mirror of
https://github.com/teamchong/pxpipe.git
synced 2026-07-22 02:02:51 +02:00
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:
+76
-7
@@ -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}`);
|
||||
|
||||
Reference in New Issue
Block a user