fix(llm): double max_tokens on finish_reason=length before retrying JSON

[smoryan]

Description

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" — not as an error. completeJson() treats it like
any malformed output and retries with the same budget, which truncates
again: the retry can never succeed.

On our local deployment this triggered routinely on L3-abstraction calls
(P95 useful output ≈ 4.2k tokens, with reasoning frequently pushing past the
configured budget), burning both attempts of completeJson (the initial call plus its one malformed-retry) and producing
LLM_OUTPUT_MALFORMED for what is a budget problem, not a formatting problem.

This PR detects the truncation and doubles max_tokens once for the retry,
capped at 32768.

Related Issue: Fixes #2320

Change

• core/llm/client.ts — completeJson():
  • after callWithFallback(), if completion.finishReason === "length",
    log max_tokens_truncated diagnostics (op, attempt, maxTokens,
    completion chars, usage.completionTokens / usage.totalTokens; the
    usage block is nullable — not every provider normalizes it);
  • on the first truncation only (truncatedOnce guard), still granted when the caller passes malformedRetries: 0 (fail-fast), rebuild the call
    with maxTokens = min(32768, max(2 × maxTokens, DEFAULT_MAX_TOKENS)).
    The detection sits before JSON parsing so the diagnosis does not
    depend on the parse failing, and the once-only guard keeps retries from
    ratcheting cost upward. If the configured budget is already ≥ 32768
    (the cap), the budget is left untouched — the guard never lowers an
    explicit high budget a caller set via llm.maxTokens (fix(plugin): declare llm.maxTokens and llm.headers as first-class config keys #2248).
  • Malformed-only retries (finish_reason="stop") keep the budget
    unchanged.
• tests/unit/llm/length-retry-budget.test.ts — new suite (5 tests) driving the
  real retry loop through a stub provider:
  doubling from the default (1024 → 2048) and the retry parses;
  cap at 32768 from 16384;
  at-most-once doubling across consecutive truncations (budget stays at
  32768, ultimately rejects with MemosError);
  budget unchanged when the only problem is malformed output
  (finish_reason="stop").

Tests

• npx vitest run tests/unit/llm/length-retry-budget.test.ts tests/unit/llm/client.test.ts → 30 passed (5 new + 25 existing)
• npx vitest run tests/unit/llm/ → 86 passed (6 files)
• npx tsc -p tsconfig.json --noEmit → clean (exit 0)

Reproduce: clone main, apply this PR, cd apps/memos-local-plugin,
npm install, run the commands above.

Related

• Builds on fix(plugin): declare llm.maxTokens and llm.headers as first-class config keys #2248, which made llm.maxTokens a first-class config key —
  this covers the retry behavior that PR did not address: with
  finish_reason="length", raising the configured budget up front helps,
  but a completion that still runs out mid-JSON must not burn its retries
  at the same doomed budget.

Type of change

• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Refactor (does not change functionality, e.g. code style improvements, linting)
• Documentation update

How Has This Been Tested?

• Unit Test — npx vitest run tests/unit/llm/ → 86 passed (6 files), incl. 5 new truncation-budget tests
• Test Script Or Test Steps — see reproduce steps above
• Pipeline Automated API Test

Checklist

• I have performed a self-review of my own code
• I have commented my code in hard-to-understand areas
• I have added tests that prove my fix is effective or that my feature works
• I have created related documentation issue/PR in MemOS-Docs (if applicable) — no public-API/config surface changed
• I have linked the issue to this PR (if applicable) — pending issue; will link once filed
• I have mentioned the person who will review this PR

Environment

• Plugin version: monorepo main (28dfb4e)
• Runtime: Node 26 (arm64 macOS), vitest 2.1.9

[Memtensor-AI]

🤖 Open Code Review

Target: PR #2321
Task: 0bb1cb61401ceb98
Base: main
Head: fix/llm-length-retry-budget

🔍 OpenCodeReview found 1 issue(s) in this PR.

---

1. apps/memos-local-plugin/core/llm/client.ts (L536-L551)

After widening maxTokens on a truncation retry, execution falls through and unconditionally attempts to parse the truncated (incomplete) JSON. This parse will always throw, emitting a spurious "malformed" warning in the log even though the failure is fully expected and already diagnosed by "max_tokens_truncated". It also burns one retry slot unnecessarily.

Add a continue immediately after the budget is updated (when the retry will actually fire) so the loop restarts with the enlarged budget without the false-positive "malformed" log entry. When the budget is already at the ceiling (truncatedOnce is already true), falling through and letting the parse fail is correct — that path should remain unchanged.

💡 Suggested Change

Before:

        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(
              LENGTH_RETRY_MAX_TOKENS_CEILING,
              Math.max(2 * (call.maxTokens ?? DEFAULT_MAX_TOKENS), DEFAULT_MAX_TOKENS),
            ),
          };
        }
      }
      lastRaw = completion.text;

After:

        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(
              LENGTH_RETRY_MAX_TOKENS_CEILING,
              Math.max(2 * (call.maxTokens ?? DEFAULT_MAX_TOKENS), DEFAULT_MAX_TOKENS),
            ),
          };
          // Skip parsing the known-truncated response; the next iteration
          // will use the widened budget and produce a complete response.
          continue;
        }
      }
      lastRaw = completion.text;

Generated by cloud-assistant via Open Code Review.

[smoryan]

Thanks @Memtensor-AI — both findings addressed in bd9bd17:

1. malformedRetries: 0 silent no-op — confirmed and fixed. The truncation upgrade now grants exactly one additional attempt (if (attempt > maxMalformedRetries) maxMalformedRetries = attempt), so the upgraded budget is always exercised at least once, even for fail-fast callers. Locked by a new test: still retries once on truncation when malformedRetries is 0.
2. Constant hoisting — done. LENGTH_RETRY_MAX_TOKENS_CEILING now lives at module scope next to DEFAULT_MAX_TOKENS.

llm domain: 86 passed (6 files), tsc clean.

[Memtensor-AI]

✅ Automated Test Results: PASSED

All tests passed (5/5 executed). memos_local_plugin/unit: 5/5. Duration: 3s [advisory, non-gating] AI-generated tests on branch test/auto-gen-0bb1cb61401ceb98-20260902092149: 23/23 passed — these do NOT affect the PR verdict; review the branch manually.

Branch: fix/llm-length-retry-budget