diff --git a/.changeset/agent-plugin-safety-costs.md b/.changeset/agent-plugin-safety-costs.md new file mode 100644 index 000000000..b33ed96b8 --- /dev/null +++ b/.changeset/agent-plugin-safety-costs.md @@ -0,0 +1,6 @@ +--- +"@aoagents/ao-plugin-agent-codex": patch +"@aoagents/ao-plugin-agent-claude-code": patch +--- + +Improve Claude Code and Codex session cost estimates to account for cached-token spend, and make Codex restore commands fall back to approval prompts for worker sessions instead of blindly reusing dangerous bypass flags. diff --git a/packages/plugins/agent-claude-code/src/index.test.ts b/packages/plugins/agent-claude-code/src/index.test.ts index 0d9c9cc3c..f283e681e 100644 --- a/packages/plugins/agent-claude-code/src/index.test.ts +++ b/packages/plugins/agent-claude-code/src/index.test.ts @@ -1,5 +1,11 @@ import { describe, it, expect, vi, beforeEach } from "vitest"; -import { createActivitySignal, type Session, type RuntimeHandle, type AgentLaunchConfig, type WorkspaceHooksConfig } from "@aoagents/ao-core"; +import { + createActivitySignal, + type Session, + type RuntimeHandle, + type AgentLaunchConfig, + type WorkspaceHooksConfig, +} from "@aoagents/ao-core"; // --------------------------------------------------------------------------- // Hoisted mocks — available inside vi.mock factories @@ -50,7 +56,13 @@ vi.mock("node:os", () => ({ homedir: mockHomedir, })); -import { create, manifest, default as defaultExport, resetPsCache, METADATA_UPDATER_SCRIPT } from "./index.js"; +import { + create, + manifest, + default as defaultExport, + resetPsCache, + METADATA_UPDATER_SCRIPT, +} from "./index.js"; // --------------------------------------------------------------------------- // Test helpers @@ -582,6 +594,17 @@ describe("getSessionInfo", () => { expect(result?.cost?.outputTokens).toBe(50); }); + it("uses model-aware pricing when cached tokens are present", async () => { + const jsonl = [ + '{"type":"assistant","model":"claude-sonnet-4-5","usage":{"input_tokens":1000,"output_tokens":100,"cache_read_input_tokens":10000,"cache_creation_input_tokens":2000}}', + ].join("\n"); + mockJsonlFiles(jsonl); + const result = await agent.getSessionInfo(makeSession()); + expect(result?.cost?.inputTokens).toBe(13000); + expect(result?.cost?.outputTokens).toBe(100); + expect(result?.cost?.estimatedCostUsd).toBeGreaterThan(0); + }); + it("uses costUSD field when present", async () => { const jsonl = [ '{"type":"user","message":{"content":"hi"}}', @@ -725,12 +748,8 @@ describe("METADATA_UPDATER_SCRIPT content", () => { it("does NOT use ^-anchored regexes directly on $command for gh/git detection", () => { // The old buggy patterns matched $command with ^ anchor. // After the fix, ^ is still used but on $clean_command (which has cd stripped). - expect(METADATA_UPDATER_SCRIPT).not.toMatch( - /"\$command"\s*=~\s*\^gh/, - ); - expect(METADATA_UPDATER_SCRIPT).not.toMatch( - /"\$command"\s*=~\s*\^git/, - ); + expect(METADATA_UPDATER_SCRIPT).not.toMatch(/"\$command"\s*=~\s*\^gh/); + expect(METADATA_UPDATER_SCRIPT).not.toMatch(/"\$command"\s*=~\s*\^git/); }); it("strips cd prefixes with both && and ; delimiters", () => { @@ -742,21 +761,15 @@ describe("METADATA_UPDATER_SCRIPT content", () => { }); it("detects gh pr create on clean_command", () => { - expect(METADATA_UPDATER_SCRIPT).toMatch( - /"\$clean_command"\s*=~\s*\^gh\[/, - ); + expect(METADATA_UPDATER_SCRIPT).toMatch(/"\$clean_command"\s*=~\s*\^gh\[/); }); it("detects git checkout -b on clean_command", () => { - expect(METADATA_UPDATER_SCRIPT).toMatch( - /"\$clean_command"\s*=~\s*\^git\[.*checkout/, - ); + expect(METADATA_UPDATER_SCRIPT).toMatch(/"\$clean_command"\s*=~\s*\^git\[.*checkout/); }); it("detects gh pr merge on clean_command", () => { - expect(METADATA_UPDATER_SCRIPT).toMatch( - /"\$clean_command"\s*=~\s*\^gh\[.*merge/, - ); + expect(METADATA_UPDATER_SCRIPT).toMatch(/"\$clean_command"\s*=~\s*\^gh\[.*merge/); }); }); diff --git a/packages/plugins/agent-claude-code/src/index.ts b/packages/plugins/agent-claude-code/src/index.ts index a6fe9e0e0..c661b6306 100644 --- a/packages/plugins/agent-claude-code/src/index.ts +++ b/packages/plugins/agent-claude-code/src/index.ts @@ -307,8 +307,7 @@ async function parseJsonlFileTail(filePath: string, maxBytes = 131_072): Promise // Skip potentially truncated first line only when we started mid-file. // If offset === 0 we read from the start so the first line is complete. const firstNewline = content.indexOf("\n"); - const safeContent = - offset > 0 && firstNewline >= 0 ? content.slice(firstNewline + 1) : content; + const safeContent = offset > 0 && firstNewline >= 0 ? content.slice(firstNewline + 1) : content; const lines: JsonlLine[] = []; for (const line of safeContent.split("\n")) { const trimmed = line.trim(); @@ -326,9 +325,7 @@ async function parseJsonlFileTail(filePath: string, maxBytes = 131_072): Promise } /** Extract auto-generated summary from JSONL (last "summary" type entry) */ -function extractSummary( - lines: JsonlLine[], -): { summary: string; isFallback: boolean } | null { +function extractSummary(lines: JsonlLine[]): { summary: string; isFallback: boolean } | null { for (let i = lines.length - 1; i >= 0; i--) { const line = lines[i]; if (line?.type === "summary" && line.summary) { @@ -358,6 +355,8 @@ function extractSummary( function extractCost(lines: JsonlLine[]): CostEstimate | undefined { let inputTokens = 0; let outputTokens = 0; + let cachedReadTokens = 0; + let cacheCreationTokens = 0; let totalCost = 0; for (const line of lines) { @@ -373,8 +372,8 @@ function extractCost(lines: JsonlLine[]): CostEstimate | undefined { // double-counting if a line contains both. if (line.usage) { inputTokens += line.usage.input_tokens ?? 0; - inputTokens += line.usage.cache_read_input_tokens ?? 0; - inputTokens += line.usage.cache_creation_input_tokens ?? 0; + cachedReadTokens += line.usage.cache_read_input_tokens ?? 0; + cacheCreationTokens += line.usage.cache_creation_input_tokens ?? 0; outputTokens += line.usage.output_tokens ?? 0; } else { if (typeof line.inputTokens === "number") { @@ -386,19 +385,29 @@ function extractCost(lines: JsonlLine[]): CostEstimate | undefined { } } - if (inputTokens === 0 && outputTokens === 0 && totalCost === 0) { + if ( + inputTokens === 0 && + outputTokens === 0 && + totalCost === 0 && + cachedReadTokens === 0 && + cacheCreationTokens === 0 + ) { return undefined; } - // Rough estimate when no direct cost data — uses Sonnet 4.5 pricing as a - // baseline. Will be inaccurate for other models (Opus, Haiku) but provides - // a useful order-of-magnitude signal. TODO: make pricing configurable or - // infer from model field in JSONL. - if (totalCost === 0 && (inputTokens > 0 || outputTokens > 0)) { - totalCost = (inputTokens / 1_000_000) * 3.0 + (outputTokens / 1_000_000) * 15.0; + if (totalCost === 0) { + totalCost = + (inputTokens / 1_000_000) * 3.0 + + (outputTokens / 1_000_000) * 15.0 + + (cachedReadTokens / 1_000_000) * 0.3 + + (cacheCreationTokens / 1_000_000) * 3.75; } - return { inputTokens, outputTokens, estimatedCostUsd: totalCost }; + return { + inputTokens: inputTokens + cachedReadTokens + cacheCreationTokens, + outputTokens, + estimatedCostUsd: totalCost, + }; } // ============================================================================= @@ -819,7 +828,11 @@ function createClaudeCodeAgent(): Agent { const parts: string[] = ["claude", "--resume", shellEscape(sessionUuid)]; const permissionMode = normalizeAgentPermissionMode(project.agentConfig?.permissions); - if (permissionMode === "permissionless" || permissionMode === "auto-edit") { + const isOrchestrator = session.metadata?.["role"] === "orchestrator"; + if ( + isOrchestrator && + (permissionMode === "permissionless" || permissionMode === "auto-edit") + ) { parts.push("--dangerously-skip-permissions"); } diff --git a/packages/plugins/agent-codex/src/index.test.ts b/packages/plugins/agent-codex/src/index.test.ts index 16ebd6502..ca547d935 100644 --- a/packages/plugins/agent-codex/src/index.test.ts +++ b/packages/plugins/agent-codex/src/index.test.ts @@ -1037,10 +1037,10 @@ describe("getSessionInfo", () => { expect(result!.summary).toBe("Codex session (o3-mini)"); expect(result!.summaryIsFallback).toBe(true); expect(result!.cost).toBeDefined(); - // cached_tokens/reasoning_tokens are subsets, not additive - // input: 1000 + 2000 = 3000 + // cached tokens count toward effective input spend + // input: 1000 + 2000 + 200 = 3200 // output: 500 + 300 = 800 - expect(result!.cost!.inputTokens).toBe(3000); + expect(result!.cost!.inputTokens).toBe(3200); expect(result!.cost!.outputTokens).toBe(800); expect(result!.cost!.estimatedCostUsd).toBeGreaterThan(0); }); @@ -1430,7 +1430,10 @@ describe("getRestoreCommand", () => { mockReadFile.mockResolvedValue(content); mockStat.mockResolvedValue({ mtimeMs: 1000 }); - const session = makeSession({ workspacePath: "/workspace/test" }); + const session = makeSession({ + workspacePath: "/workspace/test", + metadata: { role: "orchestrator" }, + }); const cmd = await agent.getRestoreCommand!( session, makeProjectConfig({ @@ -1453,7 +1456,10 @@ describe("getRestoreCommand", () => { mockReadFile.mockResolvedValue(content); mockStat.mockResolvedValue({ mtimeMs: 1000 }); - const session = makeSession({ workspacePath: "/workspace/test" }); + const session = makeSession({ + workspacePath: "/workspace/test", + metadata: { role: "orchestrator" }, + }); const cmd = await agent.getRestoreCommand!( session, makeProjectConfig({ @@ -1464,6 +1470,29 @@ describe("getRestoreCommand", () => { expect(cmd).toContain("--dangerously-bypass-approvals-and-sandbox"); }); + it("demotes worker restore permissionless mode to ask-for-approval never", async () => { + const content = jsonl( + { type: "session_meta", cwd: "/workspace/test", model: "gpt-4o" }, + { threadId: "thread-1" }, + ); + mockReaddir.mockResolvedValue(["sess.jsonl"]); + setupMockOpen(content); + setupMockStream(content); + mockReadFile.mockResolvedValue(content); + mockStat.mockResolvedValue({ mtimeMs: 1000 }); + + const session = makeSession({ workspacePath: "/workspace/test", metadata: { role: "worker" } }); + const cmd = await agent.getRestoreCommand!( + session, + makeProjectConfig({ + agentConfig: { permissions: "permissionless" }, + }), + ); + + expect(cmd).toContain("--ask-for-approval never"); + expect(cmd).not.toContain("--dangerously-bypass-approvals-and-sandbox"); + }); + it("includes --ask-for-approval never from project config", async () => { const content = jsonl( { type: "session_meta", cwd: "/workspace/test", model: "gpt-4o" }, diff --git a/packages/plugins/agent-codex/src/index.ts b/packages/plugins/agent-codex/src/index.ts index 9bbf51f9e..4ac3714f9 100644 --- a/packages/plugins/agent-codex/src/index.ts +++ b/packages/plugins/agent-codex/src/index.ts @@ -185,10 +185,7 @@ async function readJsonlPrefixLines(filePath: string, maxLines: number): Promise * session_meta entry matching the given workspace path. This avoids parsing a * truncated session_meta line when Codex embeds large base_instructions. */ -async function sessionFileMatchesCwd( - filePath: string, - workspacePath: string, -): Promise { +async function sessionFileMatchesCwd(filePath: string, workspacePath: string): Promise { try { const lines = await readJsonlPrefixLines(filePath, SESSION_MATCH_SCAN_LINE_LIMIT); for (const line of lines) { @@ -245,6 +242,8 @@ interface CodexSessionData { threadId: string | null; inputTokens: number; outputTokens: number; + cachedTokens: number; + reasoningTokens: number; } /** @@ -254,7 +253,14 @@ interface CodexSessionData { */ async function streamCodexSessionData(filePath: string): Promise { try { - const data: CodexSessionData = { model: null, threadId: null, inputTokens: 0, outputTokens: 0 }; + const data: CodexSessionData = { + model: null, + threadId: null, + inputTokens: 0, + outputTokens: 0, + cachedTokens: 0, + reasoningTokens: 0, + }; const rl = createInterface({ input: createReadStream(filePath, { encoding: "utf-8" }), crlfDelay: Infinity, @@ -319,6 +325,8 @@ async function streamCodexSessionData(filePath: string): Promise { // ============================================================================= /** Append approval-policy flags to a command parts array */ -function appendApprovalFlags(parts: string[], permissions: string | undefined): void { +function appendApprovalFlags( + parts: string[], + permissions: string | undefined, + allowDangerousBypass = true, +): void { const mode = normalizeAgentPermissionMode(permissions); if (mode === "permissionless") { - parts.push("--dangerously-bypass-approvals-and-sandbox"); + if (allowDangerousBypass) { + parts.push("--dangerously-bypass-approvals-and-sandbox"); + } else { + parts.push("--ask-for-approval", "never"); + } } else if (mode === "auto-edit") { parts.push("--ask-for-approval", "never"); } else if (mode === "suggest") { @@ -421,7 +437,10 @@ async function findCodexSessionFileCached(workspacePath: string): Promise { + async getActivityState( + session: Session, + readyThresholdMs?: number, + ): Promise { const threshold = readyThresholdMs ?? DEFAULT_READY_THRESHOLD_MS; // Check if process is running first @@ -670,15 +692,19 @@ function createCodexAgent(): Agent { const agentSessionId = basename(sessionFile, ".jsonl"); - const cost: CostEstimate | undefined = - data.inputTokens === 0 && data.outputTokens === 0 - ? undefined - : { - inputTokens: data.inputTokens, - outputTokens: data.outputTokens, - estimatedCostUsd: - (data.inputTokens / 1_000_000) * 2.5 + (data.outputTokens / 1_000_000) * 10.0, - }; + let cost: CostEstimate | undefined; + const totalInputTokens = data.inputTokens + data.cachedTokens; + if (totalInputTokens > 0 || data.outputTokens > 0 || data.reasoningTokens > 0) { + const estimatedCostUsd = + (data.inputTokens / 1_000_000) * 2.5 + + (data.cachedTokens / 1_000_000) * 0.625 + + ((data.outputTokens + data.reasoningTokens) / 1_000_000) * 10.0; + cost = { + inputTokens: totalInputTokens, + outputTokens: data.outputTokens, + estimatedCostUsd, + }; + } return { summary: data.model ? `Codex session (${data.model})` : null, @@ -707,7 +733,8 @@ function createCodexAgent(): Agent { const parts: string[] = [shellEscape(binary), "resume"]; appendNoUpdateCheckFlag(parts); - appendApprovalFlags(parts, project.agentConfig?.permissions); + const isOrchestrator = session.metadata?.["role"] === "orchestrator"; + appendApprovalFlags(parts, project.agentConfig?.permissions, isOrchestrator); const effectiveModel = (project.agentConfig?.model ?? data.model) as string | undefined; appendModelFlags(parts, effectiveModel ?? undefined);