fix(proxy): strip non-count_tokens fields before probing

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.
This commit is contained in:
teamchong
2026-05-19 21:19:22 -04:00
parent 00b7c80e84
commit e5bad22cd5
+76 -7
View File
@@ -261,6 +261,40 @@ function filterHeaders(src: Headers, strip: Set<string>): 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<string, unknown>;
const out: Record<string, unknown> = {};
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<number | null> {
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}`);