From f83885097aa197a73fcc8944de110b469cb2065b Mon Sep 17 00:00:00 2001 From: Jaayden Halko Date: Thu, 20 Aug 2026 08:30:06 +0000 Subject: [PATCH 1/3] feat(coderd): bill completed local tool batches in agent runtime Local tool execution between model steps previously billed nothing toward chat_messages.runtime_ms, the source of hb_agent_runtime_v1. Each completed local tool batch now bills one window: the union of its billed execution intervals, so parallel calls count once and serial calls count only from their own start. Sub-agent orchestration tools never extend the window because child chats bill their own runtime. The window persists as runtime_ms on the batch's first tool-result row, keeping the role-agnostic runtime sum query correct with at most one billed row per batch. Zero windows stay NULL. Interrupted batches still bill only the model window; a follow-up bills the partial tool window on cancellation rows. --- coderd/database/querier_test.go | 22 +- coderd/usage/usagetypes/events.go | 14 +- coderd/x/chatd/attempt.go | 3 + coderd/x/chatd/chatd_test.go | 9 + coderd/x/chatd/chatloop/chatloop.go | 138 ++++- .../chatloop/chatloop_run_internal_test.go | 12 + coderd/x/chatd/chatloop/runtime_test.go | 545 ++++++++++++++++++ coderd/x/chatd/generation.go | 2 + coderd/x/chatd/message_conversion.go | 18 +- coderd/x/chatd/message_conversion_test.go | 60 +- coderd/x/chatd/subagent_catalog.go | 12 + coderd/x/chatd/subagent_internal_test.go | 19 + coderd/x/chatd/tasks_test.go | 2 - 13 files changed, 810 insertions(+), 46 deletions(-) diff --git a/coderd/database/querier_test.go b/coderd/database/querier_test.go index 32d5bf24697..39e84ae1b72 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 9004fb2514c..a990b4c6f7e 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 fb967fb2398..580d113716b 100644 --- a/coderd/x/chatd/attempt.go +++ b/coderd/x/chatd/attempt.go @@ -31,6 +31,9 @@ type stepData struct { ContextLimit sql.NullInt64 Runtime time.Duration + // BatchRuntime is the local-tool batch window. Model steps use Runtime. + BatchRuntime time.Duration + 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 7d710597e21..9bcc84127c7 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) } } + + messages := chatMessages(ctx, t, db, chat.ID) + billedToolRows := 0 + for _, msg := range messages { + if msg.Role == database.ChatMessageRoleTool && msg.RuntimeMs.Valid { + billedToolRows++ + } + } + require.LessOrEqual(t, billedToolRows, 1) }) } diff --git a/coderd/x/chatd/chatloop/chatloop.go b/coderd/x/chatd/chatloop/chatloop.go index ba6c1a7d0f3..a51a8e9f23f 100644 --- a/coderd/x/chatd/chatloop/chatloop.go +++ b/coderd/x/chatd/chatloop/chatloop.go @@ -76,11 +76,12 @@ 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 // PendingDynamicToolCalls lists tool calls that target // dynamic tools. When non-empty the chatloop exits with // ErrDynamicToolCall so the caller can execute them @@ -269,6 +270,20 @@ 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 + // OnToolStart fires when each local call begins. Serial calls may start + // after concurrent siblings settle, so interrupts bill actual starts and + // skip calls that never run. dispatchIndex identifies the dispatch-order + // occurrence. + OnToolStart func(dispatchIndex int, startedAt time.Time) + // OnToolComplete fires concurrently as each local call finishes, before + // ordered results publish. Interrupt billing uses the live timestamp; + // dispatchIndex identifies the dispatch-order occurrence. + // The callback must be concurrency-safe. + OnToolComplete func(dispatchIndex int, completedAt time.Time) + PublishMessagePart func(codersdk.ChatMessageRole, codersdk.ChatMessagePart) Logger slog.Logger Metrics *Metrics @@ -620,6 +635,17 @@ func ExecuteLocalTools(ctx context.Context, opts ExecuteLocalToolsOptions) (Pers } maxResultBytes := toolResultByteBudget(opts.ContextLimit) + batchStart := clockNow(opts.Clock) + // Keep completions by occurrence so duplicate IDs cannot collapse them. + orderedCompletions := make([]time.Time, 0, len(localCalls)) + // Keep starts by occurrence. Serial calls may begin after unbilled waits. + orderedStarts := make([]time.Time, len(localCalls)) + onToolStart := func(dispatchIndex int, startedAt time.Time) { + orderedStarts[dispatchIndex] = startedAt + if opts.OnToolStart != nil { + opts.OnToolStart(dispatchIndex, startedAt) + } + } toolResults := executeTools( ctx, opts.Clock, @@ -636,7 +662,11 @@ func ExecuteLocalTools(ctx context.Context, opts ExecuteLocalToolsOptions) (Pers opts.BuiltinToolNames, maxResultBytes, opts.ToolNameAliases, + batchStart, + onToolStart, + opts.OnToolComplete, func(tr fantasy.ToolResultContent, completedAt time.Time) { + orderedCompletions = append(orderedCompletions, completedAt) recordToolResultTimestamp(&result, tr.ToolCallID, completedAt) publishToolAttachments(ctx, opts.Logger, tr, completedAt, publishMessagePart) ssePart := chatprompt.PartFromContentWithLogger(ctx, opts.Logger, tr) @@ -653,9 +683,77 @@ func ExecuteLocalTools(ctx context.Context, opts ExecuteLocalToolsOptions) (Pers return PersistedStep{ Content: result.content, ToolResultCreatedAt: result.toolResultCreatedAt, + BatchRuntime: billableBatchDuration( + batchStart, + localCalls, + orderedStarts, + orderedCompletions, + opts.UnbilledToolNames, + ), }, nil } +// billableBatchDuration returns the union of billed execution intervals. +// Concurrent calls start at batchStart; serial calls use their recorded starts. +// Unbilled tools and gaps between billed intervals do not count. +func billableBatchDuration( + batchStart time.Time, + toolCalls []fantasy.ToolCallContent, + starts []time.Time, + completions []time.Time, + unbilledToolNames map[string]bool, +) time.Duration { + intervals := make([]BilledInterval, 0, len(toolCalls)) + for i, tc := range toolCalls { + if i >= len(completions) { + break + } + if unbilledToolNames[tc.ToolName] || completions[i].IsZero() { + continue + } + start := batchStart + if i < len(starts) && !starts[i].IsZero() { + start = starts[i] + } + intervals = append(intervals, BilledInterval{Start: start, End: completions[i]}) + } + return BilledIntervalsDuration(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 +1179,9 @@ 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. +// executeTools runs non-serial calls concurrently, then SerialToolCalls in +// call order. Results publish in original order after all tools finish; +// onComplete fires as each tool finishes. func executeTools( ctx context.Context, clock quartz.Clock, @@ -1100,6 +1197,9 @@ func executeTools( builtinToolNames map[string]bool, maxResultBytes int, toolNameAliases map[string]string, + batchStart time.Time, + onStart func(dispatchIndex int, startedAt time.Time), + onComplete func(dispatchIndex int, completedAt time.Time), onResult func(fantasy.ToolResultContent, time.Time), ) []fantasy.ToolResultContent { if len(toolCalls) == 0 { @@ -1167,6 +1267,9 @@ func executeTools( // Captured per call so parallel tools get // accurate individual completion times. completedAt[i] = clockNow(clock) + if onComplete != nil { + onComplete(i, completedAt[i]) + } }() results[i] = executeSingleTool( ctx, @@ -1185,12 +1288,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 +1297,9 @@ func executeTools( serialIndexes = append(serialIndexes, i) continue } + if onStart != nil { + onStart(i, batchStart) + } wg.Add(1) go func() { defer wg.Done() @@ -1206,9 +1308,7 @@ 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. + // Reconcile concurrent results before serial tools inspect shared state. settled := make([]fantasy.ToolResultContent, 0, len(results)) for i := range results { if !slices.Contains(serialIndexes, i) { @@ -1218,6 +1318,10 @@ func executeTools( notifyStepToolResultObservers(toolMap, toolNameAliases, localToolCalls, observed, settled) for _, i := range serialIndexes { + // Stamp serial calls at launch, not batch start. + if onStart != nil { + onStart(i, clockNow(clock)) + } runCall(i, localToolCalls[i]) } diff --git a/coderd/x/chatd/chatloop/chatloop_run_internal_test.go b/coderd/x/chatd/chatloop/chatloop_run_internal_test.go index 7d35d2588e7..44c2a68db52 100644 --- a/coderd/x/chatd/chatloop/chatloop_run_internal_test.go +++ b/coderd/x/chatd/chatloop/chatloop_run_internal_test.go @@ -949,6 +949,9 @@ func TestExecuteToolsNotifiesStepToolCallObservers(t *testing.T) { map[string]bool{}, defaultToolResultBytes, map[string]string{"observer_alias": "observer_tool"}, + time.Time{}, + nil, + nil, nil, ) @@ -1020,6 +1023,9 @@ func TestExecuteToolsNotifiesStepToolResultObservers(t *testing.T) { map[string]bool{}, defaultToolResultBytes, map[string]string{"observer_alias": "observer_tool"}, + time.Time{}, + nil, + nil, nil, ) @@ -1097,6 +1103,9 @@ func TestExecuteToolsReconcilesResultsBeforeSerialCalls(t *testing.T) { map[string]bool{}, defaultToolResultBytes, nil, + time.Time{}, + nil, + nil, nil, ) @@ -1164,6 +1173,9 @@ func TestExecuteToolsSerialToolCallOrder(t *testing.T) { map[string]bool{}, defaultToolResultBytes, nil, + time.Time{}, + nil, + nil, nil, ) diff --git a/coderd/x/chatd/chatloop/runtime_test.go b/coderd/x/chatd/chatloop/runtime_test.go index b227b26b694..005054d5569 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,546 @@ 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) +} + +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) +} + +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) +} + +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) +} + +func TestExecuteLocalTools_ExecutionCallbacksFireOnlyForRuns(t *testing.T) { + t.Parallel() + + t.Run("started call records a paired lifecycle", func(t *testing.T) { + t.Parallel() + + starts := 0 + completions := 0 + startedWhenToolRan := false + tool := fantasy.NewAgentTool( + "fast_tool", + "test tool", + func(context.Context, struct{}, fantasy.ToolCall) (fantasy.ToolResponse, error) { + startedWhenToolRan = 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"}, + OnToolStart: func(int, time.Time) { + starts++ + }, + OnToolComplete: func(int, time.Time) { + completions++ + }, + ToolCalls: []fantasy.ToolCallContent{ + {ToolCallID: "call-1", ToolName: "fast_tool", Input: "{}"}, + }, + }) + require.NoError(t, err) + require.Len(t, outcome.Content, 1) + require.Equal(t, 1, starts) + require.Equal(t, 1, completions) + require.True(t, startedWhenToolRan, "the start callback must fire before the tool runs") + }) + + t.Run("canceled context records no lifecycle", func(t *testing.T) { + t.Parallel() + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + started := false + completed := false + _, err := chatloop.ExecuteLocalTools(ctx, chatloop.ExecuteLocalToolsOptions{ + Clock: quartz.NewMock(t), + OnToolStart: func(int, time.Time) { started = true }, + OnToolComplete: func(int, time.Time) { completed = true }, + ToolCalls: []fantasy.ToolCallContent{ + {ToolCallID: "call-1", ToolName: "fast_tool", Input: "{}"}, + }, + }) + require.ErrorIs(t, err, context.Canceled) + require.False(t, started) + require.False(t, completed) + }) + + t.Run("exclusive violation records no lifecycle", func(t *testing.T) { + t.Parallel() + + started := false + completed := false + outcome, err := chatloop.ExecuteLocalTools(context.Background(), chatloop.ExecuteLocalToolsOptions{ + Clock: quartz.NewMock(t), + ExclusiveToolNames: map[string]bool{"exclusive_tool": true}, + OnToolStart: func(int, time.Time) { started = true }, + OnToolComplete: func(int, time.Time) { completed = true }, + 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.False(t, started) + require.False(t, completed) + }) +} + +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_OnToolCompleteReportsLiveCompletions(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitShort) + clock := quartz.NewMock(t) + trap := clock.Trap().Now() + defer trap.Close() + + type completion struct { + dispatchIndex int + completedAt time.Time + } + completionCh := make(chan completion, 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"}, + OnToolComplete: func(dispatchIndex int, completedAt time.Time) { + completionCh <- completion{dispatchIndex: dispatchIndex, completedAt: completedAt} + }, + 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, completionCh) + require.Equal(t, 0, fast.dispatchIndex) + clock.Advance(50 * time.Second) + close(slowGo) + trap.MustWait(ctx).MustRelease(ctx) + slow := testutil.RequireReceive(ctx, t, completionCh) + 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 a3a2aeb2e96..f3b15a54a94 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,7 @@ func stepDataFromPersisted(step chatloop.PersistedStep) stepData { Usage: step.Usage, ContextLimit: step.ContextLimit, Runtime: step.Runtime, + BatchRuntime: step.BatchRuntime, 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 52d749175bd..a27229b54f1 100644 --- a/coderd/x/chatd/message_conversion.go +++ b/coderd/x/chatd/message_conversion.go @@ -61,7 +61,7 @@ func buildCommitStepMessages(input buildCommitStepMessagesInput) (stepMessagesFo messages = append(messages, assistantMessage(input.modelConfigID, contentVersion, assistantContent, input.step)) } - for _, toolResult := range toolResults { + for i, toolResult := range toolResults { part := chatprompt.PartFromContentWithLogger(context.Background(), input.logger, toolResult) applyToolMetadata(&part, input.toolNameToConfigID) if part.ToolCallID != "" && input.step.ToolResultCreatedAt != nil { @@ -73,7 +73,12 @@ func buildCommitStepMessages(input buildCommitStepMessagesInput) (stepMessagesFo if err != nil { return stepMessagesForCommit{}, xerrors.Errorf("marshal tool result: %w", err) } - messages = append(messages, baseMessage(database.ChatMessageRoleTool, database.ChatMessageVisibilityBoth, input.modelConfigID, contentVersion, content)) + msg := baseMessage(database.ChatMessageRoleTool, database.ChatMessageVisibilityBoth, input.modelConfigID, contentVersion, content) + // Usage sums runtime_ms across rows, so store the batch once. + if i == 0 { + msg.RuntimeMs = nullInt64IfNonZero(input.step.BatchRuntime.Milliseconds()) + } + messages = append(messages, msg) } return stepMessagesForCommit{ @@ -631,13 +636,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 c84935d72f3..ef0852fcb02 100644 --- a/coderd/x/chatd/message_conversion_test.go +++ b/coderd/x/chatd/message_conversion_test.go @@ -100,9 +100,63 @@ 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_BatchRuntimeLandsOnFirstToolRow(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, + }, + }) + require.NoError(t, err) + require.Len(t, got.Messages, 2) + require.Equal(t, sql.NullInt64{Int64: 10000, Valid: true}, got.Messages[0].RuntimeMs) + require.False(t, got.Messages[1].RuntimeMs.Valid) +} + +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, + }, + }) + 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.Equal(t, database.ChatMessageRoleTool, got.Messages[1].Role) + require.Equal(t, sql.NullInt64{Int64: 3000, Valid: true}, got.Messages[1].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 843f07a3fa4..9f1186e7bf5 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 c51ae159cf4..adc42aa61e9 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 9531cff927c..74cc4c760bd 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) From 0e70d5679d0647c83ef43912ae33ad0c39b6ac46 Mon Sep 17 00:00:00 2001 From: Jaayden Halko Date: Mon, 24 Aug 2026 09:15:13 +0000 Subject: [PATCH 2/3] chore: updates for PR comments --- coderd/x/chatd/chatloop/chatloop.go | 148 ++++++++---------- .../chatloop/chatloop_run_internal_test.go | 124 +++++++++++++-- coderd/x/chatd/chatloop/runtime_test.go | 93 ++++++----- 3 files changed, 230 insertions(+), 135 deletions(-) diff --git a/coderd/x/chatd/chatloop/chatloop.go b/coderd/x/chatd/chatloop/chatloop.go index a51a8e9f23f..7051ee42b3f 100644 --- a/coderd/x/chatd/chatloop/chatloop.go +++ b/coderd/x/chatd/chatloop/chatloop.go @@ -273,16 +273,11 @@ type ExecuteLocalToolsOptions struct { // UnbilledToolNames lists called tool names excluded from the batch // window. Include deprecated aliases. UnbilledToolNames map[string]bool - // OnToolStart fires when each local call begins. Serial calls may start - // after concurrent siblings settle, so interrupts bill actual starts and - // skip calls that never run. dispatchIndex identifies the dispatch-order - // occurrence. - OnToolStart func(dispatchIndex int, startedAt time.Time) - // OnToolComplete fires concurrently as each local call finishes, before - // ordered results publish. Interrupt billing uses the live timestamp; - // dispatchIndex identifies the dispatch-order occurrence. - // The callback must be concurrency-safe. - OnToolComplete func(dispatchIndex int, completedAt time.Time) + // 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 @@ -290,6 +285,17 @@ type ExecuteLocalToolsOptions struct { 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 @@ -636,17 +642,7 @@ func ExecuteLocalTools(ctx context.Context, opts ExecuteLocalToolsOptions) (Pers maxResultBytes := toolResultByteBudget(opts.ContextLimit) batchStart := clockNow(opts.Clock) - // Keep completions by occurrence so duplicate IDs cannot collapse them. - orderedCompletions := make([]time.Time, 0, len(localCalls)) - // Keep starts by occurrence. Serial calls may begin after unbilled waits. - orderedStarts := make([]time.Time, len(localCalls)) - onToolStart := func(dispatchIndex int, startedAt time.Time) { - orderedStarts[dispatchIndex] = startedAt - if opts.OnToolStart != nil { - opts.OnToolStart(dispatchIndex, startedAt) - } - } - toolResults := executeTools( + toolExecutions := executeTools( ctx, opts.Clock, opts.Tools, @@ -663,59 +659,42 @@ func ExecuteLocalTools(ctx context.Context, opts ExecuteLocalToolsOptions) (Pers maxResultBytes, opts.ToolNameAliases, batchStart, - onToolStart, - opts.OnToolComplete, - func(tr fantasy.ToolResultContent, completedAt time.Time) { - orderedCompletions = append(orderedCompletions, completedAt) - 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) - }, + 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) - } return PersistedStep{ Content: result.content, ToolResultCreatedAt: result.toolResultCreatedAt, - BatchRuntime: billableBatchDuration( - batchStart, - localCalls, - orderedStarts, - orderedCompletions, - opts.UnbilledToolNames, - ), + BatchRuntime: billableBatchDuration(toolExecutions, opts.UnbilledToolNames), }, nil } // billableBatchDuration returns the union of billed execution intervals. -// Concurrent calls start at batchStart; serial calls use their recorded starts. // Unbilled tools and gaps between billed intervals do not count. func billableBatchDuration( - batchStart time.Time, - toolCalls []fantasy.ToolCallContent, - starts []time.Time, - completions []time.Time, + executions []toolExecutionResult, unbilledToolNames map[string]bool, ) time.Duration { - intervals := make([]BilledInterval, 0, len(toolCalls)) - for i, tc := range toolCalls { - if i >= len(completions) { - break - } - if unbilledToolNames[tc.ToolName] || completions[i].IsZero() { + intervals := make([]BilledInterval, 0, len(executions)) + for _, execution := range executions { + if unbilledToolNames[execution.content.ToolName] || + execution.interval.Start.IsZero() || + execution.interval.End.IsZero() { continue } - start := batchStart - if i < len(starts) && !starts[i].IsZero() { - start = starts[i] - } - intervals = append(intervals, BilledInterval{Start: start, End: completions[i]}) + intervals = append(intervals, execution.interval) } return BilledIntervalsDuration(intervals) } @@ -1179,9 +1158,14 @@ func processStepStream( return result, nil } +type toolExecutionResult struct { + content fantasy.ToolResultContent + interval BilledInterval +} + // executeTools runs non-serial calls concurrently, then SerialToolCalls in -// call order. Results publish in original order after all tools finish; -// onComplete fires as each tool finishes. +// 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, @@ -1198,10 +1182,8 @@ func executeTools( maxResultBytes int, toolNameAliases map[string]string, batchStart time.Time, - onStart func(dispatchIndex int, startedAt time.Time), - onComplete func(dispatchIndex int, completedAt time.Time), - onResult func(fantasy.ToolResultContent, time.Time), -) []fantasy.ToolResultContent { + recorder ToolBillingRecorder, +) []toolExecutionResult { if len(toolCalls) == 0 { return nil } @@ -1250,12 +1232,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{ @@ -1266,12 +1247,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) - if onComplete != nil { - onComplete(i, completedAt[i]) + 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, @@ -1297,8 +1279,9 @@ func executeTools( serialIndexes = append(serialIndexes, i) continue } - if onStart != nil { - onStart(i, batchStart) + executions[i].interval.Start = batchStart + if recorder != nil { + recorder.RecordStart(i, batchStart) } wg.Add(1) go func() { @@ -1309,30 +1292,25 @@ func executeTools( wg.Wait() // Reconcile concurrent results before serial tools inspect shared state. - settled := make([]fantasy.ToolResultContent, 0, len(results)) - for i := range results { + 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. - if onStart != nil { - onStart(i, clockNow(clock)) + 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 44c2a68db52..393d7acbaeb 100644 --- a/coderd/x/chatd/chatloop/chatloop_run_internal_test.go +++ b/coderd/x/chatd/chatloop/chatloop_run_internal_test.go @@ -951,8 +951,6 @@ func TestExecuteToolsNotifiesStepToolCallObservers(t *testing.T) { map[string]string{"observer_alias": "observer_tool"}, time.Time{}, nil, - nil, - nil, ) require.Equal(t, []string{"observer_tool", "other_tool", "denied_tool"}, observedNames, @@ -1025,8 +1023,6 @@ func TestExecuteToolsNotifiesStepToolResultObservers(t *testing.T) { map[string]string{"observer_alias": "observer_tool"}, time.Time{}, nil, - nil, - nil, ) require.Equal(t, 1, notifications, "each called observer is notified once per step") @@ -1105,15 +1101,13 @@ func TestExecuteToolsReconcilesResultsBeforeSerialCalls(t *testing.T) { nil, time.Time{}, nil, - nil, - nil, ) require.True(t, notified) 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) { @@ -1175,8 +1169,6 @@ func TestExecuteToolsSerialToolCallOrder(t *testing.T) { nil, time.Time{}, nil, - nil, - nil, ) require.Equal(t, []string{"a:start", "a:end", "b:start", "b:end", "c:start", "c:end"}, events, @@ -1189,8 +1181,120 @@ 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) { diff --git a/coderd/x/chatd/chatloop/runtime_test.go b/coderd/x/chatd/chatloop/runtime_test.go index 005054d5569..2c4a31b4978 100644 --- a/coderd/x/chatd/chatloop/runtime_test.go +++ b/coderd/x/chatd/chatloop/runtime_test.go @@ -400,42 +400,63 @@ func TestExecuteLocalTools_DuplicateToolCallIDsKeepOccurrenceCompletions(t *test require.Equal(t, 60*time.Second, outcome.BatchRuntime) } -func TestExecuteLocalTools_ExecutionCallbacksFireOnlyForRuns(t *testing.T) { +// 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() - starts := 0 - completions := 0 + recorder := &recordingToolBillingRecorder{} startedWhenToolRan := false tool := fantasy.NewAgentTool( "fast_tool", "test tool", func(context.Context, struct{}, fantasy.ToolCall) (fantasy.ToolResponse, error) { - startedWhenToolRan = starts > 0 + 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"}, - OnToolStart: func(int, time.Time) { - starts++ - }, - OnToolComplete: func(int, time.Time) { - completions++ - }, + 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, starts) - require.Equal(t, 1, completions) - require.True(t, startedWhenToolRan, "the start callback must fire before the tool runs") + 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) { @@ -443,31 +464,27 @@ func TestExecuteLocalTools_ExecutionCallbacksFireOnlyForRuns(t *testing.T) { ctx, cancel := context.WithCancel(context.Background()) cancel() - started := false - completed := false + recorder := &recordingToolBillingRecorder{} _, err := chatloop.ExecuteLocalTools(ctx, chatloop.ExecuteLocalToolsOptions{ - Clock: quartz.NewMock(t), - OnToolStart: func(int, time.Time) { started = true }, - OnToolComplete: func(int, time.Time) { completed = true }, + Clock: quartz.NewMock(t), + BillingRecorder: recorder, ToolCalls: []fantasy.ToolCallContent{ {ToolCallID: "call-1", ToolName: "fast_tool", Input: "{}"}, }, }) require.ErrorIs(t, err, context.Canceled) - require.False(t, started) - require.False(t, completed) + require.Zero(t, recorder.starts) + require.Zero(t, recorder.completions) }) t.Run("exclusive violation records no lifecycle", func(t *testing.T) { t.Parallel() - started := false - completed := false + recorder := &recordingToolBillingRecorder{} outcome, err := chatloop.ExecuteLocalTools(context.Background(), chatloop.ExecuteLocalToolsOptions{ Clock: quartz.NewMock(t), ExclusiveToolNames: map[string]bool{"exclusive_tool": true}, - OnToolStart: func(int, time.Time) { started = true }, - OnToolComplete: func(int, time.Time) { completed = true }, + BillingRecorder: recorder, ToolCalls: []fantasy.ToolCallContent{ {ToolCallID: "call-1", ToolName: "exclusive_tool", Input: "{}"}, {ToolCallID: "call-2", ToolName: "fast_tool", Input: "{}"}, @@ -475,8 +492,8 @@ func TestExecuteLocalTools_ExecutionCallbacksFireOnlyForRuns(t *testing.T) { }) require.NoError(t, err) require.Len(t, outcome.Content, 2, "the whole batch resolves to synthesized policy errors") - require.False(t, started) - require.False(t, completed) + require.Zero(t, recorder.starts) + require.Zero(t, recorder.completions) }) } @@ -514,7 +531,7 @@ func TestExecuteLocalTools_EmptyToolCallIDStillBillsWindow(t *testing.T) { require.Equal(t, 60*time.Second, outcome.BatchRuntime) } -func TestExecuteLocalTools_OnToolCompleteReportsLiveCompletions(t *testing.T) { +func TestExecuteLocalTools_BillingRecorderReportsLiveCompletions(t *testing.T) { t.Parallel() ctx := testutil.Context(t, testutil.WaitShort) @@ -522,11 +539,9 @@ func TestExecuteLocalTools_OnToolCompleteReportsLiveCompletions(t *testing.T) { trap := clock.Trap().Now() defer trap.Close() - type completion struct { - dispatchIndex int - completedAt time.Time + recorder := &recordingToolBillingRecorder{ + completeCh: make(chan recordedToolCompletion, 2), } - completionCh := make(chan completion, 2) fastGo := make(chan struct{}) slowGo := make(chan struct{}) resultCh := executeToolBatch(t, clock, chatloop.ExecuteLocalToolsOptions{ @@ -534,10 +549,8 @@ func TestExecuteLocalTools_OnToolCompleteReportsLiveCompletions(t *testing.T) { blockingTool("fast_tool", fastGo, fantasy.NewTextResponse("done")), blockingTool("slow_tool", slowGo, fantasy.NewTextResponse("done")), }, - ActiveTools: []string{"fast_tool", "slow_tool"}, - OnToolComplete: func(dispatchIndex int, completedAt time.Time) { - completionCh <- completion{dispatchIndex: dispatchIndex, completedAt: completedAt} - }, + 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: "{}"}, @@ -548,12 +561,12 @@ func TestExecuteLocalTools_OnToolCompleteReportsLiveCompletions(t *testing.T) { clock.Advance(10 * time.Second) close(fastGo) trap.MustWait(ctx).MustRelease(ctx) - fast := testutil.RequireReceive(ctx, t, completionCh) + 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, completionCh) + slow := testutil.RequireReceive(ctx, t, recorder.completeCh) require.Equal(t, 1, slow.dispatchIndex) require.Equal(t, 50*time.Second, slow.completedAt.Sub(fast.completedAt)) From c93e039ad3b6fd1b5151069e35d2dd69b5349ced Mon Sep 17 00:00:00 2001 From: Jaayden Halko Date: Mon, 24 Aug 2026 15:07:20 +0000 Subject: [PATCH 3/3] feat(coderd/x/chatd): bill local tool batches on a dedicated usage record Usage sums runtime_ms across rows, so the batch window previously landed on the first tool-result row of the batch. Replace that arbitrary member stamping with one model-invisible tool-role row per billed batch carrying the runtime and an audit payload (billed window and interval count). Real tool results no longer carry batch-level runtime. --- coderd/x/chatd/attempt.go | 3 + coderd/x/chatd/chatd_test.go | 8 +-- coderd/x/chatd/chatloop/chatloop.go | 17 ++++-- coderd/x/chatd/chatloop/runtime_test.go | 3 + coderd/x/chatd/generation.go | 1 + coderd/x/chatd/message_conversion.go | 67 ++++++++++++++++++++--- coderd/x/chatd/message_conversion_test.go | 36 +++++++++--- 7 files changed, 111 insertions(+), 24 deletions(-) diff --git a/coderd/x/chatd/attempt.go b/coderd/x/chatd/attempt.go index 580d113716b..746df38fc4d 100644 --- a/coderd/x/chatd/attempt.go +++ b/coderd/x/chatd/attempt.go @@ -33,6 +33,9 @@ type stepData struct { // 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 diff --git a/coderd/x/chatd/chatd_test.go b/coderd/x/chatd/chatd_test.go index 9bcc84127c7..27de4c9fce6 100644 --- a/coderd/x/chatd/chatd_test.go +++ b/coderd/x/chatd/chatd_test.go @@ -6592,14 +6592,14 @@ func TestActiveServer_ToolExecutionAndPolicy(t *testing.T) { } } + // 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) - billedToolRows := 0 for _, msg := range messages { - if msg.Role == database.ChatMessageRoleTool && msg.RuntimeMs.Valid { - billedToolRows++ + if msg.Role == database.ChatMessageRoleTool { + require.False(t, msg.RuntimeMs.Valid) } } - require.LessOrEqual(t, billedToolRows, 1) }) } diff --git a/coderd/x/chatd/chatloop/chatloop.go b/coderd/x/chatd/chatloop/chatloop.go index 7051ee42b3f..f22cbc8d64b 100644 --- a/coderd/x/chatd/chatloop/chatloop.go +++ b/coderd/x/chatd/chatloop/chatloop.go @@ -82,6 +82,9 @@ type PersistedStep struct { // 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 @@ -674,19 +677,21 @@ func ExecuteLocalTools(ctx context.Context, opts ExecuteLocalToolsOptions) (Pers if ctx.Err() != nil { return PersistedStep{}, ctx.Err() } + billedIntervals := billableBatchIntervals(toolExecutions, opts.UnbilledToolNames) return PersistedStep{ Content: result.content, ToolResultCreatedAt: result.toolResultCreatedAt, - BatchRuntime: billableBatchDuration(toolExecutions, opts.UnbilledToolNames), + BatchRuntime: BilledIntervalsDuration(billedIntervals), + BatchBilledCalls: len(billedIntervals), }, nil } -// billableBatchDuration returns the union of billed execution intervals. -// Unbilled tools and gaps between billed intervals do not count. -func billableBatchDuration( +// billableBatchIntervals returns the billed execution intervals. +// Unbilled tools and calls without both stamps do not count. +func billableBatchIntervals( executions []toolExecutionResult, unbilledToolNames map[string]bool, -) time.Duration { +) []BilledInterval { intervals := make([]BilledInterval, 0, len(executions)) for _, execution := range executions { if unbilledToolNames[execution.content.ToolName] || @@ -696,7 +701,7 @@ func billableBatchDuration( } intervals = append(intervals, execution.interval) } - return BilledIntervalsDuration(intervals) + return intervals } // BilledInterval is one billed tool call's execution window. diff --git a/coderd/x/chatd/chatloop/runtime_test.go b/coderd/x/chatd/chatloop/runtime_test.go index 2c4a31b4978..62211b13f92 100644 --- a/coderd/x/chatd/chatloop/runtime_test.go +++ b/coderd/x/chatd/chatloop/runtime_test.go @@ -229,6 +229,7 @@ func TestExecuteLocalTools_BatchWindowIsMaxNotSum(t *testing.T) { 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) { @@ -325,6 +326,7 @@ func TestExecuteLocalTools_UnbilledOnlyBatchBillsNothing(t *testing.T) { outcome := testutil.RequireReceive(ctx, t, resultCh) require.Zero(t, outcome.BatchRuntime) + require.Zero(t, outcome.BatchBilledCalls) } func TestExecuteLocalTools_AliasNamesClassifyAsCalled(t *testing.T) { @@ -364,6 +366,7 @@ func TestExecuteLocalTools_AliasNamesClassifyAsCalled(t *testing.T) { 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) { diff --git a/coderd/x/chatd/generation.go b/coderd/x/chatd/generation.go index f3b15a54a94..7c948383244 100644 --- a/coderd/x/chatd/generation.go +++ b/coderd/x/chatd/generation.go @@ -1552,6 +1552,7 @@ func stepDataFromPersisted(step chatloop.PersistedStep) stepData { 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 a27229b54f1..c5fc4be8b89 100644 --- a/coderd/x/chatd/message_conversion.go +++ b/coderd/x/chatd/message_conversion.go @@ -61,7 +61,7 @@ func buildCommitStepMessages(input buildCommitStepMessagesInput) (stepMessagesFo messages = append(messages, assistantMessage(input.modelConfigID, contentVersion, assistantContent, input.step)) } - for i, toolResult := range toolResults { + for _, toolResult := range toolResults { part := chatprompt.PartFromContentWithLogger(context.Background(), input.logger, toolResult) applyToolMetadata(&part, input.toolNameToConfigID) if part.ToolCallID != "" && input.step.ToolResultCreatedAt != nil { @@ -73,12 +73,17 @@ func buildCommitStepMessages(input buildCommitStepMessagesInput) (stepMessagesFo if err != nil { return stepMessagesForCommit{}, xerrors.Errorf("marshal tool result: %w", err) } - msg := baseMessage(database.ChatMessageRoleTool, database.ChatMessageVisibilityBoth, input.modelConfigID, contentVersion, content) - // Usage sums runtime_ms across rows, so store the batch once. - if i == 0 { - msg.RuntimeMs = nullInt64IfNonZero(input.step.BatchRuntime.Milliseconds()) - } - messages = append(messages, msg) + 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{ @@ -230,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 { diff --git a/coderd/x/chatd/message_conversion_test.go b/coderd/x/chatd/message_conversion_test.go index ef0852fcb02..d3258a43faf 100644 --- a/coderd/x/chatd/message_conversion_test.go +++ b/coderd/x/chatd/message_conversion_test.go @@ -100,7 +100,7 @@ func TestBuildCommitStepMessages_LocalToolResultsBecomeToolMessages(t *testing.T require.JSONEq(t, `{"stdout":"/tmp"}`, string(toolParts[0].Result)) } -func TestBuildCommitStepMessages_BatchRuntimeLandsOnFirstToolRow(t *testing.T) { +func TestBuildCommitStepMessages_BatchRuntimeBillsDedicatedUsageRow(t *testing.T) { t.Parallel() got, err := buildCommitStepMessages(buildCommitStepMessagesInput{ @@ -120,13 +120,32 @@ func TestBuildCommitStepMessages_BatchRuntimeLandsOnFirstToolRow(t *testing.T) { Result: fantasy.ToolResultOutputContentText{Text: `{"stdout":"/tmp"}`}, }, }, - BatchRuntime: 10 * time.Second, + BatchRuntime: 10 * time.Second, + BatchBilledCalls: 2, }, }) require.NoError(t, err) - require.Len(t, got.Messages, 2) - require.Equal(t, sql.NullInt64{Int64: 10000, Valid: true}, got.Messages[0].RuntimeMs) + 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) { @@ -145,15 +164,18 @@ func TestBuildCommitStepMessages_BatchAttachmentAssistantRowStaysNull(t *testing ClientMetadata: `{"attachments":[{"file_id":"` + uuid.NewString() + `","media_type":"image/png","name":"shot.png"}]}`, }, }, - BatchRuntime: 3 * time.Second, + BatchRuntime: 3 * time.Second, + BatchBilledCalls: 1, }, }) require.NoError(t, err) - require.Len(t, got.Messages, 2) + 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.Equal(t, sql.NullInt64{Int64: 3000, Valid: true}, got.Messages[1].RuntimeMs) + 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) {