From 0ea990ef31d154596dcb1365edc1d6e578ff88bb Mon Sep 17 00:00:00 2001 From: Christiaan Arnoldus Date: Wed, 23 Sep 2026 11:15:02 +0200 Subject: [PATCH 1/6] feat(ai-gateway): support TypeSafe System One via OpenRouter --- apps/web/package.json | 1 + .../typesafe/v1/systemone/route.test.ts | 587 ++++++++++++++++++ .../gateway/typesafe/v1/systemone/route.ts | 191 ++++++ apps/web/src/lib/ai-gateway/typesafe.test.ts | 110 ++++ apps/web/src/lib/ai-gateway/typesafe.ts | 88 +++ packages/db/src/schema-types.ts | 1 + pnpm-lock.yaml | 9 + 7 files changed, 987 insertions(+) create mode 100644 apps/web/src/app/api/gateway/typesafe/v1/systemone/route.test.ts create mode 100644 apps/web/src/app/api/gateway/typesafe/v1/systemone/route.ts create mode 100644 apps/web/src/lib/ai-gateway/typesafe.test.ts create mode 100644 apps/web/src/lib/ai-gateway/typesafe.ts diff --git a/apps/web/package.json b/apps/web/package.json index 596c67f1a8..570e51f010 100644 --- a/apps/web/package.json +++ b/apps/web/package.json @@ -112,6 +112,7 @@ "@types/archiver": "7.0.0", "@types/js-yaml": "4.0.9", "@types/mdx": "2.0.13", + "@typesafe-ai/sdk": "0.6.0", "@upstash/redis": "^1.38.0", "@vercel/firewall": "1.2.1", "@vercel/functions": "3.4.6", diff --git a/apps/web/src/app/api/gateway/typesafe/v1/systemone/route.test.ts b/apps/web/src/app/api/gateway/typesafe/v1/systemone/route.test.ts new file mode 100644 index 0000000000..78470c4cbb --- /dev/null +++ b/apps/web/src/app/api/gateway/typesafe/v1/systemone/route.test.ts @@ -0,0 +1,587 @@ +import { afterEach, beforeEach, describe, expect, it, jest } from '@jest/globals'; +import { captureException, captureMessage } from '@sentry/nextjs'; +import { TypeSafeClient, choice, noul, score } from '@typesafe-ai/sdk'; +import type { User } from '@kilocode/db/schema'; +import { KILO_GATEWAY_AUDIENCE } from '@kilocode/worker-utils/internal-service-token-audiences'; +import { after, NextRequest, NextResponse } from 'next/server'; +import { getUserFromAuth } from '@/lib/user/server'; +import { getBalanceAndOrgSettings } from '@/lib/organizations/organization-usage'; +import { resolveOrganizationMemberModelDecision } from '@/lib/organizations/effective-model-access.server'; +import { + gatewayRateLimitKey, + isGatewayAccountRateLimited, +} from '@/lib/ai-gateway/gateway-account-rate-limit'; +import { + checkOrganizationModelRestrictions, + creditsBlockedResponse, + extractFraudAndProjectHeaders, + extractHeaderAndLimitLength, + modelNotAllowedResponse, + wrapInSafeNextResponse, +} from '@/lib/ai-gateway/llm-proxy-helpers'; +import { OPENROUTER } from '@/lib/ai-gateway/providers/definitions/openrouter'; +import { generateProviderSpecificHash } from '@/lib/ai-gateway/providerHash'; +import { logMicrodollarUsage } from '@/lib/ai-gateway/processUsage'; +import { TYPESAFE_MODEL } from '@/lib/ai-gateway/typesafe'; +import { EmptyFraudDetectionHeaders } from '@/lib/utils'; +import { POST } from './route'; + +jest.mock('next/server', () => ({ + ...jest.requireActual('next/server'), + after: jest.fn(), +})); +jest.mock('@sentry/nextjs', () => ({ + captureException: jest.fn(), + captureMessage: jest.fn(), +})); +jest.mock('@/lib/user/server', () => ({ getUserFromAuth: jest.fn() })); +jest.mock('@/lib/organizations/organization-usage', () => ({ + getBalanceAndOrgSettings: jest.fn(), +})); +jest.mock('@/lib/organizations/effective-model-access.server', () => ({ + resolveOrganizationMemberModelDecision: jest.fn(), +})); +jest.mock('@/lib/ai-gateway/gateway-account-rate-limit', () => ({ + gatewayRateLimitKey: jest.fn(), + isGatewayAccountRateLimited: jest.fn(), +})); +jest.mock('@/lib/ai-gateway/llm-proxy-helpers', () => ({ + checkOrganizationModelRestrictions: jest.fn(), + creditsBlockedResponse: jest.fn(), + extractFraudAndProjectHeaders: jest.fn(), + extractHeaderAndLimitLength: jest.fn(), + modelNotAllowedResponse: jest.fn(), + wrapInSafeNextResponse: jest.fn(), +})); +jest.mock('@/lib/ai-gateway/providers/definitions/openrouter', () => ({ + OPENROUTER: { + id: 'openrouter', + apiUrl: 'https://openrouter.ai/api/v1', + apiKey: 'test-platform-openrouter-key', + }, +})); +jest.mock('@/lib/ai-gateway/providerHash', () => ({ generateProviderSpecificHash: jest.fn() })); +jest.mock('@/lib/ai-gateway/processUsage', () => ({ logMicrodollarUsage: jest.fn() })); + +const routeUrl = 'http://localhost:3000/api/gateway/typesafe/v1/systemone'; +const user = { + id: 'oauth/test-user', + google_user_email: 'test@example.com', + microdollars_used: 123, +} as User; +const questions = { + billing: noul('Is this about billing?'), + category: choice('Choose a category', { billing: 'Payments', other: null }), + urgency: score('How urgent is this?', ['Low', 'High']), +}; +const requestBody = { state: { message: 'I was charged twice.' }, questions }; +const upstreamBody = { + id: 'gen-systemone-123', + model: TYPESAFE_MODEL, + provider: 'TypeSafe upstream', + answers: { + billing: { type: 'noul', noul: 0.95 }, + category: { + type: 'choice', + choice: 'billing', + confidence: 0.9, + probabilities: { billing: 0.9, other: 0.1 }, + }, + urgency: { + type: 'score', + score: 0.75, + confidence: 0.8, + probabilities: { '0': 0.25, '1': 0.75 }, + legend: { '0': 'Low', '1': 'High' }, + }, + }, + usage: { input_tokens: 23, output_tokens: 7, cost: 0.0001236, total_tokens: 30 }, + upstream_metadata: { retained: true }, +}; +const memberDecision = { + policy: { + requireModelInCurrentSnapshot: false, + organizationModelDenyList: [], + memberGrant: { mode: 'unrestricted' as const }, + policyRevision: 1, + }, + decision: { allowed: true }, +}; +const mockedFetch = jest.fn(); +const originalFetch = globalThis.fetch; + +function makeRequest(body: unknown = requestBody, headers: Record = {}) { + return new NextRequest(routeUrl, { + method: 'POST', + headers: { 'Content-Type': 'application/json', ...headers }, + body: JSON.stringify(body), + }); +} + +function setAuth(organizationId?: string) { + jest.mocked(getUserFromAuth).mockResolvedValue({ + user, + authFailedResponse: null, + organizationId, + botId: 'bot-123', + tokenSource: 'api-key', + }); +} + +function upstreamRequest() { + expect(mockedFetch).toHaveBeenCalledTimes(1); + const [url, init] = mockedFetch.mock.calls[0]; + expect(url).toBe('https://openrouter.ai/api/v1/systemone'); + expect(init?.method).toBe('POST'); + if (typeof init?.body !== 'string') throw new Error('Expected a JSON upstream body'); + return { body: JSON.parse(init.body), headers: new Headers(init.headers) }; +} + +async function runAfter() { + expect(after).toHaveBeenCalledTimes(1); + const [callback] = jest.mocked(after).mock.calls[0]; + if (typeof callback !== 'function') throw new Error('Expected deferred usage callback'); + await callback(); +} + +describe('POST /api/gateway/typesafe/v1/systemone', () => { + beforeEach(() => { + jest.resetAllMocks(); + globalThis.fetch = mockedFetch; + setAuth(); + jest.mocked(getBalanceAndOrgSettings).mockResolvedValue({ balance: 1_000_000 }); + jest.mocked(isGatewayAccountRateLimited).mockResolvedValue(false); + jest.mocked(gatewayRateLimitKey).mockReturnValue('test-rate-limit-key'); + jest.mocked(checkOrganizationModelRestrictions).mockReturnValue({ error: null }); + jest.mocked(resolveOrganizationMemberModelDecision).mockResolvedValue(memberDecision); + jest.mocked(generateProviderSpecificHash).mockReturnValue('hashed-user'); + jest.mocked(extractFraudAndProjectHeaders).mockReturnValue({ + fraudHeaders: EmptyFraudDetectionHeaders, + projectId: 'project-123', + xKiloCodeVersion: null, + numericKiloCodeVersion: 0, + }); + jest + .mocked(extractHeaderAndLimitLength) + .mockImplementation((request, name) => request.headers.get(name)); + jest + .mocked(modelNotAllowedResponse) + .mockImplementation(() => + NextResponse.json({ error_type: 'model_not_allowed' }, { status: 404 }) + ); + mockedFetch.mockImplementation(async () => Response.json(upstreamBody)); + }); + + afterEach(() => { + globalThis.fetch = originalFetch; + }); + + it('supports the TypeSafe SDK contract through a custom fetch transport', async () => { + const routeFetch = jest.fn(async (input: string, init?: RequestInit) => + POST(new NextRequest(input, init)) + ); + const client = new TypeSafeClient({ + apiKey: 'test-kilo-key', + baseURL: 'http://localhost:3000/api/gateway/typesafe', + defaultModel: 'jev-1.13', + retry: { maxRetries: 0 }, + logLevel: 'off', + fetch: routeFetch, + }); + + const { data, response } = await client.systemOne(requestBody).withResponse(); + + expect(response.status).toBe(200); + expect(data).toEqual(upstreamBody); + expect(data.answers.billing.noul).toBe(0.95); + expect(data.answers.category.choice).toBe('billing'); + expect(data.answers.urgency.score).toBe(0.75); + expect(routeFetch).toHaveBeenCalledTimes(1); + const [url, init] = routeFetch.mock.calls[0]; + expect(url).toBe(routeUrl); + expect(new Headers(init?.headers).get('authorization')).toBe('Bearer test-kilo-key'); + expect(upstreamRequest().body).toEqual({ + ...JSON.parse(JSON.stringify(requestBody)), + model: TYPESAFE_MODEL, + user: 'hashed-user', + }); + expect(getUserFromAuth).toHaveBeenCalledWith({ + adminOnly: false, + expectedAudience: KILO_GATEWAY_AUDIENCE, + }); + expect(resolveOrganizationMemberModelDecision).not.toHaveBeenCalled(); + }); + + it.each([undefined, 'jev-1.13', 'typesafe/jev-1.13'])( + 'pins model %s to OpenRouter and drops client credentials and routing overrides', + async model => { + const response = await POST( + makeRequest( + { + ...requestBody, + model, + provider: { only: ['attacker'], api_key: 'test-body-key' }, + user: 'attacker-user', + api_key: 'test-body-key', + byok: true, + user_byok: [{ providerId: 'typesafe', apiKey: 'test-byok-key' }], + base_url: 'https://attacker.invalid', + }, + { + Authorization: 'Bearer test-client-key', + 'Proxy-Authorization': 'Bearer test-proxy-key', + 'x-api-key': 'test-forwarded-key', + 'x-openrouter-api-key': 'test-user-openrouter-key', + Cookie: 'session=test-cookie', + 'HTTP-Referer': 'https://attacker.invalid', + 'X-Title': 'Attacker', + } + ) + ); + + expect(response.status).toBe(200); + const upstream = upstreamRequest(); + expect(upstream.body).toEqual({ + ...JSON.parse(JSON.stringify(requestBody)), + model: 'typesafe/jev-1.13', + user: 'hashed-user', + }); + expect(Object.fromEntries(upstream.headers)).toEqual({ + authorization: 'Bearer test-platform-openrouter-key', + 'content-type': 'application/json', + 'http-referer': 'https://kilocode.ai', + 'x-title': 'Kilo Code', + }); + expect(generateProviderSpecificHash).toHaveBeenCalledWith(user.id, OPENROUTER); + expect(checkOrganizationModelRestrictions).toHaveBeenCalledWith({ + modelId: TYPESAFE_MODEL, + settings: undefined, + organizationPlan: undefined, + }); + } + ); + + it('defers billing and converts upstream cost to rounded microdollars without a markup', async () => { + setAuth('org-123'); + const response = await POST( + makeRequest(requestBody, { + 'x-kilocode-editorname': 'vscode', + 'x-kilocode-machineid': 'machine-123', + 'x-kilocode-feature': 'cli', + 'X-KiloCode-TaskId': 'task-123', + 'x-kilocode-mode': 'code', + }) + ); + + expect(response.status).toBe(200); + expect(await response.json()).toEqual(upstreamBody); + expect(logMicrodollarUsage).not.toHaveBeenCalled(); + await runAfter(); + expect(logMicrodollarUsage).toHaveBeenCalledTimes(1); + expect(logMicrodollarUsage).toHaveBeenCalledWith( + expect.objectContaining({ + messageId: upstreamBody.id, + model: TYPESAFE_MODEL, + inference_provider: 'TypeSafe upstream', + cost_mUsd: 124, + market_cost: 124, + inputTokens: 23, + outputTokens: 7, + cacheWriteTokens: 0, + cacheHitTokens: 0, + is_byok: false, + streamed: false, + hasError: false, + status_code: 200, + responseContent: '', + latency: expect.any(Number), + }), + expect.objectContaining({ + api_kind: 'systemone', + kiloUserId: user.id, + provider: 'openrouter', + requested_model: TYPESAFE_MODEL, + organizationId: 'org-123', + prior_microdollar_usage: 123, + posthog_distinct_id: user.google_user_email, + project_id: 'project-123', + fraudHeaders: EmptyFraudDetectionHeaders, + editor_name: 'vscode', + machine_id: 'machine-123', + feature: 'cli', + session_id: 'task-123', + mode: 'code', + user_byok: false, + isStreaming: false, + botId: 'bot-123', + tokenSource: 'api-key', + ttfb_ms: expect.any(Number), + }) + ); + }); + + it('accepts zero-cost usage and falls back to TypeSafe when the provider is absent', async () => { + mockedFetch.mockResolvedValue( + Response.json({ + ...upstreamBody, + provider: undefined, + usage: { ...upstreamBody.usage, cost: 0 }, + }) + ); + + expect((await POST(makeRequest())).status).toBe(200); + await runAfter(); + + expect(logMicrodollarUsage).toHaveBeenCalledWith( + expect.objectContaining({ inference_provider: 'TypeSafe', cost_mUsd: 0, market_cost: 0 }), + expect.objectContaining({ provider: 'openrouter', user_byok: false }) + ); + }); + + it('returns the authentication failure and requests the gateway audience', async () => { + const authFailedResponse = NextResponse.json( + { success: false as const, error: 'Unauthorized' }, + { status: 401 } + ); + jest.mocked(getUserFromAuth).mockResolvedValue({ user: null, authFailedResponse }); + + expect(await POST(makeRequest())).toBe(authFailedResponse); + expect(getUserFromAuth).toHaveBeenCalledWith({ + adminOnly: false, + expectedAudience: KILO_GATEWAY_AUDIENCE, + }); + expect(getBalanceAndOrgSettings).not.toHaveBeenCalled(); + expect(mockedFetch).not.toHaveBeenCalled(); + expect(after).not.toHaveBeenCalled(); + }); + + it('rate limits before authentication or upstream work', async () => { + jest.mocked(isGatewayAccountRateLimited).mockResolvedValue(true); + const request = makeRequest(requestBody, { 'x-forwarded-for': ' 192.0.2.1, 192.0.2.2' }); + + const response = await POST(request); + + expect(response.status).toBe(429); + expect(await response.json()).toMatchObject({ error_type: 'rate_limit_exceeded' }); + expect(gatewayRateLimitKey).toHaveBeenCalledWith(request.headers, '192.0.2.1'); + expect(isGatewayAccountRateLimited).toHaveBeenCalledWith(request, 'test-rate-limit-key'); + expect(getUserFromAuth).not.toHaveBeenCalled(); + expect(mockedFetch).not.toHaveBeenCalled(); + }); + + it.each([ + ['malformed JSON', '{'], + ['another model', JSON.stringify({ ...requestBody, model: 'openai/gpt-4o' })], + ['an unpinned alias', JSON.stringify({ ...requestBody, model: 'jev-latest' })], + ['empty questions', JSON.stringify({ ...requestBody, questions: {} })], + ])('rejects %s before balance checks or upstream work', async (_name, body) => { + const response = await POST( + new NextRequest(routeUrl, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body, + }) + ); + + expect(response.status).toBe(400); + expect(await response.json()).toMatchObject({ error_type: 'invalid_request' }); + expect(getBalanceAndOrgSettings).not.toHaveBeenCalled(); + expect(mockedFetch).not.toHaveBeenCalled(); + expect(after).not.toHaveBeenCalled(); + }); + + it.each([0, -1])('blocks balance %s using the shared credits response', async balance => { + setAuth('org-123'); + jest.mocked(getBalanceAndOrgSettings).mockResolvedValue({ + balance, + balanceLimitedByUserAllowance: true, + }); + const blockedResponse = NextResponse.json({ error_type: 'credits_exhausted' }, { status: 402 }); + jest.mocked(creditsBlockedResponse).mockResolvedValue(blockedResponse); + + expect(await POST(makeRequest())).toBe(blockedResponse); + expect(getBalanceAndOrgSettings).toHaveBeenCalledWith('org-123', user); + expect(creditsBlockedResponse).toHaveBeenCalledWith({ + user, + balance, + organizationId: 'org-123', + balanceLimitedByUserAllowance: true, + }); + expect(checkOrganizationModelRestrictions).not.toHaveBeenCalled(); + expect(mockedFetch).not.toHaveBeenCalled(); + expect(after).not.toHaveBeenCalled(); + }); + + it('honors organization model restrictions before resolving member access', async () => { + setAuth('org-123'); + const settings = { model_deny_list: [TYPESAFE_MODEL] }; + jest.mocked(getBalanceAndOrgSettings).mockResolvedValue({ + balance: 1000, + settings, + plan: 'enterprise', + }); + const error = NextResponse.json({ error_type: 'model_not_allowed' }, { status: 404 }); + jest.mocked(checkOrganizationModelRestrictions).mockReturnValue({ error }); + + expect(await POST(makeRequest())).toBe(error); + expect(checkOrganizationModelRestrictions).toHaveBeenCalledWith({ + modelId: TYPESAFE_MODEL, + settings, + organizationPlan: 'enterprise', + }); + expect(resolveOrganizationMemberModelDecision).not.toHaveBeenCalled(); + expect(mockedFetch).not.toHaveBeenCalled(); + expect(after).not.toHaveBeenCalled(); + }); + + it('denies a model rejected by the effective organization member decision', async () => { + setAuth('org-123'); + jest.mocked(resolveOrganizationMemberModelDecision).mockResolvedValue({ + ...memberDecision, + decision: { allowed: false }, + }); + + expect((await POST(makeRequest())).status).toBe(404); + expect(resolveOrganizationMemberModelDecision).toHaveBeenCalledWith({ + organizationId: 'org-123', + kiloUserId: user.id, + modelId: TYPESAFE_MODEL, + providerLookup: expect.any(Function), + }); + const [{ providerLookup }] = jest.mocked(resolveOrganizationMemberModelDecision).mock.calls[0]; + if (!providerLookup) throw new Error('Expected the fixed TypeSafe provider lookup'); + await expect(providerLookup(TYPESAFE_MODEL)).resolves.toEqual(new Set(['typesafe'])); + expect(modelNotAllowedResponse).toHaveBeenCalledTimes(1); + expect(mockedFetch).not.toHaveBeenCalled(); + expect(after).not.toHaveBeenCalled(); + }); + + it.each([ + { + only: ['typesafe', 'other'], + eligible: ['typesafe', 'outside-ceiling'], + expected: ['typesafe'], + }, + { only: undefined, eligible: ['typesafe'], expected: ['typesafe'] }, + { only: ['typesafe'], eligible: undefined, expected: ['typesafe'] }, + ])('applies the provider policy intersection: %j', async ({ only, eligible, expected }) => { + setAuth('org-123'); + jest.mocked(checkOrganizationModelRestrictions).mockReturnValue({ + error: null, + providerConfig: { only, data_collection: 'deny' }, + }); + jest.mocked(resolveOrganizationMemberModelDecision).mockResolvedValue({ + ...memberDecision, + decision: { allowed: true, eligibleProviderRoutes: eligible ? new Set(eligible) : undefined }, + }); + + const response = await POST(makeRequest({ ...requestBody, provider: { only: ['attacker'] } })); + + expect(response.status).toBe(200); + expect(upstreamRequest().body.provider).toEqual({ only: expected, data_collection: 'deny' }); + }); + + it.each([ + { only: ['other'], eligible: ['typesafe'] }, + { only: [], eligible: ['typesafe'] }, + { only: undefined, eligible: [] }, + ])('rejects an empty provider policy intersection: %j', async ({ only, eligible }) => { + setAuth('org-123'); + jest.mocked(checkOrganizationModelRestrictions).mockReturnValue({ + error: null, + providerConfig: { only, data_collection: 'deny' }, + }); + jest.mocked(resolveOrganizationMemberModelDecision).mockResolvedValue({ + ...memberDecision, + decision: { allowed: true, eligibleProviderRoutes: new Set(eligible) }, + }); + + expect((await POST(makeRequest())).status).toBe(404); + expect(modelNotAllowedResponse).toHaveBeenCalledTimes(1); + expect(mockedFetch).not.toHaveBeenCalled(); + expect(after).not.toHaveBeenCalled(); + }); + + it('maps OpenRouter credit exhaustion to a service error without charging the user', async () => { + mockedFetch.mockResolvedValue( + Response.json({ error: 'upstream credits exhausted' }, { status: 402 }) + ); + + const response = await POST(makeRequest()); + + expect(response.status).toBe(503); + expect(await response.json()).toEqual({ + message: 'Service temporarily unavailable', + error_type: 'upstream_error', + }); + expect(captureMessage).toHaveBeenCalledWith('OpenRouter System One balance exhausted', { + level: 'error', + }); + expect(wrapInSafeNextResponse).not.toHaveBeenCalled(); + expect(after).not.toHaveBeenCalled(); + expect(logMicrodollarUsage).not.toHaveBeenCalled(); + }); + + it.each([400, 429, 500])( + 'delegates upstream HTTP %s to the safe response wrapper', + async status => { + const upstream = Response.json({ error: 'upstream failure' }, { status }); + const safeResponse = NextResponse.json({ error: 'upstream failure' }, { status }); + mockedFetch.mockResolvedValue(upstream); + jest.mocked(wrapInSafeNextResponse).mockReturnValue(safeResponse); + + expect(await POST(makeRequest())).toBe(safeResponse); + expect(wrapInSafeNextResponse).toHaveBeenCalledWith(upstream); + expect(after).not.toHaveBeenCalled(); + expect(logMicrodollarUsage).not.toHaveBeenCalled(); + } + ); + + it.each(['network failure', 'non-JSON success'])('handles %s without charging', async failure => { + if (failure === 'network failure') { + mockedFetch.mockRejectedValue(new Error('connection failed')); + } else { + mockedFetch.mockResolvedValue(new Response('not JSON', { status: 200 })); + } + + const response = await POST(makeRequest()); + + expect(response.status).toBe(502); + expect(await response.json()).toEqual({ + message: 'Upstream request failed', + error_type: 'upstream_error', + }); + expect(captureException).toHaveBeenCalledWith(expect.any(Error), { + tags: { source: 'systemone_proxy' }, + }); + expect(after).not.toHaveBeenCalled(); + expect(logMicrodollarUsage).not.toHaveBeenCalled(); + }); + + it.each([ + ['missing usage', undefined], + ['null usage', null], + ['missing cost', { input_tokens: 23, output_tokens: 7 }], + ['string cost', { ...upstreamBody.usage, cost: '0.01' }], + ['negative cost', { ...upstreamBody.usage, cost: -0.01 }], + ['missing input tokens', { output_tokens: 7, cost: 0.01 }], + ['negative tokens', { ...upstreamBody.usage, input_tokens: -1 }], + ['fractional tokens', { ...upstreamBody.usage, output_tokens: 1.5 }], + ])('rejects upstream %s without scheduling billing', async (_name, usage) => { + mockedFetch.mockResolvedValue(Response.json({ ...upstreamBody, usage })); + + const response = await POST(makeRequest()); + + expect(response.status).toBe(502); + expect(await response.json()).toEqual({ + message: 'Invalid upstream response', + error_type: 'upstream_error', + }); + expect(captureMessage).toHaveBeenCalledWith( + 'Invalid OpenRouter System One response or missing usage', + { level: 'error' } + ); + expect(after).not.toHaveBeenCalled(); + expect(logMicrodollarUsage).not.toHaveBeenCalled(); + }); +}); diff --git a/apps/web/src/app/api/gateway/typesafe/v1/systemone/route.ts b/apps/web/src/app/api/gateway/typesafe/v1/systemone/route.ts new file mode 100644 index 0000000000..3d8c0d0551 --- /dev/null +++ b/apps/web/src/app/api/gateway/typesafe/v1/systemone/route.ts @@ -0,0 +1,191 @@ +import { after, NextResponse, type NextRequest } from 'next/server'; +import { captureException, captureMessage } from '@sentry/nextjs'; +import { KILO_GATEWAY_AUDIENCE } from '@kilocode/worker-utils/internal-service-token-audiences'; +import { getUserFromAuth } from '@/lib/user/server'; +import { getBalanceAndOrgSettings } from '@/lib/organizations/organization-usage'; +import { resolveOrganizationMemberModelDecision } from '@/lib/organizations/effective-model-access.server'; +import { + gatewayRateLimitKey, + isGatewayAccountRateLimited, +} from '@/lib/ai-gateway/gateway-account-rate-limit'; +import { + checkOrganizationModelRestrictions, + creditsBlockedResponse, + extractFraudAndProjectHeaders, + extractHeaderAndLimitLength, + modelNotAllowedResponse, + wrapInSafeNextResponse, +} from '@/lib/ai-gateway/llm-proxy-helpers'; +import { OPENROUTER } from '@/lib/ai-gateway/providers/definitions/openrouter'; +import { ATTRIBUTION_HEADERS } from '@/lib/ai-gateway/providers/openrouter/attribution-headers'; +import { generateProviderSpecificHash } from '@/lib/ai-gateway/providerHash'; +import { logMicrodollarUsage } from '@/lib/ai-gateway/processUsage'; +import { + systemOneRequestSchema, + systemOneResponseSchema, + TYPESAFE_MODEL, +} from '@/lib/ai-gateway/typesafe'; +import { FEATURE_HEADER, validateFeatureHeader } from '@/lib/feature-detection'; +import { toMicrodollars } from '@/lib/utils'; + +export const maxDuration = 800; + +function errorResponse(message: string, error_type: string, status: number) { + return NextResponse.json({ message, error_type }, { status }); +} + +export async function POST(request: NextRequest) { + const startedAt = performance.now(); + const ipAddress = request.headers.get('x-forwarded-for')?.split(',')[0]?.trim(); + if (await isGatewayAccountRateLimited(request, gatewayRateLimitKey(request.headers, ipAddress))) { + return errorResponse('Rate limit exceeded', 'rate_limit_exceeded', 429); + } + + const { user, authFailedResponse, organizationId, botId, tokenSource } = await getUserFromAuth({ + adminOnly: false, + expectedAudience: KILO_GATEWAY_AUDIENCE, + }); + if (authFailedResponse) return authFailedResponse; + if (!user) return errorResponse('Authentication required', 'authentication_error', 401); + + let body: unknown; + try { + body = await request.json(); + } catch { + return errorResponse('Invalid JSON body', 'invalid_request', 400); + } + const parsed = systemOneRequestSchema.safeParse(body); + if (!parsed.success) { + return errorResponse( + parsed.error.issues.map(issue => `${issue.path.join('.')}: ${issue.message}`).join('; '), + 'invalid_request', + 400 + ); + } + + const { balance, settings, plan, balanceLimitedByUserAllowance } = await getBalanceAndOrgSettings( + organizationId, + user + ); + if (balance <= 0) { + return creditsBlockedResponse({ user, balance, organizationId, balanceLimitedByUserAllowance }); + } + + const { error, providerConfig } = checkOrganizationModelRestrictions({ + modelId: TYPESAFE_MODEL, + settings, + organizationPlan: plan, + }); + if (error) return error; + let providerPolicy = providerConfig; + if (organizationId) { + const { decision } = await resolveOrganizationMemberModelDecision({ + organizationId, + kiloUserId: user.id, + modelId: TYPESAFE_MODEL, + providerLookup: async () => new Set(['typesafe']), + }); + if (!decision.allowed) return modelNotAllowedResponse(); + if (decision.eligibleProviderRoutes) { + const only = providerConfig?.only + ? providerConfig.only.filter(route => decision.eligibleProviderRoutes?.has(route)) + : [...decision.eligibleProviderRoutes]; + if (only.length === 0) return modelNotAllowedResponse(); + providerPolicy = { ...providerConfig, only }; + } + } + + let response: Response; + let responseBody: unknown; + let ttfbMs: number; + try { + response = await fetch(`${OPENROUTER.apiUrl}/systemone`, { + method: 'POST', + headers: { + ...ATTRIBUTION_HEADERS, + Authorization: `Bearer ${OPENROUTER.apiKey}`, + 'Content-Type': 'application/json', + }, + body: JSON.stringify({ + ...parsed.data, + provider: providerPolicy, + user: generateProviderSpecificHash(user.id, OPENROUTER), + }), + signal: AbortSignal.any([request.signal, AbortSignal.timeout(10 * 60 * 1000)]), + }); + ttfbMs = Math.max(0, Math.round(performance.now() - startedAt)); + if (response.status === 402) { + captureMessage('OpenRouter System One balance exhausted', { level: 'error' }); + return errorResponse('Service temporarily unavailable', 'upstream_error', 503); + } + if (!response.ok) return wrapInSafeNextResponse(response); + responseBody = await response.json(); + } catch (error) { + captureException(error, { tags: { source: 'systemone_proxy' } }); + return errorResponse('Upstream request failed', 'upstream_error', 502); + } + + const result = systemOneResponseSchema.safeParse(responseBody); + if (!result.success) { + captureMessage('Invalid OpenRouter System One response or missing usage', { level: 'error' }); + return errorResponse('Invalid upstream response', 'upstream_error', 502); + } + + const { id, model, provider, usage } = result.data; + const { fraudHeaders, projectId } = extractFraudAndProjectHeaders(request); + const cost = toMicrodollars(usage.cost); + after(async () => { + await logMicrodollarUsage( + { + messageId: id, + model, + responseContent: '', + hasError: false, + inference_provider: provider ?? 'TypeSafe', + upstream_id: null, + finish_reason: null, + latency: ttfbMs, + moderation_latency: null, + generation_time: null, + streamed: false, + cancelled: false, + status_code: response.status, + cost_mUsd: cost, + market_cost: cost, + inputTokens: usage.input_tokens, + outputTokens: usage.output_tokens, + cacheWriteTokens: 0, + cacheHitTokens: 0, + is_byok: false, + }, + { + api_kind: 'systemone', + kiloUserId: user.id, + provider: 'openrouter', + requested_model: TYPESAFE_MODEL, + promptInfo: { system_prompt_prefix: '', system_prompt_length: 0, user_prompt_prefix: '' }, + max_tokens: null, + has_middle_out_transform: null, + fraudHeaders, + isStreaming: false, + organizationId, + prior_microdollar_usage: user.microdollars_used, + posthog_distinct_id: user.google_user_email, + project_id: projectId, + status_code: response.status, + editor_name: extractHeaderAndLimitLength(request, 'x-kilocode-editorname'), + machine_id: extractHeaderAndLimitLength(request, 'x-kilocode-machineid'), + user_byok: false, + has_tools: false, + botId, + tokenSource, + feature: validateFeatureHeader(request.headers.get(FEATURE_HEADER) || ''), + session_id: extractHeaderAndLimitLength(request, 'X-KiloCode-TaskId'), + mode: extractHeaderAndLimitLength(request, 'x-kilocode-mode'), + auto_model: null, + ttfb_ms: ttfbMs, + } + ); + }); + return NextResponse.json(responseBody); +} diff --git a/apps/web/src/lib/ai-gateway/typesafe.test.ts b/apps/web/src/lib/ai-gateway/typesafe.test.ts new file mode 100644 index 0000000000..fb413496ff --- /dev/null +++ b/apps/web/src/lib/ai-gateway/typesafe.test.ts @@ -0,0 +1,110 @@ +import { describe, expect, it } from '@jest/globals'; +import { choice, noul, score } from '@typesafe-ai/sdk'; +import { systemOneRequestSchema, systemOneResponseSchema, TYPESAFE_MODEL } from './typesafe'; + +describe('systemOneRequestSchema', () => { + it('accepts SDK question helpers with structured entries and strips routing overrides', () => { + const request = { + state: { text: 'A support request', metadata: [1, true, null] }, + questions: { + relevant: noul({ prompt: 'Is this relevant?' }, { true: ['Yes'], false: null }), + category: choice(null, { billing: { description: 'Payments' }, other: null }), + urgency: score(['Assess urgency'], [null, 'Medium', { description: 'High' }]), + }, + }; + + expect( + systemOneRequestSchema.parse({ + ...request, + model: 'jev-1.13', + provider: { only: ['attacker'] }, + api_key: 'test-key', + user: 'attacker', + }) + ).toEqual({ ...request, model: TYPESAFE_MODEL }); + }); + + it.each(['text', { nested: [true, 1, null] }, ['text', 1, false], null])( + 'accepts SDK state entries: %j', + state => { + expect( + systemOneRequestSchema.safeParse({ state, questions: { relevant: { type: 'noul' } } }) + .success + ).toBe(true); + } + ); + + it.each([ + ['missing state', { questions: { relevant: { type: 'noul' } } }], + ['numeric state', { state: 1, questions: { relevant: { type: 'noul' } } }], + ['boolean state', { state: true, questions: { relevant: { type: 'noul' } } }], + ['unknown question type', { state: null, questions: { q: { type: 'boolean' } } }], + ['missing choice criteria', { state: null, questions: { q: { type: 'choice' } } }], + [ + 'array choice criteria', + { state: null, questions: { q: { type: 'choice', criteria: ['a', 'b'] } } }, + ], + ['short score rubric', { state: null, questions: { q: { type: 'score', criteria: ['Low'] } } }], + [ + 'object score rubric', + { + state: null, + questions: { q: { type: 'score', criteria: { '0': 'Low', '1': 'High' } } }, + }, + ], + ])('rejects %s', (_name, request) => { + expect(systemOneRequestSchema.safeParse(request).success).toBe(false); + }); +}); + +describe('systemOneResponseSchema', () => { + const response = { + id: 'gen-123', + model: TYPESAFE_MODEL, + answers: { relevant: { type: 'noul', noul: 0.5 } }, + usage: { input_tokens: 1, output_tokens: 2, cost: 0.00001 }, + }; + + it('retains upstream metadata, usage extensions, and answer extensions', () => { + const extended = { + ...response, + provider: 'TypeSafe', + metadata: { request_id: 'upstream-123' }, + answers: { relevant: { type: 'noul', noul: 1, explanation: 'Relevant' } }, + usage: { input_tokens: 0, output_tokens: 0, cost: 0, total_tokens: 0 }, + }; + + expect(systemOneResponseSchema.parse(extended)).toEqual(extended); + }); + + it.each([NaN, Infinity, -Infinity])('rejects non-finite upstream cost %s', cost => { + expect( + systemOneResponseSchema.safeParse({ ...response, usage: { ...response.usage, cost } }).success + ).toBe(false); + }); + + it.each([ + ['missing generation ID', { ...response, id: undefined }], + ['empty generation ID', { ...response, id: '' }], + ['missing model', { ...response, model: undefined }], + ['empty model', { ...response, model: '' }], + ['missing answers', { ...response, answers: undefined }], + ])('rejects %s', (_name, value) => { + expect(systemOneResponseSchema.safeParse(value).success).toBe(false); + }); + + it.each([ + { type: 'noul', noul: -0.1 }, + { type: 'noul', noul: 1.1 }, + { type: 'noul', noul: '0.5' }, + { type: 'unknown', noul: 0.5 }, + { type: 'choice', choice: 'yes', confidence: 1.1, probabilities: { yes: 1 } }, + { type: 'choice', choice: 'yes', confidence: 1, probabilities: { yes: -0.1 } }, + { type: 'score', score: -1, confidence: 1, probabilities: { '0': 1 }, legend: { '0': 'Low' } }, + { type: 'score', score: 0, confidence: 1, probabilities: { '0': 1 } }, + ])('rejects malformed answers: %j', answer => { + expect( + systemOneResponseSchema.safeParse({ ...response, answers: { relevant: answer } }).success + ).toBe(false); + }); +}); diff --git a/apps/web/src/lib/ai-gateway/typesafe.ts b/apps/web/src/lib/ai-gateway/typesafe.ts new file mode 100644 index 0000000000..cc291ec7ea --- /dev/null +++ b/apps/web/src/lib/ai-gateway/typesafe.ts @@ -0,0 +1,88 @@ +import type { + EntryType, + Question, + Questions, + SystemOneRequestPayload, + SystemOneResult, +} from '@typesafe-ai/sdk'; +import { z } from 'zod'; + +export const TYPESAFE_MODEL = 'typesafe/jev-1.13'; + +const entrySchema = z.union([ + z.string(), + z.record(z.string(), z.json()), + z.array(z.json()), + z.null(), +]) satisfies z.ZodType; + +const questionSchema = z.discriminatedUnion('type', [ + z.object({ + type: z.literal('noul'), + instructions: entrySchema.optional(), + criteria: z + .object({ true: entrySchema.optional(), false: entrySchema.optional() }) + .nullable() + .optional(), + }), + z.object({ + type: z.literal('choice'), + instructions: entrySchema.optional(), + criteria: z.record(z.string(), entrySchema), + }), + z.object({ + type: z.literal('score'), + instructions: entrySchema.optional(), + criteria: z.tuple([entrySchema, entrySchema]).rest(entrySchema), + }), +]) satisfies z.ZodType; + +export const systemOneRequestSchema = z.object({ + model: z + .enum([TYPESAFE_MODEL, 'jev-1.13']) + .default(TYPESAFE_MODEL) + .transform(() => TYPESAFE_MODEL), + state: entrySchema, + questions: z + .record(z.string(), questionSchema) + .refine(questions => Object.keys(questions).length > 0, 'At least one question is required'), +}) satisfies z.ZodType; + +type OpenRouterSystemOneResult = SystemOneResult & { + id: string; + provider?: string; + usage: SystemOneResult['usage'] & { cost: number }; +}; + +const probability = z.number().min(0).max(1); +const probabilities = z.record(z.string(), probability); + +export const systemOneResponseSchema = z.looseObject({ + id: z.string().min(1), + model: z.string().min(1), + provider: z.string().optional(), + answers: z.record( + z.string(), + z.discriminatedUnion('type', [ + z.looseObject({ type: z.literal('noul'), noul: probability }), + z.looseObject({ + type: z.literal('choice'), + choice: z.string(), + confidence: probability, + probabilities, + }), + z.looseObject({ + type: z.literal('score'), + score: z.number().nonnegative(), + confidence: probability, + probabilities, + legend: z.record(z.string(), entrySchema), + }), + ]) + ), + usage: z.looseObject({ + input_tokens: z.number().int().nonnegative(), + output_tokens: z.number().int().nonnegative(), + cost: z.number().nonnegative(), + }), +}) satisfies z.ZodType; diff --git a/packages/db/src/schema-types.ts b/packages/db/src/schema-types.ts index 0376657896..72cc7b878b 100644 --- a/packages/db/src/schema-types.ts +++ b/packages/db/src/schema-types.ts @@ -1278,6 +1278,7 @@ export const GatewayApiKindSchema = z.enum([ 'messages', 'responses', 'audio_transcriptions', + 'systemone', ]); export type GatewayApiKind = z.infer; diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 6452671f49..b0117c28a5 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -1060,6 +1060,9 @@ importers: '@types/mdx': specifier: 2.0.13 version: 2.0.13 + '@typesafe-ai/sdk': + specifier: 0.6.0 + version: 0.6.0 '@upstash/redis': specifier: ^1.38.0 version: 1.38.0 @@ -10625,6 +10628,10 @@ packages: '@types/yargs@17.0.35': resolution: {integrity: sha512-qUHkeCyQFxMXg79wQfTtfndEC+N9ZZg76HJftDJp+qH2tV7Gj4OJi7l+PiWwJ+pWtW8GwSmqsDj/oymhrTWXjg==} + '@typesafe-ai/sdk@0.6.0': + resolution: {integrity: sha512-IddX+Q0XM+VagOUZFeP7wZjaO4SHMdvnh2zEBdrZZnXedWI3BNK1lKhMx3ayrkFWvVLbVcUHJy6AVZlY+e6Jaw==} + engines: {node: '>=20'} + '@typescript-eslint/project-service@8.67.0': resolution: {integrity: sha512-cvE8c7ulYeXN9fYuszhCeCsbzyVEXuhrRCybnBre7TUmqb5nRmBfQAwCj0O3WJFDeyAZt4VYv51vMCC9LHSdYw==} engines: {node: ^18.18.0 || ^20.9.0 || >=21.1.0} @@ -28565,6 +28572,8 @@ snapshots: dependencies: '@types/yargs-parser': 21.0.3 + '@typesafe-ai/sdk@0.6.0': {} + '@typescript-eslint/project-service@8.67.0(typescript@5.9.3)': dependencies: '@typescript-eslint/tsconfig-utils': 8.67.0(typescript@5.9.3) From bb7b126ab447c6504ff795105a5a09a50b2d9dfb Mon Sep 17 00:00:00 2001 From: Christiaan Arnoldus Date: Wed, 23 Sep 2026 11:24:38 +0200 Subject: [PATCH 2/6] test(ai-gateway): fix TypeSafe route test fixtures --- .../typesafe/v1/systemone/route.test.ts | 38 ++++++++++++++----- 1 file changed, 29 insertions(+), 9 deletions(-) diff --git a/apps/web/src/app/api/gateway/typesafe/v1/systemone/route.test.ts b/apps/web/src/app/api/gateway/typesafe/v1/systemone/route.test.ts index 78470c4cbb..f5935d0771 100644 --- a/apps/web/src/app/api/gateway/typesafe/v1/systemone/route.test.ts +++ b/apps/web/src/app/api/gateway/typesafe/v1/systemone/route.test.ts @@ -1,9 +1,10 @@ -import { afterEach, beforeEach, describe, expect, it, jest } from '@jest/globals'; +import { afterEach, beforeEach, describe, expect, it } from '@jest/globals'; import { captureException, captureMessage } from '@sentry/nextjs'; import { TypeSafeClient, choice, noul, score } from '@typesafe-ai/sdk'; import type { User } from '@kilocode/db/schema'; import { KILO_GATEWAY_AUDIENCE } from '@kilocode/worker-utils/internal-service-token-audiences'; import { after, NextRequest, NextResponse } from 'next/server'; +import type * as NextServer from 'next/server'; import { getUserFromAuth } from '@/lib/user/server'; import { getBalanceAndOrgSettings } from '@/lib/organizations/organization-usage'; import { resolveOrganizationMemberModelDecision } from '@/lib/organizations/effective-model-access.server'; @@ -27,7 +28,7 @@ import { EmptyFraudDetectionHeaders } from '@/lib/utils'; import { POST } from './route'; jest.mock('next/server', () => ({ - ...jest.requireActual('next/server'), + ...jest.requireActual('next/server'), after: jest.fn(), })); jest.mock('@sentry/nextjs', () => ({ @@ -107,7 +108,7 @@ const memberDecision = { }, decision: { allowed: true }, }; -const mockedFetch = jest.fn(); +const mockedFetch = jest.fn() as jest.MockedFunction; const originalFetch = globalThis.fetch; function makeRequest(body: unknown = requestBody, headers: Record = {}) { @@ -167,7 +168,14 @@ describe('POST /api/gateway/typesafe/v1/systemone', () => { jest .mocked(modelNotAllowedResponse) .mockImplementation(() => - NextResponse.json({ error_type: 'model_not_allowed' }, { status: 404 }) + NextResponse.json( + { + error: 'Model not allowed', + error_type: 'model_not_allowed', + message: 'Model not allowed', + }, + { status: 404 } + ) ); mockedFetch.mockImplementation(async () => Response.json(upstreamBody)); }); @@ -178,7 +186,7 @@ describe('POST /api/gateway/typesafe/v1/systemone', () => { it('supports the TypeSafe SDK contract through a custom fetch transport', async () => { const routeFetch = jest.fn(async (input: string, init?: RequestInit) => - POST(new NextRequest(input, init)) + POST(new NextRequest(new Request(input, init))) ); const client = new TypeSafeClient({ apiKey: 'test-kilo-key', @@ -396,7 +404,18 @@ describe('POST /api/gateway/typesafe/v1/systemone', () => { balance, balanceLimitedByUserAllowance: true, }); - const blockedResponse = NextResponse.json({ error_type: 'credits_exhausted' }, { status: 402 }); + const blockedResponse = NextResponse.json( + { + error_type: 'usage_limit_exceeded' as const, + error: { + title: 'Credits exhausted', + message: 'Credits exhausted', + balance, + buyCreditsUrl: '', + }, + }, + { status: 402 } + ); jest.mocked(creditsBlockedResponse).mockResolvedValue(blockedResponse); expect(await POST(makeRequest())).toBe(blockedResponse); @@ -551,9 +570,10 @@ describe('POST /api/gateway/typesafe/v1/systemone', () => { message: 'Upstream request failed', error_type: 'upstream_error', }); - expect(captureException).toHaveBeenCalledWith(expect.any(Error), { - tags: { source: 'systemone_proxy' }, - }); + expect(captureException).toHaveBeenCalledWith( + expect.objectContaining({ message: expect.any(String) }), + { tags: { source: 'systemone_proxy' } } + ); expect(after).not.toHaveBeenCalled(); expect(logMicrodollarUsage).not.toHaveBeenCalled(); }); From 48b9dff4c7ffdd4b7b5880c73218a1f8646f1040 Mon Sep 17 00:00:00 2001 From: Christiaan Arnoldus Date: Wed, 23 Sep 2026 11:26:53 +0200 Subject: [PATCH 3/6] test(ai-gateway): cover System One regional billing contract --- .../typesafe/v1/systemone/route.test.ts | 22 +++++++++---------- .../ai-gateway/usage-record-contract.test.ts | 6 +++++ 2 files changed, 16 insertions(+), 12 deletions(-) diff --git a/apps/web/src/app/api/gateway/typesafe/v1/systemone/route.test.ts b/apps/web/src/app/api/gateway/typesafe/v1/systemone/route.test.ts index f5935d0771..dcecad2755 100644 --- a/apps/web/src/app/api/gateway/typesafe/v1/systemone/route.test.ts +++ b/apps/web/src/app/api/gateway/typesafe/v1/systemone/route.test.ts @@ -165,18 +165,16 @@ describe('POST /api/gateway/typesafe/v1/systemone', () => { jest .mocked(extractHeaderAndLimitLength) .mockImplementation((request, name) => request.headers.get(name)); - jest - .mocked(modelNotAllowedResponse) - .mockImplementation(() => - NextResponse.json( - { - error: 'Model not allowed', - error_type: 'model_not_allowed', - message: 'Model not allowed', - }, - { status: 404 } - ) - ); + jest.mocked(modelNotAllowedResponse).mockImplementation(() => + NextResponse.json( + { + error: 'Model not allowed', + error_type: 'model_not_allowed', + message: 'Model not allowed', + }, + { status: 404 } + ) + ); mockedFetch.mockImplementation(async () => Response.json(upstreamBody)); }); diff --git a/apps/web/src/lib/ai-gateway/usage-record-contract.test.ts b/apps/web/src/lib/ai-gateway/usage-record-contract.test.ts index 30929276d0..1c20e86287 100644 --- a/apps/web/src/lib/ai-gateway/usage-record-contract.test.ts +++ b/apps/web/src/lib/ai-gateway/usage-record-contract.test.ts @@ -146,6 +146,12 @@ describe('UsageRecordRequestSchema', () => { expect(UsageRecordRequestSchema.safeParse(request).success).toBe(true); }); + test('accepts System One usage for primary-region billing', () => { + const request = validRequest(); + request.metadata.api_kind = 'systemone'; + expect(UsageRecordRequestSchema.parse(request).metadata.api_kind).toBe('systemone'); + }); + test('rejects an unknown api_kind', () => { const request = validRequest(); request.metadata.api_kind = 'not_a_real_kind'; From eed8fd2ff10753220b37257f4b8d454c06d5a7b8 Mon Sep 17 00:00:00 2001 From: Christiaan Arnoldus Date: Wed, 23 Sep 2026 11:38:29 +0200 Subject: [PATCH 4/6] refactor(ai-gateway): extract TypeSafe request handler --- .../typesafe/v1/systemone/route.test.ts | 40 ++-- .../gateway/typesafe/v1/systemone/route.ts | 190 +----------------- .../src/lib/ai-gateway/typesafe/handler.ts | 186 +++++++++++++++++ .../schemas.test.ts} | 2 +- .../{typesafe.ts => typesafe/schemas.ts} | 0 5 files changed, 213 insertions(+), 205 deletions(-) create mode 100644 apps/web/src/lib/ai-gateway/typesafe/handler.ts rename apps/web/src/lib/ai-gateway/{typesafe.test.ts => typesafe/schemas.test.ts} (99%) rename apps/web/src/lib/ai-gateway/{typesafe.ts => typesafe/schemas.ts} (100%) diff --git a/apps/web/src/app/api/gateway/typesafe/v1/systemone/route.test.ts b/apps/web/src/app/api/gateway/typesafe/v1/systemone/route.test.ts index dcecad2755..2dc30269c6 100644 --- a/apps/web/src/app/api/gateway/typesafe/v1/systemone/route.test.ts +++ b/apps/web/src/app/api/gateway/typesafe/v1/systemone/route.test.ts @@ -1,5 +1,6 @@ import { afterEach, beforeEach, describe, expect, it } from '@jest/globals'; -import { captureException, captureMessage } from '@sentry/nextjs'; +import { z } from 'zod'; +import { errorExceptInTest } from '@/lib/utils.server'; import { TypeSafeClient, choice, noul, score } from '@typesafe-ai/sdk'; import type { User } from '@kilocode/db/schema'; import { KILO_GATEWAY_AUDIENCE } from '@kilocode/worker-utils/internal-service-token-audiences'; @@ -23,7 +24,7 @@ import { import { OPENROUTER } from '@/lib/ai-gateway/providers/definitions/openrouter'; import { generateProviderSpecificHash } from '@/lib/ai-gateway/providerHash'; import { logMicrodollarUsage } from '@/lib/ai-gateway/processUsage'; -import { TYPESAFE_MODEL } from '@/lib/ai-gateway/typesafe'; +import { systemOneRequestSchema, TYPESAFE_MODEL } from '@/lib/ai-gateway/typesafe/schemas'; import { EmptyFraudDetectionHeaders } from '@/lib/utils'; import { POST } from './route'; @@ -31,10 +32,7 @@ jest.mock('next/server', () => ({ ...jest.requireActual('next/server'), after: jest.fn(), })); -jest.mock('@sentry/nextjs', () => ({ - captureException: jest.fn(), - captureMessage: jest.fn(), -})); +jest.mock('@/lib/utils.server', () => ({ errorExceptInTest: jest.fn() })); jest.mock('@/lib/user/server', () => ({ getUserFromAuth: jest.fn() })); jest.mock('@/lib/organizations/organization-usage', () => ({ getBalanceAndOrgSettings: jest.fn(), @@ -396,6 +394,21 @@ describe('POST /api/gateway/typesafe/v1/systemone', () => { expect(after).not.toHaveBeenCalled(); }); + it('formats validation errors with Zod while preserving the TypeSafe error shape', async () => { + const body = { ...requestBody, questions: { refund: { type: 'invalid' } } }; + const validation = systemOneRequestSchema.safeParse(body); + if (validation.success) throw new Error('Expected an invalid question'); + + const response = await POST(makeRequest(body)); + + expect(response.status).toBe(400); + expect(await response.json()).toEqual({ + message: z.prettifyError(validation.error), + error_type: 'invalid_request', + }); + expect(mockedFetch).not.toHaveBeenCalled(); + }); + it.each([0, -1])('blocks balance %s using the shared credits response', async balance => { setAuth('org-123'); jest.mocked(getBalanceAndOrgSettings).mockResolvedValue({ @@ -531,9 +544,7 @@ describe('POST /api/gateway/typesafe/v1/systemone', () => { message: 'Service temporarily unavailable', error_type: 'upstream_error', }); - expect(captureMessage).toHaveBeenCalledWith('OpenRouter System One balance exhausted', { - level: 'error', - }); + expect(errorExceptInTest).toHaveBeenCalledWith('OpenRouter System One balance exhausted'); expect(wrapInSafeNextResponse).not.toHaveBeenCalled(); expect(after).not.toHaveBeenCalled(); expect(logMicrodollarUsage).not.toHaveBeenCalled(); @@ -568,9 +579,9 @@ describe('POST /api/gateway/typesafe/v1/systemone', () => { message: 'Upstream request failed', error_type: 'upstream_error', }); - expect(captureException).toHaveBeenCalledWith( - expect.objectContaining({ message: expect.any(String) }), - { tags: { source: 'systemone_proxy' } } + expect(errorExceptInTest).toHaveBeenCalledWith( + 'OpenRouter System One request failed', + expect.objectContaining({ message: expect.any(String) }) ); expect(after).not.toHaveBeenCalled(); expect(logMicrodollarUsage).not.toHaveBeenCalled(); @@ -595,9 +606,8 @@ describe('POST /api/gateway/typesafe/v1/systemone', () => { message: 'Invalid upstream response', error_type: 'upstream_error', }); - expect(captureMessage).toHaveBeenCalledWith( - 'Invalid OpenRouter System One response or missing usage', - { level: 'error' } + expect(errorExceptInTest).toHaveBeenCalledWith( + 'Invalid OpenRouter System One response or missing usage' ); expect(after).not.toHaveBeenCalled(); expect(logMicrodollarUsage).not.toHaveBeenCalled(); diff --git a/apps/web/src/app/api/gateway/typesafe/v1/systemone/route.ts b/apps/web/src/app/api/gateway/typesafe/v1/systemone/route.ts index 3d8c0d0551..9053460883 100644 --- a/apps/web/src/app/api/gateway/typesafe/v1/systemone/route.ts +++ b/apps/web/src/app/api/gateway/typesafe/v1/systemone/route.ts @@ -1,191 +1,3 @@ -import { after, NextResponse, type NextRequest } from 'next/server'; -import { captureException, captureMessage } from '@sentry/nextjs'; -import { KILO_GATEWAY_AUDIENCE } from '@kilocode/worker-utils/internal-service-token-audiences'; -import { getUserFromAuth } from '@/lib/user/server'; -import { getBalanceAndOrgSettings } from '@/lib/organizations/organization-usage'; -import { resolveOrganizationMemberModelDecision } from '@/lib/organizations/effective-model-access.server'; -import { - gatewayRateLimitKey, - isGatewayAccountRateLimited, -} from '@/lib/ai-gateway/gateway-account-rate-limit'; -import { - checkOrganizationModelRestrictions, - creditsBlockedResponse, - extractFraudAndProjectHeaders, - extractHeaderAndLimitLength, - modelNotAllowedResponse, - wrapInSafeNextResponse, -} from '@/lib/ai-gateway/llm-proxy-helpers'; -import { OPENROUTER } from '@/lib/ai-gateway/providers/definitions/openrouter'; -import { ATTRIBUTION_HEADERS } from '@/lib/ai-gateway/providers/openrouter/attribution-headers'; -import { generateProviderSpecificHash } from '@/lib/ai-gateway/providerHash'; -import { logMicrodollarUsage } from '@/lib/ai-gateway/processUsage'; -import { - systemOneRequestSchema, - systemOneResponseSchema, - TYPESAFE_MODEL, -} from '@/lib/ai-gateway/typesafe'; -import { FEATURE_HEADER, validateFeatureHeader } from '@/lib/feature-detection'; -import { toMicrodollars } from '@/lib/utils'; +export { handleSystemOneRequest as POST } from '@/lib/ai-gateway/typesafe/handler'; export const maxDuration = 800; - -function errorResponse(message: string, error_type: string, status: number) { - return NextResponse.json({ message, error_type }, { status }); -} - -export async function POST(request: NextRequest) { - const startedAt = performance.now(); - const ipAddress = request.headers.get('x-forwarded-for')?.split(',')[0]?.trim(); - if (await isGatewayAccountRateLimited(request, gatewayRateLimitKey(request.headers, ipAddress))) { - return errorResponse('Rate limit exceeded', 'rate_limit_exceeded', 429); - } - - const { user, authFailedResponse, organizationId, botId, tokenSource } = await getUserFromAuth({ - adminOnly: false, - expectedAudience: KILO_GATEWAY_AUDIENCE, - }); - if (authFailedResponse) return authFailedResponse; - if (!user) return errorResponse('Authentication required', 'authentication_error', 401); - - let body: unknown; - try { - body = await request.json(); - } catch { - return errorResponse('Invalid JSON body', 'invalid_request', 400); - } - const parsed = systemOneRequestSchema.safeParse(body); - if (!parsed.success) { - return errorResponse( - parsed.error.issues.map(issue => `${issue.path.join('.')}: ${issue.message}`).join('; '), - 'invalid_request', - 400 - ); - } - - const { balance, settings, plan, balanceLimitedByUserAllowance } = await getBalanceAndOrgSettings( - organizationId, - user - ); - if (balance <= 0) { - return creditsBlockedResponse({ user, balance, organizationId, balanceLimitedByUserAllowance }); - } - - const { error, providerConfig } = checkOrganizationModelRestrictions({ - modelId: TYPESAFE_MODEL, - settings, - organizationPlan: plan, - }); - if (error) return error; - let providerPolicy = providerConfig; - if (organizationId) { - const { decision } = await resolveOrganizationMemberModelDecision({ - organizationId, - kiloUserId: user.id, - modelId: TYPESAFE_MODEL, - providerLookup: async () => new Set(['typesafe']), - }); - if (!decision.allowed) return modelNotAllowedResponse(); - if (decision.eligibleProviderRoutes) { - const only = providerConfig?.only - ? providerConfig.only.filter(route => decision.eligibleProviderRoutes?.has(route)) - : [...decision.eligibleProviderRoutes]; - if (only.length === 0) return modelNotAllowedResponse(); - providerPolicy = { ...providerConfig, only }; - } - } - - let response: Response; - let responseBody: unknown; - let ttfbMs: number; - try { - response = await fetch(`${OPENROUTER.apiUrl}/systemone`, { - method: 'POST', - headers: { - ...ATTRIBUTION_HEADERS, - Authorization: `Bearer ${OPENROUTER.apiKey}`, - 'Content-Type': 'application/json', - }, - body: JSON.stringify({ - ...parsed.data, - provider: providerPolicy, - user: generateProviderSpecificHash(user.id, OPENROUTER), - }), - signal: AbortSignal.any([request.signal, AbortSignal.timeout(10 * 60 * 1000)]), - }); - ttfbMs = Math.max(0, Math.round(performance.now() - startedAt)); - if (response.status === 402) { - captureMessage('OpenRouter System One balance exhausted', { level: 'error' }); - return errorResponse('Service temporarily unavailable', 'upstream_error', 503); - } - if (!response.ok) return wrapInSafeNextResponse(response); - responseBody = await response.json(); - } catch (error) { - captureException(error, { tags: { source: 'systemone_proxy' } }); - return errorResponse('Upstream request failed', 'upstream_error', 502); - } - - const result = systemOneResponseSchema.safeParse(responseBody); - if (!result.success) { - captureMessage('Invalid OpenRouter System One response or missing usage', { level: 'error' }); - return errorResponse('Invalid upstream response', 'upstream_error', 502); - } - - const { id, model, provider, usage } = result.data; - const { fraudHeaders, projectId } = extractFraudAndProjectHeaders(request); - const cost = toMicrodollars(usage.cost); - after(async () => { - await logMicrodollarUsage( - { - messageId: id, - model, - responseContent: '', - hasError: false, - inference_provider: provider ?? 'TypeSafe', - upstream_id: null, - finish_reason: null, - latency: ttfbMs, - moderation_latency: null, - generation_time: null, - streamed: false, - cancelled: false, - status_code: response.status, - cost_mUsd: cost, - market_cost: cost, - inputTokens: usage.input_tokens, - outputTokens: usage.output_tokens, - cacheWriteTokens: 0, - cacheHitTokens: 0, - is_byok: false, - }, - { - api_kind: 'systemone', - kiloUserId: user.id, - provider: 'openrouter', - requested_model: TYPESAFE_MODEL, - promptInfo: { system_prompt_prefix: '', system_prompt_length: 0, user_prompt_prefix: '' }, - max_tokens: null, - has_middle_out_transform: null, - fraudHeaders, - isStreaming: false, - organizationId, - prior_microdollar_usage: user.microdollars_used, - posthog_distinct_id: user.google_user_email, - project_id: projectId, - status_code: response.status, - editor_name: extractHeaderAndLimitLength(request, 'x-kilocode-editorname'), - machine_id: extractHeaderAndLimitLength(request, 'x-kilocode-machineid'), - user_byok: false, - has_tools: false, - botId, - tokenSource, - feature: validateFeatureHeader(request.headers.get(FEATURE_HEADER) || ''), - session_id: extractHeaderAndLimitLength(request, 'X-KiloCode-TaskId'), - mode: extractHeaderAndLimitLength(request, 'x-kilocode-mode'), - auto_model: null, - ttfb_ms: ttfbMs, - } - ); - }); - return NextResponse.json(responseBody); -} diff --git a/apps/web/src/lib/ai-gateway/typesafe/handler.ts b/apps/web/src/lib/ai-gateway/typesafe/handler.ts new file mode 100644 index 0000000000..7c03b64197 --- /dev/null +++ b/apps/web/src/lib/ai-gateway/typesafe/handler.ts @@ -0,0 +1,186 @@ +import { after, NextResponse, type NextRequest } from 'next/server'; +import { z } from 'zod'; +import { KILO_GATEWAY_AUDIENCE } from '@kilocode/worker-utils/internal-service-token-audiences'; +import { getUserFromAuth } from '@/lib/user/server'; +import { getBalanceAndOrgSettings } from '@/lib/organizations/organization-usage'; +import { resolveOrganizationMemberModelDecision } from '@/lib/organizations/effective-model-access.server'; +import { + gatewayRateLimitKey, + isGatewayAccountRateLimited, +} from '@/lib/ai-gateway/gateway-account-rate-limit'; +import { + checkOrganizationModelRestrictions, + creditsBlockedResponse, + extractFraudAndProjectHeaders, + extractHeaderAndLimitLength, + modelNotAllowedResponse, + wrapInSafeNextResponse, +} from '@/lib/ai-gateway/llm-proxy-helpers'; +import { OPENROUTER } from '@/lib/ai-gateway/providers/definitions/openrouter'; +import { ATTRIBUTION_HEADERS } from '@/lib/ai-gateway/providers/openrouter/attribution-headers'; +import { generateProviderSpecificHash } from '@/lib/ai-gateway/providerHash'; +import { logMicrodollarUsage } from '@/lib/ai-gateway/processUsage'; +import { + systemOneRequestSchema, + systemOneResponseSchema, + TYPESAFE_MODEL, +} from '@/lib/ai-gateway/typesafe/schemas'; +import { FEATURE_HEADER, validateFeatureHeader } from '@/lib/feature-detection'; +import { toMicrodollars } from '@/lib/utils'; +import { errorExceptInTest } from '@/lib/utils.server'; + +function errorResponse(message: string, error_type: string, status: number) { + return NextResponse.json({ message, error_type }, { status }); +} + +export async function handleSystemOneRequest(request: NextRequest) { + const startedAt = performance.now(); + const ipAddress = request.headers.get('x-forwarded-for')?.split(',')[0]?.trim(); + if (await isGatewayAccountRateLimited(request, gatewayRateLimitKey(request.headers, ipAddress))) { + return errorResponse('Rate limit exceeded', 'rate_limit_exceeded', 429); + } + + const { user, authFailedResponse, organizationId, botId, tokenSource } = await getUserFromAuth({ + adminOnly: false, + expectedAudience: KILO_GATEWAY_AUDIENCE, + }); + if (authFailedResponse) return authFailedResponse; + if (!user) return errorResponse('Authentication required', 'authentication_error', 401); + + let body: unknown; + try { + body = await request.json(); + } catch { + return errorResponse('Invalid JSON body', 'invalid_request', 400); + } + const parsed = systemOneRequestSchema.safeParse(body); + if (!parsed.success) { + return errorResponse(z.prettifyError(parsed.error), 'invalid_request', 400); + } + + const { balance, settings, plan, balanceLimitedByUserAllowance } = await getBalanceAndOrgSettings( + organizationId, + user + ); + if (balance <= 0) { + return creditsBlockedResponse({ user, balance, organizationId, balanceLimitedByUserAllowance }); + } + + const { error, providerConfig } = checkOrganizationModelRestrictions({ + modelId: TYPESAFE_MODEL, + settings, + organizationPlan: plan, + }); + if (error) return error; + let providerPolicy = providerConfig; + if (organizationId) { + const { decision } = await resolveOrganizationMemberModelDecision({ + organizationId, + kiloUserId: user.id, + modelId: TYPESAFE_MODEL, + providerLookup: async () => new Set(['typesafe']), + }); + if (!decision.allowed) return modelNotAllowedResponse(); + if (decision.eligibleProviderRoutes) { + const only = providerConfig?.only + ? providerConfig.only.filter(route => decision.eligibleProviderRoutes?.has(route)) + : [...decision.eligibleProviderRoutes]; + if (only.length === 0) return modelNotAllowedResponse(); + providerPolicy = { ...providerConfig, only }; + } + } + + let response: Response; + let responseBody: unknown; + let ttfbMs: number; + try { + response = await fetch(`${OPENROUTER.apiUrl}/systemone`, { + method: 'POST', + headers: { + ...ATTRIBUTION_HEADERS, + Authorization: `Bearer ${OPENROUTER.apiKey}`, + 'Content-Type': 'application/json', + }, + body: JSON.stringify({ + ...parsed.data, + provider: providerPolicy, + user: generateProviderSpecificHash(user.id, OPENROUTER), + }), + signal: AbortSignal.any([request.signal, AbortSignal.timeout(10 * 60 * 1000)]), + }); + ttfbMs = Math.max(0, Math.round(performance.now() - startedAt)); + if (response.status === 402) { + errorExceptInTest('OpenRouter System One balance exhausted'); + return errorResponse('Service temporarily unavailable', 'upstream_error', 503); + } + if (!response.ok) return wrapInSafeNextResponse(response); + responseBody = await response.json(); + } catch (error) { + errorExceptInTest('OpenRouter System One request failed', error); + return errorResponse('Upstream request failed', 'upstream_error', 502); + } + + const result = systemOneResponseSchema.safeParse(responseBody); + if (!result.success) { + errorExceptInTest('Invalid OpenRouter System One response or missing usage'); + return errorResponse('Invalid upstream response', 'upstream_error', 502); + } + + const { id, model, provider, usage } = result.data; + const { fraudHeaders, projectId } = extractFraudAndProjectHeaders(request); + const cost = toMicrodollars(usage.cost); + after(async () => { + await logMicrodollarUsage( + { + messageId: id, + model, + responseContent: '', + hasError: false, + inference_provider: provider ?? 'TypeSafe', + upstream_id: null, + finish_reason: null, + latency: ttfbMs, + moderation_latency: null, + generation_time: null, + streamed: false, + cancelled: false, + status_code: response.status, + cost_mUsd: cost, + market_cost: cost, + inputTokens: usage.input_tokens, + outputTokens: usage.output_tokens, + cacheWriteTokens: 0, + cacheHitTokens: 0, + is_byok: false, + }, + { + api_kind: 'systemone', + kiloUserId: user.id, + provider: 'openrouter', + requested_model: TYPESAFE_MODEL, + promptInfo: { system_prompt_prefix: '', system_prompt_length: 0, user_prompt_prefix: '' }, + max_tokens: null, + has_middle_out_transform: null, + fraudHeaders, + isStreaming: false, + organizationId, + prior_microdollar_usage: user.microdollars_used, + posthog_distinct_id: user.google_user_email, + project_id: projectId, + status_code: response.status, + editor_name: extractHeaderAndLimitLength(request, 'x-kilocode-editorname'), + machine_id: extractHeaderAndLimitLength(request, 'x-kilocode-machineid'), + user_byok: false, + has_tools: false, + botId, + tokenSource, + feature: validateFeatureHeader(request.headers.get(FEATURE_HEADER) || ''), + session_id: extractHeaderAndLimitLength(request, 'X-KiloCode-TaskId'), + mode: extractHeaderAndLimitLength(request, 'x-kilocode-mode'), + auto_model: null, + ttfb_ms: ttfbMs, + } + ); + }); + return NextResponse.json(responseBody); +} diff --git a/apps/web/src/lib/ai-gateway/typesafe.test.ts b/apps/web/src/lib/ai-gateway/typesafe/schemas.test.ts similarity index 99% rename from apps/web/src/lib/ai-gateway/typesafe.test.ts rename to apps/web/src/lib/ai-gateway/typesafe/schemas.test.ts index fb413496ff..7046550f43 100644 --- a/apps/web/src/lib/ai-gateway/typesafe.test.ts +++ b/apps/web/src/lib/ai-gateway/typesafe/schemas.test.ts @@ -1,6 +1,6 @@ import { describe, expect, it } from '@jest/globals'; import { choice, noul, score } from '@typesafe-ai/sdk'; -import { systemOneRequestSchema, systemOneResponseSchema, TYPESAFE_MODEL } from './typesafe'; +import { systemOneRequestSchema, systemOneResponseSchema, TYPESAFE_MODEL } from './schemas'; describe('systemOneRequestSchema', () => { it('accepts SDK question helpers with structured entries and strips routing overrides', () => { diff --git a/apps/web/src/lib/ai-gateway/typesafe.ts b/apps/web/src/lib/ai-gateway/typesafe/schemas.ts similarity index 100% rename from apps/web/src/lib/ai-gateway/typesafe.ts rename to apps/web/src/lib/ai-gateway/typesafe/schemas.ts From cabfd2f081180e9996792227328e254e2dab4c69 Mon Sep 17 00:00:00 2001 From: Christiaan Arnoldus Date: Wed, 23 Sep 2026 11:50:14 +0200 Subject: [PATCH 5/6] fix(ai-gateway): address TypeSafe review findings --- .../api/gateway/typesafe/v1/systemone/route.test.ts | 12 ++++++------ apps/web/src/lib/ai-gateway/typesafe/handler.ts | 7 ++++--- 2 files changed, 10 insertions(+), 9 deletions(-) diff --git a/apps/web/src/app/api/gateway/typesafe/v1/systemone/route.test.ts b/apps/web/src/app/api/gateway/typesafe/v1/systemone/route.test.ts index 2dc30269c6..78f523b919 100644 --- a/apps/web/src/app/api/gateway/typesafe/v1/systemone/route.test.ts +++ b/apps/web/src/app/api/gateway/typesafe/v1/systemone/route.test.ts @@ -324,7 +324,7 @@ describe('POST /api/gateway/typesafe/v1/systemone', () => { ); }); - it('accepts zero-cost usage and falls back to TypeSafe when the provider is absent', async () => { + it('accepts zero-cost usage and leaves an absent inference provider unknown', async () => { mockedFetch.mockResolvedValue( Response.json({ ...upstreamBody, @@ -337,7 +337,7 @@ describe('POST /api/gateway/typesafe/v1/systemone', () => { await runAfter(); expect(logMicrodollarUsage).toHaveBeenCalledWith( - expect.objectContaining({ inference_provider: 'TypeSafe', cost_mUsd: 0, market_cost: 0 }), + expect.objectContaining({ inference_provider: null, cost_mUsd: 0, market_cost: 0 }), expect.objectContaining({ provider: 'openrouter', user_byok: false }) ); }); @@ -532,14 +532,14 @@ describe('POST /api/gateway/typesafe/v1/systemone', () => { expect(after).not.toHaveBeenCalled(); }); - it('maps OpenRouter credit exhaustion to a service error without charging the user', async () => { - mockedFetch.mockResolvedValue( - Response.json({ error: 'upstream credits exhausted' }, { status: 402 }) - ); + it('cancels the upstream body on credit exhaustion and returns a service error without charging', async () => { + const cancel = jest.fn(); + mockedFetch.mockResolvedValue(new Response(new ReadableStream({ cancel }), { status: 402 })); const response = await POST(makeRequest()); expect(response.status).toBe(503); + expect(cancel).toHaveBeenCalledTimes(1); expect(await response.json()).toEqual({ message: 'Service temporarily unavailable', error_type: 'upstream_error', diff --git a/apps/web/src/lib/ai-gateway/typesafe/handler.ts b/apps/web/src/lib/ai-gateway/typesafe/handler.ts index 7c03b64197..be7abd0793 100644 --- a/apps/web/src/lib/ai-gateway/typesafe/handler.ts +++ b/apps/web/src/lib/ai-gateway/typesafe/handler.ts @@ -28,8 +28,9 @@ import { import { FEATURE_HEADER, validateFeatureHeader } from '@/lib/feature-detection'; import { toMicrodollars } from '@/lib/utils'; import { errorExceptInTest } from '@/lib/utils.server'; +import type { ProxyErrorType } from '@/lib/proxy-error-types'; -function errorResponse(message: string, error_type: string, status: number) { +function errorResponse(message: string, error_type: ProxyErrorType, status: number) { return NextResponse.json({ message, error_type }, { status }); } @@ -45,7 +46,6 @@ export async function handleSystemOneRequest(request: NextRequest) { expectedAudience: KILO_GATEWAY_AUDIENCE, }); if (authFailedResponse) return authFailedResponse; - if (!user) return errorResponse('Authentication required', 'authentication_error', 401); let body: unknown; try { @@ -110,6 +110,7 @@ export async function handleSystemOneRequest(request: NextRequest) { }); ttfbMs = Math.max(0, Math.round(performance.now() - startedAt)); if (response.status === 402) { + await response.body?.cancel(); errorExceptInTest('OpenRouter System One balance exhausted'); return errorResponse('Service temporarily unavailable', 'upstream_error', 503); } @@ -136,7 +137,7 @@ export async function handleSystemOneRequest(request: NextRequest) { model, responseContent: '', hasError: false, - inference_provider: provider ?? 'TypeSafe', + inference_provider: provider ?? null, upstream_id: null, finish_reason: null, latency: ttfbMs, From 688579f093dd1e3ce5c1e09918a79099ed753ce6 Mon Sep 17 00:00:00 2001 From: Christiaan Arnoldus Date: Wed, 23 Sep 2026 13:30:03 +0200 Subject: [PATCH 6/6] test(ai-gateway): colocate TypeSafe handler coverage --- .../typesafe/v1/systemone/route.test.ts | 619 +----------------- .../lib/ai-gateway/typesafe/handler.test.ts | 617 +++++++++++++++++ 2 files changed, 624 insertions(+), 612 deletions(-) create mode 100644 apps/web/src/lib/ai-gateway/typesafe/handler.test.ts diff --git a/apps/web/src/app/api/gateway/typesafe/v1/systemone/route.test.ts b/apps/web/src/app/api/gateway/typesafe/v1/systemone/route.test.ts index 78f523b919..50e9e54f9e 100644 --- a/apps/web/src/app/api/gateway/typesafe/v1/systemone/route.test.ts +++ b/apps/web/src/app/api/gateway/typesafe/v1/systemone/route.test.ts @@ -1,615 +1,10 @@ -import { afterEach, beforeEach, describe, expect, it } from '@jest/globals'; -import { z } from 'zod'; -import { errorExceptInTest } from '@/lib/utils.server'; -import { TypeSafeClient, choice, noul, score } from '@typesafe-ai/sdk'; -import type { User } from '@kilocode/db/schema'; -import { KILO_GATEWAY_AUDIENCE } from '@kilocode/worker-utils/internal-service-token-audiences'; -import { after, NextRequest, NextResponse } from 'next/server'; -import type * as NextServer from 'next/server'; -import { getUserFromAuth } from '@/lib/user/server'; -import { getBalanceAndOrgSettings } from '@/lib/organizations/organization-usage'; -import { resolveOrganizationMemberModelDecision } from '@/lib/organizations/effective-model-access.server'; -import { - gatewayRateLimitKey, - isGatewayAccountRateLimited, -} from '@/lib/ai-gateway/gateway-account-rate-limit'; -import { - checkOrganizationModelRestrictions, - creditsBlockedResponse, - extractFraudAndProjectHeaders, - extractHeaderAndLimitLength, - modelNotAllowedResponse, - wrapInSafeNextResponse, -} from '@/lib/ai-gateway/llm-proxy-helpers'; -import { OPENROUTER } from '@/lib/ai-gateway/providers/definitions/openrouter'; -import { generateProviderSpecificHash } from '@/lib/ai-gateway/providerHash'; -import { logMicrodollarUsage } from '@/lib/ai-gateway/processUsage'; -import { systemOneRequestSchema, TYPESAFE_MODEL } from '@/lib/ai-gateway/typesafe/schemas'; -import { EmptyFraudDetectionHeaders } from '@/lib/utils'; -import { POST } from './route'; +import { expect, it } from '@jest/globals'; +import { handleSystemOneRequest } from '@/lib/ai-gateway/typesafe/handler'; +import { maxDuration, POST } from './route'; -jest.mock('next/server', () => ({ - ...jest.requireActual('next/server'), - after: jest.fn(), -})); -jest.mock('@/lib/utils.server', () => ({ errorExceptInTest: jest.fn() })); -jest.mock('@/lib/user/server', () => ({ getUserFromAuth: jest.fn() })); -jest.mock('@/lib/organizations/organization-usage', () => ({ - getBalanceAndOrgSettings: jest.fn(), -})); -jest.mock('@/lib/organizations/effective-model-access.server', () => ({ - resolveOrganizationMemberModelDecision: jest.fn(), -})); -jest.mock('@/lib/ai-gateway/gateway-account-rate-limit', () => ({ - gatewayRateLimitKey: jest.fn(), - isGatewayAccountRateLimited: jest.fn(), -})); -jest.mock('@/lib/ai-gateway/llm-proxy-helpers', () => ({ - checkOrganizationModelRestrictions: jest.fn(), - creditsBlockedResponse: jest.fn(), - extractFraudAndProjectHeaders: jest.fn(), - extractHeaderAndLimitLength: jest.fn(), - modelNotAllowedResponse: jest.fn(), - wrapInSafeNextResponse: jest.fn(), -})); -jest.mock('@/lib/ai-gateway/providers/definitions/openrouter', () => ({ - OPENROUTER: { - id: 'openrouter', - apiUrl: 'https://openrouter.ai/api/v1', - apiKey: 'test-platform-openrouter-key', - }, -})); -jest.mock('@/lib/ai-gateway/providerHash', () => ({ generateProviderSpecificHash: jest.fn() })); -jest.mock('@/lib/ai-gateway/processUsage', () => ({ logMicrodollarUsage: jest.fn() })); +jest.mock('@/lib/ai-gateway/typesafe/handler', () => ({ handleSystemOneRequest: jest.fn() })); -const routeUrl = 'http://localhost:3000/api/gateway/typesafe/v1/systemone'; -const user = { - id: 'oauth/test-user', - google_user_email: 'test@example.com', - microdollars_used: 123, -} as User; -const questions = { - billing: noul('Is this about billing?'), - category: choice('Choose a category', { billing: 'Payments', other: null }), - urgency: score('How urgent is this?', ['Low', 'High']), -}; -const requestBody = { state: { message: 'I was charged twice.' }, questions }; -const upstreamBody = { - id: 'gen-systemone-123', - model: TYPESAFE_MODEL, - provider: 'TypeSafe upstream', - answers: { - billing: { type: 'noul', noul: 0.95 }, - category: { - type: 'choice', - choice: 'billing', - confidence: 0.9, - probabilities: { billing: 0.9, other: 0.1 }, - }, - urgency: { - type: 'score', - score: 0.75, - confidence: 0.8, - probabilities: { '0': 0.25, '1': 0.75 }, - legend: { '0': 'Low', '1': 'High' }, - }, - }, - usage: { input_tokens: 23, output_tokens: 7, cost: 0.0001236, total_tokens: 30 }, - upstream_metadata: { retained: true }, -}; -const memberDecision = { - policy: { - requireModelInCurrentSnapshot: false, - organizationModelDenyList: [], - memberGrant: { mode: 'unrestricted' as const }, - policyRevision: 1, - }, - decision: { allowed: true }, -}; -const mockedFetch = jest.fn() as jest.MockedFunction; -const originalFetch = globalThis.fetch; - -function makeRequest(body: unknown = requestBody, headers: Record = {}) { - return new NextRequest(routeUrl, { - method: 'POST', - headers: { 'Content-Type': 'application/json', ...headers }, - body: JSON.stringify(body), - }); -} - -function setAuth(organizationId?: string) { - jest.mocked(getUserFromAuth).mockResolvedValue({ - user, - authFailedResponse: null, - organizationId, - botId: 'bot-123', - tokenSource: 'api-key', - }); -} - -function upstreamRequest() { - expect(mockedFetch).toHaveBeenCalledTimes(1); - const [url, init] = mockedFetch.mock.calls[0]; - expect(url).toBe('https://openrouter.ai/api/v1/systemone'); - expect(init?.method).toBe('POST'); - if (typeof init?.body !== 'string') throw new Error('Expected a JSON upstream body'); - return { body: JSON.parse(init.body), headers: new Headers(init.headers) }; -} - -async function runAfter() { - expect(after).toHaveBeenCalledTimes(1); - const [callback] = jest.mocked(after).mock.calls[0]; - if (typeof callback !== 'function') throw new Error('Expected deferred usage callback'); - await callback(); -} - -describe('POST /api/gateway/typesafe/v1/systemone', () => { - beforeEach(() => { - jest.resetAllMocks(); - globalThis.fetch = mockedFetch; - setAuth(); - jest.mocked(getBalanceAndOrgSettings).mockResolvedValue({ balance: 1_000_000 }); - jest.mocked(isGatewayAccountRateLimited).mockResolvedValue(false); - jest.mocked(gatewayRateLimitKey).mockReturnValue('test-rate-limit-key'); - jest.mocked(checkOrganizationModelRestrictions).mockReturnValue({ error: null }); - jest.mocked(resolveOrganizationMemberModelDecision).mockResolvedValue(memberDecision); - jest.mocked(generateProviderSpecificHash).mockReturnValue('hashed-user'); - jest.mocked(extractFraudAndProjectHeaders).mockReturnValue({ - fraudHeaders: EmptyFraudDetectionHeaders, - projectId: 'project-123', - xKiloCodeVersion: null, - numericKiloCodeVersion: 0, - }); - jest - .mocked(extractHeaderAndLimitLength) - .mockImplementation((request, name) => request.headers.get(name)); - jest.mocked(modelNotAllowedResponse).mockImplementation(() => - NextResponse.json( - { - error: 'Model not allowed', - error_type: 'model_not_allowed', - message: 'Model not allowed', - }, - { status: 404 } - ) - ); - mockedFetch.mockImplementation(async () => Response.json(upstreamBody)); - }); - - afterEach(() => { - globalThis.fetch = originalFetch; - }); - - it('supports the TypeSafe SDK contract through a custom fetch transport', async () => { - const routeFetch = jest.fn(async (input: string, init?: RequestInit) => - POST(new NextRequest(new Request(input, init))) - ); - const client = new TypeSafeClient({ - apiKey: 'test-kilo-key', - baseURL: 'http://localhost:3000/api/gateway/typesafe', - defaultModel: 'jev-1.13', - retry: { maxRetries: 0 }, - logLevel: 'off', - fetch: routeFetch, - }); - - const { data, response } = await client.systemOne(requestBody).withResponse(); - - expect(response.status).toBe(200); - expect(data).toEqual(upstreamBody); - expect(data.answers.billing.noul).toBe(0.95); - expect(data.answers.category.choice).toBe('billing'); - expect(data.answers.urgency.score).toBe(0.75); - expect(routeFetch).toHaveBeenCalledTimes(1); - const [url, init] = routeFetch.mock.calls[0]; - expect(url).toBe(routeUrl); - expect(new Headers(init?.headers).get('authorization')).toBe('Bearer test-kilo-key'); - expect(upstreamRequest().body).toEqual({ - ...JSON.parse(JSON.stringify(requestBody)), - model: TYPESAFE_MODEL, - user: 'hashed-user', - }); - expect(getUserFromAuth).toHaveBeenCalledWith({ - adminOnly: false, - expectedAudience: KILO_GATEWAY_AUDIENCE, - }); - expect(resolveOrganizationMemberModelDecision).not.toHaveBeenCalled(); - }); - - it.each([undefined, 'jev-1.13', 'typesafe/jev-1.13'])( - 'pins model %s to OpenRouter and drops client credentials and routing overrides', - async model => { - const response = await POST( - makeRequest( - { - ...requestBody, - model, - provider: { only: ['attacker'], api_key: 'test-body-key' }, - user: 'attacker-user', - api_key: 'test-body-key', - byok: true, - user_byok: [{ providerId: 'typesafe', apiKey: 'test-byok-key' }], - base_url: 'https://attacker.invalid', - }, - { - Authorization: 'Bearer test-client-key', - 'Proxy-Authorization': 'Bearer test-proxy-key', - 'x-api-key': 'test-forwarded-key', - 'x-openrouter-api-key': 'test-user-openrouter-key', - Cookie: 'session=test-cookie', - 'HTTP-Referer': 'https://attacker.invalid', - 'X-Title': 'Attacker', - } - ) - ); - - expect(response.status).toBe(200); - const upstream = upstreamRequest(); - expect(upstream.body).toEqual({ - ...JSON.parse(JSON.stringify(requestBody)), - model: 'typesafe/jev-1.13', - user: 'hashed-user', - }); - expect(Object.fromEntries(upstream.headers)).toEqual({ - authorization: 'Bearer test-platform-openrouter-key', - 'content-type': 'application/json', - 'http-referer': 'https://kilocode.ai', - 'x-title': 'Kilo Code', - }); - expect(generateProviderSpecificHash).toHaveBeenCalledWith(user.id, OPENROUTER); - expect(checkOrganizationModelRestrictions).toHaveBeenCalledWith({ - modelId: TYPESAFE_MODEL, - settings: undefined, - organizationPlan: undefined, - }); - } - ); - - it('defers billing and converts upstream cost to rounded microdollars without a markup', async () => { - setAuth('org-123'); - const response = await POST( - makeRequest(requestBody, { - 'x-kilocode-editorname': 'vscode', - 'x-kilocode-machineid': 'machine-123', - 'x-kilocode-feature': 'cli', - 'X-KiloCode-TaskId': 'task-123', - 'x-kilocode-mode': 'code', - }) - ); - - expect(response.status).toBe(200); - expect(await response.json()).toEqual(upstreamBody); - expect(logMicrodollarUsage).not.toHaveBeenCalled(); - await runAfter(); - expect(logMicrodollarUsage).toHaveBeenCalledTimes(1); - expect(logMicrodollarUsage).toHaveBeenCalledWith( - expect.objectContaining({ - messageId: upstreamBody.id, - model: TYPESAFE_MODEL, - inference_provider: 'TypeSafe upstream', - cost_mUsd: 124, - market_cost: 124, - inputTokens: 23, - outputTokens: 7, - cacheWriteTokens: 0, - cacheHitTokens: 0, - is_byok: false, - streamed: false, - hasError: false, - status_code: 200, - responseContent: '', - latency: expect.any(Number), - }), - expect.objectContaining({ - api_kind: 'systemone', - kiloUserId: user.id, - provider: 'openrouter', - requested_model: TYPESAFE_MODEL, - organizationId: 'org-123', - prior_microdollar_usage: 123, - posthog_distinct_id: user.google_user_email, - project_id: 'project-123', - fraudHeaders: EmptyFraudDetectionHeaders, - editor_name: 'vscode', - machine_id: 'machine-123', - feature: 'cli', - session_id: 'task-123', - mode: 'code', - user_byok: false, - isStreaming: false, - botId: 'bot-123', - tokenSource: 'api-key', - ttfb_ms: expect.any(Number), - }) - ); - }); - - it('accepts zero-cost usage and leaves an absent inference provider unknown', async () => { - mockedFetch.mockResolvedValue( - Response.json({ - ...upstreamBody, - provider: undefined, - usage: { ...upstreamBody.usage, cost: 0 }, - }) - ); - - expect((await POST(makeRequest())).status).toBe(200); - await runAfter(); - - expect(logMicrodollarUsage).toHaveBeenCalledWith( - expect.objectContaining({ inference_provider: null, cost_mUsd: 0, market_cost: 0 }), - expect.objectContaining({ provider: 'openrouter', user_byok: false }) - ); - }); - - it('returns the authentication failure and requests the gateway audience', async () => { - const authFailedResponse = NextResponse.json( - { success: false as const, error: 'Unauthorized' }, - { status: 401 } - ); - jest.mocked(getUserFromAuth).mockResolvedValue({ user: null, authFailedResponse }); - - expect(await POST(makeRequest())).toBe(authFailedResponse); - expect(getUserFromAuth).toHaveBeenCalledWith({ - adminOnly: false, - expectedAudience: KILO_GATEWAY_AUDIENCE, - }); - expect(getBalanceAndOrgSettings).not.toHaveBeenCalled(); - expect(mockedFetch).not.toHaveBeenCalled(); - expect(after).not.toHaveBeenCalled(); - }); - - it('rate limits before authentication or upstream work', async () => { - jest.mocked(isGatewayAccountRateLimited).mockResolvedValue(true); - const request = makeRequest(requestBody, { 'x-forwarded-for': ' 192.0.2.1, 192.0.2.2' }); - - const response = await POST(request); - - expect(response.status).toBe(429); - expect(await response.json()).toMatchObject({ error_type: 'rate_limit_exceeded' }); - expect(gatewayRateLimitKey).toHaveBeenCalledWith(request.headers, '192.0.2.1'); - expect(isGatewayAccountRateLimited).toHaveBeenCalledWith(request, 'test-rate-limit-key'); - expect(getUserFromAuth).not.toHaveBeenCalled(); - expect(mockedFetch).not.toHaveBeenCalled(); - }); - - it.each([ - ['malformed JSON', '{'], - ['another model', JSON.stringify({ ...requestBody, model: 'openai/gpt-4o' })], - ['an unpinned alias', JSON.stringify({ ...requestBody, model: 'jev-latest' })], - ['empty questions', JSON.stringify({ ...requestBody, questions: {} })], - ])('rejects %s before balance checks or upstream work', async (_name, body) => { - const response = await POST( - new NextRequest(routeUrl, { - method: 'POST', - headers: { 'Content-Type': 'application/json' }, - body, - }) - ); - - expect(response.status).toBe(400); - expect(await response.json()).toMatchObject({ error_type: 'invalid_request' }); - expect(getBalanceAndOrgSettings).not.toHaveBeenCalled(); - expect(mockedFetch).not.toHaveBeenCalled(); - expect(after).not.toHaveBeenCalled(); - }); - - it('formats validation errors with Zod while preserving the TypeSafe error shape', async () => { - const body = { ...requestBody, questions: { refund: { type: 'invalid' } } }; - const validation = systemOneRequestSchema.safeParse(body); - if (validation.success) throw new Error('Expected an invalid question'); - - const response = await POST(makeRequest(body)); - - expect(response.status).toBe(400); - expect(await response.json()).toEqual({ - message: z.prettifyError(validation.error), - error_type: 'invalid_request', - }); - expect(mockedFetch).not.toHaveBeenCalled(); - }); - - it.each([0, -1])('blocks balance %s using the shared credits response', async balance => { - setAuth('org-123'); - jest.mocked(getBalanceAndOrgSettings).mockResolvedValue({ - balance, - balanceLimitedByUserAllowance: true, - }); - const blockedResponse = NextResponse.json( - { - error_type: 'usage_limit_exceeded' as const, - error: { - title: 'Credits exhausted', - message: 'Credits exhausted', - balance, - buyCreditsUrl: '', - }, - }, - { status: 402 } - ); - jest.mocked(creditsBlockedResponse).mockResolvedValue(blockedResponse); - - expect(await POST(makeRequest())).toBe(blockedResponse); - expect(getBalanceAndOrgSettings).toHaveBeenCalledWith('org-123', user); - expect(creditsBlockedResponse).toHaveBeenCalledWith({ - user, - balance, - organizationId: 'org-123', - balanceLimitedByUserAllowance: true, - }); - expect(checkOrganizationModelRestrictions).not.toHaveBeenCalled(); - expect(mockedFetch).not.toHaveBeenCalled(); - expect(after).not.toHaveBeenCalled(); - }); - - it('honors organization model restrictions before resolving member access', async () => { - setAuth('org-123'); - const settings = { model_deny_list: [TYPESAFE_MODEL] }; - jest.mocked(getBalanceAndOrgSettings).mockResolvedValue({ - balance: 1000, - settings, - plan: 'enterprise', - }); - const error = NextResponse.json({ error_type: 'model_not_allowed' }, { status: 404 }); - jest.mocked(checkOrganizationModelRestrictions).mockReturnValue({ error }); - - expect(await POST(makeRequest())).toBe(error); - expect(checkOrganizationModelRestrictions).toHaveBeenCalledWith({ - modelId: TYPESAFE_MODEL, - settings, - organizationPlan: 'enterprise', - }); - expect(resolveOrganizationMemberModelDecision).not.toHaveBeenCalled(); - expect(mockedFetch).not.toHaveBeenCalled(); - expect(after).not.toHaveBeenCalled(); - }); - - it('denies a model rejected by the effective organization member decision', async () => { - setAuth('org-123'); - jest.mocked(resolveOrganizationMemberModelDecision).mockResolvedValue({ - ...memberDecision, - decision: { allowed: false }, - }); - - expect((await POST(makeRequest())).status).toBe(404); - expect(resolveOrganizationMemberModelDecision).toHaveBeenCalledWith({ - organizationId: 'org-123', - kiloUserId: user.id, - modelId: TYPESAFE_MODEL, - providerLookup: expect.any(Function), - }); - const [{ providerLookup }] = jest.mocked(resolveOrganizationMemberModelDecision).mock.calls[0]; - if (!providerLookup) throw new Error('Expected the fixed TypeSafe provider lookup'); - await expect(providerLookup(TYPESAFE_MODEL)).resolves.toEqual(new Set(['typesafe'])); - expect(modelNotAllowedResponse).toHaveBeenCalledTimes(1); - expect(mockedFetch).not.toHaveBeenCalled(); - expect(after).not.toHaveBeenCalled(); - }); - - it.each([ - { - only: ['typesafe', 'other'], - eligible: ['typesafe', 'outside-ceiling'], - expected: ['typesafe'], - }, - { only: undefined, eligible: ['typesafe'], expected: ['typesafe'] }, - { only: ['typesafe'], eligible: undefined, expected: ['typesafe'] }, - ])('applies the provider policy intersection: %j', async ({ only, eligible, expected }) => { - setAuth('org-123'); - jest.mocked(checkOrganizationModelRestrictions).mockReturnValue({ - error: null, - providerConfig: { only, data_collection: 'deny' }, - }); - jest.mocked(resolveOrganizationMemberModelDecision).mockResolvedValue({ - ...memberDecision, - decision: { allowed: true, eligibleProviderRoutes: eligible ? new Set(eligible) : undefined }, - }); - - const response = await POST(makeRequest({ ...requestBody, provider: { only: ['attacker'] } })); - - expect(response.status).toBe(200); - expect(upstreamRequest().body.provider).toEqual({ only: expected, data_collection: 'deny' }); - }); - - it.each([ - { only: ['other'], eligible: ['typesafe'] }, - { only: [], eligible: ['typesafe'] }, - { only: undefined, eligible: [] }, - ])('rejects an empty provider policy intersection: %j', async ({ only, eligible }) => { - setAuth('org-123'); - jest.mocked(checkOrganizationModelRestrictions).mockReturnValue({ - error: null, - providerConfig: { only, data_collection: 'deny' }, - }); - jest.mocked(resolveOrganizationMemberModelDecision).mockResolvedValue({ - ...memberDecision, - decision: { allowed: true, eligibleProviderRoutes: new Set(eligible) }, - }); - - expect((await POST(makeRequest())).status).toBe(404); - expect(modelNotAllowedResponse).toHaveBeenCalledTimes(1); - expect(mockedFetch).not.toHaveBeenCalled(); - expect(after).not.toHaveBeenCalled(); - }); - - it('cancels the upstream body on credit exhaustion and returns a service error without charging', async () => { - const cancel = jest.fn(); - mockedFetch.mockResolvedValue(new Response(new ReadableStream({ cancel }), { status: 402 })); - - const response = await POST(makeRequest()); - - expect(response.status).toBe(503); - expect(cancel).toHaveBeenCalledTimes(1); - expect(await response.json()).toEqual({ - message: 'Service temporarily unavailable', - error_type: 'upstream_error', - }); - expect(errorExceptInTest).toHaveBeenCalledWith('OpenRouter System One balance exhausted'); - expect(wrapInSafeNextResponse).not.toHaveBeenCalled(); - expect(after).not.toHaveBeenCalled(); - expect(logMicrodollarUsage).not.toHaveBeenCalled(); - }); - - it.each([400, 429, 500])( - 'delegates upstream HTTP %s to the safe response wrapper', - async status => { - const upstream = Response.json({ error: 'upstream failure' }, { status }); - const safeResponse = NextResponse.json({ error: 'upstream failure' }, { status }); - mockedFetch.mockResolvedValue(upstream); - jest.mocked(wrapInSafeNextResponse).mockReturnValue(safeResponse); - - expect(await POST(makeRequest())).toBe(safeResponse); - expect(wrapInSafeNextResponse).toHaveBeenCalledWith(upstream); - expect(after).not.toHaveBeenCalled(); - expect(logMicrodollarUsage).not.toHaveBeenCalled(); - } - ); - - it.each(['network failure', 'non-JSON success'])('handles %s without charging', async failure => { - if (failure === 'network failure') { - mockedFetch.mockRejectedValue(new Error('connection failed')); - } else { - mockedFetch.mockResolvedValue(new Response('not JSON', { status: 200 })); - } - - const response = await POST(makeRequest()); - - expect(response.status).toBe(502); - expect(await response.json()).toEqual({ - message: 'Upstream request failed', - error_type: 'upstream_error', - }); - expect(errorExceptInTest).toHaveBeenCalledWith( - 'OpenRouter System One request failed', - expect.objectContaining({ message: expect.any(String) }) - ); - expect(after).not.toHaveBeenCalled(); - expect(logMicrodollarUsage).not.toHaveBeenCalled(); - }); - - it.each([ - ['missing usage', undefined], - ['null usage', null], - ['missing cost', { input_tokens: 23, output_tokens: 7 }], - ['string cost', { ...upstreamBody.usage, cost: '0.01' }], - ['negative cost', { ...upstreamBody.usage, cost: -0.01 }], - ['missing input tokens', { output_tokens: 7, cost: 0.01 }], - ['negative tokens', { ...upstreamBody.usage, input_tokens: -1 }], - ['fractional tokens', { ...upstreamBody.usage, output_tokens: 1.5 }], - ])('rejects upstream %s without scheduling billing', async (_name, usage) => { - mockedFetch.mockResolvedValue(Response.json({ ...upstreamBody, usage })); - - const response = await POST(makeRequest()); - - expect(response.status).toBe(502); - expect(await response.json()).toEqual({ - message: 'Invalid upstream response', - error_type: 'upstream_error', - }); - expect(errorExceptInTest).toHaveBeenCalledWith( - 'Invalid OpenRouter System One response or missing usage' - ); - expect(after).not.toHaveBeenCalled(); - expect(logMicrodollarUsage).not.toHaveBeenCalled(); - }); +it('exposes the System One handler with the gateway duration limit', () => { + expect(POST).toBe(handleSystemOneRequest); + expect(maxDuration).toBe(800); }); diff --git a/apps/web/src/lib/ai-gateway/typesafe/handler.test.ts b/apps/web/src/lib/ai-gateway/typesafe/handler.test.ts new file mode 100644 index 0000000000..40d0d70939 --- /dev/null +++ b/apps/web/src/lib/ai-gateway/typesafe/handler.test.ts @@ -0,0 +1,617 @@ +import { afterEach, beforeEach, describe, expect, it } from '@jest/globals'; +import { z } from 'zod'; +import { errorExceptInTest } from '@/lib/utils.server'; +import { TypeSafeClient, choice, noul, score } from '@typesafe-ai/sdk'; +import type { User } from '@kilocode/db/schema'; +import { KILO_GATEWAY_AUDIENCE } from '@kilocode/worker-utils/internal-service-token-audiences'; +import { after, NextRequest, NextResponse } from 'next/server'; +import type * as NextServer from 'next/server'; +import { getUserFromAuth } from '@/lib/user/server'; +import { getBalanceAndOrgSettings } from '@/lib/organizations/organization-usage'; +import { resolveOrganizationMemberModelDecision } from '@/lib/organizations/effective-model-access.server'; +import { + gatewayRateLimitKey, + isGatewayAccountRateLimited, +} from '@/lib/ai-gateway/gateway-account-rate-limit'; +import { + checkOrganizationModelRestrictions, + creditsBlockedResponse, + extractFraudAndProjectHeaders, + extractHeaderAndLimitLength, + modelNotAllowedResponse, + wrapInSafeNextResponse, +} from '@/lib/ai-gateway/llm-proxy-helpers'; +import { OPENROUTER } from '@/lib/ai-gateway/providers/definitions/openrouter'; +import { generateProviderSpecificHash } from '@/lib/ai-gateway/providerHash'; +import { logMicrodollarUsage } from '@/lib/ai-gateway/processUsage'; +import { systemOneRequestSchema, TYPESAFE_MODEL } from '@/lib/ai-gateway/typesafe/schemas'; +import { EmptyFraudDetectionHeaders } from '@/lib/utils'; +import { handleSystemOneRequest } from './handler'; + +jest.mock('next/server', () => ({ + ...jest.requireActual('next/server'), + after: jest.fn(), +})); +jest.mock('@/lib/utils.server', () => ({ errorExceptInTest: jest.fn() })); +jest.mock('@/lib/user/server', () => ({ getUserFromAuth: jest.fn() })); +jest.mock('@/lib/organizations/organization-usage', () => ({ + getBalanceAndOrgSettings: jest.fn(), +})); +jest.mock('@/lib/organizations/effective-model-access.server', () => ({ + resolveOrganizationMemberModelDecision: jest.fn(), +})); +jest.mock('@/lib/ai-gateway/gateway-account-rate-limit', () => ({ + gatewayRateLimitKey: jest.fn(), + isGatewayAccountRateLimited: jest.fn(), +})); +jest.mock('@/lib/ai-gateway/llm-proxy-helpers', () => ({ + checkOrganizationModelRestrictions: jest.fn(), + creditsBlockedResponse: jest.fn(), + extractFraudAndProjectHeaders: jest.fn(), + extractHeaderAndLimitLength: jest.fn(), + modelNotAllowedResponse: jest.fn(), + wrapInSafeNextResponse: jest.fn(), +})); +jest.mock('@/lib/ai-gateway/providers/definitions/openrouter', () => ({ + OPENROUTER: { + id: 'openrouter', + apiUrl: 'https://openrouter.ai/api/v1', + apiKey: 'test-platform-openrouter-key', + }, +})); +jest.mock('@/lib/ai-gateway/providerHash', () => ({ generateProviderSpecificHash: jest.fn() })); +jest.mock('@/lib/ai-gateway/processUsage', () => ({ logMicrodollarUsage: jest.fn() })); + +const routeUrl = 'http://localhost:3000/api/gateway/typesafe/v1/systemone'; +const user = { + id: 'oauth/test-user', + google_user_email: 'test@example.com', + microdollars_used: 123, +} as User; +const questions = { + billing: noul('Is this about billing?'), + category: choice('Choose a category', { billing: 'Payments', other: null }), + urgency: score('How urgent is this?', ['Low', 'High']), +}; +const requestBody = { state: { message: 'I was charged twice.' }, questions }; +const upstreamBody = { + id: 'gen-systemone-123', + model: TYPESAFE_MODEL, + provider: 'TypeSafe upstream', + answers: { + billing: { type: 'noul', noul: 0.95 }, + category: { + type: 'choice', + choice: 'billing', + confidence: 0.9, + probabilities: { billing: 0.9, other: 0.1 }, + }, + urgency: { + type: 'score', + score: 0.75, + confidence: 0.8, + probabilities: { '0': 0.25, '1': 0.75 }, + legend: { '0': 'Low', '1': 'High' }, + }, + }, + usage: { input_tokens: 23, output_tokens: 7, cost: 0.0001236, total_tokens: 30 }, + upstream_metadata: { retained: true }, +}; +const memberDecision = { + policy: { + requireModelInCurrentSnapshot: false, + organizationModelDenyList: [], + memberGrant: { mode: 'unrestricted' as const }, + policyRevision: 1, + }, + decision: { allowed: true }, +}; +const mockedFetch = jest.fn() as jest.MockedFunction; +const originalFetch = globalThis.fetch; + +function makeRequest(body: unknown = requestBody, headers: Record = {}) { + return new NextRequest(routeUrl, { + method: 'POST', + headers: { 'Content-Type': 'application/json', ...headers }, + body: JSON.stringify(body), + }); +} + +function setAuth(organizationId?: string) { + jest.mocked(getUserFromAuth).mockResolvedValue({ + user, + authFailedResponse: null, + organizationId, + botId: 'bot-123', + tokenSource: 'api-key', + }); +} + +function upstreamRequest() { + expect(mockedFetch).toHaveBeenCalledTimes(1); + const [url, init] = mockedFetch.mock.calls[0]; + expect(url).toBe('https://openrouter.ai/api/v1/systemone'); + expect(init?.method).toBe('POST'); + if (typeof init?.body !== 'string') throw new Error('Expected a JSON upstream body'); + return { body: JSON.parse(init.body), headers: new Headers(init.headers) }; +} + +async function runAfter() { + expect(after).toHaveBeenCalledTimes(1); + const [callback] = jest.mocked(after).mock.calls[0]; + if (typeof callback !== 'function') throw new Error('Expected deferred usage callback'); + await callback(); +} + +describe('handleSystemOneRequest', () => { + beforeEach(() => { + jest.resetAllMocks(); + globalThis.fetch = mockedFetch; + setAuth(); + jest.mocked(getBalanceAndOrgSettings).mockResolvedValue({ balance: 1_000_000 }); + jest.mocked(isGatewayAccountRateLimited).mockResolvedValue(false); + jest.mocked(gatewayRateLimitKey).mockReturnValue('test-rate-limit-key'); + jest.mocked(checkOrganizationModelRestrictions).mockReturnValue({ error: null }); + jest.mocked(resolveOrganizationMemberModelDecision).mockResolvedValue(memberDecision); + jest.mocked(generateProviderSpecificHash).mockReturnValue('hashed-user'); + jest.mocked(extractFraudAndProjectHeaders).mockReturnValue({ + fraudHeaders: EmptyFraudDetectionHeaders, + projectId: 'project-123', + xKiloCodeVersion: null, + numericKiloCodeVersion: 0, + }); + jest + .mocked(extractHeaderAndLimitLength) + .mockImplementation((request, name) => request.headers.get(name)); + jest.mocked(modelNotAllowedResponse).mockImplementation(() => + NextResponse.json( + { + error: 'Model not allowed', + error_type: 'model_not_allowed', + message: 'Model not allowed', + }, + { status: 404 } + ) + ); + mockedFetch.mockImplementation(async () => Response.json(upstreamBody)); + }); + + afterEach(() => { + globalThis.fetch = originalFetch; + }); + + it('supports the TypeSafe SDK contract through a custom fetch transport', async () => { + const routeFetch = jest.fn(async (input: string, init?: RequestInit) => + handleSystemOneRequest(new NextRequest(new Request(input, init))) + ); + const client = new TypeSafeClient({ + apiKey: 'test-kilo-key', + baseURL: 'http://localhost:3000/api/gateway/typesafe', + defaultModel: 'jev-1.13', + retry: { maxRetries: 0 }, + logLevel: 'off', + fetch: routeFetch, + }); + + const { data, response } = await client.systemOne(requestBody).withResponse(); + + expect(response.status).toBe(200); + expect(data).toEqual(upstreamBody); + expect(data.answers.billing.noul).toBe(0.95); + expect(data.answers.category.choice).toBe('billing'); + expect(data.answers.urgency.score).toBe(0.75); + expect(routeFetch).toHaveBeenCalledTimes(1); + const [url, init] = routeFetch.mock.calls[0]; + expect(url).toBe(routeUrl); + expect(new Headers(init?.headers).get('authorization')).toBe('Bearer test-kilo-key'); + expect(upstreamRequest().body).toEqual({ + ...JSON.parse(JSON.stringify(requestBody)), + model: TYPESAFE_MODEL, + user: 'hashed-user', + }); + expect(getUserFromAuth).toHaveBeenCalledWith({ + adminOnly: false, + expectedAudience: KILO_GATEWAY_AUDIENCE, + }); + expect(resolveOrganizationMemberModelDecision).not.toHaveBeenCalled(); + }); + + it.each([undefined, 'jev-1.13', 'typesafe/jev-1.13'])( + 'pins model %s to OpenRouter and drops client credentials and routing overrides', + async model => { + const response = await handleSystemOneRequest( + makeRequest( + { + ...requestBody, + model, + provider: { only: ['attacker'], api_key: 'test-body-key' }, + user: 'attacker-user', + api_key: 'test-body-key', + byok: true, + user_byok: [{ providerId: 'typesafe', apiKey: 'test-byok-key' }], + base_url: 'https://attacker.invalid', + }, + { + Authorization: 'Bearer test-client-key', + 'Proxy-Authorization': 'Bearer test-proxy-key', + 'x-api-key': 'test-forwarded-key', + 'x-openrouter-api-key': 'test-user-openrouter-key', + Cookie: 'session=test-cookie', + 'HTTP-Referer': 'https://attacker.invalid', + 'X-Title': 'Attacker', + } + ) + ); + + expect(response.status).toBe(200); + const upstream = upstreamRequest(); + expect(upstream.body).toEqual({ + ...JSON.parse(JSON.stringify(requestBody)), + model: 'typesafe/jev-1.13', + user: 'hashed-user', + }); + expect(Object.fromEntries(upstream.headers)).toEqual({ + authorization: 'Bearer test-platform-openrouter-key', + 'content-type': 'application/json', + 'http-referer': 'https://kilocode.ai', + 'x-title': 'Kilo Code', + }); + expect(generateProviderSpecificHash).toHaveBeenCalledWith(user.id, OPENROUTER); + expect(checkOrganizationModelRestrictions).toHaveBeenCalledWith({ + modelId: TYPESAFE_MODEL, + settings: undefined, + organizationPlan: undefined, + }); + } + ); + + it('defers billing and converts upstream cost to rounded microdollars without a markup', async () => { + setAuth('org-123'); + const response = await handleSystemOneRequest( + makeRequest(requestBody, { + 'x-kilocode-editorname': 'vscode', + 'x-kilocode-machineid': 'machine-123', + 'x-kilocode-feature': 'cli', + 'X-KiloCode-TaskId': 'task-123', + 'x-kilocode-mode': 'code', + }) + ); + + expect(response.status).toBe(200); + expect(await response.json()).toEqual(upstreamBody); + expect(logMicrodollarUsage).not.toHaveBeenCalled(); + await runAfter(); + expect(logMicrodollarUsage).toHaveBeenCalledTimes(1); + expect(logMicrodollarUsage).toHaveBeenCalledWith( + expect.objectContaining({ + messageId: upstreamBody.id, + model: TYPESAFE_MODEL, + inference_provider: 'TypeSafe upstream', + cost_mUsd: 124, + market_cost: 124, + inputTokens: 23, + outputTokens: 7, + cacheWriteTokens: 0, + cacheHitTokens: 0, + is_byok: false, + streamed: false, + hasError: false, + status_code: 200, + responseContent: '', + latency: expect.any(Number), + }), + expect.objectContaining({ + api_kind: 'systemone', + kiloUserId: user.id, + provider: 'openrouter', + requested_model: TYPESAFE_MODEL, + organizationId: 'org-123', + prior_microdollar_usage: 123, + posthog_distinct_id: user.google_user_email, + project_id: 'project-123', + fraudHeaders: EmptyFraudDetectionHeaders, + editor_name: 'vscode', + machine_id: 'machine-123', + feature: 'cli', + session_id: 'task-123', + mode: 'code', + user_byok: false, + isStreaming: false, + botId: 'bot-123', + tokenSource: 'api-key', + ttfb_ms: expect.any(Number), + }) + ); + }); + + it('accepts zero-cost usage and leaves an absent inference provider unknown', async () => { + mockedFetch.mockResolvedValue( + Response.json({ + ...upstreamBody, + provider: undefined, + usage: { ...upstreamBody.usage, cost: 0 }, + }) + ); + + expect((await handleSystemOneRequest(makeRequest())).status).toBe(200); + await runAfter(); + + expect(logMicrodollarUsage).toHaveBeenCalledWith( + expect.objectContaining({ inference_provider: null, cost_mUsd: 0, market_cost: 0 }), + expect.objectContaining({ provider: 'openrouter', user_byok: false }) + ); + }); + + it('returns the authentication failure and requests the gateway audience', async () => { + const authFailedResponse = NextResponse.json( + { success: false as const, error: 'Unauthorized' }, + { status: 401 } + ); + jest.mocked(getUserFromAuth).mockResolvedValue({ user: null, authFailedResponse }); + + expect(await handleSystemOneRequest(makeRequest())).toBe(authFailedResponse); + expect(getUserFromAuth).toHaveBeenCalledWith({ + adminOnly: false, + expectedAudience: KILO_GATEWAY_AUDIENCE, + }); + expect(getBalanceAndOrgSettings).not.toHaveBeenCalled(); + expect(mockedFetch).not.toHaveBeenCalled(); + expect(after).not.toHaveBeenCalled(); + }); + + it('rate limits before authentication or upstream work', async () => { + jest.mocked(isGatewayAccountRateLimited).mockResolvedValue(true); + const request = makeRequest(requestBody, { 'x-forwarded-for': ' 192.0.2.1, 192.0.2.2' }); + + const response = await handleSystemOneRequest(request); + + expect(response.status).toBe(429); + expect(await response.json()).toMatchObject({ error_type: 'rate_limit_exceeded' }); + expect(gatewayRateLimitKey).toHaveBeenCalledWith(request.headers, '192.0.2.1'); + expect(isGatewayAccountRateLimited).toHaveBeenCalledWith(request, 'test-rate-limit-key'); + expect(getUserFromAuth).not.toHaveBeenCalled(); + expect(mockedFetch).not.toHaveBeenCalled(); + }); + + it.each([ + ['malformed JSON', '{'], + ['another model', JSON.stringify({ ...requestBody, model: 'openai/gpt-4o' })], + ['an unpinned alias', JSON.stringify({ ...requestBody, model: 'jev-latest' })], + ['empty questions', JSON.stringify({ ...requestBody, questions: {} })], + ])('rejects %s before balance checks or upstream work', async (_name, body) => { + const response = await handleSystemOneRequest( + new NextRequest(routeUrl, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body, + }) + ); + + expect(response.status).toBe(400); + expect(await response.json()).toMatchObject({ error_type: 'invalid_request' }); + expect(getBalanceAndOrgSettings).not.toHaveBeenCalled(); + expect(mockedFetch).not.toHaveBeenCalled(); + expect(after).not.toHaveBeenCalled(); + }); + + it('formats validation errors with Zod while preserving the TypeSafe error shape', async () => { + const body = { ...requestBody, questions: { refund: { type: 'invalid' } } }; + const validation = systemOneRequestSchema.safeParse(body); + if (validation.success) throw new Error('Expected an invalid question'); + + const response = await handleSystemOneRequest(makeRequest(body)); + + expect(response.status).toBe(400); + expect(await response.json()).toEqual({ + message: z.prettifyError(validation.error), + error_type: 'invalid_request', + }); + expect(mockedFetch).not.toHaveBeenCalled(); + }); + + it.each([0, -1])('blocks balance %s using the shared credits response', async balance => { + setAuth('org-123'); + jest.mocked(getBalanceAndOrgSettings).mockResolvedValue({ + balance, + balanceLimitedByUserAllowance: true, + }); + const blockedResponse = NextResponse.json( + { + error_type: 'usage_limit_exceeded' as const, + error: { + title: 'Credits exhausted', + message: 'Credits exhausted', + balance, + buyCreditsUrl: '', + }, + }, + { status: 402 } + ); + jest.mocked(creditsBlockedResponse).mockResolvedValue(blockedResponse); + + expect(await handleSystemOneRequest(makeRequest())).toBe(blockedResponse); + expect(getBalanceAndOrgSettings).toHaveBeenCalledWith('org-123', user); + expect(creditsBlockedResponse).toHaveBeenCalledWith({ + user, + balance, + organizationId: 'org-123', + balanceLimitedByUserAllowance: true, + }); + expect(checkOrganizationModelRestrictions).not.toHaveBeenCalled(); + expect(mockedFetch).not.toHaveBeenCalled(); + expect(after).not.toHaveBeenCalled(); + }); + + it('honors organization model restrictions before resolving member access', async () => { + setAuth('org-123'); + const settings = { model_deny_list: [TYPESAFE_MODEL] }; + jest.mocked(getBalanceAndOrgSettings).mockResolvedValue({ + balance: 1000, + settings, + plan: 'enterprise', + }); + const error = NextResponse.json({ error_type: 'model_not_allowed' }, { status: 404 }); + jest.mocked(checkOrganizationModelRestrictions).mockReturnValue({ error }); + + expect(await handleSystemOneRequest(makeRequest())).toBe(error); + expect(checkOrganizationModelRestrictions).toHaveBeenCalledWith({ + modelId: TYPESAFE_MODEL, + settings, + organizationPlan: 'enterprise', + }); + expect(resolveOrganizationMemberModelDecision).not.toHaveBeenCalled(); + expect(mockedFetch).not.toHaveBeenCalled(); + expect(after).not.toHaveBeenCalled(); + }); + + it('denies a model rejected by the effective organization member decision', async () => { + setAuth('org-123'); + jest.mocked(resolveOrganizationMemberModelDecision).mockResolvedValue({ + ...memberDecision, + decision: { allowed: false }, + }); + + expect((await handleSystemOneRequest(makeRequest())).status).toBe(404); + expect(resolveOrganizationMemberModelDecision).toHaveBeenCalledWith({ + organizationId: 'org-123', + kiloUserId: user.id, + modelId: TYPESAFE_MODEL, + providerLookup: expect.any(Function), + }); + const [{ providerLookup }] = jest.mocked(resolveOrganizationMemberModelDecision).mock.calls[0]; + if (!providerLookup) throw new Error('Expected the fixed TypeSafe provider lookup'); + await expect(providerLookup(TYPESAFE_MODEL)).resolves.toEqual(new Set(['typesafe'])); + expect(modelNotAllowedResponse).toHaveBeenCalledTimes(1); + expect(mockedFetch).not.toHaveBeenCalled(); + expect(after).not.toHaveBeenCalled(); + }); + + it.each([ + { + only: ['typesafe', 'other'], + eligible: ['typesafe', 'outside-ceiling'], + expected: ['typesafe'], + }, + { only: undefined, eligible: ['typesafe'], expected: ['typesafe'] }, + { only: ['typesafe'], eligible: undefined, expected: ['typesafe'] }, + ])('applies the provider policy intersection: %j', async ({ only, eligible, expected }) => { + setAuth('org-123'); + jest.mocked(checkOrganizationModelRestrictions).mockReturnValue({ + error: null, + providerConfig: { only, data_collection: 'deny' }, + }); + jest.mocked(resolveOrganizationMemberModelDecision).mockResolvedValue({ + ...memberDecision, + decision: { allowed: true, eligibleProviderRoutes: eligible ? new Set(eligible) : undefined }, + }); + + const response = await handleSystemOneRequest( + makeRequest({ ...requestBody, provider: { only: ['attacker'] } }) + ); + + expect(response.status).toBe(200); + expect(upstreamRequest().body.provider).toEqual({ only: expected, data_collection: 'deny' }); + }); + + it.each([ + { only: ['other'], eligible: ['typesafe'] }, + { only: [], eligible: ['typesafe'] }, + { only: undefined, eligible: [] }, + ])('rejects an empty provider policy intersection: %j', async ({ only, eligible }) => { + setAuth('org-123'); + jest.mocked(checkOrganizationModelRestrictions).mockReturnValue({ + error: null, + providerConfig: { only, data_collection: 'deny' }, + }); + jest.mocked(resolveOrganizationMemberModelDecision).mockResolvedValue({ + ...memberDecision, + decision: { allowed: true, eligibleProviderRoutes: new Set(eligible) }, + }); + + expect((await handleSystemOneRequest(makeRequest())).status).toBe(404); + expect(modelNotAllowedResponse).toHaveBeenCalledTimes(1); + expect(mockedFetch).not.toHaveBeenCalled(); + expect(after).not.toHaveBeenCalled(); + }); + + it('cancels the upstream body on credit exhaustion and returns a service error without charging', async () => { + const cancel = jest.fn(); + mockedFetch.mockResolvedValue(new Response(new ReadableStream({ cancel }), { status: 402 })); + + const response = await handleSystemOneRequest(makeRequest()); + + expect(response.status).toBe(503); + expect(cancel).toHaveBeenCalledTimes(1); + expect(await response.json()).toEqual({ + message: 'Service temporarily unavailable', + error_type: 'upstream_error', + }); + expect(errorExceptInTest).toHaveBeenCalledWith('OpenRouter System One balance exhausted'); + expect(wrapInSafeNextResponse).not.toHaveBeenCalled(); + expect(after).not.toHaveBeenCalled(); + expect(logMicrodollarUsage).not.toHaveBeenCalled(); + }); + + it.each([400, 429, 500])( + 'delegates upstream HTTP %s to the safe response wrapper', + async status => { + const upstream = Response.json({ error: 'upstream failure' }, { status }); + const safeResponse = NextResponse.json({ error: 'upstream failure' }, { status }); + mockedFetch.mockResolvedValue(upstream); + jest.mocked(wrapInSafeNextResponse).mockReturnValue(safeResponse); + + expect(await handleSystemOneRequest(makeRequest())).toBe(safeResponse); + expect(wrapInSafeNextResponse).toHaveBeenCalledWith(upstream); + expect(after).not.toHaveBeenCalled(); + expect(logMicrodollarUsage).not.toHaveBeenCalled(); + } + ); + + it.each(['network failure', 'non-JSON success'])('handles %s without charging', async failure => { + if (failure === 'network failure') { + mockedFetch.mockRejectedValue(new Error('connection failed')); + } else { + mockedFetch.mockResolvedValue(new Response('not JSON', { status: 200 })); + } + + const response = await handleSystemOneRequest(makeRequest()); + + expect(response.status).toBe(502); + expect(await response.json()).toEqual({ + message: 'Upstream request failed', + error_type: 'upstream_error', + }); + expect(errorExceptInTest).toHaveBeenCalledWith( + 'OpenRouter System One request failed', + expect.objectContaining({ message: expect.any(String) }) + ); + expect(after).not.toHaveBeenCalled(); + expect(logMicrodollarUsage).not.toHaveBeenCalled(); + }); + + it.each([ + ['missing usage', undefined], + ['null usage', null], + ['missing cost', { input_tokens: 23, output_tokens: 7 }], + ['string cost', { ...upstreamBody.usage, cost: '0.01' }], + ['negative cost', { ...upstreamBody.usage, cost: -0.01 }], + ['missing input tokens', { output_tokens: 7, cost: 0.01 }], + ['negative tokens', { ...upstreamBody.usage, input_tokens: -1 }], + ['fractional tokens', { ...upstreamBody.usage, output_tokens: 1.5 }], + ])('rejects upstream %s without scheduling billing', async (_name, usage) => { + mockedFetch.mockResolvedValue(Response.json({ ...upstreamBody, usage })); + + const response = await handleSystemOneRequest(makeRequest()); + + expect(response.status).toBe(502); + expect(await response.json()).toEqual({ + message: 'Invalid upstream response', + error_type: 'upstream_error', + }); + expect(errorExceptInTest).toHaveBeenCalledWith( + 'Invalid OpenRouter System One response or missing usage' + ); + expect(after).not.toHaveBeenCalled(); + expect(logMicrodollarUsage).not.toHaveBeenCalled(); + }); +});