package mill

import (
	"bytes"
	"context"
	"errors"
	"io"
	"log/slog"
	"strings"
	"sync"
	"testing"
	"time"

	"tangled.org/core/api/tangled"
	"tangled.org/core/spindle/db"
	"tangled.org/core/spindle/engine"
	"tangled.org/core/spindle/models"
	"tangled.org/core/spindle/storage"

	millproto "tangled.org/core/spindle/mill/proto"
	millv1 "tangled.org/core/spindle/mill/proto/gen"
)

type scriptedEncoder func(*millproto.Message) error

func (e scriptedEncoder) Encode(msg *millproto.Message) error { return e(msg) }

func waitForCacheDeletion(t *testing.T, cache storage.Storage, index *db.DB, key string) {
	t.Helper()
	deadline := time.Now().Add(2 * time.Second)
	for time.Now().Before(deadline) {
		reader, err := cache.Get(context.Background(), key)
		if errors.Is(err, storage.ErrNotExist) {
			if index != nil {
				pending, queryErr := index.PendingCacheObjectDeletions(context.Background(), 10)
				if queryErr != nil {
					t.Fatal(queryErr)
				}
				if len(pending) == 0 {
					return
				}
			} else {
				return
			}
		} else if err == nil {
			_ = reader.Close()
		}
		time.Sleep(time.Millisecond)
	}
	t.Fatalf("cache object %q was not deleted", key)
}

func testWorkflow(name string) *models.Workflow {
	repoDid := "did:web:example.com"
	return &models.Workflow{
		Name:        name,
		Environment: map[string]string{},
		Steps:       []models.Step{remoteStep{}},
		RepoDID:     repoDid,
		Data: &millWorkflowState{
			RawWorkflow: tangled.Pipeline_Workflow{Name: name},
			RawPipeline: testPipeline(),
		},
	}
}
func TestMillEngineRequiresSharedCacheStoreID(t *testing.T) {
	twf := tangled.Pipeline_Workflow{
		Name: "build",
		Raw:  "cache:\n  - key: deps\n    paths: [deps]\n",
	}
	for _, tc := range []struct {
		name    string
		storeID string
		want    int
	}{
		{name: "unidentified store", want: 0},
		{name: "shared store", storeID: "fleet-cache", want: 1},
	} {
		t.Run(tc.name, func(t *testing.T) {
			m := New(discardLogger(), Config{CacheStoreID: tc.storeID})
			wf, err := NewEngine("microvm", m).InitWorkflow(twf, tangled.Pipeline{})
			if err != nil {
				t.Fatal(err)
			}
			if got := len(wf.Caches); got != tc.want {
				t.Fatalf("cache entries = %d, want %d", got, tc.want)
			}
		})
	}
}

func TestEngineInitWorkflowCarriesRepositoryIdentity(t *testing.T) {
	m := &Mill{l: slog.New(slog.NewTextHandler(io.Discard, nil))}
	wf, err := NewEngine("dummy", m).InitWorkflow(tangled.Pipeline_Workflow{Name: "build"}, testPipeline())
	if err != nil {
		t.Fatal(err)
	}
	if wf.OwnerDID != "did:plc:testowner" || wf.RepoDID != "did:plc:testrepo" {
		t.Fatalf("workflow identity = %q/%q, want source owner and repository", wf.OwnerDID, wf.RepoDID)
	}
}

func testWorkflowWithRunsOn(name string, runsOn []string) *models.Workflow {
	wf := testWorkflow(name)
	wf.Data.(*millWorkflowState).RawWorkflow.RunsOn = runsOn
	return wf
}

func addCandidateSession(t *testing.T, m *Mill, nodeID string, labels []string, load float64, enc messageEncoder) *millSession {
	t.Helper()
	if enc == nil {
		enc = scriptedEncoder(func(*millproto.Message) error { return nil })
	}
	sess := newSession(nodeID, "inc-"+nodeID, labels, enc, slog.New(slog.NewTextHandler(io.Discard, nil)))
	sess.snapshot = &millv1.NodeSnapshot{
		Seqno: 1,
		Engines: map[string]*millv1.EngineAvailability{
			"dummy": {Available: load < 1.0, Load: map[string]float64{"slots": load}},
		},
	}
	m.mu.Lock()
	m.sessions[nodeID] = sess
	m.mu.Unlock()
	return sess
}

func assertRankedNodes(t *testing.T, got []*millSession, want []string) {
	t.Helper()
	if len(got) != len(want) {
		t.Fatalf("rankCandidates() returned %d candidates, want %d: got %v want %v", len(got), len(want), sessionIDs(got), want)
	}
	for i := range want {
		if got[i].nodeID != want[i] {
			t.Fatalf("rankCandidates()[%d] = %q, want %q; full order got %v want %v", i, got[i].nodeID, want[i], sessionIDs(got), want)
		}
	}
}

func sessionIDs(sessions []*millSession) []string {
	out := make([]string, len(sessions))
	for i, sess := range sessions {
		out[i] = sess.nodeID
	}
	return out
}

func sameStringMultiset(a, b []string) bool {
	if len(a) != len(b) {
		return false
	}
	counts := make(map[string]int, len(a))
	for _, s := range a {
		counts[s]++
	}
	for _, s := range b {
		if counts[s] == 0 {
			return false
		}
		counts[s]--
	}
	return true
}

type reserveReply struct {
	accepted      bool
	rejectClass   millv1.RejectClass
	reason        string
	failureClass  string
	failureReason string
}

func addReplyingCandidateSession(t *testing.T, m *Mill, nodeID string, labels []string, load float64, asked chan<- string, reply reserveReply) *millSession {
	t.Helper()
	var sess *millSession
	sess = addCandidateSession(t, m, nodeID, labels, load, scriptedEncoder(func(msg *millproto.Message) error {
		rs := msg.GetReserveSeat()
		if rs == nil {
			return nil
		}
		if asked != nil {
			asked <- nodeID
		}
		sess.deliver(rs.GetLeaseId(), &millproto.Message{ReserveResult: &millv1.ReserveResult{
			LeaseId:       rs.GetLeaseId(),
			Accepted:      reply.accepted,
			RejectReason:  reply.reason,
			RejectClass:   reply.rejectClass,
			FailureClass:  reply.failureClass,
			FailureReason: reply.failureReason,
		}})
		return nil
	}))
	return sess
}

func drainAsked(ch <-chan string) []string {
	var out []string
	for {
		select {
		case nodeID := <-ch:
			out = append(out, nodeID)
		default:
			return out
		}
	}
}

func TestCommitRetriesAfterSessionCloseBeforeCommitted(t *testing.T) {
	l := slog.New(slog.NewTextHandler(io.Discard, nil))
	m := New(l, Config{BidTimeout: 25 * time.Millisecond, ReconnectGrace: time.Second})
	wf := testWorkflow("build")
	restoreID := "11111111-1111-4111-8111-111111111111"
	saveID := "22222222-2222-4222-8222-222222222222"
	wf.CacheBindings = []models.CacheBinding{{
		EntryIndex:  0,
		Hash:        "abc123",
		RestoreID:   restoreID,
		RestoreKey:  "objects/did:web:example.com/" + restoreID,
		RestoreName: "deps-abc123",
		SaveID:      saveID,
		SaveKey:     "objects/did:web:example.com/" + saveID,
	}}
	wid := models.WorkflowId{PipelineId: models.PipelineId{Knot: "k", Rkey: "r"}, Name: "build"}
	lease := newLease("lease-1", "node-1", "inc-1", "dummy")
	lease.wid = wid
	wf.Data.(*millWorkflowState).Lease = lease

	m.mu.Lock()
	m.leases[lease.id] = lease
	m.mu.Unlock()

	var sess1 *millSession
	firstCommit := make(chan struct{})
	sess1 = newSession("node-1", "inc-1", nil, scriptedEncoder(func(msg *millproto.Message) error {
		if msg.GetCommitLease() != nil {
			close(firstCommit)
			m.detachSession(sess1)
		}
		return nil
	}), l)
	m.attachSession(sess1)

	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	defer cancel()
	done := make(chan error, 1)
	go func() { done <- m.commitAndWait(ctx, wf, nil) }()

	select {
	case <-firstCommit:
	case <-ctx.Done():
		t.Fatal("first commit was not sent")
	}

	var sess2 *millSession
	sess2 = newSession("node-1", "inc-1", nil, scriptedEncoder(func(msg *millproto.Message) error {
		if msg.GetCommitLease() == nil {
			return nil
		}
		bindings := msg.GetCommitLease().GetCacheBindings()
		if len(bindings) != 1 ||
			bindings[0].GetRestoreId() != restoreID ||
			bindings[0].GetSaveId() != saveID {
			t.Errorf("CommitLease cache bindings = %+v", bindings)
		}
		leaseID := msg.GetCommitLease().GetLeaseId()
		sess2.deliver(leaseID, &millproto.Message{Committed: &millv1.Committed{LeaseId: leaseID}})
		_ = m.onEventBatch(sess2, &millv1.EventBatch{
			Epoch: sess2.epoch,
			Events: []*millv1.Event{
				{
					Seqno:   1,
					LeaseId: leaseID,
					Payload: &millv1.Event_AttemptResult{
						AttemptResult: &millv1.AttemptResult{
							Status: millv1.TerminalStatus_TERMINAL_STATUS_SUCCESS,
						},
					},
				},
			},
		})
		return nil
	}), l)
	m.attachSession(sess2)
	m.sessionReady(sess2)

	select {
	case err := <-done:
		if err != nil {
			t.Fatalf("commitAndWait() error = %v, want success after reconnect", err)
		}
	case <-ctx.Done():
		t.Fatal("commitAndWait() did not finish after reconnect")
	}
}

func TestDestroyRunningLeaseDoesNotDropCancelledTerminal(t *testing.T) {
	l := slog.New(slog.NewTextHandler(io.Discard, nil))
	m := New(l, Config{})
	wid := models.WorkflowId{PipelineId: models.PipelineId{Knot: "k", Rkey: "r"}, Name: "build"}
	lease := newLease("lease-1", "node-1", "inc-1", "dummy")
	lease.wid = wid
	lease.setState(leaseRunning)

	m.mu.Lock()
	m.leases[lease.id] = lease
	m.mu.Unlock()

	m.destroy(wid)
	if lease.getState() == leaseDone {
		t.Fatal("destroy sealed the lease before the terminal result")
	}

	lease.deliverTerminal(&millv1.AttemptResult{
		Status: millv1.TerminalStatus_TERMINAL_STATUS_CANCELLED,
	})
	res := <-lease.terminal
	if err := terminalError(res.Status); !errors.Is(err, engine.ErrWorkflowCanceled) {
		t.Fatalf("terminalError() = %v, want ErrCancelled", err)
	}
}

func TestPlaceBlocksWhenNoCapacity(t *testing.T) {
	l := slog.New(slog.NewTextHandler(io.Discard, nil))
	m := New(l, Config{})
	wf := testWorkflow("build")
	wid := models.WorkflowId{PipelineId: models.PipelineId{Knot: "k", Rkey: "r"}, Name: "build"}

	// with no executors at all, place must block until ctx expires and the user sees pending
	ctx, cancel := context.WithTimeout(context.Background(), 150*time.Millisecond)
	defer cancel()

	_, err := m.place(ctx, "dummy", wid, wf)
	if err != context.DeadlineExceeded {
		t.Fatalf("place() error = %v, want DeadlineExceeded", err)
	}
}

func TestRankCandidatesFiltersRequiredLabelsWithANDSemantics(t *testing.T) {
	m := New(slog.New(slog.NewTextHandler(io.Discard, nil)), Config{})
	addCandidateSession(t, m, "linux-high", []string{"linux"}, 0.0, nil)
	addCandidateSession(t, m, "linux-arm", []string{"linux", "arm64"}, 0.25, nil)
	addCandidateSession(t, m, "unlabeled", nil, 0.5, nil)
	addCandidateSession(t, m, "linux-arm-gpu", []string{"linux", "arm64", "gpu"}, 0.75, nil)
	addCandidateSession(t, m, "linux-arm-full", []string{"linux", "arm64"}, 1.0, nil)

	tests := []struct {
		name           string
		requiredLabels []string
		want           []string
	}{
		{
			name: "no required labels keeps old capacity ranking",
			want: []string{"linux-high", "linux-arm", "unlabeled", "linux-arm-gpu"},
		},
		{
			name:           "single required label includes every candidate carrying it",
			requiredLabels: []string{"linux"},
			want:           []string{"linux-high", "linux-arm", "linux-arm-gpu"},
		},
		{
			name:           "all required labels must be present",
			requiredLabels: []string{"linux", "arm64"},
			want:           []string{"linux-arm", "linux-arm-gpu"},
		},
		{
			name:           "one missing required label excludes the candidate",
			requiredLabels: []string{"linux", "arm64", "gpu"},
			want:           []string{"linux-arm-gpu"},
		},
		{
			name:           "unknown required label leaves no candidate",
			requiredLabels: []string{"linux", "arm64", "metal"},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			assertRankedNodes(t, m.rankCandidates("dummy", tt.requiredLabels, false), tt.want)
		})
	}
}

func TestRankCandidatesRequiresSharedCacheStore(t *testing.T) {
	m := New(discardLogger(), Config{CacheStoreID: "disk:/cache"})
	matching := addCandidateSession(t, m, "matching", []string{"linux"}, 0.5, nil)
	matching.cacheStoreID = "disk:/cache"
	other := addCandidateSession(t, m, "other", []string{"linux"}, 0.1, nil)
	other.cacheStoreID = "s3:bucket/cache"

	assertRankedNodes(t, m.rankCandidates("dummy", []string{"linux"}, true), []string{"matching"})
	assertRankedNodes(t, m.rankCandidates("dummy", []string{"linux"}, false), []string{"other", "matching"})
}

func TestPlaceWithMissingRequiredLabelsStaysPendingWithoutReserve(t *testing.T) {
	l := slog.New(slog.NewTextHandler(io.Discard, nil))
	m := New(l, Config{BidTimeout: 10 * time.Millisecond})
	reserveSent := make(chan struct{}, 1)
	addCandidateSession(t, m, "linux-only", []string{"linux"}, 0.75, scriptedEncoder(func(msg *millproto.Message) error {
		if msg.GetReserveSeat() != nil {
			select {
			case reserveSent <- struct{}{}:
			default:
			}
		}
		return nil
	}))
	wf := testWorkflowWithRunsOn("build", []string{"linux", "arm64"})
	wid := models.WorkflowId{PipelineId: models.PipelineId{Knot: "k", Rkey: "r"}, Name: "build"}

	ctx, cancel := context.WithTimeout(context.Background(), 120*time.Millisecond)
	defer cancel()
	_, err := m.place(ctx, "dummy", wid, wf)
	if err != context.DeadlineExceeded {
		t.Fatalf("place() error = %v, want DeadlineExceeded while job remains pending", err)
	}
	select {
	case <-reserveSent:
		t.Fatal("place() sent ReserveSeat to executor missing a required label")
	default:
	}
}

func TestMaxPendingRejects(t *testing.T) {
	l := slog.New(slog.NewTextHandler(io.Discard, nil))
	m := New(l, Config{MaxPending: 1})

	m.mu.Lock()
	m.pending = 1
	m.mu.Unlock()

	wf2 := testWorkflow("b")
	_, err := m.place(context.Background(), "dummy", models.WorkflowId{Name: "b"}, wf2)
	if err == nil {
		t.Fatal("place() past maxPending should error")
	}
}

func TestCancelledRunningLeaseSurvivesReleaseForReconnectReplay(t *testing.T) {
	m, bdb := restoreTestMill(t, Config{ReconnectGrace: time.Minute})
	wid := models.WorkflowId{PipelineId: models.PipelineId{Knot: "k", Rkey: "r"}, Name: "build"}
	lease := newLease("lease-1", "node-1", "inc-1", "dummy")
	lease.wid = wid
	lease.setState(leaseRunning)
	if err := m.persistLease(lease, leaseRowRunning); err != nil {
		t.Fatalf("persistLease: %v", err)
	}
	m.mu.Lock()
	m.leases[lease.id] = lease
	m.mu.Unlock()

	m.destroy(wid)
	slot := &millSlot{fleet: m, lease: lease}
	slot.Release()

	m.mu.Lock()
	_, retained := m.leases[lease.id]
	m.mu.Unlock()
	if !retained {
		t.Fatal("slot release removed a cancellation-requested running lease before its terminal result")
	}
	if rows, err := bdb.ListMillLeases(); err != nil || len(rows) != 1 {
		t.Fatalf("durable leases after slot release = %+v, err = %v; want retained lease", rows, err)
	}

	var sentMu sync.Mutex
	var sent []*millproto.Message
	sess := newSession("node-1", "inc-1", nil, scriptedEncoder(func(msg *millproto.Message) error {
		sentMu.Lock()
		sent = append(sent, msg)
		sentMu.Unlock()
		return nil
	}), discardLogger())
	if _, ok := m.attachSession(sess); !ok {
		t.Fatal("attachSession rejected reconnect")
	}
	m.sessionReady(sess)
	sentMu.Lock()
	var replayed bool
	for _, msg := range sent {
		if cancel := msg.GetCancelAttempt(); cancel != nil && cancel.GetLeaseId() == lease.id {
			replayed = true
		}
	}
	sentMu.Unlock()
	if !replayed {
		t.Fatal("reconnect did not replay CancelAttempt for retained lease")
	}

	if err := m.onEventBatch(sess, &millv1.EventBatch{
		Epoch: sess.epoch,
		Events: []*millv1.Event{
			{
				Seqno:   1,
				LeaseId: lease.id,
				Payload: &millv1.Event_AttemptResult{
					AttemptResult: &millv1.AttemptResult{
						Status: millv1.TerminalStatus_TERMINAL_STATUS_CANCELLED,
					},
				},
			},
		},
	}); err != nil {
		t.Fatalf("onEventBatch: %v", err)
	}
	m.mu.Lock()
	_, retained = m.leases[lease.id]
	m.mu.Unlock()
	if retained {
		t.Fatal("terminal result did not clean retained cancelled lease")
	}
	if rows, err := bdb.ListMillLeases(); err != nil || len(rows) != 0 {
		t.Fatalf("durable leases after terminal = %+v, err = %v; want none", rows, err)
	}
	slot.Release()
}

func TestSessionRequestCancelledBeforeRegistrationOrSend(t *testing.T) {
	sent := 0
	sess := newSession("node-1", "inc-1", nil, scriptedEncoder(func(*millproto.Message) error {
		sent++
		return nil
	}), discardLogger())
	ctx, cancel := context.WithCancel(context.Background())
	cancel()
	_, err := sess.request(ctx, "lease-1", &millproto.Message{ReleaseLease: &millv1.ReleaseLease{LeaseId: "lease-1"}})
	if !errors.Is(err, context.Canceled) {
		t.Fatalf("request error = %v, want context.Canceled", err)
	}
	if sent != 0 {
		t.Fatalf("request sent %d messages for an already-cancelled context, want 0", sent)
	}
	sess.mu.Lock()
	pending := len(sess.pending)
	sess.mu.Unlock()
	if pending != 0 {
		t.Fatalf("request left %d pending waiters, want 0", pending)
	}
}

func TestAttachSessionReplacesSilentIncumbentButRejectsActiveDuplicate(t *testing.T) {
	m := New(discardLogger(), Config{ReconnectGrace: time.Minute})
	old := newSession("node-1", "inc-old", nil, nopEncoder(), discardLogger())
	transportClosed := make(chan struct{})
	old.closeTransport = func() error {
		close(transportClosed)
		return nil
	}
	if _, ok := m.attachSession(old); !ok {
		t.Fatal("first attach rejected")
	}
	m.mu.Lock()
	old.lastSeen = time.Now().Add(-2 * m.cfg.ReconnectGrace)
	m.mu.Unlock()

	replacement := newSession("node-1", "inc-new", nil, nopEncoder(), discardLogger())
	if _, ok := m.attachSession(replacement); !ok {
		t.Fatal("silent incumbent blocked authenticated replacement")
	}
	select {
	case <-transportClosed:
	default:
		t.Fatal("replacing a silent incumbent did not close its transport")
	}
	m.mu.Lock()
	replacement.lastSeen = time.Now().Add(-2 * m.cfg.ReconnectGrace)
	m.mu.Unlock()
	if err := replacement.dispatch(m, &millproto.Message{NodeSnapshot: &millv1.NodeSnapshot{Seqno: 1}}); err != nil {
		t.Fatalf("periodic snapshot dispatch: %v", err)
	}
	if _, ok := m.attachSession(newSession("node-1", "inc-dup", nil, nopEncoder(), discardLogger())); ok {
		t.Fatal("active replacement did not reject a duplicate session")
	}
}

func TestSilentSessionReplacementWithoutSnapshotFailsLeaseAtDeadline(t *testing.T) {
	m, _ := restoreTestMill(t, Config{ReconnectGrace: time.Minute})
	old := newSession("node-1", "inc-old", nil, nopEncoder(), discardLogger())
	if _, ok := m.attachSession(old); !ok {
		t.Fatal("first attach rejected")
	}
	lease := newLease("lease-1", old.nodeID, old.epoch, "dummy")
	lease.wid = models.WorkflowId{PipelineId: models.PipelineId{Knot: "k", Rkey: "r"}, Name: "build"}
	lease.setState(leaseRunning)
	if err := m.persistLease(lease, leaseRowRunning); err != nil {
		t.Fatalf("persistLease: %v", err)
	}
	m.mu.Lock()
	m.leases[lease.id] = lease
	old.lastSeen = time.Now().Add(-2 * m.cfg.ReconnectGrace)
	m.mu.Unlock()

	replacement := newSession("node-1", "inc-new", nil, nopEncoder(), discardLogger())
	attachStarted := time.Now()
	if _, ok := m.attachSession(replacement); !ok {
		t.Fatal("silent incumbent blocked replacement")
	}
	attachFinished := time.Now()
	if !replacement.recovering || replacement.graceTimer == nil || replacement.graceDeadline.IsZero() {
		t.Fatal("replacement without snapshot has no recovery deadline")
	}
	if replacement.graceDeadline.Before(attachStarted.Add(m.cfg.ReconnectGrace)) ||
		replacement.graceDeadline.After(attachFinished.Add(m.cfg.ReconnectGrace)) {
		t.Fatalf("replacement deadline = %v, want one reconnect grace after attach", replacement.graceDeadline)
	}
	m.failLeasesAfterGrace(replacement)

	m.mu.Lock()
	_, stillActive := m.leases[lease.id]
	m.mu.Unlock()
	if stillActive {
		t.Fatal("replacement without snapshot stranded incumbent lease")
	}
}

func TestPlaceReleasesRemoteReservationWhenInitialPersistenceFails(t *testing.T) {
	m, bdb := restoreTestMill(t, Config{BidTimeout: time.Second})
	if err := bdb.Close(); err != nil {
		t.Fatalf("close db: %v", err)
	}
	released := make(chan string, 1)
	var sess *millSession
	sess = addCandidateSession(t, m, "node-1", nil, 0, scriptedEncoder(func(msg *millproto.Message) error {
		switch {
		case msg.GetReserveSeat() != nil:
			leaseID := msg.GetReserveSeat().GetLeaseId()
			sess.deliver(leaseID, &millproto.Message{ReserveResult: &millv1.ReserveResult{
				LeaseId:  leaseID,
				Accepted: true,
			}})
		case msg.GetReleaseLease() != nil:
			released <- msg.GetReleaseLease().GetLeaseId()
		}
		return nil
	}))

	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
	defer cancel()
	slot, err := m.place(
		ctx,
		"dummy",
		models.WorkflowId{PipelineId: models.PipelineId{Knot: "k", Rkey: "r"}, Name: "build"},
		testWorkflow("build"),
	)
	if err == nil {
		t.Fatal("place succeeded after reserved lease persistence failed")
	}
	if slot != nil {
		t.Fatalf("place returned slot %T after persistence failure", slot)
	}
	select {
	case leaseID := <-released:
		if leaseID == "" {
			t.Fatal("ReleaseLease had empty lease id")
		}
	case <-time.After(time.Second):
		t.Fatal("persistence failure did not compensate with ReleaseLease")
	}
	m.mu.Lock()
	leases := len(m.leases)
	m.mu.Unlock()
	if leases != 0 {
		t.Fatalf("mill published %d leases after initial persistence failure, want 0", leases)
	}
}

func TestGapsAndDuplicates(t *testing.T) {
	m, _ := restoreTestMill(t, Config{ReconnectGrace: time.Minute})
	sess := newSession("node-1", "inc-1", nil, nopEncoder(), discardLogger())
	m.attachSession(sess)

	owned := newLease("lease-1", "node-1", "inc-1", "dummy")
	owned.wid = models.WorkflowId{PipelineId: models.PipelineId{Knot: "k", Rkey: "r"}, Name: "build"}
	m.mu.Lock()
	m.leases[owned.id] = owned
	m.mu.Unlock()

	err := m.onEventBatch(sess, &millv1.EventBatch{
		Epoch: sess.epoch,
		Events: []*millv1.Event{
			{
				Seqno:   0,
				LeaseId: owned.id,
				Payload: &millv1.Event_StatusEvent{
					StatusEvent: &millv1.StatusEvent{Status: millv1.NonterminalStatus_NONTERMINAL_STATUS_RUNNING},
				},
			},
		},
	})
	if err != nil {
		t.Fatalf("expected duplicate to be skipped without error, got: %v", err)
	}

	err = m.onEventBatch(sess, &millv1.EventBatch{
		Epoch: sess.epoch,
		Events: []*millv1.Event{
			{
				Seqno:   2,
				LeaseId: owned.id,
				Payload: &millv1.Event_StatusEvent{
					StatusEvent: &millv1.StatusEvent{Status: millv1.NonterminalStatus_NONTERMINAL_STATUS_RUNNING},
				},
			},
		},
	})
	if err == nil {
		t.Fatal("expected error due to seqno gap, got nil")
	}
}

func TestAtomicBatchRollback(t *testing.T) {
	m, bdb := restoreTestMill(t, Config{ReconnectGrace: time.Minute})
	sess := newSession("node-1", "inc-1", nil, nopEncoder(), discardLogger())
	m.attachSession(sess)

	owned := newLease("lease-1", "node-1", "inc-1", "dummy")
	owned.wid = models.WorkflowId{PipelineId: models.PipelineId{Knot: "k", Rkey: "r"}, Name: "build"}
	m.mu.Lock()
	m.leases[owned.id] = owned
	m.mu.Unlock()

	if _, err := bdb.Exec(`
		create trigger reject_status_event
		before insert on events
		begin
			select raise(abort, 'forced status event failure');
		end
	`); err != nil {
		t.Fatalf("failed to create fail trigger: %v", err)
	}
	defer bdb.Exec("drop trigger reject_status_event")

	err := m.onEventBatch(sess, &millv1.EventBatch{
		Epoch: sess.epoch,
		Events: []*millv1.Event{
			{
				Seqno:   1,
				LeaseId: owned.id,
				Payload: &millv1.Event_StatusEvent{
					StatusEvent: &millv1.StatusEvent{Status: millv1.NonterminalStatus_NONTERMINAL_STATUS_RUNNING},
				},
			},
		},
	})
	if err == nil {
		t.Fatal("expected status event insertion to fail due to trigger")
	}

	m.mu.Lock()
	seqno := m.nodeSeqno["node-1/inc-1"]
	m.mu.Unlock()
	if seqno != 0 {
		t.Fatalf("expected seqno 0 due to rollback, got %d", seqno)
	}
}

func TestCacheUpdatesCommitWithStreamCursor(t *testing.T) {
	m, bdb := restoreTestMill(t, Config{ReconnectGrace: time.Minute})
	cache, err := storage.NewDisk(t.TempDir())
	if err != nil {
		t.Fatal(err)
	}
	m.AttachCache(cache)

	at := time.Date(2026, 8, 11, 20, 0, 0, 0, time.UTC)
	old := db.CacheEntry{
		ID: "old", StorageKey: "objects/old", OwnerDID: "did:plc:owner",
		RepoDID: "did:plc:repo", Engine: "dummy", CacheKey: "deps", CacheHash: "hash",
		SizeBytes: 3, State: "ready", CreatedAt: at.Add(-time.Hour), LastUsedAt: at.Add(-time.Hour),
	}
	pending := old
	pending.ID = "pending"
	pending.StorageKey = "objects/pending"
	pending.SizeBytes = 0
	pending.State = "pending"
	pending.CreatedAt = at
	pending.LastUsedAt = at
	for _, entry := range []db.CacheEntry{old, pending} {
		if err := bdb.InsertCacheEntry(context.Background(), entry); err != nil {
			t.Fatal(err)
		}
		if err := cache.Put(context.Background(), entry.StorageKey, bytes.NewReader([]byte(entry.ID))); err != nil {
			t.Fatal(err)
		}
	}

	sess := newSession("node-1", "inc-1", nil, nopEncoder(), discardLogger())
	m.attachSession(sess)
	owned := newLease("lease-1", "node-1", "inc-1", "dummy")
	owned.wid = models.WorkflowId{
		PipelineId: models.PipelineId{Knot: "k", Rkey: "r"},
		Name:       "build",
	}
	if err := m.persistLease(owned, leaseRowRunning); err != nil {
		t.Fatal(err)
	}
	if err := bdb.SaveMillCacheCapabilities(owned.id, []db.MillCacheCapability{
		{Action: "restore", CacheID: old.ID, StorageKey: old.StorageKey},
		{Action: "save", CacheID: pending.ID, StorageKey: pending.StorageKey},
	}); err != nil {
		t.Fatal(err)
	}
	m.mu.Lock()
	m.leases[owned.id] = owned
	m.mu.Unlock()

	batch := &millv1.EventBatch{
		Epoch: sess.epoch,
		Events: []*millv1.Event{
			{
				Seqno: 1, LeaseId: owned.id,
				Payload: &millv1.Event_CacheUpdate{CacheUpdate: &millv1.CacheUpdate{
					Action: millv1.CacheUpdateAction_CACHE_UPDATE_ACTION_USED, Id: old.ID,
				}},
			},
			{
				Seqno: 2, LeaseId: owned.id,
				Payload: &millv1.Event_CacheUpdate{CacheUpdate: &millv1.CacheUpdate{
					Action: millv1.CacheUpdateAction_CACHE_UPDATE_ACTION_STORED, Id: pending.ID,
					SizeBytes: int64(len(pending.ID)),
				}},
			},
		},
	}
	if err := m.onEventBatch(sess, batch); err != nil {
		t.Fatalf("onEventBatch: %v", err)
	}
	if err := m.onEventBatch(sess, batch); err != nil {
		t.Fatalf("replay onEventBatch: %v", err)
	}

	ready, err := bdb.FindCacheEntry(context.Background(), old.RepoDID, old.Engine, old.CacheKey, old.CacheHash)
	if err != nil {
		t.Fatal(err)
	}
	if ready.ID != pending.ID || ready.SizeBytes != int64(len(pending.ID)) {
		t.Fatalf("ready cache = %+v", ready)
	}

	waitForCacheDeletion(t, cache, nil, old.StorageKey)
	if cursor, err := bdb.GetExecutorCursor(sess.nodeID, sess.epoch); err != nil || cursor != 2 {
		t.Fatalf("cursor = %d, %v; want 2", cursor, err)
	}
}
func TestCacheUpdateCannotDeleteUnplannedObject(t *testing.T) {
	m, bdb := restoreTestMill(t, Config{ReconnectGrace: time.Minute})
	cache, err := storage.NewDisk(t.TempDir())
	if err != nil {
		t.Fatal(err)
	}
	m.AttachCache(cache)
	victim := "objects/did:web:victim.example/11111111-1111-4111-8111-111111111111"
	if err := cache.Put(context.Background(), victim, strings.NewReader("victim")); err != nil {
		t.Fatal(err)
	}

	sess := newSession("node-1", "inc-1", nil, nopEncoder(), discardLogger())
	m.attachSession(sess)
	lease := newLease("lease-1", sess.nodeID, sess.epoch, "dummy")
	lease.wid = models.WorkflowId{
		PipelineId: models.PipelineId{Knot: "k", Rkey: "r"},
		Name:       "build",
	}
	if err := m.persistLease(lease, leaseRowRunning); err != nil {
		t.Fatal(err)
	}
	m.mu.Lock()
	m.leases[lease.id] = lease
	m.mu.Unlock()

	err = m.onEventBatch(sess, &millv1.EventBatch{
		Epoch: sess.epoch,
		Events: []*millv1.Event{{
			Seqno:   1,
			LeaseId: lease.id,
			Payload: &millv1.Event_CacheUpdate{CacheUpdate: &millv1.CacheUpdate{
				Action: millv1.CacheUpdateAction_CACHE_UPDATE_ACTION_STORED,
				Id:     "unknown",
			}},
		}},
	})
	if err == nil {
		t.Fatal("unplanned cache update succeeded")
	}
	reader, err := cache.Get(context.Background(), victim)
	if err != nil {
		t.Fatalf("victim object was deleted: %v", err)
	}
	_ = reader.Close()
	if cursor, err := bdb.GetExecutorCursor(sess.nodeID, sess.epoch); err != nil || cursor != 0 {
		t.Fatalf("cursor = %d, %v; want 0", cursor, err)
	}
}

func TestMissingPendingRowQueuesAndDeletesPlannedObject(t *testing.T) {
	m, bdb := restoreTestMill(t, Config{ReconnectGrace: time.Minute})
	cache, err := storage.NewDisk(t.TempDir())
	if err != nil {
		t.Fatal(err)
	}
	m.AttachCache(cache)
	id := "11111111-1111-4111-8111-111111111111"
	key := "objects/did:web:example.com/" + id
	if err := cache.Put(context.Background(), key, strings.NewReader("orphan")); err != nil {
		t.Fatal(err)
	}

	sess := newSession("node-1", "inc-1", nil, nopEncoder(), discardLogger())
	m.attachSession(sess)
	lease := newLease("lease-1", sess.nodeID, sess.epoch, "dummy")
	lease.wid = models.WorkflowId{
		PipelineId: models.PipelineId{Knot: "k", Rkey: "r"},
		Name:       "build",
	}
	if err := m.persistLease(lease, leaseRowRunning); err != nil {
		t.Fatal(err)
	}
	if err := bdb.SaveMillCacheCapabilities(lease.id, []db.MillCacheCapability{{
		Action:     "save",
		CacheID:    id,
		StorageKey: key,
	}}); err != nil {
		t.Fatal(err)
	}
	m.mu.Lock()
	m.leases[lease.id] = lease
	m.mu.Unlock()

	if err := m.onEventBatch(sess, &millv1.EventBatch{
		Epoch: sess.epoch,
		Events: []*millv1.Event{{
			Seqno:   1,
			LeaseId: lease.id,
			Payload: &millv1.Event_CacheUpdate{CacheUpdate: &millv1.CacheUpdate{
				Action: millv1.CacheUpdateAction_CACHE_UPDATE_ACTION_STORED,
				Id:     id,
			}},
		}},
	}); err != nil {
		t.Fatal(err)
	}
	waitForCacheDeletion(t, cache, bdb, key)
}

func TestTerminalBeforeACK(t *testing.T) {
	m, bdb := restoreTestMill(t, Config{ReconnectGrace: time.Minute})

	ackSent := make(chan struct{})
	sess := newSession("node-1", "inc-1", nil, scriptedEncoder(func(msg *millproto.Message) error {
		if msg.GetAck() != nil {
			wid := models.WorkflowId{PipelineId: models.PipelineId{Knot: "k", Rkey: "r"}, Name: "build"}
			st, err := bdb.GetStatus(wid)
			if err != nil || st.Status != "success" {
				t.Errorf("expected terminal status success at ACK time, got status: %v, err: %v", st, err)
			}
			close(ackSent)
		}
		return nil
	}), discardLogger())
	m.attachSession(sess)

	owned := newLease("lease-1", "node-1", "inc-1", "dummy")
	owned.wid = models.WorkflowId{PipelineId: models.PipelineId{Knot: "k", Rkey: "r"}, Name: "build"}
	m.mu.Lock()
	m.leases[owned.id] = owned
	m.mu.Unlock()

	err := m.onEventBatch(sess, &millv1.EventBatch{
		Epoch: sess.epoch,
		Events: []*millv1.Event{
			{
				Seqno:   1,
				LeaseId: owned.id,
				Payload: &millv1.Event_AttemptResult{
					AttemptResult: &millv1.AttemptResult{Status: millv1.TerminalStatus_TERMINAL_STATUS_SUCCESS},
				},
			},
		},
	})
	if err != nil {
		t.Fatalf("onEventBatch: %v", err)
	}

	select {
	case <-ackSent:
	case <-time.After(2 * time.Second):
		t.Fatal("ACK was not sent")
	}
}

func TestTerminalWithIncompleteIdentityAdvancesCursor(t *testing.T) {
	m, bdb := restoreTestMill(t, Config{ReconnectGrace: time.Minute})
	sess := newSession("node-1", "inc-1", nil, scriptedEncoder(func(*millproto.Message) error { return nil }), discardLogger())
	m.attachSession(sess)

	lease := newLease("lease-1", sess.nodeID, sess.epoch, "dummy")
	lease.wid = models.WorkflowId{PipelineId: models.PipelineId{Rkey: "r"}, Name: "build"}
	m.mu.Lock()
	m.leases[lease.id] = lease
	m.mu.Unlock()

	err := m.onEventBatch(sess, &millv1.EventBatch{
		Epoch: sess.epoch,
		Events: []*millv1.Event{{
			Seqno:   1,
			LeaseId: lease.id,
			Payload: &millv1.Event_AttemptResult{AttemptResult: &millv1.AttemptResult{
				Status: millv1.TerminalStatus_TERMINAL_STATUS_SUCCESS,
				LogArtifact: &millv1.LogArtifact{
					Ref:  "logs/" + lease.id + ".log",
					Hash: "sha256:test",
				},
			}},
		}},
	})
	if err != nil {
		t.Fatalf("onEventBatch: %v", err)
	}

	m.mu.Lock()
	seqno := m.nodeSeqno[sess.nodeID+"/"+sess.epoch]
	m.mu.Unlock()
	if seqno != 1 {
		t.Fatalf("cursor = %d, want 1", seqno)
	}
	var artifacts int
	if err := bdb.QueryRow(`select count(*) from mill_artifacts`).Scan(&artifacts); err != nil {
		t.Fatal(err)
	}
	if artifacts != 0 {
		t.Fatalf("artifact rows = %d, want 0", artifacts)
	}
}

func TestExecutorRestartEmptySnapshot(t *testing.T) {
	m, _ := restoreTestMill(t, Config{ReconnectGrace: time.Minute})

	lease := newLease("lease-1", "node-1", "inc-old", "dummy")
	lease.wid = models.WorkflowId{PipelineId: models.PipelineId{Knot: "k", Rkey: "r"}, Name: "build"}
	m.mu.Lock()
	m.leases[lease.id] = lease
	m.mu.Unlock()
	if err := m.persistLease(lease, leaseRowRunning); err != nil {
		t.Fatalf("persistLease: %v", err)
	}

	sess := newSession("node-1", "inc-new", nil, nopEncoder(), discardLogger())
	m.attachSession(sess)

	err := m.onSnapshot(sess, &millv1.NodeSnapshot{
		Seqno:          1,
		ActiveLeaseIds: nil,
	})
	if err != nil {
		t.Fatalf("onSnapshot: %v", err)
	}

	m.mu.Lock()
	_, stillActive := m.leases["lease-1"]
	m.mu.Unlock()
	if stillActive {
		t.Fatal("expected old epoch lease to be reconciled and failed")
	}
}

func TestReplacementLostBeforeSnapshotFailsOldEpochLease(t *testing.T) {
	m, _ := restoreTestMill(t, Config{ReconnectGrace: time.Minute})
	lease := newLease("lease-1", "node-1", "inc-old", "dummy")
	lease.wid = models.WorkflowId{PipelineId: models.PipelineId{Knot: "k", Rkey: "r"}, Name: "build"}
	lease.setState(leaseRunning)
	if err := m.persistLease(lease, leaseRowRunning); err != nil {
		t.Fatalf("persistLease: %v", err)
	}
	m.mu.Lock()
	m.leases[lease.id] = lease
	m.mu.Unlock()

	replacement := newSession("node-1", "inc-new", nil, nopEncoder(), discardLogger())
	if _, ok := m.attachSession(replacement); !ok {
		t.Fatal("attachSession rejected replacement")
	}
	m.detachSession(replacement)
	m.failLeasesAfterGrace(replacement)

	m.mu.Lock()
	_, stillActive := m.leases[lease.id]
	m.mu.Unlock()
	if stillActive {
		t.Fatal("replacement loss stranded old-epoch lease")
	}
}

func TestSeqRegression(t *testing.T) {
	m, _ := restoreTestMill(t, Config{ReconnectGrace: time.Minute})
	sess := newSession("node-1", "inc-1", nil, nopEncoder(), discardLogger())
	m.attachSession(sess)

	err := m.onSnapshot(sess, &millv1.NodeSnapshot{
		Seqno: 5,
	})
	if err != nil {
		t.Fatalf("first snapshot: %v", err)
	}

	err = m.onSnapshot(sess, &millv1.NodeSnapshot{
		Seqno: 4,
	})
	if err == nil {
		t.Fatal("expected seqno regression to be rejected")
	}
}

func TestClaimedSweep(t *testing.T) {
	m, _ := restoreTestMill(t, Config{ReconnectGrace: time.Minute})

	lease := newLease("lease-1", "node-1", "inc-1", "dummy")
	lease.wid = models.WorkflowId{PipelineId: models.PipelineId{Knot: "k", Rkey: "r"}, Name: "build"}
	lease.orphaned = true
	lease.claimed = false
	m.mu.Lock()
	m.leases[lease.id] = lease
	m.mu.Unlock()

	sess := newSession("node-1", "inc-1", nil, nopEncoder(), discardLogger())
	m.attachSession(sess)
	err := m.onSnapshot(sess, &millv1.NodeSnapshot{
		Seqno:          1,
		ActiveLeaseIds: []string{"lease-1"},
	})
	if err != nil {
		t.Fatalf("onSnapshot: %v", err)
	}

	m.detachSession(sess)

	m.sweepUnclaimedOrphans()

	m.mu.Lock()
	_, stillRunning := m.leases["lease-1"]
	m.mu.Unlock()
	if !stillRunning {
		t.Fatal("claimed lease was incorrectly swept by startup sweep")
	}
}

func TestUnacknowledgedCancelClosesSession(t *testing.T) {
	m, bdb := restoreTestMill(t, Config{ReconnectGrace: time.Minute, CancelAckTimeout: 50 * time.Millisecond})
	registerTestExecutor(t, bdb, "node-1", HashToken("tok-1"), nil)

	sessClosed := make(chan struct{})
	sess := newSession("node-1", "inc-1", nil, nopEncoder(), discardLogger())
	sess.closeTransport = func() error {
		close(sessClosed)
		return nil
	}
	m.attachSession(sess)

	lease := newLease("lease-1", "node-1", "inc-1", "dummy")
	lease.wid = models.WorkflowId{PipelineId: models.PipelineId{Knot: "k", Rkey: "r"}, Name: "build"}
	lease.setState(leaseRunning)
	m.mu.Lock()
	m.leases[lease.id] = lease
	m.mu.Unlock()

	m.destroy(lease.wid)

	select {
	case <-sessClosed:
	case <-time.After(2 * time.Second):
		t.Fatal("session was not closed after the cancel went unacknowledged")
	}
	if _, _, ok, err := bdb.ResolveExecutorToken(HashToken("tok-1")); err != nil || !ok {
		t.Fatalf("executor identity rejected after session timeout: ok = %v, err = %v", ok, err)
	}
}

func TestPendingCancelBlocksPlacementUntilAcknowledged(t *testing.T) {
	m, _ := restoreTestMill(t, Config{ReconnectGrace: time.Minute, CancelAckTimeout: time.Minute})
	sess := addCandidateSession(t, m, "node-1", nil, 0, nil)
	lease := newLease("lease-1", "node-1", sess.epoch, "dummy")
	lease.wid = models.WorkflowId{PipelineId: models.PipelineId{Knot: "k", Rkey: "r"}, Name: "build"}
	lease.setState(leaseRunning)
	m.mu.Lock()
	m.leases[lease.id] = lease
	m.mu.Unlock()

	assertRankedNodes(t, m.rankCandidates("dummy", nil, false), []string{"node-1"})
	m.destroy(lease.wid)
	assertRankedNodes(t, m.rankCandidates("dummy", nil, false), nil)

	m.onCancelAck(sess, &millv1.CancelAck{LeaseId: lease.id})
	assertRankedNodes(t, m.rankCandidates("dummy", nil, false), []string{"node-1"})
}

func TestTerminalSettlesPendingCancelPlacementGate(t *testing.T) {
	m, _ := restoreTestMill(t, Config{ReconnectGrace: time.Minute, CancelAckTimeout: time.Minute})
	sess := addCandidateSession(t, m, "node-1", nil, 0, nil)
	lease := newLease("lease-1", sess.nodeID, sess.epoch, "dummy")
	lease.wid = models.WorkflowId{PipelineId: models.PipelineId{Knot: "k", Rkey: "r"}, Name: "build"}
	lease.setState(leaseRunning)
	m.mu.Lock()
	m.leases[lease.id] = lease
	m.mu.Unlock()

	m.destroy(lease.wid)
	assertRankedNodes(t, m.rankCandidates("dummy", nil, false), nil)
	if err := m.onEventBatch(sess, &millv1.EventBatch{
		Epoch: sess.epoch,
		Events: []*millv1.Event{{
			Seqno:   1,
			LeaseId: lease.id,
			Payload: &millv1.Event_AttemptResult{AttemptResult: &millv1.AttemptResult{
				Status: millv1.TerminalStatus_TERMINAL_STATUS_CANCELLED,
			}},
		}},
	}); err != nil {
		t.Fatalf("onEventBatch: %v", err)
	}
	assertRankedNodes(t, m.rankCandidates("dummy", nil, false), []string{"node-1"})
}

func TestCancelSentAfterTerminalDoesNotLeakPlacementGate(t *testing.T) {
	m, _ := restoreTestMill(t, Config{ReconnectGrace: time.Minute, CancelAckTimeout: time.Minute})
	sess := addCandidateSession(t, m, "node-1", nil, 0, nil)
	lease := newLease("lease-1", sess.nodeID, sess.epoch, "dummy")
	lease.wid = models.WorkflowId{PipelineId: models.PipelineId{Knot: "k", Rkey: "r"}, Name: "build"}
	lease.setState(leaseRunning)
	if err := m.persistLease(lease, leaseRowRunning); err != nil {
		t.Fatalf("persistLease: %v", err)
	}
	m.mu.Lock()
	m.leases[lease.id] = lease
	m.mu.Unlock()

	lease.requestCancel("done")
	if err := m.finishLiveLease(lease, string(models.StatusKindCancelled), "done"); err != nil {
		t.Fatalf("finishLiveLease: %v", err)
	}
	m.sendCancel(sess, lease, "done")

	m.mu.Lock()
	pending := len(sess.pendingCancels)
	m.mu.Unlock()
	if pending != 0 {
		t.Fatalf("pending cancels after terminal = %d, want 0", pending)
	}
	assertRankedNodes(t, m.rankCandidates("dummy", nil, false), []string{"node-1"})
}

func TestValidSnapshotRecoversAfterProtocolError(t *testing.T) {
	m, _ := restoreTestMill(t, Config{ReconnectGrace: time.Minute})
	first := addCandidateSession(t, m, "node-1", nil, 0, nil)
	lease := newLease("lease-1", first.nodeID, first.epoch, "dummy")
	lease.wid = models.WorkflowId{PipelineId: models.PipelineId{Knot: "k", Rkey: "r"}, Name: "build"}
	lease.setState(leaseRunning)
	m.mu.Lock()
	m.leases[lease.id] = lease
	m.mu.Unlock()

	m.noteSessionError(first, protoErrf("bad stream"))
	m.detachSession(first)
	replacement := newSession(first.nodeID, first.epoch, nil, nopEncoder(), discardLogger())
	if _, ok := m.attachSession(replacement); !ok {
		t.Fatal("reconnect was rejected")
	}
	m.sessionReady(replacement)
	if err := m.onSnapshot(replacement, &millv1.NodeSnapshot{
		Seqno:          1,
		ActiveLeaseIds: []string{lease.id},
		Engines: map[string]*millv1.EngineAvailability{
			"dummy": {Available: true},
		},
	}); err != nil {
		t.Fatalf("onSnapshot: %v", err)
	}

	m.mu.Lock()
	recovering := replacement.recovering
	m.mu.Unlock()
	if recovering {
		t.Fatal("valid reconnect snapshot did not end recovery")
	}
	assertRankedNodes(t, m.rankCandidates("dummy", nil, false), []string{"node-1"})
}

func TestReconnectRetriesKeepOriginalFailureDeadline(t *testing.T) {
	m, _ := restoreTestMill(t, Config{ReconnectGrace: time.Minute})
	first := addCandidateSession(t, m, "node-1", nil, 0, nil)
	lease := newLease("lease-1", first.nodeID, first.epoch, "dummy")
	lease.wid = models.WorkflowId{PipelineId: models.PipelineId{Knot: "k", Rkey: "r"}, Name: "build"}
	lease.setState(leaseRunning)
	if err := m.persistLease(lease, leaseRowRunning); err != nil {
		t.Fatalf("persistLease: %v", err)
	}
	m.mu.Lock()
	m.leases[lease.id] = lease
	m.mu.Unlock()

	m.detachSession(first)
	deadline := first.graceDeadline
	second := newSession(first.nodeID, first.epoch, nil, nopEncoder(), discardLogger())
	if _, ok := m.attachSession(second); !ok {
		t.Fatal("first reconnect was rejected")
	}
	m.detachSession(second)
	third := newSession(first.nodeID, first.epoch, nil, nopEncoder(), discardLogger())
	if _, ok := m.attachSession(third); !ok {
		t.Fatal("second reconnect was rejected")
	}
	if !third.graceDeadline.Equal(deadline) {
		t.Fatalf("reconnect moved failure deadline from %v to %v", deadline, third.graceDeadline)
	}

	m.failLeasesAfterGrace(third)
	m.mu.Lock()
	_, tracked := m.leases[lease.id]
	m.mu.Unlock()
	if tracked {
		t.Fatal("reconnect retries kept the lease past its original failure deadline")
	}
}

func TestUnacknowledgedCancelReconnectsReachBoundedTerminal(t *testing.T) {
	m, _ := restoreTestMill(t, Config{ReconnectGrace: time.Minute, CancelAckTimeout: time.Minute})
	first := addCandidateSession(t, m, "node-1", nil, 0, nil)
	lease := newLease("lease-1", first.nodeID, first.epoch, "dummy")
	lease.wid = models.WorkflowId{PipelineId: models.PipelineId{Knot: "k", Rkey: "r"}, Name: "build"}
	lease.setState(leaseRunning)
	if err := m.persistLease(lease, leaseRowRunning); err != nil {
		t.Fatalf("persistLease: %v", err)
	}
	m.mu.Lock()
	m.leases[lease.id] = lease
	m.mu.Unlock()

	m.destroy(lease.wid)
	m.checkCancelAck(first, lease)
	m.detachSession(first)
	deadline := first.graceDeadline

	second := newSession(first.nodeID, first.epoch, nil, nopEncoder(), discardLogger())
	if _, ok := m.attachSession(second); !ok {
		t.Fatal("first reconnect was rejected")
	}
	m.sessionReady(second)
	if err := m.onSnapshot(second, &millv1.NodeSnapshot{
		Seqno:          1,
		ActiveLeaseIds: []string{lease.id},
		Engines: map[string]*millv1.EngineAvailability{
			"dummy": {Available: true},
		},
	}); err != nil {
		t.Fatalf("onSnapshot: %v", err)
	}
	if !second.recovering {
		t.Fatal("snapshot cleared recovery while cancel was still pending")
	}
	m.checkCancelAck(second, lease)
	m.detachSession(second)

	third := newSession(first.nodeID, first.epoch, nil, nopEncoder(), discardLogger())
	if _, ok := m.attachSession(third); !ok {
		t.Fatal("second reconnect was rejected")
	}
	m.sessionReady(third)
	if !third.graceDeadline.Equal(deadline) {
		t.Fatalf("cancel reconnect moved failure deadline from %v to %v", deadline, third.graceDeadline)
	}
	m.failLeasesAfterGrace(third)

	select {
	case result := <-lease.terminal:
		if result.GetStatus() != millv1.TerminalStatus_TERMINAL_STATUS_CANCELLED {
			t.Fatalf("terminal status = %v, want CANCELLED", result.GetStatus())
		}
	case <-time.After(time.Second):
		t.Fatal("unacknowledged cancel did not reach a bounded terminal")
	}
}

func TestAckedCancelFinishesOnlyItsLeaseAfterTeardownDeadline(t *testing.T) {
	m, bdb := restoreTestMill(t, Config{
		ReconnectGrace:        time.Minute,
		CancelAckTimeout:      30 * time.Millisecond,
		CancelTeardownTimeout: 150 * time.Millisecond,
	})
	registerTestExecutor(t, bdb, "node-1", HashToken("tok-1"), nil)

	sessClosed := make(chan struct{})
	sess := newSession("node-1", "inc-1", nil, nopEncoder(), discardLogger())
	sess.closeTransport = func() error {
		close(sessClosed)
		return nil
	}
	m.attachSession(sess)

	lease := newLease("lease-1", "node-1", "inc-1", "dummy")
	lease.wid = models.WorkflowId{PipelineId: models.PipelineId{Knot: "k", Rkey: "r"}, Name: "build"}
	lease.setState(leaseRunning)
	if err := m.persistLease(lease, leaseRowRunning); err != nil {
		t.Fatalf("persistLease: %v", err)
	}
	m.mu.Lock()
	m.leases[lease.id] = lease
	m.mu.Unlock()

	m.destroy(lease.wid)
	m.onCancelAck(sess, &millv1.CancelAck{LeaseId: lease.id})

	time.Sleep(80 * time.Millisecond)
	select {
	case <-sessClosed:
		t.Fatal("session was closed while an acked cancel was still tearing down")
	default:
	}

	select {
	case res := <-lease.terminal:
		if res.GetStatus() != millv1.TerminalStatus_TERMINAL_STATUS_CANCELLED {
			t.Fatalf("terminal status = %v, want CANCELLED", res.GetStatus())
		}
	case <-time.After(2 * time.Second):
		t.Fatal("teardown deadline passed without the mill finishing the lease")
	}

	select {
	case <-sessClosed:
		t.Fatal("teardown deadline closed the session, failing the node's other jobs")
	default:
	}
	if _, _, ok, err := bdb.ResolveExecutorToken(HashToken("tok-1")); err != nil || !ok {
		t.Fatalf("executor identity rejected for a slow teardown: ok = %v, err = %v", ok, err)
	}
}

func TestLateCancelAckTimerKeepsReconnectGrace(t *testing.T) {
	m, bdb := restoreTestMill(t, Config{ReconnectGrace: time.Minute, CancelAckTimeout: time.Minute})
	registerTestExecutor(t, bdb, "node-1", HashToken("tok-1"), nil)

	sess := newSession("node-1", "inc-1", nil, nopEncoder(), discardLogger())
	m.attachSession(sess)
	if err := m.onSnapshot(sess, &millv1.NodeSnapshot{Seqno: 1}); err != nil {
		t.Fatalf("onSnapshot: %v", err)
	}

	lease := newLease("lease-1", "node-1", "inc-1", "dummy")
	lease.wid = models.WorkflowId{PipelineId: models.PipelineId{Knot: "k", Rkey: "r"}, Name: "build"}
	lease.setState(leaseRunning)
	if err := m.persistLease(lease, leaseRowRunning); err != nil {
		t.Fatalf("persistLease: %v", err)
	}
	m.mu.Lock()
	m.leases[lease.id] = lease
	m.mu.Unlock()

	m.destroy(lease.wid)
	m.onCancelAck(sess, &millv1.CancelAck{LeaseId: lease.id})
	m.detachSession(sess)
	m.checkCancelAck(sess, lease)

	m.mu.Lock()
	recovering := sess.recovering
	timerArmed := sess.graceTimer != nil
	m.mu.Unlock()
	if !recovering || !timerArmed {
		t.Fatal("late cancel timer disarmed reconnect grace")
	}
	m.failLeasesAfterGrace(sess)
	select {
	case result := <-lease.terminal:
		if result.GetStatus() != millv1.TerminalStatus_TERMINAL_STATUS_CANCELLED {
			t.Fatalf("terminal status = %v, want CANCELLED", result.GetStatus())
		}
	case <-time.After(time.Second):
		t.Fatal("reconnect grace did not finish the disconnected lease")
	}
	if _, _, ok, err := bdb.ResolveExecutorToken(HashToken("tok-1")); err != nil || !ok {
		t.Fatalf("executor identity rejected for a cancel it could not answer: ok = %v, err = %v", ok, err)
	}
}

func TestCancelAckDeadlineDoesNotCloseReplacementSession(t *testing.T) {
	m, _ := restoreTestMill(t, Config{ReconnectGrace: time.Minute, CancelAckTimeout: time.Minute})
	oldSession := newSession("node-1", "inc-1", nil, nopEncoder(), discardLogger())
	m.attachSession(oldSession)
	lease := newLease("lease-1", "node-1", oldSession.epoch, "dummy")
	lease.wid = models.WorkflowId{PipelineId: models.PipelineId{Knot: "k", Rkey: "r"}, Name: "build"}
	lease.setState(leaseRunning)
	m.mu.Lock()
	m.leases[lease.id] = lease
	m.mu.Unlock()
	m.destroy(lease.wid)
	m.detachSession(oldSession)

	replacementClosed := make(chan struct{})
	replacement := newSession("node-1", "inc-1", nil, nopEncoder(), discardLogger())
	replacement.closeTransport = func() error {
		close(replacementClosed)
		return nil
	}
	if _, ok := m.attachSession(replacement); !ok {
		t.Fatal("replacement session was rejected")
	}
	replacement.snapshot = &millv1.NodeSnapshot{
		Seqno: 1,
		Engines: map[string]*millv1.EngineAvailability{
			"dummy": {Available: true},
		},
	}
	m.sessionReady(replacement)
	assertRankedNodes(t, m.rankCandidates("dummy", nil, false), nil)

	m.checkCancelAck(oldSession, lease)
	select {
	case <-replacementClosed:
		t.Fatal("cancel timer from old session closed its replacement")
	default:
	}
}

func TestCleanupRetry(t *testing.T) {
	m, bdb := restoreTestMill(t, Config{ReconnectGrace: time.Minute})

	lease := newLease("lease-1", "node-1", "inc-1", "dummy")
	lease.wid = models.WorkflowId{PipelineId: models.PipelineId{Knot: "k", Rkey: "r"}, Name: "build"}
	if err := m.persistLease(lease, leaseRowRunning); err != nil {
		t.Fatalf("persist lease: %v", err)
	}
	m.mu.Lock()
	m.leases[lease.id] = lease
	m.mu.Unlock()

	if _, err := bdb.Exec(`
		create trigger reject_cleanup_delete
		before delete on mill_leases
		begin
			select raise(abort, 'forced delete failure');
		end
	`); err != nil {
		t.Fatalf("failed to create fail trigger: %v", err)
	}

	err := m.cleanupLease(lease)
	if err == nil {
		t.Fatal("expected cleanupLease to fail")
	}

	m.mu.Lock()
	_, stillRunning := m.leases["lease-1"]
	m.mu.Unlock()
	if !stillRunning {
		t.Fatal("lease was removed from memory despite cleanup failure")
	}

	if _, err := bdb.Exec("drop trigger reject_cleanup_delete"); err != nil {
		t.Fatalf("drop trigger: %v", err)
	}

	err = m.cleanupLease(lease)
	if err != nil {
		t.Fatalf("expected retry cleanup to succeed, got: %v", err)
	}

	m.mu.Lock()
	_, stillRunning = m.leases["lease-1"]
	m.mu.Unlock()
	if stillRunning {
		t.Fatal("lease still in memory after successful cleanup retry")
	}
}

func TestBidPreservesConsistentTypedIncompatibility(t *testing.T) {
	m := New(discardLogger(), Config{TopK: 1, BidTimeout: time.Second})
	addReplyingCandidateSession(t, m, "node-1", nil, 0, nil, reserveReply{
		rejectClass:   millv1.RejectClass_REJECT_CLASS_INCOMPATIBLE,
		reason:        "invalid workflow",
		failureClass:  string(engine.FailureClassUser),
		failureReason: string(engine.FailureReasonWorkflowInvalid),
	})

	wid := models.WorkflowId{
		PipelineId: models.PipelineId{Knot: "knot.test", Rkey: "pipeline"},
		Name:       "build",
	}
	_, err := m.bid(context.Background(), "dummy", wid, testWorkflow("build"))
	if err == nil {
		t.Fatal("bid succeeded despite incompatible executor")
	}
	var failure *engine.WorkflowFailure
	if !errors.As(err, &failure) {
		t.Fatalf("bid error %T did not preserve typed attribution", err)
	}
	if failure.Class != engine.FailureClassUser || failure.Reason != engine.FailureReasonWorkflowInvalid {
		t.Fatalf("failure attribution = %s/%s", failure.Class, failure.Reason)
	}
}

func TestBoundedBidding(t *testing.T) {
	m := New(discardLogger(), Config{TopK: 2, BidTimeout: 10 * time.Millisecond})

	addCandidateSession(t, m, "node-1", nil, 0, nil)
	addCandidateSession(t, m, "node-2", nil, 0, nil)
	addCandidateSession(t, m, "node-3", nil, 0, nil)
	addCandidateSession(t, m, "node-4", nil, 0, nil)
	addCandidateSession(t, m, "node-5", nil, 0, nil)

	ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
	defer cancel()

	lease, err := m.bid(ctx, "dummy", models.WorkflowId{}, testWorkflow("build"))
	if err != nil {
		t.Fatalf("bid: %v", err)
	}
	if lease != nil {
		t.Fatalf("did not expect a lease, got %+v", lease)
	}
}

func TestProtocolErrorsDoNotRejectExecutorIdentity(t *testing.T) {
	m, bdb := restoreTestMill(t, Config{ReconnectGrace: time.Minute})
	registerTestExecutor(t, bdb, "node-1", HashToken("tok-1"), nil)
	sess := newSession("node-1", "inc-1", nil, nopEncoder(), discardLogger())

	for range 3 {
		m.noteSessionError(sess, protoErrf("gap in stream seqnos. Expected 1, got 9"))
	}
	if _, _, ok, err := bdb.ResolveExecutorToken(HashToken("tok-1")); err != nil || !ok {
		t.Fatalf("executor identity rejected after protocol errors: ok = %v, err = %v", ok, err)
	}
}

func TestAttachSessionRereadsDatabaseCursor(t *testing.T) {
	m, bdb := restoreTestMill(t, Config{ReconnectGrace: time.Minute})
	if err := bdb.ApplyEventBatch(nil, func(tx *db.EventBatchTx) error {
		return tx.AdvanceCursor("node-1", "inc-1", 3)
	}); err != nil {
		t.Fatalf("AdvanceCursor: %v", err)
	}

	sess := newSession("node-1", "inc-1", nil, nopEncoder(), discardLogger())
	if resume, ok := m.attachSession(sess); !ok || resume != 3 {
		t.Fatalf("attachSession resume = %d, %v; want 3 from the db cursor", resume, ok)
	}

	// an operator skip-forward must take effect on the next reconnect without
	// a mill restart
	if _, err := bdb.SetExecutorCursors("node-1", 28); err != nil {
		t.Fatalf("SetExecutorCursors: %v", err)
	}
	m.mu.Lock()
	sess.disconnected = true
	m.mu.Unlock()
	replacement := newSession("node-1", "inc-1", nil, nopEncoder(), discardLogger())
	if resume, ok := m.attachSession(replacement); !ok || resume != 28 {
		t.Fatalf("attachSession resume after reset = %d, %v; want 28", resume, ok)
	}
}