diff --git a/lib/output-token-limit.ts b/lib/output-token-limit.ts index 9b37f17..bd4a987 100644 --- a/lib/output-token-limit.ts +++ b/lib/output-token-limit.ts @@ -16,6 +16,19 @@ 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 @@ -28,9 +41,26 @@ export const MAX_OUTPUT_TOKENS_LIMIT = 200000 * 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 } + 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 : "", @@ -43,18 +73,17 @@ export function parseOutputTokenLimit(error: unknown): number | null { 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 + 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) || - text.match(/lower than (\d+)/i) + text.match(/at most (\d+) completion tokens/i) - return output ? Number(output[1]) : null + return output ? usableLimit(Number(output[1])) : null } /** @@ -92,19 +121,25 @@ export function withOutputTokenLimitFallback( }) } +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. + * 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 { - 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 + return ( + validBudget(headerValue) ?? + validBudget(process.env.MAX_OUTPUT_TOKENS) ?? + DEFAULT_MAX_OUTPUT_TOKENS + ) } diff --git a/tests/unit/output-token-limit.test.ts b/tests/unit/output-token-limit.test.ts index 541565a..d549f44 100644 --- a/tests/unit/output-token-limit.test.ts +++ b/tests/unit/output-token-limit.test.ts @@ -3,6 +3,7 @@ import { DEFAULT_MAX_OUTPUT_TOKENS, parseOutputTokenLimit, resolveMaxOutputTokens, + withOutputTokenLimitFallback, } from "@/lib/output-token-limit" describe("parseOutputTokenLimit", () => { @@ -52,6 +53,44 @@ describe("parseOutputTokenLimit", () => { 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: @@ -76,4 +115,157 @@ describe("resolveMaxOutputTokens", () => { // 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" } }, + }) + }) })