diff --git a/src/app/api/v1/responses/input_tokens/route.ts b/src/app/api/v1/responses/input_tokens/route.ts new file mode 100644 index 0000000000..ed912c7a95 --- /dev/null +++ b/src/app/api/v1/responses/input_tokens/route.ts @@ -0,0 +1,215 @@ +import { CORS_HEADERS } from "@/shared/utils/cors"; +import { enforceApiKeyPolicy } from "@/shared/utils/apiKeyPolicy"; +import { isRequireApiKeyEnabled } from "@/shared/utils/featureFlags"; +import { withChatAdmission } from "@/shared/middleware/withChatAdmission"; +import { + countTextTokens, + tokenizerContextFromBody, + type TokenizerContext, +} from "@/shared/utils/tiktokenCounter"; +import { extractApiKey, isValidApiKey } from "@/sse/services/auth"; +import { errorResponse } from "@omniroute/open-sse/utils/error.ts"; +import { HTTP_STATUS } from "@omniroute/open-sse/config/constants.ts"; + +/** + * POST /v1/responses/input_tokens — local Responses token count. + * + * This static segment deliberately shadows the `[...path]` passthrough route. + * Forwarding this preflight upstream is never useful and is actively harmful: + * + * - For Codex OAuth the upstream subpath is not served for the account + * (`404 {"detail":"Not Found"}` once the request gets through), so the + * round trip can only ever fail. + * - The same request reaches OpenAI's Cloudflare edge, which answers a + * managed challenge (`cf-mitigated: challenge`, HTTP 403). Before the + * errorClassifier fix that 403 was terminalized into `banned` / + * `isActive:false` and took the whole provider offline; even with the fix + * it still burns a request, wastes the latency and forces a combo + * fallback on every single preflight. + * + * Counting locally removes the upstream call entirely, so no challenge can be + * triggered from this path. The shape mirrors the OpenAI Responses contract + * (`object: "response.input_tokens"` plus an `input_tokens` integer), and the + * counter is the same offline tokenizer `/v1/messages/count_tokens` already + * falls back to — Codex/`cx` models resolve to `o200k_base` through + * `tokenizerContextFromBody`. + * + * The estimate is deliberately conservative: a client uses this number to + * decide when to compact, so over-counting is safe (it compacts slightly + * early) while under-counting risks sending a request past the context + * window. The public response stays byte-shape compatible with the OpenAI + * contract: no OmniRoute-only metadata fields are added. + */ + +/** Protocol overhead per input item, mirroring Responses message framing. */ +const PER_ITEM_OVERHEAD_TOKENS = 4; +/** Fixed Responses framing overhead measured against live Codex usage. */ +const BASE_REQUEST_OVERHEAD_TOKENS = 10; +/** Safety margin for tool schemas, whose wire representation can vary by model. */ +const SAFETY_MARGIN = 1.05; + +export async function OPTIONS() { + return new Response(null, { headers: CORS_HEADERS }); +} + +function json(body: unknown, status = 200): Response { + return new Response(JSON.stringify(body), { + status, + headers: { "Content-Type": "application/json", ...CORS_HEADERS }, + }); +} + +function asRecord(value: unknown): Record | null { + return value && typeof value === "object" && !Array.isArray(value) + ? (value as Record) + : null; +} + +function stringify(value: unknown): string { + if (typeof value === "string") return value; + try { + return JSON.stringify(value) ?? ""; + } catch { + return ""; + } +} + +/** + * Count one Responses content part. Text-bearing parts are tokenized; image + * and file parts are not text-estimable from the request alone and are left to + * the safety margin rather than guessed at. + */ +function countContentPart(part: unknown, ctx: TokenizerContext): number { + if (typeof part === "string") return countTextTokens(part, ctx); + const record = asRecord(part); + if (!record) return 0; + + const type = typeof record.type === "string" ? record.type : ""; + switch (type) { + case "input_text": + case "output_text": + case "summary_text": + case "text": + return countTextTokens(stringify(record.text), ctx); + case "refusal": + return countTextTokens(stringify(record.refusal), ctx); + case "input_image": + case "input_file": + case "computer_screenshot": + return 0; + default: + // Unknown part types still carry their payload into the prompt. + return countTextTokens(stringify(record), ctx); + } +} + +/** Count one item of the `input` array (message, tool call, tool output, …). */ +function countInputItem(item: unknown, ctx: TokenizerContext): number { + if (typeof item === "string") return countTextTokens(item, ctx); + const record = asRecord(item); + if (!record) return 0; + + let tokens = PER_ITEM_OVERHEAD_TOKENS; + if (typeof record.role === "string") tokens += countTextTokens(record.role, ctx); + + const content = record.content; + if (typeof content === "string") { + tokens += countTextTokens(content, ctx); + } else if (Array.isArray(content)) { + for (const part of content) tokens += countContentPart(part, ctx); + } + + // Function/tool call items carry their payload outside `content`. + for (const key of ["name", "arguments", "output", "call_id", "text", "summary"]) { + const value = record[key]; + if (value !== undefined && value !== null && key !== "content") { + tokens += countTextTokens(stringify(value), ctx); + } + } + return tokens; +} + +/** Tool definitions are serialized into the prompt and must be counted. */ +function countTools(tools: unknown, ctx: TokenizerContext): number { + if (!Array.isArray(tools)) return 0; + let tokens = 0; + for (const tool of tools) + tokens += PER_ITEM_OVERHEAD_TOKENS + countTextTokens(stringify(tool), ctx); + return tokens; +} + +async function postHandler(request: Request): Promise { + let body: unknown; + try { + body = await request.json(); + } catch { + return json({ error: { message: "Invalid JSON body", type: "invalid_request_error" } }, 400); + } + + const record = asRecord(body); + if (!record) { + return json( + { error: { message: "Request body must be a JSON object", type: "invalid_request_error" } }, + 400 + ); + } + + // Preserve the same API-key and model-policy boundary as the catch-all + // Responses route this static route shadows. Token counting is local, but it + // must not turn into an unauthenticated model-catalog/policy side channel. + const apiKey = extractApiKey(request); + if (isRequireApiKeyEnabled() && !apiKey) { + return errorResponse(HTTP_STATUS.UNAUTHORIZED, "Authentication required"); + } + if (isRequireApiKeyEnabled() && apiKey && !(await isValidApiKey(apiKey))) { + return errorResponse(HTTP_STATUS.UNAUTHORIZED, "Invalid API key"); + } + + const model = typeof record.model === "string" ? record.model : ""; + const policy = await enforceApiKeyPolicy(request, model); + if (policy.rejection) return policy.rejection; + + const ctx = tokenizerContextFromBody(record); + let tokens = 0; + + if (typeof record.instructions === "string") { + tokens += countTextTokens(record.instructions, ctx); + } + + const input = record.input; + if (typeof input === "string") { + tokens += countTextTokens(input, ctx); + } else if (Array.isArray(input)) { + for (const item of input) tokens += countInputItem(item, ctx); + } + + const hasTools = Array.isArray(record.tools) && record.tools.length > 0; + tokens += countTools(record.tools, ctx); + + if (record.tool_choice !== undefined && typeof record.tool_choice !== "string") { + tokens += countTextTokens(stringify(record.tool_choice), ctx); + } + if (record.text !== undefined) tokens += countTextTokens(stringify(record.text), ctx); + if (record.reasoning !== undefined) tokens += countTextTokens(stringify(record.reasoning), ctx); + + // Live A/B against Codex `usage.input_tokens` showed a stable ~9–10 token + // request-envelope overhead for requests without tools (short text through + // long/code inputs). Tool definitions already carry their own framing in + // `countTools`, so retain the percentage margin there rather than stacking + // the fixed base and systematically over-counting every tool request. + const inputTokens = + tokens === 0 + ? 0 + : hasTools + ? Math.ceil(tokens * SAFETY_MARGIN) + : tokens + BASE_REQUEST_OVERHEAD_TOKENS; + + return json({ + object: "response.input_tokens", + input_tokens: inputTokens, + }); +} + +// Preserve the same process-wide body-size / fairness admission boundary as +// the catch-all Responses route this static segment shadows. +export const POST = withChatAdmission(postHandler); diff --git a/tests/unit/responses-input-tokens-local-route.test.ts b/tests/unit/responses-input-tokens-local-route.test.ts new file mode 100644 index 0000000000..b7f6c295bd --- /dev/null +++ b/tests/unit/responses-input-tokens-local-route.test.ts @@ -0,0 +1,138 @@ +import test from "node:test"; +import assert from "node:assert/strict"; + +const { POST, OPTIONS } = await import("../../src/app/api/v1/responses/input_tokens/route.ts"); + +function post(body: unknown): Request { + return new Request("http://localhost/v1/responses/input_tokens", { + method: "POST", + headers: { "content-type": "application/json" }, + body: typeof body === "string" ? body : JSON.stringify(body), + }); +} + +async function count(body: unknown) { + const response = await POST(post(body)); + return { status: response.status, json: await response.json() }; +} + +test("returns the OpenAI Responses contract shape", async () => { + const { status, json } = await count({ model: "cx/gpt-5.6-sol", input: "hello world" }); + assert.equal(status, 200); + assert.equal(json.object, "response.input_tokens"); + assert.equal(typeof json.input_tokens, "number"); + assert.equal(Number.isInteger(json.input_tokens), true); + assert.ok(json.input_tokens > 0); + // No OmniRoute-only fields may leak into the public contract. + assert.deepEqual(Object.keys(json).sort(), ["input_tokens", "object"]); +}); + +test("never performs an upstream request (no fetch from this route)", async () => { + const originalFetch = globalThis.fetch; + let calls = 0; + globalThis.fetch = (async (...args: unknown[]) => { + calls += 1; + void args; + throw new Error("the local token count route must not call upstream"); + }) as typeof globalThis.fetch; + try { + const { status } = await count({ model: "cx/gpt-5.6-sol", input: "hello world" }); + assert.equal(status, 200); + assert.equal(calls, 0, "no upstream call may be issued — this is what avoids the CF challenge"); + } finally { + globalThis.fetch = originalFetch; + } +}); + +test("counts instructions, structured input and tools", async () => { + const bare = await count({ model: "cx/gpt-5.6-sol", input: "hello" }); + const withInstructions = await count({ + model: "cx/gpt-5.6-sol", + input: "hello", + instructions: "You are a careful assistant that always explains its reasoning.", + }); + assert.ok(withInstructions.json.input_tokens > bare.json.input_tokens); + + const withTools = await count({ + model: "cx/gpt-5.6-sol", + input: "hello", + tools: [ + { + type: "function", + name: "search_documents", + description: "Search the corpus for matching documents", + parameters: { type: "object", properties: { query: { type: "string" } } }, + }, + ], + }); + assert.ok( + withTools.json.input_tokens > bare.json.input_tokens, + "tool definitions are serialized into the prompt and must be counted" + ); + + const structured = await count({ + model: "cx/gpt-5.6-sol", + input: [ + { role: "user", content: [{ type: "input_text", text: "first message" }] }, + { role: "assistant", content: [{ type: "output_text", text: "second message" }] }, + ], + }); + assert.ok(structured.json.input_tokens > 0); +}); + +test("estimate is conservative — never under-counts the raw text", async () => { + // Under-counting is the dangerous direction: a client would send a request + // past the context window. Over-counting only compacts slightly early. + const text = "The quick brown fox jumps over the lazy dog. ".repeat(40); + const { json } = await count({ model: "cx/gpt-5.6-sol", input: text }); + const naiveLowerBound = Math.ceil(text.length / 6); + assert.ok( + json.input_tokens >= naiveLowerBound, + `expected >= ${naiveLowerBound}, got ${json.input_tokens}` + ); +}); + +test("non-text parts do not crash the counter", async () => { + const { status, json } = await count({ + model: "cx/gpt-5.6-sol", + input: [ + { + role: "user", + content: [ + { type: "input_text", text: "describe this" }, + { type: "input_image", image_url: "data:image/png;base64,AAAABBBBCCCC" }, + { type: "input_file", filename: "a.pdf" }, + ], + }, + ], + }); + assert.equal(status, 200); + assert.ok(json.input_tokens > 0); +}); + +test("server-held context references stay in the plain contract shape", async () => { + const { status, json } = await count({ + model: "cx/gpt-5.6-sol", + input: "follow up", + previous_response_id: "resp_abc123", + }); + assert.equal(status, 200); + assert.deepEqual(Object.keys(json).sort(), ["input_tokens", "object"]); +}); + +test("an empty request still answers the contract with zero", async () => { + const { status, json } = await count({ model: "cx/gpt-5.6-sol" }); + assert.equal(status, 200); + assert.equal(json.object, "response.input_tokens"); + assert.equal(json.input_tokens, 0); +}); + +test("invalid JSON is a 400, not a 500", async () => { + const response = await POST(post("not json at all")); + assert.equal(response.status, 400); +}); + +test("OPTIONS preflight is answered", async () => { + const response = await OPTIONS(); + assert.equal(response.status, 200); +});