package executor

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"github.com/bluesky-social/indigo/atproto/syntax"
	"github.com/gorilla/websocket"
	"google.golang.org/protobuf/proto"
	"io"
	"log/slog"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"strings"
	"testing"
	"time"

	"tangled.org/core/api/tangled"
	"tangled.org/core/notifier"
	"tangled.org/core/spindle/config"
	"tangled.org/core/spindle/db"
	"tangled.org/core/spindle/engine"
	millproto "tangled.org/core/spindle/mill/proto"
	millv1 "tangled.org/core/spindle/mill/proto/gen"
	"tangled.org/core/spindle/models"
	"tangled.org/core/spindle/quota"
	"tangled.org/core/spindle/secrets"
	"tangled.org/core/spindle/storage"
)

type captureEncoder struct {
	messages chan *millproto.Message
}

func newCaptureEncoder() *captureEncoder {
	return &captureEncoder{messages: make(chan *millproto.Message, 4)}
}

func (e *captureEncoder) Encode(msg *millproto.Message) error {
	e.messages <- msg
	return nil
}

type blockingEncoder struct {
	started chan struct{}
	release chan struct{}
}

type artifactWriterFunc func(context.Context, string, io.Reader) error

func (f artifactWriterFunc) Put(ctx context.Context, ref string, r io.Reader) error {
	return f(ctx, ref, r)
}

func (e *blockingEncoder) Encode(*millproto.Message) error {
	select {
	case e.started <- struct{}{}:
	default:
	}
	<-e.release
	return nil
}

type fakeSlot struct {
	released int
	release  func()
}

func (s *fakeSlot) Release() {
	if s.release != nil {
		s.release()
	}
	s.released++
}

type fakeEngine struct {
	setupCalled    bool
	runCalled      bool
	destroyCalled  bool
	acquireCalled  bool
	acquireStarted chan struct{}
	releaseAcquire chan struct{}
	slot           engine.WorkflowSlot
	secrets        chan []secrets.UnlockedSecret
	done           chan struct{}
	initErr        error
}

func (e *fakeEngine) InitWorkflow(twf tangled.Pipeline_Workflow, tpl tangled.Pipeline) (*models.Workflow, error) {
	if e.initErr != nil {
		return nil, e.initErr
	}
	return &models.Workflow{Name: twf.Name}, nil
}
func (e *fakeEngine) SetupWorkflow(ctx context.Context, wid models.WorkflowId, wf *models.Workflow, l models.WorkflowLogger) error {
	e.setupCalled = true
	return nil
}
func (e *fakeEngine) WorkflowTimeout() time.Duration { return 7 * time.Minute }
func (e *fakeEngine) DestroyWorkflow(ctx context.Context, wid models.WorkflowId) error {
	e.destroyCalled = true
	return nil
}
func (e *fakeEngine) RunStep(ctx context.Context, wid models.WorkflowId, w *models.Workflow, idx int, s []secrets.UnlockedSecret, l models.WorkflowLogger) error {
	e.runCalled = true
	if e.secrets != nil {
		e.secrets <- s
	}
	if e.done != nil {
		close(e.done)
	}
	return nil
}
func (e *fakeEngine) AcquireWorkflowSlot(ctx context.Context, wid models.WorkflowId, wf *models.Workflow, mode engine.AcquireMode) (engine.WorkflowSlot, error) {
	e.acquireCalled = true
	if e.acquireStarted != nil {
		close(e.acquireStarted)
	}
	if e.releaseAcquire != nil {
		select {
		case <-e.releaseAcquire:
		case <-ctx.Done():
			return nil, ctx.Err()
		}
	}
	if e.slot != nil {
		return e.slot, nil
	}
	return &fakeSlot{}, nil
}

type resourceReportingEngine struct {
	*fakeEngine
	resources quota.Resources
	usage     engine.WorkflowResourceUsage
}

func (e *resourceReportingEngine) QuotaResources(*models.Workflow) quota.Resources {
	return e.resources
}

func (e *resourceReportingEngine) WorkflowResourceUsage(*models.Workflow) (engine.WorkflowResourceUsage, bool) {
	return e.usage, true
}

type fakeCacheEngine struct {
	*fakeEngine
	restored bool
	saved    bool
}

func (e *fakeCacheEngine) RestoreCache(context.Context, models.WorkflowId, *models.Workflow, storage.Storage, []models.CacheBinding, models.WorkflowLogger) error {
	e.restored = true
	return nil
}

func (e *fakeCacheEngine) SaveCache(context.Context, models.WorkflowId, *models.Workflow, storage.Storage, []models.CacheBinding, models.WorkflowLogger) error {
	e.saved = true
	return nil
}

type fakeStep struct{}

func (fakeStep) Name() string          { return "test" }
func (fakeStep) Command() string       { return "true" }
func (fakeStep) Kind() models.StepKind { return models.StepKindUser }

func testReserveSeat(t *testing.T, leaseID, engineName string) *millv1.ReserveSeat {
	t.Helper()
	twf, err := json.Marshal(tangled.Pipeline_Workflow{Name: "build"})
	if err != nil {
		t.Fatal(err)
	}
	repoDID := "did:web:example.com"
	tpl, err := json.Marshal(tangled.Pipeline{TriggerMetadata: &tangled.Pipeline_TriggerMetadata{
		Repo: &tangled.Pipeline_TriggerRepo{RepoDid: &repoDID},
	}})
	if err != nil {
		t.Fatal(err)
	}
	return &millv1.ReserveSeat{
		LeaseId:         leaseID,
		TargetEngine:    engineName,
		RawWorkflowJson: string(twf),
		RawPipelineJson: string(tpl),
		PipelineId:      "pipeline",
		RepoDid:         repoDID,
	}
}

func TestNewFailsWhenOutboxCannotInitialize(t *testing.T) {
	d := testDB(t)
	if err := d.Close(); err != nil {
		t.Fatal(err)
	}
	n := notifier.New()
	cfg := &config.Config{}
	if _, err := New(cfg, nil, d, &n, slog.New(slog.NewTextHandler(io.Discard, nil)), nil, nil); err == nil {
		t.Fatal("New succeeded with an unavailable outbox database")
	}
}

func TestReservedEngineHandsBackHeldSlotOnce(t *testing.T) {
	inner := &fakeEngine{}
	slot := &fakeSlot{}
	re := newReservedEngine(inner, slot)

	got, err := re.(engine.WorkflowSlotter).AcquireWorkflowSlot(context.Background(), models.WorkflowId{}, nil, engine.Wait)
	if err != nil {
		t.Fatal(err)
	}
	if got != engine.WorkflowSlot(slot) {
		t.Fatal("AcquireWorkflowSlot() did not return the held slot")
	}
	if inner.acquireCalled {
		t.Fatal("wrapper must not call the inner engine's AcquireWorkflowSlot")
	}

	if _, err := re.(engine.WorkflowSlotter).AcquireWorkflowSlot(context.Background(), models.WorkflowId{}, nil, engine.Wait); err == nil {
		t.Fatal("second AcquireWorkflowSlot() should error")
	}
}

func TestReservedEngineForwardsResourceReporting(t *testing.T) {
	inner := &resourceReportingEngine{
		fakeEngine: &fakeEngine{},
		resources: quota.Resources{
			quota.ResourceWorkflows: 1,
			quota.ResourceVCPUs:     4,
		},
		usage: engine.WorkflowResourceUsage{
			CPUUsec:              1234,
			MemoryPeakBytes:      4096,
			CgroupAvailable:      true,
			VolumeAllocatedBytes: 8192,
			VolumeAvailable:      true,
		},
	}
	re := newReservedEngine(inner, &fakeSlot{})
	wf := &models.Workflow{}

	resources := re.(engine.WorkflowQuotaReporter).QuotaResources(wf)
	if resources[quota.ResourceWorkflows] != 1 || resources[quota.ResourceVCPUs] != 4 {
		t.Fatalf("QuotaResources() = %v, want workflow and vCPU resources", resources)
	}

	usage, ok := re.(engine.WorkflowResourceUsageReporter).WorkflowResourceUsage(wf)
	if !ok {
		t.Fatal("WorkflowResourceUsage() unavailable")
	}
	if usage != inner.usage {
		t.Fatalf("WorkflowResourceUsage() = %+v, want %+v", usage, inner.usage)
	}
}

func TestReservedEngineForwardsCacheRunner(t *testing.T) {
	inner := &fakeCacheEngine{fakeEngine: &fakeEngine{}}
	re := newReservedEngine(inner, &fakeSlot{})
	runner, ok := re.(engine.CacheRunner)
	if !ok {
		t.Fatal("reserved engine dropped CacheRunner")
	}

	ctx := context.Background()
	if err := runner.RestoreCache(ctx, models.WorkflowId{}, nil, nil, nil, nil); err != nil {
		t.Fatal(err)
	}
	if err := runner.SaveCache(ctx, models.WorkflowId{}, nil, nil, nil, nil); err != nil {
		t.Fatal(err)
	}
	if !inner.restored || !inner.saved {
		t.Fatalf("cache calls were not forwarded: restored=%t saved=%t", inner.restored, inner.saved)
	}
}

func TestHandleCommitIsIdempotent(t *testing.T) {
	enc := newCaptureEncoder()
	e := testExecutor(t)
	e.enc = enc
	e.active["lease-1"] = &reservation{leaseID: "lease-1", committed: true}

	e.handleCommit(context.Background(), &millv1.CommitLease{LeaseId: "lease-1"})
	msg := <-enc.messages
	if got := msg.GetCommitted().GetLeaseId(); got != "lease-1" {
		t.Fatalf("Committed lease = %q, want lease-1", got)
	}
}

func TestHandleCommitRejectsMissingReservation(t *testing.T) {
	enc := newCaptureEncoder()
	e := testExecutor(t)
	e.enc = enc

	e.handleCommit(context.Background(), &millv1.CommitLease{LeaseId: "expired"})
	result := (<-enc.messages).GetReserveResult()
	if result == nil {
		t.Fatal("missing reservation commit did not receive a ReserveResult")
	}
	if result.GetLeaseId() != "expired" || result.GetAccepted() {
		t.Fatalf("ReserveResult = %+v, want correlated rejection", result)
	}
}

func TestHandleCancelFinalizesExpiredReservation(t *testing.T) {
	enc := newCaptureEncoder()
	e := testExecutor(t)
	e.enc = enc

	e.handleCancel("lease-expired")
	var ack *millv1.CancelAck
	for ack == nil {
		select {
		case msg := <-enc.messages:
			ack = msg.GetCancelAck()
		case <-time.After(time.Second):
			t.Fatal("cancel acknowledgement timed out")
		}
	}
	if ack.GetLeaseId() != "lease-expired" {
		t.Fatalf("CancelAck = %+v, want lease-expired", ack)
	}
	rows, err := e.db.ListOutboxRows()
	if err != nil {
		t.Fatal(err)
	}
	if len(rows) != 1 {
		t.Fatalf("cancel terminal outbox rows = %d, want 1", len(rows))
	}
	var entry millv1.Event
	if err := proto.Unmarshal(rows[0].Payload, &entry); err != nil {
		t.Fatal(err)
	}
	if got := entry.GetAttemptResult().GetStatus(); got != millv1.TerminalStatus_TERMINAL_STATUS_CANCELLED {
		t.Fatalf("cancel terminal = %v, want CANCELLED", got)
	}
}

func TestTerminalOutboxEntrySuppressesLaterEvents(t *testing.T) {
	e := testExecutor(t)
	if err := e.appendTerminal("lease-1", string(models.StatusKindSuccess), nil); err != nil {
		t.Fatal(err)
	}

	e.handleCancel("lease-1")
	if err := e.appendStatus("lease-1", &db.StatusRow{Status: string(models.StatusKindRunning)}); err != nil {
		t.Fatal(err)
	}

	rows, err := e.db.ListOutboxRows()
	if err != nil {
		t.Fatal(err)
	}
	if len(rows) != 1 {
		t.Fatalf("outbox rows after terminal = %d, want 1", len(rows))
	}
	var entry millv1.Event
	if err := proto.Unmarshal(rows[0].Payload, &entry); err != nil {
		t.Fatal(err)
	}
	if got := entry.GetAttemptResult().GetStatus(); got != millv1.TerminalStatus_TERMINAL_STATUS_SUCCESS {
		t.Fatalf("terminal status = %v, want SUCCESS", got)
	}
}

func TestTerminalGuardSurvivesExecutorRestart(t *testing.T) {
	d := testDB(t)
	first := &Executor{
		db:             d,
		l:              slog.New(slog.NewTextHandler(io.Discard, nil)),
		active:         make(map[string]*reservation),
		maxOutboxBytes: 10 * 1024 * 1024,
	}
	if err := first.initOutbox(); err != nil {
		t.Fatal(err)
	}
	if err := first.appendTerminal("lease-1", string(models.StatusKindSuccess), nil); err != nil {
		t.Fatal(err)
	}

	restarted := &Executor{
		db:             d,
		l:              slog.New(slog.NewTextHandler(io.Discard, nil)),
		active:         make(map[string]*reservation),
		maxOutboxBytes: 10 * 1024 * 1024,
	}
	if err := restarted.initOutbox(); err != nil {
		t.Fatal(err)
	}
	restarted.handleCancel("lease-1")

	rows, err := d.ListOutboxRows()
	if err != nil {
		t.Fatal(err)
	}
	if len(rows) != 1 {
		t.Fatalf("outbox rows after restart and duplicate cancel = %d, want 1", len(rows))
	}
}

func TestPendingArtifactRecoveryCompletesWithoutDeadlock(t *testing.T) {
	d := testDB(t)
	logDir := t.TempDir()
	wid := models.WorkflowId{
		PipelineId: models.PipelineId("3abc"),
		Name:       "build",
	}
	logPath := models.LogFilePath(logDir, wid)
	if err := os.MkdirAll(filepath.Dir(logPath), 0755); err != nil {
		t.Fatal(err)
	}
	if err := os.WriteFile(logPath, []byte("finished"), 0600); err != nil {
		t.Fatal(err)
	}
	if err := d.SavePendingArtifact(
		"lease-1", wid, string(models.StatusKindSuccess), "", 0,
		"logs/lease-1.log", "sha256:test", string(engine.FailureClassNone),
		string(engine.FailureReasonSuccess), true,
	); err != nil {
		t.Fatal(err)
	}

	uploaded := make(chan string, 1)
	e := &Executor{
		db:     d,
		cfg:    &config.Config{Server: config.Server{LogDir: logDir}},
		l:      slog.New(slog.NewTextHandler(io.Discard, nil)),
		active: make(map[string]*reservation),
		writer: artifactWriterFunc(func(_ context.Context, ref string, r io.Reader) error {
			if _, err := io.ReadAll(r); err != nil {
				return err
			}
			uploaded <- ref
			return nil
		}),
		maxOutboxBytes: 10 * 1024 * 1024,
	}
	if err := e.initOutbox(); err != nil {
		t.Fatalf("initOutbox: %v", err)
	}
	done := make(chan error, 1)
	go func() { done <- e.recoverPendingArtifacts(context.Background()) }()
	select {
	case err := <-done:
		if err != nil {
			t.Fatalf("recoverPendingArtifacts: %v", err)
		}
	case <-time.After(time.Second):
		t.Fatal("pending artifact recovery deadlocked")
	}
	select {
	case ref := <-uploaded:
		if ref != "logs/lease-1.log" {
			t.Fatalf("uploaded ref = %q", ref)
		}
	default:
		t.Fatal("pending artifact was not uploaded")
	}
	pending, err := d.ListPendingArtifacts()
	if err != nil {
		t.Fatal(err)
	}
	if len(pending) != 0 {
		t.Fatalf("pending artifacts after recovery = %d, want 0", len(pending))
	}
	rows, err := d.ListOutboxRows()
	if err != nil {
		t.Fatal(err)
	}
	if len(rows) != 1 {
		t.Fatalf("outbox rows after recovery = %d, want 1", len(rows))
	}
}

func TestPendingArtifactRecoveryKeepsRowWhenLogIsMissing(t *testing.T) {
	d := testDB(t)
	wid := models.WorkflowId{
		PipelineId: models.PipelineId("3abc"),
		Name:       "build",
	}
	if err := d.SavePendingArtifact(
		"lease-1", wid, string(models.StatusKindSuccess), "", 0,
		"logs/lease-1.log", "sha256:test", string(engine.FailureClassNone),
		string(engine.FailureReasonSuccess), true,
	); err != nil {
		t.Fatal(err)
	}
	e := &Executor{
		db:     d,
		cfg:    &config.Config{Server: config.Server{LogDir: t.TempDir()}},
		l:      slog.New(slog.NewTextHandler(io.Discard, nil)),
		active: make(map[string]*reservation),
		writer: artifactWriterFunc(func(context.Context, string, io.Reader) error {
			t.Fatal("writer called without a log file")
			return nil
		}),
		maxOutboxBytes: 10 * 1024 * 1024,
	}
	if err := e.initOutbox(); err != nil {
		t.Fatalf("initOutbox: %v", err)
	}
	if err := e.recoverPendingArtifacts(context.Background()); err != nil {
		t.Fatalf("recoverPendingArtifacts: %v", err)
	}
	pending, err := d.ListPendingArtifacts()
	if err != nil {
		t.Fatal(err)
	}
	if len(pending) != 1 {
		t.Fatalf("pending artifacts with missing log = %d, want 1", len(pending))
	}
	rows, err := d.ListOutboxRows()
	if err != nil {
		t.Fatal(err)
	}
	if len(rows) != 0 {
		t.Fatalf("outbox rows with missing log = %d, want 0", len(rows))
	}
}

func TestPendingArtifactRecoveryUsesAggregateDeadline(t *testing.T) {
	d := testDB(t)
	logDir := t.TempDir()
	wids := []models.WorkflowId{
		{PipelineId: models.PipelineId("3abc"), Name: "build"},
		{PipelineId: models.PipelineId("3def"), Name: "test"},
	}
	for i, wid := range wids {
		logPath := models.LogFilePath(logDir, wid)
		if err := os.MkdirAll(filepath.Dir(logPath), 0755); err != nil {
			t.Fatal(err)
		}
		if err := os.WriteFile(logPath, []byte("finished"), 0600); err != nil {
			t.Fatal(err)
		}
		leaseID := fmt.Sprintf("lease-%d", i+1)
		if err := d.SavePendingArtifact(
			leaseID, wid, string(models.StatusKindSuccess), "", 0,
			"logs/"+leaseID+".log", "sha256:test", string(engine.FailureClassNone),
			string(engine.FailureReasonSuccess), true,
		); err != nil {
			t.Fatal(err)
		}
	}

	var deadlines []time.Time
	e := &Executor{
		db:     d,
		cfg:    &config.Config{Server: config.Server{LogDir: logDir}},
		l:      slog.New(slog.NewTextHandler(io.Discard, nil)),
		active: make(map[string]*reservation),
		writer: artifactWriterFunc(func(ctx context.Context, _ string, _ io.Reader) error {
			deadline, ok := ctx.Deadline()
			if !ok {
				t.Fatal("artifact recovery context has no deadline")
			}
			deadlines = append(deadlines, deadline)
			return errors.New("unavailable")
		}),
		maxOutboxBytes: 10 * 1024 * 1024,
	}
	if err := e.initOutbox(); err != nil {
		t.Fatalf("initOutbox: %v", err)
	}
	started := time.Now()
	if err := e.recoverPendingArtifacts(context.Background()); err != nil {
		t.Fatalf("recoverPendingArtifacts: %v", err)
	}
	finished := time.Now()
	if len(deadlines) != len(wids) {
		t.Fatalf("upload attempts = %d, want %d", len(deadlines), len(wids))
	}
	if !deadlines[0].Equal(deadlines[1]) {
		t.Fatalf("artifact recovery deadlines differ: %v, %v", deadlines[0], deadlines[1])
	}
	wantTimeout := 2 * time.Minute
	if deadlines[0].Before(started.Add(wantTimeout)) || deadlines[0].After(finished.Add(wantTimeout)) {
		t.Fatalf("artifact recovery deadline = %v, want %v after startup", deadlines[0], wantTimeout)
	}
}

func TestHandleCommitPreservesPreauthorizedSecrets(t *testing.T) {
	d := testDB(t)
	n := notifier.New()
	enc := newCaptureEncoder()
	e := &Executor{
		cfg:            &config.Config{Server: config.Server{LogDir: t.TempDir()}},
		db:             d,
		n:              &n,
		l:              slog.New(slog.NewTextHandler(io.Discard, nil)),
		active:         make(map[string]*reservation),
		maxOutboxBytes: 10 * 1024 * 1024,
		enc:            enc,
	}
	if err := e.initOutbox(); err != nil {
		t.Fatal(err)
	}
	e.lifecycleCtx = context.Background()

	inner := &fakeEngine{secrets: make(chan []secrets.UnlockedSecret, 1), done: make(chan struct{})}
	slot := &fakeSlot{}
	repoDid, err := syntax.ParseDID("did:web:example.com")
	if err != nil {
		t.Fatal(err)
	}
	repoDidString := repoDid.String()
	res := &reservation{
		leaseID:    "lease-1",
		wid:        models.WorkflowId{PipelineId: models.PipelineId("r"), Name: "build"},
		realEngine: inner,
		slot:       slot,
		wf:         &models.Workflow{Name: "build", Steps: []models.Step{fakeStep{}}},
		pipeline: &models.Pipeline{
			RepoDid: repoDid,
			TriggerMetadata: &tangled.Pipeline_TriggerMetadata{
				Repo: &tangled.Pipeline_TriggerRepo{RepoDid: &repoDidString},
			},
			TrustedSource: true,
		},
	}
	e.active[res.leaseID] = res

	e.handleCommit(context.Background(), &millv1.CommitLease{
		LeaseId: res.leaseID,
		Secrets: []*millv1.Secret{{Key: "TOKEN", Value: "secret-value"}},
	})
	if got := (<-enc.messages).GetCommitted().GetLeaseId(); got != res.leaseID {
		t.Fatalf("Committed lease = %q, want %q", got, res.leaseID)
	}
	if got := e.maskSecrets(res, "value=secret-value"); got != "value=***" {
		t.Fatalf("masked log = %q", got)
	}
	select {
	case got := <-inner.secrets:
		if len(got) != 1 || got[0].Key != "TOKEN" || got[0].Value != "secret-value" {
			t.Fatalf("RunStep secrets = %+v", got)
		}
	case <-time.After(2 * time.Second):
		t.Fatal("RunStep did not receive CommitLease secrets")
	}
	select {
	case <-inner.done:
	case <-time.After(2 * time.Second):
		t.Fatal("workflow did not finish")
	}
	if res.stopTail != nil {
		res.stopTail()
	}
	e.jobsWG.Wait()
}

func TestConnectHandshakesWhilePendingArtifactRecoveryIsBlocked(t *testing.T) {
	handshake := make(chan error, 1)
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		conn, err := websocket.Upgrade(w, r, nil, 1024, 1024)
		if err != nil {
			handshake <- err
			return
		}
		defer conn.Close()
		stream := millproto.NewWSStream(conn)
		dec := millproto.NewDecoder(stream)
		enc := millproto.NewEncoder(stream)
		helloMsg, err := dec.Decode()
		if err != nil {
			handshake <- err
			return
		}
		hello := helloMsg.GetHello()
		if hello == nil {
			handshake <- errors.New("first executor message was not Hello")
			return
		}
		if err := enc.Encode(&millproto.Message{Resume: &millv1.Resume{Epoch: hello.GetEpoch()}}); err != nil {
			handshake <- err
			return
		}
		snapshotMsg, err := dec.Decode()
		if err != nil {
			handshake <- err
			return
		}
		if snapshotMsg.GetNodeSnapshot() == nil {
			handshake <- errors.New("executor did not send NodeSnapshot after Resume")
			return
		}
		handshake <- nil
		for {
			if _, err := dec.Decode(); err != nil {
				return
			}
		}
	}))
	t.Cleanup(srv.Close)

	e := testSessionExecutor(t, "ws"+strings.TrimPrefix(srv.URL, "http"))
	n := notifier.New()
	e.n = &n
	logDir := t.TempDir()
	e.cfg.Server.LogDir = logDir
	wid := models.WorkflowId{
		PipelineId: models.PipelineId("3abc"),
		Name:       "build",
	}
	logPath := models.LogFilePath(logDir, wid)
	if err := os.MkdirAll(filepath.Dir(logPath), 0755); err != nil {
		t.Fatal(err)
	}
	if err := os.WriteFile(logPath, []byte("finished"), 0600); err != nil {
		t.Fatal(err)
	}
	if err := e.db.SavePendingArtifact(
		"lease-1", wid, string(models.StatusKindSuccess), "", 0,
		"logs/lease-1.log", "sha256:test", string(engine.FailureClassNone),
		string(engine.FailureReasonSuccess), true,
	); err != nil {
		t.Fatal(err)
	}
	recoveryStarted := make(chan struct{})
	e.writer = artifactWriterFunc(func(ctx context.Context, _ string, _ io.Reader) error {
		close(recoveryStarted)
		<-ctx.Done()
		return ctx.Err()
	})

	ctx, cancel := context.WithCancel(context.Background())
	done := make(chan struct{})
	go func() {
		e.Connect(ctx)
		close(done)
	}()
	select {
	case <-recoveryStarted:
	case <-time.After(time.Second):
		cancel()
		t.Fatal("pending artifact recovery did not start")
	}
	select {
	case err := <-handshake:
		if err != nil {
			cancel()
			t.Fatalf("mill handshake: %v", err)
		}
	case <-time.After(time.Second):
		cancel()
		t.Fatal("blocked artifact recovery delayed the mill handshake")
	}
	cancel()
	select {
	case <-done:
	case <-time.After(2 * time.Second):
		t.Fatal("Connect did not stop after cancellation")
	}
}

func TestHandleCommitRejectsCachesAndSecretsForForkSource(t *testing.T) {
	d := testDB(t)
	n := notifier.New()
	enc := newCaptureEncoder()
	cache, err := storage.NewDisk(t.TempDir())
	if err != nil {
		t.Fatal(err)
	}
	e := &Executor{
		db:             d,
		n:              &n,
		l:              slog.New(slog.NewTextHandler(io.Discard, nil)),
		active:         make(map[string]*reservation),
		maxOutboxBytes: 10 * 1024 * 1024,
		cfg:            &config.Config{Server: config.Server{LogDir: t.TempDir()}},
		cache:          cache,
		enc:            enc,
	}
	if err := e.initOutbox(); err != nil {
		t.Fatal(err)
	}
	e.lifecycleCtx = context.Background()

	inner := &fakeCacheEngine{fakeEngine: &fakeEngine{
		secrets: make(chan []secrets.UnlockedSecret, 1),
		done:    make(chan struct{}),
	}}
	targetRepoDid, err := syntax.ParseDID("did:plc:target")
	if err != nil {
		t.Fatal(err)
	}
	sourceRepoDid := "did:plc:fork"
	cacheObjectID := "11111111-1111-4111-8111-111111111111"
	res := &reservation{
		leaseID:    "lease-1",
		wid:        models.WorkflowId{PipelineId: models.PipelineId("r"), Name: "build"},
		realEngine: inner,
		slot:       &fakeSlot{},
		wf: &models.Workflow{
			Name:   "build",
			Steps:  []models.Step{fakeStep{}},
			Caches: []models.CacheEntry{{Key: "deps", Paths: []string{"deps"}}},
		},
		pipeline: &models.Pipeline{
			RepoDid: targetRepoDid,
			TriggerMetadata: &tangled.Pipeline_TriggerMetadata{
				SourceRepo: &sourceRepoDid,
				Repo:       &tangled.Pipeline_TriggerRepo{RepoDid: &sourceRepoDid},
			},
		},
	}
	e.active[res.leaseID] = res

	e.handleCommit(context.Background(), &millv1.CommitLease{
		LeaseId: res.leaseID,
		Secrets: []*millv1.Secret{{Key: "TOKEN", Value: "secret-value"}},
		CacheBindings: []*millv1.CacheBinding{{
			EntryIndex: 0,
			RestoreId:  cacheObjectID,
			RestoreKey: "objects/" + targetRepoDid.String() + "/" + cacheObjectID,
			SaveId:     cacheObjectID,
			SaveKey:    "objects/" + targetRepoDid.String() + "/" + cacheObjectID,
		}},
	})
	if got := (<-enc.messages).GetCommitted().GetLeaseId(); got != res.leaseID {
		t.Fatalf("Committed lease = %q, want %q", got, res.leaseID)
	}
	e.jobsWG.Wait()

	if inner.restored || inner.saved {
		t.Fatalf("fork cache access = restored %t, saved %t", inner.restored, inner.saved)
	}
	if got := <-inner.secrets; len(got) != 0 {
		t.Fatalf("fork secrets = %+v, want none", got)
	}
}

func TestRunSessionCancellationClosesStalledWebsocket(t *testing.T) {
	connected := make(chan struct{})
	release := make(chan struct{})
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		conn, err := websocket.Upgrade(w, r, nil, 1024, 1024)
		if err != nil {
			return
		}
		defer conn.Close()
		close(connected)
		<-release
	}))
	t.Cleanup(func() {
		close(release)
		srv.Close()
	})

	e := testSessionExecutor(t, "ws"+strings.TrimPrefix(srv.URL, "http"))
	ctx, cancel := context.WithCancel(context.Background())
	done := make(chan error, 1)
	go func() { done <- e.runSession(ctx) }()
	<-connected
	cancel()

	select {
	case <-done:
	case <-time.After(2 * time.Second):
		t.Fatal("runSession did not return after context cancellation")
	}
}

func testSessionExecutor(t *testing.T, url string) *Executor {
	d := testDB(t)
	e := &Executor{
		millURL:        url,
		seats:          1,
		engines:        make(map[string]models.Engine),
		db:             d,
		cfg:            &config.Config{Server: config.Server{Dev: true}},
		l:              slog.New(slog.NewTextHandler(io.Discard, nil)),
		active:         make(map[string]*reservation),
		maxOutboxBytes: 10 * 1024 * 1024,
	}
	if err := e.initOutbox(); err != nil {
		t.Fatal(err)
	}
	return e
}

func testExecutor(t *testing.T) *Executor {
	d := testDB(t)
	e := &Executor{
		db:             d,
		l:              slog.New(slog.NewTextHandler(io.Discard, nil)),
		active:         make(map[string]*reservation),
		maxOutboxBytes: 10 * 1024 * 1024,
	}
	if err := e.initOutbox(); err != nil {
		t.Fatal(err)
	}
	return e
}

func TestReserveHonorsExecutorSeatCapacity(t *testing.T) {
	enc := newCaptureEncoder()
	e := testExecutor(t)
	e.enc = enc
	e.seats = 1
	e.engines = map[string]models.Engine{"dummy": &fakeEngine{}}

	e.handleReserve(context.Background(), testReserveSeat(t, "lease-1", "dummy"))
	first := (<-enc.messages).GetReserveResult()
	if first == nil || !first.GetAccepted() {
		t.Fatalf("first reserve result = %+v, want accepted", first)
	}
	if snapshot := (<-enc.messages).GetNodeSnapshot(); snapshot == nil {
		t.Fatal("accepted reservation did not publish a snapshot")
	}

	e.handleReserve(context.Background(), testReserveSeat(t, "lease-2", "dummy"))
	second := (<-enc.messages).GetReserveResult()
	if second == nil || second.GetAccepted() {
		t.Fatalf("second reserve result = %+v, want transient rejection", second)
	}
	if second.GetRejectClass() != millv1.RejectClass_REJECT_CLASS_TRANSIENT {
		t.Fatalf("second reject class = %v, want transient", second.GetRejectClass())
	}
	if second.GetRejectReason() != "no executor seats available" {
		t.Fatalf("second reject reason = %q", second.GetRejectReason())
	}

	cleanup, ok := e.takeUncommittedReservation("lease-1", true)
	if !ok {
		t.Fatal("first reservation missing during cleanup")
	}
	cleanup()
}

func testDB(t *testing.T) *db.DB {
	d, err := db.Make(context.Background(), filepath.Join(t.TempDir(), "spindle.db"))
	if err != nil {
		t.Fatal(err)
	}
	return d
}

func TestFinishJobReportsCancelledReservationAsCancelled(t *testing.T) {
	d := testDB(t)
	res := &reservation{
		leaseID:   "lease-1",
		wid:       models.WorkflowId{PipelineId: models.PipelineId("r"), Name: "build"},
		cancelled: true,
	}
	e := &Executor{
		db:             d,
		l:              slog.New(slog.NewTextHandler(io.Discard, nil)),
		active:         map[string]*reservation{res.leaseID: res},
		maxOutboxBytes: 10 * 1024 * 1024,
	}
	if err := e.initOutbox(); err != nil {
		t.Fatal(err)
	}

	e.finishJob(res, &db.StatusRow{
		Pipeline: string(res.wid.PipelineId),
		Workflow: res.wid.Name,
		Status:   string(models.StatusKindFailed),
	})

	rows, err := d.ListOutboxRows()
	if err != nil {
		t.Fatal(err)
	}
	if len(rows) != 1 {
		t.Fatalf("outbox rows = %d, want 1", len(rows))
	}

	var entry millv1.Event
	if err := proto.Unmarshal(rows[0].Payload, &entry); err != nil {
		t.Fatal(err)
	}
	got := entry.GetAttemptResult().GetStatus()
	if got != millv1.TerminalStatus_TERMINAL_STATUS_CANCELLED {
		t.Fatalf("terminal status = %v, want CANCELLED", got)
	}
}

func TestFinishJobWaitsForEngineCleanup(t *testing.T) {
	d := testDB(t)
	res := &reservation{
		leaseID: "lease-1",
		wid: models.WorkflowId{
			PipelineId: models.PipelineId("r"),
			Name:       "build",
		},
		runDone: make(chan struct{}),
	}
	e := &Executor{
		db:             d,
		l:              slog.New(slog.NewTextHandler(io.Discard, nil)),
		active:         map[string]*reservation{res.leaseID: res},
		maxOutboxBytes: 10 * 1024 * 1024,
	}
	if err := e.initOutbox(); err != nil {
		t.Fatal(err)
	}

	done := make(chan error, 1)
	go func() {
		done <- e.finishJob(res, &db.StatusRow{
			Pipeline: string(res.wid.PipelineId),
			Workflow: res.wid.Name,
			Status:   string(models.StatusKindSuccess),
		})
	}()

	select {
	case err := <-done:
		t.Fatalf("finishJob returned before engine cleanup: %v", err)
	case <-time.After(50 * time.Millisecond):
	}
	rows, err := d.ListOutboxRows()
	if err != nil {
		t.Fatal(err)
	}
	if len(rows) != 0 {
		t.Fatalf("terminal rows before engine cleanup = %d, want 0", len(rows))
	}

	close(res.runDone)
	select {
	case err := <-done:
		if err != nil {
			t.Fatal(err)
		}
	case <-time.After(time.Second):
		t.Fatal("finishJob did not resume after engine cleanup")
	}
}

func TestReplayRejectsMalformedOutboxRow(t *testing.T) {
	d := testDB(t)
	e := &Executor{
		db:  d,
		enc: newCaptureEncoder(),
		l:   slog.New(slog.NewTextHandler(io.Discard, nil)),
	}
	if err := e.initOutbox(); err != nil {
		t.Fatal(err)
	}
	if _, err := d.AppendOutboxRow([]byte("not protobuf"), true); err != nil {
		t.Fatal(err)
	}
	if err := e.replay(0); err == nil {
		t.Fatal("replay accepted a malformed row and would leave a permanent seqno gap")
	}
}

func TestReplayDropsRowsAlreadyAcknowledgedByResume(t *testing.T) {
	e := testExecutor(t)
	if err := e.appendTerminal("lease-1", string(models.StatusKindSuccess), nil); err != nil {
		t.Fatal(err)
	}
	rows, err := e.db.ListOutboxRows()
	if err != nil {
		t.Fatal(err)
	}
	if len(rows) != 1 {
		t.Fatalf("outbox rows = %d, want 1", len(rows))
	}

	e.enc = newCaptureEncoder()
	if err := e.replay(rows[0].Seqno); err != nil {
		t.Fatalf("replay: %v", err)
	}
	rows, err = e.db.ListOutboxRows()
	if err != nil {
		t.Fatal(err)
	}
	if len(rows) != 0 {
		t.Fatalf("outbox rows after resume = %d, want 0", len(rows))
	}
	if len(e.terminalSeqnos) != 0 {
		t.Fatalf("terminal guards after resume = %d, want 0", len(e.terminalSeqnos))
	}
	select {
	case <-e.outboxIdle():
	default:
		t.Fatal("resumed acknowledgement left the outbox busy")
	}
}

func TestSocketCancellationIndependence(t *testing.T) {
	d := testDB(t)
	n := notifier.New()
	e := &Executor{
		db:             d,
		n:              &n,
		l:              slog.New(slog.NewTextHandler(io.Discard, nil)),
		active:         make(map[string]*reservation),
		maxOutboxBytes: 10 * 1024 * 1024,
		cfg:            &config.Config{Server: config.Server{LogDir: t.TempDir()}},
	}
	if err := e.initOutbox(); err != nil {
		t.Fatal(err)
	}

	lifecycleCtx, cancelLifecycle := context.WithCancel(context.Background())
	defer cancelLifecycle()
	e.lifecycleCtx = lifecycleCtx

	inner := &fakeEngine{secrets: make(chan []secrets.UnlockedSecret, 1), done: make(chan struct{})}
	slot := &fakeSlot{}
	repoDid, err := syntax.ParseDID("did:web:example.com")
	if err != nil {
		t.Fatal(err)
	}
	res := &reservation{
		leaseID:    "lease-1",
		wid:        models.WorkflowId{PipelineId: models.PipelineId("r"), Name: "build"},
		realEngine: inner,
		slot:       slot,
		wf:         &models.Workflow{Name: "build", Steps: []models.Step{fakeStep{}}},
		pipeline:   &models.Pipeline{RepoDid: repoDid},
	}
	e.active[res.leaseID] = res

	sessionCtx, cancelSession := context.WithCancel(lifecycleCtx)

	e.handleCommit(sessionCtx, &millv1.CommitLease{
		LeaseId: "lease-1",
	})

	cancelSession()

	// session disconnect must not cancel the running job
	select {
	case <-inner.done:
	case <-time.After(2 * time.Second):
		t.Fatal("workflow did not complete even though websocket session was cancelled")
	}

	e.jobsWG.Wait()
}

func TestMonotonicSnapshots(t *testing.T) {
	enc := newCaptureEncoder()
	e := &Executor{
		enc:            enc,
		l:              slog.New(slog.NewTextHandler(io.Discard, nil)),
		active:         make(map[string]*reservation),
		maxOutboxBytes: 10 * 1024 * 1024,
	}

	e.pushSnapshot()
	msg1 := <-enc.messages
	seq1 := msg1.GetNodeSnapshot().GetSeqno()
	if seq1 != 1 {
		t.Fatalf("first seq = %d, want 1", seq1)
	}

	e.pushSnapshot()
	msg2 := <-enc.messages
	seq2 := msg2.GetNodeSnapshot().GetSeqno()
	if seq2 != 2 {
		t.Fatalf("second seq = %d, want 2", seq2)
	}
}

func TestDrainWaitsForActiveReservation(t *testing.T) {
	enc := newCaptureEncoder()
	e := testExecutor(t)
	e.enc = enc
	e.engines = map[string]models.Engine{"dummy": &fakeEngine{}}

	res := &reservation{leaseID: "lease-1", committed: true}
	e.active[res.leaseID] = res
	done := make(chan error, 1)
	go func() {
		done <- e.Drain(context.Background())
	}()

	snapshot := (<-enc.messages).GetNodeSnapshot()
	if snapshot == nil || snapshot.GetEngines()["dummy"].GetAvailable() {
		t.Fatalf("drain snapshot = %+v, want unavailable engine", snapshot)
	}
	if got := snapshot.GetActiveLeaseIds(); len(got) != 1 || got[0] != res.leaseID {
		t.Fatalf("active leases = %v, want [%s]", got, res.leaseID)
	}
	select {
	case err := <-done:
		t.Fatalf("Drain returned before reservation cleanup: %v", err)
	default:
	}

	e.mu.Lock()
	cleanup := e.removeReservationLocked(res, false)
	e.mu.Unlock()
	cleanup()
	select {
	case err := <-done:
		if err != nil {
			t.Fatalf("Drain: %v", err)
		}
	case <-time.After(2 * time.Second):
		t.Fatal("Drain did not return after reservation cleanup")
	}
}

func TestDrainWaitsForReservationCleanup(t *testing.T) {
	releaseStarted := make(chan struct{})
	releaseSlot := make(chan struct{})
	slot := &fakeSlot{release: func() {
		close(releaseStarted)
		<-releaseSlot
	}}
	e := testExecutor(t)
	e.enc = newCaptureEncoder()
	res := &reservation{leaseID: "lease-1", slot: slot}
	e.active[res.leaseID] = res

	cleanup, ok := e.takeUncommittedReservation(res.leaseID, true)
	if !ok {
		t.Fatal("failed to take reservation")
	}
	cleanupDone := make(chan struct{})
	go func() {
		cleanup()
		close(cleanupDone)
	}()
	<-releaseStarted

	done := make(chan error, 1)
	go func() {
		done <- e.Drain(context.Background())
	}()
	<-e.enc.(*captureEncoder).messages
	select {
	case err := <-done:
		t.Fatalf("Drain returned before slot release completed: %v", err)
	default:
	}

	close(releaseSlot)
	<-cleanupDone
	select {
	case err := <-done:
		if err != nil {
			t.Fatalf("Drain: %v", err)
		}
	case <-time.After(2 * time.Second):
		t.Fatal("Drain did not return after slot release")
	}
	if slot.released != 1 {
		t.Fatalf("slot releases = %d, want 1", slot.released)
	}
}

func TestDrainWaitsForWorkflowCleanup(t *testing.T) {
	enc := newCaptureEncoder()
	e := testExecutor(t)
	e.enc = enc
	e.jobsWG.Add(1)

	done := make(chan error, 1)
	go func() {
		done <- e.Drain(context.Background())
	}()
	<-enc.messages
	select {
	case err := <-done:
		t.Fatalf("Drain returned before workflow cleanup: %v", err)
	default:
	}

	e.jobsWG.Done()
	select {
	case err := <-done:
		if err != nil {
			t.Fatalf("Drain: %v", err)
		}
	case <-time.After(2 * time.Second):
		t.Fatal("Drain did not return after workflow cleanup")
	}
}

func TestDrainWaitsForTerminalAcknowledgement(t *testing.T) {
	enc := newCaptureEncoder()
	e := testExecutor(t)
	e.enc = enc
	if err := e.appendTerminal("lease-1", string(models.StatusKindSuccess), nil); err != nil {
		t.Fatal(err)
	}
	batch := (<-enc.messages).GetEventBatch()
	if batch == nil || len(batch.GetEvents()) != 1 {
		t.Fatalf("terminal batch = %+v, want one event", batch)
	}

	done := make(chan error, 1)
	go func() {
		done <- e.Drain(context.Background())
	}()
	<-enc.messages
	select {
	case err := <-done:
		t.Fatalf("Drain returned before terminal acknowledgement: %v", err)
	default:
	}

	e.handleAck(&millv1.Ack{Epoch: e.epoch, UpToSeqno: batch.GetEvents()[0].GetSeqno()})
	select {
	case err := <-done:
		if err != nil {
			t.Fatalf("Drain: %v", err)
		}
	case <-time.After(2 * time.Second):
		t.Fatal("Drain did not return after terminal acknowledgement")
	}
}

func TestDrainReturnsWithDurableOutboxAfterReconnectGrace(t *testing.T) {
	e := testExecutor(t)
	e.cfg = &config.Config{}
	e.cfg.Mill.ReconnectGrace = 10 * time.Millisecond
	if err := e.appendTerminal("lease-1", string(models.StatusKindSuccess), nil); err != nil {
		t.Fatal(err)
	}

	if err := e.Drain(context.Background()); err != nil {
		t.Fatalf("Drain with disconnected Mill: %v", err)
	}

	rows, err := e.db.ListOutboxRows()
	if err != nil {
		t.Fatal(err)
	}
	if len(rows) != 1 {
		t.Fatalf("durable outbox rows = %d, want 1", len(rows))
	}
}

func TestDrainAllowsReconnectToAcknowledgeOutbox(t *testing.T) {
	enc := newCaptureEncoder()
	e := testExecutor(t)
	e.cfg = &config.Config{}
	e.cfg.Mill.ReconnectGrace = 2 * time.Second
	e.enc = enc
	if err := e.appendTerminal("lease-1", string(models.StatusKindSuccess), nil); err != nil {
		t.Fatal(err)
	}
	batch := (<-enc.messages).GetEventBatch()
	if batch == nil || len(batch.GetEvents()) != 1 {
		t.Fatalf("terminal batch = %+v, want one event", batch)
	}

	done := make(chan error, 1)
	go func() {
		done <- e.Drain(context.Background())
	}()
	<-enc.messages
	e.connMu.Lock()
	e.enc = nil
	e.connMu.Unlock()
	select {
	case err := <-done:
		t.Fatalf("Drain returned before reconnect grace elapsed: %v", err)
	case <-time.After(50 * time.Millisecond):
	}

	e.connMu.Lock()
	e.enc = enc
	e.connMu.Unlock()
	e.handleAck(&millv1.Ack{Epoch: e.epoch, UpToSeqno: batch.GetEvents()[0].GetSeqno()})
	select {
	case err := <-done:
		if err != nil {
			t.Fatalf("Drain after reconnect acknowledgement: %v", err)
		}
	case <-time.After(2 * time.Second):
		t.Fatal("Drain did not accept acknowledgement during reconnect grace")
	}

	rows, err := e.db.ListOutboxRows()
	if err != nil {
		t.Fatal(err)
	}
	if len(rows) != 0 {
		t.Fatalf("durable outbox rows = %d, want 0", len(rows))
	}
}

func TestReserveCannotRacePastDrain(t *testing.T) {
	enc := newCaptureEncoder()
	e := testExecutor(t)
	e.enc = enc
	slot := &fakeSlot{}
	inner := &fakeEngine{
		acquireStarted: make(chan struct{}),
		releaseAcquire: make(chan struct{}),
		slot:           slot,
	}
	e.engines = map[string]models.Engine{"dummy": inner}

	reserve := testReserveSeat(t, "lease-1", "dummy")
	reserveDone := make(chan struct{})
	go func() {
		e.handleReserve(context.Background(), reserve)
		close(reserveDone)
	}()
	<-inner.acquireStarted

	if err := e.Drain(context.Background()); err != nil {
		t.Fatalf("Drain: %v", err)
	}
	snapshot := (<-enc.messages).GetNodeSnapshot()
	if snapshot == nil || snapshot.GetEngines()["dummy"].GetAvailable() {
		t.Fatalf("drain snapshot = %+v, want unavailable engine", snapshot)
	}
	close(inner.releaseAcquire)
	<-reserveDone

	rejected := (<-enc.messages).GetReserveResult()
	if rejected == nil || rejected.GetAccepted() || rejected.GetRejectReason() != "draining" {
		t.Fatalf("reserve result = %+v, want draining rejection", rejected)
	}
	if slot.released != 1 {
		t.Fatalf("slot releases = %d, want 1", slot.released)
	}
	if len(e.active) != 0 {
		t.Fatalf("active reservations = %d, want 0", len(e.active))
	}
}

func TestDrainCancellationIncludesUnavailableSnapshot(t *testing.T) {
	enc := &blockingEncoder{
		started: make(chan struct{}, 1),
		release: make(chan struct{}),
	}
	e := testExecutor(t)
	e.enc = enc
	ctx, cancel := context.WithCancel(context.Background())
	done := make(chan error, 1)
	go func() {
		done <- e.Drain(ctx)
	}()

	<-enc.started
	cancel()
	select {
	case err := <-done:
		if !errors.Is(err, context.Canceled) {
			t.Fatalf("Drain error = %v, want context canceled", err)
		}
	case <-time.After(2 * time.Second):
		t.Fatal("Drain ignored cancellation while sending the unavailable snapshot")
	}
	close(enc.release)
}

func TestTimerRace(t *testing.T) {
	d := testDB(t)
	e := &Executor{
		db:             d,
		l:              slog.New(slog.NewTextHandler(io.Discard, nil)),
		active:         make(map[string]*reservation),
		maxOutboxBytes: 10 * 1024 * 1024,
	}
	if err := e.initOutbox(); err != nil {
		t.Fatal(err)
	}

	twf, _ := json.Marshal(tangled.Pipeline_Workflow{Name: "build"})
	tpl, _ := json.Marshal(tangled.Pipeline{TriggerMetadata: &tangled.Pipeline_TriggerMetadata{}})

	inner := &fakeEngine{}
	e.engines = map[string]models.Engine{"microvm": inner}

	e.handleReserve(context.Background(), &millv1.ReserveSeat{
		LeaseId:         "lease-1",
		TargetEngine:    "microvm",
		RawWorkflowJson: string(twf),
		RawPipelineJson: string(tpl),
		PipelineId:      "pipeline",
		TtlSeconds:      1,
		RepoDid:         "did:web:example.com",
	})

	e.mu.Lock()
	res := e.active["lease-1"]
	e.mu.Unlock()

	if res == nil {
		t.Fatal("reservation was not added")
	}

	deadline := time.Now().Add(5 * time.Second)
	for {
		e.mu.Lock()
		activeLen := len(e.active)
		e.mu.Unlock()
		if activeLen == 0 {
			break
		}
		if time.Now().After(deadline) {
			t.Fatal("reservation was leaked and never expired")
		}
		time.Sleep(10 * time.Millisecond)
	}
}

func TestStructuredShutdown(t *testing.T) {
	d := testDB(t)
	n := notifier.New()
	e := &Executor{
		db:             d,
		n:              &n,
		l:              slog.New(slog.NewTextHandler(io.Discard, nil)),
		active:         make(map[string]*reservation),
		maxOutboxBytes: 10 * 1024 * 1024,
		cfg:            &config.Config{Server: config.Server{LogDir: t.TempDir()}},
	}
	if err := e.initOutbox(); err != nil {
		t.Fatal(err)
	}

	ctx, cancel := context.WithCancel(context.Background())
	e.lifecycleCtx = ctx

	inner := &fakeEngine{secrets: make(chan []secrets.UnlockedSecret, 1), done: make(chan struct{})}
	slot := &fakeSlot{}
	repoDid, err := syntax.ParseDID("did:web:example.com")
	if err != nil {
		t.Fatal(err)
	}
	res := &reservation{
		leaseID:    "lease-1",
		wid:        models.WorkflowId{PipelineId: models.PipelineId("r"), Name: "build"},
		realEngine: inner,
		slot:       slot,
		wf:         &models.Workflow{Name: "build", Steps: []models.Step{fakeStep{}}},
		pipeline:   &models.Pipeline{RepoDid: repoDid},
	}
	e.active[res.leaseID] = res

	e.handleCommit(ctx, &millv1.CommitLease{
		LeaseId: "lease-1",
	})

	cancel()

	e.jobsWG.Wait()

	select {
	case <-inner.done:
	default:
		t.Fatal("shutdown returned but running job did not finish")
	}
}