diff --git a/packages/web-shared/src/lib/ai-gateway/free-model-rate-limit.test.ts b/packages/web-shared/src/lib/ai-gateway/free-model-rate-limit.test.ts index a45ccaacca0..abb8c6cd7f4 100644 --- a/packages/web-shared/src/lib/ai-gateway/free-model-rate-limit.test.ts +++ b/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' }) ); }); }); diff --git a/packages/web-shared/src/lib/ai-gateway/free-model-rate-limit.ts b/packages/web-shared/src/lib/ai-gateway/free-model-rate-limit.ts index 65e3b49d7de..b70b38b40ae 100644 --- a/packages/web-shared/src/lib/ai-gateway/free-model-rate-limit.ts +++ b/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 { +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 { return isFirewallRateLimited(request, { - id: FREE_MODEL_RATE_LIMIT_ID, - key, - sentrySource: 'free_model_inference', + ...FREE_MODEL_RATE_LIMIT_RULES[subject.kind], + key: freeModelRateLimitKey(subject), }); } diff --git a/packages/web-shared/src/lib/ai-gateway/handlers/llm-proxy.test.ts b/packages/web-shared/src/lib/ai-gateway/handlers/llm-proxy.test.ts index f63f028ec83..501d927e839 100644 --- a/packages/web-shared/src/lib/ai-gateway/handlers/llm-proxy.test.ts +++ b/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>); +} + 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>); + 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); }); diff --git a/packages/web-shared/src/lib/ai-gateway/handlers/llm-proxy.ts b/packages/web-shared/src/lib/ai-gateway/handlers/llm-proxy.ts index 5b9d90f8860..a95f7714bcc 100644 --- a/packages/web-shared/src/lib/ai-gateway/handlers/llm-proxy.ts +++ b/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,7 +82,11 @@ 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, @@ -93,7 +94,6 @@ import { } 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 | 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 +): Promise | 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(); } } diff --git a/packages/web-shared/src/lib/ai-gateway/llm-proxy-helpers.ts b/packages/web-shared/src/lib/ai-gateway/llm-proxy-helpers.ts index ca135e3f2c5..18f44748a8b 100644 --- a/packages/web-shared/src/lib/ai-gateway/llm-proxy-helpers.ts +++ b/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 }, diff --git a/packages/web-shared/src/lib/bouncer/inference.ts b/packages/web-shared/src/lib/bouncer/inference.ts index 0f308fb47b8..e3932326f94 100644 --- a/packages/web-shared/src/lib/bouncer/inference.ts +++ b/packages/web-shared/src/lib/bouncer/inference.ts @@ -6,7 +6,7 @@ import { after, NextResponse, type NextRequest } from 'next/server'; import { isCloudflareIP } from '@kilocode/web-shared/lib/cloudflare-ip'; import { - isUserRateLimitedFeature, + isServerSideProductFeature, type FeatureValue, } from '@kilocode/web-shared/lib/feature-detection'; import { @@ -73,7 +73,7 @@ export function payerSharingIp( ip: string | undefined, feature: FeatureValue | null ): string | undefined { - if (ip !== undefined && isUserRateLimitedFeature(feature) && isCloudflareIP(ip)) { + if (ip !== undefined && isServerSideProductFeature(feature) && isCloudflareIP(ip)) { return undefined; } return ip; diff --git a/packages/web-shared/src/lib/feature-detection.test.ts b/packages/web-shared/src/lib/feature-detection.test.ts index 9bd8f986fc3..24988ab0d58 100644 --- a/packages/web-shared/src/lib/feature-detection.test.ts +++ b/packages/web-shared/src/lib/feature-detection.test.ts @@ -1,5 +1,5 @@ import { describe, test, expect } from '@jest/globals'; -import { isUserRateLimitedFeature, validateFeatureHeader } from './feature-detection'; +import { isServerSideProductFeature, validateFeatureHeader } from './feature-detection'; describe('validateFeatureHeader', () => { test('returns null for null input', () => { @@ -36,22 +36,22 @@ describe('validateFeatureHeader', () => { }); }); -describe('isUserRateLimitedFeature', () => { +describe('isServerSideProductFeature', () => { test('returns true for server-side products', () => { - expect(isUserRateLimitedFeature('cloud-agent')).toBe(true); - expect(isUserRateLimitedFeature('code-review')).toBe(true); - expect(isUserRateLimitedFeature('app-builder')).toBe(true); - expect(isUserRateLimitedFeature('gastown')).toBe(true); + expect(isServerSideProductFeature('cloud-agent')).toBe(true); + expect(isServerSideProductFeature('code-review')).toBe(true); + expect(isServerSideProductFeature('app-builder')).toBe(true); + expect(isServerSideProductFeature('gastown')).toBe(true); }); test('returns false for client-side products', () => { - expect(isUserRateLimitedFeature('vscode-extension')).toBe(false); - expect(isUserRateLimitedFeature('jetbrains-extension')).toBe(false); - expect(isUserRateLimitedFeature('cli')).toBe(false); - expect(isUserRateLimitedFeature('direct-gateway')).toBe(false); + expect(isServerSideProductFeature('vscode-extension')).toBe(false); + expect(isServerSideProductFeature('jetbrains-extension')).toBe(false); + expect(isServerSideProductFeature('cli')).toBe(false); + expect(isServerSideProductFeature('direct-gateway')).toBe(false); }); test('returns false for null', () => { - expect(isUserRateLimitedFeature(null)).toBe(false); + expect(isServerSideProductFeature(null)).toBe(false); }); }); diff --git a/packages/web-shared/src/lib/feature-detection.ts b/packages/web-shared/src/lib/feature-detection.ts index 1bd84758b89..4ffc71d1267 100644 --- a/packages/web-shared/src/lib/feature-detection.ts +++ b/packages/web-shared/src/lib/feature-detection.ts @@ -61,17 +61,17 @@ export function validateFeatureHeader(headerValue: string | null): FeatureValue } /** - * Server-side products that rate-limit free models per user instead of per IP. - * These products share infrastructure IPs, so IP-based limits would be too restrictive. + * Server-side products that reach the gateway from shared Kilo infrastructure + * (Cloudflare Workers), so their source IP says nothing about the user. */ -const USER_RATE_LIMITED_FEATURES: ReadonlySet = new Set([ +const SERVER_SIDE_PRODUCT_FEATURES: ReadonlySet = new Set([ 'cloud-agent', 'code-review', 'app-builder', 'gastown', ]); -export function isUserRateLimitedFeature(feature: FeatureValue | null): boolean { +export function isServerSideProductFeature(feature: FeatureValue | null): boolean { if (!feature) return false; - return USER_RATE_LIMITED_FEATURES.has(feature); + return SERVER_SIDE_PRODUCT_FEATURES.has(feature); }