mirror of
https://github.com/DayuanJiang/next-ai-draw-io.git
synced 2026-09-02 01:20:23 +08:00
aa276ca9b0
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.
272 lines
9.5 KiB
TypeScript
272 lines
9.5 KiB
TypeScript
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" } },
|
|
})
|
|
})
|
|
})
|