refactor(ai-gateway): fold request logging into rewriteModelResponse

apps/web/src/app/api/openrouter/[...path]/route.test.ts

@@ -14,7 +14,7 @@ import { fetchEfficientAutoDecision } from '@/lib/ai-gateway/auto-routing-decisi
 import { logMicrodollarUsage } from '@/lib/ai-gateway/processUsage';
 import { applyResolvedAutoModel } from '@/lib/ai-gateway/auto-model/resolution';
 import { getDirectByokModel } from '@/lib/ai-gateway/providers/direct-byok';
-import { handleRequestLogging } from '@/lib/ai-gateway/handleRequestLogging';
+import { rewriteModelResponse } from '@/lib/rewriteModelResponse';
 
 jest.mock('next/server', () => {
   return {
@@ -55,9 +55,13 @@ jest.mock('@/lib/ai-gateway/o11y/api-metrics.server', () => ({
   getToolsAvailable: jest.fn(() => false),
   getToolsUsed: jest.fn(() => false),
 }));
-jest.mock('@/lib/ai-gateway/handleRequestLogging', () => ({
-  handleRequestLogging: jest.fn(),
-}));
+jest.mock('@/lib/rewriteModelResponse', () => {
+  const actual = jest.requireActual('@/lib/rewriteModelResponse');
+  return {
+    ...actual,
+    rewriteModelResponse: jest.fn(async () => null),
+  };
+});
 jest.mock('@/lib/ai-gateway/llm-proxy-helpers', () => {
   const actual = jest.requireActual('@/lib/ai-gateway/llm-proxy-helpers');
   return {
@@ -96,7 +100,7 @@ const mockedFetchEfficientAutoDecision = jest.mocked(fetchEfficientAutoDecision)
 const mockedLogMicrodollarUsage = jest.mocked(logMicrodollarUsage);
 const mockedApplyResolvedAutoModel = jest.mocked(applyResolvedAutoModel);
 const mockedGetDirectByokModel = jest.mocked(getDirectByokModel);
-const mockedHandleRequestLogging = jest.mocked(handleRequestLogging);
+const mockedRewriteModelResponse = jest.mocked(rewriteModelResponse);
 
 const provider = {
   id: 'openrouter',
@@ -255,7 +259,11 @@ describe('POST /api/openrouter/v1/chat/completions rules-engine actions', () =>
     );
 
     expect(response.status).toBe(200);
-    expect(mockedHandleRequestLogging).toHaveBeenCalledWith(
+    expect(mockedRewriteModelResponse).toHaveBeenCalledWith(
+      expect.anything(),
+      expect.anything(),
+      expect.anything(),
+      expect.anything(),
       expect.objectContaining({ vercel_request_id: 'iad1::iad1::request-id' })
     );
   });

apps/web/src/app/api/openrouter/[...path]/route.ts

@@ -56,7 +56,7 @@ import {
 import { ProxyErrorType } from '@/lib/proxy-error-types';
 import { getBalanceAndOrgSettings } from '@/lib/organizations/organization-usage';
 import { isDataCollectionExplicitlyDisallowed } from '@/lib/ai-gateway/providers/openrouter/types';
-import { rewriteModelResponse } from '@/lib/rewriteModelResponse';
+import { rewriteModelResponse, logUnrewrittenResponse } from '@/lib/rewriteModelResponse';
 import {
   createAnonymousContext,
   isAnonymousContext,
@@ -69,7 +69,6 @@ import {
   checkPromotionLimit,
 } from '@/lib/free-model-rate-limiter';
 import { PROMOTION_MAX_REQUESTS, PROMOTION_WINDOW_HOURS } from '@/lib/constants';
-import { handleRequestLogging } from '@/lib/ai-gateway/handleRequestLogging';
 import {
   classifyAbuse,
   awaitClassifyAbuse,
@@ -957,16 +956,13 @@ export async function POST(request: NextRequest): Promise<NextResponseType<unkno
 
   accountForMicrodollarUsage(clonedReponse, usageContext, openrouterRequestSpan);
 
-  await handleRequestLogging({
-    clonedResponse: response.clone(),
+  const requestLogging = {
     user: maybeUser,
     organization_id: organizationId || null,
-    provider: effectiveProviderContext.provider.id,
-    model: effectiveModelIdLowerCased,
     session_id: usageContext.session_id,
     vercel_request_id: extractHeaderAndLimitLength(request, 'x-vercel-id'),
     request: requestBodyParsed,
-  });
+  };
 
   {
     const errorResponse = await makeErrorReadable({
@@ -977,6 +973,12 @@ export async function POST(request: NextRequest): Promise<NextResponseType<unkno
       isUserByok: !!effectiveProviderContext.userByok,
     });
     if (errorResponse) {
+      await logUnrewrittenResponse(
+        response,
+        effectiveModelIdLowerCased,
+        effectiveProviderContext.provider.id,
+        requestLogging
+      );
       return errorResponse;
     }
   }
@@ -986,7 +988,7 @@ export async function POST(request: NextRequest): Promise<NextResponseType<unkno
     effectiveModelIdLowerCased,
     effectiveProviderContext.provider.id,
     requestBodyParsed.kind,
-    organizationId
+    requestLogging
   );
   if (rewrittenResponse) {
     return rewrittenResponse;

apps/web/src/lib/rewriteModelResponse.test.ts

@@ -1,12 +1,29 @@
-import { describe, test, expect, jest } from '@jest/globals';
+import { describe, test, expect, beforeEach } from '@jest/globals';
 import {
   rewriteModelResponse_ChatCompletions,
   rewriteModelResponse_Messages,
   rewriteModelResponse_Responses,
   rewriteModelResponse,
+  type RequestLoggingParams,
 } from './rewriteModelResponse';
+import { isDynamicallyOptedIntoRequestLogging } from '@/lib/ai-gateway/request-logging-opt-ins';
 import { KILO_ORGANIZATION_ID } from '@/lib/organizations/constants';
 
+jest.mock('next/server', () => ({
+  ...(jest.requireActual('next/server') as Record<string, unknown>),
+  after: jest.fn(),
+}));
+
+jest.mock('@/lib/ai-gateway/request-logging-opt-ins', () => ({
+  isDynamicallyOptedIntoRequestLogging: jest.fn(async () => false),
+}));
+
+const mockedOptIn = jest.mocked(isDynamicallyOptedIntoRequestLogging);
+
+beforeEach(() => {
+  mockedOptIn.mockClear();
+});
+
 function jsonResponse(body: unknown, status = 200): Response {
   return new Response(JSON.stringify(body), {
     status,
@@ -478,6 +495,17 @@ describe('rewriteModelResponse_Responses', () => {
   });
 });
 
+function makeLogging(overrides?: Partial<RequestLoggingParams>): RequestLoggingParams {
+  return {
+    user: null,
+    organization_id: null,
+    session_id: null,
+    vercel_request_id: null,
+    request: { body: {} } as unknown as RequestLoggingParams['request'],
+    ...overrides,
+  };
+}
+
 describe('rewriteModelResponse', () => {
   test('rewrites paid-model Kilo organization traffic without stripping cost', async () => {
     const result = await rewriteModelResponse(
@@ -492,7 +520,7 @@ describe('rewriteModelResponse', () => {
       'openai/gpt-5',
       'openrouter',
       'chat_completions',
-      KILO_ORGANIZATION_ID
+      makeLogging({ organization_id: KILO_ORGANIZATION_ID })
     );
 
     expect(result).not.toBeNull();
@@ -511,7 +539,7 @@ describe('rewriteModelResponse', () => {
       'openai/gpt-5',
       'openrouter',
       'chat_completions',
-      '00000000-0000-0000-0000-000000000000'
+      makeLogging({ organization_id: '00000000-0000-0000-0000-000000000000' })
     );
 
     expect(result).toBeNull();
@@ -525,7 +553,8 @@ describe('rewriteModelResponse', () => {
       }),
       'google/gemma-4-26b-a4b-it:free',
       'openrouter',
-      'chat_completions'
+      'chat_completions',
+      makeLogging()
     );
 
     expect(result).not.toBeNull();
@@ -534,4 +563,114 @@ describe('rewriteModelResponse', () => {
       usage: {},
     });
   });
+
+  test('processes responses it would normally skip when request logging is enabled', async () => {
+    mockedOptIn.mockResolvedValueOnce(true);
+    const result = await rewriteModelResponse(
+      jsonResponse({ model: 'openai/gpt-5' }),
+      'openai/gpt-5',
+      'openrouter',
+      'chat_completions',
+      makeLogging({ organization_id: '00000000-0000-0000-0000-000000000000' })
+    );
+
+    expect(result).not.toBeNull();
+  });
+});
+
+function makeCapture() {
+  return { setBody: jest.fn(), setReadError: jest.fn() };
+}
+
+describe('request log capture', () => {
+  test.each(rewriters)('%s: captures the raw JSON body', async (_name, rewrite) => {
+    const capture = makeCapture();
+    const body = { model: 'upstream-model' };
+
+    const result = await rewrite(jsonResponse(body), true, capture);
+
+    expect(result.status).toBe(200);
+    expect(capture.setBody).toHaveBeenCalledTimes(1);
+    expect(capture.setBody).toHaveBeenCalledWith(JSON.stringify(body));
+    expect(capture.setReadError).not.toHaveBeenCalled();
+  });
+
+  test.each(rewriters)('%s: captures the raw event stream', async (_name, rewrite) => {
+    const capture = makeCapture();
+    const sseBody =
+      'data: {"id":"gen-1","model":"upstream-model","choices":[]}\n\n' + 'data: [DONE]\n\n';
+
+    const result = await rewrite(sseResponse(sseBody), true, capture);
+    await readOutputStream(result);
+
+    expect(capture.setBody).toHaveBeenCalledTimes(1);
+    expect(capture.setBody).toHaveBeenCalledWith(sseBody);
+    expect(capture.setReadError).not.toHaveBeenCalled();
+  });
+
+  test.each(rewriters)(
+    '%s: captures an empty body when upstream has no body',
+    async (_name, rewrite) => {
+      const capture = makeCapture();
+
+      const result = await rewrite(
+        new Response(null, { headers: { 'content-type': 'text/event-stream' } }),
+        true,
+        capture
+      );
+      await readOutputStream(result);
+
+      expect(capture.setBody).toHaveBeenCalledWith('');
+      expect(capture.setReadError).not.toHaveBeenCalled();
+    }
+  );
+
+  test.each(rewriters)('%s: records a read error when the stream fails', async (_name, rewrite) => {
+    const capture = makeCapture();
+
+    const result = await rewrite(
+      failingResponse(
+        'text/event-stream',
+        'ResponseAborted',
+        'data: {"id":"gen-1","choices":[]}\n\n'
+      ),
+      true,
+      capture
+    );
+    await readOutputStream(result);
+
+    expect(capture.setReadError).toHaveBeenCalledTimes(1);
+    expect(capture.setBody).not.toHaveBeenCalled();
+  });
+
+  test.each(rewriters)(
+    '%s: records a read error when a JSON body cannot be read',
+    async (_name, rewrite) => {
+      const capture = makeCapture();
+
+      const result = await rewrite(
+        failingResponse('application/json', 'TimeoutError'),
+        true,
+        capture
+      );
+
+      expect(result.status).toBe(503);
+      expect(capture.setReadError).toHaveBeenCalledTimes(1);
+      expect(capture.setBody).not.toHaveBeenCalled();
+    }
+  );
+
+  test('records a read error when the response stream is cancelled', async () => {
+    const capture = makeCapture();
+    const upstream = new Response(new ReadableStream<Uint8Array>({ start() {} }), {
+      headers: { 'content-type': 'text/event-stream' },
+    });
+
+    const result = await rewriteModelResponse_ChatCompletions(upstream, true, capture);
+    const reader = result.body?.getReader();
+    await reader?.cancel();
+
+    expect(capture.setReadError).toHaveBeenCalled();
+    expect(capture.setBody).not.toHaveBeenCalled();
+  });
 });