From 0a46096c41e7a6f56925643413aaccae8fc554bb Mon Sep 17 00:00:00 2001 From: Daniel Salazar Date: Thu, 9 Jul 2026 19:35:03 -0400 Subject: [PATCH] fix: ai rate limits (#3371) --- AGENTS.md | 2 +- .../puterai/PuterAIController.test.ts | 57 ++++++++++++++++++- .../controllers/puterai/PuterAIController.ts | 24 ++++++++ 3 files changed, 81 insertions(+), 2 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 7a8f33b27..cc6a34417 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -53,7 +53,7 @@ The `extension` global ([src/backend/extensions.ts](src/backend/extensions.ts)) - **Modules:** We transpile and build as needed — write ES modules, not CommonJS. - **TypeScript preferred for new files.** Existing JS is fine; convert opportunistically when you're already touching a file. -- **Reuse types before inventing them.** Search for an existing type first; extend it if close. Only define a new type when nothing fits. +- **Reuse types and code before adding duplicates.** Before defining a type, search for an existing one; extend it if close. Before writing a helper or repeating logic, search for an existing utility, helper, or implementation and reuse or extend it. Only add a new type or helper when nothing suitable exists. - **Make new types findable.** Co-locate them with the layer/module that owns them, export from the obvious entry point, and use a descriptive `PascalCase` name. Don't hide types in random files where future readers won't grep them. - **Naming:** `camelCase` for variables/functions, `PascalCase` for classes and for files containing a class (`AuthService.ts`, `KVStoreDriver.ts`). diff --git a/src/backend/controllers/puterai/PuterAIController.test.ts b/src/backend/controllers/puterai/PuterAIController.test.ts index e8ed260fd..ef2332d10 100644 --- a/src/backend/controllers/puterai/PuterAIController.test.ts +++ b/src/backend/controllers/puterai/PuterAIController.test.ts @@ -44,7 +44,9 @@ import { } from 'vitest'; import type { Actor } from '../../core/actor.js'; +import type { RouteOptions } from '../../core/http/index.js'; import type { ChatCompletionDriver } from '../../drivers/ai-chat/ChatCompletionDriver.js'; +import { AI_CONCURRENT, AI_RATE_LIMIT } from '../../drivers/util/aiLimits.js'; import { PuterServer } from '../../server.js'; import { setupTestServer } from '../../testUtil.js'; import { PuterAIController } from './PuterAIController.js'; @@ -192,7 +194,10 @@ describe('PuterAIController.registerRoutes', () => { // call puter-chat-completion directly), but admit the user's own // full-access personal access token: `requireUserActor` keeps apps out // and `allowFullAccessToken` opens the gate to a full-access PAT. Each - // proxy route carries both flags. + // proxy route carries both flags, plus the shared per-tier AI + // rate-limit / concurrency policy — these routes bypass the + // `/drivers/call` dispatch (where the driver-declared limits are + // enforced), so without the route gates they'd be unthrottled. const userOnlyPaths = [ '/puterai/openai/v1/chat/completions', '/puterai/openai/v1/completions', @@ -205,6 +210,16 @@ describe('PuterAIController.registerRoutes', () => { subdomain: 'api', requireUserActor: true, allowFullAccessToken: true, + rateLimit: { + ...AI_RATE_LIMIT.default, + scope: 'driver:puter-chat-completion:complete', + key: expect.any(Function), + }, + concurrent: { + ...AI_CONCURRENT.default, + scope: 'driver:puter-chat-completion:complete', + key: expect.any(Function), + }, }); } const modelsRoute = calls.find( @@ -215,6 +230,46 @@ describe('PuterAIController.registerRoutes', () => { requireAuth: false, }); }); + + it('keys the proxy-route AI limits by user uuid so they share the /drivers/call buckets', () => { + const calls: Array<{ path: string; opts: RouteOptions }> = []; + const router = { + post: vi.fn((path: string, opts: RouteOptions) => { + calls.push({ path, opts }); + return router; + }), + get: vi.fn(() => router), + }; + + controller.registerRoutes(router as never); + + const route = calls.find( + (c) => c.path === '/puterai/openai/v1/chat/completions', + ); + const rateLimit = route?.opts.rateLimit as { + key: (req: Request) => string; + scope: string; + }; + const concurrent = route?.opts.concurrent as { + key: (req: Request) => string; + scope: string; + }; + + // The dispatch buckets requests as `driver:::` + // where uid is the actor's user uuid; scope + key must compose to the + // identical string or wire traffic mints a second per-user budget. + const req = makeReq({ actor: makeUserActor() }); + expect(rateLimit.key(req)).toBe('u-7'); + expect(concurrent.key(req)).toBe('u-7'); + expect(rateLimit.scope).toBe('driver:puter-chat-completion:complete'); + expect(concurrent.scope).toBe('driver:puter-chat-completion:complete'); + + // No-actor fallback still yields a usable (fingerprint) key rather + // than throwing or bucketing everyone together under undefined. + const anonKey = rateLimit.key(makeReq({})); + expect(typeof anonKey).toBe('string'); + expect(anonKey.length).toBeGreaterThan(0); + }); }); // ── /openai/v1/chat/completions ───────────────────────────────────── diff --git a/src/backend/controllers/puterai/PuterAIController.ts b/src/backend/controllers/puterai/PuterAIController.ts index 0ac3a7ed2..b9e80bc5b 100644 --- a/src/backend/controllers/puterai/PuterAIController.ts +++ b/src/backend/controllers/puterai/PuterAIController.ts @@ -22,6 +22,7 @@ import crypto from 'node:crypto'; import { Readable } from 'node:stream'; import { HttpError } from '../../core/http/HttpError.js'; import { RouteOptions } from '../../core/http/index.js'; +import { computeNetworkFingerprint } from '../../core/http/middleware/rateLimit.js'; import type { PuterRouter } from '../../core/http/PuterRouter.js'; import type { ChatCompletionDriver } from '../../drivers/ai-chat/ChatCompletionDriver.js'; import type { @@ -29,6 +30,7 @@ import type { ICompleteArguments, } from '../../drivers/ai-chat/types.js'; import { isDriverStreamResult } from '../../drivers/meta.js'; +import { AI_CONCURRENT, AI_RATE_LIMIT } from '../../drivers/util/aiLimits.js'; import { PuterController } from '../types.js'; const GEMINI_DOWNLOAD_BASE = @@ -49,6 +51,18 @@ const GEMINI_DOWNLOAD_BASE = */ export class PuterAIController extends PuterController { registerRoutes(router: PuterRouter): void { + /** + * The wire routes call the chat driver directly instead of going + * through the `/drivers/call` dispatch, so the shared per-tier AI + * rate-limit / concurrency policy must be declared as route gates + * here. `scope` + `key` reproduce the dispatch's bucket key + * (`driver:::`) exactly, so wire traffic and + * `/drivers/call` traffic draw from one per-user budget rather than + * each surface minting its own. + */ + const aiPolicyScope = 'driver:puter-chat-completion:complete'; + const aiPolicyKey = (req: Request): string => + req.actor?.user?.uuid || computeNetworkFingerprint(req); const apiAuthOpts = { subdomain: 'api', requireUserActor: true, @@ -58,6 +72,16 @@ export class PuterAIController extends PuterController { // personal access tokens (CLI/MCP/scripts). Apps and scoped tokens // stay blocked. allowFullAccessToken: true, + rateLimit: { + ...AI_RATE_LIMIT.default!, + scope: aiPolicyScope, + key: aiPolicyKey, + }, + concurrent: { + ...AI_CONCURRENT.default!, + scope: aiPolicyScope, + key: aiPolicyKey, + }, } as RouteOptions; const publicOpts = { subdomain: 'api', requireAuth: false } as const;