mirror of
https://github.com/coder/coder.git
synced 2026-09-24 15:04:27 +08:00
feat: back the per-chat cost endpoint with AI Gateway data (#27328)
## Stack Context
This stack removes native chat cost tracking and native chat usage
limits, making the AI Gateway the single source of AI spend data and
budget enforcement.
1. **This PR:** re-back the per-chat cost endpoint with AI Gateway data.
2. Remove native chat usage limits end to end, rewiring the sidebar
indicator to gateway spend.
3. Remove native chat cost tracking end to end, deleting the
Analytics/Spend cost UI.
## What?
`GET /api/experimental/chats/{chat}/cost` summed
`chat_messages.total_cost_micros`, which native chat cost tracking
maintained. It now aggregates AI Gateway interception data instead, and
has no native fallback.
- New `GetAIBridgeChatCost` query, authorized through the root chat so
members can read their own chat's cost without gaining access to raw
interception rows.
- Response fields renamed: `priced_message_count` -> `request_count`,
`unpriced_messages_having_usage_count` -> `unpriced_request_count`.
- The chat summary sidebar keys its cost cache by root chat, and hides
the cost row where the AI Gateway is off or unlicensed. The root cost is
invalidated when a chat leaves an active status and when a generated
title lands, since title generation bills its own gateway request.
`GetChatModelUsageCostByChatID` and the rest of native cost tracking are
untouched here; PR 3 removes them.
## Why?
Native cost tracking duplicates what the AI Gateway already records, and
the two disagree. Repointing the endpoint first means the cost UI keeps
working while the native implementation is deleted later in the stack.
Two behaviour changes follow from gateway semantics and are intentional:
- **Requests, not messages.** The gateway records interceptions, so
counts are requests. Title-generation traffic now counts.
- **Whole-tree totals.** The gateway records the *spawning* chat's ID as
the interception session ID, so a subagent's requests are attributed to
its immediate parent, not always the root. Only a whole chat tree can be
summed, so the query resolves the root and aggregates the tree, and
every chat in a tree reports the same total. Native returned per-subtree
totals.
## Attribution and counting semantics
The aggregate groups token usage per interception before counting, so
the reported numbers are per request even though a request records one
usage row per provider response:
- `RequestCount` counts finished `Coder Agents` interceptions in the
tree, including unpriced ones.
- `UnpricedRequestCount` counts requests with at least one usage row the
gateway could not price. It is a subset of `RequestCount`.
- `TotalCostMicros` omits only unpriced usage, so a partially priced
request still contributes its priced portion. The sidebar therefore says
`Excludes unpriced usage from N request(s)` rather than claiming whole
requests were dropped.
A recorded cost of zero is a free request, not an unpriced one. Usage
without an effective group is excluded, matching what never reached
`ai_user_daily_spend`.
## Authorization
Reads go through `ExtractChatParam` plus `ResourceChat`, with no
cost-specific RBAC widening. `TestGetChatCost/MemberCanReadOwnChat`
covers a scoped `agents-access` member reading their own chat's cost,
and `MemberCannotReadOtherUsersChat` still asserts 404 for a non-owner.
Plain members without `agents-access` cannot create or read chats at
all, so they never reach this endpoint.
## Known limitation
AI Gateway data has its own retention period, 60 days by default and
configured independently of chat retention, so spend for requests older
than that is no longer reported. A chat whose gateway records have all
been purged reports zero cost, which is indistinguishable from genuinely
free usage under this contract. The endpoint documents the caveat;
#27330 documents it on the Spend Management page.
In-flight interceptions are excluded, since cost is only known once the
response is recorded. A chat's cost therefore lags the active turn by
one request.
## Rebase note
Rebased onto `main` after #27579 removed the `ai-gateway-cost-control`
experiment. The per-chat cost row is now gated on the `aibridge` feature
alone, matching how #27579 degated the other cost-control surfaces.
> Mux prepared this PR on Mike's behalf.
This commit is contained in:
Generated
+3
-3
@@ -582,7 +582,7 @@ const docTemplate = `{
|
||||
},
|
||||
"/api/experimental/chats/{chat}/cost": {
|
||||
"get": {
|
||||
"description": "Experimental: this endpoint is subject to change.",
|
||||
"description": "Experimental: this endpoint is subject to change.\n\nCost covers the whole chat tree: the root chat plus every\nsubagent chat beneath it. Requesting cost for a subagent chat\nreturns that same total.\n\nCost is derived from AI Gateway data, which is subject to its\nown retention period, 60 days by default, configured\nindependently of chat retention. Spend for requests older than\nthat period is no longer reported, so a chat whose requests\nhave all been purged reports zero cost.",
|
||||
"produces": [
|
||||
"application/json"
|
||||
],
|
||||
@@ -17504,13 +17504,13 @@ const docTemplate = `{
|
||||
"type": "string",
|
||||
"format": "uuid"
|
||||
},
|
||||
"priced_message_count": {
|
||||
"request_count": {
|
||||
"type": "integer"
|
||||
},
|
||||
"total_cost_micros": {
|
||||
"type": "integer"
|
||||
},
|
||||
"unpriced_messages_having_usage_count": {
|
||||
"unpriced_request_count": {
|
||||
"type": "integer"
|
||||
}
|
||||
}
|
||||
|
||||
Generated
+3
-3
@@ -511,7 +511,7 @@
|
||||
},
|
||||
"/api/experimental/chats/{chat}/cost": {
|
||||
"get": {
|
||||
"description": "Experimental: this endpoint is subject to change.",
|
||||
"description": "Experimental: this endpoint is subject to change.\n\nCost covers the whole chat tree: the root chat plus every\nsubagent chat beneath it. Requesting cost for a subagent chat\nreturns that same total.\n\nCost is derived from AI Gateway data, which is subject to its\nown retention period, 60 days by default, configured\nindependently of chat retention. Spend for requests older than\nthat period is no longer reported, so a chat whose requests\nhave all been purged reports zero cost.",
|
||||
"produces": ["application/json"],
|
||||
"tags": ["Chats"],
|
||||
"summary": "Get chat cost",
|
||||
@@ -15730,13 +15730,13 @@
|
||||
"type": "string",
|
||||
"format": "uuid"
|
||||
},
|
||||
"priced_message_count": {
|
||||
"request_count": {
|
||||
"type": "integer"
|
||||
},
|
||||
"total_cost_micros": {
|
||||
"type": "integer"
|
||||
},
|
||||
"unpriced_messages_having_usage_count": {
|
||||
"unpriced_request_count": {
|
||||
"type": "integer"
|
||||
}
|
||||
}
|
||||
|
||||
@@ -2793,6 +2793,16 @@ func (q *querier) FindMatchingPresetID(ctx context.Context, arg database.FindMat
|
||||
return q.db.FindMatchingPresetID(ctx, arg)
|
||||
}
|
||||
|
||||
func (q *querier) GetAIBridgeChatCost(ctx context.Context, rootChatID uuid.UUID) (database.GetAIBridgeChatCostRow, error) {
|
||||
// The aggregate covers one chat tree, so it is authorized through the
|
||||
// root chat. Members cannot read interception rows back, but they can
|
||||
// read their own chats.
|
||||
if _, err := q.GetChatByID(ctx, rootChatID); err != nil {
|
||||
return database.GetAIBridgeChatCostRow{}, err
|
||||
}
|
||||
return q.db.GetAIBridgeChatCost(ctx, rootChatID)
|
||||
}
|
||||
|
||||
func (q *querier) GetAIBridgeInterceptionByID(ctx context.Context, id uuid.UUID) (database.AIBridgeInterception, error) {
|
||||
return fetch(q.log, q.auth, q.db.GetAIBridgeInterceptionByID)(ctx, id)
|
||||
}
|
||||
|
||||
@@ -1069,6 +1069,13 @@ func (s *MethodTestSuite) TestChats() {
|
||||
dbm.EXPECT().GetChatMessagesByChatID(gomock.Any(), arg).Return(msgs, nil).AnyTimes()
|
||||
check.Args(arg).Asserts(chat, policy.ActionRead).Returns(msgs)
|
||||
}))
|
||||
s.Run("GetAIBridgeChatCost", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
|
||||
chat := testutil.Fake(s.T(), faker, database.Chat{})
|
||||
row := database.GetAIBridgeChatCostRow{TotalCostMicros: 1000, RequestCount: 2}
|
||||
dbm.EXPECT().GetChatByID(gomock.Any(), chat.ID).Return(chat, nil).AnyTimes()
|
||||
dbm.EXPECT().GetAIBridgeChatCost(gomock.Any(), chat.ID).Return(row, nil).AnyTimes()
|
||||
check.Args(chat.ID).Asserts(chat, policy.ActionRead).Returns(row)
|
||||
}))
|
||||
s.Run("GetChatModelUsageCostByChatID", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
|
||||
chat := testutil.Fake(s.T(), faker, database.Chat{})
|
||||
row := database.GetChatModelUsageCostByChatIDRow{ChatID: chat.ID, TotalCostMicros: 1000, PricedMessageCount: 2}
|
||||
|
||||
+8
@@ -1081,6 +1081,14 @@ func (m queryMetricsStore) FindMatchingPresetID(ctx context.Context, arg databas
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) GetAIBridgeChatCost(ctx context.Context, rootChatID uuid.UUID) (database.GetAIBridgeChatCostRow, error) {
|
||||
start := time.Now()
|
||||
r0, r1 := m.s.GetAIBridgeChatCost(ctx, rootChatID)
|
||||
m.queryLatencies.WithLabelValues("GetAIBridgeChatCost").Observe(time.Since(start).Seconds())
|
||||
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetAIBridgeChatCost").Inc()
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) GetAIBridgeInterceptionByID(ctx context.Context, id uuid.UUID) (database.AIBridgeInterception, error) {
|
||||
start := time.Now()
|
||||
r0, r1 := m.s.GetAIBridgeInterceptionByID(ctx, id)
|
||||
|
||||
Generated
+15
@@ -1858,6 +1858,21 @@ func (mr *MockStoreMockRecorder) FindMatchingPresetID(ctx, arg any) *gomock.Call
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "FindMatchingPresetID", reflect.TypeOf((*MockStore)(nil).FindMatchingPresetID), ctx, arg)
|
||||
}
|
||||
|
||||
// GetAIBridgeChatCost mocks base method.
|
||||
func (m *MockStore) GetAIBridgeChatCost(ctx context.Context, rootChatID uuid.UUID) (database.GetAIBridgeChatCostRow, error) {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "GetAIBridgeChatCost", ctx, rootChatID)
|
||||
ret0, _ := ret[0].(database.GetAIBridgeChatCostRow)
|
||||
ret1, _ := ret[1].(error)
|
||||
return ret0, ret1
|
||||
}
|
||||
|
||||
// GetAIBridgeChatCost indicates an expected call of GetAIBridgeChatCost.
|
||||
func (mr *MockStoreMockRecorder) GetAIBridgeChatCost(ctx, rootChatID any) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetAIBridgeChatCost", reflect.TypeOf((*MockStore)(nil).GetAIBridgeChatCost), ctx, rootChatID)
|
||||
}
|
||||
|
||||
// GetAIBridgeInterceptionByID mocks base method.
|
||||
func (m *MockStore) GetAIBridgeInterceptionByID(ctx context.Context, id uuid.UUID) (database.AIBridgeInterception, error) {
|
||||
m.ctrl.T.Helper()
|
||||
|
||||
Generated
+7
@@ -295,6 +295,13 @@ type sqlcQuerier interface {
|
||||
// The query finds presets where all preset parameters are present in the provided parameters,
|
||||
// and returns the preset with the most parameters (largest subset).
|
||||
FindMatchingPresetID(ctx context.Context, arg FindMatchingPresetIDParams) (uuid.UUID, error)
|
||||
// AI Gateway cost for one chat tree: the root chat plus every subagent
|
||||
// beneath it. The spawning chat's ID is recorded as the interception session
|
||||
// ID (see chatprovider.CoderHeaders), so a subagent's requests are attributed
|
||||
// to its parent rather than the root, and only whole trees can be summed. The
|
||||
// owner check guards against session-id collisions. Usage without an
|
||||
// effective group never reaches ai_user_daily_spend.
|
||||
GetAIBridgeChatCost(ctx context.Context, rootChatID uuid.UUID) (GetAIBridgeChatCostRow, error)
|
||||
GetAIBridgeInterceptionByID(ctx context.Context, id uuid.UUID) (AIBridgeInterception, error)
|
||||
// Look up the parent interception and the root of the thread by finding
|
||||
// which interception recorded a tool usage with the given tool call ID.
|
||||
|
||||
Generated
+53
@@ -1156,6 +1156,59 @@ func (q *sqlQuerier) DeleteOldAIBridgeRecords(ctx context.Context, beforeTime ti
|
||||
return total_deleted, err
|
||||
}
|
||||
|
||||
const getAIBridgeChatCost = `-- name: GetAIBridgeChatCost :one
|
||||
WITH per_request AS (
|
||||
-- One row per interception. A request records one token usage per provider
|
||||
-- response, so aggregating here keeps the outer counts per request and
|
||||
-- flags a request whose cost is partial because some usage was unpriced.
|
||||
-- The usage join is a LEFT JOIN so a request that ended without eligible
|
||||
-- usage, such as one that failed upstream, still counts as a request. The
|
||||
-- tu.id guard keeps that row from reading as unpriced usage, since the
|
||||
-- unmatched side is all NULL.
|
||||
SELECT
|
||||
SUM(tu.cost_micros) AS cost_micros,
|
||||
BOOL_OR(tu.id IS NOT NULL AND tu.cost_micros IS NULL) AS has_unpriced_usage
|
||||
FROM aibridge_interceptions i
|
||||
JOIN chats c ON c.id::text = i.session_id AND c.owner_id = i.initiator_id
|
||||
LEFT JOIN aibridge_token_usages tu ON tu.interception_id = i.id AND tu.effective_group_id IS NOT NULL
|
||||
WHERE (
|
||||
-- Spelled out instead of COALESCE(c.root_chat_id, c.id) so each branch
|
||||
-- stays a plain comparison against an indexed column.
|
||||
c.root_chat_id = $1::uuid
|
||||
OR (c.root_chat_id IS NULL AND c.id = $1::uuid)
|
||||
)
|
||||
-- Restrict to aibridge.ClientCoderAgents so another client's session
|
||||
-- reference cannot match a chat ID.
|
||||
AND i.client = 'Coder Agents'
|
||||
AND i.ended_at IS NOT NULL
|
||||
GROUP BY i.id
|
||||
)
|
||||
SELECT
|
||||
COALESCE(SUM(cost_micros), 0)::bigint AS total_cost_micros,
|
||||
COUNT(*)::bigint AS request_count,
|
||||
COUNT(*) FILTER (WHERE has_unpriced_usage)::bigint AS unpriced_request_count
|
||||
FROM per_request
|
||||
`
|
||||
|
||||
type GetAIBridgeChatCostRow struct {
|
||||
TotalCostMicros int64 `db:"total_cost_micros" json:"total_cost_micros"`
|
||||
RequestCount int64 `db:"request_count" json:"request_count"`
|
||||
UnpricedRequestCount int64 `db:"unpriced_request_count" json:"unpriced_request_count"`
|
||||
}
|
||||
|
||||
// AI Gateway cost for one chat tree: the root chat plus every subagent
|
||||
// beneath it. The spawning chat's ID is recorded as the interception session
|
||||
// ID (see chatprovider.CoderHeaders), so a subagent's requests are attributed
|
||||
// to its parent rather than the root, and only whole trees can be summed. The
|
||||
// owner check guards against session-id collisions. Usage without an
|
||||
// effective group never reaches ai_user_daily_spend.
|
||||
func (q *sqlQuerier) GetAIBridgeChatCost(ctx context.Context, rootChatID uuid.UUID) (GetAIBridgeChatCostRow, error) {
|
||||
row := q.db.QueryRowContext(ctx, getAIBridgeChatCost, rootChatID)
|
||||
var i GetAIBridgeChatCostRow
|
||||
err := row.Scan(&i.TotalCostMicros, &i.RequestCount, &i.UnpricedRequestCount)
|
||||
return i, err
|
||||
}
|
||||
|
||||
const getAIBridgeInterceptionByID = `-- name: GetAIBridgeInterceptionByID :one
|
||||
SELECT
|
||||
id, initiator_id, provider, model, started_at, metadata, ended_at, api_key_id, client, thread_parent_id, thread_root_id, client_session_id, session_id, provider_name, credential_kind, credential_hint, agent_firewall_session_id, agent_firewall_sequence_number, error_type, error_message
|
||||
|
||||
@@ -663,3 +663,42 @@ GROUP BY
|
||||
LIMIT COALESCE(NULLIF(@limit_::integer, 0), 100)
|
||||
OFFSET @offset_
|
||||
;
|
||||
|
||||
-- name: GetAIBridgeChatCost :one
|
||||
-- AI Gateway cost for one chat tree: the root chat plus every subagent
|
||||
-- beneath it. The spawning chat's ID is recorded as the interception session
|
||||
-- ID (see chatprovider.CoderHeaders), so a subagent's requests are attributed
|
||||
-- to its parent rather than the root, and only whole trees can be summed. The
|
||||
-- owner check guards against session-id collisions. Usage without an
|
||||
-- effective group never reaches ai_user_daily_spend.
|
||||
WITH per_request AS (
|
||||
-- One row per interception. A request records one token usage per provider
|
||||
-- response, so aggregating here keeps the outer counts per request and
|
||||
-- flags a request whose cost is partial because some usage was unpriced.
|
||||
-- The usage join is a LEFT JOIN so a request that ended without eligible
|
||||
-- usage, such as one that failed upstream, still counts as a request. The
|
||||
-- tu.id guard keeps that row from reading as unpriced usage, since the
|
||||
-- unmatched side is all NULL.
|
||||
SELECT
|
||||
SUM(tu.cost_micros) AS cost_micros,
|
||||
BOOL_OR(tu.id IS NOT NULL AND tu.cost_micros IS NULL) AS has_unpriced_usage
|
||||
FROM aibridge_interceptions i
|
||||
JOIN chats c ON c.id::text = i.session_id AND c.owner_id = i.initiator_id
|
||||
LEFT JOIN aibridge_token_usages tu ON tu.interception_id = i.id AND tu.effective_group_id IS NOT NULL
|
||||
WHERE (
|
||||
-- Spelled out instead of COALESCE(c.root_chat_id, c.id) so each branch
|
||||
-- stays a plain comparison against an indexed column.
|
||||
c.root_chat_id = @root_chat_id::uuid
|
||||
OR (c.root_chat_id IS NULL AND c.id = @root_chat_id::uuid)
|
||||
)
|
||||
-- Restrict to aibridge.ClientCoderAgents so another client's session
|
||||
-- reference cannot match a chat ID.
|
||||
AND i.client = 'Coder Agents'
|
||||
AND i.ended_at IS NOT NULL
|
||||
GROUP BY i.id
|
||||
)
|
||||
SELECT
|
||||
COALESCE(SUM(cost_micros), 0)::bigint AS total_cost_micros,
|
||||
COUNT(*)::bigint AS request_count,
|
||||
COUNT(*) FILTER (WHERE has_unpriced_usage)::bigint AS unpriced_request_count
|
||||
FROM per_request;
|
||||
|
||||
+30
-8
@@ -2472,16 +2472,38 @@ func (api *API) getChatMessages(rw http.ResponseWriter, r *http.Request) {
|
||||
// @Success 200 {object} codersdk.ChatCost
|
||||
// @Router /api/experimental/chats/{chat}/cost [get]
|
||||
// @Description Experimental: this endpoint is subject to change.
|
||||
// @Description
|
||||
// @Description Cost covers the whole chat tree: the root chat plus every
|
||||
// @Description subagent chat beneath it. Requesting cost for a subagent chat
|
||||
// @Description returns that same total.
|
||||
// @Description
|
||||
// @Description Cost is derived from AI Gateway data, which is subject to its
|
||||
// @Description own retention period, 60 days by default, configured
|
||||
// @Description independently of chat retention. Spend for requests older than
|
||||
// @Description that period is no longer reported, so a chat whose requests
|
||||
// @Description have all been purged reports zero cost.
|
||||
//
|
||||
//nolint:revive // HTTP handler writes to ResponseWriter.
|
||||
func (api *API) getChatCost(rw http.ResponseWriter, r *http.Request) {
|
||||
ctx := r.Context()
|
||||
chat := httpmw.ChatParam(r)
|
||||
|
||||
// The query rolls up the requested chat's subtree, so a root chat
|
||||
// reports itself plus all subagents while a subagent reports only
|
||||
// its own spend (plus any nested subagents it spawned).
|
||||
row, err := api.Database.GetChatModelUsageCostByChatID(ctx, chat.ID)
|
||||
// AI Gateway attributes a subagent's requests to the chat that spawned
|
||||
// it, so cost is only meaningful for a whole chat tree. Resolve the root
|
||||
// chat and report the tree total, including for subagent chats. Fall back
|
||||
// to the parent when root_chat_id is NULL, matching the
|
||||
// COALESCE(root_chat_id, parent_chat_id) resolution the chat queries use:
|
||||
// both columns are ON DELETE SET NULL, so deleting a root leaves
|
||||
// descendants with only a parent.
|
||||
rootChatID := chat.ID
|
||||
switch {
|
||||
case chat.RootChatID.Valid:
|
||||
rootChatID = chat.RootChatID.UUID
|
||||
case chat.ParentChatID.Valid:
|
||||
rootChatID = chat.ParentChatID.UUID
|
||||
}
|
||||
|
||||
row, err := api.Database.GetAIBridgeChatCost(ctx, rootChatID)
|
||||
if err != nil {
|
||||
if httpapi.Is404Error(err) {
|
||||
httpapi.ResourceNotFound(rw)
|
||||
@@ -2495,10 +2517,10 @@ func (api *API) getChatCost(rw http.ResponseWriter, r *http.Request) {
|
||||
}
|
||||
|
||||
httpapi.Write(ctx, rw, http.StatusOK, codersdk.ChatCost{
|
||||
ChatID: row.ChatID,
|
||||
TotalCostMicros: row.TotalCostMicros,
|
||||
PricedMessageCount: row.PricedMessageCount,
|
||||
UnpricedMessagesHavingUsageCount: row.UnpricedMessagesHavingUsageCount,
|
||||
ChatID: chat.ID,
|
||||
TotalCostMicros: row.TotalCostMicros,
|
||||
RequestCount: row.RequestCount,
|
||||
UnpricedRequestCount: row.UnpricedRequestCount,
|
||||
})
|
||||
}
|
||||
|
||||
|
||||
@@ -23,9 +23,9 @@ import (
|
||||
"github.com/coder/coder/v2/testutil"
|
||||
)
|
||||
|
||||
// ExtractChatParam authorizes the read, then GetChatModelUsageCostByChatID
|
||||
// authorizes it again. A denial on the second check means the ACL changed in
|
||||
// between (a read-authz race). Assert it surfaces as 404, not 500.
|
||||
// ExtractChatParam authorizes the read, then GetAIBridgeChatCost authorizes it
|
||||
// again. A denial on the second check means the ACL changed in between (a
|
||||
// read-authz race). Assert it surfaces as 404, not 500.
|
||||
func TestGetChatCostSurfacesReadAuthzRace(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
@@ -38,8 +38,8 @@ func TestGetChatCostSurfacesReadAuthzRace(t *testing.T) {
|
||||
}
|
||||
|
||||
dbm.EXPECT().GetChatByID(gomock.Any(), chat.ID).Return(chat, nil)
|
||||
dbm.EXPECT().GetChatModelUsageCostByChatID(gomock.Any(), chat.ID).Return(
|
||||
database.GetChatModelUsageCostByChatIDRow{},
|
||||
dbm.EXPECT().GetAIBridgeChatCost(gomock.Any(), chat.ID).Return(
|
||||
database.GetAIBridgeChatCostRow{},
|
||||
dbauthz.NotAuthorizedError{Err: sql.ErrNoRows},
|
||||
)
|
||||
|
||||
@@ -56,9 +56,9 @@ func TestGetChatCostSurfacesReadAuthzRace(t *testing.T) {
|
||||
require.Equal(t, http.StatusNotFound, resp.StatusCode)
|
||||
}
|
||||
|
||||
// A subagent chat's cost is scoped to its own subtree, so the handler
|
||||
// must query the requested chat ID rather than resolving to the root.
|
||||
func TestGetChatCostQueriesRequestedChat(t *testing.T) {
|
||||
// AI Gateway attributes a subagent's requests to the chat that spawned it, so
|
||||
// a subagent request must be answered with its root chat's tree cost.
|
||||
func TestGetChatCostQueriesRootChat(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
ctrl := gomock.NewController(t)
|
||||
@@ -73,11 +73,11 @@ func TestGetChatCostQueriesRequestedChat(t *testing.T) {
|
||||
}
|
||||
|
||||
dbm.EXPECT().GetChatByID(gomock.Any(), child.ID).Return(child, nil)
|
||||
dbm.EXPECT().GetChatModelUsageCostByChatID(gomock.Any(), child.ID).Return(
|
||||
database.GetChatModelUsageCostByChatIDRow{
|
||||
ChatID: child.ID,
|
||||
TotalCostMicros: 250,
|
||||
PricedMessageCount: 1,
|
||||
dbm.EXPECT().GetAIBridgeChatCost(gomock.Any(), rootID).Return(
|
||||
database.GetAIBridgeChatCostRow{
|
||||
TotalCostMicros: 250,
|
||||
RequestCount: 2,
|
||||
UnpricedRequestCount: 1,
|
||||
},
|
||||
nil,
|
||||
)
|
||||
@@ -97,7 +97,43 @@ func TestGetChatCostQueriesRequestedChat(t *testing.T) {
|
||||
require.NoError(t, json.NewDecoder(resp.Body).Decode(&cost))
|
||||
require.Equal(t, child.ID, cost.ChatID)
|
||||
require.Equal(t, int64(250), cost.TotalCostMicros)
|
||||
require.Equal(t, int64(1), cost.PricedMessageCount)
|
||||
require.Equal(t, int64(2), cost.RequestCount)
|
||||
require.Equal(t, int64(1), cost.UnpricedRequestCount)
|
||||
}
|
||||
|
||||
func TestGetChatCostFallsBackToParentChat(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
dbm := dbmock.NewMockStore(gomock.NewController(t))
|
||||
parentID := uuid.New()
|
||||
// chats.parent_chat_id and chats.root_chat_id are both ON DELETE SET NULL,
|
||||
// so deleting a root leaves descendants with only a parent.
|
||||
child := database.Chat{
|
||||
ID: uuid.New(),
|
||||
OwnerID: uuid.New(),
|
||||
ParentChatID: uuid.NullUUID{UUID: parentID, Valid: true},
|
||||
}
|
||||
|
||||
dbm.EXPECT().GetChatByID(gomock.Any(), child.ID).Return(child, nil)
|
||||
dbm.EXPECT().GetAIBridgeChatCost(gomock.Any(), parentID).Return(
|
||||
database.GetAIBridgeChatCostRow{TotalCostMicros: 125, RequestCount: 1},
|
||||
nil,
|
||||
)
|
||||
|
||||
api := &API{Options: &Options{Database: dbm}}
|
||||
rtr := chi.NewRouter()
|
||||
rtr.With(httpmw.ExtractChatParam(dbm)).Get("/chats/{chat}/cost", api.getChatCost)
|
||||
|
||||
req := httptest.NewRequest(http.MethodGet, "/chats/"+child.ID.String()+"/cost", nil)
|
||||
rec := httptest.NewRecorder()
|
||||
rtr.ServeHTTP(rec, req)
|
||||
resp := rec.Result()
|
||||
defer resp.Body.Close()
|
||||
|
||||
require.Equal(t, http.StatusOK, resp.StatusCode)
|
||||
var cost codersdk.ChatCost
|
||||
require.NoError(t, json.NewDecoder(resp.Body).Decode(&cost))
|
||||
require.Equal(t, int64(125), cost.TotalCostMicros)
|
||||
}
|
||||
|
||||
func TestEnrichMissingChatAgentIDs(t *testing.T) {
|
||||
|
||||
+334
-115
@@ -27,6 +27,7 @@ import (
|
||||
"golang.org/x/xerrors"
|
||||
|
||||
"cdr.dev/slog/v3/sloggers/slogtest"
|
||||
agplaibridge "github.com/coder/coder/v2/aibridge"
|
||||
"github.com/coder/coder/v2/coderd"
|
||||
"github.com/coder/coder/v2/coderd/aibridge"
|
||||
"github.com/coder/coder/v2/coderd/aibridgedtest"
|
||||
@@ -11617,29 +11618,40 @@ func TestChatCostSummary_AdminDrilldown(t *testing.T) {
|
||||
})
|
||||
}
|
||||
|
||||
// seedChatGatewayRequest records one finished Coder Agents gateway request
|
||||
// under sessionChatID, mirroring aibridged: the session ID is the spawning
|
||||
// chat, and each usage is one provider response within that one request.
|
||||
func seedChatGatewayRequest(t *testing.T, db database.Store, initiatorID, sessionChatID uuid.UUID, usages ...database.InsertAIBridgeTokenUsageParams) {
|
||||
t.Helper()
|
||||
|
||||
now := dbtime.Now()
|
||||
endedAt := now.Add(time.Second)
|
||||
interception := dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{
|
||||
InitiatorID: initiatorID,
|
||||
Provider: "anthropic",
|
||||
Model: "claude-4",
|
||||
StartedAt: now,
|
||||
Client: sql.NullString{String: string(agplaibridge.ClientCoderAgents), Valid: true},
|
||||
ClientSessionID: sql.NullString{String: sessionChatID.String(), Valid: true},
|
||||
}, &endedAt)
|
||||
|
||||
for _, usage := range usages {
|
||||
usage.InterceptionID = interception.ID
|
||||
usage.CreatedAt = now
|
||||
dbgen.AIBridgeTokenUsage(t, db, usage)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetChatCost(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
t.Run("BasicCost", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
f := seedChatCostFixture(t)
|
||||
ctx := testutil.Context(t, testutil.WaitLong)
|
||||
|
||||
cost, err := f.Client.GetChatCost(ctx, f.ChatID)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, f.ChatID, cost.ChatID)
|
||||
require.Equal(t, int64(1000), cost.TotalCostMicros)
|
||||
require.Equal(t, int64(2), cost.PricedMessageCount)
|
||||
require.Equal(t, int64(0), cost.UnpricedMessagesHavingUsageCount)
|
||||
})
|
||||
|
||||
t.Run("RollsUpSubtree", func(t *testing.T) {
|
||||
t.Run("RollsUpChatTree", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
client, db := newChatClientWithDatabase(t)
|
||||
firstUser := coderdtest.CreateFirstUser(t, client.Client)
|
||||
modelConfig := createChatModelConfig(t, client)
|
||||
everyoneGroup := uuid.NullUUID{UUID: firstUser.OrganizationID, Valid: true}
|
||||
|
||||
rootChat := dbgen.Chat(t, db, database.Chat{
|
||||
OrganizationID: firstUser.OrganizationID,
|
||||
@@ -11647,13 +11659,6 @@ func TestGetChatCost(t *testing.T) {
|
||||
LastModelConfigID: modelConfig.ID,
|
||||
Title: "root chat",
|
||||
})
|
||||
_ = dbgen.ChatMessage(t, db, database.ChatMessage{
|
||||
ChatID: rootChat.ID,
|
||||
ModelConfigID: uuid.NullUUID{UUID: modelConfig.ID, Valid: true},
|
||||
Role: database.ChatMessageRoleAssistant,
|
||||
TotalCostMicros: sql.NullInt64{Int64: 500, Valid: true},
|
||||
})
|
||||
|
||||
childChat := dbgen.Chat(t, db, database.Chat{
|
||||
OrganizationID: firstUser.OrganizationID,
|
||||
OwnerID: firstUser.UserID,
|
||||
@@ -11662,15 +11667,6 @@ func TestGetChatCost(t *testing.T) {
|
||||
ParentChatID: uuid.NullUUID{UUID: rootChat.ID, Valid: true},
|
||||
RootChatID: uuid.NullUUID{UUID: rootChat.ID, Valid: true},
|
||||
})
|
||||
_ = dbgen.ChatMessage(t, db, database.ChatMessage{
|
||||
ChatID: childChat.ID,
|
||||
ModelConfigID: uuid.NullUUID{UUID: modelConfig.ID, Valid: true},
|
||||
Role: database.ChatMessageRoleAssistant,
|
||||
TotalCostMicros: sql.NullInt64{Int64: 250, Valid: true},
|
||||
})
|
||||
|
||||
// root_chat_id is flattened to the top-level root at any depth,
|
||||
// so subtree traversal must follow parent_chat_id instead.
|
||||
grandchildChat := dbgen.Chat(t, db, database.Chat{
|
||||
OrganizationID: firstUser.OrganizationID,
|
||||
OwnerID: firstUser.UserID,
|
||||
@@ -11679,39 +11675,219 @@ func TestGetChatCost(t *testing.T) {
|
||||
ParentChatID: uuid.NullUUID{UUID: childChat.ID, Valid: true},
|
||||
RootChatID: uuid.NullUUID{UUID: rootChat.ID, Valid: true},
|
||||
})
|
||||
_ = dbgen.ChatMessage(t, db, database.ChatMessage{
|
||||
ChatID: grandchildChat.ID,
|
||||
ModelConfigID: uuid.NullUUID{UUID: modelConfig.ID, Valid: true},
|
||||
Role: database.ChatMessageRoleAssistant,
|
||||
TotalCostMicros: sql.NullInt64{Int64: 100, Valid: true},
|
||||
|
||||
// A subagent's requests carry the spawning chat's ID, so the child's
|
||||
// spend lands on the root session and the grandchild's on the child.
|
||||
seedChatGatewayRequest(t, db, firstUser.UserID, rootChat.ID, database.InsertAIBridgeTokenUsageParams{
|
||||
EffectiveGroupID: everyoneGroup,
|
||||
CostMicros: sql.NullInt64{Int64: 500, Valid: true},
|
||||
})
|
||||
seedChatGatewayRequest(t, db, firstUser.UserID, rootChat.ID, database.InsertAIBridgeTokenUsageParams{
|
||||
EffectiveGroupID: everyoneGroup,
|
||||
CostMicros: sql.NullInt64{Int64: 250, Valid: true},
|
||||
})
|
||||
seedChatGatewayRequest(t, db, firstUser.UserID, childChat.ID, database.InsertAIBridgeTokenUsageParams{
|
||||
EffectiveGroupID: everyoneGroup,
|
||||
CostMicros: sql.NullInt64{Int64: 100, Valid: true},
|
||||
})
|
||||
|
||||
ctx := testutil.Context(t, testutil.WaitLong)
|
||||
|
||||
// The root rolls up every descendant's cost.
|
||||
rootCost, err := client.GetChatCost(ctx, rootChat.ID)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, rootChat.ID, rootCost.ChatID)
|
||||
require.Equal(t, int64(850), rootCost.TotalCostMicros)
|
||||
require.Equal(t, int64(3), rootCost.PricedMessageCount)
|
||||
|
||||
// A subagent reports only its own subtree: itself plus the
|
||||
// nested subagents it spawned, excluding the parent's spend.
|
||||
childCost, err := client.GetChatCost(ctx, childChat.ID)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, childChat.ID, childCost.ChatID)
|
||||
require.Equal(t, int64(350), childCost.TotalCostMicros)
|
||||
require.Equal(t, int64(2), childCost.PricedMessageCount)
|
||||
|
||||
// A leaf subagent reports only its own spend.
|
||||
grandchildCost, err := client.GetChatCost(ctx, grandchildChat.ID)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, grandchildChat.ID, grandchildCost.ChatID)
|
||||
require.Equal(t, int64(100), grandchildCost.TotalCostMicros)
|
||||
require.Equal(t, int64(1), grandchildCost.PricedMessageCount)
|
||||
for _, chatID := range []uuid.UUID{rootChat.ID, childChat.ID, grandchildChat.ID} {
|
||||
cost, err := client.GetChatCost(ctx, chatID)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, chatID, cost.ChatID)
|
||||
require.Equal(t, int64(850), cost.TotalCostMicros)
|
||||
require.Equal(t, int64(3), cost.RequestCount)
|
||||
require.Equal(t, int64(0), cost.UnpricedRequestCount)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("UnpricedMessages", func(t *testing.T) {
|
||||
t.Run("UnpricedRequests", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
client, db := newChatClientWithDatabase(t)
|
||||
firstUser := coderdtest.CreateFirstUser(t, client.Client)
|
||||
modelConfig := createChatModelConfig(t, client)
|
||||
everyoneGroup := uuid.NullUUID{UUID: firstUser.OrganizationID, Valid: true}
|
||||
|
||||
chat := dbgen.Chat(t, db, database.Chat{
|
||||
OrganizationID: firstUser.OrganizationID,
|
||||
OwnerID: firstUser.UserID,
|
||||
LastModelConfigID: modelConfig.ID,
|
||||
Title: "unpriced chat",
|
||||
})
|
||||
seedChatGatewayRequest(t, db, firstUser.UserID, chat.ID, database.InsertAIBridgeTokenUsageParams{
|
||||
EffectiveGroupID: everyoneGroup,
|
||||
CostMicros: sql.NullInt64{Int64: 400, Valid: true},
|
||||
})
|
||||
seedChatGatewayRequest(t, db, firstUser.UserID, chat.ID, database.InsertAIBridgeTokenUsageParams{
|
||||
EffectiveGroupID: everyoneGroup,
|
||||
})
|
||||
|
||||
ctx := testutil.Context(t, testutil.WaitLong)
|
||||
|
||||
cost, err := client.GetChatCost(ctx, chat.ID)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, int64(400), cost.TotalCostMicros)
|
||||
require.Equal(t, int64(2), cost.RequestCount)
|
||||
require.Equal(t, int64(1), cost.UnpricedRequestCount)
|
||||
})
|
||||
|
||||
t.Run("PartiallyPricedRequest", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
client, db := newChatClientWithDatabase(t)
|
||||
firstUser := coderdtest.CreateFirstUser(t, client.Client)
|
||||
modelConfig := createChatModelConfig(t, client)
|
||||
everyoneGroup := uuid.NullUUID{UUID: firstUser.OrganizationID, Valid: true}
|
||||
|
||||
chat := dbgen.Chat(t, db, database.Chat{
|
||||
OrganizationID: firstUser.OrganizationID,
|
||||
OwnerID: firstUser.UserID,
|
||||
LastModelConfigID: modelConfig.ID,
|
||||
Title: "partially priced chat",
|
||||
})
|
||||
// One request, two provider responses, only one of them priced. The
|
||||
// priced usage still counts towards the total, and the request is
|
||||
// reported as unpriced so the total is not presented as exact.
|
||||
seedChatGatewayRequest(t, db, firstUser.UserID, chat.ID,
|
||||
database.InsertAIBridgeTokenUsageParams{
|
||||
EffectiveGroupID: everyoneGroup,
|
||||
CostMicros: sql.NullInt64{Int64: 300, Valid: true},
|
||||
},
|
||||
database.InsertAIBridgeTokenUsageParams{
|
||||
EffectiveGroupID: everyoneGroup,
|
||||
},
|
||||
)
|
||||
|
||||
ctx := testutil.Context(t, testutil.WaitLong)
|
||||
|
||||
cost, err := client.GetChatCost(ctx, chat.ID)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, int64(300), cost.TotalCostMicros)
|
||||
require.Equal(t, int64(1), cost.RequestCount)
|
||||
require.Equal(t, int64(1), cost.UnpricedRequestCount)
|
||||
})
|
||||
|
||||
t.Run("ZeroCostRequests", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
client, db := newChatClientWithDatabase(t)
|
||||
firstUser := coderdtest.CreateFirstUser(t, client.Client)
|
||||
modelConfig := createChatModelConfig(t, client)
|
||||
everyoneGroup := uuid.NullUUID{UUID: firstUser.OrganizationID, Valid: true}
|
||||
|
||||
chat := dbgen.Chat(t, db, database.Chat{
|
||||
OrganizationID: firstUser.OrganizationID,
|
||||
OwnerID: firstUser.UserID,
|
||||
LastModelConfigID: modelConfig.ID,
|
||||
Title: "free chat",
|
||||
})
|
||||
// A recorded cost of zero is a free request, not an unpriced one.
|
||||
seedChatGatewayRequest(t, db, firstUser.UserID, chat.ID, database.InsertAIBridgeTokenUsageParams{
|
||||
EffectiveGroupID: everyoneGroup,
|
||||
CostMicros: sql.NullInt64{Int64: 0, Valid: true},
|
||||
})
|
||||
|
||||
ctx := testutil.Context(t, testutil.WaitLong)
|
||||
|
||||
cost, err := client.GetChatCost(ctx, chat.ID)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, int64(0), cost.TotalCostMicros)
|
||||
require.Equal(t, int64(1), cost.RequestCount)
|
||||
require.Equal(t, int64(0), cost.UnpricedRequestCount)
|
||||
})
|
||||
|
||||
t.Run("RequestWithoutUsage", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
client, db := newChatClientWithDatabase(t)
|
||||
firstUser := coderdtest.CreateFirstUser(t, client.Client)
|
||||
modelConfig := createChatModelConfig(t, client)
|
||||
everyoneGroup := uuid.NullUUID{UUID: firstUser.OrganizationID, Valid: true}
|
||||
|
||||
chat := dbgen.Chat(t, db, database.Chat{
|
||||
OrganizationID: firstUser.OrganizationID,
|
||||
OwnerID: firstUser.UserID,
|
||||
LastModelConfigID: modelConfig.ID,
|
||||
Title: "failed request chat",
|
||||
})
|
||||
seedChatGatewayRequest(t, db, firstUser.UserID, chat.ID, database.InsertAIBridgeTokenUsageParams{
|
||||
EffectiveGroupID: everyoneGroup,
|
||||
CostMicros: sql.NullInt64{Int64: 200, Valid: true},
|
||||
})
|
||||
// A request that fails upstream still ends, but records no usage. It
|
||||
// counts as a request, adds no cost, and is not unpriced usage.
|
||||
seedChatGatewayRequest(t, db, firstUser.UserID, chat.ID)
|
||||
|
||||
ctx := testutil.Context(t, testutil.WaitLong)
|
||||
|
||||
cost, err := client.GetChatCost(ctx, chat.ID)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, int64(200), cost.TotalCostMicros)
|
||||
require.Equal(t, int64(2), cost.RequestCount)
|
||||
require.Equal(t, int64(0), cost.UnpricedRequestCount)
|
||||
})
|
||||
|
||||
t.Run("IsolatesSiblingChatTrees", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
client, db := newChatClientWithDatabase(t)
|
||||
firstUser := coderdtest.CreateFirstUser(t, client.Client)
|
||||
modelConfig := createChatModelConfig(t, client)
|
||||
everyoneGroup := uuid.NullUUID{UUID: firstUser.OrganizationID, Valid: true}
|
||||
|
||||
firstRoot := dbgen.Chat(t, db, database.Chat{
|
||||
OrganizationID: firstUser.OrganizationID,
|
||||
OwnerID: firstUser.UserID,
|
||||
LastModelConfigID: modelConfig.ID,
|
||||
Title: "first root chat",
|
||||
})
|
||||
secondRoot := dbgen.Chat(t, db, database.Chat{
|
||||
OrganizationID: firstUser.OrganizationID,
|
||||
OwnerID: firstUser.UserID,
|
||||
LastModelConfigID: modelConfig.ID,
|
||||
Title: "second root chat",
|
||||
})
|
||||
secondChild := dbgen.Chat(t, db, database.Chat{
|
||||
OrganizationID: firstUser.OrganizationID,
|
||||
OwnerID: firstUser.UserID,
|
||||
LastModelConfigID: modelConfig.ID,
|
||||
Title: "second root subagent",
|
||||
ParentChatID: uuid.NullUUID{UUID: secondRoot.ID, Valid: true},
|
||||
RootChatID: uuid.NullUUID{UUID: secondRoot.ID, Valid: true},
|
||||
})
|
||||
|
||||
seedChatGatewayRequest(t, db, firstUser.UserID, firstRoot.ID, database.InsertAIBridgeTokenUsageParams{
|
||||
EffectiveGroupID: everyoneGroup,
|
||||
CostMicros: sql.NullInt64{Int64: 500, Valid: true},
|
||||
})
|
||||
seedChatGatewayRequest(t, db, firstUser.UserID, secondRoot.ID, database.InsertAIBridgeTokenUsageParams{
|
||||
EffectiveGroupID: everyoneGroup,
|
||||
CostMicros: sql.NullInt64{Int64: 120, Valid: true},
|
||||
})
|
||||
seedChatGatewayRequest(t, db, firstUser.UserID, secondChild.ID, database.InsertAIBridgeTokenUsageParams{
|
||||
EffectiveGroupID: everyoneGroup,
|
||||
CostMicros: sql.NullInt64{Int64: 30, Valid: true},
|
||||
})
|
||||
|
||||
ctx := testutil.Context(t, testutil.WaitLong)
|
||||
|
||||
firstCost, err := client.GetChatCost(ctx, firstRoot.ID)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, int64(500), firstCost.TotalCostMicros)
|
||||
require.Equal(t, int64(1), firstCost.RequestCount)
|
||||
|
||||
for _, chatID := range []uuid.UUID{secondRoot.ID, secondChild.ID} {
|
||||
secondCost, err := client.GetChatCost(ctx, chatID)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, int64(150), secondCost.TotalCostMicros)
|
||||
require.Equal(t, int64(2), secondCost.RequestCount)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("ExcludesUnattributedUsageFromCost", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
client, db := newChatClientWithDatabase(t)
|
||||
@@ -11722,30 +11898,108 @@ func TestGetChatCost(t *testing.T) {
|
||||
OrganizationID: firstUser.OrganizationID,
|
||||
OwnerID: firstUser.UserID,
|
||||
LastModelConfigID: modelConfig.ID,
|
||||
Title: "unpriced chat",
|
||||
Title: "legacy chat",
|
||||
})
|
||||
_ = dbgen.ChatMessage(t, db, database.ChatMessage{
|
||||
ChatID: chat.ID,
|
||||
ModelConfigID: uuid.NullUUID{UUID: modelConfig.ID, Valid: true},
|
||||
Role: database.ChatMessageRoleAssistant,
|
||||
TotalCostMicros: sql.NullInt64{Int64: 400, Valid: true},
|
||||
})
|
||||
// Token usage but no cost (no model pricing) counts as unpriced.
|
||||
_ = dbgen.ChatMessage(t, db, database.ChatMessage{
|
||||
ChatID: chat.ID,
|
||||
ModelConfigID: uuid.NullUUID{UUID: modelConfig.ID, Valid: true},
|
||||
Role: database.ChatMessageRoleAssistant,
|
||||
InputTokens: sql.NullInt64{Int64: 100, Valid: true},
|
||||
OutputTokens: sql.NullInt64{Int64: 50, Valid: true},
|
||||
// Usage recorded before group attribution existed never reached
|
||||
// ai_user_daily_spend, so it must not appear as chat spend either. The
|
||||
// request itself still finished, so it stays in the request count.
|
||||
seedChatGatewayRequest(t, db, firstUser.UserID, chat.ID, database.InsertAIBridgeTokenUsageParams{
|
||||
CostMicros: sql.NullInt64{Int64: 900, Valid: true},
|
||||
})
|
||||
|
||||
ctx := testutil.Context(t, testutil.WaitLong)
|
||||
|
||||
cost, err := client.GetChatCost(ctx, chat.ID)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, int64(400), cost.TotalCostMicros)
|
||||
require.Equal(t, int64(1), cost.PricedMessageCount)
|
||||
require.Equal(t, int64(1), cost.UnpricedMessagesHavingUsageCount)
|
||||
require.Equal(t, int64(0), cost.TotalCostMicros)
|
||||
require.Equal(t, int64(1), cost.RequestCount)
|
||||
require.Equal(t, int64(0), cost.UnpricedRequestCount)
|
||||
})
|
||||
|
||||
t.Run("ExcludesForeignAndUnfinishedRequests", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
client, db := newChatClientWithDatabase(t)
|
||||
firstUser := coderdtest.CreateFirstUser(t, client.Client)
|
||||
_, otherUser := coderdtest.CreateAnotherUser(t, client.Client, firstUser.OrganizationID)
|
||||
modelConfig := createChatModelConfig(t, client)
|
||||
everyoneGroup := uuid.NullUUID{UUID: firstUser.OrganizationID, Valid: true}
|
||||
|
||||
chat := dbgen.Chat(t, db, database.Chat{
|
||||
OrganizationID: firstUser.OrganizationID,
|
||||
OwnerID: firstUser.UserID,
|
||||
LastModelConfigID: modelConfig.ID,
|
||||
Title: "owner chat",
|
||||
})
|
||||
|
||||
seedChatGatewayRequest(t, db, otherUser.ID, chat.ID, database.InsertAIBridgeTokenUsageParams{
|
||||
EffectiveGroupID: everyoneGroup,
|
||||
CostMicros: sql.NullInt64{Int64: 700, Valid: true},
|
||||
})
|
||||
|
||||
foreignEndedAt := dbtime.Now().Add(time.Second)
|
||||
foreignInterception := dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{
|
||||
InitiatorID: firstUser.UserID,
|
||||
StartedAt: dbtime.Now(),
|
||||
Client: sql.NullString{String: string(agplaibridge.ClientClaudeCode), Valid: true},
|
||||
ClientSessionID: sql.NullString{String: chat.ID.String(), Valid: true},
|
||||
}, &foreignEndedAt)
|
||||
dbgen.AIBridgeTokenUsage(t, db, database.InsertAIBridgeTokenUsageParams{
|
||||
InterceptionID: foreignInterception.ID,
|
||||
EffectiveGroupID: everyoneGroup,
|
||||
CostMicros: sql.NullInt64{Int64: 800, Valid: true},
|
||||
})
|
||||
|
||||
unfinishedInterception := dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{
|
||||
InitiatorID: firstUser.UserID,
|
||||
StartedAt: dbtime.Now(),
|
||||
Client: sql.NullString{String: string(agplaibridge.ClientCoderAgents), Valid: true},
|
||||
ClientSessionID: sql.NullString{String: chat.ID.String(), Valid: true},
|
||||
}, nil)
|
||||
dbgen.AIBridgeTokenUsage(t, db, database.InsertAIBridgeTokenUsageParams{
|
||||
InterceptionID: unfinishedInterception.ID,
|
||||
EffectiveGroupID: everyoneGroup,
|
||||
CostMicros: sql.NullInt64{Int64: 600, Valid: true},
|
||||
})
|
||||
|
||||
ctx := testutil.Context(t, testutil.WaitLong)
|
||||
|
||||
cost, err := client.GetChatCost(ctx, chat.ID)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, int64(0), cost.TotalCostMicros)
|
||||
require.Equal(t, int64(0), cost.RequestCount)
|
||||
})
|
||||
|
||||
t.Run("MemberCanReadOwnChat", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
client, db := newChatClientWithDatabase(t)
|
||||
firstUser := coderdtest.CreateFirstUser(t, client.Client)
|
||||
// agents-access is what grants ResourceChat; plain members cannot
|
||||
// create or read chats at all, so they never reach this endpoint.
|
||||
memberClientRaw, member := coderdtest.CreateAnotherUser(t, client.Client, firstUser.OrganizationID, rbac.ScopedRoleAgentsAccess(firstUser.OrganizationID))
|
||||
memberClient := codersdk.NewExperimentalClient(memberClientRaw)
|
||||
modelConfig := createChatModelConfig(t, client)
|
||||
everyoneGroup := uuid.NullUUID{UUID: firstUser.OrganizationID, Valid: true}
|
||||
|
||||
chat := dbgen.Chat(t, db, database.Chat{
|
||||
OrganizationID: firstUser.OrganizationID,
|
||||
OwnerID: member.ID,
|
||||
LastModelConfigID: modelConfig.ID,
|
||||
Title: "member chat",
|
||||
})
|
||||
seedChatGatewayRequest(t, db, member.ID, chat.ID, database.InsertAIBridgeTokenUsageParams{
|
||||
EffectiveGroupID: everyoneGroup,
|
||||
CostMicros: sql.NullInt64{Int64: 450, Valid: true},
|
||||
})
|
||||
|
||||
ctx := testutil.Context(t, testutil.WaitLong)
|
||||
|
||||
cost, err := memberClient.GetChatCost(ctx, chat.ID)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, chat.ID, cost.ChatID)
|
||||
require.Equal(t, int64(450), cost.TotalCostMicros)
|
||||
require.Equal(t, int64(1), cost.RequestCount)
|
||||
})
|
||||
|
||||
t.Run("MemberCannotReadOtherUsersChat", func(t *testing.T) {
|
||||
@@ -11773,15 +12027,15 @@ func TestGetChatCost(t *testing.T) {
|
||||
require.Equal(t, http.StatusNotFound, sdkErr.StatusCode())
|
||||
})
|
||||
|
||||
t.Run("ZeroMessages", func(t *testing.T) {
|
||||
t.Run("ZeroRequests", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
client, db := newChatClientWithDatabase(t)
|
||||
firstUser := coderdtest.CreateFirstUser(t, client.Client)
|
||||
modelConfig := createChatModelConfig(t, client)
|
||||
|
||||
// No assistant messages must still return one zero-total row; a COALESCE
|
||||
// or :one regression would surface as sql.ErrNoRows -> 500.
|
||||
// An ungrouped aggregate always returns a row, so a chat with no
|
||||
// gateway requests reports zeros instead of failing.
|
||||
chat := dbgen.Chat(t, db, database.Chat{
|
||||
OrganizationID: firstUser.OrganizationID,
|
||||
OwnerID: firstUser.UserID,
|
||||
@@ -11795,43 +12049,8 @@ func TestGetChatCost(t *testing.T) {
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, chat.ID, cost.ChatID)
|
||||
require.Equal(t, int64(0), cost.TotalCostMicros)
|
||||
require.Equal(t, int64(0), cost.PricedMessageCount)
|
||||
require.Equal(t, int64(0), cost.UnpricedMessagesHavingUsageCount)
|
||||
})
|
||||
|
||||
t.Run("ExcludesNonAssistantMessages", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
client, db := newChatClientWithDatabase(t)
|
||||
firstUser := coderdtest.CreateFirstUser(t, client.Client)
|
||||
modelConfig := createChatModelConfig(t, client)
|
||||
|
||||
chat := dbgen.Chat(t, db, database.Chat{
|
||||
OrganizationID: firstUser.OrganizationID,
|
||||
OwnerID: firstUser.UserID,
|
||||
LastModelConfigID: modelConfig.ID,
|
||||
Title: "mixed-role chat",
|
||||
})
|
||||
_ = dbgen.ChatMessage(t, db, database.ChatMessage{
|
||||
ChatID: chat.ID,
|
||||
ModelConfigID: uuid.NullUUID{UUID: modelConfig.ID, Valid: true},
|
||||
Role: database.ChatMessageRoleAssistant,
|
||||
TotalCostMicros: sql.NullInt64{Int64: 600, Valid: true},
|
||||
})
|
||||
// User-role cost must be excluded; the query bills only assistant messages.
|
||||
_ = dbgen.ChatMessage(t, db, database.ChatMessage{
|
||||
ChatID: chat.ID,
|
||||
Role: database.ChatMessageRoleUser,
|
||||
TotalCostMicros: sql.NullInt64{Int64: 999, Valid: true},
|
||||
})
|
||||
|
||||
ctx := testutil.Context(t, testutil.WaitLong)
|
||||
|
||||
cost, err := client.GetChatCost(ctx, chat.ID)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, int64(600), cost.TotalCostMicros)
|
||||
require.Equal(t, int64(1), cost.PricedMessageCount)
|
||||
require.Equal(t, int64(0), cost.UnpricedMessagesHavingUsageCount)
|
||||
require.Equal(t, int64(0), cost.RequestCount)
|
||||
require.Equal(t, int64(0), cost.UnpricedRequestCount)
|
||||
})
|
||||
}
|
||||
|
||||
|
||||
+13
-9
@@ -1990,15 +1990,18 @@ type ChatCostChatBreakdown struct {
|
||||
TotalRuntimeMs int64 `json:"total_runtime_ms"`
|
||||
}
|
||||
|
||||
// ChatCost is the cumulative cost for a selected chat's subtree: the
|
||||
// chat itself plus every descendant (subagent) chat it spawned. A root
|
||||
// chat therefore reports its whole tree, while a subagent reports only
|
||||
// its own spend plus any nested subagents.
|
||||
// ChatCost is the AI Gateway cost for the requested chat's whole tree.
|
||||
// Root and subagent chats report the same total.
|
||||
// RequestCount counts every finished request in the tree, including ones that
|
||||
// recorded no billable usage at all, such as a request that failed upstream.
|
||||
// UnpricedRequestCount counts requests with at least one usage record whose
|
||||
// model had no recorded price; RequestCount includes them and
|
||||
// TotalCostMicros omits only their unpriced usage.
|
||||
type ChatCost struct {
|
||||
ChatID uuid.UUID `json:"chat_id" format:"uuid"`
|
||||
TotalCostMicros int64 `json:"total_cost_micros"`
|
||||
PricedMessageCount int64 `json:"priced_message_count"`
|
||||
UnpricedMessagesHavingUsageCount int64 `json:"unpriced_messages_having_usage_count"`
|
||||
ChatID uuid.UUID `json:"chat_id" format:"uuid"`
|
||||
TotalCostMicros int64 `json:"total_cost_micros"`
|
||||
RequestCount int64 `json:"request_count"`
|
||||
UnpricedRequestCount int64 `json:"unpriced_request_count"`
|
||||
}
|
||||
|
||||
// ChatCostUserRollup contains per-user cost aggregation for admin views.
|
||||
@@ -2592,7 +2595,8 @@ func (c *ExperimentalClient) GetChatCostSummary(ctx context.Context, user string
|
||||
return summary, json.NewDecoder(res.Body).Decode(&summary)
|
||||
}
|
||||
|
||||
// GetChatCost returns the cumulative cost for a single chat.
|
||||
// GetChatCost returns the AI Gateway cost for the whole chat tree that
|
||||
// contains chatID.
|
||||
func (c *ExperimentalClient) GetChatCost(ctx context.Context, chatID uuid.UUID) (ChatCost, error) {
|
||||
res, err := c.Request(ctx, http.MethodGet, fmt.Sprintf("/api/experimental/chats/%s/cost", chatID), nil)
|
||||
if err != nil {
|
||||
|
||||
Generated
+12
-2
@@ -1298,6 +1298,16 @@ curl -X GET http://coder-server:8080/api/experimental/chats/{chat}/cost \
|
||||
|
||||
Experimental: this endpoint is subject to change.
|
||||
|
||||
Cost covers the whole chat tree: the root chat plus every
|
||||
subagent chat beneath it. Requesting cost for a subagent chat
|
||||
returns that same total.
|
||||
|
||||
Cost is derived from AI Gateway data, which is subject to its
|
||||
own retention period, 60 days by default, configured
|
||||
independently of chat retention. Spend for requests older than
|
||||
that period is no longer reported, so a chat whose requests
|
||||
have all been purged reports zero cost.
|
||||
|
||||
### Parameters
|
||||
|
||||
| Name | In | Type | Required | Description |
|
||||
@@ -1311,9 +1321,9 @@ Experimental: this endpoint is subject to change.
|
||||
```json
|
||||
{
|
||||
"chat_id": "efc9fe20-a1e5-4a8c-9c48-f1b30c1e4f86",
|
||||
"priced_message_count": 0,
|
||||
"request_count": 0,
|
||||
"total_cost_micros": 0,
|
||||
"unpriced_messages_having_usage_count": 0
|
||||
"unpriced_request_count": 0
|
||||
}
|
||||
```
|
||||
|
||||
|
||||
Generated
+8
-8
@@ -2558,20 +2558,20 @@ AuthorizationObject can represent a "set" of objects, such as: all workspaces in
|
||||
```json
|
||||
{
|
||||
"chat_id": "efc9fe20-a1e5-4a8c-9c48-f1b30c1e4f86",
|
||||
"priced_message_count": 0,
|
||||
"request_count": 0,
|
||||
"total_cost_micros": 0,
|
||||
"unpriced_messages_having_usage_count": 0
|
||||
"unpriced_request_count": 0
|
||||
}
|
||||
```
|
||||
|
||||
### Properties
|
||||
|
||||
| Name | Type | Required | Restrictions | Description |
|
||||
|----------------------------------------|---------|----------|--------------|-------------|
|
||||
| `chat_id` | string | false | | |
|
||||
| `priced_message_count` | integer | false | | |
|
||||
| `total_cost_micros` | integer | false | | |
|
||||
| `unpriced_messages_having_usage_count` | integer | false | | |
|
||||
| Name | Type | Required | Restrictions | Description |
|
||||
|--------------------------|---------|----------|--------------|-------------|
|
||||
| `chat_id` | string | false | | |
|
||||
| `request_count` | integer | false | | |
|
||||
| `total_cost_micros` | integer | false | | |
|
||||
| `unpriced_request_count` | integer | false | | |
|
||||
|
||||
## codersdk.ChatDiffContents
|
||||
|
||||
|
||||
@@ -1966,17 +1966,15 @@ export const chatCostSummary = (user = "me", params?: ChatCostDateParams) => ({
|
||||
staleTime: 60_000,
|
||||
});
|
||||
|
||||
export const chatCostKey = (chatId: string) =>
|
||||
[...chatsKey, chatId, "cost"] as const;
|
||||
export const chatCostKey = (rootChatId: string) =>
|
||||
[...chatsKey, rootChatId, "cost"] as const;
|
||||
|
||||
// Chat cost changes only when a new assistant message is priced, so a short
|
||||
// stale window refreshes the sidebar without refetching on every render.
|
||||
const ASSISTANT_MESSAGE_PRICING_STALE_MS = 30_000;
|
||||
const GATEWAY_REQUEST_STALE_MS = 30_000;
|
||||
|
||||
export const chatCost = (chatId: string) => ({
|
||||
queryKey: chatCostKey(chatId),
|
||||
queryFn: () => API.experimental.getChatCost(chatId),
|
||||
staleTime: ASSISTANT_MESSAGE_PRICING_STALE_MS,
|
||||
export const chatCost = (rootChatId: string) => ({
|
||||
queryKey: chatCostKey(rootChatId),
|
||||
queryFn: () => API.experimental.getChatCost(rootChatId),
|
||||
staleTime: GATEWAY_REQUEST_STALE_MS,
|
||||
});
|
||||
|
||||
interface PaginatedChatCostUsersPayload {
|
||||
|
||||
Generated
+9
-6
@@ -2155,16 +2155,19 @@ export interface ChatContextTool {
|
||||
|
||||
// From codersdk/chats.go
|
||||
/**
|
||||
* ChatCost is the cumulative cost for a selected chat's subtree: the
|
||||
* chat itself plus every descendant (subagent) chat it spawned. A root
|
||||
* chat therefore reports its whole tree, while a subagent reports only
|
||||
* its own spend plus any nested subagents.
|
||||
* ChatCost is the AI Gateway cost for the requested chat's whole tree.
|
||||
* Root and subagent chats report the same total.
|
||||
* RequestCount counts every finished request in the tree, including ones that
|
||||
* recorded no billable usage at all, such as a request that failed upstream.
|
||||
* UnpricedRequestCount counts requests with at least one usage record whose
|
||||
* model had no recorded price; RequestCount includes them and
|
||||
* TotalCostMicros omits only their unpriced usage.
|
||||
*/
|
||||
export interface ChatCost {
|
||||
readonly chat_id: string;
|
||||
readonly total_cost_micros: number;
|
||||
readonly priced_message_count: number;
|
||||
readonly unpriced_messages_having_usage_count: number;
|
||||
readonly request_count: number;
|
||||
readonly unpriced_request_count: number;
|
||||
}
|
||||
|
||||
// From codersdk/chats.go
|
||||
|
||||
@@ -196,15 +196,11 @@ const StoryAgentChatPageView: FC<StoryProps> = ({ editing, ...overrides }) => {
|
||||
const meta: Meta<typeof AgentChatPageView> = {
|
||||
title: "pages/AgentsPage/AgentChatPageView",
|
||||
component: AgentChatPageView,
|
||||
// Summary is the default tab and reads chat + cost; mock both so the sidebar renders.
|
||||
// Summary is the default tab and reads the chat, so mock it for the sidebar.
|
||||
// Cost needs no mock: these stories leave the aibridge feature off, so the
|
||||
// summary panel never requests it.
|
||||
beforeEach: () => {
|
||||
spyOn(API.experimental, "getChat").mockResolvedValue(buildChat());
|
||||
spyOn(API.experimental, "getChatCost").mockResolvedValue({
|
||||
chat_id: AGENT_ID,
|
||||
total_cost_micros: 0,
|
||||
priced_message_count: 0,
|
||||
unpriced_messages_having_usage_count: 0,
|
||||
});
|
||||
},
|
||||
decorators: [withAuthProvider, withDashboardProvider, withProxyProvider()],
|
||||
parameters: {
|
||||
|
||||
@@ -2,7 +2,7 @@ import { act, renderHook } from "@testing-library/react";
|
||||
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
|
||||
import type * as TypesGen from "#/api/typesGenerated";
|
||||
import {
|
||||
chatCostIdsToInvalidate,
|
||||
chatCostIdToInvalidate,
|
||||
shouldInvalidateFilteredChatList,
|
||||
} from "./AgentsPageLayout";
|
||||
import {
|
||||
@@ -933,30 +933,21 @@ describe(shouldInvalidateFilteredChatList.name, () => {
|
||||
});
|
||||
});
|
||||
|
||||
describe(chatCostIdsToInvalidate.name, () => {
|
||||
describe(chatCostIdToInvalidate.name, () => {
|
||||
it.each<{
|
||||
name: string;
|
||||
updatedChat: TypesGen.Chat;
|
||||
eventKind: TypesGen.ChatWatchEventKind;
|
||||
expected: readonly string[];
|
||||
expected: string | undefined;
|
||||
}>([
|
||||
{
|
||||
name: "invalidates when a status change ends active generation",
|
||||
updatedChat: chatForFilterInvalidation({ status: "waiting" }),
|
||||
eventKind: "status_change",
|
||||
expected: ["chat-1"],
|
||||
expected: "chat-1",
|
||||
},
|
||||
{
|
||||
name: "does not duplicate a self-referential root id",
|
||||
updatedChat: chatForFilterInvalidation({
|
||||
status: "waiting",
|
||||
root_chat_id: "chat-1",
|
||||
}),
|
||||
eventKind: "status_change",
|
||||
expected: ["chat-1"],
|
||||
},
|
||||
{
|
||||
name: "invalidates the root chat's rolled-up cost when a subagent finishes",
|
||||
name: "invalidates the root's tree cost when a subagent finishes",
|
||||
updatedChat: chatForFilterInvalidation({
|
||||
id: "child-1",
|
||||
parent_chat_id: "root-1",
|
||||
@@ -964,10 +955,10 @@ describe(chatCostIdsToInvalidate.name, () => {
|
||||
status: "waiting",
|
||||
}),
|
||||
eventKind: "status_change",
|
||||
expected: ["child-1", "root-1"],
|
||||
expected: "root-1",
|
||||
},
|
||||
{
|
||||
name: "invalidates the parent and root when a nested subagent finishes",
|
||||
name: "invalidates the root's tree cost when a nested subagent finishes",
|
||||
updatedChat: chatForFilterInvalidation({
|
||||
id: "grandchild-1",
|
||||
parent_chat_id: "child-1",
|
||||
@@ -975,13 +966,26 @@ describe(chatCostIdsToInvalidate.name, () => {
|
||||
status: "waiting",
|
||||
}),
|
||||
eventKind: "status_change",
|
||||
expected: ["grandchild-1", "child-1", "root-1"],
|
||||
expected: "root-1",
|
||||
},
|
||||
{
|
||||
// Deleting a root nulls root_chat_id on descendants, leaving only
|
||||
// parent_chat_id, so cost is keyed on the parent.
|
||||
name: "falls back to the parent when the root chat is gone",
|
||||
updatedChat: chatForFilterInvalidation({
|
||||
id: "grandchild-1",
|
||||
parent_chat_id: "child-1",
|
||||
root_chat_id: undefined,
|
||||
status: "waiting",
|
||||
}),
|
||||
eventKind: "status_change",
|
||||
expected: "child-1",
|
||||
},
|
||||
{
|
||||
name: "waits while the chat is still active",
|
||||
updatedChat: chatForFilterInvalidation({ status: "running" }),
|
||||
eventKind: "status_change",
|
||||
expected: [],
|
||||
expected: undefined,
|
||||
},
|
||||
{
|
||||
name: "waits while a subagent is still active",
|
||||
@@ -992,21 +996,61 @@ describe(chatCostIdsToInvalidate.name, () => {
|
||||
status: "running",
|
||||
}),
|
||||
eventKind: "status_change",
|
||||
expected: [],
|
||||
expected: undefined,
|
||||
},
|
||||
{
|
||||
name: "waits while the chat is interrupting",
|
||||
updatedChat: chatForFilterInvalidation({ status: "interrupting" }),
|
||||
eventKind: "status_change",
|
||||
expected: [],
|
||||
expected: undefined,
|
||||
},
|
||||
{
|
||||
name: "ignores non-status events",
|
||||
name: "ignores events that bill no gateway request",
|
||||
updatedChat: chatForFilterInvalidation({ status: "waiting" }),
|
||||
eventKind: "diff_status_change",
|
||||
expected: undefined,
|
||||
},
|
||||
{
|
||||
name: "invalidates when a generated title lands on an idle chat",
|
||||
updatedChat: chatForFilterInvalidation({ status: "waiting" }),
|
||||
eventKind: "title_change",
|
||||
expected: "chat-1",
|
||||
},
|
||||
{
|
||||
name: "invalidates the root's tree cost for a subagent title change",
|
||||
updatedChat: chatForFilterInvalidation({
|
||||
id: "child-1",
|
||||
parent_chat_id: "root-1",
|
||||
root_chat_id: "root-1",
|
||||
status: "running",
|
||||
}),
|
||||
eventKind: "title_change",
|
||||
expected: "root-1",
|
||||
},
|
||||
{
|
||||
name: "invalidates when a generated turn status label lands",
|
||||
updatedChat: chatForFilterInvalidation({ status: "waiting" }),
|
||||
eventKind: "summary_change",
|
||||
expected: [],
|
||||
expected: "chat-1",
|
||||
},
|
||||
{
|
||||
name: "invalidates when a generated whole-chat summary lands",
|
||||
updatedChat: chatForFilterInvalidation({ status: "waiting" }),
|
||||
eventKind: "chat_summary_change",
|
||||
expected: "chat-1",
|
||||
},
|
||||
{
|
||||
name: "invalidates the root's tree cost for a subagent summary change",
|
||||
updatedChat: chatForFilterInvalidation({
|
||||
id: "child-1",
|
||||
parent_chat_id: "root-1",
|
||||
root_chat_id: "root-1",
|
||||
status: "running",
|
||||
}),
|
||||
eventKind: "chat_summary_change",
|
||||
expected: "root-1",
|
||||
},
|
||||
])("$name", ({ updatedChat, eventKind, expected }) => {
|
||||
expect(chatCostIdsToInvalidate(updatedChat, eventKind)).toEqual(expected);
|
||||
expect(chatCostIdToInvalidate(updatedChat, eventKind)).toBe(expected);
|
||||
});
|
||||
});
|
||||
|
||||
@@ -59,6 +59,7 @@ import { cn } from "#/utils/cn";
|
||||
import { pageTitle } from "#/utils/page";
|
||||
import { createReconnectingWebSocket } from "#/utils/reconnectingWebSocket";
|
||||
import { emptyInputStorageKey } from "./components/AgentCreateForm";
|
||||
import { getChatCostTreeID } from "./components/ChatConversation/chatHelpers";
|
||||
import { isActiveChatStatus } from "./components/ChatConversation/chatStore";
|
||||
import {
|
||||
ChatsSidebar,
|
||||
@@ -126,32 +127,25 @@ export const shouldInvalidateFilteredChatList = (
|
||||
): boolean =>
|
||||
!chat.parent_chat_id && FILTER_MEMBERSHIP_EVENT_KINDS.has(eventKind);
|
||||
|
||||
// Chat IDs whose cost queries must refetch after a watch event, or an
|
||||
// empty array when the event cannot change any cost. Cost accrues while
|
||||
// a chat generates, so refetch when a status change lands in a
|
||||
// non-active status. The cost endpoint sums the requested chat's
|
||||
// subtree (GetChatModelUsageCostByChatID walks parent_chat_id), so a
|
||||
// subagent going idle must also refresh its ancestors' rolled-up
|
||||
// totals. The watch payload only carries the immediate parent and the
|
||||
// root, which covers every ancestor for nesting up to two levels deep;
|
||||
// deeper intermediate ancestors are refreshed by the query staleTime.
|
||||
export const chatCostIdsToInvalidate = (
|
||||
// Summary and title generation can bill after the turn reports a non-active
|
||||
// status, so invalidate the root-keyed cost query when those events arrive.
|
||||
const POST_TURN_BILLED_EVENT_KINDS = new Set<TypesGen.ChatWatchEventKind>([
|
||||
"chat_summary_change",
|
||||
"summary_change",
|
||||
"title_change",
|
||||
]);
|
||||
|
||||
export const chatCostIdToInvalidate = (
|
||||
chat: TypesGen.Chat,
|
||||
eventKind: TypesGen.ChatWatchEventKind,
|
||||
): readonly string[] => {
|
||||
): string | undefined => {
|
||||
if (POST_TURN_BILLED_EVENT_KINDS.has(eventKind)) {
|
||||
return getChatCostTreeID(chat);
|
||||
}
|
||||
if (eventKind !== "status_change" || isActiveChatStatus(chat.status)) {
|
||||
return [];
|
||||
return undefined;
|
||||
}
|
||||
// root_chat_id is self-referential on root chats and parent_chat_id
|
||||
// equals root_chat_id at depth one; the set dedupes both cases.
|
||||
const ids = new Set([chat.id]);
|
||||
if (chat.parent_chat_id) {
|
||||
ids.add(chat.parent_chat_id);
|
||||
}
|
||||
if (chat.root_chat_id) {
|
||||
ids.add(chat.root_chat_id);
|
||||
}
|
||||
return [...ids];
|
||||
return getChatCostTreeID(chat);
|
||||
};
|
||||
|
||||
const AgentsPageLayout: FC = () => {
|
||||
@@ -683,10 +677,11 @@ const AgentsPageLayout: FC = () => {
|
||||
if (shouldInvalidateFilteredChatList(updatedChat, chatEvent.kind)) {
|
||||
void invalidateChatListQueries(queryClient);
|
||||
}
|
||||
for (const costChatId of chatCostIdsToInvalidate(
|
||||
const costChatId = chatCostIdToInvalidate(
|
||||
updatedChat,
|
||||
chatEvent.kind,
|
||||
)) {
|
||||
);
|
||||
if (costChatId) {
|
||||
void queryClient.invalidateQueries({
|
||||
queryKey: chatCostKey(costChatId),
|
||||
exact: true,
|
||||
|
||||
@@ -70,6 +70,20 @@ export const getParentChatID = (
|
||||
return asNonEmptyString(chat?.parent_chat_id);
|
||||
};
|
||||
|
||||
/**
|
||||
* Identifies the chat tree that AI Gateway cost is aggregated over, matching
|
||||
* the server's COALESCE(root_chat_id, parent_chat_id) precedence. Both columns
|
||||
* are ON DELETE SET NULL, so deleting a root leaves descendants with only a
|
||||
* parent. Cost readers and cost invalidators must agree, or a mounted cost
|
||||
* goes stale.
|
||||
*/
|
||||
export const getChatCostTreeID = (
|
||||
chat: TypesGen.Chat | undefined,
|
||||
): string | undefined =>
|
||||
asNonEmptyString(chat?.root_chat_id) ??
|
||||
asNonEmptyString(chat?.parent_chat_id) ??
|
||||
asNonEmptyString(chat?.id);
|
||||
|
||||
export const resolveModelFromChatConfig = (
|
||||
modelConfig: unknown,
|
||||
modelOptions: readonly ModelSelectorOption[],
|
||||
|
||||
@@ -11,6 +11,7 @@ const meta: Meta<typeof ChatSummary> = {
|
||||
createdAt: "2024-05-01T12:00:00Z",
|
||||
updatedAt: "2024-05-02T15:30:00Z",
|
||||
costMicros: 1_250_000,
|
||||
showCost: true,
|
||||
},
|
||||
decorators: [
|
||||
(Story) => (
|
||||
@@ -93,13 +94,31 @@ export const CostError: Story = {
|
||||
};
|
||||
|
||||
export const PartialCost: Story = {
|
||||
args: { costMicros: 0, unpricedMessagesHavingUsageCount: 3 },
|
||||
args: { costMicros: 0, unpricedRequestCount: 3 },
|
||||
play: async ({ canvasElement }) => {
|
||||
const canvas = within(canvasElement);
|
||||
await expect(
|
||||
canvas.getByText(
|
||||
"Excludes 3 messages with usage but without model pricing.",
|
||||
),
|
||||
canvas.getByText("Excludes unpriced usage from 3 requests."),
|
||||
).toBeInTheDocument();
|
||||
},
|
||||
};
|
||||
|
||||
export const SubagentTreeCost: Story = {
|
||||
args: { isSubagent: true },
|
||||
play: async ({ canvasElement }) => {
|
||||
const canvas = within(canvasElement);
|
||||
await expect(
|
||||
canvas.getByText(/Cost covers this agent's whole chat/),
|
||||
).toBeInTheDocument();
|
||||
},
|
||||
};
|
||||
|
||||
export const CostHidden: Story = {
|
||||
args: { showCost: false },
|
||||
play: async ({ canvasElement }) => {
|
||||
const canvas = within(canvasElement);
|
||||
await expect(canvas.getByText("Updated:")).toBeInTheDocument();
|
||||
await expect(canvas.queryByText("Cost:")).not.toBeInTheDocument();
|
||||
await expect(canvas.queryByText("$1.25")).not.toBeInTheDocument();
|
||||
},
|
||||
};
|
||||
|
||||
@@ -9,12 +9,13 @@ interface ChatSummaryProps {
|
||||
summary: string | null;
|
||||
createdAt: string;
|
||||
updatedAt: string;
|
||||
/** Cumulative chat cost in microdollars (1 USD = 1,000,000). */
|
||||
/** Cost of the whole chat tree in microdollars (1 USD = 1,000,000). */
|
||||
costMicros?: number | null;
|
||||
isCostLoading?: boolean;
|
||||
costError?: boolean;
|
||||
/** Assistant messages with usage but no model pricing; when > 0 the cost is partial and a note is shown. */
|
||||
unpricedMessagesHavingUsageCount?: number;
|
||||
/** Requests with usage the gateway could not price, so the reported cost is partial. */
|
||||
unpricedRequestCount?: number;
|
||||
showCost: boolean;
|
||||
/** Subagent summaries are the agent's final report, persisted when it completes, so the empty state reads as pending rather than absent. */
|
||||
isSubagent?: boolean;
|
||||
}
|
||||
@@ -26,16 +27,15 @@ export const ChatSummary: FC<ChatSummaryProps> = ({
|
||||
costMicros,
|
||||
isCostLoading,
|
||||
costError,
|
||||
unpricedMessagesHavingUsageCount,
|
||||
unpricedRequestCount,
|
||||
showCost,
|
||||
isSubagent,
|
||||
}) => {
|
||||
const trimmedSummary = summary?.trim();
|
||||
const hasUnpricedMessages =
|
||||
!isCostLoading &&
|
||||
!costError &&
|
||||
costMicros != null &&
|
||||
unpricedMessagesHavingUsageCount != null &&
|
||||
unpricedMessagesHavingUsageCount > 0;
|
||||
const hasCost =
|
||||
showCost && !isCostLoading && !costError && costMicros != null;
|
||||
const hasUnpricedRequests =
|
||||
hasCost && unpricedRequestCount != null && unpricedRequestCount > 0;
|
||||
|
||||
return (
|
||||
<div className="flex flex-col gap-4">
|
||||
@@ -56,24 +56,32 @@ export const ChatSummary: FC<ChatSummaryProps> = ({
|
||||
<ChatSummaryRow label="Updated:">
|
||||
{formatDateTime(updatedAt, DATE_FORMAT.MEDIUM_DATE)}
|
||||
</ChatSummaryRow>
|
||||
<ChatSummaryRow label="Cost:">
|
||||
{isCostLoading ? (
|
||||
<Skeleton aria-label="Loading cost" className="my-1 h-4 w-16" />
|
||||
) : costError ? (
|
||||
<span className="text-content-secondary">Unavailable</span>
|
||||
) : costMicros != null ? (
|
||||
formatCostMicros(costMicros)
|
||||
) : (
|
||||
EMPTY_VALUE
|
||||
)}
|
||||
</ChatSummaryRow>
|
||||
{showCost && (
|
||||
<ChatSummaryRow label="Cost:">
|
||||
{isCostLoading ? (
|
||||
<Skeleton aria-label="Loading cost" className="my-1 h-4 w-16" />
|
||||
) : costError ? (
|
||||
<span className="text-content-secondary">Unavailable</span>
|
||||
) : costMicros != null ? (
|
||||
formatCostMicros(costMicros)
|
||||
) : (
|
||||
EMPTY_VALUE
|
||||
)}
|
||||
</ChatSummaryRow>
|
||||
)}
|
||||
</dl>
|
||||
|
||||
{hasUnpricedMessages && (
|
||||
{isSubagent && hasCost && (
|
||||
<p className="m-0 text-xs italic text-content-secondary">
|
||||
Excludes {unpricedMessagesHavingUsageCount} message
|
||||
{unpricedMessagesHavingUsageCount === 1 ? "" : "s"} with usage but
|
||||
without model pricing.
|
||||
Cost covers this agent's whole chat, including the chat that started
|
||||
it and any other subagents.
|
||||
</p>
|
||||
)}
|
||||
|
||||
{hasUnpricedRequests && (
|
||||
<p className="m-0 text-xs italic text-content-secondary">
|
||||
Excludes unpriced usage from {unpricedRequestCount} request
|
||||
{unpricedRequestCount === 1 ? "" : "s"}.
|
||||
</p>
|
||||
)}
|
||||
</div>
|
||||
|
||||
@@ -7,11 +7,13 @@ import { MockChat } from "#/testHelpers/chatEntities";
|
||||
import { withDashboardProvider } from "#/testHelpers/storybook";
|
||||
import { ChatSummaryPanel } from "./ChatSummaryPanel";
|
||||
|
||||
const ROOT_CHAT_ID = "root-chat-id";
|
||||
|
||||
const mockCost: TypesGen.ChatCost = {
|
||||
chat_id: MockChat.id,
|
||||
total_cost_micros: 1_250_000,
|
||||
priced_message_count: 8,
|
||||
unpriced_messages_having_usage_count: 0,
|
||||
request_count: 8,
|
||||
unpriced_request_count: 0,
|
||||
};
|
||||
|
||||
type MockRequestOptions = {
|
||||
@@ -19,6 +21,7 @@ type MockRequestOptions = {
|
||||
summary?: string | null;
|
||||
chatError?: boolean;
|
||||
parentChatId?: string;
|
||||
rootChatId?: string;
|
||||
};
|
||||
|
||||
const mockRequests = ({
|
||||
@@ -26,6 +29,7 @@ const mockRequests = ({
|
||||
summary = null,
|
||||
chatError,
|
||||
parentChatId,
|
||||
rootChatId,
|
||||
}: MockRequestOptions = {}) => {
|
||||
if (chatError) {
|
||||
spyOn(API.experimental, "getChat").mockRejectedValue(
|
||||
@@ -36,6 +40,7 @@ const mockRequests = ({
|
||||
...MockChat,
|
||||
summary,
|
||||
...(parentChatId ? { parent_chat_id: parentChatId } : {}),
|
||||
...(rootChatId ? { root_chat_id: rootChatId } : {}),
|
||||
});
|
||||
}
|
||||
|
||||
@@ -53,6 +58,7 @@ const meta: Meta<typeof ChatSummaryPanel> = {
|
||||
title: "pages/AgentsPage/ChatSummaryPanel",
|
||||
component: ChatSummaryPanel,
|
||||
decorators: [PanelFrame, withDashboardProvider],
|
||||
parameters: { features: ["aibridge"] satisfies TypesGen.FeatureName[] },
|
||||
args: {
|
||||
chatId: MockChat.id,
|
||||
isVisible: true,
|
||||
@@ -93,6 +99,26 @@ export const SubagentSummaryPending: Story = {
|
||||
},
|
||||
};
|
||||
|
||||
export const SubagentTreeCost: Story = {
|
||||
beforeEach: () =>
|
||||
mockRequests({
|
||||
parentChatId: "parent-chat-id",
|
||||
rootChatId: ROOT_CHAT_ID,
|
||||
cost: { ...mockCost, chat_id: ROOT_CHAT_ID },
|
||||
}),
|
||||
play: async ({ canvasElement }) => {
|
||||
const canvas = within(canvasElement);
|
||||
await waitFor(() => {
|
||||
expect(canvas.getByText("$1.25")).toBeInTheDocument();
|
||||
});
|
||||
expect(
|
||||
canvas.getByText(/Cost covers this agent's whole chat/),
|
||||
).toBeInTheDocument();
|
||||
expect(API.experimental.getChatCost).toHaveBeenCalledWith(ROOT_CHAT_ID);
|
||||
expect(API.experimental.getChatCost).not.toHaveBeenCalledWith(MockChat.id);
|
||||
},
|
||||
};
|
||||
|
||||
export const ChatError: Story = {
|
||||
beforeEach: () => mockRequests({ chatError: true }),
|
||||
play: async ({ canvasElement }) => {
|
||||
@@ -117,3 +143,16 @@ export const NotVisible: Story = {
|
||||
expect(canvas.queryByText("No summary yet.")).not.toBeInTheDocument();
|
||||
},
|
||||
};
|
||||
|
||||
export const GatewayUnavailable: Story = {
|
||||
parameters: { features: [] },
|
||||
beforeEach: () => mockRequests({ summary: "Gateway is off here." }),
|
||||
play: async ({ canvasElement }) => {
|
||||
const canvas = within(canvasElement);
|
||||
await waitFor(() => {
|
||||
expect(canvas.getByText("Gateway is off here.")).toBeInTheDocument();
|
||||
});
|
||||
expect(canvas.queryByText("Cost:")).not.toBeInTheDocument();
|
||||
expect(API.experimental.getChatCost).not.toHaveBeenCalled();
|
||||
},
|
||||
};
|
||||
|
||||
@@ -2,6 +2,8 @@ import type { FC, ReactNode } from "react";
|
||||
import { useQuery } from "react-query";
|
||||
import { chat, chatCost } from "#/api/queries/chats";
|
||||
import { ErrorAlert } from "#/components/Alert/ErrorAlert";
|
||||
import { useFeatureVisibility } from "#/modules/dashboard/useFeatureVisibility";
|
||||
import { getChatCostTreeID } from "./ChatConversation/chatHelpers";
|
||||
import { ChatSummary } from "./ChatSummary";
|
||||
|
||||
type ChatSummaryPanelProps = {
|
||||
@@ -14,10 +16,15 @@ export const ChatSummaryPanel: FC<ChatSummaryPanelProps> = ({
|
||||
chatId,
|
||||
isVisible,
|
||||
}) => {
|
||||
const showCost = Boolean(useFeatureVisibility().aibridge);
|
||||
const chatQuery = useQuery({ ...chat(chatId), enabled: isVisible });
|
||||
const costQuery = useQuery({ ...chatCost(chatId), enabled: isVisible });
|
||||
|
||||
const chatData = chatQuery.data;
|
||||
const rootChatId = getChatCostTreeID(chatData) ?? chatId;
|
||||
const costQuery = useQuery({
|
||||
...chatCost(rootChatId),
|
||||
enabled: isVisible && showCost && chatData !== undefined,
|
||||
});
|
||||
|
||||
let content: ReactNode = null;
|
||||
if (chatQuery.isError) {
|
||||
@@ -30,9 +37,8 @@ export const ChatSummaryPanel: FC<ChatSummaryPanelProps> = ({
|
||||
createdAt={chatData.created_at}
|
||||
updatedAt={chatData.updated_at}
|
||||
costMicros={costQuery.data?.total_cost_micros}
|
||||
unpricedMessagesHavingUsageCount={
|
||||
costQuery.data?.unpriced_messages_having_usage_count
|
||||
}
|
||||
unpricedRequestCount={costQuery.data?.unpriced_request_count}
|
||||
showCost={showCost}
|
||||
isCostLoading={costQuery.isLoading}
|
||||
costError={costQuery.isError}
|
||||
/>
|
||||
|
||||
Reference in New Issue
Block a user