fix(ai): bound large AI request bodies per user and to chat (#4420)

Follow-up to #4408.

The 503 returned when a process holds too many large bodies is load
shedding, so it records an info alarm instead of a critical one.

One user can hold at most 128 MiB of large bodies at once; past that the
request gets 429 with Retry-After, so a single account can't take a
process's whole budget, e.g. with slow uploads.

Through /drivers/call only puter-chat-completion takes a body over 50 MiB;
other drivers keep the default limit and get 413. The check reads the
parsed body's interface, since route matching ignores case and trailing
slashes.

src/backend/core/http/middleware/largeJsonBody.test.ts

@@ -29,14 +29,19 @@ import {
 const MiB = 1024 * 1024;
 
 // A complete request with no socket reads as already consumed to body-parser,
-// so the parsers call straight through; these tests are about admission.
-const makeReq = (bytes: number | undefined, signedIn = true) =>
+// so the parsers call straight through, leaving any preset `body` in place.
+const makeReq = (
+    bytes: number | undefined,
+    user: string | null = 'u1',
+    body?: unknown,
+) =>
     ({
         complete: true,
+        body,
         headers: {
             'content-type': 'application/json',
             ...(bytes === undefined ? {} : { 'content-length': String(bytes) }),
-            ...(signedIn ? { authorization: 'Bearer t' } : {}),
+            ...(user ? { authorization: user } : {}),
         },
     }) as unknown as Request;
 
@@ -50,7 +55,9 @@ const makeRes = () => {
 
 const authProbe = vi.fn<RequestHandler>((req, _res, next) => {
     if (req.headers.authorization) {
-        (req as { actor?: unknown }).actor = { user: { uuid: 'u' } };
+        (req as { actor?: unknown }).actor = {
+            user: { uuid: req.headers.authorization },
+        };
     }
     next();
 });
@@ -66,12 +73,24 @@ const run = (mw: RequestHandler, req: Request, res = makeRes()) =>
         },
     );
 
+const statusOf = (err: unknown) => (isHttpError(err) ? err.statusCode : err);
+
 describe('createLargeJsonBodyParser', () => {
-    const create = (budgetBytes = 200 * MiB) =>
+    const create = ({
+        budgetBytes = 200 * MiB,
+        perUserBytes = 128 * MiB,
+        accepts,
+    }: {
+        budgetBytes?: number;
+        perUserBytes?: number;
+        accepts?: (req: Request) => boolean;
+    } = {}) =>
         createLargeJsonBodyParser({
             authProbe,
             limitBytes: 128 * MiB,
             budgetBytes,
+            perUserBytes,
+            accepts,
         });
 
     it.each([
@@ -86,30 +105,59 @@ describe('createLargeJsonBodyParser', () => {
     });
 
     it('leaves an unauthenticated large body to the global parser', async () => {
-        const mw = create(60 * MiB);
-        const { err } = await run(mw, makeReq(60 * MiB, false));
+        const mw = create({ budgetBytes: 60 * MiB });
+        const { err } = await run(mw, makeReq(60 * MiB, null));
         expect(err).toBeUndefined();
         // Nothing was charged, so a signed-in caller still fits.
         expect((await run(mw, makeReq(60 * MiB))).err).toBeUndefined();
     });
 
     it('charges nothing for a client that left during the probe', async () => {
-        const mw = create(60 * MiB);
+        const mw = create({ budgetBytes: 60 * MiB });
         const gone = Object.assign(makeRes(), { closed: true });
         expect((await run(mw, makeReq(60 * MiB), gone)).err).toBeUndefined();
         expect((await run(mw, makeReq(60 * MiB))).err).toBeUndefined();
     });
 
     it('refuses with 503 past the budget until a response closes', async () => {
-        const mw = create(130 * MiB);
-        const first = await run(mw, makeReq(60 * MiB));
-        await run(mw, makeReq(60 * MiB));
+        const mw = create({ budgetBytes: 130 * MiB });
+        const first = await run(mw, makeReq(60 * MiB, 'u1'));
+        await run(mw, makeReq(60 * MiB, 'u2'));
 
-        const refused = await run(mw, makeReq(60 * MiB));
-        expect(isHttpError(refused.err) && refused.err.statusCode).toBe(503);
+        const refused = await run(mw, makeReq(60 * MiB, 'u3'));
+        expect(statusOf(refused.err)).toBe(503);
+        // Load shedding is recorded, not paged.
+        expect(isHttpError(refused.err) && refused.err.alarmSeverity).toBe(
+            'info',
+        );
         expect(refused.res.setHeader).toHaveBeenCalledWith('Retry-After', '5');
 
         first.res.emit('close');
-        expect((await run(mw, makeReq(60 * MiB))).err).toBeUndefined();
+        expect((await run(mw, makeReq(60 * MiB, 'u3'))).err).toBeUndefined();
+    });
+
+    it("refuses one user's bodies with 429 past their share", async () => {
+        const mw = create({ perUserBytes: 100 * MiB });
+        const first = await run(mw, makeReq(60 * MiB, 'u1'));
+
+        const refused = await run(mw, makeReq(60 * MiB, 'u1'));
+        expect(statusOf(refused.err)).toBe(429);
+        expect(refused.res.setHeader).toHaveBeenCalledWith('Retry-After', '5');
+        // Another user isn't held back by it.
+        expect((await run(mw, makeReq(60 * MiB, 'u2'))).err).toBeUndefined();
+
+        first.res.emit('close');
+        expect((await run(mw, makeReq(60 * MiB, 'u1'))).err).toBeUndefined();
+    });
+
+    it('refuses with 413 a parsed body `accepts` turns down', async () => {
+        const mw = create({
+            accepts: (req) => (req.body as { ok?: boolean }).ok === true,
+        });
+        const refused = await run(mw, makeReq(60 * MiB, 'u1', { ok: false }));
+        expect(statusOf(refused.err)).toBe(413);
+        expect(
+            (await run(mw, makeReq(60 * MiB, 'u2', { ok: true }))).err,
+        ).toBeUndefined();
     });
 });

src/backend/core/http/middleware/largeJsonBody.ts

@@ -17,7 +17,8 @@
  * along with this program.  If not, see <https://www.gnu.org/licenses/>.
  */
 
-import express, { type RequestHandler } from 'express';
+import express, { type Request, type RequestHandler } from 'express';
+import { actorUid } from '../../actor';
 import { HttpError } from '../HttpError';
 
 /** The global JSON parser's limit, which applies before auth. */
@@ -29,6 +30,10 @@ interface LargeJsonBodyOptions {
     limitBytes: number;
     /** Declared bytes of large bodies a process holds at once. */
     budgetBytes: number;
+    /** Declared bytes of large bodies one user holds at once. */
+    perUserBytes: number;
+    /** Checked once the body is parsed; a body it refuses gets a 413. */
+    accepts?: (req: Request) => boolean;
 }
 
 /**
@@ -38,12 +43,15 @@ interface LargeJsonBodyOptions {
  * limit.
  *
  * The parsed body lives as long as the request, so a body's bytes count against
- * the budget until its response closes. Over budget gets a 503.
+ * both budgets until its response closes. Over the user's budget gets a 429,
+ * over the process budget a 503.
  */
 export const createLargeJsonBodyParser = ({
     authProbe,
     limitBytes,
     budgetBytes,
+    perUserBytes,
+    accepts,
 }: LargeJsonBodyOptions): RequestHandler => {
     const json = express.json({ limit: limitBytes });
     const textAsJson = express.json({
@@ -52,6 +60,7 @@ export const createLargeJsonBodyParser = ({
             req.headers['content-type'] === 'text/plain;actually=json',
     });
     let inFlightBytes = 0;
+    const inFlightByUser = new Map<string, number>();
 
     return (req, res, next) => {
         const declared = Number(req.headers['content-length']);
@@ -66,25 +75,59 @@ export const createLargeJsonBodyParser = ({
                 next();
                 return;
             }
+            const user = req.actor.user.uuid ?? actorUid(req.actor);
+            const userBytes = inFlightByUser.get(user) ?? 0;
+            if (userBytes + declared > perUserBytes) {
+                res.setHeader('Retry-After', '5');
+                next(
+                    new HttpError(
+                        429,
+                        'Too many large requests in progress for this account',
+                        { legacyCode: 'too_many_requests' },
+                    ),
+                );
+                return;
+            }
             if (inFlightBytes + declared > budgetBytes) {
                 res.setHeader('Retry-After', '5');
                 next(
                     new HttpError(503, 'Too many large requests in progress', {
                         code: 'large_body_capacity',
+                        // Load shedding, not a fault: record it, don't page.
+                        alarmSeverity: 'info',
                     }),
                 );
                 return;
             }
             inFlightBytes += declared;
+            inFlightByUser.set(user, userBytes + declared);
             res.once('close', () => {
                 inFlightBytes -= declared;
+                const left = (inFlightByUser.get(user) ?? 0) - declared;
+                if (left > 0) inFlightByUser.set(user, left);
+                else inFlightByUser.delete(user);
             });
             json(req, res, (err?: unknown) => {
                 if (err) {
                     next(err);
                     return;
                 }
-                textAsJson(req, res, next);
+                textAsJson(req, res, (err?: unknown) => {
+                    if (err) {
+                        next(err);
+                        return;
+                    }
+                    if (accepts && !accepts(req)) {
+                        // The same answer the default limit gives.
+                        next(
+                            new HttpError(413, 'request entity too large', {
+                                legacyCode: 'client_bad_request',
+                            }),
+                        );
+                        return;
+                    }
+                    next();
+                });
             });
         });
     };

src/backend/server.test.ts

@@ -601,10 +601,10 @@ describe('PuterServer request parsing', () => {
         },
     );
 
-    it('parses JSON above 50 MiB on an AI route for a signed-in caller', async () => {
+    it('parses a chat call above 50 MiB for a signed-in caller', async () => {
         const payload = JSON.stringify({
-            interface: 'no-such-interface',
-            method: 'complete',
+            interface: 'puter-chat-completion',
+            method: 'no-such-method',
             pad: 'x'.repeat(51 * MiB),
         });
         for (const contentType of [
@@ -615,9 +615,24 @@ describe('PuterServer request parsing', () => {
                 'content-type': contentType,
                 authorization: `Bearer ${token}`,
             });
-            // The controller read `interface` out of the body.
+            // The controller read `method` out of the body.
             expect(res.status, contentType).toBe(404);
-            expect(res.body).toContain('no-such-interface');
+            expect(res.body).toContain('no-such-method');
+        }
+    });
+
+    it('keeps the default JSON limit for other drivers', async () => {
+        const payload = JSON.stringify({
+            interface: 'puter-kvstore',
+            method: 'get',
+            pad: 'x'.repeat(51 * MiB),
+        });
+        // Routing matches these spellings too.
+        for (const path of ['/drivers/call', '/Drivers/Call/']) {
+            const res = await postJson(path, payload, {
+                authorization: `Bearer ${token}`,
+            });
+            expect(res.status, path).toBe(413);
         }
     });
 

src/backend/server.ts

@@ -145,6 +145,17 @@ const LARGE_JSON_LIMIT_BYTES = 128 * 1024 * 1024;
 /** Two max-size bodies per process. */
 const LARGE_JSON_BUDGET_BYTES = 2 * LARGE_JSON_LIMIT_BYTES;
 
+/**
+ * Through `/drivers/call`, only chat takes a large body; other drivers keep the
+ * default limit. Keyed on the body rather than the path, which routing matches
+ * loosely (case, trailing slash). Wire API bodies carry no `interface`.
+ */
+const acceptsLargeJson = (req: Request): boolean => {
+    const iface = (req.body as { interface?: unknown } | null | undefined)
+        ?.interface;
+    return iface === undefined || iface === 'puter-chat-completion';
+};
+
 /** Route options a claimed subdomain's routes may use. */
 const CLAIMED_ROUTE_OPTIONS = new Set([
     'subdomain',
@@ -559,6 +570,8 @@ export class PuterServer {
                     authProbe,
                     limitBytes: LARGE_JSON_LIMIT_BYTES,
                     budgetBytes: LARGE_JSON_BUDGET_BYTES,
+                    perUserBytes: LARGE_JSON_LIMIT_BYTES,
+                    accepts: acceptsLargeJson,
                 }),
             );
         }

src/docs/src/rate-limits-and-quotas.md

@@ -146,7 +146,7 @@ The [`puter.ai.convo()`](/AI/convo/) SDK also applies these limits before sendin
 
 `/puterai/anthropic/v1/messages/count_tokens` has its own budget of 120 requests per minute per user, with no concurrency limit, and is not charged.
 
-AI requests (`/drivers/call`, which `puter.ai.*` uses, and POSTs to `/puterai/openai/v1/*`, `/puterai/anthropic/v1/messages` and `/puterai/anthropic/v1/messages/count_tokens`) accept JSON bodies up to 128 MiB (134,217,728 bytes), including `text/plain;actually=json` bodies. A body over 50 MiB needs a `Content-Length` header and a token in the `Authorization` or `x-api-key` header (or a session cookie); an `auth_token` inside the body isn't read until the body is, so it only works up to 50 MiB. Larger bodies return HTTP 413. A server handling too many large bodies at once returns `503` with `Retry-After`; retry after the given number of seconds. Other routes accept JSON bodies up to 50 MiB.
+Chat requests accept JSON bodies up to 128 MiB (134,217,728 bytes), including `text/plain;actually=json` bodies. That covers `puter-chat-completion` calls through `/drivers/call`, which `puter.ai.chat()` uses, and POSTs to `/puterai/openai/v1/*`, `/puterai/anthropic/v1/messages` and `/puterai/anthropic/v1/messages/count_tokens`. A body over 50 MiB needs a `Content-Length` header and a token in the `Authorization` or `x-api-key` header (or a session cookie); an `auth_token` inside the body isn't read until the body is, so it only works up to 50 MiB. Larger bodies, and bodies over 50 MiB for other drivers, return HTTP 413. One account can have up to 128 MiB of these large bodies in flight at once; past that, requests get `429` with `Retry-After`. A server handling too many large bodies at once returns `503` with `Retry-After`. Retry either after the given number of seconds. Other routes accept JSON bodies up to 50 MiB.
 
 Claude's server-executed tools are clamped so a request can't reserve an unbounded amount of credit: