Skip to content

Commit b4acb6b

Browse files
lpcoxCopilot
andauthored
feat: add token usage tracking to api-proxy sidecar (#1539)
* feat: add token usage tracking to api-proxy sidecar Intercept LLM API responses in the api-proxy to extract and log token usage data. Supports both streaming (SSE) and non-streaming JSON responses from OpenAI, Anthropic, and Copilot providers. - Add token-tracker.js module with SSE and JSON usage extraction - Integrate trackTokenUsage into server.js proxy response pipeline - Write JSONL logs to /var/log/api-proxy/token-usage.jsonl - Update input_tokens_total and output_tokens_total metrics - Add 32 unit tests covering all extraction and normalization paths Closes #1536 Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * fix: address PR review feedback for token tracker - Add token-tracker.js to Dockerfile COPY step - Add try/catch require guard for graceful fallback - Fix header comment to match actual architecture - Guard empty usage objects from producing 0-token log entries - Handle write backpressure in JSONL log writer - Make closeLogStream() async to flush before process.exit - Remove unused method/targetHost opts; drop hardcoded request_bytes - Redirect test log output to temp dir to avoid /var/log noise Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
1 parent e31b23e commit b4acb6b

4 files changed

Lines changed: 884 additions & 3 deletions

File tree

containers/api-proxy/Dockerfile

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@ COPY package*.json ./
1515
RUN npm ci --omit=dev
1616

1717
# Copy application files
18-
COPY server.js logging.js metrics.js rate-limiter.js ./
18+
COPY server.js logging.js metrics.js rate-limiter.js token-tracker.js ./
1919

2020
# Create non-root user
2121
RUN addgroup -S apiproxy && adduser -S apiproxy -G apiproxy

containers/api-proxy/server.js

Lines changed: 25 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,18 @@ const { HttpsProxyAgent } = require('https-proxy-agent');
1818
const { generateRequestId, sanitizeForLog, logRequest } = require('./logging');
1919
const metrics = require('./metrics');
2020
const rateLimiter = require('./rate-limiter');
21+
let trackTokenUsage;
22+
let closeLogStream;
23+
try {
24+
({ trackTokenUsage, closeLogStream } = require('./token-tracker'));
25+
} catch (err) {
26+
if (err && err.code === 'MODULE_NOT_FOUND') {
27+
trackTokenUsage = () => {};
28+
closeLogStream = () => {};
29+
} else {
30+
throw err;
31+
}
32+
}
2133

2234
// Create rate limiter from environment variables
2335
const limiter = rateLimiter.create();
@@ -423,6 +435,15 @@ function proxyRequest(req, res, targetHost, injectHeaders, provider, basePath =
423435

424436
res.writeHead(proxyRes.statusCode, resHeaders);
425437
proxyRes.pipe(res);
438+
439+
// Attach token usage tracking (non-blocking, listens on same data/end events)
440+
trackTokenUsage(proxyRes, {
441+
requestId,
442+
provider,
443+
path: sanitizeForLog(req.url),
444+
startTime,
445+
metrics,
446+
});
426447
});
427448

428449
proxyReq.on('error', (err) => {
@@ -851,13 +872,15 @@ if (require.main === module) {
851872
}
852873

853874
// Graceful shutdown
854-
process.on('SIGTERM', () => {
875+
process.on('SIGTERM', async () => {
855876
logRequest('info', 'shutdown', { message: 'Received SIGTERM, shutting down gracefully' });
877+
await closeLogStream();
856878
process.exit(0);
857879
});
858880

859-
process.on('SIGINT', () => {
881+
process.on('SIGINT', async () => {
860882
logRequest('info', 'shutdown', { message: 'Received SIGINT, shutting down gracefully' });
883+
await closeLogStream();
861884
process.exit(0);
862885
});
863886
}

0 commit comments

Comments
 (0)