fix: harden token tracker against unsupported response encodings and WS fragmentation

containers/api-proxy/proxy-utils.accept-encoding.test.js

@@ -0,0 +1,34 @@
+'use strict';
+
+const { sanitizeAcceptEncoding } = require('./proxy-utils');
+
+describe('sanitizeAcceptEncoding', () => {
+  it('passes through supported encodings unchanged', () => {
+    expect(sanitizeAcceptEncoding('gzip, deflate, br')).toBe('gzip, deflate, br');
+  });
+
+  it('strips zstd from the list', () => {
+    expect(sanitizeAcceptEncoding('gzip, br, zstd')).toBe('gzip, br');
+  });
+
+  it('strips zstd with quality value', () => {
+    expect(sanitizeAcceptEncoding('gzip;q=1.0, zstd;q=0.8, br;q=0.5')).toBe('gzip;q=1.0, br;q=0.5');
+  });
+
+  it('handles zstd-only Accept-Encoding by returning defaults', () => {
+    expect(sanitizeAcceptEncoding('zstd')).toBe('gzip, deflate, br');
+  });
+
+  it('preserves identity encoding', () => {
+    expect(sanitizeAcceptEncoding('identity, gzip, zstd')).toBe('identity, gzip');
+  });
+
+  it('returns defaults for empty/undefined input', () => {
+    expect(sanitizeAcceptEncoding('')).toBe('gzip, deflate, br');
+    expect(sanitizeAcceptEncoding(undefined)).toBe('gzip, deflate, br');
+  });
+
+  it('strips unknown encodings', () => {
+    expect(sanitizeAcceptEncoding('gzip, compress, deflate')).toBe('gzip, deflate');
+  });
+});

containers/api-proxy/proxy-utils.js

@@ -270,6 +270,32 @@ function makeUnconfiguredHealthResponse(service, error, status = 'not_configured
   return { statusCode: 503, body: { status, service, error } };
 }
 
+/**
+ * Encodings that the token tracker can decompress for response inspection.
+ * If the upstream responds with an encoding not in this set, the tracker
+ * receives raw compressed bytes and fails to extract usage data.
+ */
+const SUPPORTED_ENCODINGS = new Set(['gzip', 'deflate', 'br', 'identity']);
+
+/**
+ * Sanitize Accept-Encoding to only include encodings the token tracker can
+ * decompress. This prevents upstream APIs from responding with unsupported
+ * encodings (e.g. zstd) that the tracker cannot parse.
+ *
+ * @param {string|undefined} value - Original Accept-Encoding header value
+ * @returns {string} Filtered Accept-Encoding value
+ */
+function sanitizeAcceptEncoding(value) {
+  if (!value) return 'gzip, deflate, br';
+  const parts = value.split(',').map(p => p.trim()).filter(Boolean);
+  const supported = parts.filter(p => {
+    // Extract encoding name (strip quality value like ";q=0.5")
+    const encoding = p.split(';')[0].trim().toLowerCase();
+    return SUPPORTED_ENCODINGS.has(encoding);
+  });
+  return supported.length > 0 ? supported.join(', ') : 'gzip, deflate, br';
+}
+
 module.exports = {
   normalizeApiTarget,
   parseApiTargetAndBasePath,
@@ -280,4 +306,5 @@ module.exports = {
   composeBodyTransforms,
   makeProviderNotConfiguredResponse,
   makeUnconfiguredHealthResponse,
+  sanitizeAcceptEncoding,
 };

containers/api-proxy/request-headers.js

@@ -1,7 +1,7 @@
 'use strict';
 
 const { logRequest } = require('./logging');
-const { shouldStripHeader } = require('./proxy-utils');
+const { shouldStripHeader, sanitizeAcceptEncoding } = require('./proxy-utils');
 const { maybeStripLearnedHeaderValues } = require('./deprecated-header-tracker');
 
 /**
@@ -50,6 +50,13 @@ function buildRequestHeaders(body, inboundBytes, req, { injectHeaders, provider,
     delete headers['transfer-encoding'];
   }
 
+  // Restrict Accept-Encoding to encodings the token tracker can decompress.
+  // Without this, upstream APIs may respond with unsupported encodings (e.g.
+  // zstd) that the tracker cannot parse, causing silent token-usage data loss.
+  if (headers['accept-encoding']) {
+    headers['accept-encoding'] = sanitizeAcceptEncoding(headers['accept-encoding']);
+  }
+
   const injectedKey = Object.entries(injectHeaders).find(([k]) =>
     ['x-api-key', 'authorization', 'x-goog-api-key'].includes(k.toLowerCase())
   )?.[1];

containers/api-proxy/token-parsers.js

@@ -19,16 +19,28 @@ function isStreamingResponse(headers) {
 }
 
 /**
- * Check if a response is gzip or deflate compressed.
+ * Check if a response is compressed with a known encoding.
+ * Includes zstd detection — even though we cannot decompress it natively,
+ * recognizing it allows the tracker to skip parsing garbage bytes and log
+ * an actionable warning.
  */
 function isCompressedResponse(headers) {
   const ce = (headers['content-encoding'] || '').toLowerCase();
-  return ce === 'gzip' || ce === 'deflate' || ce === 'br';
+  return ce === 'gzip' || ce === 'deflate' || ce === 'br' || ce === 'zstd';
+}
+
+/**
+ * Check if a response uses an encoding we cannot decompress.
+ * Currently only zstd (Node.js has no built-in zstd support).
+ */
+function isUnsupportedEncoding(headers) {
+  const ce = (headers['content-encoding'] || '').toLowerCase();
+  return ce === 'zstd';
 }
 
 /**
  * Create a decompression transform stream based on content-encoding.
- * Returns null if the encoding is not supported.
+ * Returns null if the encoding is not supported (including zstd).
  */
 function createDecompressor(headers) {
   const ce = (headers['content-encoding'] || '').toLowerCase();
@@ -414,6 +426,7 @@ function normalizeUsage(usage) {
 module.exports = {
   isStreamingResponse,
   isCompressedResponse,
+  isUnsupportedEncoding,
   createDecompressor,
   extractReasoningTokens,
   extractCacheReadTokens,

containers/api-proxy/token-tracker-http.js

@@ -22,6 +22,7 @@ const { logRequest } = require('./logging');
 const {
   isStreamingResponse,
   isCompressedResponse,
+  isUnsupportedEncoding,
   createDecompressor,
   parseSseDataLines,
   extractUsageFromSseLine,
@@ -313,6 +314,18 @@ function finalizeHttpTracking(state, proxyRes, opts) {
 
   const normalized = normalizeUsage(usage);
   if (!normalized) {
+    // Log at info level so failed extraction is visible in CI without debug mode
+    logRequest('info', 'token_track_no_usage', {
+      request_id: requestId,
+      provider,
+      path: reqPath,
+      streaming,
+      total_bytes: state.totalBytes,
+      overflow: state.overflow,
+      content_encoding: contentEncoding,
+      content_type: state.contentType,
+      message: 'No token usage extracted from 2xx response',
+    });
     if (typeof onSpanEnd === 'function') onSpanEnd(proxyRes.statusCode);
     return;
   }
@@ -394,6 +407,22 @@ function trackTokenUsage(proxyRes, opts) {
 
   const state = initHttpState({ streaming, compressed, contentType, contentEncoding });
 
+  // If the response uses an encoding we cannot decompress (e.g. zstd), skip
+  // token tracking entirely and log a warning so the issue is visible in CI.
+  if (compressed && isUnsupportedEncoding(proxyRes.headers)) {
+    logRequest('warn', 'token_track_unsupported_encoding', {
+      request_id: requestId,
+      provider,
+      path: reqPath,
+      content_encoding: contentEncoding,
+      message: `Cannot decompress ${contentEncoding} responses; token usage will not be extracted. ` +
+        'The Accept-Encoding sanitizer should have prevented this — check if the client bypassed the proxy header rewrite.',
+    });
+    diag('HTTP_TRACK_UNSUPPORTED_ENCODING', { request_id: requestId, provider, path: reqPath, content_encoding: contentEncoding });
+    if (typeof opts.onSpanEnd === 'function') opts.onSpanEnd(proxyRes.statusCode);
+    return;
+  }
+
   // If the response is compressed, create a decompressor.
   // We feed raw chunks into it and listen on the decompressed output.
   // The raw proxyRes still flows to the client unchanged via pipe().

containers/api-proxy/token-tracker-ws.fragmentation.test.js

@@ -0,0 +1,122 @@
+'use strict';
+
+const { parseWebSocketFrames } = require('./token-tracker-ws');
+
+/**
+ * Build a WebSocket frame buffer.
+ * @param {string} payload - Text payload
+ * @param {object} opts
+ * @param {boolean} [opts.fin=true] - FIN bit
+ * @param {number} [opts.opcode=1] - Opcode (1=text, 0=continuation)
+ */
+function buildFrame(payload, { fin = true, opcode = 1 } = {}) {
+  const buf = Buffer.from(payload, 'utf8');
+  const len = buf.length;
+  let header;
+
+  const firstByte = (fin ? 0x80 : 0x00) | (opcode & 0x0F);
+
+  if (len < 126) {
+    header = Buffer.alloc(2);
+    header[0] = firstByte;
+    header[1] = len;
+  } else if (len < 65536) {
+    header = Buffer.alloc(4);
+    header[0] = firstByte;
+    header[1] = 126;
+    header.writeUInt16BE(len, 2);
+  } else {
+    header = Buffer.alloc(10);
+    header[0] = firstByte;
+    header[1] = 127;
+    header.writeBigUInt64BE(BigInt(len), 2);
+  }
+
+  return Buffer.concat([header, buf]);
+}
+
+describe('parseWebSocketFrames', () => {
+  describe('unfragmented messages', () => {
+    it('extracts a single text frame', () => {
+      const frame = buildFrame('{"type":"response.done"}');
+      const { messages, consumed } = parseWebSocketFrames(frame);
+      expect(messages).toEqual(['{"type":"response.done"}']);
+      expect(consumed).toBe(frame.length);
+    });
+
+    it('extracts multiple text frames from a single buffer', () => {
+      const frame1 = buildFrame('hello');
+      const frame2 = buildFrame('world');
+      const buf = Buffer.concat([frame1, frame2]);
+      const { messages } = parseWebSocketFrames(buf);
+      expect(messages).toEqual(['hello', 'world']);
+    });
+  });
+
+  describe('fragmented messages', () => {
+    it('reassembles a text message split across two frames', () => {
+      const fragments = [];
+      // First frame: FIN=0, opcode=text
+      const frame1 = buildFrame('{"type":"resp', { fin: false, opcode: 1 });
+      // Final continuation frame: FIN=1, opcode=continuation
+      const frame2 = buildFrame('onse.done"}', { fin: true, opcode: 0 });
+
+      const buf = Buffer.concat([frame1, frame2]);
+      const { messages } = parseWebSocketFrames(buf, fragments);
+      expect(messages).toEqual(['{"type":"response.done"}']);
+    });
+
+    it('reassembles a text message split across three frames', () => {
+      const fragments = [];
+      const frame1 = buildFrame('{"type"', { fin: false, opcode: 1 });
+      const frame2 = buildFrame(':"response', { fin: false, opcode: 0 });
+      const frame3 = buildFrame('.done"}', { fin: true, opcode: 0 });
+
+      const buf = Buffer.concat([frame1, frame2, frame3]);
+      const { messages } = parseWebSocketFrames(buf, fragments);
+      expect(messages).toEqual(['{"type":"response.done"}']);
+    });
+
+    it('handles fragmentation across multiple data events', () => {
+      const fragments = [];
+
+      // First data event: start of fragmented message
+      const frame1 = buildFrame('{"part":"one"', { fin: false, opcode: 1 });
+      const result1 = parseWebSocketFrames(frame1, fragments);
+      expect(result1.messages).toEqual([]);
+      expect(fragments.length).toBe(1);
+
+      // Second data event: final continuation frame
+      const frame2 = buildFrame(', "part":"two"}', { fin: true, opcode: 0 });
+      const result2 = parseWebSocketFrames(frame2, fragments);
+      expect(result2.messages).toEqual(['{"part":"one", "part":"two"}']);
+      expect(fragments.length).toBe(0);
+    });
+
+    it('handles mix of unfragmented and fragmented messages', () => {
+      const fragments = [];
+      const unfragmented = buildFrame('{"simple":true}');
+      const frag1 = buildFrame('{"frag":', { fin: false, opcode: 1 });
+      const frag2 = buildFrame('"yes"}', { fin: true, opcode: 0 });
+
+      const buf = Buffer.concat([unfragmented, frag1, frag2]);
+      const { messages } = parseWebSocketFrames(buf, fragments);
+      expect(messages).toEqual(['{"simple":true}', '{"frag":"yes"}']);
+    });
+  });
+
+  describe('backward compatibility', () => {
+    it('works without fragments parameter (unfragmented only)', () => {
+      const frame = buildFrame('hello');
+      const { messages } = parseWebSocketFrames(frame);
+      expect(messages).toEqual(['hello']);
+    });
+
+    it('silently skips fragmented messages without fragments parameter', () => {
+      // FIN=0 text frame without fragments array — should be skipped
+      const frame = buildFrame('partial', { fin: false, opcode: 1 });
+      const { messages } = parseWebSocketFrames(frame);
+      expect(messages).toEqual([]);
+    });
+  });
+});