From 8fb9ef20bd8e343e41a8073f1d7ed4260935cf05 Mon Sep 17 00:00:00 2001 From: "dayuan.jiang" Date: Sat, 22 Aug 2026 11:45:09 +0900 Subject: [PATCH] fix: raise the output budget so reasoning models reach the tool call A reasoning model spends the output budget in order: thinking first, then prose, then the tool call. With 16000 the thinking alone can consume all of it, so the turn ends with finishReason "length" before display_diagram is ever called. The canvas stays empty and nothing surfaces in the UI, because no tool call means no tool error, and the client never reads finishReason. Measured on openrouter deepseek/deepseek-v4-flash, the model from the report: - max_tokens=800 with reasoning on returns reasoning_tokens=800, empty content, finish_reason length. So reasoning is billed against this budget, not exempt. - refining an existing diagram (19k chars of XML in the input) produced 49142 chars of reasoning, zero tool calls, finishReason "length" at 16000 - the same request at 40000 finished and called edit_diagram with 12 operations 64000 cannot just be sent to every model: bedrock claude-3-haiku caps at 4096, nova-lite at 10000, and the openrouter deepseek-r1 endpoint counts input and output against one 64000 ceiling. All three name the real limit in the 400, so parse it and retry once. Verified: nova-lite logs "64000 rejected, retrying with 10000" and then completes its tool call. Also expose the budget in Settings. It is sent as a header rather than read from env only, so desktop users can raise it themselves without an env file. vercel.json goes back to the 300s it had before #238 traded it for $2-4/month. That is now Vercel's own default, and billing pauses while the function waits on the model, so the saving that motivated 120s no longer applies. edgeone.json is left alone: its 120 may be that platform's actual ceiling. --- app/api/chat/route.ts | 25 ++++-- components/chat-panel.tsx | 21 +++++ components/settings-dialog.tsx | 22 ++++++ env.example | 6 +- lib/i18n/dictionaries/en.json | 2 + lib/i18n/dictionaries/ja.json | 2 + lib/i18n/dictionaries/zh-Hant.json | 2 + lib/i18n/dictionaries/zh.json | 2 + lib/output-token-limit.ts | 110 ++++++++++++++++++++++++++ lib/storage.ts | 3 + tests/e2e/settings.spec.ts | 20 +++++ tests/unit/output-token-limit.test.ts | 79 ++++++++++++++++++ vercel.json | 2 +- 13 files changed, 288 insertions(+), 8 deletions(-) create mode 100644 lib/output-token-limit.ts create mode 100644 tests/unit/output-token-limit.test.ts diff --git a/app/api/chat/route.ts b/app/api/chat/route.ts index ca9f723..ec404eb 100644 --- a/app/api/chat/route.ts +++ b/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 { // 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 — - // enough for a small diagram, so larger ones were cut off mid-attribute. - maxOutputTokens: Number(process.env.MAX_OUTPUT_TOKENS) || 16000, + // Must be sent: unset means the provider's own default, and Bedrock's is + // 4096, enough for a small diagram, so larger ones were cut off mid-attribute. + maxOutputTokens, stopWhen: stepCountIs(5), // Repair truncated tool calls when maxOutputTokens is reached mid-JSON experimental_repairToolCall: async ({ toolCall, error }) => { diff --git a/components/chat-panel.tsx b/components/chat-panel.tsx index 792a6ac..66f44b6 100644 --- a/components/chat-panel.tsx +++ b/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(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)} /> diff --git a/components/settings-dialog.tsx b/components/settings-dialog.tsx index d3dccc4..4d89e6d 100644 --- a/components/settings-dialog.tsx +++ b/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({ /> + {/* Max Output Tokens */} + + + onMaxOutputTokensChange(e.target.value) + } + placeholder="64000" + className="h-9 w-28 text-sm" + /> + + {/* Send Shortcut */} + +/** + * 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 + +/** + * 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" + */ +export function parseOutputTokenLimit(error: unknown): number | null { + const err = error as { message?: unknown; responseBody?: unknown } + 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) + const budget = + Number(context[1]) - (input ? Number(input[1]) : 0) - 1024 + return budget > 0 ? budget : null + } + + 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) || + text.match(/lower than (\d+)/i) + + return output ? 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, + }) + } + }, + }, + }) +} + +/** + * Resolve the output budget: user setting (sent as a header so it works in the + * desktop app too), then server env, then the default. + */ +export function resolveMaxOutputTokens(headerValue: string | null): number { + const fromHeader = Number(headerValue) + if ( + Number.isInteger(fromHeader) && + fromHeader > 0 && + fromHeader <= MAX_OUTPUT_TOKENS_LIMIT + ) { + return fromHeader + } + + return Number(process.env.MAX_OUTPUT_TOKENS) || DEFAULT_MAX_OUTPUT_TOKENS +} diff --git a/lib/storage.ts b/lib/storage.ts index ddfc54c..2f61181 100644 --- a/lib/storage.ts +++ b/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", diff --git a/tests/e2e/settings.spec.ts b/tests/e2e/settings.spec.ts index ad93515..c5c9b83 100644 --- a/tests/e2e/settings.spec.ts +++ b/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) diff --git a/tests/unit/output-token-limit.test.ts b/tests/unit/output-token-limit.test.ts new file mode 100644 index 0000000..541565a --- /dev/null +++ b/tests/unit/output-token-limit.test.ts @@ -0,0 +1,79 @@ +import { describe, expect, it } from "vitest" +import { + DEFAULT_MAX_OUTPUT_TOKENS, + parseOutputTokenLimit, + resolveMaxOutputTokens, +} 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("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) + }) +}) diff --git a/vercel.json b/vercel.json index 42052bf..d2ebc80 100644 --- a/vercel.json +++ b/vercel.json @@ -2,7 +2,7 @@ "functions": { "app/api/chat/route.ts": { "memory": 512, - "maxDuration": 120 + "maxDuration": 300 }, "app/api/**/route.ts": { "memory": 256,