fix(llm): populate usage for the Claude Code local-agent path (#140)
* fix(llm): populate usage for the Claude Code local-agent path LocalAgentAdapter.chat() never set `usage` on its LLMResponse, so AgentLoop's token budget check (tokensUsed >= maxTokens) silently never fired for Claude Code missions. maxIterations was the only brake, and it bounds turn count, not actual spend. Claude Code's --output-format json (instead of text) returns exact per-call token counts, including prompt-cache creation/read tokens that a text-length estimate has no way to see. localAgentChat now requests that format for the claude agent id only. LocalAgentAdapter parses the envelope into usage.promptTokens/completionTokens, with a character-based estimate as a fallback if parsing fails, so usage never goes back to undefined. Verified directly against the CLI: a live run now reports promptTokens 24458, completionTokens 1662 instead of undefined/0. Full suite: 675/675, no regressions. Fixes #139 * test(llm): harden Claude usage fallback --------- Co-authored-by: Joseph Magly <1159087+jmagly@users.noreply.github.com>
This commit is contained in:
1 parent
9020e5fc08
commit
d471da6603
3 files changed
+118
-2
No files matched your search
@@ -34,6 +34,7 @@ const TOOLS = [{
|
|||||||
parameters: { type: 'object' as const, properties: { target: { type: 'string' } }, required: ['target'] },
|
parameters: { type: 'object' as const, properties: { target: { type: 'string' } }, required: ['target'] },
|
||||||
}];
|
}];
|
||||||
const localBackbone = () => new LLMBackbone({ provider: 'local-agent', model: 'codex' } as never);
|
const localBackbone = () => new LLMBackbone({ provider: 'local-agent', model: 'codex' } as never);
|
||||||
|
const claudeBackbone = () => new LLMBackbone({ provider: 'local-agent', model: 'claude' } as never);
|
||||||
const codexBackbone = () => new LLMBackbone({ provider: 'codex', model: 'codex-default' } as never);
|
const codexBackbone = () => new LLMBackbone({ provider: 'codex', model: 'codex-default' } as never);
|
||||||
|
|
||||||
describe('parseTextToolCalls — happy path + drift tolerance', () => {
|
describe('parseTextToolCalls — happy path + drift tolerance', () => {
|
||||||
@@ -136,6 +137,64 @@ describe('local-agent backbone surfaces toolCalls (keyless-path fix)', () => {
|
|||||||
});
|
});
|
||||||
});
|
});
|
||||||
|
|
||||||
|
describe('Claude local-agent usage accounting (#139)', () => {
|
||||||
|
it('extracts content and aggregates real input, output, and cache token usage', async () => {
|
||||||
|
cli.mockResolvedValueOnce(JSON.stringify({
|
||||||
|
result: 'done',
|
||||||
|
usage: {
|
||||||
|
input_tokens: 11,
|
||||||
|
output_tokens: 7,
|
||||||
|
cache_creation_input_tokens: 13,
|
||||||
|
cache_read_input_tokens: 17,
|
||||||
|
},
|
||||||
|
}));
|
||||||
|
|
||||||
|
const res = await claudeBackbone().chat([{ role: 'user', content: 'hello' }]);
|
||||||
|
|
||||||
|
expect(res.content).toBe('done');
|
||||||
|
expect(res.usage).toEqual({ promptTokens: 41, completionTokens: 7, totalTokens: 48 });
|
||||||
|
});
|
||||||
|
|
||||||
|
it.each([
|
||||||
|
['missing', { result: 'fallback' }],
|
||||||
|
['empty', { result: 'fallback', usage: {} }],
|
||||||
|
['zero', { result: 'fallback', usage: { input_tokens: 0, output_tokens: 0 } }],
|
||||||
|
['malformed', { result: 'fallback', usage: { input_tokens: 'many', output_tokens: 2 } }],
|
||||||
|
['negative', { result: 'fallback', usage: { input_tokens: -1, output_tokens: 2 } }],
|
||||||
|
])('retains envelope content and estimates when usage is %s', async (_case, envelope) => {
|
||||||
|
cli.mockResolvedValueOnce(JSON.stringify(envelope));
|
||||||
|
|
||||||
|
const res = await claudeBackbone().chat([{ role: 'user', content: 'hello' }]);
|
||||||
|
|
||||||
|
expect(res.content).toBe('fallback');
|
||||||
|
expect(res.usage?.promptTokens).toBeGreaterThan(0);
|
||||||
|
expect(res.usage?.completionTokens).toBe(Math.ceil('fallback'.length / 4));
|
||||||
|
expect(res.usage?.totalTokens).toBe((res.usage?.promptTokens || 0) + (res.usage?.completionTokens || 0));
|
||||||
|
});
|
||||||
|
|
||||||
|
it('estimates usage and preserves raw content when the CLI does not return JSON', async () => {
|
||||||
|
cli.mockResolvedValueOnce('plain reply');
|
||||||
|
|
||||||
|
const res = await claudeBackbone().chat([{ role: 'user', content: 'hello' }]);
|
||||||
|
|
||||||
|
expect(res.content).toBe('plain reply');
|
||||||
|
expect(res.usage?.totalTokens).toBeGreaterThan(0);
|
||||||
|
});
|
||||||
|
|
||||||
|
it('parses tool calls from the JSON result rather than the envelope', async () => {
|
||||||
|
cli.mockResolvedValueOnce(JSON.stringify({
|
||||||
|
result: '{"tool_calls":[{"name":"nmap_scan","arguments":{"target":"x"}}]}',
|
||||||
|
usage: { input_tokens: 10, output_tokens: 5 },
|
||||||
|
}));
|
||||||
|
|
||||||
|
const res = await claudeBackbone().chatWithTools([{ role: 'user', content: 'scan' }], TOOLS as never);
|
||||||
|
|
||||||
|
expect(res.toolCalls?.[0]?.name).toBe('nmap_scan');
|
||||||
|
expect(res.finishReason).toBe('tool_calls');
|
||||||
|
expect(res.usage?.totalTokens).toBe(15);
|
||||||
|
});
|
||||||
|
});
|
||||||
|
|
||||||
describe('codex backbone surfaces toolCalls (guards the CodexAdapter half of the fix)', () => {
|
describe('codex backbone surfaces toolCalls (guards the CodexAdapter half of the fix)', () => {
|
||||||
it('returns toolCalls when codex emits the contract', async () => {
|
it('returns toolCalls when codex emits the contract', async () => {
|
||||||
fileRead.mockResolvedValueOnce('```json\n{"tool_calls":[{"name":"nmap_scan","arguments":{"target":"x"}}]}\n```');
|
fileRead.mockResolvedValueOnce('```json\n{"tool_calls":[{"name":"nmap_scan","arguments":{"target":"x"}}]}\n```');
|
||||||
|
|||||||
@@ -470,7 +470,10 @@ export function localAgentChat(id: string, prompt: string, opts: { model?: strin
|
|||||||
let outFile: string | null = null;
|
let outFile: string | null = null;
|
||||||
let workDir: string | null = null;
|
let workDir: string | null = null;
|
||||||
if (id === 'claude') {
|
if (id === 'claude') {
|
||||||
args = ['-p', '--output-format', 'text', ...(model ? ['--model', model] : [])];
|
// json (not text): the envelope carries REAL per-call token usage (input/output tokens
|
||||||
|
// plus prompt-cache creation/read) that LocalAgentAdapter.chat() parses to drive
|
||||||
|
// AgentLoop's token budget check. text mode reports no usage at all.
|
||||||
|
args = ['-p', '--output-format', 'json', ...(model ? ['--model', model] : [])];
|
||||||
} else if (id === 'codex') {
|
} else if (id === 'codex') {
|
||||||
workDir = mkdtempSync(join(tmpdir(), 't3mp3st-codexllm-'));
|
workDir = mkdtempSync(join(tmpdir(), 't3mp3st-codexllm-'));
|
||||||
outFile = join(workDir, 'reply.txt');
|
outFile = join(workDir, 'reply.txt');
|
||||||
|
|||||||
+55
-1
@@ -1308,6 +1308,51 @@ class CodexAdapter implements LLMProviderAdapter {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// Claude Code's `--output-format json` envelope (requested only for agentId 'claude', see
|
||||||
|
// localAgentChat in local-agents.ts). Carries REAL per-call token accounting straight from the
|
||||||
|
// CLI — input/output tokens plus the prompt-cache creation/read tokens from Claude Code's own
|
||||||
|
// bootstrap (CLAUDE.md, skills, MCP tool manifests), which a text-length estimate has no way to
|
||||||
|
// see and which measured 4-25k tokens on a single trivial call in testing.
|
||||||
|
interface ClaudeJsonEnvelope {
|
||||||
|
result?: string;
|
||||||
|
usage?: {
|
||||||
|
input_tokens?: number;
|
||||||
|
output_tokens?: number;
|
||||||
|
cache_creation_input_tokens?: number;
|
||||||
|
cache_read_input_tokens?: number;
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
|
/** Parse the envelope; invalid usage is omitted so the caller can retain content and estimate. */
|
||||||
|
function parseClaudeJsonEnvelope(raw: string): { content: string; usage?: LLMResponse['usage'] } | null {
|
||||||
|
let json: ClaudeJsonEnvelope;
|
||||||
|
try {
|
||||||
|
json = JSON.parse(raw);
|
||||||
|
} catch {
|
||||||
|
return null;
|
||||||
|
}
|
||||||
|
if (typeof json.result !== 'string') return null;
|
||||||
|
const u = json.usage || {};
|
||||||
|
const tokenValues = [u.input_tokens, u.output_tokens, u.cache_creation_input_tokens, u.cache_read_input_tokens];
|
||||||
|
if (!tokenValues.some((v) => v !== undefined)
|
||||||
|
|| tokenValues.some((v) => v !== undefined && (!Number.isFinite(v) || v < 0))) {
|
||||||
|
return { content: json.result };
|
||||||
|
}
|
||||||
|
const promptTokens = (u.input_tokens ?? 0) + (u.cache_creation_input_tokens ?? 0) + (u.cache_read_input_tokens ?? 0);
|
||||||
|
const completionTokens = u.output_tokens ?? 0;
|
||||||
|
if (promptTokens + completionTokens === 0) return { content: json.result };
|
||||||
|
return { content: json.result, usage: { promptTokens, completionTokens, totalTokens: promptTokens + completionTokens } };
|
||||||
|
}
|
||||||
|
|
||||||
|
// Character-based fallback ONLY — used when the real envelope can't be parsed (older CLI,
|
||||||
|
// unexpected output shape). Rough on purpose: the goal is to give AgentLoop's budget check SOME
|
||||||
|
// number instead of a permanently-zero one, not to be precise.
|
||||||
|
function estimateUsage(promptText: string, completionText: string): LLMResponse['usage'] {
|
||||||
|
const promptTokens = Math.ceil(promptText.length / 4);
|
||||||
|
const completionTokens = Math.ceil(completionText.length / 4);
|
||||||
|
return { promptTokens, completionTokens, totalTokens: promptTokens + completionTokens };
|
||||||
|
}
|
||||||
|
|
||||||
// Generic adapter that drives a CONNECTED local agent CLI (Claude Code / Codex / Hermes) as the LLM
|
// Generic adapter that drives a CONNECTED local agent CLI (Claude Code / Codex / Hermes) as the LLM
|
||||||
// backend — NO API key needed; each CLI uses its own login. The agent id travels in `config.model`,
|
// backend — NO API key needed; each CLI uses its own login. The agent id travels in `config.model`,
|
||||||
// optionally with an underlying model after a `::` separator ("claude::opus", "claude::claude-opus-4-8")
|
// optionally with an underlying model after a `::` separator ("claude::opus", "claude::claude-opus-4-8")
|
||||||
@@ -1347,7 +1392,15 @@ class LocalAgentAdapter implements LLMProviderAdapter {
|
|||||||
const { agentId, agentModel } = this.parseAgentSpec();
|
const { agentId, agentModel } = this.parseAgentSpec();
|
||||||
const prompt = this.formatPrompt(messages, options);
|
const prompt = this.formatPrompt(messages, options);
|
||||||
const timeoutMs = typeof this.config.timeout === 'number' && this.config.timeout > 0 ? this.config.timeout : undefined;
|
const timeoutMs = typeof this.config.timeout === 'number' && this.config.timeout > 0 ? this.config.timeout : undefined;
|
||||||
const content = (await localAgentChat(agentId, prompt, { model: agentModel, timeoutMs })).trim();
|
const raw = (await localAgentChat(agentId, prompt, { model: agentModel, timeoutMs })).trim();
|
||||||
|
// Claude requests --output-format json (see local-agents.ts) so this parses to REAL usage.
|
||||||
|
// Anything else (parse failure, or a non-claude agent) falls back to the raw text — claude
|
||||||
|
// additionally gets a character-based usage ESTIMATE so its budget check is never blind again.
|
||||||
|
const parsed = agentId === 'claude' ? parseClaudeJsonEnvelope(raw) : null;
|
||||||
|
const content = parsed ? parsed.content : raw;
|
||||||
|
const usage = agentId === 'claude'
|
||||||
|
? (parsed?.usage ?? estimateUsage(prompt, parsed?.content ?? raw))
|
||||||
|
: undefined;
|
||||||
// Tool-calling over text: if the Arsenal was offered, parse the agent's tool requests so the
|
// Tool-calling over text: if the Arsenal was offered, parse the agent's tool requests so the
|
||||||
// ReAct loop EXECUTES them instead of treating this planning turn as the (abstaining) final answer.
|
// ReAct loop EXECUTES them instead of treating this planning turn as the (abstaining) final answer.
|
||||||
const toolCalls = options?.tools?.length ? parseTextToolCalls(content) : undefined;
|
const toolCalls = options?.tools?.length ? parseTextToolCalls(content) : undefined;
|
||||||
@@ -1356,6 +1409,7 @@ class LocalAgentAdapter implements LLMProviderAdapter {
|
|||||||
model: `local-agent:${agentId}${agentModel ? '/' + agentModel : ''}`,
|
model: `local-agent:${agentId}${agentModel ? '/' + agentModel : ''}`,
|
||||||
finishReason: toolCalls?.length ? 'tool_calls' : 'stop',
|
finishReason: toolCalls?.length ? 'tool_calls' : 'stop',
|
||||||
toolCalls,
|
toolCalls,
|
||||||
|
usage,
|
||||||
};
|
};
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
Reference in new issue
Block a user