fix: ai rate limits (#3371)

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`).
 

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:<iface>:<method>:<uid>`
+        // 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 ─────────────────────────────────────

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:<iface>:<method>:<uid>`) 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;