fix: only reinterpret an error as a budget rejection when it says so

Review of the first commit found the retry could fire on errors that have
nothing to do with the budget, which would replace a readable provider error
with a truncated response: exactly the symptom this PR exists to remove.

- Drop the generic "lower than N" pattern. For the Bedrock message it was dead
  code, since "model limit of N" matches first with the same number. Left live,
  it would read a number out of any message shaped like "must be lower than 2".
- Skip errors whose status is not 400 or 422, so auth and rate-limit failures
  are never reinterpreted.
- Require the parsed ceiling to be at least 1024. Below that a diagram cannot
  come out whole, so retrying would hide the error behind broken XML.
- Validate MAX_OUTPUT_TOKENS from env the same way as the header, so a stray
  "-1" falls back instead of reaching the provider.

Adds tests for the retry wrapper itself, which had none: it retries once with
the named ceiling, leaves a 401 alone, does not retry when the ceiling is not
smaller, propagates a second rejection, and preserves the other call options.

Re-verified against the live APIs: bedrock nova-lite still logs "64000 rejected,
retrying with 10000" and completes its tool call, and deepseek-v4-flash still
finishes normally at 64000.

app/api/chat/route.ts

@@ -34,11 +34,17 @@ import {
     setTraceOutput,
     wrapWithObserve,
 } from "@/lib/langfuse"
+import {
+    resolveMaxOutputTokens,
+    withOutputTokenLimitFallback,
+} from "@/lib/output-token-limit"
 import { findServerModelById } from "@/lib/server-model-config"
 import { getSystemPrompt } from "@/lib/system-prompts"
 import { getUserIdFromRequest } from "@/lib/user-id"
 
-export const maxDuration = 120
+// No explicit cap: a reasoning model can spend minutes planning before it emits
+// the tool call, so take whatever the host allows. Vercel's own default is 300s,
+// which is also where Node's response-body timeout on the upstream stream lands.
 
 // Helper function to create cached stream response
 function createCachedStreamResponse(xml: string): Response {
@@ -241,13 +247,22 @@ async function handleChatRequest(req: Request): Promise<Response> {
 
     // Get AI model with optional client overrides
     const {
-        model,
+        model: baseModel,
         providerOptions,
         headers,
         modelId,
         provider: resolvedProvider,
     } = getAIModel(clientOverrides)
 
+    // Retry with a smaller budget if the provider rejects the requested one
+    const model = withOutputTokenLimitFallback(baseModel)
+
+    // User setting wins over server env, so desktop users can raise it themselves
+    const maxOutputTokens = resolveMaxOutputTokens(
+        req.headers.get("x-max-output-tokens"),
+    )
+    console.log(`[maxOutputTokens] ${maxOutputTokens}`)
+
     // Check if model supports prompt caching
     const shouldCache = supportsPromptCaching(modelId)
     console.log(
@@ -493,9 +508,9 @@ IMPORTANT: The "Current diagram XML" is the SINGLE SOURCE OF TRUTH for what's on
     const result = streamText({
         model,
         abortSignal: req.signal,
-        // Must be sent: unset means the provider's own default, and Bedrock's is 4096 —
+        // Must be sent: unset means the provider's own default, and Bedrock's is
-        // enough for a small diagram, so larger ones were cut off mid-attribute.
+        // 4096, enough for a small diagram, so larger ones were cut off mid-attribute.
-        maxOutputTokens: Number(process.env.MAX_OUTPUT_TOKENS) || 16000,
+        maxOutputTokens,
         stopWhen: stepCountIs(5),
         // Repair truncated tool calls when maxOutputTokens is reached mid-JSON
         experimental_repairToolCall: async ({ toolCall, error }) => {

components/chat-panel.tsx

@@ -178,6 +178,7 @@ export default function ChatPanel({
     const [minimalStyle, setMinimalStyle] = useState(false)
     const [vlmValidationEnabled, setVlmValidationEnabled] = useState(false)
     const [customSystemMessage, setCustomSystemMessage] = useState("")
+    const [maxOutputTokens, setMaxOutputTokens] = useState("")
     const [shouldFocusInput, setShouldFocusInput] = useState(false)
 
     // Restore input from sessionStorage on mount (when ChatPanel remounts due to key change)
@@ -204,6 +205,14 @@ export default function ChatPanel({
         }
     }, [])
 
+    // Load output token budget from localStorage on mount
+    useEffect(() => {
+        const stored = localStorage.getItem(STORAGE_KEYS.maxOutputTokens)
+        if (stored !== null) {
+            setMaxOutputTokens(stored)
+        }
+    }, [])
+
     // Check config on mount
     useEffect(() => {
         fetch(getApiEndpoint("/api/config"))
@@ -320,6 +329,13 @@ export default function ChatPanel({
         localStorage.setItem(STORAGE_KEYS.customSystemMessage, value)
     }, [])
 
+    // Handler for output token budget change (empty string = use server default)
+    const handleMaxOutputTokensChange = useCallback((value: string) => {
+        const digitsOnly = value.replace(/\D/g, "")
+        setMaxOutputTokens(digitsOnly)
+        localStorage.setItem(STORAGE_KEYS.maxOutputTokens, digitsOnly)
+    }, [])
+
     // Ref to store the sendMessage function for use in callbacks
     const sendMessageRef = useRef<typeof sendMessage | null>(null)
 
@@ -1104,6 +1120,9 @@ export default function ChatPanel({
                     ...(minimalStyle && {
                         "x-minimal-style": "true",
                     }),
+                    ...(maxOutputTokens && {
+                        "x-max-output-tokens": maxOutputTokens,
+                    }),
                 },
             },
         )
@@ -1448,6 +1467,8 @@ export default function ChatPanel({
                 onVlmValidationChange={handleVlmValidationChange}
                 customSystemMessage={customSystemMessage}
                 onCustomSystemMessageChange={handleCustomSystemMessageChange}
+                maxOutputTokens={maxOutputTokens}
+                onMaxOutputTokensChange={handleMaxOutputTokensChange}
                 onOpenModelConfig={() => setShowModelConfigDialog(true)}
             />
 

components/settings-dialog.tsx

@@ -75,6 +75,8 @@ interface SettingsDialogProps {
     onOpenModelConfig?: () => void
     customSystemMessage?: string
     onCustomSystemMessageChange?: (value: string) => void
+    maxOutputTokens?: string
+    onMaxOutputTokensChange?: (value: string) => void
 }
 
 export const STORAGE_ACCESS_CODE_KEY = "next-ai-draw-io-access-code"
@@ -101,6 +103,8 @@ function SettingsContent({
     onOpenModelConfig,
     customSystemMessage = "",
     onCustomSystemMessageChange = () => {},
+    maxOutputTokens = "",
+    onMaxOutputTokensChange = () => {},
 }: SettingsDialogProps) {
     const dict = useDictionary()
     const router = useRouter()
@@ -591,6 +595,24 @@ function SettingsContent({
                         />
                     </div>
 
+                    {/* Max Output Tokens */}
+                    <SettingItem
+                        label={dict.settings.maxOutputTokens}
+                        description={dict.settings.maxOutputTokensDescription}
+                    >
+                        <Input
+                            id="max-output-tokens"
+                            type="text"
+                            inputMode="numeric"
+                            value={maxOutputTokens}
+                            onChange={(e) =>
+                                onMaxOutputTokensChange(e.target.value)
+                            }
+                            placeholder="64000"
+                            className="h-9 w-28 text-sm"
+                        />
+                    </SettingItem>
+
                     {/* Send Shortcut */}
                     <SettingItem
                         label={dict.settings.sendShortcut}

env.example

@@ -11,8 +11,10 @@ AI_PROVIDER=bedrock
 # Example: AI_MODEL=doubao-seed-1-8-251215,doubao-seed-1-6-flash,doubao-seed-1-6-pro
 AI_MODEL=global.anthropic.claude-sonnet-4-5-20250929-v1:0
 
-# Output limit, all providers (default: 16000). Raise it if large diagrams arrive cut off.
+# Output limit, all providers (default: 64000). Shared by reasoning and the diagram XML,
-# MAX_OUTPUT_TOKENS=16000
+# so a thinking model can spend it all before the tool call. Users can override it in Settings.
+# If a model's own ceiling is lower, the request is retried with that ceiling automatically.
+# MAX_OUTPUT_TOKENS=64000
 
 # AWS Bedrock Configuration
 # AWS_REGION=us-east-1

lib/i18n/dictionaries/en.json

@@ -132,6 +132,8 @@
         "customSystemMessage": "Custom System Message",
         "customSystemMessageDescription": "Add custom instructions appended to the AI's system prompt.",
         "customSystemMessagePlaceholder": "e.g., Always use blue color scheme for diagrams...",
+        "maxOutputTokens": "Max Output Tokens",
+        "maxOutputTokensDescription": "Budget for one reply, shared by thinking and the diagram XML. Raise it if the AI keeps thinking and no diagram appears. Leave empty for the default.",
         "panelVisibility": "Lobby Panels",
         "panelVisibilityDescription": "Choose which panels to show on the chat lobby.",
         "showRecentChats": "Recent Chats",

lib/i18n/dictionaries/ja.json

@@ -132,6 +132,8 @@
         "customSystemMessage": "カスタムシステムメッセージ",
         "customSystemMessageDescription": "AIのシステムプロンプトに追加されるカスタム指示を入力します。",
         "customSystemMessagePlaceholder": "例：ダイアグラムには常に青色のカラースキームを使用...",
+        "maxOutputTokens": "最大出力トークン数",
+        "maxOutputTokensDescription": "1回の応答の予算で、思考過程とダイアグラムの XML が共有します。AI が考え続けてダイアグラムが生成されない場合は大きくしてください。空欄ならデフォルト値を使います。",
         "panelVisibility": "ロビーパネル",
         "panelVisibilityDescription": "チャットロビーに表示するパネルを選択します。",
         "showRecentChats": "最近のチャット",

lib/i18n/dictionaries/zh-Hant.json

@@ -132,6 +132,8 @@
         "customSystemMessage": "自訂系統訊息",
         "customSystemMessageDescription": "新增自訂指示，將附加到 AI 的系統提示末尾。",
         "customSystemMessagePlaceholder": "例如：圖表始終使用藍色配色方案...",
+        "maxOutputTokens": "最大輸出 token 數",
+        "maxOutputTokensDescription": "單次回覆的額度，思考過程與圖表 XML 共用。若 AI 一直在思考卻沒有產生圖表，請將它調大。留空則使用預設值。",
         "panelVisibility": "大廳面板",
         "panelVisibilityDescription": "選擇在聊天大廳顯示哪些面板。",
         "showRecentChats": "最近聊天",

lib/i18n/dictionaries/zh.json

@@ -132,6 +132,8 @@
         "customSystemMessage": "自定义系统消息",
         "customSystemMessageDescription": "添加自定义指令，将附加到 AI 的系统提示末尾。",
         "customSystemMessagePlaceholder": "例如：图表始终使用蓝色配色方案...",
+        "maxOutputTokens": "最大输出 token 数",
+        "maxOutputTokensDescription": "单次回复的额度，思考过程和图表 XML 共用。如果 AI 一直在思考却没有生成图表，请把它调大。留空则使用默认值。",
         "panelVisibility": "大厅面板",
         "panelVisibilityDescription": "选择在聊天大厅显示哪些面板。",
         "showRecentChats": "最近聊天",

lib/output-token-limit.ts

@@ -0,0 +1,145 @@
+import { wrapLanguageModel } from "ai"
+
+type WrappedModel = ReturnType<typeof wrapLanguageModel>
+
+/**
+ * Default output budget for a chat turn.
+ *
+ * This has to cover thinking + prose + the tool call, because reasoning models
+ * spend it in that order. Measured on deepseek-v4-flash: refining an existing
+ * diagram burned 16000 tokens on thinking alone and the request ended with
+ * finishReason "length" before display_diagram was ever called (issue #924).
+ * 64000 leaves room for the plan and the XML in one turn.
+ */
+export const DEFAULT_MAX_OUTPUT_TOKENS = 64000
+
+/** Ceiling for the user-supplied override, to catch typos like an extra zero. */
+export const MAX_OUTPUT_TOKENS_LIMIT = 200000
+
+/**
+ * Below this a diagram cannot come out whole, so a retry would just produce
+ * truncated XML instead of the provider's error. Better to surface the error.
+ */
+const MIN_USABLE_OUTPUT_TOKENS = 1024
+
+/** Status codes that can carry a complaint about the requested budget. */
+const BUDGET_REJECTION_STATUSES = new Set([400, 422])
+
+function usableLimit(value: number): number | null {
+    return value >= MIN_USABLE_OUTPUT_TOKENS ? value : null
+}
+
+/**
+ * A budget this large exceeds what some models accept. Providers reject it with a
+ * 400 that names the real limit, so we parse the number out and retry once
+ * instead of failing the turn.
+ *
+ * Formats seen in the wild:
+ * - Bedrock: "The maximum tokens you requested exceeds the model limit of 4096."
+ * - OpenRouter: "This endpoint's maximum context length is 64000 tokens. However,
+ *   you requested about 64025 tokens (25 of text input, 64000 in the output)."
+ *   Note this one is an input+output ceiling, so the input has to be subtracted.
+ * - Anthropic: "max_tokens: 200000 > 64000, which is the maximum allowed..."
+ * - OpenAI: "This model supports at most 16384 completion tokens"
+ *
+ * Every pattern names tokens explicitly. A generic one (an earlier draft matched
+ * "lower than N") would reinterpret unrelated failures, and retrying on a bogus
+ * number turns a readable error into an empty diagram.
+ */
+export function parseOutputTokenLimit(error: unknown): number | null {
+    const err = error as {
+        message?: unknown
+        responseBody?: unknown
+        statusCode?: unknown
+    }
+
+    // An auth or rate-limit failure is not about the budget, so leave it alone.
+    if (
+        typeof err?.statusCode === "number" &&
+        !BUDGET_REJECTION_STATUSES.has(err.statusCode)
+    ) {
+        return null
+    }
+
+    const text = [
+        typeof err?.message === "string" ? err.message : "",
+        typeof err?.responseBody === "string" ? err.responseBody : "",
+    ].join(" ")
+
+    if (!text) return null
+
+    // Combined input+output ceiling: subtract the input the provider counted,
+    // plus a small margin because its estimate is approximate.
+    const context = text.match(/maximum context length is (\d+)/i)
+    if (context) {
+        const input = text.match(/(\d+) of text input/i)
+        return usableLimit(
+            Number(context[1]) - (input ? Number(input[1]) : 0) - 1024,
+        )
+    }
+
+    const output =
+        text.match(/model limit of (\d+)/i) ||
+        text.match(/> (\d+), which is the maximum/i) ||
+        text.match(/at most (\d+) completion tokens/i)
+
+    return output ? usableLimit(Number(output[1])) : null
+}
+
+/**
+ * Retry the stream once with a smaller budget when the provider rejects the
+ * requested one. Without this, raising the default breaks every model whose
+ * ceiling is below it (measured: bedrock claude-3-haiku 4096, nova-lite 10000,
+ * openrouter deepseek-r1 64000 shared with the input).
+ */
+export function withOutputTokenLimitFallback(
+    model: WrappedModel,
+): WrappedModel {
+    return wrapLanguageModel({
+        model,
+        middleware: {
+            specificationVersion: "v3",
+            async wrapStream({ doStream, params, model: inner }) {
+                try {
+                    return await doStream()
+                } catch (error) {
+                    const limit = parseOutputTokenLimit(error)
+                    const requested = params.maxOutputTokens
+
+                    if (!limit || !requested || limit >= requested) throw error
+
+                    console.warn(
+                        `[maxOutputTokens] ${requested} rejected, retrying with ${limit}`,
+                    )
+                    return await inner.doStream({
+                        ...params,
+                        maxOutputTokens: limit,
+                    })
+                }
+            },
+        },
+    })
+}
+
+function validBudget(value: string | null | undefined): number | null {
+    const parsed = Number(value)
+    return Number.isInteger(parsed) &&
+        parsed > 0 &&
+        parsed <= MAX_OUTPUT_TOKENS_LIMIT
+        ? parsed
+        : null
+}
+
+/**
+ * Resolve the output budget: user setting (sent as a header so it works in the
+ * desktop app too), then server env, then the default. Both sources go through
+ * the same validation, so a typo in either falls back instead of reaching the
+ * provider.
+ */
+export function resolveMaxOutputTokens(headerValue: string | null): number {
+    return (
+        validBudget(headerValue) ??
+        validBudget(process.env.MAX_OUTPUT_TOKENS) ??
+        DEFAULT_MAX_OUTPUT_TOKENS
+    )
+}

lib/storage.ts

@@ -31,6 +31,9 @@ export const STORAGE_KEYS = {
     // Custom system message
     customSystemMessage: "next-ai-draw-io-custom-system-message",
 
+    // Output token budget per turn (empty = server default)
+    maxOutputTokens: "next-ai-draw-io-max-output-tokens",
+
     // Panel visibility
     showRecentChats: "next-ai-draw-io-show-recent-chats",
     showMyTemplates: "next-ai-draw-io-show-my-templates",

tests/e2e/settings.spec.ts

@@ -18,6 +18,26 @@ test.describe("Settings", () => {
         await expect(dialog.locator('text="English"')).toBeVisible()
     })
 
+    test("max output tokens is editable and persists", async ({ page }) => {
+        await openSettings(page)
+
+        const input = page.locator("#max-output-tokens")
+        await expect(input).toBeVisible()
+
+        await input.fill("48000")
+        await expect
+            .poll(() =>
+                page.evaluate(() =>
+                    localStorage.getItem("next-ai-draw-io-max-output-tokens"),
+                ),
+            )
+            .toBe("48000")
+
+        // Non-digits are dropped so the header always carries a plain number
+        await input.fill("12k000")
+        await expect(input).toHaveValue("12000")
+    })
+
     test("draw.io theme toggle exists", async ({ page }) => {
         await openSettings(page)
 

tests/unit/output-token-limit.test.ts

@@ -0,0 +1,271 @@
+import { describe, expect, it } from "vitest"
+import {
+    DEFAULT_MAX_OUTPUT_TOKENS,
+    parseOutputTokenLimit,
+    resolveMaxOutputTokens,
+    withOutputTokenLimitFallback,
+} from "@/lib/output-token-limit"
+
+describe("parseOutputTokenLimit", () => {
+    it("reads the ceiling from a Bedrock rejection", () => {
+        const error = {
+            message:
+                "The maximum tokens you requested exceeds the model limit of 4096. Try again with a maximum tokens value that is lower than 4096.",
+        }
+        expect(parseOutputTokenLimit(error)).toBe(4096)
+    })
+
+    it("subtracts the input when the ceiling covers input plus output", () => {
+        const error = {
+            message:
+                "This endpoint's maximum context length is 64000 tokens. However, you requested about 64025 tokens (25 of text input, 64000 in the output).",
+        }
+        // 64000 - 25 - 1024 margin
+        expect(parseOutputTokenLimit(error)).toBe(62951)
+    })
+
+    it("reads the ceiling from an Anthropic rejection", () => {
+        const error = {
+            message:
+                "max_tokens: 200000 > 64000, which is the maximum allowed number of output tokens for claude-sonnet-4-5",
+        }
+        expect(parseOutputTokenLimit(error)).toBe(64000)
+    })
+
+    it("reads the ceiling from an OpenAI rejection", () => {
+        const error = {
+            message:
+                "max_tokens is too large: 64000. This model supports at most 16384 completion tokens",
+        }
+        expect(parseOutputTokenLimit(error)).toBe(16384)
+    })
+
+    it("looks in the response body too", () => {
+        const error = {
+            message: "Bad request",
+            responseBody: '{"message":"exceeds the model limit of 10000."}',
+        }
+        expect(parseOutputTokenLimit(error)).toBe(10000)
+    })
+
+    it("returns null for unrelated errors", () => {
+        expect(parseOutputTokenLimit({ message: "Invalid API key" })).toBeNull()
+        expect(parseOutputTokenLimit(undefined)).toBeNull()
+    })
+
+    it("ignores a number that is not about tokens", () => {
+        // An earlier draft matched "lower than N" generically, which turned any
+        // message shaped like this into a bogus budget
+        expect(
+            parseOutputTokenLimit({
+                message: "temperature must be lower than 2",
+                statusCode: 400,
+            }),
+        ).toBeNull()
+        expect(
+            parseOutputTokenLimit({
+                message: "reduce requests to lower than 60 per minute",
+                statusCode: 429,
+            }),
+        ).toBeNull()
+    })
+
+    it("skips errors whose status is not a bad request", () => {
+        const error = {
+            message: "exceeds the model limit of 4096",
+            statusCode: 429,
+        }
+        expect(parseOutputTokenLimit(error)).toBeNull()
+    })
+
+    it("rejects a ceiling too small to hold a diagram", () => {
+        expect(
+            parseOutputTokenLimit({ message: "model limit of 200" }),
+        ).toBeNull()
+        // Context ceiling that leaves almost nothing after the input
+        expect(
+            parseOutputTokenLimit({
+                message:
+                    "This endpoint's maximum context length is 64000 tokens. However, you requested about 128000 tokens (63500 of text input, 64000 in the output).",
+            }),
+        ).toBeNull()
+    })
+
+    it("returns null when the input alone fills the context", () => {
+        const error = {
+            message:
+                "This endpoint's maximum context length is 1000 tokens. However, you requested about 65000 tokens (64000 of text input, 1000 in the output).",
+        }
+        expect(parseOutputTokenLimit(error)).toBeNull()
+    })
+})
+
+describe("resolveMaxOutputTokens", () => {
+    it("uses a valid header value", () => {
+        expect(resolveMaxOutputTokens("32000")).toBe(32000)
+    })
+
+    it("falls back to the default for missing or bogus values", () => {
+        expect(resolveMaxOutputTokens(null)).toBe(DEFAULT_MAX_OUTPUT_TOKENS)
+        expect(resolveMaxOutputTokens("")).toBe(DEFAULT_MAX_OUTPUT_TOKENS)
+        expect(resolveMaxOutputTokens("abc")).toBe(DEFAULT_MAX_OUTPUT_TOKENS)
+        expect(resolveMaxOutputTokens("0")).toBe(DEFAULT_MAX_OUTPUT_TOKENS)
+        expect(resolveMaxOutputTokens("-5")).toBe(DEFAULT_MAX_OUTPUT_TOKENS)
+        expect(resolveMaxOutputTokens("1.5")).toBe(DEFAULT_MAX_OUTPUT_TOKENS)
+        // Above the sanity ceiling, e.g. an extra zero
+        expect(resolveMaxOutputTokens("640000")).toBe(DEFAULT_MAX_OUTPUT_TOKENS)
+    })
+
+    it("uses the env value when no header is sent, and validates it too", () => {
+        const original = process.env.MAX_OUTPUT_TOKENS
+        try {
+            process.env.MAX_OUTPUT_TOKENS = "24000"
+            expect(resolveMaxOutputTokens(null)).toBe(24000)
+            // Header still wins
+            expect(resolveMaxOutputTokens("8000")).toBe(8000)
+
+            process.env.MAX_OUTPUT_TOKENS = "-1"
+            expect(resolveMaxOutputTokens(null)).toBe(DEFAULT_MAX_OUTPUT_TOKENS)
+        } finally {
+            if (original === undefined) delete process.env.MAX_OUTPUT_TOKENS
+            else process.env.MAX_OUTPUT_TOKENS = original
+        }
+    })
+})
+
+/** Minimal stand-in for a v3 language model that records what it was asked for. */
+function fakeModel(
+    behaviors: Array<() => Promise<unknown>>,
+): [any, Array<Record<string, unknown>>] {
+    const calls: Array<Record<string, unknown>> = []
+    let index = 0
+    const model = {
+        specificationVersion: "v3" as const,
+        provider: "test",
+        modelId: "test-model",
+        supportedUrls: {},
+        doGenerate: async () => {
+            throw new Error("not used")
+        },
+        doStream: async (options: Record<string, unknown>) => {
+            calls.push(options)
+            const behavior = behaviors[index] ?? behaviors[behaviors.length - 1]
+            index++
+            return behavior()
+        },
+    }
+    return [model, calls]
+}
+
+const STREAM_OK = { stream: new ReadableStream() }
+
+describe("withOutputTokenLimitFallback", () => {
+    it("retries once with the ceiling named in the rejection", async () => {
+        const [model, calls] = fakeModel([
+            () =>
+                Promise.reject(
+                    Object.assign(
+                        new Error("exceeds the model limit of 4096"),
+                        { statusCode: 400 },
+                    ),
+                ),
+            () => Promise.resolve(STREAM_OK),
+        ])
+
+        const wrapped = withOutputTokenLimitFallback(model)
+        await wrapped.doStream({ prompt: [], maxOutputTokens: 64000 } as any)
+
+        expect(calls.map((c) => c.maxOutputTokens)).toEqual([64000, 4096])
+    })
+
+    it("does not retry an error it cannot attribute to the budget", async () => {
+        const [model, calls] = fakeModel([
+            () =>
+                Promise.reject(
+                    Object.assign(new Error("Invalid API key"), {
+                        statusCode: 401,
+                    }),
+                ),
+        ])
+
+        const wrapped = withOutputTokenLimitFallback(model)
+        await expect(
+            wrapped.doStream({ prompt: [], maxOutputTokens: 64000 } as any),
+        ).rejects.toThrow("Invalid API key")
+
+        expect(calls).toHaveLength(1)
+    })
+
+    it("does not retry when the ceiling is not actually smaller", async () => {
+        const [model, calls] = fakeModel([
+            () =>
+                Promise.reject(
+                    Object.assign(
+                        new Error("exceeds the model limit of 64000"),
+                        { statusCode: 400 },
+                    ),
+                ),
+        ])
+
+        const wrapped = withOutputTokenLimitFallback(model)
+        await expect(
+            wrapped.doStream({ prompt: [], maxOutputTokens: 64000 } as any),
+        ).rejects.toThrow()
+
+        expect(calls).toHaveLength(1)
+    })
+
+    it("retries at most once, so a second rejection propagates", async () => {
+        const [model, calls] = fakeModel([
+            () =>
+                Promise.reject(
+                    Object.assign(
+                        new Error("exceeds the model limit of 4096"),
+                        { statusCode: 400 },
+                    ),
+                ),
+            () =>
+                Promise.reject(
+                    Object.assign(
+                        new Error("exceeds the model limit of 2048"),
+                        { statusCode: 400 },
+                    ),
+                ),
+        ])
+
+        const wrapped = withOutputTokenLimitFallback(model)
+        await expect(
+            wrapped.doStream({ prompt: [], maxOutputTokens: 64000 } as any),
+        ).rejects.toThrow("model limit of 2048")
+
+        expect(calls).toHaveLength(2)
+    })
+
+    it("keeps the other call options when retrying", async () => {
+        const [model, calls] = fakeModel([
+            () =>
+                Promise.reject(
+                    Object.assign(
+                        new Error("exceeds the model limit of 4096"),
+                        { statusCode: 400 },
+                    ),
+                ),
+            () => Promise.resolve(STREAM_OK),
+        ])
+
+        const wrapped = withOutputTokenLimitFallback(model)
+        await wrapped.doStream({
+            prompt: [],
+            maxOutputTokens: 64000,
+            temperature: 0.4,
+            providerOptions: {
+                bedrock: { reasoningConfig: { type: "enabled" } },
+            },
+        } as any)
+
+        expect(calls[1].temperature).toBe(0.4)
+        expect(calls[1].providerOptions).toEqual({
+            bedrock: { reasoningConfig: { type: "enabled" } },
+        })
+    })
+})

vercel.json

@@ -2,7 +2,7 @@
     "functions": {
         "app/api/chat/route.ts": {
             "memory": 512,
-            "maxDuration": 120
+            "maxDuration": 300
         },
         "app/api/**/route.ts": {
             "memory": 256,