diff --git a/coderd/x/chatd/ARCHITECTURE.md b/coderd/x/chatd/ARCHITECTURE.md index b03bd9d32c5c1..de78ad7215383 100644 --- a/coderd/x/chatd/ARCHITECTURE.md +++ b/coderd/x/chatd/ARCHITECTURE.md @@ -717,6 +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. +- `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. diff --git a/coderd/x/chatd/attempt.go b/coderd/x/chatd/attempt.go index 1e1cc6b0e2d92..fb967fb239887 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 cfb9c51f44765..9d102a29536ed 100644 --- a/coderd/x/chatd/chatd_test.go +++ b/coderd/x/chatd/chatd_test.go @@ -2925,6 +2925,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) @@ -2939,12 +2940,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") @@ -5909,7 +5914,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.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") server = newActiveTestServer(t, db, ps, func(cfg *chatd.Config) { diff --git a/coderd/x/chatd/chatloop/chatloop.go b/coderd/x/chatd/chatloop/chatloop.go index 609a88484bc80..89063ff061280 100644 --- a/coderd/x/chatd/chatloop/chatloop.go +++ b/coderd/x/chatd/chatloop/chatloop.go @@ -73,10 +73,10 @@ 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. Runtime time.Duration // PendingDynamicToolCalls lists tool calls that target // dynamic tools. When non-empty the chatloop exits with @@ -219,6 +219,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 } @@ -305,6 +310,13 @@ type GenerateCompactionOptions struct { ProviderOptions fantasy.ProviderOptions PublishMessagePart func(codersdk.ChatMessageRole, codersdk.ChatMessagePart) + + // Clock measures the summary call duration. Required. + Clock quartz.Clock + + // OnModelStreamStart runs immediately before the summary model call, + // at the instant CompactionResult.Runtime starts measuring. + OnModelStreamStart func() } // ProviderTool pairs a provider-native tool definition with an @@ -400,6 +412,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 2304cadd5ee30..320edca3d9237 100644 --- a/coderd/x/chatd/chatloop/compaction.go +++ b/coderd/x/chatd/chatloop/compaction.go @@ -130,6 +130,11 @@ type CompactionResult struct { UsagePercent float64 ContextTokens int64 ContextLimit int64 + // Runtime is the wall-clock duration of the summarization 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 } // GenerateCompaction generates one context summary and returns it without @@ -140,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 @@ -171,11 +179,16 @@ func GenerateCompaction(ctx context.Context, opts GenerateCompactionOptions) (Co ) } + summaryStart := opts.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 } + summaryRuntime := opts.Clock.Since(summaryStart) if summary == "" { publishCompactionError(config, "compaction produced an empty summary") return CompactionResult{}, xerrors.New("compaction produced an empty summary") @@ -191,6 +204,7 @@ func GenerateCompaction(ctx context.Context, opts GenerateCompactionOptions) (Co UsagePercent: usagePercent, ContextTokens: contextTokens, ContextLimit: contextLimit, + Runtime: summaryRuntime, } if config.PublishMessagePart != nil && config.ToolCallID != "" { resultJSON, _ := json.Marshal(map[string]any{ diff --git a/coderd/x/chatd/chatloop/compaction_internal_test.go b/coderd/x/chatd/chatloop/compaction_internal_test.go index 422aaa5d98ca2..476531139558d 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/chatloop/runtime_test.go b/coderd/x/chatd/chatloop/runtime_test.go new file mode 100644 index 0000000000000..b227b26b6947c --- /dev/null +++ b/coderd/x/chatd/chatloop/runtime_test.go @@ -0,0 +1,161 @@ +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" +) + +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) +} + +// 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() + + 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) +} + +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 + }, + } + + var startedAt []time.Time + 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, + 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 db9f38a960a90..5fb7f00e6fd21 100644 --- a/coderd/x/chatd/generation.go +++ b/coderd/x/chatd/generation.go @@ -734,6 +734,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, @@ -946,8 +947,10 @@ func (s *taskStarter) generateCompaction( } compactionOpts.SummaryHint = preResult.GetModelContext() compactionOpts.PublishMessagePart = attempt.publish + compactionOpts.OnModelStreamStart = attempt.startModelInvocation 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. @@ -1039,6 +1042,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() @@ -1082,6 +1090,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 ddad79e239619..52d749175bd4f 100644 --- a/coderd/x/chatd/message_conversion.go +++ b/coderd/x/chatd/message_conversion.go @@ -195,9 +195,10 @@ func assistantMessage( msg.CacheReadTokens = nullInt64IfNonZero(step.Usage.CacheReadTokens) } 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 } @@ -306,6 +307,8 @@ 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) + assistantMsg.RuntimeMs = nullInt64IfNonZero(input.compaction.Runtime.Milliseconds()) messages := []chatstate.Message{ { Role: database.ChatMessageRoleUser, @@ -314,7 +317,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 { @@ -547,6 +550,12 @@ type bufferedPartsToPartialMessagesInput struct { contentVersion int16 logger slog.Logger interruptedAt time.Time + // 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 } type partialToolCall struct { @@ -597,6 +606,17 @@ func bufferedPartsToPartialMessages(input bufferedPartsToPartialMessagesInput) ( if err := state.appendSyntheticInterruptionResults(); err != nil { return nil, err } + if input.attemptRuntime > 0 && state.modelStreamedAssistant { + // 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 + } + state.messages[i].RuntimeMs = nullInt64IfNonZero(input.attemptRuntime.Milliseconds()) + break + } + } return state.messages, nil } @@ -611,6 +631,14 @@ 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. 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 } func (s *partialMessageConversionState) consume(buffered messagepartbuffer.Part) error { @@ -631,6 +659,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 5679a76e51361..c84935d72f30d 100644 --- a/coderd/x/chatd/message_conversion_test.go +++ b/coderd/x/chatd/message_conversion_test.go @@ -68,18 +68,23 @@ 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) + 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) @@ -95,6 +100,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() @@ -201,6 +234,7 @@ func TestBuildCompactionMessages_CompressedSummaryToolCallAndResult(t *testing.T UsagePercent: 81.5, ContextTokens: 815, ContextLimit: 1000, + Runtime: 1500 * time.Millisecond, }, }) require.NoError(t, err) @@ -212,10 +246,12 @@ 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) + 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) @@ -224,12 +260,41 @@ 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) 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() @@ -624,6 +689,59 @@ func TestBufferedPartsToPartialMessages_NormalizesToolCallDeltasBeforeFinal(t *t require.Equal(t, "call-1", syntheticParts[0].ToolCallID) } +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) + require.Equal(t, database.ChatMessageRoleTool, got[1].Role) + require.False(t, got[1].RuntimeMs.Valid) +} + +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) + + // 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) { t.Parallel() diff --git a/coderd/x/chatd/messagepartbuffer/message_part_buffer.go b/coderd/x/chatd/messagepartbuffer/message_part_buffer.go index 9b14c3287e400..82f41648be6a3 100644 --- a/coderd/x/chatd/messagepartbuffer/message_part_buffer.go +++ b/coderd/x/chatd/messagepartbuffer/message_part_buffer.go @@ -96,7 +96,12 @@ type Buffer struct { } type episodeState struct { - created bool + created bool + // 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 @@ -194,6 +199,25 @@ func (b *Buffer) CreateEpisode(key Key) error { return nil } +// StartModelInvocation stamps the instant the episode opens its provider +// stream, which starts the episode's billable model invocation window. +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 +} + // AddPart appends a part to an existing episode. // // Parts receive contiguous sequence numbers so stream endpoints can detect @@ -266,6 +290,19 @@ func (b *Buffer) GetParts(key Key) ([]Part, error) { return slices.Clone(episode.parts), nil } +// 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 { + return time.Time{} + } + return episode.modelStartedAt +} + // 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 46ff27e5f3565..ea489af244889 100644 --- a/coderd/x/chatd/messagepartbuffer/message_part_buffer_test.go +++ b/coderd/x/chatd/messagepartbuffer/message_part_buffer_test.go @@ -99,6 +99,55 @@ func TestBuffer_CloseEpisodeIdempotent(t *testing.T) { require.NoError(t, buffer.CloseEpisode(key)) } +func TestBuffer_ModelInvokedAt(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.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)) + invokedAt := buffer.ModelInvokedAt(key) + require.Equal(t, clock.Now(), invokedAt) + + // 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, invokedAt, buffer.ModelInvokedAt(key)) + + // Episodes that never open a provider stream, such as local tool + // 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.ModelInvokedAt(toolBatch)) + + // Episodes created implicitly by CloseEpisode never started a + // generation attempt, so they report no invocation stamp. + implicit := testEpisodeKey() + require.NoError(t, buffer.CloseEpisode(implicit)) + require.Zero(t, buffer.ModelInvokedAt(implicit)) +} + func TestBuffer_SubscribeExistingReplaysThenStreamsLiveParts(t *testing.T) { t.Parallel() diff --git a/coderd/x/chatd/tasks.go b/coderd/x/chatd/tasks.go index c1c0c840e4a39..a3738b886dbc5 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,12 +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"), + interruptedAt: interruptedAt, + attemptRuntime: attemptRuntime, }) 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 8f95da4e1933d..feedcbb382ded 100644 --- a/coderd/x/chatd/tasks_test.go +++ b/coderd/x/chatd/tasks_test.go @@ -391,6 +391,92 @@ func TestInterruptTask_BufferedPartsBecomePartialMessages(t *testing.T) { require.True(t, toolParts[0].IsError) } +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)) + // 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) + + 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 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() @@ -1244,13 +1330,20 @@ 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. +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/usage-data-reporting.md b/docs/ai-coder/usage-data-reporting.md index 7483599a827c5..3c69e27ae0b3c 100644 --- a/docs/ai-coder/usage-data-reporting.md +++ b/docs/ai-coder/usage-data-reporting.md @@ -55,3 +55,30 @@ Example of a failed request (e.g. Tallyman Server is blocked by your network): > [!NOTE] > Air-gapped deployments and/or those with legal restrictions around usage reporting can [contact us](https://coder.com/contact) to discuss alternative methods. + +## Agent runtime measurement + +Total Coder Agent runtime is 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.