diff --git a/coderd/apidoc/docs.go b/coderd/apidoc/docs.go index 6a906f903f..3f90aa8c1e 100644 --- a/coderd/apidoc/docs.go +++ b/coderd/apidoc/docs.go @@ -580,6 +580,42 @@ const docTemplate = `{ ] } }, + "/api/experimental/chats/{chat}/cost": { + "get": { + "description": "Experimental: this endpoint is subject to change.", + "produces": [ + "application/json" + ], + "tags": [ + "Chats" + ], + "summary": "Get chat cost", + "operationId": "get-chat-cost", + "parameters": [ + { + "type": "string", + "format": "uuid", + "description": "Chat ID", + "name": "chat", + "in": "path", + "required": true + } + ], + "responses": { + "200": { + "description": "OK", + "schema": { + "$ref": "#/definitions/codersdk.ChatCost" + } + } + }, + "security": [ + { + "CoderSessionToken": [] + } + ] + } + }, "/api/experimental/chats/{chat}/diff": { "get": { "description": "Experimental: this endpoint is subject to change.", @@ -17271,6 +17307,24 @@ const docTemplate = `{ } } }, + "codersdk.ChatCost": { + "type": "object", + "properties": { + "chat_id": { + "type": "string", + "format": "uuid" + }, + "priced_message_count": { + "type": "integer" + }, + "total_cost_micros": { + "type": "integer" + }, + "unpriced_messages_having_usage_count": { + "type": "integer" + } + } + }, "codersdk.ChatDiffContents": { "type": "object", "properties": { diff --git a/coderd/apidoc/swagger.json b/coderd/apidoc/swagger.json index 283dcea72e..b958b0e4a2 100644 --- a/coderd/apidoc/swagger.json +++ b/coderd/apidoc/swagger.json @@ -509,6 +509,38 @@ ] } }, + "/api/experimental/chats/{chat}/cost": { + "get": { + "description": "Experimental: this endpoint is subject to change.", + "produces": ["application/json"], + "tags": ["Chats"], + "summary": "Get chat cost", + "operationId": "get-chat-cost", + "parameters": [ + { + "type": "string", + "format": "uuid", + "description": "Chat ID", + "name": "chat", + "in": "path", + "required": true + } + ], + "responses": { + "200": { + "description": "OK", + "schema": { + "$ref": "#/definitions/codersdk.ChatCost" + } + } + }, + "security": [ + { + "CoderSessionToken": [] + } + ] + } + }, "/api/experimental/chats/{chat}/diff": { "get": { "description": "Experimental: this endpoint is subject to change.", @@ -15523,6 +15555,24 @@ } } }, + "codersdk.ChatCost": { + "type": "object", + "properties": { + "chat_id": { + "type": "string", + "format": "uuid" + }, + "priced_message_count": { + "type": "integer" + }, + "total_cost_micros": { + "type": "integer" + }, + "unpriced_messages_having_usage_count": { + "type": "integer" + } + } + }, "codersdk.ChatDiffContents": { "type": "object", "properties": { diff --git a/coderd/coderd.go b/coderd/coderd.go index 1eae3544b9..158af33929 100644 --- a/coderd/coderd.go +++ b/coderd/coderd.go @@ -1420,6 +1420,7 @@ func New(options *Options) *API { }) r.Get("/", api.getChat) r.Patch("/", api.patchChat) + r.Get("/cost", api.getChatCost) r.Get("/messages", api.getChatMessages) r.Post("/messages", api.postChatMessages) r.Patch("/messages/{message}", api.patchChatMessage) diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go index b35529b028..578fb3ab18 100644 --- a/coderd/database/dbauthz/dbauthz.go +++ b/coderd/database/dbauthz/dbauthz.go @@ -3488,6 +3488,13 @@ func (q *querier) GetChatModelConfigsForTelemetry(ctx context.Context) ([]databa return q.db.GetChatModelConfigsForTelemetry(ctx) } +func (q *querier) GetChatModelUsageCostByChatID(ctx context.Context, chatID uuid.UUID) (database.GetChatModelUsageCostByChatIDRow, error) { + if _, err := q.GetChatByID(ctx, chatID); err != nil { + return database.GetChatModelUsageCostByChatIDRow{}, err + } + return q.db.GetChatModelUsageCostByChatID(ctx, chatID) +} + func (q *querier) GetChatPersonalModelOverridesEnabled(ctx context.Context) (bool, error) { // The personal model overrides flag is a deployment-wide setting read by // authenticated chat users. We only require that an explicit actor is diff --git a/coderd/database/dbauthz/dbauthz_test.go b/coderd/database/dbauthz/dbauthz_test.go index d06e4e5c5e..8f98fd7d7b 100644 --- a/coderd/database/dbauthz/dbauthz_test.go +++ b/coderd/database/dbauthz/dbauthz_test.go @@ -938,11 +938,11 @@ func (s *MethodTestSuite) TestChats() { EndDate: time.Date(2025, 2, 1, 0, 0, 0, 0, time.UTC), } row := database.GetChatCostSummaryRow{ - TotalCostMicros: 987, - PricedMessageCount: 12, - UnpricedMessageCount: 2, - TotalInputTokens: 400, - TotalOutputTokens: 800, + TotalCostMicros: 987, + PricedMessageCount: 12, + UnpricedMessagesHavingUsageCount: 2, + TotalInputTokens: 400, + TotalOutputTokens: 800, } dbm.EXPECT().GetChatCostSummary(gomock.Any(), arg).Return(row, nil).AnyTimes() check.Args(arg).Asserts(rbac.ResourceChat.WithOwner(arg.OwnerID.String()).AnyOrganization(), policy.ActionRead).Returns(row) @@ -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("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} + dbm.EXPECT().GetChatByID(gomock.Any(), chat.ID).Return(chat, nil).AnyTimes() + dbm.EXPECT().GetChatModelUsageCostByChatID(gomock.Any(), chat.ID).Return(row, nil).AnyTimes() + check.Args(chat.ID).Asserts(chat, policy.ActionRead).Returns(row) + })) s.Run("GetChatMessagesByChatIDAscPaginated", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { chat := testutil.Fake(s.T(), faker, database.Chat{}) msgs := []database.ChatMessage{testutil.Fake(s.T(), faker, database.ChatMessage{ChatID: chat.ID})} diff --git a/coderd/database/dbmetrics/querymetrics.go b/coderd/database/dbmetrics/querymetrics.go index 37d8cc9641..07d3321747 100644 --- a/coderd/database/dbmetrics/querymetrics.go +++ b/coderd/database/dbmetrics/querymetrics.go @@ -1713,6 +1713,14 @@ func (m queryMetricsStore) GetChatModelConfigsForTelemetry(ctx context.Context) return r0, r1 } +func (m queryMetricsStore) GetChatModelUsageCostByChatID(ctx context.Context, chatID uuid.UUID) (database.GetChatModelUsageCostByChatIDRow, error) { + start := time.Now() + r0, r1 := m.s.GetChatModelUsageCostByChatID(ctx, chatID) + m.queryLatencies.WithLabelValues("GetChatModelUsageCostByChatID").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetChatModelUsageCostByChatID").Inc() + return r0, r1 +} + func (m queryMetricsStore) GetChatPersonalModelOverridesEnabled(ctx context.Context) (bool, error) { start := time.Now() r0, r1 := m.s.GetChatPersonalModelOverridesEnabled(ctx) diff --git a/coderd/database/dbmock/dbmock.go b/coderd/database/dbmock/dbmock.go index 845b3c279d..9a255e2df7 100644 --- a/coderd/database/dbmock/dbmock.go +++ b/coderd/database/dbmock/dbmock.go @@ -3163,6 +3163,21 @@ func (mr *MockStoreMockRecorder) GetChatModelConfigsForTelemetry(ctx any) *gomoc return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetChatModelConfigsForTelemetry", reflect.TypeOf((*MockStore)(nil).GetChatModelConfigsForTelemetry), ctx) } +// GetChatModelUsageCostByChatID mocks base method. +func (m *MockStore) GetChatModelUsageCostByChatID(ctx context.Context, chatID uuid.UUID) (database.GetChatModelUsageCostByChatIDRow, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetChatModelUsageCostByChatID", ctx, chatID) + ret0, _ := ret[0].(database.GetChatModelUsageCostByChatIDRow) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// GetChatModelUsageCostByChatID indicates an expected call of GetChatModelUsageCostByChatID. +func (mr *MockStoreMockRecorder) GetChatModelUsageCostByChatID(ctx, chatID any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetChatModelUsageCostByChatID", reflect.TypeOf((*MockStore)(nil).GetChatModelUsageCostByChatID), ctx, chatID) +} + // GetChatPersonalModelOverridesEnabled mocks base method. func (m *MockStore) GetChatPersonalModelOverridesEnabled(ctx context.Context) (bool, error) { m.ctrl.T.Helper() diff --git a/coderd/database/querier.go b/coderd/database/querier.go index c637c654c7..7579a25c7d 100644 --- a/coderd/database/querier.go +++ b/coderd/database/querier.go @@ -456,6 +456,11 @@ type sqlcQuerier interface { // Returns all model configurations for telemetry snapshot collection. // deleted = false guarantees ai_provider_id is non-null, so INNER JOIN is safe. GetChatModelConfigsForTelemetry(ctx context.Context) ([]GetChatModelConfigsForTelemetryRow, error) + // Assistant-message cost rolled up over the requested chat's subtree: the + // chat itself plus every descendant reachable through parent_chat_id. A + // root chat therefore reports its whole tree, while a subagent chat + // reports only its own spend plus any nested subagents it spawned. + GetChatModelUsageCostByChatID(ctx context.Context, chatID uuid.UUID) (GetChatModelUsageCostByChatIDRow, error) // GetChatPersonalModelOverridesEnabled returns whether users may configure // personal chat model overrides. It defaults to false when unset. GetChatPersonalModelOverridesEnabled(ctx context.Context) (bool, error) diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index b55924928a..f4ae41eb3c 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -7557,7 +7557,7 @@ SELECT OR cm.cache_creation_tokens IS NOT NULL OR cm.cache_read_tokens IS NOT NULL ) - )::bigint AS unpriced_message_count, + )::bigint AS unpriced_messages_having_usage_count, COALESCE(SUM(cm.input_tokens), 0)::bigint AS total_input_tokens, COALESCE(SUM(cm.output_tokens), 0)::bigint AS total_output_tokens, COALESCE(SUM(cm.cache_read_tokens), 0)::bigint AS total_cache_read_tokens, @@ -7581,14 +7581,14 @@ type GetChatCostSummaryParams struct { } type GetChatCostSummaryRow struct { - TotalCostMicros int64 `db:"total_cost_micros" json:"total_cost_micros"` - PricedMessageCount int64 `db:"priced_message_count" json:"priced_message_count"` - UnpricedMessageCount int64 `db:"unpriced_message_count" json:"unpriced_message_count"` - TotalInputTokens int64 `db:"total_input_tokens" json:"total_input_tokens"` - TotalOutputTokens int64 `db:"total_output_tokens" json:"total_output_tokens"` - TotalCacheReadTokens int64 `db:"total_cache_read_tokens" json:"total_cache_read_tokens"` - TotalCacheCreationTokens int64 `db:"total_cache_creation_tokens" json:"total_cache_creation_tokens"` - TotalRuntimeMs int64 `db:"total_runtime_ms" json:"total_runtime_ms"` + TotalCostMicros int64 `db:"total_cost_micros" json:"total_cost_micros"` + PricedMessageCount int64 `db:"priced_message_count" json:"priced_message_count"` + UnpricedMessagesHavingUsageCount int64 `db:"unpriced_messages_having_usage_count" json:"unpriced_messages_having_usage_count"` + TotalInputTokens int64 `db:"total_input_tokens" json:"total_input_tokens"` + TotalOutputTokens int64 `db:"total_output_tokens" json:"total_output_tokens"` + TotalCacheReadTokens int64 `db:"total_cache_read_tokens" json:"total_cache_read_tokens"` + TotalCacheCreationTokens int64 `db:"total_cache_creation_tokens" json:"total_cache_creation_tokens"` + TotalRuntimeMs int64 `db:"total_runtime_ms" json:"total_runtime_ms"` } // Aggregate cost summary for a single user within a date range. @@ -7599,7 +7599,7 @@ func (q *sqlQuerier) GetChatCostSummary(ctx context.Context, arg GetChatCostSumm err := row.Scan( &i.TotalCostMicros, &i.PricedMessageCount, - &i.UnpricedMessageCount, + &i.UnpricedMessagesHavingUsageCount, &i.TotalInputTokens, &i.TotalOutputTokens, &i.TotalCacheReadTokens, @@ -8363,6 +8363,67 @@ func (q *sqlQuerier) GetChatModelConfigsForTelemetry(ctx context.Context) ([]Get return items, nil } +const getChatModelUsageCostByChatID = `-- name: GetChatModelUsageCostByChatID :one +WITH RECURSIVE target AS ( + SELECT $1::uuid AS chat_id +), subtree AS ( + SELECT chat_id AS id FROM target + UNION ALL + SELECT c.id + FROM chats c + JOIN subtree s ON c.parent_chat_id = s.id +), costs AS ( + SELECT + COALESCE(SUM(cm.total_cost_micros), 0)::bigint AS total_cost_micros, + COUNT(*) FILTER ( + WHERE cm.total_cost_micros IS NOT NULL + )::bigint AS priced_message_count, + COUNT(*) FILTER ( + WHERE cm.total_cost_micros IS NULL + AND ( + cm.input_tokens IS NOT NULL + OR cm.output_tokens IS NOT NULL + OR cm.reasoning_tokens IS NOT NULL + OR cm.cache_creation_tokens IS NOT NULL + OR cm.cache_read_tokens IS NOT NULL + ) + )::bigint AS unpriced_messages_having_usage_count + FROM chat_messages cm + JOIN subtree s ON s.id = cm.chat_id + WHERE cm.role = 'assistant' +) +SELECT + t.chat_id, + costs.total_cost_micros, + costs.priced_message_count, + costs.unpriced_messages_having_usage_count +FROM target t +CROSS JOIN costs +` + +type GetChatModelUsageCostByChatIDRow struct { + ChatID uuid.UUID `db:"chat_id" json:"chat_id"` + TotalCostMicros int64 `db:"total_cost_micros" json:"total_cost_micros"` + PricedMessageCount int64 `db:"priced_message_count" json:"priced_message_count"` + UnpricedMessagesHavingUsageCount int64 `db:"unpriced_messages_having_usage_count" json:"unpriced_messages_having_usage_count"` +} + +// Assistant-message cost rolled up over the requested chat's subtree: the +// chat itself plus every descendant reachable through parent_chat_id. A +// root chat therefore reports its whole tree, while a subagent chat +// reports only its own spend plus any nested subagents it spawned. +func (q *sqlQuerier) GetChatModelUsageCostByChatID(ctx context.Context, chatID uuid.UUID) (GetChatModelUsageCostByChatIDRow, error) { + row := q.db.QueryRowContext(ctx, getChatModelUsageCostByChatID, chatID) + var i GetChatModelUsageCostByChatIDRow + err := row.Scan( + &i.ChatID, + &i.TotalCostMicros, + &i.PricedMessageCount, + &i.UnpricedMessagesHavingUsageCount, + ) + return i, err +} + const getChatQueuedMessageByID = `-- name: GetChatQueuedMessageByID :one SELECT id, chat_id, content, created_at, model_config_id, position, created_by, reasoning_effort FROM chat_queued_messages WHERE id = $1::bigint AND chat_id = $2::uuid diff --git a/coderd/database/queries/chats.sql b/coderd/database/queries/chats.sql index 2d7389630f..836d1200b3 100644 --- a/coderd/database/queries/chats.sql +++ b/coderd/database/queries/chats.sql @@ -2198,7 +2198,7 @@ SELECT OR cm.cache_creation_tokens IS NOT NULL OR cm.cache_read_tokens IS NOT NULL ) - )::bigint AS unpriced_message_count, + )::bigint AS unpriced_messages_having_usage_count, COALESCE(SUM(cm.input_tokens), 0)::bigint AS total_input_tokens, COALESCE(SUM(cm.output_tokens), 0)::bigint AS total_output_tokens, COALESCE(SUM(cm.cache_read_tokens), 0)::bigint AS total_cache_read_tokens, @@ -2295,6 +2295,47 @@ FROM chat_costs cc LEFT JOIN chats rc ON rc.id = cc.root_chat_id ORDER BY cc.total_cost_micros DESC; +-- name: GetChatModelUsageCostByChatID :one +-- Assistant-message cost rolled up over the requested chat's subtree: the +-- chat itself plus every descendant reachable through parent_chat_id. A +-- root chat therefore reports its whole tree, while a subagent chat +-- reports only its own spend plus any nested subagents it spawned. +WITH RECURSIVE target AS ( + SELECT @chat_id::uuid AS chat_id +), subtree AS ( + SELECT chat_id AS id FROM target + UNION ALL + SELECT c.id + FROM chats c + JOIN subtree s ON c.parent_chat_id = s.id +), costs AS ( + SELECT + COALESCE(SUM(cm.total_cost_micros), 0)::bigint AS total_cost_micros, + COUNT(*) FILTER ( + WHERE cm.total_cost_micros IS NOT NULL + )::bigint AS priced_message_count, + COUNT(*) FILTER ( + WHERE cm.total_cost_micros IS NULL + AND ( + cm.input_tokens IS NOT NULL + OR cm.output_tokens IS NOT NULL + OR cm.reasoning_tokens IS NOT NULL + OR cm.cache_creation_tokens IS NOT NULL + OR cm.cache_read_tokens IS NOT NULL + ) + )::bigint AS unpriced_messages_having_usage_count + FROM chat_messages cm + JOIN subtree s ON s.id = cm.chat_id + WHERE cm.role = 'assistant' +) +SELECT + t.chat_id, + costs.total_cost_micros, + costs.priced_message_count, + costs.unpriced_messages_having_usage_count +FROM target t +CROSS JOIN costs; + -- name: GetChatCostPerUser :many -- Deployment-wide per-user cost rollup within a date range. -- Only counts assistant-role messages. diff --git a/coderd/exp_chats.go b/coderd/exp_chats.go index aa95221609..35ca899bac 100644 --- a/coderd/exp_chats.go +++ b/coderd/exp_chats.go @@ -1649,18 +1649,18 @@ func (api *API) chatCostSummary(rw http.ResponseWriter, r *http.Request) { } response := codersdk.ChatCostSummary{ - StartDate: startDate, - EndDate: endDate, - TotalCostMicros: summary.TotalCostMicros, - PricedMessageCount: summary.PricedMessageCount, - UnpricedMessageCount: summary.UnpricedMessageCount, - TotalInputTokens: summary.TotalInputTokens, - TotalOutputTokens: summary.TotalOutputTokens, - TotalCacheReadTokens: summary.TotalCacheReadTokens, - TotalCacheCreationTokens: summary.TotalCacheCreationTokens, - TotalRuntimeMs: summary.TotalRuntimeMs, - ByModel: modelBreakdowns, - ByChat: chatBreakdowns, + StartDate: startDate, + EndDate: endDate, + TotalCostMicros: summary.TotalCostMicros, + PricedMessageCount: summary.PricedMessageCount, + UnpricedMessagesHavingUsageCount: summary.UnpricedMessagesHavingUsageCount, + TotalInputTokens: summary.TotalInputTokens, + TotalOutputTokens: summary.TotalOutputTokens, + TotalCacheReadTokens: summary.TotalCacheReadTokens, + TotalCacheCreationTokens: summary.TotalCacheCreationTokens, + TotalRuntimeMs: summary.TotalRuntimeMs, + ByModel: modelBreakdowns, + ByChat: chatBreakdowns, } if usageStatus != nil { response.UsageLimit = usageStatus @@ -2411,6 +2411,47 @@ func (api *API) getChatMessages(rw http.ResponseWriter, r *http.Request) { }) } +// EXPERIMENTAL: this endpoint is experimental and is subject to change. +// +// @Summary Get chat cost +// @ID get-chat-cost +// @Security CoderSessionToken +// @Tags Chats +// @Produce json +// @Param chat path string true "Chat ID" format(uuid) +// @Success 200 {object} codersdk.ChatCost +// @Router /api/experimental/chats/{chat}/cost [get] +// @Description Experimental: this endpoint is subject to change. +// +//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) + if err != nil { + if httpapi.Is404Error(err) { + httpapi.ResourceNotFound(rw) + return + } + httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{ + Message: "Failed to get chat cost.", + Detail: err.Error(), + }) + return + } + + httpapi.Write(ctx, rw, http.StatusOK, codersdk.ChatCost{ + ChatID: row.ChatID, + TotalCostMicros: row.TotalCostMicros, + PricedMessageCount: row.PricedMessageCount, + UnpricedMessagesHavingUsageCount: row.UnpricedMessagesHavingUsageCount, + }) +} + // @Summary List chat user prompts // @ID list-chat-user-prompts // @Security CoderSessionToken diff --git a/coderd/exp_chats_internal_test.go b/coderd/exp_chats_internal_test.go index 1facd9ff97..bcd7624438 100644 --- a/coderd/exp_chats_internal_test.go +++ b/coderd/exp_chats_internal_test.go @@ -1,8 +1,13 @@ package coderd import ( + "database/sql" + "encoding/json" + "net/http" + "net/http/httptest" "testing" + "github.com/go-chi/chi/v5" "github.com/google/uuid" "github.com/stretchr/testify/require" "go.uber.org/mock/gomock" @@ -11,11 +16,90 @@ import ( "cdr.dev/slog/v3" "cdr.dev/slog/v3/sloggers/slogtest" "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/coderd/database/dbauthz" "github.com/coder/coder/v2/coderd/database/dbmock" + "github.com/coder/coder/v2/coderd/httpmw" "github.com/coder/coder/v2/codersdk" "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. +func TestGetChatCostSurfacesReadAuthzRace(t *testing.T) { + t.Parallel() + + ctrl := gomock.NewController(t) + dbm := dbmock.NewMockStore(ctrl) + chat := database.Chat{ + ID: uuid.New(), + OrganizationID: uuid.New(), + OwnerID: uuid.New(), + } + + dbm.EXPECT().GetChatByID(gomock.Any(), chat.ID).Return(chat, nil) + dbm.EXPECT().GetChatModelUsageCostByChatID(gomock.Any(), chat.ID).Return( + database.GetChatModelUsageCostByChatIDRow{}, + dbauthz.NotAuthorizedError{Err: sql.ErrNoRows}, + ) + + 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/"+chat.ID.String()+"/cost", nil) + rec := httptest.NewRecorder() + rtr.ServeHTTP(rec, req) + resp := rec.Result() + defer resp.Body.Close() + + 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) { + t.Parallel() + + ctrl := gomock.NewController(t) + dbm := dbmock.NewMockStore(ctrl) + rootID := uuid.New() + child := database.Chat{ + ID: uuid.New(), + OrganizationID: uuid.New(), + OwnerID: uuid.New(), + ParentChatID: uuid.NullUUID{UUID: rootID, Valid: true}, + RootChatID: uuid.NullUUID{UUID: rootID, Valid: true}, + } + + 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, + }, + 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, child.ID, cost.ChatID) + require.Equal(t, int64(250), cost.TotalCostMicros) + require.Equal(t, int64(1), cost.PricedMessageCount) +} + func TestEnrichMissingChatAgentIDs(t *testing.T) { t.Parallel() newAPI := func(t *testing.T) (*API, *dbmock.MockStore) { diff --git a/coderd/exp_chats_test.go b/coderd/exp_chats_test.go index 09d934b341..d4762af6c3 100644 --- a/coderd/exp_chats_test.go +++ b/coderd/exp_chats_test.go @@ -11683,7 +11683,7 @@ func assertChatCostSummary(t *testing.T, summary codersdk.ChatCostSummary, model require.Equal(t, int64(1000), summary.TotalCostMicros) require.Equal(t, int64(2), summary.PricedMessageCount) - require.Equal(t, int64(0), summary.UnpricedMessageCount) + require.Equal(t, int64(0), summary.UnpricedMessagesHavingUsageCount) require.Equal(t, int64(200), summary.TotalInputTokens) require.Equal(t, int64(100), summary.TotalOutputTokens) require.Equal(t, int64(4000), summary.TotalRuntimeMs) @@ -11792,6 +11792,224 @@ func TestChatCostSummary_AdminDrilldown(t *testing.T) { }) } +func TestGetChatCost(t *testing.T) { + t.Parallel() + + t.Run("BasicCost", func(t *testing.T) { + t.Parallel() + + f := seedChatCostFixture(t) + ctx := testutil.Context(t, testutil.WaitLong) + + cost, err := f.Client.GetChatCost(ctx, f.ChatID) + require.NoError(t, err) + require.Equal(t, f.ChatID, cost.ChatID) + require.Equal(t, int64(1000), cost.TotalCostMicros) + require.Equal(t, int64(2), cost.PricedMessageCount) + require.Equal(t, int64(0), cost.UnpricedMessagesHavingUsageCount) + }) + + t.Run("RollsUpSubtree", func(t *testing.T) { + t.Parallel() + + client, db := newChatClientWithDatabase(t) + firstUser := coderdtest.CreateFirstUser(t, client.Client) + modelConfig := createChatModelConfig(t, client) + + rootChat := dbgen.Chat(t, db, database.Chat{ + OrganizationID: firstUser.OrganizationID, + OwnerID: firstUser.UserID, + 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, + LastModelConfigID: modelConfig.ID, + Title: "child chat", + 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, + LastModelConfigID: modelConfig.ID, + Title: "grandchild chat", + 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}, + }) + + 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) + }) + + t.Run("UnpricedMessages", func(t *testing.T) { + t.Parallel() + + client, db := newChatClientWithDatabase(t) + firstUser := coderdtest.CreateFirstUser(t, client.Client) + modelConfig := createChatModelConfig(t, client) + + chat := dbgen.Chat(t, db, database.Chat{ + OrganizationID: firstUser.OrganizationID, + OwnerID: firstUser.UserID, + LastModelConfigID: modelConfig.ID, + Title: "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}, + }) + // 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}, + }) + + ctx := testutil.Context(t, testutil.WaitLong) + + cost, err := client.GetChatCost(ctx, chat.ID) + require.NoError(t, err) + require.Equal(t, int64(400), cost.TotalCostMicros) + require.Equal(t, int64(1), cost.PricedMessageCount) + require.Equal(t, int64(1), cost.UnpricedMessagesHavingUsageCount) + }) + + 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: "owner chat", + }) + + 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("ZeroMessages", func(t *testing.T) { + t.Parallel() + + client, db := newChatClientWithDatabase(t) + firstUser := coderdtest.CreateFirstUser(t, client.Client) + modelConfig := createChatModelConfig(t, client) + + // No assistant messages must still return one zero-total row; a COALESCE + // or :one regression would surface as sql.ErrNoRows -> 500. + 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, chat.ID, cost.ChatID) + require.Equal(t, int64(0), cost.TotalCostMicros) + require.Equal(t, int64(0), cost.PricedMessageCount) + require.Equal(t, int64(0), cost.UnpricedMessagesHavingUsageCount) + }) + + t.Run("ExcludesNonAssistantMessages", func(t *testing.T) { + t.Parallel() + + client, db := newChatClientWithDatabase(t) + firstUser := coderdtest.CreateFirstUser(t, client.Client) + modelConfig := createChatModelConfig(t, client) + + chat := dbgen.Chat(t, db, database.Chat{ + OrganizationID: firstUser.OrganizationID, + OwnerID: firstUser.UserID, + LastModelConfigID: modelConfig.ID, + Title: "mixed-role chat", + }) + _ = dbgen.ChatMessage(t, db, database.ChatMessage{ + ChatID: chat.ID, + ModelConfigID: uuid.NullUUID{UUID: modelConfig.ID, Valid: true}, + Role: database.ChatMessageRoleAssistant, + TotalCostMicros: sql.NullInt64{Int64: 600, Valid: true}, + }) + // User-role cost must be excluded; the query bills only assistant messages. + _ = dbgen.ChatMessage(t, db, database.ChatMessage{ + ChatID: chat.ID, + Role: database.ChatMessageRoleUser, + TotalCostMicros: sql.NullInt64{Int64: 999, Valid: true}, + }) + + ctx := testutil.Context(t, testutil.WaitLong) + + cost, err := client.GetChatCost(ctx, chat.ID) + require.NoError(t, err) + require.Equal(t, int64(600), cost.TotalCostMicros) + require.Equal(t, int64(1), cost.PricedMessageCount) + require.Equal(t, int64(0), cost.UnpricedMessagesHavingUsageCount) + }) +} + func TestChatCostUsers(t *testing.T) { t.Parallel() @@ -11985,7 +12203,7 @@ func TestChatCostSummary_UnpricedMessages(t *testing.T) { require.Equal(t, int64(500), summary.TotalCostMicros) require.Equal(t, int64(1), summary.PricedMessageCount) - require.Equal(t, int64(1), summary.UnpricedMessageCount) + require.Equal(t, int64(1), summary.UnpricedMessagesHavingUsageCount) require.Equal(t, int64(300), summary.TotalInputTokens) require.Equal(t, int64(125), summary.TotalOutputTokens) } diff --git a/coderd/x/chatd/chatd.go b/coderd/x/chatd/chatd.go index 6ca6fa56db..59072a6e0a 100644 --- a/coderd/x/chatd/chatd.go +++ b/coderd/x/chatd/chatd.go @@ -4381,6 +4381,12 @@ func (p *Server) maybeFinalizeTurnStatusLabelAndPush( logger slog.Logger, ) { if chat.ParentChatID.Valid { + // Subagent chats skip turn status labels and generated + // summaries, but a successful turn's final report doubles as + // the chat summary so subagents are not summary-less. + if status == database.ChatStatusWaiting { + p.storeSubagentReportSummaryAsync(ctx, chat, logger) + } return } @@ -4505,17 +4511,6 @@ func (p *Server) dispatchSuccessfulTurnPush( p.dispatchPush(ctx, chat, pushBody, database.ChatStatusWaiting, logger) } -func (p *Server) maybeClearLastTurnSummaryAsync( - ctx context.Context, - chat database.Chat, - logger slog.Logger, -) { - if chat.ParentChatID.Valid { - return - } - p.clearLastTurnSummaryAsync(ctx, chat, logger) -} - func (p *Server) setLastTurnSummaryAsync( ctx context.Context, chat database.Chat, @@ -4624,6 +4619,15 @@ const ( chatSummaryWorkTimeout = 120 * time.Second chatSummaryGenerateTimeout = 60 * time.Second chatSummaryWriteTimeout = 5 * time.Second + + // Subagent summaries reuse the final report instead of generating + // text, so their work timeout only covers two database round trips. + subagentReportSummaryTimeout = 15 * time.Second + // Bound the extracted report snippet near the 1-3 sentence + // generated summaries that root chats get, so subagent and parent + // summary panels read the same. + subagentReportSummaryMaxRunes = 300 + subagentReportSummaryMaxSentences = 3 ) // maybeGenerateChatSummaryAsync launches best-effort whole-chat summary @@ -4809,6 +4813,45 @@ func (p *Server) updateChatSummary( p.publishChatPubsubEvent(updatedChat, codersdk.ChatWatchEventKindChatSummaryChange, nil) } +func (p *Server) storeSubagentReportSummaryAsync( + ctx context.Context, + chat database.Chat, + logger slog.Logger, +) { + summaryCtx, stopSummaryCtx := p.inflightContext(ctx) + if err := p.goInflight(func() { + defer stopSummaryCtx() + p.storeSubagentReportSummary(summaryCtx, chat, logger) + }); err != nil { + stopSummaryCtx() + logger.Debug(context.WithoutCancel(ctx), "skipped subagent report summary", + slog.F("chat_id", chat.ID), slog.Error(err)) + } +} + +func (p *Server) storeSubagentReportSummary( + ctx context.Context, + chat database.Chat, + logger slog.Logger, +) { + ctx, cancel := context.WithTimeout(ctx, subagentReportSummaryTimeout) + defer cancel() + + //nolint:gocritic // Narrow daemon access for best-effort summary writes. + authCtx := dbauthz.AsChatd(ctx) + report, err := latestSubagentAssistantMessage(authCtx, p.db, chat.ID) + if err != nil { + logger.Debug(ctx, "failed to load subagent report for summary", + slog.F("chat_id", chat.ID), slog.Error(err)) + return + } + summary := subagentReportSummarySnippet(report) + if summary == "" { + return + } + p.updateChatSummary(ctx, logger, chat, chat.HistoryVersion, summary) +} + func (p *Server) webpushConfigured() bool { return p.webpushDispatcher != nil && p.webpushDispatcher.PublicKey() != "" } diff --git a/coderd/x/chatd/chatd_internal_test.go b/coderd/x/chatd/chatd_internal_test.go index f5c7b5c633..15bb8df0f4 100644 --- a/coderd/x/chatd/chatd_internal_test.go +++ b/coderd/x/chatd/chatd_internal_test.go @@ -14,6 +14,7 @@ import ( "charm.land/fantasy" "github.com/google/uuid" + "github.com/sqlc-dev/pqtype" "github.com/stretchr/testify/require" "go.uber.org/mock/gomock" "golang.org/x/xerrors" @@ -161,6 +162,142 @@ func TestUpdateChatSummary(t *testing.T) { }) } +func assistantReportMessage(t *testing.T, chatID uuid.UUID, id int64, text string) database.ChatMessage { + t.Helper() + + parts := []codersdk.ChatMessagePart{codersdk.ChatMessageText(text)} + data, err := json.Marshal(parts) + require.NoError(t, err) + + return database.ChatMessage{ + ID: id, + ChatID: chatID, + Role: database.ChatMessageRoleAssistant, + Content: pqtype.NullRawMessage{RawMessage: data, Valid: true}, + ContentVersion: chatprompt.ContentVersionV1, + Visibility: database.ChatMessageVisibilityBoth, + } +} + +func TestStoreSubagentReportSummary(t *testing.T) { + t.Parallel() + + t.Run("PersistsFinalReport", func(t *testing.T) { + t.Parallel() + + ctrl := gomock.NewController(t) + db := dbmock.NewMockStore(ctrl) + server := &Server{db: db} + chat := database.Chat{ + ID: uuid.New(), + OwnerID: uuid.New(), + ParentChatID: uuid.NullUUID{UUID: uuid.New(), Valid: true}, + HistoryVersion: 3, + } + logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}) + + db.EXPECT().GetChatMessagesByChatID(gomock.Any(), database.GetChatMessagesByChatIDParams{ + ChatID: chat.ID, + }).Return([]database.ChatMessage{ + assistantReportMessage(t, chat.ID, 1, "intermediate progress"), + assistantReportMessage(t, chat.ID, 2, "final report"), + }, nil) + db.EXPECT().UpdateChatSummary(gomock.Any(), database.UpdateChatSummaryParams{ + ID: chat.ID, + ExpectedHistoryVersion: chat.HistoryVersion, + Summary: sql.NullString{String: "final report", Valid: true}, + }).Return(int64(1), nil) + + server.storeSubagentReportSummary(context.Background(), chat, logger) + }) + + t.Run("SkipsWhenNoVisibleReport", func(t *testing.T) { + t.Parallel() + + ctrl := gomock.NewController(t) + db := dbmock.NewMockStore(ctrl) + server := &Server{db: db} + chat := database.Chat{ + ID: uuid.New(), + OwnerID: uuid.New(), + ParentChatID: uuid.NullUUID{UUID: uuid.New(), Valid: true}, + HistoryVersion: 3, + } + logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}) + + db.EXPECT().GetChatMessagesByChatID(gomock.Any(), database.GetChatMessagesByChatIDParams{ + ChatID: chat.ID, + }).Return([]database.ChatMessage{}, nil) + + server.storeSubagentReportSummary(context.Background(), chat, logger) + }) + + t.Run("TruncatesLongReport", func(t *testing.T) { + t.Parallel() + + ctrl := gomock.NewController(t) + db := dbmock.NewMockStore(ctrl) + server := &Server{db: db} + chat := database.Chat{ + ID: uuid.New(), + OwnerID: uuid.New(), + ParentChatID: uuid.NullUUID{UUID: uuid.New(), Valid: true}, + HistoryVersion: 3, + } + logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}) + + longReport := strings.Repeat("a", subagentReportSummaryMaxRunes+100) + wantSummary := strings.Repeat("a", subagentReportSummaryMaxRunes-1) + "…" + + db.EXPECT().GetChatMessagesByChatID(gomock.Any(), database.GetChatMessagesByChatIDParams{ + ChatID: chat.ID, + }).Return([]database.ChatMessage{ + assistantReportMessage(t, chat.ID, 1, longReport), + }, nil) + db.EXPECT().UpdateChatSummary(gomock.Any(), database.UpdateChatSummaryParams{ + ID: chat.ID, + ExpectedHistoryVersion: chat.HistoryVersion, + Summary: sql.NullString{String: wantSummary, Valid: true}, + }).Return(int64(1), nil) + + server.storeSubagentReportSummary(context.Background(), chat, logger) + }) + + t.Run("StoresSnippetOfMarkdownReport", func(t *testing.T) { + t.Parallel() + + ctrl := gomock.NewController(t) + db := dbmock.NewMockStore(ctrl) + server := &Server{db: db} + chat := database.Chat{ + ID: uuid.New(), + OwnerID: uuid.New(), + ParentChatID: uuid.NullUUID{UUID: uuid.New(), Valid: true}, + HistoryVersion: 3, + } + logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}) + + report := "## Result\n\nFixed **the race** in `cache.go`. " + + "Added a regression test.\n\nLonger explanation follows here." + + db.EXPECT().GetChatMessagesByChatID(gomock.Any(), database.GetChatMessagesByChatIDParams{ + ChatID: chat.ID, + }).Return([]database.ChatMessage{ + assistantReportMessage(t, chat.ID, 1, report), + }, nil) + db.EXPECT().UpdateChatSummary(gomock.Any(), database.UpdateChatSummaryParams{ + ID: chat.ID, + ExpectedHistoryVersion: chat.HistoryVersion, + Summary: sql.NullString{ + String: "Fixed the race in cache.go. Added a regression test.", + Valid: true, + }, + }).Return(int64(1), nil) + + server.storeSubagentReportSummary(context.Background(), chat, logger) + }) +} + func TestMaybeGenerateChatSummaryAsync_CloseCancelsInflight(t *testing.T) { t.Parallel() diff --git a/coderd/x/chatd/generation_preparer.go b/coderd/x/chatd/generation_preparer.go index 06b47af4e4..fcbfbf744d 100644 --- a/coderd/x/chatd/generation_preparer.go +++ b/coderd/x/chatd/generation_preparer.go @@ -709,8 +709,8 @@ func (server *Server) afterInterruptionOutcome( chat := outcome.Chat logger := server.logger.With(slog.F("chat_id", chat.ID), slog.F("owner_id", chat.OwnerID)) - if outcome.Kind == runnerActionKindFinishInterruption { - server.maybeClearLastTurnSummaryAsync(context.WithoutCancel(ctx), chat, logger) + if outcome.Kind == runnerActionKindFinishInterruption && !chat.ParentChatID.Valid { + server.clearLastTurnSummaryAsync(context.WithoutCancel(ctx), chat, logger) } return nil } diff --git a/coderd/x/chatd/quickgen.go b/coderd/x/chatd/quickgen.go index 16e1dea454..dfbcfac6e5 100644 --- a/coderd/x/chatd/quickgen.go +++ b/coderd/x/chatd/quickgen.go @@ -6,6 +6,7 @@ import ( "errors" "fmt" "net/http" + "regexp" "slices" "strings" "time" @@ -1171,6 +1172,145 @@ func countSentenceTerminators(text string) int { return count } +// markdownLinkRe matches inline links and images so snippet extraction +// can keep the link text and drop the URL. +var markdownLinkRe = regexp.MustCompile(`!?\[([^\]]*)\]\([^)]*\)`) + +func subagentReportSummarySnippet(report string) string { + paragraph := firstReportParagraph(report) + if paragraph == "" { + return "" + } + return boundSnippetSentences( + paragraph, + subagentReportSummaryMaxSentences, + subagentReportSummaryMaxRunes, + ) +} + +func firstReportParagraph(report string) string { + var paragraph []string + inFence := false + for line := range strings.Lines(report) { + trimmed := strings.TrimSpace(line) + if strings.HasPrefix(trimmed, "```") || strings.HasPrefix(trimmed, "~~~") { + if !inFence && len(paragraph) > 0 { + break + } + inFence = !inFence + continue + } + if inFence { + continue + } + if trimmed == "" || isMarkdownStructureLine(trimmed) { + if len(paragraph) > 0 { + break + } + continue + } + content := stripInlineMarkdown(stripLineMarkers(trimmed)) + if content == "" { + if len(paragraph) > 0 { + break + } + continue + } + paragraph = append(paragraph, content) + } + return strings.TrimSpace(strings.Join(paragraph, " ")) +} + +func isMarkdownStructureLine(trimmed string) bool { + if strings.HasPrefix(trimmed, "#") || strings.HasPrefix(trimmed, "|") { + return true + } + // Horizontal rules: at least three of the same marker character + // and nothing else. + if len(trimmed) >= 3 && strings.Trim(trimmed, "-") == "" { + return true + } + if len(trimmed) >= 3 && strings.Trim(trimmed, "*") == "" { + return true + } + if len(trimmed) >= 3 && strings.Trim(trimmed, "_") == "" { + return true + } + return false +} + +func stripLineMarkers(trimmed string) string { + for { + next := trimmed + next = strings.TrimPrefix(next, ">") + if rest, ok := trimListMarker(next); ok { + next = rest + } + next = strings.TrimSpace(next) + if next == trimmed { + return trimmed + } + trimmed = next + } +} + +// trimListMarker strips one leading bullet ("- ", "* ", "+ "), ordered +// ("1. ", "1) "), or task-list ("[ ] ", "[x] ") marker. +func trimListMarker(line string) (string, bool) { + for _, marker := range []string{"- ", "* ", "+ ", "[ ] ", "[x] ", "[X] "} { + if rest, ok := strings.CutPrefix(line, marker); ok { + return rest, true + } + } + digits := 0 + for _, r := range line { + if r < '0' || r > '9' { + break + } + digits++ + } + if digits > 0 && len(line) > digits+1 && + (line[digits] == '.' || line[digits] == ')') && line[digits+1] == ' ' { + return line[digits+2:], true + } + return line, false +} + +func stripInlineMarkdown(text string) string { + text = markdownLinkRe.ReplaceAllString(text, "$1") + replacer := strings.NewReplacer("**", "", "__", "", "~~", "", "`", "") + return strings.TrimSpace(replacer.Replace(text)) +} + +func boundSnippetSentences(text string, maxSentences, maxRunes int) string { + runes := []rune(text) + sentences := 0 + lastEnd := 0 + for i, r := range runes { + if r != '.' && r != '!' && r != '?' { + continue + } + if i != len(runes)-1 && !unicode.IsSpace(runes[i+1]) { + continue + } + if i+1 > maxRunes { + break + } + lastEnd = i + 1 + sentences++ + if sentences >= maxSentences { + break + } + } + if lastEnd > 0 { + return strings.TrimSpace(string(runes[:lastEnd])) + } + if len(runes) <= maxRunes { + return text + } + return strings.TrimSpace(string(runes[:maxRunes-1])) + "…" +} + const turnStatusLabelPrompt = "You write compact chat status labels for a sidebar or push notification. " + "Given a chat title, current chat state, and the agent's latest message, populate the label field with a 2-5 word status label. " + "Describe the chat's current state, not the agent. " + diff --git a/coderd/x/chatd/summarygen_internal_test.go b/coderd/x/chatd/summarygen_internal_test.go index d12a108beb..824dbb3c6d 100644 --- a/coderd/x/chatd/summarygen_internal_test.go +++ b/coderd/x/chatd/summarygen_internal_test.go @@ -244,3 +244,87 @@ func TestCountSentenceTerminators(t *testing.T) { "Refactored pkg.cmd.server and auth.rbac.Policy in main.go and util.go. Added coverage in foo_test.go.", )) } + +func TestSubagentReportSummarySnippet(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + report string + want string + }{ + { + name: "ProseLeadKeepsFirstSentences", + report: "Done. Both fixes are pushed as separate commits. Validation passed. Extra detail here.\n\nMore paragraphs follow.", + want: "Done. Both fixes are pushed as separate commits. Validation passed.", + }, + { + name: "SkipsLeadingHeading", + report: "## Summary\n\nFixed the flaky test by pinning the clock.", + want: "Fixed the flaky test by pinning the clock.", + }, + { + name: "StripsInlineMarkdown", + report: "Fixed **the race** in `cache.go`; see [the PR](https://example.com) for details.", + want: "Fixed the race in cache.go; see the PR for details.", + }, + { + name: "JoinsWrappedLines", + report: "Fixed the race\nin the cache layer.\n\nDetails below.", + want: "Fixed the race in the cache layer.", + }, + { + name: "BulletLeadReport", + report: "- Fixed A.\n- Fixed B.\n- Fixed C.\n- Fixed D.", + want: "Fixed A. Fixed B. Fixed C.", + }, + { + name: "SkipsLeadingCodeFence", + report: "```\ngo test ./...\n```\n\nAll tests passed.", + want: "All tests passed.", + }, + { + name: "FenceEndsParagraph", + report: "Ran the suite:\n```\nok 12 packages\n```\nThen more prose.", + want: "Ran the suite:", + }, + { + name: "SkipsTableAndRule", + report: "| a | b |\n|---|---|\n\n---\n\nRolled out the migration.", + want: "Rolled out the migration.", + }, + { + name: "PreservesSnakeCaseIdentifiers", + report: "Renamed parent_chat_id to root_chat_id in the query.", + want: "Renamed parent_chat_id to root_chat_id in the query.", + }, + { + name: "TruncatesUnterminatedText", + report: strings.Repeat("a", subagentReportSummaryMaxRunes+100), + want: strings.Repeat("a", subagentReportSummaryMaxRunes-1) + "…", + }, + { + name: "DropsSentencesPastRuneCap", + report: "Short lead sentence. " + + strings.Repeat("b", subagentReportSummaryMaxRunes) + ".", + want: "Short lead sentence.", + }, + { + name: "EmptyReport", + report: " \n\t\n", + want: "", + }, + { + name: "OnlyCodeAndHeadings", + report: "## Log\n```\nstack trace\n```\n", + want: "", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + require.Equal(t, tt.want, subagentReportSummarySnippet(tt.report)) + }) + } +} diff --git a/coderd/x/chatd/turn_summary_internal_test.go b/coderd/x/chatd/turn_summary_internal_test.go index e37f79231e..d8494c4ab1 100644 --- a/coderd/x/chatd/turn_summary_internal_test.go +++ b/coderd/x/chatd/turn_summary_internal_test.go @@ -18,6 +18,7 @@ import ( "github.com/coder/coder/v2/coderd/x/chatd/chatstate" "github.com/coder/coder/v2/codersdk" "github.com/coder/coder/v2/testutil" + "github.com/coder/quartz" ) func TestUpdateLastTurnSummaryRejectsStaleWrites(t *testing.T) { @@ -112,7 +113,10 @@ func TestUpdateLastTurnSummaryRejectsStaleWrites(t *testing.T) { require.Equal(t, sql.NullString{String: "fresh summary", Valid: true}, fetched.LastTurnSummary) } -func TestSuccessfulChildChatOutcomeSkipsSummaryAndWebPush(t *testing.T) { +// A successful child chat outcome persists the subagent's final report +// as the chat summary but still skips the turn status label and web +// push, which remain parent-only. +func TestSuccessfulChildChatOutcomeStoresReportSummaryWithoutPush(t *testing.T) { t.Parallel() db, ps := dbtestutil.NewDB(t) @@ -168,6 +172,14 @@ func TestSuccessfulChildChatOutcomeSkipsSummaryAndWebPush(t *testing.T) { }) require.NoError(t, err) + const report = "Completed the delegated task." + insertAssistantMessage(t, db, child.ID, modelCfg.ID, report) + // Message inserts bump history_version via trigger; the finalize + // hook receives the post-turn chat, so mirror that here or the + // fenced summary write would be skipped as stale. + child, err = db.GetChatByID(ctx, child.ID) + require.NoError(t, err) + dispatcher := &recordingWebpushDispatcher{} server := &Server{ ctx: t.Context(), @@ -175,6 +187,11 @@ func TestSuccessfulChildChatOutcomeSkipsSummaryAndWebPush(t *testing.T) { pubsub: ps, logger: slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}), webpushDispatcher: dispatcher, + // deriveFinalTurnRunResult resolves the chat model once the + // child has an assistant message, which requires the cache and + // the clock used to mint the synthetic gateway API key. + clock: quartz.NewReal(), + configCache: newChatConfigCache(context.Background(), db, quartz.NewReal()), } require.NoError(t, server.afterGenerationOutcome(ctx, generationOutcome{ Chat: child, @@ -185,6 +202,7 @@ func TestSuccessfulChildChatOutcomeSkipsSummaryAndWebPush(t *testing.T) { fetched, err := db.GetChatByID(ctx, child.ID) require.NoError(t, err) require.False(t, fetched.LastTurnSummary.Valid) + require.Equal(t, sql.NullString{String: report, Valid: true}, fetched.Summary) require.Equal(t, int32(0), dispatcher.dispatchCount.Load()) } diff --git a/codersdk/chats.go b/codersdk/chats.go index f72c852b88..6dd5192fc8 100644 --- a/codersdk/chats.go +++ b/codersdk/chats.go @@ -1923,19 +1923,19 @@ type ChatCostUsersOptions struct { // ChatCostSummary is the response from the chat cost summary endpoint. type ChatCostSummary struct { - StartDate time.Time `json:"start_date" format:"date-time"` - EndDate time.Time `json:"end_date" format:"date-time"` - TotalCostMicros int64 `json:"total_cost_micros"` - PricedMessageCount int64 `json:"priced_message_count"` - UnpricedMessageCount int64 `json:"unpriced_message_count"` - TotalInputTokens int64 `json:"total_input_tokens"` - TotalOutputTokens int64 `json:"total_output_tokens"` - TotalCacheReadTokens int64 `json:"total_cache_read_tokens"` - TotalCacheCreationTokens int64 `json:"total_cache_creation_tokens"` - TotalRuntimeMs int64 `json:"total_runtime_ms"` - ByModel []ChatCostModelBreakdown `json:"by_model"` - ByChat []ChatCostChatBreakdown `json:"by_chat"` - UsageLimit *ChatUsageLimitStatus `json:"usage_limit,omitempty"` + StartDate time.Time `json:"start_date" format:"date-time"` + EndDate time.Time `json:"end_date" format:"date-time"` + TotalCostMicros int64 `json:"total_cost_micros"` + PricedMessageCount int64 `json:"priced_message_count"` + UnpricedMessagesHavingUsageCount int64 `json:"unpriced_messages_having_usage_count"` + TotalInputTokens int64 `json:"total_input_tokens"` + TotalOutputTokens int64 `json:"total_output_tokens"` + TotalCacheReadTokens int64 `json:"total_cache_read_tokens"` + TotalCacheCreationTokens int64 `json:"total_cache_creation_tokens"` + TotalRuntimeMs int64 `json:"total_runtime_ms"` + ByModel []ChatCostModelBreakdown `json:"by_model"` + ByChat []ChatCostChatBreakdown `json:"by_chat"` + UsageLimit *ChatUsageLimitStatus `json:"usage_limit,omitempty"` } // ChatCostModelBreakdown contains per-model cost aggregation. @@ -1966,6 +1966,17 @@ 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. +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"` +} + // ChatCostUserRollup contains per-user cost aggregation for admin views. type ChatCostUserRollup struct { UserID uuid.UUID `json:"user_id" format:"uuid"` @@ -2541,6 +2552,20 @@ 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. +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 { + return ChatCost{}, err + } + defer res.Body.Close() + if res.StatusCode != http.StatusOK { + return ChatCost{}, ReadBodyAsError(res) + } + var cost ChatCost + return cost, json.NewDecoder(res.Body).Decode(&cost) +} + // GetChatCostUsers returns a per-user cost rollup for the deployment // (admin only). Zero-valued StartDate or EndDate fields are omitted from // the request, letting the server apply its own defaults (typically the diff --git a/docs/reference/api/chats.md b/docs/reference/api/chats.md index 5e2e2c90f4..ca4de9f004 100644 --- a/docs/reference/api/chats.md +++ b/docs/reference/api/chats.md @@ -1283,6 +1283,48 @@ Experimental: this endpoint is subject to change. To perform this operation, you must be authenticated. [Learn more](authentication.md). +## Get chat cost + +### Code samples + +```sh +# Example request using curl +curl -X GET http://coder-server:8080/api/experimental/chats/{chat}/cost \ + -H 'Accept: application/json' \ + -H 'Coder-Session-Token: API_KEY' +``` + +`GET /api/experimental/chats/{chat}/cost` + +Experimental: this endpoint is subject to change. + +### Parameters + +| Name | In | Type | Required | Description | +|--------|------|--------------|----------|-------------| +| `chat` | path | string(uuid) | true | Chat ID | + +### Example responses + +> 200 Response + +```json +{ + "chat_id": "efc9fe20-a1e5-4a8c-9c48-f1b30c1e4f86", + "priced_message_count": 0, + "total_cost_micros": 0, + "unpriced_messages_having_usage_count": 0 +} +``` + +### Responses + +| Status | Meaning | Description | Schema | +|--------|---------------------------------------------------------|-------------|--------------------------------------------------| +| 200 | [OK](https://tools.ietf.org/html/rfc7231#section-6.3.1) | OK | [codersdk.ChatCost](schemas.md#codersdkchatcost) | + +To perform this operation, you must be authenticated. [Learn more](authentication.md). + ## Get chat diff contents ### Code samples diff --git a/docs/reference/api/schemas.md b/docs/reference/api/schemas.md index c118eb45b1..8739faca02 100644 --- a/docs/reference/api/schemas.md +++ b/docs/reference/api/schemas.md @@ -2517,6 +2517,26 @@ AuthorizationObject can represent a "set" of objects, such as: all workspaces in | `description` | string | false | | Description is the tool's human-readable summary; may be empty. | | `name` | string | false | | Name is the tool name with the "__" prefix the agent adds stripped, so it reads as the server exposes it. | +## codersdk.ChatCost + +```json +{ + "chat_id": "efc9fe20-a1e5-4a8c-9c48-f1b30c1e4f86", + "priced_message_count": 0, + "total_cost_micros": 0, + "unpriced_messages_having_usage_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 | | | + ## codersdk.ChatDiffContents ```json diff --git a/site/src/api/api.ts b/site/src/api/api.ts index 7cb182a144..eb0f971d10 100644 --- a/site/src/api/api.ts +++ b/site/src/api/api.ts @@ -3441,6 +3441,12 @@ class ExperimentalApiMethods { ); return response.data; }; + getChatCost = async (chatId: string): Promise => { + const response = await this.axios.get( + `/api/experimental/chats/${chatId}/cost`, + ); + return response.data; + }; getChatMessages = async ( chatId: string, opts?: { before_id?: number; after_id?: number; limit?: number }, diff --git a/site/src/api/queries/chats.test.ts b/site/src/api/queries/chats.test.ts index 64e6d25c5b..89156318d5 100644 --- a/site/src/api/queries/chats.test.ts +++ b/site/src/api/queries/chats.test.ts @@ -15,6 +15,8 @@ import { chatACLKey, chatAdvisorConfig, chatAdvisorConfigKey, + chatCost, + chatCostKey, chatCostSummary, chatCostSummaryKey, chatDebugRunsKey, @@ -59,6 +61,7 @@ vi.mock("#/api/api", () => ({ createChat: vi.fn(), deleteChatQueuedMessage: vi.fn(), getChats: vi.fn(), + getChatCost: vi.fn(), getChatCostSummary: vi.fn(), getChatCostUsers: vi.fn(), createChatMessage: vi.fn(), @@ -870,6 +873,20 @@ describe("chat cost query factories", () => { ); }); + it("builds the per-chat cost query key and forwards the chat id", async () => { + const chatId = "chat-1"; + vi.mocked(API.experimental.getChatCost).mockResolvedValue( + {} as TypesGen.ChatCost, + ); + + const query = chatCost(chatId); + + expect(chatCostKey(chatId)).toEqual(["chats", chatId, "cost"]); + expect(query.queryKey).toEqual(["chats", chatId, "cost"]); + await query.queryFn(); + expect(API.experimental.getChatCost).toHaveBeenCalledWith(chatId); + }); + it("builds paginated cost users query with correct key and coerces empty username", async () => { const payload = { start_date: "2025-01-01", diff --git a/site/src/api/queries/chats.ts b/site/src/api/queries/chats.ts index 00b26af131..a23e7efd7d 100644 --- a/site/src/api/queries/chats.ts +++ b/site/src/api/queries/chats.ts @@ -1953,6 +1953,19 @@ 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, +}); + interface PaginatedChatCostUsersPayload { username: string; start_date: string; diff --git a/site/src/api/typesGenerated.ts b/site/src/api/typesGenerated.ts index 815cef4543..8ea30011d0 100644 --- a/site/src/api/typesGenerated.ts +++ b/site/src/api/typesGenerated.ts @@ -1945,6 +1945,20 @@ export interface ChatContextTool { readonly description?: string; } +// 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. + */ +export interface ChatCost { + readonly chat_id: string; + readonly total_cost_micros: number; + readonly priced_message_count: number; + readonly unpriced_messages_having_usage_count: number; +} + // From codersdk/chats.go /** * ChatCostChatBreakdown contains per-root-chat cost aggregation. @@ -1988,7 +2002,7 @@ export interface ChatCostSummary { readonly end_date: string; readonly total_cost_micros: number; readonly priced_message_count: number; - readonly unpriced_message_count: number; + readonly unpriced_messages_having_usage_count: number; readonly total_input_tokens: number; readonly total_output_tokens: number; readonly total_cache_read_tokens: number; diff --git a/site/src/pages/AISettingsPage/SpendPage/SpendPageView.stories.tsx b/site/src/pages/AISettingsPage/SpendPage/SpendPageView.stories.tsx index 12e49c8916..47dd6ffa61 100644 --- a/site/src/pages/AISettingsPage/SpendPage/SpendPageView.stories.tsx +++ b/site/src/pages/AISettingsPage/SpendPage/SpendPageView.stories.tsx @@ -64,7 +64,7 @@ const mockCostSummary = { end_date: "2026-03-12T00:00:00Z", total_cost_micros: 2_500_000, priced_message_count: 40, - unpriced_message_count: 2, + unpriced_messages_having_usage_count: 2, total_input_tokens: 200_000, total_output_tokens: 300_000, total_cache_read_tokens: 10_000, diff --git a/site/src/pages/AISettingsPage/SpendPage/components/ChatCostSummaryView.stories.tsx b/site/src/pages/AISettingsPage/SpendPage/components/ChatCostSummaryView.stories.tsx index dbcf32f074..19d029cdbb 100644 --- a/site/src/pages/AISettingsPage/SpendPage/components/ChatCostSummaryView.stories.tsx +++ b/site/src/pages/AISettingsPage/SpendPage/components/ChatCostSummaryView.stories.tsx @@ -10,7 +10,7 @@ const buildSummary = ( end_date: "2026-03-12T00:00:00Z", total_cost_micros: 1_500_000, priced_message_count: 12, - unpriced_message_count: 0, + unpriced_messages_having_usage_count: 0, total_input_tokens: 123_456, total_output_tokens: 654_321, total_cache_read_tokens: 9_876, @@ -50,7 +50,7 @@ const buildSummary = ( const emptySummary = buildSummary({ total_cost_micros: 0, priced_message_count: 0, - unpriced_message_count: 0, + unpriced_messages_having_usage_count: 0, total_input_tokens: 0, total_output_tokens: 0, by_model: [], @@ -110,7 +110,7 @@ export const WithData: Story = { export const UnpricedWarning: Story = { args: { summary: buildSummary({ - unpriced_message_count: 2, + unpriced_messages_having_usage_count: 2, }), }, }; diff --git a/site/src/pages/AISettingsPage/SpendPage/components/ChatCostSummaryView.tsx b/site/src/pages/AISettingsPage/SpendPage/components/ChatCostSummaryView.tsx index 8d0ef16076..4f1e954f55 100644 --- a/site/src/pages/AISettingsPage/SpendPage/components/ChatCostSummaryView.tsx +++ b/site/src/pages/AISettingsPage/SpendPage/components/ChatCostSummaryView.tsx @@ -194,7 +194,8 @@ export const ChatCostSummaryView: FC = ({

{( - summary.priced_message_count + summary.unpriced_message_count + summary.priced_message_count + + summary.unpriced_messages_having_usage_count ).toLocaleString()}

@@ -253,13 +254,14 @@ export const ChatCostSummaryView: FC = ({ )} - {summary.unpriced_message_count > 0 && ( + {summary.unpriced_messages_having_usage_count > 0 && (
- {summary.unpriced_message_count} message - {summary.unpriced_message_count === 1 ? "" : "s"} could not be - priced because model pricing data was unavailable. + {summary.unpriced_messages_having_usage_count} message + {summary.unpriced_messages_having_usage_count === 1 ? "" : "s"} with + usage could not be priced because model pricing data was + unavailable.
)} diff --git a/site/src/pages/AgentsPage/AgentChatPageView.stories.tsx b/site/src/pages/AgentsPage/AgentChatPageView.stories.tsx index cc58a59997..e99359d7c4 100644 --- a/site/src/pages/AgentsPage/AgentChatPageView.stories.tsx +++ b/site/src/pages/AgentsPage/AgentChatPageView.stories.tsx @@ -196,6 +196,16 @@ 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. + beforeEach: () => { + spyOn(API.experimental, "getChat").mockResolvedValue(buildChat()); + spyOn(API.experimental, "getChatCost").mockResolvedValue({ + chat_id: AGENT_ID, + total_cost_micros: 0, + priced_message_count: 0, + unpriced_messages_having_usage_count: 0, + }); + }, decorators: [withAuthProvider, withDashboardProvider, withProxyProvider()], parameters: { layout: "fullscreen", @@ -479,6 +489,9 @@ index abc1234..def5678 100644 play: async ({ canvasElement }) => { const canvas = within(canvasElement); + // Summary is the default tab; switch to Git to view the PR diff. + await userEvent.click(canvas.getByRole("tab", { name: "Git" })); + // Wait for the initial diff fetch triggered by React Query. await waitFor(() => { expect(API.experimental.getChatDiffContents).toHaveBeenCalled(); @@ -1331,7 +1344,7 @@ export const StickyUserMessagePinsOnScroll: Story = { // which is the regression this story guards against. const sentinels = scrollContainer.querySelectorAll("[data-user-sentinel]"); expect(sentinels.length).toBeGreaterThan(0); - for (const sentinel of sentinels) { + for (const sentinel of Array.from(sentinels)) { expect(sentinel.closest("[data-testid='scroll-container']")).toBe( scrollContainer, ); @@ -1445,7 +1458,7 @@ export const StickyUserMessageClipUpdatesWhilePinned: Story = { // resize of that node reflects transcript growth. const sentinels = scrollContainer.querySelectorAll("[data-user-sentinel]"); expect(sentinels.length).toBeGreaterThan(0); - for (const sentinel of sentinels) { + for (const sentinel of Array.from(sentinels)) { expect(contentMarker?.contains(sentinel)).toBe(true); } @@ -1526,7 +1539,7 @@ export const TerminalFocusOnTabSwitch: Story = { play: async ({ canvasElement }) => { const canvas = within(canvasElement); - // The sidebar should open on the Git tab by default. + // Sidebar defaults to Summary; this story drives the Terminal tab instead. const terminalTab = await canvas.findByRole("tab", { name: "Terminal" }); // 1. Click the Terminal tab. @@ -1624,8 +1637,8 @@ export const PersistsSidebarTabClick: Story = { const canvas = within(canvasElement); await waitFor(() => { - const gitTab = canvas.getByRole("tab", { name: "Git" }); - expect(gitTab).toHaveAttribute("aria-selected", "true"); + const summaryTab = canvas.getByRole("tab", { name: "Summary" }); + expect(summaryTab).toHaveAttribute("aria-selected", "true"); }); const terminalTab = canvas.getByRole("tab", { name: "Terminal" }); @@ -1640,15 +1653,11 @@ export const PersistsSidebarTabClick: Story = { }; /** - * When localStorage holds a tab ID whose tab is not currently available - * (e.g. `"terminal"` while the workspace is stopped), the sidebar - * should fall back to the first available tab (Git) and the stored - * value must be preserved so it can be honoured once the tab reappears. - * - * This locks down the contract described in the PR: `getEffectiveTabId` - * only reads `sidebarTabId` and never writes back. A future write-back - * in the fallback path would silently break restore-after-recovery, so - * this story exists to catch that regression. + * When localStorage holds an unavailable tab ID (e.g. `"terminal"` while the + * workspace is stopped), the sidebar falls back to the first available tab + * (Summary) while preserving the stored value for when the tab reappears. + * Guards the `getEffectiveTabId` contract: it only reads `sidebarTabId`, never + * writes back, so restore-after-recovery cannot silently break. */ export const PreservesUnavailableSidebarTab: Story = { beforeEach: () => { @@ -1662,8 +1671,8 @@ export const PreservesUnavailableSidebarTab: Story = { const canvas = within(canvasElement); await waitFor(() => { - const gitTab = canvas.getByRole("tab", { name: "Git" }); - expect(gitTab).toHaveAttribute("aria-selected", "true"); + const summaryTab = canvas.getByRole("tab", { name: "Summary" }); + expect(summaryTab).toHaveAttribute("aria-selected", "true"); }); expect(canvas.queryByRole("tab", { name: "Terminal" })).toBeNull(); @@ -1702,8 +1711,8 @@ export const DoesNotPersistForArchivedChat: Story = { const canvas = within(canvasElement); await waitFor(() => { - const gitTab = canvas.getByRole("tab", { name: "Git" }); - expect(gitTab).toHaveAttribute("aria-selected", "true"); + const summaryTab = canvas.getByRole("tab", { name: "Summary" }); + expect(summaryTab).toHaveAttribute("aria-selected", "true"); }); const terminalTab = canvas.getByRole("tab", { name: "Terminal" }); diff --git a/site/src/pages/AgentsPage/AgentChatPageView.tsx b/site/src/pages/AgentsPage/AgentChatPageView.tsx index cbe48cc83a..4bb0b50221 100644 --- a/site/src/pages/AgentsPage/AgentChatPageView.tsx +++ b/site/src/pages/AgentsPage/AgentChatPageView.tsx @@ -41,6 +41,7 @@ import type { PendingAttachment } from "./components/ChatPageContent"; import { ChatPageInput, ChatPageTimeline } from "./components/ChatPageContent"; import { ChatScrollContainer } from "./components/ChatScrollContainer"; import { ChatSharingPopoverContent } from "./components/ChatSharingPopover"; +import { ChatSummaryPanel } from "./components/ChatSummaryPanel"; import { getEffectiveTabId } from "./components/ChatsSidebar/tabs/getEffectiveTabId"; import { SidebarTabView } from "./components/ChatsSidebar/tabs/SidebarTabView"; import { ChatTopBar } from "./components/ChatTopBar"; @@ -512,6 +513,7 @@ export const AgentChatPageView: FC = ({ // new tab can never be added to one without the other going out of // sync. Desktop is ordered before terminals so terminals are rightmost. const builtInSidebarTabConfigs = [ + { id: "summary", label: "Summary" }, { id: "git", label: "Git" }, ...(debugLoggingEnabled ? [{ id: "debug", label: "Debug" }] : []), ...(availableDesktopChatId ? [{ id: "desktop", label: "Desktop" }] : []), @@ -680,6 +682,13 @@ export const AgentChatPageView: FC = ({ const renderTabContent = (tabId: string): ReactNode => { switch (tabId) { + case "summary": + return ( + + ); case "git": return ( { ); }); }); + +describe(chatCostIdsToInvalidate.name, () => { + it.each<{ + name: string; + updatedChat: TypesGen.Chat; + eventKind: TypesGen.ChatWatchEventKind; + expected: readonly string[]; + }>([ + { + name: "invalidates when a status change ends active generation", + updatedChat: chatForFilterInvalidation({ status: "waiting" }), + eventKind: "status_change", + 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", + updatedChat: chatForFilterInvalidation({ + id: "child-1", + parent_chat_id: "root-1", + root_chat_id: "root-1", + status: "waiting", + }), + eventKind: "status_change", + expected: ["child-1", "root-1"], + }, + { + name: "invalidates the parent and root when a nested subagent finishes", + updatedChat: chatForFilterInvalidation({ + id: "grandchild-1", + parent_chat_id: "child-1", + root_chat_id: "root-1", + status: "waiting", + }), + eventKind: "status_change", + expected: ["grandchild-1", "child-1", "root-1"], + }, + { + name: "waits while the chat is still active", + updatedChat: chatForFilterInvalidation({ status: "running" }), + eventKind: "status_change", + expected: [], + }, + { + name: "waits while a subagent is still active", + updatedChat: chatForFilterInvalidation({ + id: "child-1", + parent_chat_id: "root-1", + root_chat_id: "root-1", + status: "running", + }), + eventKind: "status_change", + expected: [], + }, + { + name: "waits while the chat is interrupting", + updatedChat: chatForFilterInvalidation({ status: "interrupting" }), + eventKind: "status_change", + expected: [], + }, + { + name: "ignores non-status events", + updatedChat: chatForFilterInvalidation({ status: "waiting" }), + eventKind: "summary_change", + expected: [], + }, + ])("$name", ({ updatedChat, eventKind, expected }) => { + expect(chatCostIdsToInvalidate(updatedChat, eventKind)).toEqual(expected); + }); +}); diff --git a/site/src/pages/AgentsPage/AgentsPageLayout.tsx b/site/src/pages/AgentsPage/AgentsPageLayout.tsx index cc30a490d3..482bc789f8 100644 --- a/site/src/pages/AgentsPage/AgentsPageLayout.tsx +++ b/site/src/pages/AgentsPage/AgentsPageLayout.tsx @@ -20,6 +20,7 @@ import { applyChatArchiveStateToCaches, archiveChat, cancelChatListRefetches, + chatCostKey, chatDiffContentsKey, chatKey, chatModelConfigs, @@ -58,6 +59,7 @@ import { cn } from "#/utils/cn"; import { pageTitle } from "#/utils/page"; import { createReconnectingWebSocket } from "#/utils/reconnectingWebSocket"; import { emptyInputStorageKey } from "./components/AgentCreateForm"; +import { isActiveChatStatus } from "./components/ChatConversation/chatStore"; import { ChatsSidebar, isSettingsView, @@ -124,6 +126,34 @@ 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: 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); + } + if (chat.root_chat_id) { + ids.add(chat.root_chat_id); + } + return [...ids]; +}; + const AgentsPageLayout: FC = () => { useAgentsPWA(); const queryClient = useQueryClient(); @@ -653,6 +683,15 @@ const AgentsPageLayout: FC = () => { if (shouldInvalidateFilteredChatList(updatedChat, chatEvent.kind)) { void invalidateChatListQueries(queryClient); } + for (const costChatId of chatCostIdsToInvalidate( + updatedChat, + chatEvent.kind, + )) { + void queryClient.invalidateQueries({ + queryKey: chatCostKey(costChatId), + exact: true, + }); + } if (chatEvent.kind === "context_dirty") { // The watch payload carries only the lightweight // context flags (the merge above applies them); diff --git a/site/src/pages/AgentsPage/components/ChatSummary.stories.tsx b/site/src/pages/AgentsPage/components/ChatSummary.stories.tsx new file mode 100644 index 0000000000..00f93ff608 --- /dev/null +++ b/site/src/pages/AgentsPage/components/ChatSummary.stories.tsx @@ -0,0 +1,105 @@ +import type { Meta, StoryObj } from "@storybook/react-vite"; +import { expect, within } from "storybook/test"; +import { ChatSummary } from "./ChatSummary"; + +const meta: Meta = { + title: "pages/AgentsPage/ChatSummary", + component: ChatSummary, + args: { + summary: + "Investigated the flaky CI job, traced it to a race in the cache layer, and added a regression test.", + createdAt: "2024-05-01T12:00:00Z", + updatedAt: "2024-05-02T15:30:00Z", + costMicros: 1_250_000, + }, + decorators: [ + (Story) => ( +
+ +
+ ), + ], +}; + +export default meta; +type Story = StoryObj; + +export const WithSummary: Story = { + play: async ({ canvasElement }) => { + const canvas = within(canvasElement); + await expect(canvas.getByText("Created:")).toBeInTheDocument(); + await expect(canvas.getByText("Updated:")).toBeInTheDocument(); + await expect(canvas.getByText("Cost:")).toBeInTheDocument(); + await expect(canvas.getByText("May 1, 2024")).toBeInTheDocument(); + await expect(canvas.getByText("May 2, 2024")).toBeInTheDocument(); + await expect(canvas.queryByText(/12:00|15:30/)).not.toBeInTheDocument(); + // formatCostMicros is locale-pinned to en-US, so this is deterministic. + await expect(canvas.getByText("$1.25")).toBeInTheDocument(); + }, +}; + +export const NoSummary: Story = { + args: { summary: null }, + play: async ({ canvasElement }) => { + const canvas = within(canvasElement); + await expect(canvas.getByText("No summary yet.")).toBeInTheDocument(); + }, +}; + +// A subagent's summary is its final report, persisted when it +// completes, so an empty summary means the agent is still working. +export const SubagentSummaryPending: Story = { + args: { summary: null, isSubagent: true }, + play: async ({ canvasElement }) => { + const canvas = within(canvasElement); + await expect( + canvas.getByText("Summary pending agent completion."), + ).toBeInTheDocument(); + }, +}; + +export const CostLoading: Story = { + args: { isCostLoading: true, costMicros: undefined }, + play: async ({ canvasElement }) => { + const canvas = within(canvasElement); + await expect(canvas.getByLabelText("Loading cost")).toBeInTheDocument(); + }, +}; + +export const CostAbsent: Story = { + args: { costMicros: null }, + play: async ({ canvasElement }) => { + const canvas = within(canvasElement); + await expect(canvas.getByText("Cost:")).toBeInTheDocument(); + await expect(canvas.getByText("-")).toBeInTheDocument(); + }, +}; + +export const SubCentCost: Story = { + args: { costMicros: 5_000 }, + play: async ({ canvasElement }) => { + const canvas = within(canvasElement); + await expect(canvas.getByText("$0.0050")).toBeInTheDocument(); + }, +}; + +export const CostError: Story = { + args: { costMicros: undefined, costError: true }, + play: async ({ canvasElement }) => { + const canvas = within(canvasElement); + await expect(canvas.getByText("Cost:")).toBeInTheDocument(); + await expect(canvas.getByText("Unavailable")).toBeInTheDocument(); + }, +}; + +export const PartialCost: Story = { + args: { costMicros: 0, unpricedMessagesHavingUsageCount: 3 }, + play: async ({ canvasElement }) => { + const canvas = within(canvasElement); + await expect( + canvas.getByText( + "Excludes 3 messages with usage but without model pricing.", + ), + ).toBeInTheDocument(); + }, +}; diff --git a/site/src/pages/AgentsPage/components/ChatSummary.tsx b/site/src/pages/AgentsPage/components/ChatSummary.tsx new file mode 100644 index 0000000000..a7169576c1 --- /dev/null +++ b/site/src/pages/AgentsPage/components/ChatSummary.tsx @@ -0,0 +1,95 @@ +import type { FC, ReactNode } from "react"; +import { Skeleton } from "#/components/Skeleton/Skeleton"; +import { formatCostMicros } from "#/utils/currency"; +import { DATE_FORMAT, formatDateTime } from "#/utils/time"; + +const EMPTY_VALUE = "-"; + +interface ChatSummaryProps { + summary: string | null; + createdAt: string; + updatedAt: string; + /** Cumulative chat cost 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; + /** Subagent summaries are the agent's final report, persisted when it completes, so the empty state reads as pending rather than absent. */ + isSubagent?: boolean; +} + +export const ChatSummary: FC = ({ + summary, + createdAt, + updatedAt, + costMicros, + isCostLoading, + costError, + unpricedMessagesHavingUsageCount, + isSubagent, +}) => { + const trimmedSummary = summary?.trim(); + const hasUnpricedMessages = + !isCostLoading && + !costError && + costMicros != null && + unpricedMessagesHavingUsageCount != null && + unpricedMessagesHavingUsageCount > 0; + + return ( +
+ {trimmedSummary ? ( +

+ {trimmedSummary} +

+ ) : ( +

+ {isSubagent ? "Summary pending agent completion." : "No summary yet."} +

+ )} + +
+ + {formatDateTime(createdAt, DATE_FORMAT.MEDIUM_DATE)} + + + {formatDateTime(updatedAt, DATE_FORMAT.MEDIUM_DATE)} + + + {isCostLoading ? ( + + ) : costError ? ( + Unavailable + ) : costMicros != null ? ( + formatCostMicros(costMicros) + ) : ( + EMPTY_VALUE + )} + +
+ + {hasUnpricedMessages && ( +

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

+ )} +
+ ); +}; + +interface ChatSummaryRowProps { + label: string; + children: ReactNode; +} + +const ChatSummaryRow: FC = ({ label, children }) => ( +
+
{label}
+
+ {children} +
+
+); diff --git a/site/src/pages/AgentsPage/components/ChatSummaryPanel.stories.tsx b/site/src/pages/AgentsPage/components/ChatSummaryPanel.stories.tsx new file mode 100644 index 0000000000..a3ccfc11d9 --- /dev/null +++ b/site/src/pages/AgentsPage/components/ChatSummaryPanel.stories.tsx @@ -0,0 +1,119 @@ +import type { Meta, StoryObj } from "@storybook/react-vite"; +import type { FC } from "react"; +import { expect, spyOn, waitFor, within } from "storybook/test"; +import { API } from "#/api/api"; +import type * as TypesGen from "#/api/typesGenerated"; +import { MockChat } from "#/testHelpers/chatEntities"; +import { withDashboardProvider } from "#/testHelpers/storybook"; +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, +}; + +type MockRequestOptions = { + cost?: TypesGen.ChatCost; + summary?: string | null; + chatError?: boolean; + parentChatId?: string; +}; + +const mockRequests = ({ + cost = mockCost, + summary = null, + chatError, + parentChatId, +}: MockRequestOptions = {}) => { + if (chatError) { + spyOn(API.experimental, "getChat").mockRejectedValue( + new Error("Failed to load chat"), + ); + } else { + spyOn(API.experimental, "getChat").mockResolvedValue({ + ...MockChat, + summary, + ...(parentChatId ? { parent_chat_id: parentChatId } : {}), + }); + } + + spyOn(API.experimental, "getChatCost").mockResolvedValue(cost); +}; + +// The Summary tab fills the right panel, so give stories a bounded height. +const PanelFrame = (Story: FC) => ( +
+ +
+); + +const meta: Meta = { + title: "pages/AgentsPage/ChatSummaryPanel", + component: ChatSummaryPanel, + decorators: [PanelFrame, withDashboardProvider], + args: { + chatId: MockChat.id, + isVisible: true, + }, +}; + +export default meta; +type Story = StoryObj; + +export const WithSummary: Story = { + beforeEach: () => + mockRequests({ + summary: + "Investigated the flaky CI job, traced it to a cache-layer race, and added a regression test.", + }), + play: async ({ canvasElement }) => { + const canvas = within(canvasElement); + await waitFor(() => { + expect( + canvas.getByText(/traced it to a cache-layer race/), + ).toBeInTheDocument(); + expect(canvas.getByText("$1.25")).toBeInTheDocument(); + }); + }, +}; + +// A running subagent has no summary yet; its report is persisted as the +// summary when it completes, so the empty state reads as pending. +export const SubagentSummaryPending: Story = { + beforeEach: () => mockRequests({ parentChatId: "parent-chat-id" }), + play: async ({ canvasElement }) => { + const canvas = within(canvasElement); + await waitFor(() => { + expect( + canvas.getByText("Summary pending agent completion."), + ).toBeInTheDocument(); + }); + }, +}; + +export const ChatError: Story = { + beforeEach: () => mockRequests({ chatError: true }), + play: async ({ canvasElement }) => { + const canvas = within(canvasElement); + await waitFor(() => { + expect(canvas.getByText("Failed to load chat")).toBeInTheDocument(); + }); + }, +}; + +export const NotVisible: Story = { + args: { isVisible: false }, + beforeEach: () => mockRequests({ summary: "Should never be fetched." }), + play: async ({ canvasElement }) => { + const canvas = within(canvasElement); + // Gating disables both queries, so nothing renders and no API call fires. + expect(API.experimental.getChat).not.toHaveBeenCalled(); + expect(API.experimental.getChatCost).not.toHaveBeenCalled(); + expect( + canvas.queryByText("Should never be fetched."), + ).not.toBeInTheDocument(); + expect(canvas.queryByText("No summary yet.")).not.toBeInTheDocument(); + }, +}; diff --git a/site/src/pages/AgentsPage/components/ChatSummaryPanel.tsx b/site/src/pages/AgentsPage/components/ChatSummaryPanel.tsx new file mode 100644 index 0000000000..9ddddefb90 --- /dev/null +++ b/site/src/pages/AgentsPage/components/ChatSummaryPanel.tsx @@ -0,0 +1,47 @@ +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 { ChatSummary } from "./ChatSummary"; + +type ChatSummaryPanelProps = { + chatId: string; + /** Gate reads on tab visibility so the chat and cost queries don't run while the tab is hidden. */ + isVisible: boolean; +}; + +export const ChatSummaryPanel: FC = ({ + chatId, + isVisible, +}) => { + const chatQuery = useQuery({ ...chat(chatId), enabled: isVisible }); + const costQuery = useQuery({ ...chatCost(chatId), enabled: isVisible }); + + const chatData = chatQuery.data; + + let content: ReactNode = null; + if (chatQuery.isError) { + content = ; + } else if (chatData) { + content = ( + + ); + } + + return ( +
+ {content} +
+ ); +};