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 + +/** + * 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 + ) +} 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..d549f44 --- /dev/null +++ b/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>, +): [any, Array>] { + const calls: Array> = [] + 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) => { + 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" } }, + }) + }) +}) 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,