diff --git a/ENVIRONMENT.md b/ENVIRONMENT.md index c1f67705983..58358a28b6a 100644 --- a/ENVIRONMENT.md +++ b/ENVIRONMENT.md @@ -459,7 +459,7 @@ When `VERCEL_TARGET_ENV` is absent in local development or a script process, tra - `R2_API_REQUEST_LOG_BUCKET_NAME` - R2 bucket for `api_request_log` request and response bodies. [SERVER] - `R2_API_REQUEST_LOG_ACCESS_KEY_ID` - R2 access key ID used only for the `api_request_log` bucket; uses the `R2_ACCOUNT_ID` account. `[SECRET]` - `R2_API_REQUEST_LOG_SECRET_ACCESS_KEY` - R2 secret access key paired with `R2_API_REQUEST_LOG_ACCESS_KEY_ID`. `[SECRET]` -- `R2_AI_PROVIDER_METADATA_BUCKET_NAME` - Optional R2 bucket for AI provider metadata; Vercel objects use `vercel/.json` and gateway routing decisions use `kilo/.json`. [SERVER] +- `R2_AI_PROVIDER_METADATA_BUCKET_NAME` - Optional R2 bucket for AI provider metadata; `.json` contains the version-1 gateway routing decision and optional `vercel: { metadata }`. [SERVER] - `R2_AI_PROVIDER_METADATA_ACCESS_KEY_ID` - Dedicated R2 access key ID for the AI provider metadata bucket; uses the `R2_ACCOUNT_ID` account. `[SECRET]` - `R2_AI_PROVIDER_METADATA_SECRET_ACCESS_KEY` - R2 secret access key paired with `R2_AI_PROVIDER_METADATA_ACCESS_KEY_ID`. `[SECRET]` diff --git a/packages/web-shared/src/lib/ai-gateway/handlers/llm-proxy.test.ts b/packages/web-shared/src/lib/ai-gateway/handlers/llm-proxy.test.ts index 220fe05a317..cf5fa01a1de 100644 --- a/packages/web-shared/src/lib/ai-gateway/handlers/llm-proxy.test.ts +++ b/packages/web-shared/src/lib/ai-gateway/handlers/llm-proxy.test.ts @@ -892,10 +892,8 @@ describe('POST /api/openrouter/v1/chat/completions request handling', () => { expect(usageId).toMatch( /^[0-9a-f]{8}-[0-9a-f]{4}-4[0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}$/ ); - expect([...fakeR2.objects.keys()]).toEqual([`test-provider-metadata/kilo/${usageId}.json`]); - expect( - JSON.parse(fakeR2.objects.get(`test-provider-metadata/kilo/${usageId}.json`) ?? '') - ).toEqual({ + expect([...fakeR2.objects.keys()]).toEqual([`test-provider-metadata/${usageId}.json`]); + expect(JSON.parse(fakeR2.objects.get(`test-provider-metadata/${usageId}.json`) ?? '')).toEqual({ schema_version: 1, usage_id: usageId, created_at: expect.any(String), @@ -903,6 +901,177 @@ describe('POST /api/openrouter/v1/chat/completions request handling', () => { }); }); + it('archives Vercel metadata captured during accounting with the routing decision', async () => { + const { fakeR2 } = jest.requireMock( + '@kilocode/web-shared/lib/r2/create-client' + ); + fakeR2.objects.clear(); + const env = jest.replaceProperty(process, 'env', { + ...process.env, + R2_ACCOUNT_ID: 'test-metadata-account', + R2_AI_PROVIDER_METADATA_BUCKET_NAME: 'test-provider-metadata', + R2_AI_PROVIDER_METADATA_ACCESS_KEY_ID: 'test-metadata-access-key', + R2_AI_PROVIDER_METADATA_SECRET_ACCESS_KEY: 'test-metadata-secret-key', + }); + const decision: ProviderRoutingDecision = { + requested_model: 'openai/gpt-4o', + api_kind: 'chat_completions', + owner_type: 'user', + steps: [ + { + step: 'kilo', + outcome: 'selected', + gateway: 'vercel', + bounded_spending: false, + kilo_ignored_providers: [], + }, + ], + byok_candidates: [], + result: { + kind: 'provider', + route: 'kilo', + provider_type: 'vercel', + byok_provider_ids: [], + payer: 'kilo', + is_free: false, + }, + }; + const vercelFrame = { + provider_metadata: { request_id: 'vercel-request-456', region: 'iad1' }, + }; + mockedGetProvider.mockImplementation(async ({ onDecision }) => { + onDecision?.(decision); + return { + kind: 'provider', + provider: { ...provider, type: 'vercel' }, + userByok: null, + isFree: false, + bypassAccessCheck: false, + }; + }); + mockedAccountForMicrodollarUsage.mockImplementation(async (_response, usageContext) => { + const capture = usageContext.providerMetadataCapture; + capture?.captureVercelMetadata(vercelFrame); + capture?.waitFor(Promise.resolve()); + return null; + }); + const { handleLlmProxyRequest } = await import('./llm-proxy'); + + let response: Response; + try { + response = await handleLlmProxyRequest(makeRequest(makeBody()) as never); + expect(fakeR2.objects.size).toBe(0); + await Promise.all( + jest.mocked(after).mock.calls.map(([task]) => (typeof task === 'function' ? task() : task)) + ); + } finally { + env.restore(); + } + + expect(response.status).toBe(200); + expect(mockedAccountForMicrodollarUsage).toHaveBeenCalledTimes(1); + expect(mockedRewriteModelResponse).toHaveBeenCalledTimes(1); + const usageContext = mockedAccountForMicrodollarUsage.mock.calls[0]?.[1]; + expect(usageContext?.providerMetadataCapture).toEqual( + expect.objectContaining({ + captureRoutingDecision: expect.any(Function), + captureVercelMetadata: expect.any(Function), + waitFor: expect.any(Function), + }) + ); + const usageId = usageContext?.usageId; + expect(usageId).toMatch( + /^[0-9a-f]{8}-[0-9a-f]{4}-4[0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}$/ + ); + const key = `test-provider-metadata/${usageId}.json`; + expect([...fakeR2.objects.keys()]).toEqual([key]); + expect(JSON.parse(fakeR2.objects.get(key) ?? '')).toEqual({ + schema_version: 1, + usage_id: usageId, + created_at: expect.any(String), + ...decision, + vercel: { + metadata: vercelFrame.provider_metadata, + }, + }); + }); + + it('archives only the routing decision when the selected provider rejects the API kind', async () => { + const { fakeR2 } = jest.requireMock( + '@kilocode/web-shared/lib/r2/create-client' + ); + fakeR2.objects.clear(); + const env = jest.replaceProperty(process, 'env', { + ...process.env, + R2_ACCOUNT_ID: 'test-metadata-account', + R2_AI_PROVIDER_METADATA_BUCKET_NAME: 'test-provider-metadata', + R2_AI_PROVIDER_METADATA_ACCESS_KEY_ID: 'test-metadata-access-key', + R2_AI_PROVIDER_METADATA_SECRET_ACCESS_KEY: 'test-metadata-secret-key', + }); + const decision: ProviderRoutingDecision = { + requested_model: 'openai/gpt-4o', + api_kind: 'messages', + owner_type: 'user', + steps: [ + { + step: 'kilo', + outcome: 'selected', + gateway: 'openrouter', + bounded_spending: false, + kilo_ignored_providers: [], + }, + ], + byok_candidates: [], + result: { + kind: 'provider', + route: 'kilo', + provider_type: 'openrouter', + byok_provider_ids: [], + payer: 'kilo', + is_free: false, + }, + }; + mockedGetProvider.mockImplementation(async ({ onDecision }) => { + onDecision?.(decision); + return { + kind: 'provider', + provider: { ...provider, supportedChatApis: ['responses'] }, + userByok: null, + isFree: false, + bypassAccessCheck: false, + }; + }); + const { handleLlmProxyRequest } = await import('./llm-proxy'); + + let response: Response; + try { + response = await handleLlmProxyRequest( + makeRequest({ ...makeBody(), max_tokens: 100 }, undefined, '/messages') as never + ); + expect(fakeR2.objects.size).toBe(0); + await Promise.all( + jest.mocked(after).mock.calls.map(([task]) => (typeof task === 'function' ? task() : task)) + ); + } finally { + env.restore(); + } + + expect(response.status).toBe(400); + expect(mockedAccountForMicrodollarUsage).not.toHaveBeenCalled(); + expect(mockedUpstreamRequest).not.toHaveBeenCalled(); + expect(fakeR2.objects.size).toBe(1); + const [key] = [...fakeR2.objects.keys()]; + expect(key).toMatch(/^test-provider-metadata\/[0-9a-f-]+\.json$/); + const archived = JSON.parse(fakeR2.objects.get(key ?? '') ?? ''); + expect(archived).toEqual({ + schema_version: 1, + usage_id: key?.split('/').at(-1)?.replace('.json', ''), + created_at: expect.any(String), + ...decision, + }); + expect(archived).not.toHaveProperty('vercel'); + }); + it('passes the Vercel request ID to request logging', async () => { const { handleLlmProxyRequest } = await import('./llm-proxy'); diff --git a/packages/web-shared/src/lib/ai-gateway/handlers/llm-proxy.ts b/packages/web-shared/src/lib/ai-gateway/handlers/llm-proxy.ts index 4cee736c710..120d99e84c5 100644 --- a/packages/web-shared/src/lib/ai-gateway/handlers/llm-proxy.ts +++ b/packages/web-shared/src/lib/ai-gateway/handlers/llm-proxy.ts @@ -22,7 +22,7 @@ import { providerPrivacySchema, } from '@kilocode/web-shared/lib/ai-gateway/provider-privacy'; import { getProvider } from '@kilocode/web-shared/lib/ai-gateway/providers/get-provider'; -import { archiveKiloRoutingDecision } from '@kilocode/web-shared/lib/r2/kilo-routing-decision'; +import { createAiProviderMetadataCapture } from '@kilocode/web-shared/lib/r2/kilo-routing-decision'; import { getIgnoredProviders } from '@kilocode/web-shared/lib/ai-gateway/providers/apply-provider-specific-logic'; import { isNonTrialEnterpriseOrganization } from '@kilocode/web-shared/lib/organizations/non-trial-enterprise'; import { @@ -942,6 +942,7 @@ export async function handleLlmProxyRequest( const requestCostLimit = request.headers.get('x-kilo-max-cost-microdollars'); const usageId = randomUUID(); + const providerMetadataCapture = createAiProviderMetadataCapture(usageId); const providerResult = await getProvider({ requestedModel: effectiveModelIdLowerCased, request: requestBodyParsed, @@ -955,7 +956,7 @@ export async function handleLlmProxyRequest( getAllowedProviders: accessCheckResolver.getAllowedProviders, isDirectByokProviderAllowed, kiloIgnoredProviders: ignoredProviders, - onDecision: decision => archiveKiloRoutingDecision(usageId, decision), + onDecision: decision => providerMetadataCapture?.captureRoutingDecision(decision), }); if (providerResult.kind === 'chatgpt-reconnect') { // The person's enabled ChatGPT connection is terminally dead. Fail readably @@ -1121,6 +1122,7 @@ export async function handleLlmProxyRequest( requestStartedAt, clientRequestId, usageId, + providerMetadataCapture, bouncer: bouncer.usage({ occurredAt: new Date(requestStartedAtMs), clientAttributed: feature !== null || Boolean(xKiloCodeVersion), diff --git a/packages/web-shared/src/lib/ai-gateway/processUsage.messages.test.ts b/packages/web-shared/src/lib/ai-gateway/processUsage.messages.test.ts index 0abf036c025..7fc10cb9852 100644 --- a/packages/web-shared/src/lib/ai-gateway/processUsage.messages.test.ts +++ b/packages/web-shared/src/lib/ai-gateway/processUsage.messages.test.ts @@ -10,6 +10,7 @@ import { join } from 'node:path'; import { createReadStream } from 'node:fs'; import { readFile } from 'node:fs/promises'; import { Readable } from 'node:stream'; +import { createVercelProviderMetadataCollector } from '@kilocode/web-shared/lib/r2/vercel-provider-metadata'; jest.mock('@sentry/nextjs', () => ({ ...jest.requireActual('@sentry/nextjs'), @@ -309,6 +310,56 @@ describe('parseMessagesMicrodollarUsageFromStream approval tests', () => { }); describe('parseMessagesMicrodollarUsageFromStream', () => { + test('releases metadata capture on [DONE] without waiting for the connection to close', async () => { + const collector = createVercelProviderMetadataCollector(); + const { readable, writable } = new TransformStream(); + const writer = writable.getWriter(); + const metadata = { gateway: { generationId: 'gen_messages', marketCost: '0.12' } }; + let parsed = false; + const parsing = parseMessagesMicrodollarUsageFromStream( + readable, + 'fake-user-id', + undefined, + 'vercel', + 200, + false, + undefined, + { + captureRoutingDecision: jest.fn(), + captureVercelMetadata: value => collector.capture(value, 'messages'), + waitFor: jest.fn(), + } + ).then(result => { + parsed = true; + return result; + }); + let timeout: ReturnType | undefined; + try { + await writer.write( + new TextEncoder().encode( + `data: ${JSON.stringify({ type: 'message_start', message: { id: 'msg_1', model: 'model', usage: { input_tokens: 3, output_tokens: 1 }, provider_metadata: metadata } })}\n\ndata: [DONE]\n\n` + ) + ); + await Promise.race([ + collector.completed, + new Promise((_resolve, reject) => { + timeout = setTimeout(() => reject(new Error('Metadata capture waited for EOF')), 1000); + }), + ]); + expect(collector.snapshot()).toEqual(metadata); + expect(parsed).toBe(false); + } finally { + clearTimeout(timeout); + await writer.close(); + } + expect(await parsing).toMatchObject({ + cost_mUsd: 120_000, + inputTokens: 3, + outputTokens: 1, + hasError: false, + }); + }); + test.each([false, true])( 'retains Vercel start-frame cost and routing metadata with a final delta=%s', async hasDelta => { diff --git a/packages/web-shared/src/lib/ai-gateway/processUsage.messages.ts b/packages/web-shared/src/lib/ai-gateway/processUsage.messages.ts index 12a9937f4fd..62e13da2fe3 100644 --- a/packages/web-shared/src/lib/ai-gateway/processUsage.messages.ts +++ b/packages/web-shared/src/lib/ai-gateway/processUsage.messages.ts @@ -19,6 +19,7 @@ import { } from '@kilocode/web-shared/lib/ai-gateway/processUsage.shared'; import { isErrorFinishReason } from '@kilocode/web-shared/lib/ai-gateway/finishReason'; import type Anthropic from '@anthropic-ai/sdk'; +import type { AiProviderMetadataCapture } from '@kilocode/web-shared/lib/r2/kilo-routing-decision'; type MaybeHasVercelProviderMetadata = { provider_metadata?: VercelProviderMetaData; @@ -106,7 +107,8 @@ export async function parseMessagesMicrodollarUsageFromStream( provider: ProviderType, statusCode: number, usesUserCredentials: boolean, - requestStartedAt?: number + requestStartedAt?: number, + metadataCapture?: AiProviderMetadataCapture ): Promise { openrouterRequestSpan?.end(); const streamProcessingSpan = startInactiveSpan({ @@ -141,6 +143,7 @@ export async function parseMessagesMicrodollarUsageFromStream( } if (event.data === '[DONE]') { + metadataCapture?.captureVercelMetadata('[DONE]'); return; } @@ -155,6 +158,13 @@ export async function parseMessagesMicrodollarUsageFromStream( return; } + const eventProviderMetadata = ( + (json.type === 'message_start' ? json.message : json) as + | MaybeHasVercelProviderMetadata + | undefined + )?.provider_metadata; + metadataCapture?.captureVercelMetadata(json); + if (json.type === 'error') { reportedError = true; finish_reason = json.error?.error_type ?? json.error?.type ?? 'error'; @@ -177,12 +187,6 @@ export async function parseMessagesMicrodollarUsageFromStream( if (openRouterProvider) { inference_provider = openRouterProvider; } - const meta = (json.message as MaybeHasVercelProviderMetadata).provider_metadata; - if (meta) { - providerMetadata = meta; - model = getVercelResolvedModel(meta) ?? model; - inference_provider = meta.gateway?.routing?.finalProvider ?? inference_provider; - } } if ( @@ -196,12 +200,13 @@ export async function parseMessagesMicrodollarUsageFromStream( if (json.type === 'message_delta') { finish_reason = json.delta.stop_reason; deltaUsage = json.usage ?? deltaUsage; - const meta = (json as MaybeHasVercelProviderMetadata).provider_metadata; - if (meta) { - providerMetadata = meta; - model = getVercelResolvedModel(meta) ?? model; - inference_provider = meta.gateway?.routing?.finalProvider ?? inference_provider; - } + } + + if (eventProviderMetadata && (json.type === 'message_start' || json.type === 'message_delta')) { + providerMetadata = eventProviderMetadata; + model = getVercelResolvedModel(eventProviderMetadata) ?? model; + inference_provider = + eventProviderMetadata.gateway?.routing?.finalProvider ?? inference_provider; } }; @@ -244,14 +249,15 @@ export async function parseMessagesMicrodollarUsageFromStream( export function parseMessagesMicrodollarUsageFromString( fullResponse: string, statusCode: number, - provider: ProviderType + provider: ProviderType, + metadataCapture?: AiProviderMetadataCapture ): MicrodollarUsageStats { const responseJson = JSON.parse(fullResponse) as | (Anthropic.Messages.Message & MaybeHasVercelProviderMetadata & MaybeHasOpenRouterProvider) | null; - const usage = responseJson?.usage; const providerMetadata = responseJson?.provider_metadata ?? null; + metadataCapture?.captureVercelMetadata(responseJson); const inference_provider = providerMetadata?.gateway?.routing?.finalProvider ?? responseJson?.provider ?? null; diff --git a/packages/web-shared/src/lib/ai-gateway/processUsage.responses.ts b/packages/web-shared/src/lib/ai-gateway/processUsage.responses.ts index 939166fab99..3402fcd4ba5 100644 --- a/packages/web-shared/src/lib/ai-gateway/processUsage.responses.ts +++ b/packages/web-shared/src/lib/ai-gateway/processUsage.responses.ts @@ -19,6 +19,7 @@ import { selectUsageCostFields, } from '@kilocode/web-shared/lib/ai-gateway/processUsage.shared'; import { isErrorFinishReason } from '@kilocode/web-shared/lib/ai-gateway/finishReason'; +import type { AiProviderMetadataCapture } from '@kilocode/web-shared/lib/r2/kilo-routing-decision'; // OpenRouter adds cost fields to the standard Responses API usage object. // ref: https://openrouter.ai/docs/use-cases/usage-accounting#response-format @@ -86,7 +87,8 @@ export async function parseResponsesMicrodollarUsageFromStream( provider: ProviderType, statusCode: number, usesUserCredentials: boolean, - requestStartedAt?: number + requestStartedAt?: number, + metadataCapture?: AiProviderMetadataCapture ): Promise { openrouterRequestSpan?.end(); const streamProcessingSpan = startInactiveSpan({ @@ -120,6 +122,7 @@ export async function parseResponsesMicrodollarUsageFromStream( } if (event.data === '[DONE]') { + metadataCapture?.captureVercelMetadata('[DONE]'); return; } @@ -132,6 +135,9 @@ export async function parseResponsesMicrodollarUsageFromStream( return; } + const eventProviderMetadata = json.response?.provider_metadata; + metadataCapture?.captureVercelMetadata(json); + if ('error' in json && json.error) { reportedError = true; reportUpstreamStreamError({ @@ -158,10 +164,10 @@ export async function parseResponsesMicrodollarUsageFromStream( if (response.usage) { usage = response.usage as ResponsesApiUsage; } - const meta = response.provider_metadata; - if (meta) { - providerMetadata = meta; - inference_provider = meta.gateway?.routing?.finalProvider ?? inference_provider; + if (eventProviderMetadata) { + providerMetadata = eventProviderMetadata; + inference_provider = + eventProviderMetadata.gateway?.routing?.finalProvider ?? inference_provider; } finish_reason = response.status ?? finish_reason; } @@ -209,12 +215,13 @@ export async function parseResponsesMicrodollarUsageFromStream( export function parseResponsesMicrodollarUsageFromString( fullResponse: string, statusCode: number, - provider: ProviderType + provider: ProviderType, + metadataCapture?: AiProviderMetadataCapture ): MicrodollarUsageStats { const responseJson = JSON.parse(fullResponse) as ResponsesApiResponse | null; - const usage = responseJson?.usage; const providerMetadata = responseJson?.provider_metadata ?? null; + metadataCapture?.captureVercelMetadata(responseJson); const inference_provider = providerMetadata?.gateway?.routing?.finalProvider ?? null; diff --git a/packages/web-shared/src/lib/ai-gateway/processUsage.ts b/packages/web-shared/src/lib/ai-gateway/processUsage.ts index 86c82a01f39..f16e9f18b65 100644 --- a/packages/web-shared/src/lib/ai-gateway/processUsage.ts +++ b/packages/web-shared/src/lib/ai-gateway/processUsage.ts @@ -117,6 +117,7 @@ import { import { deliverBouncerUsageEventNow } from '@kilocode/web-shared/lib/bouncer/dispatch-usage-event-outbox'; import { enqueueBouncerUsageEvent } from '@kilocode/db/bouncer-usage-event-outbox'; import type { BouncerUsageEventEnqueue, UsageRecordRequest } from '@kilocode/usage-contracts'; +import type { AiProviderMetadataCapture } from '@kilocode/web-shared/lib/r2/kilo-routing-decision'; const posthogClient = PostHogClient(); @@ -1175,6 +1176,12 @@ export function countAndStoreUsage( ) { let usageStatsPromise: Promise = Promise.resolve(null); const { requestStartedAt } = usageContext; + const metadataCapture = + usageContext.provider === 'vercel' ? usageContext.providerMetadataCapture : undefined; + const contentType = clonedReponse.headers.get('content-type'); + const isStreaming = contentType?.includes('application/json') + ? false + : contentType?.includes('text/event-stream') || usageContext.isStreaming; const parseResponseText = async ( body: ReadableStream, @@ -1194,7 +1201,7 @@ export function countAndStoreUsage( if (clonedReponse.body) { if (usageContext.api_kind === 'responses') { - usageStatsPromise = usageContext.isStreaming + usageStatsPromise = isStreaming ? parseResponsesMicrodollarUsageFromStream( clonedReponse.body, usageContext.kiloUserId, @@ -1202,18 +1209,20 @@ export function countAndStoreUsage( usageContext.provider, clonedReponse.status, usesUserCredentials, - requestStartedAt + requestStartedAt, + metadataCapture ) : parseResponseText(clonedReponse.body, content => parseResponsesMicrodollarUsageFromString( content, clonedReponse.status, - usageContext.provider + usageContext.provider, + metadataCapture ) ); } if (usageContext.api_kind === 'chat_completions') { - usageStatsPromise = usageContext.isStreaming + usageStatsPromise = isStreaming ? parseMicrodollarUsageFromStream( clonedReponse.body, usageContext.kiloUserId, @@ -1221,19 +1230,21 @@ export function countAndStoreUsage( usageContext.provider, clonedReponse.status, usesUserCredentials, - requestStartedAt + requestStartedAt, + metadataCapture ) : parseResponseText(clonedReponse.body, content => parseMicrodollarUsageFromString( content, usageContext.kiloUserId, clonedReponse.status, - usageContext.provider + usageContext.provider, + metadataCapture ) ); } if (usageContext.api_kind === 'messages') { - usageStatsPromise = usageContext.isStreaming + usageStatsPromise = isStreaming ? parseMessagesMicrodollarUsageFromStream( clonedReponse.body, usageContext.kiloUserId, @@ -1241,18 +1252,21 @@ export function countAndStoreUsage( usageContext.provider, clonedReponse.status, usesUserCredentials, - requestStartedAt + requestStartedAt, + metadataCapture ) : parseResponseText(clonedReponse.body, content => parseMessagesMicrodollarUsageFromString( content, clonedReponse.status, - usageContext.provider + usageContext.provider, + metadataCapture ) ); } } + metadataCapture?.waitFor(usageStatsPromise); return usageStatsPromise.then(usageStats => processTokenData(usageStats, usageContext, servingProvider) ); @@ -1294,7 +1308,8 @@ export async function parseMicrodollarUsageFromStream( provider: ProviderType, statusCode: number, usesUserCredentials: boolean, - requestStartedAt?: number + requestStartedAt?: number, + metadataCapture?: AiProviderMetadataCapture ): Promise { // End the request span immediately as this function starts openrouterRequestSpan?.end(); @@ -1330,6 +1345,7 @@ export async function parseMicrodollarUsageFromStream( } if (event.data === '[DONE]') { + metadataCapture?.captureVercelMetadata('[DONE]'); return; } @@ -1342,6 +1358,10 @@ export async function parseMicrodollarUsageFromStream( return; } + const choice = json.choices?.[0]; + const chunkProviderMetadata = choice?.delta?.provider_metadata; + metadataCapture?.captureVercelMetadata(json); + if ('error' in json) { const error = json.error as OpenRouterError; reportedError = true; @@ -1366,8 +1386,6 @@ export async function parseMicrodollarUsageFromStream( model = json.model ?? model; messageId = json.id ?? messageId; usage = json.usage ?? usage; - const choice = json.choices?.[0]; - const chunkProviderMetadata = choice?.delta?.provider_metadata; if (chunkProviderMetadata) { vercelProviderMetadata = chunkProviderMetadata; } @@ -1422,13 +1440,17 @@ export function parseMicrodollarUsageFromString( fullResponse: string, kiloUserId: string, statusCode: number, - provider: ProviderType + provider: ProviderType, + metadataCapture?: AiProviderMetadataCapture ): MicrodollarUsageStats { const responseJson = JSON.parse(fullResponse) as | (OpenAI.Chat.Completions.ChatCompletion & MaybeHasOpenRouterUsage & MaybeHasVercelProviderMetaData) | null; + const choice = responseJson?.choices?.[0]; + const vercelProviderMetadata = choice?.message?.provider_metadata ?? null; + metadataCapture?.captureVercelMetadata(responseJson); if ( providerReportsIsByok(provider) && @@ -1440,9 +1462,7 @@ export function parseMicrodollarUsageFromString( extra: { responseJson }, }); } - const choice = responseJson?.choices?.[0]; const finish_reason = choice?.finish_reason ?? null; - const vercelProviderMetadata = choice?.message?.provider_metadata ?? null; const coreProps = { kiloUserId, messageId: responseJson?.id ?? null, diff --git a/packages/web-shared/src/lib/ai-gateway/processUsage.types.ts b/packages/web-shared/src/lib/ai-gateway/processUsage.types.ts index 0c135291932..342d8e762a4 100644 --- a/packages/web-shared/src/lib/ai-gateway/processUsage.types.ts +++ b/packages/web-shared/src/lib/ai-gateway/processUsage.types.ts @@ -5,6 +5,7 @@ import type { MicrodollarUsage, Organization } from '@kilocode/db'; import type { CustomPricingTiers } from '@kilocode/db/schema-types'; import type { GatewayApiKind } from '@kilocode/usage-contracts'; import type { OpenAI } from 'openai'; +import type { AiProviderMetadataCapture } from '@kilocode/web-shared/lib/r2/kilo-routing-decision'; export type OpenRouterUsage = { cost?: number; @@ -186,6 +187,7 @@ export type MicrodollarUsageContext = { * it: a second row with the same id would be dropped as a redelivery and go unbilled. */ usageId?: string; + providerMetadataCapture?: AiProviderMetadataCapture; }; /** diff --git a/packages/web-shared/src/lib/ai-gateway/provider-metadata-archive.test.ts b/packages/web-shared/src/lib/ai-gateway/provider-metadata-archive.test.ts new file mode 100644 index 00000000000..c78d3655f80 --- /dev/null +++ b/packages/web-shared/src/lib/ai-gateway/provider-metadata-archive.test.ts @@ -0,0 +1,733 @@ +import { PutObjectCommand } from '@aws-sdk/client-s3'; +import { setImmediate as nextTurn } from 'node:timers/promises'; +import { after } from 'next/server'; +import { countAndStoreUsage } from './processUsage'; +import type { MicrodollarUsageContext } from './processUsage.types'; +import { enqueueUsage } from './usage-publisher'; +import { fetchGeneration } from './providers/upstream-request'; +import { + KiloRoutingDecisionRecordSchema, + type ProviderRoutingDecision, +} from './providers/routing-decision'; +import { createAiProviderMetadataCapture } from '@kilocode/web-shared/lib/r2/kilo-routing-decision'; +import { getFraudDetectionHeaders } from '@kilocode/web-shared/lib/fraud-detection-headers'; +import { errorExceptInTest } from '@kilocode/web-shared/lib/utils.server'; +import type { FakeR2ClientModule } from '@kilocode/web-shared/tests/helpers/fake-r2.helper'; + +jest.mock('next/server', () => ({ + ...jest.requireActual('next/server'), + after: jest.fn(), +})); +jest.mock('@kilocode/web-shared/lib/config.server', () => ({ + ...jest.requireActual('@kilocode/web-shared/lib/config.server'), + USAGE_QUEUE_ONLY_ENABLED: true, +})); +jest.mock('./usage-publisher', () => ({ + enqueueUsage: jest.fn(async () => ({ kind: 'accepted' })), + getUsagePublisherExecutionRegion: () => 'fra1', +})); +jest.mock('./providers/upstream-request', () => ({ fetchGeneration: jest.fn() })); +jest.mock('./custom-pricing-config', () => ({ + getCustomPricingConfig: async () => ({ models: [] }), +})); +jest.mock('@kilocode/web-shared/lib/utils.server', () => ({ + ...jest.requireActual('@kilocode/web-shared/lib/utils.server'), + errorExceptInTest: jest.fn(), +})); +jest.mock('@kilocode/web-shared/lib/r2/create-client', () => + jest + .requireActual<{ + createFakeR2ClientModule: () => FakeR2ClientModule; + }>('@kilocode/web-shared/tests/helpers/fake-r2.helper') + .createFakeR2ClientModule() +); + +const { fakeR2 } = jest.requireMock( + '@kilocode/web-shared/lib/r2/create-client' +); +const mockedAfter = jest.mocked(after); +const mockedEnqueueUsage = jest.mocked(enqueueUsage); +const USAGE_ID = '7d4a3c4e-8f1b-4f63-9a0e-2b8f6f1d5c21'; +const GENERATION_ID = 'gen-provider-metadata-regression'; +const MODEL = 'synthetic/archive-regression'; +const BUCKET = 'test-provider-metadata'; +const API_KINDS = ['chat_completions', 'messages', 'responses'] as const; +type ApiKind = (typeof API_KINDS)[number]; +type JsonObject = Record; + +const metadata = { + gateway: { + generationId: 'vercel-generation-123', + cost: '0.002', + marketCost: '0.002', + routing: { finalProvider: 'Synthetic', futureRoutingField: ['opaque', { retained: true }] }, + }, + future_provider: { nullValue: null, values: [false, 7, 'unknown'], nested: { retained: true } }, + arbitrary_top_level: 'preserved verbatim', +}; + +function createRequest( + api_kind: ApiKind, + isStreaming: boolean, + provider: 'vercel' | 'openrouter' = 'vercel' +) { + const capture = createAiProviderMetadataCapture(USAGE_ID); + if (!capture) throw new Error('Expected configured provider metadata capture'); + const decision: ProviderRoutingDecision = { + requested_model: MODEL, + api_kind, + owner_type: 'user', + steps: [ + { step: 'custom-llm', outcome: 'skipped', reason: 'no-match' }, + { + step: 'kilo', + outcome: 'selected', + gateway: provider, + bounded_spending: true, + kilo_ignored_providers: ['synthetic-ignored'], + }, + ], + byok_candidates: [], + result: { + kind: 'provider', + route: 'kilo', + provider_type: provider, + byok_provider_ids: [], + payer: 'kilo', + is_free: false, + }, + }; + capture.captureRoutingDecision(decision); + const context: MicrodollarUsageContext = { + api_kind, + kiloUserId: 'oauth/archive-regression', + provider, + fraudHeaders: getFraudDetectionHeaders(new Headers()), + requested_model: MODEL, + is_free: false, + custom_llm_pricing: null, + promptInfo: { system_prompt_prefix: '', system_prompt_length: 0, user_prompt_prefix: '' }, + max_tokens: null, + has_middle_out_transform: null, + isStreaming, + prior_microdollar_usage: 0, + project_id: null, + status_code: null, + editor_name: null, + machine_id: null, + user_byok: false, + has_tools: false, + feature: null, + session_id: null, + mode: null, + auto_model: null, + reasoning_setting: null, + ttfb_ms: null, + usageId: USAGE_ID, + providerMetadataCapture: capture, + }; + return { capture, context, decision }; +} + +function jsonBody(apiKind: ApiKind, withUsage = true): JsonObject { + const common = { id: GENERATION_ID, model: MODEL, provider: 'Synthetic' }; + const cost = { cost: 0.002, is_byok: false }; + if (apiKind === 'chat_completions') { + return { + ...common, + choices: [{ message: { role: 'assistant', content: 'answer' }, finish_reason: 'stop' }], + ...(withUsage + ? { usage: { ...cost, prompt_tokens: 8, completion_tokens: 3, total_tokens: 11 } } + : {}), + }; + } + return { + ...common, + ...(apiKind === 'messages' + ? { type: 'message', content: [], stop_reason: 'end_turn' } + : { output: [], status: 'completed' }), + ...(withUsage ? { usage: { ...cost, input_tokens: 8, output_tokens: 3 } } : {}), + }; +} + +function initialFrames(apiKind: ApiKind, withUsage = true): JsonObject[] { + const body = jsonBody(apiKind, withUsage); + if (apiKind === 'messages') { + return [ + { type: 'message_start', message: body }, + { type: 'message_delta', delta: { stop_reason: 'end_turn' }, usage: body.usage }, + ]; + } + if (apiKind === 'responses') return [{ type: 'response.created', response: body }]; + return [{ ...body, choices: [{ delta: { content: 'answer' }, finish_reason: 'stop' }] }]; +} + +function metadataFrame(apiKind: ApiKind, value: unknown): JsonObject { + return { + ...(apiKind === 'messages' ? { type: 'ping' } : {}), + ...(apiKind === 'responses' ? { type: 'response.output_text.delta', delta: '' } : {}), + provider_metadata: value, + }; +} + +function sse(...frames: unknown[]): string { + return frames + .map(frame => `data: ${frame === '[DONE]' ? frame : JSON.stringify(frame)}\n\n`) + .join(''); +} + +function controlledResponse() { + const { promise, resolve } = Promise.withResolvers>(); + const response = new Response( + new ReadableStream({ start: controller => resolve(controller) }), + { headers: { 'content-type': 'text/event-stream' } } + ); + return { + response, + async write(text: string) { + (await promise).enqueue(new TextEncoder().encode(text)); + }, + async close() { + (await promise).close(); + }, + async abort() { + (await promise).error(new DOMException('Synthetic upstream interruption', 'ResponseAborted')); + }, + }; +} + +function runArchive() { + const callbacks = mockedAfter.mock.calls + .map(([work]) => work) + .filter((work): work is () => Promise => typeof work === 'function'); + expect(callbacks).toHaveLength(1); + return callbacks[0](); +} + +function expectArchive(decision: ProviderRoutingDecision, expectedMetadata?: unknown) { + expect(fakeR2.send).toHaveBeenCalledTimes(1); + expect(fakeR2.send).toHaveBeenCalledWith(expect.any(PutObjectCommand)); + const key = `${BUCKET}/${USAGE_ID}.json`; + expect([...fakeR2.objects.keys()]).toEqual([key]); + const record = JSON.parse(fakeR2.objects.get(key) ?? ''); + expect(record).toEqual({ + schema_version: 1, + usage_id: USAGE_ID, + created_at: expect.any(String), + ...decision, + ...(expectedMetadata === undefined ? {} : { vercel: { metadata: expectedMetadata } }), + }); + expect(KiloRoutingDecisionRecordSchema.parse(record)).toEqual(record); +} + +function expectBilling(hasError = false, cost = 2000) { + expect(mockedEnqueueUsage).toHaveBeenCalledTimes(1); + const payload = mockedEnqueueUsage.mock.calls[0]?.[0]; + expect(payload?.core).toEqual( + expect.objectContaining({ + id: USAGE_ID, + model: MODEL, + has_error: hasError, + cost, + input_tokens: 8, + output_tokens: 3, + }) + ); + expect(payload?.metadata.id).toBe(USAGE_ID); + expect(payload).not.toHaveProperty('providerMetadataCapture'); + expect(payload?.metadata).not.toHaveProperty('providerMetadataCapture'); +} + +beforeEach(() => { + jest.replaceProperty(process, 'env', { + ...process.env, + R2_ACCOUNT_ID: 'test-metadata-account', + R2_AI_PROVIDER_METADATA_BUCKET_NAME: BUCKET, + R2_AI_PROVIDER_METADATA_ACCESS_KEY_ID: 'test-metadata-access-key', + R2_AI_PROVIDER_METADATA_SECRET_ACCESS_KEY: 'test-metadata-secret-key', + }); + mockedAfter.mockReset(); + mockedEnqueueUsage.mockClear(); + jest.mocked(fetchGeneration).mockReset(); + jest.mocked(errorExceptInTest).mockClear(); + fakeR2.objects.clear(); + jest.spyOn(fakeR2, 'send'); + jest.spyOn(console, 'debug').mockImplementation(() => {}); + jest.spyOn(console, 'info').mockImplementation(() => {}); +}); + +afterEach(() => jest.restoreAllMocks()); + +describe('raw provider metadata archive with usage accounting', () => { + test.each(API_KINDS)( + '%s JSON combines routing and opaque metadata in one deferred put', + async apiKind => { + const { context, decision } = createRequest(apiKind, false); + await countAndStoreUsage( + Response.json({ ...jsonBody(apiKind), provider_metadata: metadata }), + context, + undefined + ); + + expect(fakeR2.send).not.toHaveBeenCalled(); + expect(fakeR2.objects.size).toBe(0); + expectBilling(); + await runArchive(); + expectArchive(decision, metadata); + } + ); + + test.each(API_KINDS)( + '%s JSON accepts camelCase metadata without changing billing', + async apiKind => { + const { context, decision } = createRequest(apiKind, false); + await countAndStoreUsage( + Response.json({ ...jsonBody(apiKind), providerMetadata: metadata }), + context, + undefined + ); + + await runArchive(); + expectArchive(decision, metadata); + expectBilling(); + } + ); + + test.each(['provider_metadata', 'providerMetadata'])( + 'chat JSON visits every choice message with %s and replaces rather than merges metadata', + async alias => { + const { context, decision } = createRequest('chat_completions', false); + await countAndStoreUsage( + Response.json({ + ...jsonBody('chat_completions'), + provider_metadata: { rootOnly: true }, + choices: [ + { + message: { content: 'first', [alias]: { firstChoiceOnly: true } }, + finish_reason: 'stop', + }, + { message: { content: 'second', [alias]: metadata }, finish_reason: 'stop' }, + ], + }), + context, + undefined + ); + + await runArchive(); + expectArchive(decision, metadata); + expectBilling(); + } + ); + + const streamLocations: { name: string; apiKind: ApiKind; frame: JsonObject }[] = [ + { + name: 'chat first choice delta', + apiKind: 'chat_completions', + frame: { choices: [{ delta: { provider_metadata: metadata } }] }, + }, + { + name: 'chat every choice delta and camelCase', + apiKind: 'chat_completions', + frame: { + choices: [ + { delta: { provider_metadata: { stale: true } } }, + { delta: { providerMetadata: metadata } }, + ], + }, + }, + { + name: 'chat root after choice deltas', + apiKind: 'chat_completions', + frame: { + choices: [{ delta: { provider_metadata: { stale: true } } }], + providerMetadata: metadata, + }, + }, + { + name: 'chat root without choices', + apiKind: 'chat_completions', + frame: { provider_metadata: metadata }, + }, + { + name: 'chat root with empty choices', + apiKind: 'chat_completions', + frame: { choices: [], providerMetadata: metadata }, + }, + { + name: 'Messages message_start nested message', + apiKind: 'messages', + frame: { + type: 'message_start', + message: { ...jsonBody('messages'), provider_metadata: metadata }, + }, + }, + { + name: 'Messages message_start camelCase nested message', + apiKind: 'messages', + frame: { + type: 'message_start', + message: { ...jsonBody('messages'), providerMetadata: metadata }, + }, + }, + { + name: 'Messages message_start top envelope after nested message', + apiKind: 'messages', + frame: { + type: 'message_start', + message: { ...jsonBody('messages'), provider_metadata: { stale: true } }, + provider_metadata: metadata, + }, + }, + { + name: 'Messages message_delta top envelope', + apiKind: 'messages', + frame: { + type: 'message_delta', + delta: { stop_reason: 'end_turn' }, + providerMetadata: metadata, + }, + }, + { + name: 'Messages non-usage top envelope', + apiKind: 'messages', + frame: { type: 'ping', provider_metadata: metadata }, + }, + ...[ + 'response.created', + 'response.in_progress', + 'response.completed', + 'response.incomplete', + 'response.failed', + 'error', + ].flatMap(type => [ + { + name: `Responses ${type} nested response`, + apiKind: 'responses' as const, + frame: { type, response: { ...jsonBody('responses'), providerMetadata: metadata } }, + }, + { + name: `Responses ${type} top envelope after nested response`, + apiKind: 'responses' as const, + frame: { + type, + response: { ...jsonBody('responses'), provider_metadata: { stale: true } }, + provider_metadata: metadata, + }, + }, + ]), + ]; + + test.each(streamLocations)( + '$name is archived from the raw stream', + async ({ apiKind, frame }) => { + const { context, decision } = createRequest(apiKind, true); + await countAndStoreUsage( + new Response(sse(...initialFrames(apiKind), frame)), + context, + undefined + ); + + expect(fakeR2.send).not.toHaveBeenCalled(); + await runArchive(); + expectArchive(decision, metadata); + expect(mockedEnqueueUsage).toHaveBeenCalledTimes(1); + } + ); + + test.each(API_KINDS)('%s retains latest metadata across separate events', async apiKind => { + const { context, decision } = createRequest(apiKind, true); + await countAndStoreUsage( + new Response( + sse( + metadataFrame(apiKind, { staleOnly: true, gateway: { staleGatewayOnly: true } }), + ...initialFrames(apiKind), + metadataFrame(apiKind, metadata), + metadataFrame(apiKind, null), + metadataFrame(apiKind, ['not an object']), + '[DONE]' + ) + ), + context, + undefined + ); + + await runArchive(); + expectArchive(decision, metadata); + expectBilling(); + }); + + test.each(API_KINDS)('%s archives metadata without usage or a response id', async apiKind => { + const { context, decision } = createRequest(apiKind, false); + const body = jsonBody(apiKind, false); + delete body.id; + await countAndStoreUsage( + Response.json({ ...body, provider_metadata: metadata }), + context, + undefined + ); + + await runArchive(); + expectArchive(decision, metadata); + expect(mockedEnqueueUsage.mock.calls[0]?.[0].core).toEqual( + expect.objectContaining({ id: USAGE_ID, input_tokens: 0, output_tokens: 0 }) + ); + }); + + test.each(API_KINDS)('%s SSE archives metadata without usage', async apiKind => { + const { context, decision } = createRequest(apiKind, true); + await countAndStoreUsage( + new Response( + sse(metadataFrame(apiKind, metadata), ...initialFrames(apiKind, false), '[DONE]') + ), + context, + undefined + ); + + await runArchive(); + expectArchive(decision, metadata); + expect(mockedEnqueueUsage.mock.calls[0]?.[0].core).toEqual( + expect.objectContaining({ id: USAGE_ID, input_tokens: 0, output_tokens: 0 }) + ); + }); + + test.each(API_KINDS)('%s preserves a route-only archive with an absent body', async apiKind => { + const { context, decision } = createRequest(apiKind, false); + await expect(countAndStoreUsage(new Response(null), context, undefined)).resolves.toBeNull(); + + await runArchive(); + expectArchive(decision); + expect(mockedEnqueueUsage).not.toHaveBeenCalled(); + }); + + test.each(API_KINDS)( + '%s preserves a route-only archive when JSON parsing fails', + async apiKind => { + const { context, decision } = createRequest(apiKind, false); + await expect( + countAndStoreUsage(new Response('{invalid JSON'), context, undefined) + ).rejects.toThrow(SyntaxError); + + await expect(runArchive()).resolves.toBeUndefined(); + expectArchive(decision); + expect(mockedEnqueueUsage).not.toHaveBeenCalled(); + } + ); + + test.each( + API_KINDS.flatMap(apiKind => + ['malformed', 'transport abort'].map(failure => ({ apiKind, failure })) + ) + )( + '$apiKind retains partial metadata and bills hasError after $failure', + async ({ apiKind, failure }) => { + const { context, decision, capture } = createRequest(apiKind, true); + const captured = jest.spyOn(capture, 'captureVercelMetadata'); + const upstream = controlledResponse(); + const accounting = countAndStoreUsage(upstream.response, context, undefined); + const partial = metadataFrame(apiKind, metadata); + await upstream.write(sse(...initialFrames(apiKind), partial)); + await nextTurn(); + expect(captured).toHaveBeenCalledWith(partial); + if (failure === 'malformed') { + await upstream.write('data: {malformed JSON\n\n'); + await upstream.close(); + } else { + await upstream.abort(); + } + + await expect(accounting).resolves.toEqual( + expect.objectContaining({ status: 'queued', usageId: USAGE_ID }) + ); + await runArchive(); + expectArchive(decision, metadata); + expectBilling(true); + } + ); + + test('client cancellation does not discard late metadata from the accounting response clone', async () => { + const { context, decision } = createRequest('chat_completions', true); + const upstream = controlledResponse(); + const accounting = countAndStoreUsage(upstream.response.clone(), context, undefined); + await upstream.write(sse(...initialFrames('chat_completions'))); + const client = upstream.response.body?.getReader(); + if (!client) throw new Error('Expected client response body'); + await client.read(); + const cancelled = client.cancel('client disconnected'); + await upstream.write(sse(metadataFrame('chat_completions', metadata), '[DONE]')); + await upstream.close(); + + await accounting; + await cancelled; + await runArchive(); + expectArchive(decision, metadata); + expectBilling(); + }); + + test.each(API_KINDS)( + '%s retains complete metadata before an unterminated final SSE event', + async apiKind => { + const { context, decision } = createRequest(apiKind, true); + await countAndStoreUsage( + new Response( + `${sse(...initialFrames(apiKind), metadataFrame(apiKind, metadata))}data: ${JSON.stringify(metadataFrame(apiKind, { incomplete: true }))}\n` + ), + context, + undefined + ); + + await runArchive(); + expectArchive(decision, metadata); + expectBilling(); + } + ); + + test.each( + API_KINDS.flatMap(apiKind => + (['vercel', 'openrouter'] as const).map(provider => ({ apiKind, provider })) + ) + )( + '$apiKind $provider protocol terminal does not leave an archive waiting for raw EOF', + async ({ apiKind, provider }) => { + const { context, decision } = createRequest(apiKind, true, provider); + const upstream = controlledResponse(); + const accounting = countAndStoreUsage(upstream.response, context, undefined); + const upload = runArchive(); + await nextTurn(); + if (provider === 'vercel') expect(fakeR2.send).not.toHaveBeenCalled(); + const terminal = + apiKind === 'chat_completions' + ? '[DONE]' + : apiKind === 'messages' + ? { type: 'message_stop', providerMetadata: metadata } + : { + type: 'response.completed', + response: jsonBody(apiKind), + providerMetadata: metadata, + }; + let timeout: ReturnType | undefined; + try { + await upstream.write( + sse( + ...initialFrames(apiKind), + metadataFrame( + apiKind, + apiKind === 'chat_completions' ? metadata : { beforeTerminalOnly: true } + ), + terminal, + metadataFrame(apiKind, { afterTerminalMustNotOverwrite: true }) + ) + ); + await Promise.race([ + upload, + new Promise((_resolve, reject) => { + timeout = setTimeout(() => reject(new Error('Archive waited for physical EOF')), 1000); + }), + ]); + expectArchive(decision, provider === 'vercel' ? metadata : undefined); + expect(mockedEnqueueUsage).not.toHaveBeenCalled(); + } finally { + clearTimeout(timeout); + await upstream.close(); + await accounting; + await upload; + } + expectBilling(); + } + ); + + test.each(API_KINDS)( + '%s stream:true request parses an HTTP JSON response and preserves its billed usage', + async apiKind => { + const { context, decision } = createRequest(apiKind, true); + await countAndStoreUsage( + Response.json({ ...jsonBody(apiKind), provider_metadata: metadata }), + context, + undefined + ); + + await runArchive(); + expectArchive(decision, metadata); + expectBilling(); + expect(mockedEnqueueUsage.mock.calls[0]?.[0].metadata.streamed).toBe(false); + } + ); + + test.each(API_KINDS)( + '%s stream:true request archives an HTTP JSON error with no usage', + async apiKind => { + const { context, decision } = createRequest(apiKind, true); + await countAndStoreUsage( + Response.json( + { error: { message: 'Synthetic error' }, providerMetadata: metadata }, + { status: 502 } + ), + context, + undefined + ); + + await runArchive(); + expectArchive(decision, metadata); + expect(mockedEnqueueUsage.mock.calls[0]?.[0].core).toEqual( + expect.objectContaining({ cost: 0, has_error: true }) + ); + expect(mockedEnqueueUsage.mock.calls[0]?.[0].metadata).toEqual( + expect.objectContaining({ status_code: 502, streamed: false }) + ); + } + ); + + test.each(API_KINDS)( + '%s stream:false request still harvests an HTTP SSE response', + async apiKind => { + const { context, decision } = createRequest(apiKind, false); + await countAndStoreUsage( + new Response(sse(...initialFrames(apiKind), metadataFrame(apiKind, metadata), '[DONE]'), { + headers: { 'content-type': 'text/event-stream' }, + }), + context, + undefined + ); + + await runArchive(); + expectArchive(decision, metadata); + expectBilling(); + expect(mockedEnqueueUsage.mock.calls[0]?.[0].metadata.streamed).toBe(true); + } + ); + + test('ignores Vercel-looking metadata from a different serving provider', async () => { + const { context, decision } = createRequest('chat_completions', false, 'openrouter'); + await countAndStoreUsage( + Response.json({ ...jsonBody('chat_completions'), provider_metadata: metadata }), + context, + undefined + ); + + await runArchive(); + expectArchive(decision); + expectBilling(); + }); + + test('an archive upload failure does not reject accounting or change its billed record', async () => { + const { context } = createRequest('chat_completions', false); + jest.mocked(fakeR2.send).mockRejectedValueOnce(new Error('Synthetic object storage failure')); + const accounting = countAndStoreUsage( + Response.json({ ...jsonBody('chat_completions'), provider_metadata: metadata }), + context, + undefined + ); + const upload = runArchive(); + + await expect(accounting).resolves.toEqual( + expect.objectContaining({ status: 'queued', usageId: USAGE_ID }) + ); + await expect(upload).resolves.toBeUndefined(); + expectBilling(); + expect(fakeR2.send).toHaveBeenCalledTimes(1); + expect(fakeR2.objects.size).toBe(0); + expect(errorExceptInTest).toHaveBeenCalledWith( + '[kilo-routing-decision] failed to archive metadata', + expect.objectContaining({ usageId: USAGE_ID, stage: 'upload' }) + ); + }); +}); diff --git a/packages/web-shared/src/lib/ai-gateway/providers/routing-decision.ts b/packages/web-shared/src/lib/ai-gateway/providers/routing-decision.ts index 169c477e86d..37c4e81b000 100644 --- a/packages/web-shared/src/lib/ai-gateway/providers/routing-decision.ts +++ b/packages/web-shared/src/lib/ai-gateway/providers/routing-decision.ts @@ -190,12 +190,17 @@ const ProviderRoutingDecisionSchema = z.object({ export type ProviderRoutingDecision = z.infer; -/** A `kilo/.json` object in the AI provider metadata bucket. */ +/** A `.json` object in the AI provider metadata bucket. */ export const KiloRoutingDecisionRecordSchema = ProviderRoutingDecisionSchema.extend({ schema_version: z.literal(1), /** The `microdollar_usage` and `microdollar_usage_metadata` id of the request. */ usage_id: z.uuid(), created_at: z.iso.datetime(), + vercel: z + .object({ + metadata: z.record(z.string(), z.unknown()), + }) + .optional(), }); export type KiloRoutingDecisionRecord = z.infer; diff --git a/packages/web-shared/src/lib/ai-gateway/rewriteModelResponse.shared.ts b/packages/web-shared/src/lib/ai-gateway/rewriteModelResponse.shared.ts index 5a5178918f4..5becc5b539f 100644 --- a/packages/web-shared/src/lib/ai-gateway/rewriteModelResponse.shared.ts +++ b/packages/web-shared/src/lib/ai-gateway/rewriteModelResponse.shared.ts @@ -16,7 +16,6 @@ import type { import { applyUsagePricing } from '@kilocode/web-shared/lib/ai-gateway/usage-cost'; import { errorExceptInTest, logExceptInTest } from '@kilocode/web-shared/lib/utils.server'; import { withRequestId } from '@kilocode/web-shared/lib/ai-gateway/request-id'; -import type { VercelProviderMetadataCapture } from '@kilocode/web-shared/lib/r2/vercel-provider-metadata'; import type { createParser } from 'eventsource-parser'; import { NextResponse } from 'next/server'; @@ -44,7 +43,6 @@ export type ResponseUsageContext = Pick< 'requested_model' | 'provider' | 'is_free' | 'user_byok' | 'api_kind' | 'custom_llm_pricing' > & { custom_pricing_config: CustomPricingConfig; - captureProviderMetadata?: VercelProviderMetadataCapture; }; export type ResponseUsage = { @@ -85,7 +83,6 @@ function isModelRewriteExempt(model: string): boolean { } export function rewriteResponseMetadata(value: ResponseMetadata, context: ResponseUsageContext) { - context.captureProviderMetadata?.(value.provider_metadata ?? value.providerMetadata, value.id); if (value.model !== undefined && !isModelRewriteExempt(context.requested_model)) { value.model = context.requested_model; } diff --git a/packages/web-shared/src/lib/ai-gateway/rewriteModelResponse.test.ts b/packages/web-shared/src/lib/ai-gateway/rewriteModelResponse.test.ts index e1f126a890a..666e2bcd366 100644 --- a/packages/web-shared/src/lib/ai-gateway/rewriteModelResponse.test.ts +++ b/packages/web-shared/src/lib/ai-gateway/rewriteModelResponse.test.ts @@ -2044,7 +2044,7 @@ async function rewriteUsageFixture( return { payload: envelope.message ?? envelope.response ?? envelope, objects }; } -describe('Vercel provider metadata archival', () => { +describe('Vercel provider metadata rewriting', () => { const { fakeR2 } = jest.requireMock( '@kilocode/web-shared/lib/r2/create-client' ); @@ -2067,6 +2067,7 @@ describe('Vercel provider metadata archival', () => { }; beforeEach(() => { + jest.spyOn(fakeR2, 'send'); jest.replaceProperty(process, 'env', { ...process.env, R2_AI_PROVIDER_METADATA_BUCKET_NAME: 'test-ai-provider-metadata', @@ -2078,22 +2079,16 @@ describe('Vercel provider metadata archival', () => { afterEach(() => jest.restoreAllMocks()); - async function uploadAndReadMetadata() { + function expectNoArchiveIO() { + expect(mockedAfter).not.toHaveBeenCalled(); + expect(fakeR2.send).not.toHaveBeenCalled(); expect(fakeR2.objects.size).toBe(0); - expect(mockedAfter).toHaveBeenCalledTimes(1); - const callback = mockedAfter.mock.calls[0]?.[0]; - if (typeof callback !== 'function') throw new Error('Expected an after callback'); - await callback(); - expect(fakeR2.objects.size).toBe(1); - return JSON.parse( - fakeR2.objects.get('test-ai-provider-metadata/vercel/gen_metadata_test.json') ?? '' - ); } describe.each(usageShapes)('$kind', shape => { describe.each([false, true])('stream=%s', stream => { test.each(['provider_metadata', 'providerMetadata'])( - 'archives %s without exposing it or storing the response body', + 'strips %s without scheduling archive IO', async metadataKey => { const payload = { id: 'gen_metadata_test', @@ -2138,14 +2133,14 @@ describe('Vercel provider metadata archival', () => { expect(output).not.toContain('provider_metadata'); expect(output).not.toContain('providerMetadata'); - expect(await uploadAndReadMetadata()).toEqual(metadata); + expectNoArchiveIO(); } ); }); }); test.each([false, true])( - 'archives nested Chat Completions metadata without usage, stream=%s', + 'strips nested Chat Completions metadata without usage or archive IO, stream=%s', async stream => { const choice = { content: 'answer', provider_metadata: metadata }; const body = { @@ -2166,31 +2161,11 @@ describe('Vercel provider metadata archival', () => { const output = stream ? await readOutputStream(result) : JSON.stringify(await result.json()); expect(output).not.toContain('provider_metadata'); - expect(await uploadAndReadMetadata()).toEqual(metadata); - } - ); - - test.each(['openrouter', 'custom'] as const)( - 'does not archive metadata from %s', - async providerId => { - await rewriteModelResponse({ - response: jsonResponse({ id: 'gen_metadata_test', provider_metadata: metadata }), - model: 'requested-model', - providerId, - kind: 'chat_completions', - isFree: false, - userByok: false, - customLlmPricing: null, - logging: makeLogging(), - responseTransforms: null, - }); - - expect(mockedAfter).not.toHaveBeenCalled(); - expect(fakeR2.objects.size).toBe(0); + expectNoArchiveIO(); } ); - test('archives observed metadata even when the upstream stream disconnects', async () => { + test('does not perform archive IO for observed metadata before an upstream disconnect', async () => { const result = await rewriteModelResponse({ response: failingResponse( 'text/event-stream', @@ -2208,7 +2183,7 @@ describe('Vercel provider metadata archival', () => { }); expect(await readOutputStream(result)).toContain('upstream_disconnect'); - expect(await uploadAndReadMetadata()).toEqual(metadata); + expectNoArchiveIO(); }); }); diff --git a/packages/web-shared/src/lib/ai-gateway/rewriteModelResponse.ts b/packages/web-shared/src/lib/ai-gateway/rewriteModelResponse.ts index 78877568d53..bbe6901e22a 100644 --- a/packages/web-shared/src/lib/ai-gateway/rewriteModelResponse.ts +++ b/packages/web-shared/src/lib/ai-gateway/rewriteModelResponse.ts @@ -13,7 +13,6 @@ import { db } from '@kilocode/web-shared/lib/drizzle'; import { logExceptInTest } from '@kilocode/web-shared/lib/utils.server'; import { sanitizeJsonbValue } from '@kilocode/web-shared/lib/sanitize-jsonb'; import { uploadApiRequestLogBlobs } from '@kilocode/web-shared/lib/r2/api-request-log'; -import { createVercelProviderMetadataCapture } from '@kilocode/web-shared/lib/r2/vercel-provider-metadata'; import { after, type NextResponse } from 'next/server'; import type { RequestLogCapture, @@ -224,8 +223,6 @@ export async function rewriteModelResponse({ user_byok: userByok, custom_llm_pricing: customLlmPricing, custom_pricing_config: await getCustomPricingConfig(), - captureProviderMetadata: - providerId === 'vercel' ? createVercelProviderMetadataCapture() : undefined, }; console.debug('[rewriteModelResponse] rewriting response for %s', model); diff --git a/packages/web-shared/src/lib/r2/ai-provider-metadata-bucket.ts b/packages/web-shared/src/lib/r2/ai-provider-metadata-bucket.ts index 0e2d4da3925..c0f83237aa6 100644 --- a/packages/web-shared/src/lib/r2/ai-provider-metadata-bucket.ts +++ b/packages/web-shared/src/lib/r2/ai-provider-metadata-bucket.ts @@ -50,7 +50,7 @@ export function scheduleAiProviderMetadataArchive({ logSource: string; logContext: ArchiveLogContext; key: string; - body: () => string; + body: () => string | Promise; }): void { const reportFailure = (stage: 'registration' | 'upload', error: unknown) => errorExceptInTest(`[${logSource}] failed to archive metadata`, { @@ -71,12 +71,13 @@ export function scheduleAiProviderMetadataArchive({ console.debug(`[${logSource}] scheduling archive`, { ...logContext, bucket }); after(async () => { try { + const serializedBody = await body(); client ??= createR2Client({ accessKeyId, secretAccessKey }); await client.send( new PutObjectCommand({ Bucket: bucket, Key: key, - Body: body(), + Body: serializedBody, ContentType: 'application/json', }) ); diff --git a/packages/web-shared/src/lib/r2/kilo-routing-decision.test.ts b/packages/web-shared/src/lib/r2/kilo-routing-decision.test.ts index 24aec0c28f9..2f0dcb8db8e 100644 --- a/packages/web-shared/src/lib/r2/kilo-routing-decision.test.ts +++ b/packages/web-shared/src/lib/r2/kilo-routing-decision.test.ts @@ -1,9 +1,10 @@ import { after } from 'next/server'; +import { PutObjectCommand, S3ServiceException } from '@aws-sdk/client-s3'; import { KiloRoutingDecisionRecordSchema, type ProviderRoutingDecision, } from '@kilocode/web-shared/lib/ai-gateway/providers/routing-decision'; -import { archiveKiloRoutingDecision } from '@kilocode/web-shared/lib/r2/kilo-routing-decision'; +import { createAiProviderMetadataCapture } from '@kilocode/web-shared/lib/r2/kilo-routing-decision'; import { errorExceptInTest } from '@kilocode/web-shared/lib/utils.server'; import type { FakeR2ClientModule } from '@kilocode/web-shared/tests/helpers/fake-r2.helper'; @@ -33,7 +34,7 @@ const decision: ProviderRoutingDecision = { { step: 'kilo', outcome: 'selected', - gateway: 'openrouter', + gateway: 'vercel', bounded_spending: false, kilo_ignored_providers: [], }, @@ -42,7 +43,7 @@ const decision: ProviderRoutingDecision = { result: { kind: 'provider', route: 'kilo', - provider_type: 'openrouter', + provider_type: 'vercel', byok_provider_ids: [], payer: 'kilo', is_free: false, @@ -55,6 +56,12 @@ async function runScheduledUpload() { await callback(); } +function archivedRecord() { + return KiloRoutingDecisionRecordSchema.parse( + JSON.parse(fakeR2.objects.get(`test-provider-metadata/${USAGE_ID}.json`) ?? '') + ); +} + beforeEach(() => { jest.replaceProperty(process, 'env', { ...process.env, @@ -67,6 +74,7 @@ beforeEach(() => { jest.spyOn(console, 'debug').mockImplementation(() => {}); jest.mocked(errorExceptInTest).mockClear(); fakeR2.objects.clear(); + jest.spyOn(fakeR2, 'send'); }); afterEach(() => jest.restoreAllMocks()); @@ -78,8 +86,9 @@ test.each([ 'R2_AI_PROVIDER_METADATA_SECRET_ACCESS_KEY', ])('does not archive without %s', name => { delete process.env[name]; - archiveKiloRoutingDecision(USAGE_ID, decision); + expect(createAiProviderMetadataCapture(USAGE_ID)).toBeUndefined(); expect(mockedAfter).not.toHaveBeenCalled(); + expect(fakeR2.send).not.toHaveBeenCalled(); expect(console.debug).toHaveBeenCalledTimes(1); expect(console.debug).toHaveBeenCalledWith( '[kilo-routing-decision] capture disabled: missing configuration', @@ -92,10 +101,13 @@ test.each([ ); }); -test('uploads the decision after the response, keyed by the usage id', async () => { - archiveKiloRoutingDecision(USAGE_ID, decision); +test('uploads a route-only decision after the response, keyed by the usage id', async () => { + const capture = createAiProviderMetadataCapture(USAGE_ID); + expect(mockedAfter).not.toHaveBeenCalled(); + capture?.captureRoutingDecision(decision); expect(mockedAfter).toHaveBeenCalledTimes(1); expect(fakeR2.objects.size).toBe(0); + expect(fakeR2.send).not.toHaveBeenCalled(); expect(console.debug).toHaveBeenLastCalledWith('[kilo-routing-decision] scheduling archive', { usageId: USAGE_ID, bucket: 'test-provider-metadata', @@ -107,12 +119,8 @@ test('uploads the decision after the response, keyed by the usage id', async () accessKeyId: 'test-metadata-access-key', secretAccessKey: 'test-metadata-secret-key', }); - expect([...fakeR2.objects.keys()]).toEqual([`test-provider-metadata/kilo/${USAGE_ID}.json`]); - expect( - KiloRoutingDecisionRecordSchema.parse( - JSON.parse(fakeR2.objects.get(`test-provider-metadata/kilo/${USAGE_ID}.json`) ?? '') - ) - ).toEqual({ + expect([...fakeR2.objects.keys()]).toEqual([`test-provider-metadata/${USAGE_ID}.json`]); + expect(archivedRecord()).toEqual({ schema_version: 1, usage_id: USAGE_ID, created_at: expect.stringMatching(/^\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}\.\d{3}Z$/), @@ -124,9 +132,169 @@ test('uploads the decision after the response, keyed by the usage id', async () }); }); +test('combines full metadata fidelity and routing fields in exactly one PutObject', async () => { + const metadata = { + gateway: { + generationId: 'vercel-generation-123', + routing: { + modelAttempts: [ + { + providerAttempts: [ + { provider: 'openai', success: false, error: { code: 'rate_limit' }, duration: 250 }, + { provider: 'amazon-bedrock', success: true, duration: 600 }, + ], + }, + ], + }, + additionalField: { retained: true }, + }, + upstream: { additionalField: 'retained', array: [null, false, 4, 'unknown'] }, + }; + const capture = createAiProviderMetadataCapture(USAGE_ID); + capture?.captureRoutingDecision(decision); + capture?.captureRoutingDecision({ ...decision, requested_model: 'ignored-repeat' }); + capture?.captureVercelMetadata({ provider_metadata: { superseded: true } }); + capture?.captureVercelMetadata({ providerMetadata: metadata }); + capture?.captureVercelMetadata({ choices: [] }); + expect(mockedAfter).toHaveBeenCalledTimes(1); + expect(fakeR2.send).not.toHaveBeenCalled(); + + await runScheduledUpload(); + + expect(fakeR2.send).toHaveBeenCalledTimes(1); + expect(fakeR2.send).toHaveBeenCalledWith(expect.any(PutObjectCommand)); + expect(archivedRecord()).toEqual({ + schema_version: 1, + usage_id: USAGE_ID, + created_at: expect.any(String), + ...decision, + vercel: { metadata }, + }); + expect([...fakeR2.objects.keys()]).toEqual([`test-provider-metadata/${USAGE_ID}.json`]); +}); + +test('does not register an archive from metadata or parser settlement without routing', async () => { + const capture = createAiProviderMetadataCapture(USAGE_ID); + capture?.captureVercelMetadata({ provider_metadata: { captured: true } }); + capture?.waitFor(Promise.resolve()); + await Promise.resolve(); + expect(mockedAfter).not.toHaveBeenCalled(); + expect(fakeR2.send).not.toHaveBeenCalled(); +}); + +test('defers serialization until the raw parser installed after routing settles', async () => { + const parsing = Promise.withResolvers(); + const capture = createAiProviderMetadataCapture(USAGE_ID); + const toJSON = jest.fn(() => ({ complete: true })); + capture?.captureRoutingDecision(decision); + capture?.waitFor(parsing.promise); + capture?.captureVercelMetadata({ provider_metadata: { toJSON } }); + expect(toJSON).not.toHaveBeenCalled(); + + const upload = runScheduledUpload(); + await Promise.resolve(); + expect(fakeR2.send).not.toHaveBeenCalled(); + expect(toJSON).not.toHaveBeenCalled(); + capture?.captureVercelMetadata({}); + parsing.resolve({ billing: 'unchanged' }); + await upload; + + expect(toJSON).toHaveBeenCalledTimes(1); + expect(archivedRecord().vercel).toEqual({ metadata: { complete: true } }); + await expect(parsing.promise).resolves.toEqual({ billing: 'unchanged' }); +}); + +test('normalizes rejected parsing without changing its rejection or dropping partial metadata', async () => { + const parsing = Promise.withResolvers(); + const capture = createAiProviderMetadataCapture(USAGE_ID); + capture?.captureRoutingDecision(decision); + capture?.waitFor(parsing.promise); + capture?.captureVercelMetadata({ provider_metadata: { partial: true } }); + const upload = runScheduledUpload(); + const error = new Error('sensitive parsing failure'); + parsing.reject(error); + + await expect(upload).resolves.toBeUndefined(); + await expect(parsing.promise).rejects.toBe(error); + expect(archivedRecord().vercel).toEqual({ metadata: { partial: true } }); + expect(errorExceptInTest).not.toHaveBeenCalled(); +}); + +test('handles an already rejected parsing promise without creating an unhandled rejection', async () => { + const capture = createAiProviderMetadataCapture(USAGE_ID); + capture?.captureRoutingDecision(decision); + capture?.waitFor(Promise.reject(new Error('parsing failed'))); + + await expect(runScheduledUpload()).resolves.toBeUndefined(); + expect(archivedRecord()).not.toHaveProperty('vercel'); +}); + +test.each([ + { api_kind: 'chat_completions', terminal: '[DONE]' }, + { api_kind: 'chat_completions', terminal: { error: { message: 'failed' } } }, + { api_kind: 'messages', terminal: { type: 'message_stop' } }, + { api_kind: 'messages', terminal: { type: 'error' } }, + { api_kind: 'responses', terminal: { type: 'response.completed' } }, + { api_kind: 'responses', terminal: { type: 'response.failed' } }, + { api_kind: 'responses', terminal: { type: 'response.incomplete' } }, + { api_kind: 'responses', terminal: { type: 'error' } }, +] as const)( + 'archives on $api_kind terminal $terminal while raw parsing remains pending', + async ({ api_kind, terminal }) => { + const parsing = Promise.withResolvers(); + const capture = createAiProviderMetadataCapture(USAGE_ID); + capture?.captureRoutingDecision({ ...decision, api_kind }); + capture?.waitFor(parsing.promise); + capture?.captureVercelMetadata({ provider_metadata: { before: true } }); + const upload = runScheduledUpload(); + await Promise.resolve(); + expect(fakeR2.send).not.toHaveBeenCalled(); + + capture?.captureVercelMetadata( + typeof terminal === 'string' + ? terminal + : { ...terminal, provider_metadata: { terminal: true } } + ); + capture?.captureVercelMetadata({ provider_metadata: { ignored: true } }); + await upload; + + expect(archivedRecord().vercel).toEqual( + typeof terminal === 'string' + ? { metadata: { before: true } } + : { metadata: { terminal: true } } + ); + expect(fakeR2.send).toHaveBeenCalledTimes(1); + parsing.reject(new Error('later accounting failure')); + await Promise.resolve(); + expect(errorExceptInTest).not.toHaveBeenCalled(); + } +); + +test('omits the Vercel envelope when metadata is invalid', async () => { + const capture = createAiProviderMetadataCapture(USAGE_ID); + capture?.captureRoutingDecision(decision); + for (const metadata of [undefined, null, false, 42, 'invalid', []]) { + capture?.captureVercelMetadata({ provider_metadata: metadata }); + } + await runScheduledUpload(); + + expect(archivedRecord()).not.toHaveProperty('vercel'); +}); + +test('archives metadata without a response id', async () => { + const capture = createAiProviderMetadataCapture(USAGE_ID); + capture?.captureRoutingDecision(decision); + capture?.captureVercelMetadata({ provider_metadata: {} }); + await runScheduledUpload(); + + expect(archivedRecord().vercel).toEqual({ metadata: {} }); +}); + test('contains upload failures without logging the decision or the upstream error text', async () => { jest.spyOn(fakeR2, 'send').mockRejectedValue(new Error('sensitive upstream error text')); - archiveKiloRoutingDecision(USAGE_ID, decision); + const capture = createAiProviderMetadataCapture(USAGE_ID); + capture?.captureRoutingDecision(decision); + capture?.captureVercelMetadata({ provider_metadata: { private: 'private metadata' } }); await expect(runScheduledUpload()).resolves.toBeUndefined(); @@ -140,15 +308,69 @@ test('contains upload failures without logging the decision or the upstream erro } ); expect(JSON.stringify(jest.mocked(errorExceptInTest).mock.calls)).not.toMatch( - /sensitive upstream|kimi|test-metadata-secret-key/ + /sensitive upstream|kimi|private metadata|test-metadata-secret-key/ ); }); -test('reports a scheduling failure without throwing', () => { +test('includes safe S3 diagnostics without logging the raw service error', async () => { + jest.spyOn(fakeR2, 'send').mockRejectedValue( + new S3ServiceException({ + name: 'AccessDenied', + $fault: 'client', + $metadata: { httpStatusCode: 403, requestId: 'test-r2-request', attempts: 1 }, + message: 'sensitive upstream error text', + }) + ); + const capture = createAiProviderMetadataCapture(USAGE_ID); + capture?.captureRoutingDecision(decision); + capture?.captureVercelMetadata({ provider_metadata: { private: 'private metadata' } }); + + await expect(runScheduledUpload()).resolves.toBeUndefined(); + + expect(errorExceptInTest).toHaveBeenCalledWith( + '[kilo-routing-decision] failed to archive metadata', + { + usageId: USAGE_ID, + bucket: 'test-provider-metadata', + stage: 'upload', + errorType: 'AccessDenied', + httpStatusCode: 403, + requestId: 'test-r2-request', + attempts: 1, + } + ); + expect(JSON.stringify(jest.mocked(errorExceptInTest).mock.calls)).not.toMatch( + /sensitive upstream|private metadata|test-metadata-secret-key|test-metadata-access-key/ + ); +}); + +test('contains deferred serialization failures without sending an object', async () => { + const metadata: Record = {}; + metadata.circular = metadata; + const capture = createAiProviderMetadataCapture(USAGE_ID); + capture?.captureRoutingDecision(decision); + capture?.captureVercelMetadata({ provider_metadata: metadata }); + + await expect(runScheduledUpload()).resolves.toBeUndefined(); + + expect(fakeR2.send).not.toHaveBeenCalled(); + expect(errorExceptInTest).toHaveBeenCalledWith( + '[kilo-routing-decision] failed to archive metadata', + expect.objectContaining({ stage: 'upload', errorType: 'TypeError' }) + ); +}); + +test('reports a scheduling failure once without throwing or retrying registration', () => { mockedAfter.mockImplementation(() => { throw new Error('after unavailable'); }); - expect(() => archiveKiloRoutingDecision(USAGE_ID, decision)).not.toThrow(); + const capture = createAiProviderMetadataCapture(USAGE_ID); + expect(() => capture?.captureRoutingDecision(decision)).not.toThrow(); + expect(() => capture?.captureRoutingDecision(decision)).not.toThrow(); + capture?.captureVercelMetadata({ provider_metadata: {} }); + expect(mockedAfter).toHaveBeenCalledTimes(1); + expect(errorExceptInTest).toHaveBeenCalledTimes(1); + expect(fakeR2.send).not.toHaveBeenCalled(); expect(errorExceptInTest).toHaveBeenCalledWith( '[kilo-routing-decision] failed to archive metadata', { diff --git a/packages/web-shared/src/lib/r2/kilo-routing-decision.ts b/packages/web-shared/src/lib/r2/kilo-routing-decision.ts index 69eaa09c814..07b08681d52 100644 --- a/packages/web-shared/src/lib/r2/kilo-routing-decision.ts +++ b/packages/web-shared/src/lib/r2/kilo-routing-decision.ts @@ -6,28 +6,55 @@ import { getAiProviderMetadataBucket, scheduleAiProviderMetadataArchive, } from '@kilocode/web-shared/lib/r2/ai-provider-metadata-bucket'; +import { createVercelProviderMetadataCollector } from '@kilocode/web-shared/lib/r2/vercel-provider-metadata'; const LOG_SOURCE = 'kilo-routing-decision'; -/** Archives the gateway's provider routing decision as `kilo/.json` after the response. */ -export function archiveKiloRoutingDecision( - usageId: string, - decision: ProviderRoutingDecision -): void { +export type AiProviderMetadataCapture = { + captureRoutingDecision: (decision: ProviderRoutingDecision) => void; + captureVercelMetadata: (value: unknown) => void; + waitFor: (promise: Promise) => void; +}; + +export function createAiProviderMetadataCapture( + usageId: string +): AiProviderMetadataCapture | undefined { const target = getAiProviderMetadataBucket(LOG_SOURCE); - if (!target) return; - const createdAt = new Date().toISOString(); - scheduleAiProviderMetadataArchive({ - target, - logSource: LOG_SOURCE, - logContext: { usageId }, - key: `kilo/${usageId}.json`, - body: () => - JSON.stringify({ + if (!target) return undefined; + + const vercel = createVercelProviderMetadataCollector(); + let record: KiloRoutingDecisionRecord | undefined; + let parsing: Promise | undefined; + + return { + captureRoutingDecision(decision) { + if (record) return; + record = { schema_version: 1, usage_id: usageId, - created_at: createdAt, + created_at: new Date().toISOString(), ...decision, - } satisfies KiloRoutingDecisionRecord), - }); + }; + scheduleAiProviderMetadataArchive({ + target, + logSource: LOG_SOURCE, + logContext: { usageId }, + key: `${usageId}.json`, + body: async () => { + if (parsing) await Promise.race([parsing, vercel.completed]); + const metadata = vercel.snapshot(); + return JSON.stringify({ ...record, ...(metadata ? { vercel: { metadata } } : {}) }); + }, + }); + }, + captureVercelMetadata(value) { + vercel.capture(value, record?.api_kind); + }, + waitFor(promise) { + parsing = promise.then( + () => {}, + () => {} + ); + }, + }; } diff --git a/packages/web-shared/src/lib/r2/vercel-provider-metadata.test.ts b/packages/web-shared/src/lib/r2/vercel-provider-metadata.test.ts index 8352ceadd06..634318f6d8d 100644 --- a/packages/web-shared/src/lib/r2/vercel-provider-metadata.test.ts +++ b/packages/web-shared/src/lib/r2/vercel-provider-metadata.test.ts @@ -1,232 +1,212 @@ -import { after } from 'next/server'; -import { S3ServiceException } from '@aws-sdk/client-s3'; -import { createVercelProviderMetadataCapture } from '@kilocode/web-shared/lib/r2/vercel-provider-metadata'; -import { errorExceptInTest } from '@kilocode/web-shared/lib/utils.server'; -import type { FakeR2ClientModule } from '@kilocode/web-shared/tests/helpers/fake-r2.helper'; - -jest.mock('next/server', () => ({ after: jest.fn() })); -jest.mock('@kilocode/web-shared/lib/utils.server', () => ({ errorExceptInTest: jest.fn() })); -jest.mock('@kilocode/web-shared/lib/r2/create-client', () => - jest - .requireActual<{ - createFakeR2ClientModule: () => FakeR2ClientModule; - }>('@kilocode/web-shared/tests/helpers/fake-r2.helper') - .createFakeR2ClientModule() -); +import { createVercelProviderMetadataCollector } from '@kilocode/web-shared/lib/r2/vercel-provider-metadata'; + +test('walks Chat JSON root then every choice message and preserves the raw object', () => { + const collector = createVercelProviderMetadataCollector(); + const json = { + provider_metadata: { root: true }, + choices: [ + { message: { provider_metadata: { first: true } } }, + { message: { providerMetadata: { final: { unknown: [null, false, 'retained'] } } } }, + ], + }; + const original = structuredClone(json); -const { fakeR2 } = jest.requireMock( - '@kilocode/web-shared/lib/r2/create-client' -); -const mockedAfter = jest.mocked(after); - -async function runScheduledUpload() { - const callback = mockedAfter.mock.calls[0]?.[0]; - if (typeof callback !== 'function') throw new Error('Expected an after callback'); - await callback(); -} - -beforeEach(() => { - jest.replaceProperty(process, 'env', { - ...process.env, - R2_ACCOUNT_ID: 'test-metadata-account', - R2_AI_PROVIDER_METADATA_BUCKET_NAME: 'test-provider-metadata', - R2_AI_PROVIDER_METADATA_ACCESS_KEY_ID: 'test-metadata-access-key', - R2_AI_PROVIDER_METADATA_SECRET_ACCESS_KEY: 'test-metadata-secret-key', - }); - mockedAfter.mockReset(); - jest.spyOn(console, 'debug').mockImplementation(() => {}); - jest.mocked(errorExceptInTest).mockClear(); - fakeR2.objects.clear(); + collector.capture(json, 'chat_completions'); + + expect(collector.snapshot()).toEqual({ final: { unknown: [null, false, 'retained'] } }); + expect(json).toEqual(original); }); -afterEach(() => jest.restoreAllMocks()); - -test.each([ - 'R2_ACCOUNT_ID', - 'R2_AI_PROVIDER_METADATA_BUCKET_NAME', - 'R2_AI_PROVIDER_METADATA_ACCESS_KEY_ID', - 'R2_AI_PROVIDER_METADATA_SECRET_ACCESS_KEY', -])('disables capture without %s', name => { - delete process.env[name]; - expect(createVercelProviderMetadataCapture()).toBeUndefined(); - expect(mockedAfter).not.toHaveBeenCalled(); - expect(console.debug).toHaveBeenCalledTimes(1); - expect(console.debug).toHaveBeenCalledWith( - '[vercel-provider-metadata] capture disabled: missing configuration', +test('walks every Chat stream delta before the root', () => { + const collector = createVercelProviderMetadataCollector(); + collector.capture( { - accountConfigured: name !== 'R2_ACCOUNT_ID', - bucketConfigured: name !== 'R2_AI_PROVIDER_METADATA_BUCKET_NAME', - accessKeyIdConfigured: name !== 'R2_AI_PROVIDER_METADATA_ACCESS_KEY_ID', - secretAccessKeyConfigured: name !== 'R2_AI_PROVIDER_METADATA_SECRET_ACCESS_KEY', - } + providerMetadata: { root: true }, + choices: [ + { delta: { provider_metadata: { first: true } } }, + { delta: { provider_metadata: { second: true } } }, + ], + }, + 'chat_completions' ); + + expect(collector.snapshot()).toEqual({ root: true }); }); -test('uploads the latest metadata once, after the response, keyed by generation ID', async () => { - const capture = createVercelProviderMetadataCapture(); - const metadata = { - gateway: { - routing: { - modelAttempts: [ - { - providerAttempts: [ - { provider: 'openai', success: false, error: { code: 'rate_limit' }, duration: 250 }, - { provider: 'amazon-bedrock', success: true, duration: 600 }, - ], - }, - ], - }, - additionalField: { retained: true }, +test('retains metadata from a later Chat choice without usage or top-level metadata', () => { + const collector = createVercelProviderMetadataCollector(); + collector.capture( + { + choices: [ + { delta: { provider_metadata: { first: true } } }, + { delta: { providerMetadata: { second: true } } }, + ], }, - upstream: { additionalField: 'retained' }, - }; + 'chat_completions' + ); - capture?.({ gateway: { cost: '0.1' } }, 'gen_test'); - capture?.(metadata); - capture?.(undefined, 'gen_test'); - expect(mockedAfter).toHaveBeenCalledTimes(1); - expect(fakeR2.objects.size).toBe(0); - expect(console.debug).toHaveBeenCalledTimes(1); - expect(console.debug).toHaveBeenLastCalledWith('[vercel-provider-metadata] scheduling archive', { - generationId: 'gen_test', - bucket: 'test-provider-metadata', - }); - - await runScheduledUpload(); - - expect(fakeR2.credentials).toEqual({ - accessKeyId: 'test-metadata-access-key', - secretAccessKey: 'test-metadata-secret-key', - }); - expect([...fakeR2.objects]).toEqual([ - ['test-provider-metadata/vercel/gen_test.json', JSON.stringify(metadata)], - ]); - expect(console.debug).toHaveBeenCalledTimes(2); - expect(console.debug).toHaveBeenLastCalledWith('[vercel-provider-metadata] archived metadata', { - generationId: 'gen_test', - bucket: 'test-provider-metadata', - }); + expect(collector.snapshot()).toEqual({ second: true }); }); -test('retains metadata received before the generation ID', async () => { - const capture = createVercelProviderMetadataCapture(); - const metadata = { gateway: { cost: '0.1' } }; - capture?.(metadata); - capture?.(metadata); - expect(mockedAfter).not.toHaveBeenCalled(); - expect(console.debug).toHaveBeenCalledTimes(1); - expect(console.debug).toHaveBeenCalledWith( - '[vercel-provider-metadata] waiting for metadata and generation ID', - { generationId: undefined, metadataReceived: true, generationIdReceived: false } +test('walks Messages start message before root and captures top-level metadata on every event', () => { + const collector = createVercelProviderMetadataCollector(); + collector.capture( + { + type: 'message_start', + message: { provider_metadata: { nested: true } }, + providerMetadata: { root: true }, + }, + 'messages' ); + expect(collector.snapshot()).toEqual({ root: true }); - capture?.(undefined, 'gen_late_id'); - await runScheduledUpload(); + for (const type of ['content_block_start', 'content_block_delta', 'message_delta']) { + collector.capture({ type, provider_metadata: { latest: type } }, 'messages'); + expect(collector.snapshot()).toEqual({ latest: type }); + } + collector.capture({ type: 'message_stop', providerMetadata: { stopped: true } }, 'messages'); + expect(collector.snapshot()).toEqual({ stopped: true }); +}); - expect(fakeR2.objects.get('test-provider-metadata/vercel/gen_late_id.json')).toBe( - JSON.stringify(metadata) +test('ignores Messages nested message outside message_start', () => { + const collector = createVercelProviderMetadataCollector(); + collector.capture( + { + type: 'message_delta', + message: { provider_metadata: { invalid: true } }, + }, + 'messages' ); + expect(collector.snapshot()).toBeUndefined(); }); -test('does not schedule an upload without an ID and valid metadata', () => { - const capture = createVercelProviderMetadataCapture(); - for (const metadata of [undefined, null, false, 42, 'invalid', []]) { - capture?.(metadata, 'gen_test'); +test('walks every Responses nested response before top-level metadata', () => { + const collector = createVercelProviderMetadataCollector(); + for (const type of ['response.created', 'response.in_progress', 'response.output_text.delta']) { + collector.capture( + { + type, + response: { providerMetadata: { nested: true } }, + provider_metadata: { latest: type }, + }, + 'responses' + ); + expect(collector.snapshot()).toEqual({ latest: type }); } - expect(mockedAfter).not.toHaveBeenCalled(); - expect(console.debug).toHaveBeenCalledTimes(1); - expect(console.debug).toHaveBeenCalledWith( - '[vercel-provider-metadata] waiting for metadata and generation ID', - { generationId: 'gen_test', metadataReceived: false, generationIdReceived: true } + collector.capture( + { + type: 'response.completed', + response: { provider_metadata: { terminal: { retained: true } } }, + }, + 'responses' ); - - const captureWithoutId = createVercelProviderMetadataCapture(); - captureWithoutId?.({ gateway: {} }, ''); - expect(mockedAfter).not.toHaveBeenCalled(); - expect(console.debug).toHaveBeenCalledTimes(2); -}); - -test('does not log empty stream chunks', () => { - const capture = createVercelProviderMetadataCapture(); - capture?.(undefined); - capture?.(undefined); - expect(console.debug).not.toHaveBeenCalled(); + expect(collector.snapshot()).toEqual({ terminal: { retained: true } }); }); -test('contains upload failures without logging metadata or upstream error text', async () => { - jest.spyOn(fakeR2, 'send').mockRejectedValue(new Error('sensitive upstream error text')); - const capture = createVercelProviderMetadataCapture(); - capture?.({ gateway: { additionalField: 'private metadata' } }, 'gen_failed'); +test.each(['messages', 'responses'] as const)( + 'captures %s non-stream JSON root metadata', + apiKind => { + const collector = createVercelProviderMetadataCollector(); + collector.capture({ provider_metadata: { unknown: 'retained' } }, apiKind); + expect(collector.snapshot()).toEqual({ unknown: 'retained' }); + } +); - await expect(runScheduledUpload()).resolves.toBeUndefined(); +test('replaces metadata instead of merging and retains it across metadata-free frames', () => { + const collector = createVercelProviderMetadataCollector(); + collector.capture({ provider_metadata: { first: true } }, 'chat_completions'); + collector.capture({ providerMetadata: { second: true } }, 'chat_completions'); + collector.capture({ choices: [] }, 'chat_completions'); + expect(collector.snapshot()).toEqual({ second: true }); +}); - expect(errorExceptInTest).toHaveBeenCalledWith( - '[vercel-provider-metadata] failed to archive metadata', - { - generationId: 'gen_failed', - bucket: 'test-provider-metadata', - stage: 'upload', - errorType: 'Error', - } - ); - expect(JSON.stringify(jest.mocked(errorExceptInTest).mock.calls)).not.toMatch( - /sensitive upstream|private metadata/ - ); - expect(console.debug).toHaveBeenCalledTimes(1); - expect(JSON.stringify(jest.mocked(console.debug).mock.calls)).not.toMatch( - /sensitive upstream|private metadata/ +test('prefers valid snake_case metadata to valid camelCase metadata', () => { + const collector = createVercelProviderMetadataCollector(); + collector.capture( + { provider_metadata: { snake: true }, providerMetadata: { camel: true } }, + 'chat_completions' ); + expect(collector.snapshot()).toEqual({ snake: true }); }); -test('includes safe S3 diagnostics without logging the raw service error', async () => { - jest.spyOn(fakeR2, 'send').mockRejectedValue( - new S3ServiceException({ - name: 'AccessDenied', - $fault: 'client', - $metadata: { httpStatusCode: 403, requestId: 'test-r2-request', attempts: 1 }, - message: 'sensitive upstream error text', - }) - ); - const capture = createVercelProviderMetadataCapture(); - capture?.({ gateway: { additionalField: 'private metadata' } }, 'gen_denied'); +test.each([undefined, null, false, 42, 'invalid', []])( + 'ignores invalid metadata %p without clearing the previous object', + invalid => { + const collector = createVercelProviderMetadataCollector(); + collector.capture({ provider_metadata: { retained: true } }, 'chat_completions'); + collector.capture( + { provider_metadata: invalid, providerMetadata: invalid }, + 'chat_completions' + ); + expect(collector.snapshot()).toEqual({ retained: true }); + } +); - await expect(runScheduledUpload()).resolves.toBeUndefined(); +test.each([undefined, null])( + 'falls back to valid camelCase metadata when snake_case is %p', + invalid => { + const collector = createVercelProviderMetadataCollector(); + collector.capture( + { provider_metadata: invalid, providerMetadata: { valid: true } }, + 'chat_completions' + ); + expect(collector.snapshot()).toEqual({ valid: true }); + } +); - expect(errorExceptInTest).toHaveBeenCalledWith( - '[vercel-provider-metadata] failed to archive metadata', +test.each([false, 42, 'invalid', []])( + 'does not fall back to camelCase when non-null snake_case is invalid: %p', + invalid => { + const collector = createVercelProviderMetadataCollector(); + collector.capture({ provider_metadata: { retained: true } }, 'chat_completions'); + collector.capture( + { provider_metadata: invalid, providerMetadata: { ignored: true } }, + 'chat_completions' + ); + expect(collector.snapshot()).toEqual({ retained: true }); + } +); + +test('ignores non-object frames, unrelated subtrees, and invalid choice nodes', () => { + const collector = createVercelProviderMetadataCollector(); + for (const value of [undefined, null, true, 42, 'not-json', []]) { + collector.capture(value, 'chat_completions'); + } + collector.capture( { - generationId: 'gen_denied', - bucket: 'test-provider-metadata', - stage: 'upload', - errorType: 'AccessDenied', - httpStatusCode: 403, - requestId: 'test-r2-request', - attempts: 1, - } - ); - expect(JSON.stringify(jest.mocked(errorExceptInTest).mock.calls)).not.toMatch( - /sensitive upstream|private metadata|test-metadata-secret-key|test-metadata-access-key/ + output: [{ provider_metadata: { ignored: true } }], + choices: [null, 42, {}, { delta: [], message: false }], + }, + 'chat_completions' ); + expect(collector.snapshot()).toBeUndefined(); }); -test('attempts registration and reports scheduling failures only once per response', () => { - mockedAfter.mockImplementation(() => { - throw new Error('after unavailable'); - }); - const capture = createVercelProviderMetadataCapture(); - expect(() => capture?.({ gateway: {} }, 'gen_test')).not.toThrow(); - expect(() => capture?.(undefined, 'gen_test')).not.toThrow(); - expect(() => capture?.({ gateway: { cost: '0.1' } }, 'gen_test')).not.toThrow(); - expect(mockedAfter).toHaveBeenCalledTimes(1); - expect(errorExceptInTest).toHaveBeenCalledTimes(1); - expect(errorExceptInTest).toHaveBeenCalledWith( - '[vercel-provider-metadata] failed to archive metadata', - { - generationId: 'gen_test', - bucket: 'test-provider-metadata', - stage: 'registration', - errorType: 'Error', - } +test('does not treat a finish_reason or mismatched protocol event as metadata completion', async () => { + const collector = createVercelProviderMetadataCollector(); + const completed = jest.fn(); + void collector.completed.then(completed); + collector.capture({ choices: [{ finish_reason: 'stop' }] }, 'chat_completions'); + collector.capture({ type: 'error' }, 'chat_completions'); + collector.capture({ type: 'response.completed' }, 'chat_completions'); + collector.capture({ type: 'message_stop' }, 'responses'); + await Promise.resolve(); + expect(completed).not.toHaveBeenCalled(); + + collector.capture({ provider_metadata: { retained: true } }, 'chat_completions'); + collector.capture('[DONE]', 'chat_completions'); + await collector.completed; + expect(completed).toHaveBeenCalledTimes(1); + expect(collector.snapshot()).toEqual({ retained: true }); +}); + +test('captures a terminal error envelope before freezing and ignores all subsequent metadata', async () => { + const collector = createVercelProviderMetadataCollector(); + collector.capture( + { type: 'error', provider_metadata: { error: { unknown: 'retained' } } }, + 'messages' ); - expect(console.debug).toHaveBeenCalledTimes(1); + collector.capture({ provider_metadata: { ignored: true } }, 'messages'); + await collector.completed; + expect(collector.snapshot()).toEqual({ error: { unknown: 'retained' } }); }); diff --git a/packages/web-shared/src/lib/r2/vercel-provider-metadata.ts b/packages/web-shared/src/lib/r2/vercel-provider-metadata.ts index 3708f38564e..74c0eb103e8 100644 --- a/packages/web-shared/src/lib/r2/vercel-provider-metadata.ts +++ b/packages/web-shared/src/lib/r2/vercel-provider-metadata.ts @@ -1,48 +1,60 @@ -import { - getAiProviderMetadataBucket, - scheduleAiProviderMetadataArchive, -} from '@kilocode/web-shared/lib/r2/ai-provider-metadata-bucket'; +import type { GatewayChatApiKind } from '@kilocode/web-shared/lib/ai-gateway/providers/types'; -export type VercelProviderMetadataCapture = (metadata: unknown, generationId?: string) => void; +function isRecord(value: unknown): value is Record { + return typeof value === 'object' && value !== null && !Array.isArray(value); +} -const LOG_SOURCE = 'vercel-provider-metadata'; +export function createVercelProviderMetadataCollector() { + let metadata: Record | undefined; + let terminal = false; + const { promise: completed, resolve: complete } = Promise.withResolvers(); -export function createVercelProviderMetadataCapture(): VercelProviderMetadataCapture | undefined { - const target = getAiProviderMetadataBucket(LOG_SOURCE); - if (!target) return undefined; + const captureNode = (value: unknown) => { + if (!isRecord(value)) return; + const candidate = value.provider_metadata ?? value.providerMetadata; + if (isRecord(candidate)) metadata = candidate; + }; - let capturedMetadata: object | undefined; - let capturedGenerationId: string | undefined; - let waitingLogged = false; - let registrationAttempted = false; + return { + completed, + capture(value: unknown, apiKind?: GatewayChatApiKind) { + if (terminal) return; + if (value === '[DONE]') { + terminal = true; + complete(); + return; + } + if (!isRecord(value)) return; - return (metadata, generationId) => { - if (typeof generationId === 'string' && generationId.length > 0) { - capturedGenerationId = generationId; - } - if (typeof metadata === 'object' && metadata !== null && !Array.isArray(metadata)) { - capturedMetadata = metadata; - } - if (registrationAttempted) return; - if (!capturedMetadata || !capturedGenerationId) { - if (!waitingLogged && (metadata !== undefined || capturedGenerationId)) { - console.debug(`[${LOG_SOURCE}] waiting for metadata and generation ID`, { - generationId: capturedGenerationId, - metadataReceived: Boolean(capturedMetadata), - generationIdReceived: Boolean(capturedGenerationId), - }); - waitingLogged = true; + if (apiKind === 'messages' || (apiKind === undefined && value.type === 'message_start')) { + if (value.type === 'message_start') captureNode(value.message); + captureNode(value); + } else if (apiKind === 'responses' || (apiKind === undefined && isRecord(value.response))) { + captureNode(value.response); + captureNode(value); + } else { + const choices = Array.isArray(value.choices) ? value.choices : []; + for (const choice of choices) { + if (isRecord(choice)) captureNode(choice.delta); + } + captureNode(value); + for (const choice of choices) { + if (isRecord(choice)) captureNode(choice.message); + } } - return; - } - registrationAttempted = true; - scheduleAiProviderMetadataArchive({ - target, - logSource: LOG_SOURCE, - logContext: { generationId: capturedGenerationId }, - key: `vercel/${capturedGenerationId}.json`, - body: () => JSON.stringify(capturedMetadata), - }); + terminal = + (apiKind === 'chat_completions' && value.error != null) || + (apiKind === 'messages' && (value.type === 'message_stop' || value.type === 'error')) || + (apiKind === 'responses' && + (value.type === 'response.completed' || + value.type === 'response.failed' || + value.type === 'response.incomplete' || + value.type === 'error')); + if (terminal) complete(); + }, + snapshot(): Record | undefined { + return metadata; + }, }; }