diff --git a/packages/broker-proxy/openai-chat-compat.js b/packages/broker-proxy/openai-chat-compat.js new file mode 100644 index 0000000..b5eedcc --- /dev/null +++ b/packages/broker-proxy/openai-chat-compat.js @@ -0,0 +1,171 @@ +'use strict'; + +const crypto = require('crypto'); + +const CURSOR_MODEL_ALIASES = Object.freeze({ + 'anthropic:sonnet': 'claude-sonnet-4-6', + 'anthropic:opus': 'claude-opus-4-6', + 'anthropic:haiku': 'claude-haiku-4-5-20251001' +}); + +function cursorModel(model) { + const requested = String(model || 'anthropic:sonnet'); + if (CURSOR_MODEL_ALIASES[requested]) return CURSOR_MODEL_ALIASES[requested]; + if (requested.startsWith('anthropic:claude-')) return requested.slice('anthropic:'.length); + return requested; +} + +function contentBlocks(content) { + if (content == null) return []; + if (typeof content === 'string') return content ? [{ type: 'text', text: content }] : []; + if (!Array.isArray(content)) return [{ type: 'text', text: String(content) }]; + const out = []; + for (const part of content) { + if (!part || typeof part !== 'object') continue; + if ((part.type === 'text' || part.type === 'input_text') && typeof part.text === 'string') { + out.push({ type: 'text', text: part.text }); + } else if (part.type === 'image_url' && part.image_url) { + const url = typeof part.image_url === 'string' ? part.image_url : part.image_url.url; + const data = typeof url === 'string' && url.match(/^data:([^;,]+);base64,(.+)$/s); + if (data) out.push({ type: 'image', source: { type: 'base64', media_type: data[1], data: data[2] } }); + else if (typeof url === 'string') out.push({ type: 'image', source: { type: 'url', url } }); + } + } + return out; +} + +function append(messages, role, blocks) { + if (!blocks.length) return; + const tail = messages[messages.length - 1]; + if (tail && tail.role === role) tail.content.push(...blocks); + else messages.push({ role, content: blocks }); +} + +function chatToAnthropic(body) { + if (!body || typeof body !== 'object' || Array.isArray(body)) throw new Error('request body must be a JSON object'); + if (!Array.isArray(body.messages)) throw new Error('messages must be an array'); + if (body.n != null && body.n !== 1) throw new Error('only n=1 is supported'); + const system = []; + const messages = []; + for (const message of body.messages) { + if (!message || typeof message !== 'object') continue; + if (message.role === 'system' || message.role === 'developer') { + system.push(...contentBlocks(message.content)); + } else if (message.role === 'assistant') { + const blocks = contentBlocks(message.content); + for (const call of message.tool_calls || []) { + if (!call || call.type !== 'function' || !call.function) continue; + let input; + try { input = JSON.parse(call.function.arguments || '{}'); } + catch (_) { throw new Error(`tool call ${call.id || call.function.name || ''} arguments must be valid JSON`); } + blocks.push({ type: 'tool_use', id: call.id || `call_${crypto.randomUUID()}`, name: call.function.name, input }); + } + append(messages, 'assistant', blocks); + } else if (message.role === 'tool') { + append(messages, 'user', [{ type: 'tool_result', tool_use_id: message.tool_call_id, content: contentBlocks(message.content) }]); + } else { + append(messages, 'user', contentBlocks(message.content)); + } + } + const out = { + model: cursorModel(body.model), + max_tokens: body.max_completion_tokens ?? body.max_tokens ?? 8192, + messages, + stream: body.stream === true + }; + if (system.length) out.system = system; + if (body.temperature != null) out.temperature = body.temperature; + if (body.top_p != null) out.top_p = body.top_p; + if (body.stop != null) out.stop_sequences = Array.isArray(body.stop) ? body.stop : [body.stop]; + if (Array.isArray(body.tools) && body.tool_choice !== 'none') { + out.tools = body.tools.filter(t => t && t.type === 'function' && t.function && t.function.name).map(t => ({ + name: t.function.name, + ...(t.function.description ? { description: t.function.description } : {}), + input_schema: t.function.parameters || { type: 'object', properties: {} } + })); + } + const choice = body.tool_choice; + if (choice === 'auto') out.tool_choice = { type: 'auto' }; + else if (choice === 'required') out.tool_choice = { type: 'any' }; + else if (choice && choice.type === 'function' && choice.function && choice.function.name) out.tool_choice = { type: 'tool', name: choice.function.name }; + return out; +} + +function finishReason(reason) { + if (reason === 'max_tokens') return 'length'; + if (reason === 'tool_use') return 'tool_calls'; + return 'stop'; +} + +function anthropicToChat(body) { + const blocks = Array.isArray(body.content) ? body.content : []; + const text = blocks.filter(b => b && b.type === 'text').map(b => b.text || '').join(''); + const calls = blocks.filter(b => b && b.type === 'tool_use').map(b => ({ + id: b.id, type: 'function', function: { name: b.name, arguments: JSON.stringify(b.input || {}) } + })); + const sourceUsage = body.usage || {}; + const message = { role: 'assistant', content: text || null }; + if (calls.length) message.tool_calls = calls; + return { + id: `chatcmpl-${String(body.id || crypto.randomUUID()).replace(/^msg_/, '')}`, + object: 'chat.completion', created: Math.floor(Date.now() / 1000), model: body.model, + choices: [{ index: 0, message, finish_reason: finishReason(body.stop_reason) }], + usage: { + prompt_tokens: sourceUsage.input_tokens || 0, + completion_tokens: sourceUsage.output_tokens || 0, + total_tokens: (sourceUsage.input_tokens || 0) + (sourceUsage.output_tokens || 0), + ...(sourceUsage.cache_read_input_tokens != null ? { prompt_tokens_details: { cached_tokens: sourceUsage.cache_read_input_tokens } } : {}) + } + }; +} + +function openAiError(status, body) { + const source = body && body.error ? body.error : body || {}; + return { error: { message: source.message || `Claude upstream returned HTTP ${status}`, type: source.type || 'api_error', param: null, code: source.type || null } }; +} + +function createSseTranslator() { + const state = { id: `chatcmpl-${crypto.randomUUID()}`, model: 'claude', created: Math.floor(Date.now() / 1000), tool: -1, input: 0, output: 0, cached: 0, finish: 'stop', buffer: '', done: false }; + const frame = (delta, finish = null, usage) => `data: ${JSON.stringify({ id: state.id, object: 'chat.completion.chunk', created: state.created, model: state.model, choices: usage ? [] : [{ index: 0, delta, finish_reason: finish }], ...(usage ? { usage } : {}) })}\n\n`; + function translate(event) { + if (event.type === 'message_start') { + const message = event.message || {}; const usage = message.usage || {}; + state.id = `chatcmpl-${String(message.id || crypto.randomUUID()).replace(/^msg_/, '')}`; state.model = message.model || state.model; + state.input = usage.input_tokens || 0; state.cached = usage.cache_read_input_tokens || 0; + return frame({ role: 'assistant', content: '' }); + } + if (event.type === 'content_block_start' && event.content_block && event.content_block.type === 'tool_use') { + state.tool++; const block = event.content_block; + return frame({ tool_calls: [{ index: state.tool, id: block.id, type: 'function', function: { name: block.name, arguments: '' } }] }); + } + if (event.type === 'content_block_delta' && event.delta) { + if (event.delta.type === 'text_delta') return frame({ content: event.delta.text || '' }); + if (event.delta.type === 'input_json_delta') return frame({ tool_calls: [{ index: state.tool, function: { arguments: event.delta.partial_json || '' } }] }); + } + if (event.type === 'message_delta') { + state.finish = finishReason(event.delta && event.delta.stop_reason); state.output = event.usage && event.usage.output_tokens || state.output; return ''; + } + if (event.type === 'error') { state.done = true; return `data: ${JSON.stringify(openAiError(502, event))}\n\ndata: [DONE]\n\n`; } + if (event.type === 'message_stop' && !state.done) { + state.done = true; + const usage = { prompt_tokens: state.input, completion_tokens: state.output, total_tokens: state.input + state.output, ...(state.cached ? { prompt_tokens_details: { cached_tokens: state.cached } } : {}) }; + return frame({}, state.finish) + frame({}, null, usage) + 'data: [DONE]\n\n'; + } + return ''; + } + return { + push(chunk) { + state.buffer += chunk.toString('utf8'); let out = ''; let at; + while ((at = state.buffer.indexOf('\n\n')) >= 0) { + const raw = state.buffer.slice(0, at); state.buffer = state.buffer.slice(at + 2); + const line = raw.split('\n').find(l => l.startsWith('data:')); if (!line) continue; + const payload = line.slice(5).trim(); if (!payload || payload === '[DONE]') continue; + try { out += translate(JSON.parse(payload)); } catch (_) {} + } + return out; + }, + end() { const out = state.done ? '' : translate({ type: 'message_stop' }); state.buffer = ''; return out; } + }; +} + +module.exports = { anthropicToChat, chatToAnthropic, createSseTranslator, cursorModel, openAiError }; diff --git a/packages/broker-proxy/openai-chat-compat.test.js b/packages/broker-proxy/openai-chat-compat.test.js new file mode 100644 index 0000000..3ce7807 --- /dev/null +++ b/packages/broker-proxy/openai-chat-compat.test.js @@ -0,0 +1,50 @@ +'use strict'; +const test = require('node:test'); +const assert = require('node:assert/strict'); +const { anthropicToChat, chatToAnthropic, createSseTranslator, cursorModel } = require('./openai-chat-compat'); + +test('maps Cursor-safe Anthropic aliases to upstream Claude models', () => { + assert.equal(cursorModel('anthropic:sonnet'), 'claude-sonnet-4-6'); + assert.equal(cursorModel('anthropic:opus'), 'claude-opus-4-6'); + assert.equal(cursorModel('anthropic:haiku'), 'claude-haiku-4-5-20251001'); + assert.equal(cursorModel('anthropic:claude-sonnet-4-6'), 'claude-sonnet-4-6'); + assert.equal(cursorModel('custom-model'), 'custom-model'); +}); + +test('converts Cursor messages and tools to Anthropic', () => { + const out = chatToAnthropic({ model: 'claude-sonnet-4-6', stream: true, messages: [ + { role: 'system', content: 'Use tools.' }, { role: 'user', content: 'Read.' }, + { role: 'assistant', tool_calls: [{ id: 'call_1', type: 'function', function: { name: 'read_file', arguments: '{"path":"README.md"}' } }] }, + { role: 'tool', tool_call_id: 'call_1', content: 'pool docs' } + ], tools: [{ type: 'function', function: { name: 'read_file', parameters: { type: 'object', properties: { path: { type: 'string' } } } } }] }); + assert.equal(out.system[0].text, 'Use tools.'); + assert.deepEqual(out.messages[1].content[0], { type: 'tool_use', id: 'call_1', name: 'read_file', input: { path: 'README.md' } }); + assert.equal(out.messages[2].content[0].tool_use_id, 'call_1'); + assert.equal(out.tools[0].name, 'read_file'); +}); + +test('converts Anthropic JSON tool calls and usage to Chat Completions', () => { + const out = anthropicToChat({ id: 'msg_1', model: 'claude', stop_reason: 'tool_use', content: [{ type: 'tool_use', id: 'toolu_1', name: 'read_file', input: { path: 'README.md' } }], usage: { input_tokens: 10, output_tokens: 3 } }); + assert.equal(out.id, 'chatcmpl-1'); + assert.equal(out.choices[0].finish_reason, 'tool_calls'); + assert.equal(out.choices[0].message.tool_calls[0].function.arguments, '{"path":"README.md"}'); + assert.deepEqual(out.usage, { prompt_tokens: 10, completion_tokens: 3, total_tokens: 13 }); +}); + +test('converts split Anthropic SSE with tool deltas and final usage', () => { + const t = createSseTranslator(); + const raw = [ + { type: 'message_start', message: { id: 'msg_s', model: 'claude', usage: { input_tokens: 12 } } }, + { type: 'content_block_start', content_block: { type: 'tool_use', id: 'toolu_1', name: 'read_file' } }, + { type: 'content_block_delta', delta: { type: 'input_json_delta', partial_json: '{"path":"README.md"}' } }, + { type: 'message_delta', delta: { stop_reason: 'tool_use' }, usage: { output_tokens: 7 } }, + { type: 'message_stop' } + ].map(e => `data: ${JSON.stringify(e)}\n\n`).join(''); + const split = Math.floor(raw.length / 2); + const out = t.push(raw.slice(0, split)) + t.push(raw.slice(split)) + t.end(); + assert.match(out, /"id":"toolu_1"/); + assert.match(out, /"finish_reason":"tool_calls"/); + assert.match(out, /"prompt_tokens":12/); + assert.match(out, /"completion_tokens":7/); + assert.equal(out.endsWith('data: [DONE]\n\n'), true); +}); diff --git a/packages/broker-proxy/proxy.js b/packages/broker-proxy/proxy.js index 4795e12..3419ffb 100644 --- a/packages/broker-proxy/proxy.js +++ b/packages/broker-proxy/proxy.js @@ -35,6 +35,7 @@ const { applyClaudeCodeProtocol, claudeCodeHeaders } = require('./claude-code-protocol'); +const { anthropicToChat, chatToAnthropic, createSseTranslator, openAiError } = require('./openai-chat-compat'); // ─── Defaults ─────────────────────────────────────────────────────────────── const DEFAULT_PORT = 18801; @@ -1660,6 +1661,7 @@ function sendUpstreamOnce(config, req, res, body, auth, reqNum, abortSignal, act return; } const lib = upstreamUrl.protocol === 'http:' ? http : https; + const isChat = req.method === 'POST' && req.url.split('?')[0] === '/v1/chat/completions'; let requestModelIsHaiku = false; try { requestModelIsHaiku = /haiku/i.test(JSON.parse(body.toString('utf8')).model || ''); } catch (_) {} const headers = buildUpstreamHeaders(req, body.length, auth.accessToken, requestModelIsHaiku); @@ -1667,7 +1669,7 @@ function sendUpstreamOnce(config, req, res, body, auth, reqNum, abortSignal, act let finished = false; const requestOptions = { protocol: upstreamUrl.protocol, - path: req.url, + path: isChat ? '/v1/messages' : req.url, method: req.method, headers, timeout: config.upstreamTimeoutMs @@ -1742,6 +1744,10 @@ function sendUpstreamOnce(config, req, res, body, auth, reqNum, abortSignal, act console.error(`[${ts}] #${reqNum} DETECTION! Body: ${body.length}b`); } errBody = reverseMap(errBody, config, activeRenames); + if (isChat) { + let parsed; try { parsed = JSON.parse(errBody); } catch (_) { parsed = { message: errBody }; } + errBody = JSON.stringify(openAiError(status, parsed)); + } const nh = { ...upRes.headers }; delete nh['transfer-encoding']; nh['content-length'] = Buffer.byteLength(errBody); @@ -1763,12 +1769,22 @@ function sendUpstreamOnce(config, req, res, body, auth, reqNum, abortSignal, act const TAIL_SIZE = 64; const decoder = new StringDecoder('utf8'); const observer = createSseUsageObserver(); + const chatTranslator = isChat ? createSseTranslator() : null; let pending = ''; let buffered = ''; - if (!config.bufferSseResponses) res.writeHead(status, sseHeaders); + if (chatTranslator) { + sseHeaders['content-type'] = 'text/event-stream; charset=utf-8'; + sseHeaders['cache-control'] = 'no-cache'; + res.writeHead(status, sseHeaders); + } else if (!config.bufferSseResponses) res.writeHead(status, sseHeaders); upRes.on('data', (chunk) => { const decoded = decoder.write(chunk); observer.push(decoded); + if (chatTranslator) { + const translated = chatTranslator.push(decoded); + if (translated) res.write(translated); + return; + } pending += decoded; if (pending.length > TAIL_SIZE) { let sliceIdx = pending.length - TAIL_SIZE; @@ -1783,7 +1799,11 @@ function sendUpstreamOnce(config, req, res, body, auth, reqNum, abortSignal, act upRes.on('end', () => { pending += decoder.end(); const observed = observer.result(); - if (config.bufferSseResponses) { + if (chatTranslator) { + const translated = chatTranslator.end(); + if (translated) res.write(translated); + res.end(); + } else if (config.bufferSseResponses) { buffered += pending; if (observed.errorCode) { const message = observed.errorMessage || `Upstream SSE terminated with ${observed.errorCode}`; @@ -1831,6 +1851,7 @@ function sendUpstreamOnce(config, req, res, body, auth, reqNum, abortSignal, act try { const parsed = JSON.parse(respBody); if (typeof parsed.model === 'string' && parsed.model.length > 0) actualModel = parsed.model; + if (isChat) respBody = JSON.stringify(anthropicToChat(parsed)); } catch (_) {} nh['x-actual-model'] = actualModel; nh['content-length'] = Buffer.byteLength(respBody); @@ -1927,6 +1948,15 @@ function createRequestHandler(config, state) { let body = Buffer.concat(chunks); let bodyStr = body.toString('utf8'); const originalSize = bodyStr.length; + const isChat = req.method === 'POST' && req.url.split('?')[0] === '/v1/chat/completions'; + if (isChat) { + try { bodyStr = JSON.stringify(chatToAnthropic(JSON.parse(bodyStr))); } + catch (e) { + const errorBody = JSON.stringify(openAiError(400, { type: 'invalid_request_error', message: e.message })); + res.writeHead(400, { 'content-type': 'application/json', 'content-length': Buffer.byteLength(errorBody) }); + res.end(errorBody); responseFinished = true; return; + } + } // Composability: determine which renames the CALLER actually originated, // from the untouched request body, BEFORE we shape it. Reverse-map will // only undo these, so native-CC / arbitrary harnesses get their own tool @@ -2023,7 +2053,11 @@ function createRequestHandler(config, state) { // above already reported this exact lease, this is a no-op instead of // a duplicate report the broker rejects with 404 unknown_lease. if (result.outcome) await reportBrokerOutcome(config, auth.lease, { ...result.outcome, latencyMs }); - const errBody = reverseMap(result.body, config, activeRenames); + let errBody = reverseMap(result.body, config, activeRenames); + if (isChat) { + let parsed; try { parsed = JSON.parse(errBody); } catch (_) { parsed = { message: errBody }; } + errBody = JSON.stringify(openAiError(result.status, parsed)); + } const nh = { ...result.headers }; delete nh['transfer-encoding']; nh['content-length'] = Buffer.byteLength(errBody); diff --git a/packages/broker-proxy/proxy.test.js b/packages/broker-proxy/proxy.test.js index 22a00f9..8ae31ab 100644 --- a/packages/broker-proxy/proxy.test.js +++ b/packages/broker-proxy/proxy.test.js @@ -171,7 +171,7 @@ function invokeProxy(config, body, headers = {}, opts = {}) { const req = Readable.from([Buffer.from(JSON.stringify(body))]); Object.assign(req, { method: 'POST', - url: '/v1/messages', + url: opts.url || '/v1/messages', headers: { 'content-type': 'application/json', 'anthropic-version': '2023-06-01', @@ -311,6 +311,26 @@ test('broker mode preserves complex Anthropic JSON fields and required headers', assert.equal(broker.state.reports[0].tokens, 18); }); +test('Cursor chat completions route uses a pooled Claude seat', async (t) => { + const broker = makeBroker([{ leaseId: 'lease-cursor', accountId: 'acct-cursor', accessToken: 'cursor-token' }]); + let sent; + installHttpMock(t, broker.handler.bind(broker), (options, body, callback) => { + sent = { options, body: JSON.parse(body.toString('utf8')) }; + emitResponse(callback, jsonResponse(200, { id: 'msg_cursor', type: 'message', model: 'claude-sonnet-4-6', content: [{ type: 'tool_use', id: 'toolu_1', name: 'read_file', input: { path: 'README.md' } }], stop_reason: 'tool_use', usage: { input_tokens: 9, output_tokens: 4 } })); + }); + const out = await invokeProxy(makeConfig(), { model: 'anthropic:sonnet', messages: [{ role: 'user', content: 'Read it.' }], tools: [{ type: 'function', function: { name: 'read_file', parameters: { type: 'object', properties: { path: { type: 'string' } } } } }] }, {}, { url: '/v1/chat/completions' }); + assert.equal(sent.options.path, '/v1/messages'); + assert.equal(sent.options.headers.authorization, 'Bearer cursor-token'); + assert.equal(sent.body.model, 'claude-sonnet-4-6'); + assert.equal(sent.body.tools[0].name, 'read_file'); + assert.match(sent.body.system[0].text, /cc_version=2\.1\.224/); + const response = JSON.parse(out.text); + assert.equal(response.object, 'chat.completion'); + assert.equal(response.choices[0].message.tool_calls[0].function.name, 'read_file'); + await waitFor(() => broker.state.reports.length >= 1); + assert.equal(broker.state.reports[0].tokens, 13); +}); + test('streaming SSE event order is preserved and usage is reported without full buffering', async (t) => { const broker = makeBroker([{ leaseId: 'lease-stream', accountId: 'acct-stream', accessToken: 'lease-token-stream' }]); const events = [ diff --git a/packages/pool-edge/src/lib/policy.js b/packages/pool-edge/src/lib/policy.js index a3a2058..ddfa44f 100644 --- a/packages/pool-edge/src/lib/policy.js +++ b/packages/pool-edge/src/lib/policy.js @@ -25,6 +25,11 @@ import { EDGE_CONFIG } from '../edge.gen.js'; export const ROUTES = Object.freeze({ // inference legs, forwarded upstream with the contributor's mapped pool key 'POST /v1/messages': { kind: 'proxy', upstream: '/v1/messages', meter: true }, + 'POST /v1/chat/completions': { + kind: 'proxy', + upstream: '/v1/chat/completions', + meter: true, + }, 'POST /v1/messages/count_tokens': { kind: 'proxy', upstream: '/v1/messages/count_tokens', diff --git a/packages/pool-edge/src/worker.js b/packages/pool-edge/src/worker.js index ae0ff78..02632aa 100644 --- a/packages/pool-edge/src/worker.js +++ b/packages/pool-edge/src/worker.js @@ -170,11 +170,11 @@ async function handleProxy(request, env, ctx, route, url) { // Model gate for restricted tiers. Only restricted tiers pay the buffering // cost; everyone else keeps a pure streaming pass-through. let body = request.body; - if (tier.models && route.upstream === '/v1/messages') { + if (tier.models && (route.upstream === '/v1/messages' || route.upstream === '/v1/chat/completions')) { const text = await request.text(); let model = null; try { - model = JSON.parse(text).model || null; + model = normalizeCursorModel(JSON.parse(text).model || null); } catch { // error-policy: an unparseable body is upstream's to reject, not ours to guess. } @@ -230,6 +230,16 @@ async function handleProxy(request, env, ctx, route, url) { return new Response(toClient, { status: response.status, headers }); } +function normalizeCursorModel(model) { + if (model === 'anthropic:sonnet') return 'claude-sonnet-4-6'; + if (model === 'anthropic:opus') return 'claude-opus-4-6'; + if (model === 'anthropic:haiku') return 'claude-haiku-4-5-20251001'; + if (typeof model === 'string' && model.startsWith('anthropic:claude-')) { + return model.slice('anthropic:'.length); + } + return model; +} + /** * Read usage off the response copy and add it to the edge counter. * SSE and non-SSE both carry `usage`; we parse whichever shape arrives. @@ -268,7 +278,7 @@ export function extractWeightedUsage(text, contentType) { if (!payload || payload === '[DONE]') continue; try { const event = JSON.parse(payload); - if (event.usage) total += weighUsage(event.usage); + if (event.usage) total += weighCompatibleUsage(event.usage); if (event.message && event.message.usage) total += weighUsage(event.message.usage); } catch { // error-policy: a malformed SSE frame is skipped, not treated as zero-cost overall. @@ -278,10 +288,21 @@ export function extractWeightedUsage(text, contentType) { } try { const parsed = JSON.parse(text); - return weighUsage(parsed.usage); + return weighCompatibleUsage(parsed.usage); } catch { return 0; } } +function weighCompatibleUsage(usage) { + if (!usage || typeof usage !== 'object') return 0; + if (usage.prompt_tokens == null && usage.completion_tokens == null) return weighUsage(usage); + const cached = Number(usage.prompt_tokens_details && usage.prompt_tokens_details.cached_tokens || 0); + return weighUsage({ + input_tokens: Math.max(0, Number(usage.prompt_tokens || 0) - cached), + output_tokens: Number(usage.completion_tokens || 0), + cache_read_input_tokens: cached, + }); +} + export { grantKey }; diff --git a/packages/pool-edge/test/policy.test.js b/packages/pool-edge/test/policy.test.js index d56ff6f..a4ebd88 100644 --- a/packages/pool-edge/test/policy.test.js +++ b/packages/pool-edge/test/policy.test.js @@ -62,6 +62,7 @@ test('ledger.json is forbidden but a lookalike prefix is not silently blocked', test('only the allowlisted routes resolve', () => { assert.ok(resolveRoute('POST', '/v1/messages')); + assert.ok(resolveRoute('POST', '/v1/chat/completions')); assert.ok(resolveRoute('GET', '/keys/status')); assert.ok(resolveRoute('GET', '/join')); assert.ok(resolveRoute('POST', '/join/start')); diff --git a/packages/pool-edge/test/worker.test.js b/packages/pool-edge/test/worker.test.js index 76601dd..969e4fa 100644 --- a/packages/pool-edge/test/worker.test.js +++ b/packages/pool-edge/test/worker.test.js @@ -233,6 +233,13 @@ test('streaming usage is extracted from SSE frames', () => { ); }); +test('Chat Completions usage is converted before edge weighting', () => { + const usage = { prompt_tokens: 100, completion_tokens: 10, prompt_tokens_details: { cached_tokens: 80 } }; + const expected = 20 * 1 + 80 * 0.1 + 10 * 5; + assert.equal(extractWeightedUsage(JSON.stringify({ usage }), 'application/json'), expected); + assert.equal(extractWeightedUsage(`data: ${JSON.stringify({ usage })}\n\ndata: [DONE]\n\n`, 'text/event-stream'), expected); +}); + test('malformed accounting payloads weigh zero rather than throwing', () => { assert.equal(extractWeightedUsage('not json', 'application/json'), 0); assert.equal(extractWeightedUsage('data: {bad\n\n', 'text/event-stream'), 0); diff --git a/packages/pool-meter/src/lib/openai-usage.js b/packages/pool-meter/src/lib/openai-usage.js index 53b0751..44b087a 100644 --- a/packages/pool-meter/src/lib/openai-usage.js +++ b/packages/pool-meter/src/lib/openai-usage.js @@ -50,7 +50,8 @@ function makeResponsesUsageParser(usage, reqMeta) { if (m) reqMeta.model = m; } const u = (resp && resp.usage) || obj.usage || null; - if (u && (u.input_tokens != null || u.output_tokens != null)) { + if (u && (u.input_tokens != null || u.output_tokens != null || + u.prompt_tokens != null || u.completion_tokens != null)) { applyOpenAiUsage(usage, u); } } diff --git a/packages/pool-meter/src/pool-meter.js b/packages/pool-meter/src/pool-meter.js index 6b271d6..0daae6b 100644 --- a/packages/pool-meter/src/pool-meter.js +++ b/packages/pool-meter/src/pool-meter.js @@ -2821,6 +2821,7 @@ function proxyRequest(req, res, info, prebuffered) { // capture request body for model/stream detection while forwarding const reqMeta = { model: null, stream: false }; const isMessages = req.url.startsWith('/v1/messages'); + const isChatCompletions = req.url.startsWith('/v1/chat/completions'); // Pooled trace capture (Feature 2): pooled usage defaults to traced, honoring // the per-key opt-out. When tracing, capture the FULL request body (bounded // by the trace clip) instead of just the 512K head we need for model/stream. @@ -2828,7 +2829,7 @@ function proxyRequest(req, res, info, prebuffered) { let reqBodyChunks = []; let reqBodyLen = 0; const REQ_CAP = wantTrace ? (2 * 1024 * 1024) : (512 * 1024); - const captureReq = isMessages || wantTrace; + const captureReq = isMessages || isChatCompletions || wantTrace; const upReq = http.request({ host: UPSTREAM_HOST, @@ -2910,7 +2911,7 @@ function proxyRequest(req, res, info, prebuffered) { }; if (isSse) { - const feed = makeSseUsageParser(usage); + const feed = isChatCompletions ? makeResponsesUsageParser(usage, reqMeta) : makeSseUsageParser(usage); upRes.on('data', (chunk) => { if (ttfb == null) ttfb = Date.now() - start; feed(chunk); if (textCollector) { try { textCollector.feed(chunk); } catch (_) {} } res.write(chunk); }); // tee: parse copy, pass original bytes upRes.on('end', () => { res.end(); finish(); }); } else { @@ -2930,10 +2931,13 @@ function proxyRequest(req, res, info, prebuffered) { if (wantTrace) nonSseRespText = raw; const body = JSON.parse(raw); if (body && body.usage) { - usage.input_tokens = body.usage.input_tokens || 0; - usage.output_tokens = body.usage.output_tokens || 0; - usage.cache_creation_input_tokens = body.usage.cache_creation_input_tokens || 0; - usage.cache_read_input_tokens = body.usage.cache_read_input_tokens || 0; + if (isChatCompletions) applyOpenAiUsage(usage, body.usage); + else { + usage.input_tokens = body.usage.input_tokens || 0; + usage.output_tokens = body.usage.output_tokens || 0; + usage.cache_creation_input_tokens = body.usage.cache_creation_input_tokens || 0; + usage.cache_read_input_tokens = body.usage.cache_read_input_tokens || 0; + } } if (!reqMeta.model && body && body.model) reqMeta.model = body.model; } catch (_) { /* non-JSON or huge body: log zeros */ } @@ -2957,7 +2961,7 @@ function proxyRequest(req, res, info, prebuffered) { // The body was already consumed by the tier model gate: replay it verbatim // instead of piping a stream that has no data left to emit. if (prebuffered) { - if (isMessages) { + if (isMessages || isChatCompletions) { try { const b = JSON.parse(prebuffered.toString('utf8')); reqMeta.model = b.model || null; diff --git a/packages/pool-meter/test/openai-usage-unit.js b/packages/pool-meter/test/openai-usage-unit.js index 6de5d14..e988813 100644 --- a/packages/pool-meter/test/openai-usage-unit.js +++ b/packages/pool-meter/test/openai-usage-unit.js @@ -90,6 +90,15 @@ t('SSE without cached details leaves input as reported', () => { assert.strictEqual(usage.cache_read_input_tokens, 0); }); +t('SSE chat-completions usage shape is parsed', () => { + const usage = { input_tokens: 0, output_tokens: 0, cache_read_input_tokens: 0 }; + const feed = makeResponsesUsageParser(usage, {}); + feed(Buffer.from('data: {"model":"claude-sonnet-4-6","choices":[],"usage":{"prompt_tokens":1200,"completion_tokens":44,"prompt_tokens_details":{"cached_tokens":1024}}}\n')); + assert.strictEqual(usage.input_tokens, 176); + assert.strictEqual(usage.cache_read_input_tokens, 1024); + assert.strictEqual(usage.output_tokens, 44); +}); + t('SSE events without usage do not clobber parsed usage', () => { const usage = { input_tokens: 0, output_tokens: 0, cache_read_input_tokens: 0 }; const feed = makeResponsesUsageParser(usage, {});