From 2644927f4fcbad3819044a8e9aa78896b18d28b8 Mon Sep 17 00:00:00 2001 From: LamzQ Date: Tue, 1 Sep 2026 21:59:22 +0800 Subject: [PATCH 1/2] fix(llm): double max_tokens on finish_reason=length before retrying JSON Thinking models share one max_tokens budget between reasoning and the final answer, so a truncated completion arrives as 200 OK with finish_reason="length". completeJson treated it like any malformed output and retried with the same max_tokens, which truncates again - the retry could never succeed. On the first truncation, log usage diagnostics (max_tokens, completion chars, token usage) and double max_tokens for the next attempt, capped at 32768. Malformed-only retries (finish_reason=stop) keep the budget unchanged. Covered by tests/unit/llm/length-retry-budget.test.ts: doubling from the default, the 32768 cap, at-most-once doubling across consecutive truncations, and no change when finish_reason is stop. --- apps/memos-local-plugin/core/llm/client.ts | 34 +++++- .../unit/llm/length-retry-budget.test.ts | 115 ++++++++++++++++++ 2 files changed, 148 insertions(+), 1 deletion(-) create mode 100644 apps/memos-local-plugin/tests/unit/llm/length-retry-budget.test.ts diff --git a/apps/memos-local-plugin/core/llm/client.ts b/apps/memos-local-plugin/core/llm/client.ts index ee456ac12..0516a1546 100644 --- a/apps/memos-local-plugin/core/llm/client.ts +++ b/apps/memos-local-plugin/core/llm/client.ts @@ -501,16 +501,48 @@ export function createLlmClientWithProvider( const messages = normalizeMessages(input); const systemHint = buildJsonSystemHint(opts.schemaHint); const msgs = ensureJsonWordInUserMessage(inject(messages, systemHint)); - const call = buildCallInput(opts, true); + let call = buildCallInput(opts, true); const op = opts.op ?? "complete.json"; const maxMalformedRetries = Math.max(0, opts.malformedRetries ?? 1); let attempt = 0; let lastRaw = ""; let lastErr: unknown = null; + // Thinking models share one max_tokens budget between reasoning and the + // final answer, so a truncated response arrives as 200 OK with + // finish_reason="length". Without an explicit check it looks like plain + // malformed JSON and the retry re-sends the same doomed budget. On the + // first truncation, log diagnostics and double max_tokens for the next + // attempt (capped). + const LENGTH_RETRY_MAX_TOKENS_CEILING = 32_768; + let truncatedOnce = false; while (attempt <= maxMalformedRetries) { attempt++; const { completion } = await callWithFallback(msgs, call, opts, op); + if (completion.finishReason === "length") { + jsonLog.warn("max_tokens_truncated", { + op, + attempt, + maxTokens: call.maxTokens, + completionChars: completion.text.length, + usage: completion.usage + ? { + completionTokens: completion.usage.completionTokens, + totalTokens: completion.usage.totalTokens, + } + : undefined, + }); + if (!truncatedOnce && (call.maxTokens ?? 0) < LENGTH_RETRY_MAX_TOKENS_CEILING) { + truncatedOnce = true; + call = { + ...call, + maxTokens: Math.min( + LENGTH_RETRY_MAX_TOKENS_CEILING, + Math.max(2 * (call.maxTokens ?? DEFAULT_MAX_TOKENS), DEFAULT_MAX_TOKENS), + ), + }; + } + } lastRaw = completion.text; try { const parsed = opts.parse diff --git a/apps/memos-local-plugin/tests/unit/llm/length-retry-budget.test.ts b/apps/memos-local-plugin/tests/unit/llm/length-retry-budget.test.ts new file mode 100644 index 000000000..870bf3e7e --- /dev/null +++ b/apps/memos-local-plugin/tests/unit/llm/length-retry-budget.test.ts @@ -0,0 +1,115 @@ +/** + * Truncation retry budget: finish_reason="length" is a budget problem, not + * plain malformed output. The retry must double max_tokens (once, capped) + * instead of re-sending the same doomed budget. + */ +import { beforeAll, describe, expect, it } from "vitest"; + +import { MemosError } from "../../../agent-contract/errors.js"; +import { createLlmClientWithProvider } from "../../../core/llm/index.js"; +import { initTestLogger } from "../../../core/logger/index.js"; +import type { + LlmConfig, + LlmMessage, + LlmProvider, + LlmProviderCtx, + LlmProviderName, + ProviderCallInput, + ProviderCompletion, +} from "../../../core/llm/types.js"; + +beforeAll(async () => { + await initTestLogger(); +}); + +function cfg(partial: Partial = {}): LlmConfig { + return { + provider: "openai_compatible", + model: "gpt-test", + endpoint: "", + apiKey: "X", + temperature: 0.3, + fallbackToHost: false, + timeoutMs: 5_000, + maxRetries: 0, + ...partial, + }; +} + +class StubProvider implements LlmProvider { + public inputs: ProviderCallInput[] = []; + public readonly name: LlmProviderName; + constructor( + name: LlmProviderName, + private readonly responder: (n: number) => ProviderCompletion, + ) { + this.name = name; + } + async complete( + _messages: LlmMessage[], + opts: ProviderCallInput, + _ctx: LlmProviderCtx, + ): Promise { + this.inputs.push(opts); + return this.responder(this.inputs.length); + } +} + +const DEFAULT_MAX_TOKENS = 1024; + +describe("completeJson truncation retry budget", () => { + it("doubles max_tokens after finish_reason=length and parses the retry", async () => { + const stub = new StubProvider("openai_compatible", (n) => + n === 1 + ? { text: "{\"x\":", durationMs: 1, finishReason: "length" as const } + : { text: "{\"x\":1}", durationMs: 1, finishReason: "stop" as const }, + ); + const client = createLlmClientWithProvider(cfg(), stub); + const r = await client.completeJson<{ x: number }>("ask", { malformedRetries: 1 }); + expect(r.value.x).toBe(1); + expect(stub.inputs).toHaveLength(2); + expect(stub.inputs[0]!.maxTokens).toBe(DEFAULT_MAX_TOKENS); + expect(stub.inputs[1]!.maxTokens).toBe(2 * DEFAULT_MAX_TOKENS); + }); + + it("caps the doubled budget at 32768", async () => { + const stub = new StubProvider("openai_compatible", (n) => + n === 1 + ? { text: "truncated", durationMs: 1, finishReason: "length" as const } + : { text: "{\"y\":2}", durationMs: 1, finishReason: "stop" as const }, + ); + const client = createLlmClientWithProvider(cfg({ maxTokens: 16_384 }), stub); + const r = await client.completeJson<{ y: number }>("ask", { malformedRetries: 1 }); + expect(r.value.y).toBe(2); + expect(stub.inputs[0]!.maxTokens).toBe(16_384); + expect(stub.inputs[1]!.maxTokens).toBe(32_768); + }); + + it("doubles at most once across consecutive truncations", async () => { + const stub = new StubProvider("openai_compatible", () => ({ + text: "truncated", + durationMs: 1, + finishReason: "length" as const, + })); + const client = createLlmClientWithProvider(cfg({ maxTokens: 16_384 }), stub); + await expect( + client.completeJson("ask", { malformedRetries: 2 }), + ).rejects.toBeInstanceOf(MemosError); + expect(stub.inputs).toHaveLength(3); + expect(stub.inputs[0]!.maxTokens).toBe(16_384); + expect(stub.inputs[1]!.maxTokens).toBe(32_768); + expect(stub.inputs[2]!.maxTokens).toBe(32_768); // no further doubling + }); + + it("keeps the budget unchanged when finish_reason is stop", async () => { + const stub = new StubProvider("openai_compatible", (n) => + n === 1 + ? { text: "not json", durationMs: 1, finishReason: "stop" as const } + : { text: "{\"z\":3}", durationMs: 1, finishReason: "stop" as const }, + ); + const client = createLlmClientWithProvider(cfg(), stub); + const r = await client.completeJson<{ z: number }>("ask", { malformedRetries: 1 }); + expect(r.value.z).toBe(3); + expect(stub.inputs[1]!.maxTokens).toBe(stub.inputs[0]!.maxTokens); + }); +}); From bd9bd176bbbd43b45c2a293734e8d7ae2f06d6f0 Mon Sep 17 00:00:00 2001 From: LamzQ Date: Wed, 2 Sep 2026 09:01:58 +0800 Subject: [PATCH 2/2] =?UTF-8?q?fix(llm):=20address=20review=20=E2=80=94=20?= =?UTF-8?q?guarantee=20the=20truncation=20retry=20one=20attempt=20when=20m?= =?UTF-8?q?alformedRetries=20is=200,=20hoist=20ceiling=20constant=20to=20m?= =?UTF-8?q?odule=20scope?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - completeJson: the upgraded budget was a silent no-op for callers that pass malformedRetries: 0 (fail-fast parsing) — the while-loop exited before the second attempt could run. The truncation upgrade now grants exactly one additional attempt, independent of the malformed-parse budget. - LENGTH_RETRY_MAX_TOKENS_CEILING hoisted to module scope alongside DEFAULT_MAX_TOKENS (review: re-declared per invocation otherwise). - tests: +1 case locking the malformedRetries: 0 truncation retry. Co-Authored-By: LamzQ --- apps/memos-local-plugin/core/llm/client.ts | 9 +++++++-- .../tests/unit/llm/length-retry-budget.test.ts | 13 +++++++++++++ 2 files changed, 20 insertions(+), 2 deletions(-) diff --git a/apps/memos-local-plugin/core/llm/client.ts b/apps/memos-local-plugin/core/llm/client.ts index 0516a1546..f44e0cb25 100644 --- a/apps/memos-local-plugin/core/llm/client.ts +++ b/apps/memos-local-plugin/core/llm/client.ts @@ -45,6 +45,8 @@ import type { } from "./types.js"; const DEFAULT_MAX_TOKENS = 1024; +// Upper bound for the one-shot truncation retry budget (see completeJson). +const LENGTH_RETRY_MAX_TOKENS_CEILING = 32_768; // ─── Factory ───────────────────────────────────────────────────────────────── @@ -503,7 +505,7 @@ export function createLlmClientWithProvider( const msgs = ensureJsonWordInUserMessage(inject(messages, systemHint)); let call = buildCallInput(opts, true); const op = opts.op ?? "complete.json"; - const maxMalformedRetries = Math.max(0, opts.malformedRetries ?? 1); + let maxMalformedRetries = Math.max(0, opts.malformedRetries ?? 1); let attempt = 0; let lastRaw = ""; let lastErr: unknown = null; @@ -513,7 +515,6 @@ export function createLlmClientWithProvider( // malformed JSON and the retry re-sends the same doomed budget. On the // first truncation, log diagnostics and double max_tokens for the next // attempt (capped). - const LENGTH_RETRY_MAX_TOKENS_CEILING = 32_768; let truncatedOnce = false; while (attempt <= maxMalformedRetries) { @@ -534,6 +535,10 @@ export function createLlmClientWithProvider( }); if (!truncatedOnce && (call.maxTokens ?? 0) < LENGTH_RETRY_MAX_TOKENS_CEILING) { truncatedOnce = true; + // The upgraded budget is useless without at least one more attempt: + // a caller may pass malformedRetries: 0 for fail-fast parsing, and + // the truncation retry must not be silently skipped then. + if (attempt > maxMalformedRetries) maxMalformedRetries = attempt; call = { ...call, maxTokens: Math.min( diff --git a/apps/memos-local-plugin/tests/unit/llm/length-retry-budget.test.ts b/apps/memos-local-plugin/tests/unit/llm/length-retry-budget.test.ts index 870bf3e7e..70a7060e6 100644 --- a/apps/memos-local-plugin/tests/unit/llm/length-retry-budget.test.ts +++ b/apps/memos-local-plugin/tests/unit/llm/length-retry-budget.test.ts @@ -72,6 +72,19 @@ describe("completeJson truncation retry budget", () => { expect(stub.inputs[1]!.maxTokens).toBe(2 * DEFAULT_MAX_TOKENS); }); + it("still retries once on truncation when malformedRetries is 0", async () => { + const stub = new StubProvider("openai_compatible", (n) => + n === 1 + ? { text: "{\"q\":", durationMs: 1, finishReason: "length" as const } + : { text: "{\"q\":2}", durationMs: 1, finishReason: "stop" as const }, + ); + const client = createLlmClientWithProvider(cfg(), stub); + const r = await client.completeJson<{ q: number }>("ask", { malformedRetries: 0 }); + expect(r.value.q).toBe(2); + expect(stub.inputs).toHaveLength(2); + expect(stub.inputs[1]!.maxTokens).toBe(2 * DEFAULT_MAX_TOKENS); + }); + it("caps the doubled budget at 32768", async () => { const stub = new StubProvider("openai_compatible", (n) => n === 1