diff --git a/coderd/database/querier_test.go b/coderd/database/querier_test.go index 2c4df715781..092d3aeeee9 100644 --- a/coderd/database/querier_test.go +++ b/coderd/database/querier_test.go @@ -11109,13 +11109,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) @@ -11124,22 +11124,24 @@ func TestGetTotalChatMessageRuntimeMsInRange(t *testing.T) { // Counted: on the inclusive start boundary, in the middle (across two // chats), soft-deleted, and just before the exclusive end boundary. - insertMessage(chat1.ID, 1, rangeStart, false) - insertMessage(chat2.ID, 2, rangeStart.Add(30*time.Minute), false) - insertMessage(chat1.ID, 4, rangeStart.Add(45*time.Minute), true) - insertMessage(chat1.ID, 8, rangeEnd.Add(-time.Second), false) + insertMessage(chat1.ID, database.ChatMessageRoleAssistant, 1, rangeStart, false) + insertMessage(chat2.ID, database.ChatMessageRoleAssistant, 2, rangeStart.Add(30*time.Minute), false) + insertMessage(chat1.ID, database.ChatMessageRoleAssistant, 4, rangeStart.Add(45*time.Minute), true) + insertMessage(chat1.ID, database.ChatMessageRoleAssistant, 8, rangeEnd.Add(-time.Second), false) + // Tool rows count because runtime totals are role-agnostic. + insertMessage(chat1.ID, database.ChatMessageRoleTool, 64, rangeStart.Add(20*time.Minute), false) // Not counted: before the range, on the exclusive end boundary, and a // NULL runtime (runtime 0 is stored as NULL). - insertMessage(chat1.ID, 16, rangeStart.Add(-time.Second), false) - insertMessage(chat1.ID, 32, rangeEnd, false) - insertMessage(chat1.ID, 0, rangeStart.Add(10*time.Minute), false) + insertMessage(chat1.ID, database.ChatMessageRoleAssistant, 16, rangeStart.Add(-time.Second), false) + insertMessage(chat1.ID, database.ChatMessageRoleAssistant, 32, rangeEnd, false) + insertMessage(chat1.ID, database.ChatMessageRoleAssistant, 0, rangeStart.Add(10*time.Minute), false) total, err = db.GetTotalChatMessageRuntimeMsInRange(ctx, database.GetTotalChatMessageRuntimeMsInRangeParams{ StartTime: rangeStart, EndTime: rangeEnd, }) require.NoError(t, err) - require.EqualValues(t, 15, total) + require.EqualValues(t, 79, total) } func TestListUsageEventCreatedAtsByTypeSince(t *testing.T) { diff --git a/coderd/usage/usagetypes/events.go b/coderd/usage/usagetypes/events.go index 9004fb2514c..a990b4c6f7e 100644 --- a/coderd/usage/usagetypes/events.go +++ b/coderd/usage/usagetypes/events.go @@ -203,11 +203,15 @@ func (e HBAISeats) Fields() map[string]any { } // HBAgentRuntime is the event associated with hb_agent_runtime_v1. RuntimeMs -// is the total agent-loop runtime in milliseconds consumed by Coder Agents -// (chats) in one UTC hour. Each measured step spans model streaming (including -// provider-executed tools) and stream retries, and ends when the model stream -// finishes. Time spent executing local tools between steps, including -// sub-agents that bill their own model calls, is excluded. +// is total Coder Agent chat runtime in milliseconds for one UTC hour. +// +// Model steps bill provider streaming. Local tool batches bill the union of +// billed execution intervals, so parallel calls count once and serial calls +// count only from their own start. +// +// Excluded: sub-agent orchestration, client and external-agent work, user or +// idle waits, and retry backoff. Server-executed tools count even when their +// work runs in a connected workspace. // // This measures the new Coder Agents (the `chats` tables), not the deprecated // Tasks counted by dc_managed_agents_v1. diff --git a/coderd/x/chatd/ARCHITECTURE.md b/coderd/x/chatd/ARCHITECTURE.md index d9f1e16f8f6..0b68adeea4e 100644 --- a/coderd/x/chatd/ARCHITECTURE.md +++ b/coderd/x/chatd/ARCHITECTURE.md @@ -724,7 +724,9 @@ 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. +- `ModelInvokedAt(chat_id, history_version, generation_attempt)`: returns the `StartModelInvocation` stamp, or zero if none exists. Interrupt handling reads it 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. @@ -930,10 +932,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. +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 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/attempt.go b/coderd/x/chatd/attempt.go index fb967fb2398..580d113716b 100644 --- a/coderd/x/chatd/attempt.go +++ b/coderd/x/chatd/attempt.go @@ -31,6 +31,9 @@ type stepData struct { ContextLimit sql.NullInt64 Runtime time.Duration + // BatchRuntime is the local-tool batch window. Model steps use Runtime. + BatchRuntime time.Duration + ToolCallCreatedAt map[string]time.Time ToolResultCreatedAt map[string]time.Time ReasoningStartedAt []time.Time diff --git a/coderd/x/chatd/chatd_test.go b/coderd/x/chatd/chatd_test.go index 11865529fb1..8b108ae766e 100644 --- a/coderd/x/chatd/chatd_test.go +++ b/coderd/x/chatd/chatd_test.go @@ -6437,6 +6437,15 @@ func TestActiveServer_ToolExecutionAndPolicy(t *testing.T) { require.False(t, result.ProviderExecuted) } } + + messages := chatMessages(ctx, t, db, chat.ID) + billedToolRows := 0 + for _, msg := range messages { + if msg.Role == database.ChatMessageRoleTool && msg.RuntimeMs.Valid { + billedToolRows++ + } + } + require.LessOrEqual(t, billedToolRows, 1) }) } diff --git a/coderd/x/chatd/chatloop/chatloop.go b/coderd/x/chatd/chatloop/chatloop.go index e7586520f1e..0a0545da5b9 100644 --- a/coderd/x/chatd/chatloop/chatloop.go +++ b/coderd/x/chatd/chatloop/chatloop.go @@ -73,11 +73,12 @@ type PersistedStep struct { Content []fantasy.Content Usage fantasy.Usage ContextLimit sql.NullInt64 - // Runtime is the wall-clock duration of the model invocation - // that produced this step's content, measured from just before - // the provider stream is opened until the stream is fully - // consumed. + // Runtime is the wall-clock duration from opening to consuming the + // model stream. Runtime time.Duration + // BatchRuntime is the union of billed local-tool execution intervals. + // Parallel calls count once and serial calls count from their own start. + BatchRuntime time.Duration // PendingDynamicToolCalls lists tool calls that target // dynamic tools. When non-empty the chatloop exits with // ErrDynamicToolCall so the caller can execute them @@ -265,17 +266,26 @@ type ExecuteLocalToolsOptions struct { // is renamed but old chat histories still reference the old name. ToolNameAliases map[string]string + // UnbilledToolNames lists called tool names excluded from the batch + // window. Include deprecated aliases. + UnbilledToolNames map[string]bool + // OnToolStart fires when each local call begins. Serial calls may start + // after concurrent siblings settle, so interrupts bill actual starts and + // skip calls that never run. dispatchIndex identifies the dispatch-order + // occurrence. + OnToolStart func(dispatchIndex int, startedAt time.Time) + // OnToolComplete fires concurrently as each local call finishes, before + // ordered results publish. Interrupt billing uses the live timestamp; + // dispatchIndex identifies the dispatch-order occurrence. + // The callback must be concurrency-safe. + OnToolComplete func(dispatchIndex int, completedAt time.Time) + PublishMessagePart func(codersdk.ChatMessageRole, codersdk.ChatMessagePart) Logger slog.Logger Metrics *Metrics Clock quartz.Clock } -// ToolExecutionOutcome is the durable tool-result content from one batch. -type ToolExecutionOutcome struct { - Step PersistedStep -} - // GenerateCompactionOptions configures one context compaction call. type GenerateCompactionOptions struct { Model fantasy.LanguageModel @@ -532,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() } @@ -554,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)) @@ -564,7 +574,7 @@ func ExecuteLocalTools(ctx context.Context, opts ExecuteLocalToolsOptions) (Tool } } if len(localCalls) == 0 { - return ToolExecutionOutcome{}, nil + return PersistedStep{}, nil } var result stepResult @@ -586,15 +596,26 @@ 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) + batchStart := clockNow(opts.Clock) + // Keep completions by occurrence so duplicate IDs cannot collapse them. + orderedCompletions := make([]time.Time, 0, len(localCalls)) + // Keep starts by occurrence. Serial calls may begin after unbilled waits. + orderedStarts := make([]time.Time, len(localCalls)) + onToolStart := func(dispatchIndex int, startedAt time.Time) { + orderedStarts[dispatchIndex] = startedAt + if opts.OnToolStart != nil { + opts.OnToolStart(dispatchIndex, startedAt) + } + } toolResults := executeTools( ctx, opts.Clock, @@ -611,7 +632,11 @@ func ExecuteLocalTools(ctx context.Context, opts ExecuteLocalToolsOptions) (Tool opts.BuiltinToolNames, maxResultBytes, opts.ToolNameAliases, + batchStart, + onToolStart, + opts.OnToolComplete, func(tr fantasy.ToolResultContent, completedAt time.Time) { + orderedCompletions = append(orderedCompletions, completedAt) recordToolResultTimestamp(&result, tr.ToolCallID, completedAt) publishToolAttachments(ctx, opts.Logger, tr, completedAt, publishMessagePart) ssePart := chatprompt.PartFromContentWithLogger(ctx, opts.Logger, tr) @@ -620,15 +645,83 @@ 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) } - return ToolExecutionOutcome{Step: PersistedStep{ + return PersistedStep{ Content: result.content, ToolResultCreatedAt: result.toolResultCreatedAt, - }}, nil + BatchRuntime: billableBatchDuration( + batchStart, + localCalls, + orderedStarts, + orderedCompletions, + opts.UnbilledToolNames, + ), + }, nil +} + +// billableBatchDuration returns the union of billed execution intervals. +// Concurrent calls start at batchStart; serial calls use their recorded starts. +// Unbilled tools and gaps between billed intervals do not count. +func billableBatchDuration( + batchStart time.Time, + toolCalls []fantasy.ToolCallContent, + starts []time.Time, + completions []time.Time, + unbilledToolNames map[string]bool, +) time.Duration { + intervals := make([]BilledInterval, 0, len(toolCalls)) + for i, tc := range toolCalls { + if i >= len(completions) { + break + } + if unbilledToolNames[tc.ToolName] || completions[i].IsZero() { + continue + } + start := batchStart + if i < len(starts) && !starts[i].IsZero() { + start = starts[i] + } + intervals = append(intervals, BilledInterval{Start: start, End: completions[i]}) + } + return BilledIntervalsDuration(intervals) +} + +// BilledInterval is one billed tool call's execution window. +type BilledInterval struct { + Start time.Time + End time.Time +} + +// BilledIntervalsDuration returns the union duration of valid intervals. +// Overlaps count once, gaps do not, and inverted intervals are ignored. +// Committed and interrupted batches share this helper. +func BilledIntervalsDuration(intervals []BilledInterval) time.Duration { + valid := slices.DeleteFunc(slices.Clone(intervals), func(iv BilledInterval) bool { + return iv.End.Before(iv.Start) + }) + if len(valid) == 0 { + return 0 + } + slices.SortFunc(valid, func(a, b BilledInterval) int { + return a.Start.Compare(b.Start) + }) + curStart, curEnd := valid[0].Start, valid[0].End + var total time.Duration + for _, iv := range valid[1:] { + if iv.Start.After(curEnd) { + total += curEnd.Sub(curStart) + curStart, curEnd = iv.Start, iv.End + continue + } + if iv.End.After(curEnd) { + curEnd = iv.End + } + } + return total + curEnd.Sub(curStart) } // prepareMessagesForRequest applies the prompt preparation pipeline used @@ -1056,10 +1149,9 @@ func processStepStream( return result, nil } -// executeTools runs all tool calls concurrently after the stream -// completes. Results are published via onResult in the original -// tool-call order after all tools finish, preserving deterministic -// event ordering for SSE subscribers. +// executeTools runs non-serial calls concurrently, then SerialToolCalls in +// call order. Results publish in original order after all tools finish; +// onComplete fires as each tool finishes. func executeTools( ctx context.Context, clock quartz.Clock, @@ -1075,6 +1167,9 @@ func executeTools( builtinToolNames map[string]bool, maxResultBytes int, toolNameAliases map[string]string, + batchStart time.Time, + onStart func(dispatchIndex int, startedAt time.Time), + onComplete func(dispatchIndex int, completedAt time.Time), onResult func(fantasy.ToolResultContent, time.Time), ) []fantasy.ToolResultContent { if len(toolCalls) == 0 { @@ -1142,6 +1237,9 @@ func executeTools( // Captured per call so parallel tools get // accurate individual completion times. completedAt[i] = clockNow(clock) + if onComplete != nil { + onComplete(i, completedAt[i]) + } }() results[i] = executeSingleTool( ctx, @@ -1160,12 +1258,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 { @@ -1173,6 +1267,9 @@ func executeTools( serialIndexes = append(serialIndexes, i) continue } + if onStart != nil { + onStart(i, batchStart) + } wg.Add(1) go func() { defer wg.Done() @@ -1181,9 +1278,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) { @@ -1193,6 +1288,10 @@ func executeTools( notifyStepToolResultObservers(toolMap, toolNameAliases, localToolCalls, observed, settled) for _, i := range serialIndexes { + // Stamp serial calls at launch, not batch start. + if onStart != nil { + onStart(i, clockNow(clock)) + } runCall(i, localToolCalls[i]) } diff --git a/coderd/x/chatd/chatloop/chatloop_run_internal_test.go b/coderd/x/chatd/chatloop/chatloop_run_internal_test.go index 7d35d2588e7..44c2a68db52 100644 --- a/coderd/x/chatd/chatloop/chatloop_run_internal_test.go +++ b/coderd/x/chatd/chatloop/chatloop_run_internal_test.go @@ -949,6 +949,9 @@ func TestExecuteToolsNotifiesStepToolCallObservers(t *testing.T) { map[string]bool{}, defaultToolResultBytes, map[string]string{"observer_alias": "observer_tool"}, + time.Time{}, + nil, + nil, nil, ) @@ -1020,6 +1023,9 @@ func TestExecuteToolsNotifiesStepToolResultObservers(t *testing.T) { map[string]bool{}, defaultToolResultBytes, map[string]string{"observer_alias": "observer_tool"}, + time.Time{}, + nil, + nil, nil, ) @@ -1097,6 +1103,9 @@ func TestExecuteToolsReconcilesResultsBeforeSerialCalls(t *testing.T) { map[string]bool{}, defaultToolResultBytes, nil, + time.Time{}, + nil, + nil, nil, ) @@ -1164,6 +1173,9 @@ func TestExecuteToolsSerialToolCallOrder(t *testing.T) { map[string]bool{}, defaultToolResultBytes, nil, + time.Time{}, + nil, + nil, nil, ) diff --git a/coderd/x/chatd/chatloop/runtime_test.go b/coderd/x/chatd/chatloop/runtime_test.go index b227b26b694..005054d5569 100644 --- a/coderd/x/chatd/chatloop/runtime_test.go +++ b/coderd/x/chatd/chatloop/runtime_test.go @@ -6,11 +6,13 @@ import ( "time" "charm.land/fantasy" + "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "golang.org/x/xerrors" "github.com/coder/coder/v2/coderd/x/chatd/chatloop" "github.com/coder/coder/v2/coderd/x/chatd/chattest" + "github.com/coder/coder/v2/testutil" "github.com/coder/quartz" ) @@ -159,3 +161,546 @@ func TestGenerateCompaction_RecordsRuntime(t *testing.T) { require.Len(t, startedAt, 1) require.Equal(t, result.Runtime, clock.Since(startedAt[0])) } + +// executeToolBatch lets tests release trapped clock events in order, so +// goroutines cannot race clock advances. +func executeToolBatch( + t *testing.T, + clock *quartz.Mock, + opts chatloop.ExecuteLocalToolsOptions, +) <-chan chatloop.PersistedStep { + t.Helper() + opts.Clock = clock + resultCh := make(chan chatloop.PersistedStep, 1) + go func() { + outcome, err := chatloop.ExecuteLocalTools(context.Background(), opts) + assert.NoError(t, err) + resultCh <- outcome + }() + return resultCh +} + +func blockingTool(name string, release <-chan struct{}, response fantasy.ToolResponse) fantasy.AgentTool { + return fantasy.NewAgentTool( + name, + "test tool that completes when released", + func(_ context.Context, _ struct{}, _ fantasy.ToolCall) (fantasy.ToolResponse, error) { + <-release + return response, nil + }, + ) +} + +type serialTool struct { + fantasy.AgentTool +} + +func (serialTool) SerialToolCalls() bool { return true } + +func TestExecuteLocalTools_BatchWindowIsMaxNotSum(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitShort) + clock := quartz.NewMock(t) + trap := clock.Trap().Now() + defer trap.Close() + + fastGo := make(chan struct{}) + slowGo := make(chan struct{}) + resultCh := executeToolBatch(t, clock, chatloop.ExecuteLocalToolsOptions{ + Tools: []fantasy.AgentTool{ + blockingTool("fast_tool", fastGo, fantasy.NewTextResponse("done")), + blockingTool("slow_tool", slowGo, fantasy.NewTextErrorResponse("blew up")), + }, + ActiveTools: []string{"fast_tool", "slow_tool"}, + ToolCalls: []fantasy.ToolCallContent{ + {ToolCallID: "call-fast", ToolName: "fast_tool", Input: "{}"}, + {ToolCallID: "call-slow", ToolName: "slow_tool", Input: "{}"}, + }, + }) + + trap.MustWait(ctx).MustRelease(ctx) + clock.Advance(10 * time.Second) + close(fastGo) + trap.MustWait(ctx).MustRelease(ctx) + clock.Advance(50 * time.Second) + close(slowGo) + trap.MustWait(ctx).MustRelease(ctx) + + outcome := testutil.RequireReceive(ctx, t, resultCh) + require.Equal(t, 60*time.Second, outcome.BatchRuntime) +} + +func TestExecuteLocalTools_SimultaneousCompletionsBillOnce(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitShort) + clock := quartz.NewMock(t) + trap := clock.Trap().Now() + defer trap.Close() + + release := make(chan struct{}) + resultCh := executeToolBatch(t, clock, chatloop.ExecuteLocalToolsOptions{ + Tools: []fantasy.AgentTool{ + blockingTool("read_tool", release, fantasy.NewTextResponse("done")), + }, + ActiveTools: []string{"read_tool"}, + ToolCalls: []fantasy.ToolCallContent{ + {ToolCallID: "call-1", ToolName: "read_tool", Input: "{}"}, + {ToolCallID: "call-2", ToolName: "read_tool", Input: "{}"}, + {ToolCallID: "call-3", ToolName: "read_tool", Input: "{}"}, + }, + }) + + trap.MustWait(ctx).MustRelease(ctx) + clock.Advance(10 * time.Second) + close(release) + for range 3 { + trap.MustWait(ctx).MustRelease(ctx) + } + + outcome := testutil.RequireReceive(ctx, t, resultCh) + require.Equal(t, 10*time.Second, outcome.BatchRuntime) +} + +func TestExecuteLocalTools_UnbilledToolNeverExtendsWindow(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitShort) + clock := quartz.NewMock(t) + trap := clock.Trap().Now() + defer trap.Close() + + executeGo := make(chan struct{}) + waitGo := make(chan struct{}) + resultCh := executeToolBatch(t, clock, chatloop.ExecuteLocalToolsOptions{ + Tools: []fantasy.AgentTool{ + blockingTool("execute", executeGo, fantasy.NewTextResponse("done")), + blockingTool("wait_agent", waitGo, fantasy.NewTextResponse("child report")), + }, + ActiveTools: []string{"execute", "wait_agent"}, + UnbilledToolNames: map[string]bool{"wait_agent": true}, + ToolCalls: []fantasy.ToolCallContent{ + {ToolCallID: "call-execute", ToolName: "execute", Input: "{}"}, + {ToolCallID: "call-wait", ToolName: "wait_agent", Input: "{}"}, + }, + }) + + trap.MustWait(ctx).MustRelease(ctx) + clock.Advance(10 * time.Second) + close(executeGo) + trap.MustWait(ctx).MustRelease(ctx) + clock.Advance(50 * time.Second) + close(waitGo) + trap.MustWait(ctx).MustRelease(ctx) + + outcome := testutil.RequireReceive(ctx, t, resultCh) + require.Equal(t, 10*time.Second, outcome.BatchRuntime) +} + +func TestExecuteLocalTools_UnbilledOnlyBatchBillsNothing(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitShort) + clock := quartz.NewMock(t) + trap := clock.Trap().Now() + defer trap.Close() + + waitGo := make(chan struct{}) + resultCh := executeToolBatch(t, clock, chatloop.ExecuteLocalToolsOptions{ + Tools: []fantasy.AgentTool{ + blockingTool("wait_agent", waitGo, fantasy.NewTextResponse("child report")), + }, + ActiveTools: []string{"wait_agent"}, + UnbilledToolNames: map[string]bool{"wait_agent": true}, + ToolCalls: []fantasy.ToolCallContent{ + {ToolCallID: "call-wait", ToolName: "wait_agent", Input: "{}"}, + }, + }) + + trap.MustWait(ctx).MustRelease(ctx) + clock.Advance(60 * time.Second) + close(waitGo) + trap.MustWait(ctx).MustRelease(ctx) + + outcome := testutil.RequireReceive(ctx, t, resultCh) + require.Zero(t, outcome.BatchRuntime) +} + +func TestExecuteLocalTools_AliasNamesClassifyAsCalled(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitShort) + clock := quartz.NewMock(t) + trap := clock.Trap().Now() + defer trap.Close() + + executeGo := make(chan struct{}) + legacyGo := make(chan struct{}) + resultCh := executeToolBatch(t, clock, chatloop.ExecuteLocalToolsOptions{ + Tools: []fantasy.AgentTool{ + blockingTool("execute", executeGo, fantasy.NewTextResponse("done")), + blockingTool("interrupt_agent", legacyGo, fantasy.NewTextResponse("stopped")), + }, + ActiveTools: []string{"execute", "interrupt_agent"}, + ToolNameAliases: map[string]string{"close_agent": "interrupt_agent"}, + UnbilledToolNames: map[string]bool{ + "interrupt_agent": true, + "close_agent": true, + }, + ToolCalls: []fantasy.ToolCallContent{ + {ToolCallID: "call-execute", ToolName: "execute", Input: "{}"}, + {ToolCallID: "call-legacy", ToolName: "close_agent", Input: "{}"}, + }, + }) + + trap.MustWait(ctx).MustRelease(ctx) + clock.Advance(10 * time.Second) + close(executeGo) + trap.MustWait(ctx).MustRelease(ctx) + clock.Advance(50 * time.Second) + close(legacyGo) + trap.MustWait(ctx).MustRelease(ctx) + + outcome := testutil.RequireReceive(ctx, t, resultCh) + require.Equal(t, 10*time.Second, outcome.BatchRuntime) +} + +func TestExecuteLocalTools_DuplicateToolCallIDsKeepOccurrenceCompletions(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitShort) + clock := quartz.NewMock(t) + trap := clock.Trap().Now() + defer trap.Close() + + slowGo := make(chan struct{}) + fastGo := make(chan struct{}) + resultCh := executeToolBatch(t, clock, chatloop.ExecuteLocalToolsOptions{ + Tools: []fantasy.AgentTool{ + blockingTool("slow_tool", slowGo, fantasy.NewTextResponse("done")), + blockingTool("fast_tool", fastGo, fantasy.NewTextResponse("done")), + }, + ActiveTools: []string{"slow_tool", "fast_tool"}, + ToolCalls: []fantasy.ToolCallContent{ + {ToolCallID: "call-dup", ToolName: "slow_tool", Input: "{}"}, + {ToolCallID: "call-dup", ToolName: "fast_tool", Input: "{}"}, + }, + }) + + trap.MustWait(ctx).MustRelease(ctx) + clock.Advance(10 * time.Second) + close(fastGo) + trap.MustWait(ctx).MustRelease(ctx) + clock.Advance(50 * time.Second) + close(slowGo) + trap.MustWait(ctx).MustRelease(ctx) + + outcome := testutil.RequireReceive(ctx, t, resultCh) + require.Equal(t, 60*time.Second, outcome.BatchRuntime) +} + +func TestExecuteLocalTools_ExecutionCallbacksFireOnlyForRuns(t *testing.T) { + t.Parallel() + + t.Run("started call records a paired lifecycle", func(t *testing.T) { + t.Parallel() + + starts := 0 + completions := 0 + startedWhenToolRan := false + tool := fantasy.NewAgentTool( + "fast_tool", + "test tool", + func(context.Context, struct{}, fantasy.ToolCall) (fantasy.ToolResponse, error) { + startedWhenToolRan = starts > 0 + return fantasy.NewTextResponse("done"), nil + }, + ) + outcome, err := chatloop.ExecuteLocalTools(context.Background(), chatloop.ExecuteLocalToolsOptions{ + Clock: quartz.NewMock(t), + Tools: []fantasy.AgentTool{tool}, + ActiveTools: []string{"fast_tool"}, + OnToolStart: func(int, time.Time) { + starts++ + }, + OnToolComplete: func(int, time.Time) { + completions++ + }, + ToolCalls: []fantasy.ToolCallContent{ + {ToolCallID: "call-1", ToolName: "fast_tool", Input: "{}"}, + }, + }) + require.NoError(t, err) + require.Len(t, outcome.Content, 1) + require.Equal(t, 1, starts) + require.Equal(t, 1, completions) + require.True(t, startedWhenToolRan, "the start callback must fire before the tool runs") + }) + + t.Run("canceled context records no lifecycle", func(t *testing.T) { + t.Parallel() + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + started := false + completed := false + _, err := chatloop.ExecuteLocalTools(ctx, chatloop.ExecuteLocalToolsOptions{ + Clock: quartz.NewMock(t), + OnToolStart: func(int, time.Time) { started = true }, + OnToolComplete: func(int, time.Time) { completed = true }, + ToolCalls: []fantasy.ToolCallContent{ + {ToolCallID: "call-1", ToolName: "fast_tool", Input: "{}"}, + }, + }) + require.ErrorIs(t, err, context.Canceled) + require.False(t, started) + require.False(t, completed) + }) + + t.Run("exclusive violation records no lifecycle", func(t *testing.T) { + t.Parallel() + + started := false + completed := false + outcome, err := chatloop.ExecuteLocalTools(context.Background(), chatloop.ExecuteLocalToolsOptions{ + Clock: quartz.NewMock(t), + ExclusiveToolNames: map[string]bool{"exclusive_tool": true}, + OnToolStart: func(int, time.Time) { started = true }, + OnToolComplete: func(int, time.Time) { completed = true }, + ToolCalls: []fantasy.ToolCallContent{ + {ToolCallID: "call-1", ToolName: "exclusive_tool", Input: "{}"}, + {ToolCallID: "call-2", ToolName: "fast_tool", Input: "{}"}, + }, + }) + require.NoError(t, err) + require.Len(t, outcome.Content, 2, "the whole batch resolves to synthesized policy errors") + require.False(t, started) + require.False(t, completed) + }) +} + +func TestExecuteLocalTools_EmptyToolCallIDStillBillsWindow(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitShort) + clock := quartz.NewMock(t) + trap := clock.Trap().Now() + defer trap.Close() + + idlessGo := make(chan struct{}) + fastGo := make(chan struct{}) + resultCh := executeToolBatch(t, clock, chatloop.ExecuteLocalToolsOptions{ + Tools: []fantasy.AgentTool{ + blockingTool("idless_tool", idlessGo, fantasy.NewTextResponse("done")), + blockingTool("fast_tool", fastGo, fantasy.NewTextResponse("done")), + }, + ActiveTools: []string{"idless_tool", "fast_tool"}, + ToolCalls: []fantasy.ToolCallContent{ + {ToolCallID: "", ToolName: "idless_tool", Input: "{}"}, + {ToolCallID: "call-fast", ToolName: "fast_tool", Input: "{}"}, + }, + }) + + trap.MustWait(ctx).MustRelease(ctx) + clock.Advance(10 * time.Second) + close(fastGo) + trap.MustWait(ctx).MustRelease(ctx) + clock.Advance(50 * time.Second) + close(idlessGo) + trap.MustWait(ctx).MustRelease(ctx) + + outcome := testutil.RequireReceive(ctx, t, resultCh) + require.Equal(t, 60*time.Second, outcome.BatchRuntime) +} + +func TestExecuteLocalTools_OnToolCompleteReportsLiveCompletions(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitShort) + clock := quartz.NewMock(t) + trap := clock.Trap().Now() + defer trap.Close() + + type completion struct { + dispatchIndex int + completedAt time.Time + } + completionCh := make(chan completion, 2) + fastGo := make(chan struct{}) + slowGo := make(chan struct{}) + resultCh := executeToolBatch(t, clock, chatloop.ExecuteLocalToolsOptions{ + Tools: []fantasy.AgentTool{ + blockingTool("fast_tool", fastGo, fantasy.NewTextResponse("done")), + blockingTool("slow_tool", slowGo, fantasy.NewTextResponse("done")), + }, + ActiveTools: []string{"fast_tool", "slow_tool"}, + OnToolComplete: func(dispatchIndex int, completedAt time.Time) { + completionCh <- completion{dispatchIndex: dispatchIndex, completedAt: completedAt} + }, + ToolCalls: []fantasy.ToolCallContent{ + {ToolCallID: "call-fast", ToolName: "fast_tool", Input: "{}"}, + {ToolCallID: "call-slow", ToolName: "slow_tool", Input: "{}"}, + }, + }) + + trap.MustWait(ctx).MustRelease(ctx) + clock.Advance(10 * time.Second) + close(fastGo) + trap.MustWait(ctx).MustRelease(ctx) + fast := testutil.RequireReceive(ctx, t, completionCh) + require.Equal(t, 0, fast.dispatchIndex) + clock.Advance(50 * time.Second) + close(slowGo) + trap.MustWait(ctx).MustRelease(ctx) + slow := testutil.RequireReceive(ctx, t, completionCh) + require.Equal(t, 1, slow.dispatchIndex) + require.Equal(t, 50*time.Second, slow.completedAt.Sub(fast.completedAt)) + + outcome := testutil.RequireReceive(ctx, t, resultCh) + require.Equal(t, map[string]time.Time{ + "call-fast": fast.completedAt, + "call-slow": slow.completedAt, + }, outcome.ToolResultCreatedAt) +} + +func TestExecuteLocalTools_SerialCallBillsFromItsOwnStart(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitShort) + clock := quartz.NewMock(t) + trap := clock.Trap().Now() + defer trap.Close() + + waitGo := make(chan struct{}) + serialGo := make(chan struct{}) + resultCh := executeToolBatch(t, clock, chatloop.ExecuteLocalToolsOptions{ + Tools: []fantasy.AgentTool{ + blockingTool("wait_agent", waitGo, fantasy.NewTextResponse("child report")), + serialTool{blockingTool("serial_tool", serialGo, fantasy.NewTextResponse("done"))}, + }, + ActiveTools: []string{"wait_agent", "serial_tool"}, + UnbilledToolNames: map[string]bool{"wait_agent": true}, + ToolCalls: []fantasy.ToolCallContent{ + {ToolCallID: "call-wait", ToolName: "wait_agent", Input: "{}"}, + {ToolCallID: "call-serial", ToolName: "serial_tool", Input: "{}"}, + }, + }) + + trap.MustWait(ctx).MustRelease(ctx) + clock.Advance(10 * time.Minute) + close(waitGo) + trap.MustWait(ctx).MustRelease(ctx) + // Release the serial start timestamp. + trap.MustWait(ctx).MustRelease(ctx) + clock.Advance(2 * time.Second) + close(serialGo) + trap.MustWait(ctx).MustRelease(ctx) + + outcome := testutil.RequireReceive(ctx, t, resultCh) + require.Equal(t, 2*time.Second, outcome.BatchRuntime, + "a serial call bills its own execution, not the unbilled wait that delayed its launch") +} + +func TestExecuteLocalTools_SerialAfterBilledSiblingBillsUnion(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitShort) + clock := quartz.NewMock(t) + trap := clock.Trap().Now() + defer trap.Close() + + execGo := make(chan struct{}) + waitGo := make(chan struct{}) + serialGo := make(chan struct{}) + resultCh := executeToolBatch(t, clock, chatloop.ExecuteLocalToolsOptions{ + Tools: []fantasy.AgentTool{ + blockingTool("execute", execGo, fantasy.NewTextResponse("done")), + blockingTool("wait_agent", waitGo, fantasy.NewTextResponse("child report")), + serialTool{blockingTool("serial_tool", serialGo, fantasy.NewTextResponse("done"))}, + }, + ActiveTools: []string{"execute", "wait_agent", "serial_tool"}, + UnbilledToolNames: map[string]bool{"wait_agent": true}, + ToolCalls: []fantasy.ToolCallContent{ + {ToolCallID: "call-execute", ToolName: "execute", Input: "{}"}, + {ToolCallID: "call-wait", ToolName: "wait_agent", Input: "{}"}, + {ToolCallID: "call-serial", ToolName: "serial_tool", Input: "{}"}, + }, + }) + + trap.MustWait(ctx).MustRelease(ctx) + clock.Advance(3 * time.Second) + close(execGo) + trap.MustWait(ctx).MustRelease(ctx) + clock.Advance(7 * time.Second) + close(waitGo) + trap.MustWait(ctx).MustRelease(ctx) + // Release the serial start timestamp. + trap.MustWait(ctx).MustRelease(ctx) + clock.Advance(2 * time.Second) + close(serialGo) + trap.MustWait(ctx).MustRelease(ctx) + + outcome := testutil.RequireReceive(ctx, t, resultCh) + require.Equal(t, 5*time.Second, outcome.BatchRuntime, + "the 3s concurrent window and the 2s serial window bill; the 7s span where only wait_agent ran does not") +} + +func TestBilledIntervalsDuration(t *testing.T) { + t.Parallel() + + base := time.Date(2024, 1, 1, 0, 0, 0, 0, time.UTC) + at := base.Add + for _, tc := range []struct { + name string + intervals []chatloop.BilledInterval + want time.Duration + }{ + {name: "empty", want: 0}, + { + name: "overlapping intervals bill once", + intervals: []chatloop.BilledInterval{ + {Start: at(0), End: at(10 * time.Second)}, + {Start: at(0), End: at(4 * time.Second)}, + }, + want: 10 * time.Second, + }, + { + name: "gap between intervals is not billed", + intervals: []chatloop.BilledInterval{ + {Start: at(0), End: at(3 * time.Second)}, + {Start: at(10 * time.Second), End: at(12 * time.Second)}, + }, + want: 5 * time.Second, + }, + { + name: "unsorted contained interval adds nothing", + intervals: []chatloop.BilledInterval{ + {Start: at(2 * time.Second), End: at(4 * time.Second)}, + {Start: at(0), End: at(10 * time.Second)}, + }, + want: 10 * time.Second, + }, + { + name: "touching intervals merge without a gap", + intervals: []chatloop.BilledInterval{ + {Start: at(0), End: at(3 * time.Second)}, + {Start: at(3 * time.Second), End: at(5 * time.Second)}, + }, + want: 5 * time.Second, + }, + { + name: "inverted interval is ignored", + intervals: []chatloop.BilledInterval{ + {Start: at(5 * time.Second), End: at(0)}, + {Start: at(0), End: at(2 * time.Second)}, + }, + want: 2 * time.Second, + }, + } { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + require.Equal(t, tc.want, chatloop.BilledIntervalsDuration(tc.intervals)) + }) + } +} diff --git a/coderd/x/chatd/generation.go b/coderd/x/chatd/generation.go index 8b1bb881458..ddd73a23c3f 100644 --- a/coderd/x/chatd/generation.go +++ b/coderd/x/chatd/generation.go @@ -802,7 +802,7 @@ func (s *taskStarter) admitStepToolCalls( countBatch() 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 { countBatch() @@ -833,10 +833,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) } // find_tools calls are counted here, at the single point every // model-emitted call passes through, because rejections upstream of @@ -860,9 +862,27 @@ 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) + var onToolComplete func(int, time.Time) + if !exclusiveRejected { + // 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) + } + 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{ Tools: prepared.Tools, ActiveTools: prepared.ActiveTools, @@ -876,6 +896,9 @@ func (s *taskStarter) executeLocalTools( ModelName: modelName, ContextLimit: prepared.ContextLimitFallback, ToolNameAliases: subagentToolNameAliases, + UnbilledToolNames: unbilledSubagentToolNames, + OnToolStart: onToolStart, + OnToolComplete: onToolComplete, PublishMessagePart: attempt.publish, Logger: s.opts.Logger, Metrics: s.server.metrics, @@ -888,18 +911,19 @@ 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) + chathooks.RestoreToolCallOrder(outcome.Content, decision.localToolCalls) + step := stepDataFromPersisted(outcome) messages, err := buildCommitStepMessages(buildCommitStepMessagesInput{ modelConfigID: prepared.ModelConfigID, - step: stepDataFromPersisted(outcome.Step), + step: step, toolNameToConfigID: prepared.ToolNameToConfigID, logger: s.opts.Logger, contentVersion: chatprompt.CurrentContentVersion, @@ -1088,6 +1112,12 @@ type generationAttempt struct { // can bill the window the step would have reported. It is always // non-nil when beginGenerationAttempt succeeds. startModelInvocation func() + // recordToolStart stamps an occurrence's actual start; serial calls may + // start after dispatch. It is always non-nil after beginGenerationAttempt. + recordToolStart func(callIndex int, startedAt time.Time) + // recordToolCompletion stamps an occurrence's completion. It is always + // non-nil after beginGenerationAttempt. + recordToolCompletion func(callIndex int, completedAt time.Time) // closeEpisode closes the attempt's buffer episode. It is always // non-nil when beginGenerationAttempt succeeds. closeEpisode func() @@ -1134,6 +1164,12 @@ func (s *taskStarter) beginGenerationAttempt( startModelInvocation: func() { _ = s.opts.MessagePartBuffer.StartModelInvocation(key) }, + recordToolStart: func(callIndex int, startedAt time.Time) { + _ = s.opts.MessagePartBuffer.RecordToolStart(key, callIndex, startedAt) + }, + recordToolCompletion: func(callIndex int, completedAt time.Time) { + _ = s.opts.MessagePartBuffer.RecordToolCompletion(key, callIndex, completedAt) + }, closeEpisode: func() { _ = s.opts.MessagePartBuffer.CloseEpisode(key) }, @@ -1543,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 52d749175bd..a27229b54f1 100644 --- a/coderd/x/chatd/message_conversion.go +++ b/coderd/x/chatd/message_conversion.go @@ -61,7 +61,7 @@ func buildCommitStepMessages(input buildCommitStepMessagesInput) (stepMessagesFo messages = append(messages, assistantMessage(input.modelConfigID, contentVersion, assistantContent, input.step)) } - for _, toolResult := range toolResults { + for i, toolResult := range toolResults { part := chatprompt.PartFromContentWithLogger(context.Background(), input.logger, toolResult) applyToolMetadata(&part, input.toolNameToConfigID) if part.ToolCallID != "" && input.step.ToolResultCreatedAt != nil { @@ -73,7 +73,12 @@ func buildCommitStepMessages(input buildCommitStepMessagesInput) (stepMessagesFo if err != nil { return stepMessagesForCommit{}, xerrors.Errorf("marshal tool result: %w", err) } - messages = append(messages, baseMessage(database.ChatMessageRoleTool, database.ChatMessageVisibilityBoth, input.modelConfigID, contentVersion, content)) + msg := baseMessage(database.ChatMessageRoleTool, database.ChatMessageVisibilityBoth, input.modelConfigID, contentVersion, content) + // Usage sums runtime_ms across rows, so store the batch once. + if i == 0 { + msg.RuntimeMs = nullInt64IfNonZero(input.step.BatchRuntime.Milliseconds()) + } + messages = append(messages, msg) } return stepMessagesForCommit{ @@ -631,13 +636,8 @@ type partialMessageConversionState struct { toolResults map[string]*partialToolResult toolResultOrder []string answered map[string]bool - // modelStreamedAssistant records whether any assistant part came - // from the model stream itself (text, reasoning, tool calls, - // sources). Tool execution also publishes assistant-role file - // parts for attachments; those alone must not attract the - // attempt's runtime, because tool batches are not billable. The - // buffer episode only carries a runtime when a provider stream - // was opened, so this is a second gate rather than the only one. + // modelStreamedAssistant distinguishes streamed content from tool + // attachment parts, which must not carry model runtime. modelStreamedAssistant bool } diff --git a/coderd/x/chatd/message_conversion_test.go b/coderd/x/chatd/message_conversion_test.go index c84935d72f3..ef0852fcb02 100644 --- a/coderd/x/chatd/message_conversion_test.go +++ b/coderd/x/chatd/message_conversion_test.go @@ -100,9 +100,63 @@ func TestBuildCommitStepMessages_LocalToolResultsBecomeToolMessages(t *testing.T require.JSONEq(t, `{"stdout":"/tmp"}`, string(toolParts[0].Result)) } -// A step with no model invocation (a local tool execution batch) must -// persist runtime_ms NULL: its wall time is not billable. -func TestBuildCommitStepMessages_ZeroRuntimeLeavesRuntimeNull(t *testing.T) { +func TestBuildCommitStepMessages_BatchRuntimeLandsOnFirstToolRow(t *testing.T) { + t.Parallel() + + got, err := buildCommitStepMessages(buildCommitStepMessagesInput{ + modelConfigID: uuid.New(), + contentVersion: chatprompt.CurrentContentVersion, + logger: slog.Make(), + step: stepData{ + Content: []fantasy.Content{ + fantasy.ToolResultContent{ + ToolCallID: "call-1", + ToolName: "read_file", + Result: fantasy.ToolResultOutputContentText{Text: `{"data":"fast"}`}, + }, + fantasy.ToolResultContent{ + ToolCallID: "call-2", + ToolName: "execute", + Result: fantasy.ToolResultOutputContentText{Text: `{"stdout":"/tmp"}`}, + }, + }, + BatchRuntime: 10 * time.Second, + }, + }) + require.NoError(t, err) + require.Len(t, got.Messages, 2) + require.Equal(t, sql.NullInt64{Int64: 10000, Valid: true}, got.Messages[0].RuntimeMs) + require.False(t, got.Messages[1].RuntimeMs.Valid) +} + +func TestBuildCommitStepMessages_BatchAttachmentAssistantRowStaysNull(t *testing.T) { + t.Parallel() + + got, err := buildCommitStepMessages(buildCommitStepMessagesInput{ + modelConfigID: uuid.New(), + contentVersion: chatprompt.CurrentContentVersion, + logger: slog.Make(), + step: stepData{ + Content: []fantasy.Content{ + fantasy.ToolResultContent{ + ToolCallID: "call-1", + ToolName: "attach_file", + Result: fantasy.ToolResultOutputContentText{Text: `{"ok":true}`}, + ClientMetadata: `{"attachments":[{"file_id":"` + uuid.NewString() + `","media_type":"image/png","name":"shot.png"}]}`, + }, + }, + BatchRuntime: 3 * time.Second, + }, + }) + require.NoError(t, err) + require.Len(t, got.Messages, 2) + require.Equal(t, database.ChatMessageRoleAssistant, got.Messages[0].Role) + require.False(t, got.Messages[0].RuntimeMs.Valid) + require.Equal(t, database.ChatMessageRoleTool, got.Messages[1].Role) + require.Equal(t, sql.NullInt64{Int64: 3000, Valid: true}, got.Messages[1].RuntimeMs) +} + +func TestBuildCommitStepMessages_ZeroBatchRuntimeLeavesRuntimeNull(t *testing.T) { t.Parallel() got, err := buildCommitStepMessages(buildCommitStepMessagesInput{ diff --git a/coderd/x/chatd/messagepartbuffer/message_part_buffer.go b/coderd/x/chatd/messagepartbuffer/message_part_buffer.go index 82f41648be6..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,12 +103,14 @@ 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{} + // toolCompletions stores started occurrences by unresolved-call position. + toolCompletions map[int]ToolCompletion + closed bool + closedAt time.Time + closedHeapItem *closedEpisodeItem + parts []Part + bytes int64 + subscribers map[*episodeSubscriber]struct{} } type closedEpisodeItem struct { @@ -218,6 +221,65 @@ func (b *Buffer) StartModelInvocation(key Key) error { return nil } +// ToolCompletion tracks a started tool-call occurrence. CompletedAt is zero +// while the call is unfinished. +type ToolCompletion struct { + StartedAt time.Time + CompletedAt time.Time +} + +// 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 { + return ErrMessagePartBufferClosed + } + episode, err := b.getEpisodeLocked(key) + if err != nil { + return err + } + if episode.closed { + return ErrEpisodeClosed + } + if callIndex < 0 { + return nil + } + if _, ok := episode.toolCompletions[callIndex]; ok { + return nil + } + if episode.toolCompletions == nil { + episode.toolCompletions = make(map[int]ToolCompletion) + } + episode.toolCompletions[callIndex] = ToolCompletion{StartedAt: startedAt} + return nil +} + +// RecordToolCompletion stamps a call as it finishes, so interrupts use the +// 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 { + return ErrMessagePartBufferClosed + } + episode, err := b.getEpisodeLocked(key) + if err != nil { + return err + } + if episode.closed { + return ErrEpisodeClosed + } + entry, ok := episode.toolCompletions[callIndex] + if ok && entry.CompletedAt.IsZero() { + entry.CompletedAt = completedAt + episode.toolCompletions[callIndex] = entry + } + return nil +} + // AddPart appends a part to an existing episode. // // Parts receive contiguous sequence numbers so stream endpoints can detect @@ -303,6 +365,19 @@ func (b *Buffer) ModelInvokedAt(key Key) time.Time { return episode.modelStartedAt } +// 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) map[int]ToolCompletion { + b.mu.Lock() + defer b.mu.Unlock() + episode := b.episodes[key] + if episode == nil { + return nil + } + return maps.Clone(episode.toolCompletions) +} + // 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..fc67ff4de3a 100644 --- a/coderd/x/chatd/messagepartbuffer/message_part_buffer_test.go +++ b/coderd/x/chatd/messagepartbuffer/message_part_buffer_test.go @@ -148,6 +148,66 @@ func TestBuffer_ModelInvokedAt(t *testing.T) { require.Zero(t, buffer.ModelInvokedAt(implicit)) } +func TestBuffer_ToolCompletions(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.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 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)) + 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, 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)) + clock.Advance(2 * time.Second) + firstCompletedAt := clock.Now() + require.NoError(t, buffer.RecordToolCompletion(key, 1, firstCompletedAt)) + completions := buffer.ToolCompletions(key) + require.Equal(t, map[int]messagepartbuffer.ToolCompletion{ + 1: {StartedAt: firstStartedAt, CompletedAt: firstCompletedAt}, + 2: {StartedAt: secondStartedAt, CompletedAt: secondCompletedAt}, + }, completions) + + completion := completions[2] + completion.CompletedAt = clock.Now() + completions[2] = completion + require.Equal(t, secondCompletedAt, buffer.ToolCompletions(key)[2].CompletedAt) + + clock.Advance(time.Second) + 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)) + require.ErrorIs(t, buffer.RecordToolStart(key, 0, clock.Now()), messagepartbuffer.ErrEpisodeClosed) + require.ErrorIs(t, buffer.RecordToolCompletion(key, 0, clock.Now()), messagepartbuffer.ErrEpisodeClosed) +} + func TestBuffer_SubscribeExistingReplaysThenStreamsLiveParts(t *testing.T) { t.Parallel() diff --git a/coderd/x/chatd/options.go b/coderd/x/chatd/options.go index b887b928baf..6c464240755 100644 --- a/coderd/x/chatd/options.go +++ b/coderd/x/chatd/options.go @@ -70,6 +70,9 @@ type chatWorkerTaskStartInput struct { DebugTurn *runnerDebugTurn SessionStart *sessionStartTracker StopNudges *stopNudgeTracker + // InterruptSnapshot carries one interrupt task's first episode snapshot + // across retries. Nil re-reads 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..d5a77b971e4 100644 --- a/coderd/x/chatd/runner.go +++ b/coderd/x/chatd/runner.go @@ -232,6 +232,9 @@ func (r *runner) spawnTaskIfNeeded(kind taskKind, state runnerStateUpdate) { SessionStart: &r.sessionStart, StopNudges: &r.stopNudges, } + if kind == taskKindInterrupt { + 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 a79bf81fa5d..18673c571ab 100644 --- a/coderd/x/chatd/subagent_catalog.go +++ b/coderd/x/chatd/subagent_catalog.go @@ -38,6 +38,18 @@ const ( "external or web research, parallel research, or tasks that may need edits." ) +// unbilledSubagentToolNames excludes parent-side orchestration because +// child chats bill their own runtime. Include deprecated aliases. +var unbilledSubagentToolNames = map[string]bool{ + spawnAgentToolName: true, + "wait_agent": true, + "message_agent": true, + "interrupt_agent": true, + "close_agent": true, + "list_agents": true, + listSubagentModelsToolName: true, +} + type spawnAgentArgs struct { Type string `json:"type"` Prompt string `json:"prompt"` diff --git a/coderd/x/chatd/subagent_internal_test.go b/coderd/x/chatd/subagent_internal_test.go index 2774d90f9cb..432bff2ce29 100644 --- a/coderd/x/chatd/subagent_internal_test.go +++ b/coderd/x/chatd/subagent_internal_test.go @@ -4262,6 +4262,25 @@ func TestAwaitSubagentCompletion(t *testing.T) { }) } +func TestUnbilledSubagentToolNamesMatchCatalog(t *testing.T) { + t.Parallel() + + db, ps := dbtestutil.NewDB(t) + server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{}) + ctx := chatdTestContext(t) + user, org, model := seedInternalChatDeps(t, db) + parent, _ := createParentChildChats(ctx, t, server, user, org, model) + + catalog := make(map[string]bool) + for _, tool := range server.subagentTools(ctx, func() database.Chat { return parent }, parent.LastModelConfigID) { + catalog[tool.Info().Name] = true + } + for alias := range subagentToolNameAliases { + catalog[alias] = true + } + require.Equal(t, catalog, unbilledSubagentToolNames) +} + func TestWaitAgentToolSchema(t *testing.T) { t.Parallel() diff --git a/coderd/x/chatd/tasks.go b/coderd/x/chatd/tasks.go index a3738b886db..c6460338049 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" @@ -240,7 +241,69 @@ func (o chatWorkerOptions) retryOptions() retryWrapperOptions { } } +// 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 + billing interruptEpisodeBilling + parts []messagepartbuffer.Part +} + +type interruptEpisodeBilling struct { + interruptedAt time.Time + modelInvokedAt time.Time + toolCompletions map[int]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), + toolCompletions: s.opts.MessagePartBuffer.ToolCompletions(key), + } + if err := s.opts.MessagePartBuffer.CloseEpisode(key); err != nil { + if ctx.Err() != nil { + return interruptEpisodeBilling{}, nil, errors.Join(errTaskExpectedExit, xerrors.Errorf("close message part episode: %w", err), ctx.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 + err = nil + } + if err != nil { + if ctx.Err() != nil { + return interruptEpisodeBilling{}, nil, errors.Join(errTaskExpectedExit, xerrors.Errorf("get message part episode: %w", err), ctx.Err()) + } + return interruptEpisodeBilling{}, 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 + // 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, + 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 { @@ -260,28 +323,28 @@ func (s *taskStarter) StartInterrupt(ctx context.Context, input chatWorkerTaskSt HistoryVersion: input.HistoryVersion, GenerationAttempt: chat.GenerationAttempt, } - modelInvokedAt := s.opts.MessagePartBuffer.ModelInvokedAt(key) - if err := s.opts.MessagePartBuffer.CloseEpisode(key); err != nil { - if ctx.Err() != nil { - return errors.Join(errTaskExpectedExit, xerrors.Errorf("close message part episode: %w", err), ctx.Err()) + 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. + episodeBilling = snapshot.billing + parts = snapshot.parts + } else { + episodeBilling, parts, err = s.closeInterruptEpisode(ctx, key) + if err != nil { + return 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()) + 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)} } - interruptedAt := s.opts.Clock.Now("chatworker", "interrupt") + interruptedAt := episodeBilling.interruptedAt 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, @@ -302,7 +365,9 @@ 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")) + // Reuse the captured interrupt instant so database delay and retries do + // not inflate billing. + committedCancels, err := committedPendingLocalToolCancellationMessages(ctx, store, chat, interruptedAt, episodeBilling.toolCompletions) if err != nil { return xerrors.Errorf("committed pending local tool cancellation messages: %w", err) } @@ -683,6 +748,7 @@ func committedPendingLocalToolCancellationMessages( store database.Store, chat database.Chat, interruptedAt time.Time, + toolCompletions map[int]messagepartbuffer.ToolCompletion, ) ([]chatstate.Message, error) { messages, err := store.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{ ChatID: chat.ID, @@ -698,8 +764,13 @@ func committedPendingLocalToolCancellationMessages( if len(localCalls) == 0 { return nil, nil } + var ( + windowEnd time.Time + windowRowIdx = -1 + intervals []chatloop.BilledInterval + ) 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) @@ -719,6 +790,32 @@ func committedPendingLocalToolCancellationMessages( ModelConfigID: uuid.NullUUID{UUID: chat.LastModelConfigID, Valid: chat.LastModelConfigID != uuid.Nil}, ContentVersion: chatprompt.CurrentContentVersion, }) + if unbilledSubagentToolNames[call.ToolName] { + continue + } + // Bill only matching started calls. Completed calls end at completion, + // running calls end at the interrupt, and ties keep the first call. + occurrence, ok := toolCompletions[i] + if !ok { + continue + } + start := occurrence.StartedAt + if start.IsZero() { + continue + } + end := occurrence.CompletedAt + if end.IsZero() { + end = interruptedAt + } + intervals = append(intervals, chatloop.BilledInterval{Start: start, End: end}) + if end.After(windowEnd) { + windowEnd = end + windowRowIdx = len(result) - 1 + } + } + // Bill the interval union once on the row whose interval ends last. + 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 a1265b270b9..30ae6037bff 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,6 +475,406 @@ func TestInterruptTask_PartialAssistantWithoutModelInvocationHasNoRuntime(t *tes require.False(t, assistant.RuntimeMs.Valid) } +type interruptedBatch struct { + chat database.Chat + starter *taskStarter + clock *quartz.Mock + key messagepartbuffer.Key + workerID uuid.UUID + runnerID uuid.UUID +} + +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() + 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{ + ChatID: b.chat.ID, + WorkerID: b.workerID, + RunnerID: b.runnerID, + 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}) + require.NoError(t, err) + return messages +} + +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{} +} + +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 + + // Keep advances below the buffer's 15-second cleanup tick. + batch.clock.Advance(2 * time.Second) + 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, batch.clock.Now())) + batch.clock.Advance(5 * time.Second) + + 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) +} + +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(`{}`)}, + }) + + // Use only the carried snapshot; the buffer has no episode state. + batch.clock.Advance(2 * time.Second) + startedAt := batch.clock.Now() + batch.clock.Advance(7 * time.Second) + snapshot := &interruptEpisodeSnapshot{ + loaded: true, + key: batch.key, + billing: interruptEpisodeBilling{ + interruptedAt: batch.clock.Now(), + toolCompletions: map[int]messagepartbuffer.ToolCompletion{0: {StartedAt: startedAt}}, + }, + } + + messages := batch.interruptWithSnapshot(t, f, snapshot) + execRow := findToolResultMessage(t, messages, execCallID) + require.Equal(t, sql.NullInt64{Int64: 7_000, Valid: true}, execRow.RuntimeMs) +} + +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.RecordToolStart(batch.key, 0, batch.clock.Now())) + batch.clock.Advance(3 * time.Second) + + 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) + + 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) + + messages := batch.interruptWithSnapshot(t, f, snapshot) + execRow := findToolResultMessage(t, messages, execCallID) + require.Equal(t, sql.NullInt64{Int64: 3_000, Valid: true}, execRow.RuntimeMs) +} + +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.RecordToolStart(batch.key, 0, batch.clock.Now())) + 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) +} + +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.RecordToolStart(batch.key, 0, batch.clock.Now())) + batch.clock.Advance(10 * time.Second) + + messages := batch.interrupt(t, f) + waitRow := findToolResultMessage(t, messages, waitCallID) + require.False(t, waitRow.RuntimeMs.Valid) +} + +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.RecordToolStart(batch.key, 0, 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") + } + } +} + +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.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, batch.clock.Now())) + batch.clock.Advance(3 * time.Second) + 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) + 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) +} + +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.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, 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") +} + +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.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, 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") +} + +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) + 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) + 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") + } + } +} + +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) + 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) + execRow := findToolResultMessage(t, messages, execCallID) + require.False(t, execRow.RuntimeMs.Valid) + waitRow := findToolResultMessage(t, messages, waitCallID) + require.False(t, waitRow.RuntimeMs.Valid) +} + +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(`{}`)}, + }) + + 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/coderd/x/chatd/toolinput.go b/coderd/x/chatd/toolinput.go index 4aa1a8df6a8..6e961f72219 100644 --- a/coderd/x/chatd/toolinput.go +++ b/coderd/x/chatd/toolinput.go @@ -15,16 +15,13 @@ 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 maps allowed calls back to the input +// order without relying on duplicate-prone IDs. 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 +31,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) }) }