diff --git a/coderd/database/querier_test.go b/coderd/database/querier_test.go index 32d5bf246979f..39e84ae1b720d 100644 --- a/coderd/database/querier_test.go +++ b/coderd/database/querier_test.go @@ -11190,13 +11190,13 @@ func TestGetTotalChatMessageRuntimeMsInRange(t *testing.T) { LastModelConfigID: mc.ID, }) - insertMessage := func(chatID uuid.UUID, runtimeMs int64, createdAt time.Time, deleted bool) { + insertMessage := func(chatID uuid.UUID, role database.ChatMessageRole, runtimeMs int64, createdAt time.Time, deleted bool) { t.Helper() msg := dbgen.ChatMessage(t, db, database.ChatMessage{ ChatID: chatID, CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, ModelConfigID: uuid.NullUUID{UUID: mc.ID, Valid: true}, - Role: database.ChatMessageRoleAssistant, + Role: role, RuntimeMs: sql.NullInt64{Int64: runtimeMs, Valid: true}, }) _, err := sqlDB.ExecContext(ctx, "UPDATE chat_messages SET created_at = $1, deleted = $2 WHERE id = $3", createdAt, deleted, msg.ID) @@ -11205,22 +11205,24 @@ func TestGetTotalChatMessageRuntimeMsInRange(t *testing.T) { // Counted: on the inclusive start boundary, in the middle (across two // chats), soft-deleted, and just before the exclusive end boundary. - insertMessage(chat1.ID, 1, rangeStart, false) - insertMessage(chat2.ID, 2, rangeStart.Add(30*time.Minute), false) - insertMessage(chat1.ID, 4, rangeStart.Add(45*time.Minute), true) - insertMessage(chat1.ID, 8, rangeEnd.Add(-time.Second), false) + insertMessage(chat1.ID, database.ChatMessageRoleAssistant, 1, rangeStart, false) + insertMessage(chat2.ID, database.ChatMessageRoleAssistant, 2, rangeStart.Add(30*time.Minute), false) + insertMessage(chat1.ID, database.ChatMessageRoleAssistant, 4, rangeStart.Add(45*time.Minute), true) + insertMessage(chat1.ID, database.ChatMessageRoleAssistant, 8, rangeEnd.Add(-time.Second), false) + // Tool rows count because runtime totals are role-agnostic. + insertMessage(chat1.ID, database.ChatMessageRoleTool, 64, rangeStart.Add(20*time.Minute), false) // Not counted: before the range, on the exclusive end boundary, and a // NULL runtime (runtime 0 is stored as NULL). - insertMessage(chat1.ID, 16, rangeStart.Add(-time.Second), false) - insertMessage(chat1.ID, 32, rangeEnd, false) - insertMessage(chat1.ID, 0, rangeStart.Add(10*time.Minute), false) + insertMessage(chat1.ID, database.ChatMessageRoleAssistant, 16, rangeStart.Add(-time.Second), false) + insertMessage(chat1.ID, database.ChatMessageRoleAssistant, 32, rangeEnd, false) + insertMessage(chat1.ID, database.ChatMessageRoleAssistant, 0, rangeStart.Add(10*time.Minute), false) total, err = db.GetTotalChatMessageRuntimeMsInRange(ctx, database.GetTotalChatMessageRuntimeMsInRangeParams{ StartTime: rangeStart, EndTime: rangeEnd, }) require.NoError(t, err) - require.EqualValues(t, 15, total) + require.EqualValues(t, 79, total) } func TestListUsageEventCreatedAtsByTypeSince(t *testing.T) { diff --git a/coderd/usage/usagetypes/events.go b/coderd/usage/usagetypes/events.go index 9004fb2514cc0..a990b4c6f7ed1 100644 --- a/coderd/usage/usagetypes/events.go +++ b/coderd/usage/usagetypes/events.go @@ -203,11 +203,15 @@ func (e HBAISeats) Fields() map[string]any { } // HBAgentRuntime is the event associated with hb_agent_runtime_v1. RuntimeMs -// is the total agent-loop runtime in milliseconds consumed by Coder Agents -// (chats) in one UTC hour. Each measured step spans model streaming (including -// provider-executed tools) and stream retries, and ends when the model stream -// finishes. Time spent executing local tools between steps, including -// sub-agents that bill their own model calls, is excluded. +// is total Coder Agent chat runtime in milliseconds for one UTC hour. +// +// Model steps bill provider streaming. Local tool batches bill the union of +// billed execution intervals, so parallel calls count once and serial calls +// count only from their own start. +// +// Excluded: sub-agent orchestration, client and external-agent work, user or +// idle waits, and retry backoff. Server-executed tools count even when their +// work runs in a connected workspace. // // This measures the new Coder Agents (the `chats` tables), not the deprecated // Tasks counted by dc_managed_agents_v1. diff --git a/coderd/x/chatd/attempt.go b/coderd/x/chatd/attempt.go index fb967fb239887..746df38fc4d9e 100644 --- a/coderd/x/chatd/attempt.go +++ b/coderd/x/chatd/attempt.go @@ -31,6 +31,12 @@ type stepData struct { ContextLimit sql.NullInt64 Runtime time.Duration + // BatchRuntime is the local-tool batch window. Model steps use Runtime. + BatchRuntime time.Duration + // BatchBilledCalls counts the calls whose intervals produced + // BatchRuntime. Audit metadata for the batch usage record. + BatchBilledCalls int + ToolCallCreatedAt map[string]time.Time ToolResultCreatedAt map[string]time.Time ReasoningStartedAt []time.Time diff --git a/coderd/x/chatd/chatd_test.go b/coderd/x/chatd/chatd_test.go index 7d710597e21c5..27de4c9fce61f 100644 --- a/coderd/x/chatd/chatd_test.go +++ b/coderd/x/chatd/chatd_test.go @@ -6591,6 +6591,15 @@ func TestActiveServer_ToolExecutionAndPolicy(t *testing.T) { require.False(t, result.ProviderExecuted) } } + + // Batch runtime bills a dedicated model-only usage row, so no + // user-visible tool row ever carries runtime. + messages := chatMessages(ctx, t, db, chat.ID) + for _, msg := range messages { + if msg.Role == database.ChatMessageRoleTool { + require.False(t, msg.RuntimeMs.Valid) + } + } }) } diff --git a/coderd/x/chatd/chatloop/chatloop.go b/coderd/x/chatd/chatloop/chatloop.go index ba6c1a7d0f3c0..f22cbc8d64b78 100644 --- a/coderd/x/chatd/chatloop/chatloop.go +++ b/coderd/x/chatd/chatloop/chatloop.go @@ -76,11 +76,15 @@ type PersistedStep struct { Content []fantasy.Content Usage fantasy.Usage ContextLimit sql.NullInt64 - // 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 is the wall-clock duration from opening to consuming the + // model stream. Runtime time.Duration + // BatchRuntime is the union of billed local-tool execution intervals. + // Parallel calls count once and serial calls count from their own start. + BatchRuntime time.Duration + // BatchBilledCalls counts the executed calls whose intervals produced + // BatchRuntime. Audit metadata for the batch usage record. + BatchBilledCalls int // PendingDynamicToolCalls lists tool calls that target // dynamic tools. When non-empty the chatloop exits with // ErrDynamicToolCall so the caller can execute them @@ -269,12 +273,32 @@ type ExecuteLocalToolsOptions struct { // is renamed but old chat histories still reference the old name. ToolNameAliases map[string]string + // UnbilledToolNames lists called tool names excluded from the batch + // window. Include deprecated aliases. + UnbilledToolNames map[string]bool + // BillingRecorder observes each local call's start and completion + // for interrupt billing. Serial calls may start after concurrent + // siblings settle, so interrupts bill actual starts and skip calls + // that never run. Optional. + BillingRecorder ToolBillingRecorder + PublishMessagePart func(codersdk.ChatMessageRole, codersdk.ChatMessagePart) Logger slog.Logger Metrics *Metrics Clock quartz.Clock } +// ToolBillingRecorder records live start and completion timestamps for +// local tool calls. Interrupt billing uses these so a cancel can bill +// work that already started and skip calls that never ran. +// dispatchIndex identifies the dispatch-order occurrence. +// RecordComplete may run from multiple tool goroutines; implementations +// must be concurrency-safe. +type ToolBillingRecorder interface { + RecordStart(dispatchIndex int, startedAt time.Time) + RecordComplete(dispatchIndex int, completedAt time.Time) +} + // GenerateCompactionOptions configures one context compaction call. type GenerateCompactionOptions struct { Model fantasy.LanguageModel @@ -620,7 +644,8 @@ func ExecuteLocalTools(ctx context.Context, opts ExecuteLocalToolsOptions) (Pers } maxResultBytes := toolResultByteBudget(opts.ContextLimit) - toolResults := executeTools( + batchStart := clockNow(opts.Clock) + toolExecutions := executeTools( ctx, opts.Clock, opts.Tools, @@ -636,26 +661,83 @@ func ExecuteLocalTools(ctx context.Context, opts ExecuteLocalToolsOptions) (Pers opts.BuiltinToolNames, maxResultBytes, opts.ToolNameAliases, - func(tr fantasy.ToolResultContent, completedAt time.Time) { - recordToolResultTimestamp(&result, tr.ToolCallID, completedAt) - publishToolAttachments(ctx, opts.Logger, tr, completedAt, publishMessagePart) - ssePart := chatprompt.PartFromContentWithLogger(ctx, opts.Logger, tr) - ssePart.CreatedAt = &completedAt - publishMessagePart(codersdk.ChatMessageRoleTool, ssePart) - }, + batchStart, + opts.BillingRecorder, ) + for _, execution := range toolExecutions { + tr := execution.content + completedAt := execution.interval.End + recordToolResultTimestamp(&result, tr.ToolCallID, completedAt) + publishToolAttachments(ctx, opts.Logger, tr, completedAt, publishMessagePart) + ssePart := chatprompt.PartFromContentWithLogger(ctx, opts.Logger, tr) + ssePart.CreatedAt = &completedAt + publishMessagePart(codersdk.ChatMessageRoleTool, ssePart) + result.content = append(result.content, tr) + } if ctx.Err() != nil { return PersistedStep{}, ctx.Err() } - for _, tr := range toolResults { - result.content = append(result.content, tr) - } + billedIntervals := billableBatchIntervals(toolExecutions, opts.UnbilledToolNames) return PersistedStep{ Content: result.content, ToolResultCreatedAt: result.toolResultCreatedAt, + BatchRuntime: BilledIntervalsDuration(billedIntervals), + BatchBilledCalls: len(billedIntervals), }, nil } +// billableBatchIntervals returns the billed execution intervals. +// Unbilled tools and calls without both stamps do not count. +func billableBatchIntervals( + executions []toolExecutionResult, + unbilledToolNames map[string]bool, +) []BilledInterval { + intervals := make([]BilledInterval, 0, len(executions)) + for _, execution := range executions { + if unbilledToolNames[execution.content.ToolName] || + execution.interval.Start.IsZero() || + execution.interval.End.IsZero() { + continue + } + intervals = append(intervals, execution.interval) + } + return intervals +} + +// BilledInterval is one billed tool call's execution window. +type BilledInterval struct { + Start time.Time + End time.Time +} + +// BilledIntervalsDuration returns the union duration of valid intervals. +// Overlaps count once, gaps do not, and inverted intervals are ignored. +// Committed and interrupted batches share this helper. +func BilledIntervalsDuration(intervals []BilledInterval) time.Duration { + valid := slices.DeleteFunc(slices.Clone(intervals), func(iv BilledInterval) bool { + return iv.End.Before(iv.Start) + }) + if len(valid) == 0 { + return 0 + } + slices.SortFunc(valid, func(a, b BilledInterval) int { + return a.Start.Compare(b.Start) + }) + curStart, curEnd := valid[0].Start, valid[0].End + var total time.Duration + for _, iv := range valid[1:] { + if iv.Start.After(curEnd) { + total += curEnd.Sub(curStart) + curStart, curEnd = iv.Start, iv.End + continue + } + if iv.End.After(curEnd) { + curEnd = iv.End + } + } + return total + curEnd.Sub(curStart) +} + // prepareMessagesForRequest applies the prompt preparation pipeline used // immediately before sending messages to a provider. It returns the // possibly updated canonical messages and an independent provider-ready @@ -1081,10 +1163,14 @@ func processStepStream( return result, nil } -// executeTools runs all tool calls concurrently after the stream -// completes. Results are published via onResult in the original -// tool-call order after all tools finish, preserving deterministic -// event ordering for SSE subscribers. +type toolExecutionResult struct { + content fantasy.ToolResultContent + interval BilledInterval +} + +// executeTools runs non-serial calls concurrently, then SerialToolCalls in +// call order. Results are returned in original order after all tools finish. +// recorder, if set, receives live start and completion timestamps. func executeTools( ctx context.Context, clock quartz.Clock, @@ -1100,8 +1186,9 @@ func executeTools( builtinToolNames map[string]bool, maxResultBytes int, toolNameAliases map[string]string, - onResult func(fantasy.ToolResultContent, time.Time), -) []fantasy.ToolResultContent { + batchStart time.Time, + recorder ToolBillingRecorder, +) []toolExecutionResult { if len(toolCalls) == 0 { return nil } @@ -1150,12 +1237,11 @@ func executeTools( } notifyStepToolCallObservers(toolMap, toolNameAliases, observed) - results := make([]fantasy.ToolResultContent, len(localToolCalls)) - completedAt := make([]time.Time, len(localToolCalls)) + executions := make([]toolExecutionResult, len(localToolCalls)) runCall := func(i int, tc fantasy.ToolCallContent) { defer func() { if r := recover(); r != nil { - results[i] = fantasy.ToolResultContent{ + executions[i].content = fantasy.ToolResultContent{ ToolCallID: tc.ToolCallID, ToolName: tc.ToolName, Result: fantasy.ToolResultOutputContentError{ @@ -1166,9 +1252,13 @@ func executeTools( // Record when this tool completed (or panicked). // Captured per call so parallel tools get // accurate individual completion times. - completedAt[i] = clockNow(clock) + completedAt := clockNow(clock) + executions[i].interval.End = completedAt + if recorder != nil { + recorder.RecordComplete(i, completedAt) + } }() - results[i] = executeSingleTool( + executions[i].content = executeSingleTool( ctx, toolMap, tc, @@ -1185,12 +1275,8 @@ func executeTools( toolNameAliases, ) } - // Calls to tools that opt in via SerialToolCalls run in tool-call - // order after every concurrent sibling has settled. The step waits - // for all calls anyway, so sequencing them last costs nothing, and - // order-sensitive shared state (for example the find_tools - // activation budget) is claimed deterministically after sibling - // outcomes are known. All other calls stay concurrent. + // SerialToolCalls run in call order after concurrent siblings settle, so + // order-sensitive state observes final sibling outcomes. var serialIndexes []int var wg sync.WaitGroup for i, tc := range localToolCalls { @@ -1198,6 +1284,10 @@ func executeTools( serialIndexes = append(serialIndexes, i) continue } + executions[i].interval.Start = batchStart + if recorder != nil { + recorder.RecordStart(i, batchStart) + } wg.Add(1) go func() { defer wg.Done() @@ -1206,29 +1296,26 @@ func executeTools( } wg.Wait() - // Reconcile settled sibling outcomes before serial tools run, so - // for example find_tools refunds reservations of errored direct - // calls before its searches admit activations. - settled := make([]fantasy.ToolResultContent, 0, len(results)) - for i := range results { + // Reconcile concurrent results before serial tools inspect shared state. + settled := make([]fantasy.ToolResultContent, 0, len(executions)) + for i := range executions { if !slices.Contains(serialIndexes, i) { - settled = append(settled, results[i]) + settled = append(settled, executions[i].content) } } notifyStepToolResultObservers(toolMap, toolNameAliases, localToolCalls, observed, settled) for _, i := range serialIndexes { + // Stamp serial calls at launch, not batch start. + startedAt := clockNow(clock) + executions[i].interval.Start = startedAt + if recorder != nil { + recorder.RecordStart(i, startedAt) + } runCall(i, localToolCalls[i]) } - // Publish results in the original tool-call order so SSE - // subscribers see a deterministic event sequence. - if onResult != nil { - for i, tr := range results { - onResult(tr, completedAt[i]) - } - } - return results + return executions } // applyExclusiveToolPolicy checks whether toolCalls violate the diff --git a/coderd/x/chatd/chatloop/chatloop_run_internal_test.go b/coderd/x/chatd/chatloop/chatloop_run_internal_test.go index 7d35d2588e708..393d7acbaeb23 100644 --- a/coderd/x/chatd/chatloop/chatloop_run_internal_test.go +++ b/coderd/x/chatd/chatloop/chatloop_run_internal_test.go @@ -949,6 +949,7 @@ func TestExecuteToolsNotifiesStepToolCallObservers(t *testing.T) { map[string]bool{}, defaultToolResultBytes, map[string]string{"observer_alias": "observer_tool"}, + time.Time{}, nil, ) @@ -1020,6 +1021,7 @@ func TestExecuteToolsNotifiesStepToolResultObservers(t *testing.T) { map[string]bool{}, defaultToolResultBytes, map[string]string{"observer_alias": "observer_tool"}, + time.Time{}, nil, ) @@ -1097,6 +1099,7 @@ func TestExecuteToolsReconcilesResultsBeforeSerialCalls(t *testing.T) { map[string]bool{}, defaultToolResultBytes, nil, + time.Time{}, nil, ) @@ -1104,7 +1107,7 @@ func TestExecuteToolsReconcilesResultsBeforeSerialCalls(t *testing.T) { require.Equal(t, []string{"failing_tool"}, erroredAtRun, "a serial tool must see settled sibling outcomes before it executes") require.Len(t, results, 2) - require.Equal(t, "1", results[0].ToolCallID, "results keep original call order") + require.Equal(t, "1", results[0].content.ToolCallID, "results keep original call order") } func TestExecuteToolsSerialToolCallOrder(t *testing.T) { @@ -1164,6 +1167,7 @@ func TestExecuteToolsSerialToolCallOrder(t *testing.T) { map[string]bool{}, defaultToolResultBytes, nil, + time.Time{}, nil, ) @@ -1177,10 +1181,122 @@ func TestExecuteToolsSerialToolCallOrder(t *testing.T) { } require.Len(t, results, len(calls)) for i, tc := range calls { - require.Equal(t, tc.ToolCallID, results[i].ToolCallID, "results keep original call order") + require.Equal(t, tc.ToolCallID, results[i].content.ToolCallID, "results keep original call order") } } +type timingEvent struct { + dispatchIndex int + at time.Time +} + +// liveToolBillingRecorder is a test ToolBillingRecorder that publishes +// start and completion timestamps as they happen. +type liveToolBillingRecorder struct { + started chan timingEvent + completed chan timingEvent +} + +func (r liveToolBillingRecorder) RecordStart(dispatchIndex int, startedAt time.Time) { + r.started <- timingEvent{dispatchIndex: dispatchIndex, at: startedAt} +} + +func (r liveToolBillingRecorder) RecordComplete(dispatchIndex int, completedAt time.Time) { + r.completed <- timingEvent{dispatchIndex: dispatchIndex, at: completedAt} +} + +func TestExecuteToolsReturnsExecutionIntervals(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitShort) + started := make(chan timingEvent, 3) + completed := make(chan timingEvent, 3) + + slowGo := make(chan struct{}) + fastGo := make(chan struct{}) + blocking := func(name string, release <-chan struct{}) fantasy.AgentTool { + return fantasy.NewAgentTool( + name, + "waits for the test to release it", + func(_ context.Context, _ struct{}, _ fantasy.ToolCall) (fantasy.ToolResponse, error) { + <-release + return fantasy.NewTextResponse("ok"), nil + }, + ) + } + slow := blocking("slow_tool", slowGo) + fast := blocking("fast_tool", fastGo) + serial := serialMarkerTool{AgentTool: fantasy.NewAgentTool( + "serial_tool", + "runs after concurrent tools", + func(_ context.Context, _ struct{}, _ fantasy.ToolCall) (fantasy.ToolResponse, error) { + return fantasy.NewTextResponse("ok"), nil + }, + )} + calls := []fantasy.ToolCallContent{ + {ToolCallID: "slow", ToolName: "slow_tool", Input: "{}"}, + {ToolCallID: "fast", ToolName: "fast_tool", Input: "{}"}, + {ToolCallID: "serial", ToolName: "serial_tool", Input: "{}"}, + } + batchStart := time.Date(2024, time.January, 1, 0, 0, 0, 0, time.UTC) + resultsCh := make(chan []toolExecutionResult, 1) + go func() { + resultsCh <- executeTools( + ctx, + quartz.NewReal(), + []fantasy.AgentTool{slow, fast, serial}, + nil, + nil, + nil, + calls, + nil, + NewMetrics(prometheus.NewRegistry()), + slog.Make(), + "fake", "fake-model", + map[string]bool{}, + defaultToolResultBytes, + nil, + batchStart, + liveToolBillingRecorder{started: started, completed: completed}, + ) + }() + + startByIndex := make([]time.Time, len(calls)) + for range 2 { + event := testutil.RequireReceive(ctx, t, started) + startByIndex[event.dispatchIndex] = event.at + } + + close(fastGo) + fastCompleted := testutil.RequireReceive(ctx, t, completed) + require.Equal(t, 1, fastCompleted.dispatchIndex) + + close(slowGo) + slowCompleted := testutil.RequireReceive(ctx, t, completed) + require.Equal(t, 0, slowCompleted.dispatchIndex) + + serialStarted := testutil.RequireReceive(ctx, t, started) + require.Equal(t, 2, serialStarted.dispatchIndex) + startByIndex[serialStarted.dispatchIndex] = serialStarted.at + serialCompleted := testutil.RequireReceive(ctx, t, completed) + require.Equal(t, 2, serialCompleted.dispatchIndex) + + endByIndex := []time.Time{slowCompleted.at, fastCompleted.at, serialCompleted.at} + results := testutil.RequireReceive(ctx, t, resultsCh) + require.Len(t, results, len(calls)) + for i, call := range calls { + require.Equal(t, call.ToolCallID, results[i].content.ToolCallID, + "results keep dispatch order when calls complete out of order") + require.Equal(t, startByIndex[i], results[i].interval.Start) + require.Equal(t, endByIndex[i], results[i].interval.End) + } + require.Equal(t, batchStart, results[0].interval.Start) + require.Equal(t, batchStart, results[1].interval.Start) + require.NotEqual(t, batchStart, results[2].interval.Start) + require.False(t, results[2].interval.Start.Before(slowCompleted.at), + "serial execution starts only after concurrent calls settle") +} + func TestExecuteSingleTool_MediaBase64Encoding(t *testing.T) { t.Parallel() diff --git a/coderd/x/chatd/chatloop/runtime_test.go b/coderd/x/chatd/chatloop/runtime_test.go index b227b26b6947c..62211b13f928e 100644 --- a/coderd/x/chatd/chatloop/runtime_test.go +++ b/coderd/x/chatd/chatloop/runtime_test.go @@ -6,11 +6,13 @@ import ( "time" "charm.land/fantasy" + "github.com/stretchr/testify/assert" "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/coder/v2/testutil" "github.com/coder/quartz" ) @@ -159,3 +161,562 @@ func TestGenerateCompaction_RecordsRuntime(t *testing.T) { require.Len(t, startedAt, 1) require.Equal(t, result.Runtime, clock.Since(startedAt[0])) } + +// executeToolBatch lets tests release trapped clock events in order, so +// goroutines cannot race clock advances. +func executeToolBatch( + t *testing.T, + clock *quartz.Mock, + opts chatloop.ExecuteLocalToolsOptions, +) <-chan chatloop.PersistedStep { + t.Helper() + opts.Clock = clock + resultCh := make(chan chatloop.PersistedStep, 1) + go func() { + outcome, err := chatloop.ExecuteLocalTools(context.Background(), opts) + assert.NoError(t, err) + resultCh <- outcome + }() + return resultCh +} + +func blockingTool(name string, release <-chan struct{}, response fantasy.ToolResponse) fantasy.AgentTool { + return fantasy.NewAgentTool( + name, + "test tool that completes when released", + func(_ context.Context, _ struct{}, _ fantasy.ToolCall) (fantasy.ToolResponse, error) { + <-release + return response, nil + }, + ) +} + +type serialTool struct { + fantasy.AgentTool +} + +func (serialTool) SerialToolCalls() bool { return true } + +func TestExecuteLocalTools_BatchWindowIsMaxNotSum(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitShort) + clock := quartz.NewMock(t) + trap := clock.Trap().Now() + defer trap.Close() + + fastGo := make(chan struct{}) + slowGo := make(chan struct{}) + resultCh := executeToolBatch(t, clock, chatloop.ExecuteLocalToolsOptions{ + Tools: []fantasy.AgentTool{ + blockingTool("fast_tool", fastGo, fantasy.NewTextResponse("done")), + blockingTool("slow_tool", slowGo, fantasy.NewTextErrorResponse("blew up")), + }, + ActiveTools: []string{"fast_tool", "slow_tool"}, + ToolCalls: []fantasy.ToolCallContent{ + {ToolCallID: "call-fast", ToolName: "fast_tool", Input: "{}"}, + {ToolCallID: "call-slow", ToolName: "slow_tool", Input: "{}"}, + }, + }) + + trap.MustWait(ctx).MustRelease(ctx) + clock.Advance(10 * time.Second) + close(fastGo) + trap.MustWait(ctx).MustRelease(ctx) + clock.Advance(50 * time.Second) + close(slowGo) + trap.MustWait(ctx).MustRelease(ctx) + + outcome := testutil.RequireReceive(ctx, t, resultCh) + require.Equal(t, 60*time.Second, outcome.BatchRuntime) + require.Equal(t, 2, outcome.BatchBilledCalls) +} + +func TestExecuteLocalTools_SimultaneousCompletionsBillOnce(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitShort) + clock := quartz.NewMock(t) + trap := clock.Trap().Now() + defer trap.Close() + + release := make(chan struct{}) + resultCh := executeToolBatch(t, clock, chatloop.ExecuteLocalToolsOptions{ + Tools: []fantasy.AgentTool{ + blockingTool("read_tool", release, fantasy.NewTextResponse("done")), + }, + ActiveTools: []string{"read_tool"}, + ToolCalls: []fantasy.ToolCallContent{ + {ToolCallID: "call-1", ToolName: "read_tool", Input: "{}"}, + {ToolCallID: "call-2", ToolName: "read_tool", Input: "{}"}, + {ToolCallID: "call-3", ToolName: "read_tool", Input: "{}"}, + }, + }) + + trap.MustWait(ctx).MustRelease(ctx) + clock.Advance(10 * time.Second) + close(release) + for range 3 { + trap.MustWait(ctx).MustRelease(ctx) + } + + outcome := testutil.RequireReceive(ctx, t, resultCh) + require.Equal(t, 10*time.Second, outcome.BatchRuntime) +} + +func TestExecuteLocalTools_UnbilledToolNeverExtendsWindow(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitShort) + clock := quartz.NewMock(t) + trap := clock.Trap().Now() + defer trap.Close() + + executeGo := make(chan struct{}) + waitGo := make(chan struct{}) + resultCh := executeToolBatch(t, clock, chatloop.ExecuteLocalToolsOptions{ + Tools: []fantasy.AgentTool{ + blockingTool("execute", executeGo, fantasy.NewTextResponse("done")), + blockingTool("wait_agent", waitGo, fantasy.NewTextResponse("child report")), + }, + ActiveTools: []string{"execute", "wait_agent"}, + UnbilledToolNames: map[string]bool{"wait_agent": true}, + ToolCalls: []fantasy.ToolCallContent{ + {ToolCallID: "call-execute", ToolName: "execute", Input: "{}"}, + {ToolCallID: "call-wait", ToolName: "wait_agent", Input: "{}"}, + }, + }) + + trap.MustWait(ctx).MustRelease(ctx) + clock.Advance(10 * time.Second) + close(executeGo) + trap.MustWait(ctx).MustRelease(ctx) + clock.Advance(50 * time.Second) + close(waitGo) + trap.MustWait(ctx).MustRelease(ctx) + + outcome := testutil.RequireReceive(ctx, t, resultCh) + require.Equal(t, 10*time.Second, outcome.BatchRuntime) +} + +func TestExecuteLocalTools_UnbilledOnlyBatchBillsNothing(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitShort) + clock := quartz.NewMock(t) + trap := clock.Trap().Now() + defer trap.Close() + + waitGo := make(chan struct{}) + resultCh := executeToolBatch(t, clock, chatloop.ExecuteLocalToolsOptions{ + Tools: []fantasy.AgentTool{ + blockingTool("wait_agent", waitGo, fantasy.NewTextResponse("child report")), + }, + ActiveTools: []string{"wait_agent"}, + UnbilledToolNames: map[string]bool{"wait_agent": true}, + ToolCalls: []fantasy.ToolCallContent{ + {ToolCallID: "call-wait", ToolName: "wait_agent", Input: "{}"}, + }, + }) + + trap.MustWait(ctx).MustRelease(ctx) + clock.Advance(60 * time.Second) + close(waitGo) + trap.MustWait(ctx).MustRelease(ctx) + + outcome := testutil.RequireReceive(ctx, t, resultCh) + require.Zero(t, outcome.BatchRuntime) + require.Zero(t, outcome.BatchBilledCalls) +} + +func TestExecuteLocalTools_AliasNamesClassifyAsCalled(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitShort) + clock := quartz.NewMock(t) + trap := clock.Trap().Now() + defer trap.Close() + + executeGo := make(chan struct{}) + legacyGo := make(chan struct{}) + resultCh := executeToolBatch(t, clock, chatloop.ExecuteLocalToolsOptions{ + Tools: []fantasy.AgentTool{ + blockingTool("execute", executeGo, fantasy.NewTextResponse("done")), + blockingTool("interrupt_agent", legacyGo, fantasy.NewTextResponse("stopped")), + }, + ActiveTools: []string{"execute", "interrupt_agent"}, + ToolNameAliases: map[string]string{"close_agent": "interrupt_agent"}, + UnbilledToolNames: map[string]bool{ + "interrupt_agent": true, + "close_agent": true, + }, + ToolCalls: []fantasy.ToolCallContent{ + {ToolCallID: "call-execute", ToolName: "execute", Input: "{}"}, + {ToolCallID: "call-legacy", ToolName: "close_agent", Input: "{}"}, + }, + }) + + trap.MustWait(ctx).MustRelease(ctx) + clock.Advance(10 * time.Second) + close(executeGo) + trap.MustWait(ctx).MustRelease(ctx) + clock.Advance(50 * time.Second) + close(legacyGo) + trap.MustWait(ctx).MustRelease(ctx) + + outcome := testutil.RequireReceive(ctx, t, resultCh) + require.Equal(t, 10*time.Second, outcome.BatchRuntime) + require.Equal(t, 1, outcome.BatchBilledCalls) +} + +func TestExecuteLocalTools_DuplicateToolCallIDsKeepOccurrenceCompletions(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitShort) + clock := quartz.NewMock(t) + trap := clock.Trap().Now() + defer trap.Close() + + slowGo := make(chan struct{}) + fastGo := make(chan struct{}) + resultCh := executeToolBatch(t, clock, chatloop.ExecuteLocalToolsOptions{ + Tools: []fantasy.AgentTool{ + blockingTool("slow_tool", slowGo, fantasy.NewTextResponse("done")), + blockingTool("fast_tool", fastGo, fantasy.NewTextResponse("done")), + }, + ActiveTools: []string{"slow_tool", "fast_tool"}, + ToolCalls: []fantasy.ToolCallContent{ + {ToolCallID: "call-dup", ToolName: "slow_tool", Input: "{}"}, + {ToolCallID: "call-dup", ToolName: "fast_tool", Input: "{}"}, + }, + }) + + trap.MustWait(ctx).MustRelease(ctx) + clock.Advance(10 * time.Second) + close(fastGo) + trap.MustWait(ctx).MustRelease(ctx) + clock.Advance(50 * time.Second) + close(slowGo) + trap.MustWait(ctx).MustRelease(ctx) + + outcome := testutil.RequireReceive(ctx, t, resultCh) + require.Equal(t, 60*time.Second, outcome.BatchRuntime) +} + +// recordingToolBillingRecorder is a test ToolBillingRecorder that +// counts start/complete calls and can publish live completions. +type recordingToolBillingRecorder struct { + starts int + completions int + completeCh chan recordedToolCompletion +} + +type recordedToolCompletion struct { + dispatchIndex int + completedAt time.Time +} + +func (r *recordingToolBillingRecorder) RecordStart(int, time.Time) { + r.starts++ +} + +func (r *recordingToolBillingRecorder) RecordComplete(dispatchIndex int, completedAt time.Time) { + r.completions++ + if r.completeCh != nil { + r.completeCh <- recordedToolCompletion{ + dispatchIndex: dispatchIndex, + completedAt: completedAt, + } + } +} + +func TestExecuteLocalTools_BillingRecorderRecordsOnlyRuns(t *testing.T) { + t.Parallel() + + t.Run("started call records a paired lifecycle", func(t *testing.T) { + t.Parallel() + + recorder := &recordingToolBillingRecorder{} + startedWhenToolRan := false + tool := fantasy.NewAgentTool( + "fast_tool", + "test tool", + func(context.Context, struct{}, fantasy.ToolCall) (fantasy.ToolResponse, error) { + startedWhenToolRan = recorder.starts > 0 + return fantasy.NewTextResponse("done"), nil + }, + ) + outcome, err := chatloop.ExecuteLocalTools(context.Background(), chatloop.ExecuteLocalToolsOptions{ + Clock: quartz.NewMock(t), + Tools: []fantasy.AgentTool{tool}, + ActiveTools: []string{"fast_tool"}, + BillingRecorder: recorder, + ToolCalls: []fantasy.ToolCallContent{ + {ToolCallID: "call-1", ToolName: "fast_tool", Input: "{}"}, + }, + }) + require.NoError(t, err) + require.Len(t, outcome.Content, 1) + require.Equal(t, 1, recorder.starts) + require.Equal(t, 1, recorder.completions) + require.True(t, startedWhenToolRan, "RecordStart must run before the tool runs") + }) + + t.Run("canceled context records no lifecycle", func(t *testing.T) { + t.Parallel() + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + recorder := &recordingToolBillingRecorder{} + _, err := chatloop.ExecuteLocalTools(ctx, chatloop.ExecuteLocalToolsOptions{ + Clock: quartz.NewMock(t), + BillingRecorder: recorder, + ToolCalls: []fantasy.ToolCallContent{ + {ToolCallID: "call-1", ToolName: "fast_tool", Input: "{}"}, + }, + }) + require.ErrorIs(t, err, context.Canceled) + require.Zero(t, recorder.starts) + require.Zero(t, recorder.completions) + }) + + t.Run("exclusive violation records no lifecycle", func(t *testing.T) { + t.Parallel() + + recorder := &recordingToolBillingRecorder{} + outcome, err := chatloop.ExecuteLocalTools(context.Background(), chatloop.ExecuteLocalToolsOptions{ + Clock: quartz.NewMock(t), + ExclusiveToolNames: map[string]bool{"exclusive_tool": true}, + BillingRecorder: recorder, + ToolCalls: []fantasy.ToolCallContent{ + {ToolCallID: "call-1", ToolName: "exclusive_tool", Input: "{}"}, + {ToolCallID: "call-2", ToolName: "fast_tool", Input: "{}"}, + }, + }) + require.NoError(t, err) + require.Len(t, outcome.Content, 2, "the whole batch resolves to synthesized policy errors") + require.Zero(t, recorder.starts) + require.Zero(t, recorder.completions) + }) +} + +func TestExecuteLocalTools_EmptyToolCallIDStillBillsWindow(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitShort) + clock := quartz.NewMock(t) + trap := clock.Trap().Now() + defer trap.Close() + + idlessGo := make(chan struct{}) + fastGo := make(chan struct{}) + resultCh := executeToolBatch(t, clock, chatloop.ExecuteLocalToolsOptions{ + Tools: []fantasy.AgentTool{ + blockingTool("idless_tool", idlessGo, fantasy.NewTextResponse("done")), + blockingTool("fast_tool", fastGo, fantasy.NewTextResponse("done")), + }, + ActiveTools: []string{"idless_tool", "fast_tool"}, + ToolCalls: []fantasy.ToolCallContent{ + {ToolCallID: "", ToolName: "idless_tool", Input: "{}"}, + {ToolCallID: "call-fast", ToolName: "fast_tool", Input: "{}"}, + }, + }) + + trap.MustWait(ctx).MustRelease(ctx) + clock.Advance(10 * time.Second) + close(fastGo) + trap.MustWait(ctx).MustRelease(ctx) + clock.Advance(50 * time.Second) + close(idlessGo) + trap.MustWait(ctx).MustRelease(ctx) + + outcome := testutil.RequireReceive(ctx, t, resultCh) + require.Equal(t, 60*time.Second, outcome.BatchRuntime) +} + +func TestExecuteLocalTools_BillingRecorderReportsLiveCompletions(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitShort) + clock := quartz.NewMock(t) + trap := clock.Trap().Now() + defer trap.Close() + + recorder := &recordingToolBillingRecorder{ + completeCh: make(chan recordedToolCompletion, 2), + } + fastGo := make(chan struct{}) + slowGo := make(chan struct{}) + resultCh := executeToolBatch(t, clock, chatloop.ExecuteLocalToolsOptions{ + Tools: []fantasy.AgentTool{ + blockingTool("fast_tool", fastGo, fantasy.NewTextResponse("done")), + blockingTool("slow_tool", slowGo, fantasy.NewTextResponse("done")), + }, + ActiveTools: []string{"fast_tool", "slow_tool"}, + BillingRecorder: recorder, + ToolCalls: []fantasy.ToolCallContent{ + {ToolCallID: "call-fast", ToolName: "fast_tool", Input: "{}"}, + {ToolCallID: "call-slow", ToolName: "slow_tool", Input: "{}"}, + }, + }) + + trap.MustWait(ctx).MustRelease(ctx) + clock.Advance(10 * time.Second) + close(fastGo) + trap.MustWait(ctx).MustRelease(ctx) + fast := testutil.RequireReceive(ctx, t, recorder.completeCh) + require.Equal(t, 0, fast.dispatchIndex) + clock.Advance(50 * time.Second) + close(slowGo) + trap.MustWait(ctx).MustRelease(ctx) + slow := testutil.RequireReceive(ctx, t, recorder.completeCh) + require.Equal(t, 1, slow.dispatchIndex) + require.Equal(t, 50*time.Second, slow.completedAt.Sub(fast.completedAt)) + + outcome := testutil.RequireReceive(ctx, t, resultCh) + require.Equal(t, map[string]time.Time{ + "call-fast": fast.completedAt, + "call-slow": slow.completedAt, + }, outcome.ToolResultCreatedAt) +} + +func TestExecuteLocalTools_SerialCallBillsFromItsOwnStart(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitShort) + clock := quartz.NewMock(t) + trap := clock.Trap().Now() + defer trap.Close() + + waitGo := make(chan struct{}) + serialGo := make(chan struct{}) + resultCh := executeToolBatch(t, clock, chatloop.ExecuteLocalToolsOptions{ + Tools: []fantasy.AgentTool{ + blockingTool("wait_agent", waitGo, fantasy.NewTextResponse("child report")), + serialTool{blockingTool("serial_tool", serialGo, fantasy.NewTextResponse("done"))}, + }, + ActiveTools: []string{"wait_agent", "serial_tool"}, + UnbilledToolNames: map[string]bool{"wait_agent": true}, + ToolCalls: []fantasy.ToolCallContent{ + {ToolCallID: "call-wait", ToolName: "wait_agent", Input: "{}"}, + {ToolCallID: "call-serial", ToolName: "serial_tool", Input: "{}"}, + }, + }) + + trap.MustWait(ctx).MustRelease(ctx) + clock.Advance(10 * time.Minute) + close(waitGo) + trap.MustWait(ctx).MustRelease(ctx) + // Release the serial start timestamp. + trap.MustWait(ctx).MustRelease(ctx) + clock.Advance(2 * time.Second) + close(serialGo) + trap.MustWait(ctx).MustRelease(ctx) + + outcome := testutil.RequireReceive(ctx, t, resultCh) + require.Equal(t, 2*time.Second, outcome.BatchRuntime, + "a serial call bills its own execution, not the unbilled wait that delayed its launch") +} + +func TestExecuteLocalTools_SerialAfterBilledSiblingBillsUnion(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitShort) + clock := quartz.NewMock(t) + trap := clock.Trap().Now() + defer trap.Close() + + execGo := make(chan struct{}) + waitGo := make(chan struct{}) + serialGo := make(chan struct{}) + resultCh := executeToolBatch(t, clock, chatloop.ExecuteLocalToolsOptions{ + Tools: []fantasy.AgentTool{ + blockingTool("execute", execGo, fantasy.NewTextResponse("done")), + blockingTool("wait_agent", waitGo, fantasy.NewTextResponse("child report")), + serialTool{blockingTool("serial_tool", serialGo, fantasy.NewTextResponse("done"))}, + }, + ActiveTools: []string{"execute", "wait_agent", "serial_tool"}, + UnbilledToolNames: map[string]bool{"wait_agent": true}, + ToolCalls: []fantasy.ToolCallContent{ + {ToolCallID: "call-execute", ToolName: "execute", Input: "{}"}, + {ToolCallID: "call-wait", ToolName: "wait_agent", Input: "{}"}, + {ToolCallID: "call-serial", ToolName: "serial_tool", Input: "{}"}, + }, + }) + + trap.MustWait(ctx).MustRelease(ctx) + clock.Advance(3 * time.Second) + close(execGo) + trap.MustWait(ctx).MustRelease(ctx) + clock.Advance(7 * time.Second) + close(waitGo) + trap.MustWait(ctx).MustRelease(ctx) + // Release the serial start timestamp. + trap.MustWait(ctx).MustRelease(ctx) + clock.Advance(2 * time.Second) + close(serialGo) + trap.MustWait(ctx).MustRelease(ctx) + + outcome := testutil.RequireReceive(ctx, t, resultCh) + require.Equal(t, 5*time.Second, outcome.BatchRuntime, + "the 3s concurrent window and the 2s serial window bill; the 7s span where only wait_agent ran does not") +} + +func TestBilledIntervalsDuration(t *testing.T) { + t.Parallel() + + base := time.Date(2024, 1, 1, 0, 0, 0, 0, time.UTC) + at := base.Add + for _, tc := range []struct { + name string + intervals []chatloop.BilledInterval + want time.Duration + }{ + {name: "empty", want: 0}, + { + name: "overlapping intervals bill once", + intervals: []chatloop.BilledInterval{ + {Start: at(0), End: at(10 * time.Second)}, + {Start: at(0), End: at(4 * time.Second)}, + }, + want: 10 * time.Second, + }, + { + name: "gap between intervals is not billed", + intervals: []chatloop.BilledInterval{ + {Start: at(0), End: at(3 * time.Second)}, + {Start: at(10 * time.Second), End: at(12 * time.Second)}, + }, + want: 5 * time.Second, + }, + { + name: "unsorted contained interval adds nothing", + intervals: []chatloop.BilledInterval{ + {Start: at(2 * time.Second), End: at(4 * time.Second)}, + {Start: at(0), End: at(10 * time.Second)}, + }, + want: 10 * time.Second, + }, + { + name: "touching intervals merge without a gap", + intervals: []chatloop.BilledInterval{ + {Start: at(0), End: at(3 * time.Second)}, + {Start: at(3 * time.Second), End: at(5 * time.Second)}, + }, + want: 5 * time.Second, + }, + { + name: "inverted interval is ignored", + intervals: []chatloop.BilledInterval{ + {Start: at(5 * time.Second), End: at(0)}, + {Start: at(0), End: at(2 * time.Second)}, + }, + want: 2 * time.Second, + }, + } { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + require.Equal(t, tc.want, chatloop.BilledIntervalsDuration(tc.intervals)) + }) + } +} diff --git a/coderd/x/chatd/generation.go b/coderd/x/chatd/generation.go index a3a2aeb2e9666..7c948383244a8 100644 --- a/coderd/x/chatd/generation.go +++ b/coderd/x/chatd/generation.go @@ -874,6 +874,7 @@ func (s *taskStarter) executeLocalTools( ModelName: modelName, ContextLimit: prepared.ContextLimitFallback, ToolNameAliases: subagentToolNameAliases, + UnbilledToolNames: unbilledSubagentToolNames, PublishMessagePart: attempt.publish, Logger: s.opts.Logger, Metrics: s.server.metrics, @@ -1550,6 +1551,8 @@ func stepDataFromPersisted(step chatloop.PersistedStep) stepData { Usage: step.Usage, ContextLimit: step.ContextLimit, Runtime: step.Runtime, + BatchRuntime: step.BatchRuntime, + BatchBilledCalls: step.BatchBilledCalls, ToolCallCreatedAt: step.ToolCallCreatedAt, ToolResultCreatedAt: step.ToolResultCreatedAt, ReasoningStartedAt: step.ReasoningStartedAt, diff --git a/coderd/x/chatd/message_conversion.go b/coderd/x/chatd/message_conversion.go index 52d749175bd4f..c5fc4be8b893f 100644 --- a/coderd/x/chatd/message_conversion.go +++ b/coderd/x/chatd/message_conversion.go @@ -76,6 +76,16 @@ func buildCommitStepMessages(input buildCommitStepMessagesInput) (stepMessagesFo messages = append(messages, baseMessage(database.ChatMessageRoleTool, database.ChatMessageVisibilityBoth, input.modelConfigID, contentVersion, content)) } + // Usage sums runtime_ms across rows, so the batch window is billed + // once on a dedicated record instead of an arbitrary member row. + stamp, ok, err := batchUsageMessage(input.modelConfigID, contentVersion, input.step.BatchRuntime, input.step.BatchBilledCalls) + if err != nil { + return stepMessagesForCommit{}, err + } + if ok { + messages = append(messages, stamp) + } + return stepMessagesForCommit{ Messages: messages, VisibleIndexes: visibleMessageIndexes(messages), @@ -225,6 +235,54 @@ func nullInt64IfNonZero(value int64) sql.NullInt64 { return sql.NullInt64{Int64: value, Valid: true} } +// toolBatchUsagePartType marks the dedicated billing record for a local +// tool batch. Internal to chatd: the row is persisted with model +// visibility so it never reaches the API or SSE, and prompt replay drops +// it because the part converts to no provider content. +const toolBatchUsagePartType codersdk.ChatMessagePartType = "tool-batch-usage" + +// toolBatchUsagePayload is the audit payload stored on the usage record. +// It duplicates the row's runtime_ms so the billed window survives in +// content for debugging, alongside how many call intervals produced it. +type toolBatchUsagePayload struct { + BilledMs int64 `json:"billed_ms"` + BilledCalls int `json:"billed_calls"` +} + +// batchUsageMessage builds the single model-invisible row that carries a +// local tool batch's billed runtime. Usage sums runtime_ms across rows, +// so a dedicated record keeps real tool results free of batch-level +// runtime. Completed and interrupted batches share this helper. Returns +// false when the batch bills no whole millisecond. +func batchUsageMessage( + modelConfigID uuid.UUID, + contentVersion int16, + runtime time.Duration, + billedCalls int, +) (chatstate.Message, bool, error) { + runtimeMs := runtime.Milliseconds() + if runtimeMs <= 0 { + return chatstate.Message{}, false, nil + } + payload, err := json.Marshal(toolBatchUsagePayload{ + BilledMs: runtimeMs, + BilledCalls: billedCalls, + }) + if err != nil { + return chatstate.Message{}, false, xerrors.Errorf("marshal tool batch usage payload: %w", err) + } + content, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{{ + Type: toolBatchUsagePartType, + Result: payload, + }}) + if err != nil { + return chatstate.Message{}, false, xerrors.Errorf("marshal tool batch usage part: %w", err) + } + msg := baseMessage(database.ChatMessageRoleTool, database.ChatMessageVisibilityModel, modelConfigID, contentVersion, content) + msg.RuntimeMs = sql.NullInt64{Int64: runtimeMs, Valid: true} + return msg, true, nil +} + func visibleMessageIndexes(messages []chatstate.Message) []int { indexes := make([]int, 0, len(messages)) for i, msg := range messages { @@ -631,13 +689,8 @@ 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 distinguishes streamed content from tool + // attachment parts, which must not carry model runtime. modelStreamedAssistant bool } diff --git a/coderd/x/chatd/message_conversion_test.go b/coderd/x/chatd/message_conversion_test.go index c84935d72f30d..d3258a43faf60 100644 --- a/coderd/x/chatd/message_conversion_test.go +++ b/coderd/x/chatd/message_conversion_test.go @@ -100,9 +100,85 @@ 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) { +func TestBuildCommitStepMessages_BatchRuntimeBillsDedicatedUsageRow(t *testing.T) { + t.Parallel() + + got, err := buildCommitStepMessages(buildCommitStepMessagesInput{ + modelConfigID: uuid.New(), + contentVersion: chatprompt.CurrentContentVersion, + logger: slog.Make(), + step: stepData{ + Content: []fantasy.Content{ + fantasy.ToolResultContent{ + ToolCallID: "call-1", + ToolName: "read_file", + Result: fantasy.ToolResultOutputContentText{Text: `{"data":"fast"}`}, + }, + fantasy.ToolResultContent{ + ToolCallID: "call-2", + ToolName: "execute", + Result: fantasy.ToolResultOutputContentText{Text: `{"stdout":"/tmp"}`}, + }, + }, + BatchRuntime: 10 * time.Second, + BatchBilledCalls: 2, + }, + }) + require.NoError(t, err) + require.Len(t, got.Messages, 3) + // Real tool results never carry batch-level runtime. + require.False(t, got.Messages[0].RuntimeMs.Valid) + require.False(t, got.Messages[1].RuntimeMs.Valid) + + stamp := got.Messages[2] + require.Equal(t, database.ChatMessageRoleTool, stamp.Role) + require.Equal(t, database.ChatMessageVisibilityModel, stamp.Visibility) + require.Equal(t, sql.NullInt64{Int64: 10000, Valid: true}, stamp.RuntimeMs) + stampParts, err := chatprompt.ParseContent(database.ChatMessage{ + Role: stamp.Role, + Content: stamp.Content, + ContentVersion: chatprompt.CurrentContentVersion, + }) + require.NoError(t, err) + require.Len(t, stampParts, 1) + require.Equal(t, toolBatchUsagePartType, stampParts[0].Type) + require.JSONEq(t, `{"billed_ms":10000,"billed_calls":2}`, string(stampParts[0].Result)) + // The usage record is model-only bookkeeping, never published to + // clients. + require.Equal(t, []int{0, 1}, got.VisibleIndexes) +} + +func TestBuildCommitStepMessages_BatchAttachmentAssistantRowStaysNull(t *testing.T) { + t.Parallel() + + got, err := buildCommitStepMessages(buildCommitStepMessagesInput{ + modelConfigID: uuid.New(), + contentVersion: chatprompt.CurrentContentVersion, + logger: slog.Make(), + step: stepData{ + Content: []fantasy.Content{ + fantasy.ToolResultContent{ + ToolCallID: "call-1", + ToolName: "attach_file", + Result: fantasy.ToolResultOutputContentText{Text: `{"ok":true}`}, + ClientMetadata: `{"attachments":[{"file_id":"` + uuid.NewString() + `","media_type":"image/png","name":"shot.png"}]}`, + }, + }, + BatchRuntime: 3 * time.Second, + BatchBilledCalls: 1, + }, + }) + require.NoError(t, err) + require.Len(t, got.Messages, 3) + require.Equal(t, database.ChatMessageRoleAssistant, got.Messages[0].Role) + require.False(t, got.Messages[0].RuntimeMs.Valid) + require.Equal(t, database.ChatMessageRoleTool, got.Messages[1].Role) + require.False(t, got.Messages[1].RuntimeMs.Valid) + require.Equal(t, database.ChatMessageVisibilityModel, got.Messages[2].Visibility) + require.Equal(t, sql.NullInt64{Int64: 3000, Valid: true}, got.Messages[2].RuntimeMs) +} + +func TestBuildCommitStepMessages_ZeroBatchRuntimeLeavesRuntimeNull(t *testing.T) { t.Parallel() got, err := buildCommitStepMessages(buildCommitStepMessagesInput{ diff --git a/coderd/x/chatd/subagent_catalog.go b/coderd/x/chatd/subagent_catalog.go index 843f07a3fa41b..9f1186e7bf5dd 100644 --- a/coderd/x/chatd/subagent_catalog.go +++ b/coderd/x/chatd/subagent_catalog.go @@ -38,6 +38,18 @@ const ( "external or web research, parallel research, or tasks that may need edits." ) +// unbilledSubagentToolNames excludes parent-side orchestration because +// child chats bill their own runtime. Include deprecated aliases. +var unbilledSubagentToolNames = map[string]bool{ + spawnAgentToolName: true, + "wait_agent": true, + "message_agent": true, + "interrupt_agent": true, + "close_agent": true, + "list_agents": true, + listSubagentModelsToolName: true, +} + type spawnAgentArgs struct { Type string `json:"type"` Prompt string `json:"prompt"` diff --git a/coderd/x/chatd/subagent_internal_test.go b/coderd/x/chatd/subagent_internal_test.go index c51ae159cf4e1..adc42aa61e9c5 100644 --- a/coderd/x/chatd/subagent_internal_test.go +++ b/coderd/x/chatd/subagent_internal_test.go @@ -4290,6 +4290,25 @@ func TestAwaitSubagentCompletion(t *testing.T) { }) } +func TestUnbilledSubagentToolNamesMatchCatalog(t *testing.T) { + t.Parallel() + + db, ps := dbtestutil.NewDB(t) + server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{}) + ctx := chatdTestContext(t) + user, org, model := seedInternalChatDeps(t, db) + parent, _ := createParentChildChats(ctx, t, server, user, org, model) + + catalog := make(map[string]bool) + for _, tool := range server.subagentTools(ctx, func() database.Chat { return parent }, parent.LastModelConfigID) { + catalog[tool.Info().Name] = true + } + for alias := range subagentToolNameAliases { + catalog[alias] = true + } + require.Equal(t, catalog, unbilledSubagentToolNames) +} + func TestWaitAgentToolSchema(t *testing.T) { t.Parallel() diff --git a/coderd/x/chatd/tasks_test.go b/coderd/x/chatd/tasks_test.go index 9531cff927cf1..74cc4c760bd6d 100644 --- a/coderd/x/chatd/tasks_test.go +++ b/coderd/x/chatd/tasks_test.go @@ -452,8 +452,6 @@ func TestInterruptTask_PartialAssistantWithoutModelInvocationHasNoRuntime(t *tes 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)