feat(coderd): bill completed local tool batches in agent runtime (#28360)

**Stack**: #28359 → **#28360 (this PR)** → #28361#28362

## Summary

Completed local tool batches now count toward Coder Agent runtime
(`chat_messages.runtime_ms`, the source of truth for
`hb_agent_runtime_v1`), excluding sub-agent orchestration tools to avoid
double counting, and without multiplying runtime for parallel tool
calls.

Second of a four-PR stack splitting #28211. Stacked on #28359. Refs
CODAGT-928
(https://linear.app/codercom/issue/CODAGT-928/track-local-tool-execution-for-agent-runtime).

## Problem

Agent runtime previously measured only model invocation wall clock
(stream open to fully consumed). Time spent executing local tools
between steps, including file operations, terminal commands, workspace
provisioning, and MCP tools, was deliberately excluded, which
undercounts the product definition of "actively processing a task".

Naive inclusion has two hazards: a batch of parallel tool calls would
bill N windows for one wall-clock wait, and `wait_agent` would re-bill
child agents that already bill their own model and tool time.

## Fix

Each completed local tool batch bills one window: the union of the
billed tools' execution intervals, persisted as `runtime_ms` on a
dedicated usage record appended after the batch's tool-result rows. The
record is a tool-role message with `visibility='model'`, so it never
reaches the API, SSE, or clients, and prompt replay drops it because its
single internal `tool-batch-usage` part converts to no provider content.
Its content carries an audit payload (`billed_ms`, `billed_calls`) so
the billed window is inspectable after the fact; real tool-result rows
never carry batch-level runtime. Concurrent calls all start at batch
start, so 5 parallel 10s reads bill 10s, not 50s; serial calls
(`SerialToolCalls`) count only from their own launch, so unbilled waits
before them do not count.

Sub-agent orchestration tools (`spawn_agent`, `wait_agent`,
`message_agent`, `interrupt_agent`, `list_agents`,
`list_subagent_models`, plus the deprecated `close_agent` alias) never
extend the window: every chat, including children, applies the same
rules to its own runtime, so a parent's `wait_agent` window would double
count. A lone `wait_agent` bills 0 and appends no usage record;
`execute` 10s in parallel with `wait_agent` 60s bills 10s. A test pins
the unbilled set to the registered sub-agent tool catalog so they cannot
drift.

`GetTotalChatMessageRuntimeMsInRange` already sums `runtime_ms`
role-agnostically, so the usage record is picked up with no schema,
query, or cron changes.

Client-executed dynamic tools, external agents, parked/idle time, and
retry backoff remain unbilled. Interrupted batches still bill only the
model window; the rest of the stack adds partial-window billing via the
same usage record.

`ExecuteLocalToolsOptions` also gains an optional `ToolBillingRecorder`
(`RecordStart`/`RecordComplete`) observing per-occurrence dispatch-order
execution stamps. This PR wires no recorder; the next PR in the stack
connects it to the message part buffer for interrupt billing.

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