mlxrunner: keep prefill snapshots when a request is cancelled mid-prompt

A long prompt records restore points during prefill, but they only
reached the prefix trie when the prefill completed; a cancelled request
closed and released everything it had captured. Agent clients routinely
cancel long prefills — their timeouts are shorter than the minutes a
40k-token prompt takes — so every retry started the whole prompt over
and never got further than the timeout allowed, which presents as the
model hanging forever.

Closing a session now attaches every snapshot the prefill crossed, so a
retry resumes from the last one and makes progress across timeouts.
Scenario tests cover retries resuming exactly where a cancelled attempt
stopped and cancellations on divergent conversation variants.

Fixes #17839

x/mlxrunner/prefix_cache.go

@@ -63,7 +63,7 @@ type cacheSession struct {
 
 	// pendingSnapshots lists offsets where snapshots should be captured
 	// during prefill, sorted by offset. Entries are scheduled on the caches
-	// before prefill and drained or discarded after.
+	// before prefill and drained when the captures are attached.
 	pendingSnapshots []pendingSnapshot
 }
 
@@ -326,29 +326,6 @@ func (s *cacheSession) schedulePrefillSnapshots(offsets []int) {
 	}
 }
 
-// discardPrefillSnapshots drains and closes the snapshots scheduled by
-// schedulePrefillSnapshots without attaching them to the trie, releasing their
-// pinned/lazy state. It is a no-op once attachPrefillSnapshots has drained the
-// schedule, so close can call it unconditionally to clean up an abandoned
-// prefill.
-func (s *cacheSession) discardPrefillSnapshots() {
-	if len(s.pendingSnapshots) == 0 {
-		return
-	}
-	s.pendingSnapshots = nil
-
-	for _, kv := range s.cache.caches {
-		if kv == nil {
-			continue
-		}
-		for _, snap := range kv.TakeSnapshots() {
-			if snap != nil {
-				snap.Close()
-			}
-		}
-	}
-}
-
 // attachPrefillSnapshots collects the snapshots captured during prefill and
 // attaches them to the trie, materializing a node at each requested offset.
 // Pending offsets are ascending and were scheduled in the same order, so the
@@ -538,13 +515,10 @@ func (c *prefixCache) minCacheOffset() int {
 
 // close saves the token state if the forward pass ran.
 func (s *cacheSession) close() {
-	// Release any prefill snapshots the session scheduled but never attached to
-	// the trie. A successful prefill drains them in attachPrefillSnapshots (so
-	// this is a no-op then); an abandoned one (e.g. cancellation between
-	// schedule and attach) leaves them in the caches, where the next request's
-	// PrepareSnapshots would overwrite the schedule without closing them,
-	// leaking the pinned/lazy snapshots and their VRAM.
-	s.discardPrefillSnapshots()
+	// A cancelled prefill never reaches the success-path attach; attaching
+	// here keeps its crossed captures for the retry and drains the schedule
+	// PrepareSnapshots would otherwise overwrite, leaking them.
+	s.attachPrefillSnapshots()
 
 	offset := s.cache.minCacheOffset()
 	if offset <= 0 {

x/mlxrunner/prefix_cache_scenario_test.go

@@ -265,3 +265,126 @@ func TestScenarioConversationTurns(t *testing.T) {
 	}
 	checkSnapshotCoverage(t, e.pc, e.kvLayers())
 }
+
+// runCancelledPrefill mirrors a prefill cancelled by a client timeout at
+// (roughly) cancelAt prompt tokens: begin, schedule, feed whole chunks until
+// cancelAt with the drafter's pairs lagging one token, settle the drafter with
+// the next prompt token, then close, which attaches the crossed captures.
+func (e *hybridEnv) runCancelledPrefill(t *testing.T, inputs []int32, chunk, cancelAt int) {
+	t.Helper()
+	pc := e.pc
+	session := pc.begin(inputs, nil)
+	session.schedulePrefillSnapshots(periodic(len(inputs)))
+
+	pos := pc.minCacheOffset()
+	for pos < cancelAt && pos < len(inputs)-1 {
+		n := min(chunk, len(inputs)-1-pos)
+		for _, c := range pc.caches {
+			if c == cache.Cache(e.draft) {
+				continue
+			}
+			if fc, ok := c.(feedableCache); ok {
+				fc.feed(inputs[pos : pos+n])
+			}
+		}
+		if d := e.draft.Offset(); d < pos+n-1 {
+			e.draft.feed(inputs[d : pos+n-1])
+		}
+		pos += n
+	}
+
+	if d := e.draft.Offset(); d < pos {
+		e.draft.feed(inputs[d:pos])
+	}
+	session.close()
+}
+
+// TestScenarioCancelledPrefills covers the cancellation invariants: a retry
+// of a cancelled prefill resumes exactly where the previous attempt stopped,
+// and the captures the cancelled attempt crossed become restore points.
+func TestScenarioCancelledPrefills(t *testing.T) {
+	logs := captureWarns(t)
+	e := newHybridEnv()
+	ts := &tokenStream{}
+	begins := &beginLog{t: t, logs: logs}
+
+	prompt := ts.fresh(47)
+
+	// Cancelled after two chunks of 9. The prefill crossed the captures
+	// scheduled at interval and 2*interval — key offsets 7 and 15 after the
+	// draft look-ahead shift — and close attached them.
+	e.runCancelledPrefill(t, prompt, 9, 2*9)
+	begins.next()
+
+	// A session diverging inside the cancelled span restores to the deepest
+	// crossed capture: 17 keys match, and the entry at 15 is the closest
+	// restore point below.
+	probe := slices.Concat(slices.Clone(prompt[:18]), ts.fresh(10))
+	pr := e.pc.begin(probe, nil)
+	if _, m, c, _ := begins.next(); m != 17 || c != 15 {
+		t.Errorf("probe into cancelled span: matched=%d cached=%d, want 17/15", m, c)
+	}
+	pr.close()
+
+	// The retry resumes exactly at the 18 tokens the first attempt recorded,
+	// and is cancelled again two chunks deeper.
+	e.runCancelledPrefill(t, prompt, 9, 2*9+2*9)
+	if _, m, c, _ := begins.next(); m != 18 || c != 18 {
+		t.Errorf("first retry: matched=%d cached=%d, want 18/18", m, c)
+	}
+
+	// The final retry resumes at 36 and completes.
+	e.runRequest(t, prompt, ts.fresh(6))
+	if _, m, c, _ := begins.next(); m != 36 || c != 36 {
+		t.Errorf("second retry: matched=%d cached=%d, want 36/36", m, c)
+	}
+
+	if out := logs(); strings.Contains(out, "failed to restore cache") {
+		t.Errorf("freeAll warn fired:\n%s", out)
+	}
+	checkSnapshotCoverage(t, e.pc, e.kvLayers())
+}
+
+// TestScenarioDivergentCancels interleaves cancelled prefills that diverge
+// from a conversation mid-history with completed requests, growing branch
+// points above and below the conversation's capture offsets, then re-requests
+// the full conversation, which must restore to within the capture cadence.
+func TestScenarioDivergentCancels(t *testing.T) {
+	logs := captureWarns(t)
+	e := newHybridEnv()
+	ts := &tokenStream{}
+	begins := &beginLog{t: t, logs: logs}
+
+	prompt := ts.fresh(24)
+	gen := ts.fresh(8)
+	e.runRequest(t, prompt, gen)
+	stream := slices.Concat(prompt, gen)
+	begins.next()
+
+	// A cancelled prefill diverging between the second and third captures.
+	e.runCancelledPrefill(t, slices.Concat(slices.Clone(stream[:2*interval+2]), ts.fresh(20)), 5, 12)
+	begins.next()
+
+	// A completed turn extends the conversation.
+	p2 := slices.Concat(stream, ts.fresh(10))
+	g2 := ts.fresh(4)
+	e.runRequest(t, p2, g2)
+	stream = slices.Concat(p2, g2)
+	begins.next()
+
+	// A cancelled prefill diverging below the first capture.
+	e.runCancelledPrefill(t, slices.Concat(slices.Clone(stream[:interval-2]), ts.fresh(16)), 5, 14)
+	begins.next()
+
+	// Re-requesting the full conversation restores to within the cadence of
+	// the last completed turn.
+	e.runRequest(t, slices.Concat(stream, ts.fresh(12)), ts.fresh(5))
+	if _, m, c, _ := begins.next(); c < m-e.restoreBound(len(g2), 0) {
+		t.Errorf("re-request: cached=%d fell more than %d below matched=%d", c, e.restoreBound(len(g2), 0), m)
+	}
+
+	if out := logs(); strings.Contains(out, "failed to restore cache") {
+		t.Errorf("freeAll warn fired:\n%s", out)
+	}
+	checkSnapshotCoverage(t, e.pc, e.kvLayers())
+}

x/mlxrunner/prefix_cache_test.go

@@ -56,12 +56,18 @@ func (p *fakePending) take() []cache.Snapshot {
 func (p *fakePending) feedCapturing(start int, tokens []int32, capture func(from, reached int) cache.Snapshot, advance func([]int32)) {
 	end := start + len(tokens)
 	captureAt := func(reached int) {
+		fired := false
 		for i, o := range p.offsets {
 			if p.captured[i] == nil && o == reached {
 				p.captured[i] = capture(p.base, reached)
+				fired = true
 			}
 		}
-		p.base = reached
+		// The base advances only when a capture fires, mirroring
+		// pendingSnapshots.captureReached.
+		if fired {
+			p.base = reached
+		}
 	}
 
 	if len(p.offsets) == 0 {
@@ -1005,12 +1011,13 @@ func TestSnapshotBeyondPrefillSkipped(t *testing.T) {
 	})
 }
 
-// TestPrefillSnapshotsDiscardedOnCancel mirrors a prefill canceled after the
-// caches captured interior snapshots but before attachPrefillSnapshots ran. The
-// abandoned captures must be released when the session closes; otherwise the
-// next request's PrepareSnapshots overwrites the schedule without closing them,
-// leaking the snapshots (caught by checkSnapshotLeaks in the env cleanup).
-func TestPrefillSnapshotsDiscardedOnCancel(t *testing.T) {
+// TestPrefillSnapshotsKeptOnCancel mirrors a prefill canceled after the caches
+// captured interior snapshots but before the success-path attach ran. Closing
+// the session attaches the crossed captures so a retry can resume from them,
+// and drains the capture schedule; otherwise the next request's
+// PrepareSnapshots would overwrite it without closing the captures, leaking
+// them (caught by checkSnapshotLeaks in the env cleanup).
+func TestPrefillSnapshotsKeptOnCancel(t *testing.T) {
 	forEachEnv(t, func(t *testing.T, env *testEnv) {
 		pc := env.pc
 		inputs := []int32{1, 2, 3, 4, 5}
@@ -1018,12 +1025,17 @@ func TestPrefillSnapshotsDiscardedOnCancel(t *testing.T) {
 		session := pc.begin(inputs, nil)
 		session.schedulePrefillSnapshots([]int{3})
 		// Cross offset 3 so the caches capture it, then close the session as a
-		// canceled prefill would, before the captures are attached to the trie.
+		// canceled prefill would, before the success-path attach.
 		feedAll(pc.caches, inputs[pc.minCacheOffset():3])
 		session.close()
 
+		// The crossed capture becomes a restore point for the retry.
+		if at := 3 - pc.draftLookahead; !nodeExistsAtOffset(pc.root, at) {
+			t.Errorf("no trie node at capture point %d after cancel", at)
+		}
+
 		// A second request re-prepares snapshots on the same caches: if the
-		// discarded ones were not closed, prepare() orphans them here.
+		// pending ones were not drained, prepare() orphans them here.
 		simulateRequest(t, pc, inputs, nil, 5)
 
 		checkTrieInvariants(t, pc.root)