[AISAGE-729] fix(ai-gateway): rate-limit signed-in free model use per account

packages/web-shared/src/lib/ai-gateway/free-model-rate-limit.test.ts

@@ -10,15 +10,31 @@ import { isFreeModelRateLimited } from './free-model-rate-limit';
 const mockCheckRateLimit = jest.mocked(checkRateLimit);
 
 describe('isFreeModelRateLimited', () => {
-  it('counts against the free-model firewall rule', async () => {
+  it('counts signed-in callers against the per-account free-model rule', async () => {
     mockCheckRateLimit.mockResolvedValue({ rateLimited: true });
     const request = new NextRequest('https://gateway.example.com/api/gateway/chat/completions');
 
-    await expect(isFreeModelRateLimited(request, 'user:user-123')).resolves.toBe(true);
+    await expect(
+      isFreeModelRateLimited(request, { kind: 'account', kiloUserId: 'user-123' })
+    ).resolves.toBe(true);
+
+    expect(mockCheckRateLimit).toHaveBeenCalledWith(
+      'free-model-inference-user',
+      expect.objectContaining({ rateLimitKey: 'free-model-inference-user:user:user-123' })
+    );
+  });
+
+  it('counts anonymous callers against the per-IP free-model rule', async () => {
+    mockCheckRateLimit.mockResolvedValue({ rateLimited: false });
+    const request = new NextRequest('https://gateway.example.com/api/gateway/chat/completions');
+
+    await expect(
+      isFreeModelRateLimited(request, { kind: 'ip', ipAddress: '203.0.113.7' })
+    ).resolves.toBe(false);
 
     expect(mockCheckRateLimit).toHaveBeenCalledWith(
       'free-model-inference',
-      expect.objectContaining({ rateLimitKey: 'free-model-inference:user:user-123' })
+      expect.objectContaining({ rateLimitKey: 'free-model-inference:ip:203.0.113.7' })
     );
   });
 });

packages/web-shared/src/lib/ai-gateway/free-model-rate-limit.ts

@@ -3,16 +3,32 @@ import type { NextRequest } from 'next/server';
 import { isFirewallRateLimited } from '@kilocode/web-shared/lib/ai-gateway/firewall-rate-limit';
 
 /**
- * Vercel Firewall rate limit shared by all free models, custom LLMs excluded.
- * It exists to stop obvious abuse, not to meter normal use. The window and
- * limit live in the project's firewall configuration.
+ * Vercel Firewall rate limits shared by all free models, custom LLMs excluded.
+ * They exist to stop abuse, not to meter normal use. Signed-in callers count
+ * per account and anonymous callers per IP, each against its own rule so the
+ * two limits can differ. Abusers rotate IPs far more cheaply than accounts, and
+ * people behind one NAT should not share an allowance. The windows and limits
+ * live in the project's firewall configuration.
  */
-const FREE_MODEL_RATE_LIMIT_ID = 'free-model-inference';
+export type FreeModelRateLimitSubject =
+  | { kind: 'account'; kiloUserId: string }
+  | { kind: 'ip'; ipAddress: string };
 
-export function isFreeModelRateLimited(request: NextRequest, key: string): Promise<boolean> {
+const FREE_MODEL_RATE_LIMIT_RULES = {
+  account: { id: 'free-model-inference-user', sentrySource: 'free_model_inference_user' },
+  ip: { id: 'free-model-inference', sentrySource: 'free_model_inference' },
+} as const;
+
+export function freeModelRateLimitKey(subject: FreeModelRateLimitSubject): string {
+  return subject.kind === 'account' ? `user:${subject.kiloUserId}` : `ip:${subject.ipAddress}`;
+}
+
+export function isFreeModelRateLimited(
+  request: NextRequest,
+  subject: FreeModelRateLimitSubject
+): Promise<boolean> {
   return isFirewallRateLimited(request, {
-    id: FREE_MODEL_RATE_LIMIT_ID,
-    key,
-    sentrySource: 'free_model_inference',
+    ...FREE_MODEL_RATE_LIMIT_RULES[subject.kind],
+    key: freeModelRateLimitKey(subject),
   });
 }

packages/web-shared/src/lib/ai-gateway/handlers/llm-proxy.test.ts

@@ -288,6 +288,14 @@ function setUserAuth() {
   });
 }
 
+function setAnonymousAuth() {
+  mockedGetUserFromAuth.mockResolvedValue({
+    user: null,
+    authFailedResponse: new Response('unauthorized', { status: 401 }),
+    organizationId: undefined,
+  } as unknown as Awaited<ReturnType<typeof getUserFromAuth>>);
+}
+
 function upstreamJsonResponse(body: unknown, status = 200) {
   return new Response(JSON.stringify(body), {
     status,
@@ -1303,52 +1311,95 @@ describe('POST /api/openrouter/v1/chat/completions request handling', () => {
   });
 
   it.each([gemma_4_26b_a4b_it_free_model.public_id, 'test/free-model:free', 'openrouter/free'])(
-    'applies free-model rate limiting per IP to the free model %s',
+    'applies free-model rate limiting per IP to anonymous callers of the free model %s',
     async modelId => {
+      setAnonymousAuth();
       mockedIsFreeModelRateLimited.mockResolvedValue(true);
 
       const { handleLlmProxyRequest } = await import('./llm-proxy');
       const response = await handleLlmProxyRequest(makeRequest(makeBody(modelId)) as never);
 
       expect(response.status).toBe(429);
+      expect(response.headers.get('x-should-retry')).toBe('false');
       expect(await response.json()).toMatchObject({
+        error: { code: 'PROMOTION_MODEL_LIMIT_REACHED' },
         error_type: 'rate_limit_exceeded',
-        message: 'Model usage limit reached. Please try again later.',
       });
-      expect(mockedIsFreeModelRateLimited).toHaveBeenCalledWith(expect.anything(), 'ip:127.0.0.1');
+      expect(mockedIsFreeModelRateLimited).toHaveBeenCalledWith(expect.anything(), {
+        kind: 'ip',
+        ipAddress: '127.0.0.1',
+      });
       expect(mockedUpstreamRequest).not.toHaveBeenCalled();
     }
   );
 
-  it('applies free-model rate limiting per user to server-side products on Cloudflare IPs', async () => {
+  it.each([
+    ['a client on its own IP', {}],
+    [
+      'a server-side product on a Cloudflare IP',
+      { 'x-forwarded-for': '104.16.0.1', 'x-kilocode-feature': 'cloud-agent' },
+    ],
+  ])('applies free-model rate limiting per account to %s', async (_caller, headers) => {
     setUserAuth();
     mockedIsFreeModelRateLimited.mockResolvedValue(true);
 
     const { handleLlmProxyRequest } = await import('./llm-proxy');
     const response = await handleLlmProxyRequest(
       makeRequest(makeBody(gemma_4_26b_a4b_it_free_model.public_id), {
-        'x-forwarded-for': '104.16.0.1',
-        'x-kilocode-feature': 'cloud-agent',
+        authorization: `Bearer ${signedToken(KILO_GATEWAY_AUDIENCE)}`,
+        ...headers,
       }) as never
     );
 
     expect(response.status).toBe(429);
-    expect(mockedIsFreeModelRateLimited).toHaveBeenCalledWith(expect.anything(), 'user:user-123');
+    expect(response.headers.get('x-should-retry')).toBe('false');
+    const body = await response.json();
+    expect(body).toMatchObject({
+      error_type: 'rate_limit_exceeded',
+      message:
+        'Free model usage limit reached. Please try again later, or switch to a paid model or your own API key.',
+    });
+    expect(JSON.stringify(body)).not.toContain('PROMOTION_MODEL_LIMIT_REACHED');
+    expect(mockedIsFreeModelRateLimited).toHaveBeenCalledWith(expect.anything(), {
+      kind: 'account',
+      kiloUserId: 'user-123',
+    });
     expect(mockedUpstreamRequest).not.toHaveBeenCalled();
   });
 
-  it('requires authentication for server-side products on Cloudflare IPs', async () => {
+  it('applies free-model rate limiting per account to session-cookie callers', async () => {
+    mockedIsFreeModelRateLimited.mockResolvedValue(true);
+
+    const { handleLlmProxyRequest } = await import('./llm-proxy');
+    const response = await handleLlmProxyRequest(
+      makeRequest(makeBody(gemma_4_26b_a4b_it_free_model.public_id)) as never
+    );
+
+    expect(response.status).toBe(429);
+    expect(JSON.stringify(await response.json())).not.toContain('PROMOTION_MODEL_LIMIT_REACHED');
+    expect(mockedIsFreeModelRateLimited).toHaveBeenCalledWith(expect.anything(), {
+      kind: 'account',
+      kiloUserId: 'user-123',
+    });
+    expect(mockedUpstreamRequest).not.toHaveBeenCalled();
+  });
+
+  it.each([
+    ['a malformed token', () => 'not-a-jwt'],
+    ['a signature-valid but revoked token', () => signedToken(KILO_GATEWAY_AUDIENCE)],
+  ])('answers %s for a free model with the invalid-token response', async (_case, token) => {
     mockedGetUserFromAuth.mockResolvedValue({
       user: null,
       authFailedResponse: new Response('unauthorized', { status: 401 }),
+      credentialsRejected: true,
       organizationId: undefined,
     } as unknown as Awaited<ReturnType<typeof getUserFromAuth>>);
+    mockedIsFreeModelRateLimited.mockResolvedValue(true);
 
     const { handleLlmProxyRequest } = await import('./llm-proxy');
     const response = await handleLlmProxyRequest(
       makeRequest(makeBody(gemma_4_26b_a4b_it_free_model.public_id), {
-        'x-forwarded-for': '104.16.0.1',
-        'x-kilocode-feature': 'cloud-agent',
+        authorization: `Bearer ${token()}`,
       }) as never
     );
 
@@ -1358,6 +1409,26 @@ describe('POST /api/openrouter/v1/chat/completions request handling', () => {
     expect(mockedUpstreamRequest).not.toHaveBeenCalled();
   });
 
+  it('counts anonymous server-side product requests on Cloudflare IPs per IP', async () => {
+    setAnonymousAuth();
+    mockedIsFreeModelRateLimited.mockResolvedValue(true);
+
+    const { handleLlmProxyRequest } = await import('./llm-proxy');
+    const response = await handleLlmProxyRequest(
+      makeRequest(makeBody(gemma_4_26b_a4b_it_free_model.public_id), {
+        'x-forwarded-for': '104.16.0.1',
+        'x-kilocode-feature': 'cloud-agent',
+      }) as never
+    );
+
+    expect(response.status).toBe(429);
+    expect(mockedIsFreeModelRateLimited).toHaveBeenCalledWith(expect.anything(), {
+      kind: 'ip',
+      ipAddress: '104.16.0.1',
+    });
+    expect(mockedUpstreamRequest).not.toHaveBeenCalled();
+  });
+
   it('serves free models under the free-model rate limit', async () => {
     mockedGetProvider.mockResolvedValueOnce({
       kind: 'provider',
@@ -1379,7 +1450,10 @@ describe('POST /api/openrouter/v1/chat/completions request handling', () => {
     );
 
     expect(response.status).toBe(200);
-    expect(mockedIsFreeModelRateLimited).toHaveBeenCalledWith(expect.anything(), 'ip:127.0.0.1');
+    expect(mockedIsFreeModelRateLimited).toHaveBeenCalledWith(expect.anything(), {
+      kind: 'ip',
+      ipAddress: '127.0.0.1',
+    });
     expect(mockedUpstreamRequest).toHaveBeenCalledTimes(1);
   });
 

packages/web-shared/src/lib/ai-gateway/handlers/llm-proxy.ts

@@ -5,12 +5,7 @@ import { type NextRequest } from 'next/server';
 import { toMicrodollars } from '@kilocode/web-shared/lib/microdollars';
 import { extractPromptInfo } from '@kilocode/web-shared/lib/ai-gateway/extractPromptInfo';
 import { determineFallbackFeature } from '@kilocode/web-shared/lib/ai-gateway/determineFallbackFeature';
-import {
-  validateFeatureHeader,
-  FEATURE_HEADER,
-  isUserRateLimitedFeature,
-  type FeatureValue,
-} from '@kilocode/web-shared/lib/feature-detection';
+import { validateFeatureHeader, FEATURE_HEADER } from '@kilocode/web-shared/lib/feature-detection';
 import type {
   OpenRouterChatCompletionRequest,
   GatewayResponsesRequest,
@@ -61,6 +56,8 @@ import {
   autoRoutingUnavailableResponse,
   invalidAutoDecisionResponse,
   paidModelAuthRequiredResponse,
+  anonymousFreeModelLimitResponse,
+  accountFreeModelLimitResponse,
   temporarilyUnavailableResponse,
   creditsBlockedResponse,
   unavailableModelResponse,
@@ -85,15 +82,18 @@ import {
   isAnonymousContext,
   type AnonymousUserContext,
 } from '@kilocode/web-shared/lib/anonymous';
-import { isFreeModelRateLimited } from '@kilocode/web-shared/lib/ai-gateway/free-model-rate-limit';
+import {
+  freeModelRateLimitKey,
+  isFreeModelRateLimited,
+  type FreeModelRateLimitSubject,
+} from '@kilocode/web-shared/lib/ai-gateway/free-model-rate-limit';
 import { emitApiMetricsForResponse } from '@kilocode/web-shared/lib/ai-gateway/o11y/api-metrics.server';
 import {
   gatewayRateLimitKey,
   isGatewayAccountRateLimited,
 } from '@kilocode/web-shared/lib/ai-gateway/gateway-account-rate-limit';
 import { normalizeModelId } from '@kilocode/web-shared/lib/ai-gateway/model-utils';
 import { isUnavailableModel } from '@kilocode/web-shared/lib/ai-gateway/unavailable-models';
-import { isCloudflareIP } from '@kilocode/web-shared/lib/cloudflare-ip';
 import {
   isKiloAutoModel,
   KILO_AUTO_BALANCED_MODEL,
@@ -194,25 +194,24 @@ function validatePath(
   return { errorResponse: invalidPathResponse() };
 }
 
-async function resolveFreeModelRateLimitKey(
-  feature: FeatureValue | null,
+/**
+ * Signed-in callers count per account, anonymous callers per IP. Uses the
+ * resolved auth rather than the token's claims, so a revoked token cannot spend
+ * its account's allowance: a rejected credential gets the invalid-token
+ * response instead. Auth has already resolved by the time free-model requests
+ * get here, because balance and settings wait on it.
+ */
+async function resolveFreeModelRateLimitSubject(
   ipAddress: string,
-  authPromise: Promise<{ user: { id: string } | null }>
-): Promise<NextResponseType<unknown> | string> {
-  if (isUserRateLimitedFeature(feature) && isCloudflareIP(ipAddress)) {
-    const { user } = await authPromise;
-    if (!user) {
-      return NextResponse.json(
-        {
-          error: 'Authentication required for this feature',
-          error_type: ProxyErrorType.authentication_required,
-        },
-        { status: 401 }
-      );
-    }
-    return `user:${user.id}`;
+  authPromise: ReturnType<typeof getUserFromAuth>
+): Promise<NextResponseType<unknown> | FreeModelRateLimitSubject> {
+  const { user, authFailedResponse, credentialsRejected } = await authPromise;
+  if (authFailedResponse && credentialsRejected) {
+    return invalidTokenResponse();
   }
-  return `ip:${ipAddress}`;
+  return user && !authFailedResponse
+    ? { kind: 'account', kiloUserId: user.id }
+    : { kind: 'ip', ipAddress };
 }
 
 export async function handleLlmProxyRequest(
@@ -601,26 +600,18 @@ export async function handleLlmProxyRequest(
   // the report-only decide call and the usage event, which share this request.
   const clientIp = bareIpLiteral(ipAddress);
 
-  // For free models: server-side products (cloud-agent, code-review, app-builder,
-  // gastown) rate-limit per user when the request comes from Cloudflare IPs (Kilo
-  // infrastructure). All other products rate-limit per IP (fast pre-auth path).
   // Custom LLMs are exempt, even when their id matches a free-model pattern.
   if (isFreeRequestedModel && !(await isCustomLlm(effectiveModelIdLowerCased))) {
-    const rateLimitKey = await resolveFreeModelRateLimitKey(feature, ipAddress, authPromise);
-    if (rateLimitKey instanceof NextResponse) return rateLimitKey;
+    const rateLimitSubject = await resolveFreeModelRateLimitSubject(ipAddress, authPromise);
+    if (rateLimitSubject instanceof NextResponse) return rateLimitSubject;
 
-    if (await isFreeModelRateLimited(request, rateLimitKey)) {
+    if (await isFreeModelRateLimited(request, rateLimitSubject)) {
       console.warn(
-        `Model rate limit exceeded, ${rateLimitKey}, model: ${effectiveModelIdLowerCased}`
-      );
-      return NextResponse.json(
-        {
-          error: 'Rate limit exceeded',
-          error_type: ProxyErrorType.rate_limit_exceeded,
-          message: 'Model usage limit reached. Please try again later.',
-        },
-        { status: 429 }
+        `Model rate limit exceeded, ${freeModelRateLimitKey(rateLimitSubject)}, model: ${effectiveModelIdLowerCased}, feature: ${feature ?? 'none'}`
       );
+      return rateLimitSubject.kind === 'account'
+        ? accountFreeModelLimitResponse()
+        : anonymousFreeModelLimitResponse();
     }
   }
 

packages/web-shared/src/lib/ai-gateway/llm-proxy-helpers.ts

@@ -61,6 +61,8 @@ import type { AutoDecisionRejection } from '@kilocode/web-shared/lib/ai-gateway/
 const fimSuffixFakeSysPrompMarkers = { begin: '[FIM_SUFFIX:', end: ']' } as const;
 
 const PAID_MODEL_AUTH_REQUIRED = 'PAID_MODEL_AUTH_REQUIRED';
+// Kilo clients show a sign-up prompt for this code and do not retry the request.
+const PROMOTION_MODEL_LIMIT_REACHED = 'PROMOTION_MODEL_LIMIT_REACHED';
 
 export function invalidPathResponse() {
   return NextResponse.json(
@@ -566,6 +568,34 @@ export function paidModelAuthRequiredResponse() {
   );
 }
 
+// The free-model limits use hour-long windows, so an immediate client retry
+// only spends more of the caller's allowance. Kilo CLI v2 and the OpenAI SDKs
+// honor this header.
+const DO_NOT_RETRY_HEADERS = { 'x-should-retry': 'false' };
+
+export function anonymousFreeModelLimitResponse() {
+  const message =
+    'Sign up for free to continue and explore 500 other models. ' +
+    'Takes 2 minutes, no credit card required. Or come back later.';
+  return NextResponse.json(
+    {
+      error: { code: PROMOTION_MODEL_LIMIT_REACHED, message },
+      error_type: ProxyErrorType.rate_limit_exceeded,
+      message,
+    },
+    { status: 429, headers: DO_NOT_RETRY_HEADERS }
+  );
+}
+
+export function accountFreeModelLimitResponse() {
+  const message =
+    'Free model usage limit reached. Please try again later, or switch to a paid model or your own API key.';
+  return NextResponse.json(
+    { error: 'Rate limit exceeded', error_type: ProxyErrorType.rate_limit_exceeded, message },
+    { status: 429, headers: DO_NOT_RETRY_HEADERS }
+  );
+}
+
 export function customRouterErrorResponse(code: CustomRouterErrorCode, message: string) {
   return NextResponse.json(
     { error: { code, message }, error_type: ProxyErrorType.custom_router, message },