fix(api-proxy): track token usage for aborted streaming responses

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

@@ -126,13 +126,32 @@ function createChunkHandler(state, { requestId, provider }) {
 /**
  * Wire data/end event listeners onto proxyRes and optional decompressor.
  *
+ * Finalization runs on a clean 'end' or, as a fallback, on a premature
+ * 'aborted'/'close' so streaming responses whose sockets are torn down without
+ * a clean end (common with SSE clients) still record their accumulated usage.
+ *
  * @param {object} proxyRes - Upstream response stream
  * @param {object|null} decompressor - Zlib decompressor stream, or null
  * @param {object} state - Mutable tracking state
  * @param {(text: string) => void} onChunk - Decoded-chunk callback
  * @param {() => void} onFinalize - Finalization callback
  */
 function wireListeners(proxyRes, decompressor, state, onChunk, onFinalize) {
+  // Finalization must run at most once. Both the clean-completion path ('end')
+  // and the premature-close fallback ('aborted'/'close') route through here so
+  // usage is never written twice.
+  let finalized = false;
+  const finalizeOnce = () => {
+    if (finalized) return;
+    finalized = true;
+    onFinalize();
+  };
+
+  // Tracks whether the upstream stream completed cleanly. When true, the
+  // premature-close fallback is a no-op so we don't finalize with a
+  // still-flushing decompressor.
+  let endedCleanly = false;
+
   if (decompressor) {
     // Feed decompressed text to our parser
     decompressor.on('data', (decompressedChunk) => {
@@ -146,20 +165,44 @@ function wireListeners(proxyRes, decompressor, state, onChunk, onFinalize) {
     });
 
     proxyRes.on('end', () => {
+      endedCleanly = true;
       try { decompressor.end(); } catch { /* ignore */ }
     });
 
     // Finalize on decompressor end
-    decompressor.on('end', onFinalize);
+    decompressor.on('end', finalizeOnce);
   } else {
     // No compression — parse raw chunks directly
     proxyRes.on('data', (chunk) => {
       state.totalBytes += chunk.length;
       onChunk(chunk.toString('utf8'));
     });
 
-    proxyRes.on('end', onFinalize);
+    proxyRes.on('end', () => {
+      endedCleanly = true;
+      finalizeOnce();
+    });
   }
+
+  // Fallback for prematurely-closed connections. Streaming (SSE) clients such
+  // as Codex/OpenAI /responses frequently tear down the socket after receiving
+  // the final event, so proxyRes emits 'aborted'/'close' but never 'end'. The
+  // usage has already been accumulated per-chunk in state, so finalize it here
+  // instead of dropping the record. Skipped when the stream ended cleanly to
+  // avoid finalizing before a decompressor has flushed.
+  let prematureCloseHandled = false;
+  const onPrematureClose = () => {
+    if (endedCleanly || prematureCloseHandled) return;
+    prematureCloseHandled = true;
+    if (decompressor) {
+      decompressor.once('error', finalizeOnce);
+      try { decompressor.end(); } catch { finalizeOnce(); }
+      return;
+    }
+    finalizeOnce();
+  };
+  proxyRes.on('aborted', onPrematureClose);
+  proxyRes.on('close', onPrematureClose);
 }
 
 /**

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

@@ -213,6 +213,94 @@ describe('trackTokenUsage', () => {
     }, 10);
   });
 
+  test('records usage from a streaming response that closes without a clean end', (done) => {
+    // SSE clients (e.g. Codex/OpenAI /responses) often tear down the socket
+    // after the final event, so proxyRes emits 'close'/'aborted' but never
+    // 'end'. The accumulated usage must still be recorded.
+    const proxyRes = new EventEmitter();
+    proxyRes.headers = { 'content-type': 'text/event-stream' };
+    proxyRes.statusCode = 200;
+
+    const metricsRef = {
+      increment: jest.fn(),
+    };
+
+    trackTokenUsage(proxyRes, {
+      requestId: 'test-openai-responses-aborted',
+      provider: 'openai',
+      path: '/v1/responses',
+      startTime: Date.now(),
+      metrics: metricsRef,
+    });
+
+    const chunk = 'event: response.completed\ndata: ' + JSON.stringify({
+      type: 'response.completed',
+      response: {
+        model: 'gpt-5',
+        usage: { input_tokens: 800, output_tokens: 120, total_tokens: 920 },
+      },
+    }) + '\n\n';
+
+    proxyRes.emit('data', Buffer.from(chunk));
+    // No 'end' — the connection is torn down mid-stream.
+    proxyRes.emit('aborted');
+    proxyRes.emit('close');
+
+    setTimeout(() => {
+      expect(metricsRef.increment).toHaveBeenCalledWith(
+        'input_tokens_total',
+        { provider: 'openai' },
+        800,
+      );
+      expect(metricsRef.increment).toHaveBeenCalledWith(
+        'output_tokens_total',
+        { provider: 'openai' },
+        120,
+      );
+      done();
+    }, 10);
+  });
+
+  test('does not double-count usage when close follows a clean end', (done) => {
+    const proxyRes = new EventEmitter();
+    proxyRes.headers = { 'content-type': 'text/event-stream' };
+    proxyRes.statusCode = 200;
+
+    const metricsRef = {
+      increment: jest.fn(),
+    };
+
+    trackTokenUsage(proxyRes, {
+      requestId: 'test-openai-responses-end-then-close',
+      provider: 'openai',
+      path: '/v1/responses',
+      startTime: Date.now(),
+      metrics: metricsRef,
+    });
+
+    const chunk = 'event: response.completed\ndata: ' + JSON.stringify({
+      type: 'response.completed',
+      response: {
+        model: 'gpt-5',
+        usage: { input_tokens: 300, output_tokens: 50, total_tokens: 350 },
+      },
+    }) + '\n\n';
+
+    proxyRes.emit('data', Buffer.from(chunk));
+    proxyRes.emit('end');
+    // Node emits 'close' after 'end' on a normal completion — must be ignored.
+    proxyRes.emit('close');
+
+    setTimeout(() => {
+      const inputCalls = metricsRef.increment.mock.calls.filter(
+        (c) => c[0] === 'input_tokens_total',
+      );
+      expect(inputCalls).toHaveLength(1);
+      expect(inputCalls[0][2]).toBe(300);
+      done();
+    }, 10);
+  });
+
   test('warns when cache-read was observed in events but rolled-up value is zero', (done) => {
     const proxyRes = new EventEmitter();
     proxyRes.headers = { 'content-type': 'text/event-stream' };