| 1 | import { describe, expect, it } from "vitest"; |
| 2 | import { COMPACT_MAX_OUTPUT, DEFAULT_MAX_OUTPUT, formatTokens, pinRefusal, prepareRequest, readBody, readUsage, usageOf, weighted } from "../../src/agentic/usage.js"; |
| 3 | |
| 4 | const stream = (chunks: string[]) => |
| 5 | new ReadableStream<Uint8Array>({ |
| 6 | start(c) { |
| 7 | for (const s of chunks) c.enqueue(new TextEncoder().encode(s)); |
| 8 | c.close(); |
| 9 | }, |
| 10 | }); |
| 11 | |
| 12 | describe("usageOf", () => { |
| 13 | it("reads Anthropic, OpenAI Responses and Chat Completions usage", () => { |
| 14 | expect(usageOf({ usage: { input_tokens: 10, output_tokens: 5, cache_read_input_tokens: 100, cache_creation_input_tokens: 20 } })).toEqual({ input: 10, output: 5, cacheRead: 100, cacheWrite: 20 }); |
| 15 | expect(usageOf({ type: "message_start", message: { usage: { input_tokens: 7, output_tokens: 1 } } })).toEqual({ input: 7, output: 1, cacheRead: 0, cacheWrite: 0 }); |
| 16 | expect(usageOf({ type: "response.completed", response: { usage: { input_tokens: 120, input_tokens_details: { cached_tokens: 100 }, output_tokens: 9 } } })).toEqual({ input: 20, output: 9, cacheRead: 100, cacheWrite: 0 }); |
| 17 | expect(usageOf({ usage: { prompt_tokens: 50, completion_tokens: 4, prompt_tokens_details: { cached_tokens: 40 } } })).toEqual({ input: 10, output: 4, cacheRead: 40, cacheWrite: 0 }); |
| 18 | expect(usageOf({ type: "ping" })).toBeNull(); |
| 19 | expect(usageOf({ usage: { input_tokens: -3, output_tokens: "9" } })).toEqual({ input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }); |
| 20 | }); |
| 21 | it("weighs cache reads as a tenth", () => { |
| 22 | expect(weighted({ input: 10, output: 5, cacheRead: 101, cacheWrite: 20 })).toBe(10 + 5 + 20 + 11); |
| 23 | expect(formatTokens(2_000_000)).toBe("2M"); |
| 24 | expect(formatTokens(120_000)).toBe("120k"); |
| 25 | expect(formatTokens(950)).toBe("950"); |
| 26 | }); |
| 27 | }); |
| 28 | |
| 29 | describe("readUsage", () => { |
| 30 | it("adds up an Anthropic stream's running totals, split across chunks", async () => { |
| 31 | const sse = [ |
| 32 | 'event: message_start\ndata: {"type":"message_start","message":{"usage":{"input_tokens":1000,"cache_read_input_tokens":5000,"output_tokens":1}}}\n\n', |
| 33 | 'event: content_block_delta\ndata: {"type":"content_block_delta","delta":{"text":"hi"}}\n\nevent: message_delta\nda', |
| 34 | 'ta: {"type":"message_delta","usage":{"output_tokens":300}}\n\nevent: message_stop\ndata: {"type":"message_stop"}', |
| 35 | ]; |
| 36 | expect(await readUsage(stream(sse), "text/event-stream; charset=utf-8")).toEqual({ input: 1000, output: 300, cacheRead: 5000, cacheWrite: 0 }); |
| 37 | }); |
| 38 | it("reads an OpenAI Responses stream and a plain JSON answer", async () => { |
| 39 | const sse = ['data: {"type":"response.created","response":{}}\n\n', 'data: {"type":"response.completed","response":{"usage":{"input_tokens":80,"input_tokens_details":{"cached_tokens":0},"output_tokens":12}}}\n\n']; |
| 40 | expect(await readUsage(stream(sse), "text/event-stream")).toEqual({ input: 80, output: 12, cacheRead: 0, cacheWrite: 0 }); |
| 41 | expect(await readUsage(stream(['{"id":"x","usage":{"prompt_tokens":3,', '"completion_tokens":2}}']), "application/json")).toEqual({ input: 3, output: 2, cacheRead: 0, cacheWrite: 0 }); |
| 42 | expect(await readUsage(stream(["not json"]), "application/json")).toBeNull(); |
| 43 | expect(await readUsage(stream(['data: {"choices":[{"delta":{"content":"hi"}}]}\n\n', "data: [DONE]\n\n"]), "text/event-stream")).toBeNull(); |
| 44 | }); |
| 45 | it("drops an SSE line too long to be a usage event, and still reads the rest", async () => { |
| 46 | const huge = "data: " + "x".repeat(1100 * 1024); |
| 47 | expect(await readUsage(stream([huge.slice(0, 600_000), huge.slice(600_000), '\ndata: {"usage":{"input_tokens":4,"output_tokens":2}}\n\n']), "text/event-stream")).toEqual({ input: 4, output: 2, cacheRead: 0, cacheWrite: 0 }); |
| 48 | }); |
| 49 | }); |
| 50 | |
| 51 | describe("prepareRequest", () => { |
| 52 | const prep = (method: string, path: string, body: string) => prepareRequest(method, path, body, new TextEncoder().encode(body).length); |
| 53 | it("estimates the most an answer may cost: a token per 3 bytes plus the output it allows", () => { |
| 54 | const body = JSON.stringify({ model: "m", max_tokens: 1000, messages: [{ role: "user", content: "é".repeat(300) }] }); |
| 55 | expect(prep("POST", "/v1/messages", body)).toEqual({ body, model: "m", estimate: { input: Math.ceil((body.length + 300) / 3), output: 1000, cacheRead: 0, cacheWrite: 0 } }); |
| 56 | expect(prep("POST", "/v1/responses", '{"max_output_tokens":50}').estimate?.output).toBe(50); |
| 57 | expect(prep("POST", "/v1/chat/completions", '{"max_completion_tokens":70}').estimate?.output).toBe(70); |
| 58 | expect(prep("POST", "/v1/messages", '{"model":"m"}').estimate?.output).toBe(DEFAULT_MAX_OUTPUT); |
| 59 | expect(prep("POST", "/v1/messages", "not json").estimate).toEqual({ input: 3, output: DEFAULT_MAX_OUTPUT, cacheRead: 0, cacheWrite: 0 }); |
| 60 | // Each endpoint's own cap fields only; Chat Completions takes either. |
| 61 | expect(prep("POST", "/v1/messages", '{"max_output_tokens":50}').estimate?.output).toBe(DEFAULT_MAX_OUTPUT); |
| 62 | expect(prep("POST", "/v1/responses", '{"max_tokens":50}').estimate?.output).toBe(DEFAULT_MAX_OUTPUT); |
| 63 | expect(prep("POST", "/v1/chat/completions", '{"max_tokens":60}').estimate?.output).toBe(60); |
| 64 | expect(prep("POST", "/v1/completions", '{"prompt":"x","max_tokens":80}').estimate?.output).toBe(80); |
| 65 | // A compaction can't be capped: it books the most a model emits. |
| 66 | expect(prep("POST", "/v1/responses/compact", '{"input":[],"max_output_tokens":5}').estimate?.output).toBe(COMPACT_MAX_OUTPUT); |
| 67 | expect(prep("POST", "/v1/responses/compact", '{"input":[]}').reviewBody).toBeUndefined(); |
| 68 | // Requests that generate nothing are never estimated. |
| 69 | expect(prep("POST", "/v1/messages/count_tokens", "{}").estimate).toBeNull(); |
| 70 | expect(prep("GET", "/v1/models", "").estimate).toBeNull(); |
| 71 | }); |
| 72 | it("asks streamed Chat Completions for usage", () => { |
| 73 | const p = prep("POST", "/v1/chat/completions", '{"model":"m","stream":true,"stream_options":{"include_usage":false,"x":1}}'); |
| 74 | expect(JSON.parse(p.body).stream_options).toEqual({ include_usage: true, x: 1 }); |
| 75 | expect(JSON.parse(prep("POST", "/v1/chat/completions", '{"stream":true}').body).stream_options).toEqual({ include_usage: true }); |
| 76 | const unchanged = '{"stream":false}'; |
| 77 | expect(prep("POST", "/v1/chat/completions", unchanged).body).toBe(unchanged); |
| 78 | }); |
| 79 | it("names what a review may not send: background responses and context stored at the provider", () => { |
| 80 | expect(prep("POST", "/v1/responses", '{"background":true}').reviewRefuse).toMatch(/background/); |
| 81 | expect(prep("POST", "/v1/responses", '{"previous_response_id":"resp_1","input":"x"}').reviewRefuse).toMatch(/^previous_response_id/); |
| 82 | expect(prep("POST", "/v1/responses/compact", '{"conversation":"conv_1"}').reviewRefuse).toMatch(/^conversation/); |
| 83 | expect(prep("POST", "/v1/responses", '{"prompt":{"id":"pmpt_1"}}').reviewRefuse).toMatch(/^prompt/); |
| 84 | expect(prep("POST", "/v1/responses", '{"previous_response_id":null,"input":"x"}').reviewRefuse).toBeUndefined(); |
| 85 | expect(prep("POST", "/v1/messages", '{"prompt":"x","max_tokens":5}').reviewRefuse).toBeUndefined(); |
| 86 | expect(prep("POST", "/v1/chat/completions", '{"n":128,"max_completion_tokens":16000}').reviewRefuse).toMatch(/\(n\)/); |
| 87 | expect(prep("POST", "/v1/chat/completions", '{"n":"2"}').reviewRefuse).toMatch(/\(n\)/); |
| 88 | expect(prep("POST", "/v1/chat/completions", '{"n":1,"max_completion_tokens":5}').reviewRefuse).toBeUndefined(); |
| 89 | }); |
| 90 | it("books the largest of a request's caps, and a review sends every cap as the one booked", () => { |
| 91 | const p = prep("POST", "/v1/chat/completions", '{"max_completion_tokens":1,"max_tokens":100000}'); |
| 92 | expect(p.estimate?.output).toBe(100_000); |
| 93 | expect(JSON.parse(p.reviewBody!)).toEqual({ max_completion_tokens: 100_000, max_tokens: 100_000 }); |
| 94 | expect(prep("POST", "/v1/chat/completions", '{"max_completion_tokens":7,"max_tokens":7}').reviewBody).toBeUndefined(); |
| 95 | }); |
| 96 | it("pins a review's or a pinned run's model and routing", () => { |
| 97 | const req = (o: object) => prep("POST", "/v1/chat/completions", JSON.stringify({ max_tokens: 5, ...o })); |
| 98 | const pinned = { review: false, model: "anthropic/claude-sonnet-4.5" }; |
| 99 | expect(pinRefusal(req({ model: "anthropic/claude-sonnet-4.5" }), pinned)).toBeUndefined(); |
| 100 | expect(pinRefusal(req({ model: "openai/o1-pro" }), pinned)).toMatch(/runs on anthropic\/claude-sonnet-4\.5, not openai\/o1-pro/); |
| 101 | expect(pinRefusal(req({}), pinned)).toMatch(/runs on anthropic\/claude-sonnet-4\.5$/); |
| 102 | for (const k of ["models", "route", "provider", "plugins", "web_search_options", "speed", "best_of"]) { |
| 103 | expect(req({ model: "anthropic/claude-sonnet-4.5", [k]: k === "models" ? ["x/y"] : {} }).routing).toBe(k); |
| 104 | expect(pinRefusal(req({ model: "anthropic/claude-sonnet-4.5", [k]: {} }), pinned)).toMatch(new RegExp(`^${k} isn't available to this task`)); |
| 105 | expect(pinRefusal(req({ model: "m", [k]: {} }), { review: true, model: null })).toMatch(new RegExp(`^${k} isn't available to agent reviews`)); |
| 106 | } |
| 107 | // A run that isn't pinned (a custom command, or a direct provider) may choose: only reviews refuse routing. |
| 108 | expect(pinRefusal(req({ model: "any/model", plugins: [] }), { review: false, model: null })).toBeUndefined(); |
| 109 | expect(pinRefusal(req({ model: "any/model" }), { review: true, model: null })).toBeUndefined(); |
| 110 | // Requests that generate nothing aren't pinned. |
| 111 | expect(pinRefusal(prep("POST", "/v1/messages/count_tokens", '{"model":"other"}'), pinned)).toBeUndefined(); |
| 112 | }); |
| 113 | it("strips paid provider-side tools from a review's request, and keeps the rest", () => { |
| 114 | const fn = { name: "Read", input_schema: {} }; |
| 115 | const a = prep("POST", "/v1/messages", JSON.stringify({ max_tokens: 9, tools: [fn, { type: "web_search_20250305", name: "web_search" }, { type: "web_fetch_20250910", name: "web_fetch" }, { type: "code_execution_20250825", name: "code_execution" }, { type: "mcp_toolset", mcp_server_name: "x" }], mcp_servers: [{ url: "https://x" }], tool_choice: { type: "auto" } })); |
| 116 | expect(JSON.parse(a.reviewBody!)).toEqual({ max_tokens: 9, tools: [fn], tool_choice: { type: "auto" } }); |
| 117 | const local = { type: "local_shell" }; |
| 118 | const o = prep("POST", "/v1/responses", JSON.stringify({ max_output_tokens: 9, tools: [local, { type: "web_search" }, { type: "web_search_preview_2025_03_11" }, { type: "file_search" }, { type: "code_interpreter" }, { type: "image_generation" }, { type: "mcp", server_url: "https://x" }], tool_choice: { type: "web_search" } })); |
| 119 | expect(JSON.parse(o.reviewBody!)).toEqual({ max_output_tokens: 9, tools: [local] }); |
| 120 | // An allowlist: dated variants and tool types nobody has invented yet are dropped too; client-run tools stay. |
| 121 | const custom = { type: "custom", name: "grep" }; |
| 122 | const patch = { type: "apply_patch" }; |
| 123 | const f = { type: "function", name: "Read", parameters: {} }; |
| 124 | const d = prep("POST", "/v1/responses", JSON.stringify({ max_output_tokens: 9, tools: [f, custom, patch, { type: "web_search_2025_08_26" }, { type: "deep_research_2027_01_01" }, { type: 7 }], tool_choice: { type: "function", name: "Read" } })); |
| 125 | expect(JSON.parse(d.reviewBody!)).toEqual({ max_output_tokens: 9, tools: [f, custom, patch], tool_choice: { type: "function", name: "Read" } }); |
| 126 | const bash = { type: "bash_20250124", name: "bash" }; |
| 127 | const editor = { type: "text_editor_20250728", name: "str_replace_based_edit_tool" }; |
| 128 | const m = prep("POST", "/v1/messages", JSON.stringify({ max_tokens: 9, tools: [bash, editor, { type: "memory_20250818", name: "memory" }, { type: "web_search_20260101", name: "web_search" }, { type: "tool_search_tool_regex_20251119", name: "tool_search" }], tool_choice: { type: "tool", name: "web_search" } })); |
| 129 | expect(JSON.parse(m.reviewBody!)).toEqual({ max_tokens: 9, tools: [bash, editor, { type: "memory_20250818", name: "memory" }] }); |
| 130 | const cf = { type: "function", function: { name: "Read", parameters: {} } }; |
| 131 | const c = prep("POST", "/v1/chat/completions", JSON.stringify({ max_tokens: 9, tools: [cf, { type: "web_search" }], tool_choice: { type: "function", function: { name: "Read" } } })); |
| 132 | expect(JSON.parse(c.reviewBody!)).toEqual({ max_tokens: 9, tools: [cf], tool_choice: { type: "function", function: { name: "Read" } } }); |
| 133 | // With no tools left, neither tools nor a tool choice is sent. |
| 134 | expect(JSON.parse(prep("POST", "/v1/responses", JSON.stringify({ max_output_tokens: 9, tools: [{ type: "web_search" }], tool_choice: "required" })).reviewBody!)).toEqual({ max_output_tokens: 9 }); |
| 135 | // Nothing to strip and nothing to cap: the request goes as it is. |
| 136 | expect(prep("POST", "/v1/messages", JSON.stringify({ max_tokens: 9, tools: [fn] })).reviewBody).toBeUndefined(); |
| 137 | // Ordinary tasks keep their tools: only the review body changes. |
| 138 | expect(JSON.parse(a.body).tools).toHaveLength(5); |
| 139 | }); |
| 140 | it("gives reviews a body capped at the estimate's output when the request sets no cap", () => { |
| 141 | expect(JSON.parse(prep("POST", "/v1/responses", '{"input":"x"}').reviewBody!)).toEqual({ input: "x", max_output_tokens: DEFAULT_MAX_OUTPUT }); |
| 142 | expect(JSON.parse(prep("POST", "/v1/messages", '{"messages":[]}').reviewBody!)).toEqual({ messages: [], max_tokens: DEFAULT_MAX_OUTPUT }); |
| 143 | const streamed = JSON.parse(prep("POST", "/v1/chat/completions", '{"stream":true}').reviewBody!); |
| 144 | expect(streamed).toEqual({ stream: true, stream_options: { include_usage: true }, max_completion_tokens: DEFAULT_MAX_OUTPUT }); |
| 145 | // A request that caps its output, any way the endpoint takes it, is sent as it is. |
| 146 | expect(prep("POST", "/v1/chat/completions", '{"max_tokens":9}').reviewBody).toBeUndefined(); |
| 147 | expect(prep("POST", "/v1/responses", '{"max_output_tokens":9}').reviewBody).toBeUndefined(); |
| 148 | expect(prep("POST", "/v1/messages/count_tokens", "{}").reviewBody).toBeUndefined(); |
| 149 | expect(JSON.parse(prep("POST", "/v1/completions", '{"prompt":"x"}').reviewBody!)).toEqual({ prompt: "x", max_tokens: DEFAULT_MAX_OUTPUT }); |
| 150 | // OpenRouter's text completions take a prompt, which isn't stored context there. |
| 151 | expect(prep("POST", "/v1/completions", '{"prompt":"x","max_tokens":5}').reviewRefuse).toBeUndefined(); |
| 152 | }); |
| 153 | }); |
| 154 | |
| 155 | describe("readBody", () => { |
| 156 | const chunks = (parts: Uint8Array[]) => new ReadableStream<Uint8Array>({ start(c) { for (const p of parts) c.enqueue(p); c.close(); } }); |
| 157 | it("reads a body whole, and refuses one past the limit by its bytes", async () => { |
| 158 | const e = new TextEncoder(); |
| 159 | expect(new TextDecoder().decode((await readBody(chunks([e.encode("ab"), e.encode("çd")]), 5))!)).toBe("abçd"); |
| 160 | expect(await readBody(chunks([e.encode("ab"), e.encode("çde")]), 5)).toBeNull(); |
| 161 | expect(await readBody(null, 5)).toEqual(new Uint8Array()); |
| 162 | }); |
| 163 | }); |