refactor(gateway): centralize free model response rewrites (#3641)

* refactor(gateway): centralize free model response rewrites

* style(gateway): format response rewrite helper

---------

Co-authored-by: kiloconnect[bot] <240665456+kiloconnect[bot]@users.noreply.github.com>

Author: kilo-code-bot[bot]

apps/web/src/app/api/openrouter/[...path]/route.ts

@@ -28,7 +28,6 @@ import {
   isDeadFreeModel,
   isExcludedForFeature,
   isKiloExclusiveFreeModel,
-  isKiloStealthModel,
   isKiloExclusiveModelRequiringDataCollection,
 } from '@/lib/ai-gateway/models';
 import { isFreeModel } from '@/lib/ai-gateway/is-free-model';
@@ -57,11 +56,7 @@ import {
 import { ProxyErrorType } from '@/lib/proxy-error-types';
 import { getBalanceAndOrgSettings } from '@/lib/organizations/organization-usage';
 import { isDataCollectionExplicitlyDisallowed } from '@/lib/ai-gateway/providers/openrouter/types';
-import {
-  rewriteFreeModelResponse_ChatCompletions,
-  rewriteFreeModelResponse_Messages,
-  rewriteFreeModelResponse_Responses,
-} from '@/lib/rewriteModelResponse';
+import { rewriteFreeModelResponse } from '@/lib/rewriteModelResponse';
 import {
   createAnonymousContext,
   isAnonymousContext,
@@ -799,24 +794,14 @@ export async function POST(request: NextRequest): Promise<NextResponseType<unkno
     }
   }
 
-  const isFreeModelRequiringCostRemoval =
-    (effectiveProviderContext.provider.id === 'openrouter' ||
-      effectiveProviderContext.provider.id === 'vercel') &&
-    isKiloExclusiveFreeModel(effectiveModelIdLowerCased);
-  const isStealthModelRequiringNameRemoval =
-    effectiveProviderContext.provider.id !== 'martian' &&
-    isKiloStealthModel(effectiveModelIdLowerCased);
-
-  if (isFreeModelRequiringCostRemoval || isStealthModelRequiringNameRemoval) {
-    if (requestBodyParsed.kind === 'chat_completions') {
-      return rewriteFreeModelResponse_ChatCompletions(response, effectiveModelIdLowerCased);
-    }
-    if (requestBodyParsed.kind === 'responses') {
-      return rewriteFreeModelResponse_Responses(response, effectiveModelIdLowerCased);
-    }
-    if (requestBodyParsed.kind === 'messages') {
-      return rewriteFreeModelResponse_Messages(response, effectiveModelIdLowerCased);
-    }
+  const rewrittenResponse = await rewriteFreeModelResponse(
+    response,
+    effectiveModelIdLowerCased,
+    effectiveProviderContext.provider.id,
+    requestBodyParsed.kind
+  );
+  if (rewrittenResponse) {
+    return rewrittenResponse;
   }
 
   return wrapInSafeNextResponse(response);

apps/web/src/lib/rewriteModelResponse.ts

@@ -1,3 +1,6 @@
+import { isKiloExclusiveFreeModel, isKiloStealthModel } from '@/lib/ai-gateway/models';
+import type { GatewayRequest } from '@/lib/ai-gateway/providers/openrouter/types';
+import type { ProviderId } from '@/lib/ai-gateway/providers/types';
 import { getOutputHeaders } from '@/lib/ai-gateway/llm-proxy-helpers';
 import type { ChatCompletionChunk, OpenRouterUsage } from '@/lib/ai-gateway/processUsage.types';
 import type { EventSourceMessage } from 'eventsource-parser';
@@ -338,3 +341,30 @@ export async function rewriteFreeModelResponse_Responses(response: Response, mod
     headers,
   });
 }
+
+export async function rewriteFreeModelResponse(
+  response: Response,
+  model: string,
+  providerId: ProviderId,
+  kind: GatewayRequest['kind']
+): Promise<NextResponse | null> {
+  const isFreeModelRequiringCostRemoval =
+    (providerId === 'openrouter' || providerId === 'vercel') && isKiloExclusiveFreeModel(model);
+  const isStealthModelRequiringNameRemoval = providerId !== 'martian' && isKiloStealthModel(model);
+
+  if (!isFreeModelRequiringCostRemoval && !isStealthModelRequiringNameRemoval) {
+    return null;
+  }
+
+  if (kind === 'chat_completions') {
+    return rewriteFreeModelResponse_ChatCompletions(response, model);
+  }
+  if (kind === 'responses') {
+    return rewriteFreeModelResponse_Responses(response, model);
+  }
+  if (kind === 'messages') {
+    return rewriteFreeModelResponse_Messages(response, model);
+  }
+
+  return null;
+}