package dispatcher

import (
	"context"
	"errors"
	"os"
	"path/filepath"
	"sync"
	"testing"
	"time"

	gittransport "github.com/go-git/go-git/v5/plumbing/transport/http"
	"github.com/pingcap/failpoint"
	"github.com/stretchr/testify/require"

	"github.com/tidbcloud/AgentTeam/desktop/internal/agent"
	"github.com/tidbcloud/AgentTeam/desktop/internal/agent/drivercontract"
	"github.com/tidbcloud/AgentTeam/desktop/internal/agent/runtimeconfig"
	"github.com/tidbcloud/AgentTeam/desktop/internal/memorywatcher"
	"github.com/tidbcloud/AgentTeam/desktop/internal/profile"
	"github.com/tidbcloud/AgentTeam/desktop/internal/types"
)

// blockMemoryWatcherPublication installs the narrow production seam immediately
// after Watcher.Start and immediately before the dispatcher publication fence.
// Keeping the block there makes these tests cover the real provider lookup,
// clone/rebind, watcher construction, and shutdown entry points rather than a
// lifecycle helper in isolation.
func blockMemoryWatcherPublication(t *testing.T, d *Dispatcher) (<-chan struct{}, func()) {
	t.Helper()
	entered := make(chan struct{})
	releaseCh := make(chan struct{})
	var enteredOnce sync.Once
	var releaseOnce sync.Once
	release := func() { releaseOnce.Do(func() { close(releaseCh) }) }
	d.memoryWatcherBeforePublish = func(context.Context, *interactiveSession) {
		enteredOnce.Do(func() { close(entered) })
		<-releaseCh
	}
	t.Cleanup(release)
	return entered, release
}

func blockMemoryWatcherShutdownAfterDetach(
	t *testing.T,
	d *Dispatcher,
	target *interactiveSession,
) (<-chan struct{}, func()) {
	t.Helper()
	entered := make(chan struct{})
	releaseCh := make(chan struct{})
	var enteredOnce sync.Once
	var releaseOnce sync.Once
	release := func() { releaseOnce.Do(func() { close(releaseCh) }) }
	d.memoryWatcherShutdownAfterDetach = func(sess *interactiveSession) {
		if sess != target {
			return
		}
		enteredOnce.Do(func() { close(entered) })
		<-releaseCh
	}
	t.Cleanup(release)
	return entered, release
}

func blockMemoryWatcherFinalRemovalAfterFence(
	t *testing.T,
	d *Dispatcher,
	target *interactiveSession,
) (<-chan struct{}, func()) {
	t.Helper()
	entered := make(chan struct{})
	releaseCh := make(chan struct{})
	var enteredOnce sync.Once
	var releaseOnce sync.Once
	release := func() { releaseOnce.Do(func() { close(releaseCh) }) }
	d.memoryWatcherFinalRemovalAfterFence = func(sess *interactiveSession) {
		if sess != target {
			return
		}
		enteredOnce.Do(func() { close(entered) })
		<-releaseCh
	}
	t.Cleanup(release)
	return entered, release
}

func waitForMemoryWatcherTestSignal(t *testing.T, signal <-chan struct{}, what string) {
	t.Helper()
	select {
	case <-signal:
	case <-time.After(5 * time.Second):
		t.Fatalf("timed out waiting for %s", what)
	}
}

func memoryWatcherSessionFields(sess *interactiveSession) (watcherPresent, clonePresent, authPresent bool) {
	if sess == nil {
		return false, false, false
	}
	sess.mu.Lock()
	defer sess.mu.Unlock()
	return sess.watcher != nil, sess.memoryClone != nil, sess.memoryPushAuth != nil
}

func memorySearchAdmissionClosed(d *Dispatcher) bool {
	d.memorySearchMu.RLock()
	defer d.memorySearchMu.RUnlock()
	return d.memorySearchClosed
}

func TestDispatcherCloseFencesBlockedUnpublishedMemoryWatcherAttach(t *testing.T) {
	origin := makeBareOriginRepoForMemory(t, filepath.Join(t.TempDir(), "origin"))
	workDir := t.TempDir()
	agentID := "agt_blocked_attach_close"
	provider := &fakeMemoryAccessProvider{}
	provider.set(agentID, MemoryAccess{
		MemoryID:  "mem_blocked_attach_close",
		OriginURL: origin,
		Branch:    "main",
		Auth:      &gittransport.BasicAuth{Username: "memory", Password: "attach-close"},
	})
	d := newDispatcherWithMemoryProvider(t, provider)
	sess := &interactiveSession{
		agentID: agentID,
		agent:   agent.Config{Name: "Blocked Attach"},
		workDir: workDir,
	}
	t.Cleanup(func() { stopMemoryWatcher(sess) })

	entered, release := blockMemoryWatcherPublication(t, d)
	attachDone := make(chan struct{})
	go func() {
		d.attachMemoryWatcher(context.Background(), sess, agentID, workDir)
		close(attachDone)
	}()
	waitForMemoryWatcherTestSignal(t, entered, "blocked watcher attach publication")

	closeDone := make(chan struct{})
	go func() {
		d.Close()
		close(closeDone)
	}()

	closeReturnedBeforeRelease := false
	select {
	case <-closeDone:
		closeReturnedBeforeRelease = true
	case <-time.After(75 * time.Millisecond):
	}
	release()
	waitForMemoryWatcherTestSignal(t, attachDone, "blocked watcher attach cleanup")
	waitForMemoryWatcherTestSignal(t, closeDone, "dispatcher close after watcher owner cleanup")

	require.False(t, closeReturnedBeforeRelease,
		"Close must wait for an admitted watcher mutation that has not reached publication")
	watcherPresent, clonePresent, authPresent := memoryWatcherSessionFields(sess)
	require.False(t, watcherPresent, "shutdown must reject late watcher publication")
	require.False(t, clonePresent, "shutdown must reject late clone publication")
	require.False(t, authPresent, "shutdown must reject late push-auth publication")
}

func TestDispatcherCloseFencesBlockedMemoryWatcherRebind(t *testing.T) {
	originA := makeBareOriginRepoForMemory(t, filepath.Join(t.TempDir(), "origin-a"))
	originB := makeBareOriginRepoForMemory(t, filepath.Join(t.TempDir(), "origin-b"))
	workDir := t.TempDir()
	agentID := "agt_blocked_rebind_close"
	provider := &fakeMemoryAccessProvider{}
	provider.set(agentID, MemoryAccess{
		MemoryID:  "mem_blocked_rebind_a",
		OriginURL: originA,
		Branch:    "main",
		Auth:      &gittransport.BasicAuth{Username: "memory", Password: "rebind-a"},
	})
	d := newDispatcherWithMemoryProvider(t, provider)
	sess := &interactiveSession{
		agentID: agentID,
		agent:   agent.Config{Name: "Blocked Rebind"},
		workDir: workDir,
	}
	d.sessionManager.Set(agentID, sess)
	d.attachMemoryWatcher(context.Background(), sess, agentID, workDir)
	t.Cleanup(func() { stopMemoryWatcher(sess) })
	watcherPresent, clonePresent, authPresent := memoryWatcherSessionFields(sess)
	require.True(t, watcherPresent, "precondition: initial watcher must be published")
	require.True(t, clonePresent, "precondition: initial clone must be published")
	require.True(t, authPresent, "precondition: initial push auth must be published")

	provider.set(agentID, MemoryAccess{
		MemoryID:                    "mem_blocked_rebind_b",
		OriginURL:                   originB,
		Branch:                      "main",
		ReplaceLocalContentOnSwitch: true,
		Auth:                        &gittransport.BasicAuth{Username: "memory", Password: "rebind-b"},
	})
	entered, release := blockMemoryWatcherPublication(t, d)
	refreshDone := make(chan struct{})
	go func() {
		d.RefreshMemoryAccess(context.Background(), []string{agentID})
		close(refreshDone)
	}()
	waitForMemoryWatcherTestSignal(t, entered, "blocked watcher rebind publication")

	closeDone := make(chan struct{})
	go func() {
		d.Close()
		close(closeDone)
	}()
	closeReturnedBeforeRelease := false
	select {
	case <-closeDone:
		closeReturnedBeforeRelease = true
	case <-time.After(75 * time.Millisecond):
	}
	release()
	waitForMemoryWatcherTestSignal(t, refreshDone, "blocked watcher rebind cleanup")
	waitForMemoryWatcherTestSignal(t, closeDone, "dispatcher close after rebind owner cleanup")

	require.False(t, closeReturnedBeforeRelease,
		"Close must wait for an admitted watcher rebind that has not reached publication")
	require.True(t, sess.retired.Load(), "session manager close must retire the rebind target")
	watcherPresent, clonePresent, authPresent = memoryWatcherSessionFields(sess)
	require.False(t, watcherPresent, "retired session must not receive the replacement watcher")
	require.False(t, clonePresent, "retired session must not receive the replacement clone")
	require.False(t, authPresent, "retired session must not receive replacement push auth")
}

func TestDispatcherCloseJoinsDetachedWatcherOnlySessionDrain(t *testing.T) {
	origin := makeBareOriginRepoForMemory(t, filepath.Join(t.TempDir(), "origin"))
	workDir := t.TempDir()
	agentID := "agt_close_detached_watcher_only_drain"
	provider := &fakeMemoryAccessProvider{}
	provider.set(agentID, MemoryAccess{
		MemoryID:  "mem_close_detached_watcher_only_drain",
		OriginURL: origin,
		Branch:    "main",
	})
	d := newDispatcherWithMemoryProvider(t, provider)
	sess := &interactiveSession{
		sessionKey: agentID,
		agentID:    agentID,
		agent:      agent.Config{Name: "Detached Watcher-Only Drain"},
		workDir:    workDir,
	}
	d.sessionManager.Set(agentID, sess)
	d.attachMemoryWatcher(context.Background(), sess, agentID, workDir)
	t.Cleanup(func() { stopMemoryWatcher(sess) })

	drainEntered := make(chan struct{})
	releaseDrain := make(chan struct{})
	var enteredOnce sync.Once
	var releaseOnce sync.Once
	release := func() { releaseOnce.Do(func() { close(releaseDrain) }) }
	t.Cleanup(release)
	drainDone := make(chan struct{})
	go func() {
		_ = stopMemoryWatcherWithTimeoutResultAfterDetach(sess, 60*time.Second, func(*interactiveSession) {
			enteredOnce.Do(func() { close(drainEntered) })
			<-releaseDrain
		})
		close(drainDone)
	}()
	waitForMemoryWatcherTestSignal(t, drainEntered, "detached watcher-only drain")

	watcherPresent, clonePresent, authPresent := memoryWatcherSessionFields(sess)
	require.False(t, watcherPresent, "precondition: the active drain must have detached the watcher")
	require.False(t, clonePresent, "precondition: the active drain must have detached the clone")
	require.False(t, authPresent, "precondition: the active drain must have detached push auth")

	closeDone := make(chan struct{})
	go func() {
		d.Close()
		close(closeDone)
	}()
	select {
	case <-closeDone:
		t.Fatal("Close returned without joining the detached watcher-only drain")
	case <-time.After(75 * time.Millisecond):
	}
	require.False(t, memorySearchAdmissionClosed(d),
		"search admission must remain open until the detached watcher drain completes")

	release()
	waitForMemoryWatcherTestSignal(t, drainDone, "detached watcher-only drain completion")
	waitForMemoryWatcherTestSignal(t, closeDone, "Close after detached watcher-only drain")
	require.True(t, memorySearchAdmissionClosed(d),
		"Close must disable search after joining the detached watcher drain")
}

func TestStopForShutdownDeadlineContinuesMemoryWatcherFenceAndDrain(t *testing.T) {
	origin := makeBareOriginRepoForMemory(t, filepath.Join(t.TempDir(), "origin"))
	workDir := t.TempDir()
	agentID := "agt_blocked_attach_bounded_shutdown"
	provider := &fakeMemoryAccessProvider{}
	provider.set(agentID, MemoryAccess{
		MemoryID:  "mem_blocked_attach_bounded_shutdown",
		OriginURL: origin,
		Branch:    "main",
		Auth:      &gittransport.BasicAuth{Username: "memory", Password: "bounded-shutdown"},
	})
	d := newDispatcherWithMemoryProvider(t, provider)
	sess := &interactiveSession{
		agentID: agentID,
		agent:   agent.Config{Name: "Bounded Attach"},
		workDir: workDir,
	}
	d.sessionManager.Set(agentID, sess)
	t.Cleanup(func() { stopMemoryWatcher(sess) })

	entered, release := blockMemoryWatcherPublication(t, d)
	refreshDone := make(chan struct{})
	go func() {
		d.RefreshMemoryAccess(context.Background(), []string{agentID})
		close(refreshDone)
	}()
	waitForMemoryWatcherTestSignal(t, entered, "blocked watcher attach before bounded shutdown")

	ctx, cancel := context.WithTimeout(context.Background(), 80*time.Millisecond)
	defer cancel()
	startedAt := time.Now()
	d.StopForShutdown(ctx)
	elapsed := time.Since(startedAt)
	require.ErrorIs(t, ctx.Err(), context.DeadlineExceeded)
	require.GreaterOrEqual(t, elapsed, 50*time.Millisecond,
		"bounded shutdown should remain pending on the admitted watcher owner until its deadline")
	require.Less(t, elapsed, time.Second,
		"the admitted watcher owner must not extend the caller's shutdown deadline")
	require.False(t, memorySearchAdmissionClosed(d),
		"memory search must stay enabled while an admitted watcher can still reach publication cleanup")

	release()
	waitForMemoryWatcherTestSignal(t, refreshDone, "late watcher publication rejection")
	require.Eventually(t, func() bool {
		return memorySearchAdmissionClosed(d)
	}, 5*time.Second, 10*time.Millisecond,
		"best-effort shutdown must disable memory search after owner cleanup, even after the caller deadline")

	require.True(t, sess.retired.Load(), "bounded shutdown must retire the attach target")
	watcherPresent, clonePresent, authPresent := memoryWatcherSessionFields(sess)
	require.False(t, watcherPresent, "release after the shutdown deadline must not publish a watcher")
	require.False(t, clonePresent, "release after the shutdown deadline must not publish a clone")
	require.False(t, authPresent, "release after the shutdown deadline must not publish push auth")
}

func TestStopForShutdownDoesNotDisableSearchWhileDetachedWatcherDrainIsBlocked(t *testing.T) {
	origin := makeBareOriginRepoForMemory(t, filepath.Join(t.TempDir(), "origin"))
	workDir := t.TempDir()
	agentID := "agt_blocked_detached_drain"
	provider := &fakeMemoryAccessProvider{}
	provider.set(agentID, MemoryAccess{
		MemoryID:  "mem_blocked_detached_drain",
		OriginURL: origin,
		Branch:    "main",
		Auth:      &gittransport.BasicAuth{Username: "memory", Password: "detached-drain"},
	})
	d := newDispatcherWithMemoryProvider(t, provider)
	process, driver := newShutdownDeadlineProcess(t, agentID)
	sess := &interactiveSession{
		agentID: agentID,
		agent:   agent.Config{Name: "Blocked Detached Drain"},
		workDir: workDir,
		process: process,
	}
	d.sessionManager.Set(agentID, sess)
	d.attachMemoryWatcher(context.Background(), sess, agentID, workDir)
	t.Cleanup(func() { stopMemoryWatcher(sess) })
	watcherPresent, clonePresent, authPresent := memoryWatcherSessionFields(sess)
	require.True(t, watcherPresent, "precondition: watcher must be published")
	require.True(t, clonePresent, "precondition: clone must be published")
	require.True(t, authPresent, "precondition: push auth must be published")

	// Hold an unpublished session-start owner across shutdown. The authoritative
	// memory task must wait for this owner before taking its final snapshot. A
	// process task that independently drained the watcher would reach the seam
	// early, detach the session fields, and let the final snapshot falsely look
	// clean — the exact detached-drain race this regression test protects.
	startOwner, startInvalidated, ok := d.sessionManager.BeginStart("blocked-unpublished-start", nil)
	require.True(t, ok, "precondition: unpublished start owner must be admitted")
	t.Cleanup(func() { d.sessionManager.FinishStart(startOwner) })

	drainEntered, releaseDrain := blockMemoryWatcherShutdownAfterDetach(t, d, sess)
	ctx, cancel := context.WithTimeout(context.Background(), 220*time.Millisecond)
	defer cancel()
	stopDone := make(chan struct{})
	startedAt := time.Now()
	go func() {
		d.StopForShutdown(ctx)
		close(stopDone)
	}()
	waitForMemoryWatcherTestSignal(t, startInvalidated, "session-start admission invalidation")

	select {
	case <-drainEntered:
		t.Fatal("an interactive process task detached the watcher before session-start owners drained")
	case <-time.After(75 * time.Millisecond):
	}
	d.sessionManager.FinishStart(startOwner)
	waitForMemoryWatcherTestSignal(t, drainEntered, "authoritative watcher detach after owner drain")

	waitForMemoryWatcherTestSignal(t, stopDone, "bounded shutdown deadline while watcher drain is parked")
	require.ErrorIs(t, ctx.Err(), context.DeadlineExceeded)
	require.GreaterOrEqual(t, time.Since(startedAt), 150*time.Millisecond,
		"parked watcher drain should remain under the caller's total shutdown deadline")
	require.False(t, memorySearchAdmissionClosed(d),
		"memory search must stay enabled after session fields detach and until FlushAndClose completes")
	watcherPresent, clonePresent, authPresent = memoryWatcherSessionFields(sess)
	require.False(t, watcherPresent, "authoritative shutdown drain should detach the watcher before flushing")
	require.False(t, clonePresent, "authoritative shutdown drain should detach the clone before flushing")
	require.False(t, authPresent, "authoritative shutdown drain should detach push auth before flushing")

	releaseDrain()
	driver.allowExit()
	require.Eventually(t, func() bool {
		return memorySearchAdmissionClosed(d)
	}, 5*time.Second, 10*time.Millisecond,
		"memory search must disable only after the parked detached watcher finishes draining")
}

func TestHandleProcessExitDrainsWatcherBeforeFinalSessionRemoval(t *testing.T) {
	origin := makeBareOriginRepoForMemory(t, filepath.Join(t.TempDir(), "origin"))
	workDir := t.TempDir()
	agentID := "agt_process_exit_watcher_drain"
	provider := &fakeMemoryAccessProvider{}
	provider.set(agentID, MemoryAccess{
		MemoryID:  "mem_process_exit_watcher_drain",
		OriginURL: origin,
		Branch:    "main",
	})
	d := newDispatcherWithMemoryProvider(t, provider)
	process := &agent.Process{AgentID: agentID}
	sess := &interactiveSession{
		sessionKey: agentID,
		agentID:    agentID,
		agent:      agent.Config{Name: "Process Exit Drain"},
		workDir:    workDir,
		process:    process,
		ephemeral:  true,
	}
	d.sessionManager.Set(agentID, sess)
	d.attachMemoryWatcher(context.Background(), sess, agentID, workDir)
	t.Cleanup(func() { stopMemoryWatcher(sess) })
	watcherPresent, clonePresent, _ := memoryWatcherSessionFields(sess)
	require.True(t, watcherPresent, "precondition: watcher must be published")
	require.True(t, clonePresent, "precondition: clone must be published")

	d.handleProcessExit(sess, process)

	_, stillPublished := d.sessionManager.Get(agentID)
	require.False(t, stillPublished, "terminal process exit must remove the session")
	watcherPresent, clonePresent, authPresent := memoryWatcherSessionFields(sess)
	require.False(t, watcherPresent, "terminal process exit must drain the watcher before removal")
	require.False(t, clonePresent, "terminal process exit must clear the clone before removal")
	require.False(t, authPresent, "terminal process exit must clear push auth before removal")
}

func TestFinalSessionRemovalRejectsBlockedReplacementWatcherPublication(t *testing.T) {
	originA := makeBareOriginRepoForMemory(t, filepath.Join(t.TempDir(), "origin-a"))
	originB := makeBareOriginRepoForMemory(t, filepath.Join(t.TempDir(), "origin-b"))
	workDir := t.TempDir()
	agentID := "agt_final_remove_blocks_replacement"
	provider := &fakeMemoryAccessProvider{}
	provider.set(agentID, MemoryAccess{
		MemoryID:  "mem_final_remove_a",
		OriginURL: originA,
		Branch:    "main",
	})
	d := newDispatcherWithMemoryProvider(t, provider)
	process := &agent.Process{AgentID: agentID}
	sess := &interactiveSession{
		sessionKey: agentID,
		agentID:    agentID,
		agent:      agent.Config{Name: "Final Remove Fence"},
		workDir:    workDir,
		process:    process,
		ephemeral:  true,
	}
	d.sessionManager.Set(agentID, sess)
	d.attachMemoryWatcher(context.Background(), sess, agentID, workDir)
	t.Cleanup(func() { stopMemoryWatcher(sess) })

	provider.set(agentID, MemoryAccess{
		MemoryID:                    "mem_final_remove_b",
		OriginURL:                   originB,
		Branch:                      "main",
		ReplaceLocalContentOnSwitch: true,
	})
	publicationEntered, releasePublication := blockMemoryWatcherPublication(t, d)
	refreshDone := make(chan struct{})
	go func() {
		d.RefreshMemoryAccess(context.Background(), []string{agentID})
		close(refreshDone)
	}()
	waitForMemoryWatcherTestSignal(t, publicationEntered, "blocked replacement watcher publication")

	// The replacement watcher has started but is not published. Final removal
	// must retire the session before its drain, otherwise the drain observes nil
	// and the blocked refresh can publish into the removal gap.
	fenceEntered, releaseRemoval := blockMemoryWatcherFinalRemovalAfterFence(t, d, sess)
	exitDone := make(chan struct{})
	go func() {
		d.handleProcessExit(sess, process)
		close(exitDone)
	}()
	waitForMemoryWatcherTestSignal(t, fenceEntered, "final-removal watcher fence")
	_, stillPublished := d.sessionManager.Get(agentID)
	require.True(t, stillPublished,
		"test must release watcher publication after retirement but before manager removal")
	sess.mu.Lock()
	removalFenced := sess.memoryWatcherTerminalFenced
	sess.mu.Unlock()
	require.True(t, removalFenced, "final removal must fence watcher publication before draining")
	require.False(t, sess.retired.Load(),
		"slow watcher drain must not retire a manager-visible session before atomic removal")

	releasePublication()
	waitForMemoryWatcherTestSignal(t, refreshDone, "replacement watcher rejection after final removal")
	watcherPresent, clonePresent, authPresent := memoryWatcherSessionFields(sess)
	require.False(t, watcherPresent, "replacement watcher must not publish after final removal")
	require.False(t, clonePresent, "replacement clone must not publish after final removal")
	require.False(t, authPresent, "replacement auth must not publish after final removal")

	releaseRemoval()
	waitForMemoryWatcherTestSignal(t, exitDone, "final manager removal")
	_, stillPublished = d.sessionManager.Get(agentID)
	require.False(t, stillPublished, "terminal process exit must remove the retired session")
}

func TestResetSessionDrainsWatcherOnlySessionBeforeWorkspaceRemoval(t *testing.T) {
	origin := makeBareOriginRepoForMemory(t, filepath.Join(t.TempDir(), "origin"))
	workDir := filepath.Join(t.TempDir(), "agent-workdir")
	agentID := "agt_reset_watcher_only"
	provider := &fakeMemoryAccessProvider{}
	provider.set(agentID, MemoryAccess{
		MemoryID:  "mem_reset_watcher_only",
		OriginURL: origin,
		Branch:    "main",
	})
	d := newDispatcherWithMemoryProvider(t, provider)
	sess := &interactiveSession{
		sessionKey: agentID,
		agentID:    agentID,
		agent:      agent.Config{Name: "Reset Watcher Only"},
		workDir:    workDir,
	}
	d.sessionManager.Set(agentID, sess)
	d.attachMemoryWatcher(context.Background(), sess, agentID, workDir)
	t.Cleanup(func() { stopMemoryWatcher(sess) })
	watcherPresent, clonePresent, _ := memoryWatcherSessionFields(sess)
	require.True(t, watcherPresent, "precondition: watcher must be published")
	require.True(t, clonePresent, "precondition: clone must be published")

	require.NoError(t, d.resetSession(sess))

	_, stillPublished := d.sessionManager.Get(agentID)
	require.False(t, stillPublished, "reset must remove the watcher-only session")
	watcherPresent, clonePresent, authPresent := memoryWatcherSessionFields(sess)
	require.False(t, watcherPresent, "reset must drain a watcher-only session")
	require.False(t, clonePresent, "reset must clear a watcher-only clone")
	require.False(t, authPresent, "reset must clear watcher-only push auth")
	_, statErr := os.Stat(workDir)
	require.ErrorIs(t, statErr, os.ErrNotExist,
		"reset must remove the workspace only after the watcher has drained")
}

func TestResetSessionDoesNotRetryFailedFinalPushAfterWorkspaceRemoval(t *testing.T) {
	originDir := filepath.Join(t.TempDir(), "origin")
	origin := makeBareOriginRepoForMemory(t, originDir)
	workDir := filepath.Join(t.TempDir(), "agent-workdir")
	agentID := "agt_reset_failed_push"
	provider := &fakeMemoryAccessProvider{}
	provider.set(agentID, MemoryAccess{
		MemoryID:  "mem_reset_failed_push",
		OriginURL: origin,
		Branch:    "main",
	})
	d := newDispatcherWithMemoryProvider(t, provider)
	sess := &interactiveSession{
		sessionKey: agentID,
		agentID:    agentID,
		agent:      agent.Config{Name: "Reset Failed Push"},
		workDir:    workDir,
	}
	d.sessionManager.Set(agentID, sess)
	d.attachMemoryWatcher(t.Context(), sess, agentID, workDir)
	t.Cleanup(func() { stopMemoryWatcher(sess) })

	offlineOrigin := originDir + "-offline"
	require.NoError(t, os.Rename(originDir, offlineOrigin))
	require.NoError(t, os.WriteFile(filepath.Join(workDir, "memory", "pending-reset.md"), []byte("pending reset\n"), 0o644))
	require.Eventually(t, func() bool {
		status, ok := d.MemorySyncStatus(agentID)
		return ok && status.State == memorywatcher.StatusPendingSync
	}, 3*time.Second, 25*time.Millisecond)
	d.setMemoryLookupStatus(agentID, degradedMemorySyncStatus(memorywatcher.StatusReasonLookupFailed))

	require.NoError(t, d.resetSession(sess))
	_, ok := d.MemorySyncStatus(agentID)
	require.False(t, ok, "full reset must clear transient Memory sync state")
	d.memorySyncStatusMu.RLock()
	_, syncStatusPresent := d.memorySyncStatuses[agentID]
	_, lookupStatusPresent := d.memoryLookupStatuses[agentID]
	d.memorySyncStatusMu.RUnlock()
	require.False(t, syncStatusPresent, "full reset must clear cached watcher status")
	require.False(t, lookupStatusPresent, "full reset must clear cached lookup status")
	d.RetryMemorySync(t.Context(), agentID)
	require.Never(t, func() bool {
		_, err := os.Stat(workDir)
		return err == nil
	}, 2*idleMemoryRetryMin, 25*time.Millisecond,
		"full reset must not let a stale retry recreate the deleted workspace")
}

func TestReplacementAttachDrainsOldWatcherBeforeNewWatcherPublication(t *testing.T) {
	origin := makeBareOriginRepoForMemory(t, filepath.Join(t.TempDir(), "origin"))
	workDir := t.TempDir()
	agentID := "agt_replacement_drain_before_attach"
	provider := &fakeMemoryAccessProvider{}
	provider.set(agentID, MemoryAccess{
		MemoryID:  "mem_replacement_drain_before_attach",
		OriginURL: origin,
		Branch:    "main",
	})
	d := newDispatcherWithMemoryProvider(t, provider)
	oldSession := &interactiveSession{
		sessionKey: agentID,
		agentID:    agentID,
		agent:      agent.Config{ID: agentID, Name: "Old Replacement Session"},
		workDir:    workDir,
	}
	newSession := &interactiveSession{
		sessionKey: agentID,
		agentID:    agentID,
		agent:      agent.Config{ID: agentID, Name: "New Replacement Session"},
		workDir:    workDir,
	}
	d.attachMemoryWatcher(context.Background(), oldSession, agentID, workDir)
	t.Cleanup(func() {
		stopMemoryWatcher(newSession)
		stopMemoryWatcher(oldSession)
	})
	oldWatcher, oldClone, _ := memoryWatcherSessionFields(oldSession)
	require.True(t, oldWatcher, "precondition: old watcher must be published")
	require.True(t, oldClone, "precondition: old clone must be published")

	oldDetached := make(chan struct{})
	releaseOldDrain := make(chan struct{})
	var detachedOnce sync.Once
	var releaseOnce sync.Once
	releaseDrain := func() { releaseOnce.Do(func() { close(releaseOldDrain) }) }
	t.Cleanup(releaseDrain)
	d.memoryWatcherReplacementAfterOldDetach = func(sess *interactiveSession) {
		if sess != oldSession {
			return
		}
		detachedOnce.Do(func() { close(oldDetached) })
		<-releaseOldDrain
	}
	newBeforePublish := make(chan struct{})
	var publishOnce sync.Once
	d.memoryWatcherBeforePublish = func(_ context.Context, sess *interactiveSession) {
		if sess == newSession {
			publishOnce.Do(func() { close(newBeforePublish) })
		}
	}

	replacementDone := make(chan *memoryWatcherReplacement, 1)
	go func() {
		replacementDone <- d.attachMemoryWatcherReplacing(
			context.Background(), newSession, oldSession, agentID, workDir,
		)
	}()
	waitForMemoryWatcherTestSignal(t, oldDetached, "old replacement watcher detach")

	select {
	case <-newBeforePublish:
		t.Fatal("new watcher reached publication before the old watcher drain completed")
	case <-time.After(100 * time.Millisecond):
	}
	newWatcher, newClone, newAuth := memoryWatcherSessionFields(newSession)
	require.False(t, newWatcher, "new watcher must remain unpublished while old drain is parked")
	require.False(t, newClone, "new clone must remain unpublished while old drain is parked")
	require.False(t, newAuth, "new auth must remain unpublished while old drain is parked")

	releaseDrain()
	waitForMemoryWatcherTestSignal(t, newBeforePublish, "new watcher publication after old drain")
	var replacement *memoryWatcherReplacement
	select {
	case replacement = <-replacementDone:
	case <-time.After(5 * time.Second):
		t.Fatal("timed out waiting for replacement attach")
	}
	require.NotNil(t, replacement)
	require.Same(t, oldSession, replacement.previous)

	oldSession.mu.Lock()
	oldDrain := oldSession.memoryWatcherDrain
	oldFenced := oldSession.memoryWatcherReplacementFenced
	oldSession.mu.Unlock()
	oldWatcher, oldClone, oldAuth := memoryWatcherSessionFields(oldSession)
	require.False(t, oldWatcher, "successful replacement must leave the old watcher empty")
	require.False(t, oldClone, "successful replacement must leave the old clone empty")
	require.False(t, oldAuth, "successful replacement must leave the old auth empty")
	require.Nil(t, oldDrain, "successful replacement must fully join the old watcher drain")
	require.True(t, oldFenced, "old session must stay fenced after successful replacement attach")
	newWatcher, newClone, _ = memoryWatcherSessionFields(newSession)
	require.True(t, newWatcher, "new watcher must publish after the old drain completes")
	require.True(t, newClone, "new clone must publish after the old drain completes")
}

func TestStartSessionReplacementFailureRestoresOldWatcher(t *testing.T) {
	for _, failureStage := range []string{"resolve", "launch"} {
		t.Run(failureStage, func(t *testing.T) {
			t.Setenv(profile.ConfigDirEnvKey, t.TempDir())
			_ = initTokenForTest(t)
			stubRuntimeBinaries(t, "codex")

			origin := makeBareOriginRepoForMemory(t, filepath.Join(t.TempDir(), "origin"))
			agentID := "agt_replacement_restore_" + failureStage
			workDir := sessionWorkDir(agentID)
			provider := &fakeMemoryAccessProvider{}
			provider.set(agentID, MemoryAccess{
				MemoryID:  "mem_replacement_restore_" + failureStage,
				OriginURL: origin,
				Branch:    "main",
			})
			d := newDispatcherWithMemoryProvider(t, provider)
			process, _ := newShutdownDeadlineProcess(t, agentID)
			agentConfig := agent.Config{ID: agentID, Name: "Replacement Restore", Runtime: "codex"}
			runtimeConfig := runtimeconfig.Runtime{Type: runtimeconfig.TypeCodex, Binary: "codex"}
			oldSession := &interactiveSession{
				sessionKey: agentID,
				agentID:    agentID,
				agent:      agentConfig,
				runtime:    runtimeConfig,
				workDir:    workDir,
				process:    process,
			}
			d.sessionManager.Set(agentID, oldSession)
			d.attachMemoryWatcher(context.Background(), oldSession, agentID, workDir)
			t.Cleanup(func() { stopMemoryWatcher(oldSession) })
			oldWatcher, oldClone, _ := memoryWatcherSessionFields(oldSession)
			require.True(t, oldWatcher, "precondition: old watcher must be published")
			require.True(t, oldClone, "precondition: old clone must be published")

			reservation, _, owns := d.sessionManager.BeginStart(agentID, oldSession)
			require.True(t, owns, "precondition: replacement start must own its generation")
			t.Cleanup(func() { d.sessionManager.FinishStart(reservation) })

			injectedErr := errors.New("injected replacement " + failureStage + " failure")
			require.NoError(t, failpoint.Enable("dispatcher-new-agent-driver", "return(true)"))
			if failureStage == "resolve" {
				SetNewAgentDriverHook(func(runtimeconfig.RuntimeType) (drivercontract.Driver, error) {
					return nil, injectedErr
				})
			} else {
				SetNewAgentDriverHook(func(runtimeconfig.RuntimeType) (drivercontract.Driver, error) {
					return &stubDriver{}, nil
				})
				require.NoError(t, failpoint.Enable("dispatcher-new-agent-process", "return(true)"))
				SetNewAgentProcessHook(func(
					context.Context,
					drivercontract.Driver,
					drivercontract.LaunchConfig,
					agent.ProcessOptions,
				) (*agent.Process, error) {
					return nil, injectedErr
				})
			}
			t.Cleanup(func() {
				SetNewAgentDriverHook(nil)
				SetNewAgentProcessHook(nil)
				require.NoError(t, failpoint.Disable("dispatcher-new-agent-driver"))
				if failureStage == "launch" {
					require.NoError(t, failpoint.Disable("dispatcher-new-agent-process"))
				}
			})

			_, err := d.startSession(
				context.Background(),
				agentID,
				agentConfig,
				runtimeConfig,
				conversationBinding{channelID: "memory-replacement"},
				types.InboxMessage{
					MessageID:   "msg-replacement-" + failureStage,
					Target:      "#memory-replacement",
					ChannelType: "channel",
					ChannelName: "memory-replacement",
					AuthorProfile: types.AuthorProfile{
						ActorType:   "human",
						DisplayName: "alice",
						TargetID:    "alice",
					},
					Content:   "replace the running session",
					Timestamp: time.Now().UTC().Format(time.RFC3339),
				},
				startSessionOptions{startReservation: reservation},
			)
			require.ErrorIs(t, err, injectedErr)

			current, ok := d.sessionManager.Get(agentID)
			require.True(t, ok)
			require.Same(t, oldSession, current, "failed replacement must leave the old session current")
			oldSession.mu.Lock()
			oldFenced := oldSession.memoryWatcherReplacementFenced
			oldDrain := oldSession.memoryWatcherDrain
			oldSession.mu.Unlock()
			require.False(t, oldFenced, "best-effort restore must reopen the old watcher admission")
			require.Nil(t, oldDrain, "best-effort restore must finish before startSession returns")
			oldWatcher, oldClone, _ = memoryWatcherSessionFields(oldSession)
			require.True(t, oldWatcher, "failed replacement must restore the old watcher")
			require.True(t, oldClone, "failed replacement must restore the old clone")
			require.Equal(t, "running", process.State(), "restoration requires and must preserve the old runtime")
			provider.mu.Lock()
			lookupCalls := provider.calls
			provider.mu.Unlock()
			require.GreaterOrEqual(t, lookupCalls, 3,
				"old attach, replacement attach, and best-effort restore must each resolve access")
		})
	}
}

func TestRefreshMemoryAccessWaitsOutsideLookupLockForSessionStart(t *testing.T) {
	t.Setenv(profile.ConfigDirEnvKey, t.TempDir())
	stubRuntimeBinaries(t, "codex")
	origin := makeBareOriginRepoForMemory(t, filepath.Join(t.TempDir(), "origin"))
	agentID := "agt_refresh_real_start_interleave"
	provider := &fakeMemoryAccessProvider{}
	provider.set(agentID, MemoryAccess{
		MemoryID:                   "mem_refresh_real_start_interleave",
		OriginURL:                  origin,
		Branch:                     "main",
		ImportExistingLocalContent: true,
	})
	d := newDispatcherWithMemoryProvider(t, provider)

	reservation, _, owns := d.sessionManager.BeginStart(agentID, nil)
	require.True(t, owns)
	t.Cleanup(func() { d.sessionManager.FinishStart(reservation) })

	waiting := make(chan struct{})
	var waitingOnce sync.Once
	d.memoryRefreshBeforeWaitForStart = func() { waitingOnce.Do(func() { close(waiting) }) }
	refreshDone := make(chan struct{})
	ctx := t.Context()
	go func() {
		d.RefreshMemoryAccess(ctx, []string{agentID})
		close(refreshDone)
	}()
	<-waiting

	injectedErr := errors.New("injected start after real memory attach")
	require.NoError(t, failpoint.Enable("dispatcher-new-agent-driver", "return(true)"))
	SetNewAgentDriverHook(func(runtimeconfig.RuntimeType) (drivercontract.Driver, error) {
		return nil, injectedErr
	})
	t.Cleanup(func() {
		SetNewAgentDriverHook(nil)
		require.NoError(t, failpoint.Disable("dispatcher-new-agent-driver"))
	})
	startDone := make(chan error, 1)
	go func() {
		_, err := d.startSession(
			ctx,
			agentID,
			agent.Config{ID: agentID, Name: "Refresh Interleave", Runtime: "codex"},
			runtimeconfig.Runtime{Type: runtimeconfig.TypeCodex, Binary: "codex"},
			conversationBinding{channelID: "memory-refresh-interleave"},
			types.InboxMessage{
				MessageID:   "msg-refresh-interleave",
				Target:      "#memory-refresh-interleave",
				ChannelType: "channel",
				ChannelName: "memory-refresh-interleave",
				Content:     "start while memory access refresh is waiting",
				Timestamp:   time.Now().UTC().Format(time.RFC3339),
			},
			startSessionOptions{startReservation: reservation},
		)
		startDone <- err
	}()

	select {
	case err := <-startDone:
		require.ErrorIs(t, err, injectedErr)
	case <-time.After(5 * time.Second):
		t.Fatal("startSession could not acquire the lookup lock while refresh waited for its generation")
	}
	d.sessionManager.FinishStart(reservation)
	select {
	case <-refreshDone:
	case <-time.After(5 * time.Second):
		t.Fatal("refresh did not resume after the real start generation finished")
	}
	provider.mu.Lock()
	lookupCalls := provider.calls
	provider.mu.Unlock()
	require.GreaterOrEqual(t, lookupCalls, 2, "real start attach and resumed refresh must both perform authoritative lookup")
}

func TestStartSessionFinalizesLegacyPreferencesBeforeSpawnWritesProfile(t *testing.T) {
	t.Setenv(profile.ConfigDirEnvKey, t.TempDir())
	stubRuntimeBinaries(t, "codex")
	origin := makeBareOriginRepoForMemory(t, filepath.Join(t.TempDir(), "origin"))
	agentID := "agt_legacy_before_spawn_write"
	workDir := sessionWorkDir(agentID)
	access := MemoryAccess{
		MemoryID:                   "mem_legacy_before_spawn_write",
		Protocol:                   MemorySyncProtocolLegacyGitV1,
		OriginURL:                  origin,
		Branch:                     "main",
		ImportExistingLocalContent: true,
	}
	provider := &fakeMemoryAccessProvider{}
	provider.set(agentID, access)
	d := newDispatcherWithMemoryProvider(t, provider)
	_, migration, _, ok := d.ensureMemoryCloneForAccess(t.Context(), agentID, workDir, access)
	require.True(t, ok)
	require.NotNil(t, migration)
	legacyPath := filepath.Join(workDir, "kb", "user-preferences.md")
	require.NoError(t, os.MkdirAll(filepath.Dir(legacyPath), 0o755))
	require.NoError(t, os.WriteFile(legacyPath, []byte("# User Preferences\n\n- Legacy preference.\n"), 0o644))

	injectedErr := errors.New("injected spawn failure after active profile write")
	require.NoError(t, failpoint.Enable("dispatcher-new-agent-driver", "return(true)"))
	require.NoError(t, failpoint.Enable("dispatcher-new-agent-process", "return(true)"))
	SetNewAgentDriverHook(func(runtimeconfig.RuntimeType) (drivercontract.Driver, error) {
		return &stubDriver{}, nil
	})
	SetNewAgentProcessHook(func(
		_ context.Context,
		_ drivercontract.Driver,
		_ drivercontract.LaunchConfig,
		_ agent.ProcessOptions,
	) (*agent.Process, error) {
		profilePath := filepath.Join(workDir, "memory", "profiles", "user.md")
		profile, err := os.ReadFile(profilePath)
		require.NoError(t, err)
		require.Contains(t, string(profile), "Legacy preference.",
			"legacy merge must finish before Spawn can dispatch the initial turn")
		profile = append(profile, []byte("\n## Active Spawn Write\n\n- New preference.\n")...)
		require.NoError(t, os.WriteFile(profilePath, profile, 0o644))
		return nil, injectedErr
	})
	t.Cleanup(func() {
		SetNewAgentDriverHook(nil)
		SetNewAgentProcessHook(nil)
		require.NoError(t, failpoint.Disable("dispatcher-new-agent-driver"))
		require.NoError(t, failpoint.Disable("dispatcher-new-agent-process"))
	})

	reservation, _, owns := d.sessionManager.BeginStart(agentID, nil)
	require.True(t, owns)
	_, err := d.startSession(
		t.Context(),
		agentID,
		agent.Config{ID: agentID, Name: "Legacy Before Spawn", Runtime: "codex"},
		runtimeconfig.Runtime{Type: runtimeconfig.TypeCodex, Binary: "codex"},
		conversationBinding{channelID: "memory-legacy-before-spawn"},
		types.InboxMessage{
			MessageID:   "msg-legacy-before-spawn",
			Target:      "#memory-legacy-before-spawn",
			ChannelType: "channel",
			ChannelName: "memory-legacy-before-spawn",
			Content:     "write active preferences during Spawn",
			Timestamp:   time.Now().UTC().Format(time.RFC3339),
		},
		startSessionOptions{startReservation: reservation},
	)
	d.sessionManager.FinishStart(reservation)
	require.ErrorIs(t, err, injectedErr)
	profile := readFileString(t, filepath.Join(workDir, "memory", "profiles", "user.md"))
	require.Contains(t, profile, "Legacy preference.")
	require.Contains(t, profile, "New preference.")
	require.FileExists(t, legacyUserPreferencesMigrationStatePath(
		workDir, agentID, access.MemoryID, access.Branch, legacyUserPreferencesMigrationDoneSuffix,
	))
}

func TestStartSessionStopsLiveFallbackBeforeManagedSpawn(t *testing.T) {
	t.Setenv(profile.ConfigDirEnvKey, t.TempDir())
	stubRuntimeBinaries(t, "codex")
	origin := makeBareOriginRepoForMemory(t, filepath.Join(t.TempDir(), "origin"))
	agentID := "agt_live_fallback_before_spawn"
	workDir := sessionWorkDir(agentID)
	access := MemoryAccess{
		MemoryID:                   "mem_live_fallback_before_spawn",
		OriginURL:                  origin,
		Branch:                     "main",
		ImportExistingLocalContent: true,
	}
	provider := &fakeMemoryAccessProvider{}
	provider.set(agentID, access)
	d := newDispatcherWithMemoryProvider(t, provider)

	legacyPath := filepath.Join(workDir, "kb", "user-preferences.md")
	require.NoError(t, os.MkdirAll(filepath.Dir(legacyPath), 0o755))
	require.NoError(t, os.WriteFile(legacyPath, []byte("# User Preferences\n\n- Live fallback preference.\n"), 0o644))
	oldProcess := &agent.Process{}
	oldSession := &interactiveSession{
		sessionKey:                     agentID,
		agentID:                        agentID,
		workDir:                        workDir,
		process:                        oldProcess,
		processUsesWorkspaceRootMemory: true,
	}
	d.sessionManager.sessions[agentID] = oldSession

	injectedErr := errors.New("injected spawn failure after live fallback cutover")
	require.NoError(t, failpoint.Enable("dispatcher-new-agent-driver", "return(true)"))
	require.NoError(t, failpoint.Enable("dispatcher-new-agent-process", "return(true)"))
	SetNewAgentDriverHook(func(runtimeconfig.RuntimeType) (drivercontract.Driver, error) {
		return &stubDriver{}, nil
	})
	SetNewAgentProcessHook(func(
		_ context.Context,
		_ drivercontract.Driver,
		_ drivercontract.LaunchConfig,
		_ agent.ProcessOptions,
	) (*agent.Process, error) {
		require.Equal(t, "stopped", oldProcess.State(),
			"the fallback generation must reach its Stop barrier before managed Spawn")
		profilePath := filepath.Join(workDir, "memory", "profiles", "user.md")
		profile := readFileString(t, profilePath)
		require.Contains(t, profile, "Live fallback preference.")
		profile += "\n## Active Spawn Write\n\n- New active preference.\n"
		require.NoError(t, os.WriteFile(profilePath, []byte(profile), 0o644))
		return nil, injectedErr
	})
	t.Cleanup(func() {
		SetNewAgentDriverHook(nil)
		SetNewAgentProcessHook(nil)
		require.NoError(t, failpoint.Disable("dispatcher-new-agent-driver"))
		require.NoError(t, failpoint.Disable("dispatcher-new-agent-process"))
	})

	reservation, _, owns := d.sessionManager.BeginStart(agentID, oldSession)
	require.True(t, owns)
	_, err := d.startSession(
		t.Context(),
		agentID,
		agent.Config{ID: agentID, Name: "Live Fallback", Runtime: "codex"},
		runtimeconfig.Runtime{Type: runtimeconfig.TypeCodex, Binary: "codex"},
		conversationBinding{channelID: "memory-live-fallback"},
		types.InboxMessage{
			MessageID:   "msg-live-fallback",
			Target:      "#memory-live-fallback",
			ChannelType: "channel",
			ChannelName: "memory-live-fallback",
			Content:     "replace the fallback process",
			Timestamp:   time.Now().UTC().Format(time.RFC3339),
		},
		startSessionOptions{startReservation: reservation},
	)
	d.sessionManager.FinishStart(reservation)
	require.ErrorIs(t, err, injectedErr)
	profile := readFileString(t, filepath.Join(workDir, "memory", "profiles", "user.md"))
	require.Contains(t, profile, "Live fallback preference.")
	require.Contains(t, profile, "New active preference.")
	require.FileExists(t, legacyUserPreferencesMigrationStatePath(
		workDir, agentID, access.MemoryID, access.Branch, legacyUserPreferencesMigrationDoneSuffix,
	))
	require.NoFileExists(t, legacyUserPreferencesMigrationStatePath(
		workDir, agentID, access.MemoryID, access.Branch, legacyUserPreferencesMigrationReadySuffix,
	))
}

func TestStartSessionReplacementFailureRacingShutdownDoesNotRestoreOldWatcher(t *testing.T) {
	t.Setenv(profile.ConfigDirEnvKey, t.TempDir())
	stubRuntimeBinaries(t, "codex")

	origin := makeBareOriginRepoForMemory(t, filepath.Join(t.TempDir(), "origin"))
	agentID := "agt_replacement_failure_shutdown"
	workDir := sessionWorkDir(agentID)
	provider := &fakeMemoryAccessProvider{}
	provider.set(agentID, MemoryAccess{
		MemoryID:  "mem_replacement_failure_shutdown",
		OriginURL: origin,
		Branch:    "main",
	})
	d := newDispatcherWithMemoryProvider(t, provider)
	process, driver := newShutdownDeadlineProcess(t, agentID)
	agentConfig := agent.Config{ID: agentID, Name: "Replacement Shutdown", Runtime: "codex"}
	runtimeConfig := runtimeconfig.Runtime{Type: runtimeconfig.TypeCodex, Binary: "codex"}
	oldSession := &interactiveSession{
		sessionKey: agentID,
		agentID:    agentID,
		agent:      agentConfig,
		runtime:    runtimeConfig,
		workDir:    workDir,
		process:    process,
	}
	d.sessionManager.Set(agentID, oldSession)
	d.attachMemoryWatcher(context.Background(), oldSession, agentID, workDir)
	t.Cleanup(func() { stopMemoryWatcher(oldSession) })

	reservation, _, owns := d.sessionManager.BeginStart(agentID, oldSession)
	require.True(t, owns, "precondition: replacement start must own its generation")
	injectedErr := errors.New("injected replacement resolve failure during shutdown")
	require.NoError(t, failpoint.Enable("dispatcher-new-agent-driver", "return(true)"))
	SetNewAgentDriverHook(func(runtimeconfig.RuntimeType) (drivercontract.Driver, error) {
		return nil, injectedErr
	})
	t.Cleanup(func() {
		SetNewAgentDriverHook(nil)
		require.NoError(t, failpoint.Disable("dispatcher-new-agent-driver"))
	})

	newBeforePublish := make(chan struct{})
	releaseNewPublication := make(chan struct{})
	var enteredOnce sync.Once
	var releaseOnce sync.Once
	releasePublication := func() { releaseOnce.Do(func() { close(releaseNewPublication) }) }
	t.Cleanup(releasePublication)
	d.memoryWatcherBeforePublish = func(_ context.Context, sess *interactiveSession) {
		if sess == oldSession {
			return
		}
		enteredOnce.Do(func() { close(newBeforePublish) })
		<-releaseNewPublication
	}

	startDone := make(chan error, 1)
	go func() {
		_, err := d.startSession(
			context.Background(),
			agentID,
			agentConfig,
			runtimeConfig,
			conversationBinding{channelID: "memory-replacement"},
			types.InboxMessage{
				MessageID:   "msg-replacement-shutdown",
				Target:      "#memory-replacement",
				ChannelType: "channel",
				ChannelName: "memory-replacement",
				AuthorProfile: types.AuthorProfile{
					ActorType:   "human",
					DisplayName: "alice",
					TargetID:    "alice",
				},
				Content:   "replace while shutting down",
				Timestamp: time.Now().UTC().Format(time.RFC3339),
			},
			startSessionOptions{startReservation: reservation},
		)
		d.sessionManager.FinishStart(reservation)
		startDone <- err
	}()
	waitForMemoryWatcherTestSignal(t, newBeforePublish, "blocked replacement watcher publication")

	closeDone := make(chan struct{})
	go func() {
		d.Close()
		close(closeDone)
	}()
	require.Eventually(t, func() bool {
		d.memoryWatcherMutationMu.Lock()
		closed := d.memoryWatcherMutationClosed
		d.memoryWatcherMutationMu.Unlock()
		return closed && oldSession.retired.Load()
	}, 5*time.Second, 10*time.Millisecond,
		"shutdown must close watcher mutation admission and retire the old session before release")

	releasePublication()
	var startErr error
	select {
	case startErr = <-startDone:
	case <-time.After(5 * time.Second):
		t.Fatal("timed out waiting for failed replacement start to unwind")
	}
	require.ErrorIs(t, startErr, injectedErr)
	driver.allowExit()
	waitForMemoryWatcherTestSignal(t, closeDone, "dispatcher close after failed replacement")

	oldSession.mu.Lock()
	oldFenced := oldSession.memoryWatcherReplacementFenced
	oldDrain := oldSession.memoryWatcherDrain
	oldSession.mu.Unlock()
	oldWatcher, oldClone, oldAuth := memoryWatcherSessionFields(oldSession)
	require.True(t, oldFenced,
		"shutdown-owned failed replacement must not clear the old session's replacement fence")
	require.Nil(t, oldDrain, "shutdown must finish the old watcher drain")
	require.False(t, oldWatcher, "failed replacement must not revive an old watcher after shutdown")
	require.False(t, oldClone, "failed replacement must not revive an old clone after shutdown")
	require.False(t, oldAuth, "failed replacement must not revive old push auth after shutdown")
	provider.mu.Lock()
	lookupCalls := provider.calls
	provider.mu.Unlock()
	require.Equal(t, 2, lookupCalls,
		"shutdown must reject restore before a third provider lookup can recreate the old watcher")
}

func TestProcessExitLosingReplacementGenerationRejectsBlockedWatcherRestore(t *testing.T) {
	t.Setenv(profile.ConfigDirEnvKey, t.TempDir())
	_ = initTokenForTest(t)
	stubRuntimeBinaries(t, "codex")

	origin := makeBareOriginRepoForMemory(t, filepath.Join(t.TempDir(), "origin"))
	agentID := "agt_replacement_restore_process_exit"
	workDir := sessionWorkDir(agentID)
	provider := &fakeMemoryAccessProvider{}
	provider.set(agentID, MemoryAccess{
		MemoryID:  "mem_replacement_restore_process_exit",
		OriginURL: origin,
		Branch:    "main",
	})
	d := newDispatcherWithMemoryProvider(t, provider)

	exitHandled := make(chan struct{})
	var oldSession *interactiveSession
	process, driver := newShutdownCallbackProcess(t, agentID, func(ap *agent.Process) {
		d.handleProcessExit(oldSession, ap)
		close(exitHandled)
	})
	t.Cleanup(func() {
		process.Stop()
		driver.allowExit()
		require.True(t, process.WaitForExit(5*time.Second))
	})
	agentConfig := agent.Config{ID: agentID, Name: "Replacement Exit Fence", Runtime: "codex"}
	runtimeConfig := runtimeconfig.Runtime{Type: runtimeconfig.TypeCodex, Binary: "codex"}
	oldSession = &interactiveSession{
		sessionKey: agentID,
		agentID:    agentID,
		agent:      agentConfig,
		runtime:    runtimeConfig,
		workDir:    workDir,
		process:    process,
	}
	d.sessionManager.Set(agentID, oldSession)
	d.attachMemoryWatcher(context.Background(), oldSession, agentID, workDir)
	t.Cleanup(func() { stopMemoryWatcher(oldSession) })
	oldWatcher, oldClone, _ := memoryWatcherSessionFields(oldSession)
	require.True(t, oldWatcher, "precondition: old watcher must be published")
	require.True(t, oldClone, "precondition: old clone must be published")

	reservation, _, owns := d.sessionManager.BeginStart(agentID, oldSession)
	require.True(t, owns, "precondition: failed replacement must own the start generation")
	injectedErr := errors.New("injected replacement resolve failure before restore")
	require.NoError(t, failpoint.Enable("dispatcher-new-agent-driver", "return(true)"))
	SetNewAgentDriverHook(func(runtimeconfig.RuntimeType) (drivercontract.Driver, error) {
		return nil, injectedErr
	})
	t.Cleanup(func() {
		SetNewAgentDriverHook(nil)
		require.NoError(t, failpoint.Disable("dispatcher-new-agent-driver"))
	})

	restoreBeforePublish := make(chan struct{})
	releaseRestorePublication := make(chan struct{})
	var restoreEnteredOnce sync.Once
	var releaseOnce sync.Once
	releaseRestore := func() { releaseOnce.Do(func() { close(releaseRestorePublication) }) }
	t.Cleanup(releaseRestore)
	d.memoryWatcherBeforePublish = func(_ context.Context, sess *interactiveSession) {
		if sess != oldSession {
			return
		}
		restoreEnteredOnce.Do(func() { close(restoreBeforePublish) })
		<-releaseRestorePublication
	}

	startDone := make(chan error, 1)
	go func() {
		_, err := d.startSession(
			context.Background(),
			agentID,
			agentConfig,
			runtimeConfig,
			conversationBinding{channelID: "memory-replacement"},
			types.InboxMessage{
				MessageID:   "msg-replacement-restore-process-exit",
				Target:      "#memory-replacement",
				ChannelType: "channel",
				ChannelName: "memory-replacement",
				AuthorProfile: types.AuthorProfile{
					ActorType:   "human",
					DisplayName: "alice",
					TargetID:    "alice",
				},
				Content:   "fail replacement while the old runtime exits",
				Timestamp: time.Now().UTC().Format(time.RFC3339),
			},
			startSessionOptions{startReservation: reservation},
		)
		d.sessionManager.FinishStart(reservation)
		startDone <- err
	}()
	waitForMemoryWatcherTestSignal(t, restoreBeforePublish, "blocked old watcher restore publication")

	// The failed replacement still owns BeginStart while restore is parked at
	// the last publication boundary. The real Process onExit callback therefore
	// loses the generation. That early-return path must still install a terminal
	// watcher fence; otherwise releasing restore publishes onto process=nil.
	process.Stop()
	driver.allowExit()
	waitForMemoryWatcherTestSignal(t, exitHandled, "old process exit callback losing replacement generation")
	oldSession.mu.Lock()
	processAfterExit := oldSession.process
	terminalFenced := oldSession.memoryWatcherTerminalFenced
	oldSession.mu.Unlock()
	require.Nil(t, processAfterExit, "process exit callback must clear the old current process")
	require.True(t, terminalFenced,
		"process exit that loses BeginStart must fence the concurrently restoring watcher")

	releaseRestore()
	var startErr error
	select {
	case startErr = <-startDone:
	case <-time.After(5 * time.Second):
		t.Fatal("timed out waiting for failed replacement restore rejection")
	}
	require.ErrorIs(t, startErr, injectedErr)

	current, ok := d.sessionManager.Get(agentID)
	require.True(t, ok, "losing callback must leave removal to the competing start owner")
	require.Same(t, oldSession, current)
	oldSession.mu.Lock()
	oldDrain := oldSession.memoryWatcherDrain
	terminalFenced = oldSession.memoryWatcherTerminalFenced
	oldSession.mu.Unlock()
	oldWatcher, oldClone, oldAuth := memoryWatcherSessionFields(oldSession)
	require.True(t, terminalFenced, "terminal watcher fence must remain permanent")
	require.Nil(t, oldDrain, "rejected restore cleanup must finish before startSession returns")
	require.False(t, oldWatcher, "blocked restore must not publish onto the process-less old session")
	require.False(t, oldClone, "blocked restore must not publish its clone onto the old session")
	require.False(t, oldAuth, "blocked restore must not publish push auth onto the old session")
}