[AISAGE-733] perf(gateway): combine provider metadata uploads

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/<generation_id>.json` and gateway routing decisions use `kilo/<microdollar_usage id>.json`. [SERVER]
+- `R2_AI_PROVIDER_METADATA_BUCKET_NAME` - Optional R2 bucket for AI provider metadata; `<microdollar_usage id>.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]`
 

packages/web-shared/src/lib/ai-gateway/handlers/llm-proxy.test.ts

@@ -892,17 +892,186 @@ 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),
       ...decision,
     });
   });
 
+  it('archives Vercel metadata captured during accounting with the routing decision', async () => {
+    const { fakeR2 } = jest.requireMock<FakeR2ClientModule>(
+      '@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<FakeR2ClientModule>(
+      '@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');
 

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),

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<object>('@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<Uint8Array, Uint8Array>();
+    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<typeof setTimeout> | 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<never>((_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 => {

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<MicrodollarUsageStats> {
   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;