mirror of
https://github.com/coder/coder.git
synced 2026-09-01 14:53:15 +08:00
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:
@@ -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) {
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
|
||||
@@ -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()
|
||||
|
||||
|
||||
@@ -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))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
|
||||
@@ -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{
|
||||
|
||||
@@ -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"`
|
||||
|
||||
@@ -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()
|
||||
|
||||
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user