[solderzzc]

Fixes the remaining cases in #245 (a tool-call-only turn was fixed by #250).

Problem

Since the mlx-swift-lm b620ae8 bump the prompt cache is saved when the producer emits .info. A request that ends on a stop sequence or a client disconnect leaves the stream before .info, so nothing is saved and the next turn re-prefills the whole history.

Reproduced on current main with Qwen2.5-0.5B-Instruct-4bit (greedy, ~690-token system prompt, follow-up repeats turn 1 plus a short exchange). Tokens reused by the follow-up:

turn 1 ended on	main	this PR
the model's own end of turn	673/736	672/736
max_tokens	671/690	670/690
a custom stop sequence	4/692	672/692
client disconnect	4/688	671/688

The built-in stop strings (<|im_end|> etc.) do not cause it: the EOS token arrives first, so ordinary turns were fine. Only requests that set stop or leave early were affected.

Fix

For the full-attention path (not hybrid, rotating, draft model, MTP or TurboKV), prefill the prompt except its last token synchronously, save, then let generation feed the last token. This is what the hybrid and rotating paths already do. No decode step has started when the cache is read, so the race that moved the save to .info cannot happen, and the save no longer depends on how the request ends.

The saved entry is one token shorter than the prompt; restore already replays the last token on a full match. Turn-1 output text is identical to main for all four scenarios.

Tests

• tests/test-contract.sh against this build, Qwen2.5-0.5B-Instruct-4bit: 26 passed, 0 failed, 2 skipped. New Test 16 covers a stop-sequence turn (511 of 530 tokens reused) and a client-disconnect turn (509 of 528).
• Same suite against main: Test 16 fails (3 of 530 and 4 of 528 reused).
• Test 15 now accepts a reuse of prompt tokens minus one, because of the shorter entry.
• Not run locally: the XCTest target (the changed code is in startGeneration, which has no unit test) and the other integration legs; CI covers them.

Not covered

• Hybrid (Qwen3.5/3.6) and rotating (Gemma 4 sliding window) paths already saved at their boundary and are unchanged. MTP, draft-model and TurboKV paths still save at .info and keep the old early-exit behaviour.
• /v1/completions does not use the prompt cache at all, so it is unaffected.

🤖 Generated with Claude Code