Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
32 changes: 32 additions & 0 deletions apps/memos-local-plugin/core/llm/client.ts
Original file line number Diff line number Diff line change
Expand Up @@ -507,11 +507,13 @@ export function createLlmClientWithProvider(
let attempt = 0;
let lastRaw = "";
let lastErr: unknown = null;
let lastCompletion: LlmCompletion | null = null;

while (attempt <= maxMalformedRetries) {
attempt++;
const { completion } = await callWithFallback(msgs, call, opts, op);
lastRaw = completion.text;
lastCompletion = completion;
try {
const parsed = opts.parse
? opts.parse(completion.text)
Expand All @@ -534,6 +536,35 @@ export function createLlmClientWithProvider(
};
} catch (err) {
lastErr = err;
// Thinking models share `max_tokens` between reasoning and answer.
// When the budget is exhausted, providers return 200 OK with
// `finish_reason="length"` and an unfinished JSON body. Parsing
// fails, but retrying at the same `maxTokens` will truncate at the
// same offset — it is a budget problem, not a formatting one. Fail
// fast so callers see `truncated:true` and can raise `maxTokens`
// instead of burning another paid request.
if (completion.finishReason === "length") {
jsonLog.warn("truncated", {
op,
attempt,
maxTokens: call.maxTokens,
finishReason: completion.finishReason,
usage: completion.usage,
});
throw new MemosError(
ERROR_CODES.LLM_OUTPUT_MALFORMED,
"LLM JSON truncated: response hit max_tokens (finish_reason=length); raise maxTokens",
{
provider: provider.name,
op,
truncated: true,
finishReason: completion.finishReason,
maxTokens: call.maxTokens,
usage: completion.usage,
rawPreview: completion.text.slice(0, 512),
},
);
}
jsonLog.warn("malformed", {
op,
attempt,
Expand All @@ -552,6 +583,7 @@ export function createLlmClientWithProvider(
provider: provider.name,
op,
rawPreview: lastRaw.slice(0, 512),
finishReason: lastCompletion?.finishReason,
});
}

Expand Down
48 changes: 48 additions & 0 deletions apps/memos-local-plugin/tests/unit/llm/client.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -165,6 +165,54 @@ describe("llm/client", () => {
expect(fake.invocations).toBe(2);
});

it("completeJson fails fast on finish_reason=length instead of retrying at same budget", async () => {
// Thinking models share max_tokens between reasoning and answer, so a
// truncated completion arrives as 200 OK with finish_reason="length"
// and an unparseable JSON body. Retrying at the same budget truncates
// again — the retry should be skipped.
const fake = new FakeProvider("openai_compatible", () => ({
text: '{"partial":',
finishReason: "length",
durationMs: 1,
usage: { promptTokens: 10, completionTokens: 512, totalTokens: 522 },
}));
const client = createLlmClientWithProvider(cfg(), fake);

try {
await client.completeJson("ask", { malformedRetries: 3, maxTokens: 512 });
throw new Error("should have thrown");
} catch (err) {
expect(err).toBeInstanceOf(MemosError);
const memos = err as MemosError;
expect(memos.code).toBe(ERROR_CODES.LLM_OUTPUT_MALFORMED);
expect(memos.details).toMatchObject({
truncated: true,
finishReason: "length",
maxTokens: 512,
});
expect(memos.message).toMatch(/truncat|max_tokens|length/i);
}
// Key assertion: exactly one provider call, not one + retries.
expect(fake.invocations).toBe(1);
expect(client.stats().retries).toBe(0);
});

it("completeJson still retries on malformed output that is not a length truncation", async () => {
// Regression guard: the length short-circuit must not swallow genuine
// malformed outputs (finishReason=stop|undefined), which are still
// worth retrying — a resample may produce valid JSON.
const fake = new FakeProvider("openai_compatible", (n) => ({
text: n === 1 ? "not json" : '{"ok":1}',
finishReason: "stop",
durationMs: 1,
}));
const client = createLlmClientWithProvider(cfg(), fake);
const r = await client.completeJson<{ ok: number }>("ask", { malformedRetries: 1 });
expect(r.value.ok).toBe(1);
expect(fake.invocations).toBe(2);
expect(client.stats().retries).toBe(1);
});

it("stream passes provider-native chunks through", async () => {
const client = createLlmClientWithProvider(cfg(), new StreamingProvider());
const chunks: LlmStreamChunk[] = [];
Expand Down
Loading