From bb26857eb39b8b1e1fce7169937b243ac51c92b9 Mon Sep 17 00:00:00 2001 From: Jaayden Halko Date: Fri, 14 Aug 2026 14:47:56 +0000 Subject: [PATCH 01/20] feat: bill local tool execution time in chat agent runtime Local tool batches now persist their billable wall-clock window to chat_messages.runtime_ms, the source of truth for hb_agent_runtime_v1. Each batch bills one window, from the batch start to the last billed tool's completion, stored on that tool's message row. Tools in a batch run in parallel, so the window equals the union of billed execution intervals rather than a per-tool sum. Sub-agent orchestration tools (spawn_agent, wait_agent, message_agent, interrupt_agent, list_agents, list_subagent_models, and the deprecated close_agent alias) never extend the window: every chat, including children, bills its own runtime, so a parent's wait_agent would double count. Client-executed dynamic tools, external agents, parked time, and retry backoff remain unbilled. Interrupted batches bill the partial window: the buffer episode stamps the batch start, and the interrupt task places the window on the window-defining synthesized cancellation row. Batches without a live attempt (crash recovery, state promotion) bill nothing. The usage query already sums runtime_ms role-agnostically, so no schema, query, or cron changes are needed. Reported agent runtime increases from deploy forward. --- coderd/database/querier_test.go | 25 +- coderd/usage/usagetypes/events.go | 14 +- coderd/x/chatd/ARCHITECTURE.md | 2 + coderd/x/chatd/attempt.go | 7 + coderd/x/chatd/chatd_test.go | 12 + coderd/x/chatd/chatloop/chatloop.go | 79 +++++- coderd/x/chatd/chatloop/runtime_test.go | 232 ++++++++++++++++++ coderd/x/chatd/generation.go | 17 +- coderd/x/chatd/message_conversion.go | 18 +- coderd/x/chatd/message_conversion_test.go | 71 +++++- .../messagepartbuffer/message_part_buffer.go | 49 +++- .../message_part_buffer_test.go | 37 +++ coderd/x/chatd/subagent_catalog.go | 19 ++ coderd/x/chatd/subagent_internal_test.go | 22 ++ coderd/x/chatd/tasks.go | 78 +++++- coderd/x/chatd/tasks_test.go | 196 +++++++++++++++ docs/ai-coder/usage-data-reporting.md | 33 ++- 17 files changed, 865 insertions(+), 46 deletions(-) diff --git a/coderd/database/querier_test.go b/coderd/database/querier_test.go index 44fc384b619..059920be5b2 100644 --- a/coderd/database/querier_test.go +++ b/coderd/database/querier_test.go @@ -10996,13 +10996,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) @@ -11010,23 +11010,26 @@ 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) + // chats), soft-deleted, just before the exclusive end boundary, and a + // tool-role row carrying a local tool batch window (the sum is + // role-agnostic). + 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) + 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..c4b67dad9eb 100644 --- a/coderd/usage/usagetypes/events.go +++ b/coderd/usage/usagetypes/events.go @@ -204,10 +204,16 @@ 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. +// (chats) in one UTC hour. Two kinds of windows are measured. Model steps +// span model streaming (including provider-executed tools) and stream +// retries, ending when the model stream finishes. Local tool batches span +// the start of the batch until the last billed tool completes; tools in a +// batch run in parallel, so the batch bills one window rather than a sum. +// Sub-agent orchestration tools (spawn_agent, wait_agent, and the rest of +// that category) are excluded because each sub-agent chat bills its own +// runtime. Also excluded: client-executed (dynamic) tools and external +// agents, time parked waiting for user action, idle time between turns, +// and retry backoff between attempts. // // 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/ARCHITECTURE.md b/coderd/x/chatd/ARCHITECTURE.md index bef7b57116b..ec9c8187c20 100644 --- a/coderd/x/chatd/ARCHITECTURE.md +++ b/coderd/x/chatd/ARCHITECTURE.md @@ -725,6 +725,8 @@ The buffer exposes the following API: - `GetParts(chat_id, history_version, generation_attempt)`: returns the message parts for an episode. Returns a predefined error if the episode is not found. - `StartModelInvocation(chat_id, history_version, generation_attempt)`: stamps the instant the episode opens its provider stream. Returns a predefined error if the episode is not found or already closed. Episodes that never invoke a model, such as local tool execution batches, are never stamped. - `ModelInvokedAt(chat_id, history_version, generation_attempt)`: returns the instant stamped by `StartModelInvocation`, or the zero time when the episode is unknown or never opened a provider stream. It must be read before `CloseEpisode`, because closed episodes are garbage collected and reading afterwards races the cleanup loop. The interrupt goroutine reads it just before closing the episode and uses the span between that instant and the interrupt as the interrupted attempt's billable runtime. +- `StartToolBatch(chat_id, history_version, generation_attempt)`: stamps the instant the episode begins executing its local tool batch. Returns a predefined error if the episode is not found or already closed. Episodes that never execute local tools, such as model invocations without tool calls, are never stamped. +- `ToolBatchStartedAt(chat_id, history_version, generation_attempt)`: returns the instant stamped by `StartToolBatch`, or the zero time when the episode is unknown or never started a tool batch. Like `ModelInvokedAt`, it must be read before `CloseEpisode`. The interrupt goroutine reads it just before closing the episode and uses it as the start of the interrupted tool batch's billable window. - `SubscribeToEpisode(chat_id, history_version, generation_attempt)`: returns a go channel that will receive all message parts for the episode. It spawns a goroutine that delivers parts to the channel. It's live until the episode is closed or until a subscriber requests that the channel be closed. Once the goroutine delivers all message parts for a closed episode, it closes the channel and exits. If the episode is already closed at the time of the call, the goroutine delivers all message parts for the episode, closes the channel, and exits. `SubscribeToEpisode` does not return an error if the episode is not found: it waits for it to be created instead. Closed episodes are garbage collected after at least 15 seconds since they were closed and when they have no active subscribers. The message part buffer maintains a garbage collection goroutine. diff --git a/coderd/x/chatd/attempt.go b/coderd/x/chatd/attempt.go index fb967fb2398..28de29971c1 100644 --- a/coderd/x/chatd/attempt.go +++ b/coderd/x/chatd/attempt.go @@ -31,6 +31,13 @@ type stepData struct { ContextLimit sql.NullInt64 Runtime time.Duration + // BatchRuntime is the billable window of a local tool batch, + // persisted as runtime_ms on the tool message row identified by + // BatchRuntimeToolCallID. Zero for model-invocation steps, whose + // billable window is Runtime on the assistant row instead. + BatchRuntime time.Duration + BatchRuntimeToolCallID string + 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 53e4a4fd087..77d9af09f4c 100644 --- a/coderd/x/chatd/chatd_test.go +++ b/coderd/x/chatd/chatd_test.go @@ -6433,6 +6433,18 @@ func TestActiveServer_ToolExecutionAndPolicy(t *testing.T) { require.False(t, result.ProviderExecuted) } } + + // The parallel batch bills at most one window: whatever wall + // time elapsed, only the window-defining tool row may carry + // runtime_ms, never one per parallel call. + 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 89063ff0612..00ef8d4589e 100644 --- a/coderd/x/chatd/chatloop/chatloop.go +++ b/coderd/x/chatd/chatloop/chatloop.go @@ -76,7 +76,8 @@ type PersistedStep struct { // Runtime is the wall-clock duration of the model invocation // that produced this step's content, measured from just before // the provider stream is opened until the stream is fully - // consumed. + // consumed. Local tool batches report their billable window + // through ToolExecutionOutcome.BatchRuntime instead. Runtime time.Duration // PendingDynamicToolCalls lists tool calls that target // dynamic tools. When non-empty the chatloop exits with @@ -259,6 +260,12 @@ type ExecuteLocalToolsOptions struct { // is renamed but old chat histories still reference the old name. ToolNameAliases map[string]string + // UnbilledToolNames lists tool names whose execution never extends + // the batch's billable runtime window. Classification uses the + // name as called, so deprecated aliases must be listed alongside + // their canonical names. + UnbilledToolNames map[string]bool + PublishMessagePart func(codersdk.ChatMessageRole, codersdk.ChatMessagePart) Logger slog.Logger Metrics *Metrics @@ -268,6 +275,18 @@ type ExecuteLocalToolsOptions struct { // ToolExecutionOutcome is the durable tool-result content from one batch. type ToolExecutionOutcome struct { Step PersistedStep + // BatchRuntime is the billable local-tool counterpart of + // PersistedStep.Runtime: the wall-clock window from just before the + // batch's tools start until the last billed tool completes. All + // calls in a batch start together, so this equals the union of the + // billed tools' execution intervals; parallel calls are never + // summed. Zero when no billed tool produced a result. + BatchRuntime time.Duration + // BatchRuntimeToolCallID identifies the billed tool call whose + // completion ends the batch window, with ties broken by call + // order. The persistence layer stores BatchRuntime on that call's + // tool message row. Empty when BatchRuntime is zero. + BatchRuntimeToolCallID string } // GenerateCompactionOptions configures one context compaction call. @@ -589,6 +608,7 @@ func ExecuteLocalTools(ctx context.Context, opts ExecuteLocalToolsOptions) (Tool } maxResultBytes := toolResultByteBudget(opts.ContextLimit) + batchStart := clockNow(opts.Clock) toolResults := executeTools( ctx, opts.Clock, @@ -617,10 +637,59 @@ func ExecuteLocalTools(ctx context.Context, opts ExecuteLocalToolsOptions) (Tool for _, tr := range toolResults { result.content = append(result.content, tr) } - return ToolExecutionOutcome{Step: PersistedStep{ - Content: result.content, - ToolResultCreatedAt: result.toolResultCreatedAt, - }}, nil + batchRuntime, batchRuntimeToolCallID := billableBatchWindow( + batchStart, + localCalls, + result.toolResultCreatedAt, + opts.UnbilledToolNames, + ) + return ToolExecutionOutcome{ + Step: PersistedStep{ + Content: result.content, + ToolResultCreatedAt: result.toolResultCreatedAt, + }, + BatchRuntime: batchRuntime, + BatchRuntimeToolCallID: batchRuntimeToolCallID, + }, nil +} + +// billableBatchWindow computes one local tool batch's billable runtime: +// the span from batchStart to the latest completion among billed tools. +// Because all calls in a batch start together, that span equals the +// union of the billed tools' execution intervals, so parallel calls are +// billed once rather than summed, and unbilled tools (for example +// sub-agent orchestration) never extend the window even when they run +// longest. Returns the tool call whose completion ends the window, with +// ties broken by call order; (0, "") when no billed tool completed or +// the window rounds to nothing. +func billableBatchWindow( + batchStart time.Time, + toolCalls []fantasy.ToolCallContent, + completedAt map[string]time.Time, + unbilledToolNames map[string]bool, +) (time.Duration, string) { + var ( + windowEnd time.Time + toolCallID string + ) + for _, tc := range toolCalls { + if unbilledToolNames[tc.ToolName] { + continue + } + end, ok := completedAt[tc.ToolCallID] + if !ok { + continue + } + // Strictly-after keeps the earliest call on ties. + if end.After(windowEnd) { + windowEnd = end + toolCallID = tc.ToolCallID + } + } + if toolCallID == "" || !windowEnd.After(batchStart) { + return 0, "" + } + return windowEnd.Sub(batchStart), toolCallID } // prepareMessagesForRequest applies the prompt preparation pipeline used diff --git a/coderd/x/chatd/chatloop/runtime_test.go b/coderd/x/chatd/chatloop/runtime_test.go index b227b26b694..6c832dd6d12 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,233 @@ func TestGenerateCompaction_RecordsRuntime(t *testing.T) { require.Len(t, startedAt, 1) require.Equal(t, result.Runtime, clock.Since(startedAt[0])) } + +// executeToolBatch runs ExecuteLocalTools in a goroutine. Callers trap the +// mock clock's Now: the first trapped call is the batch start and each +// subsequent one is a tool completion, released one tool at a time so +// parallel tool goroutines never race the clock advances. +func executeToolBatch( + t *testing.T, + clock *quartz.Mock, + opts chatloop.ExecuteLocalToolsOptions, +) <-chan chatloop.ToolExecutionOutcome { + t.Helper() + opts.Clock = clock + resultCh := make(chan chatloop.ToolExecutionOutcome, 1) + go func() { + outcome, err := chatloop.ExecuteLocalTools(context.Background(), opts) + assert.NoError(t, err) + resultCh <- outcome + }() + return resultCh +} + +// blockingTool returns a tool that parks until release is closed, so the +// test controls exactly when its completion timestamp is recorded. +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 + }, + ) +} + +// Parallel billed tools bill one shared window ending at the slowest +// tool's completion, never the sum of their durations. The slower tool +// returns an error result: errored tools bill their wall clock too. +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: "{}"}, + }, + }) + + // Batch start. + trap.MustWait(ctx).MustRelease(ctx) + // The fast tool completes 10 seconds in. + clock.Advance(10 * time.Second) + close(fastGo) + trap.MustWait(ctx).MustRelease(ctx) + // The slow tool errors out at 60 seconds. + 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, "call-slow", outcome.BatchRuntimeToolCallID) +} + +// Simultaneous completions tie-break to the earliest call in call order, +// and N parallel calls of the same duration bill that duration once. +func TestExecuteLocalTools_SimultaneousCompletionsBillOnceByCallOrder(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: "{}"}, + }, + }) + + // Batch start. + trap.MustWait(ctx).MustRelease(ctx) + // All three calls complete together 10 seconds in. + 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) + require.Equal(t, "call-1", outcome.BatchRuntimeToolCallID) +} + +// An unbilled tool never extends the window, even when it runs longest: +// the batch bills up to the last billed tool's completion. +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: "{}"}, + }, + }) + + // Batch start. + trap.MustWait(ctx).MustRelease(ctx) + // execute completes 10 seconds in. + clock.Advance(10 * time.Second) + close(executeGo) + trap.MustWait(ctx).MustRelease(ctx) + // wait_agent keeps blocking on its child until 60 seconds. + 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) + require.Equal(t, "call-execute", outcome.BatchRuntimeToolCallID) +} + +// A batch of only unbilled tools bills nothing. +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: "{}"}, + }, + }) + + // Batch start. + 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.Empty(t, outcome.BatchRuntimeToolCallID) +} + +// Billing classifies on the name as called: a deprecated alias listed in +// UnbilledToolNames stays unbilled even though dispatch resolves it to +// its canonical tool through ToolNameAliases. +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: "{}"}, + }, + }) + + // Batch start. + trap.MustWait(ctx).MustRelease(ctx) + // execute completes 10 seconds in. + clock.Advance(10 * time.Second) + close(executeGo) + trap.MustWait(ctx).MustRelease(ctx) + // The aliased call completes at 60 seconds and must not bill. + 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, "call-execute", outcome.BatchRuntimeToolCallID) +} diff --git a/coderd/x/chatd/generation.go b/coderd/x/chatd/generation.go index 5fb7f00e6fd..cde2ee1f532 100644 --- a/coderd/x/chatd/generation.go +++ b/coderd/x/chatd/generation.go @@ -824,6 +824,9 @@ func (s *taskStarter) executeLocalTools( var outcome chatloop.ToolExecutionOutcome var spawnDispatchErr error if len(allowed) > 0 { + // Stamp the batch start on the buffer episode so an interrupt + // can bill the partial window this step would have reported. + attempt.startToolBatch() outcome, err = chatloop.ExecuteLocalTools(ctx, chatloop.ExecuteLocalToolsOptions{ Tools: prepared.Tools, ActiveTools: prepared.ActiveTools, @@ -835,6 +838,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, @@ -856,9 +860,12 @@ func (s *taskStarter) executeLocalTools( outcome.Step.Content = append(outcome.Step.Content, result) } chathooks.RestoreToolCallOrder(outcome.Step.Content, decision.localToolCalls) + step := stepDataFromPersisted(outcome.Step) + step.BatchRuntime = outcome.BatchRuntime + step.BatchRuntimeToolCallID = outcome.BatchRuntimeToolCallID messages, err := buildCommitStepMessages(buildCommitStepMessagesInput{ modelConfigID: prepared.ModelConfigID, - step: stepDataFromPersisted(outcome.Step), + step: step, toolNameToConfigID: prepared.ToolNameToConfigID, logger: s.opts.Logger, contentVersion: chatprompt.CurrentContentVersion, @@ -1047,6 +1054,11 @@ type generationAttempt struct { // can bill the window the step would have reported. It is always // non-nil when beginGenerationAttempt succeeds. startModelInvocation func() + // startToolBatch marks the start of the attempt's billable local + // tool batch window on the buffer episode, so an interrupt can + // bill the window the step would have reported. It is always + // non-nil when beginGenerationAttempt succeeds. + startToolBatch func() // closeEpisode closes the attempt's buffer episode. It is always // non-nil when beginGenerationAttempt succeeds. closeEpisode func() @@ -1093,6 +1105,9 @@ func (s *taskStarter) beginGenerationAttempt( startModelInvocation: func() { _ = s.opts.MessagePartBuffer.StartModelInvocation(key) }, + startToolBatch: func() { + _ = s.opts.MessagePartBuffer.StartToolBatch(key) + }, closeEpisode: func() { _ = s.opts.MessagePartBuffer.CloseEpisode(key) }, diff --git a/coderd/x/chatd/message_conversion.go b/coderd/x/chatd/message_conversion.go index 52d749175bd..d766343b1fd 100644 --- a/coderd/x/chatd/message_conversion.go +++ b/coderd/x/chatd/message_conversion.go @@ -73,7 +73,16 @@ 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) + // The batch's billable window lands on the single tool row whose + // completion ended it; every other row in the batch stays NULL so + // usage reporting, which sums runtime_ms across rows, bills the + // batch exactly once. Zero maps to NULL, so a sub-millisecond + // window persists the same way an unmeasured one does. + if toolResult.ToolCallID != "" && toolResult.ToolCallID == input.step.BatchRuntimeToolCallID { + msg.RuntimeMs = nullInt64IfNonZero(input.step.BatchRuntime.Milliseconds()) + } + messages = append(messages, msg) } return stepMessagesForCommit{ @@ -635,9 +644,10 @@ type partialMessageConversionState struct { // from the model stream itself (text, reasoning, tool calls, // sources). Tool execution also publishes assistant-role file // parts for attachments; those alone must not attract the - // attempt's runtime, because tool batches are not billable. The - // buffer episode only carries a runtime when a provider stream - // was opened, so this is a second gate rather than the only one. + // attempt's model-invocation runtime, because tool batches bill + // their window on tool rows instead. The buffer episode only + // carries a model runtime when a provider stream was opened, so + // this is a second gate rather than the only one. modelStreamedAssistant bool } diff --git a/coderd/x/chatd/message_conversion_test.go b/coderd/x/chatd/message_conversion_test.go index c84935d72f3..0b3e3063506 100644 --- a/coderd/x/chatd/message_conversion_test.go +++ b/coderd/x/chatd/message_conversion_test.go @@ -100,9 +100,74 @@ 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) { +// A local tool batch bills its window on the single tool row whose +// completion ended it; the batch's other rows stay NULL so summing +// runtime_ms across rows bills the batch exactly once. +func TestBuildCommitStepMessages_BatchRuntimeLandsOnWindowDefiningToolRow(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, + BatchRuntimeToolCallID: "call-2", + }, + }) + require.NoError(t, err) + require.Len(t, got.Messages, 2) + require.Equal(t, database.ChatMessageRoleTool, 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: 10000, Valid: true}, got.Messages[1].RuntimeMs) +} + +// Assistant rows synthesized from a tool batch (attachment file parts) +// never carry the batch runtime: it belongs to the tool row alone. +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, + BatchRuntimeToolCallID: "call-1", + }, + }) + 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) +} + +// A batch whose billable window is empty (for example only sub-agent +// orchestration tools ran) must persist runtime_ms NULL on every row. +func TestBuildCommitStepMessages_ZeroBatchRuntimeLeavesRuntimeNull(t *testing.T) { t.Parallel() got, err := buildCommitStepMessages(buildCommitStepMessagesInput{ diff --git a/coderd/x/chatd/messagepartbuffer/message_part_buffer.go b/coderd/x/chatd/messagepartbuffer/message_part_buffer.go index 82f41648be6..16996e24105 100644 --- a/coderd/x/chatd/messagepartbuffer/message_part_buffer.go +++ b/coderd/x/chatd/messagepartbuffer/message_part_buffer.go @@ -102,12 +102,17 @@ type episodeState struct { // that never invoke a model, such as local tool execution // batches. modelStartedAt time.Time - closed bool - closedAt time.Time - closedHeapItem *closedEpisodeItem - parts []Part - bytes int64 - subscribers map[*episodeSubscriber]struct{} + // toolBatchStartedAt is stamped by StartToolBatch when the + // episode begins executing its local tool batch. It is zero for + // episodes that never execute local tools, such as model + // invocations that finish without tool calls. + toolBatchStartedAt time.Time + closed bool + closedAt time.Time + closedHeapItem *closedEpisodeItem + parts []Part + bytes int64 + subscribers map[*episodeSubscriber]struct{} } type closedEpisodeItem struct { @@ -218,6 +223,25 @@ func (b *Buffer) StartModelInvocation(key Key) error { return nil } +// StartToolBatch stamps the instant the episode begins executing its local +// tool batch, which starts the batch's billable runtime window. +func (b *Buffer) StartToolBatch(key Key) error { + b.mu.Lock() + defer b.mu.Unlock() + if b.closed { + return ErrMessagePartBufferClosed + } + episode, err := b.getEpisodeLocked(key) + if err != nil { + return err + } + if episode.closed { + return ErrEpisodeClosed + } + episode.toolBatchStartedAt = b.opts.Clock.Now("message-part-buffer", "tool-batch-start") + return nil +} + // AddPart appends a part to an existing episode. // // Parts receive contiguous sequence numbers so stream endpoints can detect @@ -303,6 +327,19 @@ func (b *Buffer) ModelInvokedAt(key Key) time.Time { return episode.modelStartedAt } +// ToolBatchStartedAt returns the instant stamped by StartToolBatch, or the +// zero time if there is none. Read it before CloseEpisode: closed episodes +// are garbage collected, so reading afterwards races the cleanup loop. +func (b *Buffer) ToolBatchStartedAt(key Key) time.Time { + b.mu.Lock() + defer b.mu.Unlock() + episode := b.episodes[key] + if episode == nil { + return time.Time{} + } + return episode.toolBatchStartedAt +} + // SubscribeToEpisode replays existing parts and streams new parts. // // Subscribers may attach before CreateEpisode is called. In that case the diff --git a/coderd/x/chatd/messagepartbuffer/message_part_buffer_test.go b/coderd/x/chatd/messagepartbuffer/message_part_buffer_test.go index ea489af2448..4671bccb64b 100644 --- a/coderd/x/chatd/messagepartbuffer/message_part_buffer_test.go +++ b/coderd/x/chatd/messagepartbuffer/message_part_buffer_test.go @@ -148,6 +148,43 @@ func TestBuffer_ModelInvokedAt(t *testing.T) { require.Zero(t, buffer.ModelInvokedAt(implicit)) } +func TestBuffer_ToolBatchStartedAt(t *testing.T) { + t.Parallel() + + clock := quartz.NewMock(t) + buffer := messagepartbuffer.New(messagepartbuffer.Options{Clock: clock}) + defer buffer.Close() + + key := testEpisodeKey() + require.Zero(t, buffer.ToolBatchStartedAt(key), "unknown episode has no batch stamp") + require.ErrorIs(t, buffer.StartToolBatch(key), messagepartbuffer.ErrEpisodeNotFound) + + require.NoError(t, buffer.CreateEpisode(key)) + require.Zero(t, buffer.ToolBatchStartedAt(key), "episode without a tool batch has no batch stamp") + // Attempt setup happens before tools start executing and is not + // billable. + clock.Advance(time.Second) + require.NoError(t, buffer.StartToolBatch(key)) + startedAt := buffer.ToolBatchStartedAt(key) + require.Equal(t, clock.Now(), startedAt) + + // Closing must not move the recorded stamp, and a closed episode + // no longer accepts a batch start. + clock.Advance(1500 * time.Millisecond) + require.NoError(t, buffer.CloseEpisode(key)) + require.ErrorIs(t, buffer.StartToolBatch(key), messagepartbuffer.ErrEpisodeClosed) + require.Equal(t, startedAt, buffer.ToolBatchStartedAt(key)) + + // Episodes that never execute local tools, such as model + // invocations without tool calls, report no batch stamp. + modelOnly := testEpisodeKey() + require.NoError(t, buffer.CreateEpisode(modelOnly)) + require.NoError(t, buffer.StartModelInvocation(modelOnly)) + clock.Advance(time.Second) + require.NoError(t, buffer.CloseEpisode(modelOnly)) + require.Zero(t, buffer.ToolBatchStartedAt(modelOnly)) +} + func TestBuffer_SubscribeExistingReplaysThenStreamsLiveParts(t *testing.T) { t.Parallel() diff --git a/coderd/x/chatd/subagent_catalog.go b/coderd/x/chatd/subagent_catalog.go index a79bf81fa5d..937adc5fa0c 100644 --- a/coderd/x/chatd/subagent_catalog.go +++ b/coderd/x/chatd/subagent_catalog.go @@ -38,6 +38,25 @@ const ( "external or web research, parallel research, or tasks that may need edits." ) +// unbilledSubagentToolNames lists the sub-agent orchestration tools whose +// execution time is excluded from the local-tool runtime persisted to +// chat_messages.runtime_ms. Every chat bills its own model and tool time, +// including child agents, so a parent's wait_agent window would count each +// child's already-billed runtime a second time. The remaining orchestration +// tools are millisecond-scale bookkeeping and are excluded with it so the +// billing rule stays one explainable category. Classification uses the tool +// name as called, so the deprecated close_agent alias is listed alongside +// interrupt_agent. +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 9f1e55c54fd..6d91de351fc 100644 --- a/coderd/x/chatd/subagent_internal_test.go +++ b/coderd/x/chatd/subagent_internal_test.go @@ -4261,6 +4261,28 @@ func TestAwaitSubagentCompletion(t *testing.T) { }) } +// The unbilled set must track the sub-agent orchestration catalog exactly: +// every orchestration tool and deprecated alias is excluded from local tool +// runtime billing, and nothing outside the catalog is. +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.go b/coderd/x/chatd/tasks.go index a3738b886db..ece4bee2739 100644 --- a/coderd/x/chatd/tasks.go +++ b/coderd/x/chatd/tasks.go @@ -261,6 +261,7 @@ func (s *taskStarter) StartInterrupt(ctx context.Context, input chatWorkerTaskSt GenerationAttempt: chat.GenerationAttempt, } modelInvokedAt := s.opts.MessagePartBuffer.ModelInvokedAt(key) + toolBatchStartedAt := s.opts.MessagePartBuffer.ToolBatchStartedAt(key) if err := s.opts.MessagePartBuffer.CloseEpisode(key); err != nil { if ctx.Err() != nil { return errors.Join(errTaskExpectedExit, xerrors.Errorf("close message part episode: %w", err), ctx.Err()) @@ -302,7 +303,10 @@ func (s *taskStarter) StartInterrupt(ctx context.Context, input chatWorkerTaskSt return xerrors.Errorf("load chat for task: %w", err) } messages := partialMessages - committedCancels, err := committedPendingLocalToolCancellationMessages(ctx, store, chat, s.opts.Clock.Now("chatworker", "interrupt")) + committedCancels, err := committedPendingLocalToolCancellationMessages(ctx, store, chat, s.opts.Clock.Now("chatworker", "interrupt"), interruptedToolBatchBilling{ + batchStartedAt: toolBatchStartedAt, + bufferedParts: parts, + }) if err != nil { return xerrors.Errorf("committed pending local tool cancellation messages: %w", err) } @@ -678,11 +682,29 @@ func dynamicToolNamesFromChat(chat database.Chat) map[string]bool { return names } +// interruptedToolBatchBilling carries what the interrupt task knows about +// the live tool batch it is canceling, so the synthesized cancellation +// rows can bill the partial window the batch would have reported. +type interruptedToolBatchBilling struct { + // batchStartedAt is the StartToolBatch stamp from the interrupted + // attempt's buffer episode. Zero when no batch was live at the + // interrupt (crash recovery, state promotion), in which case the + // cancellation rows carry no runtime. + batchStartedAt time.Time + // bufferedParts are the interrupted episode's buffered parts. Tool + // results buffered before the interrupt carry the completion times + // of tools that finished early, so a batch whose billed tools all + // completed does not bill the longer window of a still-running + // unbilled tool such as wait_agent. + bufferedParts []messagepartbuffer.Part +} + func committedPendingLocalToolCancellationMessages( ctx context.Context, store database.Store, chat database.Chat, interruptedAt time.Time, + billing interruptedToolBatchBilling, ) ([]chatstate.Message, error) { messages, err := store.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{ ChatID: chat.ID, @@ -698,6 +720,11 @@ func committedPendingLocalToolCancellationMessages( if len(localCalls) == 0 { return nil, nil } + completions := bufferedToolResultCompletions(billing.bufferedParts) + var ( + windowEnd time.Time + windowRowIdx = -1 + ) result := make([]chatstate.Message, 0, len(localCalls)) for _, call := range localCalls { payload, err := json.Marshal(map[string]string{"error": interruptedToolResultErrorMessage}) @@ -719,6 +746,55 @@ func committedPendingLocalToolCancellationMessages( ModelConfigID: uuid.NullUUID{UUID: chat.LastModelConfigID, Valid: chat.LastModelConfigID != uuid.Nil}, ContentVersion: chatprompt.CurrentContentVersion, }) + if billing.batchStartedAt.IsZero() || unbilledSubagentToolNames[call.ToolName] { + continue + } + // A billed tool that buffered a durable result completed at that + // instant; one without a buffered result was still running, so + // its window ends at the interrupt. Strictly-after keeps the + // earliest call on ties, matching billableBatchWindow. + end, ok := completions[call.ToolCallID] + if !ok { + end = interruptedAt + } + if end.After(windowEnd) { + windowEnd = end + windowRowIdx = len(result) - 1 + } + } + // Mirror the committed-batch policy: bill the partial window once, on + // the cancellation row of the billed tool call that defines it. + if windowRowIdx >= 0 && windowEnd.After(billing.batchStartedAt) { + result[windowRowIdx].RuntimeMs = nullInt64IfNonZero(windowEnd.Sub(billing.batchStartedAt).Milliseconds()) } return result, nil } + +// bufferedToolResultCompletions maps tool call IDs to the completion +// instants of durable tool results buffered before an interrupt. +// Streaming deltas and resets are not completions, and only the first +// durable result per call counts, matching the buffered-part conversion +// rules in bufferedPartsToPartialMessages. +func bufferedToolResultCompletions(parts []messagepartbuffer.Part) map[string]time.Time { + completions := make(map[string]time.Time) + for _, buffered := range parts { + if buffered.Role != codersdk.ChatMessageRoleTool { + continue + } + part := buffered.MessagePart + if part.Type != codersdk.ChatMessagePartTypeToolResult || part.ToolCallID == "" { + continue + } + if part.ResultDelta != "" || part.ResultReset || len(part.Result) == 0 { + continue + } + if part.CreatedAt == nil { + continue + } + if _, ok := completions[part.ToolCallID]; ok { + continue + } + completions[part.ToolCallID] = *part.CreatedAt + } + return completions +} diff --git a/coderd/x/chatd/tasks_test.go b/coderd/x/chatd/tasks_test.go index a1265b270b9..d12357024a6 100644 --- a/coderd/x/chatd/tasks_test.go +++ b/coderd/x/chatd/tasks_test.go @@ -477,6 +477,202 @@ func TestInterruptTask_PartialAssistantWithoutModelInvocationHasNoRuntime(t *tes require.False(t, assistant.RuntimeMs.Valid) } +// interruptedBatch is a chat with committed unresolved local tool calls +// and a live attempt buffer episode, ready for StartInterrupt to +// synthesize cancellation rows. +type interruptedBatch struct { + chat database.Chat + starter *taskStarter + clock *quartz.Mock + key messagepartbuffer.Key + workerID uuid.UUID + runnerID uuid.UUID +} + +// interruptedBatchFixture commits an assistant message with the given +// unresolved local tool calls and prepares the live attempt's buffer +// episode, so StartInterrupt synthesizes cancellation rows for them. +func interruptedBatchFixture( + t *testing.T, + f *taskTestFixture, + calls []codersdk.ChatMessagePart, +) interruptedBatch { + t.Helper() + chat := f.createRunningChat(t) + raw, err := chatprompt.MarshalParts(calls) + require.NoError(t, err) + machine := chatstate.NewChatMachine(f.db, f.pubsub, chat.ID) + require.NoError(t, machine.Update(testutil.Context(t, testutil.WaitShort), func(tx *chatstate.Tx, _ database.Store) error { + _, err := tx.CommitStep(chatstate.CommitStepInput{Messages: []chatstate.Message{{ + Role: database.ChatMessageRoleAssistant, + Content: raw, + Visibility: database.ChatMessageVisibilityBoth, + ContentVersion: chatprompt.CurrentContentVersion, + ModelConfigID: uuid.NullUUID{UUID: f.model.ID, Valid: true}, + }}}) + return err + })) + workerID := uuid.New() + runnerID := uuid.New() + acquired := f.acquireChat(t, chat.ID, workerID, runnerID) + recorder := newTaskSideEffectRecorder() + clock := quartz.NewMock(t) + starter := newTestTaskStarterWithClock(t, f, recorder, clock) + key := messagepartbuffer.Key{ + ChatID: chat.ID, + HistoryVersion: acquired.HistoryVersion, + GenerationAttempt: acquired.GenerationAttempt, + } + require.NoError(t, starter.opts.MessagePartBuffer.CreateEpisode(key)) + return interruptedBatch{ + chat: chat, + starter: starter, + clock: clock, + key: key, + workerID: workerID, + runnerID: runnerID, + } +} + +func (b interruptedBatch) interrupt(t *testing.T, f *taskTestFixture) []database.ChatMessage { + t.Helper() + interrupting := f.interruptChat(t, b.chat.ID) + err := b.starter.StartInterrupt(testutil.Context(t, testutil.WaitLong), chatWorkerTaskStartInput{ + ChatID: b.chat.ID, + WorkerID: b.workerID, + RunnerID: b.runnerID, + HistoryVersion: interrupting.HistoryVersion, + GenerationAttempt: interrupting.GenerationAttempt, + Status: database.ChatStatusInterrupting, + }) + require.NoError(t, err) + messages, err := f.db.GetChatMessagesByChatID(testutil.Context(t, testutil.WaitShort), database.GetChatMessagesByChatIDParams{ChatID: b.chat.ID}) + require.NoError(t, err) + return messages +} + +// findToolResultMessage returns the tool-role message answering the given +// tool call ID. +func findToolResultMessage(t *testing.T, messages []database.ChatMessage, toolCallID string) database.ChatMessage { + t.Helper() + for _, msg := range messages { + if msg.Role != database.ChatMessageRoleTool { + continue + } + parts, err := chatprompt.ParseContent(msg) + require.NoError(t, err) + for _, part := range parts { + if part.Type == codersdk.ChatMessagePartTypeToolResult && part.ToolCallID == toolCallID { + return msg + } + } + } + t.Fatalf("no tool result message for call %s", toolCallID) + return database.ChatMessage{} +} + +// An interrupt mid tool batch bills the partial window on the cancellation +// row of the billed tool that defines it. A billed tool that buffered a +// result before the interrupt ends the window at its completion, so a +// still-running unbilled wait_agent does not stretch the bill to the +// interrupt instant. +func TestInterruptTask_ToolBatchBillsPartialWindowOnCancellationRow(t *testing.T) { + t.Parallel() + + f := newTaskTestFixture(t) + execCallID := "call_" + uuid.NewString() + waitCallID := "call_" + uuid.NewString() + batch := interruptedBatchFixture(t, f, []codersdk.ChatMessagePart{ + {Type: codersdk.ChatMessagePartTypeToolCall, ToolCallID: execCallID, ToolName: "execute", Args: json.RawMessage(`{}`)}, + {Type: codersdk.ChatMessagePartTypeToolCall, ToolCallID: waitCallID, ToolName: "wait_agent", Args: json.RawMessage(`{}`)}, + }) + buffer := batch.starter.opts.MessagePartBuffer + + // Advances stay under the buffer's 15s cleanup tick, which shares + // this mock clock. + // Attempt setup happens before the tools start and is not billable. + batch.clock.Advance(2 * time.Second) + require.NoError(t, buffer.StartToolBatch(batch.key)) + // execute completes 3 seconds into the batch and buffers its + // durable result. + batch.clock.Advance(3 * time.Second) + execCompletedAt := batch.clock.Now() + execResult := codersdk.ChatMessageToolResult(execCallID, "execute", json.RawMessage(`{"stdout":"/tmp"}`), false, false) + execResult.CreatedAt = &execCompletedAt + require.NoError(t, buffer.AddPart(batch.key, codersdk.ChatMessageRoleTool, execResult)) + // wait_agent is still blocked on its child when the interrupt lands + // 5 seconds later. + batch.clock.Advance(5 * time.Second) + + messages := batch.interrupt(t, f) + execRow := findToolResultMessage(t, messages, execCallID) + require.Equal(t, sql.NullInt64{Int64: 3_000, Valid: true}, execRow.RuntimeMs) + waitRow := findToolResultMessage(t, messages, waitCallID) + require.False(t, waitRow.RuntimeMs.Valid) +} + +// A billed tool still running at the interrupt bills up to the interrupt +// instant. +func TestInterruptTask_RunningBilledToolBillsUpToInterrupt(t *testing.T) { + t.Parallel() + + f := newTaskTestFixture(t) + execCallID := "call_" + uuid.NewString() + batch := interruptedBatchFixture(t, f, []codersdk.ChatMessagePart{ + {Type: codersdk.ChatMessagePartTypeToolCall, ToolCallID: execCallID, ToolName: "execute", Args: json.RawMessage(`{}`)}, + }) + + batch.clock.Advance(2 * time.Second) + require.NoError(t, batch.starter.opts.MessagePartBuffer.StartToolBatch(batch.key)) + // The tool is still running when the interrupt lands 7 seconds in. + batch.clock.Advance(7 * time.Second) + + messages := batch.interrupt(t, f) + execRow := findToolResultMessage(t, messages, execCallID) + require.Equal(t, sql.NullInt64{Int64: 7_000, Valid: true}, execRow.RuntimeMs) +} + +// An interrupted batch of only unbilled sub-agent orchestration tools +// bills nothing: an interrupted lone wait_agent stays free. +func TestInterruptTask_UnbilledOnlyBatchBillsNothingOnInterrupt(t *testing.T) { + t.Parallel() + + f := newTaskTestFixture(t) + waitCallID := "call_" + uuid.NewString() + batch := interruptedBatchFixture(t, f, []codersdk.ChatMessagePart{ + {Type: codersdk.ChatMessagePartTypeToolCall, ToolCallID: waitCallID, ToolName: "wait_agent", Args: json.RawMessage(`{}`)}, + }) + + batch.clock.Advance(2 * time.Second) + require.NoError(t, batch.starter.opts.MessagePartBuffer.StartToolBatch(batch.key)) + batch.clock.Advance(10 * time.Second) + + messages := batch.interrupt(t, f) + waitRow := findToolResultMessage(t, messages, waitCallID) + require.False(t, waitRow.RuntimeMs.Valid) +} + +// Cancellation rows synthesized without a live tool batch (crash +// recovery: the episode never stamped a batch start) carry no runtime, +// consistent with generation losing in-flight runtime on a crash. +func TestInterruptTask_ToolCancellationWithoutLiveBatchHasNoRuntime(t *testing.T) { + t.Parallel() + + f := newTaskTestFixture(t) + execCallID := "call_" + uuid.NewString() + batch := interruptedBatchFixture(t, f, []codersdk.ChatMessagePart{ + {Type: codersdk.ChatMessagePartTypeToolCall, ToolCallID: execCallID, ToolName: "execute", Args: json.RawMessage(`{}`)}, + }) + + // No StartToolBatch stamp: the batch never went live on this + // attempt. + batch.clock.Advance(10 * time.Second) + + messages := batch.interrupt(t, f) + execRow := findToolResultMessage(t, messages, execCallID) + require.False(t, execRow.RuntimeMs.Valid) +} + func TestRequiresActionTimeout_ExpiredCancelsOnly(t *testing.T) { t.Parallel() diff --git a/docs/ai-coder/usage-data-reporting.md b/docs/ai-coder/usage-data-reporting.md index 3c69e27ae0b..4e109956722 100644 --- a/docs/ai-coder/usage-data-reporting.md +++ b/docs/ai-coder/usage-data-reporting.md @@ -58,27 +58,38 @@ Example of a failed request (e.g. Tallyman Server is blocked by your network): ## Agent runtime measurement -Total Coder Agent runtime is summed from per-message generation time +Total Coder Agent runtime is summed from per-message runtime (`runtime_ms` on chat messages). -A message's runtime is the wall-clock duration of the model invocation that -produced its content, measured from just before the request to the model -provider opens until the response is fully consumed. +An assistant message's runtime is the wall-clock duration of the model +invocation that produced its content, measured from just before the request +to the model provider opens until the response is fully consumed. A tool +message's runtime is the wall-clock duration of the local tool batch that +produced it, measured from the start of the batch until the last counted +tool finishes. Tools in a batch run in parallel, so each batch records one +window on one tool message rather than a per-tool sum. What counts: - Assistant generation steps, in both top-level chats and sub-agent chats. - Context compaction (summarization) model calls. -- Interrupted generation: the time streamed before the interrupt is kept on - the partial assistant message. +- Local tool execution: file, terminal, and process tools, workspace + lifecycle operations, MCP tools, and other server-executed tools. +- Interrupted generation: the time streamed or spent executing tools before + the interrupt is kept on the partial messages. What does not count: -- Local tool execution, including waiting on sub-agents. A sub-agent is its - own chat and records its own model invocations, so counting the parent's - wait would double count. +- Sub-agent orchestration tools, such as spawning and waiting on + sub-agents. A sub-agent is its own chat and records its own runtime, so + counting the parent's wait would double count. Waiting on a sub-agent + never extends a tool batch's window, even when other tools in the batch + do count. +- Client-executed (dynamic) tools and external agents: the server cannot + measure work it does not execute. - Idle time: chats waiting for user input or external tool results. -- Failed model calls whose output was discarded. Retried and errored - attempts persist no content, so they record no runtime. +- Failed model calls whose output was discarded, and the backoff between + retried attempts. Retried and errored attempts persist no content, so + they record no runtime. - Ancillary model calls that produce no chat messages, such as title generation. From a88b8a4d9a1678a48d3d2e6a9b932f1de84dc6fd Mon Sep 17 00:00:00 2001 From: Jaayden Halko Date: Tue, 18 Aug 2026 05:06:52 +0000 Subject: [PATCH 02/20] fix(coderd/x/chatd): bill interrupted tool batches with live completion times Address Codex review feedback: - Record each local tool call's completion instant on the message part buffer as the tool finishes (chatloop OnToolComplete), because tool results are published only after the whole batch completes. The interrupt task now reads these live stamps instead of inferring completions from buffered results, which never exist while a sibling tool is still running, so an interrupted batch no longer bills a finished tool through to the interrupt instant. - Assign the batch runtime to only the first tool row matching the window-defining tool call ID, so duplicate tool call IDs cannot multiply the billed sum. - Replace the agent-authored ARCHITECTURE.md buffer API entries with a TODO item for the PR author, per the chatd documentation rule. --- coderd/x/chatd/ARCHITECTURE.md | 3 +- coderd/x/chatd/chatloop/chatloop.go | 22 ++++++- coderd/x/chatd/chatloop/runtime_test.go | 60 +++++++++++++++++++ coderd/x/chatd/generation.go | 10 ++++ coderd/x/chatd/message_conversion.go | 11 +++- coderd/x/chatd/message_conversion_test.go | 33 ++++++++++ .../messagepartbuffer/message_part_buffer.go | 59 ++++++++++++++++-- .../message_part_buffer_test.go | 38 ++++++++++++ coderd/x/chatd/tasks.go | 58 +++++------------- coderd/x/chatd/tasks_test.go | 15 +++-- 10 files changed, 246 insertions(+), 63 deletions(-) diff --git a/coderd/x/chatd/ARCHITECTURE.md b/coderd/x/chatd/ARCHITECTURE.md index ec9c8187c20..7880d4fe4f3 100644 --- a/coderd/x/chatd/ARCHITECTURE.md +++ b/coderd/x/chatd/ARCHITECTURE.md @@ -725,8 +725,7 @@ The buffer exposes the following API: - `GetParts(chat_id, history_version, generation_attempt)`: returns the message parts for an episode. Returns a predefined error if the episode is not found. - `StartModelInvocation(chat_id, history_version, generation_attempt)`: stamps the instant the episode opens its provider stream. Returns a predefined error if the episode is not found or already closed. Episodes that never invoke a model, such as local tool execution batches, are never stamped. - `ModelInvokedAt(chat_id, history_version, generation_attempt)`: returns the instant stamped by `StartModelInvocation`, or the zero time when the episode is unknown or never opened a provider stream. It must be read before `CloseEpisode`, because closed episodes are garbage collected and reading afterwards races the cleanup loop. The interrupt goroutine reads it just before closing the episode and uses the span between that instant and the interrupt as the interrupted attempt's billable runtime. -- `StartToolBatch(chat_id, history_version, generation_attempt)`: stamps the instant the episode begins executing its local tool batch. Returns a predefined error if the episode is not found or already closed. Episodes that never execute local tools, such as model invocations without tool calls, are never stamped. -- `ToolBatchStartedAt(chat_id, history_version, generation_attempt)`: returns the instant stamped by `StartToolBatch`, or the zero time when the episode is unknown or never started a tool batch. Like `ModelInvokedAt`, it must be read before `CloseEpisode`. The interrupt goroutine reads it just before closing the episode and uses it as the start of the interrupted tool batch's billable window. +- TODO(CODAGT-928): document the tool-batch billing methods `StartToolBatch`, `ToolBatchStartedAt`, `RecordToolCompletion`, and `ToolCompletionsAt`, the local-tool counterparts of `StartModelInvocation`/`ModelInvokedAt` that the interrupt task reads before `CloseEpisode` to bill an interrupted tool batch's partial window. - `SubscribeToEpisode(chat_id, history_version, generation_attempt)`: returns a go channel that will receive all message parts for the episode. It spawns a goroutine that delivers parts to the channel. It's live until the episode is closed or until a subscriber requests that the channel be closed. Once the goroutine delivers all message parts for a closed episode, it closes the channel and exits. If the episode is already closed at the time of the call, the goroutine delivers all message parts for the episode, closes the channel, and exits. `SubscribeToEpisode` does not return an error if the episode is not found: it waits for it to be created instead. Closed episodes are garbage collected after at least 15 seconds since they were closed and when they have no active subscribers. The message part buffer maintains a garbage collection goroutine. diff --git a/coderd/x/chatd/chatloop/chatloop.go b/coderd/x/chatd/chatloop/chatloop.go index 00ef8d4589e..f642e590208 100644 --- a/coderd/x/chatd/chatloop/chatloop.go +++ b/coderd/x/chatd/chatloop/chatloop.go @@ -265,6 +265,16 @@ type ExecuteLocalToolsOptions struct { // name as called, so deprecated aliases must be listed alongside // their canonical names. UnbilledToolNames map[string]bool + // OnToolComplete, when set, is invoked with each local tool call's + // completion instant as the tool finishes, the same instant + // PersistedStep.ToolResultCreatedAt later carries. Tool results are + // published only after the whole batch completes, to keep event + // ordering deterministic, so this callback is the only live signal + // that a tool already finished while siblings are still running; + // the interrupt path uses it to bill an interrupted batch's + // partial window. It is called concurrently from tool goroutines + // and must be safe for concurrent use. + OnToolComplete func(toolCallID string, completedAt time.Time) PublishMessagePart func(codersdk.ChatMessageRole, codersdk.ChatMessagePart) Logger slog.Logger @@ -592,6 +602,9 @@ func ExecuteLocalTools(ctx context.Context, opts ExecuteLocalToolsOptions) (Tool now := clockNow(opts.Clock) for _, tr := range policyResults { recordToolResultTimestamp(&result, tr.ToolCallID, now) + if opts.OnToolComplete != nil { + opts.OnToolComplete(tr.ToolCallID, now) + } publishToolAttachments(ctx, opts.Logger, tr, now, publishMessagePart) ssePart := chatprompt.PartFromContentWithLogger(ctx, opts.Logger, tr) ssePart.CreatedAt = &now @@ -623,6 +636,7 @@ func ExecuteLocalTools(ctx context.Context, opts ExecuteLocalToolsOptions) (Tool opts.BuiltinToolNames, maxResultBytes, opts.ToolNameAliases, + opts.OnToolComplete, func(tr fantasy.ToolResultContent, completedAt time.Time) { recordToolResultTimestamp(&result, tr.ToolCallID, completedAt) publishToolAttachments(ctx, opts.Logger, tr, completedAt, publishMessagePart) @@ -1120,7 +1134,9 @@ func processStepStream( // 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. +// event ordering for SSE subscribers. onComplete, in contrast, fires +// from each tool's goroutine the instant that tool finishes, so +// callers can observe completions while slower siblings still run. func executeTools( ctx context.Context, clock quartz.Clock, @@ -1134,6 +1150,7 @@ func executeTools( builtinToolNames map[string]bool, maxResultBytes int, toolNameAliases map[string]string, + onComplete func(toolCallID string, completedAt time.Time), onResult func(fantasy.ToolResultContent, time.Time), ) []fantasy.ToolResultContent { if len(toolCalls) == 0 { @@ -1199,6 +1216,9 @@ func executeTools( // Captured per-goroutine so parallel tools get // accurate individual completion times. completedAt[i] = clockNow(clock) + if onComplete != nil { + onComplete(tc.ToolCallID, completedAt[i]) + } }() results[i] = executeSingleTool( ctx, diff --git a/coderd/x/chatd/chatloop/runtime_test.go b/coderd/x/chatd/chatloop/runtime_test.go index 6c832dd6d12..a4a79a0eec7 100644 --- a/coderd/x/chatd/chatloop/runtime_test.go +++ b/coderd/x/chatd/chatloop/runtime_test.go @@ -391,3 +391,63 @@ func TestExecuteLocalTools_AliasNamesClassifyAsCalled(t *testing.T) { require.Equal(t, 10*time.Second, outcome.BatchRuntime) require.Equal(t, "call-execute", outcome.BatchRuntimeToolCallID) } + +// OnToolComplete reports each tool's completion the instant it finishes, +// while slower siblings are still running, with the same instants the +// outcome's ToolResultCreatedAt later carries. The interrupt path +// depends on this live signal: results publish only after the whole +// batch finishes, so without it an interrupt could not tell finished +// tools from still-running ones. +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 { + toolCallID string + 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(toolCallID string, completedAt time.Time) { + completionCh <- completion{toolCallID: toolCallID, completedAt: completedAt} + }, + ToolCalls: []fantasy.ToolCallContent{ + {ToolCallID: "call-fast", ToolName: "fast_tool", Input: "{}"}, + {ToolCallID: "call-slow", ToolName: "slow_tool", Input: "{}"}, + }, + }) + + // Batch start. + trap.MustWait(ctx).MustRelease(ctx) + // The fast tool completes 10 seconds in. Its completion arrives + // while the slow tool is still parked on its release channel. + clock.Advance(10 * time.Second) + close(fastGo) + trap.MustWait(ctx).MustRelease(ctx) + fast := testutil.RequireReceive(ctx, t, completionCh) + require.Equal(t, "call-fast", fast.toolCallID) + // The slow tool completes at 60 seconds. + clock.Advance(50 * time.Second) + close(slowGo) + trap.MustWait(ctx).MustRelease(ctx) + slow := testutil.RequireReceive(ctx, t, completionCh) + require.Equal(t, "call-slow", slow.toolCallID) + 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.Step.ToolResultCreatedAt) +} diff --git a/coderd/x/chatd/generation.go b/coderd/x/chatd/generation.go index cde2ee1f532..52630d4fc5a 100644 --- a/coderd/x/chatd/generation.go +++ b/coderd/x/chatd/generation.go @@ -839,6 +839,7 @@ func (s *taskStarter) executeLocalTools( ContextLimit: prepared.ContextLimitFallback, ToolNameAliases: subagentToolNameAliases, UnbilledToolNames: unbilledSubagentToolNames, + OnToolComplete: attempt.recordToolCompletion, PublishMessagePart: attempt.publish, Logger: s.opts.Logger, Metrics: s.server.metrics, @@ -1059,6 +1060,12 @@ type generationAttempt struct { // bill the window the step would have reported. It is always // non-nil when beginGenerationAttempt succeeds. startToolBatch func() + // recordToolCompletion records a tool call's completion instant on + // the buffer episode as the batch executes, so an interrupt can + // end an already-finished tool's billable window at its real + // completion instead of the interrupt instant. It is always + // non-nil when beginGenerationAttempt succeeds. + recordToolCompletion func(toolCallID string, completedAt time.Time) // closeEpisode closes the attempt's buffer episode. It is always // non-nil when beginGenerationAttempt succeeds. closeEpisode func() @@ -1108,6 +1115,9 @@ func (s *taskStarter) beginGenerationAttempt( startToolBatch: func() { _ = s.opts.MessagePartBuffer.StartToolBatch(key) }, + recordToolCompletion: func(toolCallID string, completedAt time.Time) { + _ = s.opts.MessagePartBuffer.RecordToolCompletion(key, toolCallID, completedAt) + }, closeEpisode: func() { _ = s.opts.MessagePartBuffer.CloseEpisode(key) }, diff --git a/coderd/x/chatd/message_conversion.go b/coderd/x/chatd/message_conversion.go index d766343b1fd..458ba3e26e7 100644 --- a/coderd/x/chatd/message_conversion.go +++ b/coderd/x/chatd/message_conversion.go @@ -61,6 +61,7 @@ func buildCommitStepMessages(input buildCommitStepMessagesInput) (stepMessagesFo messages = append(messages, assistantMessage(input.modelConfigID, contentVersion, assistantContent, input.step)) } + batchRuntimeAssigned := false for _, toolResult := range toolResults { part := chatprompt.PartFromContentWithLogger(context.Background(), input.logger, toolResult) applyToolMetadata(&part, input.toolNameToConfigID) @@ -77,10 +78,14 @@ func buildCommitStepMessages(input buildCommitStepMessagesInput) (stepMessagesFo // The batch's billable window lands on the single tool row whose // completion ended it; every other row in the batch stays NULL so // usage reporting, which sums runtime_ms across rows, bills the - // batch exactly once. Zero maps to NULL, so a sub-millisecond - // window persists the same way an unmeasured one does. - if toolResult.ToolCallID != "" && toolResult.ToolCallID == input.step.BatchRuntimeToolCallID { + // batch exactly once. Only the first row with the window-defining + // ID carries it, because providers can emit duplicate tool call + // IDs and billing every duplicate would multiply the sum. Zero + // maps to NULL, so a sub-millisecond window persists the same + // way an unmeasured one does. + if !batchRuntimeAssigned && toolResult.ToolCallID != "" && toolResult.ToolCallID == input.step.BatchRuntimeToolCallID { msg.RuntimeMs = nullInt64IfNonZero(input.step.BatchRuntime.Milliseconds()) + batchRuntimeAssigned = true } messages = append(messages, msg) } diff --git a/coderd/x/chatd/message_conversion_test.go b/coderd/x/chatd/message_conversion_test.go index 0b3e3063506..0d7515265fb 100644 --- a/coderd/x/chatd/message_conversion_test.go +++ b/coderd/x/chatd/message_conversion_test.go @@ -135,6 +135,39 @@ func TestBuildCommitStepMessages_BatchRuntimeLandsOnWindowDefiningToolRow(t *tes require.Equal(t, sql.NullInt64{Int64: 10000, Valid: true}, got.Messages[1].RuntimeMs) } +// Duplicate tool call IDs, which providers can emit and which admission +// does not always reject, must not multiply the bill: only the first row +// with the window-defining ID carries the batch runtime. +func TestBuildCommitStepMessages_DuplicateToolCallIDsBillOnce(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: "execute", + Result: fantasy.ToolResultOutputContentText{Text: `{"stdout":"first"}`}, + }, + fantasy.ToolResultContent{ + ToolCallID: "call-1", + ToolName: "execute", + Result: fantasy.ToolResultOutputContentText{Text: `{"stdout":"second"}`}, + }, + }, + BatchRuntime: 10 * time.Second, + BatchRuntimeToolCallID: "call-1", + }, + }) + 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) +} + // Assistant rows synthesized from a tool batch (attachment file parts) // never carry the batch runtime: it belongs to the tool row alone. func TestBuildCommitStepMessages_BatchAttachmentAssistantRowStaysNull(t *testing.T) { diff --git a/coderd/x/chatd/messagepartbuffer/message_part_buffer.go b/coderd/x/chatd/messagepartbuffer/message_part_buffer.go index 16996e24105..1a482b60cf9 100644 --- a/coderd/x/chatd/messagepartbuffer/message_part_buffer.go +++ b/coderd/x/chatd/messagepartbuffer/message_part_buffer.go @@ -19,6 +19,7 @@ import ( "container/heap" "context" "encoding/json" + "maps" "slices" "sync" "time" @@ -107,12 +108,18 @@ type episodeState struct { // episodes that never execute local tools, such as model // invocations that finish without tool calls. toolBatchStartedAt time.Time - closed bool - closedAt time.Time - closedHeapItem *closedEpisodeItem - parts []Part - bytes int64 - subscribers map[*episodeSubscriber]struct{} + // toolCompletedAt maps tool call IDs to the completion instants + // recorded by RecordToolCompletion as the batch's tools finish. + // Tool results are published only after the whole batch + // completes, so these per-tool stamps are the only live view of + // which tools already finished when an interrupt lands. + toolCompletedAt map[string]time.Time + closed bool + closedAt time.Time + closedHeapItem *closedEpisodeItem + parts []Part + bytes int64 + subscribers map[*episodeSubscriber]struct{} } type closedEpisodeItem struct { @@ -242,6 +249,31 @@ func (b *Buffer) StartToolBatch(key Key) error { return nil } +// RecordToolCompletion records the instant a local tool call in the +// episode's tool batch finished. Tool goroutines report completions as +// they happen, so an interrupt can bill tools that already finished up +// to their real completion instead of treating every canceled call as +// still running. +func (b *Buffer) RecordToolCompletion(key Key, toolCallID string, completedAt time.Time) error { + b.mu.Lock() + defer b.mu.Unlock() + if b.closed { + return ErrMessagePartBufferClosed + } + episode, err := b.getEpisodeLocked(key) + if err != nil { + return err + } + if episode.closed { + return ErrEpisodeClosed + } + if episode.toolCompletedAt == nil { + episode.toolCompletedAt = make(map[string]time.Time) + } + episode.toolCompletedAt[toolCallID] = completedAt + return nil +} + // AddPart appends a part to an existing episode. // // Parts receive contiguous sequence numbers so stream endpoints can detect @@ -340,6 +372,21 @@ func (b *Buffer) ToolBatchStartedAt(key Key) time.Time { return episode.toolBatchStartedAt } +// ToolCompletionsAt returns a copy of the completion instants recorded +// by RecordToolCompletion, keyed by tool call ID, or nil when the +// episode is unknown or recorded none. Read it before CloseEpisode: +// closed episodes are garbage collected, so reading afterwards races +// the cleanup loop. +func (b *Buffer) ToolCompletionsAt(key Key) map[string]time.Time { + b.mu.Lock() + defer b.mu.Unlock() + episode := b.episodes[key] + if episode == nil { + return nil + } + return maps.Clone(episode.toolCompletedAt) +} + // SubscribeToEpisode replays existing parts and streams new parts. // // Subscribers may attach before CreateEpisode is called. In that case the diff --git a/coderd/x/chatd/messagepartbuffer/message_part_buffer_test.go b/coderd/x/chatd/messagepartbuffer/message_part_buffer_test.go index 4671bccb64b..9d4f1845312 100644 --- a/coderd/x/chatd/messagepartbuffer/message_part_buffer_test.go +++ b/coderd/x/chatd/messagepartbuffer/message_part_buffer_test.go @@ -185,6 +185,44 @@ func TestBuffer_ToolBatchStartedAt(t *testing.T) { require.Zero(t, buffer.ToolBatchStartedAt(modelOnly)) } +func TestBuffer_ToolCompletionsAt(t *testing.T) { + t.Parallel() + + clock := quartz.NewMock(t) + buffer := messagepartbuffer.New(messagepartbuffer.Options{Clock: clock}) + defer buffer.Close() + + key := testEpisodeKey() + require.Nil(t, buffer.ToolCompletionsAt(key), "unknown episode has no completions") + require.ErrorIs(t, buffer.RecordToolCompletion(key, "call-1", clock.Now()), messagepartbuffer.ErrEpisodeNotFound) + + require.NoError(t, buffer.CreateEpisode(key)) + require.Empty(t, buffer.ToolCompletionsAt(key), "episode without recorded completions has none") + // Tools complete at different instants; each keeps its own stamp. + clock.Advance(time.Second) + firstCompletedAt := clock.Now() + require.NoError(t, buffer.RecordToolCompletion(key, "call-1", firstCompletedAt)) + clock.Advance(2 * time.Second) + secondCompletedAt := clock.Now() + require.NoError(t, buffer.RecordToolCompletion(key, "call-2", secondCompletedAt)) + completions := buffer.ToolCompletionsAt(key) + require.Equal(t, map[string]time.Time{ + "call-1": firstCompletedAt, + "call-2": secondCompletedAt, + }, completions) + + // The returned map is a copy: mutating it must not corrupt the + // episode's recorded completions. + completions["call-3"] = clock.Now() + require.Len(t, buffer.ToolCompletionsAt(key), 2) + + // A closed episode keeps its recorded completions but accepts no + // more, matching the batch-start stamp's lifecycle. + require.NoError(t, buffer.CloseEpisode(key)) + require.ErrorIs(t, buffer.RecordToolCompletion(key, "call-3", clock.Now()), messagepartbuffer.ErrEpisodeClosed) + require.Len(t, buffer.ToolCompletionsAt(key), 2) +} + func TestBuffer_SubscribeExistingReplaysThenStreamsLiveParts(t *testing.T) { t.Parallel() diff --git a/coderd/x/chatd/tasks.go b/coderd/x/chatd/tasks.go index ece4bee2739..d2cc31cc18d 100644 --- a/coderd/x/chatd/tasks.go +++ b/coderd/x/chatd/tasks.go @@ -262,6 +262,7 @@ func (s *taskStarter) StartInterrupt(ctx context.Context, input chatWorkerTaskSt } modelInvokedAt := s.opts.MessagePartBuffer.ModelInvokedAt(key) toolBatchStartedAt := s.opts.MessagePartBuffer.ToolBatchStartedAt(key) + toolCompletions := s.opts.MessagePartBuffer.ToolCompletionsAt(key) if err := s.opts.MessagePartBuffer.CloseEpisode(key); err != nil { if ctx.Err() != nil { return errors.Join(errTaskExpectedExit, xerrors.Errorf("close message part episode: %w", err), ctx.Err()) @@ -304,8 +305,8 @@ func (s *taskStarter) StartInterrupt(ctx context.Context, input chatWorkerTaskSt } messages := partialMessages committedCancels, err := committedPendingLocalToolCancellationMessages(ctx, store, chat, s.opts.Clock.Now("chatworker", "interrupt"), interruptedToolBatchBilling{ - batchStartedAt: toolBatchStartedAt, - bufferedParts: parts, + batchStartedAt: toolBatchStartedAt, + toolCompletions: toolCompletions, }) if err != nil { return xerrors.Errorf("committed pending local tool cancellation messages: %w", err) @@ -691,12 +692,13 @@ type interruptedToolBatchBilling struct { // interrupt (crash recovery, state promotion), in which case the // cancellation rows carry no runtime. batchStartedAt time.Time - // bufferedParts are the interrupted episode's buffered parts. Tool - // results buffered before the interrupt carry the completion times - // of tools that finished early, so a batch whose billed tools all - // completed does not bill the longer window of a still-running - // unbilled tool such as wait_agent. - bufferedParts []messagepartbuffer.Part + // toolCompletions maps tool call IDs to the completion instants the + // live batch recorded as each tool finished. Tools with a recorded + // completion end their billable window there, so a batch whose + // billed tools all finished early does not bill the longer window + // of a still-running unbilled tool such as wait_agent. Tools + // without one were still running when the interrupt landed. + toolCompletions map[string]time.Time } func committedPendingLocalToolCancellationMessages( @@ -720,7 +722,6 @@ func committedPendingLocalToolCancellationMessages( if len(localCalls) == 0 { return nil, nil } - completions := bufferedToolResultCompletions(billing.bufferedParts) var ( windowEnd time.Time windowRowIdx = -1 @@ -749,11 +750,11 @@ func committedPendingLocalToolCancellationMessages( if billing.batchStartedAt.IsZero() || unbilledSubagentToolNames[call.ToolName] { continue } - // A billed tool that buffered a durable result completed at that - // instant; one without a buffered result was still running, so - // its window ends at the interrupt. Strictly-after keeps the - // earliest call on ties, matching billableBatchWindow. - end, ok := completions[call.ToolCallID] + // A billed tool with a recorded completion finished at that + // instant; one without was still running, so its window ends + // at the interrupt. Strictly-after keeps the earliest call on + // ties, matching billableBatchWindow. + end, ok := billing.toolCompletions[call.ToolCallID] if !ok { end = interruptedAt } @@ -769,32 +770,3 @@ func committedPendingLocalToolCancellationMessages( } return result, nil } - -// bufferedToolResultCompletions maps tool call IDs to the completion -// instants of durable tool results buffered before an interrupt. -// Streaming deltas and resets are not completions, and only the first -// durable result per call counts, matching the buffered-part conversion -// rules in bufferedPartsToPartialMessages. -func bufferedToolResultCompletions(parts []messagepartbuffer.Part) map[string]time.Time { - completions := make(map[string]time.Time) - for _, buffered := range parts { - if buffered.Role != codersdk.ChatMessageRoleTool { - continue - } - part := buffered.MessagePart - if part.Type != codersdk.ChatMessagePartTypeToolResult || part.ToolCallID == "" { - continue - } - if part.ResultDelta != "" || part.ResultReset || len(part.Result) == 0 { - continue - } - if part.CreatedAt == nil { - continue - } - if _, ok := completions[part.ToolCallID]; ok { - continue - } - completions[part.ToolCallID] = *part.CreatedAt - } - return completions -} diff --git a/coderd/x/chatd/tasks_test.go b/coderd/x/chatd/tasks_test.go index d12357024a6..0947bb42e87 100644 --- a/coderd/x/chatd/tasks_test.go +++ b/coderd/x/chatd/tasks_test.go @@ -572,8 +572,8 @@ func findToolResultMessage(t *testing.T, messages []database.ChatMessage, toolCa } // An interrupt mid tool batch bills the partial window on the cancellation -// row of the billed tool that defines it. A billed tool that buffered a -// result before the interrupt ends the window at its completion, so a +// row of the billed tool that defines it. A billed tool that recorded its +// completion before the interrupt ends the window there, so a // still-running unbilled wait_agent does not stretch the bill to the // interrupt instant. func TestInterruptTask_ToolBatchBillsPartialWindowOnCancellationRow(t *testing.T) { @@ -593,13 +593,12 @@ func TestInterruptTask_ToolBatchBillsPartialWindowOnCancellationRow(t *testing.T // Attempt setup happens before the tools start and is not billable. batch.clock.Advance(2 * time.Second) require.NoError(t, buffer.StartToolBatch(batch.key)) - // execute completes 3 seconds into the batch and buffers its - // durable result. + // execute completes 3 seconds into the batch and records its + // completion, the way the tool goroutine's completion callback + // does. Its result is not published: results publish only after + // the whole batch finishes. batch.clock.Advance(3 * time.Second) - execCompletedAt := batch.clock.Now() - execResult := codersdk.ChatMessageToolResult(execCallID, "execute", json.RawMessage(`{"stdout":"/tmp"}`), false, false) - execResult.CreatedAt = &execCompletedAt - require.NoError(t, buffer.AddPart(batch.key, codersdk.ChatMessageRoleTool, execResult)) + require.NoError(t, buffer.RecordToolCompletion(batch.key, execCallID, batch.clock.Now())) // wait_agent is still blocked on its child when the interrupt lands // 5 seconds later. batch.clock.Advance(5 * time.Second) From 589bfce2a19632c4209b8f494efd1deb00e61714 Mon Sep 17 00:00:00 2001 From: Jaayden Halko Date: Tue, 18 Aug 2026 05:26:27 +0000 Subject: [PATCH 03/20] fix(coderd/x/chatd): track batch billing per occurrence and per dispatched call Address Codex review feedback: - Compute the batch's billable window from completion instants aligned with the tool calls by occurrence instead of the ID-keyed timestamp map. Duplicate tool call IDs, which reach execution when lifecycle hooks are disabled, previously overwrote each other there, letting a short duplicate shrink the window a longer one should have defined. - Seed the buffer episode's tool batch with the dispatched call IDs so an interrupt can tell a call that was still running (seeded, no completion) from one rejected before execution (absent). Rejected calls previously looked like still-running work and billed the whole window up to the interrupt even when only unbilled tools actually ran. --- coderd/x/chatd/chatloop/chatloop.go | 31 +++++++++---- coderd/x/chatd/chatloop/runtime_test.go | 45 +++++++++++++++++++ coderd/x/chatd/generation.go | 27 +++++++---- .../messagepartbuffer/message_part_buffer.go | 38 +++++++++++----- .../message_part_buffer_test.go | 22 ++++++--- coderd/x/chatd/tasks.go | 33 +++++++++----- coderd/x/chatd/tasks_test.go | 35 +++++++++++++-- 7 files changed, 182 insertions(+), 49 deletions(-) diff --git a/coderd/x/chatd/chatloop/chatloop.go b/coderd/x/chatd/chatloop/chatloop.go index f642e590208..d475214b3fb 100644 --- a/coderd/x/chatd/chatloop/chatloop.go +++ b/coderd/x/chatd/chatloop/chatloop.go @@ -622,6 +622,12 @@ func ExecuteLocalTools(ctx context.Context, opts ExecuteLocalToolsOptions) (Tool maxResultBytes := toolResultByteBudget(opts.ContextLimit) batchStart := clockNow(opts.Clock) + // Completion instants aligned with localCalls by occurrence. + // billableBatchWindow reads these instead of the ID-keyed + // ToolResultCreatedAt map, where duplicate tool call IDs, which + // reach execution when lifecycle hooks are disabled, would + // overwrite each other and corrupt the window. + orderedCompletions := make([]time.Time, 0, len(localCalls)) toolResults := executeTools( ctx, opts.Clock, @@ -638,6 +644,10 @@ func ExecuteLocalTools(ctx context.Context, opts ExecuteLocalToolsOptions) (Tool opts.ToolNameAliases, opts.OnToolComplete, func(tr fantasy.ToolResultContent, completedAt time.Time) { + // onResult fires once per local call in call order, so + // appending keeps orderedCompletions aligned with + // localCalls even when tool call IDs collide. + orderedCompletions = append(orderedCompletions, completedAt) recordToolResultTimestamp(&result, tr.ToolCallID, completedAt) publishToolAttachments(ctx, opts.Logger, tr, completedAt, publishMessagePart) ssePart := chatprompt.PartFromContentWithLogger(ctx, opts.Logger, tr) @@ -654,7 +664,7 @@ func ExecuteLocalTools(ctx context.Context, opts ExecuteLocalToolsOptions) (Tool batchRuntime, batchRuntimeToolCallID := billableBatchWindow( batchStart, localCalls, - result.toolResultCreatedAt, + orderedCompletions, opts.UnbilledToolNames, ) return ToolExecutionOutcome{ @@ -673,25 +683,30 @@ func ExecuteLocalTools(ctx context.Context, opts ExecuteLocalToolsOptions) (Tool // union of the billed tools' execution intervals, so parallel calls are // billed once rather than summed, and unbilled tools (for example // sub-agent orchestration) never extend the window even when they run -// longest. Returns the tool call whose completion ends the window, with -// ties broken by call order; (0, "") when no billed tool completed or -// the window rounds to nothing. +// longest. completions is aligned with toolCalls by occurrence, so +// duplicate tool call IDs each keep their own completion. Returns the +// tool call whose completion ends the window, with ties broken by call +// order; (0, "") when no billed tool completed or the window rounds to +// nothing. func billableBatchWindow( batchStart time.Time, toolCalls []fantasy.ToolCallContent, - completedAt map[string]time.Time, + completions []time.Time, unbilledToolNames map[string]bool, ) (time.Duration, string) { var ( windowEnd time.Time toolCallID string ) - for _, tc := range toolCalls { + for i, tc := range toolCalls { + if i >= len(completions) { + break + } if unbilledToolNames[tc.ToolName] { continue } - end, ok := completedAt[tc.ToolCallID] - if !ok { + end := completions[i] + if end.IsZero() { continue } // Strictly-after keeps the earliest call on ties. diff --git a/coderd/x/chatd/chatloop/runtime_test.go b/coderd/x/chatd/chatloop/runtime_test.go index a4a79a0eec7..28748dbc965 100644 --- a/coderd/x/chatd/chatloop/runtime_test.go +++ b/coderd/x/chatd/chatloop/runtime_test.go @@ -392,6 +392,51 @@ func TestExecuteLocalTools_AliasNamesClassifyAsCalled(t *testing.T) { require.Equal(t, "call-execute", outcome.BatchRuntimeToolCallID) } +// Duplicate tool call IDs, which reach execution when lifecycle hooks +// are disabled, must not corrupt the window: completions are tracked per +// occurrence, so a later short duplicate cannot overwrite an earlier +// long one and shrink the bill. +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"}, + // Both calls share one ID: an ID-keyed completion map would + // let the fast occurrence overwrite the slow one. + ToolCalls: []fantasy.ToolCallContent{ + {ToolCallID: "call-dup", ToolName: "slow_tool", Input: "{}"}, + {ToolCallID: "call-dup", ToolName: "fast_tool", Input: "{}"}, + }, + }) + + // Batch start. + trap.MustWait(ctx).MustRelease(ctx) + // The fast occurrence completes at 10 seconds. + clock.Advance(10 * time.Second) + close(fastGo) + trap.MustWait(ctx).MustRelease(ctx) + // The slow occurrence completes at 60 seconds and must define the + // window even though the fast occurrence shares its ID. + 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, "call-dup", outcome.BatchRuntimeToolCallID) +} + // OnToolComplete reports each tool's completion the instant it finishes, // while slower siblings are still running, with the same instants the // outcome's ToolResultCreatedAt later carries. The interrupt path diff --git a/coderd/x/chatd/generation.go b/coderd/x/chatd/generation.go index 52630d4fc5a..2d595c246ba 100644 --- a/coderd/x/chatd/generation.go +++ b/coderd/x/chatd/generation.go @@ -824,9 +824,16 @@ func (s *taskStarter) executeLocalTools( var outcome chatloop.ToolExecutionOutcome var spawnDispatchErr error if len(allowed) > 0 { - // Stamp the batch start on the buffer episode so an interrupt - // can bill the partial window this step would have reported. - attempt.startToolBatch() + // Stamp the batch start and the dispatched call IDs on the + // buffer episode so an interrupt can bill the partial window + // this step would have reported. Only allowed calls are + // listed: denied calls never run, so an interrupt must not + // treat their missing completions as still-running work. + allowedCallIDs := make([]string, 0, len(allowed)) + for _, tc := range allowed { + allowedCallIDs = append(allowedCallIDs, tc.ToolCallID) + } + attempt.startToolBatch(allowedCallIDs) outcome, err = chatloop.ExecuteLocalTools(ctx, chatloop.ExecuteLocalToolsOptions{ Tools: prepared.Tools, ActiveTools: prepared.ActiveTools, @@ -1056,10 +1063,12 @@ type generationAttempt struct { // non-nil when beginGenerationAttempt succeeds. startModelInvocation func() // startToolBatch marks the start of the attempt's billable local - // tool batch window on the buffer episode, so an interrupt can - // bill the window the step would have reported. It is always - // non-nil when beginGenerationAttempt succeeds. - startToolBatch func() + // tool batch window on the buffer episode and records which tool + // calls were actually dispatched, so an interrupt can bill the + // window the step would have reported without charging calls that + // were rejected before execution. It is always non-nil when + // beginGenerationAttempt succeeds. + startToolBatch func(toolCallIDs []string) // recordToolCompletion records a tool call's completion instant on // the buffer episode as the batch executes, so an interrupt can // end an already-finished tool's billable window at its real @@ -1112,8 +1121,8 @@ func (s *taskStarter) beginGenerationAttempt( startModelInvocation: func() { _ = s.opts.MessagePartBuffer.StartModelInvocation(key) }, - startToolBatch: func() { - _ = s.opts.MessagePartBuffer.StartToolBatch(key) + startToolBatch: func(toolCallIDs []string) { + _ = s.opts.MessagePartBuffer.StartToolBatch(key, toolCallIDs) }, recordToolCompletion: func(toolCallID string, completedAt time.Time) { _ = s.opts.MessagePartBuffer.RecordToolCompletion(key, toolCallID, completedAt) diff --git a/coderd/x/chatd/messagepartbuffer/message_part_buffer.go b/coderd/x/chatd/messagepartbuffer/message_part_buffer.go index 1a482b60cf9..0ef5db96c71 100644 --- a/coderd/x/chatd/messagepartbuffer/message_part_buffer.go +++ b/coderd/x/chatd/messagepartbuffer/message_part_buffer.go @@ -108,11 +108,12 @@ type episodeState struct { // episodes that never execute local tools, such as model // invocations that finish without tool calls. toolBatchStartedAt time.Time - // toolCompletedAt maps tool call IDs to the completion instants - // recorded by RecordToolCompletion as the batch's tools finish. - // Tool results are published only after the whole batch - // completes, so these per-tool stamps are the only live view of - // which tools already finished when an interrupt lands. + // toolCompletedAt holds the batch's dispatched tool call IDs, + // seeded with zero times by StartToolBatch and stamped by + // RecordToolCompletion as each tool finishes. Tool results are + // published only after the whole batch completes, so these + // per-tool stamps are the only live view of which tools were + // dispatched and which already finished when an interrupt lands. toolCompletedAt map[string]time.Time closed bool closedAt time.Time @@ -231,8 +232,12 @@ func (b *Buffer) StartModelInvocation(key Key) error { } // StartToolBatch stamps the instant the episode begins executing its local -// tool batch, which starts the batch's billable runtime window. -func (b *Buffer) StartToolBatch(key Key) error { +// tool batch, which starts the batch's billable runtime window, and seeds +// the batch's dispatched tool call IDs with zero completions. Readers can +// then distinguish a dispatched call that is still running (present, zero) +// from one never dispatched at all (absent), such as a call denied by a +// lifecycle hook or rejected as ambiguous before execution. +func (b *Buffer) StartToolBatch(key Key, toolCallIDs []string) error { b.mu.Lock() defer b.mu.Unlock() if b.closed { @@ -246,6 +251,14 @@ func (b *Buffer) StartToolBatch(key Key) error { return ErrEpisodeClosed } episode.toolBatchStartedAt = b.opts.Clock.Now("message-part-buffer", "tool-batch-start") + if episode.toolCompletedAt == nil { + episode.toolCompletedAt = make(map[string]time.Time, len(toolCallIDs)) + } + for _, id := range toolCallIDs { + if _, ok := episode.toolCompletedAt[id]; !ok { + episode.toolCompletedAt[id] = time.Time{} + } + } return nil } @@ -372,11 +385,12 @@ func (b *Buffer) ToolBatchStartedAt(key Key) time.Time { return episode.toolBatchStartedAt } -// ToolCompletionsAt returns a copy of the completion instants recorded -// by RecordToolCompletion, keyed by tool call ID, or nil when the -// episode is unknown or recorded none. Read it before CloseEpisode: -// closed episodes are garbage collected, so reading afterwards races -// the cleanup loop. +// ToolCompletionsAt returns a copy of the tool batch's completion map, +// keyed by tool call ID, or nil when the episode is unknown or never +// started a batch. Calls seeded by StartToolBatch but not yet stamped +// by RecordToolCompletion carry the zero time: they were dispatched and +// are still running. Read it before CloseEpisode: closed episodes are +// garbage collected, so reading afterwards races the cleanup loop. func (b *Buffer) ToolCompletionsAt(key Key) map[string]time.Time { b.mu.Lock() defer b.mu.Unlock() diff --git a/coderd/x/chatd/messagepartbuffer/message_part_buffer_test.go b/coderd/x/chatd/messagepartbuffer/message_part_buffer_test.go index 9d4f1845312..d81a7a509a5 100644 --- a/coderd/x/chatd/messagepartbuffer/message_part_buffer_test.go +++ b/coderd/x/chatd/messagepartbuffer/message_part_buffer_test.go @@ -157,14 +157,14 @@ func TestBuffer_ToolBatchStartedAt(t *testing.T) { key := testEpisodeKey() require.Zero(t, buffer.ToolBatchStartedAt(key), "unknown episode has no batch stamp") - require.ErrorIs(t, buffer.StartToolBatch(key), messagepartbuffer.ErrEpisodeNotFound) + require.ErrorIs(t, buffer.StartToolBatch(key, []string{"call-1"}), messagepartbuffer.ErrEpisodeNotFound) require.NoError(t, buffer.CreateEpisode(key)) require.Zero(t, buffer.ToolBatchStartedAt(key), "episode without a tool batch has no batch stamp") // Attempt setup happens before tools start executing and is not // billable. clock.Advance(time.Second) - require.NoError(t, buffer.StartToolBatch(key)) + require.NoError(t, buffer.StartToolBatch(key, []string{"call-1"})) startedAt := buffer.ToolBatchStartedAt(key) require.Equal(t, clock.Now(), startedAt) @@ -172,7 +172,7 @@ func TestBuffer_ToolBatchStartedAt(t *testing.T) { // no longer accepts a batch start. clock.Advance(1500 * time.Millisecond) require.NoError(t, buffer.CloseEpisode(key)) - require.ErrorIs(t, buffer.StartToolBatch(key), messagepartbuffer.ErrEpisodeClosed) + require.ErrorIs(t, buffer.StartToolBatch(key, []string{"call-1"}), messagepartbuffer.ErrEpisodeClosed) require.Equal(t, startedAt, buffer.ToolBatchStartedAt(key)) // Episodes that never execute local tools, such as model @@ -197,11 +197,23 @@ func TestBuffer_ToolCompletionsAt(t *testing.T) { require.ErrorIs(t, buffer.RecordToolCompletion(key, "call-1", clock.Now()), messagepartbuffer.ErrEpisodeNotFound) require.NoError(t, buffer.CreateEpisode(key)) - require.Empty(t, buffer.ToolCompletionsAt(key), "episode without recorded completions has none") - // Tools complete at different instants; each keeps its own stamp. + require.Empty(t, buffer.ToolCompletionsAt(key), "episode without a tool batch has no completions") + // Starting the batch seeds every dispatched call with a zero + // completion, marking it dispatched but still running. + require.NoError(t, buffer.StartToolBatch(key, []string{"call-1", "call-2"})) + require.Equal(t, map[string]time.Time{ + "call-1": {}, + "call-2": {}, + }, buffer.ToolCompletionsAt(key)) + // Tools complete at different instants; each keeps its own stamp + // while the still-running call stays zero. clock.Advance(time.Second) firstCompletedAt := clock.Now() require.NoError(t, buffer.RecordToolCompletion(key, "call-1", firstCompletedAt)) + require.Equal(t, map[string]time.Time{ + "call-1": firstCompletedAt, + "call-2": {}, + }, buffer.ToolCompletionsAt(key)) clock.Advance(2 * time.Second) secondCompletedAt := clock.Now() require.NoError(t, buffer.RecordToolCompletion(key, "call-2", secondCompletedAt)) diff --git a/coderd/x/chatd/tasks.go b/coderd/x/chatd/tasks.go index d2cc31cc18d..1912d3fab6b 100644 --- a/coderd/x/chatd/tasks.go +++ b/coderd/x/chatd/tasks.go @@ -692,12 +692,16 @@ type interruptedToolBatchBilling struct { // interrupt (crash recovery, state promotion), in which case the // cancellation rows carry no runtime. batchStartedAt time.Time - // toolCompletions maps tool call IDs to the completion instants the - // live batch recorded as each tool finished. Tools with a recorded - // completion end their billable window there, so a batch whose - // billed tools all finished early does not bill the longer window - // of a still-running unbilled tool such as wait_agent. Tools - // without one were still running when the interrupt landed. + // toolCompletions holds the live batch's dispatched tool call IDs: + // seeded with zero times when the batch started and stamped with + // completion instants as each tool finished. A stamped tool ends + // its billable window at its completion, so a batch whose billed + // tools all finished early does not bill the longer window of a + // still-running unbilled tool such as wait_agent. A seeded but + // unstamped tool was still running when the interrupt landed. A + // tool absent from the map was never dispatched, such as a call + // denied by a lifecycle hook or rejected as ambiguous, and bills + // nothing even though it too receives a cancellation row. toolCompletions map[string]time.Time } @@ -750,12 +754,17 @@ func committedPendingLocalToolCancellationMessages( if billing.batchStartedAt.IsZero() || unbilledSubagentToolNames[call.ToolName] { continue } - // A billed tool with a recorded completion finished at that - // instant; one without was still running, so its window ends - // at the interrupt. Strictly-after keeps the earliest call on - // ties, matching billableBatchWindow. - end, ok := billing.toolCompletions[call.ToolCallID] - if !ok { + // Only dispatched calls bill: a call absent from the batch's + // completion map was rejected before execution. A dispatched + // call with a stamped completion finished at that instant; a + // seeded but unstamped one was still running, so its window + // ends at the interrupt. Strictly-after keeps the earliest + // call on ties, matching billableBatchWindow. + end, dispatched := billing.toolCompletions[call.ToolCallID] + if !dispatched { + continue + } + if end.IsZero() { end = interruptedAt } if end.After(windowEnd) { diff --git a/coderd/x/chatd/tasks_test.go b/coderd/x/chatd/tasks_test.go index 0947bb42e87..a8b89193a56 100644 --- a/coderd/x/chatd/tasks_test.go +++ b/coderd/x/chatd/tasks_test.go @@ -592,7 +592,7 @@ func TestInterruptTask_ToolBatchBillsPartialWindowOnCancellationRow(t *testing.T // this mock clock. // Attempt setup happens before the tools start and is not billable. batch.clock.Advance(2 * time.Second) - require.NoError(t, buffer.StartToolBatch(batch.key)) + require.NoError(t, buffer.StartToolBatch(batch.key, []string{execCallID, waitCallID})) // execute completes 3 seconds into the batch and records its // completion, the way the tool goroutine's completion callback // does. Its result is not published: results publish only after @@ -622,7 +622,7 @@ func TestInterruptTask_RunningBilledToolBillsUpToInterrupt(t *testing.T) { }) batch.clock.Advance(2 * time.Second) - require.NoError(t, batch.starter.opts.MessagePartBuffer.StartToolBatch(batch.key)) + require.NoError(t, batch.starter.opts.MessagePartBuffer.StartToolBatch(batch.key, []string{execCallID})) // The tool is still running when the interrupt lands 7 seconds in. batch.clock.Advance(7 * time.Second) @@ -643,7 +643,7 @@ func TestInterruptTask_UnbilledOnlyBatchBillsNothingOnInterrupt(t *testing.T) { }) batch.clock.Advance(2 * time.Second) - require.NoError(t, batch.starter.opts.MessagePartBuffer.StartToolBatch(batch.key)) + require.NoError(t, batch.starter.opts.MessagePartBuffer.StartToolBatch(batch.key, []string{waitCallID})) batch.clock.Advance(10 * time.Second) messages := batch.interrupt(t, f) @@ -651,6 +651,35 @@ func TestInterruptTask_UnbilledOnlyBatchBillsNothingOnInterrupt(t *testing.T) { require.False(t, waitRow.RuntimeMs.Valid) } +// A billed call rejected before execution (hook denial, ambiguous-call +// rejection) is absent from the batch's dispatched set, so its missing +// completion is not evidence that it ran: an interrupted batch whose +// only dispatched call is an unbilled wait_agent bills nothing even +// though the rejected execute call also receives a cancellation row. +func TestInterruptTask_RejectedCallBillsNothingOnInterrupt(t *testing.T) { + t.Parallel() + + f := newTaskTestFixture(t) + execCallID := "call_" + uuid.NewString() + waitCallID := "call_" + uuid.NewString() + batch := interruptedBatchFixture(t, f, []codersdk.ChatMessagePart{ + {Type: codersdk.ChatMessagePartTypeToolCall, ToolCallID: execCallID, ToolName: "execute", Args: json.RawMessage(`{}`)}, + {Type: codersdk.ChatMessagePartTypeToolCall, ToolCallID: waitCallID, ToolName: "wait_agent", Args: json.RawMessage(`{}`)}, + }) + + batch.clock.Advance(2 * time.Second) + // Only wait_agent was dispatched: execute was rejected before + // execution and never ran. + require.NoError(t, batch.starter.opts.MessagePartBuffer.StartToolBatch(batch.key, []string{waitCallID})) + batch.clock.Advance(10 * time.Second) + + messages := batch.interrupt(t, f) + execRow := findToolResultMessage(t, messages, execCallID) + require.False(t, execRow.RuntimeMs.Valid) + waitRow := findToolResultMessage(t, messages, waitCallID) + require.False(t, waitRow.RuntimeMs.Valid) +} + // Cancellation rows synthesized without a live tool batch (crash // recovery: the episode never stamped a batch start) carry no runtime, // consistent with generation losing in-flight runtime on a crash. From 9ff5f50314b3ced5a5d67faec0cd10349c73d465 Mon Sep 17 00:00:00 2001 From: Jaayden Halko Date: Tue, 18 Aug 2026 05:53:50 +0000 Subject: [PATCH 04/20] fix(coderd/x/chatd): snapshot interrupt billing atomically per call occurrence Address Codex review feedback: - Store the tool batch's dispatched calls on the buffer episode as an occurrence list instead of an ID-keyed map, stamping completions onto the first still-running occurrence of the reported ID. Duplicate tool call IDs, which reach execution when lifecycle hooks are disabled, previously collapsed into one shared state, so one duplicate finishing made its still-running twin look finished and ended the interrupted batch's window early. The interrupt path now consumes per-ID occurrence queues in the same dispatch order the unresolved history rows walk. - Close the buffer episode and snapshot its billing stamps in one critical section (CloseEpisodeForBilling) instead of reading ModelInvokedAt, ToolBatchStartedAt, and the completions before a separate CloseEpisode. Stamps recorded in those gaps went missing from the snapshot, billing a finished tool through to the interrupt or dropping a just-started batch's window entirely. --- coderd/x/chatd/ARCHITECTURE.md | 2 +- .../messagepartbuffer/message_part_buffer.go | 119 +++++++++++++----- .../message_part_buffer_test.go | 103 +++++++++++---- coderd/x/chatd/tasks.go | 67 ++++++---- coderd/x/chatd/tasks_test.go | 33 +++++ 5 files changed, 239 insertions(+), 85 deletions(-) diff --git a/coderd/x/chatd/ARCHITECTURE.md b/coderd/x/chatd/ARCHITECTURE.md index 7880d4fe4f3..5fb8aa47959 100644 --- a/coderd/x/chatd/ARCHITECTURE.md +++ b/coderd/x/chatd/ARCHITECTURE.md @@ -725,7 +725,7 @@ The buffer exposes the following API: - `GetParts(chat_id, history_version, generation_attempt)`: returns the message parts for an episode. Returns a predefined error if the episode is not found. - `StartModelInvocation(chat_id, history_version, generation_attempt)`: stamps the instant the episode opens its provider stream. Returns a predefined error if the episode is not found or already closed. Episodes that never invoke a model, such as local tool execution batches, are never stamped. - `ModelInvokedAt(chat_id, history_version, generation_attempt)`: returns the instant stamped by `StartModelInvocation`, or the zero time when the episode is unknown or never opened a provider stream. It must be read before `CloseEpisode`, because closed episodes are garbage collected and reading afterwards races the cleanup loop. The interrupt goroutine reads it just before closing the episode and uses the span between that instant and the interrupt as the interrupted attempt's billable runtime. -- TODO(CODAGT-928): document the tool-batch billing methods `StartToolBatch`, `ToolBatchStartedAt`, `RecordToolCompletion`, and `ToolCompletionsAt`, the local-tool counterparts of `StartModelInvocation`/`ModelInvokedAt` that the interrupt task reads before `CloseEpisode` to bill an interrupted tool batch's partial window. +- TODO(CODAGT-928): document the tool-batch billing methods `StartToolBatch`, `ToolBatchStartedAt`, `RecordToolCompletion`, and `ToolCompletions`, the local-tool counterparts of `StartModelInvocation`/`ModelInvokedAt`, and `CloseEpisodeForBilling`, which the interrupt task uses to close the episode and snapshot its billing stamps in one atomic step so it can bill an interrupted attempt's partial window. - `SubscribeToEpisode(chat_id, history_version, generation_attempt)`: returns a go channel that will receive all message parts for the episode. It spawns a goroutine that delivers parts to the channel. It's live until the episode is closed or until a subscriber requests that the channel be closed. Once the goroutine delivers all message parts for a closed episode, it closes the channel and exits. If the episode is already closed at the time of the call, the goroutine delivers all message parts for the episode, closes the channel, and exits. `SubscribeToEpisode` does not return an error if the episode is not found: it waits for it to be created instead. Closed episodes are garbage collected after at least 15 seconds since they were closed and when they have no active subscribers. The message part buffer maintains a garbage collection goroutine. diff --git a/coderd/x/chatd/messagepartbuffer/message_part_buffer.go b/coderd/x/chatd/messagepartbuffer/message_part_buffer.go index 0ef5db96c71..7a8cced2403 100644 --- a/coderd/x/chatd/messagepartbuffer/message_part_buffer.go +++ b/coderd/x/chatd/messagepartbuffer/message_part_buffer.go @@ -19,7 +19,6 @@ import ( "container/heap" "context" "encoding/json" - "maps" "slices" "sync" "time" @@ -108,13 +107,16 @@ type episodeState struct { // episodes that never execute local tools, such as model // invocations that finish without tool calls. toolBatchStartedAt time.Time - // toolCompletedAt holds the batch's dispatched tool call IDs, - // seeded with zero times by StartToolBatch and stamped by - // RecordToolCompletion as each tool finishes. Tool results are - // published only after the whole batch completes, so these - // per-tool stamps are the only live view of which tools were - // dispatched and which already finished when an interrupt lands. - toolCompletedAt map[string]time.Time + // toolCompletions holds one entry per dispatched tool call + // occurrence, seeded by StartToolBatch in dispatch order and + // stamped by RecordToolCompletion as each tool finishes. Tool + // results are published only after the whole batch completes, so + // these per-occurrence stamps are the only live view of which + // tools were dispatched and which already finished when an + // interrupt lands. Keyed storage would collapse duplicate tool + // call IDs, which reach execution when lifecycle hooks are + // disabled, into one shared state. + toolCompletions []ToolCompletion closed bool closedAt time.Time closedHeapItem *closedEpisodeItem @@ -231,12 +233,21 @@ func (b *Buffer) StartModelInvocation(key Key) error { return nil } +// ToolCompletion tracks one dispatched tool call occurrence in an +// episode's local tool batch. CompletedAt is zero while the call is +// still running. +type ToolCompletion struct { + ToolCallID string + CompletedAt time.Time +} + // StartToolBatch stamps the instant the episode begins executing its local // tool batch, which starts the batch's billable runtime window, and seeds -// the batch's dispatched tool call IDs with zero completions. Readers can -// then distinguish a dispatched call that is still running (present, zero) -// from one never dispatched at all (absent), such as a call denied by a -// lifecycle hook or rejected as ambiguous before execution. +// one still-running entry per dispatched tool call occurrence. Readers can +// then distinguish a dispatched call that is still running (seeded, zero +// completion) from one never dispatched at all (absent), such as a call +// denied by a lifecycle hook or rejected as ambiguous before execution, +// and duplicate tool call IDs keep distinct per-occurrence states. func (b *Buffer) StartToolBatch(key Key, toolCallIDs []string) error { b.mu.Lock() defer b.mu.Unlock() @@ -251,13 +262,9 @@ func (b *Buffer) StartToolBatch(key Key, toolCallIDs []string) error { return ErrEpisodeClosed } episode.toolBatchStartedAt = b.opts.Clock.Now("message-part-buffer", "tool-batch-start") - if episode.toolCompletedAt == nil { - episode.toolCompletedAt = make(map[string]time.Time, len(toolCallIDs)) - } + episode.toolCompletions = make([]ToolCompletion, 0, len(toolCallIDs)) for _, id := range toolCallIDs { - if _, ok := episode.toolCompletedAt[id]; !ok { - episode.toolCompletedAt[id] = time.Time{} - } + episode.toolCompletions = append(episode.toolCompletions, ToolCompletion{ToolCallID: id}) } return nil } @@ -266,7 +273,11 @@ func (b *Buffer) StartToolBatch(key Key, toolCallIDs []string) error { // episode's tool batch finished. Tool goroutines report completions as // they happen, so an interrupt can bill tools that already finished up // to their real completion instead of treating every canceled call as -// still running. +// still running. The stamp lands on the first still-running occurrence +// with the given ID: same-ID occurrences are indistinguishable to +// callers, and assigning them in seeded order keeps the batch's set of +// completion states correct. A completion whose ID was never seeded is +// appended, so an executed call is never dropped. func (b *Buffer) RecordToolCompletion(key Key, toolCallID string, completedAt time.Time) error { b.mu.Lock() defer b.mu.Unlock() @@ -280,10 +291,17 @@ func (b *Buffer) RecordToolCompletion(key Key, toolCallID string, completedAt ti if episode.closed { return ErrEpisodeClosed } - if episode.toolCompletedAt == nil { - episode.toolCompletedAt = make(map[string]time.Time) + for i := range episode.toolCompletions { + entry := &episode.toolCompletions[i] + if entry.ToolCallID == toolCallID && entry.CompletedAt.IsZero() { + entry.CompletedAt = completedAt + return nil + } } - episode.toolCompletedAt[toolCallID] = completedAt + episode.toolCompletions = append(episode.toolCompletions, ToolCompletion{ + ToolCallID: toolCallID, + CompletedAt: completedAt, + }) return nil } @@ -385,20 +403,61 @@ func (b *Buffer) ToolBatchStartedAt(key Key) time.Time { return episode.toolBatchStartedAt } -// ToolCompletionsAt returns a copy of the tool batch's completion map, -// keyed by tool call ID, or nil when the episode is unknown or never -// started a batch. Calls seeded by StartToolBatch but not yet stamped -// by RecordToolCompletion carry the zero time: they were dispatched and -// are still running. Read it before CloseEpisode: closed episodes are -// garbage collected, so reading afterwards races the cleanup loop. -func (b *Buffer) ToolCompletionsAt(key Key) map[string]time.Time { +// ToolCompletions returns a copy of the tool batch's dispatched call +// occurrences in dispatch order, or nil when the episode is unknown or +// never started a batch. Occurrences seeded by StartToolBatch but not +// yet stamped by RecordToolCompletion carry the zero time: they were +// dispatched and are still running. Interrupt handling must use +// CloseEpisodeForBilling instead: a separate read-then-close would let +// completions land in the gap and go missing from the snapshot. +func (b *Buffer) ToolCompletions(key Key) []ToolCompletion { b.mu.Lock() defer b.mu.Unlock() episode := b.episodes[key] if episode == nil { return nil } - return maps.Clone(episode.toolCompletedAt) + return slices.Clone(episode.toolCompletions) +} + +// EpisodeBilling is the billing state an episode accumulated before it +// closed. +type EpisodeBilling struct { + // ModelInvokedAt is the StartModelInvocation stamp, or zero when + // the episode never opened a provider stream. + ModelInvokedAt time.Time + // ToolBatchStartedAt is the StartToolBatch stamp, or zero when the + // episode never started a local tool batch. + ToolBatchStartedAt time.Time + // ToolCompletions are the tool batch's dispatched call occurrences + // in dispatch order; see Buffer.ToolCompletions. + ToolCompletions []ToolCompletion +} + +// CloseEpisodeForBilling closes the episode like CloseEpisode and +// returns its billing stamps from the same critical section. The +// interrupt task uses it so every stamp accepted before closure is in +// the snapshot and none can be recorded afterwards: reading and closing +// in separate steps would let a tool completion or batch start land in +// the gap, billing a finished tool as still running or losing a live +// batch's window entirely. Closing an unknown episode creates it +// closed, mirroring CloseEpisode, and reports empty billing. +func (b *Buffer) CloseEpisodeForBilling(key Key) (EpisodeBilling, error) { + b.mu.Lock() + defer b.mu.Unlock() + if b.closed { + return EpisodeBilling{}, ErrMessagePartBufferClosed + } + episode := b.getOrCreateEpisodeLocked(key) + if episode.close(b.opts.Clock.Now("message-part-buffer", "close")) { + b.queueClosedEpisodeLocked(key, episode) + episode.notifySubscribers() + } + return EpisodeBilling{ + ModelInvokedAt: episode.modelStartedAt, + ToolBatchStartedAt: episode.toolBatchStartedAt, + ToolCompletions: slices.Clone(episode.toolCompletions), + }, nil } // SubscribeToEpisode replays existing parts and streams new parts. diff --git a/coderd/x/chatd/messagepartbuffer/message_part_buffer_test.go b/coderd/x/chatd/messagepartbuffer/message_part_buffer_test.go index d81a7a509a5..edbb64003f2 100644 --- a/coderd/x/chatd/messagepartbuffer/message_part_buffer_test.go +++ b/coderd/x/chatd/messagepartbuffer/message_part_buffer_test.go @@ -185,7 +185,7 @@ func TestBuffer_ToolBatchStartedAt(t *testing.T) { require.Zero(t, buffer.ToolBatchStartedAt(modelOnly)) } -func TestBuffer_ToolCompletionsAt(t *testing.T) { +func TestBuffer_ToolCompletions(t *testing.T) { t.Parallel() clock := quartz.NewMock(t) @@ -193,46 +193,95 @@ func TestBuffer_ToolCompletionsAt(t *testing.T) { defer buffer.Close() key := testEpisodeKey() - require.Nil(t, buffer.ToolCompletionsAt(key), "unknown episode has no completions") + require.Nil(t, buffer.ToolCompletions(key), "unknown episode has no completions") require.ErrorIs(t, buffer.RecordToolCompletion(key, "call-1", clock.Now()), messagepartbuffer.ErrEpisodeNotFound) require.NoError(t, buffer.CreateEpisode(key)) - require.Empty(t, buffer.ToolCompletionsAt(key), "episode without a tool batch has no completions") - // Starting the batch seeds every dispatched call with a zero - // completion, marking it dispatched but still running. - require.NoError(t, buffer.StartToolBatch(key, []string{"call-1", "call-2"})) - require.Equal(t, map[string]time.Time{ - "call-1": {}, - "call-2": {}, - }, buffer.ToolCompletionsAt(key)) - // Tools complete at different instants; each keeps its own stamp - // while the still-running call stays zero. + require.Empty(t, buffer.ToolCompletions(key), "episode without a tool batch has no completions") + // Starting the batch seeds one still-running occurrence per + // dispatched call, in dispatch order. Duplicate IDs keep distinct + // occurrences instead of collapsing into one shared state. + require.NoError(t, buffer.StartToolBatch(key, []string{"call-1", "call-dup", "call-dup"})) + require.Equal(t, []messagepartbuffer.ToolCompletion{ + {ToolCallID: "call-1"}, + {ToolCallID: "call-dup"}, + {ToolCallID: "call-dup"}, + }, buffer.ToolCompletions(key)) + // A completion stamps the first still-running occurrence with its + // ID; the duplicate's second occurrence stays running. clock.Advance(time.Second) firstCompletedAt := clock.Now() - require.NoError(t, buffer.RecordToolCompletion(key, "call-1", firstCompletedAt)) - require.Equal(t, map[string]time.Time{ - "call-1": firstCompletedAt, - "call-2": {}, - }, buffer.ToolCompletionsAt(key)) + require.NoError(t, buffer.RecordToolCompletion(key, "call-dup", firstCompletedAt)) + require.Equal(t, []messagepartbuffer.ToolCompletion{ + {ToolCallID: "call-1"}, + {ToolCallID: "call-dup", CompletedAt: firstCompletedAt}, + {ToolCallID: "call-dup"}, + }, buffer.ToolCompletions(key)) clock.Advance(2 * time.Second) secondCompletedAt := clock.Now() - require.NoError(t, buffer.RecordToolCompletion(key, "call-2", secondCompletedAt)) - completions := buffer.ToolCompletionsAt(key) - require.Equal(t, map[string]time.Time{ - "call-1": firstCompletedAt, - "call-2": secondCompletedAt, + require.NoError(t, buffer.RecordToolCompletion(key, "call-dup", secondCompletedAt)) + completions := buffer.ToolCompletions(key) + require.Equal(t, []messagepartbuffer.ToolCompletion{ + {ToolCallID: "call-1"}, + {ToolCallID: "call-dup", CompletedAt: firstCompletedAt}, + {ToolCallID: "call-dup", CompletedAt: secondCompletedAt}, }, completions) - // The returned map is a copy: mutating it must not corrupt the + // The returned slice is a copy: mutating it must not corrupt the // episode's recorded completions. - completions["call-3"] = clock.Now() - require.Len(t, buffer.ToolCompletionsAt(key), 2) + completions[0].CompletedAt = clock.Now() + require.True(t, buffer.ToolCompletions(key)[0].CompletedAt.IsZero()) // A closed episode keeps its recorded completions but accepts no // more, matching the batch-start stamp's lifecycle. require.NoError(t, buffer.CloseEpisode(key)) - require.ErrorIs(t, buffer.RecordToolCompletion(key, "call-3", clock.Now()), messagepartbuffer.ErrEpisodeClosed) - require.Len(t, buffer.ToolCompletionsAt(key), 2) + require.ErrorIs(t, buffer.RecordToolCompletion(key, "call-1", clock.Now()), messagepartbuffer.ErrEpisodeClosed) + require.True(t, buffer.ToolCompletions(key)[0].CompletedAt.IsZero()) +} + +func TestBuffer_CloseEpisodeForBilling(t *testing.T) { + t.Parallel() + + clock := quartz.NewMock(t) + buffer := messagepartbuffer.New(messagepartbuffer.Options{Clock: clock}) + defer buffer.Close() + + // Closing an unknown episode creates it closed, like CloseEpisode, + // and reports empty billing. + unknown := testEpisodeKey() + billing, err := buffer.CloseEpisodeForBilling(unknown) + require.NoError(t, err) + require.Zero(t, billing.ModelInvokedAt) + require.Zero(t, billing.ToolBatchStartedAt) + require.Empty(t, billing.ToolCompletions) + + // The snapshot carries every stamp accepted before closure, and + // stamps are rejected afterwards, so nothing can land in a gap + // between reading and closing. + key := testEpisodeKey() + require.NoError(t, buffer.CreateEpisode(key)) + clock.Advance(time.Second) + batchStartedAt := clock.Now() + require.NoError(t, buffer.StartToolBatch(key, []string{"call-1", "call-2"})) + clock.Advance(time.Second) + completedAt := clock.Now() + require.NoError(t, buffer.RecordToolCompletion(key, "call-1", completedAt)) + billing, err = buffer.CloseEpisodeForBilling(key) + require.NoError(t, err) + require.Zero(t, billing.ModelInvokedAt) + require.Equal(t, batchStartedAt, billing.ToolBatchStartedAt) + require.Equal(t, []messagepartbuffer.ToolCompletion{ + {ToolCallID: "call-1", CompletedAt: completedAt}, + {ToolCallID: "call-2"}, + }, billing.ToolCompletions) + require.ErrorIs(t, buffer.RecordToolCompletion(key, "call-2", clock.Now()), messagepartbuffer.ErrEpisodeClosed) + + // Closing an already-closed episode still reports its billing, so + // an interrupt racing the generation task's own close loses + // nothing. + again, err := buffer.CloseEpisodeForBilling(key) + require.NoError(t, err) + require.Equal(t, billing, again) } func TestBuffer_SubscribeExistingReplaysThenStreamsLiveParts(t *testing.T) { diff --git a/coderd/x/chatd/tasks.go b/coderd/x/chatd/tasks.go index 1912d3fab6b..40973bfdeb6 100644 --- a/coderd/x/chatd/tasks.go +++ b/coderd/x/chatd/tasks.go @@ -260,10 +260,12 @@ func (s *taskStarter) StartInterrupt(ctx context.Context, input chatWorkerTaskSt HistoryVersion: input.HistoryVersion, GenerationAttempt: chat.GenerationAttempt, } - modelInvokedAt := s.opts.MessagePartBuffer.ModelInvokedAt(key) - toolBatchStartedAt := s.opts.MessagePartBuffer.ToolBatchStartedAt(key) - toolCompletions := s.opts.MessagePartBuffer.ToolCompletionsAt(key) - if err := s.opts.MessagePartBuffer.CloseEpisode(key); err != nil { + // Closing and snapshotting billing state must be one atomic step: + // the generation goroutine records batch starts and tool + // completions concurrently, so a read-then-close would let stamps + // land in the gap and go missing from the snapshot. + episodeBilling, err := s.opts.MessagePartBuffer.CloseEpisodeForBilling(key) + if err != nil { if ctx.Err() != nil { return errors.Join(errTaskExpectedExit, xerrors.Errorf("close message part episode: %w", err), ctx.Err()) } @@ -282,8 +284,8 @@ func (s *taskStarter) StartInterrupt(ctx context.Context, input chatWorkerTaskSt } interruptedAt := s.opts.Clock.Now("chatworker", "interrupt") var attemptRuntime time.Duration - if !modelInvokedAt.IsZero() { - attemptRuntime = interruptedAt.Sub(modelInvokedAt) + if !episodeBilling.ModelInvokedAt.IsZero() { + attemptRuntime = interruptedAt.Sub(episodeBilling.ModelInvokedAt) } partialMessages, err := bufferedPartsToPartialMessages(bufferedPartsToPartialMessagesInput{ parts: parts, @@ -305,8 +307,8 @@ func (s *taskStarter) StartInterrupt(ctx context.Context, input chatWorkerTaskSt } messages := partialMessages committedCancels, err := committedPendingLocalToolCancellationMessages(ctx, store, chat, s.opts.Clock.Now("chatworker", "interrupt"), interruptedToolBatchBilling{ - batchStartedAt: toolBatchStartedAt, - toolCompletions: toolCompletions, + batchStartedAt: episodeBilling.ToolBatchStartedAt, + toolCompletions: episodeBilling.ToolCompletions, }) if err != nil { return xerrors.Errorf("committed pending local tool cancellation messages: %w", err) @@ -692,17 +694,18 @@ type interruptedToolBatchBilling struct { // interrupt (crash recovery, state promotion), in which case the // cancellation rows carry no runtime. batchStartedAt time.Time - // toolCompletions holds the live batch's dispatched tool call IDs: - // seeded with zero times when the batch started and stamped with - // completion instants as each tool finished. A stamped tool ends - // its billable window at its completion, so a batch whose billed - // tools all finished early does not bill the longer window of a - // still-running unbilled tool such as wait_agent. A seeded but - // unstamped tool was still running when the interrupt landed. A - // tool absent from the map was never dispatched, such as a call - // denied by a lifecycle hook or rejected as ambiguous, and bills - // nothing even though it too receives a cancellation row. - toolCompletions map[string]time.Time + // toolCompletions holds the live batch's dispatched tool call + // occurrences in dispatch order: seeded with zero completions when + // the batch started and stamped as each tool finished. A stamped + // occurrence ends its billable window at its completion, so a + // batch whose billed tools all finished early does not bill the + // longer window of a still-running unbilled tool such as + // wait_agent. A seeded but unstamped occurrence was still running + // when the interrupt landed. A call with no matching occurrence + // was never dispatched, such as a call denied by a lifecycle hook + // or rejected as ambiguous, and bills nothing even though it too + // receives a cancellation row. + toolCompletions []messagepartbuffer.ToolCompletion } func committedPendingLocalToolCancellationMessages( @@ -726,6 +729,14 @@ func committedPendingLocalToolCancellationMessages( if len(localCalls) == 0 { return nil, nil } + // Per-ID queues of dispatched occurrences, in dispatch order. + // Unresolved calls walk the same assistant part order the batch was + // dispatched from, so each history row consumes its own occurrence + // and duplicate tool call IDs never share one completion state. + pendingOccurrences := make(map[string][]time.Time, len(billing.toolCompletions)) + for _, completion := range billing.toolCompletions { + pendingOccurrences[completion.ToolCallID] = append(pendingOccurrences[completion.ToolCallID], completion.CompletedAt) + } var ( windowEnd time.Time windowRowIdx = -1 @@ -754,16 +765,18 @@ func committedPendingLocalToolCancellationMessages( if billing.batchStartedAt.IsZero() || unbilledSubagentToolNames[call.ToolName] { continue } - // Only dispatched calls bill: a call absent from the batch's - // completion map was rejected before execution. A dispatched - // call with a stamped completion finished at that instant; a - // seeded but unstamped one was still running, so its window - // ends at the interrupt. Strictly-after keeps the earliest - // call on ties, matching billableBatchWindow. - end, dispatched := billing.toolCompletions[call.ToolCallID] - if !dispatched { + // Only dispatched calls bill: a call with no remaining + // occurrence was rejected before execution. A stamped + // occurrence finished at that instant; a seeded but unstamped + // one was still running, so its window ends at the interrupt. + // Strictly-after keeps the earliest call on ties, matching + // billableBatchWindow. + queue := pendingOccurrences[call.ToolCallID] + if len(queue) == 0 { continue } + end := queue[0] + pendingOccurrences[call.ToolCallID] = queue[1:] if end.IsZero() { end = interruptedAt } diff --git a/coderd/x/chatd/tasks_test.go b/coderd/x/chatd/tasks_test.go index a8b89193a56..e3c52820362 100644 --- a/coderd/x/chatd/tasks_test.go +++ b/coderd/x/chatd/tasks_test.go @@ -651,6 +651,39 @@ func TestInterruptTask_UnbilledOnlyBatchBillsNothingOnInterrupt(t *testing.T) { require.False(t, waitRow.RuntimeMs.Valid) } +// Two dispatched billed calls sharing one tool call ID keep distinct +// occurrence states: one completing early must not make the other look +// finished, so the interrupted batch still bills through to the +// interrupt, once. +func TestInterruptTask_DuplicateCallIDsKeepOccurrenceStates(t *testing.T) { + t.Parallel() + + f := newTaskTestFixture(t) + dupCallID := "call_" + uuid.NewString() + batch := interruptedBatchFixture(t, f, []codersdk.ChatMessagePart{ + {Type: codersdk.ChatMessagePartTypeToolCall, ToolCallID: dupCallID, ToolName: "execute", Args: json.RawMessage(`{}`)}, + {Type: codersdk.ChatMessagePartTypeToolCall, ToolCallID: dupCallID, ToolName: "execute", Args: json.RawMessage(`{}`)}, + }) + buffer := batch.starter.opts.MessagePartBuffer + + batch.clock.Advance(2 * time.Second) + require.NoError(t, buffer.StartToolBatch(batch.key, []string{dupCallID, dupCallID})) + // One occurrence completes 3 seconds in; the other keeps running + // until the interrupt lands 3 seconds later, defining the window. + batch.clock.Advance(3 * time.Second) + require.NoError(t, buffer.RecordToolCompletion(batch.key, dupCallID, batch.clock.Now())) + batch.clock.Advance(3 * time.Second) + + messages := batch.interrupt(t, f) + var billed []int64 + for _, msg := range messages { + if msg.Role == database.ChatMessageRoleTool && msg.RuntimeMs.Valid { + billed = append(billed, msg.RuntimeMs.Int64) + } + } + require.Equal(t, []int64{6_000}, billed, "the still-running occurrence bills the full window exactly once") +} + // A billed call rejected before execution (hook denial, ambiguous-call // rejection) is absent from the batch's dispatched set, so its missing // completion is not evidence that it ran: an interrupted batch whose From 68dee097733f8b4144fbf47784371d3cef8c5483 Mon Sep 17 00:00:00 2001 From: Jaayden Halko Date: Tue, 18 Aug 2026 06:09:03 +0000 Subject: [PATCH 05/20] fix(coderd/x/chatd): stamp tool completions by occurrence and reuse the interrupt instant Address Codex review feedback: - Thread the call's occurrence index through OnToolComplete and RecordToolCompletion so a completion stamps its own occurrence instead of the first still-running occurrence sharing its tool call ID. Same-ID calls with different billing classifications, such as an unbilled wait_agent finishing before a billed execute, previously ended the billed occurrence's window early. The interrupt path also consumes an occurrence for every dispatched call, including unbilled ones, to keep same-ID queues aligned with the history walk. - Pass the interrupt instant captured at episode close into the cancellation billing instead of taking a fresh clock reading inside machine.Update, so database contention or transaction retries no longer inflate a still-running call's billed window past the actual interrupt. --- coderd/x/chatd/chatloop/chatloop.go | 20 ++++++---- coderd/x/chatd/chatloop/runtime_test.go | 7 +++- coderd/x/chatd/generation.go | 18 +++++---- .../messagepartbuffer/message_part_buffer.go | 21 +++++++--- .../message_part_buffer_test.go | 27 ++++++------- coderd/x/chatd/tasks.go | 24 ++++++++---- coderd/x/chatd/tasks_test.go | 38 ++++++++++++++++++- 7 files changed, 110 insertions(+), 45 deletions(-) diff --git a/coderd/x/chatd/chatloop/chatloop.go b/coderd/x/chatd/chatloop/chatloop.go index d475214b3fb..3570f6ea378 100644 --- a/coderd/x/chatd/chatloop/chatloop.go +++ b/coderd/x/chatd/chatloop/chatloop.go @@ -272,9 +272,13 @@ type ExecuteLocalToolsOptions struct { // ordering deterministic, so this callback is the only live signal // that a tool already finished while siblings are still running; // the interrupt path uses it to bill an interrupted batch's - // partial window. It is called concurrently from tool goroutines - // and must be safe for concurrent use. - OnToolComplete func(toolCallID string, completedAt time.Time) + // partial window. callIndex is the call's position among the + // batch's local (non provider-executed) calls in dispatch order; + // it identifies the exact occurrence because duplicate tool call + // IDs, which reach execution when lifecycle hooks are disabled, + // make the ID alone ambiguous. It is called concurrently from + // tool goroutines and must be safe for concurrent use. + OnToolComplete func(callIndex int, toolCallID string, completedAt time.Time) PublishMessagePart func(codersdk.ChatMessageRole, codersdk.ChatMessagePart) Logger slog.Logger @@ -600,10 +604,12 @@ func ExecuteLocalTools(ctx context.Context, opts ExecuteLocalToolsOptions) (Tool ) if exclusiveViolation { now := clockNow(opts.Clock) - for _, tr := range policyResults { + for i, tr := range policyResults { recordToolResultTimestamp(&result, tr.ToolCallID, now) if opts.OnToolComplete != nil { - opts.OnToolComplete(tr.ToolCallID, now) + // Policy results are index-aligned with localCalls, so + // i is the call's occurrence index. + opts.OnToolComplete(i, tr.ToolCallID, now) } publishToolAttachments(ctx, opts.Logger, tr, now, publishMessagePart) ssePart := chatprompt.PartFromContentWithLogger(ctx, opts.Logger, tr) @@ -1165,7 +1171,7 @@ func executeTools( builtinToolNames map[string]bool, maxResultBytes int, toolNameAliases map[string]string, - onComplete func(toolCallID string, completedAt time.Time), + onComplete func(callIndex int, toolCallID string, completedAt time.Time), onResult func(fantasy.ToolResultContent, time.Time), ) []fantasy.ToolResultContent { if len(toolCalls) == 0 { @@ -1232,7 +1238,7 @@ func executeTools( // accurate individual completion times. completedAt[i] = clockNow(clock) if onComplete != nil { - onComplete(tc.ToolCallID, completedAt[i]) + onComplete(i, tc.ToolCallID, completedAt[i]) } }() results[i] = executeSingleTool( diff --git a/coderd/x/chatd/chatloop/runtime_test.go b/coderd/x/chatd/chatloop/runtime_test.go index 28748dbc965..c3dfa607852 100644 --- a/coderd/x/chatd/chatloop/runtime_test.go +++ b/coderd/x/chatd/chatloop/runtime_test.go @@ -452,6 +452,7 @@ func TestExecuteLocalTools_OnToolCompleteReportsLiveCompletions(t *testing.T) { defer trap.Close() type completion struct { + callIndex int toolCallID string completedAt time.Time } @@ -464,8 +465,8 @@ func TestExecuteLocalTools_OnToolCompleteReportsLiveCompletions(t *testing.T) { blockingTool("slow_tool", slowGo, fantasy.NewTextResponse("done")), }, ActiveTools: []string{"fast_tool", "slow_tool"}, - OnToolComplete: func(toolCallID string, completedAt time.Time) { - completionCh <- completion{toolCallID: toolCallID, completedAt: completedAt} + OnToolComplete: func(callIndex int, toolCallID string, completedAt time.Time) { + completionCh <- completion{callIndex: callIndex, toolCallID: toolCallID, completedAt: completedAt} }, ToolCalls: []fantasy.ToolCallContent{ {ToolCallID: "call-fast", ToolName: "fast_tool", Input: "{}"}, @@ -482,12 +483,14 @@ func TestExecuteLocalTools_OnToolCompleteReportsLiveCompletions(t *testing.T) { trap.MustWait(ctx).MustRelease(ctx) fast := testutil.RequireReceive(ctx, t, completionCh) require.Equal(t, "call-fast", fast.toolCallID) + require.Equal(t, 0, fast.callIndex, "callIndex is the call's position in dispatch order") // The slow tool completes at 60 seconds. clock.Advance(50 * time.Second) close(slowGo) trap.MustWait(ctx).MustRelease(ctx) slow := testutil.RequireReceive(ctx, t, completionCh) require.Equal(t, "call-slow", slow.toolCallID) + require.Equal(t, 1, slow.callIndex, "callIndex is the call's position in dispatch order") require.Equal(t, 50*time.Second, slow.completedAt.Sub(fast.completedAt)) outcome := testutil.RequireReceive(ctx, t, resultCh) diff --git a/coderd/x/chatd/generation.go b/coderd/x/chatd/generation.go index 2d595c246ba..83868c102ce 100644 --- a/coderd/x/chatd/generation.go +++ b/coderd/x/chatd/generation.go @@ -1069,12 +1069,14 @@ type generationAttempt struct { // were rejected before execution. It is always non-nil when // beginGenerationAttempt succeeds. startToolBatch func(toolCallIDs []string) - // recordToolCompletion records a tool call's completion instant on - // the buffer episode as the batch executes, so an interrupt can - // end an already-finished tool's billable window at its real - // completion instead of the interrupt instant. It is always - // non-nil when beginGenerationAttempt succeeds. - recordToolCompletion func(toolCallID string, completedAt time.Time) + // recordToolCompletion records a tool call occurrence's completion + // instant on the buffer episode as the batch executes, so an + // interrupt can end an already-finished tool's billable window at + // its real completion instead of the interrupt instant. callIndex + // addresses the occurrence within the dispatched batch, matching + // the order startToolBatch seeded. It is always non-nil when + // beginGenerationAttempt succeeds. + recordToolCompletion func(callIndex int, toolCallID string, completedAt time.Time) // closeEpisode closes the attempt's buffer episode. It is always // non-nil when beginGenerationAttempt succeeds. closeEpisode func() @@ -1124,8 +1126,8 @@ func (s *taskStarter) beginGenerationAttempt( startToolBatch: func(toolCallIDs []string) { _ = s.opts.MessagePartBuffer.StartToolBatch(key, toolCallIDs) }, - recordToolCompletion: func(toolCallID string, completedAt time.Time) { - _ = s.opts.MessagePartBuffer.RecordToolCompletion(key, toolCallID, completedAt) + recordToolCompletion: func(callIndex int, toolCallID string, completedAt time.Time) { + _ = s.opts.MessagePartBuffer.RecordToolCompletion(key, callIndex, toolCallID, completedAt) }, closeEpisode: func() { _ = s.opts.MessagePartBuffer.CloseEpisode(key) diff --git a/coderd/x/chatd/messagepartbuffer/message_part_buffer.go b/coderd/x/chatd/messagepartbuffer/message_part_buffer.go index 7a8cced2403..e2f4f2e3ba2 100644 --- a/coderd/x/chatd/messagepartbuffer/message_part_buffer.go +++ b/coderd/x/chatd/messagepartbuffer/message_part_buffer.go @@ -273,12 +273,14 @@ func (b *Buffer) StartToolBatch(key Key, toolCallIDs []string) error { // episode's tool batch finished. Tool goroutines report completions as // they happen, so an interrupt can bill tools that already finished up // to their real completion instead of treating every canceled call as -// still running. The stamp lands on the first still-running occurrence -// with the given ID: same-ID occurrences are indistinguishable to -// callers, and assigning them in seeded order keeps the batch's set of -// completion states correct. A completion whose ID was never seeded is -// appended, so an executed call is never dropped. -func (b *Buffer) RecordToolCompletion(key Key, toolCallID string, completedAt time.Time) error { +// still running. callIndex addresses the exact occurrence seeded by +// StartToolBatch: duplicate tool call IDs make the ID alone ambiguous, +// and stamping the wrong same-ID occurrence would let a finished call +// mark its still-running twin as done. When callIndex does not match +// the seeded occurrence, the stamp falls back to the first +// still-running occurrence with the ID, or is appended, so an executed +// call is never dropped. +func (b *Buffer) RecordToolCompletion(key Key, callIndex int, toolCallID string, completedAt time.Time) error { b.mu.Lock() defer b.mu.Unlock() if b.closed { @@ -291,6 +293,13 @@ func (b *Buffer) RecordToolCompletion(key Key, toolCallID string, completedAt ti if episode.closed { return ErrEpisodeClosed } + if callIndex >= 0 && callIndex < len(episode.toolCompletions) { + entry := &episode.toolCompletions[callIndex] + if entry.ToolCallID == toolCallID && entry.CompletedAt.IsZero() { + entry.CompletedAt = completedAt + return nil + } + } for i := range episode.toolCompletions { entry := &episode.toolCompletions[i] if entry.ToolCallID == toolCallID && entry.CompletedAt.IsZero() { diff --git a/coderd/x/chatd/messagepartbuffer/message_part_buffer_test.go b/coderd/x/chatd/messagepartbuffer/message_part_buffer_test.go index edbb64003f2..697074b0334 100644 --- a/coderd/x/chatd/messagepartbuffer/message_part_buffer_test.go +++ b/coderd/x/chatd/messagepartbuffer/message_part_buffer_test.go @@ -194,7 +194,7 @@ func TestBuffer_ToolCompletions(t *testing.T) { key := testEpisodeKey() require.Nil(t, buffer.ToolCompletions(key), "unknown episode has no completions") - require.ErrorIs(t, buffer.RecordToolCompletion(key, "call-1", clock.Now()), messagepartbuffer.ErrEpisodeNotFound) + require.ErrorIs(t, buffer.RecordToolCompletion(key, 0, "call-1", clock.Now()), messagepartbuffer.ErrEpisodeNotFound) require.NoError(t, buffer.CreateEpisode(key)) require.Empty(t, buffer.ToolCompletions(key), "episode without a tool batch has no completions") @@ -207,24 +207,25 @@ func TestBuffer_ToolCompletions(t *testing.T) { {ToolCallID: "call-dup"}, {ToolCallID: "call-dup"}, }, buffer.ToolCompletions(key)) - // A completion stamps the first still-running occurrence with its - // ID; the duplicate's second occurrence stays running. + // A completion stamps the occurrence addressed by its call index, + // not the first occurrence with a matching ID: the duplicate's + // second occurrence finishing must leave the first still running. clock.Advance(time.Second) - firstCompletedAt := clock.Now() - require.NoError(t, buffer.RecordToolCompletion(key, "call-dup", firstCompletedAt)) + secondDupCompletedAt := clock.Now() + require.NoError(t, buffer.RecordToolCompletion(key, 2, "call-dup", secondDupCompletedAt)) require.Equal(t, []messagepartbuffer.ToolCompletion{ {ToolCallID: "call-1"}, - {ToolCallID: "call-dup", CompletedAt: firstCompletedAt}, {ToolCallID: "call-dup"}, + {ToolCallID: "call-dup", CompletedAt: secondDupCompletedAt}, }, buffer.ToolCompletions(key)) clock.Advance(2 * time.Second) - secondCompletedAt := clock.Now() - require.NoError(t, buffer.RecordToolCompletion(key, "call-dup", secondCompletedAt)) + firstDupCompletedAt := clock.Now() + require.NoError(t, buffer.RecordToolCompletion(key, 1, "call-dup", firstDupCompletedAt)) completions := buffer.ToolCompletions(key) require.Equal(t, []messagepartbuffer.ToolCompletion{ {ToolCallID: "call-1"}, - {ToolCallID: "call-dup", CompletedAt: firstCompletedAt}, - {ToolCallID: "call-dup", CompletedAt: secondCompletedAt}, + {ToolCallID: "call-dup", CompletedAt: firstDupCompletedAt}, + {ToolCallID: "call-dup", CompletedAt: secondDupCompletedAt}, }, completions) // The returned slice is a copy: mutating it must not corrupt the @@ -235,7 +236,7 @@ func TestBuffer_ToolCompletions(t *testing.T) { // A closed episode keeps its recorded completions but accepts no // more, matching the batch-start stamp's lifecycle. require.NoError(t, buffer.CloseEpisode(key)) - require.ErrorIs(t, buffer.RecordToolCompletion(key, "call-1", clock.Now()), messagepartbuffer.ErrEpisodeClosed) + require.ErrorIs(t, buffer.RecordToolCompletion(key, 0, "call-1", clock.Now()), messagepartbuffer.ErrEpisodeClosed) require.True(t, buffer.ToolCompletions(key)[0].CompletedAt.IsZero()) } @@ -265,7 +266,7 @@ func TestBuffer_CloseEpisodeForBilling(t *testing.T) { require.NoError(t, buffer.StartToolBatch(key, []string{"call-1", "call-2"})) clock.Advance(time.Second) completedAt := clock.Now() - require.NoError(t, buffer.RecordToolCompletion(key, "call-1", completedAt)) + require.NoError(t, buffer.RecordToolCompletion(key, 0, "call-1", completedAt)) billing, err = buffer.CloseEpisodeForBilling(key) require.NoError(t, err) require.Zero(t, billing.ModelInvokedAt) @@ -274,7 +275,7 @@ func TestBuffer_CloseEpisodeForBilling(t *testing.T) { {ToolCallID: "call-1", CompletedAt: completedAt}, {ToolCallID: "call-2"}, }, billing.ToolCompletions) - require.ErrorIs(t, buffer.RecordToolCompletion(key, "call-2", clock.Now()), messagepartbuffer.ErrEpisodeClosed) + require.ErrorIs(t, buffer.RecordToolCompletion(key, 1, "call-2", clock.Now()), messagepartbuffer.ErrEpisodeClosed) // Closing an already-closed episode still reports its billing, so // an interrupt racing the generation task's own close loses diff --git a/coderd/x/chatd/tasks.go b/coderd/x/chatd/tasks.go index 40973bfdeb6..5ac3847a9c9 100644 --- a/coderd/x/chatd/tasks.go +++ b/coderd/x/chatd/tasks.go @@ -306,7 +306,12 @@ func (s *taskStarter) StartInterrupt(ctx context.Context, input chatWorkerTaskSt return xerrors.Errorf("load chat for task: %w", err) } messages := partialMessages - committedCancels, err := committedPendingLocalToolCancellationMessages(ctx, store, chat, s.opts.Clock.Now("chatworker", "interrupt"), interruptedToolBatchBilling{ + // Reuse the interrupt instant captured when the episode closed: + // a fresh clock read here would run while this transaction + // waits on the database (and again on retries), inflating the + // billed window of a still-running call past the actual + // interrupt. + committedCancels, err := committedPendingLocalToolCancellationMessages(ctx, store, chat, interruptedAt, interruptedToolBatchBilling{ batchStartedAt: episodeBilling.ToolBatchStartedAt, toolCompletions: episodeBilling.ToolCompletions, }) @@ -762,21 +767,26 @@ func committedPendingLocalToolCancellationMessages( ModelConfigID: uuid.NullUUID{UUID: chat.LastModelConfigID, Valid: chat.LastModelConfigID != uuid.Nil}, ContentVersion: chatprompt.CurrentContentVersion, }) - if billing.batchStartedAt.IsZero() || unbilledSubagentToolNames[call.ToolName] { + if billing.batchStartedAt.IsZero() { continue } // Only dispatched calls bill: a call with no remaining - // occurrence was rejected before execution. A stamped - // occurrence finished at that instant; a seeded but unstamped - // one was still running, so its window ends at the interrupt. - // Strictly-after keeps the earliest call on ties, matching - // billableBatchWindow. + // occurrence was rejected before execution. Every dispatched + // call consumes its occurrence, including unbilled ones, so + // same-ID occurrences stay aligned with the history walk. A + // stamped occurrence finished at that instant; a seeded but + // unstamped one was still running, so its window ends at the + // interrupt. Strictly-after keeps the earliest call on ties, + // matching billableBatchWindow. queue := pendingOccurrences[call.ToolCallID] if len(queue) == 0 { continue } end := queue[0] pendingOccurrences[call.ToolCallID] = queue[1:] + if unbilledSubagentToolNames[call.ToolName] { + continue + } if end.IsZero() { end = interruptedAt } diff --git a/coderd/x/chatd/tasks_test.go b/coderd/x/chatd/tasks_test.go index e3c52820362..966539b5fae 100644 --- a/coderd/x/chatd/tasks_test.go +++ b/coderd/x/chatd/tasks_test.go @@ -598,7 +598,7 @@ func TestInterruptTask_ToolBatchBillsPartialWindowOnCancellationRow(t *testing.T // does. Its result is not published: results publish only after // the whole batch finishes. batch.clock.Advance(3 * time.Second) - require.NoError(t, buffer.RecordToolCompletion(batch.key, execCallID, batch.clock.Now())) + require.NoError(t, buffer.RecordToolCompletion(batch.key, 0, execCallID, batch.clock.Now())) // wait_agent is still blocked on its child when the interrupt lands // 5 seconds later. batch.clock.Advance(5 * time.Second) @@ -671,7 +671,7 @@ func TestInterruptTask_DuplicateCallIDsKeepOccurrenceStates(t *testing.T) { // One occurrence completes 3 seconds in; the other keeps running // until the interrupt lands 3 seconds later, defining the window. batch.clock.Advance(3 * time.Second) - require.NoError(t, buffer.RecordToolCompletion(batch.key, dupCallID, batch.clock.Now())) + require.NoError(t, buffer.RecordToolCompletion(batch.key, 0, dupCallID, batch.clock.Now())) batch.clock.Advance(3 * time.Second) messages := batch.interrupt(t, f) @@ -684,6 +684,40 @@ func TestInterruptTask_DuplicateCallIDsKeepOccurrenceStates(t *testing.T) { require.Equal(t, []int64{6_000}, billed, "the still-running occurrence bills the full window exactly once") } +// Same-ID occurrences with different billing classifications stay +// correlated: an unbilled wait_agent occurrence finishing early stamps +// its own occurrence, not the still-running billed execute sharing its +// ID, so the execute occurrence still bills through to the interrupt. +func TestInterruptTask_DuplicateCallIDCompletionStampsOwnOccurrence(t *testing.T) { + t.Parallel() + + f := newTaskTestFixture(t) + dupCallID := "call_" + uuid.NewString() + batch := interruptedBatchFixture(t, f, []codersdk.ChatMessagePart{ + {Type: codersdk.ChatMessagePartTypeToolCall, ToolCallID: dupCallID, ToolName: "execute", Args: json.RawMessage(`{}`)}, + {Type: codersdk.ChatMessagePartTypeToolCall, ToolCallID: dupCallID, ToolName: "wait_agent", Args: json.RawMessage(`{}`)}, + }) + buffer := batch.starter.opts.MessagePartBuffer + + batch.clock.Advance(2 * time.Second) + require.NoError(t, buffer.StartToolBatch(batch.key, []string{dupCallID, dupCallID})) + // The unbilled wait_agent occurrence (index 1) completes 3 seconds + // in; the billed execute occurrence (index 0) keeps running until + // the interrupt 3 seconds later. + batch.clock.Advance(3 * time.Second) + require.NoError(t, buffer.RecordToolCompletion(batch.key, 1, dupCallID, batch.clock.Now())) + batch.clock.Advance(3 * time.Second) + + messages := batch.interrupt(t, f) + var billed []int64 + for _, msg := range messages { + if msg.Role == database.ChatMessageRoleTool && msg.RuntimeMs.Valid { + billed = append(billed, msg.RuntimeMs.Int64) + } + } + require.Equal(t, []int64{6_000}, billed, "the running execute occurrence bills to the interrupt; wait_agent's early completion must not end it at 3s") +} + // A billed call rejected before execution (hook denial, ambiguous-call // rejection) is absent from the batch's dispatched set, so its missing // completion is not evidence that it ran: an interrupted batch whose From 2f6e743565b703d5f41c3e43bdf14221ab5998bb Mon Sep 17 00:00:00 2001 From: Jaayden Halko Date: Tue, 18 Aug 2026 06:30:33 +0000 Subject: [PATCH 06/20] fix(coderd/x/chatd): correlate interrupt billing by unresolved call position Address Codex review feedback: - Seed each dispatched occurrence with its position in the step's unresolved tool-call order, the order interrupt reconstruction walks, and match occurrences to cancellation rows positionally instead of through per-ID queues. A call rejected before execution that shared an ID with a dispatched call previously consumed the dispatched occurrence and billed the whole window on a row that never ran. partitionAmbiguousToolCalls now reports each allowed call's input position so the seed does not re-derive it from IDs. - Skip starting the billing batch for exclusive-policy violations. chatloop synthesizes error results for the whole batch without dispatching any tool, so seeding those calls let an interrupt racing the synthetic results bill never-run work from the batch stamp to the interrupt. --- coderd/x/chatd/generation.go | 53 +++++++++------ .../messagepartbuffer/message_part_buffer.go | 49 ++++++++++---- .../message_part_buffer_test.go | 51 +++++++++----- coderd/x/chatd/tasks.go | 66 +++++++++---------- coderd/x/chatd/tasks_test.go | 54 +++++++++++++-- coderd/x/chatd/toolinput.go | 16 ++--- coderd/x/chatd/toolinput_internal_test.go | 8 ++- 7 files changed, 196 insertions(+), 101 deletions(-) diff --git a/coderd/x/chatd/generation.go b/coderd/x/chatd/generation.go index 83868c102ce..6d1fee6cb26 100644 --- a/coderd/x/chatd/generation.go +++ b/coderd/x/chatd/generation.go @@ -786,7 +786,7 @@ func (s *taskStarter) admitStepToolCalls( if err := chathooks.RejectDuplicateToolUseIDs(toolCalls); err != nil { return chathooks.PreToolUseExecutionResult{}, chathooks.GenerationDispatchError(agenthooks.EventPreToolUse, err) } - unambiguous, ambiguous := partitionAmbiguousToolCalls(prepared, toolCalls) + unambiguous, _, ambiguous := partitionAmbiguousToolCalls(prepared, toolCalls) preflight, err := s.server.hooks.PreflightPendingToolCalls(ctx, chathooks.ChatFor(prepared.Chat, input.hookTurnID()), unambiguous) if err != nil { return chathooks.PreToolUseExecutionResult{}, chathooks.GenerationDispatchError(agenthooks.EventPreToolUse, err) @@ -805,10 +805,12 @@ func (s *taskStarter) executeLocalTools( prepared generationPrepared, decision generationDecision, ) error { + exclusiveRejected := exclusiveBatchRejected(decision.localToolCalls, prepared.ExclusiveToolNames) allowed := decision.localToolCalls + var allowedIndexes []int var denied []fantasy.ToolResultContent - if !exclusiveBatchRejected(decision.localToolCalls, prepared.ExclusiveToolNames) { - allowed, denied = partitionAmbiguousToolCalls(prepared, decision.localToolCalls) + if !exclusiveRejected { + allowed, allowedIndexes, denied = partitionAmbiguousToolCalls(prepared, decision.localToolCalls) } attempt, err := s.beginGenerationAttempt(ctx, machine, input) if err != nil { @@ -824,16 +826,27 @@ func (s *taskStarter) executeLocalTools( var outcome chatloop.ToolExecutionOutcome var spawnDispatchErr error if len(allowed) > 0 { - // Stamp the batch start and the dispatched call IDs on the - // buffer episode so an interrupt can bill the partial window - // this step would have reported. Only allowed calls are - // listed: denied calls never run, so an interrupt must not - // treat their missing completions as still-running work. - allowedCallIDs := make([]string, 0, len(allowed)) - for _, tc := range allowed { - allowedCallIDs = append(allowedCallIDs, tc.ToolCallID) - } - attempt.startToolBatch(allowedCallIDs) + // Stamp the batch start and the dispatched calls on the buffer + // episode so an interrupt can bill the partial window this step + // would have reported. Only dispatched calls are seeded, keyed + // by their position in the unresolved call order the interrupt + // walks: rejected calls never run, so an interrupt must not + // treat their missing completions as still-running work, and a + // rejected call must not consume a same-ID dispatched + // occurrence. An exclusive-policy violation dispatches nothing + // (chatloop synthesizes error results for the whole batch), so + // no billable batch starts and an interrupt racing those + // synthetic results bills nothing. + if !exclusiveRejected { + dispatched := make([]messagepartbuffer.DispatchedToolCall, 0, len(allowed)) + for j, tc := range allowed { + dispatched = append(dispatched, messagepartbuffer.DispatchedToolCall{ + CallIndex: allowedIndexes[j], + ToolCallID: tc.ToolCallID, + }) + } + attempt.startToolBatch(dispatched) + } outcome, err = chatloop.ExecuteLocalTools(ctx, chatloop.ExecuteLocalToolsOptions{ Tools: prepared.Tools, ActiveTools: prepared.ActiveTools, @@ -1064,11 +1077,11 @@ type generationAttempt struct { startModelInvocation func() // startToolBatch marks the start of the attempt's billable local // tool batch window on the buffer episode and records which tool - // calls were actually dispatched, so an interrupt can bill the - // window the step would have reported without charging calls that - // were rejected before execution. It is always non-nil when - // beginGenerationAttempt succeeds. - startToolBatch func(toolCallIDs []string) + // call occurrences were actually dispatched, so an interrupt can + // bill the window the step would have reported without charging + // calls that were rejected before execution. It is always non-nil + // when beginGenerationAttempt succeeds. + startToolBatch func(calls []messagepartbuffer.DispatchedToolCall) // recordToolCompletion records a tool call occurrence's completion // instant on the buffer episode as the batch executes, so an // interrupt can end an already-finished tool's billable window at @@ -1123,8 +1136,8 @@ func (s *taskStarter) beginGenerationAttempt( startModelInvocation: func() { _ = s.opts.MessagePartBuffer.StartModelInvocation(key) }, - startToolBatch: func(toolCallIDs []string) { - _ = s.opts.MessagePartBuffer.StartToolBatch(key, toolCallIDs) + startToolBatch: func(calls []messagepartbuffer.DispatchedToolCall) { + _ = s.opts.MessagePartBuffer.StartToolBatch(key, calls) }, recordToolCompletion: func(callIndex int, toolCallID string, completedAt time.Time) { _ = s.opts.MessagePartBuffer.RecordToolCompletion(key, callIndex, toolCallID, completedAt) diff --git a/coderd/x/chatd/messagepartbuffer/message_part_buffer.go b/coderd/x/chatd/messagepartbuffer/message_part_buffer.go index e2f4f2e3ba2..5810c9deded 100644 --- a/coderd/x/chatd/messagepartbuffer/message_part_buffer.go +++ b/coderd/x/chatd/messagepartbuffer/message_part_buffer.go @@ -233,10 +233,26 @@ func (b *Buffer) StartModelInvocation(key Key) error { return nil } +// DispatchedToolCall identifies one tool call occurrence dispatched in an +// episode's local tool batch. +type DispatchedToolCall struct { + // CallIndex is the caller-assigned position of this occurrence in the + // step's full unresolved tool-call order, which is the order interrupt + // reconstruction walks. It differs from the occurrence's position in + // the dispatched batch when calls were rejected before execution, and + // it disambiguates duplicate tool call IDs. + CallIndex int + ToolCallID string +} + // ToolCompletion tracks one dispatched tool call occurrence in an // episode's local tool batch. CompletedAt is zero while the call is // still running. type ToolCompletion struct { + // CallIndex mirrors DispatchedToolCall.CallIndex. It is -1 for + // completions recorded without a matching seeded occurrence, which + // never correlate to an unresolved call. + CallIndex int ToolCallID string CompletedAt time.Time } @@ -248,7 +264,7 @@ type ToolCompletion struct { // completion) from one never dispatched at all (absent), such as a call // denied by a lifecycle hook or rejected as ambiguous before execution, // and duplicate tool call IDs keep distinct per-occurrence states. -func (b *Buffer) StartToolBatch(key Key, toolCallIDs []string) error { +func (b *Buffer) StartToolBatch(key Key, calls []DispatchedToolCall) error { b.mu.Lock() defer b.mu.Unlock() if b.closed { @@ -262,9 +278,12 @@ func (b *Buffer) StartToolBatch(key Key, toolCallIDs []string) error { return ErrEpisodeClosed } episode.toolBatchStartedAt = b.opts.Clock.Now("message-part-buffer", "tool-batch-start") - episode.toolCompletions = make([]ToolCompletion, 0, len(toolCallIDs)) - for _, id := range toolCallIDs { - episode.toolCompletions = append(episode.toolCompletions, ToolCompletion{ToolCallID: id}) + episode.toolCompletions = make([]ToolCompletion, 0, len(calls)) + for _, call := range calls { + episode.toolCompletions = append(episode.toolCompletions, ToolCompletion{ + CallIndex: call.CallIndex, + ToolCallID: call.ToolCallID, + }) } return nil } @@ -273,14 +292,15 @@ func (b *Buffer) StartToolBatch(key Key, toolCallIDs []string) error { // episode's tool batch finished. Tool goroutines report completions as // they happen, so an interrupt can bill tools that already finished up // to their real completion instead of treating every canceled call as -// still running. callIndex addresses the exact occurrence seeded by -// StartToolBatch: duplicate tool call IDs make the ID alone ambiguous, -// and stamping the wrong same-ID occurrence would let a finished call -// mark its still-running twin as done. When callIndex does not match -// the seeded occurrence, the stamp falls back to the first -// still-running occurrence with the ID, or is appended, so an executed -// call is never dropped. -func (b *Buffer) RecordToolCompletion(key Key, callIndex int, toolCallID string, completedAt time.Time) error { +// still running. dispatchIndex addresses the exact occurrence seeded by +// StartToolBatch, as the occurrence's position within the dispatched +// batch: duplicate tool call IDs make the ID alone ambiguous, and +// stamping the wrong same-ID occurrence would let a finished call mark +// its still-running twin as done. When dispatchIndex does not match the +// seeded occurrence, the stamp falls back to the first still-running +// occurrence with the ID, or is appended with CallIndex -1, so an +// executed call is never dropped. +func (b *Buffer) RecordToolCompletion(key Key, dispatchIndex int, toolCallID string, completedAt time.Time) error { b.mu.Lock() defer b.mu.Unlock() if b.closed { @@ -293,8 +313,8 @@ func (b *Buffer) RecordToolCompletion(key Key, callIndex int, toolCallID string, if episode.closed { return ErrEpisodeClosed } - if callIndex >= 0 && callIndex < len(episode.toolCompletions) { - entry := &episode.toolCompletions[callIndex] + if dispatchIndex >= 0 && dispatchIndex < len(episode.toolCompletions) { + entry := &episode.toolCompletions[dispatchIndex] if entry.ToolCallID == toolCallID && entry.CompletedAt.IsZero() { entry.CompletedAt = completedAt return nil @@ -308,6 +328,7 @@ func (b *Buffer) RecordToolCompletion(key Key, callIndex int, toolCallID string, } } episode.toolCompletions = append(episode.toolCompletions, ToolCompletion{ + CallIndex: -1, ToolCallID: toolCallID, CompletedAt: completedAt, }) diff --git a/coderd/x/chatd/messagepartbuffer/message_part_buffer_test.go b/coderd/x/chatd/messagepartbuffer/message_part_buffer_test.go index 697074b0334..615f356b73a 100644 --- a/coderd/x/chatd/messagepartbuffer/message_part_buffer_test.go +++ b/coderd/x/chatd/messagepartbuffer/message_part_buffer_test.go @@ -157,14 +157,14 @@ func TestBuffer_ToolBatchStartedAt(t *testing.T) { key := testEpisodeKey() require.Zero(t, buffer.ToolBatchStartedAt(key), "unknown episode has no batch stamp") - require.ErrorIs(t, buffer.StartToolBatch(key, []string{"call-1"}), messagepartbuffer.ErrEpisodeNotFound) + require.ErrorIs(t, buffer.StartToolBatch(key, []messagepartbuffer.DispatchedToolCall{{CallIndex: 0, ToolCallID: "call-1"}}), messagepartbuffer.ErrEpisodeNotFound) require.NoError(t, buffer.CreateEpisode(key)) require.Zero(t, buffer.ToolBatchStartedAt(key), "episode without a tool batch has no batch stamp") // Attempt setup happens before tools start executing and is not // billable. clock.Advance(time.Second) - require.NoError(t, buffer.StartToolBatch(key, []string{"call-1"})) + require.NoError(t, buffer.StartToolBatch(key, []messagepartbuffer.DispatchedToolCall{{CallIndex: 0, ToolCallID: "call-1"}})) startedAt := buffer.ToolBatchStartedAt(key) require.Equal(t, clock.Now(), startedAt) @@ -172,7 +172,7 @@ func TestBuffer_ToolBatchStartedAt(t *testing.T) { // no longer accepts a batch start. clock.Advance(1500 * time.Millisecond) require.NoError(t, buffer.CloseEpisode(key)) - require.ErrorIs(t, buffer.StartToolBatch(key, []string{"call-1"}), messagepartbuffer.ErrEpisodeClosed) + require.ErrorIs(t, buffer.StartToolBatch(key, []messagepartbuffer.DispatchedToolCall{{CallIndex: 0, ToolCallID: "call-1"}}), messagepartbuffer.ErrEpisodeClosed) require.Equal(t, startedAt, buffer.ToolBatchStartedAt(key)) // Episodes that never execute local tools, such as model @@ -201,11 +201,15 @@ func TestBuffer_ToolCompletions(t *testing.T) { // Starting the batch seeds one still-running occurrence per // dispatched call, in dispatch order. Duplicate IDs keep distinct // occurrences instead of collapsing into one shared state. - require.NoError(t, buffer.StartToolBatch(key, []string{"call-1", "call-dup", "call-dup"})) + require.NoError(t, buffer.StartToolBatch(key, []messagepartbuffer.DispatchedToolCall{ + {CallIndex: 0, ToolCallID: "call-1"}, + {CallIndex: 1, ToolCallID: "call-dup"}, + {CallIndex: 2, ToolCallID: "call-dup"}, + })) require.Equal(t, []messagepartbuffer.ToolCompletion{ - {ToolCallID: "call-1"}, - {ToolCallID: "call-dup"}, - {ToolCallID: "call-dup"}, + {CallIndex: 0, ToolCallID: "call-1"}, + {CallIndex: 1, ToolCallID: "call-dup"}, + {CallIndex: 2, ToolCallID: "call-dup"}, }, buffer.ToolCompletions(key)) // A completion stamps the occurrence addressed by its call index, // not the first occurrence with a matching ID: the duplicate's @@ -214,18 +218,18 @@ func TestBuffer_ToolCompletions(t *testing.T) { secondDupCompletedAt := clock.Now() require.NoError(t, buffer.RecordToolCompletion(key, 2, "call-dup", secondDupCompletedAt)) require.Equal(t, []messagepartbuffer.ToolCompletion{ - {ToolCallID: "call-1"}, - {ToolCallID: "call-dup"}, - {ToolCallID: "call-dup", CompletedAt: secondDupCompletedAt}, + {CallIndex: 0, ToolCallID: "call-1"}, + {CallIndex: 1, ToolCallID: "call-dup"}, + {CallIndex: 2, ToolCallID: "call-dup", CompletedAt: secondDupCompletedAt}, }, buffer.ToolCompletions(key)) clock.Advance(2 * time.Second) firstDupCompletedAt := clock.Now() require.NoError(t, buffer.RecordToolCompletion(key, 1, "call-dup", firstDupCompletedAt)) completions := buffer.ToolCompletions(key) require.Equal(t, []messagepartbuffer.ToolCompletion{ - {ToolCallID: "call-1"}, - {ToolCallID: "call-dup", CompletedAt: firstDupCompletedAt}, - {ToolCallID: "call-dup", CompletedAt: secondDupCompletedAt}, + {CallIndex: 0, ToolCallID: "call-1"}, + {CallIndex: 1, ToolCallID: "call-dup", CompletedAt: firstDupCompletedAt}, + {CallIndex: 2, ToolCallID: "call-dup", CompletedAt: secondDupCompletedAt}, }, completions) // The returned slice is a copy: mutating it must not corrupt the @@ -233,6 +237,18 @@ func TestBuffer_ToolCompletions(t *testing.T) { completions[0].CompletedAt = clock.Now() require.True(t, buffer.ToolCompletions(key)[0].CompletedAt.IsZero()) + // A completion whose ID was never seeded is appended with + // CallIndex -1: the executed call is not dropped, but it never + // correlates to an unresolved call. + clock.Advance(time.Second) + unseededAt := clock.Now() + require.NoError(t, buffer.RecordToolCompletion(key, 5, "call-unseeded", unseededAt)) + require.Equal(t, messagepartbuffer.ToolCompletion{ + CallIndex: -1, + ToolCallID: "call-unseeded", + CompletedAt: unseededAt, + }, buffer.ToolCompletions(key)[3]) + // A closed episode keeps its recorded completions but accepts no // more, matching the batch-start stamp's lifecycle. require.NoError(t, buffer.CloseEpisode(key)) @@ -263,7 +279,10 @@ func TestBuffer_CloseEpisodeForBilling(t *testing.T) { require.NoError(t, buffer.CreateEpisode(key)) clock.Advance(time.Second) batchStartedAt := clock.Now() - require.NoError(t, buffer.StartToolBatch(key, []string{"call-1", "call-2"})) + require.NoError(t, buffer.StartToolBatch(key, []messagepartbuffer.DispatchedToolCall{ + {CallIndex: 0, ToolCallID: "call-1"}, + {CallIndex: 1, ToolCallID: "call-2"}, + })) clock.Advance(time.Second) completedAt := clock.Now() require.NoError(t, buffer.RecordToolCompletion(key, 0, "call-1", completedAt)) @@ -272,8 +291,8 @@ func TestBuffer_CloseEpisodeForBilling(t *testing.T) { require.Zero(t, billing.ModelInvokedAt) require.Equal(t, batchStartedAt, billing.ToolBatchStartedAt) require.Equal(t, []messagepartbuffer.ToolCompletion{ - {ToolCallID: "call-1", CompletedAt: completedAt}, - {ToolCallID: "call-2"}, + {CallIndex: 0, ToolCallID: "call-1", CompletedAt: completedAt}, + {CallIndex: 1, ToolCallID: "call-2"}, }, billing.ToolCompletions) require.ErrorIs(t, buffer.RecordToolCompletion(key, 1, "call-2", clock.Now()), messagepartbuffer.ErrEpisodeClosed) diff --git a/coderd/x/chatd/tasks.go b/coderd/x/chatd/tasks.go index 5ac3847a9c9..e6ef6ada775 100644 --- a/coderd/x/chatd/tasks.go +++ b/coderd/x/chatd/tasks.go @@ -700,16 +700,17 @@ type interruptedToolBatchBilling struct { // cancellation rows carry no runtime. batchStartedAt time.Time // toolCompletions holds the live batch's dispatched tool call - // occurrences in dispatch order: seeded with zero completions when - // the batch started and stamped as each tool finished. A stamped - // occurrence ends its billable window at its completion, so a - // batch whose billed tools all finished early does not bill the - // longer window of a still-running unbilled tool such as - // wait_agent. A seeded but unstamped occurrence was still running - // when the interrupt landed. A call with no matching occurrence - // was never dispatched, such as a call denied by a lifecycle hook - // or rejected as ambiguous, and bills nothing even though it too - // receives a cancellation row. + // occurrences, each keyed by its position in the step's unresolved + // call order: seeded with zero completions when the batch started + // and stamped as each tool finished. A stamped occurrence ends its + // billable window at its completion, so a batch whose billed tools + // all finished early does not bill the longer window of a + // still-running unbilled tool such as wait_agent. A seeded but + // unstamped occurrence was still running when the interrupt + // landed. A call with no occurrence at its position was never + // dispatched, such as a call rejected as malformed or ambiguous + // before execution, and bills nothing even though it too receives + // a cancellation row. toolCompletions []messagepartbuffer.ToolCompletion } @@ -734,20 +735,24 @@ func committedPendingLocalToolCancellationMessages( if len(localCalls) == 0 { return nil, nil } - // Per-ID queues of dispatched occurrences, in dispatch order. - // Unresolved calls walk the same assistant part order the batch was - // dispatched from, so each history row consumes its own occurrence - // and duplicate tool call IDs never share one completion state. - pendingOccurrences := make(map[string][]time.Time, len(billing.toolCompletions)) + // Dispatched occurrences keyed by their position in the unresolved + // call order, the same order this loop walks. Positional matching + // keeps a call rejected before execution from consuming a same-ID + // dispatched occurrence and keeps duplicate tool call IDs from + // sharing one completion state. + dispatched := make(map[int]messagepartbuffer.ToolCompletion, len(billing.toolCompletions)) for _, completion := range billing.toolCompletions { - pendingOccurrences[completion.ToolCallID] = append(pendingOccurrences[completion.ToolCallID], completion.CompletedAt) + if completion.CallIndex < 0 { + continue + } + dispatched[completion.CallIndex] = completion } var ( windowEnd time.Time windowRowIdx = -1 ) result := make([]chatstate.Message, 0, len(localCalls)) - for _, call := range localCalls { + for i, call := range localCalls { payload, err := json.Marshal(map[string]string{"error": interruptedToolResultErrorMessage}) if err != nil { return nil, xerrors.Errorf("marshal interrupted tool result: %w", err) @@ -767,26 +772,21 @@ func committedPendingLocalToolCancellationMessages( ModelConfigID: uuid.NullUUID{UUID: chat.LastModelConfigID, Valid: chat.LastModelConfigID != uuid.Nil}, ContentVersion: chatprompt.CurrentContentVersion, }) - if billing.batchStartedAt.IsZero() { - continue - } - // Only dispatched calls bill: a call with no remaining - // occurrence was rejected before execution. Every dispatched - // call consumes its occurrence, including unbilled ones, so - // same-ID occurrences stay aligned with the history walk. A - // stamped occurrence finished at that instant; a seeded but - // unstamped one was still running, so its window ends at the - // interrupt. Strictly-after keeps the earliest call on ties, - // matching billableBatchWindow. - queue := pendingOccurrences[call.ToolCallID] - if len(queue) == 0 { + if billing.batchStartedAt.IsZero() || unbilledSubagentToolNames[call.ToolName] { continue } - end := queue[0] - pendingOccurrences[call.ToolCallID] = queue[1:] - if unbilledSubagentToolNames[call.ToolName] { + // Only dispatched calls bill: a call with no occurrence at its + // position was rejected before execution, and an ID mismatch + // means the seed does not describe this call. A stamped + // occurrence finished at that instant; a seeded but unstamped + // one was still running, so its window ends at the interrupt. + // Strictly-after keeps the earliest call on ties, matching + // billableBatchWindow. + occurrence, ok := dispatched[i] + if !ok || occurrence.ToolCallID != call.ToolCallID { continue } + end := occurrence.CompletedAt if end.IsZero() { end = interruptedAt } diff --git a/coderd/x/chatd/tasks_test.go b/coderd/x/chatd/tasks_test.go index 966539b5fae..e7a9ea980d4 100644 --- a/coderd/x/chatd/tasks_test.go +++ b/coderd/x/chatd/tasks_test.go @@ -592,7 +592,10 @@ func TestInterruptTask_ToolBatchBillsPartialWindowOnCancellationRow(t *testing.T // this mock clock. // Attempt setup happens before the tools start and is not billable. batch.clock.Advance(2 * time.Second) - require.NoError(t, buffer.StartToolBatch(batch.key, []string{execCallID, waitCallID})) + require.NoError(t, buffer.StartToolBatch(batch.key, []messagepartbuffer.DispatchedToolCall{ + {CallIndex: 0, ToolCallID: execCallID}, + {CallIndex: 1, ToolCallID: waitCallID}, + })) // execute completes 3 seconds into the batch and records its // completion, the way the tool goroutine's completion callback // does. Its result is not published: results publish only after @@ -622,7 +625,7 @@ func TestInterruptTask_RunningBilledToolBillsUpToInterrupt(t *testing.T) { }) batch.clock.Advance(2 * time.Second) - require.NoError(t, batch.starter.opts.MessagePartBuffer.StartToolBatch(batch.key, []string{execCallID})) + require.NoError(t, batch.starter.opts.MessagePartBuffer.StartToolBatch(batch.key, []messagepartbuffer.DispatchedToolCall{{CallIndex: 0, ToolCallID: execCallID}})) // The tool is still running when the interrupt lands 7 seconds in. batch.clock.Advance(7 * time.Second) @@ -643,7 +646,7 @@ func TestInterruptTask_UnbilledOnlyBatchBillsNothingOnInterrupt(t *testing.T) { }) batch.clock.Advance(2 * time.Second) - require.NoError(t, batch.starter.opts.MessagePartBuffer.StartToolBatch(batch.key, []string{waitCallID})) + require.NoError(t, batch.starter.opts.MessagePartBuffer.StartToolBatch(batch.key, []messagepartbuffer.DispatchedToolCall{{CallIndex: 0, ToolCallID: waitCallID}})) batch.clock.Advance(10 * time.Second) messages := batch.interrupt(t, f) @@ -667,7 +670,10 @@ func TestInterruptTask_DuplicateCallIDsKeepOccurrenceStates(t *testing.T) { buffer := batch.starter.opts.MessagePartBuffer batch.clock.Advance(2 * time.Second) - require.NoError(t, buffer.StartToolBatch(batch.key, []string{dupCallID, dupCallID})) + require.NoError(t, buffer.StartToolBatch(batch.key, []messagepartbuffer.DispatchedToolCall{ + {CallIndex: 0, ToolCallID: dupCallID}, + {CallIndex: 1, ToolCallID: dupCallID}, + })) // One occurrence completes 3 seconds in; the other keeps running // until the interrupt lands 3 seconds later, defining the window. batch.clock.Advance(3 * time.Second) @@ -700,7 +706,10 @@ func TestInterruptTask_DuplicateCallIDCompletionStampsOwnOccurrence(t *testing.T buffer := batch.starter.opts.MessagePartBuffer batch.clock.Advance(2 * time.Second) - require.NoError(t, buffer.StartToolBatch(batch.key, []string{dupCallID, dupCallID})) + require.NoError(t, buffer.StartToolBatch(batch.key, []messagepartbuffer.DispatchedToolCall{ + {CallIndex: 0, ToolCallID: dupCallID}, + {CallIndex: 1, ToolCallID: dupCallID}, + })) // The unbilled wait_agent occurrence (index 1) completes 3 seconds // in; the billed execute occurrence (index 0) keeps running until // the interrupt 3 seconds later. @@ -718,6 +727,36 @@ func TestInterruptTask_DuplicateCallIDCompletionStampsOwnOccurrence(t *testing.T require.Equal(t, []int64{6_000}, billed, "the running execute occurrence bills to the interrupt; wait_agent's early completion must not end it at 3s") } +// A rejected call must not steal a same-ID dispatched occurrence: with +// a rejected billed execute preceding a dispatched unbilled wait_agent +// that shares its ID, positional matching leaves the execute row +// without an occurrence, so the interrupted batch bills nothing instead +// of charging the whole wait window to a call that never ran. +func TestInterruptTask_RejectedDuplicateIDDoesNotStealDispatchedOccurrence(t *testing.T) { + t.Parallel() + + f := newTaskTestFixture(t) + dupCallID := "call_" + uuid.NewString() + batch := interruptedBatchFixture(t, f, []codersdk.ChatMessagePart{ + {Type: codersdk.ChatMessagePartTypeToolCall, ToolCallID: dupCallID, ToolName: "execute", Args: json.RawMessage(`{}`)}, + {Type: codersdk.ChatMessagePartTypeToolCall, ToolCallID: dupCallID, ToolName: "wait_agent", Args: json.RawMessage(`{}`)}, + }) + + batch.clock.Advance(2 * time.Second) + // Only the wait_agent occurrence (unresolved position 1) was + // dispatched: the execute occurrence sharing its ID was rejected + // as malformed before execution. + require.NoError(t, batch.starter.opts.MessagePartBuffer.StartToolBatch(batch.key, []messagepartbuffer.DispatchedToolCall{{CallIndex: 1, ToolCallID: dupCallID}})) + batch.clock.Advance(10 * time.Second) + + messages := batch.interrupt(t, f) + for _, msg := range messages { + if msg.Role == database.ChatMessageRoleTool { + require.False(t, msg.RuntimeMs.Valid, "no cancellation row may bill: only the unbilled wait_agent occurrence ran") + } + } +} + // A billed call rejected before execution (hook denial, ambiguous-call // rejection) is absent from the batch's dispatched set, so its missing // completion is not evidence that it ran: an interrupted batch whose @@ -736,8 +775,9 @@ func TestInterruptTask_RejectedCallBillsNothingOnInterrupt(t *testing.T) { batch.clock.Advance(2 * time.Second) // Only wait_agent was dispatched: execute was rejected before - // execution and never ran. - require.NoError(t, batch.starter.opts.MessagePartBuffer.StartToolBatch(batch.key, []string{waitCallID})) + // execution and never ran, so only the occurrence at wait_agent's + // unresolved position (1) is seeded. + require.NoError(t, batch.starter.opts.MessagePartBuffer.StartToolBatch(batch.key, []messagepartbuffer.DispatchedToolCall{{CallIndex: 1, ToolCallID: waitCallID}})) batch.clock.Advance(10 * time.Second) messages := batch.interrupt(t, f) diff --git a/coderd/x/chatd/toolinput.go b/coderd/x/chatd/toolinput.go index 4aa1a8df6a8..d1b15b91ce6 100644 --- a/coderd/x/chatd/toolinput.go +++ b/coderd/x/chatd/toolinput.go @@ -15,16 +15,15 @@ import ( // reject before pre_tool_use so a hook consumer is never asked to authorize // bytes whose meaning depends on which reader resolves them, and so input that // cannot be carried in a hook payload fails as a retryable tool error instead -// of a dispatch failure. +// of a dispatch failure. allowedIndexes carries each allowed call's position +// in the input slice, so callers that need to correlate the dispatched subset +// back to the full unresolved call order (tool batch billing) do not have to +// re-derive it from tool call IDs, which duplicates make ambiguous. func partitionAmbiguousToolCalls( prepared generationPrepared, toolCalls []fantasy.ToolCallContent, -) ([]fantasy.ToolCallContent, []fantasy.ToolResultContent) { - var ( - allowed []fantasy.ToolCallContent - rejected []fantasy.ToolResultContent - ) - for _, toolCall := range toolCalls { +) (allowed []fantasy.ToolCallContent, allowedIndexes []int, rejected []fantasy.ToolResultContent) { + for i, toolCall := range toolCalls { if !json.Valid([]byte(toolCall.Input)) { rejected = append(rejected, malformedToolResult(toolCall)) continue @@ -34,8 +33,9 @@ func partitionAmbiguousToolCalls( continue } allowed = append(allowed, toolCall) + allowedIndexes = append(allowedIndexes, i) } - return allowed, rejected + return allowed, allowedIndexes, rejected } // validateOverriddenToolInputs rechecks the inputs a pre_tool_use consumer diff --git a/coderd/x/chatd/toolinput_internal_test.go b/coderd/x/chatd/toolinput_internal_test.go index e19bbe48a9d..807c2fa4870 100644 --- a/coderd/x/chatd/toolinput_internal_test.go +++ b/coderd/x/chatd/toolinput_internal_test.go @@ -40,20 +40,22 @@ func TestPartitionAmbiguousToolCallsGatesOnBuiltins(t *testing.T) { Tools: []fantasy.AgentTool{fetch}, BuiltinToolNames: map[string]bool{"fetch": true}, } - allowed, rejected := partitionAmbiguousToolCalls(prepared, []fantasy.ToolCallContent{ambiguous, clean}) + allowed, allowedIndexes, rejected := partitionAmbiguousToolCalls(prepared, []fantasy.ToolCallContent{ambiguous, clean}) require.Len(t, rejected, 1) require.Equal(t, "call_ambiguous", rejected[0].ToolCallID) require.Len(t, allowed, 1) require.Equal(t, "call_clean", allowed[0].ToolCallID) + require.Equal(t, []int{1}, allowedIndexes, "allowed indexes point at positions in the input slice") }) t.Run("non-builtin", func(t *testing.T) { t.Parallel() prepared := generationPrepared{Tools: []fantasy.AgentTool{fetch}} - allowed, rejected := partitionAmbiguousToolCalls(prepared, []fantasy.ToolCallContent{ambiguous, clean}) + allowed, allowedIndexes, rejected := partitionAmbiguousToolCalls(prepared, []fantasy.ToolCallContent{ambiguous, clean}) require.Empty(t, rejected) require.Len(t, allowed, 2) + require.Equal(t, []int{0, 1}, allowedIndexes) }) // Execution resolves a deprecated name to its canonical tool, so @@ -85,7 +87,7 @@ func TestPartitionAmbiguousToolCallsGatesOnBuiltins(t *testing.T) { Tools: []fantasy.AgentTool{tool}, BuiltinToolNames: map[string]bool{canonical: true}, } - _, rejected := partitionAmbiguousToolCalls(prepared, []fantasy.ToolCallContent{aliased}) + _, _, rejected := partitionAmbiguousToolCalls(prepared, []fantasy.ToolCallContent{aliased}) require.Len(t, rejected, 1) }) } From 45f962c4a694bda1bcd5d57bc66586f40d7d3dbf Mon Sep 17 00:00:00 2001 From: Jaayden Halko Date: Tue, 18 Aug 2026 06:40:27 +0000 Subject: [PATCH 07/20] docs(docs/ai-coder): reformat agent runtime section to one sentence per line Address Codex review feedback: the docs prose style guide requires each sentence on its own source line and the full text of a bullet item on a single line, reformatting the entire paragraph when any line in it is edited. The section kept fixed-column wrapping. --- docs/ai-coder/usage-data-reporting.md | 35 ++++++++------------------- 1 file changed, 10 insertions(+), 25 deletions(-) diff --git a/docs/ai-coder/usage-data-reporting.md b/docs/ai-coder/usage-data-reporting.md index 4e109956722..6973a76ad90 100644 --- a/docs/ai-coder/usage-data-reporting.md +++ b/docs/ai-coder/usage-data-reporting.md @@ -58,38 +58,23 @@ Example of a failed request (e.g. Tallyman Server is blocked by your network): ## Agent runtime measurement -Total Coder Agent runtime is summed from per-message runtime -(`runtime_ms` on chat messages). +Total Coder Agent runtime is summed from per-message runtime (`runtime_ms` on chat messages). -An assistant message's runtime is the wall-clock duration of the model -invocation that produced its content, measured from just before the request -to the model provider opens until the response is fully consumed. A tool -message's runtime is the wall-clock duration of the local tool batch that -produced it, measured from the start of the batch until the last counted -tool finishes. Tools in a batch run in parallel, so each batch records one -window on one tool message rather than a per-tool sum. +An assistant message's runtime is the wall-clock duration of the model invocation that produced its content, measured from just before the request to the model provider opens until the response is fully consumed. +A tool message's runtime is the wall-clock duration of the local tool batch that produced it, measured from the start of the batch until the last counted tool finishes. +Tools in a batch run in parallel, so each batch records one window on one tool message rather than a per-tool sum. What counts: - Assistant generation steps, in both top-level chats and sub-agent chats. - Context compaction (summarization) model calls. -- Local tool execution: file, terminal, and process tools, workspace - lifecycle operations, MCP tools, and other server-executed tools. -- Interrupted generation: the time streamed or spent executing tools before - the interrupt is kept on the partial messages. +- Local tool execution: file, terminal, and process tools, workspace lifecycle operations, MCP tools, and other server-executed tools. +- Interrupted generation: the time streamed or spent executing tools before the interrupt is kept on the partial messages. What does not count: -- Sub-agent orchestration tools, such as spawning and waiting on - sub-agents. A sub-agent is its own chat and records its own runtime, so - counting the parent's wait would double count. Waiting on a sub-agent - never extends a tool batch's window, even when other tools in the batch - do count. -- Client-executed (dynamic) tools and external agents: the server cannot - measure work it does not execute. +- Sub-agent orchestration tools, such as spawning and waiting on sub-agents. A sub-agent is its own chat and records its own runtime, so counting the parent's wait would double count. Waiting on a sub-agent never extends a tool batch's window, even when other tools in the batch do count. +- Client-executed (dynamic) tools and external agents: the server cannot measure work it does not execute. - Idle time: chats waiting for user input or external tool results. -- Failed model calls whose output was discarded, and the backoff between - retried attempts. Retried and errored attempts persist no content, so - they record no runtime. -- Ancillary model calls that produce no chat messages, such as title - generation. +- Failed model calls whose output was discarded, and the backoff between retried attempts. Retried and errored attempts persist no content, so they record no runtime. +- Ancillary model calls that produce no chat messages, such as title generation. From 167d9f8922cf3cf04235e7a47332567d5243af1a Mon Sep 17 00:00:00 2001 From: Jaayden Halko Date: Tue, 18 Aug 2026 06:56:51 +0000 Subject: [PATCH 08/20] fix(coderd/x/chatd): keep interrupted billing stable across task retries Address Codex review feedback: - Return the episode's first-close instant in the billing snapshot and use it as the interrupt instant. Each retried interrupt attempt previously took a fresh clock reading, billing still-running tools (and the model window) through every retry of a transient database failure; every billing input now comes from the episode's stable first-close state, so retries recompute identical windows. - Push the episode's eviction deadline out on every re-close, so a retry loop outlasting the buffer's retention window keeps its billing snapshot and buffered parts through a database outage instead of losing them to the cleanup loop; retry backoff (max 5s) stays well under the retention window (15s). The superseded heap item is skipped by the cleanup loop's existing identity check. --- coderd/x/chatd/ARCHITECTURE.md | 2 +- .../messagepartbuffer/message_part_buffer.go | 30 +++++++++- .../message_part_buffer_test.go | 59 +++++++++++++++++-- coderd/x/chatd/tasks.go | 7 ++- 4 files changed, 91 insertions(+), 7 deletions(-) diff --git a/coderd/x/chatd/ARCHITECTURE.md b/coderd/x/chatd/ARCHITECTURE.md index 5fb8aa47959..59a98a4a0a0 100644 --- a/coderd/x/chatd/ARCHITECTURE.md +++ b/coderd/x/chatd/ARCHITECTURE.md @@ -725,7 +725,7 @@ The buffer exposes the following API: - `GetParts(chat_id, history_version, generation_attempt)`: returns the message parts for an episode. Returns a predefined error if the episode is not found. - `StartModelInvocation(chat_id, history_version, generation_attempt)`: stamps the instant the episode opens its provider stream. Returns a predefined error if the episode is not found or already closed. Episodes that never invoke a model, such as local tool execution batches, are never stamped. - `ModelInvokedAt(chat_id, history_version, generation_attempt)`: returns the instant stamped by `StartModelInvocation`, or the zero time when the episode is unknown or never opened a provider stream. It must be read before `CloseEpisode`, because closed episodes are garbage collected and reading afterwards races the cleanup loop. The interrupt goroutine reads it just before closing the episode and uses the span between that instant and the interrupt as the interrupted attempt's billable runtime. -- TODO(CODAGT-928): document the tool-batch billing methods `StartToolBatch`, `ToolBatchStartedAt`, `RecordToolCompletion`, and `ToolCompletions`, the local-tool counterparts of `StartModelInvocation`/`ModelInvokedAt`, and `CloseEpisodeForBilling`, which the interrupt task uses to close the episode and snapshot its billing stamps in one atomic step so it can bill an interrupted attempt's partial window. +- TODO(CODAGT-928): document the tool-batch billing methods `StartToolBatch`, `ToolBatchStartedAt`, `RecordToolCompletion`, and `ToolCompletions`, the local-tool counterparts of `StartModelInvocation`/`ModelInvokedAt`, and `CloseEpisodeForBilling`, which the interrupt task uses to close the episode and snapshot its billing stamps (including the stable first-close instant used as the interrupt instant) in one atomic step; repeat calls return the identical snapshot and push the eviction deadline out, so retried interrupt tasks bill the same window and keep their snapshot through a database outage. - `SubscribeToEpisode(chat_id, history_version, generation_attempt)`: returns a go channel that will receive all message parts for the episode. It spawns a goroutine that delivers parts to the channel. It's live until the episode is closed or until a subscriber requests that the channel be closed. Once the goroutine delivers all message parts for a closed episode, it closes the channel and exits. If the episode is already closed at the time of the call, the goroutine delivers all message parts for the episode, closes the channel, and exits. `SubscribeToEpisode` does not return an error if the episode is not found: it waits for it to be created instead. Closed episodes are garbage collected after at least 15 seconds since they were closed and when they have no active subscribers. The message part buffer maintains a garbage collection goroutine. diff --git a/coderd/x/chatd/messagepartbuffer/message_part_buffer.go b/coderd/x/chatd/messagepartbuffer/message_part_buffer.go index 5810c9deded..175f3d9616d 100644 --- a/coderd/x/chatd/messagepartbuffer/message_part_buffer.go +++ b/coderd/x/chatd/messagepartbuffer/message_part_buffer.go @@ -453,6 +453,12 @@ func (b *Buffer) ToolCompletions(key Key) []ToolCompletion { // EpisodeBilling is the billing state an episode accumulated before it // closed. type EpisodeBilling struct { + // ClosedAt is the instant the episode first closed. Interrupt + // handling uses it as the interrupt instant: it is stable across + // repeat closes, so a retried interrupt task bills the same window + // every attempt instead of one that grows with each retry's later + // clock reading. + ClosedAt time.Time // ModelInvokedAt is the StartModelInvocation stamp, or zero when // the episode never opened a provider stream. ModelInvokedAt time.Time @@ -472,6 +478,14 @@ type EpisodeBilling struct { // the gap, billing a finished tool as still running or losing a live // batch's window entirely. Closing an unknown episode creates it // closed, mirroring CloseEpisode, and reports empty billing. +// +// Calling it again on an already-closed episode returns the same +// snapshot and pushes the episode's eviction deadline out by the +// retention window. Interrupt task retries re-read the snapshot on +// every attempt with backoff well under the retention window, so a +// retrying interrupt keeps its billing state and buffered parts alive +// through a database outage instead of losing them to the cleanup +// loop mid-retry. func (b *Buffer) CloseEpisodeForBilling(key Key) (EpisodeBilling, error) { b.mu.Lock() defer b.mu.Unlock() @@ -479,11 +493,15 @@ func (b *Buffer) CloseEpisodeForBilling(key Key) (EpisodeBilling, error) { return EpisodeBilling{}, ErrMessagePartBufferClosed } episode := b.getOrCreateEpisodeLocked(key) - if episode.close(b.opts.Clock.Now("message-part-buffer", "close")) { + now := b.opts.Clock.Now("message-part-buffer", "close") + if episode.close(now) { b.queueClosedEpisodeLocked(key, episode) episode.notifySubscribers() + } else { + b.refreshClosedEpisodeEvictionLocked(key, episode, now) } return EpisodeBilling{ + ClosedAt: episode.closedAt, ModelInvokedAt: episode.modelStartedAt, ToolBatchStartedAt: episode.toolBatchStartedAt, ToolCompletions: slices.Clone(episode.toolCompletions), @@ -611,6 +629,16 @@ func (b *Buffer) queueClosedEpisodeLocked(key Key, episode *episodeState) { heap.Push(&b.closedEpisodes, item) } +// refreshClosedEpisodeEvictionLocked pushes a closed episode's eviction +// deadline out to evictAt plus the retention window by queueing a fresh +// heap item. The superseded item stays in the heap until the cleanup +// loop pops it and skips it via the closedHeapItem identity check. +func (b *Buffer) refreshClosedEpisodeEvictionLocked(key Key, episode *episodeState, evictAt time.Time) { + item := &closedEpisodeItem{key: key, closedAt: evictAt} + episode.closedHeapItem = item + heap.Push(&b.closedEpisodes, item) +} + func (b *Buffer) getOrCreateEpisodeLocked(key Key) *episodeState { episode := b.episodes[key] if episode != nil { diff --git a/coderd/x/chatd/messagepartbuffer/message_part_buffer_test.go b/coderd/x/chatd/messagepartbuffer/message_part_buffer_test.go index 615f356b73a..f7aeaac05ba 100644 --- a/coderd/x/chatd/messagepartbuffer/message_part_buffer_test.go +++ b/coderd/x/chatd/messagepartbuffer/message_part_buffer_test.go @@ -264,10 +264,11 @@ func TestBuffer_CloseEpisodeForBilling(t *testing.T) { defer buffer.Close() // Closing an unknown episode creates it closed, like CloseEpisode, - // and reports empty billing. + // and reports empty billing stamped with the close instant. unknown := testEpisodeKey() billing, err := buffer.CloseEpisodeForBilling(unknown) require.NoError(t, err) + require.Equal(t, clock.Now(), billing.ClosedAt) require.Zero(t, billing.ModelInvokedAt) require.Zero(t, billing.ToolBatchStartedAt) require.Empty(t, billing.ToolCompletions) @@ -286,8 +287,11 @@ func TestBuffer_CloseEpisodeForBilling(t *testing.T) { clock.Advance(time.Second) completedAt := clock.Now() require.NoError(t, buffer.RecordToolCompletion(key, 0, "call-1", completedAt)) + clock.Advance(time.Second) + closedAt := clock.Now() billing, err = buffer.CloseEpisodeForBilling(key) require.NoError(t, err) + require.Equal(t, closedAt, billing.ClosedAt) require.Zero(t, billing.ModelInvokedAt) require.Equal(t, batchStartedAt, billing.ToolBatchStartedAt) require.Equal(t, []messagepartbuffer.ToolCompletion{ @@ -296,14 +300,61 @@ func TestBuffer_CloseEpisodeForBilling(t *testing.T) { }, billing.ToolCompletions) require.ErrorIs(t, buffer.RecordToolCompletion(key, 1, "call-2", clock.Now()), messagepartbuffer.ErrEpisodeClosed) - // Closing an already-closed episode still reports its billing, so - // an interrupt racing the generation task's own close loses - // nothing. + // Closing an already-closed episode reports the identical snapshot, + // including the original close instant: a retried interrupt task + // bills the same window every attempt, and an interrupt racing the + // generation task's own close loses nothing. + clock.Advance(time.Second) again, err := buffer.CloseEpisodeForBilling(key) require.NoError(t, err) require.Equal(t, billing, again) } +// A retrying interrupt task re-reads its billing snapshot on every +// attempt; each re-read pushes the episode's eviction deadline out, so +// a retry loop outlasting the original retention window keeps the +// snapshot instead of losing it to the cleanup loop mid-outage. +func TestBuffer_CloseEpisodeForBillingRefreshesEviction(t *testing.T) { + t.Parallel() + + clock := quartz.NewMock(t) + buffer := messagepartbuffer.New(messagepartbuffer.Options{ + Clock: clock, + ClosedEpisodeRetention: time.Minute, + }) + defer buffer.Close() + ctx := testutil.Context(t, testutil.WaitShort) + + key := testEpisodeKey() + require.NoError(t, buffer.CreateEpisode(key)) + require.NoError(t, buffer.StartToolBatch(key, []messagepartbuffer.DispatchedToolCall{{CallIndex: 0, ToolCallID: "call-1"}})) + first, err := buffer.CloseEpisodeForBilling(key) + require.NoError(t, err) + + // A retry 45 seconds in re-reads the snapshot and refreshes the + // deadline. + clock.Advance(45 * time.Second).MustWait(ctx) + again, err := buffer.CloseEpisodeForBilling(key) + require.NoError(t, err) + require.Equal(t, first, again) + + // The cleanup tick after the original one-minute deadline must not + // collect the refreshed episode: the snapshot survives. + clock.Advance(15 * time.Second).MustWait(ctx) + again, err = buffer.CloseEpisodeForBilling(key) + require.NoError(t, err) + require.Equal(t, first, again) + + // Once retries stop refreshing it, the episode ages out and a later + // close reports empty billing again. + clock.Advance(time.Minute).MustWait(ctx) + clock.Advance(time.Minute).MustWait(ctx) + expired, err := buffer.CloseEpisodeForBilling(key) + require.NoError(t, err) + require.Zero(t, expired.ToolBatchStartedAt) + require.Empty(t, expired.ToolCompletions) +} + func TestBuffer_SubscribeExistingReplaysThenStreamsLiveParts(t *testing.T) { t.Parallel() diff --git a/coderd/x/chatd/tasks.go b/coderd/x/chatd/tasks.go index e6ef6ada775..b91bddec642 100644 --- a/coderd/x/chatd/tasks.go +++ b/coderd/x/chatd/tasks.go @@ -282,7 +282,12 @@ func (s *taskStarter) StartInterrupt(ctx context.Context, input chatWorkerTaskSt } return taskRetryableError{err: xerrors.Errorf("get message part episode: %w", err)} } - interruptedAt := s.opts.Clock.Now("chatworker", "interrupt") + // The interrupt instant is the episode's first close, which is + // stable across repeat closes: a retried interrupt task (transient + // database errors) recomputes identical partial messages and + // billing windows instead of billing still-running tools through + // each retry's later clock reading. + interruptedAt := episodeBilling.ClosedAt var attemptRuntime time.Duration if !episodeBilling.ModelInvokedAt.IsZero() { attemptRuntime = interruptedAt.Sub(episodeBilling.ModelInvokedAt) From 9fd2445eb5d5b09f0214a867f3c5d2a140776c73 Mon Sep 17 00:00:00 2001 From: Jaayden Halko Date: Tue, 18 Aug 2026 07:10:28 +0000 Subject: [PATCH 09/20] fix(coderd/x/chatd): carry the interrupt episode snapshot across task retries Address Codex review feedback: one interrupt attempt can outlive the buffer's closed-episode retention (a machine.Update stalled on a database outage runs up to the task timeout), and the eviction refresh only happens when an attempt begins. A retry after such a stall found a blank recreated episode and lost the interrupt instant, partial parts, and tool completions, underbilling the interrupted work and dropping the partial messages. The runner now attaches a snapshot holder to the interrupt task's input, which every retry attempt of the task instance shares. The first attempt to read the episode stores the billing state and buffered parts there, and later attempts reuse the carried snapshot instead of re-reading the evictable buffer. --- coderd/x/chatd/options.go | 5 +++ coderd/x/chatd/runner.go | 6 +++ coderd/x/chatd/tasks.go | 71 ++++++++++++++++++++++++++---------- coderd/x/chatd/tasks_test.go | 51 +++++++++++++++++++++++++- 4 files changed, 113 insertions(+), 20 deletions(-) diff --git a/coderd/x/chatd/options.go b/coderd/x/chatd/options.go index f7570a9035b..49758ebab77 100644 --- a/coderd/x/chatd/options.go +++ b/coderd/x/chatd/options.go @@ -69,6 +69,11 @@ type chatWorkerTaskStartInput struct { DebugTurn *runnerDebugTurn SessionStart *sessionStartTracker StopNudges *stopNudgeTracker + // InterruptSnapshot, set for interrupt tasks, is shared by every + // retry attempt of the task instance so the first attempt's episode + // snapshot survives buffer eviction during a stalled attempt. Nil + // makes each attempt re-read the buffer. + InterruptSnapshot *interruptEpisodeSnapshot } func (i chatWorkerTaskStartInput) hookTurnID() *uuid.UUID { diff --git a/coderd/x/chatd/runner.go b/coderd/x/chatd/runner.go index 72b916a38a6..bcc4bd222ce 100644 --- a/coderd/x/chatd/runner.go +++ b/coderd/x/chatd/runner.go @@ -232,6 +232,12 @@ func (r *runner) spawnTaskIfNeeded(kind taskKind, state runnerStateUpdate) { SessionStart: &r.sessionStart, StopNudges: &r.stopNudges, } + if kind == taskKindInterrupt { + // Shared by every retry attempt of this task instance so the + // first attempt's episode snapshot survives buffer eviction + // during a stalled attempt. + input.InterruptSnapshot = &interruptEpisodeSnapshot{} + } go r.runTask(taskCtx, kind, key, input, done) } diff --git a/coderd/x/chatd/tasks.go b/coderd/x/chatd/tasks.go index b91bddec642..711d778260e 100644 --- a/coderd/x/chatd/tasks.go +++ b/coderd/x/chatd/tasks.go @@ -240,6 +240,21 @@ func (o chatWorkerOptions) retryOptions() retryWrapperOptions { } } +// interruptEpisodeSnapshot carries the interrupt task's one-time episode +// snapshot across retry attempts of the same task instance. One attempt +// can outlive the buffer's closed-episode retention (a machine.Update +// stalled on a database outage runs up to the task timeout), so +// re-reading the buffer on a later attempt could find a blank recreated +// episode and underbill the interrupted work or drop the partial +// messages. Attempts of one task run sequentially, so no locking is +// needed. +type interruptEpisodeSnapshot struct { + loaded bool + key messagepartbuffer.Key + billing messagepartbuffer.EpisodeBilling + parts []messagepartbuffer.Part +} + func (s *taskStarter) StartInterrupt(ctx context.Context, input chatWorkerTaskStartInput) error { machine := chatstate.NewChatMachine(s.opts.Store, s.opts.Pubsub, input.ChatID) var chat database.Chat @@ -260,27 +275,45 @@ func (s *taskStarter) StartInterrupt(ctx context.Context, input chatWorkerTaskSt HistoryVersion: input.HistoryVersion, GenerationAttempt: chat.GenerationAttempt, } - // Closing and snapshotting billing state must be one atomic step: - // the generation goroutine records batch starts and tool - // completions concurrently, so a read-then-close would let stamps - // land in the gap and go missing from the snapshot. - episodeBilling, err := s.opts.MessagePartBuffer.CloseEpisodeForBilling(key) - if err != nil { - if ctx.Err() != nil { - return errors.Join(errTaskExpectedExit, xerrors.Errorf("close message part episode: %w", err), ctx.Err()) + var episodeBilling messagepartbuffer.EpisodeBilling + var parts []messagepartbuffer.Part + snapshot := input.InterruptSnapshot + if snapshot != nil && snapshot.loaded && snapshot.key == key { + // A previous attempt of this task already snapshotted the + // episode. Reuse it: the episode may have been evicted from + // the buffer while that attempt stalled, and re-reading would + // find a blank recreated episode. + episodeBilling = snapshot.billing + parts = snapshot.parts + } else { + // Closing and snapshotting billing state must be one atomic + // step: the generation goroutine records batch starts and tool + // completions concurrently, so a read-then-close would let + // stamps land in the gap and go missing from the snapshot. + episodeBilling, err = s.opts.MessagePartBuffer.CloseEpisodeForBilling(key) + if err != nil { + if ctx.Err() != nil { + return errors.Join(errTaskExpectedExit, xerrors.Errorf("close message part episode: %w", err), ctx.Err()) + } + return taskRetryableError{err: xerrors.Errorf("close message part episode: %w", err)} } - return taskRetryableError{err: xerrors.Errorf("close message part episode: %w", err)} - } - parts, err := s.opts.MessagePartBuffer.GetParts(key) - if errors.Is(err, messagepartbuffer.ErrEpisodeNotFound) { - parts = nil - err = nil - } - if err != nil { - if ctx.Err() != nil { - return errors.Join(errTaskExpectedExit, xerrors.Errorf("get message part episode: %w", err), ctx.Err()) + parts, err = s.opts.MessagePartBuffer.GetParts(key) + if errors.Is(err, messagepartbuffer.ErrEpisodeNotFound) { + parts = nil + err = nil + } + if err != nil { + if ctx.Err() != nil { + return errors.Join(errTaskExpectedExit, xerrors.Errorf("get message part episode: %w", err), ctx.Err()) + } + return taskRetryableError{err: xerrors.Errorf("get message part episode: %w", err)} + } + if snapshot != nil { + snapshot.loaded = true + snapshot.key = key + snapshot.billing = episodeBilling + snapshot.parts = parts } - return taskRetryableError{err: xerrors.Errorf("get message part episode: %w", err)} } // The interrupt instant is the episode's first close, which is // stable across repeat closes: a retried interrupt task (transient diff --git a/coderd/x/chatd/tasks_test.go b/coderd/x/chatd/tasks_test.go index e7a9ea980d4..10200f21b5a 100644 --- a/coderd/x/chatd/tasks_test.go +++ b/coderd/x/chatd/tasks_test.go @@ -535,6 +535,11 @@ func interruptedBatchFixture( } func (b interruptedBatch) interrupt(t *testing.T, f *taskTestFixture) []database.ChatMessage { + t.Helper() + return b.interruptWithSnapshot(t, f, nil) +} + +func (b interruptedBatch) interruptWithSnapshot(t *testing.T, f *taskTestFixture, snapshot *interruptEpisodeSnapshot) []database.ChatMessage { t.Helper() interrupting := f.interruptChat(t, b.chat.ID) err := b.starter.StartInterrupt(testutil.Context(t, testutil.WaitLong), chatWorkerTaskStartInput{ @@ -544,6 +549,7 @@ func (b interruptedBatch) interrupt(t *testing.T, f *taskTestFixture) []database HistoryVersion: interrupting.HistoryVersion, GenerationAttempt: interrupting.GenerationAttempt, Status: database.ChatStatusInterrupting, + InterruptSnapshot: snapshot, }) require.NoError(t, err) messages, err := f.db.GetChatMessagesByChatID(testutil.Context(t, testutil.WaitShort), database.GetChatMessagesByChatIDParams{ChatID: b.chat.ID}) @@ -606,11 +612,54 @@ func TestInterruptTask_ToolBatchBillsPartialWindowOnCancellationRow(t *testing.T // 5 seconds later. batch.clock.Advance(5 * time.Second) - messages := batch.interrupt(t, f) + // The first attempt stores the episode snapshot it read, so a + // retry after buffer eviction reuses it. + snapshot := &interruptEpisodeSnapshot{} + messages := batch.interruptWithSnapshot(t, f, snapshot) execRow := findToolResultMessage(t, messages, execCallID) require.Equal(t, sql.NullInt64{Int64: 3_000, Valid: true}, execRow.RuntimeMs) waitRow := findToolResultMessage(t, messages, waitCallID) require.False(t, waitRow.RuntimeMs.Valid) + require.True(t, snapshot.loaded) + require.Equal(t, batch.key, snapshot.key) + require.Len(t, snapshot.billing.ToolCompletions, 2) +} + +// A retry attempt reuses the snapshot its first attempt captured: after +// a stalled attempt outlives the buffer's retention and the episode is +// evicted, the blank recreated episode must not replace the snapshot's +// billing state, so the interrupted batch still bills its window. +func TestInterruptTask_RetrySnapshotOutlivesEpisodeEviction(t *testing.T) { + t.Parallel() + + f := newTaskTestFixture(t) + execCallID := "call_" + uuid.NewString() + batch := interruptedBatchFixture(t, f, []codersdk.ChatMessagePart{ + {Type: codersdk.ChatMessagePartTypeToolCall, ToolCallID: execCallID, ToolName: "execute", Args: json.RawMessage(`{}`)}, + }) + + // Simulate a retry: the first attempt snapshotted a live batch + // (started at 2s, closed at 9s, execute still running) and then + // stalled past the buffer's retention. The buffer's episode + // carries none of that state, like the blank episode a retry + // recreates after eviction, so billing can only come from the + // carried snapshot. + batch.clock.Advance(2 * time.Second) + batchStartedAt := batch.clock.Now() + batch.clock.Advance(7 * time.Second) + snapshot := &interruptEpisodeSnapshot{ + loaded: true, + key: batch.key, + billing: messagepartbuffer.EpisodeBilling{ + ClosedAt: batch.clock.Now(), + ToolBatchStartedAt: batchStartedAt, + ToolCompletions: []messagepartbuffer.ToolCompletion{{CallIndex: 0, ToolCallID: execCallID}}, + }, + } + + messages := batch.interruptWithSnapshot(t, f, snapshot) + execRow := findToolResultMessage(t, messages, execCallID) + require.Equal(t, sql.NullInt64{Int64: 7_000, Valid: true}, execRow.RuntimeMs) } // A billed tool still running at the interrupt bills up to the interrupt From dd4384d73c782bb9c7509c70541de6086ad828f8 Mon Sep 17 00:00:00 2001 From: Jaayden Halko Date: Tue, 18 Aug 2026 07:24:36 +0000 Subject: [PATCH 10/20] fix(coderd/x/chatd): bill windows defined by calls without tool call IDs Address Codex review feedback: - billableBatchWindow used an empty tool call ID as its no-window sentinel, discarding the window when the defining call had no ID, which providers can emit and which still executes. A found flag now tracks whether a billed completion exists, and the commit path gates the row match on a positive window instead of a non-empty ID, so an ID-less window lands on the first ID-less tool row. - Add a TODO to the architecture document's interrupt goroutine section, whose close-then-read flow description is stale after the atomic billing snapshot changes, for the PR author to update. - Deflake the eviction-refresh buffer test: assert collection through GetParts, which garbage collects due episodes synchronously, instead of racing the cleanup goroutine's handling of delivered ticks. --- coderd/x/chatd/ARCHITECTURE.md | 2 + coderd/x/chatd/chatloop/chatloop.go | 12 ++++-- coderd/x/chatd/chatloop/runtime_test.go | 41 +++++++++++++++++++ coderd/x/chatd/message_conversion.go | 11 +++-- coderd/x/chatd/message_conversion_test.go | 33 +++++++++++++++ .../message_part_buffer_test.go | 6 ++- 6 files changed, 96 insertions(+), 9 deletions(-) diff --git a/coderd/x/chatd/ARCHITECTURE.md b/coderd/x/chatd/ARCHITECTURE.md index 59a98a4a0a0..59d6ee39fc4 100644 --- a/coderd/x/chatd/ARCHITECTURE.md +++ b/coderd/x/chatd/ARCHITECTURE.md @@ -936,6 +936,8 @@ The goroutine does the following in order: 3. It reads the buffered parts for that episode by calling the `GetParts` method on the message part buffer. 4. It applies the `FinishInterruption(partial?)` transition on the core state machine. If there are no buffered parts for that episode, or the episode is not found, it passes `nil` as the `partial` argument. +TODO(CODAGT-928): update steps 2 and 3 above. The interrupt goroutine now closes the episode and snapshots its billing stamps in one atomic step via `CloseEpisodeForBilling`, uses the episode's first-close instant as the interrupt instant, carries the snapshot and buffered parts across task retry attempts so buffer eviction during a stalled attempt cannot lose them, and synthesizes tool cancellation rows whose `runtime_ms` bills the interrupted tool batch's partial window. + #### Dynamic tools timeout goroutine The dynamic tools timeout goroutine is responsible for waiting for the dynamic tool timeout to pass, which is determined by the `requires_action_deadline_at` field on the chat. It is spawned when the event indicates the core state machine is in `A0` or `A1` (status is `requires_action`). The goroutine fetches the deadline value from the database. When the timeout passes, it applies the `CancelRequiresAction` transition on the core state machine. diff --git a/coderd/x/chatd/chatloop/chatloop.go b/coderd/x/chatd/chatloop/chatloop.go index 3570f6ea378..11348ed5c0b 100644 --- a/coderd/x/chatd/chatloop/chatloop.go +++ b/coderd/x/chatd/chatloop/chatloop.go @@ -691,9 +691,11 @@ func ExecuteLocalTools(ctx context.Context, opts ExecuteLocalToolsOptions) (Tool // sub-agent orchestration) never extend the window even when they run // longest. completions is aligned with toolCalls by occurrence, so // duplicate tool call IDs each keep their own completion. Returns the -// tool call whose completion ends the window, with ties broken by call -// order; (0, "") when no billed tool completed or the window rounds to -// nothing. +// ID of the tool call whose completion ends the window, with ties +// broken by call order; a zero duration means no billed tool completed +// or the window rounds to nothing. The ID can be empty even when the +// window exists, because providers can emit calls without IDs and +// those still execute and bill. func billableBatchWindow( batchStart time.Time, toolCalls []fantasy.ToolCallContent, @@ -701,6 +703,7 @@ func billableBatchWindow( unbilledToolNames map[string]bool, ) (time.Duration, string) { var ( + found bool windowEnd time.Time toolCallID string ) @@ -717,11 +720,12 @@ func billableBatchWindow( } // Strictly-after keeps the earliest call on ties. if end.After(windowEnd) { + found = true windowEnd = end toolCallID = tc.ToolCallID } } - if toolCallID == "" || !windowEnd.After(batchStart) { + if !found || !windowEnd.After(batchStart) { return 0, "" } return windowEnd.Sub(batchStart), toolCallID diff --git a/coderd/x/chatd/chatloop/runtime_test.go b/coderd/x/chatd/chatloop/runtime_test.go index c3dfa607852..0ba5002bf79 100644 --- a/coderd/x/chatd/chatloop/runtime_test.go +++ b/coderd/x/chatd/chatloop/runtime_test.go @@ -437,6 +437,47 @@ func TestExecuteLocalTools_DuplicateToolCallIDsKeepOccurrenceCompletions(t *test require.Equal(t, "call-dup", outcome.BatchRuntimeToolCallID) } +// A call without a tool call ID, which providers can emit, still +// executes: its completion defines the window like any other billed +// call instead of being discarded as a missing result. +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: "{}"}, + }, + }) + + // Batch start. + trap.MustWait(ctx).MustRelease(ctx) + // The identified call completes at 10 seconds. + clock.Advance(10 * time.Second) + close(fastGo) + trap.MustWait(ctx).MustRelease(ctx) + // The ID-less call completes at 60 seconds and defines the window. + 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) + require.Empty(t, outcome.BatchRuntimeToolCallID) +} + // OnToolComplete reports each tool's completion the instant it finishes, // while slower siblings are still running, with the same instants the // outcome's ToolResultCreatedAt later carries. The interrupt path diff --git a/coderd/x/chatd/message_conversion.go b/coderd/x/chatd/message_conversion.go index 458ba3e26e7..f430e2ee28a 100644 --- a/coderd/x/chatd/message_conversion.go +++ b/coderd/x/chatd/message_conversion.go @@ -80,10 +80,13 @@ func buildCommitStepMessages(input buildCommitStepMessagesInput) (stepMessagesFo // usage reporting, which sums runtime_ms across rows, bills the // batch exactly once. Only the first row with the window-defining // ID carries it, because providers can emit duplicate tool call - // IDs and billing every duplicate would multiply the sum. Zero - // maps to NULL, so a sub-millisecond window persists the same - // way an unmeasured one does. - if !batchRuntimeAssigned && toolResult.ToolCallID != "" && toolResult.ToolCallID == input.step.BatchRuntimeToolCallID { + // IDs and billing every duplicate would multiply the sum. A + // positive window gates the match instead of a non-empty ID, + // because calls without IDs also execute and bill; their window + // lands on the first ID-less tool row. Zero maps to NULL, so a + // sub-millisecond window persists the same way an unmeasured + // one does. + if !batchRuntimeAssigned && input.step.BatchRuntime > 0 && toolResult.ToolCallID == input.step.BatchRuntimeToolCallID { msg.RuntimeMs = nullInt64IfNonZero(input.step.BatchRuntime.Milliseconds()) batchRuntimeAssigned = true } diff --git a/coderd/x/chatd/message_conversion_test.go b/coderd/x/chatd/message_conversion_test.go index 0d7515265fb..5c22606a4c9 100644 --- a/coderd/x/chatd/message_conversion_test.go +++ b/coderd/x/chatd/message_conversion_test.go @@ -168,6 +168,39 @@ func TestBuildCommitStepMessages_DuplicateToolCallIDsBillOnce(t *testing.T) { require.False(t, got.Messages[1].RuntimeMs.Valid) } +// A window defined by a call without a tool call ID still persists: the +// positive window gates the match, so the runtime lands on the first +// ID-less tool row instead of being dropped with the empty-ID sentinel. +func TestBuildCommitStepMessages_EmptyIDWindowLandsOnIDLessRow(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-fast", + ToolName: "fast_tool", + Result: fantasy.ToolResultOutputContentText{Text: `{"stdout":"fast"}`}, + }, + fantasy.ToolResultContent{ + ToolCallID: "", + ToolName: "idless_tool", + Result: fantasy.ToolResultOutputContentText{Text: `{"stdout":"slow"}`}, + }, + }, + BatchRuntime: 60 * time.Second, + BatchRuntimeToolCallID: "", + }, + }) + require.NoError(t, err) + require.Len(t, got.Messages, 2) + require.False(t, got.Messages[0].RuntimeMs.Valid, "the identified row is not the window-defining one") + require.Equal(t, sql.NullInt64{Int64: 60_000, Valid: true}, got.Messages[1].RuntimeMs) +} + // Assistant rows synthesized from a tool batch (attachment file parts) // never carry the batch runtime: it belongs to the tool row alone. func TestBuildCommitStepMessages_BatchAttachmentAssistantRowStaysNull(t *testing.T) { diff --git a/coderd/x/chatd/messagepartbuffer/message_part_buffer_test.go b/coderd/x/chatd/messagepartbuffer/message_part_buffer_test.go index f7aeaac05ba..84dd9eb204f 100644 --- a/coderd/x/chatd/messagepartbuffer/message_part_buffer_test.go +++ b/coderd/x/chatd/messagepartbuffer/message_part_buffer_test.go @@ -346,9 +346,13 @@ func TestBuffer_CloseEpisodeForBillingRefreshesEviction(t *testing.T) { require.Equal(t, first, again) // Once retries stop refreshing it, the episode ages out and a later - // close reports empty billing again. + // close reports empty billing again. GetParts collects due episodes + // synchronously, so the assertions cannot race the cleanup + // goroutine's handling of the delivered ticks. clock.Advance(time.Minute).MustWait(ctx) clock.Advance(time.Minute).MustWait(ctx) + _, err = buffer.GetParts(key) + require.ErrorIs(t, err, messagepartbuffer.ErrEpisodeNotFound) expired, err := buffer.CloseEpisodeForBilling(key) require.NoError(t, err) require.Zero(t, expired.ToolBatchStartedAt) From 3041968eb5e2130272b946c48fbf936c7b4d9fe2 Mon Sep 17 00:00:00 2001 From: Jaayden Halko Date: Tue, 18 Aug 2026 07:36:20 +0000 Subject: [PATCH 11/20] docs(coderd/usage): clarify that workspace tools bill regardless of agent type Address Codex review feedback: the runtime measurement docs excluded "external agents" without qualification, which reads as if tools run against a workspace with an externally managed agent were unbilled. The exclusion is about chats handed off to external coding agents, which park while the work happens outside the server. Server-executed tools dispatch to a connected workspace and wait on it during the turn for every workspace agent type, Coder-provisioned or external, and bill deliberately; excluding externally managed workspace agents would make tool billing trivially avoidable. Reword the usage-data-reporting page and the HBAgentRuntime doc comment to say so. --- coderd/usage/usagetypes/events.go | 10 +++++++--- docs/ai-coder/usage-data-reporting.md | 4 ++-- 2 files changed, 9 insertions(+), 5 deletions(-) diff --git a/coderd/usage/usagetypes/events.go b/coderd/usage/usagetypes/events.go index c4b67dad9eb..2b59238c614 100644 --- a/coderd/usage/usagetypes/events.go +++ b/coderd/usage/usagetypes/events.go @@ -211,9 +211,13 @@ func (e HBAISeats) Fields() map[string]any { // batch run in parallel, so the batch bills one window rather than a sum. // Sub-agent orchestration tools (spawn_agent, wait_agent, and the rest of // that category) are excluded because each sub-agent chat bills its own -// runtime. Also excluded: client-executed (dynamic) tools and external -// agents, time parked waiting for user action, idle time between turns, -// and retry backoff between attempts. +// runtime. Also excluded: client-executed (dynamic) tools and chats handed +// off to external coding agents (the chat parks while that work happens +// outside the server, which cannot measure it), time parked waiting for +// user action, idle time between turns, and retry backoff between +// attempts. Server-executed tools count even when their work runs in a +// connected workspace, whatever the workspace's agent type: the server +// dispatches and waits on them during the turn. // // This measures the new Coder Agents (the `chats` tables), not the deprecated // Tasks counted by dc_managed_agents_v1. diff --git a/docs/ai-coder/usage-data-reporting.md b/docs/ai-coder/usage-data-reporting.md index 6973a76ad90..cfcaf116cd6 100644 --- a/docs/ai-coder/usage-data-reporting.md +++ b/docs/ai-coder/usage-data-reporting.md @@ -68,13 +68,13 @@ What counts: - Assistant generation steps, in both top-level chats and sub-agent chats. - Context compaction (summarization) model calls. -- Local tool execution: file, terminal, and process tools, workspace lifecycle operations, MCP tools, and other server-executed tools. +- Local tool execution: file, terminal, and process tools, workspace lifecycle operations, MCP tools, and other server-executed tools. A tool's work may run in a connected workspace, but the server dispatches and waits on it during the turn, so it counts regardless of how the workspace's agent is managed. - Interrupted generation: the time streamed or spent executing tools before the interrupt is kept on the partial messages. What does not count: - Sub-agent orchestration tools, such as spawning and waiting on sub-agents. A sub-agent is its own chat and records its own runtime, so counting the parent's wait would double count. Waiting on a sub-agent never extends a tool batch's window, even when other tools in the batch do count. -- Client-executed (dynamic) tools and external agents: the server cannot measure work it does not execute. +- Client-executed (dynamic) tools and chats handed off to external coding agents: the chat parks while that work happens outside the server, which cannot measure it. - Idle time: chats waiting for user input or external tool results. - Failed model calls whose output was discarded, and the backoff between retried attempts. Retried and errored attempts persist no content, so they record no runtime. - Ancillary model calls that produce no chat messages, such as title generation. From 8d4760a5d2a913a47af3536472ff985c8c7c2dd7 Mon Sep 17 00:00:00 2001 From: Jaayden Halko Date: Tue, 18 Aug 2026 07:46:48 +0000 Subject: [PATCH 12/20] fix(coderd/x/chatd): seed tool batch billing only once dispatch is guaranteed Address Codex review feedback: generation seeded the buffer episode's tool batch before calling chatloop.ExecuteLocalTools, whose cancellation check runs inside. An interrupt whose cancellation landed in that gap dispatched no tools yet snapshotted seeded zero-completion entries, billing never-run calls from the batch stamp to the episode close. ExecuteLocalTools now takes an OnBatchStart callback invoked after its cancellation and exclusive-policy checks, immediately before the tool goroutines launch, where dispatch is unconditional; generation seeds the batch there. Canceled and policy-rejected batches never seed, so a racing interrupt finds no batch and bills nothing. --- coderd/x/chatd/chatloop/chatloop.go | 11 ++++ coderd/x/chatd/chatloop/runtime_test.go | 71 +++++++++++++++++++++++++ coderd/x/chatd/generation.go | 24 +++++---- 3 files changed, 97 insertions(+), 9 deletions(-) diff --git a/coderd/x/chatd/chatloop/chatloop.go b/coderd/x/chatd/chatloop/chatloop.go index 11348ed5c0b..a2e8982019e 100644 --- a/coderd/x/chatd/chatloop/chatloop.go +++ b/coderd/x/chatd/chatloop/chatloop.go @@ -265,6 +265,14 @@ type ExecuteLocalToolsOptions struct { // name as called, so deprecated aliases must be listed alongside // their canonical names. UnbilledToolNames map[string]bool + // OnBatchStart, when set, is invoked once cancellation and + // exclusive-policy checks have passed, immediately before the + // batch's tool goroutines launch. Dispatch is unconditional after + // it fires. The interrupt path uses it to seed the batch's billing + // state on the buffer episode: seeding any earlier would let an + // interrupt whose cancellation lands before dispatch bill calls + // that never ran as if they were still running. + OnBatchStart func() // OnToolComplete, when set, is invoked with each local tool call's // completion instant as the tool finishes, the same instant // PersistedStep.ToolResultCreatedAt later carries. Tool results are @@ -627,6 +635,9 @@ func ExecuteLocalTools(ctx context.Context, opts ExecuteLocalToolsOptions) (Tool } maxResultBytes := toolResultByteBudget(opts.ContextLimit) + if opts.OnBatchStart != nil { + opts.OnBatchStart() + } batchStart := clockNow(opts.Clock) // Completion instants aligned with localCalls by occurrence. // billableBatchWindow reads these instead of the ID-keyed diff --git a/coderd/x/chatd/chatloop/runtime_test.go b/coderd/x/chatd/chatloop/runtime_test.go index 0ba5002bf79..76489b72a72 100644 --- a/coderd/x/chatd/chatloop/runtime_test.go +++ b/coderd/x/chatd/chatloop/runtime_test.go @@ -437,6 +437,77 @@ func TestExecuteLocalTools_DuplicateToolCallIDsKeepOccurrenceCompletions(t *test require.Equal(t, "call-dup", outcome.BatchRuntimeToolCallID) } +// OnBatchStart seeds billing state only once dispatch is guaranteed: it +// fires before the tool goroutines launch, and never fires when +// cancellation or the exclusive policy stops the batch, so an interrupt +// racing those paths finds no batch to bill. +func TestExecuteLocalTools_OnBatchStartFiresOnlyOnDispatch(t *testing.T) { + t.Parallel() + + t.Run("dispatched batch seeds before tools run", func(t *testing.T) { + t.Parallel() + + starts := 0 + seededWhenToolRan := false + tool := fantasy.NewAgentTool( + "fast_tool", + "test tool", + func(context.Context, struct{}, fantasy.ToolCall) (fantasy.ToolResponse, error) { + seededWhenToolRan = 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"}, + OnBatchStart: func() { starts++ }, + ToolCalls: []fantasy.ToolCallContent{ + {ToolCallID: "call-1", ToolName: "fast_tool", Input: "{}"}, + }, + }) + require.NoError(t, err) + require.Len(t, outcome.Step.Content, 1) + require.Equal(t, 1, starts) + require.True(t, seededWhenToolRan, "the batch must be seeded before any tool goroutine runs") + }) + + t.Run("canceled context never seeds", func(t *testing.T) { + t.Parallel() + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + started := false + _, err := chatloop.ExecuteLocalTools(ctx, chatloop.ExecuteLocalToolsOptions{ + Clock: quartz.NewMock(t), + OnBatchStart: func() { started = true }, + ToolCalls: []fantasy.ToolCallContent{ + {ToolCallID: "call-1", ToolName: "fast_tool", Input: "{}"}, + }, + }) + require.ErrorIs(t, err, context.Canceled) + require.False(t, started, "a canceled batch dispatches nothing and must not seed billing state") + }) + + t.Run("exclusive violation never seeds", func(t *testing.T) { + t.Parallel() + + started := false + outcome, err := chatloop.ExecuteLocalTools(context.Background(), chatloop.ExecuteLocalToolsOptions{ + Clock: quartz.NewMock(t), + ExclusiveToolNames: map[string]bool{"exclusive_tool": true}, + OnBatchStart: func() { started = 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.Step.Content, 2, "the whole batch resolves to synthesized policy errors") + require.False(t, started, "a policy-rejected batch dispatches nothing and must not seed billing state") + }) +} + // A call without a tool call ID, which providers can emit, still // executes: its completion defines the window like any other billed // call instead of being discarded as a missing result. diff --git a/coderd/x/chatd/generation.go b/coderd/x/chatd/generation.go index 6d1fee6cb26..64ef5918b28 100644 --- a/coderd/x/chatd/generation.go +++ b/coderd/x/chatd/generation.go @@ -826,17 +826,20 @@ func (s *taskStarter) executeLocalTools( var outcome chatloop.ToolExecutionOutcome var spawnDispatchErr error if len(allowed) > 0 { - // Stamp the batch start and the dispatched calls on the buffer + // Seed the batch start and the dispatched calls on the buffer // episode so an interrupt can bill the partial window this step - // would have reported. Only dispatched calls are seeded, keyed - // by their position in the unresolved call order the interrupt - // walks: rejected calls never run, so an interrupt must not - // treat their missing completions as still-running work, and a + // would have reported. Chatloop invokes the callback only once + // dispatch is guaranteed, so an interrupt whose cancellation + // lands before the tools launch finds no batch and bills + // nothing. Only dispatched calls are seeded, keyed by their + // position in the unresolved call order the interrupt walks: + // rejected calls never run, so an interrupt must not treat + // their missing completions as still-running work, and a // rejected call must not consume a same-ID dispatched // occurrence. An exclusive-policy violation dispatches nothing - // (chatloop synthesizes error results for the whole batch), so - // no billable batch starts and an interrupt racing those - // synthetic results bills nothing. + // (chatloop synthesizes error results before the callback), so + // no billable batch starts there either. + var onBatchStart func() if !exclusiveRejected { dispatched := make([]messagepartbuffer.DispatchedToolCall, 0, len(allowed)) for j, tc := range allowed { @@ -845,7 +848,9 @@ func (s *taskStarter) executeLocalTools( ToolCallID: tc.ToolCallID, }) } - attempt.startToolBatch(dispatched) + onBatchStart = func() { + attempt.startToolBatch(dispatched) + } } outcome, err = chatloop.ExecuteLocalTools(ctx, chatloop.ExecuteLocalToolsOptions{ Tools: prepared.Tools, @@ -859,6 +864,7 @@ func (s *taskStarter) executeLocalTools( ContextLimit: prepared.ContextLimitFallback, ToolNameAliases: subagentToolNameAliases, UnbilledToolNames: unbilledSubagentToolNames, + OnBatchStart: onBatchStart, OnToolComplete: attempt.recordToolCompletion, PublishMessagePart: attempt.publish, Logger: s.opts.Logger, From ca3ae1aeb6f6252b7a434c1c122272adf4af4a7d Mon Sep 17 00:00:00 2001 From: Jaayden Halko Date: Tue, 18 Aug 2026 08:04:29 +0000 Subject: [PATCH 13/20] fix(coderd/x/chatd): capture the interrupt episode snapshot before the chat read Address Codex review feedback: the interrupt task read the chat row before snapshotting the buffer episode, so an initial read stalled past the buffer's retention let the cleanup loop evict the episode first, and the eventual snapshot found a blank recreated episode with no billing stamps or partial parts. StartInterrupt now captures the snapshot before its first database read, keyed by the task input's attempt number. Generation attempts only advance while the chat is running, so the pre-read key matches the row while the chat stays interrupting; if the post-read attempt number ever disagrees, the snapshot is retaken with the row's key. The close-and-read sequence moved into a closeInterruptEpisode helper shared by both paths. --- coderd/x/chatd/ARCHITECTURE.md | 2 +- coderd/x/chatd/tasks.go | 101 ++++++++++++++++++++++----------- coderd/x/chatd/tasks_test.go | 52 +++++++++++++++++ 3 files changed, 122 insertions(+), 33 deletions(-) diff --git a/coderd/x/chatd/ARCHITECTURE.md b/coderd/x/chatd/ARCHITECTURE.md index 59d6ee39fc4..6e9c8d48011 100644 --- a/coderd/x/chatd/ARCHITECTURE.md +++ b/coderd/x/chatd/ARCHITECTURE.md @@ -936,7 +936,7 @@ The goroutine does the following in order: 3. It reads the buffered parts for that episode by calling the `GetParts` method on the message part buffer. 4. It applies the `FinishInterruption(partial?)` transition on the core state machine. If there are no buffered parts for that episode, or the episode is not found, it passes `nil` as the `partial` argument. -TODO(CODAGT-928): update steps 2 and 3 above. The interrupt goroutine now closes the episode and snapshots its billing stamps in one atomic step via `CloseEpisodeForBilling`, uses the episode's first-close instant as the interrupt instant, carries the snapshot and buffered parts across task retry attempts so buffer eviction during a stalled attempt cannot lose them, and synthesizes tool cancellation rows whose `runtime_ms` bills the interrupted tool batch's partial window. +TODO(CODAGT-928): update steps 2 and 3 above. The interrupt goroutine now closes the episode and snapshots its billing stamps in one atomic step via `CloseEpisodeForBilling`, uses the episode's first-close instant as the interrupt instant, captures the snapshot and buffered parts before its first database read and carries them across task retry attempts so buffer eviction during a stalled read cannot lose them, and synthesizes tool cancellation rows whose `runtime_ms` bills the interrupted tool batch's partial window. #### Dynamic tools timeout goroutine diff --git a/coderd/x/chatd/tasks.go b/coderd/x/chatd/tasks.go index 711d778260e..d5e117ddf46 100644 --- a/coderd/x/chatd/tasks.go +++ b/coderd/x/chatd/tasks.go @@ -241,13 +241,13 @@ func (o chatWorkerOptions) retryOptions() retryWrapperOptions { } // interruptEpisodeSnapshot carries the interrupt task's one-time episode -// snapshot across retry attempts of the same task instance. One attempt -// can outlive the buffer's closed-episode retention (a machine.Update -// stalled on a database outage runs up to the task timeout), so -// re-reading the buffer on a later attempt could find a blank recreated -// episode and underbill the interrupted work or drop the partial -// messages. Attempts of one task run sequentially, so no locking is -// needed. +// snapshot across retry attempts of the same task instance. It is +// captured before the task's first database read: one stalled read (a +// database outage holds an attempt up to the task timeout) can outlive +// the buffer's closed-episode retention, after which the buffer only +// offers a blank recreated episode that would underbill the interrupted +// work and drop the partial messages. Attempts of one task run +// sequentially, so no locking is needed. type interruptEpisodeSnapshot struct { loaded bool key messagepartbuffer.Key @@ -255,7 +255,59 @@ type interruptEpisodeSnapshot struct { parts []messagepartbuffer.Part } +// closeInterruptEpisode closes the interrupted attempt's buffer episode +// and returns its billing snapshot and buffered parts. Closing and +// snapshotting billing state must be one atomic step: the generation +// goroutine records batch starts and tool completions concurrently, so +// a read-then-close would let stamps land in the gap and go missing +// from the snapshot. Unknown episodes close blank, so interruption +// converges even when the worker exited before publishing parts. +func (s *taskStarter) closeInterruptEpisode(ctx context.Context, key messagepartbuffer.Key) (messagepartbuffer.EpisodeBilling, []messagepartbuffer.Part, error) { + billing, err := s.opts.MessagePartBuffer.CloseEpisodeForBilling(key) + if err != nil { + if ctx.Err() != nil { + return messagepartbuffer.EpisodeBilling{}, nil, errors.Join(errTaskExpectedExit, xerrors.Errorf("close message part episode: %w", err), ctx.Err()) + } + return messagepartbuffer.EpisodeBilling{}, nil, taskRetryableError{err: xerrors.Errorf("close message part episode: %w", err)} + } + parts, err := s.opts.MessagePartBuffer.GetParts(key) + if errors.Is(err, messagepartbuffer.ErrEpisodeNotFound) { + parts = nil + err = nil + } + if err != nil { + if ctx.Err() != nil { + return messagepartbuffer.EpisodeBilling{}, nil, errors.Join(errTaskExpectedExit, xerrors.Errorf("get message part episode: %w", err), ctx.Err()) + } + return messagepartbuffer.EpisodeBilling{}, nil, taskRetryableError{err: xerrors.Errorf("get message part episode: %w", err)} + } + return billing, parts, nil +} + func (s *taskStarter) StartInterrupt(ctx context.Context, input chatWorkerTaskStartInput) error { + snapshot := input.InterruptSnapshot + // Capture the episode before the chat read below, which can stall + // on a database outage past the buffer's retention and let the + // cleanup loop evict the episode before a post-read capture. The + // input carries the interrupted attempt number, and generation + // attempts only advance while the chat is running, so the pre-read + // key matches the row while it stays interrupting; if the read + // below disagrees, the snapshot is retaken with the row's key. + if snapshot != nil && !snapshot.loaded { + earlyKey := messagepartbuffer.Key{ + ChatID: input.ChatID, + HistoryVersion: input.HistoryVersion, + GenerationAttempt: input.GenerationAttempt, + } + billing, parts, err := s.closeInterruptEpisode(ctx, earlyKey) + if err != nil { + return err + } + snapshot.loaded = true + snapshot.key = earlyKey + snapshot.billing = billing + snapshot.parts = parts + } machine := chatstate.NewChatMachine(s.opts.Store, s.opts.Pubsub, input.ChatID) var chat database.Chat err := machine.ReadLock(ctx, func(store database.Store) error { @@ -277,36 +329,21 @@ func (s *taskStarter) StartInterrupt(ctx context.Context, input chatWorkerTaskSt } var episodeBilling messagepartbuffer.EpisodeBilling var parts []messagepartbuffer.Part - snapshot := input.InterruptSnapshot if snapshot != nil && snapshot.loaded && snapshot.key == key { - // A previous attempt of this task already snapshotted the - // episode. Reuse it: the episode may have been evicted from - // the buffer while that attempt stalled, and re-reading would - // find a blank recreated episode. + // This task already snapshotted the episode, before its first + // chat read or on a previous attempt. Reuse it: the episode + // may have been evicted from the buffer while an attempt + // stalled, and re-reading would find a blank recreated + // episode. episodeBilling = snapshot.billing parts = snapshot.parts } else { - // Closing and snapshotting billing state must be one atomic - // step: the generation goroutine records batch starts and tool - // completions concurrently, so a read-then-close would let - // stamps land in the gap and go missing from the snapshot. - episodeBilling, err = s.opts.MessagePartBuffer.CloseEpisodeForBilling(key) + // No usable snapshot: either the caller passed none (tests) or + // the row's attempt number disagrees with the pre-read key, so + // the pre-read snapshot describes the wrong episode. + episodeBilling, parts, err = s.closeInterruptEpisode(ctx, key) if err != nil { - if ctx.Err() != nil { - return errors.Join(errTaskExpectedExit, xerrors.Errorf("close message part episode: %w", err), ctx.Err()) - } - return taskRetryableError{err: xerrors.Errorf("close message part episode: %w", err)} - } - parts, err = s.opts.MessagePartBuffer.GetParts(key) - if errors.Is(err, messagepartbuffer.ErrEpisodeNotFound) { - parts = nil - err = nil - } - if err != nil { - if ctx.Err() != nil { - return errors.Join(errTaskExpectedExit, xerrors.Errorf("get message part episode: %w", err), ctx.Err()) - } - return taskRetryableError{err: xerrors.Errorf("get message part episode: %w", err)} + return err } if snapshot != nil { snapshot.loaded = true diff --git a/coderd/x/chatd/tasks_test.go b/coderd/x/chatd/tasks_test.go index 10200f21b5a..c5f7583430f 100644 --- a/coderd/x/chatd/tasks_test.go +++ b/coderd/x/chatd/tasks_test.go @@ -662,6 +662,58 @@ func TestInterruptTask_RetrySnapshotOutlivesEpisodeEviction(t *testing.T) { require.Equal(t, sql.NullInt64{Int64: 7_000, Valid: true}, execRow.RuntimeMs) } +// The interrupt task captures its episode snapshot before its first +// database read: an attempt whose chat read fails (or stalls past the +// buffer's retention, evicting the episode) has already stored the +// billing state, so a later attempt bills the original window even +// after the buffer forgets the episode. +func TestInterruptTask_SnapshotCapturedBeforeChatRead(t *testing.T) { + t.Parallel() + + f := newTaskTestFixture(t) + execCallID := "call_" + uuid.NewString() + batch := interruptedBatchFixture(t, f, []codersdk.ChatMessagePart{ + {Type: codersdk.ChatMessagePartTypeToolCall, ToolCallID: execCallID, ToolName: "execute", Args: json.RawMessage(`{}`)}, + }) + buffer := batch.starter.opts.MessagePartBuffer + + batch.clock.Advance(2 * time.Second) + require.NoError(t, buffer.StartToolBatch(batch.key, []messagepartbuffer.DispatchedToolCall{{CallIndex: 0, ToolCallID: execCallID}})) + batch.clock.Advance(3 * time.Second) + + // The first attempt's chat read fails: the chat is still running, + // so the interrupting fence rejects it. The snapshot must already + // be captured by then, with the interrupt instant at this failed + // attempt's episode close. + snapshot := &interruptEpisodeSnapshot{} + err := batch.starter.StartInterrupt(testutil.Context(t, testutil.WaitShort), chatWorkerTaskStartInput{ + ChatID: batch.chat.ID, + WorkerID: batch.workerID, + RunnerID: batch.runnerID, + HistoryVersion: batch.key.HistoryVersion, + GenerationAttempt: batch.key.GenerationAttempt, + Status: database.ChatStatusInterrupting, + InterruptSnapshot: snapshot, + }) + require.Error(t, err) + require.True(t, snapshot.loaded, "the snapshot must be captured before the chat read") + require.Equal(t, batch.key, snapshot.key) + + // The attempt outlives the buffer retention: cleanup ticks evict + // the closed episode, so only the carried snapshot remains. + ctx := testutil.Context(t, testutil.WaitShort) + batch.clock.Advance(10 * time.Second).MustWait(ctx) + batch.clock.Advance(15 * time.Second).MustWait(ctx) + _, err = buffer.GetParts(batch.key) + require.ErrorIs(t, err, messagepartbuffer.ErrEpisodeNotFound) + + // The retry bills the pre-read window: batch start at 2s to the + // first attempt's close at 5s. + messages := batch.interruptWithSnapshot(t, f, snapshot) + execRow := findToolResultMessage(t, messages, execCallID) + require.Equal(t, sql.NullInt64{Int64: 3_000, Valid: true}, execRow.RuntimeMs) +} + // A billed tool still running at the interrupt bills up to the interrupt // instant. func TestInterruptTask_RunningBilledToolBillsUpToInterrupt(t *testing.T) { From d2d9a23a6443ee93c5255d55fb53fe96879c9be8 Mon Sep 17 00:00:00 2001 From: Jaayden Halko Date: Wed, 19 Aug 2026 09:18:03 +0000 Subject: [PATCH 14/20] fix(coderd/x/chatd): bill serial tool calls from their actual start Serial tool calls (SerialToolCalls) launch only after every concurrent sibling settles, so measuring them from the shared batch start charged the whole concurrent phase, including waits on unbilled sub-agent orchestration tools, and an interrupt treated a dispatched serial call that never launched as running since batch start. The batch window is now the union of the billed calls' execution intervals: chatloop stamps each call's start (concurrent calls at batch start, serial calls at their launch) and reports it through the new OnToolStart callback, the buffer records the marks per occurrence via RecordToolStart, and both the committed path and the interrupt path bill with the shared BilledIntervalsDuration union, so a span where only unbilled tools were running never bills and a dispatched call without a start mark bills nothing. --- coderd/usage/usagetypes/events.go | 8 +- coderd/x/chatd/ARCHITECTURE.md | 2 +- coderd/x/chatd/chatloop/chatloop.go | 134 ++++++++++++-- .../chatloop/chatloop_run_internal_test.go | 8 + coderd/x/chatd/chatloop/runtime_test.go | 170 ++++++++++++++++++ coderd/x/chatd/generation.go | 13 ++ .../messagepartbuffer/message_part_buffer.go | 58 +++++- .../message_part_buffer_test.go | 56 ++++++ coderd/x/chatd/tasks.go | 58 ++++-- coderd/x/chatd/tasks_test.go | 102 ++++++++++- docs/ai-coder/usage-data-reporting.md | 3 +- 11 files changed, 566 insertions(+), 46 deletions(-) diff --git a/coderd/usage/usagetypes/events.go b/coderd/usage/usagetypes/events.go index 2b59238c614..3bb24f033d2 100644 --- a/coderd/usage/usagetypes/events.go +++ b/coderd/usage/usagetypes/events.go @@ -206,9 +206,11 @@ func (e HBAISeats) Fields() map[string]any { // is the total agent-loop runtime in milliseconds consumed by Coder Agents // (chats) in one UTC hour. Two kinds of windows are measured. Model steps // span model streaming (including provider-executed tools) and stream -// retries, ending when the model stream finishes. Local tool batches span -// the start of the batch until the last billed tool completes; tools in a -// batch run in parallel, so the batch bills one window rather than a sum. +// retries, ending when the model stream finishes. Local tool batches bill +// the union of the billed tools' execution windows: tools in a batch run +// in parallel, so the batch bills one window rather than a sum, and a +// tool that runs serially after its siblings is measured from its own +// start. // Sub-agent orchestration tools (spawn_agent, wait_agent, and the rest of // that category) are excluded because each sub-agent chat bills its own // runtime. Also excluded: client-executed (dynamic) tools and chats handed diff --git a/coderd/x/chatd/ARCHITECTURE.md b/coderd/x/chatd/ARCHITECTURE.md index f83b853ca84..be9cc88a47d 100644 --- a/coderd/x/chatd/ARCHITECTURE.md +++ b/coderd/x/chatd/ARCHITECTURE.md @@ -725,7 +725,7 @@ The buffer exposes the following API: - `GetParts(chat_id, history_version, generation_attempt)`: returns the message parts for an episode. Returns a predefined error if the episode is not found. - `StartModelInvocation(chat_id, history_version, generation_attempt)`: stamps the instant the episode opens its provider stream. Returns a predefined error if the episode is not found or already closed. Episodes that never invoke a model, such as local tool execution batches, are never stamped. - `ModelInvokedAt(chat_id, history_version, generation_attempt)`: returns the instant stamped by `StartModelInvocation`, or the zero time when the episode is unknown or never opened a provider stream. It must be read before `CloseEpisode`, because closed episodes are garbage collected and reading afterwards races the cleanup loop. The interrupt goroutine reads it just before closing the episode and uses the span between that instant and the interrupt as the interrupted attempt's billable runtime. -- TODO(CODAGT-928): document the tool-batch billing methods `StartToolBatch`, `ToolBatchStartedAt`, `RecordToolCompletion`, and `ToolCompletions`, the local-tool counterparts of `StartModelInvocation`/`ModelInvokedAt`, and `CloseEpisodeForBilling`, which the interrupt task uses to close the episode and snapshot its billing stamps (including the stable first-close instant used as the interrupt instant) in one atomic step; repeat calls return the identical snapshot and push the eviction deadline out, so retried interrupt tasks bill the same window and keep their snapshot through a database outage. +- TODO(CODAGT-928): document the tool-batch billing methods `StartToolBatch`, `ToolBatchStartedAt`, `RecordToolStart`, `RecordToolCompletion`, and `ToolCompletions`, the local-tool counterparts of `StartModelInvocation`/`ModelInvokedAt`, and `CloseEpisodeForBilling`, which the interrupt task uses to close the episode and snapshot its billing stamps (including the stable first-close instant used as the interrupt instant) in one atomic step; repeat calls return the identical snapshot and push the eviction deadline out, so retried interrupt tasks bill the same window and keep their snapshot through a database outage. `RecordToolStart` marks when a dispatched call begins executing: serial tool calls launch only after every concurrent sibling settles, so the interrupt bills each call from its own start and bills nothing for a dispatched call that never launched. - `SubscribeToEpisode(chat_id, history_version, generation_attempt)`: returns a go channel that will receive all message parts for the episode. It spawns a goroutine that delivers parts to the channel. It's live until the episode is closed or until a subscriber requests that the channel be closed. Once the goroutine delivers all message parts for a closed episode, it closes the channel and exits. If the episode is already closed at the time of the call, the goroutine delivers all message parts for the episode, closes the channel, and exits. `SubscribeToEpisode` does not return an error if the episode is not found: it waits for it to be created instead. Closed episodes are garbage collected after at least 15 seconds since they were closed and when they have no active subscribers. The message part buffer maintains a garbage collection goroutine. diff --git a/coderd/x/chatd/chatloop/chatloop.go b/coderd/x/chatd/chatloop/chatloop.go index 692a7f8bd23..39dd3b3e91b 100644 --- a/coderd/x/chatd/chatloop/chatloop.go +++ b/coderd/x/chatd/chatloop/chatloop.go @@ -279,6 +279,18 @@ type ExecuteLocalToolsOptions struct { // interrupt whose cancellation lands before dispatch bill calls // that never ran as if they were still running. OnBatchStart func() + // OnToolStart, when set, is invoked with each local tool call's + // occurrence index, ID, and the instant the call begins executing. + // Concurrent calls start together when the batch launches, but + // calls whose tools opt in to SerialToolCalls launch only after + // every concurrent sibling has settled, so their start marks + // arrive later. The interrupt path uses the marks to bill each + // interrupted call from its actual start rather than the batch + // start, and to avoid billing a dispatched serial call that never + // began executing. Serial calls invoke it from the batch goroutine + // after concurrent siblings finished, so implementations must be + // safe for use from a different goroutine than OnBatchStart. + OnToolStart func(callIndex int, toolCallID string, startedAt time.Time) // OnToolComplete, when set, is invoked with each local tool call's // completion instant as the tool finishes, the same instant // PersistedStep.ToolResultCreatedAt later carries. Tool results are @@ -651,6 +663,18 @@ func ExecuteLocalTools(ctx context.Context, opts ExecuteLocalToolsOptions) (Tool // reach execution when lifecycle hooks are disabled, would // overwrite each other and corrupt the window. orderedCompletions := make([]time.Time, 0, len(localCalls)) + // Start instants aligned with localCalls by occurrence. Concurrent + // calls start at batchStart, but serial calls launch only after + // every concurrent sibling settles, so billing them from + // batchStart would charge the whole concurrent phase, including + // waits on unbilled tools. + orderedStarts := make([]time.Time, len(localCalls)) + onToolStart := func(callIndex int, toolCallID string, startedAt time.Time) { + orderedStarts[callIndex] = startedAt + if opts.OnToolStart != nil { + opts.OnToolStart(callIndex, toolCallID, startedAt) + } + } toolResults := executeTools( ctx, opts.Clock, @@ -667,6 +691,8 @@ func ExecuteLocalTools(ctx context.Context, opts ExecuteLocalToolsOptions) (Tool opts.BuiltinToolNames, maxResultBytes, opts.ToolNameAliases, + batchStart, + onToolStart, opts.OnToolComplete, func(tr fantasy.ToolResultContent, completedAt time.Time) { // onResult fires once per local call in call order, so @@ -689,6 +715,7 @@ func ExecuteLocalTools(ctx context.Context, opts ExecuteLocalToolsOptions) (Tool batchRuntime, batchRuntimeToolCallID := billableBatchWindow( batchStart, localCalls, + orderedStarts, orderedCompletions, opts.UnbilledToolNames, ) @@ -703,21 +730,25 @@ func ExecuteLocalTools(ctx context.Context, opts ExecuteLocalToolsOptions) (Tool } // billableBatchWindow computes one local tool batch's billable runtime: -// the span from batchStart to the latest completion among billed tools. -// Because all calls in a batch start together, that span equals the -// union of the billed tools' execution intervals, so parallel calls are -// billed once rather than summed, and unbilled tools (for example -// sub-agent orchestration) never extend the window even when they run -// longest. completions is aligned with toolCalls by occurrence, so -// duplicate tool call IDs each keep their own completion. Returns the -// ID of the tool call whose completion ends the window, with ties -// broken by call order; a zero duration means no billed tool completed -// or the window rounds to nothing. The ID can be empty even when the -// window exists, because providers can emit calls without IDs and -// those still execute and bill. +// the union of the billed tools' execution intervals. Concurrent calls +// all start together at batchStart, so parallel calls are billed once +// rather than summed, and unbilled tools (for example sub-agent +// orchestration) never extend the window even when they run longest. +// Calls whose tools opt in to SerialToolCalls launch only after every +// concurrent sibling settles, so each is measured from its own instant +// in starts, and a span where only unbilled tools were running never +// bills. starts and completions are aligned with toolCalls by +// occurrence, so duplicate tool call IDs each keep their own window; a +// zero start means the call launched with the batch. Returns the ID of +// the tool call whose completion ends the window, with ties broken by +// call order; a zero duration means no billed tool completed or the +// window rounds to nothing. The ID can be empty even when the window +// exists, because providers can emit calls without IDs and those still +// execute and bill. func billableBatchWindow( batchStart time.Time, toolCalls []fantasy.ToolCallContent, + starts []time.Time, completions []time.Time, unbilledToolNames map[string]bool, ) (time.Duration, string) { @@ -725,6 +756,7 @@ func billableBatchWindow( found bool windowEnd time.Time toolCallID string + intervals []BilledInterval ) for i, tc := range toolCalls { if i >= len(completions) { @@ -737,6 +769,11 @@ func billableBatchWindow( if end.IsZero() { continue } + start := batchStart + if i < len(starts) && !starts[i].IsZero() { + start = starts[i] + } + intervals = append(intervals, BilledInterval{Start: start, End: end}) // Strictly-after keeps the earliest call on ties. if end.After(windowEnd) { found = true @@ -744,10 +781,65 @@ func billableBatchWindow( toolCallID = tc.ToolCallID } } - if !found || !windowEnd.After(batchStart) { + if !found { return 0, "" } - return windowEnd.Sub(batchStart), toolCallID + runtime := BilledIntervalsDuration(intervals) + if runtime <= 0 { + return 0, "" + } + return runtime, toolCallID +} + +// BilledInterval is one billed tool call's execution window. +type BilledInterval struct { + Start time.Time + End time.Time +} + +// BilledIntervalsDuration returns the total length of the union of the +// given intervals. Overlapping windows, such as concurrent tool calls +// sharing a batch start, bill once rather than summing, while gaps, +// such as a span where only unbilled tools were running before a +// serial call launched, bill nothing. Intervals whose End precedes +// their Start are ignored. The interrupt path shares this helper so a +// canceled batch bills the same window the committed step would have +// reported. +func BilledIntervalsDuration(intervals []BilledInterval) time.Duration { + valid := make([]BilledInterval, 0, len(intervals)) + for _, iv := range intervals { + if iv.End.Before(iv.Start) { + continue + } + valid = append(valid, iv) + } + slices.SortFunc(valid, func(a, b BilledInterval) int { + return a.Start.Compare(b.Start) + }) + var ( + total time.Duration + curStart time.Time + curEnd time.Time + open bool + ) + for _, iv := range valid { + if !open { + curStart, curEnd, open = iv.Start, iv.End, true + continue + } + if iv.Start.After(curEnd) { + total += curEnd.Sub(curStart) + curStart, curEnd = iv.Start, iv.End + continue + } + if iv.End.After(curEnd) { + curEnd = iv.End + } + } + if open { + total += curEnd.Sub(curStart) + } + return total } // prepareMessagesForRequest applies the prompt preparation pipeline used @@ -1196,6 +1288,8 @@ func executeTools( builtinToolNames map[string]bool, maxResultBytes int, toolNameAliases map[string]string, + batchStart time.Time, + onStart func(callIndex int, toolCallID string, startedAt time.Time), onComplete func(callIndex int, toolCallID string, completedAt time.Time), onResult func(fantasy.ToolResultContent, time.Time), ) []fantasy.ToolResultContent { @@ -1298,6 +1392,11 @@ func executeTools( serialIndexes = append(serialIndexes, i) continue } + // Concurrent calls all start executing now, at the batch + // start already stamped by the caller. + if onStart != nil { + onStart(i, tc.ToolCallID, batchStart) + } wg.Add(1) go func() { defer wg.Done() @@ -1318,6 +1417,13 @@ func executeTools( notifyStepToolResultObservers(toolMap, toolNameAliases, localToolCalls, observed, settled) for _, i := range serialIndexes { + // A serial call starts executing only here, after every + // concurrent sibling settled, so its start mark is stamped at + // launch rather than at batchStart: billing it from + // batchStart would charge the whole concurrent phase. + if onStart != nil { + onStart(i, localToolCalls[i].ToolCallID, 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 bd7b5699e60..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,8 @@ func TestExecuteToolsNotifiesStepToolCallObservers(t *testing.T) { map[string]bool{}, defaultToolResultBytes, map[string]string{"observer_alias": "observer_tool"}, + time.Time{}, + nil, nil, nil, ) @@ -1021,6 +1023,8 @@ func TestExecuteToolsNotifiesStepToolResultObservers(t *testing.T) { map[string]bool{}, defaultToolResultBytes, map[string]string{"observer_alias": "observer_tool"}, + time.Time{}, + nil, nil, nil, ) @@ -1099,6 +1103,8 @@ func TestExecuteToolsReconcilesResultsBeforeSerialCalls(t *testing.T) { map[string]bool{}, defaultToolResultBytes, nil, + time.Time{}, + nil, nil, nil, ) @@ -1167,6 +1173,8 @@ 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 76489b72a72..0dd9f5ded24 100644 --- a/coderd/x/chatd/chatloop/runtime_test.go +++ b/coderd/x/chatd/chatloop/runtime_test.go @@ -195,6 +195,15 @@ func blockingTool(name string, release <-chan struct{}, response fantasy.ToolRes ) } +// serialTool wraps a tool so its calls run serially after every +// concurrent sibling settles, matching tools that opt in via +// SerialToolCalls. +type serialTool struct { + fantasy.AgentTool +} + +func (serialTool) SerialToolCalls() bool { return true } + // Parallel billed tools bill one shared window ending at the slowest // tool's completion, never the sum of their durations. The slower tool // returns an error result: errored tools bill their wall clock too. @@ -611,3 +620,164 @@ func TestExecuteLocalTools_OnToolCompleteReportsLiveCompletions(t *testing.T) { "call-slow": slow.completedAt, }, outcome.Step.ToolResultCreatedAt) } + +// A billed serial tool queued behind a long-running unbilled tool bills +// only its own execution: it launches only after every concurrent +// sibling settles, so measuring it from the batch start would charge +// the whole unbilled wait that delayed its launch. +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: "{}"}, + }, + }) + + // Batch start. + trap.MustWait(ctx).MustRelease(ctx) + // The unbilled wait completes 10 minutes in. + clock.Advance(10 * time.Minute) + close(waitGo) + trap.MustWait(ctx).MustRelease(ctx) + // The serial call launches only now and stamps its start. + trap.MustWait(ctx).MustRelease(ctx) + // It completes 2 seconds later. + 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") + require.Equal(t, "call-serial", outcome.BatchRuntimeToolCallID) +} + +// A batch mixing billed concurrent and billed serial calls bills the +// union of their execution intervals: the concurrent window and the +// serial call's own window sum, while the gap where only the unbilled +// tool was running is never charged. +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: "{}"}, + }, + }) + + // Batch start. + trap.MustWait(ctx).MustRelease(ctx) + // execute completes 3 seconds in. + clock.Advance(3 * time.Second) + close(execGo) + trap.MustWait(ctx).MustRelease(ctx) + // wait_agent keeps blocking until 10 seconds. + clock.Advance(7 * time.Second) + close(waitGo) + trap.MustWait(ctx).MustRelease(ctx) + // The serial call launches at 10 seconds. + trap.MustWait(ctx).MustRelease(ctx) + // It completes at 12 seconds. + 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") + require.Equal(t, "call-serial", outcome.BatchRuntimeToolCallID, + "the serial call's completion ends the window") +} + +// BilledIntervalsDuration merges overlapping windows and skips gaps, so +// parallel calls bill once and spans without billed work bill nothing. +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 acecb8b65d0..ba8aeafb352 100644 --- a/coderd/x/chatd/generation.go +++ b/coderd/x/chatd/generation.go @@ -906,6 +906,7 @@ func (s *taskStarter) executeLocalTools( ToolNameAliases: subagentToolNameAliases, UnbilledToolNames: unbilledSubagentToolNames, OnBatchStart: onBatchStart, + OnToolStart: attempt.recordToolStart, OnToolComplete: attempt.recordToolCompletion, PublishMessagePart: attempt.publish, Logger: s.opts.Logger, @@ -1129,6 +1130,15 @@ type generationAttempt struct { // calls that were rejected before execution. It is always non-nil // when beginGenerationAttempt succeeds. startToolBatch func(calls []messagepartbuffer.DispatchedToolCall) + // recordToolStart records the instant a tool call occurrence began + // executing on the buffer episode. Concurrent calls start with the + // batch, but serial calls launch only after every concurrent + // sibling settles, so an interrupt needs the marks to bill each + // call from its actual start and to skip a dispatched call that + // never launched. callIndex addresses the occurrence within the + // dispatched batch, matching the order startToolBatch seeded. It + // is always non-nil when beginGenerationAttempt succeeds. + recordToolStart func(callIndex int, toolCallID string, startedAt time.Time) // recordToolCompletion records a tool call occurrence's completion // instant on the buffer episode as the batch executes, so an // interrupt can end an already-finished tool's billable window at @@ -1186,6 +1196,9 @@ func (s *taskStarter) beginGenerationAttempt( startToolBatch: func(calls []messagepartbuffer.DispatchedToolCall) { _ = s.opts.MessagePartBuffer.StartToolBatch(key, calls) }, + recordToolStart: func(callIndex int, toolCallID string, startedAt time.Time) { + _ = s.opts.MessagePartBuffer.RecordToolStart(key, callIndex, toolCallID, startedAt) + }, recordToolCompletion: func(callIndex int, toolCallID string, completedAt time.Time) { _ = s.opts.MessagePartBuffer.RecordToolCompletion(key, callIndex, toolCallID, completedAt) }, diff --git a/coderd/x/chatd/messagepartbuffer/message_part_buffer.go b/coderd/x/chatd/messagepartbuffer/message_part_buffer.go index 175f3d9616d..a4d48d89d87 100644 --- a/coderd/x/chatd/messagepartbuffer/message_part_buffer.go +++ b/coderd/x/chatd/messagepartbuffer/message_part_buffer.go @@ -246,24 +246,28 @@ type DispatchedToolCall struct { } // ToolCompletion tracks one dispatched tool call occurrence in an -// episode's local tool batch. CompletedAt is zero while the call is -// still running. +// episode's local tool batch. StartedAt is zero until the call begins +// executing, which for a serial tool call happens only after every +// concurrent sibling has settled; CompletedAt is zero while the call +// has not finished. type ToolCompletion struct { // CallIndex mirrors DispatchedToolCall.CallIndex. It is -1 for // completions recorded without a matching seeded occurrence, which // never correlate to an unresolved call. CallIndex int ToolCallID string + StartedAt time.Time CompletedAt time.Time } // StartToolBatch stamps the instant the episode begins executing its local // tool batch, which starts the batch's billable runtime window, and seeds -// one still-running entry per dispatched tool call occurrence. Readers can -// then distinguish a dispatched call that is still running (seeded, zero -// completion) from one never dispatched at all (absent), such as a call -// denied by a lifecycle hook or rejected as ambiguous before execution, -// and duplicate tool call IDs keep distinct per-occurrence states. +// one not-yet-started entry per dispatched tool call occurrence. Readers +// can then distinguish a dispatched call awaiting launch (seeded, zero +// start), one that is executing (started, zero completion), and one never +// dispatched at all (absent), such as a call denied by a lifecycle hook or +// rejected as ambiguous before execution, and duplicate tool call IDs keep +// distinct per-occurrence states. func (b *Buffer) StartToolBatch(key Key, calls []DispatchedToolCall) error { b.mu.Lock() defer b.mu.Unlock() @@ -288,6 +292,46 @@ func (b *Buffer) StartToolBatch(key Key, calls []DispatchedToolCall) error { return nil } +// RecordToolStart records the instant a dispatched local tool call +// occurrence began executing. Concurrent calls start with the batch, but +// calls whose tools run serially launch only after every concurrent +// sibling has settled, so an interrupt must bill each call from its own +// start and must not treat a dispatched call that never launched as +// running work. dispatchIndex addresses the exact occurrence seeded by +// StartToolBatch, mirroring RecordToolCompletion; when it does not match, +// the stamp falls back to the first unstarted occurrence with the ID. A +// start with no seeded occurrence is dropped: it cannot correlate to an +// unresolved call, so it could never bill. +func (b *Buffer) RecordToolStart(key Key, dispatchIndex int, toolCallID string, startedAt time.Time) error { + b.mu.Lock() + defer b.mu.Unlock() + if b.closed { + return ErrMessagePartBufferClosed + } + episode, err := b.getEpisodeLocked(key) + if err != nil { + return err + } + if episode.closed { + return ErrEpisodeClosed + } + if dispatchIndex >= 0 && dispatchIndex < len(episode.toolCompletions) { + entry := &episode.toolCompletions[dispatchIndex] + if entry.ToolCallID == toolCallID && entry.StartedAt.IsZero() { + entry.StartedAt = startedAt + return nil + } + } + for i := range episode.toolCompletions { + entry := &episode.toolCompletions[i] + if entry.ToolCallID == toolCallID && entry.StartedAt.IsZero() { + entry.StartedAt = startedAt + return nil + } + } + return nil +} + // RecordToolCompletion records the instant a local tool call in the // episode's tool batch finished. Tool goroutines report completions as // they happen, so an interrupt can bill tools that already finished up diff --git a/coderd/x/chatd/messagepartbuffer/message_part_buffer_test.go b/coderd/x/chatd/messagepartbuffer/message_part_buffer_test.go index 84dd9eb204f..0f604270e70 100644 --- a/coderd/x/chatd/messagepartbuffer/message_part_buffer_test.go +++ b/coderd/x/chatd/messagepartbuffer/message_part_buffer_test.go @@ -256,6 +256,62 @@ func TestBuffer_ToolCompletions(t *testing.T) { require.True(t, buffer.ToolCompletions(key)[0].CompletedAt.IsZero()) } +func TestBuffer_RecordToolStart(t *testing.T) { + t.Parallel() + + clock := quartz.NewMock(t) + buffer := messagepartbuffer.New(messagepartbuffer.Options{Clock: clock}) + defer buffer.Close() + + key := testEpisodeKey() + require.ErrorIs(t, buffer.RecordToolStart(key, 0, "call-1", clock.Now()), messagepartbuffer.ErrEpisodeNotFound) + + require.NoError(t, buffer.CreateEpisode(key)) + // Seeded occurrences carry no start: a dispatched serial call + // waits behind its concurrent siblings, so seeding must not make + // it look like running work. + require.NoError(t, buffer.StartToolBatch(key, []messagepartbuffer.DispatchedToolCall{ + {CallIndex: 0, ToolCallID: "call-1"}, + {CallIndex: 1, ToolCallID: "call-dup"}, + {CallIndex: 2, ToolCallID: "call-dup"}, + })) + for _, completion := range buffer.ToolCompletions(key) { + require.True(t, completion.StartedAt.IsZero(), "seeding must not mark occurrences as started") + } + + // A start stamps the occurrence addressed by its dispatch index, + // not the first occurrence with a matching ID: the duplicate's + // second occurrence launching must leave the first unstarted. + clock.Advance(time.Second) + secondDupStartedAt := clock.Now() + require.NoError(t, buffer.RecordToolStart(key, 2, "call-dup", secondDupStartedAt)) + require.Equal(t, []messagepartbuffer.ToolCompletion{ + {CallIndex: 0, ToolCallID: "call-1"}, + {CallIndex: 1, ToolCallID: "call-dup"}, + {CallIndex: 2, ToolCallID: "call-dup", StartedAt: secondDupStartedAt}, + }, buffer.ToolCompletions(key)) + + // A start whose dispatch index does not match falls back to the + // first unstarted occurrence with the ID. + clock.Advance(time.Second) + firstDupStartedAt := clock.Now() + require.NoError(t, buffer.RecordToolStart(key, 7, "call-dup", firstDupStartedAt)) + require.Equal(t, firstDupStartedAt, buffer.ToolCompletions(key)[1].StartedAt) + + // A start whose ID was never seeded is dropped rather than + // appended: it cannot correlate to an unresolved call, so it could + // never bill. + require.NoError(t, buffer.RecordToolStart(key, 9, "call-unseeded", clock.Now())) + require.Len(t, buffer.ToolCompletions(key), 3) + + // The billing snapshot carries start marks through the close. + billing, err := buffer.CloseEpisodeForBilling(key) + require.NoError(t, err) + require.Equal(t, secondDupStartedAt, billing.ToolCompletions[2].StartedAt) + require.ErrorIs(t, buffer.RecordToolStart(key, 0, "call-1", clock.Now()), messagepartbuffer.ErrEpisodeClosed) + require.True(t, buffer.ToolCompletions(key)[0].StartedAt.IsZero()) +} + func TestBuffer_CloseEpisodeForBilling(t *testing.T) { t.Parallel() diff --git a/coderd/x/chatd/tasks.go b/coderd/x/chatd/tasks.go index d5e117ddf46..35a4b2820fd 100644 --- a/coderd/x/chatd/tasks.go +++ b/coderd/x/chatd/tasks.go @@ -15,6 +15,7 @@ import ( "github.com/coder/coder/v2/coderd/database" coderdpubsub "github.com/coder/coder/v2/coderd/pubsub" "github.com/coder/coder/v2/coderd/x/chatd/chatdebug" + "github.com/coder/coder/v2/coderd/x/chatd/chatloop" "github.com/coder/coder/v2/coderd/x/chatd/chatprompt" "github.com/coder/coder/v2/coderd/x/chatd/chatstate" "github.com/coder/coder/v2/coderd/x/chatd/messagepartbuffer" @@ -776,16 +777,18 @@ type interruptedToolBatchBilling struct { batchStartedAt time.Time // toolCompletions holds the live batch's dispatched tool call // occurrences, each keyed by its position in the step's unresolved - // call order: seeded with zero completions when the batch started - // and stamped as each tool finished. A stamped occurrence ends its - // billable window at its completion, so a batch whose billed tools - // all finished early does not bill the longer window of a - // still-running unbilled tool such as wait_agent. A seeded but - // unstamped occurrence was still running when the interrupt - // landed. A call with no occurrence at its position was never - // dispatched, such as a call rejected as malformed or ambiguous - // before execution, and bills nothing even though it too receives - // a cancellation row. + // call order: seeded when the batch started, marked as each call + // began executing, and stamped as each tool finished. A completed + // occurrence bills the interval from its start to its completion, + // so a batch whose billed tools all finished early does not bill + // the longer window of a still-running unbilled tool such as + // wait_agent. A started but uncompleted occurrence was still + // running when the interrupt landed. A seeded occurrence that + // never started is a serial call that was still waiting behind + // its siblings, and bills nothing. A call with no occurrence at + // its position was never dispatched, such as a call rejected as + // malformed or ambiguous before execution, and bills nothing even + // though it too receives a cancellation row. toolCompletions []messagepartbuffer.ToolCompletion } @@ -825,6 +828,7 @@ func committedPendingLocalToolCancellationMessages( var ( windowEnd time.Time windowRowIdx = -1 + intervals []chatloop.BilledInterval ) result := make([]chatstate.Message, 0, len(localCalls)) for i, call := range localCalls { @@ -852,28 +856,44 @@ func committedPendingLocalToolCancellationMessages( } // Only dispatched calls bill: a call with no occurrence at its // position was rejected before execution, and an ID mismatch - // means the seed does not describe this call. A stamped - // occurrence finished at that instant; a seeded but unstamped - // one was still running, so its window ends at the interrupt. - // Strictly-after keeps the earliest call on ties, matching - // billableBatchWindow. + // means the seed does not describe this call. A completed + // occurrence bills its execution interval; a started but + // uncompleted one was still running, so its window ends at the + // interrupt. A dispatched occurrence that never started is a + // serial call still waiting behind its siblings when the + // interrupt landed, and bills nothing. Strictly-after keeps + // the earliest call on ties, matching billableBatchWindow. occurrence, ok := dispatched[i] if !ok || occurrence.ToolCallID != call.ToolCallID { continue } + start := occurrence.StartedAt end := occurrence.CompletedAt if end.IsZero() { + if start.IsZero() { + continue + } end = interruptedAt } + if start.IsZero() { + // Live execution always marks a start before a + // completion; fall back to the batch start rather than + // dropping a completed call's billed work. + start = billing.batchStartedAt + } + intervals = append(intervals, chatloop.BilledInterval{Start: start, End: end}) if end.After(windowEnd) { windowEnd = end windowRowIdx = len(result) - 1 } } - // Mirror the committed-batch policy: bill the partial window once, on - // the cancellation row of the billed tool call that defines it. - if windowRowIdx >= 0 && windowEnd.After(billing.batchStartedAt) { - result[windowRowIdx].RuntimeMs = nullInt64IfNonZero(windowEnd.Sub(billing.batchStartedAt).Milliseconds()) + // Mirror the committed-batch policy: bill the union of the billed + // calls' execution intervals once, on the cancellation row of the + // billed tool call whose window ends last. The union never charges + // a span where only unbilled tools were running, such as a + // sub-agent wait that delayed a serial call's launch. + if windowRowIdx >= 0 { + result[windowRowIdx].RuntimeMs = nullInt64IfNonZero(chatloop.BilledIntervalsDuration(intervals).Milliseconds()) } return result, nil } diff --git a/coderd/x/chatd/tasks_test.go b/coderd/x/chatd/tasks_test.go index c5f7583430f..18af7c49597 100644 --- a/coderd/x/chatd/tasks_test.go +++ b/coderd/x/chatd/tasks_test.go @@ -602,6 +602,10 @@ func TestInterruptTask_ToolBatchBillsPartialWindowOnCancellationRow(t *testing.T {CallIndex: 0, ToolCallID: execCallID}, {CallIndex: 1, ToolCallID: waitCallID}, })) + // Both concurrent calls launch with the batch, the way the launch + // loop's start callback records them. + require.NoError(t, buffer.RecordToolStart(batch.key, 0, execCallID, batch.clock.Now())) + require.NoError(t, buffer.RecordToolStart(batch.key, 1, waitCallID, batch.clock.Now())) // execute completes 3 seconds into the batch and records its // completion, the way the tool goroutine's completion callback // does. Its result is not published: results publish only after @@ -653,7 +657,7 @@ func TestInterruptTask_RetrySnapshotOutlivesEpisodeEviction(t *testing.T) { billing: messagepartbuffer.EpisodeBilling{ ClosedAt: batch.clock.Now(), ToolBatchStartedAt: batchStartedAt, - ToolCompletions: []messagepartbuffer.ToolCompletion{{CallIndex: 0, ToolCallID: execCallID}}, + ToolCompletions: []messagepartbuffer.ToolCompletion{{CallIndex: 0, ToolCallID: execCallID, StartedAt: batchStartedAt}}, }, } @@ -679,6 +683,7 @@ func TestInterruptTask_SnapshotCapturedBeforeChatRead(t *testing.T) { batch.clock.Advance(2 * time.Second) require.NoError(t, buffer.StartToolBatch(batch.key, []messagepartbuffer.DispatchedToolCall{{CallIndex: 0, ToolCallID: execCallID}})) + require.NoError(t, buffer.RecordToolStart(batch.key, 0, execCallID, batch.clock.Now())) batch.clock.Advance(3 * time.Second) // The first attempt's chat read fails: the chat is still running, @@ -727,6 +732,7 @@ func TestInterruptTask_RunningBilledToolBillsUpToInterrupt(t *testing.T) { batch.clock.Advance(2 * time.Second) require.NoError(t, batch.starter.opts.MessagePartBuffer.StartToolBatch(batch.key, []messagepartbuffer.DispatchedToolCall{{CallIndex: 0, ToolCallID: execCallID}})) + require.NoError(t, batch.starter.opts.MessagePartBuffer.RecordToolStart(batch.key, 0, execCallID, batch.clock.Now())) // The tool is still running when the interrupt lands 7 seconds in. batch.clock.Advance(7 * time.Second) @@ -748,6 +754,7 @@ func TestInterruptTask_UnbilledOnlyBatchBillsNothingOnInterrupt(t *testing.T) { batch.clock.Advance(2 * time.Second) require.NoError(t, batch.starter.opts.MessagePartBuffer.StartToolBatch(batch.key, []messagepartbuffer.DispatchedToolCall{{CallIndex: 0, ToolCallID: waitCallID}})) + require.NoError(t, batch.starter.opts.MessagePartBuffer.RecordToolStart(batch.key, 0, waitCallID, batch.clock.Now())) batch.clock.Advance(10 * time.Second) messages := batch.interrupt(t, f) @@ -755,6 +762,93 @@ func TestInterruptTask_UnbilledOnlyBatchBillsNothingOnInterrupt(t *testing.T) { require.False(t, waitRow.RuntimeMs.Valid) } +// A dispatched call that never began executing bills nothing on +// interrupt: a billed serial tool stays queued until every concurrent +// sibling settles, so an interrupt landing during an unbilled +// wait_agent must not treat the waiting serial call as running from +// the batch start. +func TestInterruptTask_UnstartedSerialCallBillsNothingOnInterrupt(t *testing.T) { + t.Parallel() + + f := newTaskTestFixture(t) + waitCallID := "call_" + uuid.NewString() + serialCallID := "call_" + uuid.NewString() + batch := interruptedBatchFixture(t, f, []codersdk.ChatMessagePart{ + {Type: codersdk.ChatMessagePartTypeToolCall, ToolCallID: waitCallID, ToolName: "wait_agent", Args: json.RawMessage(`{}`)}, + {Type: codersdk.ChatMessagePartTypeToolCall, ToolCallID: serialCallID, ToolName: "execute", Args: json.RawMessage(`{}`)}, + }) + buffer := batch.starter.opts.MessagePartBuffer + + batch.clock.Advance(2 * time.Second) + require.NoError(t, buffer.StartToolBatch(batch.key, []messagepartbuffer.DispatchedToolCall{ + {CallIndex: 0, ToolCallID: waitCallID}, + {CallIndex: 1, ToolCallID: serialCallID}, + })) + // Only the concurrent wait_agent launched; the serial call stays + // queued behind it, so no start mark arrives before the interrupt + // lands 10 seconds later. + require.NoError(t, buffer.RecordToolStart(batch.key, 0, waitCallID, batch.clock.Now())) + batch.clock.Advance(10 * time.Second) + + messages := batch.interrupt(t, f) + for _, msg := range messages { + if msg.Role == database.ChatMessageRoleTool { + require.False(t, msg.RuntimeMs.Valid, + "no cancellation row may bill: the billed call never began executing") + } + } +} + +// A billed serial call that launched bills from its own start: the +// interrupt window unions the early billed sibling's interval with the +// serial call's, so the span where only the unbilled wait_agent was +// running is never charged. +func TestInterruptTask_StartedSerialCallBillsFromItsOwnStart(t *testing.T) { + t.Parallel() + + f := newTaskTestFixture(t) + execCallID := "call_" + uuid.NewString() + waitCallID := "call_" + uuid.NewString() + serialCallID := "call_" + uuid.NewString() + batch := interruptedBatchFixture(t, f, []codersdk.ChatMessagePart{ + {Type: codersdk.ChatMessagePartTypeToolCall, ToolCallID: execCallID, ToolName: "execute", Args: json.RawMessage(`{}`)}, + {Type: codersdk.ChatMessagePartTypeToolCall, ToolCallID: waitCallID, ToolName: "wait_agent", Args: json.RawMessage(`{}`)}, + {Type: codersdk.ChatMessagePartTypeToolCall, ToolCallID: serialCallID, ToolName: "execute", Args: json.RawMessage(`{}`)}, + }) + buffer := batch.starter.opts.MessagePartBuffer + + batch.clock.Advance(2 * time.Second) + require.NoError(t, buffer.StartToolBatch(batch.key, []messagepartbuffer.DispatchedToolCall{ + {CallIndex: 0, ToolCallID: execCallID}, + {CallIndex: 1, ToolCallID: waitCallID}, + {CallIndex: 2, ToolCallID: serialCallID}, + })) + // The two concurrent calls launch with the batch. + require.NoError(t, buffer.RecordToolStart(batch.key, 0, execCallID, batch.clock.Now())) + require.NoError(t, buffer.RecordToolStart(batch.key, 1, waitCallID, batch.clock.Now())) + // execute completes 3 seconds in. + batch.clock.Advance(3 * time.Second) + require.NoError(t, buffer.RecordToolCompletion(batch.key, 0, execCallID, batch.clock.Now())) + // wait_agent completes at 6 seconds, letting the serial call + // launch; it is still running when the interrupt lands at 8 + // seconds. + batch.clock.Advance(3 * time.Second) + require.NoError(t, buffer.RecordToolCompletion(batch.key, 1, waitCallID, batch.clock.Now())) + require.NoError(t, buffer.RecordToolStart(batch.key, 2, serialCallID, batch.clock.Now())) + batch.clock.Advance(2 * time.Second) + + messages := batch.interrupt(t, f) + // The union bills execute's 3s and the serial call's 2s, not the + // 3s gap where only wait_agent ran, and lands once on the serial + // call's row: its window ends last. + serialRow := findToolResultMessage(t, messages, serialCallID) + require.Equal(t, sql.NullInt64{Int64: 5_000, Valid: true}, serialRow.RuntimeMs) + execRow := findToolResultMessage(t, messages, execCallID) + require.False(t, execRow.RuntimeMs.Valid) + waitRow := findToolResultMessage(t, messages, waitCallID) + require.False(t, waitRow.RuntimeMs.Valid) +} + // Two dispatched billed calls sharing one tool call ID keep distinct // occurrence states: one completing early must not make the other look // finished, so the interrupted batch still bills through to the @@ -775,6 +869,8 @@ func TestInterruptTask_DuplicateCallIDsKeepOccurrenceStates(t *testing.T) { {CallIndex: 0, ToolCallID: dupCallID}, {CallIndex: 1, ToolCallID: dupCallID}, })) + require.NoError(t, buffer.RecordToolStart(batch.key, 0, dupCallID, batch.clock.Now())) + require.NoError(t, buffer.RecordToolStart(batch.key, 1, dupCallID, batch.clock.Now())) // One occurrence completes 3 seconds in; the other keeps running // until the interrupt lands 3 seconds later, defining the window. batch.clock.Advance(3 * time.Second) @@ -811,6 +907,8 @@ func TestInterruptTask_DuplicateCallIDCompletionStampsOwnOccurrence(t *testing.T {CallIndex: 0, ToolCallID: dupCallID}, {CallIndex: 1, ToolCallID: dupCallID}, })) + require.NoError(t, buffer.RecordToolStart(batch.key, 0, dupCallID, batch.clock.Now())) + require.NoError(t, buffer.RecordToolStart(batch.key, 1, dupCallID, batch.clock.Now())) // The unbilled wait_agent occurrence (index 1) completes 3 seconds // in; the billed execute occurrence (index 0) keeps running until // the interrupt 3 seconds later. @@ -848,6 +946,7 @@ func TestInterruptTask_RejectedDuplicateIDDoesNotStealDispatchedOccurrence(t *te // dispatched: the execute occurrence sharing its ID was rejected // as malformed before execution. require.NoError(t, batch.starter.opts.MessagePartBuffer.StartToolBatch(batch.key, []messagepartbuffer.DispatchedToolCall{{CallIndex: 1, ToolCallID: dupCallID}})) + require.NoError(t, batch.starter.opts.MessagePartBuffer.RecordToolStart(batch.key, 0, dupCallID, batch.clock.Now())) batch.clock.Advance(10 * time.Second) messages := batch.interrupt(t, f) @@ -879,6 +978,7 @@ func TestInterruptTask_RejectedCallBillsNothingOnInterrupt(t *testing.T) { // execution and never ran, so only the occurrence at wait_agent's // unresolved position (1) is seeded. require.NoError(t, batch.starter.opts.MessagePartBuffer.StartToolBatch(batch.key, []messagepartbuffer.DispatchedToolCall{{CallIndex: 1, ToolCallID: waitCallID}})) + require.NoError(t, batch.starter.opts.MessagePartBuffer.RecordToolStart(batch.key, 0, waitCallID, batch.clock.Now())) batch.clock.Advance(10 * time.Second) messages := batch.interrupt(t, f) diff --git a/docs/ai-coder/usage-data-reporting.md b/docs/ai-coder/usage-data-reporting.md index cfcaf116cd6..cffe7194bbb 100644 --- a/docs/ai-coder/usage-data-reporting.md +++ b/docs/ai-coder/usage-data-reporting.md @@ -61,8 +61,9 @@ Example of a failed request (e.g. Tallyman Server is blocked by your network): Total Coder Agent runtime is summed from per-message runtime (`runtime_ms` on chat messages). An assistant message's runtime is the wall-clock duration of the model invocation that produced its content, measured from just before the request to the model provider opens until the response is fully consumed. -A tool message's runtime is the wall-clock duration of the local tool batch that produced it, measured from the start of the batch until the last counted tool finishes. +A tool message's runtime is the wall-clock duration of the local tool batch that produced it: the combined time during which at least one counted tool was executing. Tools in a batch run in parallel, so each batch records one window on one tool message rather than a per-tool sum. +A tool that runs serially after the rest of the batch is measured from its own start, so time when only excluded tools were running is never counted. What counts: From da26bcdec0f128a6f9ac0d441900b8ef621639d5 Mon Sep 17 00:00:00 2001 From: Jaayden Halko Date: Wed, 19 Aug 2026 09:32:49 +0000 Subject: [PATCH 15/20] docs(docs/ai-coder): revert usage-data-reporting changes Restore docs/ai-coder/usage-data-reporting.md to its state on main, removing this branch's edits to the agent runtime measurement section. --- docs/ai-coder/usage-data-reporting.md | 25 ++++++++++++++----------- 1 file changed, 14 insertions(+), 11 deletions(-) diff --git a/docs/ai-coder/usage-data-reporting.md b/docs/ai-coder/usage-data-reporting.md index cffe7194bbb..3c69e27ae0b 100644 --- a/docs/ai-coder/usage-data-reporting.md +++ b/docs/ai-coder/usage-data-reporting.md @@ -58,24 +58,27 @@ Example of a failed request (e.g. Tallyman Server is blocked by your network): ## Agent runtime measurement -Total Coder Agent runtime is summed from per-message runtime (`runtime_ms` on chat messages). +Total Coder Agent runtime is summed from per-message generation time +(`runtime_ms` on chat messages). -An assistant message's runtime is the wall-clock duration of the model invocation that produced its content, measured from just before the request to the model provider opens until the response is fully consumed. -A tool message's runtime is the wall-clock duration of the local tool batch that produced it: the combined time during which at least one counted tool was executing. -Tools in a batch run in parallel, so each batch records one window on one tool message rather than a per-tool sum. -A tool that runs serially after the rest of the batch is measured from its own start, so time when only excluded tools were running is never counted. +A message's runtime is the wall-clock duration of the model invocation that +produced its content, measured from just before the request to the model +provider opens until the response is fully consumed. What counts: - Assistant generation steps, in both top-level chats and sub-agent chats. - Context compaction (summarization) model calls. -- Local tool execution: file, terminal, and process tools, workspace lifecycle operations, MCP tools, and other server-executed tools. A tool's work may run in a connected workspace, but the server dispatches and waits on it during the turn, so it counts regardless of how the workspace's agent is managed. -- Interrupted generation: the time streamed or spent executing tools before the interrupt is kept on the partial messages. +- Interrupted generation: the time streamed before the interrupt is kept on + the partial assistant message. What does not count: -- Sub-agent orchestration tools, such as spawning and waiting on sub-agents. A sub-agent is its own chat and records its own runtime, so counting the parent's wait would double count. Waiting on a sub-agent never extends a tool batch's window, even when other tools in the batch do count. -- Client-executed (dynamic) tools and chats handed off to external coding agents: the chat parks while that work happens outside the server, which cannot measure it. +- Local tool execution, including waiting on sub-agents. A sub-agent is its + own chat and records its own model invocations, so counting the parent's + wait would double count. - Idle time: chats waiting for user input or external tool results. -- Failed model calls whose output was discarded, and the backoff between retried attempts. Retried and errored attempts persist no content, so they record no runtime. -- Ancillary model calls that produce no chat messages, such as title generation. +- Failed model calls whose output was discarded. Retried and errored + attempts persist no content, so they record no runtime. +- Ancillary model calls that produce no chat messages, such as title + generation. From 38715616797aacefcc65f543dde42ff286d9f441 Mon Sep 17 00:00:00 2001 From: Jaayden Halko Date: Wed, 19 Aug 2026 09:42:01 +0000 Subject: [PATCH 16/20] style(coderd): trim agent runtime comments --- coderd/database/querier_test.go | 5 +- coderd/usage/usagetypes/events.go | 26 ++- coderd/x/chatd/ARCHITECTURE.md | 6 +- coderd/x/chatd/attempt.go | 6 +- coderd/x/chatd/chatd_test.go | 3 - coderd/x/chatd/chatloop/chatloop.go | 150 +++++------------- coderd/x/chatd/chatloop/runtime_test.go | 85 +--------- coderd/x/chatd/generation.go | 42 +---- coderd/x/chatd/message_conversion.go | 24 +-- coderd/x/chatd/message_conversion_test.go | 13 -- .../messagepartbuffer/message_part_buffer.go | 140 +++++----------- .../message_part_buffer_test.go | 53 +------ coderd/x/chatd/options.go | 6 +- coderd/x/chatd/runner.go | 3 - coderd/x/chatd/subagent_catalog.go | 11 +- coderd/x/chatd/subagent_internal_test.go | 3 - coderd/x/chatd/tasks.go | 107 +++---------- coderd/x/chatd/tasks_test.go | 113 +------------ coderd/x/chatd/toolinput.go | 6 +- 19 files changed, 145 insertions(+), 657 deletions(-) diff --git a/coderd/database/querier_test.go b/coderd/database/querier_test.go index 20dfe4bab03..092d3aeeee9 100644 --- a/coderd/database/querier_test.go +++ b/coderd/database/querier_test.go @@ -11123,13 +11123,12 @@ func TestGetTotalChatMessageRuntimeMsInRange(t *testing.T) { } // Counted: on the inclusive start boundary, in the middle (across two - // chats), soft-deleted, just before the exclusive end boundary, and a - // tool-role row carrying a local tool batch window (the sum is - // role-agnostic). + // chats), soft-deleted, and just before the exclusive end boundary. 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). diff --git a/coderd/usage/usagetypes/events.go b/coderd/usage/usagetypes/events.go index 3bb24f033d2..a990b4c6f7e 100644 --- a/coderd/usage/usagetypes/events.go +++ b/coderd/usage/usagetypes/events.go @@ -203,23 +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. Two kinds of windows are measured. Model steps -// span model streaming (including provider-executed tools) and stream -// retries, ending when the model stream finishes. Local tool batches bill -// the union of the billed tools' execution windows: tools in a batch run -// in parallel, so the batch bills one window rather than a sum, and a -// tool that runs serially after its siblings is measured from its own -// start. -// Sub-agent orchestration tools (spawn_agent, wait_agent, and the rest of -// that category) are excluded because each sub-agent chat bills its own -// runtime. Also excluded: client-executed (dynamic) tools and chats handed -// off to external coding agents (the chat parks while that work happens -// outside the server, which cannot measure it), time parked waiting for -// user action, idle time between turns, and retry backoff between -// attempts. Server-executed tools count even when their work runs in a -// connected workspace, whatever the workspace's agent type: the server -// dispatches and waits on them during the turn. +// 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/ARCHITECTURE.md b/coderd/x/chatd/ARCHITECTURE.md index be9cc88a47d..033902a9821 100644 --- a/coderd/x/chatd/ARCHITECTURE.md +++ b/coderd/x/chatd/ARCHITECTURE.md @@ -725,7 +725,8 @@ The buffer exposes the following API: - `GetParts(chat_id, history_version, generation_attempt)`: returns the message parts for an episode. Returns a predefined error if the episode is not found. - `StartModelInvocation(chat_id, history_version, generation_attempt)`: stamps the instant the episode opens its provider stream. Returns a predefined error if the episode is not found or already closed. Episodes that never invoke a model, such as local tool execution batches, are never stamped. - `ModelInvokedAt(chat_id, history_version, generation_attempt)`: returns the instant stamped by `StartModelInvocation`, or the zero time when the episode is unknown or never opened a provider stream. It must be read before `CloseEpisode`, because closed episodes are garbage collected and reading afterwards races the cleanup loop. The interrupt goroutine reads it just before closing the episode and uses the span between that instant and the interrupt as the interrupted attempt's billable runtime. -- TODO(CODAGT-928): document the tool-batch billing methods `StartToolBatch`, `ToolBatchStartedAt`, `RecordToolStart`, `RecordToolCompletion`, and `ToolCompletions`, the local-tool counterparts of `StartModelInvocation`/`ModelInvokedAt`, and `CloseEpisodeForBilling`, which the interrupt task uses to close the episode and snapshot its billing stamps (including the stable first-close instant used as the interrupt instant) in one atomic step; repeat calls return the identical snapshot and push the eviction deadline out, so retried interrupt tasks bill the same window and keep their snapshot through a database outage. `RecordToolStart` marks when a dispatched call begins executing: serial tool calls launch only after every concurrent sibling settles, so the interrupt bills each call from its own start and bills nothing for a dispatched call that never launched. +- TODO(CODAGT-928): document tool-batch billing APIs, including serial + starts, duplicate IDs, atomic snapshots, stable close time, and retries. - `SubscribeToEpisode(chat_id, history_version, generation_attempt)`: returns a go channel that will receive all message parts for the episode. It spawns a goroutine that delivers parts to the channel. It's live until the episode is closed or until a subscriber requests that the channel be closed. Once the goroutine delivers all message parts for a closed episode, it closes the channel and exits. If the episode is already closed at the time of the call, the goroutine delivers all message parts for the episode, closes the channel, and exits. `SubscribeToEpisode` does not return an error if the episode is not found: it waits for it to be created instead. Closed episodes are garbage collected after at least 15 seconds since they were closed and when they have no active subscribers. The message part buffer maintains a garbage collection goroutine. @@ -936,7 +937,8 @@ The goroutine does the following in order: 3. It reads the buffered parts for that episode by calling the `GetParts` method on the message part buffer. 4. It applies the `FinishInterruption(partial?)` transition on the core state machine. If there are no buffered parts for that episode, or the episode is not found, it passes `nil` as the `partial` argument. -TODO(CODAGT-928): update steps 2 and 3 above. The interrupt goroutine now closes the episode and snapshots its billing stamps in one atomic step via `CloseEpisodeForBilling`, uses the episode's first-close instant as the interrupt instant, captures the snapshot and buffered parts before its first database read and carries them across task retry attempts so buffer eviction during a stalled read cannot lose them, and synthesizes tool cancellation rows whose `runtime_ms` bills the interrupted tool batch's partial window. +TODO(CODAGT-928): update interrupt steps for atomic snapshots, pre-read +retry carry, stable close time, and partial tool-batch billing. #### Dynamic tools timeout goroutine diff --git a/coderd/x/chatd/attempt.go b/coderd/x/chatd/attempt.go index 28de29971c1..7804a25c0d7 100644 --- a/coderd/x/chatd/attempt.go +++ b/coderd/x/chatd/attempt.go @@ -31,10 +31,8 @@ type stepData struct { ContextLimit sql.NullInt64 Runtime time.Duration - // BatchRuntime is the billable window of a local tool batch, - // persisted as runtime_ms on the tool message row identified by - // BatchRuntimeToolCallID. Zero for model-invocation steps, whose - // billable window is Runtime on the assistant row instead. + // BatchRuntime is the local-tool batch window persisted on the + // BatchRuntimeToolCallID tool row. Model steps use Runtime instead. BatchRuntime time.Duration BatchRuntimeToolCallID string diff --git a/coderd/x/chatd/chatd_test.go b/coderd/x/chatd/chatd_test.go index 04d915cb6f0..8b108ae766e 100644 --- a/coderd/x/chatd/chatd_test.go +++ b/coderd/x/chatd/chatd_test.go @@ -6438,9 +6438,6 @@ func TestActiveServer_ToolExecutionAndPolicy(t *testing.T) { } } - // The parallel batch bills at most one window: whatever wall - // time elapsed, only the window-defining tool row may carry - // runtime_ms, never one per parallel call. messages := chatMessages(ctx, t, db, chat.ID) billedToolRows := 0 for _, msg := range messages { diff --git a/coderd/x/chatd/chatloop/chatloop.go b/coderd/x/chatd/chatloop/chatloop.go index 39dd3b3e91b..48a67474418 100644 --- a/coderd/x/chatd/chatloop/chatloop.go +++ b/coderd/x/chatd/chatloop/chatloop.go @@ -73,11 +73,8 @@ 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. Local tool batches report their billable window - // through ToolExecutionOutcome.BatchRuntime instead. + // Runtime is the wall-clock duration from opening to consuming the + // model stream. Local tool batches use ToolExecutionOutcome.BatchRuntime. Runtime time.Duration // PendingDynamicToolCalls lists tool calls that target // dynamic tools. When non-empty the chatloop exits with @@ -266,44 +263,21 @@ type ExecuteLocalToolsOptions struct { // is renamed but old chat histories still reference the old name. ToolNameAliases map[string]string - // UnbilledToolNames lists tool names whose execution never extends - // the batch's billable runtime window. Classification uses the - // name as called, so deprecated aliases must be listed alongside - // their canonical names. + // UnbilledToolNames lists called tool names excluded from the batch + // window. Include deprecated aliases. UnbilledToolNames map[string]bool - // OnBatchStart, when set, is invoked once cancellation and - // exclusive-policy checks have passed, immediately before the - // batch's tool goroutines launch. Dispatch is unconditional after - // it fires. The interrupt path uses it to seed the batch's billing - // state on the buffer episode: seeding any earlier would let an - // interrupt whose cancellation lands before dispatch bill calls - // that never ran as if they were still running. + // OnBatchStart fires immediately before dispatch. Interrupt billing uses + // it to avoid charging pre-dispatch cancellations. OnBatchStart func() - // OnToolStart, when set, is invoked with each local tool call's - // occurrence index, ID, and the instant the call begins executing. - // Concurrent calls start together when the batch launches, but - // calls whose tools opt in to SerialToolCalls launch only after - // every concurrent sibling has settled, so their start marks - // arrive later. The interrupt path uses the marks to bill each - // interrupted call from its actual start rather than the batch - // start, and to avoid billing a dispatched serial call that never - // began executing. Serial calls invoke it from the batch goroutine - // after concurrent siblings finished, so implementations must be - // safe for use from a different goroutine than OnBatchStart. + // OnToolStart fires when each local call begins. Serial calls may start + // after concurrent siblings settle, so interrupts bill actual starts and + // skip dispatched calls that never run. callIndex identifies the + // dispatch-order occurrence. OnToolStart func(callIndex int, toolCallID string, startedAt time.Time) - // OnToolComplete, when set, is invoked with each local tool call's - // completion instant as the tool finishes, the same instant - // PersistedStep.ToolResultCreatedAt later carries. Tool results are - // published only after the whole batch completes, to keep event - // ordering deterministic, so this callback is the only live signal - // that a tool already finished while siblings are still running; - // the interrupt path uses it to bill an interrupted batch's - // partial window. callIndex is the call's position among the - // batch's local (non provider-executed) calls in dispatch order; - // it identifies the exact occurrence because duplicate tool call - // IDs, which reach execution when lifecycle hooks are disabled, - // make the ID alone ambiguous. It is called concurrently from - // tool goroutines and must be safe for concurrent use. + // OnToolComplete fires concurrently as each local call finishes, before + // ordered results publish. Interrupt billing uses the live timestamp; + // callIndex identifies the dispatch-order occurrence when IDs collide. + // The callback must be concurrency-safe. OnToolComplete func(callIndex int, toolCallID string, completedAt time.Time) PublishMessagePart func(codersdk.ChatMessageRole, codersdk.ChatMessagePart) @@ -315,17 +289,12 @@ type ExecuteLocalToolsOptions struct { // ToolExecutionOutcome is the durable tool-result content from one batch. type ToolExecutionOutcome struct { Step PersistedStep - // BatchRuntime is the billable local-tool counterpart of - // PersistedStep.Runtime: the wall-clock window from just before the - // batch's tools start until the last billed tool completes. All - // calls in a batch start together, so this equals the union of the - // billed tools' execution intervals; parallel calls are never - // summed. Zero when no billed tool produced a result. + // BatchRuntime is the union of billed tool execution intervals. Parallel + // calls count once and serial calls count only from their own start. Zero + // means no billed tool produced a result. BatchRuntime time.Duration - // BatchRuntimeToolCallID identifies the billed tool call whose - // completion ends the batch window, with ties broken by call - // order. The persistence layer stores BatchRuntime on that call's - // tool message row. Empty when BatchRuntime is zero. + // BatchRuntimeToolCallID is the ID on the billed interval ending last. + // Ties use call order; the ID can be empty or non-unique. BatchRuntimeToolCallID string } @@ -633,8 +602,6 @@ func ExecuteLocalTools(ctx context.Context, opts ExecuteLocalToolsOptions) (Tool for i, tr := range policyResults { recordToolResultTimestamp(&result, tr.ToolCallID, now) if opts.OnToolComplete != nil { - // Policy results are index-aligned with localCalls, so - // i is the call's occurrence index. opts.OnToolComplete(i, tr.ToolCallID, now) } publishToolAttachments(ctx, opts.Logger, tr, now, publishMessagePart) @@ -657,17 +624,9 @@ func ExecuteLocalTools(ctx context.Context, opts ExecuteLocalToolsOptions) (Tool opts.OnBatchStart() } batchStart := clockNow(opts.Clock) - // Completion instants aligned with localCalls by occurrence. - // billableBatchWindow reads these instead of the ID-keyed - // ToolResultCreatedAt map, where duplicate tool call IDs, which - // reach execution when lifecycle hooks are disabled, would - // overwrite each other and corrupt the window. + // Keep completions by occurrence so duplicate IDs cannot collapse them. orderedCompletions := make([]time.Time, 0, len(localCalls)) - // Start instants aligned with localCalls by occurrence. Concurrent - // calls start at batchStart, but serial calls launch only after - // every concurrent sibling settles, so billing them from - // batchStart would charge the whole concurrent phase, including - // waits on unbilled tools. + // Keep starts by occurrence. Serial calls may begin after unbilled waits. orderedStarts := make([]time.Time, len(localCalls)) onToolStart := func(callIndex int, toolCallID string, startedAt time.Time) { orderedStarts[callIndex] = startedAt @@ -695,9 +654,6 @@ func ExecuteLocalTools(ctx context.Context, opts ExecuteLocalToolsOptions) (Tool onToolStart, opts.OnToolComplete, func(tr fantasy.ToolResultContent, completedAt time.Time) { - // onResult fires once per local call in call order, so - // appending keeps orderedCompletions aligned with - // localCalls even when tool call IDs collide. orderedCompletions = append(orderedCompletions, completedAt) recordToolResultTimestamp(&result, tr.ToolCallID, completedAt) publishToolAttachments(ctx, opts.Logger, tr, completedAt, publishMessagePart) @@ -729,22 +685,14 @@ func ExecuteLocalTools(ctx context.Context, opts ExecuteLocalToolsOptions) (Tool }, nil } -// billableBatchWindow computes one local tool batch's billable runtime: -// the union of the billed tools' execution intervals. Concurrent calls -// all start together at batchStart, so parallel calls are billed once -// rather than summed, and unbilled tools (for example sub-agent -// orchestration) never extend the window even when they run longest. -// Calls whose tools opt in to SerialToolCalls launch only after every -// concurrent sibling settles, so each is measured from its own instant -// in starts, and a span where only unbilled tools were running never -// bills. starts and completions are aligned with toolCalls by -// occurrence, so duplicate tool call IDs each keep their own window; a -// zero start means the call launched with the batch. Returns the ID of -// the tool call whose completion ends the window, with ties broken by -// call order; a zero duration means no billed tool completed or the -// window rounds to nothing. The ID can be empty even when the window -// exists, because providers can emit calls without IDs and those still -// execute and bill. +// billableBatchWindow returns the union of billed execution intervals and +// the call ID whose interval ends last. Concurrent calls start at +// batchStart; serial calls use their recorded starts. Unbilled tools and +// gaps between billed intervals do not count. +// +// Starts and completions align with toolCalls by occurrence, so duplicate or +// empty IDs stay distinct. A zero start means the call used batchStart. Ties +// keep the earliest call. func billableBatchWindow( batchStart time.Time, toolCalls []fantasy.ToolCallContent, @@ -774,7 +722,6 @@ func billableBatchWindow( start = starts[i] } intervals = append(intervals, BilledInterval{Start: start, End: end}) - // Strictly-after keeps the earliest call on ties. if end.After(windowEnd) { found = true windowEnd = end @@ -797,14 +744,9 @@ type BilledInterval struct { End time.Time } -// BilledIntervalsDuration returns the total length of the union of the -// given intervals. Overlapping windows, such as concurrent tool calls -// sharing a batch start, bill once rather than summing, while gaps, -// such as a span where only unbilled tools were running before a -// serial call launched, bill nothing. Intervals whose End precedes -// their Start are ignored. The interrupt path shares this helper so a -// canceled batch bills the same window the committed step would have -// reported. +// 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 := make([]BilledInterval, 0, len(intervals)) for _, iv := range intervals { @@ -1267,12 +1209,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. onComplete, in contrast, fires -// from each tool's goroutine the instant that tool finishes, so -// callers can observe completions while slower siblings still run. +// 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, @@ -1379,12 +1318,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 { @@ -1392,8 +1327,6 @@ func executeTools( serialIndexes = append(serialIndexes, i) continue } - // Concurrent calls all start executing now, at the batch - // start already stamped by the caller. if onStart != nil { onStart(i, tc.ToolCallID, batchStart) } @@ -1405,9 +1338,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) { @@ -1417,10 +1348,7 @@ func executeTools( notifyStepToolResultObservers(toolMap, toolNameAliases, localToolCalls, observed, settled) for _, i := range serialIndexes { - // A serial call starts executing only here, after every - // concurrent sibling settled, so its start mark is stamped at - // launch rather than at batchStart: billing it from - // batchStart would charge the whole concurrent phase. + // Stamp serial calls at launch, not batch start. if onStart != nil { onStart(i, localToolCalls[i].ToolCallID, clockNow(clock)) } diff --git a/coderd/x/chatd/chatloop/runtime_test.go b/coderd/x/chatd/chatloop/runtime_test.go index 0dd9f5ded24..b542dd69c67 100644 --- a/coderd/x/chatd/chatloop/runtime_test.go +++ b/coderd/x/chatd/chatloop/runtime_test.go @@ -162,10 +162,8 @@ func TestGenerateCompaction_RecordsRuntime(t *testing.T) { require.Equal(t, result.Runtime, clock.Since(startedAt[0])) } -// executeToolBatch runs ExecuteLocalTools in a goroutine. Callers trap the -// mock clock's Now: the first trapped call is the batch start and each -// subsequent one is a tool completion, released one tool at a time so -// parallel tool goroutines never race the clock advances. +// executeToolBatch lets tests release trapped clock events in order, so +// goroutines cannot race clock advances. func executeToolBatch( t *testing.T, clock *quartz.Mock, @@ -182,8 +180,6 @@ func executeToolBatch( return resultCh } -// blockingTool returns a tool that parks until release is closed, so the -// test controls exactly when its completion timestamp is recorded. func blockingTool(name string, release <-chan struct{}, response fantasy.ToolResponse) fantasy.AgentTool { return fantasy.NewAgentTool( name, @@ -195,18 +191,12 @@ func blockingTool(name string, release <-chan struct{}, response fantasy.ToolRes ) } -// serialTool wraps a tool so its calls run serially after every -// concurrent sibling settles, matching tools that opt in via -// SerialToolCalls. type serialTool struct { fantasy.AgentTool } func (serialTool) SerialToolCalls() bool { return true } -// Parallel billed tools bill one shared window ending at the slowest -// tool's completion, never the sum of their durations. The slower tool -// returns an error result: errored tools bill their wall clock too. func TestExecuteLocalTools_BatchWindowIsMaxNotSum(t *testing.T) { t.Parallel() @@ -229,13 +219,10 @@ func TestExecuteLocalTools_BatchWindowIsMaxNotSum(t *testing.T) { }, }) - // Batch start. trap.MustWait(ctx).MustRelease(ctx) - // The fast tool completes 10 seconds in. clock.Advance(10 * time.Second) close(fastGo) trap.MustWait(ctx).MustRelease(ctx) - // The slow tool errors out at 60 seconds. clock.Advance(50 * time.Second) close(slowGo) trap.MustWait(ctx).MustRelease(ctx) @@ -245,8 +232,6 @@ func TestExecuteLocalTools_BatchWindowIsMaxNotSum(t *testing.T) { require.Equal(t, "call-slow", outcome.BatchRuntimeToolCallID) } -// Simultaneous completions tie-break to the earliest call in call order, -// and N parallel calls of the same duration bill that duration once. func TestExecuteLocalTools_SimultaneousCompletionsBillOnceByCallOrder(t *testing.T) { t.Parallel() @@ -268,9 +253,7 @@ func TestExecuteLocalTools_SimultaneousCompletionsBillOnceByCallOrder(t *testing }, }) - // Batch start. trap.MustWait(ctx).MustRelease(ctx) - // All three calls complete together 10 seconds in. clock.Advance(10 * time.Second) close(release) for range 3 { @@ -282,8 +265,6 @@ func TestExecuteLocalTools_SimultaneousCompletionsBillOnceByCallOrder(t *testing require.Equal(t, "call-1", outcome.BatchRuntimeToolCallID) } -// An unbilled tool never extends the window, even when it runs longest: -// the batch bills up to the last billed tool's completion. func TestExecuteLocalTools_UnbilledToolNeverExtendsWindow(t *testing.T) { t.Parallel() @@ -307,13 +288,10 @@ func TestExecuteLocalTools_UnbilledToolNeverExtendsWindow(t *testing.T) { }, }) - // Batch start. trap.MustWait(ctx).MustRelease(ctx) - // execute completes 10 seconds in. clock.Advance(10 * time.Second) close(executeGo) trap.MustWait(ctx).MustRelease(ctx) - // wait_agent keeps blocking on its child until 60 seconds. clock.Advance(50 * time.Second) close(waitGo) trap.MustWait(ctx).MustRelease(ctx) @@ -323,7 +301,6 @@ func TestExecuteLocalTools_UnbilledToolNeverExtendsWindow(t *testing.T) { require.Equal(t, "call-execute", outcome.BatchRuntimeToolCallID) } -// A batch of only unbilled tools bills nothing. func TestExecuteLocalTools_UnbilledOnlyBatchBillsNothing(t *testing.T) { t.Parallel() @@ -344,7 +321,6 @@ func TestExecuteLocalTools_UnbilledOnlyBatchBillsNothing(t *testing.T) { }, }) - // Batch start. trap.MustWait(ctx).MustRelease(ctx) clock.Advance(60 * time.Second) close(waitGo) @@ -355,9 +331,6 @@ func TestExecuteLocalTools_UnbilledOnlyBatchBillsNothing(t *testing.T) { require.Empty(t, outcome.BatchRuntimeToolCallID) } -// Billing classifies on the name as called: a deprecated alias listed in -// UnbilledToolNames stays unbilled even though dispatch resolves it to -// its canonical tool through ToolNameAliases. func TestExecuteLocalTools_AliasNamesClassifyAsCalled(t *testing.T) { t.Parallel() @@ -385,13 +358,10 @@ func TestExecuteLocalTools_AliasNamesClassifyAsCalled(t *testing.T) { }, }) - // Batch start. trap.MustWait(ctx).MustRelease(ctx) - // execute completes 10 seconds in. clock.Advance(10 * time.Second) close(executeGo) trap.MustWait(ctx).MustRelease(ctx) - // The aliased call completes at 60 seconds and must not bill. clock.Advance(50 * time.Second) close(legacyGo) trap.MustWait(ctx).MustRelease(ctx) @@ -401,10 +371,6 @@ func TestExecuteLocalTools_AliasNamesClassifyAsCalled(t *testing.T) { require.Equal(t, "call-execute", outcome.BatchRuntimeToolCallID) } -// Duplicate tool call IDs, which reach execution when lifecycle hooks -// are disabled, must not corrupt the window: completions are tracked per -// occurrence, so a later short duplicate cannot overwrite an earlier -// long one and shrink the bill. func TestExecuteLocalTools_DuplicateToolCallIDsKeepOccurrenceCompletions(t *testing.T) { t.Parallel() @@ -421,22 +387,16 @@ func TestExecuteLocalTools_DuplicateToolCallIDsKeepOccurrenceCompletions(t *test blockingTool("fast_tool", fastGo, fantasy.NewTextResponse("done")), }, ActiveTools: []string{"slow_tool", "fast_tool"}, - // Both calls share one ID: an ID-keyed completion map would - // let the fast occurrence overwrite the slow one. ToolCalls: []fantasy.ToolCallContent{ {ToolCallID: "call-dup", ToolName: "slow_tool", Input: "{}"}, {ToolCallID: "call-dup", ToolName: "fast_tool", Input: "{}"}, }, }) - // Batch start. trap.MustWait(ctx).MustRelease(ctx) - // The fast occurrence completes at 10 seconds. clock.Advance(10 * time.Second) close(fastGo) trap.MustWait(ctx).MustRelease(ctx) - // The slow occurrence completes at 60 seconds and must define the - // window even though the fast occurrence shares its ID. clock.Advance(50 * time.Second) close(slowGo) trap.MustWait(ctx).MustRelease(ctx) @@ -446,10 +406,6 @@ func TestExecuteLocalTools_DuplicateToolCallIDsKeepOccurrenceCompletions(t *test require.Equal(t, "call-dup", outcome.BatchRuntimeToolCallID) } -// OnBatchStart seeds billing state only once dispatch is guaranteed: it -// fires before the tool goroutines launch, and never fires when -// cancellation or the exclusive policy stops the batch, so an interrupt -// racing those paths finds no batch to bill. func TestExecuteLocalTools_OnBatchStartFiresOnlyOnDispatch(t *testing.T) { t.Parallel() @@ -517,9 +473,6 @@ func TestExecuteLocalTools_OnBatchStartFiresOnlyOnDispatch(t *testing.T) { }) } -// A call without a tool call ID, which providers can emit, still -// executes: its completion defines the window like any other billed -// call instead of being discarded as a missing result. func TestExecuteLocalTools_EmptyToolCallIDStillBillsWindow(t *testing.T) { t.Parallel() @@ -542,13 +495,10 @@ func TestExecuteLocalTools_EmptyToolCallIDStillBillsWindow(t *testing.T) { }, }) - // Batch start. trap.MustWait(ctx).MustRelease(ctx) - // The identified call completes at 10 seconds. clock.Advance(10 * time.Second) close(fastGo) trap.MustWait(ctx).MustRelease(ctx) - // The ID-less call completes at 60 seconds and defines the window. clock.Advance(50 * time.Second) close(idlessGo) trap.MustWait(ctx).MustRelease(ctx) @@ -558,12 +508,6 @@ func TestExecuteLocalTools_EmptyToolCallIDStillBillsWindow(t *testing.T) { require.Empty(t, outcome.BatchRuntimeToolCallID) } -// OnToolComplete reports each tool's completion the instant it finishes, -// while slower siblings are still running, with the same instants the -// outcome's ToolResultCreatedAt later carries. The interrupt path -// depends on this live signal: results publish only after the whole -// batch finishes, so without it an interrupt could not tell finished -// tools from still-running ones. func TestExecuteLocalTools_OnToolCompleteReportsLiveCompletions(t *testing.T) { t.Parallel() @@ -595,17 +539,13 @@ func TestExecuteLocalTools_OnToolCompleteReportsLiveCompletions(t *testing.T) { }, }) - // Batch start. trap.MustWait(ctx).MustRelease(ctx) - // The fast tool completes 10 seconds in. Its completion arrives - // while the slow tool is still parked on its release channel. clock.Advance(10 * time.Second) close(fastGo) trap.MustWait(ctx).MustRelease(ctx) fast := testutil.RequireReceive(ctx, t, completionCh) require.Equal(t, "call-fast", fast.toolCallID) require.Equal(t, 0, fast.callIndex, "callIndex is the call's position in dispatch order") - // The slow tool completes at 60 seconds. clock.Advance(50 * time.Second) close(slowGo) trap.MustWait(ctx).MustRelease(ctx) @@ -621,10 +561,6 @@ func TestExecuteLocalTools_OnToolCompleteReportsLiveCompletions(t *testing.T) { }, outcome.Step.ToolResultCreatedAt) } -// A billed serial tool queued behind a long-running unbilled tool bills -// only its own execution: it launches only after every concurrent -// sibling settles, so measuring it from the batch start would charge -// the whole unbilled wait that delayed its launch. func TestExecuteLocalTools_SerialCallBillsFromItsOwnStart(t *testing.T) { t.Parallel() @@ -648,15 +584,12 @@ func TestExecuteLocalTools_SerialCallBillsFromItsOwnStart(t *testing.T) { }, }) - // Batch start. trap.MustWait(ctx).MustRelease(ctx) - // The unbilled wait completes 10 minutes in. clock.Advance(10 * time.Minute) close(waitGo) trap.MustWait(ctx).MustRelease(ctx) - // The serial call launches only now and stamps its start. + // Release the serial start timestamp. trap.MustWait(ctx).MustRelease(ctx) - // It completes 2 seconds later. clock.Advance(2 * time.Second) close(serialGo) trap.MustWait(ctx).MustRelease(ctx) @@ -667,10 +600,6 @@ func TestExecuteLocalTools_SerialCallBillsFromItsOwnStart(t *testing.T) { require.Equal(t, "call-serial", outcome.BatchRuntimeToolCallID) } -// A batch mixing billed concurrent and billed serial calls bills the -// union of their execution intervals: the concurrent window and the -// serial call's own window sum, while the gap where only the unbilled -// tool was running is never charged. func TestExecuteLocalTools_SerialAfterBilledSiblingBillsUnion(t *testing.T) { t.Parallel() @@ -697,19 +626,15 @@ func TestExecuteLocalTools_SerialAfterBilledSiblingBillsUnion(t *testing.T) { }, }) - // Batch start. trap.MustWait(ctx).MustRelease(ctx) - // execute completes 3 seconds in. clock.Advance(3 * time.Second) close(execGo) trap.MustWait(ctx).MustRelease(ctx) - // wait_agent keeps blocking until 10 seconds. clock.Advance(7 * time.Second) close(waitGo) trap.MustWait(ctx).MustRelease(ctx) - // The serial call launches at 10 seconds. + // Release the serial start timestamp. trap.MustWait(ctx).MustRelease(ctx) - // It completes at 12 seconds. clock.Advance(2 * time.Second) close(serialGo) trap.MustWait(ctx).MustRelease(ctx) @@ -721,8 +646,6 @@ func TestExecuteLocalTools_SerialAfterBilledSiblingBillsUnion(t *testing.T) { "the serial call's completion ends the window") } -// BilledIntervalsDuration merges overlapping windows and skips gaps, so -// parallel calls bill once and spans without billed work bill nothing. func TestBilledIntervalsDuration(t *testing.T) { t.Parallel() diff --git a/coderd/x/chatd/generation.go b/coderd/x/chatd/generation.go index ba8aeafb352..a1379422def 100644 --- a/coderd/x/chatd/generation.go +++ b/coderd/x/chatd/generation.go @@ -865,19 +865,8 @@ func (s *taskStarter) executeLocalTools( var outcome chatloop.ToolExecutionOutcome var spawnDispatchErr error if len(allowed) > 0 { - // Seed the batch start and the dispatched calls on the buffer - // episode so an interrupt can bill the partial window this step - // would have reported. Chatloop invokes the callback only once - // dispatch is guaranteed, so an interrupt whose cancellation - // lands before the tools launch finds no batch and bills - // nothing. Only dispatched calls are seeded, keyed by their - // position in the unresolved call order the interrupt walks: - // rejected calls never run, so an interrupt must not treat - // their missing completions as still-running work, and a - // rejected call must not consume a same-ID dispatched - // occurrence. An exclusive-policy violation dispatches nothing - // (chatloop synthesizes error results before the callback), so - // no billable batch starts there either. + // Seed only calls that will dispatch, after policy checks. Positional + // indexes skip rejected calls and keep duplicate IDs distinct. var onBatchStart func() if !exclusiveRejected { dispatched := make([]messagepartbuffer.DispatchedToolCall, 0, len(allowed)) @@ -1123,29 +1112,14 @@ type generationAttempt struct { // can bill the window the step would have reported. It is always // non-nil when beginGenerationAttempt succeeds. startModelInvocation func() - // startToolBatch marks the start of the attempt's billable local - // tool batch window on the buffer episode and records which tool - // call occurrences were actually dispatched, so an interrupt can - // bill the window the step would have reported without charging - // calls that were rejected before execution. It is always non-nil - // when beginGenerationAttempt succeeds. + // startToolBatch stamps dispatch and seeds tool-call occurrences. It is + // always non-nil after beginGenerationAttempt. startToolBatch func(calls []messagepartbuffer.DispatchedToolCall) - // recordToolStart records the instant a tool call occurrence began - // executing on the buffer episode. Concurrent calls start with the - // batch, but serial calls launch only after every concurrent - // sibling settles, so an interrupt needs the marks to bill each - // call from its actual start and to skip a dispatched call that - // never launched. callIndex addresses the occurrence within the - // dispatched batch, matching the order startToolBatch seeded. It - // is always non-nil when beginGenerationAttempt succeeds. + // recordToolStart stamps an occurrence's actual start; serial calls may + // start after dispatch. It is always non-nil after beginGenerationAttempt. recordToolStart func(callIndex int, toolCallID string, startedAt time.Time) - // recordToolCompletion records a tool call occurrence's completion - // instant on the buffer episode as the batch executes, so an - // interrupt can end an already-finished tool's billable window at - // its real completion instead of the interrupt instant. callIndex - // addresses the occurrence within the dispatched batch, matching - // the order startToolBatch seeded. It is always non-nil when - // beginGenerationAttempt succeeds. + // recordToolCompletion stamps an occurrence's completion. It is always + // non-nil after beginGenerationAttempt. recordToolCompletion func(callIndex int, toolCallID string, completedAt time.Time) // closeEpisode closes the attempt's buffer episode. It is always // non-nil when beginGenerationAttempt succeeds. diff --git a/coderd/x/chatd/message_conversion.go b/coderd/x/chatd/message_conversion.go index f430e2ee28a..b6fd8c12d0e 100644 --- a/coderd/x/chatd/message_conversion.go +++ b/coderd/x/chatd/message_conversion.go @@ -75,17 +75,9 @@ func buildCommitStepMessages(input buildCommitStepMessagesInput) (stepMessagesFo return stepMessagesForCommit{}, xerrors.Errorf("marshal tool result: %w", err) } msg := baseMessage(database.ChatMessageRoleTool, database.ChatMessageVisibilityBoth, input.modelConfigID, contentVersion, content) - // The batch's billable window lands on the single tool row whose - // completion ended it; every other row in the batch stays NULL so - // usage reporting, which sums runtime_ms across rows, bills the - // batch exactly once. Only the first row with the window-defining - // ID carries it, because providers can emit duplicate tool call - // IDs and billing every duplicate would multiply the sum. A - // positive window gates the match instead of a non-empty ID, - // because calls without IDs also execute and bill; their window - // lands on the first ID-less tool row. Zero maps to NULL, so a - // sub-millisecond window persists the same way an unmeasured - // one does. + // Assign the batch window to one matching tool row because usage sums + // runtime_ms. The first match handles duplicate and ID-less calls; + // zero stays NULL. if !batchRuntimeAssigned && input.step.BatchRuntime > 0 && toolResult.ToolCallID == input.step.BatchRuntimeToolCallID { msg.RuntimeMs = nullInt64IfNonZero(input.step.BatchRuntime.Milliseconds()) batchRuntimeAssigned = true @@ -648,14 +640,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 model-invocation runtime, because tool batches bill - // their window on tool rows instead. The buffer episode only - // carries a model 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 5c22606a4c9..e47c9e49a3f 100644 --- a/coderd/x/chatd/message_conversion_test.go +++ b/coderd/x/chatd/message_conversion_test.go @@ -100,9 +100,6 @@ func TestBuildCommitStepMessages_LocalToolResultsBecomeToolMessages(t *testing.T require.JSONEq(t, `{"stdout":"/tmp"}`, string(toolParts[0].Result)) } -// A local tool batch bills its window on the single tool row whose -// completion ended it; the batch's other rows stay NULL so summing -// runtime_ms across rows bills the batch exactly once. func TestBuildCommitStepMessages_BatchRuntimeLandsOnWindowDefiningToolRow(t *testing.T) { t.Parallel() @@ -135,9 +132,6 @@ func TestBuildCommitStepMessages_BatchRuntimeLandsOnWindowDefiningToolRow(t *tes require.Equal(t, sql.NullInt64{Int64: 10000, Valid: true}, got.Messages[1].RuntimeMs) } -// Duplicate tool call IDs, which providers can emit and which admission -// does not always reject, must not multiply the bill: only the first row -// with the window-defining ID carries the batch runtime. func TestBuildCommitStepMessages_DuplicateToolCallIDsBillOnce(t *testing.T) { t.Parallel() @@ -168,9 +162,6 @@ func TestBuildCommitStepMessages_DuplicateToolCallIDsBillOnce(t *testing.T) { require.False(t, got.Messages[1].RuntimeMs.Valid) } -// A window defined by a call without a tool call ID still persists: the -// positive window gates the match, so the runtime lands on the first -// ID-less tool row instead of being dropped with the empty-ID sentinel. func TestBuildCommitStepMessages_EmptyIDWindowLandsOnIDLessRow(t *testing.T) { t.Parallel() @@ -201,8 +192,6 @@ func TestBuildCommitStepMessages_EmptyIDWindowLandsOnIDLessRow(t *testing.T) { require.Equal(t, sql.NullInt64{Int64: 60_000, Valid: true}, got.Messages[1].RuntimeMs) } -// Assistant rows synthesized from a tool batch (attachment file parts) -// never carry the batch runtime: it belongs to the tool row alone. func TestBuildCommitStepMessages_BatchAttachmentAssistantRowStaysNull(t *testing.T) { t.Parallel() @@ -231,8 +220,6 @@ func TestBuildCommitStepMessages_BatchAttachmentAssistantRowStaysNull(t *testing require.Equal(t, sql.NullInt64{Int64: 3000, Valid: true}, got.Messages[1].RuntimeMs) } -// A batch whose billable window is empty (for example only sub-agent -// orchestration tools ran) must persist runtime_ms NULL on every row. func TestBuildCommitStepMessages_ZeroBatchRuntimeLeavesRuntimeNull(t *testing.T) { t.Parallel() diff --git a/coderd/x/chatd/messagepartbuffer/message_part_buffer.go b/coderd/x/chatd/messagepartbuffer/message_part_buffer.go index a4d48d89d87..379993db1c9 100644 --- a/coderd/x/chatd/messagepartbuffer/message_part_buffer.go +++ b/coderd/x/chatd/messagepartbuffer/message_part_buffer.go @@ -101,21 +101,11 @@ type episodeState struct { // episode's provider stream is opened. It is zero for episodes // that never invoke a model, such as local tool execution // batches. - modelStartedAt time.Time - // toolBatchStartedAt is stamped by StartToolBatch when the - // episode begins executing its local tool batch. It is zero for - // episodes that never execute local tools, such as model - // invocations that finish without tool calls. + modelStartedAt time.Time toolBatchStartedAt time.Time - // toolCompletions holds one entry per dispatched tool call - // occurrence, seeded by StartToolBatch in dispatch order and - // stamped by RecordToolCompletion as each tool finishes. Tool - // results are published only after the whole batch completes, so - // these per-occurrence stamps are the only live view of which - // tools were dispatched and which already finished when an - // interrupt lands. Keyed storage would collapse duplicate tool - // call IDs, which reach execution when lifecycle hooks are - // disabled, into one shared state. + // toolCompletions stores per-occurrence start and completion stamps for + // interrupts. Positional storage preserves queued serial calls and + // duplicate IDs. toolCompletions []ToolCompletion closed bool closedAt time.Time @@ -233,41 +223,27 @@ func (b *Buffer) StartModelInvocation(key Key) error { return nil } -// DispatchedToolCall identifies one tool call occurrence dispatched in an -// episode's local tool batch. +// DispatchedToolCall identifies a dispatched tool-call occurrence. type DispatchedToolCall struct { - // CallIndex is the caller-assigned position of this occurrence in the - // step's full unresolved tool-call order, which is the order interrupt - // reconstruction walks. It differs from the occurrence's position in - // the dispatched batch when calls were rejected before execution, and - // it disambiguates duplicate tool call IDs. + // CallIndex is its position in the unresolved call order, used to + // distinguish rejected and duplicate-ID calls. CallIndex int ToolCallID string } -// ToolCompletion tracks one dispatched tool call occurrence in an -// episode's local tool batch. StartedAt is zero until the call begins -// executing, which for a serial tool call happens only after every -// concurrent sibling has settled; CompletedAt is zero while the call -// has not finished. +// ToolCompletion tracks a tool-call occurrence. StartedAt is zero when the +// start is unknown; CompletedAt is zero while unfinished. type ToolCompletion struct { - // CallIndex mirrors DispatchedToolCall.CallIndex. It is -1 for - // completions recorded without a matching seeded occurrence, which - // never correlate to an unresolved call. + // CallIndex is -1 when no seeded occurrence matched. CallIndex int ToolCallID string StartedAt time.Time CompletedAt time.Time } -// StartToolBatch stamps the instant the episode begins executing its local -// tool batch, which starts the batch's billable runtime window, and seeds -// one not-yet-started entry per dispatched tool call occurrence. Readers -// can then distinguish a dispatched call awaiting launch (seeded, zero -// start), one that is executing (started, zero completion), and one never -// dispatched at all (absent), such as a call denied by a lifecycle hook or -// rejected as ambiguous before execution, and duplicate tool call IDs keep -// distinct per-occurrence states. +// StartToolBatch stamps dispatch and seeds one entry per occurrence. A zero +// start is queued; a started entry with zero completion is running. Absent +// calls were not dispatched, and duplicate IDs remain distinct. func (b *Buffer) StartToolBatch(key Key, calls []DispatchedToolCall) error { b.mu.Lock() defer b.mu.Unlock() @@ -292,16 +268,10 @@ func (b *Buffer) StartToolBatch(key Key, calls []DispatchedToolCall) error { return nil } -// RecordToolStart records the instant a dispatched local tool call -// occurrence began executing. Concurrent calls start with the batch, but -// calls whose tools run serially launch only after every concurrent -// sibling has settled, so an interrupt must bill each call from its own -// start and must not treat a dispatched call that never launched as -// running work. dispatchIndex addresses the exact occurrence seeded by -// StartToolBatch, mirroring RecordToolCompletion; when it does not match, -// the stamp falls back to the first unstarted occurrence with the ID. A -// start with no seeded occurrence is dropped: it cannot correlate to an -// unresolved call, so it could never bill. +// RecordToolStart stamps a dispatched call's actual start. Concurrent calls +// start with the batch; serial calls may start later. If dispatchIndex does +// not match, the first unstarted same-ID occurrence is used. Unknown starts +// are dropped because they cannot correlate to unresolved calls. func (b *Buffer) RecordToolStart(key Key, dispatchIndex int, toolCallID string, startedAt time.Time) error { b.mu.Lock() defer b.mu.Unlock() @@ -332,18 +302,10 @@ func (b *Buffer) RecordToolStart(key Key, dispatchIndex int, toolCallID string, return nil } -// RecordToolCompletion records the instant a local tool call in the -// episode's tool batch finished. Tool goroutines report completions as -// they happen, so an interrupt can bill tools that already finished up -// to their real completion instead of treating every canceled call as -// still running. dispatchIndex addresses the exact occurrence seeded by -// StartToolBatch, as the occurrence's position within the dispatched -// batch: duplicate tool call IDs make the ID alone ambiguous, and -// stamping the wrong same-ID occurrence would let a finished call mark -// its still-running twin as done. When dispatchIndex does not match the -// seeded occurrence, the stamp falls back to the first still-running -// occurrence with the ID, or is appended with CallIndex -1, so an -// executed call is never dropped. +// RecordToolCompletion stamps a call as it finishes, so interrupts use the +// actual completion. dispatchIndex selects the seeded occurrence; otherwise +// the first unfinished same-ID occurrence is used. Unmatched completions +// append with CallIndex -1. func (b *Buffer) RecordToolCompletion(key Key, dispatchIndex int, toolCallID string, completedAt time.Time) error { b.mu.Lock() defer b.mu.Unlock() @@ -464,9 +426,8 @@ func (b *Buffer) ModelInvokedAt(key Key) time.Time { return episode.modelStartedAt } -// ToolBatchStartedAt returns the instant stamped by StartToolBatch, or the -// zero time if there is none. Read it before CloseEpisode: closed episodes -// are garbage collected, so reading afterwards races the cleanup loop. +// ToolBatchStartedAt returns the StartToolBatch stamp, or zero if absent. +// Read it before CloseEpisode because closed episodes are garbage collected. func (b *Buffer) ToolBatchStartedAt(key Key) time.Time { b.mu.Lock() defer b.mu.Unlock() @@ -477,13 +438,9 @@ func (b *Buffer) ToolBatchStartedAt(key Key) time.Time { return episode.toolBatchStartedAt } -// ToolCompletions returns a copy of the tool batch's dispatched call -// occurrences in dispatch order, or nil when the episode is unknown or -// never started a batch. Occurrences seeded by StartToolBatch but not -// yet stamped by RecordToolCompletion carry the zero time: they were -// dispatched and are still running. Interrupt handling must use -// CloseEpisodeForBilling instead: a separate read-then-close would let -// completions land in the gap and go missing from the snapshot. +// ToolCompletions returns copied occurrence state. A zero start means not +// launched; a zero completion means unfinished. Interrupts must use +// CloseEpisodeForBilling to avoid a read-close race. func (b *Buffer) ToolCompletions(key Key) []ToolCompletion { b.mu.Lock() defer b.mu.Unlock() @@ -494,42 +451,23 @@ func (b *Buffer) ToolCompletions(key Key) []ToolCompletion { return slices.Clone(episode.toolCompletions) } -// EpisodeBilling is the billing state an episode accumulated before it -// closed. +// EpisodeBilling is the billing snapshot captured when an episode closes. type EpisodeBilling struct { - // ClosedAt is the instant the episode first closed. Interrupt - // handling uses it as the interrupt instant: it is stable across - // repeat closes, so a retried interrupt task bills the same window - // every attempt instead of one that grows with each retry's later - // clock reading. + // ClosedAt is the first close instant, used as the retry-stable + // interrupt time. ClosedAt time.Time - // ModelInvokedAt is the StartModelInvocation stamp, or zero when - // the episode never opened a provider stream. + // ModelInvokedAt is the StartModelInvocation stamp, or zero if absent. ModelInvokedAt time.Time - // ToolBatchStartedAt is the StartToolBatch stamp, or zero when the - // episode never started a local tool batch. + // ToolBatchStartedAt is the StartToolBatch stamp, or zero if absent. ToolBatchStartedAt time.Time - // ToolCompletions are the tool batch's dispatched call occurrences - // in dispatch order; see Buffer.ToolCompletions. + // ToolCompletions is the close-time occurrence snapshot. ToolCompletions []ToolCompletion } -// CloseEpisodeForBilling closes the episode like CloseEpisode and -// returns its billing stamps from the same critical section. The -// interrupt task uses it so every stamp accepted before closure is in -// the snapshot and none can be recorded afterwards: reading and closing -// in separate steps would let a tool completion or batch start land in -// the gap, billing a finished tool as still running or losing a live -// batch's window entirely. Closing an unknown episode creates it -// closed, mirroring CloseEpisode, and reports empty billing. -// -// Calling it again on an already-closed episode returns the same -// snapshot and pushes the episode's eviction deadline out by the -// retention window. Interrupt task retries re-read the snapshot on -// every attempt with backoff well under the retention window, so a -// retrying interrupt keeps its billing state and buffered parts alive -// through a database outage instead of losing them to the cleanup -// loop mid-retry. +// CloseEpisodeForBilling closes the episode and returns billing stamps from +// the same critical section, preventing a read-close race. Unknown episodes +// close blank. Re-closing returns the original snapshot and refreshes +// retention so retries keep the same billing state and buffered parts. func (b *Buffer) CloseEpisodeForBilling(key Key) (EpisodeBilling, error) { b.mu.Lock() defer b.mu.Unlock() @@ -673,10 +611,8 @@ func (b *Buffer) queueClosedEpisodeLocked(key Key, episode *episodeState) { heap.Push(&b.closedEpisodes, item) } -// refreshClosedEpisodeEvictionLocked pushes a closed episode's eviction -// deadline out to evictAt plus the retention window by queueing a fresh -// heap item. The superseded item stays in the heap until the cleanup -// loop pops it and skips it via the closedHeapItem identity check. +// refreshClosedEpisodeEvictionLocked refreshes retention with a new heap +// item; cleanup skips superseded items by identity. func (b *Buffer) refreshClosedEpisodeEvictionLocked(key Key, episode *episodeState, evictAt time.Time) { item := &closedEpisodeItem{key: key, closedAt: evictAt} episode.closedHeapItem = item diff --git a/coderd/x/chatd/messagepartbuffer/message_part_buffer_test.go b/coderd/x/chatd/messagepartbuffer/message_part_buffer_test.go index 0f604270e70..d79dffe47f3 100644 --- a/coderd/x/chatd/messagepartbuffer/message_part_buffer_test.go +++ b/coderd/x/chatd/messagepartbuffer/message_part_buffer_test.go @@ -161,22 +161,16 @@ func TestBuffer_ToolBatchStartedAt(t *testing.T) { require.NoError(t, buffer.CreateEpisode(key)) require.Zero(t, buffer.ToolBatchStartedAt(key), "episode without a tool batch has no batch stamp") - // Attempt setup happens before tools start executing and is not - // billable. clock.Advance(time.Second) require.NoError(t, buffer.StartToolBatch(key, []messagepartbuffer.DispatchedToolCall{{CallIndex: 0, ToolCallID: "call-1"}})) startedAt := buffer.ToolBatchStartedAt(key) require.Equal(t, clock.Now(), startedAt) - // Closing must not move the recorded stamp, and a closed episode - // no longer accepts a batch start. clock.Advance(1500 * time.Millisecond) require.NoError(t, buffer.CloseEpisode(key)) require.ErrorIs(t, buffer.StartToolBatch(key, []messagepartbuffer.DispatchedToolCall{{CallIndex: 0, ToolCallID: "call-1"}}), messagepartbuffer.ErrEpisodeClosed) require.Equal(t, startedAt, buffer.ToolBatchStartedAt(key)) - // Episodes that never execute local tools, such as model - // invocations without tool calls, report no batch stamp. modelOnly := testEpisodeKey() require.NoError(t, buffer.CreateEpisode(modelOnly)) require.NoError(t, buffer.StartModelInvocation(modelOnly)) @@ -198,9 +192,6 @@ func TestBuffer_ToolCompletions(t *testing.T) { require.NoError(t, buffer.CreateEpisode(key)) require.Empty(t, buffer.ToolCompletions(key), "episode without a tool batch has no completions") - // Starting the batch seeds one still-running occurrence per - // dispatched call, in dispatch order. Duplicate IDs keep distinct - // occurrences instead of collapsing into one shared state. require.NoError(t, buffer.StartToolBatch(key, []messagepartbuffer.DispatchedToolCall{ {CallIndex: 0, ToolCallID: "call-1"}, {CallIndex: 1, ToolCallID: "call-dup"}, @@ -211,9 +202,6 @@ func TestBuffer_ToolCompletions(t *testing.T) { {CallIndex: 1, ToolCallID: "call-dup"}, {CallIndex: 2, ToolCallID: "call-dup"}, }, buffer.ToolCompletions(key)) - // A completion stamps the occurrence addressed by its call index, - // not the first occurrence with a matching ID: the duplicate's - // second occurrence finishing must leave the first still running. clock.Advance(time.Second) secondDupCompletedAt := clock.Now() require.NoError(t, buffer.RecordToolCompletion(key, 2, "call-dup", secondDupCompletedAt)) @@ -232,14 +220,9 @@ func TestBuffer_ToolCompletions(t *testing.T) { {CallIndex: 2, ToolCallID: "call-dup", CompletedAt: secondDupCompletedAt}, }, completions) - // The returned slice is a copy: mutating it must not corrupt the - // episode's recorded completions. completions[0].CompletedAt = clock.Now() require.True(t, buffer.ToolCompletions(key)[0].CompletedAt.IsZero()) - // A completion whose ID was never seeded is appended with - // CallIndex -1: the executed call is not dropped, but it never - // correlates to an unresolved call. clock.Advance(time.Second) unseededAt := clock.Now() require.NoError(t, buffer.RecordToolCompletion(key, 5, "call-unseeded", unseededAt)) @@ -249,8 +232,6 @@ func TestBuffer_ToolCompletions(t *testing.T) { CompletedAt: unseededAt, }, buffer.ToolCompletions(key)[3]) - // A closed episode keeps its recorded completions but accepts no - // more, matching the batch-start stamp's lifecycle. require.NoError(t, buffer.CloseEpisode(key)) require.ErrorIs(t, buffer.RecordToolCompletion(key, 0, "call-1", clock.Now()), messagepartbuffer.ErrEpisodeClosed) require.True(t, buffer.ToolCompletions(key)[0].CompletedAt.IsZero()) @@ -267,9 +248,6 @@ func TestBuffer_RecordToolStart(t *testing.T) { require.ErrorIs(t, buffer.RecordToolStart(key, 0, "call-1", clock.Now()), messagepartbuffer.ErrEpisodeNotFound) require.NoError(t, buffer.CreateEpisode(key)) - // Seeded occurrences carry no start: a dispatched serial call - // waits behind its concurrent siblings, so seeding must not make - // it look like running work. require.NoError(t, buffer.StartToolBatch(key, []messagepartbuffer.DispatchedToolCall{ {CallIndex: 0, ToolCallID: "call-1"}, {CallIndex: 1, ToolCallID: "call-dup"}, @@ -279,9 +257,6 @@ func TestBuffer_RecordToolStart(t *testing.T) { require.True(t, completion.StartedAt.IsZero(), "seeding must not mark occurrences as started") } - // A start stamps the occurrence addressed by its dispatch index, - // not the first occurrence with a matching ID: the duplicate's - // second occurrence launching must leave the first unstarted. clock.Advance(time.Second) secondDupStartedAt := clock.Now() require.NoError(t, buffer.RecordToolStart(key, 2, "call-dup", secondDupStartedAt)) @@ -291,20 +266,14 @@ func TestBuffer_RecordToolStart(t *testing.T) { {CallIndex: 2, ToolCallID: "call-dup", StartedAt: secondDupStartedAt}, }, buffer.ToolCompletions(key)) - // A start whose dispatch index does not match falls back to the - // first unstarted occurrence with the ID. clock.Advance(time.Second) firstDupStartedAt := clock.Now() require.NoError(t, buffer.RecordToolStart(key, 7, "call-dup", firstDupStartedAt)) require.Equal(t, firstDupStartedAt, buffer.ToolCompletions(key)[1].StartedAt) - // A start whose ID was never seeded is dropped rather than - // appended: it cannot correlate to an unresolved call, so it could - // never bill. require.NoError(t, buffer.RecordToolStart(key, 9, "call-unseeded", clock.Now())) require.Len(t, buffer.ToolCompletions(key), 3) - // The billing snapshot carries start marks through the close. billing, err := buffer.CloseEpisodeForBilling(key) require.NoError(t, err) require.Equal(t, secondDupStartedAt, billing.ToolCompletions[2].StartedAt) @@ -319,8 +288,6 @@ func TestBuffer_CloseEpisodeForBilling(t *testing.T) { buffer := messagepartbuffer.New(messagepartbuffer.Options{Clock: clock}) defer buffer.Close() - // Closing an unknown episode creates it closed, like CloseEpisode, - // and reports empty billing stamped with the close instant. unknown := testEpisodeKey() billing, err := buffer.CloseEpisodeForBilling(unknown) require.NoError(t, err) @@ -329,9 +296,6 @@ func TestBuffer_CloseEpisodeForBilling(t *testing.T) { require.Zero(t, billing.ToolBatchStartedAt) require.Empty(t, billing.ToolCompletions) - // The snapshot carries every stamp accepted before closure, and - // stamps are rejected afterwards, so nothing can land in a gap - // between reading and closing. key := testEpisodeKey() require.NoError(t, buffer.CreateEpisode(key)) clock.Advance(time.Second) @@ -356,20 +320,12 @@ func TestBuffer_CloseEpisodeForBilling(t *testing.T) { }, billing.ToolCompletions) require.ErrorIs(t, buffer.RecordToolCompletion(key, 1, "call-2", clock.Now()), messagepartbuffer.ErrEpisodeClosed) - // Closing an already-closed episode reports the identical snapshot, - // including the original close instant: a retried interrupt task - // bills the same window every attempt, and an interrupt racing the - // generation task's own close loses nothing. clock.Advance(time.Second) again, err := buffer.CloseEpisodeForBilling(key) require.NoError(t, err) require.Equal(t, billing, again) } -// A retrying interrupt task re-reads its billing snapshot on every -// attempt; each re-read pushes the episode's eviction deadline out, so -// a retry loop outlasting the original retention window keeps the -// snapshot instead of losing it to the cleanup loop mid-outage. func TestBuffer_CloseEpisodeForBillingRefreshesEviction(t *testing.T) { t.Parallel() @@ -387,24 +343,17 @@ func TestBuffer_CloseEpisodeForBillingRefreshesEviction(t *testing.T) { first, err := buffer.CloseEpisodeForBilling(key) require.NoError(t, err) - // A retry 45 seconds in re-reads the snapshot and refreshes the - // deadline. clock.Advance(45 * time.Second).MustWait(ctx) again, err := buffer.CloseEpisodeForBilling(key) require.NoError(t, err) require.Equal(t, first, again) - // The cleanup tick after the original one-minute deadline must not - // collect the refreshed episode: the snapshot survives. clock.Advance(15 * time.Second).MustWait(ctx) again, err = buffer.CloseEpisodeForBilling(key) require.NoError(t, err) require.Equal(t, first, again) - // Once retries stop refreshing it, the episode ages out and a later - // close reports empty billing again. GetParts collects due episodes - // synchronously, so the assertions cannot race the cleanup - // goroutine's handling of the delivered ticks. + // GetParts later runs cleanup synchronously, avoiding a goroutine race. clock.Advance(time.Minute).MustWait(ctx) clock.Advance(time.Minute).MustWait(ctx) _, err = buffer.GetParts(key) diff --git a/coderd/x/chatd/options.go b/coderd/x/chatd/options.go index a2a35165fed..6c464240755 100644 --- a/coderd/x/chatd/options.go +++ b/coderd/x/chatd/options.go @@ -70,10 +70,8 @@ type chatWorkerTaskStartInput struct { DebugTurn *runnerDebugTurn SessionStart *sessionStartTracker StopNudges *stopNudgeTracker - // InterruptSnapshot, set for interrupt tasks, is shared by every - // retry attempt of the task instance so the first attempt's episode - // snapshot survives buffer eviction during a stalled attempt. Nil - // makes each attempt re-read the buffer. + // InterruptSnapshot carries one interrupt task's first episode snapshot + // across retries. Nil re-reads the buffer. InterruptSnapshot *interruptEpisodeSnapshot } diff --git a/coderd/x/chatd/runner.go b/coderd/x/chatd/runner.go index bcc4bd222ce..d5a77b971e4 100644 --- a/coderd/x/chatd/runner.go +++ b/coderd/x/chatd/runner.go @@ -233,9 +233,6 @@ func (r *runner) spawnTaskIfNeeded(kind taskKind, state runnerStateUpdate) { StopNudges: &r.stopNudges, } if kind == taskKindInterrupt { - // Shared by every retry attempt of this task instance so the - // first attempt's episode snapshot survives buffer eviction - // during a stalled attempt. input.InterruptSnapshot = &interruptEpisodeSnapshot{} } go r.runTask(taskCtx, kind, key, input, done) diff --git a/coderd/x/chatd/subagent_catalog.go b/coderd/x/chatd/subagent_catalog.go index 937adc5fa0c..18673c571ab 100644 --- a/coderd/x/chatd/subagent_catalog.go +++ b/coderd/x/chatd/subagent_catalog.go @@ -38,15 +38,8 @@ const ( "external or web research, parallel research, or tasks that may need edits." ) -// unbilledSubagentToolNames lists the sub-agent orchestration tools whose -// execution time is excluded from the local-tool runtime persisted to -// chat_messages.runtime_ms. Every chat bills its own model and tool time, -// including child agents, so a parent's wait_agent window would count each -// child's already-billed runtime a second time. The remaining orchestration -// tools are millisecond-scale bookkeeping and are excluded with it so the -// billing rule stays one explainable category. Classification uses the tool -// name as called, so the deprecated close_agent alias is listed alongside -// interrupt_agent. +// 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, diff --git a/coderd/x/chatd/subagent_internal_test.go b/coderd/x/chatd/subagent_internal_test.go index 61a5d8c9551..432bff2ce29 100644 --- a/coderd/x/chatd/subagent_internal_test.go +++ b/coderd/x/chatd/subagent_internal_test.go @@ -4262,9 +4262,6 @@ func TestAwaitSubagentCompletion(t *testing.T) { }) } -// The unbilled set must track the sub-agent orchestration catalog exactly: -// every orchestration tool and deprecated alias is excluded from local tool -// runtime billing, and nothing outside the catalog is. func TestUnbilledSubagentToolNamesMatchCatalog(t *testing.T) { t.Parallel() diff --git a/coderd/x/chatd/tasks.go b/coderd/x/chatd/tasks.go index 35a4b2820fd..f9e9b109b3c 100644 --- a/coderd/x/chatd/tasks.go +++ b/coderd/x/chatd/tasks.go @@ -241,14 +241,9 @@ func (o chatWorkerOptions) retryOptions() retryWrapperOptions { } } -// interruptEpisodeSnapshot carries the interrupt task's one-time episode -// snapshot across retry attempts of the same task instance. It is -// captured before the task's first database read: one stalled read (a -// database outage holds an attempt up to the task timeout) can outlive -// the buffer's closed-episode retention, after which the buffer only -// offers a blank recreated episode that would underbill the interrupted -// work and drop the partial messages. Attempts of one task run -// sequentially, so no locking is needed. +// interruptEpisodeSnapshot carries one interrupt task's episode snapshot +// across retries. Capturing it before the first database read prevents a +// stalled read from losing billing state to buffer eviction. type interruptEpisodeSnapshot struct { loaded bool key messagepartbuffer.Key @@ -256,13 +251,8 @@ type interruptEpisodeSnapshot struct { parts []messagepartbuffer.Part } -// closeInterruptEpisode closes the interrupted attempt's buffer episode -// and returns its billing snapshot and buffered parts. Closing and -// snapshotting billing state must be one atomic step: the generation -// goroutine records batch starts and tool completions concurrently, so -// a read-then-close would let stamps land in the gap and go missing -// from the snapshot. Unknown episodes close blank, so interruption -// converges even when the worker exited before publishing parts. +// closeInterruptEpisode atomically closes the episode and snapshots billing. +// Unknown episodes close blank so interruption converges. func (s *taskStarter) closeInterruptEpisode(ctx context.Context, key messagepartbuffer.Key) (messagepartbuffer.EpisodeBilling, []messagepartbuffer.Part, error) { billing, err := s.opts.MessagePartBuffer.CloseEpisodeForBilling(key) if err != nil { @@ -287,13 +277,8 @@ func (s *taskStarter) closeInterruptEpisode(ctx context.Context, key messagepart func (s *taskStarter) StartInterrupt(ctx context.Context, input chatWorkerTaskStartInput) error { snapshot := input.InterruptSnapshot - // Capture the episode before the chat read below, which can stall - // on a database outage past the buffer's retention and let the - // cleanup loop evict the episode before a post-read capture. The - // input carries the interrupted attempt number, and generation - // attempts only advance while the chat is running, so the pre-read - // key matches the row while it stays interrupting; if the read - // below disagrees, the snapshot is retaken with the row's key. + // Snapshot before the chat read, which can outlive buffer retention. + // Retake it if the stored attempt differs from the input key. if snapshot != nil && !snapshot.loaded { earlyKey := messagepartbuffer.Key{ ChatID: input.ChatID, @@ -331,17 +316,10 @@ func (s *taskStarter) StartInterrupt(ctx context.Context, input chatWorkerTaskSt var episodeBilling messagepartbuffer.EpisodeBilling var parts []messagepartbuffer.Part if snapshot != nil && snapshot.loaded && snapshot.key == key { - // This task already snapshotted the episode, before its first - // chat read or on a previous attempt. Reuse it: the episode - // may have been evicted from the buffer while an attempt - // stalled, and re-reading would find a blank recreated - // episode. + // Reuse the snapshot because the buffer may have evicted the episode. episodeBilling = snapshot.billing parts = snapshot.parts } else { - // No usable snapshot: either the caller passed none (tests) or - // the row's attempt number disagrees with the pre-read key, so - // the pre-read snapshot describes the wrong episode. episodeBilling, parts, err = s.closeInterruptEpisode(ctx, key) if err != nil { return err @@ -353,11 +331,7 @@ func (s *taskStarter) StartInterrupt(ctx context.Context, input chatWorkerTaskSt snapshot.parts = parts } } - // The interrupt instant is the episode's first close, which is - // stable across repeat closes: a retried interrupt task (transient - // database errors) recomputes identical partial messages and - // billing windows instead of billing still-running tools through - // each retry's later clock reading. + // The first close is a retry-stable interrupt instant. interruptedAt := episodeBilling.ClosedAt var attemptRuntime time.Duration if !episodeBilling.ModelInvokedAt.IsZero() { @@ -382,11 +356,8 @@ func (s *taskStarter) StartInterrupt(ctx context.Context, input chatWorkerTaskSt return xerrors.Errorf("load chat for task: %w", err) } messages := partialMessages - // Reuse the interrupt instant captured when the episode closed: - // a fresh clock read here would run while this transaction - // waits on the database (and again on retries), inflating the - // billed window of a still-running call past the actual - // interrupt. + // Reuse the close instant so database delay and retries do not inflate + // billing. committedCancels, err := committedPendingLocalToolCancellationMessages(ctx, store, chat, interruptedAt, interruptedToolBatchBilling{ batchStartedAt: episodeBilling.ToolBatchStartedAt, toolCompletions: episodeBilling.ToolCompletions, @@ -766,29 +737,14 @@ func dynamicToolNamesFromChat(chat database.Chat) map[string]bool { return names } -// interruptedToolBatchBilling carries what the interrupt task knows about -// the live tool batch it is canceling, so the synthesized cancellation -// rows can bill the partial window the batch would have reported. +// interruptedToolBatchBilling is the live batch state used to bill +// synthesized cancellation rows. type interruptedToolBatchBilling struct { - // batchStartedAt is the StartToolBatch stamp from the interrupted - // attempt's buffer episode. Zero when no batch was live at the - // interrupt (crash recovery, state promotion), in which case the - // cancellation rows carry no runtime. + // batchStartedAt is zero when no local tool batch was live. batchStartedAt time.Time - // toolCompletions holds the live batch's dispatched tool call - // occurrences, each keyed by its position in the step's unresolved - // call order: seeded when the batch started, marked as each call - // began executing, and stamped as each tool finished. A completed - // occurrence bills the interval from its start to its completion, - // so a batch whose billed tools all finished early does not bill - // the longer window of a still-running unbilled tool such as - // wait_agent. A started but uncompleted occurrence was still - // running when the interrupt landed. A seeded occurrence that - // never started is a serial call that was still waiting behind - // its siblings, and bills nothing. A call with no occurrence at - // its position was never dispatched, such as a call rejected as - // malformed or ambiguous before execution, and bills nothing even - // though it too receives a cancellation row. + // toolCompletions uses unresolved-call positions. Completed calls bill + // to completion, running calls bill to the interrupt, and queued or + // undispatched calls bill nothing. toolCompletions []messagepartbuffer.ToolCompletion } @@ -813,11 +769,8 @@ func committedPendingLocalToolCancellationMessages( if len(localCalls) == 0 { return nil, nil } - // Dispatched occurrences keyed by their position in the unresolved - // call order, the same order this loop walks. Positional matching - // keeps a call rejected before execution from consuming a same-ID - // dispatched occurrence and keeps duplicate tool call IDs from - // sharing one completion state. + // Match by unresolved-call position so rejected and duplicate-ID calls + // cannot share occurrence state. dispatched := make(map[int]messagepartbuffer.ToolCompletion, len(billing.toolCompletions)) for _, completion := range billing.toolCompletions { if completion.CallIndex < 0 { @@ -854,15 +807,9 @@ func committedPendingLocalToolCancellationMessages( if billing.batchStartedAt.IsZero() || unbilledSubagentToolNames[call.ToolName] { continue } - // Only dispatched calls bill: a call with no occurrence at its - // position was rejected before execution, and an ID mismatch - // means the seed does not describe this call. A completed - // occurrence bills its execution interval; a started but - // uncompleted one was still running, so its window ends at the - // interrupt. A dispatched occurrence that never started is a - // serial call still waiting behind its siblings when the - // interrupt landed, and bills nothing. Strictly-after keeps - // the earliest call on ties, matching billableBatchWindow. + // Bill only matching dispatched calls. Completed calls end at + // completion, running calls end at the interrupt, and queued serial + // calls are skipped. Ties keep the first call. occurrence, ok := dispatched[i] if !ok || occurrence.ToolCallID != call.ToolCallID { continue @@ -876,9 +823,7 @@ func committedPendingLocalToolCancellationMessages( end = interruptedAt } if start.IsZero() { - // Live execution always marks a start before a - // completion; fall back to the batch start rather than - // dropping a completed call's billed work. + // Completed calls should have starts; batchStart preserves legacy data. start = billing.batchStartedAt } intervals = append(intervals, chatloop.BilledInterval{Start: start, End: end}) @@ -887,11 +832,7 @@ func committedPendingLocalToolCancellationMessages( windowRowIdx = len(result) - 1 } } - // Mirror the committed-batch policy: bill the union of the billed - // calls' execution intervals once, on the cancellation row of the - // billed tool call whose window ends last. The union never charges - // a span where only unbilled tools were running, such as a - // sub-agent wait that delayed a serial call's launch. + // Bill the interval union once on the row whose interval ends last. if windowRowIdx >= 0 { result[windowRowIdx].RuntimeMs = nullInt64IfNonZero(chatloop.BilledIntervalsDuration(intervals).Milliseconds()) } diff --git a/coderd/x/chatd/tasks_test.go b/coderd/x/chatd/tasks_test.go index 18af7c49597..49d594ffce5 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) @@ -477,9 +475,6 @@ func TestInterruptTask_PartialAssistantWithoutModelInvocationHasNoRuntime(t *tes require.False(t, assistant.RuntimeMs.Valid) } -// interruptedBatch is a chat with committed unresolved local tool calls -// and a live attempt buffer episode, ready for StartInterrupt to -// synthesize cancellation rows. type interruptedBatch struct { chat database.Chat starter *taskStarter @@ -489,9 +484,6 @@ type interruptedBatch struct { runnerID uuid.UUID } -// interruptedBatchFixture commits an assistant message with the given -// unresolved local tool calls and prepares the live attempt's buffer -// episode, so StartInterrupt synthesizes cancellation rows for them. func interruptedBatchFixture( t *testing.T, f *taskTestFixture, @@ -557,8 +549,6 @@ func (b interruptedBatch) interruptWithSnapshot(t *testing.T, f *taskTestFixture return messages } -// findToolResultMessage returns the tool-role message answering the given -// tool call ID. func findToolResultMessage(t *testing.T, messages []database.ChatMessage, toolCallID string) database.ChatMessage { t.Helper() for _, msg := range messages { @@ -577,11 +567,6 @@ func findToolResultMessage(t *testing.T, messages []database.ChatMessage, toolCa return database.ChatMessage{} } -// An interrupt mid tool batch bills the partial window on the cancellation -// row of the billed tool that defines it. A billed tool that recorded its -// completion before the interrupt ends the window there, so a -// still-running unbilled wait_agent does not stretch the bill to the -// interrupt instant. func TestInterruptTask_ToolBatchBillsPartialWindowOnCancellationRow(t *testing.T) { t.Parallel() @@ -594,30 +579,19 @@ func TestInterruptTask_ToolBatchBillsPartialWindowOnCancellationRow(t *testing.T }) buffer := batch.starter.opts.MessagePartBuffer - // Advances stay under the buffer's 15s cleanup tick, which shares - // this mock clock. - // Attempt setup happens before the tools start and is not billable. + // Keep advances below the buffer's 15-second cleanup tick. batch.clock.Advance(2 * time.Second) require.NoError(t, buffer.StartToolBatch(batch.key, []messagepartbuffer.DispatchedToolCall{ {CallIndex: 0, ToolCallID: execCallID}, {CallIndex: 1, ToolCallID: waitCallID}, })) - // Both concurrent calls launch with the batch, the way the launch - // loop's start callback records them. require.NoError(t, buffer.RecordToolStart(batch.key, 0, execCallID, batch.clock.Now())) require.NoError(t, buffer.RecordToolStart(batch.key, 1, waitCallID, batch.clock.Now())) - // execute completes 3 seconds into the batch and records its - // completion, the way the tool goroutine's completion callback - // does. Its result is not published: results publish only after - // the whole batch finishes. + // Record a live completion without publishing a tool result. batch.clock.Advance(3 * time.Second) require.NoError(t, buffer.RecordToolCompletion(batch.key, 0, execCallID, batch.clock.Now())) - // wait_agent is still blocked on its child when the interrupt lands - // 5 seconds later. batch.clock.Advance(5 * time.Second) - // The first attempt stores the episode snapshot it read, so a - // retry after buffer eviction reuses it. snapshot := &interruptEpisodeSnapshot{} messages := batch.interruptWithSnapshot(t, f, snapshot) execRow := findToolResultMessage(t, messages, execCallID) @@ -629,10 +603,6 @@ func TestInterruptTask_ToolBatchBillsPartialWindowOnCancellationRow(t *testing.T require.Len(t, snapshot.billing.ToolCompletions, 2) } -// A retry attempt reuses the snapshot its first attempt captured: after -// a stalled attempt outlives the buffer's retention and the episode is -// evicted, the blank recreated episode must not replace the snapshot's -// billing state, so the interrupted batch still bills its window. func TestInterruptTask_RetrySnapshotOutlivesEpisodeEviction(t *testing.T) { t.Parallel() @@ -642,12 +612,7 @@ func TestInterruptTask_RetrySnapshotOutlivesEpisodeEviction(t *testing.T) { {Type: codersdk.ChatMessagePartTypeToolCall, ToolCallID: execCallID, ToolName: "execute", Args: json.RawMessage(`{}`)}, }) - // Simulate a retry: the first attempt snapshotted a live batch - // (started at 2s, closed at 9s, execute still running) and then - // stalled past the buffer's retention. The buffer's episode - // carries none of that state, like the blank episode a retry - // recreates after eviction, so billing can only come from the - // carried snapshot. + // Use only the carried snapshot; the buffer has no episode state. batch.clock.Advance(2 * time.Second) batchStartedAt := batch.clock.Now() batch.clock.Advance(7 * time.Second) @@ -666,11 +631,6 @@ func TestInterruptTask_RetrySnapshotOutlivesEpisodeEviction(t *testing.T) { require.Equal(t, sql.NullInt64{Int64: 7_000, Valid: true}, execRow.RuntimeMs) } -// The interrupt task captures its episode snapshot before its first -// database read: an attempt whose chat read fails (or stalls past the -// buffer's retention, evicting the episode) has already stored the -// billing state, so a later attempt bills the original window even -// after the buffer forgets the episode. func TestInterruptTask_SnapshotCapturedBeforeChatRead(t *testing.T) { t.Parallel() @@ -686,10 +646,6 @@ func TestInterruptTask_SnapshotCapturedBeforeChatRead(t *testing.T) { require.NoError(t, buffer.RecordToolStart(batch.key, 0, execCallID, batch.clock.Now())) batch.clock.Advance(3 * time.Second) - // The first attempt's chat read fails: the chat is still running, - // so the interrupting fence rejects it. The snapshot must already - // be captured by then, with the interrupt instant at this failed - // attempt's episode close. snapshot := &interruptEpisodeSnapshot{} err := batch.starter.StartInterrupt(testutil.Context(t, testutil.WaitShort), chatWorkerTaskStartInput{ ChatID: batch.chat.ID, @@ -704,23 +660,17 @@ func TestInterruptTask_SnapshotCapturedBeforeChatRead(t *testing.T) { require.True(t, snapshot.loaded, "the snapshot must be captured before the chat read") require.Equal(t, batch.key, snapshot.key) - // The attempt outlives the buffer retention: cleanup ticks evict - // the closed episode, so only the carried snapshot remains. ctx := testutil.Context(t, testutil.WaitShort) batch.clock.Advance(10 * time.Second).MustWait(ctx) batch.clock.Advance(15 * time.Second).MustWait(ctx) _, err = buffer.GetParts(batch.key) require.ErrorIs(t, err, messagepartbuffer.ErrEpisodeNotFound) - // The retry bills the pre-read window: batch start at 2s to the - // first attempt's close at 5s. messages := batch.interruptWithSnapshot(t, f, snapshot) execRow := findToolResultMessage(t, messages, execCallID) require.Equal(t, sql.NullInt64{Int64: 3_000, Valid: true}, execRow.RuntimeMs) } -// A billed tool still running at the interrupt bills up to the interrupt -// instant. func TestInterruptTask_RunningBilledToolBillsUpToInterrupt(t *testing.T) { t.Parallel() @@ -733,7 +683,6 @@ func TestInterruptTask_RunningBilledToolBillsUpToInterrupt(t *testing.T) { batch.clock.Advance(2 * time.Second) require.NoError(t, batch.starter.opts.MessagePartBuffer.StartToolBatch(batch.key, []messagepartbuffer.DispatchedToolCall{{CallIndex: 0, ToolCallID: execCallID}})) require.NoError(t, batch.starter.opts.MessagePartBuffer.RecordToolStart(batch.key, 0, execCallID, batch.clock.Now())) - // The tool is still running when the interrupt lands 7 seconds in. batch.clock.Advance(7 * time.Second) messages := batch.interrupt(t, f) @@ -741,8 +690,6 @@ func TestInterruptTask_RunningBilledToolBillsUpToInterrupt(t *testing.T) { require.Equal(t, sql.NullInt64{Int64: 7_000, Valid: true}, execRow.RuntimeMs) } -// An interrupted batch of only unbilled sub-agent orchestration tools -// bills nothing: an interrupted lone wait_agent stays free. func TestInterruptTask_UnbilledOnlyBatchBillsNothingOnInterrupt(t *testing.T) { t.Parallel() @@ -762,11 +709,6 @@ func TestInterruptTask_UnbilledOnlyBatchBillsNothingOnInterrupt(t *testing.T) { require.False(t, waitRow.RuntimeMs.Valid) } -// A dispatched call that never began executing bills nothing on -// interrupt: a billed serial tool stays queued until every concurrent -// sibling settles, so an interrupt landing during an unbilled -// wait_agent must not treat the waiting serial call as running from -// the batch start. func TestInterruptTask_UnstartedSerialCallBillsNothingOnInterrupt(t *testing.T) { t.Parallel() @@ -784,9 +726,6 @@ func TestInterruptTask_UnstartedSerialCallBillsNothingOnInterrupt(t *testing.T) {CallIndex: 0, ToolCallID: waitCallID}, {CallIndex: 1, ToolCallID: serialCallID}, })) - // Only the concurrent wait_agent launched; the serial call stays - // queued behind it, so no start mark arrives before the interrupt - // lands 10 seconds later. require.NoError(t, buffer.RecordToolStart(batch.key, 0, waitCallID, batch.clock.Now())) batch.clock.Advance(10 * time.Second) @@ -799,10 +738,6 @@ func TestInterruptTask_UnstartedSerialCallBillsNothingOnInterrupt(t *testing.T) } } -// A billed serial call that launched bills from its own start: the -// interrupt window unions the early billed sibling's interval with the -// serial call's, so the span where only the unbilled wait_agent was -// running is never charged. func TestInterruptTask_StartedSerialCallBillsFromItsOwnStart(t *testing.T) { t.Parallel() @@ -823,24 +758,16 @@ func TestInterruptTask_StartedSerialCallBillsFromItsOwnStart(t *testing.T) { {CallIndex: 1, ToolCallID: waitCallID}, {CallIndex: 2, ToolCallID: serialCallID}, })) - // The two concurrent calls launch with the batch. require.NoError(t, buffer.RecordToolStart(batch.key, 0, execCallID, batch.clock.Now())) require.NoError(t, buffer.RecordToolStart(batch.key, 1, waitCallID, batch.clock.Now())) - // execute completes 3 seconds in. batch.clock.Advance(3 * time.Second) require.NoError(t, buffer.RecordToolCompletion(batch.key, 0, execCallID, batch.clock.Now())) - // wait_agent completes at 6 seconds, letting the serial call - // launch; it is still running when the interrupt lands at 8 - // seconds. batch.clock.Advance(3 * time.Second) require.NoError(t, buffer.RecordToolCompletion(batch.key, 1, waitCallID, batch.clock.Now())) require.NoError(t, buffer.RecordToolStart(batch.key, 2, serialCallID, batch.clock.Now())) batch.clock.Advance(2 * time.Second) messages := batch.interrupt(t, f) - // The union bills execute's 3s and the serial call's 2s, not the - // 3s gap where only wait_agent ran, and lands once on the serial - // call's row: its window ends last. serialRow := findToolResultMessage(t, messages, serialCallID) require.Equal(t, sql.NullInt64{Int64: 5_000, Valid: true}, serialRow.RuntimeMs) execRow := findToolResultMessage(t, messages, execCallID) @@ -849,10 +776,6 @@ func TestInterruptTask_StartedSerialCallBillsFromItsOwnStart(t *testing.T) { require.False(t, waitRow.RuntimeMs.Valid) } -// Two dispatched billed calls sharing one tool call ID keep distinct -// occurrence states: one completing early must not make the other look -// finished, so the interrupted batch still bills through to the -// interrupt, once. func TestInterruptTask_DuplicateCallIDsKeepOccurrenceStates(t *testing.T) { t.Parallel() @@ -871,8 +794,6 @@ func TestInterruptTask_DuplicateCallIDsKeepOccurrenceStates(t *testing.T) { })) require.NoError(t, buffer.RecordToolStart(batch.key, 0, dupCallID, batch.clock.Now())) require.NoError(t, buffer.RecordToolStart(batch.key, 1, dupCallID, batch.clock.Now())) - // One occurrence completes 3 seconds in; the other keeps running - // until the interrupt lands 3 seconds later, defining the window. batch.clock.Advance(3 * time.Second) require.NoError(t, buffer.RecordToolCompletion(batch.key, 0, dupCallID, batch.clock.Now())) batch.clock.Advance(3 * time.Second) @@ -887,10 +808,6 @@ func TestInterruptTask_DuplicateCallIDsKeepOccurrenceStates(t *testing.T) { require.Equal(t, []int64{6_000}, billed, "the still-running occurrence bills the full window exactly once") } -// Same-ID occurrences with different billing classifications stay -// correlated: an unbilled wait_agent occurrence finishing early stamps -// its own occurrence, not the still-running billed execute sharing its -// ID, so the execute occurrence still bills through to the interrupt. func TestInterruptTask_DuplicateCallIDCompletionStampsOwnOccurrence(t *testing.T) { t.Parallel() @@ -909,9 +826,6 @@ func TestInterruptTask_DuplicateCallIDCompletionStampsOwnOccurrence(t *testing.T })) require.NoError(t, buffer.RecordToolStart(batch.key, 0, dupCallID, batch.clock.Now())) require.NoError(t, buffer.RecordToolStart(batch.key, 1, dupCallID, batch.clock.Now())) - // The unbilled wait_agent occurrence (index 1) completes 3 seconds - // in; the billed execute occurrence (index 0) keeps running until - // the interrupt 3 seconds later. batch.clock.Advance(3 * time.Second) require.NoError(t, buffer.RecordToolCompletion(batch.key, 1, dupCallID, batch.clock.Now())) batch.clock.Advance(3 * time.Second) @@ -926,11 +840,6 @@ func TestInterruptTask_DuplicateCallIDCompletionStampsOwnOccurrence(t *testing.T require.Equal(t, []int64{6_000}, billed, "the running execute occurrence bills to the interrupt; wait_agent's early completion must not end it at 3s") } -// A rejected call must not steal a same-ID dispatched occurrence: with -// a rejected billed execute preceding a dispatched unbilled wait_agent -// that shares its ID, positional matching leaves the execute row -// without an occurrence, so the interrupted batch bills nothing instead -// of charging the whole wait window to a call that never ran. func TestInterruptTask_RejectedDuplicateIDDoesNotStealDispatchedOccurrence(t *testing.T) { t.Parallel() @@ -942,9 +851,6 @@ func TestInterruptTask_RejectedDuplicateIDDoesNotStealDispatchedOccurrence(t *te }) batch.clock.Advance(2 * time.Second) - // Only the wait_agent occurrence (unresolved position 1) was - // dispatched: the execute occurrence sharing its ID was rejected - // as malformed before execution. require.NoError(t, batch.starter.opts.MessagePartBuffer.StartToolBatch(batch.key, []messagepartbuffer.DispatchedToolCall{{CallIndex: 1, ToolCallID: dupCallID}})) require.NoError(t, batch.starter.opts.MessagePartBuffer.RecordToolStart(batch.key, 0, dupCallID, batch.clock.Now())) batch.clock.Advance(10 * time.Second) @@ -957,11 +863,6 @@ func TestInterruptTask_RejectedDuplicateIDDoesNotStealDispatchedOccurrence(t *te } } -// A billed call rejected before execution (hook denial, ambiguous-call -// rejection) is absent from the batch's dispatched set, so its missing -// completion is not evidence that it ran: an interrupted batch whose -// only dispatched call is an unbilled wait_agent bills nothing even -// though the rejected execute call also receives a cancellation row. func TestInterruptTask_RejectedCallBillsNothingOnInterrupt(t *testing.T) { t.Parallel() @@ -974,9 +875,6 @@ func TestInterruptTask_RejectedCallBillsNothingOnInterrupt(t *testing.T) { }) batch.clock.Advance(2 * time.Second) - // Only wait_agent was dispatched: execute was rejected before - // execution and never ran, so only the occurrence at wait_agent's - // unresolved position (1) is seeded. require.NoError(t, batch.starter.opts.MessagePartBuffer.StartToolBatch(batch.key, []messagepartbuffer.DispatchedToolCall{{CallIndex: 1, ToolCallID: waitCallID}})) require.NoError(t, batch.starter.opts.MessagePartBuffer.RecordToolStart(batch.key, 0, waitCallID, batch.clock.Now())) batch.clock.Advance(10 * time.Second) @@ -988,9 +886,6 @@ func TestInterruptTask_RejectedCallBillsNothingOnInterrupt(t *testing.T) { require.False(t, waitRow.RuntimeMs.Valid) } -// Cancellation rows synthesized without a live tool batch (crash -// recovery: the episode never stamped a batch start) carry no runtime, -// consistent with generation losing in-flight runtime on a crash. func TestInterruptTask_ToolCancellationWithoutLiveBatchHasNoRuntime(t *testing.T) { t.Parallel() @@ -1000,8 +895,6 @@ func TestInterruptTask_ToolCancellationWithoutLiveBatchHasNoRuntime(t *testing.T {Type: codersdk.ChatMessagePartTypeToolCall, ToolCallID: execCallID, ToolName: "execute", Args: json.RawMessage(`{}`)}, }) - // No StartToolBatch stamp: the batch never went live on this - // attempt. batch.clock.Advance(10 * time.Second) messages := batch.interrupt(t, f) diff --git a/coderd/x/chatd/toolinput.go b/coderd/x/chatd/toolinput.go index d1b15b91ce6..6e961f72219 100644 --- a/coderd/x/chatd/toolinput.go +++ b/coderd/x/chatd/toolinput.go @@ -15,10 +15,8 @@ import ( // reject before pre_tool_use so a hook consumer is never asked to authorize // bytes whose meaning depends on which reader resolves them, and so input that // cannot be carried in a hook payload fails as a retryable tool error instead -// of a dispatch failure. allowedIndexes carries each allowed call's position -// in the input slice, so callers that need to correlate the dispatched subset -// back to the full unresolved call order (tool batch billing) do not have to -// re-derive it from tool call IDs, which duplicates make ambiguous. +// of a dispatch failure. allowedIndexes maps allowed calls back to the input +// order without relying on duplicate-prone IDs. func partitionAmbiguousToolCalls( prepared generationPrepared, toolCalls []fantasy.ToolCallContent, From a2d17782ccf4e3b8c96ffd640242047774cbfb9b Mon Sep 17 00:00:00 2001 From: Jaayden Halko Date: Wed, 19 Aug 2026 10:28:09 +0000 Subject: [PATCH 17/20] docs(coderd/x/chatd): document tool runtime billing --- coderd/x/chatd/ARCHITECTURE.md | 19 +++++++++---------- 1 file changed, 9 insertions(+), 10 deletions(-) diff --git a/coderd/x/chatd/ARCHITECTURE.md b/coderd/x/chatd/ARCHITECTURE.md index 033902a9821..afd5d45e190 100644 --- a/coderd/x/chatd/ARCHITECTURE.md +++ b/coderd/x/chatd/ARCHITECTURE.md @@ -724,9 +724,11 @@ The buffer exposes the following API: - `AddPart(chat_id, history_version, generation_attempt, content)`: adds a message part to the buffer. Returns a predefined error if the episode is not found or the array is full. - `GetParts(chat_id, history_version, generation_attempt)`: returns the message parts for an episode. Returns a predefined error if the episode is not found. - `StartModelInvocation(chat_id, history_version, generation_attempt)`: stamps the instant the episode opens its provider stream. Returns a predefined error if the episode is not found or already closed. Episodes that never invoke a model, such as local tool execution batches, are never stamped. -- `ModelInvokedAt(chat_id, history_version, generation_attempt)`: returns the instant stamped by `StartModelInvocation`, or the zero time when the episode is unknown or never opened a provider stream. It must be read before `CloseEpisode`, because closed episodes are garbage collected and reading afterwards races the cleanup loop. The interrupt goroutine reads it just before closing the episode and uses the span between that instant and the interrupt as the interrupted attempt's billable runtime. -- TODO(CODAGT-928): document tool-batch billing APIs, including serial - starts, duplicate IDs, atomic snapshots, stable close time, and retries. +- `ModelInvokedAt(chat_id, history_version, generation_attempt)`: returns the `StartModelInvocation` stamp, or zero if none exists. Interrupt handling reads it through `CloseEpisodeForBilling` to avoid a read-close race. +- `StartToolBatch(chat_id, history_version, generation_attempt, calls)`: stamps the batch start and seeds one entry per dispatched call occurrence. +- `RecordToolStart(chat_id, history_version, generation_attempt, call_index, tool_call_id, started_at)` and `RecordToolCompletion(chat_id, history_version, generation_attempt, call_index, tool_call_id, completed_at)`: stamp when each call occurrence starts and finishes. An unstarted occurrence represents a queued serial call and does not bill on interrupt. +- `ToolBatchStartedAt(chat_id, history_version, generation_attempt)` and `ToolCompletions(chat_id, history_version, generation_attempt)`: return the batch start and per-occurrence execution stamps. Interrupt handling reads them through `CloseEpisodeForBilling`. +- `CloseEpisodeForBilling(chat_id, history_version, generation_attempt)`: atomically closes the episode and returns its stable first-close time plus model and tool billing stamps. Repeated calls return the same snapshot and refresh its retention deadline. - `SubscribeToEpisode(chat_id, history_version, generation_attempt)`: returns a go channel that will receive all message parts for the episode. It spawns a goroutine that delivers parts to the channel. It's live until the episode is closed or until a subscriber requests that the channel be closed. Once the goroutine delivers all message parts for a closed episode, it closes the channel and exits. If the episode is already closed at the time of the call, the goroutine delivers all message parts for the episode, closes the channel, and exits. `SubscribeToEpisode` does not return an error if the episode is not found: it waits for it to be created instead. Closed episodes are garbage collected after at least 15 seconds since they were closed and when they have no active subscribers. The message part buffer maintains a garbage collection goroutine. @@ -932,13 +934,10 @@ The interrupt goroutine is responsible for handling interrupts. It is spawned wh The goroutine does the following in order: -1. It fetches the generation attempt number from the database. -2. It closes the episode corresponding to its history version and generation attempt by calling the `CloseEpisode` method on the [Message part buffer](#message-part-buffer). -3. It reads the buffered parts for that episode by calling the `GetParts` method on the message part buffer. -4. It applies the `FinishInterruption(partial?)` transition on the core state machine. If there are no buffered parts for that episode, or the episode is not found, it passes `nil` as the `partial` argument. - -TODO(CODAGT-928): update interrupt steps for atomic snapshots, pre-read -retry carry, stable close time, and partial tool-batch billing. +1. Before its first database read, it closes the expected episode with `CloseEpisodeForBilling`, reads its buffered parts, and retains both snapshots across task retries. +2. It loads the chat and verifies the generation attempt, resnapshotting the matching episode if the key differs. +3. It uses the episode's stable first-close time to convert buffered parts and synthesize tool cancellation rows. Started, billable local tools contribute the union of their partial execution intervals; unstarted or excluded tools do not. +4. It applies the `FinishInterruption(partial?)` transition on the core state machine, passing `nil` when there are no partial messages. #### Dynamic tools timeout goroutine From fd35fd8a8eac6e7a76f3de99282750f248d60981 Mon Sep 17 00:00:00 2001 From: Jaayden Halko Date: Thu, 20 Aug 2026 07:01:59 +0000 Subject: [PATCH 18/20] chore: cleanup for comments --- coderd/x/chatd/ARCHITECTURE.md | 11 +- coderd/x/chatd/chatloop/chatloop.go | 26 ++-- coderd/x/chatd/chatloop/runtime_test.go | 15 +-- coderd/x/chatd/generation.go | 12 +- .../messagepartbuffer/message_part_buffer.go | 103 +++------------ .../message_part_buffer_test.go | 119 +++--------------- coderd/x/chatd/tasks.go | 49 +++++--- coderd/x/chatd/tasks_test.go | 50 ++++---- 8 files changed, 120 insertions(+), 265 deletions(-) diff --git a/coderd/x/chatd/ARCHITECTURE.md b/coderd/x/chatd/ARCHITECTURE.md index afd5d45e190..f961e65f569 100644 --- a/coderd/x/chatd/ARCHITECTURE.md +++ b/coderd/x/chatd/ARCHITECTURE.md @@ -724,11 +724,10 @@ The buffer exposes the following API: - `AddPart(chat_id, history_version, generation_attempt, content)`: adds a message part to the buffer. Returns a predefined error if the episode is not found or the array is full. - `GetParts(chat_id, history_version, generation_attempt)`: returns the message parts for an episode. Returns a predefined error if the episode is not found. - `StartModelInvocation(chat_id, history_version, generation_attempt)`: stamps the instant the episode opens its provider stream. Returns a predefined error if the episode is not found or already closed. Episodes that never invoke a model, such as local tool execution batches, are never stamped. -- `ModelInvokedAt(chat_id, history_version, generation_attempt)`: returns the `StartModelInvocation` stamp, or zero if none exists. Interrupt handling reads it through `CloseEpisodeForBilling` to avoid a read-close race. +- `ModelInvokedAt(chat_id, history_version, generation_attempt)`: returns the `StartModelInvocation` stamp, or zero if none exists. Interrupt handling reads it before closing the episode. - `StartToolBatch(chat_id, history_version, generation_attempt, calls)`: stamps the batch start and seeds one entry per dispatched call occurrence. -- `RecordToolStart(chat_id, history_version, generation_attempt, call_index, tool_call_id, started_at)` and `RecordToolCompletion(chat_id, history_version, generation_attempt, call_index, tool_call_id, completed_at)`: stamp when each call occurrence starts and finishes. An unstarted occurrence represents a queued serial call and does not bill on interrupt. -- `ToolBatchStartedAt(chat_id, history_version, generation_attempt)` and `ToolCompletions(chat_id, history_version, generation_attempt)`: return the batch start and per-occurrence execution stamps. Interrupt handling reads them through `CloseEpisodeForBilling`. -- `CloseEpisodeForBilling(chat_id, history_version, generation_attempt)`: atomically closes the episode and returns its stable first-close time plus model and tool billing stamps. Repeated calls return the same snapshot and refresh its retention deadline. +- `RecordToolStart(chat_id, history_version, generation_attempt, dispatch_index, started_at)` and `RecordToolCompletion(chat_id, history_version, generation_attempt, dispatch_index, completed_at)`: stamp when each call occurrence starts and finishes. An unstarted occurrence represents a queued serial call and does not bill on interrupt. +- `ToolBatchStartedAt(chat_id, history_version, generation_attempt)` and `ToolCompletions(chat_id, history_version, generation_attempt)`: return the batch start and per-occurrence execution stamps. Interrupt handling reads them before closing the episode. - `SubscribeToEpisode(chat_id, history_version, generation_attempt)`: returns a go channel that will receive all message parts for the episode. It spawns a goroutine that delivers parts to the channel. It's live until the episode is closed or until a subscriber requests that the channel be closed. Once the goroutine delivers all message parts for a closed episode, it closes the channel and exits. If the episode is already closed at the time of the call, the goroutine delivers all message parts for the episode, closes the channel, and exits. `SubscribeToEpisode` does not return an error if the episode is not found: it waits for it to be created instead. Closed episodes are garbage collected after at least 15 seconds since they were closed and when they have no active subscribers. The message part buffer maintains a garbage collection goroutine. @@ -934,9 +933,9 @@ The interrupt goroutine is responsible for handling interrupts. It is spawned wh The goroutine does the following in order: -1. Before its first database read, it closes the expected episode with `CloseEpisodeForBilling`, reads its buffered parts, and retains both snapshots across task retries. +1. Before its first database read, it reads the expected episode's billing stamps, closes it with `CloseEpisode`, reads its buffered parts, and retains the snapshot across task retries. 2. It loads the chat and verifies the generation attempt, resnapshotting the matching episode if the key differs. -3. It uses the episode's stable first-close time to convert buffered parts and synthesize tool cancellation rows. Started, billable local tools contribute the union of their partial execution intervals; unstarted or excluded tools do not. +3. It uses the snapshot's interrupt time to convert buffered parts and synthesize tool cancellation rows. Started, billable local tools contribute the union of their partial execution intervals; unstarted or excluded tools do not. 4. It applies the `FinishInterruption(partial?)` transition on the core state machine, passing `nil` when there are no partial messages. #### Dynamic tools timeout goroutine diff --git a/coderd/x/chatd/chatloop/chatloop.go b/coderd/x/chatd/chatloop/chatloop.go index 48a67474418..5f1fd04d8e6 100644 --- a/coderd/x/chatd/chatloop/chatloop.go +++ b/coderd/x/chatd/chatloop/chatloop.go @@ -271,14 +271,14 @@ type ExecuteLocalToolsOptions struct { OnBatchStart func() // OnToolStart fires when each local call begins. Serial calls may start // after concurrent siblings settle, so interrupts bill actual starts and - // skip dispatched calls that never run. callIndex identifies the + // skip dispatched calls that never run. dispatchIndex identifies the // dispatch-order occurrence. - OnToolStart func(callIndex int, toolCallID string, startedAt time.Time) + 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; - // callIndex identifies the dispatch-order occurrence when IDs collide. + // dispatchIndex identifies the dispatch-order occurrence. // The callback must be concurrency-safe. - OnToolComplete func(callIndex int, toolCallID string, completedAt time.Time) + OnToolComplete func(dispatchIndex int, completedAt time.Time) PublishMessagePart func(codersdk.ChatMessageRole, codersdk.ChatMessagePart) Logger slog.Logger @@ -602,7 +602,7 @@ func ExecuteLocalTools(ctx context.Context, opts ExecuteLocalToolsOptions) (Tool for i, tr := range policyResults { recordToolResultTimestamp(&result, tr.ToolCallID, now) if opts.OnToolComplete != nil { - opts.OnToolComplete(i, tr.ToolCallID, now) + opts.OnToolComplete(i, now) } publishToolAttachments(ctx, opts.Logger, tr, now, publishMessagePart) ssePart := chatprompt.PartFromContentWithLogger(ctx, opts.Logger, tr) @@ -628,10 +628,10 @@ func ExecuteLocalTools(ctx context.Context, opts ExecuteLocalToolsOptions) (Tool 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(callIndex int, toolCallID string, startedAt time.Time) { - orderedStarts[callIndex] = startedAt + onToolStart := func(dispatchIndex int, startedAt time.Time) { + orderedStarts[dispatchIndex] = startedAt if opts.OnToolStart != nil { - opts.OnToolStart(callIndex, toolCallID, startedAt) + opts.OnToolStart(dispatchIndex, startedAt) } } toolResults := executeTools( @@ -1228,8 +1228,8 @@ func executeTools( maxResultBytes int, toolNameAliases map[string]string, batchStart time.Time, - onStart func(callIndex int, toolCallID string, startedAt time.Time), - onComplete func(callIndex int, toolCallID string, completedAt 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 { @@ -1298,7 +1298,7 @@ func executeTools( // accurate individual completion times. completedAt[i] = clockNow(clock) if onComplete != nil { - onComplete(i, tc.ToolCallID, completedAt[i]) + onComplete(i, completedAt[i]) } }() results[i] = executeSingleTool( @@ -1328,7 +1328,7 @@ func executeTools( continue } if onStart != nil { - onStart(i, tc.ToolCallID, batchStart) + onStart(i, batchStart) } wg.Add(1) go func() { @@ -1350,7 +1350,7 @@ func executeTools( for _, i := range serialIndexes { // Stamp serial calls at launch, not batch start. if onStart != nil { - onStart(i, localToolCalls[i].ToolCallID, clockNow(clock)) + onStart(i, clockNow(clock)) } runCall(i, localToolCalls[i]) } diff --git a/coderd/x/chatd/chatloop/runtime_test.go b/coderd/x/chatd/chatloop/runtime_test.go index b542dd69c67..625234ea97c 100644 --- a/coderd/x/chatd/chatloop/runtime_test.go +++ b/coderd/x/chatd/chatloop/runtime_test.go @@ -517,9 +517,8 @@ func TestExecuteLocalTools_OnToolCompleteReportsLiveCompletions(t *testing.T) { defer trap.Close() type completion struct { - callIndex int - toolCallID string - completedAt time.Time + dispatchIndex int + completedAt time.Time } completionCh := make(chan completion, 2) fastGo := make(chan struct{}) @@ -530,8 +529,8 @@ func TestExecuteLocalTools_OnToolCompleteReportsLiveCompletions(t *testing.T) { blockingTool("slow_tool", slowGo, fantasy.NewTextResponse("done")), }, ActiveTools: []string{"fast_tool", "slow_tool"}, - OnToolComplete: func(callIndex int, toolCallID string, completedAt time.Time) { - completionCh <- completion{callIndex: callIndex, toolCallID: toolCallID, completedAt: completedAt} + OnToolComplete: func(dispatchIndex int, completedAt time.Time) { + completionCh <- completion{dispatchIndex: dispatchIndex, completedAt: completedAt} }, ToolCalls: []fantasy.ToolCallContent{ {ToolCallID: "call-fast", ToolName: "fast_tool", Input: "{}"}, @@ -544,14 +543,12 @@ func TestExecuteLocalTools_OnToolCompleteReportsLiveCompletions(t *testing.T) { close(fastGo) trap.MustWait(ctx).MustRelease(ctx) fast := testutil.RequireReceive(ctx, t, completionCh) - require.Equal(t, "call-fast", fast.toolCallID) - require.Equal(t, 0, fast.callIndex, "callIndex is the call's position in dispatch order") + 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, "call-slow", slow.toolCallID) - require.Equal(t, 1, slow.callIndex, "callIndex is the call's position in dispatch order") + require.Equal(t, 1, slow.dispatchIndex) require.Equal(t, 50*time.Second, slow.completedAt.Sub(fast.completedAt)) outcome := testutil.RequireReceive(ctx, t, resultCh) diff --git a/coderd/x/chatd/generation.go b/coderd/x/chatd/generation.go index a1379422def..1ef11b03f11 100644 --- a/coderd/x/chatd/generation.go +++ b/coderd/x/chatd/generation.go @@ -1117,10 +1117,10 @@ type generationAttempt struct { startToolBatch func(calls []messagepartbuffer.DispatchedToolCall) // recordToolStart stamps an occurrence's actual start; serial calls may // start after dispatch. It is always non-nil after beginGenerationAttempt. - recordToolStart func(callIndex int, toolCallID string, startedAt time.Time) + recordToolStart func(dispatchIndex int, startedAt time.Time) // recordToolCompletion stamps an occurrence's completion. It is always // non-nil after beginGenerationAttempt. - recordToolCompletion func(callIndex int, toolCallID string, completedAt time.Time) + recordToolCompletion func(dispatchIndex int, completedAt time.Time) // closeEpisode closes the attempt's buffer episode. It is always // non-nil when beginGenerationAttempt succeeds. closeEpisode func() @@ -1170,11 +1170,11 @@ func (s *taskStarter) beginGenerationAttempt( startToolBatch: func(calls []messagepartbuffer.DispatchedToolCall) { _ = s.opts.MessagePartBuffer.StartToolBatch(key, calls) }, - recordToolStart: func(callIndex int, toolCallID string, startedAt time.Time) { - _ = s.opts.MessagePartBuffer.RecordToolStart(key, callIndex, toolCallID, startedAt) + recordToolStart: func(dispatchIndex int, startedAt time.Time) { + _ = s.opts.MessagePartBuffer.RecordToolStart(key, dispatchIndex, startedAt) }, - recordToolCompletion: func(callIndex int, toolCallID string, completedAt time.Time) { - _ = s.opts.MessagePartBuffer.RecordToolCompletion(key, callIndex, toolCallID, completedAt) + recordToolCompletion: func(dispatchIndex int, completedAt time.Time) { + _ = s.opts.MessagePartBuffer.RecordToolCompletion(key, dispatchIndex, completedAt) }, closeEpisode: func() { _ = s.opts.MessagePartBuffer.CloseEpisode(key) diff --git a/coderd/x/chatd/messagepartbuffer/message_part_buffer.go b/coderd/x/chatd/messagepartbuffer/message_part_buffer.go index 379993db1c9..f13cc4961d3 100644 --- a/coderd/x/chatd/messagepartbuffer/message_part_buffer.go +++ b/coderd/x/chatd/messagepartbuffer/message_part_buffer.go @@ -234,7 +234,7 @@ type DispatchedToolCall struct { // ToolCompletion tracks a tool-call occurrence. StartedAt is zero when the // start is unknown; CompletedAt is zero while unfinished. type ToolCompletion struct { - // CallIndex is -1 when no seeded occurrence matched. + // CallIndex is the occurrence's position in the unresolved call order. CallIndex int ToolCallID string StartedAt time.Time @@ -269,10 +269,9 @@ func (b *Buffer) StartToolBatch(key Key, calls []DispatchedToolCall) error { } // RecordToolStart stamps a dispatched call's actual start. Concurrent calls -// start with the batch; serial calls may start later. If dispatchIndex does -// not match, the first unstarted same-ID occurrence is used. Unknown starts +// start with the batch; serial calls may start later. Unknown dispatch indexes // are dropped because they cannot correlate to unresolved calls. -func (b *Buffer) RecordToolStart(key Key, dispatchIndex int, toolCallID string, startedAt time.Time) error { +func (b *Buffer) RecordToolStart(key Key, dispatchIndex int, startedAt time.Time) error { b.mu.Lock() defer b.mu.Unlock() if b.closed { @@ -285,28 +284,19 @@ func (b *Buffer) RecordToolStart(key Key, dispatchIndex int, toolCallID string, if episode.closed { return ErrEpisodeClosed } - if dispatchIndex >= 0 && dispatchIndex < len(episode.toolCompletions) { - entry := &episode.toolCompletions[dispatchIndex] - if entry.ToolCallID == toolCallID && entry.StartedAt.IsZero() { - entry.StartedAt = startedAt - return nil - } + if dispatchIndex < 0 || dispatchIndex >= len(episode.toolCompletions) { + return nil } - for i := range episode.toolCompletions { - entry := &episode.toolCompletions[i] - if entry.ToolCallID == toolCallID && entry.StartedAt.IsZero() { - entry.StartedAt = startedAt - return nil - } + entry := &episode.toolCompletions[dispatchIndex] + if entry.StartedAt.IsZero() { + entry.StartedAt = startedAt } return nil } // RecordToolCompletion stamps a call as it finishes, so interrupts use the -// actual completion. dispatchIndex selects the seeded occurrence; otherwise -// the first unfinished same-ID occurrence is used. Unmatched completions -// append with CallIndex -1. -func (b *Buffer) RecordToolCompletion(key Key, dispatchIndex int, toolCallID string, completedAt time.Time) error { +// actual completion. Unknown dispatch indexes are dropped. +func (b *Buffer) RecordToolCompletion(key Key, dispatchIndex int, completedAt time.Time) error { b.mu.Lock() defer b.mu.Unlock() if b.closed { @@ -319,25 +309,13 @@ func (b *Buffer) RecordToolCompletion(key Key, dispatchIndex int, toolCallID str if episode.closed { return ErrEpisodeClosed } - if dispatchIndex >= 0 && dispatchIndex < len(episode.toolCompletions) { - entry := &episode.toolCompletions[dispatchIndex] - if entry.ToolCallID == toolCallID && entry.CompletedAt.IsZero() { - entry.CompletedAt = completedAt - return nil - } + if dispatchIndex < 0 || dispatchIndex >= len(episode.toolCompletions) { + return nil } - for i := range episode.toolCompletions { - entry := &episode.toolCompletions[i] - if entry.ToolCallID == toolCallID && entry.CompletedAt.IsZero() { - entry.CompletedAt = completedAt - return nil - } + entry := &episode.toolCompletions[dispatchIndex] + if entry.CompletedAt.IsZero() { + entry.CompletedAt = completedAt } - episode.toolCompletions = append(episode.toolCompletions, ToolCompletion{ - CallIndex: -1, - ToolCallID: toolCallID, - CompletedAt: completedAt, - }) return nil } @@ -439,8 +417,8 @@ func (b *Buffer) ToolBatchStartedAt(key Key) time.Time { } // ToolCompletions returns copied occurrence state. A zero start means not -// launched; a zero completion means unfinished. Interrupts must use -// CloseEpisodeForBilling to avoid a read-close race. +// launched; a zero completion means unfinished. Read it before CloseEpisode +// because closed episodes are garbage collected. func (b *Buffer) ToolCompletions(key Key) []ToolCompletion { b.mu.Lock() defer b.mu.Unlock() @@ -451,45 +429,6 @@ func (b *Buffer) ToolCompletions(key Key) []ToolCompletion { return slices.Clone(episode.toolCompletions) } -// EpisodeBilling is the billing snapshot captured when an episode closes. -type EpisodeBilling struct { - // ClosedAt is the first close instant, used as the retry-stable - // interrupt time. - ClosedAt time.Time - // ModelInvokedAt is the StartModelInvocation stamp, or zero if absent. - ModelInvokedAt time.Time - // ToolBatchStartedAt is the StartToolBatch stamp, or zero if absent. - ToolBatchStartedAt time.Time - // ToolCompletions is the close-time occurrence snapshot. - ToolCompletions []ToolCompletion -} - -// CloseEpisodeForBilling closes the episode and returns billing stamps from -// the same critical section, preventing a read-close race. Unknown episodes -// close blank. Re-closing returns the original snapshot and refreshes -// retention so retries keep the same billing state and buffered parts. -func (b *Buffer) CloseEpisodeForBilling(key Key) (EpisodeBilling, error) { - b.mu.Lock() - defer b.mu.Unlock() - if b.closed { - return EpisodeBilling{}, ErrMessagePartBufferClosed - } - episode := b.getOrCreateEpisodeLocked(key) - now := b.opts.Clock.Now("message-part-buffer", "close") - if episode.close(now) { - b.queueClosedEpisodeLocked(key, episode) - episode.notifySubscribers() - } else { - b.refreshClosedEpisodeEvictionLocked(key, episode, now) - } - return EpisodeBilling{ - ClosedAt: episode.closedAt, - ModelInvokedAt: episode.modelStartedAt, - ToolBatchStartedAt: episode.toolBatchStartedAt, - ToolCompletions: slices.Clone(episode.toolCompletions), - }, nil -} - // SubscribeToEpisode replays existing parts and streams new parts. // // Subscribers may attach before CreateEpisode is called. In that case the @@ -611,14 +550,6 @@ func (b *Buffer) queueClosedEpisodeLocked(key Key, episode *episodeState) { heap.Push(&b.closedEpisodes, item) } -// refreshClosedEpisodeEvictionLocked refreshes retention with a new heap -// item; cleanup skips superseded items by identity. -func (b *Buffer) refreshClosedEpisodeEvictionLocked(key Key, episode *episodeState, evictAt time.Time) { - item := &closedEpisodeItem{key: key, closedAt: evictAt} - episode.closedHeapItem = item - heap.Push(&b.closedEpisodes, item) -} - func (b *Buffer) getOrCreateEpisodeLocked(key Key) *episodeState { episode := b.episodes[key] if episode != nil { diff --git a/coderd/x/chatd/messagepartbuffer/message_part_buffer_test.go b/coderd/x/chatd/messagepartbuffer/message_part_buffer_test.go index d79dffe47f3..d0e7b281b74 100644 --- a/coderd/x/chatd/messagepartbuffer/message_part_buffer_test.go +++ b/coderd/x/chatd/messagepartbuffer/message_part_buffer_test.go @@ -188,7 +188,7 @@ func TestBuffer_ToolCompletions(t *testing.T) { key := testEpisodeKey() require.Nil(t, buffer.ToolCompletions(key), "unknown episode has no completions") - require.ErrorIs(t, buffer.RecordToolCompletion(key, 0, "call-1", clock.Now()), messagepartbuffer.ErrEpisodeNotFound) + require.ErrorIs(t, buffer.RecordToolCompletion(key, 0, clock.Now()), messagepartbuffer.ErrEpisodeNotFound) require.NoError(t, buffer.CreateEpisode(key)) require.Empty(t, buffer.ToolCompletions(key), "episode without a tool batch has no completions") @@ -204,7 +204,7 @@ func TestBuffer_ToolCompletions(t *testing.T) { }, buffer.ToolCompletions(key)) clock.Advance(time.Second) secondDupCompletedAt := clock.Now() - require.NoError(t, buffer.RecordToolCompletion(key, 2, "call-dup", secondDupCompletedAt)) + require.NoError(t, buffer.RecordToolCompletion(key, 2, secondDupCompletedAt)) require.Equal(t, []messagepartbuffer.ToolCompletion{ {CallIndex: 0, ToolCallID: "call-1"}, {CallIndex: 1, ToolCallID: "call-dup"}, @@ -212,7 +212,7 @@ func TestBuffer_ToolCompletions(t *testing.T) { }, buffer.ToolCompletions(key)) clock.Advance(2 * time.Second) firstDupCompletedAt := clock.Now() - require.NoError(t, buffer.RecordToolCompletion(key, 1, "call-dup", firstDupCompletedAt)) + require.NoError(t, buffer.RecordToolCompletion(key, 1, firstDupCompletedAt)) completions := buffer.ToolCompletions(key) require.Equal(t, []messagepartbuffer.ToolCompletion{ {CallIndex: 0, ToolCallID: "call-1"}, @@ -225,15 +225,11 @@ func TestBuffer_ToolCompletions(t *testing.T) { clock.Advance(time.Second) unseededAt := clock.Now() - require.NoError(t, buffer.RecordToolCompletion(key, 5, "call-unseeded", unseededAt)) - require.Equal(t, messagepartbuffer.ToolCompletion{ - CallIndex: -1, - ToolCallID: "call-unseeded", - CompletedAt: unseededAt, - }, buffer.ToolCompletions(key)[3]) + require.NoError(t, buffer.RecordToolCompletion(key, 5, unseededAt)) + require.Len(t, buffer.ToolCompletions(key), 3, "unknown dispatch index must be dropped") require.NoError(t, buffer.CloseEpisode(key)) - require.ErrorIs(t, buffer.RecordToolCompletion(key, 0, "call-1", clock.Now()), messagepartbuffer.ErrEpisodeClosed) + require.ErrorIs(t, buffer.RecordToolCompletion(key, 0, clock.Now()), messagepartbuffer.ErrEpisodeClosed) require.True(t, buffer.ToolCompletions(key)[0].CompletedAt.IsZero()) } @@ -245,7 +241,7 @@ func TestBuffer_RecordToolStart(t *testing.T) { defer buffer.Close() key := testEpisodeKey() - require.ErrorIs(t, buffer.RecordToolStart(key, 0, "call-1", clock.Now()), messagepartbuffer.ErrEpisodeNotFound) + require.ErrorIs(t, buffer.RecordToolStart(key, 0, clock.Now()), messagepartbuffer.ErrEpisodeNotFound) require.NoError(t, buffer.CreateEpisode(key)) require.NoError(t, buffer.StartToolBatch(key, []messagepartbuffer.DispatchedToolCall{ @@ -259,111 +255,32 @@ func TestBuffer_RecordToolStart(t *testing.T) { clock.Advance(time.Second) secondDupStartedAt := clock.Now() - require.NoError(t, buffer.RecordToolStart(key, 2, "call-dup", secondDupStartedAt)) + require.NoError(t, buffer.RecordToolStart(key, 2, secondDupStartedAt)) require.Equal(t, []messagepartbuffer.ToolCompletion{ {CallIndex: 0, ToolCallID: "call-1"}, {CallIndex: 1, ToolCallID: "call-dup"}, {CallIndex: 2, ToolCallID: "call-dup", StartedAt: secondDupStartedAt}, }, buffer.ToolCompletions(key)) + clock.Advance(time.Second) + require.NoError(t, buffer.RecordToolStart(key, 7, clock.Now())) + require.True(t, buffer.ToolCompletions(key)[1].StartedAt.IsZero(), "unknown dispatch index must be dropped") + clock.Advance(time.Second) firstDupStartedAt := clock.Now() - require.NoError(t, buffer.RecordToolStart(key, 7, "call-dup", firstDupStartedAt)) + require.NoError(t, buffer.RecordToolStart(key, 1, firstDupStartedAt)) require.Equal(t, firstDupStartedAt, buffer.ToolCompletions(key)[1].StartedAt) - require.NoError(t, buffer.RecordToolStart(key, 9, "call-unseeded", clock.Now())) + require.NoError(t, buffer.RecordToolStart(key, 9, clock.Now())) require.Len(t, buffer.ToolCompletions(key), 3) - billing, err := buffer.CloseEpisodeForBilling(key) - require.NoError(t, err) - require.Equal(t, secondDupStartedAt, billing.ToolCompletions[2].StartedAt) - require.ErrorIs(t, buffer.RecordToolStart(key, 0, "call-1", clock.Now()), messagepartbuffer.ErrEpisodeClosed) + completions := buffer.ToolCompletions(key) + require.NoError(t, buffer.CloseEpisode(key)) + require.Equal(t, secondDupStartedAt, completions[2].StartedAt) + require.ErrorIs(t, buffer.RecordToolStart(key, 0, clock.Now()), messagepartbuffer.ErrEpisodeClosed) require.True(t, buffer.ToolCompletions(key)[0].StartedAt.IsZero()) } -func TestBuffer_CloseEpisodeForBilling(t *testing.T) { - t.Parallel() - - clock := quartz.NewMock(t) - buffer := messagepartbuffer.New(messagepartbuffer.Options{Clock: clock}) - defer buffer.Close() - - unknown := testEpisodeKey() - billing, err := buffer.CloseEpisodeForBilling(unknown) - require.NoError(t, err) - require.Equal(t, clock.Now(), billing.ClosedAt) - require.Zero(t, billing.ModelInvokedAt) - require.Zero(t, billing.ToolBatchStartedAt) - require.Empty(t, billing.ToolCompletions) - - key := testEpisodeKey() - require.NoError(t, buffer.CreateEpisode(key)) - clock.Advance(time.Second) - batchStartedAt := clock.Now() - require.NoError(t, buffer.StartToolBatch(key, []messagepartbuffer.DispatchedToolCall{ - {CallIndex: 0, ToolCallID: "call-1"}, - {CallIndex: 1, ToolCallID: "call-2"}, - })) - clock.Advance(time.Second) - completedAt := clock.Now() - require.NoError(t, buffer.RecordToolCompletion(key, 0, "call-1", completedAt)) - clock.Advance(time.Second) - closedAt := clock.Now() - billing, err = buffer.CloseEpisodeForBilling(key) - require.NoError(t, err) - require.Equal(t, closedAt, billing.ClosedAt) - require.Zero(t, billing.ModelInvokedAt) - require.Equal(t, batchStartedAt, billing.ToolBatchStartedAt) - require.Equal(t, []messagepartbuffer.ToolCompletion{ - {CallIndex: 0, ToolCallID: "call-1", CompletedAt: completedAt}, - {CallIndex: 1, ToolCallID: "call-2"}, - }, billing.ToolCompletions) - require.ErrorIs(t, buffer.RecordToolCompletion(key, 1, "call-2", clock.Now()), messagepartbuffer.ErrEpisodeClosed) - - clock.Advance(time.Second) - again, err := buffer.CloseEpisodeForBilling(key) - require.NoError(t, err) - require.Equal(t, billing, again) -} - -func TestBuffer_CloseEpisodeForBillingRefreshesEviction(t *testing.T) { - t.Parallel() - - clock := quartz.NewMock(t) - buffer := messagepartbuffer.New(messagepartbuffer.Options{ - Clock: clock, - ClosedEpisodeRetention: time.Minute, - }) - defer buffer.Close() - ctx := testutil.Context(t, testutil.WaitShort) - - key := testEpisodeKey() - require.NoError(t, buffer.CreateEpisode(key)) - require.NoError(t, buffer.StartToolBatch(key, []messagepartbuffer.DispatchedToolCall{{CallIndex: 0, ToolCallID: "call-1"}})) - first, err := buffer.CloseEpisodeForBilling(key) - require.NoError(t, err) - - clock.Advance(45 * time.Second).MustWait(ctx) - again, err := buffer.CloseEpisodeForBilling(key) - require.NoError(t, err) - require.Equal(t, first, again) - - clock.Advance(15 * time.Second).MustWait(ctx) - again, err = buffer.CloseEpisodeForBilling(key) - require.NoError(t, err) - require.Equal(t, first, again) - - // GetParts later runs cleanup synchronously, avoiding a goroutine race. - clock.Advance(time.Minute).MustWait(ctx) - clock.Advance(time.Minute).MustWait(ctx) - _, err = buffer.GetParts(key) - require.ErrorIs(t, err, messagepartbuffer.ErrEpisodeNotFound) - expired, err := buffer.CloseEpisodeForBilling(key) - require.NoError(t, err) - require.Zero(t, expired.ToolBatchStartedAt) - require.Empty(t, expired.ToolCompletions) -} - func TestBuffer_SubscribeExistingReplaysThenStreamsLiveParts(t *testing.T) { t.Parallel() diff --git a/coderd/x/chatd/tasks.go b/coderd/x/chatd/tasks.go index f9e9b109b3c..91598a44b68 100644 --- a/coderd/x/chatd/tasks.go +++ b/coderd/x/chatd/tasks.go @@ -247,20 +247,32 @@ func (o chatWorkerOptions) retryOptions() retryWrapperOptions { type interruptEpisodeSnapshot struct { loaded bool key messagepartbuffer.Key - billing messagepartbuffer.EpisodeBilling + billing interruptEpisodeBilling parts []messagepartbuffer.Part } -// closeInterruptEpisode atomically closes the episode and snapshots billing. -// Unknown episodes close blank so interruption converges. -func (s *taskStarter) closeInterruptEpisode(ctx context.Context, key messagepartbuffer.Key) (messagepartbuffer.EpisodeBilling, []messagepartbuffer.Part, error) { - billing, err := s.opts.MessagePartBuffer.CloseEpisodeForBilling(key) - if err != nil { +type interruptEpisodeBilling struct { + interruptedAt time.Time + modelInvokedAt time.Time + toolBatchStartedAt time.Time + toolCompletions []messagepartbuffer.ToolCompletion +} + +// closeInterruptEpisode snapshots billing, closes the episode, and returns its +// buffered parts. Unknown episodes close blank so interruption converges. +func (s *taskStarter) closeInterruptEpisode(ctx context.Context, key messagepartbuffer.Key) (interruptEpisodeBilling, []messagepartbuffer.Part, error) { + billing := interruptEpisodeBilling{ + modelInvokedAt: s.opts.MessagePartBuffer.ModelInvokedAt(key), + toolBatchStartedAt: s.opts.MessagePartBuffer.ToolBatchStartedAt(key), + toolCompletions: s.opts.MessagePartBuffer.ToolCompletions(key), + } + if err := s.opts.MessagePartBuffer.CloseEpisode(key); err != nil { if ctx.Err() != nil { - return messagepartbuffer.EpisodeBilling{}, nil, errors.Join(errTaskExpectedExit, xerrors.Errorf("close message part episode: %w", err), ctx.Err()) + return interruptEpisodeBilling{}, nil, errors.Join(errTaskExpectedExit, xerrors.Errorf("close message part episode: %w", err), ctx.Err()) } - return messagepartbuffer.EpisodeBilling{}, nil, taskRetryableError{err: xerrors.Errorf("close message part episode: %w", err)} + return interruptEpisodeBilling{}, nil, taskRetryableError{err: xerrors.Errorf("close message part episode: %w", err)} } + billing.interruptedAt = s.opts.Clock.Now("chatworker", "interrupt") parts, err := s.opts.MessagePartBuffer.GetParts(key) if errors.Is(err, messagepartbuffer.ErrEpisodeNotFound) { parts = nil @@ -268,9 +280,9 @@ func (s *taskStarter) closeInterruptEpisode(ctx context.Context, key messagepart } if err != nil { if ctx.Err() != nil { - return messagepartbuffer.EpisodeBilling{}, nil, errors.Join(errTaskExpectedExit, xerrors.Errorf("get message part episode: %w", err), ctx.Err()) + return interruptEpisodeBilling{}, nil, errors.Join(errTaskExpectedExit, xerrors.Errorf("get message part episode: %w", err), ctx.Err()) } - return messagepartbuffer.EpisodeBilling{}, nil, taskRetryableError{err: xerrors.Errorf("get message part episode: %w", err)} + return interruptEpisodeBilling{}, nil, taskRetryableError{err: xerrors.Errorf("get message part episode: %w", err)} } return billing, parts, nil } @@ -313,7 +325,7 @@ func (s *taskStarter) StartInterrupt(ctx context.Context, input chatWorkerTaskSt HistoryVersion: input.HistoryVersion, GenerationAttempt: chat.GenerationAttempt, } - var episodeBilling messagepartbuffer.EpisodeBilling + var episodeBilling interruptEpisodeBilling var parts []messagepartbuffer.Part if snapshot != nil && snapshot.loaded && snapshot.key == key { // Reuse the snapshot because the buffer may have evicted the episode. @@ -331,11 +343,10 @@ func (s *taskStarter) StartInterrupt(ctx context.Context, input chatWorkerTaskSt snapshot.parts = parts } } - // The first close is a retry-stable interrupt instant. - interruptedAt := episodeBilling.ClosedAt + interruptedAt := episodeBilling.interruptedAt var attemptRuntime time.Duration - if !episodeBilling.ModelInvokedAt.IsZero() { - attemptRuntime = interruptedAt.Sub(episodeBilling.ModelInvokedAt) + if !episodeBilling.modelInvokedAt.IsZero() { + attemptRuntime = interruptedAt.Sub(episodeBilling.modelInvokedAt) } partialMessages, err := bufferedPartsToPartialMessages(bufferedPartsToPartialMessagesInput{ parts: parts, @@ -356,11 +367,11 @@ func (s *taskStarter) StartInterrupt(ctx context.Context, input chatWorkerTaskSt return xerrors.Errorf("load chat for task: %w", err) } messages := partialMessages - // Reuse the close instant so database delay and retries do not inflate - // billing. + // Reuse the captured interrupt instant so database delay and retries do + // not inflate billing. committedCancels, err := committedPendingLocalToolCancellationMessages(ctx, store, chat, interruptedAt, interruptedToolBatchBilling{ - batchStartedAt: episodeBilling.ToolBatchStartedAt, - toolCompletions: episodeBilling.ToolCompletions, + batchStartedAt: episodeBilling.toolBatchStartedAt, + toolCompletions: episodeBilling.toolCompletions, }) if err != nil { return xerrors.Errorf("committed pending local tool cancellation messages: %w", err) diff --git a/coderd/x/chatd/tasks_test.go b/coderd/x/chatd/tasks_test.go index 49d594ffce5..55dafe0663d 100644 --- a/coderd/x/chatd/tasks_test.go +++ b/coderd/x/chatd/tasks_test.go @@ -585,11 +585,11 @@ func TestInterruptTask_ToolBatchBillsPartialWindowOnCancellationRow(t *testing.T {CallIndex: 0, ToolCallID: execCallID}, {CallIndex: 1, ToolCallID: waitCallID}, })) - require.NoError(t, buffer.RecordToolStart(batch.key, 0, execCallID, batch.clock.Now())) - require.NoError(t, buffer.RecordToolStart(batch.key, 1, waitCallID, batch.clock.Now())) + require.NoError(t, buffer.RecordToolStart(batch.key, 0, batch.clock.Now())) + require.NoError(t, buffer.RecordToolStart(batch.key, 1, batch.clock.Now())) // Record a live completion without publishing a tool result. batch.clock.Advance(3 * time.Second) - require.NoError(t, buffer.RecordToolCompletion(batch.key, 0, execCallID, batch.clock.Now())) + require.NoError(t, buffer.RecordToolCompletion(batch.key, 0, batch.clock.Now())) batch.clock.Advance(5 * time.Second) snapshot := &interruptEpisodeSnapshot{} @@ -600,7 +600,7 @@ func TestInterruptTask_ToolBatchBillsPartialWindowOnCancellationRow(t *testing.T require.False(t, waitRow.RuntimeMs.Valid) require.True(t, snapshot.loaded) require.Equal(t, batch.key, snapshot.key) - require.Len(t, snapshot.billing.ToolCompletions, 2) + require.Len(t, snapshot.billing.toolCompletions, 2) } func TestInterruptTask_RetrySnapshotOutlivesEpisodeEviction(t *testing.T) { @@ -619,10 +619,10 @@ func TestInterruptTask_RetrySnapshotOutlivesEpisodeEviction(t *testing.T) { snapshot := &interruptEpisodeSnapshot{ loaded: true, key: batch.key, - billing: messagepartbuffer.EpisodeBilling{ - ClosedAt: batch.clock.Now(), - ToolBatchStartedAt: batchStartedAt, - ToolCompletions: []messagepartbuffer.ToolCompletion{{CallIndex: 0, ToolCallID: execCallID, StartedAt: batchStartedAt}}, + billing: interruptEpisodeBilling{ + interruptedAt: batch.clock.Now(), + toolBatchStartedAt: batchStartedAt, + toolCompletions: []messagepartbuffer.ToolCompletion{{CallIndex: 0, ToolCallID: execCallID, StartedAt: batchStartedAt}}, }, } @@ -643,7 +643,7 @@ func TestInterruptTask_SnapshotCapturedBeforeChatRead(t *testing.T) { batch.clock.Advance(2 * time.Second) require.NoError(t, buffer.StartToolBatch(batch.key, []messagepartbuffer.DispatchedToolCall{{CallIndex: 0, ToolCallID: execCallID}})) - require.NoError(t, buffer.RecordToolStart(batch.key, 0, execCallID, batch.clock.Now())) + require.NoError(t, buffer.RecordToolStart(batch.key, 0, batch.clock.Now())) batch.clock.Advance(3 * time.Second) snapshot := &interruptEpisodeSnapshot{} @@ -682,7 +682,7 @@ func TestInterruptTask_RunningBilledToolBillsUpToInterrupt(t *testing.T) { batch.clock.Advance(2 * time.Second) require.NoError(t, batch.starter.opts.MessagePartBuffer.StartToolBatch(batch.key, []messagepartbuffer.DispatchedToolCall{{CallIndex: 0, ToolCallID: execCallID}})) - require.NoError(t, batch.starter.opts.MessagePartBuffer.RecordToolStart(batch.key, 0, execCallID, batch.clock.Now())) + require.NoError(t, batch.starter.opts.MessagePartBuffer.RecordToolStart(batch.key, 0, batch.clock.Now())) batch.clock.Advance(7 * time.Second) messages := batch.interrupt(t, f) @@ -701,7 +701,7 @@ func TestInterruptTask_UnbilledOnlyBatchBillsNothingOnInterrupt(t *testing.T) { batch.clock.Advance(2 * time.Second) require.NoError(t, batch.starter.opts.MessagePartBuffer.StartToolBatch(batch.key, []messagepartbuffer.DispatchedToolCall{{CallIndex: 0, ToolCallID: waitCallID}})) - require.NoError(t, batch.starter.opts.MessagePartBuffer.RecordToolStart(batch.key, 0, waitCallID, batch.clock.Now())) + require.NoError(t, batch.starter.opts.MessagePartBuffer.RecordToolStart(batch.key, 0, batch.clock.Now())) batch.clock.Advance(10 * time.Second) messages := batch.interrupt(t, f) @@ -726,7 +726,7 @@ func TestInterruptTask_UnstartedSerialCallBillsNothingOnInterrupt(t *testing.T) {CallIndex: 0, ToolCallID: waitCallID}, {CallIndex: 1, ToolCallID: serialCallID}, })) - require.NoError(t, buffer.RecordToolStart(batch.key, 0, waitCallID, batch.clock.Now())) + require.NoError(t, buffer.RecordToolStart(batch.key, 0, batch.clock.Now())) batch.clock.Advance(10 * time.Second) messages := batch.interrupt(t, f) @@ -758,13 +758,13 @@ func TestInterruptTask_StartedSerialCallBillsFromItsOwnStart(t *testing.T) { {CallIndex: 1, ToolCallID: waitCallID}, {CallIndex: 2, ToolCallID: serialCallID}, })) - require.NoError(t, buffer.RecordToolStart(batch.key, 0, execCallID, batch.clock.Now())) - require.NoError(t, buffer.RecordToolStart(batch.key, 1, waitCallID, batch.clock.Now())) + require.NoError(t, buffer.RecordToolStart(batch.key, 0, batch.clock.Now())) + require.NoError(t, buffer.RecordToolStart(batch.key, 1, batch.clock.Now())) batch.clock.Advance(3 * time.Second) - require.NoError(t, buffer.RecordToolCompletion(batch.key, 0, execCallID, batch.clock.Now())) + require.NoError(t, buffer.RecordToolCompletion(batch.key, 0, batch.clock.Now())) batch.clock.Advance(3 * time.Second) - require.NoError(t, buffer.RecordToolCompletion(batch.key, 1, waitCallID, batch.clock.Now())) - require.NoError(t, buffer.RecordToolStart(batch.key, 2, serialCallID, batch.clock.Now())) + require.NoError(t, buffer.RecordToolCompletion(batch.key, 1, batch.clock.Now())) + require.NoError(t, buffer.RecordToolStart(batch.key, 2, batch.clock.Now())) batch.clock.Advance(2 * time.Second) messages := batch.interrupt(t, f) @@ -792,10 +792,10 @@ func TestInterruptTask_DuplicateCallIDsKeepOccurrenceStates(t *testing.T) { {CallIndex: 0, ToolCallID: dupCallID}, {CallIndex: 1, ToolCallID: dupCallID}, })) - require.NoError(t, buffer.RecordToolStart(batch.key, 0, dupCallID, batch.clock.Now())) - require.NoError(t, buffer.RecordToolStart(batch.key, 1, dupCallID, batch.clock.Now())) + require.NoError(t, buffer.RecordToolStart(batch.key, 0, batch.clock.Now())) + require.NoError(t, buffer.RecordToolStart(batch.key, 1, batch.clock.Now())) batch.clock.Advance(3 * time.Second) - require.NoError(t, buffer.RecordToolCompletion(batch.key, 0, dupCallID, batch.clock.Now())) + require.NoError(t, buffer.RecordToolCompletion(batch.key, 0, batch.clock.Now())) batch.clock.Advance(3 * time.Second) messages := batch.interrupt(t, f) @@ -824,10 +824,10 @@ func TestInterruptTask_DuplicateCallIDCompletionStampsOwnOccurrence(t *testing.T {CallIndex: 0, ToolCallID: dupCallID}, {CallIndex: 1, ToolCallID: dupCallID}, })) - require.NoError(t, buffer.RecordToolStart(batch.key, 0, dupCallID, batch.clock.Now())) - require.NoError(t, buffer.RecordToolStart(batch.key, 1, dupCallID, batch.clock.Now())) + require.NoError(t, buffer.RecordToolStart(batch.key, 0, batch.clock.Now())) + require.NoError(t, buffer.RecordToolStart(batch.key, 1, batch.clock.Now())) batch.clock.Advance(3 * time.Second) - require.NoError(t, buffer.RecordToolCompletion(batch.key, 1, dupCallID, batch.clock.Now())) + require.NoError(t, buffer.RecordToolCompletion(batch.key, 1, batch.clock.Now())) batch.clock.Advance(3 * time.Second) messages := batch.interrupt(t, f) @@ -852,7 +852,7 @@ func TestInterruptTask_RejectedDuplicateIDDoesNotStealDispatchedOccurrence(t *te batch.clock.Advance(2 * time.Second) require.NoError(t, batch.starter.opts.MessagePartBuffer.StartToolBatch(batch.key, []messagepartbuffer.DispatchedToolCall{{CallIndex: 1, ToolCallID: dupCallID}})) - require.NoError(t, batch.starter.opts.MessagePartBuffer.RecordToolStart(batch.key, 0, dupCallID, batch.clock.Now())) + require.NoError(t, batch.starter.opts.MessagePartBuffer.RecordToolStart(batch.key, 0, batch.clock.Now())) batch.clock.Advance(10 * time.Second) messages := batch.interrupt(t, f) @@ -876,7 +876,7 @@ func TestInterruptTask_RejectedCallBillsNothingOnInterrupt(t *testing.T) { batch.clock.Advance(2 * time.Second) require.NoError(t, batch.starter.opts.MessagePartBuffer.StartToolBatch(batch.key, []messagepartbuffer.DispatchedToolCall{{CallIndex: 1, ToolCallID: waitCallID}})) - require.NoError(t, batch.starter.opts.MessagePartBuffer.RecordToolStart(batch.key, 0, waitCallID, batch.clock.Now())) + require.NoError(t, batch.starter.opts.MessagePartBuffer.RecordToolStart(batch.key, 0, batch.clock.Now())) batch.clock.Advance(10 * time.Second) messages := batch.interrupt(t, f) From 82dcbcf13ec1fab39187d08dbd4cd645a26b663c Mon Sep 17 00:00:00 2001 From: Jaayden Halko Date: Thu, 20 Aug 2026 07:16:55 +0000 Subject: [PATCH 19/20] chore: simplifications --- coderd/x/chatd/ARCHITECTURE.md | 5 +- coderd/x/chatd/chatloop/chatloop.go | 15 +- coderd/x/chatd/chatloop/runtime_test.go | 45 +++--- coderd/x/chatd/generation.go | 48 +++---- .../messagepartbuffer/message_part_buffer.go | 106 +++++--------- .../message_part_buffer_test.go | 134 +++++------------- coderd/x/chatd/tasks.go | 46 +++--- coderd/x/chatd/tasks_test.go | 37 +---- 8 files changed, 144 insertions(+), 292 deletions(-) diff --git a/coderd/x/chatd/ARCHITECTURE.md b/coderd/x/chatd/ARCHITECTURE.md index f961e65f569..0b68adeea4e 100644 --- a/coderd/x/chatd/ARCHITECTURE.md +++ b/coderd/x/chatd/ARCHITECTURE.md @@ -725,9 +725,8 @@ The buffer exposes the following API: - `GetParts(chat_id, history_version, generation_attempt)`: returns the message parts for an episode. Returns a predefined error if the episode is not found. - `StartModelInvocation(chat_id, history_version, generation_attempt)`: stamps the instant the episode opens its provider stream. Returns a predefined error if the episode is not found or already closed. Episodes that never invoke a model, such as local tool execution batches, are never stamped. - `ModelInvokedAt(chat_id, history_version, generation_attempt)`: returns the `StartModelInvocation` stamp, or zero if none exists. Interrupt handling reads it before closing the episode. -- `StartToolBatch(chat_id, history_version, generation_attempt, calls)`: stamps the batch start and seeds one entry per dispatched call occurrence. -- `RecordToolStart(chat_id, history_version, generation_attempt, dispatch_index, started_at)` and `RecordToolCompletion(chat_id, history_version, generation_attempt, dispatch_index, completed_at)`: stamp when each call occurrence starts and finishes. An unstarted occurrence represents a queued serial call and does not bill on interrupt. -- `ToolBatchStartedAt(chat_id, history_version, generation_attempt)` and `ToolCompletions(chat_id, history_version, generation_attempt)`: return the batch start and per-occurrence execution stamps. Interrupt handling reads them before closing the episode. +- `RecordToolStart(chat_id, history_version, generation_attempt, call_index, started_at)` and `RecordToolCompletion(chat_id, history_version, generation_attempt, call_index, completed_at)`: record when each call occurrence starts and finishes. Calls are keyed by unresolved-call position so rejected gaps and duplicate IDs remain distinct. Queued and undispatched calls are absent because neither has started. +- `ToolCompletions(chat_id, history_version, generation_attempt)`: returns the started per-occurrence execution stamps. Interrupt handling reads them before closing the episode. - `SubscribeToEpisode(chat_id, history_version, generation_attempt)`: returns a go channel that will receive all message parts for the episode. It spawns a goroutine that delivers parts to the channel. It's live until the episode is closed or until a subscriber requests that the channel be closed. Once the goroutine delivers all message parts for a closed episode, it closes the channel and exits. If the episode is already closed at the time of the call, the goroutine delivers all message parts for the episode, closes the channel, and exits. `SubscribeToEpisode` does not return an error if the episode is not found: it waits for it to be created instead. Closed episodes are garbage collected after at least 15 seconds since they were closed and when they have no active subscribers. The message part buffer maintains a garbage collection goroutine. diff --git a/coderd/x/chatd/chatloop/chatloop.go b/coderd/x/chatd/chatloop/chatloop.go index 5f1fd04d8e6..0a679361773 100644 --- a/coderd/x/chatd/chatloop/chatloop.go +++ b/coderd/x/chatd/chatloop/chatloop.go @@ -266,13 +266,10 @@ type ExecuteLocalToolsOptions struct { // UnbilledToolNames lists called tool names excluded from the batch // window. Include deprecated aliases. UnbilledToolNames map[string]bool - // OnBatchStart fires immediately before dispatch. Interrupt billing uses - // it to avoid charging pre-dispatch cancellations. - OnBatchStart func() // OnToolStart fires when each local call begins. Serial calls may start // after concurrent siblings settle, so interrupts bill actual starts and - // skip dispatched calls that never run. dispatchIndex identifies the - // dispatch-order occurrence. + // 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; @@ -599,11 +596,8 @@ func ExecuteLocalTools(ctx context.Context, opts ExecuteLocalToolsOptions) (Tool ) if exclusiveViolation { now := clockNow(opts.Clock) - for i, tr := range policyResults { + for _, tr := range policyResults { recordToolResultTimestamp(&result, tr.ToolCallID, now) - if opts.OnToolComplete != nil { - opts.OnToolComplete(i, now) - } publishToolAttachments(ctx, opts.Logger, tr, now, publishMessagePart) ssePart := chatprompt.PartFromContentWithLogger(ctx, opts.Logger, tr) ssePart.CreatedAt = &now @@ -620,9 +614,6 @@ func ExecuteLocalTools(ctx context.Context, opts ExecuteLocalToolsOptions) (Tool } maxResultBytes := toolResultByteBudget(opts.ContextLimit) - if opts.OnBatchStart != nil { - opts.OnBatchStart() - } batchStart := clockNow(opts.Clock) // Keep completions by occurrence so duplicate IDs cannot collapse them. orderedCompletions := make([]time.Time, 0, len(localCalls)) diff --git a/coderd/x/chatd/chatloop/runtime_test.go b/coderd/x/chatd/chatloop/runtime_test.go index 625234ea97c..17921cd6620 100644 --- a/coderd/x/chatd/chatloop/runtime_test.go +++ b/coderd/x/chatd/chatloop/runtime_test.go @@ -406,27 +406,33 @@ func TestExecuteLocalTools_DuplicateToolCallIDsKeepOccurrenceCompletions(t *test require.Equal(t, "call-dup", outcome.BatchRuntimeToolCallID) } -func TestExecuteLocalTools_OnBatchStartFiresOnlyOnDispatch(t *testing.T) { +func TestExecuteLocalTools_ExecutionCallbacksFireOnlyForRuns(t *testing.T) { t.Parallel() - t.Run("dispatched batch seeds before tools run", func(t *testing.T) { + t.Run("started call records a paired lifecycle", func(t *testing.T) { t.Parallel() starts := 0 - seededWhenToolRan := false + completions := 0 + startedWhenToolRan := false tool := fantasy.NewAgentTool( "fast_tool", "test tool", func(context.Context, struct{}, fantasy.ToolCall) (fantasy.ToolResponse, error) { - seededWhenToolRan = starts > 0 + 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"}, - OnBatchStart: func() { starts++ }, + 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: "{}"}, }, @@ -434,34 +440,40 @@ func TestExecuteLocalTools_OnBatchStartFiresOnlyOnDispatch(t *testing.T) { require.NoError(t, err) require.Len(t, outcome.Step.Content, 1) require.Equal(t, 1, starts) - require.True(t, seededWhenToolRan, "the batch must be seeded before any tool goroutine runs") + require.Equal(t, 1, completions) + require.True(t, startedWhenToolRan, "the start callback must fire before the tool runs") }) - t.Run("canceled context never seeds", func(t *testing.T) { + 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), - OnBatchStart: func() { started = true }, + 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, "a canceled batch dispatches nothing and must not seed billing state") + require.False(t, started) + require.False(t, completed) }) - t.Run("exclusive violation never seeds", func(t *testing.T) { + 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}, - OnBatchStart: func() { started = 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: "{}"}, @@ -469,7 +481,8 @@ func TestExecuteLocalTools_OnBatchStartFiresOnlyOnDispatch(t *testing.T) { }) require.NoError(t, err) require.Len(t, outcome.Step.Content, 2, "the whole batch resolves to synthesized policy errors") - require.False(t, started, "a policy-rejected batch dispatches nothing and must not seed billing state") + require.False(t, started) + require.False(t, completed) }) } diff --git a/coderd/x/chatd/generation.go b/coderd/x/chatd/generation.go index 1ef11b03f11..5ced8e7e298 100644 --- a/coderd/x/chatd/generation.go +++ b/coderd/x/chatd/generation.go @@ -865,19 +865,22 @@ func (s *taskStarter) executeLocalTools( var outcome chatloop.ToolExecutionOutcome var spawnDispatchErr error if len(allowed) > 0 { - // Seed only calls that will dispatch, after policy checks. Positional - // indexes skip rejected calls and keep duplicate IDs distinct. - var onBatchStart func() + var onToolStart func(int, time.Time) + var onToolComplete func(int, time.Time) if !exclusiveRejected { - dispatched := make([]messagepartbuffer.DispatchedToolCall, 0, len(allowed)) - for j, tc := range allowed { - dispatched = append(dispatched, messagepartbuffer.DispatchedToolCall{ - CallIndex: allowedIndexes[j], - ToolCallID: tc.ToolCallID, - }) + // Translate dispatch-order callbacks back to unresolved-call + // positions so rejected gaps and duplicate IDs remain distinct. + onToolStart = func(dispatchIndex int, startedAt time.Time) { + if dispatchIndex < 0 || dispatchIndex >= len(allowedIndexes) { + return + } + attempt.recordToolStart(allowedIndexes[dispatchIndex], startedAt) } - onBatchStart = func() { - attempt.startToolBatch(dispatched) + onToolComplete = func(dispatchIndex int, completedAt time.Time) { + if dispatchIndex < 0 || dispatchIndex >= len(allowedIndexes) { + return + } + attempt.recordToolCompletion(allowedIndexes[dispatchIndex], completedAt) } } outcome, err = chatloop.ExecuteLocalTools(ctx, chatloop.ExecuteLocalToolsOptions{ @@ -894,9 +897,8 @@ func (s *taskStarter) executeLocalTools( ContextLimit: prepared.ContextLimitFallback, ToolNameAliases: subagentToolNameAliases, UnbilledToolNames: unbilledSubagentToolNames, - OnBatchStart: onBatchStart, - OnToolStart: attempt.recordToolStart, - OnToolComplete: attempt.recordToolCompletion, + OnToolStart: onToolStart, + OnToolComplete: onToolComplete, PublishMessagePart: attempt.publish, Logger: s.opts.Logger, Metrics: s.server.metrics, @@ -1112,15 +1114,12 @@ type generationAttempt struct { // can bill the window the step would have reported. It is always // non-nil when beginGenerationAttempt succeeds. startModelInvocation func() - // startToolBatch stamps dispatch and seeds tool-call occurrences. It is - // always non-nil after beginGenerationAttempt. - startToolBatch func(calls []messagepartbuffer.DispatchedToolCall) // recordToolStart stamps an occurrence's actual start; serial calls may // start after dispatch. It is always non-nil after beginGenerationAttempt. - recordToolStart func(dispatchIndex int, startedAt time.Time) + recordToolStart func(callIndex int, startedAt time.Time) // recordToolCompletion stamps an occurrence's completion. It is always // non-nil after beginGenerationAttempt. - recordToolCompletion func(dispatchIndex int, completedAt time.Time) + recordToolCompletion func(callIndex int, completedAt time.Time) // closeEpisode closes the attempt's buffer episode. It is always // non-nil when beginGenerationAttempt succeeds. closeEpisode func() @@ -1167,14 +1166,11 @@ func (s *taskStarter) beginGenerationAttempt( startModelInvocation: func() { _ = s.opts.MessagePartBuffer.StartModelInvocation(key) }, - startToolBatch: func(calls []messagepartbuffer.DispatchedToolCall) { - _ = s.opts.MessagePartBuffer.StartToolBatch(key, calls) - }, - recordToolStart: func(dispatchIndex int, startedAt time.Time) { - _ = s.opts.MessagePartBuffer.RecordToolStart(key, dispatchIndex, startedAt) + recordToolStart: func(callIndex int, startedAt time.Time) { + _ = s.opts.MessagePartBuffer.RecordToolStart(key, callIndex, startedAt) }, - recordToolCompletion: func(dispatchIndex int, completedAt time.Time) { - _ = s.opts.MessagePartBuffer.RecordToolCompletion(key, dispatchIndex, completedAt) + recordToolCompletion: func(callIndex int, completedAt time.Time) { + _ = s.opts.MessagePartBuffer.RecordToolCompletion(key, callIndex, completedAt) }, closeEpisode: func() { _ = s.opts.MessagePartBuffer.CloseEpisode(key) diff --git a/coderd/x/chatd/messagepartbuffer/message_part_buffer.go b/coderd/x/chatd/messagepartbuffer/message_part_buffer.go index f13cc4961d3..39da1751482 100644 --- a/coderd/x/chatd/messagepartbuffer/message_part_buffer.go +++ b/coderd/x/chatd/messagepartbuffer/message_part_buffer.go @@ -101,11 +101,9 @@ type episodeState struct { // episode's provider stream is opened. It is zero for episodes // that never invoke a model, such as local tool execution // batches. - modelStartedAt time.Time - toolBatchStartedAt time.Time - // toolCompletions stores per-occurrence start and completion stamps for - // interrupts. Positional storage preserves queued serial calls and - // duplicate IDs. + modelStartedAt time.Time + // toolCompletions stores started occurrences for interrupts. CallIndex + // distinguishes rejected-call gaps and duplicate IDs. toolCompletions []ToolCompletion closed bool closedAt time.Time @@ -223,55 +221,19 @@ func (b *Buffer) StartModelInvocation(key Key) error { return nil } -// DispatchedToolCall identifies a dispatched tool-call occurrence. -type DispatchedToolCall struct { - // CallIndex is its position in the unresolved call order, used to - // distinguish rejected and duplicate-ID calls. - CallIndex int - ToolCallID string -} - -// ToolCompletion tracks a tool-call occurrence. StartedAt is zero when the -// start is unknown; CompletedAt is zero while unfinished. +// ToolCompletion tracks a started tool-call occurrence. CompletedAt is zero +// while the call is unfinished. type ToolCompletion struct { // CallIndex is the occurrence's position in the unresolved call order. CallIndex int - ToolCallID string StartedAt time.Time CompletedAt time.Time } -// StartToolBatch stamps dispatch and seeds one entry per occurrence. A zero -// start is queued; a started entry with zero completion is running. Absent -// calls were not dispatched, and duplicate IDs remain distinct. -func (b *Buffer) StartToolBatch(key Key, calls []DispatchedToolCall) error { - b.mu.Lock() - defer b.mu.Unlock() - if b.closed { - return ErrMessagePartBufferClosed - } - episode, err := b.getEpisodeLocked(key) - if err != nil { - return err - } - if episode.closed { - return ErrEpisodeClosed - } - episode.toolBatchStartedAt = b.opts.Clock.Now("message-part-buffer", "tool-batch-start") - episode.toolCompletions = make([]ToolCompletion, 0, len(calls)) - for _, call := range calls { - episode.toolCompletions = append(episode.toolCompletions, ToolCompletion{ - CallIndex: call.CallIndex, - ToolCallID: call.ToolCallID, - }) - } - return nil -} - -// RecordToolStart stamps a dispatched call's actual start. Concurrent calls -// start with the batch; serial calls may start later. Unknown dispatch indexes -// are dropped because they cannot correlate to unresolved calls. -func (b *Buffer) RecordToolStart(key Key, dispatchIndex int, startedAt time.Time) error { +// RecordToolStart adds a tool-call occurrence when it begins execution. +// CallIndex preserves rejected-call gaps and keeps duplicate IDs distinct. +// Repeated starts keep the first timestamp. +func (b *Buffer) RecordToolStart(key Key, callIndex int, startedAt time.Time) error { b.mu.Lock() defer b.mu.Unlock() if b.closed { @@ -284,19 +246,24 @@ func (b *Buffer) RecordToolStart(key Key, dispatchIndex int, startedAt time.Time if episode.closed { return ErrEpisodeClosed } - if dispatchIndex < 0 || dispatchIndex >= len(episode.toolCompletions) { + if callIndex < 0 { return nil } - entry := &episode.toolCompletions[dispatchIndex] - if entry.StartedAt.IsZero() { - entry.StartedAt = startedAt + for _, entry := range episode.toolCompletions { + if entry.CallIndex == callIndex { + return nil + } } + episode.toolCompletions = append(episode.toolCompletions, ToolCompletion{ + CallIndex: callIndex, + StartedAt: startedAt, + }) return nil } // RecordToolCompletion stamps a call as it finishes, so interrupts use the -// actual completion. Unknown dispatch indexes are dropped. -func (b *Buffer) RecordToolCompletion(key Key, dispatchIndex int, completedAt time.Time) error { +// actual completion. Calls without a recorded start are dropped. +func (b *Buffer) RecordToolCompletion(key Key, callIndex int, completedAt time.Time) error { b.mu.Lock() defer b.mu.Unlock() if b.closed { @@ -309,12 +276,15 @@ func (b *Buffer) RecordToolCompletion(key Key, dispatchIndex int, completedAt ti if episode.closed { return ErrEpisodeClosed } - if dispatchIndex < 0 || dispatchIndex >= len(episode.toolCompletions) { - return nil - } - entry := &episode.toolCompletions[dispatchIndex] - if entry.CompletedAt.IsZero() { - entry.CompletedAt = completedAt + for i := range episode.toolCompletions { + entry := &episode.toolCompletions[i] + if entry.CallIndex != callIndex { + continue + } + if entry.CompletedAt.IsZero() { + entry.CompletedAt = completedAt + } + break } return nil } @@ -404,21 +374,9 @@ func (b *Buffer) ModelInvokedAt(key Key) time.Time { return episode.modelStartedAt } -// ToolBatchStartedAt returns the StartToolBatch stamp, or zero if absent. -// Read it before CloseEpisode because closed episodes are garbage collected. -func (b *Buffer) ToolBatchStartedAt(key Key) time.Time { - b.mu.Lock() - defer b.mu.Unlock() - episode := b.episodes[key] - if episode == nil { - return time.Time{} - } - return episode.toolBatchStartedAt -} - -// ToolCompletions returns copied occurrence state. A zero start means not -// launched; a zero completion means unfinished. Read it before CloseEpisode -// because closed episodes are garbage collected. +// ToolCompletions returns copied started-occurrence state. A zero completion +// means unfinished. Read it before CloseEpisode because closed episodes are +// garbage collected. func (b *Buffer) ToolCompletions(key Key) []ToolCompletion { b.mu.Lock() defer b.mu.Unlock() diff --git a/coderd/x/chatd/messagepartbuffer/message_part_buffer_test.go b/coderd/x/chatd/messagepartbuffer/message_part_buffer_test.go index d0e7b281b74..3f341ea2f20 100644 --- a/coderd/x/chatd/messagepartbuffer/message_part_buffer_test.go +++ b/coderd/x/chatd/messagepartbuffer/message_part_buffer_test.go @@ -148,37 +148,6 @@ func TestBuffer_ModelInvokedAt(t *testing.T) { require.Zero(t, buffer.ModelInvokedAt(implicit)) } -func TestBuffer_ToolBatchStartedAt(t *testing.T) { - t.Parallel() - - clock := quartz.NewMock(t) - buffer := messagepartbuffer.New(messagepartbuffer.Options{Clock: clock}) - defer buffer.Close() - - key := testEpisodeKey() - require.Zero(t, buffer.ToolBatchStartedAt(key), "unknown episode has no batch stamp") - require.ErrorIs(t, buffer.StartToolBatch(key, []messagepartbuffer.DispatchedToolCall{{CallIndex: 0, ToolCallID: "call-1"}}), messagepartbuffer.ErrEpisodeNotFound) - - require.NoError(t, buffer.CreateEpisode(key)) - require.Zero(t, buffer.ToolBatchStartedAt(key), "episode without a tool batch has no batch stamp") - clock.Advance(time.Second) - require.NoError(t, buffer.StartToolBatch(key, []messagepartbuffer.DispatchedToolCall{{CallIndex: 0, ToolCallID: "call-1"}})) - startedAt := buffer.ToolBatchStartedAt(key) - require.Equal(t, clock.Now(), startedAt) - - clock.Advance(1500 * time.Millisecond) - require.NoError(t, buffer.CloseEpisode(key)) - require.ErrorIs(t, buffer.StartToolBatch(key, []messagepartbuffer.DispatchedToolCall{{CallIndex: 0, ToolCallID: "call-1"}}), messagepartbuffer.ErrEpisodeClosed) - require.Equal(t, startedAt, buffer.ToolBatchStartedAt(key)) - - modelOnly := testEpisodeKey() - require.NoError(t, buffer.CreateEpisode(modelOnly)) - require.NoError(t, buffer.StartModelInvocation(modelOnly)) - clock.Advance(time.Second) - require.NoError(t, buffer.CloseEpisode(modelOnly)) - require.Zero(t, buffer.ToolBatchStartedAt(modelOnly)) -} - func TestBuffer_ToolCompletions(t *testing.T) { t.Parallel() @@ -188,97 +157,60 @@ func TestBuffer_ToolCompletions(t *testing.T) { key := testEpisodeKey() require.Nil(t, buffer.ToolCompletions(key), "unknown episode has no completions") + require.ErrorIs(t, buffer.RecordToolStart(key, 0, clock.Now()), messagepartbuffer.ErrEpisodeNotFound) require.ErrorIs(t, buffer.RecordToolCompletion(key, 0, clock.Now()), messagepartbuffer.ErrEpisodeNotFound) require.NoError(t, buffer.CreateEpisode(key)) - require.Empty(t, buffer.ToolCompletions(key), "episode without a tool batch has no completions") - require.NoError(t, buffer.StartToolBatch(key, []messagepartbuffer.DispatchedToolCall{ - {CallIndex: 0, ToolCallID: "call-1"}, - {CallIndex: 1, ToolCallID: "call-dup"}, - {CallIndex: 2, ToolCallID: "call-dup"}, - })) + require.Empty(t, buffer.ToolCompletions(key), "episode without started tools has no completions") + require.NoError(t, buffer.RecordToolCompletion(key, 2, clock.Now())) + require.Empty(t, buffer.ToolCompletions(key), "completion without a start must be dropped") + + clock.Advance(time.Second) + secondStartedAt := clock.Now() + require.NoError(t, buffer.RecordToolStart(key, 2, secondStartedAt)) + clock.Advance(time.Second) + firstStartedAt := clock.Now() + require.NoError(t, buffer.RecordToolStart(key, 1, firstStartedAt)) require.Equal(t, []messagepartbuffer.ToolCompletion{ - {CallIndex: 0, ToolCallID: "call-1"}, - {CallIndex: 1, ToolCallID: "call-dup"}, - {CallIndex: 2, ToolCallID: "call-dup"}, + {CallIndex: 2, StartedAt: secondStartedAt}, + {CallIndex: 1, StartedAt: firstStartedAt}, }, buffer.ToolCompletions(key)) + + clock.Advance(time.Second) + require.NoError(t, buffer.RecordToolStart(key, 2, clock.Now())) + require.NoError(t, buffer.RecordToolStart(key, -1, clock.Now())) + require.Equal(t, []messagepartbuffer.ToolCompletion{ + {CallIndex: 2, StartedAt: secondStartedAt}, + {CallIndex: 1, StartedAt: firstStartedAt}, + }, buffer.ToolCompletions(key), "repeated and invalid starts must not replace or append occurrences") + clock.Advance(time.Second) - secondDupCompletedAt := clock.Now() - require.NoError(t, buffer.RecordToolCompletion(key, 2, secondDupCompletedAt)) + secondCompletedAt := clock.Now() + require.NoError(t, buffer.RecordToolCompletion(key, 2, secondCompletedAt)) require.Equal(t, []messagepartbuffer.ToolCompletion{ - {CallIndex: 0, ToolCallID: "call-1"}, - {CallIndex: 1, ToolCallID: "call-dup"}, - {CallIndex: 2, ToolCallID: "call-dup", CompletedAt: secondDupCompletedAt}, + {CallIndex: 2, StartedAt: secondStartedAt, CompletedAt: secondCompletedAt}, + {CallIndex: 1, StartedAt: firstStartedAt}, }, buffer.ToolCompletions(key)) clock.Advance(2 * time.Second) - firstDupCompletedAt := clock.Now() - require.NoError(t, buffer.RecordToolCompletion(key, 1, firstDupCompletedAt)) + firstCompletedAt := clock.Now() + require.NoError(t, buffer.RecordToolCompletion(key, 1, firstCompletedAt)) completions := buffer.ToolCompletions(key) require.Equal(t, []messagepartbuffer.ToolCompletion{ - {CallIndex: 0, ToolCallID: "call-1"}, - {CallIndex: 1, ToolCallID: "call-dup", CompletedAt: firstDupCompletedAt}, - {CallIndex: 2, ToolCallID: "call-dup", CompletedAt: secondDupCompletedAt}, + {CallIndex: 2, StartedAt: secondStartedAt, CompletedAt: secondCompletedAt}, + {CallIndex: 1, StartedAt: firstStartedAt, CompletedAt: firstCompletedAt}, }, completions) completions[0].CompletedAt = clock.Now() - require.True(t, buffer.ToolCompletions(key)[0].CompletedAt.IsZero()) + require.Equal(t, secondCompletedAt, buffer.ToolCompletions(key)[0].CompletedAt) clock.Advance(time.Second) unseededAt := clock.Now() require.NoError(t, buffer.RecordToolCompletion(key, 5, unseededAt)) - require.Len(t, buffer.ToolCompletions(key), 3, "unknown dispatch index must be dropped") - - require.NoError(t, buffer.CloseEpisode(key)) - require.ErrorIs(t, buffer.RecordToolCompletion(key, 0, clock.Now()), messagepartbuffer.ErrEpisodeClosed) - require.True(t, buffer.ToolCompletions(key)[0].CompletedAt.IsZero()) -} - -func TestBuffer_RecordToolStart(t *testing.T) { - t.Parallel() - - clock := quartz.NewMock(t) - buffer := messagepartbuffer.New(messagepartbuffer.Options{Clock: clock}) - defer buffer.Close() - - key := testEpisodeKey() - require.ErrorIs(t, buffer.RecordToolStart(key, 0, clock.Now()), messagepartbuffer.ErrEpisodeNotFound) - - require.NoError(t, buffer.CreateEpisode(key)) - require.NoError(t, buffer.StartToolBatch(key, []messagepartbuffer.DispatchedToolCall{ - {CallIndex: 0, ToolCallID: "call-1"}, - {CallIndex: 1, ToolCallID: "call-dup"}, - {CallIndex: 2, ToolCallID: "call-dup"}, - })) - for _, completion := range buffer.ToolCompletions(key) { - require.True(t, completion.StartedAt.IsZero(), "seeding must not mark occurrences as started") - } + require.Len(t, buffer.ToolCompletions(key), 2, "completion without a start must be dropped") - clock.Advance(time.Second) - secondDupStartedAt := clock.Now() - require.NoError(t, buffer.RecordToolStart(key, 2, secondDupStartedAt)) - require.Equal(t, []messagepartbuffer.ToolCompletion{ - {CallIndex: 0, ToolCallID: "call-1"}, - {CallIndex: 1, ToolCallID: "call-dup"}, - {CallIndex: 2, ToolCallID: "call-dup", StartedAt: secondDupStartedAt}, - }, buffer.ToolCompletions(key)) - - clock.Advance(time.Second) - require.NoError(t, buffer.RecordToolStart(key, 7, clock.Now())) - require.True(t, buffer.ToolCompletions(key)[1].StartedAt.IsZero(), "unknown dispatch index must be dropped") - - clock.Advance(time.Second) - firstDupStartedAt := clock.Now() - require.NoError(t, buffer.RecordToolStart(key, 1, firstDupStartedAt)) - require.Equal(t, firstDupStartedAt, buffer.ToolCompletions(key)[1].StartedAt) - - require.NoError(t, buffer.RecordToolStart(key, 9, clock.Now())) - require.Len(t, buffer.ToolCompletions(key), 3) - - completions := buffer.ToolCompletions(key) require.NoError(t, buffer.CloseEpisode(key)) - require.Equal(t, secondDupStartedAt, completions[2].StartedAt) require.ErrorIs(t, buffer.RecordToolStart(key, 0, clock.Now()), messagepartbuffer.ErrEpisodeClosed) - require.True(t, buffer.ToolCompletions(key)[0].StartedAt.IsZero()) + require.ErrorIs(t, buffer.RecordToolCompletion(key, 0, clock.Now()), messagepartbuffer.ErrEpisodeClosed) } func TestBuffer_SubscribeExistingReplaysThenStreamsLiveParts(t *testing.T) { diff --git a/coderd/x/chatd/tasks.go b/coderd/x/chatd/tasks.go index 91598a44b68..f21e21b9d8b 100644 --- a/coderd/x/chatd/tasks.go +++ b/coderd/x/chatd/tasks.go @@ -252,19 +252,17 @@ type interruptEpisodeSnapshot struct { } type interruptEpisodeBilling struct { - interruptedAt time.Time - modelInvokedAt time.Time - toolBatchStartedAt time.Time - toolCompletions []messagepartbuffer.ToolCompletion + interruptedAt time.Time + modelInvokedAt time.Time + toolCompletions []messagepartbuffer.ToolCompletion } // closeInterruptEpisode snapshots billing, closes the episode, and returns its // buffered parts. Unknown episodes close blank so interruption converges. func (s *taskStarter) closeInterruptEpisode(ctx context.Context, key messagepartbuffer.Key) (interruptEpisodeBilling, []messagepartbuffer.Part, error) { billing := interruptEpisodeBilling{ - modelInvokedAt: s.opts.MessagePartBuffer.ModelInvokedAt(key), - toolBatchStartedAt: s.opts.MessagePartBuffer.ToolBatchStartedAt(key), - toolCompletions: s.opts.MessagePartBuffer.ToolCompletions(key), + modelInvokedAt: s.opts.MessagePartBuffer.ModelInvokedAt(key), + toolCompletions: s.opts.MessagePartBuffer.ToolCompletions(key), } if err := s.opts.MessagePartBuffer.CloseEpisode(key); err != nil { if ctx.Err() != nil { @@ -370,7 +368,6 @@ func (s *taskStarter) StartInterrupt(ctx context.Context, input chatWorkerTaskSt // Reuse the captured interrupt instant so database delay and retries do // not inflate billing. committedCancels, err := committedPendingLocalToolCancellationMessages(ctx, store, chat, interruptedAt, interruptedToolBatchBilling{ - batchStartedAt: episodeBilling.toolBatchStartedAt, toolCompletions: episodeBilling.toolCompletions, }) if err != nil { @@ -751,11 +748,9 @@ func dynamicToolNamesFromChat(chat database.Chat) map[string]bool { // interruptedToolBatchBilling is the live batch state used to bill // synthesized cancellation rows. type interruptedToolBatchBilling struct { - // batchStartedAt is zero when no local tool batch was live. - batchStartedAt time.Time - // toolCompletions uses unresolved-call positions. Completed calls bill - // to completion, running calls bill to the interrupt, and queued or - // undispatched calls bill nothing. + // toolCompletions contains started occurrences at unresolved-call + // positions. Completed calls bill to completion and running calls bill + // to the interrupt. Absent calls never started and bill nothing. toolCompletions []messagepartbuffer.ToolCompletion } @@ -782,12 +777,12 @@ func committedPendingLocalToolCancellationMessages( } // Match by unresolved-call position so rejected and duplicate-ID calls // cannot share occurrence state. - dispatched := make(map[int]messagepartbuffer.ToolCompletion, len(billing.toolCompletions)) + started := make(map[int]messagepartbuffer.ToolCompletion, len(billing.toolCompletions)) for _, completion := range billing.toolCompletions { if completion.CallIndex < 0 { continue } - dispatched[completion.CallIndex] = completion + started[completion.CallIndex] = completion } var ( windowEnd time.Time @@ -815,28 +810,23 @@ func committedPendingLocalToolCancellationMessages( ModelConfigID: uuid.NullUUID{UUID: chat.LastModelConfigID, Valid: chat.LastModelConfigID != uuid.Nil}, ContentVersion: chatprompt.CurrentContentVersion, }) - if billing.batchStartedAt.IsZero() || unbilledSubagentToolNames[call.ToolName] { + if unbilledSubagentToolNames[call.ToolName] { continue } - // Bill only matching dispatched calls. Completed calls end at - // completion, running calls end at the interrupt, and queued serial - // calls are skipped. Ties keep the first call. - occurrence, ok := dispatched[i] - if !ok || occurrence.ToolCallID != call.ToolCallID { + // Bill only matching started calls. Completed calls end at completion, + // running calls end at the interrupt, and ties keep the first call. + occurrence, ok := started[i] + if !ok { continue } start := occurrence.StartedAt + if start.IsZero() { + continue + } end := occurrence.CompletedAt if end.IsZero() { - if start.IsZero() { - continue - } end = interruptedAt } - if start.IsZero() { - // Completed calls should have starts; batchStart preserves legacy data. - start = billing.batchStartedAt - } intervals = append(intervals, chatloop.BilledInterval{Start: start, End: end}) if end.After(windowEnd) { windowEnd = end diff --git a/coderd/x/chatd/tasks_test.go b/coderd/x/chatd/tasks_test.go index 55dafe0663d..215340c23a0 100644 --- a/coderd/x/chatd/tasks_test.go +++ b/coderd/x/chatd/tasks_test.go @@ -581,10 +581,6 @@ func TestInterruptTask_ToolBatchBillsPartialWindowOnCancellationRow(t *testing.T // Keep advances below the buffer's 15-second cleanup tick. batch.clock.Advance(2 * time.Second) - require.NoError(t, buffer.StartToolBatch(batch.key, []messagepartbuffer.DispatchedToolCall{ - {CallIndex: 0, ToolCallID: execCallID}, - {CallIndex: 1, ToolCallID: waitCallID}, - })) require.NoError(t, buffer.RecordToolStart(batch.key, 0, batch.clock.Now())) require.NoError(t, buffer.RecordToolStart(batch.key, 1, batch.clock.Now())) // Record a live completion without publishing a tool result. @@ -614,15 +610,14 @@ func TestInterruptTask_RetrySnapshotOutlivesEpisodeEviction(t *testing.T) { // Use only the carried snapshot; the buffer has no episode state. batch.clock.Advance(2 * time.Second) - batchStartedAt := batch.clock.Now() + startedAt := batch.clock.Now() batch.clock.Advance(7 * time.Second) snapshot := &interruptEpisodeSnapshot{ loaded: true, key: batch.key, billing: interruptEpisodeBilling{ - interruptedAt: batch.clock.Now(), - toolBatchStartedAt: batchStartedAt, - toolCompletions: []messagepartbuffer.ToolCompletion{{CallIndex: 0, ToolCallID: execCallID, StartedAt: batchStartedAt}}, + interruptedAt: batch.clock.Now(), + toolCompletions: []messagepartbuffer.ToolCompletion{{CallIndex: 0, StartedAt: startedAt}}, }, } @@ -642,7 +637,6 @@ func TestInterruptTask_SnapshotCapturedBeforeChatRead(t *testing.T) { buffer := batch.starter.opts.MessagePartBuffer batch.clock.Advance(2 * time.Second) - require.NoError(t, buffer.StartToolBatch(batch.key, []messagepartbuffer.DispatchedToolCall{{CallIndex: 0, ToolCallID: execCallID}})) require.NoError(t, buffer.RecordToolStart(batch.key, 0, batch.clock.Now())) batch.clock.Advance(3 * time.Second) @@ -681,7 +675,6 @@ func TestInterruptTask_RunningBilledToolBillsUpToInterrupt(t *testing.T) { }) batch.clock.Advance(2 * time.Second) - require.NoError(t, batch.starter.opts.MessagePartBuffer.StartToolBatch(batch.key, []messagepartbuffer.DispatchedToolCall{{CallIndex: 0, ToolCallID: execCallID}})) require.NoError(t, batch.starter.opts.MessagePartBuffer.RecordToolStart(batch.key, 0, batch.clock.Now())) batch.clock.Advance(7 * time.Second) @@ -700,7 +693,6 @@ func TestInterruptTask_UnbilledOnlyBatchBillsNothingOnInterrupt(t *testing.T) { }) batch.clock.Advance(2 * time.Second) - require.NoError(t, batch.starter.opts.MessagePartBuffer.StartToolBatch(batch.key, []messagepartbuffer.DispatchedToolCall{{CallIndex: 0, ToolCallID: waitCallID}})) require.NoError(t, batch.starter.opts.MessagePartBuffer.RecordToolStart(batch.key, 0, batch.clock.Now())) batch.clock.Advance(10 * time.Second) @@ -722,10 +714,6 @@ func TestInterruptTask_UnstartedSerialCallBillsNothingOnInterrupt(t *testing.T) buffer := batch.starter.opts.MessagePartBuffer batch.clock.Advance(2 * time.Second) - require.NoError(t, buffer.StartToolBatch(batch.key, []messagepartbuffer.DispatchedToolCall{ - {CallIndex: 0, ToolCallID: waitCallID}, - {CallIndex: 1, ToolCallID: serialCallID}, - })) require.NoError(t, buffer.RecordToolStart(batch.key, 0, batch.clock.Now())) batch.clock.Advance(10 * time.Second) @@ -753,11 +741,6 @@ func TestInterruptTask_StartedSerialCallBillsFromItsOwnStart(t *testing.T) { buffer := batch.starter.opts.MessagePartBuffer batch.clock.Advance(2 * time.Second) - require.NoError(t, buffer.StartToolBatch(batch.key, []messagepartbuffer.DispatchedToolCall{ - {CallIndex: 0, ToolCallID: execCallID}, - {CallIndex: 1, ToolCallID: waitCallID}, - {CallIndex: 2, ToolCallID: serialCallID}, - })) require.NoError(t, buffer.RecordToolStart(batch.key, 0, batch.clock.Now())) require.NoError(t, buffer.RecordToolStart(batch.key, 1, batch.clock.Now())) batch.clock.Advance(3 * time.Second) @@ -788,10 +771,6 @@ func TestInterruptTask_DuplicateCallIDsKeepOccurrenceStates(t *testing.T) { buffer := batch.starter.opts.MessagePartBuffer batch.clock.Advance(2 * time.Second) - require.NoError(t, buffer.StartToolBatch(batch.key, []messagepartbuffer.DispatchedToolCall{ - {CallIndex: 0, ToolCallID: dupCallID}, - {CallIndex: 1, ToolCallID: dupCallID}, - })) require.NoError(t, buffer.RecordToolStart(batch.key, 0, batch.clock.Now())) require.NoError(t, buffer.RecordToolStart(batch.key, 1, batch.clock.Now())) batch.clock.Advance(3 * time.Second) @@ -820,10 +799,6 @@ func TestInterruptTask_DuplicateCallIDCompletionStampsOwnOccurrence(t *testing.T buffer := batch.starter.opts.MessagePartBuffer batch.clock.Advance(2 * time.Second) - require.NoError(t, buffer.StartToolBatch(batch.key, []messagepartbuffer.DispatchedToolCall{ - {CallIndex: 0, ToolCallID: dupCallID}, - {CallIndex: 1, ToolCallID: dupCallID}, - })) require.NoError(t, buffer.RecordToolStart(batch.key, 0, batch.clock.Now())) require.NoError(t, buffer.RecordToolStart(batch.key, 1, batch.clock.Now())) batch.clock.Advance(3 * time.Second) @@ -851,8 +826,7 @@ func TestInterruptTask_RejectedDuplicateIDDoesNotStealDispatchedOccurrence(t *te }) batch.clock.Advance(2 * time.Second) - require.NoError(t, batch.starter.opts.MessagePartBuffer.StartToolBatch(batch.key, []messagepartbuffer.DispatchedToolCall{{CallIndex: 1, ToolCallID: dupCallID}})) - require.NoError(t, batch.starter.opts.MessagePartBuffer.RecordToolStart(batch.key, 0, batch.clock.Now())) + require.NoError(t, batch.starter.opts.MessagePartBuffer.RecordToolStart(batch.key, 1, batch.clock.Now())) batch.clock.Advance(10 * time.Second) messages := batch.interrupt(t, f) @@ -875,8 +849,7 @@ func TestInterruptTask_RejectedCallBillsNothingOnInterrupt(t *testing.T) { }) batch.clock.Advance(2 * time.Second) - require.NoError(t, batch.starter.opts.MessagePartBuffer.StartToolBatch(batch.key, []messagepartbuffer.DispatchedToolCall{{CallIndex: 1, ToolCallID: waitCallID}})) - require.NoError(t, batch.starter.opts.MessagePartBuffer.RecordToolStart(batch.key, 0, batch.clock.Now())) + require.NoError(t, batch.starter.opts.MessagePartBuffer.RecordToolStart(batch.key, 1, batch.clock.Now())) batch.clock.Advance(10 * time.Second) messages := batch.interrupt(t, f) From d52fd1a50afcbcb3b8be33aabca25338698a9391 Mon Sep 17 00:00:00 2001 From: Jaayden Halko Date: Thu, 20 Aug 2026 07:47:35 +0000 Subject: [PATCH 20/20] refactor(coderd/x/chatd): simplify tool runtime billing --- coderd/x/chatd/attempt.go | 6 +- coderd/x/chatd/chatloop/chatloop.go | 129 ++++++------------ coderd/x/chatd/chatloop/runtime_test.go | 22 +-- coderd/x/chatd/generation.go | 15 +- coderd/x/chatd/message_conversion.go | 10 +- coderd/x/chatd/message_conversion_test.go | 70 +--------- .../messagepartbuffer/message_part_buffer.go | 39 ++---- .../message_part_buffer_test.go | 33 ++--- coderd/x/chatd/tasks.go | 28 +--- coderd/x/chatd/tasks_test.go | 2 +- 10 files changed, 94 insertions(+), 260 deletions(-) diff --git a/coderd/x/chatd/attempt.go b/coderd/x/chatd/attempt.go index 7804a25c0d7..580d113716b 100644 --- a/coderd/x/chatd/attempt.go +++ b/coderd/x/chatd/attempt.go @@ -31,10 +31,8 @@ type stepData struct { ContextLimit sql.NullInt64 Runtime time.Duration - // BatchRuntime is the local-tool batch window persisted on the - // BatchRuntimeToolCallID tool row. Model steps use Runtime instead. - BatchRuntime time.Duration - BatchRuntimeToolCallID string + // BatchRuntime is the local-tool batch window. Model steps use Runtime. + BatchRuntime time.Duration ToolCallCreatedAt map[string]time.Time ToolResultCreatedAt map[string]time.Time diff --git a/coderd/x/chatd/chatloop/chatloop.go b/coderd/x/chatd/chatloop/chatloop.go index 0a679361773..0a0545da5b9 100644 --- a/coderd/x/chatd/chatloop/chatloop.go +++ b/coderd/x/chatd/chatloop/chatloop.go @@ -74,8 +74,11 @@ type PersistedStep struct { Usage fantasy.Usage ContextLimit sql.NullInt64 // Runtime is the wall-clock duration from opening to consuming the - // model stream. Local tool batches use ToolExecutionOutcome.BatchRuntime. + // 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 @@ -283,18 +286,6 @@ type ExecuteLocalToolsOptions struct { Clock quartz.Clock } -// ToolExecutionOutcome is the durable tool-result content from one batch. -type ToolExecutionOutcome struct { - Step PersistedStep - // BatchRuntime is the union of billed tool execution intervals. Parallel - // calls count once and serial calls count only from their own start. Zero - // means no billed tool produced a result. - BatchRuntime time.Duration - // BatchRuntimeToolCallID is the ID on the billed interval ending last. - // Ties use call order; the ID can be empty or non-unique. - BatchRuntimeToolCallID string -} - // GenerateCompactionOptions configures one context compaction call. type GenerateCompactionOptions struct { Model fantasy.LanguageModel @@ -551,7 +542,7 @@ func contentFilterError(provider string, metadata fantasy.ProviderMetadata) erro // ExecuteLocalTools runs local tool calls and returns durable tool results. It // does not retry or persist. -func ExecuteLocalTools(ctx context.Context, opts ExecuteLocalToolsOptions) (ToolExecutionOutcome, error) { +func ExecuteLocalTools(ctx context.Context, opts ExecuteLocalToolsOptions) (PersistedStep, error) { if opts.Metrics == nil { opts.Metrics = NopMetrics() } @@ -573,7 +564,7 @@ func ExecuteLocalTools(ctx context.Context, opts ExecuteLocalToolsOptions) (Tool // without capturing the publisher at construction time. ctx = WithMessagePartPublisher(ctx, opts.PublishMessagePart) if ctx.Err() != nil { - return ToolExecutionOutcome{}, ctx.Err() + return PersistedStep{}, ctx.Err() } localCalls := make([]fantasy.ToolCallContent, 0, len(opts.ToolCalls)) @@ -583,7 +574,7 @@ func ExecuteLocalTools(ctx context.Context, opts ExecuteLocalToolsOptions) (Tool } } if len(localCalls) == 0 { - return ToolExecutionOutcome{}, nil + return PersistedStep{}, nil } var result stepResult @@ -605,12 +596,12 @@ func ExecuteLocalTools(ctx context.Context, opts ExecuteLocalToolsOptions) (Tool result.content = append(result.content, tr) } if ctx.Err() != nil { - return ToolExecutionOutcome{}, ctx.Err() + return PersistedStep{}, ctx.Err() } - return ToolExecutionOutcome{Step: PersistedStep{ + return PersistedStep{ Content: result.content, ToolResultCreatedAt: result.toolResultCreatedAt, - }}, nil + }, nil } maxResultBytes := toolResultByteBudget(opts.ContextLimit) @@ -654,79 +645,49 @@ func ExecuteLocalTools(ctx context.Context, opts ExecuteLocalToolsOptions) (Tool }, ) if ctx.Err() != nil { - return ToolExecutionOutcome{}, ctx.Err() + return PersistedStep{}, ctx.Err() } for _, tr := range toolResults { result.content = append(result.content, tr) } - batchRuntime, batchRuntimeToolCallID := billableBatchWindow( - batchStart, - localCalls, - orderedStarts, - orderedCompletions, - opts.UnbilledToolNames, - ) - return ToolExecutionOutcome{ - Step: PersistedStep{ - Content: result.content, - ToolResultCreatedAt: result.toolResultCreatedAt, - }, - BatchRuntime: batchRuntime, - BatchRuntimeToolCallID: batchRuntimeToolCallID, + return PersistedStep{ + Content: result.content, + ToolResultCreatedAt: result.toolResultCreatedAt, + BatchRuntime: billableBatchDuration( + batchStart, + localCalls, + orderedStarts, + orderedCompletions, + opts.UnbilledToolNames, + ), }, nil } -// billableBatchWindow returns the union of billed execution intervals and -// the call ID whose interval ends last. Concurrent calls start at -// batchStart; serial calls use their recorded starts. Unbilled tools and -// gaps between billed intervals do not count. -// -// Starts and completions align with toolCalls by occurrence, so duplicate or -// empty IDs stay distinct. A zero start means the call used batchStart. Ties -// keep the earliest call. -func billableBatchWindow( +// 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, string) { - var ( - found bool - windowEnd time.Time - toolCallID string - intervals []BilledInterval - ) +) time.Duration { + intervals := make([]BilledInterval, 0, len(toolCalls)) for i, tc := range toolCalls { if i >= len(completions) { break } - if unbilledToolNames[tc.ToolName] { - continue - } - end := completions[i] - if end.IsZero() { + 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: end}) - if end.After(windowEnd) { - found = true - windowEnd = end - toolCallID = tc.ToolCallID - } - } - if !found { - return 0, "" + intervals = append(intervals, BilledInterval{Start: start, End: completions[i]}) } - runtime := BilledIntervalsDuration(intervals) - if runtime <= 0 { - return 0, "" - } - return runtime, toolCallID + return BilledIntervalsDuration(intervals) } // BilledInterval is one billed tool call's execution window. @@ -739,27 +700,18 @@ type BilledInterval struct { // 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 := make([]BilledInterval, 0, len(intervals)) - for _, iv := range intervals { - if iv.End.Before(iv.Start) { - continue - } - valid = append(valid, iv) + 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) }) - var ( - total time.Duration - curStart time.Time - curEnd time.Time - open bool - ) - for _, iv := range valid { - if !open { - curStart, curEnd, open = iv.Start, iv.End, true - continue - } + 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 @@ -769,10 +721,7 @@ func BilledIntervalsDuration(intervals []BilledInterval) time.Duration { curEnd = iv.End } } - if open { - total += curEnd.Sub(curStart) - } - return total + return total + curEnd.Sub(curStart) } // prepareMessagesForRequest applies the prompt preparation pipeline used diff --git a/coderd/x/chatd/chatloop/runtime_test.go b/coderd/x/chatd/chatloop/runtime_test.go index 17921cd6620..005054d5569 100644 --- a/coderd/x/chatd/chatloop/runtime_test.go +++ b/coderd/x/chatd/chatloop/runtime_test.go @@ -168,10 +168,10 @@ func executeToolBatch( t *testing.T, clock *quartz.Mock, opts chatloop.ExecuteLocalToolsOptions, -) <-chan chatloop.ToolExecutionOutcome { +) <-chan chatloop.PersistedStep { t.Helper() opts.Clock = clock - resultCh := make(chan chatloop.ToolExecutionOutcome, 1) + resultCh := make(chan chatloop.PersistedStep, 1) go func() { outcome, err := chatloop.ExecuteLocalTools(context.Background(), opts) assert.NoError(t, err) @@ -229,10 +229,9 @@ func TestExecuteLocalTools_BatchWindowIsMaxNotSum(t *testing.T) { outcome := testutil.RequireReceive(ctx, t, resultCh) require.Equal(t, 60*time.Second, outcome.BatchRuntime) - require.Equal(t, "call-slow", outcome.BatchRuntimeToolCallID) } -func TestExecuteLocalTools_SimultaneousCompletionsBillOnceByCallOrder(t *testing.T) { +func TestExecuteLocalTools_SimultaneousCompletionsBillOnce(t *testing.T) { t.Parallel() ctx := testutil.Context(t, testutil.WaitShort) @@ -262,7 +261,6 @@ func TestExecuteLocalTools_SimultaneousCompletionsBillOnceByCallOrder(t *testing outcome := testutil.RequireReceive(ctx, t, resultCh) require.Equal(t, 10*time.Second, outcome.BatchRuntime) - require.Equal(t, "call-1", outcome.BatchRuntimeToolCallID) } func TestExecuteLocalTools_UnbilledToolNeverExtendsWindow(t *testing.T) { @@ -298,7 +296,6 @@ func TestExecuteLocalTools_UnbilledToolNeverExtendsWindow(t *testing.T) { outcome := testutil.RequireReceive(ctx, t, resultCh) require.Equal(t, 10*time.Second, outcome.BatchRuntime) - require.Equal(t, "call-execute", outcome.BatchRuntimeToolCallID) } func TestExecuteLocalTools_UnbilledOnlyBatchBillsNothing(t *testing.T) { @@ -328,7 +325,6 @@ func TestExecuteLocalTools_UnbilledOnlyBatchBillsNothing(t *testing.T) { outcome := testutil.RequireReceive(ctx, t, resultCh) require.Zero(t, outcome.BatchRuntime) - require.Empty(t, outcome.BatchRuntimeToolCallID) } func TestExecuteLocalTools_AliasNamesClassifyAsCalled(t *testing.T) { @@ -368,7 +364,6 @@ func TestExecuteLocalTools_AliasNamesClassifyAsCalled(t *testing.T) { outcome := testutil.RequireReceive(ctx, t, resultCh) require.Equal(t, 10*time.Second, outcome.BatchRuntime) - require.Equal(t, "call-execute", outcome.BatchRuntimeToolCallID) } func TestExecuteLocalTools_DuplicateToolCallIDsKeepOccurrenceCompletions(t *testing.T) { @@ -403,7 +398,6 @@ func TestExecuteLocalTools_DuplicateToolCallIDsKeepOccurrenceCompletions(t *test outcome := testutil.RequireReceive(ctx, t, resultCh) require.Equal(t, 60*time.Second, outcome.BatchRuntime) - require.Equal(t, "call-dup", outcome.BatchRuntimeToolCallID) } func TestExecuteLocalTools_ExecutionCallbacksFireOnlyForRuns(t *testing.T) { @@ -438,7 +432,7 @@ func TestExecuteLocalTools_ExecutionCallbacksFireOnlyForRuns(t *testing.T) { }, }) require.NoError(t, err) - require.Len(t, outcome.Step.Content, 1) + 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") @@ -480,7 +474,7 @@ func TestExecuteLocalTools_ExecutionCallbacksFireOnlyForRuns(t *testing.T) { }, }) require.NoError(t, err) - require.Len(t, outcome.Step.Content, 2, "the whole batch resolves to synthesized policy errors") + require.Len(t, outcome.Content, 2, "the whole batch resolves to synthesized policy errors") require.False(t, started) require.False(t, completed) }) @@ -518,7 +512,6 @@ func TestExecuteLocalTools_EmptyToolCallIDStillBillsWindow(t *testing.T) { outcome := testutil.RequireReceive(ctx, t, resultCh) require.Equal(t, 60*time.Second, outcome.BatchRuntime) - require.Empty(t, outcome.BatchRuntimeToolCallID) } func TestExecuteLocalTools_OnToolCompleteReportsLiveCompletions(t *testing.T) { @@ -568,7 +561,7 @@ func TestExecuteLocalTools_OnToolCompleteReportsLiveCompletions(t *testing.T) { require.Equal(t, map[string]time.Time{ "call-fast": fast.completedAt, "call-slow": slow.completedAt, - }, outcome.Step.ToolResultCreatedAt) + }, outcome.ToolResultCreatedAt) } func TestExecuteLocalTools_SerialCallBillsFromItsOwnStart(t *testing.T) { @@ -607,7 +600,6 @@ func TestExecuteLocalTools_SerialCallBillsFromItsOwnStart(t *testing.T) { 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") - require.Equal(t, "call-serial", outcome.BatchRuntimeToolCallID) } func TestExecuteLocalTools_SerialAfterBilledSiblingBillsUnion(t *testing.T) { @@ -652,8 +644,6 @@ func TestExecuteLocalTools_SerialAfterBilledSiblingBillsUnion(t *testing.T) { 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") - require.Equal(t, "call-serial", outcome.BatchRuntimeToolCallID, - "the serial call's completion ends the window") } func TestBilledIntervalsDuration(t *testing.T) { diff --git a/coderd/x/chatd/generation.go b/coderd/x/chatd/generation.go index 5ced8e7e298..ddd73a23c3f 100644 --- a/coderd/x/chatd/generation.go +++ b/coderd/x/chatd/generation.go @@ -862,7 +862,7 @@ func (s *taskStarter) executeLocalTools( provider = prepared.Model.Provider() modelName = prepared.Model.ModelID() } - var outcome chatloop.ToolExecutionOutcome + var outcome chatloop.PersistedStep var spawnDispatchErr error if len(allowed) > 0 { var onToolStart func(int, time.Time) @@ -911,18 +911,16 @@ func (s *taskStarter) executeLocalTools( // the tool run; its failure surfaces as a tool result error. The // step still commits so a sibling tool that already ran keeps its // result and is not re-executed, and the turn fails afterwards. - if hookErr := chathooks.DispatchFailureFromResults(outcome.Step.Content); hookErr != nil { + if hookErr := chathooks.DispatchFailureFromResults(outcome.Content); hookErr != nil { spawnDispatchErr = chathooks.GenerationDispatchError(agenthooks.EventUserPromptSubmit, hookErr) } } - postResults, postDispatchErr := s.server.hooks.PostToolUseResults(ctx, chathooks.ChatFor(prepared.Chat, input.hookTurnID()), outcome.Step.Content) + postResults, postDispatchErr := s.server.hooks.PostToolUseResults(ctx, chathooks.ChatFor(prepared.Chat, input.hookTurnID()), outcome.Content) for _, result := range denied { - outcome.Step.Content = append(outcome.Step.Content, result) + outcome.Content = append(outcome.Content, result) } - chathooks.RestoreToolCallOrder(outcome.Step.Content, decision.localToolCalls) - step := stepDataFromPersisted(outcome.Step) - step.BatchRuntime = outcome.BatchRuntime - step.BatchRuntimeToolCallID = outcome.BatchRuntimeToolCallID + chathooks.RestoreToolCallOrder(outcome.Content, decision.localToolCalls) + step := stepDataFromPersisted(outcome) messages, err := buildCommitStepMessages(buildCommitStepMessagesInput{ modelConfigID: prepared.ModelConfigID, step: step, @@ -1581,6 +1579,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 b6fd8c12d0e..a27229b54f1 100644 --- a/coderd/x/chatd/message_conversion.go +++ b/coderd/x/chatd/message_conversion.go @@ -61,8 +61,7 @@ func buildCommitStepMessages(input buildCommitStepMessagesInput) (stepMessagesFo messages = append(messages, assistantMessage(input.modelConfigID, contentVersion, assistantContent, input.step)) } - batchRuntimeAssigned := false - 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 { @@ -75,12 +74,9 @@ func buildCommitStepMessages(input buildCommitStepMessagesInput) (stepMessagesFo return stepMessagesForCommit{}, xerrors.Errorf("marshal tool result: %w", err) } msg := baseMessage(database.ChatMessageRoleTool, database.ChatMessageVisibilityBoth, input.modelConfigID, contentVersion, content) - // Assign the batch window to one matching tool row because usage sums - // runtime_ms. The first match handles duplicate and ID-less calls; - // zero stays NULL. - if !batchRuntimeAssigned && input.step.BatchRuntime > 0 && toolResult.ToolCallID == input.step.BatchRuntimeToolCallID { + // Usage sums runtime_ms across rows, so store the batch once. + if i == 0 { msg.RuntimeMs = nullInt64IfNonZero(input.step.BatchRuntime.Milliseconds()) - batchRuntimeAssigned = true } messages = append(messages, msg) } diff --git a/coderd/x/chatd/message_conversion_test.go b/coderd/x/chatd/message_conversion_test.go index e47c9e49a3f..ef0852fcb02 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_BatchRuntimeLandsOnWindowDefiningToolRow(t *testing.T) { +func TestBuildCommitStepMessages_BatchRuntimeLandsOnFirstToolRow(t *testing.T) { t.Parallel() got, err := buildCommitStepMessages(buildCommitStepMessagesInput{ @@ -120,40 +120,7 @@ func TestBuildCommitStepMessages_BatchRuntimeLandsOnWindowDefiningToolRow(t *tes Result: fantasy.ToolResultOutputContentText{Text: `{"stdout":"/tmp"}`}, }, }, - BatchRuntime: 10 * time.Second, - BatchRuntimeToolCallID: "call-2", - }, - }) - require.NoError(t, err) - require.Len(t, got.Messages, 2) - require.Equal(t, database.ChatMessageRoleTool, 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: 10000, Valid: true}, got.Messages[1].RuntimeMs) -} - -func TestBuildCommitStepMessages_DuplicateToolCallIDsBillOnce(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: "execute", - Result: fantasy.ToolResultOutputContentText{Text: `{"stdout":"first"}`}, - }, - fantasy.ToolResultContent{ - ToolCallID: "call-1", - ToolName: "execute", - Result: fantasy.ToolResultOutputContentText{Text: `{"stdout":"second"}`}, - }, - }, - BatchRuntime: 10 * time.Second, - BatchRuntimeToolCallID: "call-1", + BatchRuntime: 10 * time.Second, }, }) require.NoError(t, err) @@ -162,36 +129,6 @@ func TestBuildCommitStepMessages_DuplicateToolCallIDsBillOnce(t *testing.T) { require.False(t, got.Messages[1].RuntimeMs.Valid) } -func TestBuildCommitStepMessages_EmptyIDWindowLandsOnIDLessRow(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-fast", - ToolName: "fast_tool", - Result: fantasy.ToolResultOutputContentText{Text: `{"stdout":"fast"}`}, - }, - fantasy.ToolResultContent{ - ToolCallID: "", - ToolName: "idless_tool", - Result: fantasy.ToolResultOutputContentText{Text: `{"stdout":"slow"}`}, - }, - }, - BatchRuntime: 60 * time.Second, - BatchRuntimeToolCallID: "", - }, - }) - require.NoError(t, err) - require.Len(t, got.Messages, 2) - require.False(t, got.Messages[0].RuntimeMs.Valid, "the identified row is not the window-defining one") - require.Equal(t, sql.NullInt64{Int64: 60_000, Valid: true}, got.Messages[1].RuntimeMs) -} - func TestBuildCommitStepMessages_BatchAttachmentAssistantRowStaysNull(t *testing.T) { t.Parallel() @@ -208,8 +145,7 @@ func TestBuildCommitStepMessages_BatchAttachmentAssistantRowStaysNull(t *testing ClientMetadata: `{"attachments":[{"file_id":"` + uuid.NewString() + `","media_type":"image/png","name":"shot.png"}]}`, }, }, - BatchRuntime: 3 * time.Second, - BatchRuntimeToolCallID: "call-1", + BatchRuntime: 3 * time.Second, }, }) require.NoError(t, err) diff --git a/coderd/x/chatd/messagepartbuffer/message_part_buffer.go b/coderd/x/chatd/messagepartbuffer/message_part_buffer.go index 39da1751482..67115f5358d 100644 --- a/coderd/x/chatd/messagepartbuffer/message_part_buffer.go +++ b/coderd/x/chatd/messagepartbuffer/message_part_buffer.go @@ -19,6 +19,7 @@ import ( "container/heap" "context" "encoding/json" + "maps" "slices" "sync" "time" @@ -102,9 +103,8 @@ type episodeState struct { // that never invoke a model, such as local tool execution // batches. modelStartedAt time.Time - // toolCompletions stores started occurrences for interrupts. CallIndex - // distinguishes rejected-call gaps and duplicate IDs. - toolCompletions []ToolCompletion + // toolCompletions stores started occurrences by unresolved-call position. + toolCompletions map[int]ToolCompletion closed bool closedAt time.Time closedHeapItem *closedEpisodeItem @@ -224,8 +224,6 @@ func (b *Buffer) StartModelInvocation(key Key) error { // ToolCompletion tracks a started tool-call occurrence. CompletedAt is zero // while the call is unfinished. type ToolCompletion struct { - // CallIndex is the occurrence's position in the unresolved call order. - CallIndex int StartedAt time.Time CompletedAt time.Time } @@ -249,15 +247,13 @@ func (b *Buffer) RecordToolStart(key Key, callIndex int, startedAt time.Time) er if callIndex < 0 { return nil } - for _, entry := range episode.toolCompletions { - if entry.CallIndex == callIndex { - return nil - } + if _, ok := episode.toolCompletions[callIndex]; ok { + return nil + } + if episode.toolCompletions == nil { + episode.toolCompletions = make(map[int]ToolCompletion) } - episode.toolCompletions = append(episode.toolCompletions, ToolCompletion{ - CallIndex: callIndex, - StartedAt: startedAt, - }) + episode.toolCompletions[callIndex] = ToolCompletion{StartedAt: startedAt} return nil } @@ -276,15 +272,10 @@ func (b *Buffer) RecordToolCompletion(key Key, callIndex int, completedAt time.T if episode.closed { return ErrEpisodeClosed } - for i := range episode.toolCompletions { - entry := &episode.toolCompletions[i] - if entry.CallIndex != callIndex { - continue - } - if entry.CompletedAt.IsZero() { - entry.CompletedAt = completedAt - } - break + entry, ok := episode.toolCompletions[callIndex] + if ok && entry.CompletedAt.IsZero() { + entry.CompletedAt = completedAt + episode.toolCompletions[callIndex] = entry } return nil } @@ -377,14 +368,14 @@ func (b *Buffer) ModelInvokedAt(key Key) time.Time { // ToolCompletions returns copied started-occurrence state. A zero completion // means unfinished. Read it before CloseEpisode because closed episodes are // garbage collected. -func (b *Buffer) ToolCompletions(key Key) []ToolCompletion { +func (b *Buffer) ToolCompletions(key Key) map[int]ToolCompletion { b.mu.Lock() defer b.mu.Unlock() episode := b.episodes[key] if episode == nil { return nil } - return slices.Clone(episode.toolCompletions) + return maps.Clone(episode.toolCompletions) } // SubscribeToEpisode replays existing parts and streams new parts. diff --git a/coderd/x/chatd/messagepartbuffer/message_part_buffer_test.go b/coderd/x/chatd/messagepartbuffer/message_part_buffer_test.go index 3f341ea2f20..fc67ff4de3a 100644 --- a/coderd/x/chatd/messagepartbuffer/message_part_buffer_test.go +++ b/coderd/x/chatd/messagepartbuffer/message_part_buffer_test.go @@ -171,41 +171,36 @@ func TestBuffer_ToolCompletions(t *testing.T) { clock.Advance(time.Second) firstStartedAt := clock.Now() require.NoError(t, buffer.RecordToolStart(key, 1, firstStartedAt)) - require.Equal(t, []messagepartbuffer.ToolCompletion{ - {CallIndex: 2, StartedAt: secondStartedAt}, - {CallIndex: 1, StartedAt: firstStartedAt}, - }, buffer.ToolCompletions(key)) + started := map[int]messagepartbuffer.ToolCompletion{ + 1: {StartedAt: firstStartedAt}, + 2: {StartedAt: secondStartedAt}, + } + require.Equal(t, started, buffer.ToolCompletions(key)) clock.Advance(time.Second) require.NoError(t, buffer.RecordToolStart(key, 2, clock.Now())) require.NoError(t, buffer.RecordToolStart(key, -1, clock.Now())) - require.Equal(t, []messagepartbuffer.ToolCompletion{ - {CallIndex: 2, StartedAt: secondStartedAt}, - {CallIndex: 1, StartedAt: firstStartedAt}, - }, buffer.ToolCompletions(key), "repeated and invalid starts must not replace or append occurrences") + require.Equal(t, started, buffer.ToolCompletions(key), "repeated and invalid starts must not replace or append occurrences") clock.Advance(time.Second) secondCompletedAt := clock.Now() require.NoError(t, buffer.RecordToolCompletion(key, 2, secondCompletedAt)) - require.Equal(t, []messagepartbuffer.ToolCompletion{ - {CallIndex: 2, StartedAt: secondStartedAt, CompletedAt: secondCompletedAt}, - {CallIndex: 1, StartedAt: firstStartedAt}, - }, buffer.ToolCompletions(key)) clock.Advance(2 * time.Second) firstCompletedAt := clock.Now() require.NoError(t, buffer.RecordToolCompletion(key, 1, firstCompletedAt)) completions := buffer.ToolCompletions(key) - require.Equal(t, []messagepartbuffer.ToolCompletion{ - {CallIndex: 2, StartedAt: secondStartedAt, CompletedAt: secondCompletedAt}, - {CallIndex: 1, StartedAt: firstStartedAt, CompletedAt: firstCompletedAt}, + require.Equal(t, map[int]messagepartbuffer.ToolCompletion{ + 1: {StartedAt: firstStartedAt, CompletedAt: firstCompletedAt}, + 2: {StartedAt: secondStartedAt, CompletedAt: secondCompletedAt}, }, completions) - completions[0].CompletedAt = clock.Now() - require.Equal(t, secondCompletedAt, buffer.ToolCompletions(key)[0].CompletedAt) + completion := completions[2] + completion.CompletedAt = clock.Now() + completions[2] = completion + require.Equal(t, secondCompletedAt, buffer.ToolCompletions(key)[2].CompletedAt) clock.Advance(time.Second) - unseededAt := clock.Now() - require.NoError(t, buffer.RecordToolCompletion(key, 5, unseededAt)) + require.NoError(t, buffer.RecordToolCompletion(key, 5, clock.Now())) require.Len(t, buffer.ToolCompletions(key), 2, "completion without a start must be dropped") require.NoError(t, buffer.CloseEpisode(key)) diff --git a/coderd/x/chatd/tasks.go b/coderd/x/chatd/tasks.go index f21e21b9d8b..c6460338049 100644 --- a/coderd/x/chatd/tasks.go +++ b/coderd/x/chatd/tasks.go @@ -254,7 +254,7 @@ type interruptEpisodeSnapshot struct { type interruptEpisodeBilling struct { interruptedAt time.Time modelInvokedAt time.Time - toolCompletions []messagepartbuffer.ToolCompletion + toolCompletions map[int]messagepartbuffer.ToolCompletion } // closeInterruptEpisode snapshots billing, closes the episode, and returns its @@ -367,9 +367,7 @@ func (s *taskStarter) StartInterrupt(ctx context.Context, input chatWorkerTaskSt messages := partialMessages // Reuse the captured interrupt instant so database delay and retries do // not inflate billing. - committedCancels, err := committedPendingLocalToolCancellationMessages(ctx, store, chat, interruptedAt, interruptedToolBatchBilling{ - toolCompletions: episodeBilling.toolCompletions, - }) + committedCancels, err := committedPendingLocalToolCancellationMessages(ctx, store, chat, interruptedAt, episodeBilling.toolCompletions) if err != nil { return xerrors.Errorf("committed pending local tool cancellation messages: %w", err) } @@ -745,21 +743,12 @@ func dynamicToolNamesFromChat(chat database.Chat) map[string]bool { return names } -// interruptedToolBatchBilling is the live batch state used to bill -// synthesized cancellation rows. -type interruptedToolBatchBilling struct { - // toolCompletions contains started occurrences at unresolved-call - // positions. Completed calls bill to completion and running calls bill - // to the interrupt. Absent calls never started and bill nothing. - toolCompletions []messagepartbuffer.ToolCompletion -} - func committedPendingLocalToolCancellationMessages( ctx context.Context, store database.Store, chat database.Chat, interruptedAt time.Time, - billing interruptedToolBatchBilling, + toolCompletions map[int]messagepartbuffer.ToolCompletion, ) ([]chatstate.Message, error) { messages, err := store.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{ ChatID: chat.ID, @@ -775,15 +764,6 @@ func committedPendingLocalToolCancellationMessages( if len(localCalls) == 0 { return nil, nil } - // Match by unresolved-call position so rejected and duplicate-ID calls - // cannot share occurrence state. - started := make(map[int]messagepartbuffer.ToolCompletion, len(billing.toolCompletions)) - for _, completion := range billing.toolCompletions { - if completion.CallIndex < 0 { - continue - } - started[completion.CallIndex] = completion - } var ( windowEnd time.Time windowRowIdx = -1 @@ -815,7 +795,7 @@ func committedPendingLocalToolCancellationMessages( } // Bill only matching started calls. Completed calls end at completion, // running calls end at the interrupt, and ties keep the first call. - occurrence, ok := started[i] + occurrence, ok := toolCompletions[i] if !ok { continue } diff --git a/coderd/x/chatd/tasks_test.go b/coderd/x/chatd/tasks_test.go index 215340c23a0..30ae6037bff 100644 --- a/coderd/x/chatd/tasks_test.go +++ b/coderd/x/chatd/tasks_test.go @@ -617,7 +617,7 @@ func TestInterruptTask_RetrySnapshotOutlivesEpisodeEviction(t *testing.T) { key: batch.key, billing: interruptEpisodeBilling{ interruptedAt: batch.clock.Now(), - toolCompletions: []messagepartbuffer.ToolCompletion{{CallIndex: 0, StartedAt: startedAt}}, + toolCompletions: map[int]messagepartbuffer.ToolCompletion{0: {StartedAt: startedAt}}, }, }