From 7584dffc4bde72e6c0186f12e20a7ea9b0fc044c Mon Sep 17 00:00:00 2001
From: Daniel Salazar <daniel.salazar@puter.com>
Date: Tue, 1 Sep 2026 09:11:04 -0700
Subject: [PATCH] fix: metering for ai (#3703)

---
 .../controllers/puterai/PuterAIController.ts  | 26 ++++++++++++-------
 .../drivers/ai-chat/ChatCompletionDriver.ts   | 13 ++++++++--
 2 files changed, 27 insertions(+), 12 deletions(-)

diff --git a/src/backend/controllers/puterai/PuterAIController.ts b/src/backend/controllers/puterai/PuterAIController.ts
index f338a7271..da0e4b346 100644
--- a/src/backend/controllers/puterai/PuterAIController.ts
+++ b/src/backend/controllers/puterai/PuterAIController.ts
@@ -340,9 +340,7 @@ export class PuterAIController extends PuterController {
             ...(body.temperature !== undefined
                 ? { temperature: Number(body.temperature) }
                 : {}),
-            ...(body.max_tokens !== undefined
-                ? { max_tokens: Number(body.max_tokens) }
-                : {}),
+            ...(finiteMaxTokens(body.max_tokens) ?? {}),
             ...(body.provider
                 ? { provider: toStringOrEmpty(body.provider) }
                 : { provider: DEFAULTS.openaiChat }),
@@ -500,9 +498,7 @@ export class PuterAIController extends PuterController {
             ...(body.temperature !== undefined
                 ? { temperature: Number(body.temperature) }
                 : {}),
-            ...(body.max_tokens !== undefined
-                ? { max_tokens: Number(body.max_tokens) }
-                : {}),
+            ...(finiteMaxTokens(body.max_tokens) ?? {}),
             ...(body.provider
                 ? { provider: toStringOrEmpty(body.provider) }
                 : { provider: DEFAULTS.openaiCompletion }),
@@ -593,7 +589,8 @@ export class PuterAIController extends PuterController {
                     text: extractTextContent(
                         (
                             messageResult.message as
-                                Record<string, unknown> | undefined
+                                | Record<string, unknown>
+                                | undefined
                         )?.content,
                     ),
                     index: 0,
@@ -984,9 +981,7 @@ export class PuterAIController extends PuterController {
             ...(body.temperature !== undefined
                 ? { temperature: Number(body.temperature) }
                 : {}),
-            ...(body.max_tokens !== undefined
-                ? { max_tokens: Number(body.max_tokens) }
-                : {}),
+            ...(finiteMaxTokens(body.max_tokens) ?? {}),
             ...(body.context_management !== undefined
                 ? { context_management: body.context_management }
                 : {}),
@@ -1223,6 +1218,17 @@ const asRecord = (value: unknown): Record<string, unknown> => {
 const toStringOrEmpty = (v: unknown): string =>
     typeof v === 'string' ? v : '';
 
+// A user-supplied max_tokens must coerce to a finite number. A non-numeric
+// value (e.g. the string "NaN") becomes NaN, which slips through the credit
+// gate's `?? Infinity` and every `< 1` comparison, disabling the output cap.
+// Drop it instead so the request runs with no client-requested cap rather than
+// a poisoned one.
+const finiteMaxTokens = (v: unknown): { max_tokens: number } | undefined => {
+    if (v === undefined) return undefined;
+    const n = Number(v);
+    return Number.isFinite(n) ? { max_tokens: n } : undefined;
+};
+
 const setSseHeaders = (res: Response): void => {
     res.setHeader('Content-Type', 'text/event-stream; charset=utf-8');
     res.setHeader('Cache-Control', 'no-cache, no-transform');
diff --git a/src/backend/drivers/ai-chat/ChatCompletionDriver.ts b/src/backend/drivers/ai-chat/ChatCompletionDriver.ts
index eff5194ee..36e460c8b 100644
--- a/src/backend/drivers/ai-chat/ChatCompletionDriver.ts
+++ b/src/backend/drivers/ai-chat/ChatCompletionDriver.ts
@@ -381,7 +381,13 @@ export class ChatCompletionDriver extends PuterDriver {
         // `args.max_tokens`, so the user's requested value has to be kept
         // apart from what the previous attempt was capped to.
         const promptTokenEstimate = estimatePromptTokens(args.messages ?? []);
-        const requestedMaxTokens = args.max_tokens;
+        // Direct driver calls reach here without the controller's coercion, so
+        // a non-finite max_tokens (e.g. NaN from the string "NaN") could still
+        // arrive. Drop it to undefined — a poisoned value would disable the
+        // output cap below and skip the credit hold entirely.
+        const requestedMaxTokens = Number.isFinite(args.max_tokens as number)
+            ? args.max_tokens
+            : undefined;
 
         const completionId = crypto
             .randomUUID()
@@ -902,7 +908,10 @@ export class ChatCompletionDriver extends PuterDriver {
             // 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) {
+            // `!(cap >= 1)` rather than `cap < 1` so a non-finite cap is caught
+            // too: a NaN requested max_tokens poisons the Math.min above, and
+            // `NaN < 1` is false — letting an uncapped request through.
+            if (!(cap >= 1)) {
                 throw new HttpError(402, 'No usage left for request.', {
                     legacyCode: 'insufficient_funds',
                 });