diff --git a/packages/loop-js/src/engine/claude.ts b/packages/loop-js/src/engine/claude.ts index 64f46ff..b336464 100644 --- a/packages/loop-js/src/engine/claude.ts +++ b/packages/loop-js/src/engine/claude.ts @@ -130,6 +130,17 @@ export type TokenUsage = { cache_creation?: { ephemeral_5m_input_tokens?: number | null; ephemeral_1h_input_tokens?: number | null } | null } +function hasTokenUsage(u: TokenUsage): boolean { + return ( + (u.input_tokens ?? 0) !== 0 || + (u.output_tokens ?? 0) !== 0 || + (u.cache_creation_input_tokens ?? 0) !== 0 || + (u.cache_read_input_tokens ?? 0) !== 0 || + (u.cache_creation?.ephemeral_5m_input_tokens ?? 0) !== 0 || + (u.cache_creation?.ephemeral_1h_input_tokens ?? 0) !== 0 + ) +} + export function stepUsage(u: TokenUsage, model: string): StepUsage { const inputTokens = u.input_tokens ?? 0 const outputTokens = u.output_tokens ?? 0 @@ -146,7 +157,7 @@ export function stepUsage(u: TokenUsage, model: string): StepUsage { cachedInputTokens * p.input * CACHE_READ + outputTokens * p.output) / 1e6 - : 0 // an unpriced model derives nothing; the result's `total_cost_usd` reconciles it + : 0 // an unpriced model derives nothing; drainSession fails closed before this can bypass the guard return { inputTokens, outputTokens, cachedInputTokens, usd } } @@ -181,6 +192,9 @@ export async function* drainSession( // repeating the turn's cumulative usage (contract §Mapping 1). Cost a turn once, on its id. if (m.message.id !== turn) { turn = m.message.id + if (!PRICES[model] && hasTokenUsage(m.message.usage)) { + throw new Interruption("budget", `unpriced model '${model}' produced token usage`) + } const usage = stepUsage(m.message.usage, model) derived += usage.usd yield { kind: "cost", usage } diff --git a/packages/loop-js/src/engine/claude.unpriced.test.ts b/packages/loop-js/src/engine/claude.unpriced.test.ts new file mode 100644 index 0000000..2ca0784 --- /dev/null +++ b/packages/loop-js/src/engine/claude.unpriced.test.ts @@ -0,0 +1,30 @@ +import { expect, test } from "bun:test" +import type { SDKMessage } from "@anthropic-ai/claude-agent-sdk" +import { drainSession, stepUsage } from "./claude.ts" +import { Interruption } from "./executor.ts" + +async function* feed(...messages: SDKMessage[]): AsyncGenerator { + for (const message of messages) yield message +} + +test("unpriced non-zero streamed usage fails closed at the cost boundary", async () => { + const assistant = { + type: "assistant", + message: { id: "turn_1", content: [{ type: "text", text: "partial" }], usage: { input_tokens: 1 } }, + } as unknown as SDKMessage + const gen = drainSession(feed(assistant), "some-future-model") + + expect(await gen.next()).toEqual({ done: false, value: { kind: "text", text: "partial" } }) + try { + await gen.next() + throw new Error("expected unpriced usage to fail closed") + } catch (err) { + expect(err).toBeInstanceOf(Interruption) + expect((err as Interruption).cause).toBe("budget") + expect((err as Interruption).detail).toContain("unpriced model") + } +}) + +test("stepUsage keeps zero as its best-effort value for an unpriced model", () => { + expect(stepUsage({ input_tokens: 1 }, "some-future-model").usd).toBe(0) +})