diff --git a/containers/api-proxy/token-tracker.js b/containers/api-proxy/token-tracker.js index 9cdb862fe..16e58a217 100644 --- a/containers/api-proxy/token-tracker.js +++ b/containers/api-proxy/token-tracker.js @@ -84,11 +84,76 @@ function getLogStream() { } } +/** + * Validate a token usage record against the token-usage/v1 schema contract. + * + * Checks that all required fields are present and have the expected types. + * Logs a warning and returns false if the record is non-conformant; does + * not throw, so a bad record is dropped rather than crashing the proxy. + * + * @param {object} record - The record to validate + * @returns {boolean} true if the record is valid, false otherwise + */ +function validateTokenUsageRecord(record) { + if (!record || typeof record !== 'object') { + logRequest('warn', 'token_record_schema_violation', { + field: 'record', + expected: 'object', + actual: record === null ? 'null' : typeof record, + }); + return false; + } + + const required = [ + ['_schema', 'string'], + ['timestamp', 'string'], + ['request_id', 'string'], + ['provider', 'string'], + ['model', 'string'], + ['path', 'string'], + ['status', 'number'], + ['streaming', 'boolean'], + ['input_tokens', 'number'], + ['output_tokens', 'number'], + ['cache_read_tokens', 'number'], + ['cache_write_tokens', 'number'], + ['duration_ms', 'number'], + ]; + + for (const [field, expectedType] of required) { + // eslint-disable-next-line valid-typeof + if (typeof record[field] !== expectedType) { + logRequest('warn', 'token_record_schema_violation', { + request_id: record.request_id, + field, + expected: expectedType, + actual: typeof record[field], + }); + return false; + } + } + + if (record._schema !== 'token-usage/v1') { + logRequest('warn', 'token_record_schema_violation', { + request_id: record.request_id, + field: '_schema', + expected: 'token-usage/v1', + actual: record._schema, + }); + return false; + } + + return true; +} + /** * Write a token usage record to the JSONL log file. + * Validates the record against the token-usage/v1 schema before writing. * Handles backpressure by dropping writes when the stream buffer is full. */ function writeTokenUsage(record) { + if (!validateTokenUsageRecord(record)) return; + const stream = getLogStream(); if (stream && !stream.writableEnded) { const ok = stream.write(JSON.stringify(record) + '\n'); @@ -478,6 +543,7 @@ function trackTokenUsage(proxyRes, opts) { // Build log record const record = { + _schema: 'token-usage/v1', timestamp: new Date().toISOString(), request_id: requestId, provider, @@ -693,6 +759,7 @@ function trackWebSocketTokenUsage(upstreamSocket, opts) { } const record = { + _schema: 'token-usage/v1', timestamp: new Date().toISOString(), request_id: requestId, provider, @@ -759,6 +826,7 @@ module.exports = { normalizeUsage, isStreamingResponse, isCompressedResponse, + validateTokenUsageRecord, writeTokenUsage, TOKEN_LOG_FILE, }; diff --git a/containers/api-proxy/token-tracker.test.js b/containers/api-proxy/token-tracker.test.js index 737ab25c2..2bf042513 100644 --- a/containers/api-proxy/token-tracker.test.js +++ b/containers/api-proxy/token-tracker.test.js @@ -12,6 +12,9 @@ const { isCompressedResponse, trackTokenUsage, trackWebSocketTokenUsage, + validateTokenUsageRecord, + writeTokenUsage, + closeLogStream, } = require('./token-tracker'); const { EventEmitter } = require('events'); const os = require('os'); @@ -1044,3 +1047,208 @@ describe('trackWebSocketTokenUsage', () => { }, 10); }); }); + +// ── validateTokenUsageRecord ───────────────────────────────────────── + +describe('validateTokenUsageRecord', () => { + const validRecord = { + _schema: 'token-usage/v1', + timestamp: '2025-01-01T00:00:00.000Z', + request_id: 'req-123', + provider: 'anthropic', + model: 'claude-sonnet-4-20250514', + path: '/v1/messages', + status: 200, + streaming: false, + input_tokens: 100, + output_tokens: 50, + cache_read_tokens: 0, + cache_write_tokens: 0, + duration_ms: 1234, + }; + + test('accepts a valid record', () => { + expect(validateTokenUsageRecord(validRecord)).toBe(true); + }); + + test('accepts a record with optional response_bytes', () => { + expect(validateTokenUsageRecord({ ...validRecord, response_bytes: 512 })).toBe(true); + }); + + test('rejects a record with wrong _schema', () => { + expect(validateTokenUsageRecord({ ...validRecord, _schema: 'wrong/v99' })).toBe(false); + }); + + test('rejects a record missing _schema', () => { + const { _schema, ...noSchema } = validRecord; + expect(validateTokenUsageRecord(noSchema)).toBe(false); + }); + + test('rejects a record with non-string timestamp', () => { + expect(validateTokenUsageRecord({ ...validRecord, timestamp: 1234567890 })).toBe(false); + }); + + test('rejects a record with non-number input_tokens', () => { + expect(validateTokenUsageRecord({ ...validRecord, input_tokens: '100' })).toBe(false); + }); + + test('rejects a record with non-boolean streaming', () => { + expect(validateTokenUsageRecord({ ...validRecord, streaming: 'true' })).toBe(false); + }); + + test('rejects a record missing a required field', () => { + const { model, ...noModel } = validRecord; + expect(validateTokenUsageRecord(noModel)).toBe(false); + }); + + test('rejects null without throwing', () => { + expect(validateTokenUsageRecord(null)).toBe(false); + }); + + test('rejects undefined without throwing', () => { + expect(validateTokenUsageRecord(undefined)).toBe(false); + }); + + test('rejects a non-object primitive without throwing', () => { + expect(validateTokenUsageRecord('not-an-object')).toBe(false); + expect(validateTokenUsageRecord(42)).toBe(false); + }); +}); + +// ── JSONL records include _schema field ─────────────────────────────── + +/** + * Build a writable mock stream that captures all written chunks. + * The `written` getter parses the accumulated JSONL and returns records. + */ +function makeMockStream() { + const chunks = []; + const stream = { + writableEnded: false, + write: jest.fn((chunk) => { chunks.push(chunk); return true; }), + end: jest.fn((cb) => { stream.writableEnded = true; if (cb) cb(); }), + on: jest.fn(), + get writtenRecords() { + return chunks.map(c => JSON.parse(c.trim())); + }, + }; + return stream; +} + +describe('token-usage JSONL record schema field', () => { + let mockStream; + let mkdirSyncSpy; + let createWriteStreamSpy; + + beforeEach(async () => { + // Close any open log stream so the next getLogStream() call creates a fresh one. + await closeLogStream(); + + mockStream = makeMockStream(); + + // Redirect fs.mkdirSync and fs.createWriteStream so the module writes to our + // in-memory stream rather than the unwritable /var/log/api-proxy path. + mkdirSyncSpy = jest.spyOn(fs, 'mkdirSync').mockReturnValue(undefined); + createWriteStreamSpy = jest.spyOn(fs, 'createWriteStream').mockReturnValue(mockStream); + }); + + afterEach(async () => { + mkdirSyncSpy.mockRestore(); + createWriteStreamSpy.mockRestore(); + await closeLogStream(); + }); + + test('writeTokenUsage serializes _schema:"token-usage/v1" into the JSONL stream', () => { + const record = { + _schema: 'token-usage/v1', + timestamp: new Date().toISOString(), + request_id: 'direct-write-test', + provider: 'openai', + model: 'gpt-4o', + path: '/v1/chat/completions', + status: 200, + streaming: false, + input_tokens: 1, + output_tokens: 1, + cache_read_tokens: 0, + cache_write_tokens: 0, + duration_ms: 10, + }; + + writeTokenUsage(record); + + expect(mockStream.write).toHaveBeenCalledTimes(1); + const parsed = mockStream.writtenRecords[0]; + expect(parsed._schema).toBe('token-usage/v1'); + expect(parsed.request_id).toBe('direct-write-test'); + }); + + test('trackTokenUsage HTTP path writes _schema:"token-usage/v1" to the stream', (done) => { + const proxyRes = new EventEmitter(); + proxyRes.headers = { 'content-type': 'application/json' }; + proxyRes.statusCode = 200; + + trackTokenUsage(proxyRes, { + requestId: 'schema-field-http', + provider: 'openai', + path: '/v1/chat/completions', + startTime: Date.now(), + metrics: null, + }); + + proxyRes.emit('data', Buffer.from(JSON.stringify({ + model: 'gpt-4o', + usage: { prompt_tokens: 10, completion_tokens: 5, total_tokens: 15 }, + }))); + proxyRes.emit('end'); + + setTimeout(() => { + expect(mockStream.write).toHaveBeenCalledTimes(1); + const parsed = mockStream.writtenRecords[0]; + expect(parsed._schema).toBe('token-usage/v1'); + expect(parsed.request_id).toBe('schema-field-http'); + done(); + }, 20); + }); + + test('trackWebSocketTokenUsage path writes _schema:"token-usage/v1" to the stream', (done) => { + const socket = new EventEmitter(); + + function buildFrame(text) { + const payload = Buffer.from(text, 'utf8'); + const header = Buffer.alloc(2); + header[0] = 0x81; + header[1] = payload.length; + return Buffer.concat([header, payload]); + } + + const httpHeader = Buffer.from('HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\n\r\n'); + const frame1 = buildFrame(JSON.stringify({ + type: 'message_start', + message: { model: 'claude-sonnet-4-20250514', usage: { input_tokens: 20, output_tokens: 0 } }, + })); + const frame2 = buildFrame(JSON.stringify({ + type: 'message_delta', + usage: { output_tokens: 8 }, + })); + + trackWebSocketTokenUsage(socket, { + requestId: 'schema-field-ws', + provider: 'anthropic', + path: '/v1/messages', + startTime: Date.now(), + metrics: null, + }); + + socket.emit('data', Buffer.concat([httpHeader, frame1, frame2])); + socket.emit('close'); + + setTimeout(() => { + expect(mockStream.write).toHaveBeenCalledTimes(1); + const parsed = mockStream.writtenRecords[0]; + expect(parsed._schema).toBe('token-usage/v1'); + expect(parsed.request_id).toBe('schema-field-ws'); + done(); + }, 20); + }); +}); diff --git a/samples/audit/audit.jsonl b/samples/audit/audit.jsonl index b3adce217..821b1382a 100644 --- a/samples/audit/audit.jsonl +++ b/samples/audit/audit.jsonl @@ -1,3 +1,3 @@ -{"ts":1774290908.910,"client":"172.30.0.20","host":"api.github.com:443","dest":"140.82.116.5:443","method":"CONNECT","status":200,"decision":"TCP_TUNNEL","url":"api.github.com:443"} -{"ts":1774290909.180,"client":"172.30.0.20","host":"api.github.com:443","dest":"140.82.116.5:443","method":"CONNECT","status":200,"decision":"TCP_TUNNEL","url":"api.github.com:443"} -{"ts":1774290909.186,"client":"172.30.0.20","host":"evil.example.com:443","dest":"-:-","method":"CONNECT","status":403,"decision":"TCP_DENIED","url":"evil.example.com:443"} +{"_schema":"audit/v1","ts":1774290908.910,"client":"172.30.0.20","host":"api.github.com:443","dest":"140.82.116.5:443","method":"CONNECT","status":200,"decision":"TCP_TUNNEL","url":"api.github.com:443"} +{"_schema":"audit/v1","ts":1774290909.180,"client":"172.30.0.20","host":"api.github.com:443","dest":"140.82.116.5:443","method":"CONNECT","status":200,"decision":"TCP_TUNNEL","url":"api.github.com:443"} +{"_schema":"audit/v1","ts":1774290909.186,"client":"172.30.0.20","host":"evil.example.com:443","dest":"-:-","method":"CONNECT","status":403,"decision":"TCP_DENIED","url":"evil.example.com:443"} diff --git a/schemas/README.md b/schemas/README.md new file mode 100644 index 000000000..d9ccd303e --- /dev/null +++ b/schemas/README.md @@ -0,0 +1,40 @@ +# AWF JSONL Schemas + +This directory contains versioned [JSON Schema](https://json-schema.org/) files for the JSONL artifact files emitted by AWF at runtime. + +## Files + +| Schema file | JSONL file | Writer | +|---|---|---| +| [`token-usage.v1.schema.json`](token-usage.v1.schema.json) | `token-usage.jsonl` | `containers/api-proxy/token-tracker.js` | +| [`audit.v1.schema.json`](audit.v1.schema.json) | `audit.jsonl` | Squid proxy (`src/squid-config.ts`) | + +## Schema versioning policy + +- **Additive changes** (new optional fields) → update the existing `v1` schema, no version bump required. +- **Breaking changes** (field removal, rename, type change, new required field) → create a new `v2` schema file and bump the `_schema` value in the writer. + +## Record identification + +Every JSONL record includes a `_schema` field that identifies the schema name and version: + +```json +{ "_schema": "token-usage/v1", "timestamp": "2025-01-01T00:00:00.000Z", ... } +{ "_schema": "audit/v1", "ts": 1761074374.646, ... } +``` + +Consumers should check `_schema` before parsing fields so they can handle future versions gracefully. + +## Validation + +You can validate a JSONL file against its schema using any JSON Schema validator. Example using [`ajv-cli`](https://github.com/ajv-validator/ajv-cli): + +```bash +# Install validator +npm install -g ajv-cli + +# Validate all records in audit.jsonl +while IFS= read -r line; do + echo "$line" | ajv validate -s schemas/audit.v1.schema.json -d /dev/stdin +done < /path/to/audit.jsonl +``` diff --git a/schemas/audit.v1.schema.json b/schemas/audit.v1.schema.json new file mode 100644 index 000000000..a803848ba --- /dev/null +++ b/schemas/audit.v1.schema.json @@ -0,0 +1,60 @@ +{ + "$schema": "https://json-schema.org/draft/2020-12/schema", + "$id": "https://github.com/github/gh-aw-firewall/schemas/audit.v1.schema.json", + "title": "AWF Audit Log Record", + "description": "A single L7 HTTP/HTTPS traffic decision record emitted to audit.jsonl by the AWF Squid proxy.", + "version": "1", + "type": "object", + "required": [ + "_schema", + "ts", + "client", + "host", + "dest", + "method", + "status", + "decision", + "url" + ], + "additionalProperties": true, + "properties": { + "_schema": { + "type": "string", + "const": "audit/v1", + "description": "Schema identifier and version for this record." + }, + "ts": { + "type": "number", + "description": "Unix timestamp with millisecond precision (e.g. 1761074374.646)." + }, + "client": { + "type": "string", + "description": "Client IP address that originated the request (e.g. '172.30.0.20')." + }, + "host": { + "type": "string", + "description": "HTTP Host header value or CONNECT target (e.g. 'api.github.com:443')." + }, + "dest": { + "type": "string", + "description": "Destination IP address and port resolved by Squid (e.g. '140.82.114.22:443'). '-:-' when the connection was denied before upstream resolution." + }, + "method": { + "type": "string", + "description": "HTTP method used by the client (e.g. 'CONNECT', 'GET', 'POST')." + }, + "status": { + "type": "integer", + "minimum": 0, + "description": "HTTP response status code. 200 = allowed, 403 = denied." + }, + "decision": { + "type": "string", + "description": "Squid cache/hierarchy decision code. 'TCP_TUNNEL' = allowed HTTPS CONNECT, 'TCP_DENIED' = blocked, 'TCP_MISS' = allowed cache miss." + }, + "url": { + "type": "string", + "description": "Request URL (for CONNECT: the domain:port tunnel target; for plain HTTP: the full URL)." + } + } +} diff --git a/schemas/token-usage.v1.schema.json b/schemas/token-usage.v1.schema.json new file mode 100644 index 000000000..ebfc02211 --- /dev/null +++ b/schemas/token-usage.v1.schema.json @@ -0,0 +1,93 @@ +{ + "$schema": "https://json-schema.org/draft/2020-12/schema", + "$id": "https://github.com/github/gh-aw-firewall/schemas/token-usage.v1.schema.json", + "title": "AWF Token Usage Record", + "description": "A single per-API-call token usage record emitted to token-usage.jsonl by the AWF api-proxy sidecar.", + "version": "1", + "type": "object", + "required": [ + "_schema", + "timestamp", + "request_id", + "provider", + "model", + "path", + "status", + "streaming", + "input_tokens", + "output_tokens", + "cache_read_tokens", + "cache_write_tokens", + "duration_ms" + ], + "additionalProperties": true, + "properties": { + "_schema": { + "type": "string", + "const": "token-usage/v1", + "description": "Schema identifier and version for this record." + }, + "timestamp": { + "type": "string", + "format": "date-time", + "description": "ISO 8601 timestamp of when the API call completed." + }, + "request_id": { + "type": "string", + "description": "Unique identifier for this API request (UUID or provider-assigned ID)." + }, + "provider": { + "type": "string", + "enum": ["anthropic", "openai", "copilot", "gemini", "opencode", "unknown"], + "description": "LLM provider that handled this request." + }, + "model": { + "type": "string", + "description": "Model name returned by the provider (e.g. 'gpt-4o', 'claude-sonnet-4-20250514')." + }, + "path": { + "type": "string", + "description": "API endpoint path for this request (e.g. '/v1/messages', '/v1/chat/completions')." + }, + "status": { + "type": "integer", + "minimum": 100, + "maximum": 599, + "description": "HTTP response status code returned by the provider." + }, + "streaming": { + "type": "boolean", + "description": "Whether the response was a streaming (SSE) response." + }, + "input_tokens": { + "type": "integer", + "minimum": 0, + "description": "Number of input/prompt tokens consumed." + }, + "output_tokens": { + "type": "integer", + "minimum": 0, + "description": "Number of output/completion tokens generated." + }, + "cache_read_tokens": { + "type": "integer", + "minimum": 0, + "description": "Number of tokens read from the provider's prompt cache (Anthropic cache_read_input_tokens or OpenAI cached_tokens)." + }, + "cache_write_tokens": { + "type": "integer", + "minimum": 0, + "description": "Number of tokens written to the provider's prompt cache (Anthropic cache_creation_input_tokens)." + }, + "duration_ms": { + "type": "number", + "minimum": 0, + "description": "End-to-end request duration in milliseconds." + }, + "response_bytes": { + "type": "integer", + "minimum": 0, + "description": "Total number of bytes in the response body (optional, omitted for WebSocket upgrades)." + } + } +} diff --git a/src/logs/log-parser.test.ts b/src/logs/log-parser.test.ts index e497aa8a8..3f5e126d7 100644 --- a/src/logs/log-parser.test.ts +++ b/src/logs/log-parser.test.ts @@ -289,5 +289,19 @@ describe('log-parser', () => { expect(parseAuditJsonlLine('not json')).toBeNull(); expect(parseAuditJsonlLine('{broken')).toBeNull(); }); + + it('should parse records that include the _schema field', () => { + const line = '{"_schema":"audit/v1","ts":1774290908.910,"client":"172.30.0.20","host":"api.github.com:443","dest":"140.82.116.5:443","method":"CONNECT","status":200,"decision":"TCP_TUNNEL","url":"api.github.com:443"}'; + const entry = parseAuditJsonlLine(line); + + expect(entry).not.toBeNull(); + expect(entry!.timestamp).toBeCloseTo(1774290908.910); + expect(entry!.clientIp).toBe('172.30.0.20'); + expect(entry!.method).toBe('CONNECT'); + expect(entry!.statusCode).toBe(200); + expect(entry!.decision).toBe('TCP_TUNNEL'); + expect(entry!.domain).toBe('api.github.com'); + expect(entry!.isAllowed).toBe(true); + }); }); }); diff --git a/src/squid-config.test.ts b/src/squid-config.test.ts index 3a3d3f7c9..b800274f8 100644 --- a/src/squid-config.test.ts +++ b/src/squid-config.test.ts @@ -569,6 +569,46 @@ describe('generateSquidConfig', () => { expect(aclIndex).toBeGreaterThan(-1); expect(accessLogIndex).toBeGreaterThan(aclIndex); }); + + it('should include JSONL audit log format (audit_jsonl)', () => { + const config: SquidConfig = { + domains: ['example.com'], + port: defaultPort, + }; + const result = generateSquidConfig(config); + expect(result).toContain('logformat audit_jsonl'); + expect(result).toContain('access_log /var/log/squid/audit.jsonl audit_jsonl'); + }); + + it('audit_jsonl logformat should include _schema:"audit/v1" field', () => { + const config: SquidConfig = { + domains: ['example.com'], + port: defaultPort, + }; + const result = generateSquidConfig(config); + // The audit_jsonl logformat line must embed the schema identifier so that + // every emitted record carries the version tag. + expect(result).toMatch(/logformat audit_jsonl \{.*"_schema":"audit\/v1"/); + }); + + it('audit_jsonl logformat should include all required fields', () => { + const config: SquidConfig = { + domains: ['example.com'], + port: defaultPort, + }; + const result = generateSquidConfig(config); + // Required fields per audit.v1.schema.json + const auditLine = result.split('\n').find(l => l.startsWith('logformat audit_jsonl')); + expect(auditLine).toBeDefined(); + expect(auditLine).toContain('"ts":'); + expect(auditLine).toContain('"client":'); + expect(auditLine).toContain('"host":'); + expect(auditLine).toContain('"dest":'); + expect(auditLine).toContain('"method":'); + expect(auditLine).toContain('"status":'); + expect(auditLine).toContain('"decision":'); + expect(auditLine).toContain('"url":'); + }); }); describe('Streaming/Long-lived Connection Support', () => { diff --git a/src/squid-config.ts b/src/squid-config.ts index 1452866f8..f9ed3adf2 100644 --- a/src/squid-config.ts +++ b/src/squid-config.ts @@ -600,7 +600,7 @@ logformat firewall_detailed %ts.%03tu %>a:%>p %{Host}>h %Hs %Ss # Structured JSONL audit log for machine-readable analysis # Note: Squid logformat does not JSON-escape strings, so fields like User-Agent # could break JSON parsing. We omit User-Agent to reduce breakage risk. -logformat audit_jsonl {"ts":%ts.%03tu,"client":"%>a","host":"%{Host}>h","dest":"%Hs,"decision":"%Ss","url":"%ru"} +logformat audit_jsonl {"_schema":"audit/v1","ts":%ts.%03tu,"client":"%>a","host":"%{Host}>h","dest":"%Hs,"decision":"%Ss","url":"%ru"} # Access log and cache configuration # Don't log healthcheck probes from localhost (using ACL filter on access_log)