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 identity', () => {
+    expect(sanitizeAcceptEncoding('zstd')).toBe('identity');
+  });
+
+  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(', ') : 'identity';
+}
+
 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-persistence.js

@@ -16,6 +16,7 @@ const { logRequest } = require('./logging');
 const TOKEN_LOG_DIR = process.env.AWF_TOKEN_LOG_DIR || '/var/log/api-proxy';
 const TOKEN_LOG_FILE = path.join(TOKEN_LOG_DIR, 'token-usage.jsonl');
 const DIAG_LOG_FILE = path.join(TOKEN_LOG_DIR, 'token-diag.jsonl');
+const AUDIT_LOG_FILE = path.join(TOKEN_LOG_DIR, 'token-tracker-audit.jsonl');
 const DIAG_ENABLED = process.env.AWF_DEBUG_TOKENS === '1';
 
 // AWF version used to identify schema version in JSONL records.
@@ -33,6 +34,7 @@ const TOKEN_DIAG_SCHEMA = `token-diag/v${AWF_VERSION || '0.0.0-dev'}`;
 
 let logStream = null;
 let diagStream = null;
+let auditStream = null;
 
 /**
  * Write a diagnostic line to the diagnostics log file.
@@ -72,6 +74,25 @@ function diag(msg, data) {
   } catch { /* best-effort */ }
 }
 
+/**
+ * Always-on audit trail for token tracking lifecycle.
+ *
+ * Writes minimal structured events to token-tracker-audit.jsonl so that CI
+ * failures can be diagnosed without AWF_DEBUG_TOKENS=1.  Each tracked
+ * request emits exactly two lines: TRACK_START and TRACK_END (with result).
+ */
+function auditTrack(event, data) {
+  try {
+    if (!auditStream) {
+      fs.mkdirSync(TOKEN_LOG_DIR, { recursive: true });
+      auditStream = fs.createWriteStream(AUDIT_LOG_FILE, { flags: 'a', mode: 0o644 });
+      auditStream.on('error', () => { auditStream = null; });
+    }
+    const line = { ts: Date.now(), event, ...data };
+    auditStream.write(JSON.stringify(line) + '\n');
+  } catch { /* best-effort */ }
+}
+
 /**
  * Get or create the JSONL append stream for token usage logs.
  * Uses a lazy singleton — created on first write.
@@ -258,6 +279,10 @@ function closeLogStream() {
         pending++;
         diagStream.end(() => { diagStream = null; pending--; check(); });
       }
+      if (auditStream) {
+        pending++;
+        auditStream.end(() => { auditStream = null; pending--; check(); });
+      }
       if (pending === 0) resolve();
     }),
     closeBlockedRequestDiagStream(),
@@ -269,6 +294,7 @@ module.exports = {
   TOKEN_USAGE_SCHEMA,
   TOKEN_DIAG_SCHEMA,
   diag,
+  auditTrack,
   buildTokenDiagRecord,
   buildTokenUsageRecord,
   incrementTokenMetrics,

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

@@ -22,6 +22,7 @@ const { logRequest } = require('./logging');
 const {
   isStreamingResponse,
   isCompressedResponse,
+  isUnsupportedEncoding,
   createDecompressor,
   parseSseDataLines,
   extractUsageFromSseLine,
@@ -33,6 +34,7 @@ const {
   buildTokenUsageRecord,
   incrementTokenMetrics,
   diag,
+  auditTrack,
 } = require('./token-persistence');
 const { warnCacheReadRollupMismatch, mergeBudgetFields } = require('./token-tracker-shared');
 
@@ -285,6 +287,7 @@ function finalizeHttpTracking(state, proxyRes, opts) {
 
   // Only process successful responses (2xx)
   if (proxyRes.statusCode < 200 || proxyRes.statusCode >= 300) {
+    auditTrack('TRACK_END', { rid: requestId, result: 'skip_status', status: proxyRes.statusCode });
     logRequest('debug', 'token_track_skip_status', {
       request_id: requestId,
       provider,
@@ -313,9 +316,25 @@ function finalizeHttpTracking(state, proxyRes, opts) {
 
   const normalized = normalizeUsage(usage);
   if (!normalized) {
+    auditTrack('TRACK_END', { rid: requestId, result: 'no_usage', streaming, bytes: state.totalBytes, overflow: state.overflow, ct: state.contentType, ce: contentEncoding });
+    // 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;
   }
+
+  auditTrack('TRACK_END', { rid: requestId, result: 'ok', streaming, model: model || requestModel, input: normalized.input_tokens, output: normalized.output_tokens, cache_read: normalized.cache_read_tokens });
+
   if (state.observedCacheReadTokens > 0 && normalized.cache_read_tokens === 0) {
     warnCacheReadRollupMismatch({ logRequest, diag, requestId, provider, model, observedCacheReadTokens: state.observedCacheReadTokens, normalizedCacheReadTokens: normalized.cache_read_tokens, streaming });
   }
@@ -381,6 +400,8 @@ function trackTokenUsage(proxyRes, opts) {
   const contentEncoding = proxyRes.headers['content-encoding'] || '(none)';
   const compressed = isCompressedResponse(proxyRes.headers);
 
+  auditTrack('TRACK_START', { rid: requestId, provider, path: reqPath, streaming, ct: contentType, ce: contentEncoding, status: proxyRes.statusCode });
+
   logRequest('debug', 'token_track_start', {
     request_id: requestId,
     provider,
@@ -394,6 +415,23 @@ 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)) {
+    auditTrack('TRACK_SKIP_ENCODING', { rid: requestId, ce: contentEncoding });
+    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().