refactor(exa): share provider dispatch and usage recording

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

@@ -1,157 +1,47 @@
-import { NextResponse } from 'next/server';
-import { type NextRequest } from 'next/server';
+import { NextResponse, after, type NextRequest } from 'next/server';
 import { getUserFromAuth } from '@/lib/user/server';
-import { EXA_API_KEY } from '@/lib/config.server';
-import { after } from 'next/server';
 import { wrapInSafeNextResponse } from '@/lib/ai-gateway/llm-proxy-helpers';
-import {
-  getExaMonthlyUsage,
-  getExaFreeAllowanceMicrodollars,
-  recordExaUsage,
-} from '@/lib/exa-usage';
-import { getBalanceAndOrgSettings } from '@/lib/organizations/organization-usage';
-import { readDb } from '@/lib/drizzle';
 import { captureException } from '@sentry/nextjs';
 import { validateFeatureHeader, FEATURE_HEADER } from '@/lib/feature-detection';
 import { EXA_ALLOWED_PATHS, isExaAllowedPath } from '@/lib/exa-paths';
-import { z } from 'zod';
-
-const EXA_BASE_URL = 'https://api.exa.ai';
-const MICRODOLLARS_PER_DOLLAR = 1_000_000;
-const ExaCostResponseSchema = z.object({
-  costDollars: z
-    .object({
-      total: z.number().finite().optional(),
-    })
-    .optional(),
-});
-
-function extractExaPath(url: URL): string | null {
-  const prefix = '/api/exa';
-  if (!url.pathname.startsWith(prefix)) return null;
-  const path = url.pathname.slice(prefix.length);
-  return isExaAllowedPath(path) ? path : null;
-}
-
-function extractCostMicrodollars(responseBody: unknown): number | undefined {
-  const costDollars = ExaCostResponseSchema.parse(responseBody).costDollars?.total;
-  if (costDollars === undefined || costDollars === 0) return undefined;
-  if (costDollars < 0) {
-    throw new Error('Exa response costDollars.total must be positive.');
-  }
-
-  const costMicrodollars = Math.round(costDollars * MICRODOLLARS_PER_DOLLAR);
-  if (!Number.isSafeInteger(costMicrodollars) || costMicrodollars <= 0) {
-    throw new Error('Exa response cost must convert to a positive safe integer.');
-  }
-  return costMicrodollars;
-}
+import { extractExaCostMicrodollars, prepareExaRequest } from '@/lib/exa-provider';
 
 export async function POST(request: NextRequest) {
-  const { user, authFailedResponse, organizationId } = await getUserFromAuth({
-    adminOnly: false,
-  });
+  const { user, authFailedResponse, organizationId } = await getUserFromAuth({ adminOnly: false });
   if (authFailedResponse) return authFailedResponse;
-
   const url = new URL(request.url);
-  const exaPath = extractExaPath(url);
-  if (!exaPath) {
+  const prefix = '/api/exa';
+  const exaPath = url.pathname.startsWith(prefix) ? url.pathname.slice(prefix.length) : '';
+  if (!isExaAllowedPath(exaPath)) {
     return NextResponse.json(
       { error: `Invalid path. Allowed: ${EXA_ALLOWED_PATHS.join(', ')}` },
       { status: 400 }
     );
   }
-
-  if (!EXA_API_KEY) {
-    captureException(new Error('EXA_API_KEY is not configured'));
-
-    return NextResponse.json({ error: 'Internal Server Error' }, { status: 500 });
-  }
-
-  // Check monthly allowance and balance.
-  // freeAllowance is the stored value from the first request of the month;
-  // null means no row yet, so we compute from the helper.
-  // Use read replica for monthly usage check - this is a read-only operation that can tolerate
-  // slight replication lag, and provides lower latency for US users
-  const { usage: monthlyUsage, freeAllowance: storedAllowance } = await getExaMonthlyUsage(
-    user.id,
-    readDb
-  );
-  const allowance = storedAllowance ?? getExaFreeAllowanceMicrodollars(new Date(), user);
-  const isPaidRequest = monthlyUsage >= allowance;
-
-  if (isPaidRequest) {
-    const { balance } = await getBalanceAndOrgSettings(organizationId, user, readDb);
-    if (balance <= 0) {
-      return NextResponse.json(
-        {
-          error: 'Exa free allowance exhausted and no credit balance available',
-          monthlyAllowance: `$${(allowance / 1_000_000).toFixed(2)}`,
-          used: `$${(monthlyUsage / 1_000_000).toFixed(2)}`,
-        },
-        { status: 402 }
-      );
-    }
-  }
-
-  // Strip `stream` to guarantee JSON responses with costDollars for billing
+  const provider = await prepareExaRequest(user, organizationId);
+  if (provider instanceof Response) return provider;
+  // Old proxy callers retain their body, response, and asynchronous billing contracts until retirement.
   const requestBody: Record<string, unknown> = await request.json();
   delete requestBody.stream;
-
   const featureId = validateFeatureHeader(request.headers.get(FEATURE_HEADER)) ?? undefined;
   const type = typeof requestBody.type === 'string' ? requestBody.type : undefined;
-
-  const response = await fetch(`${EXA_BASE_URL}${exaPath}`, {
-    method: 'POST',
-    headers: {
-      'Content-Type': 'application/json',
-      'x-api-key': EXA_API_KEY,
-    },
-    body: JSON.stringify(requestBody),
-    signal: request.signal,
-  });
-
-  if (response.status >= 400) {
-    console.error(
-      `[exa] upstream error: status=${response.status} user=${user.id} path=${exaPath}`
-    );
-  }
-
-  // Record cost asynchronously after sending the response
+  const response = await provider.send(exaPath, requestBody, request.signal);
   const cloned = response.clone();
   after(async () => {
-    if (response.status >= 400) {
-      return;
-    }
-
+    if (response.status >= 400) return;
     try {
-      const body: unknown = await cloned.json();
-      const costMicrodollars = extractCostMicrodollars(body);
-      if (costMicrodollars === undefined) return;
-
-      await recordExaUsage({
-        userId: user.id,
-        organizationId,
-        path: exaPath,
-        costMicrodollars,
-        chargedToBalance: isPaidRequest,
-        freeAllowanceMicrodollars: allowance,
+      await provider.record(
+        exaPath,
+        extractExaCostMicrodollars(await cloned.json()),
         featureId,
-        type,
-      });
+        type
+      );
     } catch (error) {
       captureException(error, {
-        tags: {
-          route: '/api/exa/[...path]',
-          exaPath,
-        },
-        extra: {
-          userId: user.id,
-          responseStatus: response.status,
-        },
+        tags: { route: '/api/exa/[...path]', exaPath },
+        extra: { userId: user.id, responseStatus: response.status },
       });
     }
   });
-
   return wrapInSafeNextResponse(response);
 }

apps/web/src/lib/exa-provider.test.ts

@@ -0,0 +1,137 @@
+import { afterAll, beforeEach, expect, it, jest } from '@jest/globals';
+import type * as Route from '@/app/api/exa/[...path]/route';
+import type { User } from '@kilocode/db/schema';
+import type * as Provider from './exa-provider';
+
+let response: Response, upstreamRequest: Request, usage: number, balance: number;
+let ledger: Record<string, unknown>[],
+  sent: { url: string; body: unknown }[],
+  after: (() => Promise<void>)[];
+jest.mock('@/lib/config.server', () => ({ EXA_API_KEY: 'provider-secret' }));
+jest.mock('@/lib/drizzle', () => ({ readDb: {} }));
+jest.mock('@sentry/nextjs', () => ({ captureException: jest.fn() }));
+jest.mock('@/lib/exa-usage', () => ({
+  getExaMonthlyUsage: async () => ({ usage, freeAllowance: 1000 }),
+  getExaFreeAllowanceMicrodollars: () => 1000,
+  recordExaUsage: async (entry: Record<string, unknown>) => {
+    ledger.push(entry);
+  },
+}));
+jest.mock('@/lib/organizations/organization-usage', () => ({
+  getBalanceAndOrgSettings: async () => ({ balance }),
+}));
+jest.mock('@/lib/user/server', () => ({
+  getUserFromAuth: async () => ({
+    user: { id: 'oauth/owner' },
+    organizationId: 'organization',
+    authFailedResponse: null,
+  }),
+}));
+jest.mock('@/lib/ai-gateway/llm-proxy-helpers', () => ({
+  wrapInSafeNextResponse: (value: Response) => value,
+}));
+jest.mock('next/server', () => ({
+  ...jest.requireActual<object>('next/server'),
+  after: (callback: () => Promise<void>) => {
+    after.push(callback);
+  },
+}));
+const { POST } = jest.requireActual<typeof Route>('@/app/api/exa/[...path]/route');
+const { prepareExaRequest, extractExaCostMicrodollars } =
+  jest.requireActual<typeof Provider>('./exa-provider');
+const originalFetch = globalThis.fetch;
+beforeEach(() => {
+  response = Response.json({ results: [], costDollars: { total: 0.002 } });
+  usage = 0;
+  balance = 1;
+  ledger = [];
+  sent = [];
+  after = [];
+  globalThis.fetch = async (url, init) => {
+    upstreamRequest = new Request(String(url), init);
+    sent.push({ url: String(url), body: JSON.parse(String(init?.body)) });
+    return response;
+  };
+});
+afterAll(() => {
+  globalThis.fetch = originalFetch;
+});
+
+it.each(['/search', '/contents', '/findSimilar', '/answer', '/context'])(
+  'preserves legacy %s responses and delayed billing',
+  async path => {
+    const result = await POST(
+      new Request(`https://kilo.example/api/exa${path}`, {
+        method: 'POST',
+        body: JSON.stringify({ query: 'old', stream: true }),
+      }) as never
+    );
+    expect(result.status).toBe(200);
+    expect(await result.json()).toEqual({ results: [], costDollars: { total: 0.002 } });
+    expect(sent[0]).toEqual({ url: `https://api.exa.ai${path}`, body: { query: 'old' } });
+    expect(upstreamRequest.redirect).toBe('follow');
+    expect(upstreamRequest.headers.get('accept')).toBeNull();
+    expect(ledger).toEqual([]);
+    for (const callback of after) await callback();
+    expect(ledger).toHaveLength(1);
+    expect(ledger[0]).toMatchObject({ path, costMicrodollars: 2000, chargedToBalance: false });
+  }
+);
+
+it.each([
+  [0, false],
+  [1000, true],
+] as const)(
+  'shares actual-cost billing with JSON-only callers at usage %s',
+  async (monthly, paid) => {
+    usage = monthly;
+    const provider = await prepareExaRequest({ id: 'oauth/owner' } as User, 'organization');
+    if (provider instanceof Response) throw new Error(`Unexpected status: ${provider.status}`);
+    const result = await provider.send(
+      '/contents',
+      { ids: ['https://example.com'] },
+      new AbortController().signal,
+      true
+    );
+    expect(upstreamRequest.redirect).toBe('error');
+    expect(upstreamRequest.headers.get('accept')).toBe('application/json');
+    expect(ledger).toEqual([]);
+    await provider.record(
+      '/contents',
+      extractExaCostMicrodollars(await result.json()),
+      'quick-chat'
+    );
+    expect(ledger).toEqual([
+      {
+        userId: 'oauth/owner',
+        organizationId: 'organization',
+        path: '/contents',
+        costMicrodollars: 2000,
+        chargedToBalance: paid,
+        freeAllowanceMicrodollars: 1000,
+        featureId: 'quick-chat',
+        type: undefined,
+      },
+    ]);
+  }
+);
+
+it.each([undefined, 0])('preserves cost %s without inventing a charge', async total => {
+  const provider = await prepareExaRequest({ id: 'oauth/owner' } as User, 'organization');
+  if (provider instanceof Response) throw new Error(`Unexpected status: ${provider.status}`);
+  const cost = extractExaCostMicrodollars({ costDollars: { total } });
+  expect(cost).toBe(total);
+  await provider.record('/search', cost);
+  expect(ledger).toEqual([]);
+});
+
+it.each([429, 503])('preserves retryable provider status %s without billing', async status => {
+  response = Response.json({ error: 'Retry later', costDollars: { total: 0.002 } }, { status });
+  const result = await POST(
+    new Request('https://kilo.example/api/exa/search', { method: 'POST', body: '{}' }) as never
+  );
+  expect(result.status).toBe(status);
+  expect(await result.json()).toEqual({ error: 'Retry later', costDollars: { total: 0.002 } });
+  for (const callback of after) await callback();
+  expect(ledger).toEqual([]);
+});