diff --git a/apps/web/next.config.mjs b/apps/web/next.config.mjs index f45ba76570..58b87813b8 100644 --- a/apps/web/next.config.mjs +++ b/apps/web/next.config.mjs @@ -60,6 +60,10 @@ const nextConfig = { source: '/api/fim/completions', destination: 'https://global-api.kilo.ai/api/fim/completions', }, + { + source: '/api/edit/completions', + destination: 'https://global-api.kilo.ai/api/edit/completions', + }, { source: '/api/exa/:path*', destination: 'https://global-api.kilo.ai/api/exa/:path*', diff --git a/apps/web/src/app/api/edit/completions/route.test.ts b/apps/web/src/app/api/edit/completions/route.test.ts new file mode 100644 index 0000000000..03e5efa674 --- /dev/null +++ b/apps/web/src/app/api/edit/completions/route.test.ts @@ -0,0 +1,310 @@ +import { describe, it, expect, beforeEach, afterAll } from '@jest/globals'; +import type { User } from '@kilocode/db/schema'; +import type { OrganizationSettings } from '@/lib/organizations/organization-types'; +import { ProxyErrorType } from '@/lib/proxy-error-types'; +import { getUserFromAuth } from '@/lib/user/server'; +import { getBalanceAndOrgSettings } from '@/lib/organizations/organization-usage'; +import { getBYOKforOrganization, getBYOKforUser } from '@/lib/ai-gateway/byok'; +import type { + MicrodollarUsageContext, + MicrodollarUsageStats, +} from '@/lib/ai-gateway/processUsage.types'; + +jest.mock('@/lib/config.server', () => ({ + INCEPTION_API_KEY: 'system-inception-key', +})); +jest.mock('@/lib/user/server'); +jest.mock('@/lib/organizations/organization-usage'); +jest.mock('@/lib/ai-gateway/byok'); +jest.mock('@/lib/debugUtils', () => ({ + debugSaveProxyRequest: jest.fn(), + debugSaveProxyResponseStream: jest.fn(), +})); + +// Run `after()` work inline so we can assert on the usage write the same tick. +jest.mock('next/server', () => ({ + ...(jest.requireActual('next/server') as Record), + after: jest.fn((work: Promise | (() => Promise)) => { + void (typeof work === 'function' ? work() : work); + }), +})); + +// Capture the persisted usage row instead of round-tripping through the DB. +// The route exercises the real `countAndStoreEditUsage`, which means we get +// end-to-end coverage of the BYOK / cache-discount zeroing path. +const mockedLogMicrodollarUsage = jest.fn( + async (_stats: MicrodollarUsageStats, _ctx: MicrodollarUsageContext) => null +); +jest.mock('@/lib/ai-gateway/processUsage', () => ({ + ...(jest.requireActual('@/lib/ai-gateway/processUsage') as Record), + logMicrodollarUsage: (stats: MicrodollarUsageStats, ctx: MicrodollarUsageContext) => + mockedLogMicrodollarUsage(stats, ctx), +})); + +const mockedGetUserFromAuth = jest.mocked(getUserFromAuth); +const mockedGetBalanceAndOrgSettings = jest.mocked(getBalanceAndOrgSettings); +const mockedGetBYOKforOrganization = jest.mocked(getBYOKforOrganization); +const mockedGetBYOKforUser = jest.mocked(getBYOKforUser); +const mockedFetch = jest.fn() as jest.MockedFunction; +const originalFetch = globalThis.fetch; + +function makeRequest(body: unknown) { + return new Request('http://localhost:3000/api/edit/completions', { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + 'x-forwarded-for': '127.0.0.1', + }, + body: JSON.stringify(body), + }); +} + +function setOrganizationAuth(settings?: OrganizationSettings) { + mockedGetUserFromAuth.mockResolvedValue({ + user: { + id: 'user-123', + google_user_email: 'test@example.com', + microdollars_used: 0, + } as User, + authFailedResponse: null, + organizationId: 'org-123', + }); + mockedGetBalanceAndOrgSettings.mockResolvedValue({ + balance: 1000, + settings, + plan: 'teams', + }); + mockedGetBYOKforOrganization.mockResolvedValue(null); + mockedGetBYOKforUser.mockResolvedValue(null); +} + +function setBYOKAuth() { + mockedGetUserFromAuth.mockResolvedValue({ + user: { + id: 'user-byok', + google_user_email: 'byok@example.com', + microdollars_used: 0, + } as User, + authFailedResponse: null, + organizationId: undefined, + }); + mockedGetBalanceAndOrgSettings.mockResolvedValue({ + balance: 0, + settings: undefined, + plan: undefined, + }); + mockedGetBYOKforOrganization.mockResolvedValue(null); + mockedGetBYOKforUser.mockResolvedValue([ + { providerId: 'inception', decryptedAPIKey: 'user-supplied-key' }, + ] as never); +} + +function makeValidRequestBody() { + return { + model: 'inception/mercury-edit-2', + messages: [{ role: 'user', content: '<|code_to_edit|>const a = 1<|/code_to_edit|>' }], + max_tokens: 100, + }; +} + +function makeUpstreamResponse(payload?: unknown) { + return new Response( + JSON.stringify( + payload ?? { + id: 'edit-test', + model: 'mercury-edit-2', + usage: { + prompt_tokens: 100_000, + cached_input_tokens: 90_000, + completion_tokens: 0, + total_tokens: 100_000, + }, + choices: [{ message: { role: 'assistant', content: 'edited' } }], + } + ), + { + status: 200, + headers: { 'content-type': 'application/json' }, + } + ); +} + +async function flushAfter() { + // `after()` invocations are scheduled as microtasks; let them settle before + // asserting on what `logMicrodollarUsage` received. + await new Promise(resolve => setImmediate(resolve)); +} + +describe('POST /api/edit/completions', () => { + beforeEach(() => { + jest.resetAllMocks(); + globalThis.fetch = mockedFetch; + mockedLogMicrodollarUsage.mockResolvedValue(null); + }); + + afterAll(() => { + globalThis.fetch = originalFetch; + }); + + it('rejects unsupported edit models with the dedicated error type', async () => { + setOrganizationAuth(); + + const { POST } = await import('./route'); + const response = await POST( + makeRequest({ ...makeValidRequestBody(), model: 'mistralai/codestral' }) as never + ); + + expect(response.status).toBe(400); + expect(await response.json()).toMatchObject({ + error_type: ProxyErrorType.unsupported_edit_model, + }); + expect(mockedFetch).not.toHaveBeenCalled(); + }); + + it('rejects requests with non-positive max_tokens', async () => { + setOrganizationAuth(); + + const { POST } = await import('./route'); + const response = await POST( + makeRequest({ ...makeValidRequestBody(), max_tokens: -1 }) as never + ); + + // -1 fails the schema's `.positive()` so the route returns invalid_request. + expect(response.status).toBe(400); + expect(mockedFetch).not.toHaveBeenCalled(); + }); + + it('rejects direct Inception requests when organization data collection is denied', async () => { + setOrganizationAuth({ data_collection: 'deny' } satisfies OrganizationSettings); + + const { POST } = await import('./route'); + const response = await POST(makeRequest(makeValidRequestBody()) as never); + + expect(response.status).toBe(400); + expect(await response.json()).toMatchObject({ + error_type: ProxyErrorType.data_collection_required, + }); + expect(mockedFetch).not.toHaveBeenCalled(); + }); + + it.each([ + { messages: [{ role: 'system', content: 'Do not forward system prompts' }] }, + { messages: [{ role: 'assistant', content: 'Do not forward assistant content' }] }, + { + messages: [ + { role: 'user', content: 'First message' }, + { role: 'user', content: 'Second message' }, + ], + }, + ])('rejects unsupported edit messages before proxying', async ({ messages }) => { + setOrganizationAuth(); + + const { POST } = await import('./route'); + const response = await POST(makeRequest({ ...makeValidRequestBody(), messages }) as never); + + expect(response.status).toBe(400); + expect(mockedFetch).not.toHaveBeenCalled(); + }); + + it('rejects requests when balance is exhausted and the user has no BYOK', async () => { + setOrganizationAuth(); + mockedGetBalanceAndOrgSettings.mockResolvedValue({ + balance: 0, + settings: undefined, + plan: 'teams', + }); + + const { POST } = await import('./route'); + const response = await POST(makeRequest(makeValidRequestBody()) as never); + + expect(response.status).toBe(402); + expect(await response.json()).toMatchObject({ + error_type: ProxyErrorType.insufficient_credits, + }); + expect(mockedFetch).not.toHaveBeenCalled(); + }); + + it('forwards a single user message to Inception', async () => { + setOrganizationAuth(); + mockedFetch.mockResolvedValue(makeUpstreamResponse()); + + const { POST } = await import('./route'); + const requestBody = makeValidRequestBody(); + const response = await POST(makeRequest(requestBody) as never); + + expect(response.status).toBe(200); + const [url, init] = mockedFetch.mock.calls[0]; + expect(url).toBe('https://api.inceptionlabs.ai/v1/edit/completions'); + const upstreamBody = JSON.parse(init?.body as string); + expect(upstreamBody.messages).toEqual(requestBody.messages); + // Provider prefix is stripped before forwarding upstream. + expect(upstreamBody.model).toBe('mercury-edit-2'); + }); + + it('persists computed cost and cache discount for paid (non-BYOK) requests', async () => { + setOrganizationAuth(); + mockedFetch.mockResolvedValue(makeUpstreamResponse()); + + const { POST } = await import('./route'); + const response = await POST(makeRequest(makeValidRequestBody()) as never); + expect(response.status).toBe(200); + + await flushAfter(); + + expect(mockedLogMicrodollarUsage).toHaveBeenCalledTimes(1); + const [stats, ctx] = mockedLogMicrodollarUsage.mock.calls[0]; + expect(ctx.api_kind).toBe('edit_completions'); + expect(ctx.user_byok).toBe(false); + expect(stats.cost_mUsd).toBe(4_750); + expect(stats.cacheDiscount_mUsd).toBe(20_250); + expect(stats.market_cost).toBe(4_750); + }); + + it('zeroes both cost and cache discount on the persisted row for BYOK requests', async () => { + setBYOKAuth(); + mockedFetch.mockResolvedValue(makeUpstreamResponse()); + + const { POST } = await import('./route'); + const response = await POST(makeRequest(makeValidRequestBody()) as never); + expect(response.status).toBe(200); + + await flushAfter(); + + expect(mockedLogMicrodollarUsage).toHaveBeenCalledTimes(1); + const [stats, ctx] = mockedLogMicrodollarUsage.mock.calls[0]; + expect(ctx.user_byok).toBe(true); + expect(stats.cost_mUsd).toBe(0); + expect(stats.cacheDiscount_mUsd).toBe(0); + // The original cost is preserved in market_cost for reporting. + expect(stats.market_cost).toBe(4_750); + + // Sanity: BYOK requests use the user's API key, not the system one. + const [, init] = mockedFetch.mock.calls[0]; + const headers = init?.headers as Record; + expect(headers.Authorization).toBe('Bearer user-supplied-key'); + }); + + it('persists a zero-cost row when the upstream response omits usage', async () => { + setOrganizationAuth(); + mockedFetch.mockResolvedValue( + makeUpstreamResponse({ + id: 'edit-no-usage', + model: 'mercury-edit-2', + choices: [{ message: { role: 'assistant', content: 'edited' } }], + }) + ); + + const { POST } = await import('./route'); + const response = await POST(makeRequest(makeValidRequestBody()) as never); + expect(response.status).toBe(200); + + await flushAfter(); + + expect(mockedLogMicrodollarUsage).toHaveBeenCalledTimes(1); + const [stats] = mockedLogMicrodollarUsage.mock.calls[0]; + expect(stats.cost_mUsd).toBe(0); + expect(stats.cacheDiscount_mUsd).toBeUndefined(); + expect(stats.inputTokens).toBe(0); + expect(stats.outputTokens).toBe(0); + }); +}); diff --git a/apps/web/src/app/api/edit/completions/route.ts b/apps/web/src/app/api/edit/completions/route.ts new file mode 100644 index 0000000000..8b54aeaf95 --- /dev/null +++ b/apps/web/src/app/api/edit/completions/route.ts @@ -0,0 +1,306 @@ +import { INCEPTION_API_KEY } from '@/lib/config.server'; +import type { NextRequest } from 'next/server'; +import { NextResponse } from 'next/server'; +import z from 'zod'; +import { captureException, setTag, startInactiveSpan } from '@sentry/nextjs'; +import type { MicrodollarUsageContext } from '@/lib/ai-gateway/processUsage.types'; +import { validateFeatureHeader, FEATURE_HEADER } from '@/lib/feature-detection'; +import { isFreeModel } from '@/lib/ai-gateway/is-free-model'; +import { sentryRootSpan } from '@/lib/getRootSpan'; +import { getUserFromAuth } from '@/lib/user/server'; +import { + checkOrganizationModelRestrictions, + countAndStoreEditUsage, + extractEditPromptInfo, + extractFraudAndProjectHeaders, + invalidRequestResponse, + dataCollectionRequiredResponse, + temporarilyUnavailableResponse, + wrapInSafeNextResponse, + captureProxyError, + extractHeaderAndLimitLength, +} from '@/lib/ai-gateway/llm-proxy-helpers'; +import { ProxyErrorType } from '@/lib/proxy-error-types'; +import { getBalanceAndOrgSettings } from '@/lib/organizations/organization-usage'; +import { readDb } from '@/lib/drizzle'; +import { debugSaveProxyRequest } from '@/lib/debugUtils'; +import { sentryLogger } from '@/lib/utils.server'; +import { getBYOKforOrganization, getBYOKforUser } from '@/lib/ai-gateway/byok'; +import type { UserByokProviderId } from '@/lib/ai-gateway/providers/openrouter/inference-provider-id'; + +// Inception's edit endpoint mirrors a chat completion shape but is hosted at +// a separate path. It accepts a single `role: "user"` message; the system prompt +// is baked in server-side and the endpoint returns 400 for any `role: "system"` +// message. See https://docs.inceptionlabs.ai/api-reference/edit/create-a-code-edit-completion +const INCEPTION_EDIT_URL = 'https://api.inceptionlabs.ai/v1/edit/completions'; +const EDIT_MAX_TOKENS_LIMIT = 1000; + +type EditProvider = 'inception'; + +function resolveEditProvider(model: string): { + provider: EditProvider; + upstreamModel: string; +} | null { + if (model.startsWith('inception/')) { + return { + provider: 'inception', + upstreamModel: model.slice('inception/'.length), + }; + } + return null; +} + +function getSystemApiKey(provider: EditProvider): string | null { + switch (provider) { + case 'inception': + return INCEPTION_API_KEY || null; + } +} + +const EditMessage = z.object({ + role: z.literal('user'), + content: z.string(), +}); + +const EditRequestBody = z.object({ + model: z.string(), + messages: z.array(EditMessage).length(1), + max_tokens: z.number().int().positive().optional(), + stop: z.string().array().optional(), + // Streaming is not supported by Inception's edit endpoint today; reject if requested. + stream: z.literal(false).optional(), +}); + +type EditRequestBody = z.infer; + +export async function POST(request: NextRequest) { + const requestStartedAt = performance.now(); + const requestBodyTextPromise = request.text(); + + const authSpan = startInactiveSpan({ name: 'auth-check' }); + const { + user: maybeUser, + authFailedResponse, + organizationId, + } = await getUserFromAuth({ adminOnly: false }); + authSpan.end(); + if (authFailedResponse) return authFailedResponse; + + const user = maybeUser; + const requestBodyText = await requestBodyTextPromise; + debugSaveProxyRequest(requestBodyText); + + let requestBody: EditRequestBody; + try { + const { success, data, error } = EditRequestBody.safeParse(JSON.parse(requestBodyText)); + if (!success) { + if (error.issues.some(issue => issue.path[0] === 'stream')) { + return NextResponse.json( + { + error: 'Streaming is not supported for edit completions', + error_type: ProxyErrorType.unsupported_field, + }, + { status: 400 } + ); + } + sentryLogger('edit-proxy')('request failed to parse', { + extra: { kiloUserId: user.id, error, organizationId }, + tags: { source: 'edit-proxy' }, + user: { id: user.id }, + }); + return invalidRequestResponse(); + } + requestBody = data; + } catch (e) { + captureException(e, { + extra: { kiloUserId: user.id }, + tags: { source: 'edit-proxy' }, + user: { id: user.id }, + }); + return invalidRequestResponse(); + } + + const resolved = resolveEditProvider(requestBody.model); + if (!resolved) { + return NextResponse.json( + { + error: requestBody.model + ' is not a supported edit model', + error_type: ProxyErrorType.unsupported_edit_model, + }, + { status: 400 } + ); + } + const { provider: editProvider, upstreamModel } = resolved; + + if (!requestBody.max_tokens || requestBody.max_tokens > EDIT_MAX_TOKENS_LIMIT) { + console.warn(`SECURITY: Edit max tokens limit exceeded or missing: ${user.id}`, { + maxTokens: requestBody.max_tokens, + }); + return temporarilyUnavailableResponse(); + } + + const { fraudHeaders, projectId } = extractFraudAndProjectHeaders(request); + const taskId = extractHeaderAndLimitLength(request, 'x-kilocode-taskid') ?? undefined; + + const promptInfo = extractEditPromptInfo(requestBody); + + const byokProviderKey: UserByokProviderId = 'inception'; + + const userByok = organizationId + ? await getBYOKforOrganization(readDb, organizationId, [byokProviderKey]) + : await getBYOKforUser(readDb, user.id, [byokProviderKey]); + + const usageContext: MicrodollarUsageContext = { + api_kind: 'edit_completions', + kiloUserId: user.id, + provider: editProvider, + requested_model: requestBody.model, + promptInfo, + max_tokens: requestBody.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: null, + editor_name: extractHeaderAndLimitLength(request, 'x-kilocode-editorname'), + machine_id: extractHeaderAndLimitLength(request, 'x-kilocode-machineid'), + user_byok: !!userByok, + has_tools: false, + feature: validateFeatureHeader(request.headers.get(FEATURE_HEADER)), + session_id: taskId ?? null, + mode: null, + auto_model: null, + ttfb_ms: null, + }; + + setTag('ui.ai_model', requestBody.model); + + // Use read replica for balance check - this is a read-only operation that can tolerate + // slight replication lag, and provides lower latency for US users. + const { balance, settings, plan } = await getBalanceAndOrgSettings(organizationId, user, readDb); + + if (balance <= 0 && !(await isFreeModel(requestBody.model)) && !userByok) { + return NextResponse.json( + { + error: { message: 'Insufficient credits' }, + error_type: ProxyErrorType.insufficient_credits, + }, + { status: 402 } + ); + } + + const { error: modelRestrictionError, providerConfig } = checkOrganizationModelRestrictions({ + modelId: requestBody.model, + settings, + organizationPlan: plan, + }); + if (modelRestrictionError) return modelRestrictionError; + + // Org-level "do not collect my data" opt-out. The OpenRouter/Vercel paths + // honor this by setting `provider.data_collection = 'deny'` on the upstream + // request body, which causes the gateway to route to a sub-provider with a + // no-training/no-retention contract. This route bypasses both gateways and + // POSTs straight to Inception, and Inception's edit endpoint exposes no + // per-request opt-out flag. Their public privacy policy + // (https://www.inceptionlabs.ai/docs/privacy-policy) lists "Personal Data + // contained in prompts, inputs and uploaded content processed by our models" + // among the data they collect for purposes including "training and refining + // our models"; "no training / no retention" is only offered as an enterprise + // feature (https://www.inceptionlabs.ai/enterprise), not as the default for + // the standard API tier we call here. Until we sign an enterprise agreement + // or Inception adds a per-request flag, refusing is the only way to honor + // the org's stated intent. + if (providerConfig?.data_collection === 'deny') { + return dataCollectionRequiredResponse(); + } + + if (providerConfig?.only && !providerConfig.only.includes(editProvider)) { + return NextResponse.json( + { + error: 'Provider not allowed for your team.', + error_type: ProxyErrorType.provider_not_allowed, + message: `The provider "${editProvider}" is not allowed for your team.`, + }, + { status: 403 } + ); + } + + if (providerConfig?.ignore?.includes(editProvider)) { + return NextResponse.json( + { + error: 'Provider not allowed for your team.', + error_type: ProxyErrorType.provider_not_allowed, + message: `The provider "${editProvider}" is not allowed for your team.`, + }, + { status: 403 } + ); + } + + const systemKey = getSystemApiKey(editProvider); + const userByokEntry = userByok?.at(0); + const apiKey = userByokEntry?.decryptedAPIKey ?? systemKey; + + if (!apiKey) { + return NextResponse.json( + { + error: 'This model requires a BYOK API key. Please configure your API key in settings.', + error_type: ProxyErrorType.byok_key_required, + }, + { status: 400 } + ); + } + + sentryRootSpan()?.setAttribute( + 'edit.time_to_request_start_ms', + performance.now() - requestStartedAt + ); + + const requestSpan = startInactiveSpan({ + name: 'edit-request-start', + op: 'http.client', + }); + + const bodyForUpstream = { ...requestBody, model: upstreamModel }; + + const proxyRes = await fetch(INCEPTION_EDIT_URL, { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + Authorization: `Bearer ${apiKey}`, + }, + body: JSON.stringify(bodyForUpstream), + }); + usageContext.ttfb_ms = Math.max(0, Math.round(performance.now() - requestStartedAt)); + usageContext.status_code = proxyRes.status; + + if (!proxyRes.body) { + return NextResponse.json( + { + error: 'No body returned from upstream', + error_type: ProxyErrorType.upstream_error, + }, + { status: 500 } + ); + } + + if (proxyRes.status >= 400) { + await captureProxyError({ + user, + request: bodyForUpstream, + response: proxyRes, + organizationId, + model: requestBody.model, + errorMessage: `Edit provider returned error ${proxyRes.status}`, + trackInSentry: proxyRes.status >= 500, + }); + } + + const clonedResponse = proxyRes.clone(); + + countAndStoreEditUsage(clonedResponse, usageContext, requestSpan); + + return wrapInSafeNextResponse(proxyRes); +} diff --git a/apps/web/src/lib/ai-gateway/llm-proxy-helpers.test.ts b/apps/web/src/lib/ai-gateway/llm-proxy-helpers.test.ts index e7d588fd96..1c6e65ee67 100644 --- a/apps/web/src/lib/ai-gateway/llm-proxy-helpers.test.ts +++ b/apps/web/src/lib/ai-gateway/llm-proxy-helpers.test.ts @@ -1,9 +1,36 @@ -import { describe, it, expect } from '@jest/globals'; +import { describe, it, expect, beforeEach } from '@jest/globals'; +import type { MicrodollarUsageContext, MicrodollarUsageStats } from './processUsage.types'; + +// `countAndStoreEditUsage` schedules the usage write through `next/server`'s +// `after()` post-response hook, which only works in a request context. Replace +// it with an immediate invocation so the test can await the work synchronously. +jest.mock('next/server', () => ({ + ...(jest.requireActual('next/server') as Record), + after: jest.fn((work: Promise | (() => Promise)) => { + void (typeof work === 'function' ? work() : work); + }), +})); + +// Capture writes that would otherwise hit the database. The helper passes +// the final, post-zeroing `usageStats` to `logMicrodollarUsage`, so spying +// here lets us assert on the persisted billing shape directly. +const mockedLogMicrodollarUsage = jest.fn( + async (_stats: MicrodollarUsageStats, _ctx: MicrodollarUsageContext) => null +); +jest.mock('./processUsage', () => ({ + ...(jest.requireActual('./processUsage') as Record), + logMicrodollarUsage: (stats: MicrodollarUsageStats, ctx: MicrodollarUsageContext) => + mockedLogMicrodollarUsage(stats, ctx), +})); + import { checkOrganizationModelRestrictions, + countAndStoreEditUsage, + extractEditPromptInfo, extractEmbeddingPromptInfo, makeErrorReadable, parseEmbeddingUsageFromResponse, + parseEditUsageFromResponse, parseTranscriptionUsageFromResponse, } from './llm-proxy-helpers'; @@ -272,6 +299,227 @@ describe('extractEmbeddingPromptInfo', () => { }); }); +describe('extractEditPromptInfo', () => { + it('uses zero system prompt length because edit has no explicit system prompt', () => { + const result = extractEditPromptInfo({ + messages: [{ role: 'user', content: '<|code_to_edit|>const a = 1<|/code_to_edit|>' }], + }); + + expect(result.system_prompt_prefix).toBe(''); + expect(result.system_prompt_length).toBe(0); + expect(result.user_prompt_prefix).toBe('<|code_to_edit|>const a = 1<|/code_to_edit|>'); + }); +}); + +describe('parseEditUsageFromResponse', () => { + it('prices cached Inception input tokens at the discounted rate', () => { + const result = parseEditUsageFromResponse( + JSON.stringify({ + id: 'edit-123', + model: 'mercury-edit-2', + usage: { + prompt_tokens: 100_000, + cached_input_tokens: 90_000, + completion_tokens: 0, + total_tokens: 100_000, + }, + choices: [], + }), + 'inception', + 200 + ); + + expect(result.inputTokens).toBe(100_000); + expect(result.cacheHitTokens).toBe(90_000); + expect(result.cost_mUsd).toBe(4_750); + expect(result.cacheDiscount_mUsd).toBe(20_250); + }); + + it('falls back to flat Inception pricing when cached_input_tokens is absent', () => { + const result = parseEditUsageFromResponse( + JSON.stringify({ + id: 'edit-456', + model: 'mercury-edit-2', + usage: { + prompt_tokens: 1_000, + completion_tokens: 100, + total_tokens: 1_100, + }, + choices: [{ message: { role: 'assistant', content: 'edited' } }], + }), + 'inception', + 200 + ); + + expect(result.cacheHitTokens).toBe(0); + expect(result.cost_mUsd).toBe(Math.round(1_000 * 0.25 + 100 * 0.75)); + expect(result.cacheDiscount_mUsd).toBe(0); + expect(result.hasError).toBe(false); + }); + + it('returns zero cost when usage is absent', () => { + const result = parseEditUsageFromResponse( + JSON.stringify({ + id: 'edit-789', + model: 'mercury-edit-2', + choices: [], + }), + 'inception', + 200 + ); + + expect(result.inputTokens).toBe(0); + expect(result.outputTokens).toBe(0); + expect(result.cacheHitTokens).toBe(0); + expect(result.cost_mUsd).toBe(0); + expect(result.cacheDiscount_mUsd).toBeUndefined(); + }); + + it('flags an error and zero cost on upstream 4xx responses', () => { + const result = parseEditUsageFromResponse( + JSON.stringify({ error: { message: 'bad request' } }), + 'inception', + 400 + ); + + expect(result.hasError).toBe(true); + expect(result.cost_mUsd).toBe(0); + expect(result.model).toBeNull(); + }); + + it('clamps cached_input_tokens that exceed prompt_tokens', () => { + const result = parseEditUsageFromResponse( + JSON.stringify({ + id: 'edit-clamp', + model: 'mercury-edit-2', + usage: { + prompt_tokens: 1_000, + cached_input_tokens: 5_000, + completion_tokens: 0, + total_tokens: 1_000, + }, + choices: [], + }), + 'inception', + 200 + ); + + // Without the clamp, uncachedInputTokens would be negative and produce a + // negative cost. The clamp pins cacheHitTokens at prompt_tokens. + expect(result.cacheHitTokens).toBe(1_000); + expect(result.cost_mUsd).toBe(25); + }); +}); + +describe('countAndStoreEditUsage', () => { + function makeUsageContext( + overrides: Partial = {} + ): MicrodollarUsageContext { + return { + api_kind: 'edit_completions', + kiloUserId: 'user-edit-test', + provider: 'inception', + requested_model: 'inception/mercury-edit-2', + promptInfo: { + system_prompt_prefix: '', + system_prompt_length: 0, + user_prompt_prefix: '', + }, + max_tokens: 100, + has_middle_out_transform: null, + fraudHeaders: {}, + isStreaming: false, + organizationId: undefined, + prior_microdollar_usage: 0, + posthog_distinct_id: undefined, + project_id: null, + status_code: 200, + editor_name: null, + machine_id: null, + user_byok: false, + has_tools: false, + feature: null, + session_id: null, + mode: null, + auto_model: null, + ttfb_ms: null, + ...overrides, + } as MicrodollarUsageContext; + } + + function makeUpstreamResponse(payload: unknown): Response { + return new Response(JSON.stringify(payload), { + status: 200, + headers: { 'content-type': 'application/json' }, + }); + } + + beforeEach(() => { + mockedLogMicrodollarUsage.mockClear(); + mockedLogMicrodollarUsage.mockResolvedValue(null); + }); + + it('zeros both cost_mUsd and cacheDiscount_mUsd for BYOK requests', async () => { + const response = makeUpstreamResponse({ + id: 'edit-byok', + model: 'mercury-edit-2', + usage: { + prompt_tokens: 100_000, + cached_input_tokens: 90_000, + completion_tokens: 0, + total_tokens: 100_000, + }, + choices: [], + }); + + countAndStoreEditUsage(response, makeUsageContext({ user_byok: true }), undefined); + + // Allow the async usage parse + after() callback to settle. + await new Promise(resolve => setImmediate(resolve)); + + expect(mockedLogMicrodollarUsage).toHaveBeenCalledTimes(1); + const [stats] = mockedLogMicrodollarUsage.mock.calls[0]; + expect(stats.cost_mUsd).toBe(0); + expect(stats.cacheDiscount_mUsd).toBe(0); + // The pre-zeroed value is preserved in `market_cost` for reporting. + expect(stats.market_cost).toBe(4_750); + }); + + it('preserves cost_mUsd and cacheDiscount_mUsd for non-BYOK requests', async () => { + const response = makeUpstreamResponse({ + id: 'edit-paid', + model: 'mercury-edit-2', + usage: { + prompt_tokens: 100_000, + cached_input_tokens: 90_000, + completion_tokens: 0, + total_tokens: 100_000, + }, + choices: [], + }); + + countAndStoreEditUsage(response, makeUsageContext({ user_byok: false }), undefined); + + await new Promise(resolve => setImmediate(resolve)); + + expect(mockedLogMicrodollarUsage).toHaveBeenCalledTimes(1); + const [stats] = mockedLogMicrodollarUsage.mock.calls[0]; + expect(stats.cost_mUsd).toBe(4_750); + expect(stats.cacheDiscount_mUsd).toBe(20_250); + expect(stats.market_cost).toBe(4_750); + }); + + it('does not log usage when the upstream body is missing', async () => { + const bodylessResponse = new Response(null, { status: 502 }); + + countAndStoreEditUsage(bodylessResponse, makeUsageContext({ status_code: 502 }), undefined); + + await new Promise(resolve => setImmediate(resolve)); + + expect(mockedLogMicrodollarUsage).not.toHaveBeenCalled(); + }); +}); + describe('parseEmbeddingUsageFromResponse', () => { function makeResponse(overrides: Record = {}) { return JSON.stringify({ diff --git a/apps/web/src/lib/ai-gateway/llm-proxy-helpers.ts b/apps/web/src/lib/ai-gateway/llm-proxy-helpers.ts index 96461b82d8..0085cd24e6 100644 --- a/apps/web/src/lib/ai-gateway/llm-proxy-helpers.ts +++ b/apps/web/src/lib/ai-gateway/llm-proxy-helpers.ts @@ -696,6 +696,147 @@ export function countAndStoreFimUsage( ); } +// ============================================================================ +// Edit-Specific Code +// ============================================================================ + +type EditMessage = { role: string; content: string }; + +export function extractEditPromptInfo(body: { messages: EditMessage[] }): PromptInfo { + const lastUser = [...body.messages].reverse().find(m => m.role === 'user'); + const content = lastUser?.content ?? ''; + return { + system_prompt_prefix: '', // /v1/edit/completions bakes its system prompt in server-side + system_prompt_length: 0, + user_prompt_prefix: content.slice(0, 100), + }; +} + +type EditUsage = FimUsage & { + cached_input_tokens?: number; +}; + +type MercuryEditCompletionResponse = { + id?: string; + model?: string; + usage?: EditUsage; + choices?: Array<{ + index?: number; + message?: { role?: string; content?: string }; + finish_reason?: string | null; + }>; +}; + +function getEditCacheHitTokens(usage: EditUsage): number { + return Math.min(usage.prompt_tokens, Math.max(usage.cached_input_tokens ?? 0, 0)); +} + +function computeEditMicrodollarCost(usage: EditUsage, provider: ProviderId): number { + switch (provider) { + case 'inception': { + // Inception Mercury Edit 2 published rates (per 1M tokens): + // $0.25 input → 0.25 mUSD/token + // $0.025 cached input → 0.025 mUSD/token + // $0.75 output → 0.75 mUSD/token + // Sources: + // https://www.inceptionlabs.ai/models + // https://www.inceptionlabs.ai/blog/introducing-mercury-edit-2 + // Mercury 2 (chat) shares the same per-token rates. + const cacheHitTokens = getEditCacheHitTokens(usage); + const uncachedInputTokens = usage.prompt_tokens - cacheHitTokens; + return Math.round( + uncachedInputTokens * 0.25 + cacheHitTokens * 0.025 + usage.completion_tokens * 0.75 + ); + } + default: + console.error('Unknown provider for edit cost calculation', provider); + return 0; + } +} + +export function parseEditUsageFromResponse( + response: string, + provider: ProviderId, + statusCode: number +): MicrodollarUsageStats { + const json: MercuryEditCompletionResponse = JSON.parse(response); + const usage = json.usage; + const cacheHitTokens = usage ? getEditCacheHitTokens(usage) : 0; + return { + messageId: json.id ?? null, + model: json.model ?? null, + responseContent: json.choices?.[0]?.message?.content || '', + hasError: !json.model || statusCode >= 400, + inference_provider: provider, + inputTokens: usage?.prompt_tokens ?? 0, + outputTokens: usage?.completion_tokens ?? 0, + cacheHitTokens, + cacheWriteTokens: 0, + cost_mUsd: usage ? computeEditMicrodollarCost(usage, provider) : 0, + cacheDiscount_mUsd: + usage && provider === 'inception' ? Math.round(cacheHitTokens * (0.25 - 0.025)) : undefined, + is_byok: null, + upstream_id: null, + finish_reason: null, + latency: null, + moderation_latency: null, + generation_time: null, + streamed: null, + cancelled: null, + status_code: statusCode, + }; +} + +export function countAndStoreEditUsage( + clonedResponse: Response, + usageContext: MicrodollarUsageContext, + requestSpan: Span | undefined +) { + debugSaveProxyResponseStream(clonedResponse, '.log.resp.json'); + + const statusCode = usageContext.status_code ?? 0; + const usageStatsPromise = !clonedResponse.body + ? Promise.resolve(null) + : clonedResponse + .text() + .then(content => parseEditUsageFromResponse(content, usageContext.provider, statusCode)) + .catch(error => { + captureException(error, { + tags: { source: 'edit_usage_processing' }, + extra: { statusCode }, + }); + return null; + }); + + after( + usageStatsPromise.then(usageStats => { + requestSpan?.end(); + if (!usageStats) { + captureMessage('SUSPICIOUS: No edit usage information', { + level: 'error', + tags: { source: 'edit_usage_processing' }, + extra: { usageContext }, + }); + return; + } + + usageStats.market_cost = usageStats.cost_mUsd; + + // Mirror the canonical chat path in `processOpenRouterUsage`: when the + // request is BYOK we don't bill the user, so the cache discount we + // would otherwise have given them must be zeroed too. Otherwise the + // usage row would claim a discount on spend that never happened and + // distort "money saved by caching" reporting. + if (usageContext.user_byok) { + usageStats.cost_mUsd = 0; + usageStats.cacheDiscount_mUsd = 0; + } + + return logMicrodollarUsage(usageStats, usageContext); + }) + ); +} + // ============================================================================ // Embedding-Specific Code // ============================================================================ diff --git a/apps/web/src/lib/proxy-error-types.ts b/apps/web/src/lib/proxy-error-types.ts index 1c9a03f519..5e19e6d720 100644 --- a/apps/web/src/lib/proxy-error-types.ts +++ b/apps/web/src/lib/proxy-error-types.ts @@ -22,6 +22,7 @@ export const proxyErrorTypeSchema = z.enum([ 'paid_model_auth_required', 'promotion_limit_reached', 'unsupported_fim_model', + 'unsupported_edit_model', 'insufficient_credits', 'provider_not_allowed', 'byok_key_required', diff --git a/packages/db/src/schema-types.ts b/packages/db/src/schema-types.ts index 2e615e1e47..983bb76762 100644 --- a/packages/db/src/schema-types.ts +++ b/packages/db/src/schema-types.ts @@ -775,6 +775,7 @@ export const GatewayApiKindSchema = z.enum([ 'chat_completions', 'embeddings', 'fim_completions', + 'edit_completions', 'messages', 'responses', 'audio_transcriptions',