From a3b2283b9ac648998ce3cc91af45cdc2053ca238 Mon Sep 17 00:00:00 2001 From: Jaayden Halko Date: Thu, 23 Jul 2026 10:01:39 +0000 Subject: [PATCH 01/11] feat: harden chat generation runtime instrumentation for billing chat_messages.runtime_ms is the billing source of truth for Coder Agents runtime, but it was built for debugging. Define the billable metric as the wall-clock duration of the model invocation that produced the persisted message content, and close the paths that dropped it: - Compaction summarization calls now record runtime on the compaction assistant message. - Interrupted attempts persist their episode's lifetime as runtime_ms on the partial assistant message committed by FinishInterruption. - Tool execution batches deliberately record no runtime: tool wall time includes idle waits such as wait_agent polling a sub-agent chat that already bills its own model invocations. - Failed model calls whose output is discarded bill nothing. Document the definition on chat_messages.runtime_ms (COMMENT ON COLUMN), PersistedStep.Runtime, the chatd architecture doc, and the Spend Management docs page. No new index: the existing idx_chat_messages_created_at already serves the hourly runtime range scan. Co-Authored-By: Claude Fable 5 --- coderd/database/dump.sql | 2 + ..._chat_messages_runtime_ms_comment.down.sql | 1 + ...51_chat_messages_runtime_ms_comment.up.sql | 1 + coderd/database/models.go | 9 +- coderd/x/chatd/ARCHITECTURE.md | 3 +- coderd/x/chatd/attempt.go | 5 +- coderd/x/chatd/chatd_test.go | 3 + coderd/x/chatd/chatloop/chatloop.go | 18 ++- coderd/x/chatd/chatloop/compaction.go | 13 ++ coderd/x/chatd/chatloop/runtime_test.go | 122 ++++++++++++++++++ coderd/x/chatd/generation.go | 1 + coderd/x/chatd/message_conversion.go | 27 +++- coderd/x/chatd/message_conversion_test.go | 72 ++++++++++- .../messagepartbuffer/message_part_buffer.go | 24 +++- .../message_part_buffer_test.go | 29 +++++ coderd/x/chatd/tasks.go | 1 + coderd/x/chatd/tasks_test.go | 57 +++++++- .../platform-controls/usage-insights.md | 27 ++++ 18 files changed, 393 insertions(+), 22 deletions(-) create mode 100644 coderd/database/migrations/000551_chat_messages_runtime_ms_comment.down.sql create mode 100644 coderd/database/migrations/000551_chat_messages_runtime_ms_comment.up.sql create mode 100644 coderd/x/chatd/chatloop/runtime_test.go diff --git a/coderd/database/dump.sql b/coderd/database/dump.sql index b7cd44e1dbf..7ef51d99226 100644 --- a/coderd/database/dump.sql +++ b/coderd/database/dump.sql @@ -1968,6 +1968,8 @@ CREATE TABLE chat_messages ( search_tsv tsvector ); +COMMENT ON COLUMN chat_messages.runtime_ms IS 'Wall-clock milliseconds of the model invocation that produced the message content: assistant steps and compaction summaries, including interrupted partials. NULL when no model invocation produced the row (user messages, tool results; local tool execution is not counted). Billing source of truth for Coder Agents runtime: usage reporting sums it over created_at ranges.'; + COMMENT ON COLUMN chat_messages.reasoning_effort IS 'Stores the selected effort for the turn triggered by this message.'; COMMENT ON COLUMN chat_messages.search_tsv IS 'Used for full text search. NULL initially, populated async via background job.'; diff --git a/coderd/database/migrations/000551_chat_messages_runtime_ms_comment.down.sql b/coderd/database/migrations/000551_chat_messages_runtime_ms_comment.down.sql new file mode 100644 index 00000000000..841ebe58b43 --- /dev/null +++ b/coderd/database/migrations/000551_chat_messages_runtime_ms_comment.down.sql @@ -0,0 +1 @@ +COMMENT ON COLUMN chat_messages.runtime_ms IS NULL; diff --git a/coderd/database/migrations/000551_chat_messages_runtime_ms_comment.up.sql b/coderd/database/migrations/000551_chat_messages_runtime_ms_comment.up.sql new file mode 100644 index 00000000000..0c2bcc420c8 --- /dev/null +++ b/coderd/database/migrations/000551_chat_messages_runtime_ms_comment.up.sql @@ -0,0 +1 @@ +COMMENT ON COLUMN chat_messages.runtime_ms IS 'Wall-clock milliseconds of the model invocation that produced the message content: assistant steps and compaction summaries, including interrupted partials. NULL when no model invocation produced the row (user messages, tool results; local tool execution is not counted). Billing source of truth for Coder Agents runtime: usage reporting sums it over created_at ranges.'; diff --git a/coderd/database/models.go b/coderd/database/models.go index 7a96121b021..680de116a8d 100644 --- a/coderd/database/models.go +++ b/coderd/database/models.go @@ -5117,10 +5117,11 @@ type ChatMessage struct { CreatedBy uuid.NullUUID `db:"created_by" json:"created_by"` ContentVersion int16 `db:"content_version" json:"content_version"` TotalCostMicros sql.NullInt64 `db:"total_cost_micros" json:"total_cost_micros"` - RuntimeMs sql.NullInt64 `db:"runtime_ms" json:"runtime_ms"` - Deleted bool `db:"deleted" json:"deleted"` - ProviderResponseID sql.NullString `db:"provider_response_id" json:"provider_response_id"` - Revision int64 `db:"revision" json:"revision"` + // Wall-clock milliseconds of the model invocation that produced the message content: assistant steps and compaction summaries, including interrupted partials. NULL when no model invocation produced the row (user messages, tool results; local tool execution is not counted). Billing source of truth for Coder Agents runtime: usage reporting sums it over created_at ranges. + RuntimeMs sql.NullInt64 `db:"runtime_ms" json:"runtime_ms"` + Deleted bool `db:"deleted" json:"deleted"` + ProviderResponseID sql.NullString `db:"provider_response_id" json:"provider_response_id"` + Revision int64 `db:"revision" json:"revision"` // Stores the selected effort for the turn triggered by this message. ReasoningEffort NullChatReasoningEffort `db:"reasoning_effort" json:"reasoning_effort"` // Used for full text search. NULL initially, populated async via background job. diff --git a/coderd/x/chatd/ARCHITECTURE.md b/coderd/x/chatd/ARCHITECTURE.md index c144316bb8c..5b2ded5ed21 100644 --- a/coderd/x/chatd/ARCHITECTURE.md +++ b/coderd/x/chatd/ARCHITECTURE.md @@ -716,6 +716,7 @@ The buffer exposes the following API: - `CloseEpisode(chat_id, history_version, generation_attempt)`: closes an episode, preventing further parts from being added to it. May be called multiple times for a given episode, subsequent calls will be no-ops. Calling it on a non-existent episode creates the episode and closes it immediately. Concurrent parts of the system may race to create the episode and close it, so creating and closing in one operation prevents race conditions. - `AddPart(chat_id, history_version, generation_attempt, content)`: adds a message part to the buffer. Returns a predefined error if the episode is not found or the array is full. - `GetParts(chat_id, history_version, generation_attempt)`: returns the message parts for an episode. Returns a predefined error if the episode is not found. +- `EpisodeDuration(chat_id, history_version, generation_attempt)`: returns the wall-clock span between `CreateEpisode` and `CloseEpisode`, or zero when the episode is unknown, was created implicitly by `CloseEpisode`, or is still open. The interrupt goroutine uses it as the interrupted generation attempt's runtime. - `SubscribeToEpisode(chat_id, history_version, generation_attempt)`: returns a go channel that will receive all message parts for the episode. It spawns a goroutine that delivers parts to the channel. It's live until the episode is closed or until a subscriber requests that the channel be closed. Once the goroutine delivers all message parts for a closed episode, it closes the channel and exits. If the episode is already closed at the time of the call, the goroutine delivers all message parts for the episode, closes the channel, and exits. `SubscribeToEpisode` does not return an error if the episode is not found: it waits for it to be created instead. Closed episodes are garbage collected after at least 15 seconds since they were closed and when they have no active subscribers. The message part buffer maintains a garbage collection goroutine. @@ -875,7 +876,7 @@ The goroutine does the following in order: 1. It fetches the generation attempt number from the database. 2. It closes the episode corresponding to its history version and generation attempt by calling the `CloseEpisode` method on the [Message part buffer](#message-part-buffer). 3. It reads the buffered parts for that episode by calling the `GetParts` method on the message part buffer. -4. It applies the `FinishInterruption(partial?)` transition on the core state machine. If there are no buffered parts for that episode, or the episode is not found, it passes `nil` as the `partial` argument. +4. It applies the `FinishInterruption(partial?)` transition on the core state machine. If there are no buffered parts for that episode, or the episode is not found, it passes `nil` as the `partial` argument. When the partial suffix contains an assistant message, the episode's lifetime (`EpisodeDuration`) is persisted on that message as `runtime_ms`, so the interrupted attempt's billable generation time is not lost. #### Dynamic tools timeout goroutine diff --git a/coderd/x/chatd/attempt.go b/coderd/x/chatd/attempt.go index 1e1cc6b0e2d..fb967fb2398 100644 --- a/coderd/x/chatd/attempt.go +++ b/coderd/x/chatd/attempt.go @@ -44,7 +44,9 @@ type pendingDynamicToolCall struct { Args string } -// compactionOutcome contains a generated context summary. +// compactionOutcome contains a generated context summary. It must stay +// field-compatible with chatloop.CompactionResult; generateCompaction +// converts between the two directly. type compactionOutcome struct { SystemSummary string SummaryReport string @@ -53,6 +55,7 @@ type compactionOutcome struct { UsagePercent float64 ContextTokens int64 ContextLimit int64 + Runtime time.Duration } type compactionStatus int diff --git a/coderd/x/chatd/chatd_test.go b/coderd/x/chatd/chatd_test.go index 4d7d31629b3..7a206e47757 100644 --- a/coderd/x/chatd/chatd_test.go +++ b/coderd/x/chatd/chatd_test.go @@ -6340,6 +6340,9 @@ func TestActiveServer_BasicAssistantGenerationAndPromptPreparation(t *testing.T) require.Equal(t, database.ChatMessageRoleAssistant, last.Role) require.True(t, last.ContextLimit.Valid) require.Equal(t, int64(4096), last.ContextLimit.Int64) + // Assistant steps must persist the model invocation's runtime; it is + // the billing source of truth for Coder Agents runtime. + require.True(t, last.RuntimeMs.Valid) require.GreaterOrEqual(t, last.RuntimeMs.Int64, int64(0)) requireTextPart(t, last, "done") diff --git a/coderd/x/chatd/chatloop/chatloop.go b/coderd/x/chatd/chatloop/chatloop.go index ff659f1021a..50662b402aa 100644 --- a/coderd/x/chatd/chatloop/chatloop.go +++ b/coderd/x/chatd/chatloop/chatloop.go @@ -73,10 +73,16 @@ type PersistedStep struct { Content []fantasy.Content Usage fantasy.Usage ContextLimit sql.NullInt64 - // Runtime is the wall-clock duration of this step, - // covering LLM streaming, tool execution, and retries. - // Zero indicates the duration was not measured (e.g. - // interrupted steps). + // Runtime is the wall-clock duration of the model invocation + // that produced this step's content, measured from just before + // the provider stream is opened until the stream is fully + // consumed. It is persisted as chat_messages.runtime_ms, the + // billable "active generation" time that usage reporting sums. + // Steps without a model invocation (local tool execution + // batches) leave it zero, which persists as NULL: tool wall + // time includes idle waits such as wait_agent polling a + // sub-agent chat that already bills its own model invocations, + // so billing it would double count. Runtime time.Duration // PendingDynamicToolCalls lists tool calls that target // dynamic tools. When non-empty the chatloop exits with @@ -304,6 +310,10 @@ type GenerateCompactionOptions struct { ProviderOptions fantasy.ProviderOptions PublishMessagePart func(codersdk.ChatMessageRole, codersdk.ChatMessagePart) + + // Clock measures the summary call duration reported as + // CompactionResult.Runtime. Nil uses a real clock. + Clock quartz.Clock } // ProviderTool pairs a provider-native tool definition with an diff --git a/coderd/x/chatd/chatloop/compaction.go b/coderd/x/chatd/chatloop/compaction.go index e0bcef5bb9c..e375ee9c81b 100644 --- a/coderd/x/chatd/chatloop/compaction.go +++ b/coderd/x/chatd/chatloop/compaction.go @@ -12,6 +12,7 @@ import ( "github.com/coder/coder/v2/coderd/x/chatd/chatdebug" "github.com/coder/coder/v2/codersdk" + "github.com/coder/quartz" ) const ( @@ -129,6 +130,11 @@ type CompactionResult struct { UsagePercent float64 ContextTokens int64 ContextLimit int64 + // Runtime is the wall-clock duration of the summarization model + // call. Compaction is a billable model invocation like a regular + // assistant step; see PersistedStep.Runtime. Zero when the run + // was gated off before calling the model. + Runtime time.Duration } // GenerateCompaction generates one context summary and returns it without @@ -170,11 +176,17 @@ func GenerateCompaction(ctx context.Context, opts GenerateCompactionOptions) (Co ) } + clock := opts.Clock + if clock == nil { + clock = quartz.NewReal() + } + summaryStart := clock.Now() summary, err := generateCompactionSummary(ctx, opts.Model, opts.Messages, config) if err != nil { publishCompactionError(config, "failed to generate compaction summary") return CompactionResult{}, err } + runtime := clock.Since(summaryStart) if summary == "" { publishCompactionError(config, "compaction produced an empty summary") return CompactionResult{}, xerrors.New("compaction produced an empty summary") @@ -190,6 +202,7 @@ func GenerateCompaction(ctx context.Context, opts GenerateCompactionOptions) (Co UsagePercent: usagePercent, ContextTokens: contextTokens, ContextLimit: contextLimit, + Runtime: runtime, } if config.PublishMessagePart != nil && config.ToolCallID != "" { resultJSON, _ := json.Marshal(map[string]any{ diff --git a/coderd/x/chatd/chatloop/runtime_test.go b/coderd/x/chatd/chatloop/runtime_test.go new file mode 100644 index 00000000000..6bf2009e116 --- /dev/null +++ b/coderd/x/chatd/chatloop/runtime_test.go @@ -0,0 +1,122 @@ +package chatloop_test + +import ( + "context" + "testing" + "time" + + "charm.land/fantasy" + "github.com/stretchr/testify/require" + "golang.org/x/xerrors" + + "github.com/coder/coder/v2/coderd/x/chatd/chatloop" + "github.com/coder/coder/v2/coderd/x/chatd/chattest" + "github.com/coder/quartz" +) + +// TestGenerateAssistant_RecordsModelInvocationRuntime pins the billable +// runtime definition: Step.Runtime spans the model invocation, from just +// before the stream opens until it is fully consumed. +func TestGenerateAssistant_RecordsModelInvocationRuntime(t *testing.T) { + t.Parallel() + + clock := quartz.NewMock(t) + model := &chattest.FakeModel{ + ProviderName: "test-provider", + ModelName: "test-model", + StreamFn: func(_ context.Context, _ fantasy.Call) (fantasy.StreamResponse, error) { + return func(yield func(fantasy.StreamPart) bool) { + if !yield(fantasy.StreamPart{Type: fantasy.StreamPartTypeTextStart, ID: "t"}) { + return + } + if !yield(fantasy.StreamPart{Type: fantasy.StreamPartTypeTextDelta, ID: "t", Delta: "hello"}) { + return + } + if !yield(fantasy.StreamPart{Type: fantasy.StreamPartTypeTextEnd, ID: "t"}) { + return + } + clock.Advance(1500 * time.Millisecond) + yield(fantasy.StreamPart{ + Type: fantasy.StreamPartTypeFinish, + FinishReason: fantasy.FinishReasonStop, + }) + }, nil + }, + } + + outcome, err := chatloop.GenerateAssistant(context.Background(), chatloop.GenerateAssistantOptions{ + Model: model, + Clock: clock, + }) + require.NoError(t, err) + require.Equal(t, 1500*time.Millisecond, outcome.Step.Runtime) +} + +// TestGenerateAssistant_ErroredStreamReturnsNoStep pins that a failed +// model invocation produces no step: its content is discarded, so no +// runtime is billed for it either. +func TestGenerateAssistant_ErroredStreamReturnsNoStep(t *testing.T) { + t.Parallel() + + clock := quartz.NewMock(t) + model := &chattest.FakeModel{ + ProviderName: "test-provider", + ModelName: "test-model", + StreamFn: func(_ context.Context, _ fantasy.Call) (fantasy.StreamResponse, error) { + return func(yield func(fantasy.StreamPart) bool) { + if !yield(fantasy.StreamPart{Type: fantasy.StreamPartTypeTextStart, ID: "t"}) { + return + } + clock.Advance(1500 * time.Millisecond) + yield(fantasy.StreamPart{ + Type: fantasy.StreamPartTypeError, + Error: xerrors.New("stream blew up"), + }) + }, nil + }, + } + + outcome, err := chatloop.GenerateAssistant(context.Background(), chatloop.GenerateAssistantOptions{ + Model: model, + Clock: clock, + }) + require.Error(t, err) + require.Zero(t, outcome.Step.Runtime) + require.Empty(t, outcome.Step.Content) +} + +// TestGenerateCompaction_RecordsRuntime verifies the summarization model +// call reports its wall-clock duration, the compaction step's billable +// runtime. +func TestGenerateCompaction_RecordsRuntime(t *testing.T) { + t.Parallel() + + clock := quartz.NewMock(t) + model := &chattest.FakeModel{ + ProviderName: "test-provider", + ModelName: "test-model", + GenerateFn: func(_ context.Context, _ fantasy.Call) (*fantasy.Response, error) { + clock.Advance(1500 * time.Millisecond) + return &fantasy.Response{ + Content: []fantasy.Content{ + fantasy.TextContent{Text: "summary"}, + }, + }, nil + }, + } + + result, err := chatloop.GenerateCompaction(context.Background(), chatloop.GenerateCompactionOptions{ + Model: model, + Messages: []fantasy.Message{{ + Role: fantasy.MessageRoleUser, + Content: []fantasy.MessagePart{fantasy.TextPart{Text: "hello"}}, + }}, + ThresholdPercent: 70, + ContextLimit: 100, + StepUsage: fantasy.Usage{InputTokens: 90}, + Clock: clock, + }) + require.NoError(t, err) + require.Equal(t, "summary", result.SummaryReport) + require.Equal(t, 1500*time.Millisecond, result.Runtime) +} diff --git a/coderd/x/chatd/generation.go b/coderd/x/chatd/generation.go index 443909157df..6f6d93d76b9 100644 --- a/coderd/x/chatd/generation.go +++ b/coderd/x/chatd/generation.go @@ -754,6 +754,7 @@ func (s *taskStarter) generateCompaction( compactionOpts.PublishMessagePart = attempt.publish compactionOpts.Source = source compactionOpts.Force = source == chatloop.CompactionSourceManual + compactionOpts.Clock = s.opts.Clock // Attach the turn debug run so the compaction call records a child // debug run; without it startCompactionDebugRun finds no parent and // skips debug instrumentation entirely. diff --git a/coderd/x/chatd/message_conversion.go b/coderd/x/chatd/message_conversion.go index c70f1efcd8c..d53312e289d 100644 --- a/coderd/x/chatd/message_conversion.go +++ b/coderd/x/chatd/message_conversion.go @@ -319,6 +319,12 @@ func buildCompactionMessages(input buildCompactionMessagesInput) (compactionMess return compactionMessagesForCommit{}, xerrors.Errorf("marshal compaction tool result: %w", err) } + assistantMsg := baseMessage(database.ChatMessageRoleAssistant, database.ChatMessageVisibilityUser, input.modelConfigID, contentVersion, assistantContent) + // Compaction is a billable model invocation; its runtime is + // persisted on the assistant message like a regular step's. + if input.compaction.Runtime > 0 { + assistantMsg.RuntimeMs = sql.NullInt64{Int64: input.compaction.Runtime.Milliseconds(), Valid: true} + } messages := []chatstate.Message{ { Role: database.ChatMessageRoleUser, @@ -327,7 +333,7 @@ func buildCompactionMessages(input buildCompactionMessagesInput) (compactionMess ModelConfigID: uuid.NullUUID{UUID: input.modelConfigID, Valid: input.modelConfigID != uuid.Nil}, ContentVersion: contentVersion, }, - baseMessage(database.ChatMessageRoleAssistant, database.ChatMessageVisibilityUser, input.modelConfigID, contentVersion, assistantContent), + assistantMsg, baseMessage(database.ChatMessageRoleTool, database.ChatMessageVisibilityBoth, input.modelConfigID, contentVersion, toolContent), } for i := range messages { @@ -560,6 +566,13 @@ type bufferedPartsToPartialMessagesInput struct { contentVersion int16 logger slog.Logger interruptedAt time.Time + // attemptRuntime is the wall-clock duration of the interrupted + // generation attempt (its message part episode's lifetime). It is + // persisted as runtime_ms on the first partial assistant message + // so interruption does not lose billable generation time. When the + // interrupted attempt streamed no assistant content (for example a + // tool execution batch, which is not billable), it is dropped. + attemptRuntime time.Duration } type partialToolCall struct { @@ -610,6 +623,18 @@ func bufferedPartsToPartialMessages(input bufferedPartsToPartialMessagesInput) ( if err := state.appendSyntheticInterruptionResults(); err != nil { return nil, err } + if input.attemptRuntime > 0 { + // The whole attempt's runtime goes on the first assistant + // message; usage reporting sums runtime_ms across rows, so + // placement within the suffix does not matter. + for i := range state.messages { + if state.messages[i].Role != database.ChatMessageRoleAssistant { + continue + } + state.messages[i].RuntimeMs = sql.NullInt64{Int64: input.attemptRuntime.Milliseconds(), Valid: true} + break + } + } return state.messages, nil } diff --git a/coderd/x/chatd/message_conversion_test.go b/coderd/x/chatd/message_conversion_test.go index 4a14d3cd34a..daf144f3198 100644 --- a/coderd/x/chatd/message_conversion_test.go +++ b/coderd/x/chatd/message_conversion_test.go @@ -69,19 +69,27 @@ func TestBuildCommitStepMessages_LocalToolResultsBecomeToolMessages(t *testing.T modelConfigID: modelConfigID, contentVersion: chatprompt.CurrentContentVersion, logger: slog.Make(), - step: stepData{Content: []fantasy.Content{ - fantasy.ToolCallContent{ToolCallID: "call-1", ToolName: "execute", Input: `{"cmd":"pwd"}`}, - fantasy.ToolResultContent{ - ToolCallID: "call-1", - ToolName: "execute", - Result: fantasy.ToolResultOutputContentText{Text: `{"stdout":"/tmp"}`}, + step: stepData{ + Content: []fantasy.Content{ + fantasy.ToolCallContent{ToolCallID: "call-1", ToolName: "execute", Input: `{"cmd":"pwd"}`}, + fantasy.ToolResultContent{ + ToolCallID: "call-1", + ToolName: "execute", + Result: fantasy.ToolResultOutputContentText{Text: `{"stdout":"/tmp"}`}, + }, }, - }}, + Runtime: 1500 * time.Millisecond, + }, }) require.NoError(t, err) require.Len(t, got.Messages, 2) require.Equal(t, []int{0, 1}, got.VisibleIndexes) + // The step's model-invocation runtime lands on the assistant + // message only; tool result rows are never billed. + require.Equal(t, sql.NullInt64{Int64: 1500, Valid: true}, got.Messages[0].RuntimeMs) + require.False(t, got.Messages[1].RuntimeMs.Valid) + assistantParts := parseMessageParts(t, got.Messages[0].Role, got.Messages[0].Content) require.Len(t, assistantParts, 1) require.Equal(t, codersdk.ChatMessagePartTypeToolCall, assistantParts[0].Type) @@ -212,6 +220,7 @@ func TestBuildCompactionMessages_CompressedSummaryToolCallAndResult(t *testing.T UsagePercent: 81.5, ContextTokens: 815, ContextLimit: 1000, + Runtime: 1500 * time.Millisecond, }, }) require.NoError(t, err) @@ -223,10 +232,14 @@ func TestBuildCompactionMessages_CompressedSummaryToolCallAndResult(t *testing.T require.True(t, got.Messages[0].Compressed) require.Equal(t, uuid.NullUUID{UUID: modelConfigID, Valid: true}, got.Messages[0].ModelConfigID) require.Equal(t, "system summary", parseMessageParts(t, got.Messages[0].Role, got.Messages[0].Content)[0].Text) + require.False(t, got.Messages[0].RuntimeMs.Valid) require.Equal(t, database.ChatMessageRoleAssistant, got.Messages[1].Role) require.Equal(t, database.ChatMessageVisibilityUser, got.Messages[1].Visibility) require.True(t, got.Messages[1].Compressed) + // The summarization call's runtime is billed on the assistant + // message, mirroring regular generation steps. + require.Equal(t, sql.NullInt64{Int64: 1500, Valid: true}, got.Messages[1].RuntimeMs) callPart := parseMessageParts(t, got.Messages[1].Role, got.Messages[1].Content)[0] require.Equal(t, codersdk.ChatMessagePartTypeToolCall, callPart.Type) require.Equal(t, "summary-1", callPart.ToolCallID) @@ -235,6 +248,7 @@ func TestBuildCompactionMessages_CompressedSummaryToolCallAndResult(t *testing.T require.Equal(t, database.ChatMessageRoleTool, got.Messages[2].Role) require.Equal(t, database.ChatMessageVisibilityBoth, got.Messages[2].Visibility) require.True(t, got.Messages[2].Compressed) + require.False(t, got.Messages[2].RuntimeMs.Valid) resultPart := parseMessageParts(t, got.Messages[2].Role, got.Messages[2].Content)[0] require.Equal(t, codersdk.ChatMessagePartTypeToolResult, resultPart.Type) require.Equal(t, "summary-1", resultPart.ToolCallID) @@ -603,6 +617,50 @@ func TestBufferedPartsToPartialMessages_NormalizesToolCallDeltasBeforeFinal(t *t require.Equal(t, "call-1", syntheticParts[0].ToolCallID) } +// TestBufferedPartsToPartialMessages_AttachesAttemptRuntime verifies an +// interrupted attempt's runtime is persisted on the first partial +// assistant message, so interruption does not lose billable generation +// time. +func TestBufferedPartsToPartialMessages_AttachesAttemptRuntime(t *testing.T) { + t.Parallel() + + parts := []messagepartbuffer.Part{ + {Seq: 1, Role: codersdk.ChatMessageRoleAssistant, MessagePart: codersdk.ChatMessageText("partial ")}, + {Seq: 2, Role: codersdk.ChatMessageRoleAssistant, MessagePart: codersdk.ChatMessageToolCall("call-1", "execute", json.RawMessage(`{"cmd":"pwd"}`))}, + } + got, err := bufferedPartsToPartialMessages(bufferedPartsToPartialMessagesInput{ + parts: parts, + modelConfigID: uuid.New(), + contentVersion: chatprompt.CurrentContentVersion, + logger: slog.Make(), + attemptRuntime: 1500 * time.Millisecond, + }) + require.NoError(t, err) + require.Len(t, got, 2) + require.Equal(t, database.ChatMessageRoleAssistant, got[0].Role) + require.Equal(t, sql.NullInt64{Int64: 1500, Valid: true}, got[0].RuntimeMs) + // The synthetic interruption tool result is not billed. + require.Equal(t, database.ChatMessageRoleTool, got[1].Role) + require.False(t, got[1].RuntimeMs.Valid) +} + +// TestBufferedPartsToPartialMessages_DropsRuntimeWithoutAssistantContent +// verifies attempts that streamed no assistant content (for example an +// interrupted tool execution batch) do not bill their episode span. +func TestBufferedPartsToPartialMessages_DropsRuntimeWithoutAssistantContent(t *testing.T) { + t.Parallel() + + got, err := bufferedPartsToPartialMessages(bufferedPartsToPartialMessagesInput{ + parts: nil, + modelConfigID: uuid.New(), + contentVersion: chatprompt.CurrentContentVersion, + logger: slog.Make(), + attemptRuntime: 1500 * time.Millisecond, + }) + require.NoError(t, err) + require.Empty(t, got) +} + func TestBufferedPartsToPartialMessages_MergesToolCallDeltasWithoutFinal(t *testing.T) { t.Parallel() diff --git a/coderd/x/chatd/messagepartbuffer/message_part_buffer.go b/coderd/x/chatd/messagepartbuffer/message_part_buffer.go index 9b14c3287e4..bd98fa2b436 100644 --- a/coderd/x/chatd/messagepartbuffer/message_part_buffer.go +++ b/coderd/x/chatd/messagepartbuffer/message_part_buffer.go @@ -96,7 +96,11 @@ type Buffer struct { } type episodeState struct { - created bool + created bool + // createdAt is set only by CreateEpisode, not by the implicit + // creation in CloseEpisode or by subscriber placeholders, so it + // marks when the generation attempt actually started. + createdAt time.Time closed bool closedAt time.Time closedHeapItem *closedEpisodeItem @@ -185,12 +189,14 @@ func (b *Buffer) CreateEpisode(key Key) error { if b.closed { return ErrMessagePartBufferClosed } - b.gcClosedEpisodesLocked(b.opts.Clock.Now("message-part-buffer", "create")) + now := b.opts.Clock.Now("message-part-buffer", "create") + b.gcClosedEpisodesLocked(now) episode := b.getOrCreateEpisodeLocked(key) if episode.created { return ErrEpisodeExists } episode.markCreated() + episode.createdAt = now return nil } @@ -266,6 +272,20 @@ func (b *Buffer) GetParts(key Key) ([]Part, error) { return slices.Clone(episode.parts), nil } +// EpisodeDuration returns the wall-clock span between CreateEpisode and +// CloseEpisode. It returns 0 when the episode is unknown, was created +// implicitly by CloseEpisode, or is not closed yet. Interruption handling +// uses this as the runtime of the interrupted generation attempt. +func (b *Buffer) EpisodeDuration(key Key) time.Duration { + b.mu.Lock() + defer b.mu.Unlock() + episode := b.episodes[key] + if episode == nil || episode.createdAt.IsZero() || !episode.closed { + return 0 + } + return episode.closedAt.Sub(episode.createdAt) +} + // SubscribeToEpisode replays existing parts and streams new parts. // // Subscribers may attach before CreateEpisode is called. In that case the diff --git a/coderd/x/chatd/messagepartbuffer/message_part_buffer_test.go b/coderd/x/chatd/messagepartbuffer/message_part_buffer_test.go index 46ff27e5f35..beadc32bb33 100644 --- a/coderd/x/chatd/messagepartbuffer/message_part_buffer_test.go +++ b/coderd/x/chatd/messagepartbuffer/message_part_buffer_test.go @@ -99,6 +99,35 @@ func TestBuffer_CloseEpisodeIdempotent(t *testing.T) { require.NoError(t, buffer.CloseEpisode(key)) } +func TestBuffer_EpisodeDuration(t *testing.T) { + t.Parallel() + + clock := quartz.NewMock(t) + buffer := messagepartbuffer.New(messagepartbuffer.Options{Clock: clock}) + defer buffer.Close() + + key := testEpisodeKey() + require.Zero(t, buffer.EpisodeDuration(key), "unknown episode has no duration") + + require.NoError(t, buffer.CreateEpisode(key)) + require.Zero(t, buffer.EpisodeDuration(key), "open episode has no duration") + + clock.Advance(1500 * time.Millisecond) + require.NoError(t, buffer.CloseEpisode(key)) + require.Equal(t, 1500*time.Millisecond, buffer.EpisodeDuration(key)) + + // A second close must not move the recorded span. + clock.Advance(time.Second) + require.NoError(t, buffer.CloseEpisode(key)) + require.Equal(t, 1500*time.Millisecond, buffer.EpisodeDuration(key)) + + // Episodes created implicitly by CloseEpisode never started a + // generation attempt, so they report no duration. + implicit := testEpisodeKey() + require.NoError(t, buffer.CloseEpisode(implicit)) + require.Zero(t, buffer.EpisodeDuration(implicit)) +} + func TestBuffer_SubscribeExistingReplaysThenStreamsLiveParts(t *testing.T) { t.Parallel() diff --git a/coderd/x/chatd/tasks.go b/coderd/x/chatd/tasks.go index c1c0c840e4a..f2e8185f392 100644 --- a/coderd/x/chatd/tasks.go +++ b/coderd/x/chatd/tasks.go @@ -283,6 +283,7 @@ func (s *taskStarter) StartInterrupt(ctx context.Context, input chatWorkerTaskSt contentVersion: chatprompt.CurrentContentVersion, logger: s.opts.Logger, interruptedAt: s.opts.Clock.Now("chatworker", "interrupt"), + attemptRuntime: s.opts.MessagePartBuffer.EpisodeDuration(key), }) if err != nil { return xerrors.Errorf("convert buffered parts: %w", err) diff --git a/coderd/x/chatd/tasks_test.go b/coderd/x/chatd/tasks_test.go index 8f95da4e193..fd1c3436c5c 100644 --- a/coderd/x/chatd/tasks_test.go +++ b/coderd/x/chatd/tasks_test.go @@ -391,6 +391,51 @@ func TestInterruptTask_BufferedPartsBecomePartialMessages(t *testing.T) { require.True(t, toolParts[0].IsError) } +// TestInterruptTask_PartialAssistantKeepsAttemptRuntime verifies an +// interrupted generation attempt's runtime (the episode's lifetime) is +// persisted as runtime_ms on the partial assistant message, so billable +// generation time survives interruption. +func TestInterruptTask_PartialAssistantKeepsAttemptRuntime(t *testing.T) { + t.Parallel() + + f := newTaskTestFixture(t) + chat := f.createRunningChat(t) + workerID := uuid.New() + runnerID := uuid.New() + acquired := f.acquireChat(t, chat.ID, workerID, runnerID) + recorder := newTaskSideEffectRecorder() + clock := quartz.NewMock(t) + starter := newTestTaskStarterWithClock(t, f, recorder, clock) + buffer := starter.opts.MessagePartBuffer + key := messagepartbuffer.Key{ + ChatID: chat.ID, + HistoryVersion: acquired.HistoryVersion, + GenerationAttempt: acquired.GenerationAttempt, + } + require.NoError(t, buffer.CreateEpisode(key)) + require.NoError(t, buffer.AddPart(key, codersdk.ChatMessageRoleAssistant, codersdk.ChatMessageText("partial answer"))) + // The attempt ran for 1500ms before the user interrupted it. + clock.Advance(1500 * time.Millisecond) + interrupting := f.interruptChat(t, chat.ID) + + err := starter.StartInterrupt(testutil.Context(t, testutil.WaitLong), chatWorkerTaskStartInput{ + ChatID: chat.ID, + WorkerID: workerID, + RunnerID: runnerID, + HistoryVersion: interrupting.HistoryVersion, + GenerationAttempt: interrupting.GenerationAttempt, + Status: database.ChatStatusInterrupting, + }) + require.NoError(t, err) + + messages, err := f.db.GetChatMessagesByChatID(testutil.Context(t, testutil.WaitShort), database.GetChatMessagesByChatIDParams{ChatID: chat.ID}) + require.NoError(t, err) + require.GreaterOrEqual(t, len(messages), 3) + assistant := messages[len(messages)-2] + require.Equal(t, database.ChatMessageRoleAssistant, assistant.Role) + require.Equal(t, sql.NullInt64{Int64: 1500, Valid: true}, assistant.RuntimeMs) +} + func TestRequiresActionTimeout_ExpiredCancelsOnly(t *testing.T) { t.Parallel() @@ -1244,13 +1289,21 @@ func (r *taskSideEffectRecorder) requireInterruptionOutcome(t *testing.T, chatID func newTestTaskStarter(t *testing.T, f *taskTestFixture, recorder *taskSideEffectRecorder) *taskStarter { t.Helper() - buffer := messagepartbuffer.New(messagepartbuffer.Options{}) + return newTestTaskStarterWithClock(t, f, recorder, quartz.NewReal()) +} + +// newTestTaskStarterWithClock shares the clock between the starter and its +// message part buffer, mirroring production wiring so episode durations are +// measured on the same clock the tasks use. +func newTestTaskStarterWithClock(t *testing.T, f *taskTestFixture, recorder *taskSideEffectRecorder, clock quartz.Clock) *taskStarter { + t.Helper() + buffer := messagepartbuffer.New(messagepartbuffer.Options{Clock: clock}) t.Cleanup(buffer.Close) starter, err := newTaskStarter(newUnstartedServer(t, f.rawPS, f.db), chatWorkerOptions{ Store: f.db, Pubsub: f.pubsub, Logger: slog.Make(), - Clock: quartz.NewReal(), + Clock: clock, MessagePartBuffer: buffer, TaskRetryInitialBackoff: time.Millisecond, TaskRetryMaxBackoff: time.Millisecond, diff --git a/docs/ai-coder/agents/platform-controls/usage-insights.md b/docs/ai-coder/agents/platform-controls/usage-insights.md index b6b2d1e5db1..0348604f8b6 100644 --- a/docs/ai-coder/agents/platform-controls/usage-insights.md +++ b/docs/ai-coder/agents/platform-controls/usage-insights.md @@ -88,3 +88,30 @@ Select a user to see: > Automatic title generation uses lightweight models, such as Claude Haiku or GPT-4o > Mini. Its token usage is not counted towards usage limits or shown in usage > summaries. + +## Generation runtime + +Cost summaries and usage reporting include agent runtime, summed from +per-message generation time (`runtime_ms` on chat messages). + +A message's runtime is the wall-clock duration of the model invocation that +produced its content, measured from just before the request to the model +provider opens until the response is fully consumed. + +What counts: + +- Assistant generation steps, in both top-level chats and sub-agent chats. +- Context compaction (summarization) model calls. +- Interrupted generation: the time streamed before the interrupt is kept on + the partial assistant message. + +What does not count: + +- Local tool execution, including waiting on sub-agents. A sub-agent is its + own chat and records its own model invocations, so counting the parent's + wait would double count. +- Idle time: chats waiting for user input or external tool results. +- Failed model calls whose output was discarded. Retried and errored + attempts persist no content, so they record no runtime. +- Ancillary model calls that produce no chat messages, such as title + generation. From 7c2486b79b0ef9807f08ebf93736f76d77db14a3 Mon Sep 17 00:00:00 2001 From: Jaayden Halko Date: Thu, 23 Jul 2026 10:10:07 +0000 Subject: [PATCH 02/11] fix: do not bill tool-batch episode spans on interrupted attachments Tool execution publishes assistant-role file parts for attachments, so an interrupt landing between a tool batch and its commit could produce a file-only partial assistant message and attract the batch's episode span as runtime_ms. Gate the runtime attachment on the suffix containing model-streamed assistant content (text, reasoning, tool calls, sources). Co-Authored-By: Claude Fable 5 --- coderd/x/chatd/message_conversion.go | 16 +++++++++++++--- coderd/x/chatd/message_conversion_test.go | 22 ++++++++++++++++++++-- 2 files changed, 33 insertions(+), 5 deletions(-) diff --git a/coderd/x/chatd/message_conversion.go b/coderd/x/chatd/message_conversion.go index d53312e289d..632afe5417b 100644 --- a/coderd/x/chatd/message_conversion.go +++ b/coderd/x/chatd/message_conversion.go @@ -570,8 +570,9 @@ type bufferedPartsToPartialMessagesInput struct { // generation attempt (its message part episode's lifetime). It is // persisted as runtime_ms on the first partial assistant message // so interruption does not lose billable generation time. When the - // interrupted attempt streamed no assistant content (for example a - // tool execution batch, which is not billable), it is dropped. + // interrupted attempt streamed no model-generated assistant content + // it is dropped: tool execution batches are not billable, including + // ones that published assistant-role file parts for attachments. attemptRuntime time.Duration } @@ -623,7 +624,7 @@ func bufferedPartsToPartialMessages(input bufferedPartsToPartialMessagesInput) ( if err := state.appendSyntheticInterruptionResults(); err != nil { return nil, err } - if input.attemptRuntime > 0 { + if input.attemptRuntime > 0 && state.modelStreamedAssistant { // The whole attempt's runtime goes on the first assistant // message; usage reporting sums runtime_ms across rows, so // placement within the suffix does not matter. @@ -649,6 +650,12 @@ type partialMessageConversionState struct { toolResults map[string]*partialToolResult toolResultOrder []string answered map[string]bool + // modelStreamedAssistant records whether any assistant part came + // from the model stream itself (text, reasoning, tool calls, + // sources). Tool execution also publishes assistant-role file + // parts for attachments; those alone must not attract the + // attempt's runtime, because tool batches are not billable. + modelStreamedAssistant bool } func (s *partialMessageConversionState) consume(buffered messagepartbuffer.Part) error { @@ -669,6 +676,9 @@ func (s *partialMessageConversionState) consumeAssistantPart(buffered messagepar s.logSkippedPart(buffered, "empty buffered assistant part type") return } + if part.Type != codersdk.ChatMessagePartTypeFile { + s.modelStreamedAssistant = true + } if part.Type != codersdk.ChatMessagePartTypeToolCall { if part.Type == codersdk.ChatMessagePartTypeReasoning && !s.input.interruptedAt.IsZero() { diff --git a/coderd/x/chatd/message_conversion_test.go b/coderd/x/chatd/message_conversion_test.go index daf144f3198..98a6df02556 100644 --- a/coderd/x/chatd/message_conversion_test.go +++ b/coderd/x/chatd/message_conversion_test.go @@ -645,8 +645,9 @@ func TestBufferedPartsToPartialMessages_AttachesAttemptRuntime(t *testing.T) { } // TestBufferedPartsToPartialMessages_DropsRuntimeWithoutAssistantContent -// verifies attempts that streamed no assistant content (for example an -// interrupted tool execution batch) do not bill their episode span. +// verifies attempts that streamed no model-generated assistant content +// (for example an interrupted tool execution batch) do not bill their +// episode span. func TestBufferedPartsToPartialMessages_DropsRuntimeWithoutAssistantContent(t *testing.T) { t.Parallel() @@ -659,6 +660,23 @@ func TestBufferedPartsToPartialMessages_DropsRuntimeWithoutAssistantContent(t *t }) require.NoError(t, err) require.Empty(t, got) + + // Tool execution publishes assistant-role file parts for + // attachments. A suffix containing only those is a tool batch, + // not model generation, so its span is not billed. + got, err = bufferedPartsToPartialMessages(bufferedPartsToPartialMessagesInput{ + parts: []messagepartbuffer.Part{ + {Seq: 1, Role: codersdk.ChatMessageRoleAssistant, MessagePart: codersdk.ChatMessageFile(uuid.New(), "image/png", "screenshot.png")}, + }, + modelConfigID: uuid.New(), + contentVersion: chatprompt.CurrentContentVersion, + logger: slog.Make(), + attemptRuntime: 1500 * time.Millisecond, + }) + require.NoError(t, err) + require.Len(t, got, 1) + require.Equal(t, database.ChatMessageRoleAssistant, got[0].Role) + require.False(t, got[0].RuntimeMs.Valid) } func TestBufferedPartsToPartialMessages_MergesToolCallDeltasWithoutFinal(t *testing.T) { From 525b0ad89d79a77e33eb0f0df9bebed34bcfe2bd Mon Sep 17 00:00:00 2001 From: Jaayden Halko Date: Thu, 23 Jul 2026 10:31:39 +0000 Subject: [PATCH 03/11] refactor(coderd/x/chatd): trim comments to edge cases and invariants Co-Authored-By: Claude Fable 5 --- coderd/x/chatd/chatd_test.go | 2 -- coderd/x/chatd/chatloop/chatloop.go | 3 +-- coderd/x/chatd/chatloop/compaction.go | 6 +++--- coderd/x/chatd/chatloop/runtime_test.go | 9 --------- coderd/x/chatd/message_conversion.go | 18 ++++++------------ coderd/x/chatd/message_conversion_test.go | 14 -------------- .../messagepartbuffer/message_part_buffer.go | 3 +-- coderd/x/chatd/tasks_test.go | 8 +------- 8 files changed, 12 insertions(+), 51 deletions(-) diff --git a/coderd/x/chatd/chatd_test.go b/coderd/x/chatd/chatd_test.go index 7a206e47757..765824045f8 100644 --- a/coderd/x/chatd/chatd_test.go +++ b/coderd/x/chatd/chatd_test.go @@ -6340,8 +6340,6 @@ func TestActiveServer_BasicAssistantGenerationAndPromptPreparation(t *testing.T) require.Equal(t, database.ChatMessageRoleAssistant, last.Role) require.True(t, last.ContextLimit.Valid) require.Equal(t, int64(4096), last.ContextLimit.Int64) - // Assistant steps must persist the model invocation's runtime; it is - // the billing source of truth for Coder Agents runtime. require.True(t, last.RuntimeMs.Valid) require.GreaterOrEqual(t, last.RuntimeMs.Int64, int64(0)) requireTextPart(t, last, "done") diff --git a/coderd/x/chatd/chatloop/chatloop.go b/coderd/x/chatd/chatloop/chatloop.go index 50662b402aa..84a3046c1ab 100644 --- a/coderd/x/chatd/chatloop/chatloop.go +++ b/coderd/x/chatd/chatloop/chatloop.go @@ -311,8 +311,7 @@ type GenerateCompactionOptions struct { PublishMessagePart func(codersdk.ChatMessageRole, codersdk.ChatMessagePart) - // Clock measures the summary call duration reported as - // CompactionResult.Runtime. Nil uses a real clock. + // Clock measures the summary call duration. Nil uses a real clock. Clock quartz.Clock } diff --git a/coderd/x/chatd/chatloop/compaction.go b/coderd/x/chatd/chatloop/compaction.go index e375ee9c81b..62c2bcfbf3e 100644 --- a/coderd/x/chatd/chatloop/compaction.go +++ b/coderd/x/chatd/chatloop/compaction.go @@ -131,9 +131,9 @@ type CompactionResult struct { ContextTokens int64 ContextLimit int64 // Runtime is the wall-clock duration of the summarization model - // call. Compaction is a billable model invocation like a regular - // assistant step; see PersistedStep.Runtime. Zero when the run - // was gated off before calling the model. + // call, the compaction step's billable runtime (see + // PersistedStep.Runtime). Zero when the run was gated off before + // calling the model. Runtime time.Duration } diff --git a/coderd/x/chatd/chatloop/runtime_test.go b/coderd/x/chatd/chatloop/runtime_test.go index 6bf2009e116..c42e86797aa 100644 --- a/coderd/x/chatd/chatloop/runtime_test.go +++ b/coderd/x/chatd/chatloop/runtime_test.go @@ -14,9 +14,6 @@ import ( "github.com/coder/quartz" ) -// TestGenerateAssistant_RecordsModelInvocationRuntime pins the billable -// runtime definition: Step.Runtime spans the model invocation, from just -// before the stream opens until it is fully consumed. func TestGenerateAssistant_RecordsModelInvocationRuntime(t *testing.T) { t.Parallel() @@ -52,9 +49,6 @@ func TestGenerateAssistant_RecordsModelInvocationRuntime(t *testing.T) { require.Equal(t, 1500*time.Millisecond, outcome.Step.Runtime) } -// TestGenerateAssistant_ErroredStreamReturnsNoStep pins that a failed -// model invocation produces no step: its content is discarded, so no -// runtime is billed for it either. func TestGenerateAssistant_ErroredStreamReturnsNoStep(t *testing.T) { t.Parallel() @@ -85,9 +79,6 @@ func TestGenerateAssistant_ErroredStreamReturnsNoStep(t *testing.T) { require.Empty(t, outcome.Step.Content) } -// TestGenerateCompaction_RecordsRuntime verifies the summarization model -// call reports its wall-clock duration, the compaction step's billable -// runtime. func TestGenerateCompaction_RecordsRuntime(t *testing.T) { t.Parallel() diff --git a/coderd/x/chatd/message_conversion.go b/coderd/x/chatd/message_conversion.go index 632afe5417b..9a55ebc2939 100644 --- a/coderd/x/chatd/message_conversion.go +++ b/coderd/x/chatd/message_conversion.go @@ -320,8 +320,6 @@ func buildCompactionMessages(input buildCompactionMessagesInput) (compactionMess } assistantMsg := baseMessage(database.ChatMessageRoleAssistant, database.ChatMessageVisibilityUser, input.modelConfigID, contentVersion, assistantContent) - // Compaction is a billable model invocation; its runtime is - // persisted on the assistant message like a regular step's. if input.compaction.Runtime > 0 { assistantMsg.RuntimeMs = sql.NullInt64{Int64: input.compaction.Runtime.Milliseconds(), Valid: true} } @@ -566,13 +564,10 @@ type bufferedPartsToPartialMessagesInput struct { contentVersion int16 logger slog.Logger interruptedAt time.Time - // attemptRuntime is the wall-clock duration of the interrupted - // generation attempt (its message part episode's lifetime). It is - // persisted as runtime_ms on the first partial assistant message - // so interruption does not lose billable generation time. When the - // interrupted attempt streamed no model-generated assistant content - // it is dropped: tool execution batches are not billable, including - // ones that published assistant-role file parts for attachments. + // attemptRuntime is the interrupted generation attempt's wall-clock + // duration (its message part episode's lifetime), persisted as + // runtime_ms on the first partial assistant message when the attempt + // streamed model-generated assistant content. attemptRuntime time.Duration } @@ -625,9 +620,8 @@ func bufferedPartsToPartialMessages(input bufferedPartsToPartialMessagesInput) ( return nil, err } if input.attemptRuntime > 0 && state.modelStreamedAssistant { - // The whole attempt's runtime goes on the first assistant - // message; usage reporting sums runtime_ms across rows, so - // placement within the suffix does not matter. + // Usage reporting sums runtime_ms across rows, so placing the + // whole span on the first assistant message is sufficient. for i := range state.messages { if state.messages[i].Role != database.ChatMessageRoleAssistant { continue diff --git a/coderd/x/chatd/message_conversion_test.go b/coderd/x/chatd/message_conversion_test.go index 98a6df02556..59a9eab3fdf 100644 --- a/coderd/x/chatd/message_conversion_test.go +++ b/coderd/x/chatd/message_conversion_test.go @@ -84,9 +84,6 @@ func TestBuildCommitStepMessages_LocalToolResultsBecomeToolMessages(t *testing.T require.NoError(t, err) require.Len(t, got.Messages, 2) require.Equal(t, []int{0, 1}, got.VisibleIndexes) - - // The step's model-invocation runtime lands on the assistant - // message only; tool result rows are never billed. require.Equal(t, sql.NullInt64{Int64: 1500, Valid: true}, got.Messages[0].RuntimeMs) require.False(t, got.Messages[1].RuntimeMs.Valid) @@ -237,8 +234,6 @@ func TestBuildCompactionMessages_CompressedSummaryToolCallAndResult(t *testing.T require.Equal(t, database.ChatMessageRoleAssistant, got.Messages[1].Role) require.Equal(t, database.ChatMessageVisibilityUser, got.Messages[1].Visibility) require.True(t, got.Messages[1].Compressed) - // The summarization call's runtime is billed on the assistant - // message, mirroring regular generation steps. require.Equal(t, sql.NullInt64{Int64: 1500, Valid: true}, got.Messages[1].RuntimeMs) callPart := parseMessageParts(t, got.Messages[1].Role, got.Messages[1].Content)[0] require.Equal(t, codersdk.ChatMessagePartTypeToolCall, callPart.Type) @@ -617,10 +612,6 @@ func TestBufferedPartsToPartialMessages_NormalizesToolCallDeltasBeforeFinal(t *t require.Equal(t, "call-1", syntheticParts[0].ToolCallID) } -// TestBufferedPartsToPartialMessages_AttachesAttemptRuntime verifies an -// interrupted attempt's runtime is persisted on the first partial -// assistant message, so interruption does not lose billable generation -// time. func TestBufferedPartsToPartialMessages_AttachesAttemptRuntime(t *testing.T) { t.Parallel() @@ -639,15 +630,10 @@ func TestBufferedPartsToPartialMessages_AttachesAttemptRuntime(t *testing.T) { require.Len(t, got, 2) require.Equal(t, database.ChatMessageRoleAssistant, got[0].Role) require.Equal(t, sql.NullInt64{Int64: 1500, Valid: true}, got[0].RuntimeMs) - // The synthetic interruption tool result is not billed. require.Equal(t, database.ChatMessageRoleTool, got[1].Role) require.False(t, got[1].RuntimeMs.Valid) } -// TestBufferedPartsToPartialMessages_DropsRuntimeWithoutAssistantContent -// verifies attempts that streamed no model-generated assistant content -// (for example an interrupted tool execution batch) do not bill their -// episode span. func TestBufferedPartsToPartialMessages_DropsRuntimeWithoutAssistantContent(t *testing.T) { t.Parallel() diff --git a/coderd/x/chatd/messagepartbuffer/message_part_buffer.go b/coderd/x/chatd/messagepartbuffer/message_part_buffer.go index bd98fa2b436..a21c3566a2d 100644 --- a/coderd/x/chatd/messagepartbuffer/message_part_buffer.go +++ b/coderd/x/chatd/messagepartbuffer/message_part_buffer.go @@ -274,8 +274,7 @@ func (b *Buffer) GetParts(key Key) ([]Part, error) { // EpisodeDuration returns the wall-clock span between CreateEpisode and // CloseEpisode. It returns 0 when the episode is unknown, was created -// implicitly by CloseEpisode, or is not closed yet. Interruption handling -// uses this as the runtime of the interrupted generation attempt. +// implicitly by CloseEpisode, or is not closed yet. func (b *Buffer) EpisodeDuration(key Key) time.Duration { b.mu.Lock() defer b.mu.Unlock() diff --git a/coderd/x/chatd/tasks_test.go b/coderd/x/chatd/tasks_test.go index fd1c3436c5c..c5cacb2f3ee 100644 --- a/coderd/x/chatd/tasks_test.go +++ b/coderd/x/chatd/tasks_test.go @@ -391,10 +391,6 @@ func TestInterruptTask_BufferedPartsBecomePartialMessages(t *testing.T) { require.True(t, toolParts[0].IsError) } -// TestInterruptTask_PartialAssistantKeepsAttemptRuntime verifies an -// interrupted generation attempt's runtime (the episode's lifetime) is -// persisted as runtime_ms on the partial assistant message, so billable -// generation time survives interruption. func TestInterruptTask_PartialAssistantKeepsAttemptRuntime(t *testing.T) { t.Parallel() @@ -414,7 +410,6 @@ func TestInterruptTask_PartialAssistantKeepsAttemptRuntime(t *testing.T) { } require.NoError(t, buffer.CreateEpisode(key)) require.NoError(t, buffer.AddPart(key, codersdk.ChatMessageRoleAssistant, codersdk.ChatMessageText("partial answer"))) - // The attempt ran for 1500ms before the user interrupted it. clock.Advance(1500 * time.Millisecond) interrupting := f.interruptChat(t, chat.ID) @@ -1293,8 +1288,7 @@ func newTestTaskStarter(t *testing.T, f *taskTestFixture, recorder *taskSideEffe } // newTestTaskStarterWithClock shares the clock between the starter and its -// message part buffer, mirroring production wiring so episode durations are -// measured on the same clock the tasks use. +// message part buffer, mirroring production wiring. func newTestTaskStarterWithClock(t *testing.T, f *taskTestFixture, recorder *taskSideEffectRecorder, clock quartz.Clock) *taskStarter { t.Helper() buffer := messagepartbuffer.New(messagepartbuffer.Options{Clock: clock}) From 359c93a11d83d6ea07b41b6ae9de3cfabec8afe5 Mon Sep 17 00:00:00 2001 From: Jaayden Halko Date: Mon, 27 Jul 2026 12:56:07 +0000 Subject: [PATCH 04/11] fix(coderd): bill interrupted turns for the model invocation window Interrupted attempts billed the whole buffer episode lifetime, which starts in beginGenerationAttempt before prompt preparation, so they over-counted against the documented "stream open until fully consumed" definition (review finding CRF-1, option A). The buffer now stamps the provider stream open instant via StartModelInvocation, wired from a new chatloop OnModelStreamStart hook that fires at the same instant PersistedStep.Runtime starts measuring. ModelInvocationDuration replaces EpisodeDuration, so episodes that never invoke a model (local tool execution batches) report no runtime at all. Also from review: - Drop the flaky runtime assertion in the prompt preparation test. It is the actual red Flake Check: InsertChatMessages maps runtime_ms 0 to NULL, and the in-process stream regularly finishes in under a millisecond (CRF-4, CRF-8). - Pin the zero-runtime exclusion for tool batches and compaction, and cover the tool-execution episode end to end (CRF-2, CRF-3). - Route runtime conversions through nullInt64IfNonZero (CRF-10) and rename the shadowing local `runtime` (CRF-9). - Point the duplicated definition at the column comment as canonical, and record the sub-millisecond NULL behavior there (CRF-7). --- coderd/database/dump.sql | 2 +- ...51_chat_messages_runtime_ms_comment.up.sql | 2 +- coderd/database/models.go | 2 +- coderd/x/chatd/ARCHITECTURE.md | 4 +- coderd/x/chatd/chatd_test.go | 13 ++++- coderd/x/chatd/chatloop/chatloop.go | 29 +++++++--- coderd/x/chatd/chatloop/compaction.go | 7 ++- coderd/x/chatd/chatloop/runtime_test.go | 48 ++++++++++++++++ coderd/x/chatd/generation.go | 10 ++++ coderd/x/chatd/message_conversion.go | 26 +++++---- coderd/x/chatd/message_conversion_test.go | 56 +++++++++++++++++++ .../messagepartbuffer/message_part_buffer.go | 50 ++++++++++++----- .../message_part_buffer_test.go | 29 +++++++--- coderd/x/chatd/tasks.go | 2 +- coderd/x/chatd/tasks_test.go | 46 +++++++++++++++ 15 files changed, 277 insertions(+), 49 deletions(-) diff --git a/coderd/database/dump.sql b/coderd/database/dump.sql index 7ef51d99226..084fe03c5a1 100644 --- a/coderd/database/dump.sql +++ b/coderd/database/dump.sql @@ -1968,7 +1968,7 @@ CREATE TABLE chat_messages ( search_tsv tsvector ); -COMMENT ON COLUMN chat_messages.runtime_ms IS 'Wall-clock milliseconds of the model invocation that produced the message content: assistant steps and compaction summaries, including interrupted partials. NULL when no model invocation produced the row (user messages, tool results; local tool execution is not counted). Billing source of truth for Coder Agents runtime: usage reporting sums it over created_at ranges.'; +COMMENT ON COLUMN chat_messages.runtime_ms IS 'Wall-clock milliseconds of the model invocation that produced the message content: assistant steps and compaction summaries, including interrupted partials, which measure from the provider stream opening until the interrupt closed it. NULL when no model invocation produced the row (user messages, tool results; local tool execution is not counted) or when the invocation rounded down to zero milliseconds. Billing source of truth for Coder Agents runtime: usage reporting sums it over created_at ranges.'; COMMENT ON COLUMN chat_messages.reasoning_effort IS 'Stores the selected effort for the turn triggered by this message.'; diff --git a/coderd/database/migrations/000551_chat_messages_runtime_ms_comment.up.sql b/coderd/database/migrations/000551_chat_messages_runtime_ms_comment.up.sql index 0c2bcc420c8..85bc1f27e31 100644 --- a/coderd/database/migrations/000551_chat_messages_runtime_ms_comment.up.sql +++ b/coderd/database/migrations/000551_chat_messages_runtime_ms_comment.up.sql @@ -1 +1 @@ -COMMENT ON COLUMN chat_messages.runtime_ms IS 'Wall-clock milliseconds of the model invocation that produced the message content: assistant steps and compaction summaries, including interrupted partials. NULL when no model invocation produced the row (user messages, tool results; local tool execution is not counted). Billing source of truth for Coder Agents runtime: usage reporting sums it over created_at ranges.'; +COMMENT ON COLUMN chat_messages.runtime_ms IS 'Wall-clock milliseconds of the model invocation that produced the message content: assistant steps and compaction summaries, including interrupted partials, which measure from the provider stream opening until the interrupt closed it. NULL when no model invocation produced the row (user messages, tool results; local tool execution is not counted) or when the invocation rounded down to zero milliseconds. Billing source of truth for Coder Agents runtime: usage reporting sums it over created_at ranges.'; diff --git a/coderd/database/models.go b/coderd/database/models.go index 680de116a8d..8f98857d7c2 100644 --- a/coderd/database/models.go +++ b/coderd/database/models.go @@ -5117,7 +5117,7 @@ type ChatMessage struct { CreatedBy uuid.NullUUID `db:"created_by" json:"created_by"` ContentVersion int16 `db:"content_version" json:"content_version"` TotalCostMicros sql.NullInt64 `db:"total_cost_micros" json:"total_cost_micros"` - // Wall-clock milliseconds of the model invocation that produced the message content: assistant steps and compaction summaries, including interrupted partials. NULL when no model invocation produced the row (user messages, tool results; local tool execution is not counted). Billing source of truth for Coder Agents runtime: usage reporting sums it over created_at ranges. + // Wall-clock milliseconds of the model invocation that produced the message content: assistant steps and compaction summaries, including interrupted partials, which measure from the provider stream opening until the interrupt closed it. NULL when no model invocation produced the row (user messages, tool results; local tool execution is not counted) or when the invocation rounded down to zero milliseconds. Billing source of truth for Coder Agents runtime: usage reporting sums it over created_at ranges. RuntimeMs sql.NullInt64 `db:"runtime_ms" json:"runtime_ms"` Deleted bool `db:"deleted" json:"deleted"` ProviderResponseID sql.NullString `db:"provider_response_id" json:"provider_response_id"` diff --git a/coderd/x/chatd/ARCHITECTURE.md b/coderd/x/chatd/ARCHITECTURE.md index 5b2ded5ed21..ac3445c983f 100644 --- a/coderd/x/chatd/ARCHITECTURE.md +++ b/coderd/x/chatd/ARCHITECTURE.md @@ -716,7 +716,7 @@ The buffer exposes the following API: - `CloseEpisode(chat_id, history_version, generation_attempt)`: closes an episode, preventing further parts from being added to it. May be called multiple times for a given episode, subsequent calls will be no-ops. Calling it on a non-existent episode creates the episode and closes it immediately. Concurrent parts of the system may race to create the episode and close it, so creating and closing in one operation prevents race conditions. - `AddPart(chat_id, history_version, generation_attempt, content)`: adds a message part to the buffer. Returns a predefined error if the episode is not found or the array is full. - `GetParts(chat_id, history_version, generation_attempt)`: returns the message parts for an episode. Returns a predefined error if the episode is not found. -- `EpisodeDuration(chat_id, history_version, generation_attempt)`: returns the wall-clock span between `CreateEpisode` and `CloseEpisode`, or zero when the episode is unknown, was created implicitly by `CloseEpisode`, or is still open. The interrupt goroutine uses it as the interrupted generation attempt's runtime. +- `EpisodeModelInvocation(chat_id, history_version, generation_attempt)`: `StartModelInvocation` stamps the instant the episode opens its provider stream, and `ModelInvocationDuration` returns the span between that stamp and `CloseEpisode`. The duration is zero when the episode is unknown, never opened a provider stream (such as a local tool execution batch), or is still open. The interrupt goroutine uses it as the interrupted attempt's billable runtime. - `SubscribeToEpisode(chat_id, history_version, generation_attempt)`: returns a go channel that will receive all message parts for the episode. It spawns a goroutine that delivers parts to the channel. It's live until the episode is closed or until a subscriber requests that the channel be closed. Once the goroutine delivers all message parts for a closed episode, it closes the channel and exits. If the episode is already closed at the time of the call, the goroutine delivers all message parts for the episode, closes the channel, and exits. `SubscribeToEpisode` does not return an error if the episode is not found: it waits for it to be created instead. Closed episodes are garbage collected after at least 15 seconds since they were closed and when they have no active subscribers. The message part buffer maintains a garbage collection goroutine. @@ -876,7 +876,7 @@ The goroutine does the following in order: 1. It fetches the generation attempt number from the database. 2. It closes the episode corresponding to its history version and generation attempt by calling the `CloseEpisode` method on the [Message part buffer](#message-part-buffer). 3. It reads the buffered parts for that episode by calling the `GetParts` method on the message part buffer. -4. It applies the `FinishInterruption(partial?)` transition on the core state machine. If there are no buffered parts for that episode, or the episode is not found, it passes `nil` as the `partial` argument. When the partial suffix contains an assistant message, the episode's lifetime (`EpisodeDuration`) is persisted on that message as `runtime_ms`, so the interrupted attempt's billable generation time is not lost. +4. It applies the `FinishInterruption(partial?)` transition on the core state machine. If there are no buffered parts for that episode, or the episode is not found, it passes `nil` as the `partial` argument. When the partial suffix contains an assistant message, the episode's model invocation window (`ModelInvocationDuration`: provider stream open until the interrupt closed the episode) is persisted on that message as `runtime_ms`, so the interrupted attempt's billable generation time is not lost. #### Dynamic tools timeout goroutine diff --git a/coderd/x/chatd/chatd_test.go b/coderd/x/chatd/chatd_test.go index 765824045f8..d12eed64f33 100644 --- a/coderd/x/chatd/chatd_test.go +++ b/coderd/x/chatd/chatd_test.go @@ -3256,6 +3256,7 @@ func TestActiveServer_InterruptionBehavior(t *testing.T) { messages := chatMessages(ctx, t, db, chat.ID) var userTexts []string var foundPartial bool + var partialRuntime sql.NullInt64 for _, msg := range messages { parts, parseErr := chatprompt.ParseContent(msg) require.NoError(t, parseErr) @@ -3270,12 +3271,16 @@ func TestActiveServer_InterruptionBehavior(t *testing.T) { for _, part := range parts { if part.Type == codersdk.ChatMessagePartTypeText && strings.Contains(part.Text, "partial assistant output") { foundPartial = true + partialRuntime = msg.RuntimeMs } } } } require.Equal(t, []string{"start and call a tool", "queued after interrupt"}, userTexts) require.True(t, foundPartial) + // The interrupted attempt bills the model invocation window it + // opened, so the partial assistant row keeps a runtime. + require.True(t, partialRuntime.Valid) parts := chatToolParts(ctx, t, db, chat.ID) call := requireToolCallPart(t, parts, "read_file") @@ -6340,8 +6345,12 @@ func TestActiveServer_BasicAssistantGenerationAndPromptPreparation(t *testing.T) require.Equal(t, database.ChatMessageRoleAssistant, last.Role) require.True(t, last.ContextLimit.Valid) require.Equal(t, int64(4096), last.ContextLimit.Int64) - require.True(t, last.RuntimeMs.Valid) - require.GreaterOrEqual(t, last.RuntimeMs.Int64, int64(0)) + // runtime_ms is not asserted here: this stream is served in + // process and routinely finishes in under a millisecond, which + // InsertChatMessages stores as NULL. Runtime measurement and + // persistence are pinned deterministically in + // chatloop.TestGenerateAssistant_RecordsModelInvocationRuntime and + // TestInterruptTask_PartialAssistantKeepsAttemptRuntime. requireTextPart(t, last, "done") requests = newAnthropicRequestRecorder() diff --git a/coderd/x/chatd/chatloop/chatloop.go b/coderd/x/chatd/chatloop/chatloop.go index 84a3046c1ab..893e58383ea 100644 --- a/coderd/x/chatd/chatloop/chatloop.go +++ b/coderd/x/chatd/chatloop/chatloop.go @@ -76,13 +76,15 @@ type PersistedStep struct { // Runtime is the wall-clock duration of the model invocation // that produced this step's content, measured from just before // the provider stream is opened until the stream is fully - // consumed. It is persisted as chat_messages.runtime_ms, the - // billable "active generation" time that usage reporting sums. - // Steps without a model invocation (local tool execution - // batches) leave it zero, which persists as NULL: tool wall - // time includes idle waits such as wait_agent polling a - // sub-agent chat that already bills its own model invocations, - // so billing it would double count. + // consumed. Interrupted attempts bill the same window, ending + // where the interrupt closed the stream. It is persisted as + // chat_messages.runtime_ms, the billable "active generation" + // time that usage reporting sums; that column's comment is the + // canonical definition. Steps without a model invocation (local + // tool execution batches) leave it zero, which persists as + // NULL: tool wall time includes idle waits such as wait_agent + // polling a sub-agent chat that already bills its own model + // invocations, so billing it would double count. Runtime time.Duration // PendingDynamicToolCalls lists tool calls that target // dynamic tools. When non-empty the chatloop exits with @@ -225,6 +227,11 @@ type GenerateAssistantOptions struct { ProviderOptions fantasy.ProviderOptions PublishMessagePart func(codersdk.ChatMessageRole, codersdk.ChatMessagePart) + // OnModelStreamStart runs immediately before the provider stream is + // opened, at the instant PersistedStep.Runtime starts measuring. It + // lets callers record the billable window's start out of band, so an + // interrupted attempt bills the same window a completed step reports. + OnModelStreamStart func() Logger slog.Logger Metrics *Metrics } @@ -313,6 +320,11 @@ type GenerateCompactionOptions struct { // Clock measures the summary call duration. Nil uses a real clock. Clock quartz.Clock + + // OnModelStreamStart runs immediately before the summary model call, + // at the instant CompactionResult.Runtime starts measuring. See + // GenerateAssistantOptions.OnModelStreamStart. + OnModelStreamStart func() } // ProviderTool pairs a provider-native tool definition with an @@ -408,6 +420,9 @@ func GenerateAssistant(ctx context.Context, opts GenerateAssistantOptions) (Assi } stepStart := opts.Clock.Now() + if opts.OnModelStreamStart != nil { + opts.OnModelStreamStart() + } stepCtx := chatdebug.ReuseStep(ctx) attempt, streamErr := guardedStream( stepCtx, diff --git a/coderd/x/chatd/chatloop/compaction.go b/coderd/x/chatd/chatloop/compaction.go index 62c2bcfbf3e..8442e0c13f8 100644 --- a/coderd/x/chatd/chatloop/compaction.go +++ b/coderd/x/chatd/chatloop/compaction.go @@ -181,12 +181,15 @@ func GenerateCompaction(ctx context.Context, opts GenerateCompactionOptions) (Co clock = quartz.NewReal() } summaryStart := clock.Now() + if opts.OnModelStreamStart != nil { + opts.OnModelStreamStart() + } summary, err := generateCompactionSummary(ctx, opts.Model, opts.Messages, config) if err != nil { publishCompactionError(config, "failed to generate compaction summary") return CompactionResult{}, err } - runtime := clock.Since(summaryStart) + summaryRuntime := clock.Since(summaryStart) if summary == "" { publishCompactionError(config, "compaction produced an empty summary") return CompactionResult{}, xerrors.New("compaction produced an empty summary") @@ -202,7 +205,7 @@ func GenerateCompaction(ctx context.Context, opts GenerateCompactionOptions) (Co UsagePercent: usagePercent, ContextTokens: contextTokens, ContextLimit: contextLimit, - Runtime: runtime, + Runtime: summaryRuntime, } if config.PublishMessagePart != nil && config.ToolCallID != "" { resultJSON, _ := json.Marshal(map[string]any{ diff --git a/coderd/x/chatd/chatloop/runtime_test.go b/coderd/x/chatd/chatloop/runtime_test.go index c42e86797aa..b227b26b694 100644 --- a/coderd/x/chatd/chatloop/runtime_test.go +++ b/coderd/x/chatd/chatloop/runtime_test.go @@ -49,6 +49,48 @@ func TestGenerateAssistant_RecordsModelInvocationRuntime(t *testing.T) { require.Equal(t, 1500*time.Millisecond, outcome.Step.Runtime) } +// The interrupt path bills the window OnModelStreamStart opens, so that +// hook must fire at the instant PersistedStep.Runtime starts measuring. +func TestGenerateAssistant_ModelStreamStartMatchesRuntimeWindow(t *testing.T) { + t.Parallel() + + clock := quartz.NewMock(t) + model := &chattest.FakeModel{ + ProviderName: "test-provider", + ModelName: "test-model", + StreamFn: func(_ context.Context, _ fantasy.Call) (fantasy.StreamResponse, error) { + return func(yield func(fantasy.StreamPart) bool) { + if !yield(fantasy.StreamPart{Type: fantasy.StreamPartTypeTextStart, ID: "t"}) { + return + } + if !yield(fantasy.StreamPart{Type: fantasy.StreamPartTypeTextDelta, ID: "t", Delta: "hello"}) { + return + } + if !yield(fantasy.StreamPart{Type: fantasy.StreamPartTypeTextEnd, ID: "t"}) { + return + } + clock.Advance(1500 * time.Millisecond) + yield(fantasy.StreamPart{ + Type: fantasy.StreamPartTypeFinish, + FinishReason: fantasy.FinishReasonStop, + }) + }, nil + }, + } + + var startedAt []time.Time + outcome, err := chatloop.GenerateAssistant(context.Background(), chatloop.GenerateAssistantOptions{ + Model: model, + Clock: clock, + OnModelStreamStart: func() { + startedAt = append(startedAt, clock.Now()) + }, + }) + require.NoError(t, err) + require.Len(t, startedAt, 1) + require.Equal(t, outcome.Step.Runtime, clock.Since(startedAt[0])) +} + func TestGenerateAssistant_ErroredStreamReturnsNoStep(t *testing.T) { t.Parallel() @@ -96,6 +138,7 @@ func TestGenerateCompaction_RecordsRuntime(t *testing.T) { }, } + var startedAt []time.Time result, err := chatloop.GenerateCompaction(context.Background(), chatloop.GenerateCompactionOptions{ Model: model, Messages: []fantasy.Message{{ @@ -106,8 +149,13 @@ func TestGenerateCompaction_RecordsRuntime(t *testing.T) { ContextLimit: 100, StepUsage: fantasy.Usage{InputTokens: 90}, Clock: clock, + OnModelStreamStart: func() { + startedAt = append(startedAt, clock.Now()) + }, }) require.NoError(t, err) require.Equal(t, "summary", result.SummaryReport) require.Equal(t, 1500*time.Millisecond, result.Runtime) + require.Len(t, startedAt, 1) + require.Equal(t, result.Runtime, clock.Since(startedAt[0])) } diff --git a/coderd/x/chatd/generation.go b/coderd/x/chatd/generation.go index 6f6d93d76b9..7efacc9112b 100644 --- a/coderd/x/chatd/generation.go +++ b/coderd/x/chatd/generation.go @@ -626,6 +626,7 @@ func (s *taskStarter) generateAssistant( ModelConfig: prepared.ModelConfig, ProviderOptions: prepared.ProviderOptions, PublishMessagePart: attempt.publish, + OnModelStreamStart: attempt.startModelInvocation, Logger: s.opts.Logger, Clock: s.opts.Clock, Metrics: s.server.metrics, @@ -752,6 +753,7 @@ func (s *taskStarter) generateCompaction( ) } compactionOpts.PublishMessagePart = attempt.publish + compactionOpts.OnModelStreamStart = attempt.startModelInvocation compactionOpts.Source = source compactionOpts.Force = source == chatloop.CompactionSourceManual compactionOpts.Clock = s.opts.Clock @@ -823,6 +825,11 @@ type generationAttempt struct { number int64 // publish streams a message part into the attempt's buffer episode. publish func(codersdk.ChatMessageRole, codersdk.ChatMessagePart) + // startModelInvocation marks the start of the attempt's billable + // model invocation window on the buffer episode, so an interrupt + // can bill the window the step would have reported. It is always + // non-nil when beginGenerationAttempt succeeds. + startModelInvocation func() // closeEpisode closes the attempt's buffer episode. It is always // non-nil when beginGenerationAttempt succeeds. closeEpisode func() @@ -866,6 +873,9 @@ func (s *taskStarter) beginGenerationAttempt( publish: func(role codersdk.ChatMessageRole, part codersdk.ChatMessagePart) { _ = s.opts.MessagePartBuffer.AddPart(key, role, part) }, + startModelInvocation: func() { + _ = s.opts.MessagePartBuffer.StartModelInvocation(key) + }, closeEpisode: func() { _ = s.opts.MessagePartBuffer.CloseEpisode(key) }, diff --git a/coderd/x/chatd/message_conversion.go b/coderd/x/chatd/message_conversion.go index 9a55ebc2939..9f96a442443 100644 --- a/coderd/x/chatd/message_conversion.go +++ b/coderd/x/chatd/message_conversion.go @@ -201,9 +201,10 @@ func assistantMessage( } } msg.ContextLimit = step.ContextLimit - if step.Runtime > 0 { - msg.RuntimeMs = sql.NullInt64{Int64: step.Runtime.Milliseconds(), Valid: true} - } + // InsertChatMessages maps a zero runtime to NULL, so a model + // invocation shorter than a millisecond persists the same way an + // unmeasured one does. + msg.RuntimeMs = nullInt64IfNonZero(step.Runtime.Milliseconds()) return msg } @@ -320,9 +321,7 @@ func buildCompactionMessages(input buildCompactionMessagesInput) (compactionMess } assistantMsg := baseMessage(database.ChatMessageRoleAssistant, database.ChatMessageVisibilityUser, input.modelConfigID, contentVersion, assistantContent) - if input.compaction.Runtime > 0 { - assistantMsg.RuntimeMs = sql.NullInt64{Int64: input.compaction.Runtime.Milliseconds(), Valid: true} - } + assistantMsg.RuntimeMs = nullInt64IfNonZero(input.compaction.Runtime.Milliseconds()) messages := []chatstate.Message{ { Role: database.ChatMessageRoleUser, @@ -564,10 +563,11 @@ type bufferedPartsToPartialMessagesInput struct { contentVersion int16 logger slog.Logger interruptedAt time.Time - // attemptRuntime is the interrupted generation attempt's wall-clock - // duration (its message part episode's lifetime), persisted as - // runtime_ms on the first partial assistant message when the attempt - // streamed model-generated assistant content. + // attemptRuntime is the interrupted attempt's billable model + // invocation window: the span from the provider stream opening to + // the interrupt closing its buffer episode. It is persisted as + // runtime_ms on the first partial assistant message when the + // attempt streamed model-generated assistant content. attemptRuntime time.Duration } @@ -626,7 +626,7 @@ func bufferedPartsToPartialMessages(input bufferedPartsToPartialMessagesInput) ( if state.messages[i].Role != database.ChatMessageRoleAssistant { continue } - state.messages[i].RuntimeMs = sql.NullInt64{Int64: input.attemptRuntime.Milliseconds(), Valid: true} + state.messages[i].RuntimeMs = nullInt64IfNonZero(input.attemptRuntime.Milliseconds()) break } } @@ -648,7 +648,9 @@ type partialMessageConversionState struct { // from the model stream itself (text, reasoning, tool calls, // sources). Tool execution also publishes assistant-role file // parts for attachments; those alone must not attract the - // attempt's runtime, because tool batches are not billable. + // attempt's runtime, because tool batches are not billable. The + // buffer episode only carries a runtime when a provider stream + // was opened, so this is a second gate rather than the only one. modelStreamedAssistant bool } diff --git a/coderd/x/chatd/message_conversion_test.go b/coderd/x/chatd/message_conversion_test.go index 59a9eab3fdf..c153e37aa1c 100644 --- a/coderd/x/chatd/message_conversion_test.go +++ b/coderd/x/chatd/message_conversion_test.go @@ -101,6 +101,34 @@ func TestBuildCommitStepMessages_LocalToolResultsBecomeToolMessages(t *testing.T require.JSONEq(t, `{"stdout":"/tmp"}`, string(toolParts[0].Result)) } +// A step with no model invocation (a local tool execution batch) must +// persist runtime_ms NULL: its wall time is not billable. +func TestBuildCommitStepMessages_ZeroRuntimeLeavesRuntimeNull(t *testing.T) { + t.Parallel() + + got, err := buildCommitStepMessages(buildCommitStepMessagesInput{ + modelConfigID: uuid.New(), + contentVersion: chatprompt.CurrentContentVersion, + logger: slog.Make(), + step: stepData{ + Content: []fantasy.Content{ + fantasy.ToolCallContent{ToolCallID: "call-1", ToolName: "execute", Input: `{"cmd":"pwd"}`}, + fantasy.ToolResultContent{ + ToolCallID: "call-1", + ToolName: "execute", + Result: fantasy.ToolResultOutputContentText{Text: `{"stdout":"/tmp"}`}, + }, + }, + Runtime: 0, + }, + }) + require.NoError(t, err) + require.Len(t, got.Messages, 2) + require.Equal(t, database.ChatMessageRoleAssistant, got.Messages[0].Role) + require.False(t, got.Messages[0].RuntimeMs.Valid) + require.False(t, got.Messages[1].RuntimeMs.Valid) +} + func TestBuildCommitStepMessages_ProviderExecutedResultsStayAssistantContent(t *testing.T) { t.Parallel() @@ -250,6 +278,34 @@ func TestBuildCompactionMessages_CompressedSummaryToolCallAndResult(t *testing.T require.JSONEq(t, `{"summary":"user report","source":"automatic","threshold_percent":70,"usage_percent":81.5,"context_tokens":815,"context_limit_tokens":1000}`, string(resultPart.Result)) } +// A compaction that never reached the summary model call carries no +// runtime, so its assistant row must persist runtime_ms NULL. +func TestBuildCompactionMessages_ZeroRuntimeLeavesRuntimeNull(t *testing.T) { + t.Parallel() + + got, err := buildCompactionMessages(buildCompactionMessagesInput{ + modelConfigID: uuid.New(), + contentVersion: chatprompt.CurrentContentVersion, + toolCallID: "summary-1", + toolName: "chat_summarized", + compaction: compactionOutcome{ + SystemSummary: "system summary", + SummaryReport: "user report", + ThresholdPercent: 70, + UsagePercent: 81.5, + ContextTokens: 815, + ContextLimit: 1000, + Runtime: 0, + }, + }) + require.NoError(t, err) + require.Len(t, got.Messages, 3) + require.Equal(t, database.ChatMessageRoleAssistant, got.Messages[1].Role) + for i := range got.Messages { + require.False(t, got.Messages[i].RuntimeMs.Valid) + } +} + func TestCurrentTurnStepCount_ExcludesCompressedCompactionMessages(t *testing.T) { t.Parallel() diff --git a/coderd/x/chatd/messagepartbuffer/message_part_buffer.go b/coderd/x/chatd/messagepartbuffer/message_part_buffer.go index a21c3566a2d..12e9ed3d522 100644 --- a/coderd/x/chatd/messagepartbuffer/message_part_buffer.go +++ b/coderd/x/chatd/messagepartbuffer/message_part_buffer.go @@ -97,10 +97,11 @@ type Buffer struct { type episodeState struct { created bool - // createdAt is set only by CreateEpisode, not by the implicit - // creation in CloseEpisode or by subscriber placeholders, so it - // marks when the generation attempt actually started. - createdAt time.Time + // modelStartedAt is stamped by StartModelInvocation when the + // episode's provider stream is opened. It is zero for episodes + // that never invoke a model, such as local tool execution + // batches. + modelStartedAt time.Time closed bool closedAt time.Time closedHeapItem *closedEpisodeItem @@ -189,14 +190,37 @@ func (b *Buffer) CreateEpisode(key Key) error { if b.closed { return ErrMessagePartBufferClosed } - now := b.opts.Clock.Now("message-part-buffer", "create") - b.gcClosedEpisodesLocked(now) + b.gcClosedEpisodesLocked(b.opts.Clock.Now("message-part-buffer", "create")) episode := b.getOrCreateEpisodeLocked(key) if episode.created { return ErrEpisodeExists } episode.markCreated() - episode.createdAt = now + return nil +} + +// StartModelInvocation stamps the instant the episode opens its provider +// stream, which starts the episode's billable model invocation window. +// +// Callers stamp it at the same instant a completed step starts measuring +// its runtime, so an interrupted attempt bills the same window the step +// would have reported had it finished. A repeat call re-stamps the start: +// only the most recent invocation is billed, which errs toward +// undercounting. +func (b *Buffer) StartModelInvocation(key Key) error { + b.mu.Lock() + defer b.mu.Unlock() + if b.closed { + return ErrMessagePartBufferClosed + } + episode, err := b.getEpisodeLocked(key) + if err != nil { + return err + } + if episode.closed { + return ErrEpisodeClosed + } + episode.modelStartedAt = b.opts.Clock.Now("message-part-buffer", "model-invocation-start") return nil } @@ -272,17 +296,17 @@ func (b *Buffer) GetParts(key Key) ([]Part, error) { return slices.Clone(episode.parts), nil } -// EpisodeDuration returns the wall-clock span between CreateEpisode and -// CloseEpisode. It returns 0 when the episode is unknown, was created -// implicitly by CloseEpisode, or is not closed yet. -func (b *Buffer) EpisodeDuration(key Key) time.Duration { +// ModelInvocationDuration returns the wall-clock span between +// StartModelInvocation and CloseEpisode. It returns 0 when the episode is +// unknown, never opened a provider stream, or is not closed yet. +func (b *Buffer) ModelInvocationDuration(key Key) time.Duration { b.mu.Lock() defer b.mu.Unlock() episode := b.episodes[key] - if episode == nil || episode.createdAt.IsZero() || !episode.closed { + if episode == nil || episode.modelStartedAt.IsZero() || !episode.closed { return 0 } - return episode.closedAt.Sub(episode.createdAt) + return episode.closedAt.Sub(episode.modelStartedAt) } // SubscribeToEpisode replays existing parts and streams new parts. diff --git a/coderd/x/chatd/messagepartbuffer/message_part_buffer_test.go b/coderd/x/chatd/messagepartbuffer/message_part_buffer_test.go index beadc32bb33..00b30ecc5a7 100644 --- a/coderd/x/chatd/messagepartbuffer/message_part_buffer_test.go +++ b/coderd/x/chatd/messagepartbuffer/message_part_buffer_test.go @@ -99,7 +99,7 @@ func TestBuffer_CloseEpisodeIdempotent(t *testing.T) { require.NoError(t, buffer.CloseEpisode(key)) } -func TestBuffer_EpisodeDuration(t *testing.T) { +func TestBuffer_ModelInvocationDuration(t *testing.T) { t.Parallel() clock := quartz.NewMock(t) @@ -107,25 +107,40 @@ func TestBuffer_EpisodeDuration(t *testing.T) { defer buffer.Close() key := testEpisodeKey() - require.Zero(t, buffer.EpisodeDuration(key), "unknown episode has no duration") + require.Zero(t, buffer.ModelInvocationDuration(key), "unknown episode has no duration") + require.ErrorIs(t, buffer.StartModelInvocation(key), messagepartbuffer.ErrEpisodeNotFound) require.NoError(t, buffer.CreateEpisode(key)) - require.Zero(t, buffer.EpisodeDuration(key), "open episode has no duration") + // Attempt setup happens before the provider stream opens and is + // not billable. + clock.Advance(time.Second) + require.NoError(t, buffer.StartModelInvocation(key)) + require.Zero(t, buffer.ModelInvocationDuration(key), "open episode has no duration") clock.Advance(1500 * time.Millisecond) require.NoError(t, buffer.CloseEpisode(key)) - require.Equal(t, 1500*time.Millisecond, buffer.EpisodeDuration(key)) + require.Equal(t, 1500*time.Millisecond, buffer.ModelInvocationDuration(key)) - // A second close must not move the recorded span. + // A second close must not move the recorded span, and a closed + // episode no longer accepts an invocation start. clock.Advance(time.Second) require.NoError(t, buffer.CloseEpisode(key)) - require.Equal(t, 1500*time.Millisecond, buffer.EpisodeDuration(key)) + require.ErrorIs(t, buffer.StartModelInvocation(key), messagepartbuffer.ErrEpisodeClosed) + require.Equal(t, 1500*time.Millisecond, buffer.ModelInvocationDuration(key)) + + // Episodes that never open a provider stream, such as local tool + // execution batches, report no duration. + toolBatch := testEpisodeKey() + require.NoError(t, buffer.CreateEpisode(toolBatch)) + clock.Advance(time.Second) + require.NoError(t, buffer.CloseEpisode(toolBatch)) + require.Zero(t, buffer.ModelInvocationDuration(toolBatch)) // Episodes created implicitly by CloseEpisode never started a // generation attempt, so they report no duration. implicit := testEpisodeKey() require.NoError(t, buffer.CloseEpisode(implicit)) - require.Zero(t, buffer.EpisodeDuration(implicit)) + require.Zero(t, buffer.ModelInvocationDuration(implicit)) } func TestBuffer_SubscribeExistingReplaysThenStreamsLiveParts(t *testing.T) { diff --git a/coderd/x/chatd/tasks.go b/coderd/x/chatd/tasks.go index f2e8185f392..6fcb384f4f1 100644 --- a/coderd/x/chatd/tasks.go +++ b/coderd/x/chatd/tasks.go @@ -283,7 +283,7 @@ func (s *taskStarter) StartInterrupt(ctx context.Context, input chatWorkerTaskSt contentVersion: chatprompt.CurrentContentVersion, logger: s.opts.Logger, interruptedAt: s.opts.Clock.Now("chatworker", "interrupt"), - attemptRuntime: s.opts.MessagePartBuffer.EpisodeDuration(key), + attemptRuntime: s.opts.MessagePartBuffer.ModelInvocationDuration(key), }) if err != nil { return xerrors.Errorf("convert buffered parts: %w", err) diff --git a/coderd/x/chatd/tasks_test.go b/coderd/x/chatd/tasks_test.go index c5cacb2f3ee..feedcbb382d 100644 --- a/coderd/x/chatd/tasks_test.go +++ b/coderd/x/chatd/tasks_test.go @@ -409,6 +409,10 @@ func TestInterruptTask_PartialAssistantKeepsAttemptRuntime(t *testing.T) { GenerationAttempt: acquired.GenerationAttempt, } require.NoError(t, buffer.CreateEpisode(key)) + // Prompt preparation and attempt bookkeeping run before the + // provider stream opens and must not be billed. + clock.Advance(3 * time.Second) + require.NoError(t, buffer.StartModelInvocation(key)) require.NoError(t, buffer.AddPart(key, codersdk.ChatMessageRoleAssistant, codersdk.ChatMessageText("partial answer"))) clock.Advance(1500 * time.Millisecond) interrupting := f.interruptChat(t, chat.ID) @@ -431,6 +435,48 @@ func TestInterruptTask_PartialAssistantKeepsAttemptRuntime(t *testing.T) { require.Equal(t, sql.NullInt64{Int64: 1500, Valid: true}, assistant.RuntimeMs) } +func TestInterruptTask_PartialAssistantWithoutModelInvocationHasNoRuntime(t *testing.T) { + t.Parallel() + + f := newTaskTestFixture(t) + chat := f.createRunningChat(t) + workerID := uuid.New() + runnerID := uuid.New() + acquired := f.acquireChat(t, chat.ID, workerID, runnerID) + recorder := newTaskSideEffectRecorder() + clock := quartz.NewMock(t) + starter := newTestTaskStarterWithClock(t, f, recorder, clock) + buffer := starter.opts.MessagePartBuffer + key := messagepartbuffer.Key{ + ChatID: chat.ID, + HistoryVersion: acquired.HistoryVersion, + GenerationAttempt: acquired.GenerationAttempt, + } + // A local tool execution batch never opens a provider stream, so + // its wall time is not billable even though it publishes parts. + require.NoError(t, buffer.CreateEpisode(key)) + require.NoError(t, buffer.AddPart(key, codersdk.ChatMessageRoleAssistant, codersdk.ChatMessageText("partial answer"))) + clock.Advance(1500 * time.Millisecond) + interrupting := f.interruptChat(t, chat.ID) + + err := starter.StartInterrupt(testutil.Context(t, testutil.WaitLong), chatWorkerTaskStartInput{ + ChatID: chat.ID, + WorkerID: workerID, + RunnerID: runnerID, + HistoryVersion: interrupting.HistoryVersion, + GenerationAttempt: interrupting.GenerationAttempt, + Status: database.ChatStatusInterrupting, + }) + require.NoError(t, err) + + messages, err := f.db.GetChatMessagesByChatID(testutil.Context(t, testutil.WaitShort), database.GetChatMessagesByChatIDParams{ChatID: chat.ID}) + require.NoError(t, err) + require.GreaterOrEqual(t, len(messages), 3) + assistant := messages[len(messages)-2] + require.Equal(t, database.ChatMessageRoleAssistant, assistant.Role) + require.False(t, assistant.RuntimeMs.Valid) +} + func TestRequiresActionTimeout_ExpiredCancelsOnly(t *testing.T) { t.Parallel() From 4b6c128a5c0c4c7909a06eda9656b7e73c71e99e Mon Sep 17 00:00:00 2001 From: Jaayden Halko Date: Mon, 27 Jul 2026 13:07:19 +0000 Subject: [PATCH 05/11] fix(coderd/database/migrations): renumber runtime_ms comment migration main added 000551_chat_summary, so the merge carried two 000551 migrations and golang-migrate refused to init ("duplicate migration file"), failing gen, lint, sqlc-vet, and the e2e job. --- ....down.sql => 000552_chat_messages_runtime_ms_comment.down.sql} | 0 ...ment.up.sql => 000552_chat_messages_runtime_ms_comment.up.sql} | 0 2 files changed, 0 insertions(+), 0 deletions(-) rename coderd/database/migrations/{000551_chat_messages_runtime_ms_comment.down.sql => 000552_chat_messages_runtime_ms_comment.down.sql} (100%) rename coderd/database/migrations/{000551_chat_messages_runtime_ms_comment.up.sql => 000552_chat_messages_runtime_ms_comment.up.sql} (100%) diff --git a/coderd/database/migrations/000551_chat_messages_runtime_ms_comment.down.sql b/coderd/database/migrations/000552_chat_messages_runtime_ms_comment.down.sql similarity index 100% rename from coderd/database/migrations/000551_chat_messages_runtime_ms_comment.down.sql rename to coderd/database/migrations/000552_chat_messages_runtime_ms_comment.down.sql diff --git a/coderd/database/migrations/000551_chat_messages_runtime_ms_comment.up.sql b/coderd/database/migrations/000552_chat_messages_runtime_ms_comment.up.sql similarity index 100% rename from coderd/database/migrations/000551_chat_messages_runtime_ms_comment.up.sql rename to coderd/database/migrations/000552_chat_messages_runtime_ms_comment.up.sql From 2ec3cb8bbfa1728e8c2d18df1039922fb7bbe612 Mon Sep 17 00:00:00 2001 From: Jaayden Halko Date: Thu, 6 Aug 2026 13:49:50 +0700 Subject: [PATCH 06/11] Update coderd/x/chatd/ARCHITECTURE.md Co-authored-by: Hugo Dutka --- coderd/x/chatd/ARCHITECTURE.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/coderd/x/chatd/ARCHITECTURE.md b/coderd/x/chatd/ARCHITECTURE.md index ab7de3b6c81..fe1af2307fd 100644 --- a/coderd/x/chatd/ARCHITECTURE.md +++ b/coderd/x/chatd/ARCHITECTURE.md @@ -926,7 +926,7 @@ The goroutine does the following in order: 1. It fetches the generation attempt number from the database. 2. It closes the episode corresponding to its history version and generation attempt by calling the `CloseEpisode` method on the [Message part buffer](#message-part-buffer). 3. It reads the buffered parts for that episode by calling the `GetParts` method on the message part buffer. -4. It applies the `FinishInterruption(partial?)` transition on the core state machine. If there are no buffered parts for that episode, or the episode is not found, it passes `nil` as the `partial` argument. When the partial suffix contains an assistant message, the episode's model invocation window (`ModelInvocationDuration`: provider stream open until the interrupt closed the episode) is persisted on that message as `runtime_ms`, so the interrupted attempt's billable generation time is not lost. +4. It applies the `FinishInterruption(partial?)` transition on the core state machine. If there are no buffered parts for that episode, or the episode is not found, it passes `nil` as the `partial` argument. #### Dynamic tools timeout goroutine From 7d1a13feb80702762e15bc1762efb725596e2ce1 Mon Sep 17 00:00:00 2001 From: Jaayden Halko Date: Thu, 6 Aug 2026 13:50:43 +0700 Subject: [PATCH 07/11] Update coderd/x/chatd/chatloop/chatloop.go Co-authored-by: Hugo Dutka --- coderd/x/chatd/chatloop/chatloop.go | 10 +--------- 1 file changed, 1 insertion(+), 9 deletions(-) diff --git a/coderd/x/chatd/chatloop/chatloop.go b/coderd/x/chatd/chatloop/chatloop.go index aed96b0dfd8..87944b96ea2 100644 --- a/coderd/x/chatd/chatloop/chatloop.go +++ b/coderd/x/chatd/chatloop/chatloop.go @@ -76,15 +76,7 @@ type PersistedStep struct { // Runtime is the wall-clock duration of the model invocation // that produced this step's content, measured from just before // the provider stream is opened until the stream is fully - // consumed. Interrupted attempts bill the same window, ending - // where the interrupt closed the stream. It is persisted as - // chat_messages.runtime_ms, the billable "active generation" - // time that usage reporting sums; that column's comment is the - // canonical definition. Steps without a model invocation (local - // tool execution batches) leave it zero, which persists as - // NULL: tool wall time includes idle waits such as wait_agent - // polling a sub-agent chat that already bills its own model - // invocations, so billing it would double count. + // consumed. Runtime time.Duration // PendingDynamicToolCalls lists tool calls that target // dynamic tools. When non-empty the chatloop exits with From e7abe6743992c80639a0cb05d2c9f61f892c8c6d Mon Sep 17 00:00:00 2001 From: Jaayden Halko Date: Thu, 6 Aug 2026 13:51:02 +0700 Subject: [PATCH 08/11] Update coderd/x/chatd/chatloop/chatloop.go Co-authored-by: Hugo Dutka --- coderd/x/chatd/chatloop/chatloop.go | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/coderd/x/chatd/chatloop/chatloop.go b/coderd/x/chatd/chatloop/chatloop.go index 87944b96ea2..06a03004234 100644 --- a/coderd/x/chatd/chatloop/chatloop.go +++ b/coderd/x/chatd/chatloop/chatloop.go @@ -315,8 +315,7 @@ type GenerateCompactionOptions struct { Clock quartz.Clock // OnModelStreamStart runs immediately before the summary model call, - // at the instant CompactionResult.Runtime starts measuring. See - // GenerateAssistantOptions.OnModelStreamStart. + // at the instant CompactionResult.Runtime starts measuring. OnModelStreamStart func() } From 8a043cbe14cf14b4b03d365ab27a60042063bc8d Mon Sep 17 00:00:00 2001 From: Jaayden Halko Date: Thu, 6 Aug 2026 13:51:14 +0700 Subject: [PATCH 09/11] Update coderd/x/chatd/messagepartbuffer/message_part_buffer.go Co-authored-by: Hugo Dutka --- coderd/x/chatd/messagepartbuffer/message_part_buffer.go | 6 ------ 1 file changed, 6 deletions(-) diff --git a/coderd/x/chatd/messagepartbuffer/message_part_buffer.go b/coderd/x/chatd/messagepartbuffer/message_part_buffer.go index 12e9ed3d522..8ba621bb681 100644 --- a/coderd/x/chatd/messagepartbuffer/message_part_buffer.go +++ b/coderd/x/chatd/messagepartbuffer/message_part_buffer.go @@ -201,12 +201,6 @@ func (b *Buffer) CreateEpisode(key Key) error { // StartModelInvocation stamps the instant the episode opens its provider // stream, which starts the episode's billable model invocation window. -// -// Callers stamp it at the same instant a completed step starts measuring -// its runtime, so an interrupted attempt bills the same window the step -// would have reported had it finished. A repeat call re-stamps the start: -// only the most recent invocation is billed, which errs toward -// undercounting. func (b *Buffer) StartModelInvocation(key Key) error { b.mu.Lock() defer b.mu.Unlock() From 1b85a756cc7e528db5cfcc661ca76c95686bae83 Mon Sep 17 00:00:00 2001 From: Jaayden Halko Date: Thu, 6 Aug 2026 07:41:33 +0000 Subject: [PATCH 10/11] updates for PR review --- coderd/database/dump.sql | 2 -- ..._chat_messages_runtime_ms_comment.down.sql | 1 - ...62_chat_messages_runtime_ms_comment.up.sql | 1 - coderd/database/models.go | 9 +++-- coderd/x/chatd/chatloop/chatloop.go | 2 +- coderd/x/chatd/chatloop/compaction.go | 12 +++---- .../chatloop/compaction_internal_test.go | 15 +++++++++ .../messagepartbuffer/message_part_buffer.go | 14 ++++---- .../message_part_buffer_test.go | 33 +++++++++++-------- coderd/x/chatd/tasks.go | 10 ++++-- 10 files changed, 59 insertions(+), 40 deletions(-) delete mode 100644 coderd/database/migrations/000562_chat_messages_runtime_ms_comment.down.sql delete mode 100644 coderd/database/migrations/000562_chat_messages_runtime_ms_comment.up.sql diff --git a/coderd/database/dump.sql b/coderd/database/dump.sql index bc9e60f5626..b63f1193958 100644 --- a/coderd/database/dump.sql +++ b/coderd/database/dump.sql @@ -1977,8 +1977,6 @@ CREATE TABLE chat_messages ( search_tsv tsvector ); -COMMENT ON COLUMN chat_messages.runtime_ms IS 'Wall-clock milliseconds of the model invocation that produced the message content: assistant steps and compaction summaries, including interrupted partials, which measure from the provider stream opening until the interrupt closed it. NULL when no model invocation produced the row (user messages, tool results; local tool execution is not counted) or when the invocation rounded down to zero milliseconds. Billing source of truth for Coder Agents runtime: usage reporting sums it over created_at ranges.'; - COMMENT ON COLUMN chat_messages.reasoning_effort IS 'Stores the selected effort for the turn triggered by this message.'; COMMENT ON COLUMN chat_messages.search_tsv IS 'Used for full text search. NULL initially, populated async via background job.'; diff --git a/coderd/database/migrations/000562_chat_messages_runtime_ms_comment.down.sql b/coderd/database/migrations/000562_chat_messages_runtime_ms_comment.down.sql deleted file mode 100644 index 841ebe58b43..00000000000 --- a/coderd/database/migrations/000562_chat_messages_runtime_ms_comment.down.sql +++ /dev/null @@ -1 +0,0 @@ -COMMENT ON COLUMN chat_messages.runtime_ms IS NULL; diff --git a/coderd/database/migrations/000562_chat_messages_runtime_ms_comment.up.sql b/coderd/database/migrations/000562_chat_messages_runtime_ms_comment.up.sql deleted file mode 100644 index 85bc1f27e31..00000000000 --- a/coderd/database/migrations/000562_chat_messages_runtime_ms_comment.up.sql +++ /dev/null @@ -1 +0,0 @@ -COMMENT ON COLUMN chat_messages.runtime_ms IS 'Wall-clock milliseconds of the model invocation that produced the message content: assistant steps and compaction summaries, including interrupted partials, which measure from the provider stream opening until the interrupt closed it. NULL when no model invocation produced the row (user messages, tool results; local tool execution is not counted) or when the invocation rounded down to zero milliseconds. Billing source of truth for Coder Agents runtime: usage reporting sums it over created_at ranges.'; diff --git a/coderd/database/models.go b/coderd/database/models.go index eaeef9e5117..4a979e94c8b 100644 --- a/coderd/database/models.go +++ b/coderd/database/models.go @@ -5125,11 +5125,10 @@ type ChatMessage struct { CreatedBy uuid.NullUUID `db:"created_by" json:"created_by"` ContentVersion int16 `db:"content_version" json:"content_version"` TotalCostMicros sql.NullInt64 `db:"total_cost_micros" json:"total_cost_micros"` - // Wall-clock milliseconds of the model invocation that produced the message content: assistant steps and compaction summaries, including interrupted partials, which measure from the provider stream opening until the interrupt closed it. NULL when no model invocation produced the row (user messages, tool results; local tool execution is not counted) or when the invocation rounded down to zero milliseconds. Billing source of truth for Coder Agents runtime: usage reporting sums it over created_at ranges. - RuntimeMs sql.NullInt64 `db:"runtime_ms" json:"runtime_ms"` - Deleted bool `db:"deleted" json:"deleted"` - ProviderResponseID sql.NullString `db:"provider_response_id" json:"provider_response_id"` - Revision int64 `db:"revision" json:"revision"` + RuntimeMs sql.NullInt64 `db:"runtime_ms" json:"runtime_ms"` + Deleted bool `db:"deleted" json:"deleted"` + ProviderResponseID sql.NullString `db:"provider_response_id" json:"provider_response_id"` + Revision int64 `db:"revision" json:"revision"` // Stores the selected effort for the turn triggered by this message. ReasoningEffort NullChatReasoningEffort `db:"reasoning_effort" json:"reasoning_effort"` // Used for full text search. NULL initially, populated async via background job. diff --git a/coderd/x/chatd/chatloop/chatloop.go b/coderd/x/chatd/chatloop/chatloop.go index 06a03004234..89063ff0612 100644 --- a/coderd/x/chatd/chatloop/chatloop.go +++ b/coderd/x/chatd/chatloop/chatloop.go @@ -311,7 +311,7 @@ type GenerateCompactionOptions struct { PublishMessagePart func(codersdk.ChatMessageRole, codersdk.ChatMessagePart) - // Clock measures the summary call duration. Nil uses a real clock. + // Clock measures the summary call duration. Required. Clock quartz.Clock // OnModelStreamStart runs immediately before the summary model call, diff --git a/coderd/x/chatd/chatloop/compaction.go b/coderd/x/chatd/chatloop/compaction.go index 7df5ce89382..320edca3d92 100644 --- a/coderd/x/chatd/chatloop/compaction.go +++ b/coderd/x/chatd/chatloop/compaction.go @@ -12,7 +12,6 @@ import ( "github.com/coder/coder/v2/coderd/x/chatd/chatdebug" "github.com/coder/coder/v2/codersdk" - "github.com/coder/quartz" ) const ( @@ -146,6 +145,9 @@ func GenerateCompaction(ctx context.Context, opts GenerateCompactionOptions) (Co if opts.Model == nil { return CompactionResult{}, xerrors.New("chat model is required") } + if opts.Clock == nil { + return CompactionResult{}, xerrors.New("clock is required") + } config, ok := normalizedCompactionGenerateConfig(opts) if !ok { return CompactionResult{}, nil @@ -177,11 +179,7 @@ func GenerateCompaction(ctx context.Context, opts GenerateCompactionOptions) (Co ) } - clock := opts.Clock - if clock == nil { - clock = quartz.NewReal() - } - summaryStart := clock.Now() + summaryStart := opts.Clock.Now() if opts.OnModelStreamStart != nil { opts.OnModelStreamStart() } @@ -190,7 +188,7 @@ func GenerateCompaction(ctx context.Context, opts GenerateCompactionOptions) (Co publishCompactionError(config, "failed to generate compaction summary") return CompactionResult{}, err } - summaryRuntime := clock.Since(summaryStart) + summaryRuntime := opts.Clock.Since(summaryStart) if summary == "" { publishCompactionError(config, "compaction produced an empty summary") return CompactionResult{}, xerrors.New("compaction produced an empty summary") diff --git a/coderd/x/chatd/chatloop/compaction_internal_test.go b/coderd/x/chatd/chatloop/compaction_internal_test.go index 422aaa5d98c..47653113955 100644 --- a/coderd/x/chatd/chatloop/compaction_internal_test.go +++ b/coderd/x/chatd/chatloop/compaction_internal_test.go @@ -18,6 +18,7 @@ import ( "github.com/coder/coder/v2/coderd/x/chatd/chatdebug" "github.com/coder/coder/v2/coderd/x/chatd/chattest" "github.com/coder/coder/v2/testutil" + "github.com/coder/quartz" ) func TestStartCompactionDebugRun_DoesNotReportDebugErrors(t *testing.T) { @@ -326,6 +327,7 @@ func TestGenerateCompaction_ForceBypassesThresholdGates(t *testing.T) { opts := tc.opts opts.Model = newModel(&calls) opts.Messages = messages + opts.Clock = quartz.NewMock(t) result, err := GenerateCompaction(context.Background(), opts) require.NoError(t, err) require.Empty(t, result.SummaryReport) @@ -365,8 +367,21 @@ func TestGenerateCompaction_DefaultSourceAutomatic(t *testing.T) { ThresholdPercent: 70, ContextLimit: 100, StepUsage: fantasy.Usage{InputTokens: 90}, + Clock: quartz.NewMock(t), }) require.NoError(t, err) require.Equal(t, "auto summary", result.SummaryReport) require.Equal(t, CompactionSourceAutomatic, result.Source) } + +// TestGenerateCompaction_RequiresClock verifies a nil clock is +// rejected instead of silently falling back to a real clock; tests +// must supply their own. +func TestGenerateCompaction_RequiresClock(t *testing.T) { + t.Parallel() + + _, err := GenerateCompaction(context.Background(), GenerateCompactionOptions{ + Model: &chattest.FakeModel{ProviderName: "fake", ModelName: "fake-model"}, + }) + require.ErrorContains(t, err, "clock is required") +} diff --git a/coderd/x/chatd/messagepartbuffer/message_part_buffer.go b/coderd/x/chatd/messagepartbuffer/message_part_buffer.go index 8ba621bb681..82f41648be6 100644 --- a/coderd/x/chatd/messagepartbuffer/message_part_buffer.go +++ b/coderd/x/chatd/messagepartbuffer/message_part_buffer.go @@ -290,17 +290,17 @@ func (b *Buffer) GetParts(key Key) ([]Part, error) { return slices.Clone(episode.parts), nil } -// ModelInvocationDuration returns the wall-clock span between -// StartModelInvocation and CloseEpisode. It returns 0 when the episode is -// unknown, never opened a provider stream, or is not closed yet. -func (b *Buffer) ModelInvocationDuration(key Key) time.Duration { +// ModelInvokedAt returns the instant stamped by StartModelInvocation, or the +// zero time if there is none. Read it before CloseEpisode: closed episodes +// are garbage collected, so reading afterwards races the cleanup loop. +func (b *Buffer) ModelInvokedAt(key Key) time.Time { b.mu.Lock() defer b.mu.Unlock() episode := b.episodes[key] - if episode == nil || episode.modelStartedAt.IsZero() || !episode.closed { - return 0 + if episode == nil { + return time.Time{} } - return episode.closedAt.Sub(episode.modelStartedAt) + return episode.modelStartedAt } // SubscribeToEpisode replays existing parts and streams new parts. diff --git a/coderd/x/chatd/messagepartbuffer/message_part_buffer_test.go b/coderd/x/chatd/messagepartbuffer/message_part_buffer_test.go index 00b30ecc5a7..ea489af2448 100644 --- a/coderd/x/chatd/messagepartbuffer/message_part_buffer_test.go +++ b/coderd/x/chatd/messagepartbuffer/message_part_buffer_test.go @@ -99,7 +99,7 @@ func TestBuffer_CloseEpisodeIdempotent(t *testing.T) { require.NoError(t, buffer.CloseEpisode(key)) } -func TestBuffer_ModelInvocationDuration(t *testing.T) { +func TestBuffer_ModelInvokedAt(t *testing.T) { t.Parallel() clock := quartz.NewMock(t) @@ -107,40 +107,45 @@ func TestBuffer_ModelInvocationDuration(t *testing.T) { defer buffer.Close() key := testEpisodeKey() - require.Zero(t, buffer.ModelInvocationDuration(key), "unknown episode has no duration") + require.Zero(t, buffer.ModelInvokedAt(key), "unknown episode has no invocation stamp") require.ErrorIs(t, buffer.StartModelInvocation(key), messagepartbuffer.ErrEpisodeNotFound) require.NoError(t, buffer.CreateEpisode(key)) + require.Zero(t, buffer.ModelInvokedAt(key), "episode without a provider stream has no invocation stamp") // Attempt setup happens before the provider stream opens and is // not billable. clock.Advance(time.Second) require.NoError(t, buffer.StartModelInvocation(key)) - require.Zero(t, buffer.ModelInvocationDuration(key), "open episode has no duration") + invokedAt := buffer.ModelInvokedAt(key) + require.Equal(t, clock.Now(), invokedAt) - clock.Advance(1500 * time.Millisecond) - require.NoError(t, buffer.CloseEpisode(key)) - require.Equal(t, 1500*time.Millisecond, buffer.ModelInvocationDuration(key)) - - // A second close must not move the recorded span, and a closed - // episode no longer accepts an invocation start. + // A repeat call re-stamps the start so only the most recent + // invocation is billed. clock.Advance(time.Second) + require.NoError(t, buffer.StartModelInvocation(key)) + require.Equal(t, clock.Now(), buffer.ModelInvokedAt(key)) + + // Closing must not move the recorded stamp, and a closed episode + // no longer accepts an invocation start. + invokedAt = buffer.ModelInvokedAt(key) + clock.Advance(1500 * time.Millisecond) require.NoError(t, buffer.CloseEpisode(key)) require.ErrorIs(t, buffer.StartModelInvocation(key), messagepartbuffer.ErrEpisodeClosed) - require.Equal(t, 1500*time.Millisecond, buffer.ModelInvocationDuration(key)) + require.Equal(t, invokedAt, buffer.ModelInvokedAt(key)) // Episodes that never open a provider stream, such as local tool - // execution batches, report no duration. + // execution batches, report no invocation stamp. toolBatch := testEpisodeKey() require.NoError(t, buffer.CreateEpisode(toolBatch)) clock.Advance(time.Second) require.NoError(t, buffer.CloseEpisode(toolBatch)) - require.Zero(t, buffer.ModelInvocationDuration(toolBatch)) + require.Zero(t, buffer.ModelInvokedAt(toolBatch)) // Episodes created implicitly by CloseEpisode never started a - // generation attempt, so they report no duration. + // generation attempt, so they report no invocation stamp. implicit := testEpisodeKey() require.NoError(t, buffer.CloseEpisode(implicit)) - require.Zero(t, buffer.ModelInvocationDuration(implicit)) + require.Zero(t, buffer.ModelInvokedAt(implicit)) } func TestBuffer_SubscribeExistingReplaysThenStreamsLiveParts(t *testing.T) { diff --git a/coderd/x/chatd/tasks.go b/coderd/x/chatd/tasks.go index 6fcb384f4f1..a3738b886db 100644 --- a/coderd/x/chatd/tasks.go +++ b/coderd/x/chatd/tasks.go @@ -260,6 +260,7 @@ func (s *taskStarter) StartInterrupt(ctx context.Context, input chatWorkerTaskSt HistoryVersion: input.HistoryVersion, GenerationAttempt: chat.GenerationAttempt, } + modelInvokedAt := s.opts.MessagePartBuffer.ModelInvokedAt(key) if err := s.opts.MessagePartBuffer.CloseEpisode(key); err != nil { if ctx.Err() != nil { return errors.Join(errTaskExpectedExit, xerrors.Errorf("close message part episode: %w", err), ctx.Err()) @@ -277,13 +278,18 @@ func (s *taskStarter) StartInterrupt(ctx context.Context, input chatWorkerTaskSt } return taskRetryableError{err: xerrors.Errorf("get message part episode: %w", err)} } + interruptedAt := s.opts.Clock.Now("chatworker", "interrupt") + var attemptRuntime time.Duration + if !modelInvokedAt.IsZero() { + attemptRuntime = interruptedAt.Sub(modelInvokedAt) + } partialMessages, err := bufferedPartsToPartialMessages(bufferedPartsToPartialMessagesInput{ parts: parts, modelConfigID: chat.LastModelConfigID, contentVersion: chatprompt.CurrentContentVersion, logger: s.opts.Logger, - interruptedAt: s.opts.Clock.Now("chatworker", "interrupt"), - attemptRuntime: s.opts.MessagePartBuffer.ModelInvocationDuration(key), + interruptedAt: interruptedAt, + attemptRuntime: attemptRuntime, }) if err != nil { return xerrors.Errorf("convert buffered parts: %w", err) From 9d1374d8259518a7ad252a5b7586f85e3f09e2f4 Mon Sep 17 00:00:00 2001 From: Jaayden Halko Date: Thu, 6 Aug 2026 08:53:31 +0000 Subject: [PATCH 11/11] chore: update for buffer interface changes --- coderd/x/chatd/ARCHITECTURE.md | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/coderd/x/chatd/ARCHITECTURE.md b/coderd/x/chatd/ARCHITECTURE.md index fe1af2307fd..de78ad72153 100644 --- a/coderd/x/chatd/ARCHITECTURE.md +++ b/coderd/x/chatd/ARCHITECTURE.md @@ -717,7 +717,8 @@ The buffer exposes the following API: - `CloseEpisode(chat_id, history_version, generation_attempt)`: closes an episode, preventing further parts from being added to it. May be called multiple times for a given episode, subsequent calls will be no-ops. Calling it on a non-existent episode creates the episode and closes it immediately. Concurrent parts of the system may race to create the episode and close it, so creating and closing in one operation prevents race conditions. - `AddPart(chat_id, history_version, generation_attempt, content)`: adds a message part to the buffer. Returns a predefined error if the episode is not found or the array is full. - `GetParts(chat_id, history_version, generation_attempt)`: returns the message parts for an episode. Returns a predefined error if the episode is not found. -- `EpisodeModelInvocation(chat_id, history_version, generation_attempt)`: `StartModelInvocation` stamps the instant the episode opens its provider stream, and `ModelInvocationDuration` returns the span between that stamp and `CloseEpisode`. The duration is zero when the episode is unknown, never opened a provider stream (such as a local tool execution batch), or is still open. The interrupt goroutine uses it as the interrupted attempt's billable runtime. +- `StartModelInvocation(chat_id, history_version, generation_attempt)`: stamps the instant the episode opens its provider stream. Returns a predefined error if the episode is not found or already closed. Episodes that never invoke a model, such as local tool execution batches, are never stamped. +- `ModelInvokedAt(chat_id, history_version, generation_attempt)`: returns the instant stamped by `StartModelInvocation`, or the zero time when the episode is unknown or never opened a provider stream. It must be read before `CloseEpisode`, because closed episodes are garbage collected and reading afterwards races the cleanup loop. The interrupt goroutine reads it just before closing the episode and uses the span between that instant and the interrupt as the interrupted attempt's billable runtime. - `SubscribeToEpisode(chat_id, history_version, generation_attempt)`: returns a go channel that will receive all message parts for the episode. It spawns a goroutine that delivers parts to the channel. It's live until the episode is closed or until a subscriber requests that the channel be closed. Once the goroutine delivers all message parts for a closed episode, it closes the channel and exits. If the episode is already closed at the time of the call, the goroutine delivers all message parts for the episode, closes the channel, and exits. `SubscribeToEpisode` does not return an error if the episode is not found: it waits for it to be created instead. Closed episodes are garbage collected after at least 15 seconds since they were closed and when they have no active subscribers. The message part buffer maintains a garbage collection goroutine.