mlx: preserve generation panics during cleanup

A failure during decoding can trigger deferred drain and cache close, which may panic again and hide the original error. Skip stateful cleanup while a panic is unwinding so the runner reports the initial failure.

Add regression tests for the masked panic paths and normal cache saving.

mlxrunner/pipeline.go

@@ -90,6 +90,15 @@ func (r *Runner) Prepare(request *Request) (err error) {
 // The runner serializes requests today so we just use a fixed slot ID.
 const pipelineSlot = 0
 
+// skipCleanupOnPanic keeps a backend failure from being replaced by cleanup
+// that evaluates or saves partially updated cache state.
+func skipCleanupOnPanic(cleanup func()) {
+	if failure := recover(); failure != nil {
+		panic(failure)
+	}
+	cleanup()
+}
+
 func (r *Runner) TextGenerationPipeline(ctx context.Context, request Request) (err error) {
 	mlx.ResetPeakMemory()
 	mlx.Scoped(func() { err = r.generate(ctx, request) })
@@ -106,16 +115,16 @@ func (r *Runner) generate(ctx context.Context, request Request) error {
 	inputs := request.Tokens
 
 	session := r.cache.begin(inputs, request.MediaItems)
-	defer session.close()
+	defer skipCleanupOnPanic(session.close)
 	caches := session.caches
 
 	media := r.openMedia(request)
-	defer media.close()
+	defer skipCleanupOnPanic(media.close)
 
 	// Built before prefill so a drafter with draft caches follows the prompt
 	// through prefill alongside the target.
 	spec := r.spec.open(request, media.rowLayout())
-	defer spec.close()
+	defer skipCleanupOnPanic(spec.close)
 
 	seed, position, promptEval, err := r.prefill(ctx, session, spec, media)
 	if err != nil {
@@ -137,7 +146,7 @@ func (r *Runner) generate(ctx context.Context, request Request) error {
 	} else {
 		d = r.pipelinedDecoder(nil, caches, seed.ExpandDims(-1), position, media.rowLayout(), grammar)
 	}
-	defer d.close()
+	defer skipCleanupOnPanic(d.close)
 	return r.decode(ctx, request, session, d, promptEval)
 }
 
@@ -253,12 +262,12 @@ type decoder interface {
 // not streamed or counted.
 func (r *Runner) decode(ctx context.Context, request Request, session *cacheSession, d decoder, promptEval time.Duration) error {
 	// A sampled-but-undelivered result is still a produced token; record it.
-	defer func() {
+	defer skipCleanupOnPanic(func() {
 		results, _, _ := d.drain()
 		for _, res := range results {
 			session.outputs = append(session.outputs, res.Token.Int())
 		}
-	}()
+	})
 
 	detok := detokenizer{
 		tokenizer:       r.Tokenizer,

mlxrunner/pipeline_panic_test.go

@@ -0,0 +1,99 @@
+package mlxrunner
+
+import (
+	"context"
+	"testing"
+
+	"github.com/ollama/ollama/api"
+	"github.com/ollama/ollama/mlx"
+	"github.com/ollama/ollama/mlx/mlxtest"
+	"github.com/ollama/ollama/mlxrunner/batch"
+	"github.com/ollama/ollama/mlxrunner/cache"
+	sampler "github.com/ollama/ollama/mlxrunner/sample"
+)
+
+type panicDecoder struct {
+	drained bool
+}
+
+func (*panicDecoder) next(int) ([]sampler.Result, error) { panic("decode failed") }
+func (d *panicDecoder) drain() ([]sampler.Result, int, error) {
+	d.drained = true
+	panic("second drain")
+}
+func (*panicDecoder) close() {}
+
+// In #18856, a second drain replaced the original MLX panic.
+func TestDecodePreservesPanic(t *testing.T) {
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		d := &panicDecoder{}
+		request := Request{CompletionRequest: CompletionRequest{Options: api.Options{NumPredict: 1}}}
+		var failure any
+		func() {
+			defer func() { failure = recover() }()
+			_ = (&Runner{}).decode(context.Background(), request, &cacheSession{inputs: []int32{1}}, d, 0)
+		}()
+		if failure != "decode failed" {
+			t.Fatalf("panic = %v, want original decode failure", failure)
+		}
+		if d.drained {
+			t.Fatal("decoder drained after a panic")
+		}
+	})
+}
+
+type panicForwardModel struct{ textOnlyModel }
+
+func (panicForwardModel) Forward(_ *batch.Batch, caches []cache.Cache) (*mlx.Array, *mlx.Array) {
+	caches[0].(*fakeRewindableCache).feed([]int32{1, 2, 3})
+	panic("forward failed")
+}
+
+// In #18856, cache close replaced the original panic with an offset error.
+func TestGeneratePreservesPanicWithoutSavingCache(t *testing.T) {
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		kv := &fakeRewindableCache{tracker: &snapshotTracker{}}
+		pc := newPrefixCache([]cache.Cache{kv})
+		r := &Runner{Model: panicForwardModel{}, cache: pc}
+		var failure any
+		func() {
+			defer func() { failure = recover() }()
+			mlx.Scoped(func() {
+				_ = r.generate(context.Background(), Request{Tokens: []int32{1, 2}})
+			})
+		}()
+		if failure != "forward failed" {
+			t.Fatalf("panic = %v, want original forward failure", failure)
+		}
+		if len(pc.root.children) != 0 {
+			t.Fatalf("saved %d cache paths after failed forward", len(pc.root.children))
+		}
+	})
+}
+
+func TestGenerateStillSavesCacheOnSuccess(t *testing.T) {
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		kv := &fakeRewindableCache{tracker: &snapshotTracker{}}
+		pc := newPrefixCache([]cache.Cache{kv})
+		tok := newTestTokenizer(t, []int32{7})
+		r := &Runner{
+			Model:     &fakeMTPModel{predict: map[int32]int32{1: 2, 2: 3, 3: 7}, tok: tok},
+			Tokenizer: tok,
+			Sampler:   sampler.New(4096),
+			cache:     pc,
+		}
+		request := Request{
+			Tokens:            []int32{1, 2},
+			Responses:         make(chan CompletionResponse, 4),
+			CompletionRequest: CompletionRequest{Options: api.Options{NumPredict: 1}},
+		}
+		mlx.Scoped(func() {
+			if err := r.generate(context.Background(), request); err != nil {
+				t.Fatalf("generate: %v", err)
+			}
+		})
+		if len(pc.root.children) != 1 {
+			t.Fatalf("saved cache paths = %d, want 1", len(pc.root.children))
+		}
+	})
+}