From 9a100c7cc4707f013438b13969c6671bc0a33c4c Mon Sep 17 00:00:00 2001 From: Armin Ronacher Date: Tue, 29 Sep 2026 16:21:39 +0200 Subject: [PATCH] feat(coding-agent): report codemode classifier and nested tool usage in session cost Codemode sums the usage of models.classify() calls into its tool result and shows each call's cost. Usage of tools called through ctx.executeTool() is added to the calling tool's result usage instead of being dropped. --- packages/coding-agent/CHANGELOG.md | 2 + packages/coding-agent/docs/extensions.md | 2 +- packages/coding-agent/docs/models.md | 2 + .../coding-agent/src/core/agent-session.ts | 12 ++-- .../src/core/compaction/compaction.ts | 24 +------- .../src/core/nested-tool-calls.ts | 31 ++++++++-- .../coding-agent/src/core/usage-totals.ts | 24 ++++++++ .../src/extensions/codemode/execute.ts | 21 ++++++- .../src/extensions/codemode/renderer.ts | 15 ++++- .../src/extensions/codemode/tool.ts | 4 ++ .../test/codemode-renderer.test.ts | 22 +++++++ .../test/nested-tool-calls.test.ts | 52 +++++++++++++++- .../test/suite/agent-session-codemode.test.ts | 60 ++++++++++++++++++- 13 files changed, 230 insertions(+), 41 deletions(-) diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index aec0bd26b..3558f3536 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -14,6 +14,7 @@ - Added inherited Claude Sonnet 5.5 support for Anthropic with adaptive thinking and a 1M context window. - Added a Built-in section in `pi config` to disable the built-in `mcp`, `llama.cpp`, `codemode`, and `tool-search` extensions globally or per project, stored as `-builtin:` in the `extensions` setting. SDK inline extensions opt in with `builtin: true`. - Added `+name` and `-name` entries to the `defaultTools` setting to add or remove tools without repeating the defaults, for example `"defaultTools": ["+codemode"]`. Project entries of this form apply on top of the user setting. Documented how to enable `codemode` without MCP and how to use classifier models such as Jev from codemode scripts. +- Added the token usage and cost of codemode `models.classify()` calls to the codemode tool result, so they count toward the session cost; the codemode result shows each call's cost. ### Changed @@ -36,6 +37,7 @@ - Fixed the OpenCode Go default model pointing at the removed Kimi K2.6 model; it now defaults to Kimi K3. - Fixed the Together default model pointing at the removed Kimi K2.6 model; it now defaults to Kimi K3. - Reduced CPU use while streaming in long sessions and when previewing themes: the footer caches session usage totals, collapsed bash results cache their preview, and `sanitizeBinaryOutput()` no longer splits output into per-character arrays. +- Fixed the usage of tools called through `ctx.executeTool()`, for example from codemode scripts, being dropped from the session cost; it is now added to the calling tool's result usage. ## [0.87.1] - 2026-09-22 diff --git a/packages/coding-agent/docs/extensions.md b/packages/coding-agent/docs/extensions.md index d4a53d392..afc64470c 100644 --- a/packages/coding-agent/docs/extensions.md +++ b/packages/coding-agent/docs/extensions.md @@ -145,7 +145,7 @@ Truncate large model-facing results and tell the model where to read the complet Declare `outputSchema` and return a matching `structuredContent` when the result is data. The model still receives `content`; programmatic callers such as codemode scripts receive `structuredContent` instead of the text. Tools without `outputSchema` are passed to scripts as their text content. To report a failure that still carries data, return the result with `isError: true` instead of throwing: the model sees an error, and scripts still receive `structuredContent`. -A tool can run other tools with `ctx.executeTool(name, args, { signal, onUpdate })`. Nested calls go through argument validation and the `tool_call` and `tool_result` handlers like model-issued calls, and emit `tool_execution_start`, `tool_execution_update`, and `tool_execution_end`; all of these events carry `parentToolCallId`, and their `toolCallId` is assigned by pi as `/`. These ids do not appear as tool calls or tool results in the transcript. Nested calls do not add transcript entries: their results only reach the calling tool, which reports them itself, for example through `onUpdate` and `details`. The session keeps a bounded record of them (name, arguments, status, duration, error; never results) as `nestedCalls` on the calling tool's result message. It is used for compaction file lists and shown in HTML exports. Arguments over 8 KiB per call or 32 KiB per tool result are omitted, at most 256 calls are kept, and `complete: false` marks a record that lost anything. `ctx.tools` lists the tools `ctx.executeTool()` can call. `tool_result` handlers that redact `content` should also replace `structuredContent`; replacing only `content` drops it. +A tool can run other tools with `ctx.executeTool(name, args, { signal, onUpdate })`. Nested calls go through argument validation and the `tool_call` and `tool_result` handlers like model-issued calls, and emit `tool_execution_start`, `tool_execution_update`, and `tool_execution_end`; all of these events carry `parentToolCallId`, and their `toolCallId` is assigned by pi as `/`. These ids do not appear as tool calls or tool results in the transcript. Nested calls do not add transcript entries: their results only reach the calling tool, which reports them itself, for example through `onUpdate` and `details`. The session keeps a bounded record of them (name, arguments, status, duration, error; never results) as `nestedCalls` on the calling tool's result message. It is used for compaction file lists and shown in HTML exports. Arguments over 8 KiB per call or 32 KiB per tool result are omitted, at most 256 calls are kept, and `complete: false` marks a record that lost anything. The `usage` of nested results, at every depth, is added to the calling tool's result `usage`, so a tool reports only its own usage, not that of the tools it called. `ctx.tools` lists the tools `ctx.executeTool()` can call. `tool_result` handlers that redact `content` should also replace `structuredContent`; replacing only `content` drops it. See [`hello.ts`](../examples/extensions/hello.ts), [`todo.ts`](../examples/extensions/todo.ts), [`dynamic-tools.ts`](../examples/extensions/dynamic-tools.ts), and [`truncated-tool.ts`](../examples/extensions/truncated-tool.ts). diff --git a/packages/coding-agent/docs/models.md b/packages/coding-agent/docs/models.md index bae53947d..da2cfe237 100644 --- a/packages/coding-agent/docs/models.md +++ b/packages/coding-agent/docs/models.md @@ -131,6 +131,8 @@ const result = await models.classify(jev, { return result.answers; ``` +When the service reports token counts, as all System One services do, `result.usage` carries them with their cost. Pi adds the usage of a script's classifier calls to the `codemode` tool result, so it counts toward the session cost in the footer and `/session`. The cost uses the model's catalog price; models without one, such as TypeSafe's direct `jev-latest`, report tokens at no cost. + Extensions call classifiers through `ctx.modelRegistry.classify()`, without codemode. [Virtual models](virtual-models.md#route-requests) can use them to route requests; see the `jev-router.ts` example. ## Add a custom provider diff --git a/packages/coding-agent/src/core/agent-session.ts b/packages/coding-agent/src/core/agent-session.ts index 1ab92f177..4aaf3c248 100644 --- a/packages/coding-agent/src/core/agent-session.ts +++ b/packages/coding-agent/src/core/agent-session.ts @@ -142,7 +142,7 @@ import { import { type BashOperations, createLocalBashOperations } from "./tools/bash.ts"; import { createAllToolDefinitions } from "./tools/index.ts"; import { createToolDefinitionFromAgentTool } from "./tools/tool-definition-wrapper.ts"; -import { addUsageToTotals, createUsageTotals } from "./usage-totals.ts"; +import { addUsageToTotals, combineUsage, createUsageTotals } from "./usage-totals.ts"; import { findLatestResponse, getBranchSelection, @@ -1059,11 +1059,15 @@ export class AgentSession { /** Internal handler for agent events - shared by subscribe and reconnect */ private _handleAgentEvent = async (event: AgentEvent): Promise => { - // Record the calls a tool made through ctx.executeTool() on its result message. + // Record the calls a tool made through ctx.executeTool() and their usage on its result message. if (this._nestedToolCalls) { if (event.type === "message_start" && event.message.role === "toolResult") { - const nestedCalls = this._nestedToolCalls.takeRecord(event.message.toolCallId); - if (nestedCalls) event.message.nestedCalls = nestedCalls; + const message = event.message; + const summary = this._nestedToolCalls.takeRecord(message.toolCallId); + if (summary?.calls) message.nestedCalls = summary.calls; + if (summary?.usage) { + message.usage = message.usage ? combineUsage(message.usage, summary.usage) : summary.usage; + } } else if (event.type === "agent_end") { this._nestedToolCalls.clear(); } diff --git a/packages/coding-agent/src/core/compaction/compaction.ts b/packages/coding-agent/src/core/compaction/compaction.ts index 1757c535c..6bcfbb2b5 100644 --- a/packages/coding-agent/src/core/compaction/compaction.ts +++ b/packages/coding-agent/src/core/compaction/compaction.ts @@ -33,6 +33,7 @@ import { type SessionProjection, sessionEntryToContextMessages, } from "../session-manager.ts"; +import { combineUsage } from "../usage-totals.ts"; import { computeFileLists, createFileOps, @@ -112,29 +113,6 @@ export interface CompactionResult { details?: T; } -function combineUsage(first: Usage, second: Usage): Usage { - return { - input: first.input + second.input, - output: first.output + second.output, - cacheRead: first.cacheRead + second.cacheRead, - cacheWrite: first.cacheWrite + second.cacheWrite, - ...(first.cacheWrite1h !== undefined || second.cacheWrite1h !== undefined - ? { cacheWrite1h: (first.cacheWrite1h ?? 0) + (second.cacheWrite1h ?? 0) } - : {}), - ...(first.reasoning !== undefined || second.reasoning !== undefined - ? { reasoning: (first.reasoning ?? 0) + (second.reasoning ?? 0) } - : {}), - totalTokens: first.totalTokens + second.totalTokens, - cost: { - input: first.cost.input + second.cost.input, - output: first.cost.output + second.cost.output, - cacheRead: first.cost.cacheRead + second.cost.cacheRead, - cacheWrite: first.cost.cacheWrite + second.cost.cacheWrite, - total: first.cost.total + second.cost.total, - }, - }; -} - // ============================================================================ // Types // ============================================================================ diff --git a/packages/coding-agent/src/core/nested-tool-calls.ts b/packages/coding-agent/src/core/nested-tool-calls.ts index 3ebf62036..b014dc0da 100644 --- a/packages/coding-agent/src/core/nested-tool-calls.ts +++ b/packages/coding-agent/src/core/nested-tool-calls.ts @@ -2,7 +2,8 @@ * Tool calls that a tool makes while it runs (`ctx.executeTool()`), for example from codemode * scripts. The agent loop does not know about them: the session runs each one through the agent's * tool pipeline (`runToolCall`) with its own hooks, emits `tool_execution_*` events with - * `parentToolCallId`, and records the calls on the model-issued call's tool result message. + * `parentToolCallId`, and records the calls and their usage on the model-issued call's tool result + * message. * * Nothing here runs until a tool calls `ctx.executeTool()`. */ @@ -14,7 +15,8 @@ import type { AgentToolResult, AgentToolUpdateCallback, } from "@earendil-works/pi-agent-core"; -import type { JsonObject, NestedToolCallRecord, NestedToolCalls, TextContent } from "@earendil-works/pi-ai"; +import type { JsonObject, NestedToolCallRecord, NestedToolCalls, TextContent, Usage } from "@earendil-works/pi-ai"; +import { combineUsage } from "./usage-totals.ts"; /** * Limits of the nested-call record on a tool result: arguments @@ -30,6 +32,14 @@ export const NESTED_CALL_LIMITS = { const encoder = new TextEncoder(); +/** What the nested calls of one model-issued tool call leave on its tool result message. */ +export interface NestedCallSummary { + /** Becomes `nestedCalls`. Undefined when no nested call was made. */ + calls: NestedToolCalls | undefined; + /** Summed `usage` of the nested results, added to the message's `usage`. */ + usage: Usage | undefined; +} + /** * Collects the nested calls of one model-issued tool call, including calls made by nested tools. * The snapshot becomes `nestedCalls` on the tool result message. @@ -39,6 +49,8 @@ export class NestedCallRecorder { private readonly startedAt = new Map(); private complete = true; private argumentBytes = 0; + /** Summed usage of every nested result, including calls dropped from the record. */ + private usage: Usage | undefined; /** Record a call as it starts. Returns undefined when the call is dropped. */ start(toolCall: AgentToolCall): NestedToolCallRecord | undefined { @@ -72,6 +84,14 @@ export class NestedCallRecorder { if (isError && errorText) record.error = errorText.slice(0, NESTED_CALL_LIMITS.maxErrorChars); } + addUsage(usage: Usage): void { + this.usage = this.usage ? combineUsage(this.usage, usage) : usage; + } + + get totalUsage(): Usage | undefined { + return this.usage; + } + /** Copy of the record so far, or undefined when no nested call was made. */ snapshot(): NestedToolCalls | undefined { if (this.calls.length === 0 && this.complete) return undefined; @@ -214,6 +234,8 @@ export class NestedToolCallRunner { } scope.recorder.finish(record, outcome.isError, textOf(outcome.result)); + // Nested results are not persisted, so their usage is only counted through the recorder. + if (outcome.result.usage) scope.recorder.addUsage(outcome.result.usage); await this.host.emit({ type: "tool_execution_end", toolCallId: toolCall.id, @@ -226,10 +248,11 @@ export class NestedToolCallRunner { } /** Remove and return the record of the nested calls a model-issued call made. */ - takeRecord(toolCallId: string): NestedToolCalls | undefined { + takeRecord(toolCallId: string): NestedCallSummary | undefined { const scope = this.scopes.get(toolCallId); this.scopes.delete(toolCallId); - return scope?.recorder.snapshot(); + if (!scope) return undefined; + return { calls: scope.recorder.snapshot(), usage: scope.recorder.totalUsage }; } clear(): void { diff --git a/packages/coding-agent/src/core/usage-totals.ts b/packages/coding-agent/src/core/usage-totals.ts index 24de1a36c..6bcafbd47 100644 --- a/packages/coding-agent/src/core/usage-totals.ts +++ b/packages/coding-agent/src/core/usage-totals.ts @@ -27,6 +27,30 @@ export function addUsageToTotals(totals: UsageTotals, usage: Usage): void { totals.cost += usage.cost.total; } +/** Sum of two usages, keeping the optional token splits when either side reports them. */ +export function combineUsage(first: Usage, second: Usage): Usage { + return { + input: first.input + second.input, + output: first.output + second.output, + cacheRead: first.cacheRead + second.cacheRead, + cacheWrite: first.cacheWrite + second.cacheWrite, + ...(first.cacheWrite1h !== undefined || second.cacheWrite1h !== undefined + ? { cacheWrite1h: (first.cacheWrite1h ?? 0) + (second.cacheWrite1h ?? 0) } + : {}), + ...(first.reasoning !== undefined || second.reasoning !== undefined + ? { reasoning: (first.reasoning ?? 0) + (second.reasoning ?? 0) } + : {}), + totalTokens: first.totalTokens + second.totalTokens, + cost: { + input: first.cost.input + second.cost.input, + output: first.cost.output + second.cost.output, + cacheRead: first.cost.cacheRead + second.cost.cacheRead, + cacheWrite: first.cost.cacheWrite + second.cost.cacheWrite, + total: first.cost.total + second.cost.total, + }, + }; +} + export interface UsageCostBreakdownEntry { key: string; cost: number; diff --git a/packages/coding-agent/src/extensions/codemode/execute.ts b/packages/coding-agent/src/extensions/codemode/execute.ts index b6cb3dcd1..0b41d84c9 100644 --- a/packages/coding-agent/src/extensions/codemode/execute.ts +++ b/packages/coding-agent/src/extensions/codemode/execute.ts @@ -8,7 +8,7 @@ import { writeFile } from "node:fs/promises"; import { tmpdir } from "node:os"; import { join } from "node:path"; import type { AgentTool, AgentToolCallOutcome, AgentToolResult } from "@earendil-works/pi-agent-core"; -import type { AnyModel, ClassifierContext, ImageContent, ModelType, TextContent } from "@earendil-works/pi-ai"; +import type { AnyModel, ClassifierContext, ImageContent, ModelType, TextContent, Usage } from "@earendil-works/pi-ai"; import { type CodemodeResult, CodemodeSandbox, @@ -21,6 +21,7 @@ import { import { getCodemodeWorkerUrl, getQuickJSWasmPath } from "../../config.ts"; import type { ExtensionToolContext } from "../../core/extensions/types.ts"; import type { SessionEntry } from "../../core/session-manager.ts"; +import { combineUsage } from "../../core/usage-totals.ts"; import { Bm25Ranker, createToolSearchDocument, DEFAULT_TOOL_SEARCH_LIMIT } from "../tool-search/tool.ts"; import { CODEMODE_STORE_ENTRY_TYPE, @@ -224,6 +225,11 @@ export async function executeCodemode( const startedAt = performance.now(); const { code, options: sourceOptions } = parseCodemodeSource(input.code); const calls: CodemodeNestedCall[] = []; + // Usage of the script's `models.*` calls. Nested tool calls report theirs through the session. + let modelUsage: Usage | undefined; + const addModelUsage = (usage: Usage) => { + modelUsage = modelUsage ? combineUsage(modelUsage, usage) : usage; + }; const snapshot = (): CodemodeToolDetails => ({ calls: calls.map((call) => ({ ...call })) }); const publish = () => onUpdate?.({ content: [], details: snapshot() }); @@ -264,7 +270,9 @@ export async function executeCodemode( tools: sandboxTools, globals: [ ...createDiscoveryGlobals(callable, samples, options), - ...(options.models && ctx ? createModelGlobals(ctx.modelRegistry, toolCallId, calls, publish) : []), + ...(options.models && ctx + ? createModelGlobals(ctx.modelRegistry, toolCallId, calls, publish, addModelUsage) + : []), ], timeoutMs: sourceOptions.timeoutMs ?? Number.POSITIVE_INFINITY, memoryLimitBytes: CODEMODE_MEMORY_LIMIT_BYTES, @@ -306,6 +314,7 @@ export async function executeCodemode( return { content: [{ type: "text", text: header }, ...truncated.items], details, + ...(modelUsage ? { usage: modelUsage } : {}), ...(result.ok ? {} : { isError: true }), }; } @@ -358,13 +367,15 @@ function createDiscoveryGlobals( /** * `models.*` for scripts: the model registry methods declared in {@link MODEL_GLOBAL_DECLARATIONS}. - * Classifier calls appear as nested call rows so the renderer shows them. + * Classifier calls appear as nested call rows so the renderer shows them, and their usage goes to + * `addUsage`. */ function createModelGlobals( models: CodemodeModelRuntime, toolCallId: string, calls: CodemodeNestedCall[], publish: () => void, + addUsage: (usage: Usage) => void, ): CodemodeTool[] { const limit = createLimiter(MAX_CONCURRENT_MODEL_CALLS); let classifyCount = 0; @@ -416,6 +427,10 @@ function createModelGlobals( record.durationMs = performance.now() - startedAt; record.status = result.stopReason === "stop" ? "ok" : result.stopReason === "aborted" ? "cancelled" : "error"; if (result.errorMessage) record.error = truncateText(result.errorMessage, ERROR_PREVIEW_CHARS); + if (result.usage) { + record.cost = result.usage.cost.total; + addUsage(result.usage); + } publish(); return result; }, diff --git a/packages/coding-agent/src/extensions/codemode/renderer.ts b/packages/coding-agent/src/extensions/codemode/renderer.ts index 684ed4b24..2a76c50fe 100644 --- a/packages/coding-agent/src/extensions/codemode/renderer.ts +++ b/packages/coding-agent/src/extensions/codemode/renderer.ts @@ -2,7 +2,8 @@ * Presentation for the codemode tool. * * The call shows the script; the result lists the nested tool calls with their status as they - * run, followed by the script output without the "Script completed" header. Nested calls are not + * run and the cost of its model calls, followed by the script output without the "Script completed" + * header. Nested calls are not * separate tool rows because they never reach the model as tool calls. */ @@ -28,6 +29,11 @@ function formatDuration(ms: number | undefined): string { return ms < 1000 ? `${Math.round(ms)}ms` : `${(ms / 1000).toFixed(1)}s`; } +/** Cents for larger amounts, two significant digits for the fractions of a cent classifier calls cost. */ +function formatCost(cost: number): string { + return `$${cost >= 0.01 ? cost.toFixed(2) : cost.toPrecision(2)}`; +} + function statusIcon(call: CodemodeNestedCall, theme: Theme): string { switch (call.status) { case "running": @@ -50,6 +56,7 @@ function formatCall(call: CodemodeNestedCall, theme: Theme, expanded: boolean): let line = `${statusIcon(call, theme)} ${theme.fg("toolTitle", call.name)}`; if (args) line += ` ${theme.fg("muted", args)}`; if (duration) line += ` ${theme.fg("dim", duration)}`; + if (call.cost) line += ` ${theme.fg("dim", formatCost(call.cost))}`; if (expanded && call.error) line += `\n ${theme.fg("error", call.error.split("\n").join("\n "))}`; return line; } @@ -85,6 +92,12 @@ export const codemodeRenderers: Pick< `${theme.fg("muted", `... (${calls.length - shown.length} earlier calls,`)} ${keyHint("app.tools.expand", "to expand")}${theme.fg("muted", ")")}`, ); } + // Collapsed rows hide earlier calls, so the total covers every call. + const priced = calls.filter((call) => call.cost); + if (priced.length > 1) { + const total = priced.reduce((sum, call) => sum + (call.cost ?? 0), 0); + lines.push(theme.fg("muted", `Model calls: ${formatCost(total)}`)); + } sections.push(lines.join("\n")); } diff --git a/packages/coding-agent/src/extensions/codemode/tool.ts b/packages/coding-agent/src/extensions/codemode/tool.ts index eca4cbdd2..831237107 100644 --- a/packages/coding-agent/src/extensions/codemode/tool.ts +++ b/packages/coding-agent/src/extensions/codemode/tool.ts @@ -111,6 +111,8 @@ export interface CodemodeNestedCall { durationMs?: number; /** Error text, truncated for display. */ error?: string; + /** Cost in USD of a `models.*` call that reported usage. */ + cost?: number; } export interface CodemodeToolDetails { @@ -182,6 +184,8 @@ interface ClassifierResult { provider: string; model: string; answers: Record; + /** Set when the service reports token counts. Cost is in USD. */ + usage?: { input: number; output: number; totalTokens: number; cost: { total: number } }; stopReason: "stop" | "error" | "aborted"; errorMessage?: string; timestamp: number; diff --git a/packages/coding-agent/test/codemode-renderer.test.ts b/packages/coding-agent/test/codemode-renderer.test.ts index 21b2761eb..122fc363d 100644 --- a/packages/coding-agent/test/codemode-renderer.test.ts +++ b/packages/coding-agent/test/codemode-renderer.test.ts @@ -49,6 +49,28 @@ describe("codemode renderer", () => { expect(text).toBe('✓ read {"path":"a"} 5ms\n\nhello'); }); + it("shows the cost of model calls and their total", () => { + const call = { name: "models.classify", args: "scorer/judge", status: "ok" as const, durationMs: 5 }; + const text = render({ + content: [{ type: "text", text: "Script completed\nWall time 0.1 seconds\nOutput:\n" }], + details: { + calls: [ + { ...call, id: "call/models.classify/1", cost: 0.000012936 }, + { ...call, id: "call/models.classify/2", cost: 0.02 }, + { ...call, id: "call/models.classify/3" }, + ], + }, + }); + expect(text).toBe( + [ + "✓ models.classify scorer/judge 5ms $0.000013", + "✓ models.classify scorer/judge 5ms $0.02", + "✓ models.classify scorer/judge 5ms", + "Model calls: $0.02", + ].join("\n"), + ); + }); + it("shows results without a header, such as rejected options", () => { const text = render( { diff --git a/packages/coding-agent/test/nested-tool-calls.test.ts b/packages/coding-agent/test/nested-tool-calls.test.ts index e6857436f..2ff64121f 100644 --- a/packages/coding-agent/test/nested-tool-calls.test.ts +++ b/packages/coding-agent/test/nested-tool-calls.test.ts @@ -1,4 +1,5 @@ import type { AgentTool, AgentToolCall } from "@earendil-works/pi-agent-core"; +import type { Usage } from "@earendil-works/pi-ai"; import { Type } from "typebox"; import { describe, expect, it } from "vitest"; import { @@ -9,6 +10,17 @@ import { type NestedToolExecutionEvent, } from "../src/core/nested-tool-calls.ts"; +function usage(input: number, cost: number): Usage { + return { + input, + output: 0, + cacheRead: 0, + cacheWrite: 0, + totalTokens: input, + cost: { input: cost, output: 0, cacheRead: 0, cacheWrite: 0, total: cost }, + }; +} + function createRunner(tools: AgentTool[], options: { sequential?: boolean } = {}) { const events: NestedToolExecutionEvent[] = []; const host: NestedToolCallHost = { @@ -66,7 +78,7 @@ describe("NestedToolCallRunner", () => { ["tool_execution_start", "call/2", "call"], ["tool_execution_end", "call/2", "call"], ]); - expect(runner.takeRecord("call")).toEqual({ + expect(runner.takeRecord("call")?.calls).toEqual({ calls: [ { id: "call/1", name: "echo", arguments: { a: 1 }, status: "ok", durationMs: expect.any(Number) }, { @@ -110,7 +122,43 @@ describe("NestedToolCallRunner", () => { await runner.execute("call", "middle", {}); - expect(runner.takeRecord("call")?.calls.map((call) => call.id)).toEqual(["call/1", "call/1/1"]); + expect(runner.takeRecord("call")?.calls?.calls.map((call) => call.id)).toEqual(["call/1", "call/1/1"]); + }); + + it("sums the usage of nested results at every depth", async () => { + const leaf: AgentTool = { + name: "leaf", + label: "Leaf", + description: "Leaf", + parameters: Type.Object({}), + async execute() { + return { content: [], details: {}, usage: usage(10, 0.01) }; + }, + }; + const plain: AgentTool = { ...leaf, name: "plain", execute: async () => ({ content: [], details: {} }) }; + const tools: AgentTool[] = [leaf, plain]; + const { runner } = createRunner(tools); + tools.push({ + name: "middle", + label: "Middle", + description: "Calls leaf", + parameters: Type.Object({}), + async execute(toolCallId) { + await runner.execute(toolCallId, "leaf", {}); + // Its own usage only: the leaf's usage is counted once, by the recorder. + return { content: [], details: {}, usage: usage(5, 0.005) }; + }, + }); + + await runner.execute("call", "middle", {}); + await runner.execute("call", "leaf", {}); + await runner.execute("call", "plain", {}); + await runner.execute("free", "plain", {}); + + const summary = runner.takeRecord("call"); + expect(summary?.usage?.input).toBe(25); + expect(summary?.usage?.cost.total).toBeCloseTo(0.025, 10); + expect(runner.takeRecord("free")).toMatchObject({ calls: { complete: true }, usage: undefined }); }); it("serializes concurrent calls to sequential tools", async () => { diff --git a/packages/coding-agent/test/suite/agent-session-codemode.test.ts b/packages/coding-agent/test/suite/agent-session-codemode.test.ts index 2fc230d92..e9a4680c5 100644 --- a/packages/coding-agent/test/suite/agent-session-codemode.test.ts +++ b/packages/coding-agent/test/suite/agent-session-codemode.test.ts @@ -8,7 +8,7 @@ import { getCurrentTools, type TranscriptContext, } from "@earendil-works/pi-ai"; -import type { ToolResultMessage } from "@earendil-works/pi-ai/compat"; +import type { ToolResultMessage, Usage } from "@earendil-works/pi-ai/compat"; import { Type } from "typebox"; import { afterEach, describe, expect, it } from "vitest"; import type { ExtensionAPI } from "../../src/core/extensions/types.ts"; @@ -65,6 +65,17 @@ const screenshotTool: AgentTool = { }), }; +function usage(input: number, cost: number): Usage { + return { + input, + output: 0, + cacheRead: 0, + cacheWrite: 0, + totalTokens: input, + cost: { input: cost, output: 0, cacheRead: 0, cacheWrite: 0, total: cost }, + }; +} + const codemodeResult = (harness: Harness) => getToolResult(harness, "codemode"); /** The output after the script header, which is checked on the way. */ @@ -225,6 +236,41 @@ describe("AgentSession codemode tool", () => { expect(details.calls.map((call) => call.status)).toEqual(["error", "ok"]); }); + it("adds the usage of nested results to the codemode result", async () => { + const billedTool: AgentTool = { + name: "billed", + label: "Billed", + description: "Run a model", + parameters: Type.Object({}), + execute: async () => ({ content: [{ type: "text", text: "ran" }], details: {}, usage: usage(100, 0.25) }), + }; + const harness = await setup([(pi) => pi.registerTool(createToolDefinitionFromAgentTool(billedTool))]); + harness.setResponses([ + fauxAssistantMessage( + [ + fauxToolCall("codemode", { + code: `await tools.billed({}); await tools.billed({}); await tools.echo({ text: "x" });`, + }), + ], + { stopReason: "toolUse" }, + ), + fauxAssistantMessage("done"), + ]); + + await harness.session.prompt("go"); + + const result = codemodeResult(harness); + expect(result.usage).toMatchObject({ input: 200, totalTokens: 200, cost: { total: 0.5 } }); + // The usage is persisted with the result, so session totals count it. + const persisted = harness.sessionManager + .getEntries() + .find((entry) => entry.type === "message" && entry.message.role === "toolResult"); + expect( + persisted?.type === "message" && persisted.message.role === "toolResult" && persisted.message.usage, + ).toEqual(result.usage); + expect(harness.session.getSessionStats().cost).toBe(0.5); + }); + it("keeps structured content that tool_result handlers replace along with the content", async () => { const harness = await setup([ (pi) => { @@ -536,6 +582,7 @@ describe("codemode models", () => { provider: model.provider, model: model.id, answers: { approved: { type: "bool", probability: text === "good" ? 0.9 : 0.1 } }, + usage: usage(300, 0.001), stopReason: "stop", timestamp: 0, }; @@ -592,6 +639,7 @@ describe("codemode models", () => { same: same.id, missing: (await models.getModelOfType("classifier", "scorer", "nope")) === undefined, probabilities: results.map((r) => r.answers.approved.probability), + cost: results[0].usage.cost.total, }; `, ); @@ -603,6 +651,7 @@ describe("codemode models", () => { same: "judge", missing: true, probabilities: [0.9, 0.1, 0.9, 0.1, 0.9, 0.1], + cost: 0.001, }); expect(observed).toHaveLength(6); expect( @@ -611,9 +660,13 @@ describe("codemode models", () => { // Six classifications with at most four in flight. expect(maxActive()).toBe(4); const details = result.details as unknown as CodemodeToolDetails; - expect(details.calls.map((call) => [call.name, call.args, call.status])).toEqual( - Array.from({ length: 6 }, () => ["models.classify", "scorer/judge", "ok"]), + expect(details.calls.map((call) => [call.name, call.args, call.status, call.cost])).toEqual( + Array.from({ length: 6 }, () => ["models.classify", "scorer/judge", "ok", 0.001]), ); + // The classifications' usage becomes the codemode result's usage. + expect(result.usage?.input).toBe(1800); + expect(result.usage?.cost.total).toBeCloseTo(0.006, 10); + expect(harness.session.getSessionStats().cost).toBeCloseTo(0.006, 10); }); it("reports provider errors as results and invalid arguments as exceptions", async () => { @@ -642,5 +695,6 @@ describe("codemode models", () => { expect(details.calls.map((call) => [call.name, call.status, call.error])).toEqual([ ["models.classify", "error", "classifier exploded"], ]); + expect(result.usage).toBeUndefined(); }); });