bkinnightskytw
/
agent-orchestrator
mirror of https://github.com/ComposioHQ/agent-orchestrator.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

fix: improve agent plugin cost accounting and restore safety

This commit is contained in:
harshitsinghbhandari 2026-04-18 13:56:59 +05:30
parent 365c149887
commit b0d0994efd
5 changed files with 145 additions and 57 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
.changeset/agent-plugin-safety-costs.md Normal file
View File

@ -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.

47
packages/plugins/agent-claude-code/src/index.test.ts
View File

@ -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/);
});
});

45
packages/plugins/agent-claude-code/src/index.ts
View File

@ -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");
}

39
packages/plugins/agent-codex/src/index.test.ts
View File

@ -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" },

65
packages/plugins/agent-codex/src/index.ts
View File

@ -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<boolean> {
async function sessionFileMatchesCwd(filePath: string, workspacePath: string): Promise<boolean> {
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<CodexSessionData | null> {
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<CodexSessionDat
if (entry.type === "event_msg" && entry.msg?.type === "token_count") {
data.inputTokens += entry.msg.input_tokens ?? 0;
data.outputTokens += entry.msg.output_tokens ?? 0;
data.cachedTokens += entry.msg.cached_tokens ?? 0;
data.reasoningTokens += entry.msg.reasoning_tokens ?? 0;
}
} catch {
// Skip malformed lines
@ -377,10 +385,18 @@ export async function resolveCodexBinary(): Promise<string> {
// =============================================================================
/** 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<string
return cached.path;
}
const result = await findCodexSessionFile(workspacePath);
sessionFileCache.set(workspacePath, { path: result, expiry: Date.now() + SESSION_FILE_CACHE_TTL_MS });
sessionFileCache.set(workspacePath, {
path: result,
expiry: Date.now() + SESSION_FILE_CACHE_TTL_MS,
});
return result;
}
@ -498,7 +517,10 @@ function createCodexAgent(): Agent {
return "active";
},
async getActivityState(session: Session, readyThresholdMs?: number): Promise<ActivityDetection | null> {
async getActivityState(
session: Session,
readyThresholdMs?: number,
): Promise<ActivityDetection | null> {
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);