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,