diff --git a/apps/web/src/lib/rewriteModelResponse.test.ts b/apps/web/src/lib/rewriteModelResponse.test.ts index 02ab61de29..935f161e95 100644 --- a/apps/web/src/lib/rewriteModelResponse.test.ts +++ b/apps/web/src/lib/rewriteModelResponse.test.ts @@ -7,6 +7,7 @@ import { type RequestLoggingParams, } from './rewriteModelResponse'; import { isDynamicallyOptedIntoRequestLogging } from '@/lib/ai-gateway/request-logging-opt-ins'; +import { QWEN37_PLUS_MODEL_ID } from '@/lib/ai-gateway/custom-pricing'; import { KILO_ORGANIZATION_ID } from '@/lib/organizations/constants'; jest.mock('next/server', () => ({ @@ -102,7 +103,7 @@ describe.each(rewriters)('%s response read errors', (_name, rewrite) => { ['ResponseAborted', 'upstream_disconnect', 'disconnected'], ['TimeoutError', 'timeout', 'timed out'], ])('returns structured JSON for %s', async (errorName, errorType, messageFragment) => { - const result = await rewrite(failingResponse('application/json', errorName)); + const result = await rewrite(failingResponse('application/json', errorName), true, null, null); expect(result.status).toBe(503); expect(await result.json()).toEqual({ @@ -150,7 +151,12 @@ describe.each(rewriters)('%s response read errors', (_name, rewrite) => { }); test('omits the request id suffix when no vercel request id is available', async () => { - const result = await rewrite(failingResponse('text/event-stream', 'ResponseAborted')); + const result = await rewrite( + failingResponse('text/event-stream', 'ResponseAborted'), + true, + null, + null + ); const events = dataObjects(await readOutputStream(result)) as { error: { message: string; vercel_request_id?: string }; }[]; @@ -178,7 +184,7 @@ describe('rewriteModelResponse_ChatCompletions', () => { }, }); - const result = await rewriteModelResponse_ChatCompletions(upstream); + const result = await rewriteModelResponse_ChatCompletions(upstream, true, null, null); const json = await result.json(); expect(json.model).toBe('upstream-model'); @@ -201,7 +207,7 @@ describe('rewriteModelResponse_ChatCompletions', () => { }, }); - const result = await rewriteModelResponse_ChatCompletions(upstream); + const result = await rewriteModelResponse_ChatCompletions(upstream, true, null, null); const json = await result.json(); expect(json.usage.prompt_tokens_details.cached_tokens).toBe(0); @@ -214,7 +220,7 @@ describe('rewriteModelResponse_ChatCompletions', () => { headers: { 'content-type': 'application/json' }, }); - const result = await rewriteModelResponse_ChatCompletions(upstream); + const result = await rewriteModelResponse_ChatCompletions(upstream, true, null, null); expect(result.status).toBe(502); expect(await result.text()).toBe('not-json{'); @@ -242,7 +248,7 @@ describe('rewriteModelResponse_ChatCompletions', () => { ); try { - const result = await rewriteModelResponse_ChatCompletions(upstream); + const result = await rewriteModelResponse_ChatCompletions(upstream, true, null, null); const reader = result.body?.getReader(); expect(reader).toBeDefined(); await reader?.read(); @@ -267,7 +273,7 @@ describe('rewriteModelResponse_ChatCompletions', () => { 'data: {"id":"gen-chat","model":"upstream-model","choices":[]}\n\n' ); - const result = await rewriteModelResponse_ChatCompletions(upstream); + const result = await rewriteModelResponse_ChatCompletions(upstream, true, null, null); const sse = await readOutputStream(result); const events = dataObjects(sse) as Array<{ error?: { code: number; type: string } }>; @@ -283,7 +289,7 @@ describe('rewriteModelResponse_ChatCompletions', () => { 'data: [DONE]\n\n' ); - const result = await rewriteModelResponse_ChatCompletions(upstream); + const result = await rewriteModelResponse_ChatCompletions(upstream, true, null, null); const sse = await readOutputStream(result); const [chunk] = dataObjects(sse) as Array<{ model: string; @@ -301,7 +307,7 @@ describe('rewriteModelResponse_ChatCompletions', () => { 'data: {"model":"upstream-model","usage":{"cost":1,"is_byok":true,"prompt_tokens":4,"completion_tokens":2,"total_tokens":6,"prompt_tokens_details":{}}}\n\n' ); - const result = await rewriteModelResponse_ChatCompletions(upstream); + const result = await rewriteModelResponse_ChatCompletions(upstream, true, null, null); const sse = await readOutputStream(result); const [chunk] = dataObjects(sse) as Array<{ model: string; @@ -325,7 +331,7 @@ describe('rewriteModelResponse_ChatCompletions', () => { ': openrouter heartbeat\n\n' + 'data: {"model":"upstream-model","choices":[]}\n\n' ); - const result = await rewriteModelResponse_ChatCompletions(upstream); + const result = await rewriteModelResponse_ChatCompletions(upstream, true, null, null); const sse = await readOutputStream(result); expect(sse).toContain(': KILO PROCESSING'); @@ -337,7 +343,7 @@ describe('rewriteModelResponse_ChatCompletions', () => { headers: { 'content-type': 'text/event-stream' }, }); - const result = await rewriteModelResponse_ChatCompletions(upstream); + const result = await rewriteModelResponse_ChatCompletions(upstream, true, null, null); expect(await readOutputStream(result)).toBe(''); }); @@ -354,7 +360,10 @@ describe('rewriteModelResponse_Messages', () => { 'text/event-stream', errorName, 'data: {"type":"message_start","message":{"id":"gen-message","usage":{"input_tokens":1,"output_tokens":0}}}\n\n' - ) + ), + true, + null, + null ); const sse = await readOutputStream(result); @@ -392,7 +401,7 @@ describe('rewriteModelResponse_Messages', () => { }, }); - const result = await rewriteModelResponse_Messages(upstream); + const result = await rewriteModelResponse_Messages(upstream, true, null, null); const json = await result.json(); expect(json.model).toBe('upstream-model'); @@ -408,7 +417,7 @@ describe('rewriteModelResponse_Messages', () => { headers: { 'content-type': 'application/json' }, }); - const result = await rewriteModelResponse_Messages(upstream); + const result = await rewriteModelResponse_Messages(upstream, true, null, null); expect(result.status).toBe(500); expect(await result.text()).toBe('}{'); @@ -421,7 +430,7 @@ describe('rewriteModelResponse_Messages', () => { 'data: [DONE]\n\n' ); - const result = await rewriteModelResponse_Messages(upstream); + const result = await rewriteModelResponse_Messages(upstream, true, null, null); const sse = await readOutputStream(result); const events = dataObjects(sse) as Array<{ type: string; @@ -450,7 +459,7 @@ describe('rewriteModelResponse_Messages', () => { 'data: {"type":"message_delta","usage":{"output_tokens":9},"delta":{}}\n\n' ); - const result = await rewriteModelResponse_Messages(upstream); + const result = await rewriteModelResponse_Messages(upstream, true, null, null); const sse = await readOutputStream(result); expect(dataPayloads(sse)).not.toContain('[DONE]'); @@ -467,7 +476,10 @@ describe('rewriteModelResponse_Responses', () => { 'text/event-stream', errorName, 'data: {"type":"response.created","sequence_number":4,"response":{"id":"gen-response"}}\n\n' - ) + ), + true, + null, + null ); const sse = await readOutputStream(result); @@ -505,7 +517,7 @@ describe('rewriteModelResponse_Responses', () => { }, }); - const result = await rewriteModelResponse_Responses(upstream); + const result = await rewriteModelResponse_Responses(upstream, true, null, null); const json = await result.json(); expect(json.model).toBe('upstream-model'); @@ -521,7 +533,7 @@ describe('rewriteModelResponse_Responses', () => { 'data: [DONE]\n\n' ); - const result = await rewriteModelResponse_Responses(upstream); + const result = await rewriteModelResponse_Responses(upstream, true, null, null); const sse = await readOutputStream(result); const [event] = dataObjects(sse) as Array<{ type: string; @@ -626,6 +638,26 @@ describe('rewriteModelResponse', () => { }); }); + test('strips cost for models with custom pricing', async () => { + const result = await rewriteModelResponse( + jsonResponse({ + model: QWEN37_PLUS_MODEL_ID, + usage: { cost: 0.5, cost_details: { upstream_inference_cost: 0.4 }, is_byok: false }, + }), + QWEN37_PLUS_MODEL_ID, + 'openrouter', + 'chat_completions', + makeLogging() + ); + + // The upstream-reported cost does not reflect the custom pricing, so it + // must be removed just like for free models. + expect(await result.json()).toEqual({ + model: QWEN37_PLUS_MODEL_ID, + usage: {}, + }); + }); + test('processes paid-model responses when request logging is enabled', async () => { mockedOptIn.mockResolvedValueOnce(true); const result = await rewriteModelResponse( @@ -649,7 +681,7 @@ describe('request log capture', () => { const capture = makeCapture(); const body = { model: 'upstream-model' }; - const result = await rewrite(jsonResponse(body), true, capture); + const result = await rewrite(jsonResponse(body), true, capture, null); expect(result.status).toBe(200); expect(capture.setBody).toHaveBeenCalledTimes(1); @@ -662,7 +694,7 @@ describe('request log capture', () => { const sseBody = 'data: {"id":"gen-1","model":"upstream-model","choices":[]}\n\n' + 'data: [DONE]\n\n'; - const result = await rewrite(sseResponse(sseBody), true, capture); + const result = await rewrite(sseResponse(sseBody), true, capture, null); await readOutputStream(result); expect(capture.setBody).toHaveBeenCalledTimes(1); @@ -678,7 +710,8 @@ describe('request log capture', () => { const result = await rewrite( new Response(null, { headers: { 'content-type': 'text/event-stream' } }), true, - capture + capture, + null ); await readOutputStream(result); @@ -696,7 +729,8 @@ describe('request log capture', () => { const result = await rewrite( failingResponse('text/event-stream', 'ResponseAborted', receivedChunks), true, - capture + capture, + null ); await readOutputStream(result); @@ -714,7 +748,8 @@ describe('request log capture', () => { const result = await rewrite( failingResponse('text/event-stream', 'ResponseAborted'), true, - capture + capture, + null ); await readOutputStream(result); @@ -732,7 +767,8 @@ describe('request log capture', () => { const result = await rewrite( failingResponse('application/json', 'TimeoutError'), true, - capture + capture, + null ); expect(result.status).toBe(503); @@ -747,7 +783,7 @@ describe('request log capture', () => { headers: { 'content-type': 'text/event-stream' }, }); - const result = await rewriteModelResponse_ChatCompletions(upstream, true, capture); + const result = await rewriteModelResponse_ChatCompletions(upstream, true, capture, null); const reader = result.body?.getReader(); await reader?.cancel(); diff --git a/apps/web/src/lib/rewriteModelResponse.ts b/apps/web/src/lib/rewriteModelResponse.ts index b67dad8597..724dd2ff60 100644 --- a/apps/web/src/lib/rewriteModelResponse.ts +++ b/apps/web/src/lib/rewriteModelResponse.ts @@ -1,5 +1,6 @@ import { api_request_log, type User } from '@kilocode/db/schema'; import { isKiloExclusiveFreeModel } from '@/lib/ai-gateway/models'; +import { getCustomPricing } from '@/lib/ai-gateway/custom-pricing'; import { detectToolCallArgumentErrors } from '@/lib/ai-gateway/api-request-log-errors'; import type { GatewayRequest } from '@/lib/ai-gateway/providers/openrouter/types'; import type { ProviderId } from '@/lib/ai-gateway/providers/types'; @@ -218,7 +219,7 @@ async function readResponseText( response: Response, headers: Headers, vercelRequestId: string | null | undefined, - capture?: RequestLogCapture | null + capture: RequestLogCapture | null ): Promise<{ text: string } | { error: unknown; errorResponse: NextResponse }> { try { return { text: await response.text() }; @@ -256,7 +257,7 @@ async function rewriteSseStream( serializeError: (error: ResponseReadError) => string, onFinally: () => void, vercelRequestId: string | null | undefined, - capture?: RequestLogCapture | null + capture: RequestLogCapture | null ) { const decoder = new TextDecoder(); // Accumulate the raw upstream text for request logging while the stream is @@ -323,9 +324,9 @@ function rewriteUsage(usage: OpenRouterUsage, removeCost: boolean) { export async function rewriteModelResponse_ChatCompletions( response: Response, - removeCost = true, - capture?: RequestLogCapture | null, - vercelRequestId?: string | null + removeCost: boolean, + capture: RequestLogCapture | null, + vercelRequestId: string | null ) { const headers = getOutputHeaders(response); @@ -478,9 +479,9 @@ function rewriteMessagesUsage(usage: MessagesApiUsage, removeCost: boolean) { export async function rewriteModelResponse_Messages( response: Response, - removeCost = true, - capture?: RequestLogCapture | null, - vercelRequestId?: string | null + removeCost: boolean, + capture: RequestLogCapture | null, + vercelRequestId: string | null ) { const headers = getOutputHeaders(response); @@ -614,9 +615,9 @@ type ResponsesApiEvent = { export async function rewriteModelResponse_Responses( response: Response, - removeCost = true, - capture?: RequestLogCapture | null, - vercelRequestId?: string | null + removeCost: boolean, + capture: RequestLogCapture | null, + vercelRequestId: string | null ) { const headers = getOutputHeaders(response); @@ -742,34 +743,25 @@ export async function rewriteModelResponse( logging: RequestLoggingParams ): Promise { const capture = await createRequestLogCapture(response, model, providerId, logging); - const isFreeModelRequiringCostRemoval = - (providerId === 'openrouter' || providerId === 'vercel') && isKiloExclusiveFreeModel(model); + const requiresCostRemoval = + (providerId === 'openrouter' || providerId === 'vercel') && + (isKiloExclusiveFreeModel(model) || getCustomPricing(model) !== undefined); console.debug('[rewriteModelResponse] rewriting response for %s', model); const { vercel_request_id: vercelRequestId } = logging; if (kind === 'chat_completions') { return rewriteModelResponse_ChatCompletions( response, - isFreeModelRequiringCostRemoval, + requiresCostRemoval, capture, vercelRequestId ); } if (kind === 'responses') { - return rewriteModelResponse_Responses( - response, - isFreeModelRequiringCostRemoval, - capture, - vercelRequestId - ); + return rewriteModelResponse_Responses(response, requiresCostRemoval, capture, vercelRequestId); } if (kind === 'messages') { - return rewriteModelResponse_Messages( - response, - isFreeModelRequiringCostRemoval, - capture, - vercelRequestId - ); + return rewriteModelResponse_Messages(response, requiresCostRemoval, capture, vercelRequestId); } const error = new Error(`implementation error: unrecognized API kind ${kind}`);