From e5bad22cd5fc6b85eb479c599e380896ac3d7e9e Mon Sep 17 00:00:00 2001 From: teamchong <25894545+teamchong@users.noreply.github.com> Date: Tue, 19 May 2026 21:19:22 -0400 Subject: [PATCH] fix(proxy): strip non-count_tokens fields before probing MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit count_tokens accepts only model/messages/system/tools/tool_choice/thinking/ mcp_servers. Forwarding the verbatim /v1/messages body included stream, max_tokens, temperature, metadata etc, which the endpoint rejected with 400. Probe returned null silently → measured_events stuck at 0. Strip to the accepted fields before posting. Now both probes succeed, the dashboard headline can use saved_pct_measured = (Σ baseline − Σ actual) / Σ baseline — exact, no regression, no priors. Diagnostic logging kept on the probe path: failures surface in stderr instead of disappearing. --- src/core/proxy.ts | 83 +++++++++++++++++++++++++++++++++++++++++++---- 1 file changed, 76 insertions(+), 7 deletions(-) diff --git a/src/core/proxy.ts b/src/core/proxy.ts index 4e3a394..323d697 100644 --- a/src/core/proxy.ts +++ b/src/core/proxy.ts @@ -261,6 +261,40 @@ function filterHeaders(src: Headers, strip: Set): Headers { return out; } +/** /v1/messages/count_tokens accepts a strict subset of /v1/messages params. + * Anything else (`stream`, `max_tokens`, `temperature`, `top_p`, `top_k`, + * `stop_sequences`, `metadata`, `service_tier`) makes it 400 with + * "Unknown parameter". This was the silent-null bug in the probe path — + * we forwarded the verbatim /v1/messages body and got 400s. + * + * Returns a fresh JSON Uint8Array with only the accepted fields, or null + * if the input can't be parsed (probe gets skipped, dashboard falls back + * to estimate). */ +const COUNT_TOKENS_FIELDS = new Set([ + 'model', + 'messages', + 'system', + 'tools', + 'tool_choice', + 'thinking', + 'mcp_servers', +]); + +function buildCountTokensBody(bytes: Uint8Array): Uint8Array | null { + try { + const obj = JSON.parse(new TextDecoder().decode(bytes)) as Record; + const out: Record = {}; + for (const k of Object.keys(obj)) { + if (COUNT_TOKENS_FIELDS.has(k)) out[k] = obj[k]; + } + // model is required by the endpoint; refuse to probe without it + if (typeof out.model !== 'string' || !Array.isArray(out.messages)) return null; + return new TextEncoder().encode(JSON.stringify(out)); + } catch { + return null; + } +} + /** Ask the upstream /v1/messages/count_tokens endpoint to tokenize a body * using the same auth + headers we'd send to /v1/messages. Returns the * exact input_tokens count or null on any failure (4xx, 5xx, network @@ -274,6 +308,7 @@ async function countTokensUpstream( upstream: string, body: BodyInit, headers: Headers, + label: string, ): Promise { try { const res = await fetch(upstream + '/v1/messages/count_tokens', { @@ -282,11 +317,31 @@ async function countTokensUpstream( body, ...(body instanceof ReadableStream ? { duplex: 'half' } : {}), } as RequestInit); - if (!res.ok) return null; + if (!res.ok) { + // Diagnostic: surface the actual failure mode. count_tokens has + // returned null silently in production; this exposes whether it's + // auth (401/403), bad request (400), rate limit (429), or upstream. + let snippet = ''; + try { + snippet = (await res.text()).slice(0, 300); + } catch { + /* body read failed — keep snippet empty */ + } + // eslint-disable-next-line no-console + console.warn(`[count_tokens:${label}] HTTP ${res.status}: ${snippet}`); + return null; + } const text = await res.text(); const json = JSON.parse(text) as { input_tokens?: unknown }; - return typeof json.input_tokens === 'number' ? json.input_tokens : null; - } catch { + const n = typeof json.input_tokens === 'number' ? json.input_tokens : null; + if (n === null) { + // eslint-disable-next-line no-console + console.warn(`[count_tokens:${label}] missing input_tokens in: ${text.slice(0, 300)}`); + } + return n; + } catch (e) { + // eslint-disable-next-line no-console + console.warn(`[count_tokens:${label}] threw: ${(e as Error).message}`); return null; } } @@ -397,16 +452,30 @@ export function createProxy(config: ProxyConfig = {}) { reqBodySha8 = await sha8Bytes(r.body); } - // Kick off the two count_tokens probes BEFORE forwarding /v1/messages - // so they run in parallel with the main request. Headers are filtered + // Kick off the count_tokens probes BEFORE forwarding /v1/messages so + // they run in parallel with the main request. Headers are filtered // exactly like the main request — same auth, same model, same // anthropic-version. Failure is silent: a null result drops the // measurement on this event without affecting the forwarded response. + // + // CRITICAL: /v1/messages/count_tokens accepts ONLY a subset of the + // /v1/messages fields. Forwarding the verbatim body causes a 400 + // ("Unknown parameter: stream" etc.) and the probe returns null. + // We strip to the accepted fields here. content-length is also + // already in STRIP_REQ_HEADERS — fetch recomputes it from the new + // body — so we don't need to touch headers further. const ctHeaders = filterHeaders(req.headers, STRIP_REQ_HEADERS); + ctHeaders.set('content-type', 'application/json'); if (config.apiKey) ctHeaders.set('x-api-key', config.apiKey); + const baselineBody = buildCountTokensBody(bodyIn); + const actualBody = buildCountTokensBody(r.body); measurePromise = Promise.all([ - countTokensUpstream(upstream, bodyIn, ctHeaders), - countTokensUpstream(upstream, r.body as unknown as BodyInit, ctHeaders), + baselineBody + ? countTokensUpstream(upstream, baselineBody as unknown as BodyInit, ctHeaders, 'baseline') + : Promise.resolve(null), + actualBody + ? countTokensUpstream(upstream, actualBody as unknown as BodyInit, ctHeaders, 'actual') + : Promise.resolve(null), ]).then(([baseline, actual]) => ({ baseline, actual })); } catch (e) { fire(502, undefined, `transform_error: ${(e as Error).message}`);