From 5cee9d43f930086c4c3657d02b90ce64bef2784b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Igor=20=C5=A0=C4=87eki=C4=87?= Date: Sun, 30 Aug 2026 07:33:20 +0200 Subject: [PATCH] refactor(exa): share provider dispatch and usage recording --- apps/web/src/app/api/exa/[...path]/route.ts | 146 +++----------------- apps/web/src/lib/exa-provider.test.ts | 137 ++++++++++++++++++ apps/web/src/lib/exa-provider.ts | 98 +++++++++++++ 3 files changed, 253 insertions(+), 128 deletions(-) create mode 100644 apps/web/src/lib/exa-provider.test.ts create mode 100644 apps/web/src/lib/exa-provider.ts diff --git a/apps/web/src/app/api/exa/[...path]/route.ts b/apps/web/src/app/api/exa/[...path]/route.ts index 729510b68e..9b252f6b72 100644 --- a/apps/web/src/app/api/exa/[...path]/route.ts +++ b/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 = 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); } diff --git a/apps/web/src/lib/exa-provider.test.ts b/apps/web/src/lib/exa-provider.test.ts new file mode 100644 index 0000000000..5408605f58 --- /dev/null +++ b/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[], + sent: { url: string; body: unknown }[], + after: (() => Promise)[]; +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) => { + 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('next/server'), + after: (callback: () => Promise) => { + after.push(callback); + }, +})); +const { POST } = jest.requireActual('@/app/api/exa/[...path]/route'); +const { prepareExaRequest, extractExaCostMicrodollars } = + jest.requireActual('./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([]); +}); diff --git a/apps/web/src/lib/exa-provider.ts b/apps/web/src/lib/exa-provider.ts new file mode 100644 index 0000000000..2ebc65852c --- /dev/null +++ b/apps/web/src/lib/exa-provider.ts @@ -0,0 +1,98 @@ +import 'server-only'; +import { NextResponse } from 'next/server'; +import type { User } from '@kilocode/db/schema'; +import { captureException } from '@sentry/nextjs'; +import { z } from 'zod'; +import { EXA_API_KEY } from '@/lib/config.server'; +import { readDb } from '@/lib/drizzle'; +import { getBalanceAndOrgSettings } from '@/lib/organizations/organization-usage'; +import { getExaMonthlyUsage, getExaFreeAllowanceMicrodollars, recordExaUsage } from './exa-usage'; +import type { ExaAllowedPath } from './exa-paths'; + +const ExaCostResponseSchema = z.object({ + costDollars: z.object({ total: z.number().finite().optional() }).optional(), +}); +export function extractExaCostMicrodollars(responseBody: unknown): number | undefined { + const costDollars = ExaCostResponseSchema.parse(responseBody).costDollars?.total; + // Keep explicit zero distinct from unknown cost; neither creates a legacy usage charge. + if (costDollars === undefined || costDollars === 0) return costDollars; + if (costDollars < 0) throw new Error('Exa response costDollars.total must be positive.'); + const costMicrodollars = Math.round(costDollars * 1_000_000); + if (!Number.isSafeInteger(costMicrodollars) || costMicrodollars <= 0) { + throw new Error('Exa response cost must convert to a positive safe integer.'); + } + return costMicrodollars; +} + +/** Call only with a freshly authorized user and context, never model-supplied identity. */ +export async function prepareExaRequest(user: User, organizationId: string | undefined) { + if (!EXA_API_KEY) { + captureException(new Error('EXA_API_KEY is not configured')); + return NextResponse.json({ error: 'Internal Server Error' }, { status: 500 }); + } + // Preserve the proxy's replica-based monthly allowance and balance checks. + 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 } + ); + } + } + return { + async send( + path: ExaAllowedPath, + body: Record, + signal: AbortSignal, + jsonOnly = false + ) { + const response = await fetch(`https://api.exa.ai${path}`, { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + 'x-api-key': EXA_API_KEY, + ...(jsonOnly ? { Accept: 'application/json' } : {}), + }, + body: JSON.stringify(body), + signal, + // Old proxy callers retain redirect behavior until that public contract retires. + ...(jsonOnly ? { redirect: 'error' as const } : {}), + }); + if (response.status >= 400) { + console.error( + `[exa] upstream error: status=${response.status} user=${user.id} path=${path}` + ); + } + return response; + }, + async record( + path: ExaAllowedPath, + costMicrodollars: number | undefined, + featureId?: string, + type?: string + ) { + if (!costMicrodollars) return; + await recordExaUsage({ + userId: user.id, + organizationId, + path, + costMicrodollars, + chargedToBalance: isPaidRequest, + freeAllowanceMicrodollars: allowance, + featureId, + type, + }); + }, + }; +}