diff --git a/apps/api/src/handlers/brain-inference/__tests__/brain-inference.test.ts b/apps/api/src/handlers/brain-inference/__tests__/brain-inference.test.ts index 9c8fc6f9c..c0665a911 100644 --- a/apps/api/src/handlers/brain-inference/__tests__/brain-inference.test.ts +++ b/apps/api/src/handlers/brain-inference/__tests__/brain-inference.test.ts @@ -6,16 +6,19 @@ const { mockGetBrainGatewayToken, mockResolveBrainInferenceProvider, mockMapBrainModelName, + mockRecordLlmUsage, } = vi.hoisted(() => ({ mockGetBrainGatewayToken: vi.fn(), mockResolveBrainInferenceProvider: vi.fn(), mockMapBrainModelName: vi.fn(), + mockRecordLlmUsage: vi.fn(), })); vi.mock('@roomote/sdk/server', () => ({ getBrainGatewayToken: mockGetBrainGatewayToken, resolveBrainInferenceProvider: mockResolveBrainInferenceProvider, mapBrainModelName: mockMapBrainModelName, + recordLlmUsage: mockRecordLlmUsage, })); const { brainInference } = await import('../index'); @@ -56,6 +59,7 @@ beforeEach(() => { vi.unstubAllGlobals(); mockGetBrainGatewayToken.mockReturnValue(GATEWAY_TOKEN); mockResolveBrainInferenceProvider.mockResolvedValue(OPENROUTER); + mockRecordLlmUsage.mockResolvedValue({ recorded: true }); mockMapBrainModelName.mockImplementation((requested: string) => requested === 'text-embedding-3-small' ? 'openai/text-embedding-3-small' @@ -139,6 +143,57 @@ describe('brain inference gateway', () => { }); }); + it('records embedding tokens, provider cost, and request diagnostics', async () => { + vi.stubGlobal( + 'fetch', + vi.fn(async () => + Response.json({ + data: [], + usage: { + prompt_tokens: 12, + total_tokens: 12, + cost: 0.0000042, + cost_details: { upstream_inference_cost: 0.000004 }, + }, + model: 'openai/text-embedding-3-small-2026-08-01', + }), + ), + ); + + const response = await post('/v1/embeddings', { + token: GATEWAY_TOKEN, + body: { model: 'text-embedding-3-small', input: ['hello'] }, + }); + + expect(response.status).toBe(200); + await vi.waitFor(() => { + expect(mockRecordLlmUsage).toHaveBeenCalledWith( + expect.objectContaining({ + eventKey: expect.stringMatching(/^brain-inference-gateway:/u), + source: 'brain-inference-gateway', + usageType: 'embedding', + providerId: 'openrouter', + modelId: 'openai/text-embedding-3-small-2026-08-01', + inputTokens: 12, + totalTokens: 12, + costMicroUsd: 4, + costSource: 'provider_response', + pricingMetadata: { + costDetails: { upstream_inference_cost: 0.000004 }, + }, + details: expect.objectContaining({ + operation: 'embeddings', + upstreamPath: '/v1/embeddings', + status: 200, + latencyMs: expect.any(Number), + usageMetadataAvailable: true, + metadataReadFailed: false, + }), + }), + ); + }); + }); + it('routes reranking through OpenRouter without exposing its key to gbrain', async () => { const fetchMock = vi.fn( async (_url: string, _init: RequestInit) => @@ -166,6 +221,40 @@ describe('brain inference gateway', () => { expect(JSON.parse(init.body as string)).toEqual(body); }); + it('records metadata-free reranking with missing cost instead of dropping it', async () => { + vi.stubGlobal( + 'fetch', + vi.fn(async () => Response.json({ results: [] })), + ); + + const response = await post('/v1/rerank', { + token: GATEWAY_TOKEN, + body: { + model: 'cohere/rerank-v3.5', + query: 'query', + documents: ['document'], + }, + }); + + expect(response.status).toBe(200); + await vi.waitFor(() => { + expect(mockRecordLlmUsage).toHaveBeenCalledWith( + expect.objectContaining({ + usageType: 'rerank', + providerId: 'openrouter', + modelId: 'cohere/rerank-v3.5', + costMicroUsd: null, + costSource: 'missing', + details: expect.objectContaining({ + operation: 'rerank', + status: 200, + usageMetadataAvailable: false, + }), + }), + ); + }); + }); + it('reports reranking as unavailable when only OpenAI is configured', async () => { mockResolveBrainInferenceProvider.mockResolvedValue({ providerId: 'openai', @@ -204,6 +293,147 @@ describe('brain inference gateway', () => { }); expect(response.status).toBe(502); + await vi.waitFor(() => { + expect(mockRecordLlmUsage).toHaveBeenCalledWith( + expect.objectContaining({ + usageType: 'embedding', + providerId: 'openrouter', + modelId: 'openai/text-embedding-3-small', + costMicroUsd: null, + costSource: 'missing', + details: expect.objectContaining({ + operation: 'embeddings', + status: 502, + usageMetadataAvailable: false, + }), + }), + ); + }); + }); + + it('records an upstream provider failure with its actual status', async () => { + vi.stubGlobal( + 'fetch', + vi.fn(async () => + Response.json({ error: 'rate limited' }, { status: 429 }), + ), + ); + + const response = await post('/v1/chat/completions', { + token: GATEWAY_TOKEN, + body: { model: 'gpt-5.6-luna', messages: [] }, + }); + + expect(response.status).toBe(429); + await vi.waitFor(() => { + expect(mockRecordLlmUsage).toHaveBeenCalledWith( + expect.objectContaining({ + usageType: 'inference', + providerId: 'openrouter', + modelId: 'gpt-5.6-luna', + costSource: 'missing', + details: expect.objectContaining({ + operation: 'chat_completions', + status: 429, + usageMetadataAvailable: false, + }), + }), + ); + }); + }); + + it('normalizes response API streaming token details as inference usage', async () => { + vi.stubGlobal( + 'fetch', + vi.fn( + async () => + new Response( + 'data: {"type":"response.completed","response":{"usage":{"input_tokens":30,"output_tokens":8,"total_tokens":38,"input_tokens_details":{"cached_tokens":10},"output_tokens_details":{"reasoning_tokens":3}}}}\n\ndata: [DONE]\n\n', + { headers: { 'content-type': 'text/event-stream' } }, + ), + ), + ); + + const response = await post('/v1/responses', { + token: GATEWAY_TOKEN, + body: { model: 'gpt-5.6-luna', input: 'hello', stream: true }, + }); + + expect(response.status).toBe(200); + await vi.waitFor(() => { + expect(mockRecordLlmUsage).toHaveBeenCalledWith( + expect.objectContaining({ + usageType: 'inference', + modelId: 'gpt-5.6-luna', + inputTokens: 30, + outputTokens: 8, + reasoningTokens: 3, + cacheReadTokens: 10, + totalTokens: 38, + details: expect.objectContaining({ operation: 'responses' }), + }), + ); + }); + }); + + it('requests usage metadata for OpenAI-compatible chat streams', async () => { + const fetchMock = vi.fn(async (_url: string, _init: RequestInit) => + Response.json({ choices: [] }), + ); + vi.stubGlobal('fetch', fetchMock); + + await post('/v1/chat/completions', { + token: GATEWAY_TOKEN, + body: { + model: 'gpt-5.6-luna', + messages: [], + stream: true, + stream_options: { custom_option: true }, + }, + }); + + const init = fetchMock.mock.calls[0]![1]; + expect(JSON.parse(init.body as string)).toMatchObject({ + stream_options: { custom_option: true, include_usage: true }, + }); + }); + + it('records the final usage chunk from a chat completions stream', async () => { + vi.stubGlobal( + 'fetch', + vi.fn( + async () => + new Response( + 'data: {"choices":[],"model":"openai/gpt-5.6-luna","usage":{"prompt_tokens":21,"completion_tokens":5,"total_tokens":26,"cost":0.0009}}\n\ndata: [DONE]\n\n', + { headers: { 'content-type': 'text/event-stream' } }, + ), + ), + ); + + const response = await post('/v1/chat/completions', { + token: GATEWAY_TOKEN, + body: { model: 'gpt-5.6-luna', messages: [], stream: true }, + }); + + expect(response.status).toBe(200); + await vi.waitFor(() => { + expect(mockRecordLlmUsage).toHaveBeenCalledWith( + expect.objectContaining({ + usageType: 'inference', + modelId: 'openai/gpt-5.6-luna', + inputTokens: 21, + outputTokens: 5, + totalTokens: 26, + costMicroUsd: 900, + costSource: 'provider_response', + details: expect.objectContaining({ + operation: 'chat_completions', + status: 200, + usageMetadataAvailable: true, + }), + }), + ); + }); }); it('forwards an operator-chosen model when one is configured', async () => { diff --git a/apps/api/src/handlers/brain-inference/index.ts b/apps/api/src/handlers/brain-inference/index.ts index 4075b411f..309612c7d 100644 --- a/apps/api/src/handlers/brain-inference/index.ts +++ b/apps/api/src/handlers/brain-inference/index.ts @@ -9,6 +9,7 @@ import { import { getBrainGatewayToken, mapBrainModelName, + recordLlmUsage, resolveBrainInferenceProvider, type ResolvedBrainInference, } from '@roomote/sdk/server'; @@ -60,6 +61,187 @@ const RESPONSE_HEADER_DENYLIST = new Set([ 'connection', ]); +type BrainInferenceOperation = { + name: 'embeddings' | 'rerank' | 'chat_completions' | 'responses'; + usageType: 'embedding' | 'rerank' | 'inference'; +}; + +type ProviderUsage = { + modelId: string | null; + inputTokens: number | null; + outputTokens: number | null; + reasoningTokens: number | null; + cacheReadTokens: number | null; + cacheWriteTokens: number | null; + totalTokens: number | null; + costUsd: number | null; + costDetails: Record | null; +}; + +function resolveOperation(upstreamPath: string): BrainInferenceOperation { + if (upstreamPath === '/v1/embeddings') { + return { name: 'embeddings', usageType: 'embedding' }; + } + + if (upstreamPath === '/v1/rerank') { + return { name: 'rerank', usageType: 'rerank' }; + } + + return { + name: upstreamPath === '/v1/responses' ? 'responses' : 'chat_completions', + usageType: 'inference', + }; +} + +function finiteNonNegative(value: unknown): number | null { + return typeof value === 'number' && Number.isFinite(value) && value >= 0 + ? value + : null; +} + +function objectValue(value: unknown): Record | null { + return typeof value === 'object' && value !== null && !Array.isArray(value) + ? (value as Record) + : null; +} + +function findUsageEnvelope(body: string): Record | null { + const candidates = [body]; + + if (body.includes('\ndata:')) { + candidates.push( + ...body + .split('\n') + .filter((line) => line.startsWith('data:')) + .map((line) => line.slice('data:'.length).trim()) + .filter((line) => line && line !== '[DONE]'), + ); + } + + for (const candidate of candidates.reverse()) { + try { + const envelope = objectValue(JSON.parse(candidate)); + const response = objectValue(envelope?.response); + + if (objectValue(envelope?.usage)) { + return envelope; + } + + if (objectValue(response?.usage)) { + return response; + } + } catch { + // Provider errors and streaming sentinels are not necessarily JSON. + } + } + + return null; +} + +function parseProviderUsage(body: string): ProviderUsage | null { + const envelope = findUsageEnvelope(body); + const usage = objectValue(envelope?.usage); + + if (!usage) { + return null; + } + + const inputDetails = + objectValue(usage.prompt_tokens_details) ?? + objectValue(usage.input_tokens_details); + const outputDetails = + objectValue(usage.completion_tokens_details) ?? + objectValue(usage.output_tokens_details); + + return { + modelId: typeof envelope?.model === 'string' ? envelope.model : null, + inputTokens: + finiteNonNegative(usage.prompt_tokens) ?? + finiteNonNegative(usage.input_tokens), + outputTokens: + finiteNonNegative(usage.completion_tokens) ?? + finiteNonNegative(usage.output_tokens), + reasoningTokens: + finiteNonNegative(usage.reasoning_tokens) ?? + finiteNonNegative(outputDetails?.reasoning_tokens), + cacheReadTokens: + finiteNonNegative(usage.cache_read_tokens) ?? + finiteNonNegative(inputDetails?.cached_tokens), + cacheWriteTokens: finiteNonNegative(usage.cache_write_tokens), + totalTokens: finiteNonNegative(usage.total_tokens), + costUsd: finiteNonNegative(usage.cost), + costDetails: objectValue(usage.cost_details), + }; +} + +function recordBrainInferenceUsage(input: { + requestId: string; + providerId: ResolvedBrainInference['providerId']; + modelId: string | null; + operation: BrainInferenceOperation; + upstreamPath: string; + status: number; + startedAt: number; + response?: Response; +}): void { + const persist = async () => { + let usage: ProviderUsage | null = null; + let metadataReadFailed = false; + + if (input.response) { + try { + usage = parseProviderUsage(await input.response.text()); + } catch { + metadataReadFailed = true; + } + } + + const costMicroUsd = + usage?.costUsd === null || usage?.costUsd === undefined + ? null + : Math.round(usage.costUsd * 1_000_000); + + await recordLlmUsage({ + eventKey: `brain-inference-gateway:${input.requestId}`, + source: 'brain-inference-gateway', + usageType: input.operation.usageType, + providerId: input.providerId, + modelId: usage?.modelId ?? input.modelId, + inputTokens: usage?.inputTokens, + outputTokens: usage?.outputTokens, + reasoningTokens: usage?.reasoningTokens, + cacheReadTokens: usage?.cacheReadTokens, + cacheWriteTokens: usage?.cacheWriteTokens, + totalTokens: usage?.totalTokens, + contextTokens: usage?.inputTokens, + costMicroUsd, + costSource: costMicroUsd === null ? 'missing' : 'provider_response', + pricingMetadata: usage?.costDetails + ? { costDetails: usage.costDetails } + : undefined, + details: { + operation: input.operation.name, + upstreamPath: input.upstreamPath, + status: input.status, + latencyMs: Date.now() - input.startedAt, + usageMetadataAvailable: usage !== null, + metadataReadFailed, + }, + }); + }; + + void persist().catch((error) => { + console.warn( + formatSingleLineLog(`${LOG_PREFIX} Failed to record usage`, { + requestId: input.requestId, + providerId: input.providerId, + upstreamPath: input.upstreamPath, + error: error instanceof Error ? error.message : String(error), + }), + ); + }); +} + function presentedToken(header: string | undefined): string | null { if (!header) { return null; @@ -94,6 +276,7 @@ function tokenMatches(presented: string, expected: string): boolean { async function rewriteBody( rawBody: string, resolved: ResolvedBrainInference, + upstreamPath: string, ): Promise { if (!rawBody) { return rawBody; @@ -106,10 +289,22 @@ async function rewriteBody( return rawBody; } - return JSON.stringify({ + const rewritten: Record = { ...parsed, model: mapBrainModelName(parsed.model, resolved), - }); + }; + + // OpenAI-compatible chat streams omit usage unless explicitly requested. + // Preserve any existing stream options while asking supported providers to + // include the final usage-only chunk for asynchronous accounting. + if (upstreamPath === '/v1/chat/completions' && parsed.stream === true) { + rewritten.stream_options = { + ...(objectValue(parsed.stream_options) ?? {}), + include_usage: true, + }; + } + + return JSON.stringify(rewritten); } catch { // Not JSON we understand; forward untouched rather than failing the call. return rawBody; @@ -218,8 +413,18 @@ brainInference.post('/*', async (c) => { : resolved.apiKey, ); - const body = await rewriteBody(await c.req.text(), resolved); + const body = await rewriteBody(await c.req.text(), resolved, upstreamPath); const startedAt = Date.now(); + const requestId = crypto.randomUUID(); + const operation = resolveOperation(upstreamPath); + let modelId: string | null = null; + + try { + const parsedBody = JSON.parse(body) as Record; + modelId = typeof parsedBody.model === 'string' ? parsedBody.model : null; + } catch { + // An unparseable body is still forwarded and accounted for without a model. + } let upstream: Response; @@ -236,6 +441,15 @@ brainInference.post('/*', async (c) => { }, ); } catch (error) { + recordBrainInferenceUsage({ + requestId, + providerId: resolved.providerId, + modelId, + operation, + upstreamPath, + status: 502, + startedAt, + }); console.warn( formatSingleLineLog(`${LOG_PREFIX} Upstream request failed`, { providerId: resolved.providerId, @@ -258,6 +472,17 @@ brainInference.post('/*', async (c) => { ); } + recordBrainInferenceUsage({ + requestId, + providerId: resolved.providerId, + modelId, + operation, + upstreamPath, + status: upstream.status, + startedAt, + response: upstream.clone(), + }); + const responseHeaders = new Headers(); for (const [name, value] of upstream.headers.entries()) { diff --git a/packages/db/src/lib/llm-usage.ts b/packages/db/src/lib/llm-usage.ts index 0b5c7e70b..4695aa6c1 100644 --- a/packages/db/src/lib/llm-usage.ts +++ b/packages/db/src/lib/llm-usage.ts @@ -127,6 +127,7 @@ export async function recordLlmUsage( costSource: costSource as | 'opencode_message' | 'litellm_gateway' + | 'provider_response' | 'missing', messageCreatedAt: input.messageCreatedAt ?? null, messageCompletedAt: input.messageCompletedAt ?? null, diff --git a/packages/db/src/schema.ts b/packages/db/src/schema.ts index 5b8773b6c..0e3f4b822 100644 --- a/packages/db/src/schema.ts +++ b/packages/db/src/schema.ts @@ -1627,7 +1627,9 @@ export const llmUsageEvents = pgTable( .default(0), costSource: text('cost_source') .notNull() - .$type<'opencode_message' | 'litellm_gateway' | 'missing'>(), + .$type< + 'opencode_message' | 'litellm_gateway' | 'provider_response' | 'missing' + >(), pricingMetadata: jsonb('pricing_metadata') .notNull() .default({}) diff --git a/packages/types/src/llm-usage.ts b/packages/types/src/llm-usage.ts index 05e4ef759..cb37fbc39 100644 --- a/packages/types/src/llm-usage.ts +++ b/packages/types/src/llm-usage.ts @@ -1,6 +1,7 @@ export const LLM_USAGE_COST_SOURCES = [ 'opencode_message', 'litellm_gateway', + 'provider_response', 'missing', ] as const;