From 93be9d544c787aebe577f39e4132c04417099780 Mon Sep 17 00:00:00 2001 From: Anh Tran <161911430+anhtran-ai@users.noreply.github.com> Date: Fri, 18 Sep 2026 22:23:55 +0700 Subject: [PATCH] fix(responses): count input tokens locally for Codex OAuth (#13167) The ChatGPT subscription backend does not serve /backend-api/codex/responses/input_tokens for the affected account. Native requests to that path are intercepted by an OpenAI Cloudflare managed challenge, while the same path over the bundled Chrome transport returns 404 Not Found. Forwarding the client preflight can therefore never return a useful count and, before the companion classifier fix, permanently disabled the healthy Codex connection on the first challenge. Add a static /v1/responses/input_tokens route that shadows the generic Responses passthrough, uses the existing offline o200k_base token counter, and returns the standard response.input_tokens contract without issuing any upstream request. Count instructions, structured input, tool definitions and config; apply a conservative five-percent margin so the failure mode is earlier client compaction rather than a context-window overflow. Preserve the API-key and model-policy boundary from the catch-all Responses path. Tests pin the public schema, prove fetch is never called, cover text, instructions, structured input, tools, non-text parts, server-held context ids, invalid JSON, OPTIONS, and the conservative lower bound. Co-authored-by: anhth2 --- .../api/v1/responses/input_tokens/route.ts | 215 ++++++++++++++++++ ...responses-input-tokens-local-route.test.ts | 138 +++++++++++ 2 files changed, 353 insertions(+) create mode 100644 src/app/api/v1/responses/input_tokens/route.ts create mode 100644 tests/unit/responses-input-tokens-local-route.test.ts 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); +});