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..50e9e54f9e --- /dev/null +++ b/apps/web/src/app/api/gateway/typesafe/v1/systemone/route.test.ts @@ -0,0 +1,10 @@ +import { expect, it } from '@jest/globals'; +import { handleSystemOneRequest } from '@/lib/ai-gateway/typesafe/handler'; +import { maxDuration, POST } from './route'; + +jest.mock('@/lib/ai-gateway/typesafe/handler', () => ({ handleSystemOneRequest: jest.fn() })); + +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/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..9053460883 --- /dev/null +++ b/apps/web/src/app/api/gateway/typesafe/v1/systemone/route.ts @@ -0,0 +1,3 @@ +export { handleSystemOneRequest as POST } from '@/lib/ai-gateway/typesafe/handler'; + +export const maxDuration = 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(); + }); +}); 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..be7abd0793 --- /dev/null +++ b/apps/web/src/lib/ai-gateway/typesafe/handler.ts @@ -0,0 +1,187 @@ +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'; +import type { ProxyErrorType } from '@/lib/proxy-error-types'; + +function errorResponse(message: string, error_type: ProxyErrorType, 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; + + 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) { + await response.body?.cancel(); + 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 ?? null, + 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/schemas.test.ts b/apps/web/src/lib/ai-gateway/typesafe/schemas.test.ts new file mode 100644 index 0000000000..7046550f43 --- /dev/null +++ b/apps/web/src/lib/ai-gateway/typesafe/schemas.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 './schemas'; + +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/schemas.ts b/apps/web/src/lib/ai-gateway/typesafe/schemas.ts new file mode 100644 index 0000000000..cc291ec7ea --- /dev/null +++ b/apps/web/src/lib/ai-gateway/typesafe/schemas.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/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'; diff --git a/packages/db/src/schema-types.ts b/packages/db/src/schema-types.ts index a9771a4828..868bc09196 100644 --- a/packages/db/src/schema-types.ts +++ b/packages/db/src/schema-types.ts @@ -1266,6 +1266,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 4b8362c39e..f0968647d9 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -1063,6 +1063,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 @@ -10628,6 +10631,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} @@ -28568,6 +28575,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)