From d52521b9153bd8a86340841fe6f1f5459f7c36e8 Mon Sep 17 00:00:00 2001 From: li-xiu-qi Date: Sun, 25 Oct 2026 01:33:13 +0800 Subject: [PATCH] fix(session): include delegated agent usage in totals --- docs/orchestration-lifecycle.md | 16 +- packages/agent-core/src/agent-loop.ts | 17 +- packages/agent-core/test/agent-loop.test.ts | 51 ++++ packages/coding-agent/docs/account-usage.md | 2 +- packages/coding-agent/docs/extensions.md | 2 +- packages/coding-agent/docs/rpc.md | 2 +- packages/coding-agent/docs/session-format.md | 2 +- packages/coding-agent/docs/usage.md | 2 +- .../coding-agent/src/core/usage-totals.ts | 47 ++++ .../src/features/step-subagent.ts | 6 +- .../src/features/subagent/execute.ts | 22 +- .../src/features/subagent/lane-lifecycle.ts | 72 ++++- .../src/features/workflow/agent-runner.ts | 2 + .../src/features/workflow/runtime.ts | 20 ++ .../src/features/workflow/types.ts | 6 + .../src/step/usage/activity-engine.ts | 67 ++++- .../src/step/usage/json-projection.ts | 13 +- .../test/agent-session-stats.test.ts | 30 +++ .../test/step-subagent-events.test.ts | 248 ++++++++++++++++- packages/coding-agent/test/step-usage.test.ts | 93 ++++++- .../test/subagent-chain-result.test.ts | 30 +++ .../agent-session-delegated-usage.test.ts | 254 ++++++++++++++++++ 22 files changed, 968 insertions(+), 36 deletions(-) create mode 100644 packages/coding-agent/test/suite/agent-session-delegated-usage.test.ts diff --git a/docs/orchestration-lifecycle.md b/docs/orchestration-lifecycle.md index fd8ee27d..a8f5aa5c 100644 --- a/docs/orchestration-lifecycle.md +++ b/docs/orchestration-lifecycle.md @@ -105,9 +105,13 @@ not hide active tasks. Phase and log updates use the same live update path. Each `onUpdate` has readable text and a complete `WorkflowProgress` snapshot in `details`, including for RPC consumers. Snapshots are independent of later -mutations. The final result remains `WorkflowRunResult`; the current tool row -retains its last progress snapshot when displaying that result or an error. -Updates from late child cleanup are ignored once the tool call has settled. +mutations. The final result remains `WorkflowRunResult`; its `usage` field +contains the run's aggregate delegated usage, and `childSessionIds` identifies +the child transcripts covered by that parent-side total. Terminal workflow +errors carry the same fields on the thrown error, which the agent loop preserves +on the error tool result. The current tool row retains its last progress +snapshot when displaying that result or an error. Updates from late child +cleanup are ignored once the tool call has settled. This projection describes assigned tasks and lifecycle states; child tool and model output remain in the child trajectories. @@ -171,7 +175,11 @@ it finishes, fails, or is stopped. It wakes an idle parent. The body opens with summary such as `Lane lane_1a2b3c4d finished: 2 completed, 1 failed.`. That summary is followed by duration, usage, child session ids and JSONL paths, any dialogs the child raised that were auto-cancelled, failure reasons, and the -output. The text passes through the shared [tool-output retention +output. Usage accounting is independent of these notifications: each settled +lane turn writes one `step-subagent-usage` entry with its aggregate usage and +covered child session ids. Queued turns, interrupted turns and shutdown settle +usage without requiring a terminal notification, and the lane/turn key prevents +duplicate accounting. The text passes through the shared [tool-output retention boundary](../packages/coding-agent/docs/tool-output-retention.md): oversized output is saved under the session's `tool-output` directory and replaced by a preview with its path. `subscribe: "progress"` adds throttled progress messages diff --git a/packages/agent-core/src/agent-loop.ts b/packages/agent-core/src/agent-loop.ts index ef7ae859..0582758b 100644 --- a/packages/agent-core/src/agent-loop.ts +++ b/packages/agent-core/src/agent-loop.ts @@ -10,6 +10,7 @@ import { classifyProviderError, EventStream, type ToolResultMessage, + type Usage, validateToolArguments, } from "@step-harness/providers"; import { getDefaultStreamFn } from "./stream-fn.ts"; @@ -775,7 +776,7 @@ async function prepareToolCall( } catch (error) { return { kind: "immediate", - result: createErrorToolResult(error instanceof Error ? error.message : String(error)), + result: createThrownToolResult(error), isError: true, }; } @@ -816,7 +817,7 @@ async function executePreparedToolCall( acceptingUpdates = false; await Promise.all(updateEvents); return { - result: createErrorToolResult(error instanceof Error ? error.message : String(error)), + result: createThrownToolResult(error), isError: true, }; } finally { @@ -859,7 +860,7 @@ async function finalizeExecutedToolCall( isError = afterResult.isError ?? isError; } } catch (error) { - result = createErrorToolResult(error instanceof Error ? error.message : String(error)); + result = createThrownToolResult(error); isError = true; } } @@ -871,10 +872,16 @@ async function finalizeExecutedToolCall( }; } -function createErrorToolResult(message: string): AgentToolResult { +function createThrownToolResult(error: unknown): AgentToolResult { + const fields = error && typeof error === "object" ? (error as { usage?: Usage; details?: unknown }) : undefined; + return createErrorToolResult(error instanceof Error ? error.message : String(error), fields?.usage, fields?.details); +} + +function createErrorToolResult(message: string, usage?: Usage, details: unknown = {}): AgentToolResult { return { content: [{ type: "text", text: message }], - details: {}, + details, + ...(usage ? { usage } : {}), }; } diff --git a/packages/agent-core/test/agent-loop.test.ts b/packages/agent-core/test/agent-loop.test.ts index 39630f65..c0a2b23b 100644 --- a/packages/agent-core/test/agent-loop.test.ts +++ b/packages/agent-core/test/agent-loop.test.ts @@ -561,6 +561,57 @@ describe("agentLoop with AgentMessage", () => { }); }); + it("preserves usage and details from a tool that throws after nested model work", async () => { + const usage = { + input: 10, + output: 20, + cacheRead: 2, + cacheWrite: 3, + totalTokens: 35, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0.5 }, + }; + const tool: AgentTool = { + name: "workflow", + label: "Workflow", + description: "Run nodes", + parameters: Type.Object({}), + execute: async () => { + throw Object.assign(new Error("workflow failed"), { usage, details: { childSessionIds: ["child-1"] } }); + }, + }; + let calls = 0; + const stream = agentLoop( + [createUserMessage("run")], + { systemPrompt: "", messages: [], tools: [tool] }, + { model: createModel(), convertToLlm: identityConverter }, + undefined, + () => { + const response = new MockAssistantStream(); + response.push({ + type: "done", + reason: calls++ === 0 ? "toolUse" : "stop", + message: + calls === 1 + ? createAssistantMessage( + [{ type: "toolCall", id: "wf", name: "workflow", arguments: {} }], + "toolUse", + ) + : createAssistantMessage([{ type: "text", text: "failed" }]), + }); + return response; + }, + ); + for await (const _event of stream) { + // consume + } + expect((await stream.result()).find((message) => message.role === "toolResult")).toMatchObject({ + isError: true, + usage, + details: { childSessionIds: ["child-1"] }, + content: [{ type: "text", text: "workflow failed" }], + }); + }); + it("should not execute tool calls from a length-truncated assistant message", async () => { const toolSchema = Type.Object({ value: Type.String() }); const executed: string[] = []; diff --git a/packages/coding-agent/docs/account-usage.md b/packages/coding-agent/docs/account-usage.md index ecca18a1..e2322761 100644 --- a/packages/coding-agent/docs/account-usage.md +++ b/packages/coding-agent/docs/account-usage.md @@ -22,7 +22,7 @@ Account queries the personal Step Plan usage endpoint with the stored CN/Global CODING plans show 5-hour and weekly percentages **remaining** and reset times. Nonpositive upstream limits mean Unlimited. TOKEN plans show the remaining percentage for each subscription/top-up Credit bucket and reset/expiry times, without displaying raw Credit balances. Account and Status use compact 12-column bars with reduced-height blocks to represent remaining quota. A missing percentage, including a zero-total bucket, is shown as unavailable rather than 0% or 100%. Credit quantities remain exact decimal strings in the snapshot; percentages use the API's validated `used_percent` without converting Credit quantities to floating point. Unknown future bucket types are marked unsupported; partial bucket errors are not displayed as zero balance. -The current main-session summary is local and includes restored/inherited history and incremental compression usage. It excludes separate child-session files. Cache hit is cache reads divided by complete input; an empty denominator is Unknown. +The current main-session summary is local and includes restored/inherited history, incremental compression usage, and delegated usage recorded on parent tool results and per-turn background usage entries. It does not scan separate child-session files for this summary. Cache hit is cache reads divided by complete input; an empty denominator is Unknown. Account responses have a 30-second memory cache and a 5-minute display lifetime. Identity/profile/endpoint/key changes invalidate it. Requests are single-flight, time out after 12 seconds and are cancelled when the panel closes or changes tabs/sessions. Attempt generations reject late responses. TTL/cooldown use a monotonic clock; HTTP Retry-After supports seconds and HTTP-date. Visible temporary failures retain a stale snapshot; 401/403 clear the snapshot without deleting credentials. No background polling runs. diff --git a/packages/coding-agent/docs/extensions.md b/packages/coding-agent/docs/extensions.md index fdfe9f67..7082f540 100644 --- a/packages/coding-agent/docs/extensions.md +++ b/packages/coding-agent/docs/extensions.md @@ -2035,7 +2035,7 @@ pi.registerTool({ }); ``` -**Usage accounting:** If a tool makes nested LLM calls, return their combined `Usage` as `usage`. Step persists it on the tool result and includes it in footer, `/session`, and RPC session totals. `tool_result` handlers can inspect or replace this value. +**Usage accounting:** If a tool makes nested LLM calls, return their combined `Usage` as `usage`. Step persists it on the tool result and includes it in footer, `/session`, and RPC session totals. `tool_result` handlers can inspect or replace this value. If nested work has already started when `execute` throws, the error object may carry `usage` and `details`; the agent loop preserves them on the error tool result. **Signaling errors:** To mark a tool execution as failed (sets `isError: true` on the result and reports it to the LLM), throw an error from `execute`. Returning a value never sets the error flag regardless of what properties you include in the return object. diff --git a/packages/coding-agent/docs/rpc.md b/packages/coding-agent/docs/rpc.md index 303f08aa..c2c6b8fc 100644 --- a/packages/coding-agent/docs/rpc.md +++ b/packages/coding-agent/docs/rpc.md @@ -590,7 +590,7 @@ Response: } ``` -`tokens` and `cost` include assistant messages, usage reported by tools, and compaction/branch-summary generation across the full session. `contextUsage` contains the actual current context-window estimate used for compaction and footer display. +`tokens` and `cost` include assistant messages, usage reported by tools, per-turn background subagent usage entries, and compaction/branch-summary generation across the full session. `contextUsage` contains the actual current context-window estimate used for compaction and footer display. `contextUsage` is omitted when no model or context window is available. `contextUsage.tokens` and `contextUsage.percent` are `null` immediately after compaction until a fresh post-compaction assistant response provides valid usage data. diff --git a/packages/coding-agent/docs/session-format.md b/packages/coding-agent/docs/session-format.md index 87e74fb4..f44b785c 100644 --- a/packages/coding-agent/docs/session-format.md +++ b/packages/coding-agent/docs/session-format.md @@ -119,7 +119,7 @@ interface Usage { The exported `StopReason` type also includes `"pending"`, but that value is reserved for partial messages in streaming events. Terminal `done`/`error` messages replace it with a completion reason before step persists the assistant message, so `"pending"` should never appear in session JSONL. -Session totals and the cost breakdown use the shared [`usage-totals.ts`](../src/core/usage-totals.ts) ledger. It includes assistant, tool, compaction, branch-summary and incremental `step-auto-clm-usage` records; cumulative maintenance summaries are excluded. Token increments must contain nonnegative safe integers and must not overflow the accumulated total. A missing/invalid cost does not hide valid tokens. `SessionStats.tokens.input` excludes caches, which are reported separately; the [Account summary](account-usage.md) combines them into complete input. +Session totals and the cost breakdown use the shared [`usage-totals.ts`](../src/core/usage-totals.ts) ledger. It includes assistant, tool, compaction, branch-summary, per-turn `step-subagent-usage` and incremental `step-auto-clm-usage` records; cumulative maintenance summaries are excluded. Delegated runners may report only aggregate cost: in that case `cost.total` is authoritative and zero component fields mean the split was not provided, not that the total cost is zero. Token increments must contain nonnegative safe integers and must not overflow the accumulated total. A missing/invalid cost does not hide valid tokens. `SessionStats.tokens.input` excludes caches, which are reported separately; the [Account summary](account-usage.md) combines them into complete input. ### Extended Message Types (from the coding agent) diff --git a/packages/coding-agent/docs/usage.md b/packages/coding-agent/docs/usage.md index f0a081d5..266a357f 100644 --- a/packages/coding-agent/docs/usage.md +++ b/packages/coding-agent/docs/usage.md @@ -11,7 +11,7 @@ The interface has four main areas: - **Startup header** - shortcuts, loaded context files, prompt templates, skills, and extensions - **Messages** - user messages, assistant responses, tool calls, tool results, notifications, errors, and extension UI - **Editor** - where you type; border color indicates the current thinking level -- **Footer** - working directory, session name, token/cache usage, cost, context usage, and current model. Totals include assistant responses, usage reported by tools, and summary generation. +- **Footer** - working directory, session name, token/cache usage, cost, context usage, and current model. Totals include assistant responses, delegated usage reported by subagents and workflows, and summary generation. The editor can be replaced temporarily by built-in UI such as `/settings` or by custom extension UI. diff --git a/packages/coding-agent/src/core/usage-totals.ts b/packages/coding-agent/src/core/usage-totals.ts index 35c533e1..64db1e40 100644 --- a/packages/coding-agent/src/core/usage-totals.ts +++ b/packages/coding-agent/src/core/usage-totals.ts @@ -24,6 +24,49 @@ export interface UsageTotals { cost: number; } +/** Common usage shape reported by delegated agents and workflows. */ +export interface DelegatedUsage { + input: number; + output: number; + cacheRead: number; + cacheWrite: number; + cost: number; +} + +/** Sum usage defensively because restored and test records may predate delegated usage. */ +export function sumDelegatedUsage(records: readonly { usage?: DelegatedUsage }[]): DelegatedUsage { + return records.reduce( + (total, record) => { + if (!record.usage) return total; + return { + input: total.input + record.usage.input, + output: total.output + record.usage.output, + cacheRead: total.cacheRead + record.usage.cacheRead, + cacheWrite: total.cacheWrite + record.usage.cacheWrite, + cost: total.cost + record.usage.cost, + }; + }, + { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, cost: 0 }, + ); +} + +export function hasDelegatedUsage(usage: DelegatedUsage): boolean { + return usage.input + usage.output + usage.cacheRead + usage.cacheWrite > 0 || usage.cost > 0; +} + +/** Adapt delegated usage to the provider Usage persisted on tool results and custom entries. */ +export function delegatedUsageToUsage(usage: DelegatedUsage): Usage { + return { + input: usage.input, + output: usage.output, + cacheRead: usage.cacheRead, + cacheWrite: usage.cacheWrite, + totalTokens: usage.input + usage.output + usage.cacheRead + usage.cacheWrite, + // Runners report aggregate cost only; preserve that total without inventing a component split. + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: usage.cost }, + }; +} + export function createUsageTotals(): UsageTotals { return { input: 0, @@ -65,6 +108,10 @@ export function entryUsage(value: unknown): { raw?: unknown; expected: boolean; return { raw: message.usage, expected: message.usage !== undefined, assistant: false }; if (entry.type === "compaction" || entry.type === "branch_summary") return { raw: entry.usage, expected: entry.usage !== undefined, assistant: false }; + if (entry.type === "custom" && entry.customType === "step-subagent-usage") { + const data = entry.data as Record | undefined; + return { raw: data?.usage, expected: data?.usage !== undefined, assistant: false }; + } if (entry.type === "custom" && entry.customType === "step-auto-clm-usage") { const data = entry.data as Record | undefined; return { raw: data?.usage, expected: true, assistant: false }; diff --git a/packages/coding-agent/src/features/step-subagent.ts b/packages/coding-agent/src/features/step-subagent.ts index 6c6d74c9..ec910b49 100644 --- a/packages/coding-agent/src/features/step-subagent.ts +++ b/packages/coding-agent/src/features/step-subagent.ts @@ -17,6 +17,7 @@ import { CONFIG_DIR_NAME } from "../config.ts"; import { CHILD_SESSION_ORIGIN } from "../core/child-session.ts"; import type { ExtensionAPI, ExtensionContext, ExtensionFactory, InlineExtension } from "../core/extensions/types.ts"; import type { NewSessionOptions } from "../core/session-manager.ts"; +import { delegatedUsageToUsage, sumDelegatedUsage } from "../core/usage-totals.ts"; import { resolveStepAgentDir, resolveStepConfigDir } from "../step/environment.ts"; import { isValidSubagentBackgroundTimeout, @@ -548,6 +549,7 @@ export function makeToolResult( return { content: [{ type: "text", text }], details, + ...(details.results.length > 0 ? { usage: delegatedUsageToUsage(sumDelegatedUsage(details.results)) } : {}), ...(isError || details.results.some((record) => record.status === "failed" || record.status === "aborted") ? { isError: true } : {}), @@ -872,8 +874,7 @@ export function createStepSubagentExtension(options: StepSubagentExtensionOption ); } const details: StepSubagentDetails = { ...lane.details, results: [...lane.details.results] }; - return makeToolResult( - details, + return controlResult( [ `Started background lane ${lane.id} (${preflight.run.mode}, ${preflight.run.tasks.length} task${preflight.run.tasks.length === 1 ? "" : "s"}; ${Math.round(lane.turnTimeoutMs / 1000)}s limit per turn).`, // Without this the parent, told to "wait for the lane", has no wait @@ -886,6 +887,7 @@ export function createStepSubagentExtension(options: StepSubagentExtensionOption ] .filter(Boolean) .join(" "), + details, ); }, renderCall: (params, theme) => { diff --git a/packages/coding-agent/src/features/subagent/execute.ts b/packages/coding-agent/src/features/subagent/execute.ts index 3d1a26f3..759167e8 100644 --- a/packages/coding-agent/src/features/subagent/execute.ts +++ b/packages/coding-agent/src/features/subagent/execute.ts @@ -535,7 +535,27 @@ export async function runPreparedSubagent( } catch (error) { if (worktree && task.retainWorktree === false) await worktree.cleanup().catch(() => undefined); const message = error instanceof Error ? error.message : String(error); - const failed = makeFailureResult(agent.name, task.task, message, agent.source); + const observed = records[index]; + const failed = observed + ? resultRecord( + agent.name, + task.task, + { + ...observed, + stderr: message, + exitCode: 1, + stopReason: "error", + errorMessage: message, + }, + { + agentSource: agent.source, + worktreePath: worktree?.path, + worktreeBranch: worktree?.branch, + }, + ) + : makeFailureResult(agent.name, task.task, message, agent.source); + // The observed record carries its previous status through resultRecord; the catch is the terminal failure. + failed.status = "failed"; records[index] = failed; emit(mode); reportFinished(failed, index); diff --git a/packages/coding-agent/src/features/subagent/lane-lifecycle.ts b/packages/coding-agent/src/features/subagent/lane-lifecycle.ts index c813635c..0f817203 100644 --- a/packages/coding-agent/src/features/subagent/lane-lifecycle.ts +++ b/packages/coding-agent/src/features/subagent/lane-lifecycle.ts @@ -9,6 +9,7 @@ import { randomBytes, randomUUID } from "node:crypto"; import type { AgentToolResult, AgentToolUpdateCallback } from "@step-harness/agent-core"; import type { ExtensionAPI, ExtensionContext } from "../../core/extensions/types.ts"; +import { delegatedUsageToUsage, hasDelegatedUsage, sumDelegatedUsage } from "../../core/usage-totals.ts"; import type { StepSubagentDetails, StepSubagentRpcSession, SubagentParams } from "../step-subagent.ts"; import type { PreparedSubagentRun, SubagentPreflight } from "./execute.ts"; import { SUBAGENT_SESSION_ID_PREFIX } from "./helpers.ts"; @@ -179,6 +180,48 @@ export function createLaneLifecycle(host: LaneLifecycleHost): LaneLifecycle { const isCurrent = (lane: BackgroundAgentLane, generation: number): boolean => laneEpochs.get(lane) === epoch && lane.runGeneration === generation; + const recordedLaneTurns = new WeakMap>(), + laneTurnDetails = new WeakMap>(); + const setLaneTurnDetails = (lane: BackgroundAgentLane, generation: number, details: StepSubagentDetails): void => { + let turns = laneTurnDetails.get(lane); + if (!turns) { + turns = new Map(); + laneTurnDetails.set(lane, turns); + } + turns.set(generation, details); + }; + /** Persist one independent usage entry per settled lane turn, including turns without a notification. */ + const recordLaneUsage = ( + lane: BackgroundAgentLane, + generation: number, + result?: AgentToolResult, + ): void => { + let recorded = recordedLaneTurns.get(lane); + if (!recorded) { + recorded = new Set(); + recordedLaneTurns.set(lane, recorded); + } + if (recorded.has(generation)) return; + recorded.add(generation); + if (result?.details) setLaneTurnDetails(lane, generation, result.details); + const records = laneTurnDetails.get(lane)?.get(generation)?.results ?? [], + usage = sumDelegatedUsage(records); + if (!hasDelegatedUsage(usage)) return; + try { + host.pi.appendEntry("step-subagent-usage", { + version: 1, + laneId: lane.id, + turn: generation, + childSessionIds: records + .map((record) => record.sessionId) + .filter((id): id is string => typeof id === "string"), + usage: delegatedUsageToUsage(usage), + }); + } catch { + // Accounting must not reopen or interrupt teardown; the result still retains its own usage. + } + }; + const clearIdleTimer = (lane: BackgroundAgentLane): void => { if (lane.idleTimer) clearTimeout(lane.idleTimer); lane.idleTimer = undefined; @@ -245,9 +288,11 @@ export function createLaneLifecycle(host: LaneLifecycleHost): LaneLifecycle { lane.updatedAt = Date.now(); lane.lastProgressNotifyAt = Date.now(); const signal = lane.controller.signal; + setLaneTurnDetails(lane, generation, { ...lane.details, results: [] }); const onUpdate = (partial: AgentToolResult): void => { if (!isCurrent(lane, generation)) return; if (partial.details) { + setLaneTurnDetails(lane, generation, partial.details); lane.details = { ...partial.details, agentId: lane.id, @@ -310,6 +355,7 @@ export function createLaneLifecycle(host: LaneLifecycleHost): LaneLifecycle { }; lane.status = lane.details.status ?? "completed"; lane.updatedAt = Date.now(); + recordLaneUsage(lane, generation, result); const next = lane.status === "aborted" ? undefined : lane.queuedPrompts.shift(); if (next) { startLane(lane, { params: followUpParams(lane, next) }); @@ -329,6 +375,7 @@ export function createLaneLifecycle(host: LaneLifecycleHost): LaneLifecycle { `Background lane ${lane.id} failed: ${error instanceof Error ? error.message : String(error)}`, lane.details, ); + recordLaneUsage(lane, generation); updateLaneWidget(lane); armIdleTimer(lane); void deliverFinal(lane, generation); @@ -398,7 +445,8 @@ export function createLaneLifecycle(host: LaneLifecycleHost): LaneLifecycle { return `Lane ${lane.id} is ${lane.status}; shut down ${processes}`; } lane.queuedPrompts.length = 0; - const generation = lane.runGeneration + 1; + const interruptedGeneration = lane.runGeneration, + generation = interruptedGeneration + 1; lane.runGeneration = generation; const interrupted = lane.promise; lane.controller.abort(); @@ -412,9 +460,11 @@ export function createLaneLifecycle(host: LaneLifecycleHost): LaneLifecycle { // Notify only after the aborted run has actually wound down, with what it // produced. The generation bump above keeps the run's own .then from // double-firing. + let interruptedResult: AgentToolResult | undefined; void interrupted .then( (result) => { + interruptedResult = result; if (!isCurrent(lane, generation)) return; lane.result = result; if (result.details) { @@ -424,7 +474,10 @@ export function createLaneLifecycle(host: LaneLifecycleHost): LaneLifecycle { () => undefined, ) .then(() => { - if (isCurrent(lane, generation) && lane.status === "aborted") return deliverFinal(lane, generation); + if (isCurrent(lane, generation) && lane.status === "aborted") { + recordLaneUsage(lane, interruptedGeneration, interruptedResult); + return deliverFinal(lane, generation); + } pendingFinal.delete(lane); }); return `Stopped lane ${lane.id}; you will get one notification once it has wound down`; @@ -465,8 +518,14 @@ export function createLaneLifecycle(host: LaneLifecycleHost): LaneLifecycle { lane.queuedPrompts.push(prompt); return `Queued follow-up for lane ${lane.id}`; } + const previousGeneration = lane.runGeneration, + previous = lane.promise; lane.controller.abort(); lane.queuedPrompts.length = 0; + void previous.then( + (result) => recordLaneUsage(lane, previousGeneration, result), + () => recordLaneUsage(lane, previousGeneration), + ); startLane(lane, { params: followUpParams(lane, prompt) }); return `Interrupted lane ${lane.id} and started the follow-up`; } @@ -513,11 +572,18 @@ export function createLaneLifecycle(host: LaneLifecycleHost): LaneLifecycle { for (const lane of owned) { clearIdleTimer(lane); lane.queuedPrompts.length = 0; + const generation = lane.runGeneration, + previous = lane.promise; lane.runGeneration += 1; if (lane.status === "running") { lane.controller.abort(); lane.status = "aborted"; - runs.push(lane.promise.catch(() => undefined)); + runs.push( + previous.then( + (result) => recordLaneUsage(lane, generation, result), + () => recordLaneUsage(lane, generation), + ), + ); } for (const session of getLaneSessions(lane.sessionId, false)) { session.stop(); diff --git a/packages/coding-agent/src/features/workflow/agent-runner.ts b/packages/coding-agent/src/features/workflow/agent-runner.ts index 2fcfd9c3..bf06423a 100644 --- a/packages/coding-agent/src/features/workflow/agent-runner.ts +++ b/packages/coding-agent/src/features/workflow/agent-runner.ts @@ -80,6 +80,7 @@ export function createDefaultWorkflowAgentRunner(options: DefaultWorkflowAgentRu // A killed or crashed child can lose a tool event in its unread output. toolActivityObserved: result.exitCode === 0 && result.stopReason !== "aborted", ...(result.model ? { model: result.model } : {}), + ...(result.sessionId ? { sessionId: result.sessionId } : {}), }; } const text = finalOutput(result.messages) || resultText(result); @@ -89,6 +90,7 @@ export function createDefaultWorkflowAgentRunner(options: DefaultWorkflowAgentRu status: "completed", toolActivityObserved: true, ...(result.model ? { model: result.model } : {}), + ...(result.sessionId ? { sessionId: result.sessionId } : {}), }; }; } diff --git a/packages/coding-agent/src/features/workflow/runtime.ts b/packages/coding-agent/src/features/workflow/runtime.ts index d419d473..66976e54 100644 --- a/packages/coding-agent/src/features/workflow/runtime.ts +++ b/packages/coding-agent/src/features/workflow/runtime.ts @@ -1,6 +1,8 @@ import path from "node:path"; import type { AgentToolResult } from "@step-harness/agent-core"; +import type { Usage } from "@step-harness/providers"; import type { ExtensionContext } from "../../core/extensions/types.ts"; +import { delegatedUsageToUsage, hasDelegatedUsage } from "../../core/usage-totals.ts"; import { type StepTelemetryKnownEventName, type StepTelemetryProperties, @@ -135,6 +137,7 @@ export class WorkflowRuntime { private sequence = 0; private agentCount = 0; private cacheHits = 0; + private readonly childSessions = new Set(); /** * Set once a live call in this run may have changed the environment. Cached * results recorded against the previous environment are then stale unless the @@ -284,6 +287,8 @@ export class WorkflowRuntime { startedAt: this.startedAt, finishedAt, spentTokens: this.budget.spent(), + usage: this.budget.usageSnapshot(), + childSessionIds: [...this.childSessions], cacheHits: this.cacheHits, agentCalls: this.agentCount, phases: [...this.phases], @@ -317,6 +322,18 @@ export class WorkflowRuntime { spent_tokens: this.budget.spent(), }); await Promise.all([this.journal.flush(), this.progress.flush()]).catch(() => undefined); + const usage = this.budget.usageSnapshot(); + if (hasDelegatedUsage(usage)) { + const withUsage = + error && typeof error === "object" && Object.isExtensible(error) + ? error + : new Error(errorMessage(error), { cause: error }); + Object.assign(withUsage, { + usage: delegatedUsageToUsage(usage), + details: { childSessionIds: [...this.childSessions] }, + } satisfies { usage: Usage; details: { childSessionIds: string[] } }); + throw withUsage; + } throw error; } finally { this.lifecycle = "closed"; @@ -596,6 +613,7 @@ export class WorkflowRuntime { }, }); raw = await Promise.race([runnerPromise, cancelled]); + if (raw.sessionId) this.childSessions.add(raw.sessionId); if (raw.toolActivityObserved !== true) toolActivityObserved = false; } catch (error: unknown) { // A thrown attempt cannot vouch that every tool start was reported. @@ -604,6 +622,7 @@ export class WorkflowRuntime { // them when it settles, after the race has already been decided. const late = await settleRunner(); if (late) { + if (late.sessionId) this.childSessions.add(late.sessionId); const lateUsage = normalizeUsage(late.usage); totalUsage = mergeWorkflowUsage(totalUsage, lateUsage); try { @@ -1227,6 +1246,7 @@ export function workflowToolResult(value: WorkflowRunResult): AgentToolResult typeof item === "string") : []; +} + +/** Parent-side delegated usage and the child transcripts it replaces in Overview. */ +function delegatedRecord(entry: Row): { usage: unknown; childSessionIds: string[] } | undefined { + const message = entry.message as Row | undefined; + if ( + entry.type === "message" && + message?.role === "toolResult" && + ["subagent", "workflow"].includes(String(message.toolName)) + ) { + const details = message.details as Row | undefined, + results = Array.isArray(details?.results) ? details.results : []; + return { + usage: message.usage, + childSessionIds: [ + ...stringValues(details?.childSessionIds), + ...results.map((record) => (record as Row).sessionId).filter((id): id is string => typeof id === "string"), + ], + }; + } + if (entry.type === "custom" && entry.customType === "step-subagent-usage") { + const data = entry.data as Row | undefined; + return { usage: data?.usage, childSessionIds: stringValues(data?.childSessionIds) }; + } + return undefined; +} + function validID(id: unknown): id is string { return typeof id === "string" && /^[a-zA-Z0-9][a-zA-Z0-9_-]{0,127}$/u.test(id); } @@ -477,6 +507,17 @@ function buildOverview( resolving.delete(key); return result; }; + // A parent delegated record is authoritative only for the child sessions it names. + // Other child transcripts, including those written before parent accounting existed, + // remain the fallback rather than being suppressed at the whole-parent-session level. + const coveredChildSessions = new Set(); + for (const file of files.values()) { + for (const entry of file.entries) { + const delegated = delegatedRecord(entry); + if (!delegated || !normalizeUsage(delegated.usage)) continue; + for (const sessionId of delegated.childSessionIds) coveredChildSessions.add(sessionId); + } + } const recordedSince = new Map(); for (const file of files.values()) if (file.header && validID(file.header.id)) @@ -528,23 +569,31 @@ function buildOverview( } const picked = entryUsage(entry), usage = normalizeUsage(picked.raw); + const coveredByParent = + origin.role === "subagent" && + picked.assistant && + (coveredChildSessions.has(String(file.header.id)) || coveredChildSessions.has(origin.sessionId)); if (picked.expected && !usage) limited.add(picked.raw === undefined ? "Missing usage" : "Invalid usage"); if (usage && !dedup.has(key)) { if (entry.customType === "step-auto-clm-usage" && data?.sessionId !== origin.sessionId) { limited.add("Usage origin conflict"); continue; } - if (!Number.isSafeInteger(globalTokens + usage.total)) { - limited.add("Token sum exceeds exact numeric range"); - continue; + if (!coveredByParent) { + if (!Number.isSafeInteger(globalTokens + usage.total)) { + limited.add("Token sum exceeds exact numeric range"); + continue; + } + day.tokens += usage.total; + globalTokens += usage.total; + if (picked.assistant && origin.role === "user") { + day.chats.add(origin.sessionId); + chats.add(origin.sessionId); + } } + // Token suppression and reasoning coverage are independent: parent-side accounting + // must not erase the child's dispatch metadata or duplicate it through a fork. dedup.add(key); - day.tokens += usage.total; - globalTokens += usage.total; - if (picked.assistant && origin.role === "user") { - day.chats.add(origin.sessionId); - chats.add(origin.sessionId); - } if (picked.assistant) { const level = provenance.get(id) ?? levels.get(id) ?? "unknown"; if (!provenance.has(id)) limited.add("Reasoning coverage is limited"); diff --git a/packages/coding-agent/src/step/usage/json-projection.ts b/packages/coding-agent/src/step/usage/json-projection.ts index 6326d3d8..2c2beb99 100644 --- a/packages/coding-agent/src/step/usage/json-projection.ts +++ b/packages/coding-agent/src/step/usage/json-projection.ts @@ -26,6 +26,7 @@ const dataKeys = new Set([ "originSessionId", "sourceLeafId", "usageEntryId", + "childSessionIds", "reasoningAtDispatch", "requestedProvider", "requestedModel", @@ -66,8 +67,16 @@ function keepPath(path: string[]): boolean { if (path[0] === "usage") return path.length === 2 && usageKeys.has(path[1]); if (path[0] === "message") { if (path.length === 2) - return ["role", "usage", "provider", "model", "responseModel", "timestamp"].includes(path[1]); - return path[1] === "usage" && path.length === 3 && usageKeys.has(path[2]); + return ["role", "usage", "provider", "model", "responseModel", "timestamp", "toolName", "details"].includes( + path[1], + ); + if (path[1] === "usage") return path.length === 3 && usageKeys.has(path[2]); + if (path[1] === "details") { + if (path.length === 2) return true; + if (path[2] === "childSessionIds") return path.length === 3; + if (path[2] === "results") return path.length === 3 || (path[3] === "sessionId" && path.length === 4); + } + return false; } if (path[0] === "data") { if (!dataKeys.has(path[1])) return false; diff --git a/packages/coding-agent/test/agent-session-stats.test.ts b/packages/coding-agent/test/agent-session-stats.test.ts index 191e7c4b..b49f153f 100644 --- a/packages/coding-agent/test/agent-session-stats.test.ts +++ b/packages/coding-agent/test/agent-session-stats.test.ts @@ -252,6 +252,36 @@ describe("AgentSession.getSessionStats", () => { session.dispose(); } }); + it("includes background subagent usage entries in session totals", async () => { + const { session, sessionManager } = await createSession(); + const usage: Usage = { + input: 10, + output: 20, + cacheRead: 30, + cacheWrite: 40, + totalTokens: 100, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 1 }, + }; + try { + sessionManager.appendCustomMessageEntry("agent-notification", "progress", false, { + event: "background_progress", + }); + sessionManager.appendCustomEntry("step-subagent-usage", { + version: 1, + laneId: "lane_00000001", + turn: 1, + childSessionIds: ["child-1"], + usage, + }); + + const stats = session.getSessionStats(); + expect(stats.tokens).toEqual({ input: 10, output: 20, cacheRead: 30, cacheWrite: 40, total: 100 }); + expect(stats.cost).toBe(1); + } finally { + session.dispose(); + } + }); + it("counts each automatic CLM request once including rejected and late responses", async () => { const { session, sessionManager } = await createSession(); const usage: Usage = { diff --git a/packages/coding-agent/test/step-subagent-events.test.ts b/packages/coding-agent/test/step-subagent-events.test.ts index eadf3690..dbf5b33d 100644 --- a/packages/coding-agent/test/step-subagent-events.test.ts +++ b/packages/coding-agent/test/step-subagent-events.test.ts @@ -2,16 +2,20 @@ import { mkdir, mkdtemp, rm, writeFile } from "node:fs/promises"; import os, { tmpdir } from "node:os"; import { join } from "node:path"; import type { AgentToolResult } from "@step-harness/agent-core"; +import type { Usage } from "@step-harness/providers/compat"; import { expect, test, vi } from "vitest"; import type { ExtensionAPI, ExtensionContext, ToolDefinition } from "../src/core/extensions/types.ts"; import { createStepSubagentExtension, escapeXmlAttr, + makeToolResult, routeSubagentRpcLine, + type StepSubagentDetails, type StepSubagentRunInput, type StepSubagentRunResult, } from "../src/features/step-subagent.ts"; import { builtinStepAgents, resolveStepAgent, type StepAgentConfig } from "../src/features/step-subagent-agents.ts"; +import { createLaneLifecycle } from "../src/features/subagent/lane-lifecycle.ts"; interface SentMessage { customType: string; @@ -81,10 +85,12 @@ function createApi(): { api: ExtensionAPI; tools: Map; sent: SentMessage[]; + entries: Array<{ customType: string; data?: unknown }>; handlers: Map; } { const tools = new Map(); const sent: SentMessage[] = []; + const entries: Array<{ customType: string; data?: unknown }> = []; const handlers = new Map(); const api = { registerTool(tool: ToolDefinition) { @@ -100,13 +106,52 @@ function createApi(): { getActiveTools: () => [], setActiveTools: () => {}, getFlag: () => false, - appendEntry: () => {}, + appendEntry: (customType: string, data?: unknown) => { + entries.push({ customType, data }); + }, sendMessage: (message: SentMessage, options?: { deliverAs?: string; triggerTurn?: boolean }) => { sent.push({ ...message, deliverAs: options?.deliverAs, triggerTurn: options?.triggerTurn }); }, sendUserMessage: () => {}, } as unknown as ExtensionAPI; - return { api, tools, sent, handlers }; + return { api, tools, sent, entries, handlers }; +} + +function usageRecords(entries: Array<{ customType: string; data?: unknown }>): Array<{ + turn: number; + childSessionIds: string[]; + usage: Usage; +}> { + return entries + .filter((entry) => entry.customType === "step-subagent-usage") + .map((entry) => entry.data as { turn: number; childSessionIds: string[]; usage: Usage }); +} + +function preparedTurn(task: string) { + return { + mode: "single", + scope: "user", + discovery: { agents: builtinStepAgents, userAgentsDir: "/agents", projectAgentsDir: null }, + tasks: [{ agent: "general", task }], + } as never; +} + +function completedTurnDetails(task: string): StepSubagentDetails { + return { + mode: "single", + agentScope: "user", + userAgentsDir: "/agents", + projectAgentsDir: null, + results: [ + { + ...textResult(task), + agent: "general", + agentSource: "builtin", + task, + status: "completed", + }, + ], + }; } function laneIdOf(result: AgentToolResult): string { @@ -252,6 +297,34 @@ test("accepts the descriptive general-purpose name as an alias for the built-in expect(result.content[0]).toMatchObject({ type: "text", text: expect.stringContaining("alias output") }); }); +test("a runner failure preserves the last reported usage", async () => { + const { api, tools } = createApi(); + createStepSubagentExtension({ + includeBuiltinAgents: true, + agentDir: "/tmp/step-agent-test", + runner: async (input) => { + input.onUpdate?.({ + ...textResult("partial"), + usage: { input: 7, output: 1, cacheRead: 2, cacheWrite: 0, cost: 0.1, contextTokens: 10, turns: 1 }, + }); + throw new Error("runner failed after usage"); + }, + })(api); + const result = await tools + .get("subagent")! + .execute( + "call", + { agent: "general", task: "inspect" } as never, + undefined, + undefined, + createContext("/workspace"), + ); + + expect(result.isError).toBe(true); + expect(result.usage).toMatchObject({ input: 7, output: 1, cacheRead: 2, totalTokens: 10 }); + expect(result.usage?.cost.total).toBeCloseTo(0.1); +}); + test("a ten-task batch starts eight and queues two on a ten-core host", async () => { const cpu = os.cpus()[0]!; const spy = vi.spyOn(os, "cpus").mockReturnValue(Array(10).fill(cpu)); @@ -710,6 +783,173 @@ test("a single-lane reply runs only the follow-up task for the lane's agent", as ]); }); +test("queued background turns each persist usage once", async () => { + const { api, tools, entries } = createApi(); + const gate = deferred(); + createStepSubagentExtension({ + agentDir: "/tmp/step-agent-test", + runner: async (input) => { + if (input.task === "first") { + input.onUpdate?.({ + ...textResult(""), + exitCode: -1, + usage: { input: 7, output: 1, cacheRead: 0, cacheWrite: 0, cost: 0.1, contextTokens: 8, turns: 1 }, + }); + await gate.promise; + return { + ...textResult("first"), + usage: { input: 7, output: 1, cacheRead: 0, cacheWrite: 0, cost: 0.1, contextTokens: 8, turns: 1 }, + }; + } + return { + ...textResult("second"), + usage: { input: 11, output: 2, cacheRead: 0, cacheWrite: 0, cost: 0.2, contextTokens: 13, turns: 1 }, + }; + }, + })(api); + const lane = laneIdOf(await startLane(tools, { agent: "general", task: "first" })); + await agentSend(tools, { lane_id: lane, action: "reply", prompt: "second" }); + gate.resolve(); + await waitFor(() => usageRecords(entries).length === 2); + const records = usageRecords(entries); + expect(records.map((record) => record.turn)).toEqual([1, 2]); + expect(records.map((record) => record.usage.input)).toEqual([7, 11]); + expect(records.map((record) => record.usage.cost.total)).toEqual([0.1, 0.2]); +}); + +test("interrupting a background turn records both the interrupted and replacement turns", async () => { + const { api, tools, entries } = createApi(); + createStepSubagentExtension({ + agentDir: "/tmp/step-agent-test", + runner: async (input) => { + if (input.task === "first") { + input.onUpdate?.({ + ...textResult(""), + exitCode: -1, + usage: { input: 7, output: 1, cacheRead: 0, cacheWrite: 0, cost: 0.1, contextTokens: 8, turns: 1 }, + }); + await new Promise((resolve) => + input.signal?.addEventListener("abort", () => resolve(), { once: true }), + ); + return { + ...textResult("interrupted"), + stopReason: "aborted", + usage: { input: 7, output: 1, cacheRead: 0, cacheWrite: 0, cost: 0.1, contextTokens: 8, turns: 1 }, + }; + } + return { + ...textResult("replacement"), + usage: { input: 11, output: 2, cacheRead: 0, cacheWrite: 0, cost: 0.2, contextTokens: 13, turns: 1 }, + }; + }, + })(api); + const lane = laneIdOf(await startLane(tools, { agent: "general", task: "first" })); + await agentSend(tools, { lane_id: lane, action: "reply", prompt: "second", interrupt: true }); + await waitFor(() => usageRecords(entries).length === 2); + expect(usageRecords(entries).map((record) => record.turn)).toEqual([1, 2]); + expect(usageRecords(entries).map((record) => record.usage.input)).toEqual([7, 11]); +}); + +test("session shutdown persists the latest observed usage without sending a notification", async () => { + const { api, tools, entries, sent, handlers } = createApi(); + let started = false; + createStepSubagentExtension({ + agentDir: "/tmp/step-agent-test", + runner: async (input) => { + input.onUpdate?.({ + ...textResult(""), + exitCode: -1, + usage: { input: 7, output: 1, cacheRead: 0, cacheWrite: 0, cost: 0.1, contextTokens: 8, turns: 1 }, + }); + started = true; + await new Promise((resolve) => input.signal?.addEventListener("abort", () => resolve(), { once: true })); + return { + ...textResult("aborted"), + stopReason: "aborted", + usage: { input: 7, output: 1, cacheRead: 0, cacheWrite: 0, cost: 0.1, contextTokens: 8, turns: 1 }, + }; + }, + })(api); + await startLane(tools, { agent: "general", task: "shutdown" }); + await waitFor(() => started); + await handlers.get("session_shutdown")![0]!( + { type: "session_shutdown", reason: "new" }, + createContext("/workspace"), + ); + const records = usageRecords(entries); + expect(records).toHaveLength(1); + expect(records[0]?.usage).toMatchObject({ input: 7, output: 1, totalTokens: 8 }); + expect(records[0]?.childSessionIds).toHaveLength(1); + expect(sent).toHaveLength(0); +}); + +test("a later preflight failure does not record the previous turn's usage", async () => { + const { api, entries, sent } = createApi(); + let second = false; + const lifecycle = createLaneLifecycle({ + pi: api, + lanes: new Map(), + prepare: async () => { + if (second) throw new Error("second preflight failed"); + return { ok: true, run: preparedTurn("first") }; + }, + run: async () => makeToolResult(completedTurnDetails("first"), "done"), + outputDirectory: () => "/tmp", + updateLaneWidget: () => {}, + clearLaneWidget: () => {}, + }); + const lane = lifecycle.launchLane( + { agent: "general", task: "first" }, + createContext("/workspace"), + preparedTurn("first"), + 1000, + ); + if (typeof lane === "string") throw new Error(lane); + await waitFor(() => usageRecords(entries).length === 1); + second = true; + lifecycle.replyToLane(lane, "second", false); + await waitFor(() => sent.some((message) => message.details?.event === "background_failed")); + expect(usageRecords(entries)).toHaveLength(1); +}); + +test("a rejected later turn does not reuse the previous turn's usage when stopped", async () => { + const { api, entries } = createApi(); + let second = false, + secondStarted = false, + rejectSecond!: (error: Error) => void; + const lifecycle = createLaneLifecycle({ + pi: api, + lanes: new Map(), + prepare: async () => ({ ok: true, run: preparedTurn(second ? "second" : "first") }), + run: async () => { + if (!second) return makeToolResult(completedTurnDetails("first"), "done"); + secondStarted = true; + return new Promise((_resolve, reject) => { + rejectSecond = reject; + }); + }, + outputDirectory: () => "/tmp", + updateLaneWidget: () => {}, + clearLaneWidget: () => {}, + }); + const lane = lifecycle.launchLane( + { agent: "general", task: "first" }, + createContext("/workspace"), + preparedTurn("first"), + 1000, + ); + if (typeof lane === "string") throw new Error(lane); + await waitFor(() => usageRecords(entries).length === 1); + second = true; + const reply = lifecycle.replyToLane(lane, "second", false); + expect(reply).toContain("Started follow-up"); + await waitFor(() => secondStarted); + lifecycle.stopLane(lane); + rejectSecond(new Error("second rejected")); + await new Promise((resolve) => setTimeout(resolve, 50)); + expect(usageRecords(entries)).toHaveLength(1); +}); + test("preflight failures return an ordinary error and create no lane", async () => { const { api, tools, sent } = createApi(); let runs = 0; @@ -964,7 +1204,7 @@ test("session shutdown stops this session's lanes silently and leaves another se }); test("stop sends a single interrupted notification carrying the partial run", async () => { - const { api, tools, sent } = createApi(); + const { api, tools, entries, sent } = createApi(); createStepSubagentExtension({ agentDir: "/tmp/step-agent-test", runner: async (input) => { @@ -980,6 +1220,8 @@ test("stop sends a single interrupted notification carrying the partial run", as expect(sent).toHaveLength(1); expect(sent[0]!.details?.event).toBe("background_interrupted"); expect(sent[0]!.content).toContain(`Lane ${lane} was stopped: 1 stopped.`); + expect(usageRecords(entries)).toHaveLength(1); + expect(usageRecords(entries)[0]?.usage).toMatchObject({ input: 1, output: 2, totalTokens: 3 }); }); test("idle keep-alive children are stopped after the TTL and a reply resumes quietly", async () => { diff --git a/packages/coding-agent/test/step-usage.test.ts b/packages/coding-agent/test/step-usage.test.ts index 1929dd91..6b6a2394 100644 --- a/packages/coding-agent/test/step-usage.test.ts +++ b/packages/coding-agent/test/step-usage.test.ts @@ -45,7 +45,7 @@ function row( }, }; } -function caps(id: string, role = "user", parentCwd?: string): Record { +function caps(id: string, role = "user", parentCwd?: string, parentSessionId?: string): Record { return { type: "custom", id: `cap-${id}`, @@ -57,12 +57,24 @@ function caps(id: string, role = "user", parentCwd?: string): Record { + return { + type: "custom", + id: `provenance-${id}`, + parentId: null, + timestamp, + customType: "step.usage_provenance", + data: { version: 1, sessionId: id, originSessionId: id, usageEntryId, reasoningAtDispatch: level }, + }; +} + function mainRows(): Record[] { return [ caps("main"), @@ -187,6 +199,79 @@ describe("annual local activity and cache", () => { expect((await scanActivity(options(path))).tokens).toBe(235); }); + it("counts delegated usage once across parent and child files, including cache hits", async () => { + const path = root(), + main = join(path, "sessions", "A", "main.jsonl"), + child = join(path, "sessions", "C", "child.jsonl"), + workflowChild = join(path, "sessions", "E", "workflow-child.jsonl"), + backgroundChild = join(path, "sessions", "D", "background-child.jsonl"); + write(main, [ + header("main"), + caps("main"), + row("parent", 10, 1), + { + type: "message", + id: "delegated", + parentId: null, + timestamp, + message: { + role: "toolResult", + toolName: "subagent", + usage: { input: 20, output: 2, cacheRead: 0, cacheWrite: 0 }, + details: { results: [{ sessionId: "child" }] }, + }, + }, + { + type: "message", + id: "workflow-usage", + parentId: null, + timestamp, + message: { + role: "toolResult", + toolName: "workflow", + usage: { input: 40, output: 4, cacheRead: 0, cacheWrite: 0 }, + details: { childSessionIds: ["workflow-child"] }, + }, + }, + { + type: "custom", + id: "background-usage", + parentId: null, + timestamp, + customType: "step-subagent-usage", + data: { + usage: { input: 30, output: 3, cacheRead: 0, cacheWrite: 0 }, + childSessionIds: ["background-child"], + }, + }, + ]); + write(child, [ + header("child", "/work/child"), + caps("child", "subagent", "/work/A", "main"), + provenance("child", "child-call"), + row("child-call", 20, 2), + ]); + write(workflowChild, [ + header("workflow-child", "/work/workflow"), + caps("workflow-child", "subagent", "/work/A", "main"), + provenance("workflow-child", "workflow-call"), + row("workflow-call", 40, 4), + ]); + write(backgroundChild, [ + header("background-child", "/work/background"), + caps("background-child", "subagent", "/work/A", "main"), + provenance("background-child", "background-call"), + row("background-call", 30, 3), + ]); + const opts = options(path), + first = await scanActivity(opts), + second = await scanActivity(opts); + expect(first.tokens).toBe(110); + expect(second.tokens).toBe(110); + expect(first.reasoning.filter(([level]) => level === "high")).toEqual([["high", 3]]); + expect(second.reasoning.filter(([level]) => level === "high")).toEqual([["high", 3]]); + }); + it("deduplicates forks and excludes child chats, including project attribution", async () => { const path = root(), main = join(path, "sessions", "A", "main.jsonl"), @@ -194,7 +279,11 @@ describe("annual local activity and cache", () => { child = join(path, "sessions", "C", "child.jsonl"); write(main, [header("main"), ...mainRows()]); write(fork, [header("fork", "/work/B", main), ...mainRows(), caps("fork"), row("new", 20, 2)]); - write(child, [header("child", "/work/isolated"), caps("child", "subagent", "/work/A"), row("child-call", 30, 3)]); + write(child, [ + header("child", "/work/isolated"), + caps("child", "subagent", "/work/A", "main"), + row("child-call", 30, 3), + ]); const checksum = createHash("sha256").update(readFileSync(main)).digest("hex"); const all = await scanActivity(options(path)); expect(all).toMatchObject({ tokens: 290, chats: 2 }); diff --git a/packages/coding-agent/test/subagent-chain-result.test.ts b/packages/coding-agent/test/subagent-chain-result.test.ts index 8e9a2a4f..a0812132 100644 --- a/packages/coding-agent/test/subagent-chain-result.test.ts +++ b/packages/coding-agent/test/subagent-chain-result.test.ts @@ -305,3 +305,33 @@ test("admitted parallel subagents ask the runner to hold their slot until the ch }); expect(seen).toEqual([true, undefined]); }); + +test("worktree cleanup failure preserves the completed child's usage", async () => { + const result = await executeSubagent( + { agent: "explore", task: "finish then clean", isolateWorkspace: true, retainWorktree: false }, + undefined, + undefined, + ctx, + { + agentDir: "/tmp/step-agent-test", + configDirName: ".step", + includeBuiltinAgents: true, + maxParallelTasks: 8, + maxConcurrency: 4, + worktreeManager: { + allocate: async () => ({ + path: "/tmp/test-worktree", + branch: "test-worktree", + cleanup: async () => { + throw new Error("cleanup failed"); + }, + }), + }, + runner: async () => runResult({ text: "completed" }), + }, + ); + + expect(result.isError).toBe(true); + expect(result.usage).toMatchObject({ input: 1, output: 2, totalTokens: 3 }); + expect(result.details?.results[0]).toMatchObject({ status: "failed", errorMessage: "cleanup failed" }); +}); diff --git a/packages/coding-agent/test/suite/agent-session-delegated-usage.test.ts b/packages/coding-agent/test/suite/agent-session-delegated-usage.test.ts new file mode 100644 index 00000000..fee99c80 --- /dev/null +++ b/packages/coding-agent/test/suite/agent-session-delegated-usage.test.ts @@ -0,0 +1,254 @@ +import { fauxAssistantMessage, fauxToolCall, type Usage } from "@step-harness/providers/compat"; +import { expect, test } from "vitest"; +import type { SessionEntry } from "../../src/core/session-manager.ts"; +import { createStepSubagentExtension, type StepSubagentRunResult } from "../../src/features/step-subagent.ts"; +import { createStepWorkflowExtension } from "../../src/features/workflow/step-workflow.ts"; +import { createHarness, type Harness } from "./harness.ts"; + +const childUsage = { + input: 120, + output: 30, + cacheRead: 10, + cacheWrite: 5, + cost: 0.25, + contextTokens: 165, + turns: 2, +}; + +const expectedToolUsage: Usage = { + input: 120, + output: 30, + cacheRead: 10, + cacheWrite: 5, + totalTokens: 165, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0.25 }, +}; + +function subagentResult(): StepSubagentRunResult { + return { + messages: [ + { + role: "assistant", + content: [{ type: "text", text: "subagent finished" }], + api: "faux", + provider: "faux", + model: "faux-1", + usage: expectedToolUsage, + stopReason: "stop", + timestamp: Date.now(), + }, + ], + stderr: "", + exitCode: 0, + usage: childUsage, + model: "faux/faux-1", + stopReason: "stop", + }; +} + +function parentTotals(entries: readonly SessionEntry[]) { + return entries.reduce( + (total, entry) => { + if (entry.type !== "message" || entry.message.role !== "assistant") return total; + const usage = entry.message.usage; + return { + input: total.input + usage.input, + output: total.output + usage.output, + cacheRead: total.cacheRead + usage.cacheRead, + cacheWrite: total.cacheWrite + usage.cacheWrite, + total: total.total + usage.input + usage.output + usage.cacheRead + usage.cacheWrite, + cost: total.cost + usage.cost.total, + }; + }, + { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0, cost: 0 }, + ); +} + +async function runDelegatedTool(harness: Harness, toolName: string, args: Record): Promise { + harness.setResponses([ + fauxAssistantMessage([fauxToolCall(toolName, args)], { stopReason: "toolUse" }), + fauxAssistantMessage("parent finished"), + ]); + await harness.session.prompt(`Run ${toolName}.`); +} + +async function waitFor(predicate: () => boolean, timeoutMs = 3_000): Promise { + const started = Date.now(); + while (!predicate()) { + if (Date.now() - started > timeoutMs) throw new Error("Timed out waiting for delegated usage"); + await new Promise((resolve) => setTimeout(resolve, 10)); + } +} + +function expectSessionTotals(harness: Harness): void { + const parent = parentTotals(harness.sessionManager.getEntries()); + const stats = harness.session.getSessionStats(); + expect(stats.tokens).toEqual({ + input: parent.input + childUsage.input, + output: parent.output + childUsage.output, + cacheRead: parent.cacheRead + childUsage.cacheRead, + cacheWrite: parent.cacheWrite + childUsage.cacheWrite, + total: parent.total + 165, + }); + expect(stats.cost).toBeCloseTo(parent.cost + childUsage.cost); +} + +function expectDelegatedUsage(harness: Harness, toolName: string, expectError = false): void { + const entries = harness.sessionManager.getEntries(); + const toolMessage = entries.find( + (entry): entry is Extract => + entry.type === "message" && entry.message.role === "toolResult" && entry.message.toolName === toolName, + )?.message; + expect(toolMessage?.role).toBe("toolResult"); + if (toolMessage?.role !== "toolResult") throw new Error(`${toolName} result was not persisted`); + expect(Boolean(toolMessage.isError)).toBe(expectError); + expect(toolMessage.usage).toEqual(expectedToolUsage); + if (toolName === "workflow") expect(toolMessage.details).toMatchObject({ childSessionIds: ["workflow-child-1"] }); + + expectSessionTotals(harness); +} + +test("foreground subagent usage joins the parent session totals", async () => { + const harness = await createHarness({ + extensionFactories: [ + { + factory: createStepSubagentExtension({ + includeBuiltinAgents: true, + agentDir: ".step-code/agents", + runner: async () => subagentResult(), + }), + }, + ], + initialActiveToolNames: ["subagent"], + }); + try { + await runDelegatedTool(harness, "subagent", { + agent: "general", + task: "run one child", + run_in_background: false, + }); + expectDelegatedUsage(harness, "subagent"); + } finally { + harness.cleanup(); + } +}); + +test("background subagent usage is counted once with synchronous progress", async () => { + const harness = await createHarness({ + extensionFactories: [ + { + factory: createStepSubagentExtension({ + includeBuiltinAgents: true, + agentDir: ".step-code/agents", + runner: async (input) => { + input.onUpdate?.({ ...subagentResult(), exitCode: -1 }); + return subagentResult(); + }, + }), + }, + ], + initialActiveToolNames: ["subagent"], + }); + try { + await harness.session.bindExtensions({ mode: "tui" }); + harness.setResponses([ + fauxAssistantMessage( + [fauxToolCall("subagent", { agent: "general", task: "background child", run_in_background: true })], + { stopReason: "toolUse" }, + ), + fauxAssistantMessage("parent finished"), + fauxAssistantMessage("background handled"), + ]); + await harness.session.prompt("Run background subagent."); + await waitFor(() => + harness.sessionManager + .getEntries() + .some((entry) => entry.type === "custom" && entry.customType === "step-subagent-usage"), + ); + const entries = harness.sessionManager.getEntries(); + const launchResult = entries.find( + (entry) => + entry.type === "message" && entry.message.role === "toolResult" && entry.message.toolName === "subagent", + ); + expect( + launchResult?.type === "message" && launchResult.message.role === "toolResult" + ? launchResult.message.usage + : undefined, + ).toBeUndefined(); + expect( + entries.filter((entry) => entry.type === "custom" && entry.customType === "step-subagent-usage"), + ).toHaveLength(1); + expectSessionTotals(harness); + } finally { + harness.cleanup(); + } +}); + +test("workflow usage joins the parent session totals", async () => { + const harness = await createHarness({ + extensionFactories: [ + { + factory: createStepWorkflowExtension({ + runner: async () => ({ + text: "workflow child finished", + usage: childUsage, + sessionId: "workflow-child-1", + status: "completed", + }), + }), + }, + ], + initialActiveToolNames: ["workflow"], + }); + try { + await runDelegatedTool(harness, "workflow", { + script: 'const value = await agent("run workflow child"); return value;', + }); + expectDelegatedUsage(harness, "workflow"); + } finally { + harness.cleanup(); + } +}); + +test.each([ + { + name: "failed agent", + script: 'return await agent("fail workflow child");', + args: {}, + status: "failed" as const, + }, + { + name: "budget exhaustion", + script: 'return await agent("exceed workflow budget");', + args: { budget: 1 }, + status: "completed" as const, + }, + { + name: "top-level script failure", + script: 'await agent("spend then throw"); throw new Error("script exploded");', + args: {}, + status: "completed" as const, + }, +])("workflow $name still joins the parent session totals", async (scenario) => { + const harness = await createHarness({ + extensionFactories: [ + { + factory: createStepWorkflowExtension({ + runner: async () => ({ + text: "workflow child failed", + usage: childUsage, + sessionId: "workflow-child-1", + status: scenario.status, + }), + }), + }, + ], + initialActiveToolNames: ["workflow"], + }); + try { + await runDelegatedTool(harness, "workflow", { script: scenario.script, ...scenario.args }); + expectDelegatedUsage(harness, "workflow", true); + } finally { + harness.cleanup(); + } +});