From eafbc148e1102ecac78f8d634d6e1d4279354b40 Mon Sep 17 00:00:00 2001 From: Michael Suchacz <203725896+ibetitsmike@users.noreply.github.com> Date: Tue, 28 Jul 2026 09:38:29 +0000 Subject: [PATCH 01/10] feat: back the per-chat cost endpoint with AI Gateway data The per-chat cost endpoint summed chat_messages.total_cost_micros, which native chat cost tracking maintained. Repoint it at AI Gateway interception data so the AI Gateway becomes the single source of AI spend. AI Gateway attributes a subagent's requests to the chat that spawned it, so cost is only meaningful for a whole chat tree. The endpoint now resolves the root chat and reports the tree total, and the sidebar keys its cache by root. Counts become requests rather than messages, which is what the gateway records, so priced_message_count becomes request_count and unpriced_messages_having_usage_count becomes unpriced_request_count. The cost row is hidden where the AI Gateway is off or unlicensed, since no interception data exists there. --- coderd/apidoc/docs.go | 6 +- coderd/apidoc/swagger.json | 6 +- coderd/database/dbauthz/dbauthz.go | 10 + coderd/database/dbauthz/dbauthz_test.go | 7 + coderd/database/dbmetrics/querymetrics.go | 8 + coderd/database/dbmock/dbmock.go | 15 ++ coderd/database/querier.go | 8 + coderd/database/queries.sql.go | 40 +++ coderd/database/queries/aibridge.sql | 26 ++ coderd/exp_chats.go | 26 +- coderd/exp_chats_internal_test.go | 27 +- coderd/exp_chats_test.go | 238 ++++++++++-------- codersdk/chats.go | 19 +- docs/reference/api/chats.md | 9 +- docs/reference/api/schemas.md | 16 +- site/src/api/queries/chats.ts | 23 +- site/src/api/typesGenerated.ts | 12 +- .../AgentsPage/AgentChatPageView.stories.tsx | 4 +- .../pages/AgentsPage/AgentsPageLayout.test.ts | 52 ++-- .../src/pages/AgentsPage/AgentsPageLayout.tsx | 41 ++- .../components/ChatSummary.stories.tsx | 26 +- .../AgentsPage/components/ChatSummary.tsx | 59 +++-- .../components/ChatSummaryPanel.stories.tsx | 26 +- .../components/ChatSummaryPanel.tsx | 23 +- 24 files changed, 476 insertions(+), 251 deletions(-) diff --git a/coderd/apidoc/docs.go b/coderd/apidoc/docs.go index 388c35f1144..1f2af9f179a 100644 --- a/coderd/apidoc/docs.go +++ b/coderd/apidoc/docs.go @@ -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 a subagent chat returns\nthat same tree total, because AI Gateway attributes a\nsubagent's requests to the chat that spawned it.", "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" } } diff --git a/coderd/apidoc/swagger.json b/coderd/apidoc/swagger.json index 4553c63fc92..34e3ffc376a 100644 --- a/coderd/apidoc/swagger.json +++ b/coderd/apidoc/swagger.json @@ -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 a subagent chat returns\nthat same tree total, because AI Gateway attributes a\nsubagent's requests to the chat that spawned it.", "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" } } diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go index de153b24c2e..80557d8bb81 100644 --- a/coderd/database/dbauthz/dbauthz.go +++ b/coderd/database/dbauthz/dbauthz.go @@ -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) } diff --git a/coderd/database/dbauthz/dbauthz_test.go b/coderd/database/dbauthz/dbauthz_test.go index 07278858db4..a8bc4efaa1d 100644 --- a/coderd/database/dbauthz/dbauthz_test.go +++ b/coderd/database/dbauthz/dbauthz_test.go @@ -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} diff --git a/coderd/database/dbmetrics/querymetrics.go b/coderd/database/dbmetrics/querymetrics.go index 26f47ff51dd..0480e54a329 100644 --- a/coderd/database/dbmetrics/querymetrics.go +++ b/coderd/database/dbmetrics/querymetrics.go @@ -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) diff --git a/coderd/database/dbmock/dbmock.go b/coderd/database/dbmock/dbmock.go index ac848e0f993..d21ee4b2262 100644 --- a/coderd/database/dbmock/dbmock.go +++ b/coderd/database/dbmock/dbmock.go @@ -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() diff --git a/coderd/database/querier.go b/coderd/database/querier.go index eb46ff6a8f8..fb7b49261f4 100644 --- a/coderd/database/querier.go +++ b/coderd/database/querier.go @@ -295,6 +295,14 @@ 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 chat + // beneath it. Coder Agents traffic records the spawning chat's ID as the + // interception session ID (chatprovider.CoderHeaders), so a subagent's + // requests are attributed to its parent and only the whole tree can be + // summed. The owner check guards against session-id collisions from other + // users. Usage without an effective group never reaches ai_user_daily_spend, + // so excluding it keeps this total consistent with AI budget 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. diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index 3fb67221f0f..5a4a8b78c9d 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -1156,6 +1156,46 @@ func (q *sqlQuerier) DeleteOldAIBridgeRecords(ctx context.Context, beforeTime ti return total_deleted, err } +const getAIBridgeChatCost = `-- name: GetAIBridgeChatCost :one +SELECT + COALESCE(SUM(tu.cost_micros), 0)::bigint AS total_cost_micros, + COUNT(DISTINCT i.id)::bigint AS request_count, + COUNT(DISTINCT i.id) FILTER (WHERE tu.cost_micros IS NULL)::bigint AS unpriced_request_count +FROM aibridge_interceptions i +JOIN chats c ON c.id::text = i.session_id AND c.owner_id = i.initiator_id +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 the planner + -- can use idx_chats_root_chat_id and the chats primary key. + c.root_chat_id = $1::uuid + OR (c.root_chat_id IS NULL AND c.id = $1::uuid) + ) + -- aibridge.ClientCoderAgents. Restricting the client keeps another + -- client's session reference from matching a chat ID. + AND i.client = 'Coder Agents' + AND i.ended_at IS NOT NULL +` + +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 chat +// beneath it. Coder Agents traffic records the spawning chat's ID as the +// interception session ID (chatprovider.CoderHeaders), so a subagent's +// requests are attributed to its parent and only the whole tree can be +// summed. The owner check guards against session-id collisions from other +// users. Usage without an effective group never reaches ai_user_daily_spend, +// so excluding it keeps this total consistent with AI budget 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 diff --git a/coderd/database/queries/aibridge.sql b/coderd/database/queries/aibridge.sql index 63635b6ae22..e06a3a105f5 100644 --- a/coderd/database/queries/aibridge.sql +++ b/coderd/database/queries/aibridge.sql @@ -663,3 +663,29 @@ 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 chat +-- beneath it. Coder Agents traffic records the spawning chat's ID as the +-- interception session ID (chatprovider.CoderHeaders), so a subagent's +-- requests are attributed to its parent and only the whole tree can be +-- summed. The owner check guards against session-id collisions from other +-- users. Usage without an effective group never reaches ai_user_daily_spend, +-- so excluding it keeps this total consistent with AI budget spend. +SELECT + COALESCE(SUM(tu.cost_micros), 0)::bigint AS total_cost_micros, + COUNT(DISTINCT i.id)::bigint AS request_count, + COUNT(DISTINCT i.id) FILTER (WHERE tu.cost_micros IS NULL)::bigint AS unpriced_request_count +FROM aibridge_interceptions i +JOIN chats c ON c.id::text = i.session_id AND c.owner_id = i.initiator_id +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 the planner + -- can use idx_chats_root_chat_id and the chats primary key. + c.root_chat_id = @root_chat_id::uuid + OR (c.root_chat_id IS NULL AND c.id = @root_chat_id::uuid) + ) + -- aibridge.ClientCoderAgents. Restricting the client keeps another + -- client's session reference from matching a chat ID. + AND i.client = 'Coder Agents' + AND i.ended_at IS NOT NULL; diff --git a/coderd/exp_chats.go b/coderd/exp_chats.go index a6f54bf80d3..0b21e6c7077 100644 --- a/coderd/exp_chats.go +++ b/coderd/exp_chats.go @@ -2472,16 +2472,26 @@ 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 a subagent chat returns +// @Description that same tree total, because AI Gateway attributes a +// @Description subagent's requests to the chat that spawned it. // //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. + rootChatID := chat.ID + if chat.RootChatID.Valid { + rootChatID = chat.RootChatID.UUID + } + + row, err := api.Database.GetAIBridgeChatCost(ctx, rootChatID) if err != nil { if httpapi.Is404Error(err) { httpapi.ResourceNotFound(rw) @@ -2495,10 +2505,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, }) } diff --git a/coderd/exp_chats_internal_test.go b/coderd/exp_chats_internal_test.go index bcd7624438e..e147bd0665e 100644 --- a/coderd/exp_chats_internal_test.go +++ b/coderd/exp_chats_internal_test.go @@ -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,10 @@ 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: 1, }, nil, ) @@ -97,7 +96,7 @@ 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(1), cost.RequestCount) } func TestEnrichMissingChatAgentIDs(t *testing.T) { diff --git a/coderd/exp_chats_test.go b/coderd/exp_chats_test.go index 07a491a5ae5..aec38670c4c 100644 --- a/coderd/exp_chats_test.go +++ b/coderd/exp_chats_test.go @@ -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,38 @@ func TestChatCostSummary_AdminDrilldown(t *testing.T) { }) } -func TestGetChatCost(t *testing.T) { - t.Parallel() - - t.Run("BasicCost", func(t *testing.T) { - t.Parallel() +// seedChatGatewayRequest records one finished Coder Agents gateway request +// under sessionChatID, mirroring what aibridged persists for chatd traffic: +// the interception's session ID is the chat that spawned the request. +func seedChatGatewayRequest(t *testing.T, db database.Store, initiatorID, sessionChatID uuid.UUID, usage database.InsertAIBridgeTokenUsageParams) { + t.Helper() - f := seedChatCostFixture(t) - ctx := testutil.Context(t, testutil.WaitLong) + 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) + + usage.InterceptionID = interception.ID + usage.CreatedAt = now + dbgen.AIBridgeTokenUsage(t, db, usage) +} - 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) - }) +func TestGetChatCost(t *testing.T) { + t.Parallel() - 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 +11657,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 +11665,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,44 +11673,41 @@ 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, @@ -11724,19 +11715,12 @@ func TestGetChatCost(t *testing.T) { LastModelConfigID: modelConfig.ID, Title: "unpriced 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}, + seedChatGatewayRequest(t, db, firstUser.UserID, chat.ID, database.InsertAIBridgeTokenUsageParams{ + EffectiveGroupID: everyoneGroup, + CostMicros: 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}, + seedChatGatewayRequest(t, db, firstUser.UserID, chat.ID, database.InsertAIBridgeTokenUsageParams{ + EffectiveGroupID: everyoneGroup, }) ctx := testutil.Context(t, testutil.WaitLong) @@ -11744,94 +11728,140 @@ func TestGetChatCost(t *testing.T) { 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(2), cost.RequestCount) + require.Equal(t, int64(1), cost.UnpricedRequestCount) }) - t.Run("MemberCannotReadOtherUsersChat", func(t *testing.T) { + t.Run("ExcludesUnattributedUsage", func(t *testing.T) { t.Parallel() client, db := newChatClientWithDatabase(t) firstUser := coderdtest.CreateFirstUser(t, client.Client) - memberClientRaw, _ := coderdtest.CreateAnotherUser(t, client.Client, firstUser.OrganizationID) - memberClient := codersdk.NewExperimentalClient(memberClientRaw) modelConfig := createChatModelConfig(t, client) chat := dbgen.Chat(t, db, database.Chat{ OrganizationID: firstUser.OrganizationID, OwnerID: firstUser.UserID, LastModelConfigID: modelConfig.ID, - Title: "owner chat", + Title: "legacy chat", + }) + // Usage recorded before group attribution existed never reached + // ai_user_daily_spend, so it must not appear as chat spend either. + seedChatGatewayRequest(t, db, firstUser.UserID, chat.ID, database.InsertAIBridgeTokenUsageParams{ + CostMicros: sql.NullInt64{Int64: 900, Valid: true}, }) ctx := testutil.Context(t, testutil.WaitLong) - _, err := memberClient.GetChatCost(ctx, chat.ID) - require.Error(t, err) - var sdkErr *codersdk.Error - require.ErrorAs(t, err, &sdkErr) - require.Equal(t, http.StatusNotFound, sdkErr.StatusCode()) + 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("ZeroMessages", func(t *testing.T) { + 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} - // No assistant messages must still return one zero-total row; a COALESCE - // or :one regression would surface as sql.ErrNoRows -> 500. chat := dbgen.Chat(t, db, database.Chat{ OrganizationID: firstUser.OrganizationID, OwnerID: firstUser.UserID, LastModelConfigID: modelConfig.ID, - Title: "empty chat", + 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, 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) + require.Equal(t, int64(0), cost.RequestCount) }) - t.Run("ExcludesNonAssistantMessages", func(t *testing.T) { + t.Run("MemberCannotReadOtherUsersChat", func(t *testing.T) { t.Parallel() client, db := newChatClientWithDatabase(t) firstUser := coderdtest.CreateFirstUser(t, client.Client) + memberClientRaw, _ := coderdtest.CreateAnotherUser(t, client.Client, firstUser.OrganizationID) + memberClient := codersdk.NewExperimentalClient(memberClientRaw) 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}, + Title: "owner chat", }) - // 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) + + _, err := memberClient.GetChatCost(ctx, chat.ID) + require.Error(t, err) + var sdkErr *codersdk.Error + require.ErrorAs(t, err, &sdkErr) + require.Equal(t, http.StatusNotFound, sdkErr.StatusCode()) + }) + + t.Run("ZeroRequests", func(t *testing.T) { + t.Parallel() + + client, db := newChatClientWithDatabase(t) + firstUser := coderdtest.CreateFirstUser(t, client.Client) + modelConfig := createChatModelConfig(t, client) + + // 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, + LastModelConfigID: modelConfig.ID, + Title: "empty chat", }) 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, chat.ID, cost.ChatID) + require.Equal(t, int64(0), cost.TotalCostMicros) + require.Equal(t, int64(0), cost.RequestCount) + require.Equal(t, int64(0), cost.UnpricedRequestCount) }) } diff --git a/codersdk/chats.go b/codersdk/chats.go index 3afbaff70da..0049dde0505 100644 --- a/codersdk/chats.go +++ b/codersdk/chats.go @@ -1990,15 +1990,15 @@ 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. AI +// Gateway attributes a subagent's requests to the chat that spawned it, so +// every chat in a tree reports the same total. UnpricedRequestCount counts +// requests whose model had no recorded price, so the total excludes them. 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 +2592,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 { diff --git a/docs/reference/api/chats.md b/docs/reference/api/chats.md index 2362d3dc701..9ecbf10b2ec 100644 --- a/docs/reference/api/chats.md +++ b/docs/reference/api/chats.md @@ -1298,6 +1298,11 @@ 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 a subagent chat returns +that same tree total, because AI Gateway attributes a +subagent's requests to the chat that spawned it. + ### Parameters | Name | In | Type | Required | Description | @@ -1311,9 +1316,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 } ``` diff --git a/docs/reference/api/schemas.md b/docs/reference/api/schemas.md index f8295c7d1ba..3983f0a41e4 100644 --- a/docs/reference/api/schemas.md +++ b/docs/reference/api/schemas.md @@ -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 diff --git a/site/src/api/queries/chats.ts b/site/src/api/queries/chats.ts index c382e11e7bd..660702e483e 100644 --- a/site/src/api/queries/chats.ts +++ b/site/src/api/queries/chats.ts @@ -1966,17 +1966,18 @@ export const chatCostSummary = (user = "me", params?: ChatCostDateParams) => ({ staleTime: 60_000, }); -export const chatCostKey = (chatId: string) => - [...chatsKey, chatId, "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; - -export const chatCost = (chatId: string) => ({ - queryKey: chatCostKey(chatId), - queryFn: () => API.experimental.getChatCost(chatId), - staleTime: ASSISTANT_MESSAGE_PRICING_STALE_MS, +// Cost covers the whole chat tree, so callers key this by the root chat. +export const chatCostKey = (rootChatId: string) => + [...chatsKey, rootChatId, "cost"] as const; + +// Chat cost changes only when a gateway request completes, so a short stale +// window refreshes the sidebar without refetching on every render. +const GATEWAY_REQUEST_STALE_MS = 30_000; + +export const chatCost = (rootChatId: string) => ({ + queryKey: chatCostKey(rootChatId), + queryFn: () => API.experimental.getChatCost(rootChatId), + staleTime: GATEWAY_REQUEST_STALE_MS, }); interface PaginatedChatCostUsersPayload { diff --git a/site/src/api/typesGenerated.ts b/site/src/api/typesGenerated.ts index 5dadf27b2b2..ec2e791b67c 100644 --- a/site/src/api/typesGenerated.ts +++ b/site/src/api/typesGenerated.ts @@ -2155,16 +2155,16 @@ 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. AI + * Gateway attributes a subagent's requests to the chat that spawned it, so + * every chat in a tree reports the same total. UnpricedRequestCount counts + * requests whose model had no recorded price, so the total excludes them. */ 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 diff --git a/site/src/pages/AgentsPage/AgentChatPageView.stories.tsx b/site/src/pages/AgentsPage/AgentChatPageView.stories.tsx index e99359d7c4b..b7ff1c9f611 100644 --- a/site/src/pages/AgentsPage/AgentChatPageView.stories.tsx +++ b/site/src/pages/AgentsPage/AgentChatPageView.stories.tsx @@ -202,8 +202,8 @@ const meta: Meta = { spyOn(API.experimental, "getChatCost").mockResolvedValue({ chat_id: AGENT_ID, total_cost_micros: 0, - priced_message_count: 0, - unpriced_messages_having_usage_count: 0, + request_count: 0, + unpriced_request_count: 0, }); }, decorators: [withAuthProvider, withDashboardProvider, withProxyProvider()], diff --git a/site/src/pages/AgentsPage/AgentsPageLayout.test.ts b/site/src/pages/AgentsPage/AgentsPageLayout.test.ts index 0e7d76430d0..76744b29a99 100644 --- a/site/src/pages/AgentsPage/AgentsPageLayout.test.ts +++ b/site/src/pages/AgentsPage/AgentsPageLayout.test.ts @@ -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,13 @@ describe(chatCostIdsToInvalidate.name, () => { status: "waiting", }), eventKind: "status_change", - expected: ["grandchild-1", "child-1", "root-1"], + expected: "root-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 +983,38 @@ 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", updatedChat: chatForFilterInvalidation({ status: "waiting" }), eventKind: "summary_change", - expected: [], + 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", ({ updatedChat, eventKind, expected }) => { - expect(chatCostIdsToInvalidate(updatedChat, eventKind)).toEqual(expected); + expect(chatCostIdToInvalidate(updatedChat, eventKind)).toBe(expected); }); }); diff --git a/site/src/pages/AgentsPage/AgentsPageLayout.tsx b/site/src/pages/AgentsPage/AgentsPageLayout.tsx index 482bc789f85..fd205a3006d 100644 --- a/site/src/pages/AgentsPage/AgentsPageLayout.tsx +++ b/site/src/pages/AgentsPage/AgentsPageLayout.tsx @@ -126,32 +126,24 @@ 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 = ( +// Chat ID whose cost query must refetch after a watch event, or undefined +// 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. +// Title generation bills its own gateway request and can land while the chat +// is idle, so a title change refetches regardless of status. +// The cost endpoint reports the whole chat tree and the sidebar keys that +// query by root, so the root covers subagents at any depth. +export const chatCostIdToInvalidate = ( chat: TypesGen.Chat, eventKind: TypesGen.ChatWatchEventKind, -): readonly string[] => { - if (eventKind !== "status_change" || isActiveChatStatus(chat.status)) { - return []; - } - // 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); +): string | undefined => { + if (eventKind === "title_change") { + return chat.root_chat_id ?? chat.id; } - if (chat.root_chat_id) { - ids.add(chat.root_chat_id); + if (eventKind !== "status_change" || isActiveChatStatus(chat.status)) { + return undefined; } - return [...ids]; + return chat.root_chat_id ?? chat.id; }; const AgentsPageLayout: FC = () => { @@ -683,10 +675,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, diff --git a/site/src/pages/AgentsPage/components/ChatSummary.stories.tsx b/site/src/pages/AgentsPage/components/ChatSummary.stories.tsx index 00f93ff6089..d48d657d33d 100644 --- a/site/src/pages/AgentsPage/components/ChatSummary.stories.tsx +++ b/site/src/pages/AgentsPage/components/ChatSummary.stories.tsx @@ -93,13 +93,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 3 requests without model pricing."), ).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(); + }, +}; diff --git a/site/src/pages/AgentsPage/components/ChatSummary.tsx b/site/src/pages/AgentsPage/components/ChatSummary.tsx index a7169576c16..a4bfab58fe6 100644 --- a/site/src/pages/AgentsPage/components/ChatSummary.tsx +++ b/site/src/pages/AgentsPage/components/ChatSummary.tsx @@ -9,12 +9,14 @@ 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 whose model had no recorded price; when > 0 the cost is partial and a note is shown. */ + unpricedRequestCount?: number; + /** Cost comes from AI Gateway, so the row is hidden where the gateway is unavailable. */ + 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 +28,15 @@ export const ChatSummary: FC = ({ costMicros, isCostLoading, costError, - unpricedMessagesHavingUsageCount, + unpricedRequestCount, + showCost = true, 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 (
@@ -56,24 +57,32 @@ export const ChatSummary: FC = ({ {formatDateTime(updatedAt, DATE_FORMAT.MEDIUM_DATE)} - - {isCostLoading ? ( - - ) : costError ? ( - Unavailable - ) : costMicros != null ? ( - formatCostMicros(costMicros) - ) : ( - EMPTY_VALUE - )} - + {showCost && ( + + {isCostLoading ? ( + + ) : costError ? ( + Unavailable + ) : costMicros != null ? ( + formatCostMicros(costMicros) + ) : ( + EMPTY_VALUE + )} + + )} - {hasUnpricedMessages && ( + {isSubagent && hasCost && ( +

+ Cost covers this agent's whole chat, including the chat that started + it and any other subagents. +

+ )} + + {hasUnpricedRequests && (

- Excludes {unpricedMessagesHavingUsageCount} message - {unpricedMessagesHavingUsageCount === 1 ? "" : "s"} with usage but - without model pricing. + Excludes {unpricedRequestCount} request + {unpricedRequestCount === 1 ? "" : "s"} without model pricing.

)}
diff --git a/site/src/pages/AgentsPage/components/ChatSummaryPanel.stories.tsx b/site/src/pages/AgentsPage/components/ChatSummaryPanel.stories.tsx index a3ccfc11d98..737fe1cbd10 100644 --- a/site/src/pages/AgentsPage/components/ChatSummaryPanel.stories.tsx +++ b/site/src/pages/AgentsPage/components/ChatSummaryPanel.stories.tsx @@ -10,8 +10,16 @@ import { ChatSummaryPanel } from "./ChatSummaryPanel"; 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, +}; + +const aiCostControl: { + features: TypesGen.FeatureName[]; + experiments: TypesGen.Experiment[]; +} = { + features: ["aibridge"], + experiments: ["ai-gateway-cost-control"], }; type MockRequestOptions = { @@ -53,6 +61,7 @@ const meta: Meta = { title: "pages/AgentsPage/ChatSummaryPanel", component: ChatSummaryPanel, decorators: [PanelFrame, withDashboardProvider], + parameters: aiCostControl, args: { chatId: MockChat.id, isVisible: true, @@ -117,3 +126,16 @@ export const NotVisible: Story = { expect(canvas.queryByText("No summary yet.")).not.toBeInTheDocument(); }, }; + +export const GatewayUnavailable: Story = { + parameters: { features: [], experiments: [] }, + 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(); + }, +}; diff --git a/site/src/pages/AgentsPage/components/ChatSummaryPanel.tsx b/site/src/pages/AgentsPage/components/ChatSummaryPanel.tsx index 9ddddefb90f..4de90059f3a 100644 --- a/site/src/pages/AgentsPage/components/ChatSummaryPanel.tsx +++ b/site/src/pages/AgentsPage/components/ChatSummaryPanel.tsx @@ -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 { useDashboard } from "#/modules/dashboard/useDashboard"; +import { useFeatureVisibility } from "#/modules/dashboard/useFeatureVisibility"; import { ChatSummary } from "./ChatSummary"; type ChatSummaryPanelProps = { @@ -14,10 +16,24 @@ export const ChatSummaryPanel: FC = ({ chatId, isVisible, }) => { + const { experiments } = useDashboard(); + // Cost is derived from AI Gateway interception data, so it is unavailable + // when the gateway is off or unlicensed. + // TODO(AIGOV-443): drop the experiment gate once cost control is stable. + const showCost = + Boolean(useFeatureVisibility().aibridge) && + experiments.includes("ai-gateway-cost-control"); const chatQuery = useQuery({ ...chat(chatId), enabled: isVisible }); - const costQuery = useQuery({ ...chatCost(chatId), enabled: isVisible }); const chatData = chatQuery.data; + // Cost covers the whole chat tree, so every chat in a tree shares one + // cache entry keyed by the root. Waiting for the chat keeps a subagent + // from caching the tree total under its own id. + const rootChatId = chatData?.root_chat_id ?? chatId; + const costQuery = useQuery({ + ...chatCost(rootChatId), + enabled: isVisible && showCost && chatData !== undefined, + }); let content: ReactNode = null; if (chatQuery.isError) { @@ -30,9 +46,8 @@ export const ChatSummaryPanel: FC = ({ 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} /> From a7b7c34235e451ccd7fdedcb398fb9beb2adc581 Mon Sep 17 00:00:00 2001 From: Michael Suchacz <203725896+ibetitsmike@users.noreply.github.com> Date: Tue, 28 Jul 2026 20:36:48 +0000 Subject: [PATCH 02/10] test(site/src/pages/AgentsPage): cover the root-keyed chat cost query The summary panel keys the cost query by the chat's root so every chat in a tree shares one cache entry. No story exercised that, so add one where the root id differs from the chat id and assert the request uses the root and never the subagent's own id. AgentChatPageView stories never enable the cost-control experiment, so the panel cannot request cost there. Drop the unreachable getChatCost mock and correct the comment above it. --- coderd/apidoc/docs.go | 2 +- coderd/apidoc/swagger.json | 2 +- coderd/database/querier.go | 13 +++++---- coderd/database/queries.sql.go | 13 +++++---- coderd/database/queries/aibridge.sql | 13 +++++---- coderd/exp_chats.go | 5 ++-- codersdk/chats.go | 8 +++--- docs/reference/api/chats.md | 5 ++-- site/src/api/queries/chats.ts | 1 - site/src/api/typesGenerated.ts | 8 +++--- .../AgentsPage/AgentChatPageView.stories.tsx | 10 +++---- .../pages/AgentsPage/AgentsPageLayout.test.ts | 27 +++++++++++++++++-- .../src/pages/AgentsPage/AgentsPageLayout.tsx | 18 +++++++------ .../components/ChatSummary.stories.tsx | 1 + .../AgentsPage/components/ChatSummary.tsx | 5 ++-- .../components/ChatSummaryPanel.stories.tsx | 25 +++++++++++++++++ .../components/ChatSummaryPanel.tsx | 5 ---- 17 files changed, 98 insertions(+), 63 deletions(-) diff --git a/coderd/apidoc/docs.go b/coderd/apidoc/docs.go index 1f2af9f179a..56600062ef5 100644 --- a/coderd/apidoc/docs.go +++ b/coderd/apidoc/docs.go @@ -582,7 +582,7 @@ const docTemplate = `{ }, "/api/experimental/chats/{chat}/cost": { "get": { - "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 a subagent chat returns\nthat same tree total, because AI Gateway attributes a\nsubagent's requests to the chat that spawned it.", + "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.", "produces": [ "application/json" ], diff --git a/coderd/apidoc/swagger.json b/coderd/apidoc/swagger.json index 34e3ffc376a..6351b4d59e8 100644 --- a/coderd/apidoc/swagger.json +++ b/coderd/apidoc/swagger.json @@ -511,7 +511,7 @@ }, "/api/experimental/chats/{chat}/cost": { "get": { - "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 a subagent chat returns\nthat same tree total, because AI Gateway attributes a\nsubagent's requests to the chat that spawned it.", + "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.", "produces": ["application/json"], "tags": ["Chats"], "summary": "Get chat cost", diff --git a/coderd/database/querier.go b/coderd/database/querier.go index fb7b49261f4..cae5e09642a 100644 --- a/coderd/database/querier.go +++ b/coderd/database/querier.go @@ -295,13 +295,12 @@ 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 chat - // beneath it. Coder Agents traffic records the spawning chat's ID as the - // interception session ID (chatprovider.CoderHeaders), so a subagent's - // requests are attributed to its parent and only the whole tree can be - // summed. The owner check guards against session-id collisions from other - // users. Usage without an effective group never reaches ai_user_daily_spend, - // so excluding it keeps this total consistent with AI budget spend. + // 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 diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index 5a4a8b78c9d..9670af318e0 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -1182,13 +1182,12 @@ type GetAIBridgeChatCostRow struct { UnpricedRequestCount int64 `db:"unpriced_request_count" json:"unpriced_request_count"` } -// AI Gateway cost for one chat tree: the root chat plus every subagent chat -// beneath it. Coder Agents traffic records the spawning chat's ID as the -// interception session ID (chatprovider.CoderHeaders), so a subagent's -// requests are attributed to its parent and only the whole tree can be -// summed. The owner check guards against session-id collisions from other -// users. Usage without an effective group never reaches ai_user_daily_spend, -// so excluding it keeps this total consistent with AI budget spend. +// 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 diff --git a/coderd/database/queries/aibridge.sql b/coderd/database/queries/aibridge.sql index e06a3a105f5..15aaea1d944 100644 --- a/coderd/database/queries/aibridge.sql +++ b/coderd/database/queries/aibridge.sql @@ -665,13 +665,12 @@ OFFSET @offset_ ; -- name: GetAIBridgeChatCost :one --- AI Gateway cost for one chat tree: the root chat plus every subagent chat --- beneath it. Coder Agents traffic records the spawning chat's ID as the --- interception session ID (chatprovider.CoderHeaders), so a subagent's --- requests are attributed to its parent and only the whole tree can be --- summed. The owner check guards against session-id collisions from other --- users. Usage without an effective group never reaches ai_user_daily_spend, --- so excluding it keeps this total consistent with AI budget spend. +-- 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. SELECT COALESCE(SUM(tu.cost_micros), 0)::bigint AS total_cost_micros, COUNT(DISTINCT i.id)::bigint AS request_count, diff --git a/coderd/exp_chats.go b/coderd/exp_chats.go index 0b21e6c7077..ec80d19cfe7 100644 --- a/coderd/exp_chats.go +++ b/coderd/exp_chats.go @@ -2474,9 +2474,8 @@ func (api *API) getChatMessages(rw http.ResponseWriter, r *http.Request) { // @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 a subagent chat returns -// @Description that same tree total, because AI Gateway attributes a -// @Description subagent's requests to the chat that spawned it. +// @Description subagent chat beneath it. Requesting cost for a subagent chat +// @Description returns that same total. // //nolint:revive // HTTP handler writes to ResponseWriter. func (api *API) getChatCost(rw http.ResponseWriter, r *http.Request) { diff --git a/codersdk/chats.go b/codersdk/chats.go index 0049dde0505..d03e6ea9400 100644 --- a/codersdk/chats.go +++ b/codersdk/chats.go @@ -1990,10 +1990,10 @@ type ChatCostChatBreakdown struct { TotalRuntimeMs int64 `json:"total_runtime_ms"` } -// ChatCost is the AI Gateway cost for the requested chat's whole tree. AI -// Gateway attributes a subagent's requests to the chat that spawned it, so -// every chat in a tree reports the same total. UnpricedRequestCount counts -// requests whose model had no recorded price, so the total excludes them. +// ChatCost is the AI Gateway cost for the requested chat's whole tree. +// Both root and leaves in a tree report the same total. +// UnpricedRequestCount counts requests whose model had no recorded price. +// RequestCount includes them; TotalCostMicros does not. type ChatCost struct { ChatID uuid.UUID `json:"chat_id" format:"uuid"` TotalCostMicros int64 `json:"total_cost_micros"` diff --git a/docs/reference/api/chats.md b/docs/reference/api/chats.md index 9ecbf10b2ec..16092486390 100644 --- a/docs/reference/api/chats.md +++ b/docs/reference/api/chats.md @@ -1299,9 +1299,8 @@ 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 a subagent chat returns -that same tree total, because AI Gateway attributes a -subagent's requests to the chat that spawned it. +subagent chat beneath it. Requesting cost for a subagent chat +returns that same total. ### Parameters diff --git a/site/src/api/queries/chats.ts b/site/src/api/queries/chats.ts index 660702e483e..5f04fbf9572 100644 --- a/site/src/api/queries/chats.ts +++ b/site/src/api/queries/chats.ts @@ -1966,7 +1966,6 @@ export const chatCostSummary = (user = "me", params?: ChatCostDateParams) => ({ staleTime: 60_000, }); -// Cost covers the whole chat tree, so callers key this by the root chat. export const chatCostKey = (rootChatId: string) => [...chatsKey, rootChatId, "cost"] as const; diff --git a/site/src/api/typesGenerated.ts b/site/src/api/typesGenerated.ts index ec2e791b67c..d650d0d367d 100644 --- a/site/src/api/typesGenerated.ts +++ b/site/src/api/typesGenerated.ts @@ -2155,10 +2155,10 @@ export interface ChatContextTool { // From codersdk/chats.go /** - * ChatCost is the AI Gateway cost for the requested chat's whole tree. AI - * Gateway attributes a subagent's requests to the chat that spawned it, so - * every chat in a tree reports the same total. UnpricedRequestCount counts - * requests whose model had no recorded price, so the total excludes them. + * ChatCost is the AI Gateway cost for the requested chat's whole tree. + * Both root and leaves in a tree report the same total. + * UnpricedRequestCount counts requests whose model had no recorded price. + * RequestCount includes them; TotalCostMicros does not. */ export interface ChatCost { readonly chat_id: string; diff --git a/site/src/pages/AgentsPage/AgentChatPageView.stories.tsx b/site/src/pages/AgentsPage/AgentChatPageView.stories.tsx index b7ff1c9f611..9afa0d6335e 100644 --- a/site/src/pages/AgentsPage/AgentChatPageView.stories.tsx +++ b/site/src/pages/AgentsPage/AgentChatPageView.stories.tsx @@ -196,15 +196,11 @@ const StoryAgentChatPageView: FC = ({ editing, ...overrides }) => { const meta: Meta = { 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 do not enable the cost-control + // experiment, 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, - request_count: 0, - unpriced_request_count: 0, - }); }, decorators: [withAuthProvider, withDashboardProvider, withProxyProvider()], parameters: { diff --git a/site/src/pages/AgentsPage/AgentsPageLayout.test.ts b/site/src/pages/AgentsPage/AgentsPageLayout.test.ts index 76744b29a99..e82d0fe85bb 100644 --- a/site/src/pages/AgentsPage/AgentsPageLayout.test.ts +++ b/site/src/pages/AgentsPage/AgentsPageLayout.test.ts @@ -992,9 +992,9 @@ describe(chatCostIdToInvalidate.name, () => { expected: undefined, }, { - name: "ignores non-status events", + name: "ignores events that bill no gateway request", updatedChat: chatForFilterInvalidation({ status: "waiting" }), - eventKind: "summary_change", + eventKind: "diff_status_change", expected: undefined, }, { @@ -1014,6 +1014,29 @@ describe(chatCostIdToInvalidate.name, () => { eventKind: "title_change", expected: "root-1", }, + { + name: "invalidates when a generated turn status label lands", + updatedChat: chatForFilterInvalidation({ status: "waiting" }), + eventKind: "summary_change", + 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(chatCostIdToInvalidate(updatedChat, eventKind)).toBe(expected); }); diff --git a/site/src/pages/AgentsPage/AgentsPageLayout.tsx b/site/src/pages/AgentsPage/AgentsPageLayout.tsx index fd205a3006d..1035137f9dc 100644 --- a/site/src/pages/AgentsPage/AgentsPageLayout.tsx +++ b/site/src/pages/AgentsPage/AgentsPageLayout.tsx @@ -126,18 +126,20 @@ export const shouldInvalidateFilteredChatList = ( ): boolean => !chat.parent_chat_id && FILTER_MEMBERSHIP_EVENT_KINDS.has(eventKind); -// Chat ID whose cost query must refetch after a watch event, or undefined -// 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. -// Title generation bills its own gateway request and can land while the chat -// is idle, so a title change refetches regardless of status. -// The cost endpoint reports the whole chat tree and the sidebar keys that -// query by root, so the root covers subagents at any depth. +// Titles, turn status labels, and whole-chat summaries are generated after the +// turn already reported a non-active status, so their gateway spend lands after +// a status-driven refetch would have run. +const POST_TURN_BILLED_EVENT_KINDS = new Set([ + "chat_summary_change", + "summary_change", + "title_change", +]); + export const chatCostIdToInvalidate = ( chat: TypesGen.Chat, eventKind: TypesGen.ChatWatchEventKind, ): string | undefined => { - if (eventKind === "title_change") { + if (POST_TURN_BILLED_EVENT_KINDS.has(eventKind)) { return chat.root_chat_id ?? chat.id; } if (eventKind !== "status_change" || isActiveChatStatus(chat.status)) { diff --git a/site/src/pages/AgentsPage/components/ChatSummary.stories.tsx b/site/src/pages/AgentsPage/components/ChatSummary.stories.tsx index d48d657d33d..4c3d31a2c1b 100644 --- a/site/src/pages/AgentsPage/components/ChatSummary.stories.tsx +++ b/site/src/pages/AgentsPage/components/ChatSummary.stories.tsx @@ -11,6 +11,7 @@ const meta: Meta = { createdAt: "2024-05-01T12:00:00Z", updatedAt: "2024-05-02T15:30:00Z", costMicros: 1_250_000, + showCost: true, }, decorators: [ (Story) => ( diff --git a/site/src/pages/AgentsPage/components/ChatSummary.tsx b/site/src/pages/AgentsPage/components/ChatSummary.tsx index a4bfab58fe6..00ed19af6a6 100644 --- a/site/src/pages/AgentsPage/components/ChatSummary.tsx +++ b/site/src/pages/AgentsPage/components/ChatSummary.tsx @@ -15,8 +15,7 @@ interface ChatSummaryProps { costError?: boolean; /** Requests whose model had no recorded price; when > 0 the cost is partial and a note is shown. */ unpricedRequestCount?: number; - /** Cost comes from AI Gateway, so the row is hidden where the gateway is unavailable. */ - showCost?: boolean; + 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; } @@ -29,7 +28,7 @@ export const ChatSummary: FC = ({ isCostLoading, costError, unpricedRequestCount, - showCost = true, + showCost, isSubagent, }) => { const trimmedSummary = summary?.trim(); diff --git a/site/src/pages/AgentsPage/components/ChatSummaryPanel.stories.tsx b/site/src/pages/AgentsPage/components/ChatSummaryPanel.stories.tsx index 737fe1cbd10..a094a68dee0 100644 --- a/site/src/pages/AgentsPage/components/ChatSummaryPanel.stories.tsx +++ b/site/src/pages/AgentsPage/components/ChatSummaryPanel.stories.tsx @@ -7,6 +7,8 @@ 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, @@ -27,6 +29,7 @@ type MockRequestOptions = { summary?: string | null; chatError?: boolean; parentChatId?: string; + rootChatId?: string; }; const mockRequests = ({ @@ -34,6 +37,7 @@ const mockRequests = ({ summary = null, chatError, parentChatId, + rootChatId, }: MockRequestOptions = {}) => { if (chatError) { spyOn(API.experimental, "getChat").mockRejectedValue( @@ -44,6 +48,7 @@ const mockRequests = ({ ...MockChat, summary, ...(parentChatId ? { parent_chat_id: parentChatId } : {}), + ...(rootChatId ? { root_chat_id: rootChatId } : {}), }); } @@ -102,6 +107,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 }) => { diff --git a/site/src/pages/AgentsPage/components/ChatSummaryPanel.tsx b/site/src/pages/AgentsPage/components/ChatSummaryPanel.tsx index 4de90059f3a..712e80f721d 100644 --- a/site/src/pages/AgentsPage/components/ChatSummaryPanel.tsx +++ b/site/src/pages/AgentsPage/components/ChatSummaryPanel.tsx @@ -17,8 +17,6 @@ export const ChatSummaryPanel: FC = ({ isVisible, }) => { const { experiments } = useDashboard(); - // Cost is derived from AI Gateway interception data, so it is unavailable - // when the gateway is off or unlicensed. // TODO(AIGOV-443): drop the experiment gate once cost control is stable. const showCost = Boolean(useFeatureVisibility().aibridge) && @@ -26,9 +24,6 @@ export const ChatSummaryPanel: FC = ({ const chatQuery = useQuery({ ...chat(chatId), enabled: isVisible }); const chatData = chatQuery.data; - // Cost covers the whole chat tree, so every chat in a tree shares one - // cache entry keyed by the root. Waiting for the chat keeps a subagent - // from caching the tree total under its own id. const rootChatId = chatData?.root_chat_id ?? chatId; const costQuery = useQuery({ ...chatCost(rootChatId), From 37eba26710865a2839cf6342eb6c6c697b8394c0 Mon Sep 17 00:00:00 2001 From: Michael Suchacz <203725896+ibetitsmike@users.noreply.github.com> Date: Wed, 29 Jul 2026 14:52:29 +0000 Subject: [PATCH 03/10] fix: aggregate chat cost per request and correct unpriced wording GetAIBridgeChatCost now aggregates token usage per interception in a CTE instead of relying on COUNT(DISTINCT) over the joined rows. The reported counts are unchanged, but a request that records several provider responses is now visibly aggregated once, and the unpriced flag reads as 'this request had usage we could not price'. Correct the SDK doc and the sidebar copy to match: a partially priced request still contributes its priced usage to the total, so the note now says the total excludes unpriced usage rather than whole requests. Add coverage for a partially priced request, a request priced at zero, a sibling root chat tree, and a member reading the cost of a chat they own. --- coderd/database/queries.sql.go | 42 +++-- coderd/database/queries/aibridge.sql | 42 +++-- coderd/exp_chats_test.go | 168 +++++++++++++++++- codersdk/chats.go | 5 +- site/src/api/queries/chats.ts | 2 - site/src/api/typesGenerated.ts | 5 +- .../src/pages/AgentsPage/AgentsPageLayout.tsx | 5 +- .../components/ChatSummary.stories.tsx | 2 +- .../AgentsPage/components/ChatSummary.tsx | 6 +- .../components/ChatSummaryPanel.tsx | 1 - 10 files changed, 227 insertions(+), 51 deletions(-) diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index 9670af318e0..88da7b43205 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -1157,23 +1157,33 @@ func (q *sqlQuerier) DeleteOldAIBridgeRecords(ctx context.Context, beforeTime ti } 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. + SELECT + SUM(tu.cost_micros) AS cost_micros, + BOOL_OR(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 + 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 the planner + -- can use idx_chats_root_chat_id and the chats primary key. + c.root_chat_id = $1::uuid + OR (c.root_chat_id IS NULL AND c.id = $1::uuid) + ) + -- aibridge.ClientCoderAgents. Restricting the client keeps another + -- client's session reference from matching a chat ID. + AND i.client = 'Coder Agents' + AND i.ended_at IS NOT NULL + GROUP BY i.id +) SELECT - COALESCE(SUM(tu.cost_micros), 0)::bigint AS total_cost_micros, - COUNT(DISTINCT i.id)::bigint AS request_count, - COUNT(DISTINCT i.id) FILTER (WHERE tu.cost_micros IS NULL)::bigint AS unpriced_request_count -FROM aibridge_interceptions i -JOIN chats c ON c.id::text = i.session_id AND c.owner_id = i.initiator_id -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 the planner - -- can use idx_chats_root_chat_id and the chats primary key. - c.root_chat_id = $1::uuid - OR (c.root_chat_id IS NULL AND c.id = $1::uuid) - ) - -- aibridge.ClientCoderAgents. Restricting the client keeps another - -- client's session reference from matching a chat ID. - AND i.client = 'Coder Agents' - AND i.ended_at IS NOT NULL + 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 { diff --git a/coderd/database/queries/aibridge.sql b/coderd/database/queries/aibridge.sql index 15aaea1d944..b5072eda5e4 100644 --- a/coderd/database/queries/aibridge.sql +++ b/coderd/database/queries/aibridge.sql @@ -671,20 +671,30 @@ OFFSET @offset_ -- 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. + SELECT + SUM(tu.cost_micros) AS cost_micros, + BOOL_OR(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 + 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 the planner + -- can use idx_chats_root_chat_id and the chats primary key. + c.root_chat_id = @root_chat_id::uuid + OR (c.root_chat_id IS NULL AND c.id = @root_chat_id::uuid) + ) + -- aibridge.ClientCoderAgents. Restricting the client keeps another + -- client's session reference from matching a chat ID. + AND i.client = 'Coder Agents' + AND i.ended_at IS NOT NULL + GROUP BY i.id +) SELECT - COALESCE(SUM(tu.cost_micros), 0)::bigint AS total_cost_micros, - COUNT(DISTINCT i.id)::bigint AS request_count, - COUNT(DISTINCT i.id) FILTER (WHERE tu.cost_micros IS NULL)::bigint AS unpriced_request_count -FROM aibridge_interceptions i -JOIN chats c ON c.id::text = i.session_id AND c.owner_id = i.initiator_id -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 the planner - -- can use idx_chats_root_chat_id and the chats primary key. - c.root_chat_id = @root_chat_id::uuid - OR (c.root_chat_id IS NULL AND c.id = @root_chat_id::uuid) - ) - -- aibridge.ClientCoderAgents. Restricting the client keeps another - -- client's session reference from matching a chat ID. - AND i.client = 'Coder Agents' - AND i.ended_at IS NOT NULL; + 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; diff --git a/coderd/exp_chats_test.go b/coderd/exp_chats_test.go index aec38670c4c..68db982f697 100644 --- a/coderd/exp_chats_test.go +++ b/coderd/exp_chats_test.go @@ -11620,8 +11620,10 @@ func TestChatCostSummary_AdminDrilldown(t *testing.T) { // seedChatGatewayRequest records one finished Coder Agents gateway request // under sessionChatID, mirroring what aibridged persists for chatd traffic: -// the interception's session ID is the chat that spawned the request. -func seedChatGatewayRequest(t *testing.T, db database.Store, initiatorID, sessionChatID uuid.UUID, usage database.InsertAIBridgeTokenUsageParams) { +// the interception's session ID is the chat that spawned the request. Every +// usage lands on that one request, as aibridged records one per provider +// response. +func seedChatGatewayRequest(t *testing.T, db database.Store, initiatorID, sessionChatID uuid.UUID, usages ...database.InsertAIBridgeTokenUsageParams) { t.Helper() now := dbtime.Now() @@ -11635,9 +11637,11 @@ func seedChatGatewayRequest(t *testing.T, db database.Store, initiatorID, sessio ClientSessionID: sql.NullString{String: sessionChatID.String(), Valid: true}, }, &endedAt) - usage.InterceptionID = interception.ID - usage.CreatedAt = now - dbgen.AIBridgeTokenUsage(t, db, usage) + for _, usage := range usages { + usage.InterceptionID = interception.ID + usage.CreatedAt = now + dbgen.AIBridgeTokenUsage(t, db, usage) + } } func TestGetChatCost(t *testing.T) { @@ -11732,6 +11736,128 @@ func TestGetChatCost(t *testing.T) { 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("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("ExcludesUnattributedUsage", func(t *testing.T) { t.Parallel() @@ -11813,6 +11939,38 @@ func TestGetChatCost(t *testing.T) { 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) { t.Parallel() diff --git a/codersdk/chats.go b/codersdk/chats.go index d03e6ea9400..f3317f32436 100644 --- a/codersdk/chats.go +++ b/codersdk/chats.go @@ -1992,8 +1992,9 @@ type ChatCostChatBreakdown struct { // ChatCost is the AI Gateway cost for the requested chat's whole tree. // Both root and leaves in a tree report the same total. -// UnpricedRequestCount counts requests whose model had no recorded price. -// RequestCount includes them; TotalCostMicros does not. +// 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"` diff --git a/site/src/api/queries/chats.ts b/site/src/api/queries/chats.ts index 5f04fbf9572..6d33afcfb1c 100644 --- a/site/src/api/queries/chats.ts +++ b/site/src/api/queries/chats.ts @@ -1969,8 +1969,6 @@ export const chatCostSummary = (user = "me", params?: ChatCostDateParams) => ({ export const chatCostKey = (rootChatId: string) => [...chatsKey, rootChatId, "cost"] as const; -// Chat cost changes only when a gateway request completes, so a short stale -// window refreshes the sidebar without refetching on every render. const GATEWAY_REQUEST_STALE_MS = 30_000; export const chatCost = (rootChatId: string) => ({ diff --git a/site/src/api/typesGenerated.ts b/site/src/api/typesGenerated.ts index d650d0d367d..e2ac4737f20 100644 --- a/site/src/api/typesGenerated.ts +++ b/site/src/api/typesGenerated.ts @@ -2157,8 +2157,9 @@ export interface ChatContextTool { /** * ChatCost is the AI Gateway cost for the requested chat's whole tree. * Both root and leaves in a tree report the same total. - * UnpricedRequestCount counts requests whose model had no recorded price. - * RequestCount includes them; TotalCostMicros does not. + * 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; diff --git a/site/src/pages/AgentsPage/AgentsPageLayout.tsx b/site/src/pages/AgentsPage/AgentsPageLayout.tsx index 1035137f9dc..d1ce1775816 100644 --- a/site/src/pages/AgentsPage/AgentsPageLayout.tsx +++ b/site/src/pages/AgentsPage/AgentsPageLayout.tsx @@ -126,9 +126,8 @@ export const shouldInvalidateFilteredChatList = ( ): boolean => !chat.parent_chat_id && FILTER_MEMBERSHIP_EVENT_KINDS.has(eventKind); -// Titles, turn status labels, and whole-chat summaries are generated after the -// turn already reported a non-active status, so their gateway spend lands after -// a status-driven refetch would have run. +// 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([ "chat_summary_change", "summary_change", diff --git a/site/src/pages/AgentsPage/components/ChatSummary.stories.tsx b/site/src/pages/AgentsPage/components/ChatSummary.stories.tsx index 4c3d31a2c1b..d680df898e7 100644 --- a/site/src/pages/AgentsPage/components/ChatSummary.stories.tsx +++ b/site/src/pages/AgentsPage/components/ChatSummary.stories.tsx @@ -98,7 +98,7 @@ export const PartialCost: Story = { play: async ({ canvasElement }) => { const canvas = within(canvasElement); await expect( - canvas.getByText("Excludes 3 requests without model pricing."), + canvas.getByText("Excludes unpriced usage from 3 requests."), ).toBeInTheDocument(); }, }; diff --git a/site/src/pages/AgentsPage/components/ChatSummary.tsx b/site/src/pages/AgentsPage/components/ChatSummary.tsx index 00ed19af6a6..41ca7328bcf 100644 --- a/site/src/pages/AgentsPage/components/ChatSummary.tsx +++ b/site/src/pages/AgentsPage/components/ChatSummary.tsx @@ -13,7 +13,7 @@ interface ChatSummaryProps { costMicros?: number | null; isCostLoading?: boolean; costError?: boolean; - /** Requests whose model had no recorded price; when > 0 the cost is partial and a note is shown. */ + /** 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. */ @@ -80,8 +80,8 @@ export const ChatSummary: FC = ({ {hasUnpricedRequests && (

- Excludes {unpricedRequestCount} request - {unpricedRequestCount === 1 ? "" : "s"} without model pricing. + Excludes unpriced usage from {unpricedRequestCount} request + {unpricedRequestCount === 1 ? "" : "s"}.

)} diff --git a/site/src/pages/AgentsPage/components/ChatSummaryPanel.tsx b/site/src/pages/AgentsPage/components/ChatSummaryPanel.tsx index 712e80f721d..04ba7d3afb7 100644 --- a/site/src/pages/AgentsPage/components/ChatSummaryPanel.tsx +++ b/site/src/pages/AgentsPage/components/ChatSummaryPanel.tsx @@ -17,7 +17,6 @@ export const ChatSummaryPanel: FC = ({ isVisible, }) => { const { experiments } = useDashboard(); - // TODO(AIGOV-443): drop the experiment gate once cost control is stable. const showCost = Boolean(useFeatureVisibility().aibridge) && experiments.includes("ai-gateway-cost-control"); From bd9accb31d18042c7c6fa2e259d1b010ff92f0c5 Mon Sep 17 00:00:00 2001 From: Michael Suchacz <203725896+ibetitsmike@users.noreply.github.com> Date: Wed, 29 Jul 2026 15:49:05 +0000 Subject: [PATCH 04/10] fix(codersdk): clarify whole-tree cost wording --- coderd/database/queries.sql.go | 4 ++-- coderd/database/queries/aibridge.sql | 4 ++-- codersdk/chats.go | 2 +- site/src/api/typesGenerated.ts | 2 +- 4 files changed, 6 insertions(+), 6 deletions(-) diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index 88da7b43205..6309ab0237e 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -1168,8 +1168,8 @@ WITH per_request AS ( JOIN chats c ON c.id::text = i.session_id AND c.owner_id = i.initiator_id 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 the planner - -- can use idx_chats_root_chat_id and the chats primary key. + -- 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) ) diff --git a/coderd/database/queries/aibridge.sql b/coderd/database/queries/aibridge.sql index b5072eda5e4..3e76da2a0eb 100644 --- a/coderd/database/queries/aibridge.sql +++ b/coderd/database/queries/aibridge.sql @@ -682,8 +682,8 @@ WITH per_request AS ( JOIN chats c ON c.id::text = i.session_id AND c.owner_id = i.initiator_id 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 the planner - -- can use idx_chats_root_chat_id and the chats primary key. + -- 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) ) diff --git a/codersdk/chats.go b/codersdk/chats.go index f3317f32436..a35f051175d 100644 --- a/codersdk/chats.go +++ b/codersdk/chats.go @@ -1991,7 +1991,7 @@ type ChatCostChatBreakdown struct { } // ChatCost is the AI Gateway cost for the requested chat's whole tree. -// Both root and leaves in a tree report the same total. +// Root and subagent chats report the same total. // 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. diff --git a/site/src/api/typesGenerated.ts b/site/src/api/typesGenerated.ts index e2ac4737f20..de25ff4b5bb 100644 --- a/site/src/api/typesGenerated.ts +++ b/site/src/api/typesGenerated.ts @@ -2156,7 +2156,7 @@ export interface ChatContextTool { // From codersdk/chats.go /** * ChatCost is the AI Gateway cost for the requested chat's whole tree. - * Both root and leaves in a tree report the same total. + * Root and subagent chats report the same total. * 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. From 42f2edeb12338d7ca4aadf70d8ff8339468cfff2 Mon Sep 17 00:00:00 2001 From: Michael Suchacz <203725896+ibetitsmike@users.noreply.github.com> Date: Wed, 29 Jul 2026 16:02:52 +0000 Subject: [PATCH 05/10] fix(coderd): tighten chat cost query and fixture comments --- coderd/database/queries.sql.go | 4 ++-- coderd/database/queries/aibridge.sql | 4 ++-- coderd/exp_chats_test.go | 6 ++---- 3 files changed, 6 insertions(+), 8 deletions(-) diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index 6309ab0237e..dd40f5423f0 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -1173,8 +1173,8 @@ WITH per_request AS ( c.root_chat_id = $1::uuid OR (c.root_chat_id IS NULL AND c.id = $1::uuid) ) - -- aibridge.ClientCoderAgents. Restricting the client keeps another - -- client's session reference from matching a chat ID. + -- 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 diff --git a/coderd/database/queries/aibridge.sql b/coderd/database/queries/aibridge.sql index 3e76da2a0eb..400eb0b5897 100644 --- a/coderd/database/queries/aibridge.sql +++ b/coderd/database/queries/aibridge.sql @@ -687,8 +687,8 @@ WITH per_request AS ( c.root_chat_id = @root_chat_id::uuid OR (c.root_chat_id IS NULL AND c.id = @root_chat_id::uuid) ) - -- aibridge.ClientCoderAgents. Restricting the client keeps another - -- client's session reference from matching a chat ID. + -- 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 diff --git a/coderd/exp_chats_test.go b/coderd/exp_chats_test.go index 68db982f697..869513820e3 100644 --- a/coderd/exp_chats_test.go +++ b/coderd/exp_chats_test.go @@ -11619,10 +11619,8 @@ func TestChatCostSummary_AdminDrilldown(t *testing.T) { } // seedChatGatewayRequest records one finished Coder Agents gateway request -// under sessionChatID, mirroring what aibridged persists for chatd traffic: -// the interception's session ID is the chat that spawned the request. Every -// usage lands on that one request, as aibridged records one per provider -// response. +// 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() From e7a7fbad8577ab2286344948c58aad8cb593d541 Mon Sep 17 00:00:00 2001 From: Michael Suchacz <203725896+ibetitsmike@users.noreply.github.com> Date: Wed, 29 Jul 2026 16:38:55 +0000 Subject: [PATCH 06/10] fix(site/src/pages/AgentsPage/components): drop the removed cost-control experiment gate The ai-gateway-cost-control experiment was removed upstream, so gate the per-chat cost row on the aibridge feature alone. --- site/src/pages/AgentsPage/AgentChatPageView.stories.tsx | 4 ++-- .../AgentsPage/components/ChatSummaryPanel.stories.tsx | 8 ++------ site/src/pages/AgentsPage/components/ChatSummaryPanel.tsx | 6 +----- 3 files changed, 5 insertions(+), 13 deletions(-) diff --git a/site/src/pages/AgentsPage/AgentChatPageView.stories.tsx b/site/src/pages/AgentsPage/AgentChatPageView.stories.tsx index 9afa0d6335e..f8affb25624 100644 --- a/site/src/pages/AgentsPage/AgentChatPageView.stories.tsx +++ b/site/src/pages/AgentsPage/AgentChatPageView.stories.tsx @@ -197,8 +197,8 @@ const meta: Meta = { title: "pages/AgentsPage/AgentChatPageView", component: AgentChatPageView, // Summary is the default tab and reads the chat, so mock it for the sidebar. - // Cost needs no mock: these stories do not enable the cost-control - // experiment, so the summary panel never requests it. + // 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()); }, diff --git a/site/src/pages/AgentsPage/components/ChatSummaryPanel.stories.tsx b/site/src/pages/AgentsPage/components/ChatSummaryPanel.stories.tsx index a094a68dee0..c00c240d53c 100644 --- a/site/src/pages/AgentsPage/components/ChatSummaryPanel.stories.tsx +++ b/site/src/pages/AgentsPage/components/ChatSummaryPanel.stories.tsx @@ -16,12 +16,8 @@ const mockCost: TypesGen.ChatCost = { unpriced_request_count: 0, }; -const aiCostControl: { - features: TypesGen.FeatureName[]; - experiments: TypesGen.Experiment[]; -} = { +const aiCostControl: { features: TypesGen.FeatureName[] } = { features: ["aibridge"], - experiments: ["ai-gateway-cost-control"], }; type MockRequestOptions = { @@ -153,7 +149,7 @@ export const NotVisible: Story = { }; export const GatewayUnavailable: Story = { - parameters: { features: [], experiments: [] }, + parameters: { features: [] }, beforeEach: () => mockRequests({ summary: "Gateway is off here." }), play: async ({ canvasElement }) => { const canvas = within(canvasElement); diff --git a/site/src/pages/AgentsPage/components/ChatSummaryPanel.tsx b/site/src/pages/AgentsPage/components/ChatSummaryPanel.tsx index 04ba7d3afb7..585172794bf 100644 --- a/site/src/pages/AgentsPage/components/ChatSummaryPanel.tsx +++ b/site/src/pages/AgentsPage/components/ChatSummaryPanel.tsx @@ -2,7 +2,6 @@ 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 { useDashboard } from "#/modules/dashboard/useDashboard"; import { useFeatureVisibility } from "#/modules/dashboard/useFeatureVisibility"; import { ChatSummary } from "./ChatSummary"; @@ -16,10 +15,7 @@ export const ChatSummaryPanel: FC = ({ chatId, isVisible, }) => { - const { experiments } = useDashboard(); - const showCost = - Boolean(useFeatureVisibility().aibridge) && - experiments.includes("ai-gateway-cost-control"); + const showCost = Boolean(useFeatureVisibility().aibridge); const chatQuery = useQuery({ ...chat(chatId), enabled: isVisible }); const chatData = chatQuery.data; From e225bccacfaebecd3bd0bf14c4e0dc920a583acc Mon Sep 17 00:00:00 2001 From: Michael Suchacz <203725896+ibetitsmike@users.noreply.github.com> Date: Thu, 30 Jul 2026 09:04:35 +0000 Subject: [PATCH 07/10] fix(coderd/database): count finished chat requests without token usage A request that fails upstream still ends but records no token usage, so the inner join dropped it from request_count. Left-join eligible usage instead, and guard the unpriced check on the usage row existing so a request with no usage is not reported as unpriced. --- coderd/database/queries.sql.go | 8 ++++-- coderd/database/queries/aibridge.sql | 8 ++++-- coderd/exp_chats_test.go | 39 +++++++++++++++++++++++++--- codersdk/chats.go | 2 ++ site/src/api/typesGenerated.ts | 2 ++ 5 files changed, 52 insertions(+), 7 deletions(-) diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index dd40f5423f0..4d52bcb438f 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -1161,12 +1161,16 @@ 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.cost_micros IS NULL) AS has_unpriced_usage + 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 - JOIN aibridge_token_usages tu ON tu.interception_id = i.id AND tu.effective_group_id IS NOT NULL + 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. diff --git a/coderd/database/queries/aibridge.sql b/coderd/database/queries/aibridge.sql index 400eb0b5897..2736d579b3d 100644 --- a/coderd/database/queries/aibridge.sql +++ b/coderd/database/queries/aibridge.sql @@ -675,12 +675,16 @@ 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.cost_micros IS NULL) AS has_unpriced_usage + 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 - JOIN aibridge_token_usages tu ON tu.interception_id = i.id AND tu.effective_group_id IS NOT NULL + 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. diff --git a/coderd/exp_chats_test.go b/coderd/exp_chats_test.go index 869513820e3..70763961168 100644 --- a/coderd/exp_chats_test.go +++ b/coderd/exp_chats_test.go @@ -11799,6 +11799,37 @@ func TestGetChatCost(t *testing.T) { 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() @@ -11856,7 +11887,7 @@ func TestGetChatCost(t *testing.T) { } }) - t.Run("ExcludesUnattributedUsage", func(t *testing.T) { + t.Run("ExcludesUnattributedUsageFromCost", func(t *testing.T) { t.Parallel() client, db := newChatClientWithDatabase(t) @@ -11870,7 +11901,8 @@ func TestGetChatCost(t *testing.T) { Title: "legacy chat", }) // Usage recorded before group attribution existed never reached - // ai_user_daily_spend, so it must not appear as chat spend either. + // 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}, }) @@ -11880,7 +11912,8 @@ func TestGetChatCost(t *testing.T) { 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) + require.Equal(t, int64(1), cost.RequestCount) + require.Equal(t, int64(0), cost.UnpricedRequestCount) }) t.Run("ExcludesForeignAndUnfinishedRequests", func(t *testing.T) { diff --git a/codersdk/chats.go b/codersdk/chats.go index a35f051175d..7c47bead863 100644 --- a/codersdk/chats.go +++ b/codersdk/chats.go @@ -1992,6 +1992,8 @@ type ChatCostChatBreakdown struct { // 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. diff --git a/site/src/api/typesGenerated.ts b/site/src/api/typesGenerated.ts index de25ff4b5bb..adb5337079b 100644 --- a/site/src/api/typesGenerated.ts +++ b/site/src/api/typesGenerated.ts @@ -2157,6 +2157,8 @@ export interface ChatContextTool { /** * 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. From 6a706b5753773b572182e346261551f258c9c05a Mon Sep 17 00:00:00 2001 From: Michael Suchacz <203725896+ibetitsmike@users.noreply.github.com> Date: Thu, 30 Jul 2026 09:28:24 +0000 Subject: [PATCH 08/10] docs(coderd): document the gateway retention boundary on chat cost The cost endpoint only totals retained AI Gateway records, so state the retention caveat where the endpoint is introduced rather than later in the stack. Also assert UnpricedRequestCount in the root-resolve test and inline the single-use story parameters constant. --- coderd/apidoc/docs.go | 2 +- coderd/apidoc/swagger.json | 2 +- coderd/exp_chats.go | 6 ++++++ coderd/exp_chats_internal_test.go | 8 +++++--- docs/reference/api/chats.md | 6 ++++++ .../AgentsPage/components/ChatSummaryPanel.stories.tsx | 6 +----- 6 files changed, 20 insertions(+), 10 deletions(-) diff --git a/coderd/apidoc/docs.go b/coderd/apidoc/docs.go index 56600062ef5..bcfd58dd35f 100644 --- a/coderd/apidoc/docs.go +++ b/coderd/apidoc/docs.go @@ -582,7 +582,7 @@ const docTemplate = `{ }, "/api/experimental/chats/{chat}/cost": { "get": { - "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.", + "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" ], diff --git a/coderd/apidoc/swagger.json b/coderd/apidoc/swagger.json index 6351b4d59e8..68af2728093 100644 --- a/coderd/apidoc/swagger.json +++ b/coderd/apidoc/swagger.json @@ -511,7 +511,7 @@ }, "/api/experimental/chats/{chat}/cost": { "get": { - "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.", + "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", diff --git a/coderd/exp_chats.go b/coderd/exp_chats.go index ec80d19cfe7..c5716730b7c 100644 --- a/coderd/exp_chats.go +++ b/coderd/exp_chats.go @@ -2476,6 +2476,12 @@ func (api *API) getChatMessages(rw http.ResponseWriter, r *http.Request) { // @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) { diff --git a/coderd/exp_chats_internal_test.go b/coderd/exp_chats_internal_test.go index e147bd0665e..3b51213cf30 100644 --- a/coderd/exp_chats_internal_test.go +++ b/coderd/exp_chats_internal_test.go @@ -75,8 +75,9 @@ func TestGetChatCostQueriesRootChat(t *testing.T) { dbm.EXPECT().GetChatByID(gomock.Any(), child.ID).Return(child, nil) dbm.EXPECT().GetAIBridgeChatCost(gomock.Any(), rootID).Return( database.GetAIBridgeChatCostRow{ - TotalCostMicros: 250, - RequestCount: 1, + TotalCostMicros: 250, + RequestCount: 2, + UnpricedRequestCount: 1, }, nil, ) @@ -96,7 +97,8 @@ func TestGetChatCostQueriesRootChat(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.RequestCount) + require.Equal(t, int64(2), cost.RequestCount) + require.Equal(t, int64(1), cost.UnpricedRequestCount) } func TestEnrichMissingChatAgentIDs(t *testing.T) { diff --git a/docs/reference/api/chats.md b/docs/reference/api/chats.md index 16092486390..aa836872bb4 100644 --- a/docs/reference/api/chats.md +++ b/docs/reference/api/chats.md @@ -1302,6 +1302,12 @@ 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 | diff --git a/site/src/pages/AgentsPage/components/ChatSummaryPanel.stories.tsx b/site/src/pages/AgentsPage/components/ChatSummaryPanel.stories.tsx index c00c240d53c..896ee03d98f 100644 --- a/site/src/pages/AgentsPage/components/ChatSummaryPanel.stories.tsx +++ b/site/src/pages/AgentsPage/components/ChatSummaryPanel.stories.tsx @@ -16,10 +16,6 @@ const mockCost: TypesGen.ChatCost = { unpriced_request_count: 0, }; -const aiCostControl: { features: TypesGen.FeatureName[] } = { - features: ["aibridge"], -}; - type MockRequestOptions = { cost?: TypesGen.ChatCost; summary?: string | null; @@ -62,7 +58,7 @@ const meta: Meta = { title: "pages/AgentsPage/ChatSummaryPanel", component: ChatSummaryPanel, decorators: [PanelFrame, withDashboardProvider], - parameters: aiCostControl, + parameters: { features: ["aibridge"] satisfies TypesGen.FeatureName[] }, args: { chatId: MockChat.id, isVisible: true, From 432b4b16b53c62170f3ab8aa3430b1bfe030b6c4 Mon Sep 17 00:00:00 2001 From: Michael Suchacz <203725896+ibetitsmike@users.noreply.github.com> Date: Thu, 30 Jul 2026 09:49:25 +0000 Subject: [PATCH 09/10] fix(coderd): fall back to the parent chat when resolving cost tree root 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. Resolve the cost tree with the same COALESCE(root_chat_id, parent_chat_id) precedence the chat queries use, on both the endpoint and the frontend cache key. --- coderd/exp_chats.go | 11 ++++-- coderd/exp_chats_internal_test.go | 35 +++++++++++++++++++ .../components/ChatSummaryPanel.tsx | 5 ++- 3 files changed, 48 insertions(+), 3 deletions(-) diff --git a/coderd/exp_chats.go b/coderd/exp_chats.go index c5716730b7c..8a2a65b06d2 100644 --- a/coderd/exp_chats.go +++ b/coderd/exp_chats.go @@ -2490,10 +2490,17 @@ func (api *API) getChatCost(rw http.ResponseWriter, r *http.Request) { // 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. + // 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 - if chat.RootChatID.Valid { + switch { + case chat.RootChatID.Valid: rootChatID = chat.RootChatID.UUID + case chat.ParentChatID.Valid: + rootChatID = chat.ParentChatID.UUID } row, err := api.Database.GetAIBridgeChatCost(ctx, rootChatID) diff --git a/coderd/exp_chats_internal_test.go b/coderd/exp_chats_internal_test.go index 3b51213cf30..11eb147c764 100644 --- a/coderd/exp_chats_internal_test.go +++ b/coderd/exp_chats_internal_test.go @@ -101,6 +101,41 @@ func TestGetChatCostQueriesRootChat(t *testing.T) { 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) { t.Parallel() newAPI := func(t *testing.T) (*API, *dbmock.MockStore) { diff --git a/site/src/pages/AgentsPage/components/ChatSummaryPanel.tsx b/site/src/pages/AgentsPage/components/ChatSummaryPanel.tsx index 585172794bf..404b4fecde9 100644 --- a/site/src/pages/AgentsPage/components/ChatSummaryPanel.tsx +++ b/site/src/pages/AgentsPage/components/ChatSummaryPanel.tsx @@ -19,7 +19,10 @@ export const ChatSummaryPanel: FC = ({ const chatQuery = useQuery({ ...chat(chatId), enabled: isVisible }); const chatData = chatQuery.data; - const rootChatId = chatData?.root_chat_id ?? chatId; + // Mirrors the server's COALESCE(root_chat_id, parent_chat_id) resolution so + // the cost cache key matches the tree the server aggregates. + const rootChatId = + chatData?.root_chat_id ?? chatData?.parent_chat_id ?? chatId; const costQuery = useQuery({ ...chatCost(rootChatId), enabled: isVisible && showCost && chatData !== undefined, From 5ce000c0bb5fe2e09df54fb6609fcc2700cc59ac Mon Sep 17 00:00:00 2001 From: Michael Suchacz <203725896+ibetitsmike@users.noreply.github.com> Date: Thu, 30 Jul 2026 10:05:42 +0000 Subject: [PATCH 10/10] fix(site/src/pages/AgentsPage): key cost invalidation on the same chat tree The cost cache key and the watch-event invalidation key resolved the tree differently, so an event on a chat whose root was deleted left a mounted cost stale. Both now use one getChatCostTreeID helper, which also trims blank IDs that a nullish check would accept. --- site/src/pages/AgentsPage/AgentsPageLayout.test.ts | 13 +++++++++++++ site/src/pages/AgentsPage/AgentsPageLayout.tsx | 5 +++-- .../components/ChatConversation/chatHelpers.ts | 14 ++++++++++++++ .../AgentsPage/components/ChatSummaryPanel.tsx | 6 ++---- 4 files changed, 32 insertions(+), 6 deletions(-) diff --git a/site/src/pages/AgentsPage/AgentsPageLayout.test.ts b/site/src/pages/AgentsPage/AgentsPageLayout.test.ts index e82d0fe85bb..ea672c68f80 100644 --- a/site/src/pages/AgentsPage/AgentsPageLayout.test.ts +++ b/site/src/pages/AgentsPage/AgentsPageLayout.test.ts @@ -968,6 +968,19 @@ describe(chatCostIdToInvalidate.name, () => { eventKind: "status_change", 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" }), diff --git a/site/src/pages/AgentsPage/AgentsPageLayout.tsx b/site/src/pages/AgentsPage/AgentsPageLayout.tsx index d1ce1775816..d3a16c2e2bc 100644 --- a/site/src/pages/AgentsPage/AgentsPageLayout.tsx +++ b/site/src/pages/AgentsPage/AgentsPageLayout.tsx @@ -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, @@ -139,12 +140,12 @@ export const chatCostIdToInvalidate = ( eventKind: TypesGen.ChatWatchEventKind, ): string | undefined => { if (POST_TURN_BILLED_EVENT_KINDS.has(eventKind)) { - return chat.root_chat_id ?? chat.id; + return getChatCostTreeID(chat); } if (eventKind !== "status_change" || isActiveChatStatus(chat.status)) { return undefined; } - return chat.root_chat_id ?? chat.id; + return getChatCostTreeID(chat); }; const AgentsPageLayout: FC = () => { diff --git a/site/src/pages/AgentsPage/components/ChatConversation/chatHelpers.ts b/site/src/pages/AgentsPage/components/ChatConversation/chatHelpers.ts index 90d8102174c..e20a7c7d39f 100644 --- a/site/src/pages/AgentsPage/components/ChatConversation/chatHelpers.ts +++ b/site/src/pages/AgentsPage/components/ChatConversation/chatHelpers.ts @@ -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[], diff --git a/site/src/pages/AgentsPage/components/ChatSummaryPanel.tsx b/site/src/pages/AgentsPage/components/ChatSummaryPanel.tsx index 404b4fecde9..7f4e0b0bede 100644 --- a/site/src/pages/AgentsPage/components/ChatSummaryPanel.tsx +++ b/site/src/pages/AgentsPage/components/ChatSummaryPanel.tsx @@ -3,6 +3,7 @@ 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 = { @@ -19,10 +20,7 @@ export const ChatSummaryPanel: FC = ({ const chatQuery = useQuery({ ...chat(chatId), enabled: isVisible }); const chatData = chatQuery.data; - // Mirrors the server's COALESCE(root_chat_id, parent_chat_id) resolution so - // the cost cache key matches the tree the server aggregates. - const rootChatId = - chatData?.root_chat_id ?? chatData?.parent_chat_id ?? chatId; + const rootChatId = getChatCostTreeID(chatData) ?? chatId; const costQuery = useQuery({ ...chatCost(rootChatId), enabled: isVisible && showCost && chatData !== undefined,