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(); }); });