From 7ed87f760b1ad9e295958d2f7a22faa0794d9304 Mon Sep 17 00:00:00 2001
From: ProgrammerIn-wonderland
 <30693865+ProgrammerIn-wonderland@users.noreply.github.com>
Date: Tue, 9 Jun 2026 14:03:13 -0400
Subject: [PATCH] make max tokens less lenient (#3240)

---
 .../ai-chat/ChatCompletionDriver.test.ts      | 40 +++++++++++++++++++
 .../drivers/ai-chat/ChatCompletionDriver.ts   | 27 ++++++++-----
 2 files changed, 58 insertions(+), 9 deletions(-)

diff --git a/src/backend/drivers/ai-chat/ChatCompletionDriver.test.ts b/src/backend/drivers/ai-chat/ChatCompletionDriver.test.ts
index 1e57910f3..251774acc 100644
--- a/src/backend/drivers/ai-chat/ChatCompletionDriver.test.ts
+++ b/src/backend/drivers/ai-chat/ChatCompletionDriver.test.ts
@@ -547,6 +547,46 @@ describe('ChatCompletionDriver.complete credit gate and max_tokens cap', () => {
         expect(passed.max_tokens!).toBeGreaterThan(0);
     });
 
+    it('throws 402 instead of leaving `max_tokens` unset when credits cannot afford one output token', async () => {
+        // Regression: previously a sub-1 cap set max_tokens to `undefined`,
+        // which let the provider run to the model's full output limit and
+        // overdraw the account. It must reject instead.
+        vi.spyOn(FakeChatProvider.prototype, 'models').mockResolvedValueOnce([
+            {
+                id: 'capme',
+                aliases: [],
+                costs_currency: 'usd-cents',
+                costs: { input_tokens: 1000, output_tokens: 2000 },
+                max_tokens: 8192,
+            },
+        ]);
+        const d = await makeDriver();
+
+        // Pass the cheap pre-flight gate but leave a balance too small to
+        // afford a single 2000-microcent output token.
+        vi.spyOn(server.services.metering, 'hasEnoughCredits').mockResolvedValue(
+            true,
+        );
+        vi.spyOn(server.services.metering, 'getRemainingUsage').mockResolvedValue(
+            100,
+        );
+
+        const completeSpy = vi.spyOn(FakeChatProvider.prototype, 'complete');
+
+        await expect(
+            withTestActor(() =>
+                d.complete({
+                    model: 'capme',
+                    messages: [{ role: 'user', content: 'hi' }],
+                }),
+            ),
+        ).rejects.toMatchObject({
+            statusCode: 402,
+            legacyCode: 'insufficient_funds',
+        });
+        expect(completeSpy).not.toHaveBeenCalled();
+    });
+
     it('rejects subscriber-only models for the default free subscription', async () => {
         vi.spyOn(FakeChatProvider.prototype, 'models').mockResolvedValueOnce([
             {
diff --git a/src/backend/drivers/ai-chat/ChatCompletionDriver.ts b/src/backend/drivers/ai-chat/ChatCompletionDriver.ts
index ae88fbe76..fad6af01a 100644
--- a/src/backend/drivers/ai-chat/ChatCompletionDriver.ts
+++ b/src/backend/drivers/ai-chat/ChatCompletionDriver.ts
@@ -400,16 +400,25 @@ export class ChatCompletionDriver extends PuterDriver {
                     remainingCredits - approximateInputCost;
                 const maxAllowedOutputTokens =
                     maxAllowedOutputUcents / outputTokenCost;
-                if (maxAllowedOutputTokens) {
-                    const cap = Math.floor(
-                        Math.min(
-                            args.max_tokens ?? Number.POSITIVE_INFINITY,
-                            maxAllowedOutputTokens,
-                            model.max_tokens - approximateTokenCount,
-                        ),
-                    );
-                    args.max_tokens = cap < 1 ? undefined : cap;
+                const cap = Math.floor(
+                    Math.min(
+                        args.max_tokens ?? Number.POSITIVE_INFINITY,
+                        maxAllowedOutputTokens,
+                        model.max_tokens - approximateTokenCount,
+                    ),
+                );
+                // `cap` is the credit-bounded ceiling on output tokens. When
+                // it drops below 1 the user can't afford even a single output
+                // token, so reject the request. Crucially we must NOT leave
+                // `max_tokens` unset here: an undefined max_tokens lets the
+                // provider run to the model's full output limit (e.g. 128k for
+                // Claude), billing far past the user's remaining balance.
+                if (cap < 1) {
+                    throw new HttpError(402, 'No usage left for request.', {
+                        legacyCode: 'insufficient_funds',
+                    });
                 }
+                args.max_tokens = cap;
             }
         }