feat(home): enhance client recovery and dispatch handling with new state management and error reporting

internal/home/client.go

@@ -93,6 +93,15 @@ var (
 	ErrDispatchFenced        = errors.New("home auth dispatch is fenced")
 )
 
+// IsMembershipTakeoverUnavailableError reports whether Home cannot preserve the previous membership state.
+func IsMembershipTakeoverUnavailableError(err error) bool {
+	if err == nil {
+		return false
+	}
+	message := strings.ToLower(err.Error())
+	return strings.Contains(message, "membership_takeover_unavailable") || strings.Contains(message, "wrong number of arguments for 'subscribe' command")
+}
+
 type clusterNode struct {
 	IP          string    `json:"ip"`
 	Port        int       `json:"port"`
@@ -127,6 +136,15 @@ type subscriptionCloser interface {
 	Close() error
 }
 
+type recoveryState uint32
+
+const (
+	recoveryStateStable recoveryState = iota
+	recoveryStateTakeoverEligible
+	recoveryStateSwitching
+	recoveryStateSwitchingTakeover
+)
+
 type Client struct {
 	mu sync.Mutex
 
@@ -139,12 +157,15 @@ type Client struct {
 	sub         *redis.Client
 	release     *redis.Client
 	connections map[*homeDispatchConn]struct{}
+	closing     chan struct{}
 	lifecycle   config.CredentialConcurrencyConfig
 	limiter     atomic.Pointer[config.CredentialConcurrencyConfig]
 	managed     bool
 
 	heartbeatOK       atomic.Bool
 	dispatchFenced    atomic.Bool
+	ambiguousDispatch atomic.Bool
+	recoveryState     atomic.Uint32
 	clusterNodes      []clusterNode
 	reconnectFailures int
 }
@@ -164,13 +185,15 @@ func (c *Client) NewLifetime() *Client {
 	}
 	c.mu.Lock()
 	defer c.mu.Unlock()
-	return &Client{
+	next := &Client{
 		homeCfg:           c.homeCfg,
 		seedHost:          c.seedHost,
 		seedPort:          c.seedPort,
 		clusterNodes:      append([]clusterNode(nil), c.clusterNodes...),
 		reconnectFailures: c.reconnectFailures,
 	}
+	next.recoveryState.Store(c.recoveryState.Load())
+	return next
 }
 
 func (c *Client) Enabled() bool {
@@ -203,11 +226,15 @@ func (c *Client) Close() {
 	commandClient, subscriptionClient, connections := c.detachClientsLocked()
 	releaseClient := c.release
 	c.release = nil
+	closing := c.closing
 	c.mu.Unlock()
 	closeDetachedClients(commandClient, subscriptionClient, connections)
 	if releaseClient != nil {
 		_ = releaseClient.Close()
 	}
+	if closing != nil {
+		<-closing
+	}
 }
 
 // closeBootstrapPools replaces private bootstrap pools without ending the client lifetime.
@@ -227,6 +254,7 @@ func (c *Client) AbortAmbiguousDispatch() {
 	if c == nil {
 		return
 	}
+	c.ambiguousDispatch.Store(true)
 	c.dispatchFenced.Store(true)
 	c.heartbeatOK.Store(false)
 	c.mu.Lock()
@@ -254,6 +282,21 @@ func (c *Client) AbortAmbiguousDispatch() {
 	}
 }
 
+// AmbiguousDispatch reports whether this lifetime observed an issued dispatch with an unknown delivery result.
+func (c *Client) AmbiguousDispatch() bool {
+	return c != nil && c.ambiguousDispatch.Load()
+}
+
+// SuppressTakeover forces the next subscriber lifetime through normal membership recovery.
+func (c *Client) SuppressTakeover() {
+	if c == nil {
+		return
+	}
+	if !c.recoveryState.CompareAndSwap(uint32(recoveryStateTakeoverEligible), uint32(recoveryStateStable)) {
+		c.recoveryState.CompareAndSwap(uint32(recoveryStateSwitchingTakeover), uint32(recoveryStateSwitching))
+	}
+}
+
 func (c *Client) detachClientsLocked() (*redis.Client, *redis.Client, []*homeDispatchConn) {
 	connections := make([]*homeDispatchConn, 0, len(c.connections))
 	for conn := range c.connections {
@@ -284,7 +327,14 @@ func (c *Client) closeClientsLocked() {
 	commandClient, subscriptionClient, connections := c.detachClientsLocked()
 	releaseClient := c.release
 	c.release = nil
+	previousClosing := c.closing
+	done := make(chan struct{})
+	c.closing = done
 	go func() {
+		defer close(done)
+		if previousClosing != nil {
+			<-previousClosing
+		}
 		closeDetachedClients(commandClient, subscriptionClient, connections)
 		if releaseClient != nil {
 			_ = releaseClient.Close()
@@ -292,6 +342,25 @@ func (c *Client) closeClientsLocked() {
 	}()
 }
 
+func (c *Client) waitForClientsClosed() {
+	for {
+		c.mu.Lock()
+		closing := c.closing
+		c.mu.Unlock()
+		if closing == nil {
+			return
+		}
+		<-closing
+		c.mu.Lock()
+		if c.closing == closing {
+			c.closing = nil
+			c.mu.Unlock()
+			return
+		}
+		c.mu.Unlock()
+	}
+}
+
 // SetManagedLifetime defers client shutdown to the Service lifetime owner.
 func (c *Client) SetManagedLifetime(managed bool) {
 	if c == nil {
@@ -341,6 +410,7 @@ func (c *Client) ensureClients() error {
 	if !c.Enabled() {
 		return ErrDisabled
 	}
+	c.waitForClientsClosed()
 	c.mu.Lock()
 	defer c.mu.Unlock()
 	if c.dispatchFenced.Load() {
@@ -685,6 +755,9 @@ func (c *Client) switchToNodeLocked(node clusterNode) bool {
 	}
 	c.homeCfg.Host = host
 	c.homeCfg.Port = node.Port
+	if !c.recoveryState.CompareAndSwap(uint32(recoveryStateStable), uint32(recoveryStateSwitching)) {
+		c.recoveryState.CompareAndSwap(uint32(recoveryStateTakeoverEligible), uint32(recoveryStateSwitchingTakeover))
+	}
 	c.closeClientsLocked()
 	return true
 }
@@ -1233,6 +1306,10 @@ func (c *Client) concurrencyReleaseClient() (*redis.Client, error) {
 	if c == nil || c.dispatchFenced.Load() {
 		return nil, ErrDispatchFenced
 	}
+	state := recoveryState(c.recoveryState.Load())
+	if state == recoveryStateSwitching || state == recoveryStateSwitchingTakeover {
+		return nil, ErrNotConnected
+	}
 	if !c.Enabled() {
 		return nil, ErrDisabled
 	}
@@ -1242,6 +1319,10 @@ func (c *Client) concurrencyReleaseClient() (*redis.Client, error) {
 	if c.dispatchFenced.Load() {
 		return nil, ErrDispatchFenced
 	}
+	state = recoveryState(c.recoveryState.Load())
+	if state == recoveryStateSwitching || state == recoveryStateSwitchingTakeover {
+		return nil, ErrNotConnected
+	}
 	if c.release != nil {
 		return c.release, nil
 	}
@@ -1525,13 +1606,21 @@ func (c *Client) subscriptionParameters() ([]string, time.Duration) {
 	args := []string{redisChannelConfig}
 	if cfg.LifecycleConfigRevision > 0 {
 		args = append(args, strconv.FormatInt(cfg.LifecycleConfigRevision, 10))
+		state := recoveryState(c.recoveryState.Load())
+		if state == recoveryStateTakeoverEligible || state == recoveryStateSwitchingTakeover {
+			args = append(args, "takeover")
+		}
 	}
 	return args, cfg.CPAHeartbeatTimeout
 }
 
 func (c *Client) rebuildCommandPoolAndProbe(ctx context.Context) error {
 	c.promoteSubscription()
-	return c.Ping(ctx)
+	if errPing := c.Ping(ctx); errPing != nil {
+		return errPing
+	}
+	c.recoveryState.Store(uint32(recoveryStateStable))
+	return nil
 }
 
 func (c *Client) promoteSubscription() {
@@ -1641,6 +1730,11 @@ func (c *Client) RunConfigSubscriberLifetime(ctx context.Context, onConfig func(
 		event, errReceive := pubsub.ReceiveTimeout(ctx, receiveTimeout)
 		if errReceive != nil {
 			if ctx.Err() == nil {
+				if c.heartbeatOK.Load() {
+					if !c.recoveryState.CompareAndSwap(uint32(recoveryStateStable), uint32(recoveryStateTakeoverEligible)) {
+						c.recoveryState.CompareAndSwap(uint32(recoveryStateSwitching), uint32(recoveryStateSwitchingTakeover))
+					}
+				}
 				if isTimeoutError(errReceive) {
 					c.markSubscriptionTimeout()
 				} else {

internal/home/client_test.go

@@ -217,6 +217,75 @@ func TestNewLifetimePreservesClusterFailoverState(t *testing.T) {
 	}
 }
 
+func TestEnsureClientsWaitsForPreviousTargetClose(t *testing.T) {
+	client := New(config.HomeConfig{Enabled: true, Host: "next.example.com", Port: 8327})
+	closing := make(chan struct{})
+	client.closing = closing
+	done := make(chan error, 1)
+	go func() {
+		done <- client.ensureClients()
+	}()
+
+	select {
+	case errEnsure := <-done:
+		t.Fatalf("ensureClients() returned before previous target closed: %v", errEnsure)
+	case <-time.After(20 * time.Millisecond):
+	}
+	close(closing)
+	select {
+	case errEnsure := <-done:
+		if errEnsure != nil {
+			t.Fatal(errEnsure)
+		}
+	case <-time.After(time.Second):
+		t.Fatal("ensureClients() did not continue after previous target closed")
+	}
+	client.Close()
+}
+
+func TestConcurrencyReleaseDoesNotOpenSwitchingTarget(t *testing.T) {
+	client := New(config.HomeConfig{Enabled: true, Host: "next.example.com", Port: 8327})
+	client.recoveryState.Store(uint32(recoveryStateSwitching))
+	errRelease := client.PushConcurrencyRelease(context.Background(), ConcurrencyReleaseFrame{CredentialID: "cred-a", Model: "model-a", ReleaseSeq: 1})
+	if !errors.Is(errRelease, ErrNotConnected) {
+		t.Fatalf("PushConcurrencyRelease() error = %v, want %v", errRelease, ErrNotConnected)
+	}
+	client.mu.Lock()
+	releaseClient := client.release
+	client.mu.Unlock()
+	if releaseClient != nil {
+		t.Fatal("release client was opened before the switched target became ready")
+	}
+}
+
+func TestAmbiguousDispatchSuppressesTakeoverForNextLifetime(t *testing.T) {
+	client := New(config.HomeConfig{Enabled: true, Host: "next.example.com", Port: 8327})
+	client.recoveryState.Store(uint32(recoveryStateSwitchingTakeover))
+	client.AbortAmbiguousDispatch()
+	if !client.AmbiguousDispatch() {
+		t.Fatal("ambiguous dispatch was not recorded")
+	}
+	client.SuppressTakeover()
+	next := client.NewLifetime()
+	if got := recoveryState(next.recoveryState.Load()); got != recoveryStateSwitching {
+		t.Fatalf("next recovery state = %d, want %d", got, recoveryStateSwitching)
+	}
+}
+
+func TestMembershipTakeoverUnavailableError(t *testing.T) {
+	for _, message := range []string{
+		"ERR membership_takeover_unavailable",
+		"ERR wrong number of arguments for 'subscribe' command",
+	} {
+		if !IsMembershipTakeoverUnavailableError(errors.New(message)) {
+			t.Fatalf("takeover unavailable error %q was not recognized", message)
+		}
+	}
+	if IsMembershipTakeoverUnavailableError(errors.New("ERR connection refused")) {
+		t.Fatal("unrelated error was recognized as takeover unavailable")
+	}
+}
+
 func TestBuildKVSetArgs(t *testing.T) {
 	args, errArgs := buildKVSetArgs("key", []byte("value"), KVSetOptions{EX: 2 * time.Second, NX: true})
 	if errArgs != nil {
@@ -1105,6 +1174,12 @@ func TestConfigSubscriberUsesAppliedLifecycleRevisionAndRebuildsCommands(t *test
 	if timeout != 4*time.Second {
 		t.Fatalf("receive timeout = %s", timeout)
 	}
+	client.recoveryState.Store(uint32(recoveryStateSwitchingTakeover))
+	args, _ = client.subscriptionParameters()
+	if !reflect.DeepEqual(args, []string{"config", "9", "takeover"}) {
+		t.Fatalf("takeover subscribe args = %#v", args)
+	}
+	client.recoveryState.Store(uint32(recoveryStateStable))
 	client.promoteSubscription()
 	client.mu.Lock()
 	commandClient := client.cmd
@@ -1177,8 +1252,9 @@ func TestRunConfigSubscriberLifetimeReturnsAfterHeartbeatLoss(t *testing.T) {
 	client.mu.Lock()
 	client.clusterNodes = []clusterNode{{IP: "failover.example.com", Port: 8327}}
 	client.mu.Unlock()
+	client.recoveryState.Store(uint32(recoveryStateSwitchingTakeover))
 
-	ready := make(chan struct{}, 1)
+	ready := make(chan bool, 1)
 	errRun := client.RunConfigSubscriberLifetime(context.Background(), func(raw []byte) error {
 		parsed, errParse := config.ParseConfigBytes(raw)
 		if errParse != nil {
@@ -1188,12 +1264,15 @@ func TestRunConfigSubscriberLifetimeReturnsAfterHeartbeatLoss(t *testing.T) {
 			return errSet
 		}
 		return nil
-	}, func() { ready <- struct{}{} })
+	}, func() { ready <- recoveryState(client.recoveryState.Load()) == recoveryStateStable })
 	if errRun == nil {
 		t.Fatal("RunConfigSubscriberLifetime() error = nil after heartbeat loss")
 	}
 	select {
-	case <-ready:
+	case cleared := <-ready:
+		if !cleared {
+			t.Fatal("successful subscription ACK and command probe did not clear takeover state")
+		}
 	default:
 		t.Fatalf("RunConfigSubscriberLifetime() did not invoke onReady after subscription ACK: %v; commands=%#v", errRun, commands.All())
 	}
@@ -1203,6 +1282,9 @@ func TestRunConfigSubscriberLifetimeReturnsAfterHeartbeatLoss(t *testing.T) {
 	if got, _ := client.addr(); got != "failover.example.com:8327" {
 		t.Fatalf("addr() = %q, want failover.example.com:8327 after heartbeat timeout", got)
 	}
+	if got := recoveryState(client.recoveryState.Load()); got != recoveryStateSwitchingTakeover {
+		t.Fatalf("recovery state = %d, want %d", got, recoveryStateSwitchingTakeover)
+	}
 	client.mu.Lock()
 	commandClient, subscriptionClient := client.cmd, client.sub
 	client.mu.Unlock()
@@ -1215,8 +1297,8 @@ func TestRunConfigSubscriberLifetimeReturnsAfterHeartbeatLoss(t *testing.T) {
 	if count := commands.CountCommandKey("SUBSCRIBE", redisChannelConfig); count != 1 {
 		t.Fatalf("SUBSCRIBE config count = %d, want 1", count)
 	}
-	if got := findRedisCommand(commands.All(), "SUBSCRIBE"); !reflect.DeepEqual(got, []string{"subscribe", "config", "1"}) {
-		t.Fatalf("SUBSCRIBE wire command = %#v, want []string{\"subscribe\", \"config\", \"1\"}", got)
+	if got := findRedisCommand(commands.All(), "SUBSCRIBE"); !reflect.DeepEqual(got, []string{"subscribe", "config", "1", "takeover"}) {
+		t.Fatalf("SUBSCRIBE wire command = %#v, want []string{\"subscribe\", \"config\", \"1\", \"takeover\"}", got)
 	}
 }
 

internal/home/concurrency_release.go

@@ -61,6 +61,17 @@ func (f *releaseFlusher) SetConfigProvider(provider func() internalconfig.Creden
 	f.signal()
 }
 
+// SetSender replaces the Home lifetime used for subsequent release attempts.
+func (f *releaseFlusher) SetSender(send func(context.Context, ConcurrencyReleaseFrame) error) {
+	if f == nil {
+		return
+	}
+	f.mu.Lock()
+	f.send = send
+	f.mu.Unlock()
+	f.signal()
+}
+
 // MarkDirty records the latest cumulative sequence for one release group and
 // returns a ticket completed when Home acknowledges that sequence.
 func (f *releaseFlusher) MarkDirty(group executionregistry.ReleaseGroup, sequence int64) *executionregistry.ReleaseTicket {
@@ -174,11 +185,12 @@ func (f *releaseFlusher) timings() releaseFlusherTimings {
 }
 
 func (f *releaseFlusher) flush(ctx context.Context) bool {
-	if f == nil || f.send == nil {
+	if f == nil {
 		return false
 	}
 
 	f.mu.Lock()
+	send := f.send
 	pending := make(map[executionregistry.ReleaseGroup]int64, len(f.groups))
 	for group, state := range f.groups {
 		if state.Latest > state.Acked {
@@ -186,10 +198,13 @@ func (f *releaseFlusher) flush(ctx context.Context) bool {
 		}
 	}
 	f.mu.Unlock()
+	if send == nil {
+		return false
+	}
 
 	failed := false
 	for group, sequence := range pending {
-		errSend := f.send(ctx, ConcurrencyReleaseFrame{
+		errSend := send(ctx, ConcurrencyReleaseFrame{
 			CredentialID: group.CredentialID,
 			Model:        group.Model,
 			ReleaseSeq:   sequence,

internal/home/concurrency_release_test.go

@@ -474,3 +474,32 @@ func TestScopeEndBlocksDrainUntilReleaseSinkFlushesFinalSequence(t *testing.T) {
 		t.Fatalf("final flushed sequence = %d, want 1", got)
 	}
 }
+
+func TestReleaseFlusherSenderReplacementPreservesTicket(t *testing.T) {
+	flusher := newReleaseFlusher(time.Hour, time.Hour, func(context.Context, ConcurrencyReleaseFrame) error {
+		return errors.New("old Home unavailable")
+	})
+	group := executionregistry.ReleaseGroup{CredentialID: "cred-1", Model: "gpt"}
+	ticket := flusher.MarkDirty(group, 1)
+	if ticket == nil {
+		t.Fatal("MarkDirty() ticket = nil")
+	}
+	if failed := flusher.flush(context.Background()); !failed {
+		t.Fatal("old sender release attempt did not fail")
+	}
+
+	flusher.SetSender(func(_ context.Context, frame ConcurrencyReleaseFrame) error {
+		if frame.CredentialID != group.CredentialID || frame.Model != group.Model || frame.ReleaseSeq != 1 {
+			t.Fatalf("replacement sender frame = %#v", frame)
+		}
+		return nil
+	})
+	if failed := flusher.flush(context.Background()); failed {
+		t.Fatal("replacement sender release attempt failed")
+	}
+	waitCtx, cancelWait := context.WithTimeout(context.Background(), time.Second)
+	defer cancelWait()
+	if errWait := ticket.Wait(waitCtx); errWait != nil {
+		t.Fatalf("ticket did not survive sender replacement: %v", errWait)
+	}
+}

sdk/cliproxy/executionregistry/registry.go

@@ -154,6 +154,30 @@ func (r *Registry) BeginDispatch() (*PendingDispatch, error) {
 	return pending, nil
 }
 
+// WaitPending waits until every dispatch with an unresolved Home response has ended or been installed.
+func (r *Registry) WaitPending(ctx context.Context) error {
+	if r == nil {
+		return ErrRegistryClosed
+	}
+	if ctx == nil {
+		ctx = context.Background()
+	}
+
+	r.mu.Lock()
+	for len(r.pending) != 0 {
+		changed := r.changed
+		r.mu.Unlock()
+		select {
+		case <-ctx.Done():
+			return ctx.Err()
+		case <-changed:
+		}
+		r.mu.Lock()
+	}
+	r.mu.Unlock()
+	return nil
+}
+
 // End releases a dispatch token that was not installed.
 func (p *PendingDispatch) End() {
 	if p == nil || p.registry == nil {

sdk/cliproxy/executionregistry/registry_test.go

@@ -91,6 +91,42 @@ func TestDrainWaitsForPendingDispatch(t *testing.T) {
 	}
 }
 
+func TestWaitPendingDoesNotDrainActiveScope(t *testing.T) {
+	registry := New()
+	activePending, errBegin := registry.BeginDispatch()
+	if errBegin != nil {
+		t.Fatal(errBegin)
+	}
+	scope, errInstall := registry.Install(activePending, ScopeSpec{})
+	if errInstall != nil {
+		t.Fatal(errInstall)
+	}
+	defer scope.End("test cleanup")
+	pending, errBegin := registry.BeginDispatch()
+	if errBegin != nil {
+		t.Fatal(errBegin)
+	}
+
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
+	defer cancel()
+	done := make(chan error, 1)
+	go func() { done <- registry.WaitPending(ctx) }()
+	select {
+	case errWait := <-done:
+		t.Fatalf("WaitPending() returned before pending dispatch ended: %v", errWait)
+	case <-time.After(20 * time.Millisecond):
+	}
+	pending.End()
+	if errWait := <-done; errWait != nil {
+		t.Fatalf("WaitPending() error = %v", errWait)
+	}
+	nextPending, errNext := registry.BeginDispatch()
+	if errNext != nil {
+		t.Fatalf("WaitPending() stopped registry acceptance: %v", errNext)
+	}
+	nextPending.End()
+}
+
 func TestDrainReturnsWhenBlockingResourceCloseExceedsContext(t *testing.T) {
 	registry := New()
 	pending, errBegin := registry.BeginDispatch()

sdk/cliproxy/service_executionregistry_test.go

@@ -1250,7 +1250,7 @@ func TestServiceExplicitReplacementCancelsRunWhenDrainTimesOut(t *testing.T) {
 	scope.End("test cleanup")
 }
 
-func TestServiceReplacesRegistryOnlyAfterNewSubscriptionAck(t *testing.T) {
+func TestServiceKeepsRegistryAcrossHeartbeatFailoverAndExposesOnlyAfterNewACK(t *testing.T) {
 	listener, errListen := net.Listen("tcp", "127.0.0.1:0")
 	if errListen != nil {
 		t.Fatalf("listen: %v", errListen)
@@ -1325,15 +1325,15 @@ func TestServiceReplacesRegistryOnlyAfterNewSubscriptionAck(t *testing.T) {
 
 	close(allowSecondAck)
 	secondRegistry := waitForServiceRegistry(t, service, time.Second)
-	if secondRegistry == firstRegistry {
-		t.Fatal("replacement subscription reused the old registry")
+	if secondRegistry != firstRegistry {
+		t.Fatal("heartbeat failover replaced the execution registry")
 	}
 	if home.Current() == nil {
 		t.Fatal("replacement client was not exposed after the replacement ACK")
 	}
 }
 
-func TestServiceDrainsBeforePreAckRetriesAndExposesOnlyAfterNewAck(t *testing.T) {
+func TestServicePreservesActiveScopeDuringPreACKFailoverRetries(t *testing.T) {
 	listener, errListen := net.Listen("tcp", "127.0.0.1:0")
 	if errListen != nil {
 		t.Fatalf("listen: %v", errListen)
@@ -1393,16 +1393,9 @@ func TestServiceDrainsBeforePreAckRetriesAndExposesOnlyAfterNewAck(t *testing.T)
 	close(loseFirst)
 	select {
 	case <-resourceClosed:
-	case <-time.After(time.Second):
-		t.Fatal("heartbeat loss did not close the active scope resource")
-	}
-	select {
-	case <-preAckAttempts:
-		t.Fatal("pre-ACK retry started before the active scope owner ended")
+		t.Fatal("heartbeat failover drained the active scope")
 	case <-time.After(50 * time.Millisecond):
 	}
-	scope.End("canceled")
-
 	firstPreAck := <-preAckAttempts
 	secondPreAck := <-preAckAttempts
 	if retryDelay := secondPreAck.Sub(firstPreAck); retryDelay < 75*time.Millisecond {
@@ -1423,12 +1416,13 @@ func TestServiceDrainsBeforePreAckRetriesAndExposesOnlyAfterNewAck(t *testing.T)
 
 	close(allowFinalAck)
 	secondRegistry := waitForServiceRegistry(t, service, time.Second)
-	if secondRegistry == firstRegistry || home.Current() == nil {
+	if secondRegistry != firstRegistry || home.Current() == nil {
 		t.Fatal("new Home lifetime was not exposed only after its subscription ACK")
 	}
+	scope.End("completed")
 }
 
-func TestServiceCancelsRunWhenBlockingScopeExceedsDrainBound(t *testing.T) {
+func TestServiceHeartbeatFailoverDoesNotDrainBlockingScope(t *testing.T) {
 	listener, errListen := net.Listen("tcp", "127.0.0.1:0")
 	if errListen != nil {
 		t.Fatalf("listen: %v", errListen)
@@ -1507,30 +1501,243 @@ func TestServiceCancelsRunWhenBlockingScopeExceedsDrainBound(t *testing.T) {
 	close(loseFirst)
 	select {
 	case <-started:
-	case <-time.After(time.Second):
-		t.Fatal("drain did not start closing the blocking scope")
+		t.Fatal("heartbeat failover started draining the blocking scope")
+	case <-time.After(50 * time.Millisecond):
 	}
 	select {
 	case <-secondSubscribe:
-		t.Fatal("new subscription started before the old registry drained")
-	case <-time.After(50 * time.Millisecond):
+	case <-time.After(time.Second):
+		t.Fatal("new subscription did not start while the old scope remained active")
 	}
 	service.homeMu.Lock()
 	exposedRegistry := service.homeRegistry
 	service.homeMu.Unlock()
 	if exposedRegistry != nil {
-		t.Fatal("new registry was exposed before the old registry drained")
+		t.Fatal("registry was exposed before the replacement ACK")
+	}
+	close(allowSecondAck)
+	if nextRegistry := waitForServiceRegistry(t, service, time.Second); nextRegistry != registry {
+		t.Fatal("heartbeat failover replaced the registry containing the active scope")
 	}
 	select {
 	case <-serviceCtx.Done():
-	case <-time.After(time.Second):
-		t.Fatal("service run was not canceled after drain timeout")
+		t.Fatal("heartbeat failover canceled the service run")
+	case <-time.After(50 * time.Millisecond):
 	}
 
 	close(release)
 	scope.End("test cleanup")
 }
 
+func TestServiceShutdownDrainsDetachedRegistryDuringRetry(t *testing.T) {
+	listener, errListen := net.Listen("tcp", "127.0.0.1:0")
+	if errListen != nil {
+		t.Fatalf("listen: %v", errListen)
+	}
+	firstAck := make(chan struct{})
+	loseFirst := make(chan struct{})
+	secondSubscribe := make(chan struct{})
+	var secondSubscribeOnce sync.Once
+	allowSecondAck := make(chan struct{})
+	stop := make(chan struct{})
+	var subscriptionMu sync.Mutex
+	subscriptions := 0
+	serverDone := make(chan struct{})
+	go func() {
+		defer close(serverDone)
+		for {
+			conn, errAccept := listener.Accept()
+			if errAccept != nil {
+				return
+			}
+			go serveRegistryTestHomeConnection(conn, &subscriptionMu, &subscriptions, firstAck, loseFirst, secondSubscribe, &secondSubscribeOnce, allowSecondAck, stop)
+		}
+	}()
+	t.Cleanup(func() {
+		close(stop)
+		_ = listener.Close()
+		<-serverDone
+		home.ClearCurrent()
+	})
+
+	service := newRegistryTestService(t, listener)
+	serviceCtx, cancelService := context.WithCancel(context.Background())
+	t.Cleanup(cancelService)
+	service.homeMu.Lock()
+	service.runCancel = cancelService
+	service.homeMu.Unlock()
+	service.startHomeSubscriber(serviceCtx)
+
+	select {
+	case <-firstAck:
+	case <-time.After(time.Second):
+		t.Fatal("first subscription was not acknowledged")
+	}
+	registry := waitForServiceRegistry(t, service, time.Second)
+	pendingRetry, errBegin := registry.BeginDispatch()
+	if errBegin != nil {
+		t.Fatal(errBegin)
+	}
+	pendingScope, errBegin := registry.BeginDispatch()
+	if errBegin != nil {
+		t.Fatal(errBegin)
+	}
+	scope, errInstall := registry.Install(pendingScope, executionregistry.ScopeSpec{})
+	if errInstall != nil {
+		t.Fatal(errInstall)
+	}
+	resourceClosed := make(chan struct{})
+	if errBind := scope.Bind(func() error {
+		close(resourceClosed)
+		go scope.End("shutdown")
+		return nil
+	}); errBind != nil {
+		t.Fatal(errBind)
+	}
+	t.Cleanup(func() {
+		pendingRetry.End()
+		scope.End("test cleanup")
+	})
+
+	service.homeMu.Lock()
+	client := service.homeClient
+	service.homeMu.Unlock()
+	if client == nil {
+		t.Fatal("ready Home client is unavailable")
+	}
+	close(loseFirst)
+	deadline := time.After(time.Second)
+	for {
+		errRelease := client.PushConcurrencyRelease(context.Background(), home.ConcurrencyReleaseFrame{CredentialID: "cred-a", Model: "model-a", ReleaseSeq: 1})
+		if errors.Is(errRelease, home.ErrDispatchFenced) {
+			break
+		}
+		select {
+		case <-deadline:
+			t.Fatal("subscriber retry did not close the previous Home client")
+		case <-time.After(time.Millisecond):
+		}
+	}
+
+	shutdownDone := make(chan error, 1)
+	go func() {
+		shutdownDone <- service.Shutdown(context.Background())
+	}()
+	pendingRetry.End()
+
+	select {
+	case <-resourceClosed:
+	case <-time.After(time.Second):
+		t.Fatal("shutdown did not drain the detached execution registry")
+	}
+	select {
+	case errShutdown := <-shutdownDone:
+		if errShutdown != nil {
+			t.Fatalf("Shutdown() error = %v", errShutdown)
+		}
+	case <-time.After(time.Second):
+		t.Fatal("Shutdown() did not complete after draining the detached registry")
+	}
+}
+
+func TestServiceAmbiguousDispatchDrainsRegistryBeforeRetry(t *testing.T) {
+	listener, errListen := net.Listen("tcp", "127.0.0.1:0")
+	if errListen != nil {
+		t.Fatalf("listen: %v", errListen)
+	}
+	firstAck := make(chan struct{})
+	loseFirst := make(chan struct{})
+	secondSubscribe := make(chan struct{})
+	var secondSubscribeOnce sync.Once
+	allowSecondAck := make(chan struct{})
+	stop := make(chan struct{})
+	var subscriptionMu sync.Mutex
+	subscriptions := 0
+	serverDone := make(chan struct{})
+	go func() {
+		defer close(serverDone)
+		for {
+			conn, errAccept := listener.Accept()
+			if errAccept != nil {
+				return
+			}
+			go serveRegistryTestHomeConnection(conn, &subscriptionMu, &subscriptions, firstAck, loseFirst, secondSubscribe, &secondSubscribeOnce, allowSecondAck, stop)
+		}
+	}()
+	t.Cleanup(func() {
+		close(stop)
+		_ = listener.Close()
+		<-serverDone
+		home.ClearCurrent()
+	})
+
+	service := newRegistryTestService(t, listener)
+	serviceCtx, cancelService := context.WithCancel(context.Background())
+	t.Cleanup(cancelService)
+	service.homeMu.Lock()
+	service.runCancel = cancelService
+	service.homeMu.Unlock()
+	service.startHomeSubscriber(serviceCtx)
+
+	select {
+	case <-firstAck:
+	case <-time.After(time.Second):
+		t.Fatal("first subscription was not acknowledged")
+	}
+	registry := waitForServiceRegistry(t, service, time.Second)
+	pending, errBegin := registry.BeginDispatch()
+	if errBegin != nil {
+		t.Fatal(errBegin)
+	}
+	scope, errInstall := registry.Install(pending, executionregistry.ScopeSpec{})
+	if errInstall != nil {
+		t.Fatal(errInstall)
+	}
+	resourceClosed := make(chan struct{})
+	if errBind := scope.Bind(func() error {
+		close(resourceClosed)
+		go scope.End("ambiguous dispatch")
+		return nil
+	}); errBind != nil {
+		t.Fatal(errBind)
+	}
+
+	service.homeMu.Lock()
+	client := service.homeClient
+	service.homeMu.Unlock()
+	if client == nil {
+		t.Fatal("ready Home client is unavailable")
+	}
+	client.AbortAmbiguousDispatch()
+	select {
+	case <-resourceClosed:
+	case <-time.After(time.Second):
+		t.Fatal("ambiguous dispatch did not drain the active registry")
+	}
+	select {
+	case <-secondSubscribe:
+	case <-time.After(time.Second):
+		t.Fatal("subscriber did not retry after ambiguous dispatch drain")
+	}
+	service.homeMu.Lock()
+	exposedRegistry := service.homeRegistry
+	service.homeMu.Unlock()
+	if exposedRegistry != nil {
+		t.Fatal("replacement registry was exposed before its subscription ACK")
+	}
+
+	close(allowSecondAck)
+	nextRegistry := waitForServiceRegistry(t, service, time.Second)
+	if nextRegistry == registry {
+		t.Fatal("ambiguous dispatch reused the drained execution registry")
+	}
+	select {
+	case <-serviceCtx.Done():
+		t.Fatal("successful ambiguous dispatch recovery canceled the service run")
+	case <-time.After(50 * time.Millisecond):
+	}
+}
+
 func TestServiceBacksOffAfterRepeatedPreAckFailures(t *testing.T) {
 	listener, errListen := net.Listen("tcp", "127.0.0.1:0")
 	if errListen != nil {
@@ -1585,7 +1792,7 @@ func TestServiceBacksOffAfterRepeatedPreAckFailures(t *testing.T) {
 	}
 }
 
-func TestServiceHeartbeatLossCancelsBlockedConfigFinalizationBeforeDrain(t *testing.T) {
+func TestServiceHeartbeatLossCancelsBlockedConfigFinalizationWithoutDrainingRegistry(t *testing.T) {
 	listener, errListen := net.Listen("tcp", "127.0.0.1:0")
 	if errListen != nil {
 		t.Fatalf("listen: %v", errListen)
@@ -1648,8 +1855,8 @@ func TestServiceHeartbeatLossCancelsBlockedConfigFinalizationBeforeDrain(t *test
 	}
 	select {
 	case <-resourceClosed:
-	case <-time.After(500 * time.Millisecond):
-		t.Fatal("heartbeat loss did not cancel the worker and drain the active execution")
+		t.Fatal("heartbeat loss drained the active execution")
+	case <-time.After(200 * time.Millisecond):
 	}
 	select {
 	case <-secondConfig:
@@ -1663,6 +1870,7 @@ func TestServiceHeartbeatLossCancelsBlockedConfigFinalizationBeforeDrain(t *test
 	if currentRegistry != nil || currentClient != nil || home.Current() != nil {
 		t.Fatal("heartbeat-lost lifetime left a published Home client or registry")
 	}
+	scope.End("completed")
 }
 
 func TestServiceConfigWorkerFinalizesRapidUpdatesInOrder(t *testing.T) {

sdk/cliproxy/service_home.go

@@ -492,6 +492,25 @@ func (s *Service) runHomeSubscriber(homeCtx context.Context, parentCtx context.C
 	}()
 
 	var previousClient *home.Client
+	registry := executionregistry.New()
+	cancelBound := atomic.Int64{}
+	cancelBound.Store(int64(internalconfig.CredentialConcurrencyConfig{}.WithDefaults().CPACancelBound))
+	releaseFlusher := home.NewReleaseFlusher(nil, nil)
+	registry.SetReleaseSink(releaseFlusher.MarkDirty)
+	defer func() {
+		registry.SetReleaseSink(nil)
+		drainBound := time.Duration(cancelBound.Load())
+		if drainBound <= 0 {
+			drainBound = internalconfig.CredentialConcurrencyConfig{}.WithDefaults().CPACancelBound
+		}
+		drainCtx, cancelDrain := context.WithTimeout(context.WithoutCancel(parentCtx), drainBound)
+		errDrain := registry.Drain(drainCtx)
+		cancelDrain()
+		if errDrain != nil && !errors.Is(errDrain, executionregistry.ErrRegistryClosed) && parentCtx.Err() == nil {
+			log.WithError(errDrain).Error("failed to drain detached Home execution registry")
+			s.cancelServiceRun()
+		}
+	}()
 	for homeCtx.Err() == nil {
 		supervisor.setPublisherCompletion(nil)
 		client := previousClient
@@ -501,18 +520,15 @@ func (s *Service) runHomeSubscriber(homeCtx context.Context, parentCtx context.C
 			client = client.NewLifetime()
 		}
 		client.SetManagedLifetime(true)
-		registry := executionregistry.New()
 		releaseCtx, releaseCancel := context.WithCancel(context.WithoutCancel(homeCtx))
-		releaseFlusher := home.NewReleaseFlusher(client.LimiterConfig, client.PushConcurrencyRelease)
-		registry.SetReleaseSink(releaseFlusher.MarkDirty)
+		releaseFlusher.SetConfigProvider(client.LimiterConfig)
+		releaseFlusher.SetSender(client.PushConcurrencyRelease)
 		releaseDone := make(chan struct{})
 		go func() {
 			defer close(releaseDone)
 			releaseFlusher.Run(releaseCtx)
 		}()
 		lifetimeCtx, lifetimeCancel := context.WithCancel(homeCtx)
-		cancelBound := atomic.Int64{}
-		cancelBound.Store(int64(internalconfig.CredentialConcurrencyConfig{}.WithDefaults().CPACancelBound))
 		queue := newHomeConfigWorkQueue()
 		ready := make(chan struct{})
 		var readyOnce sync.Once
@@ -552,6 +568,44 @@ func (s *Service) runHomeSubscriber(homeCtx context.Context, parentCtx context.C
 		}
 
 		s.detachHomeSubscriberLifetime(client, registry)
+		retry := errRun != nil && homeCtx.Err() == nil
+		if retry {
+			releaseCancel()
+			<-releaseDone
+			client.Close()
+
+			settleBound := time.Duration(cancelBound.Load())
+			settleCtx, cancelSettle := context.WithTimeout(context.WithoutCancel(parentCtx), settleBound)
+			errPending := registry.WaitPending(settleCtx)
+			cancelSettle()
+			if errPending != nil {
+				log.WithError(errPending).Error("failed to settle pending Home dispatches before subscriber replacement")
+				s.cancelServiceRun()
+				return
+			}
+			if client.AmbiguousDispatch() || home.IsMembershipTakeoverUnavailableError(errRun) {
+				registry.SetReleaseSink(nil)
+				drainCtx, cancelDrain := context.WithTimeout(context.WithoutCancel(parentCtx), settleBound)
+				errDrain := registry.Drain(drainCtx)
+				cancelDrain()
+				if errDrain != nil {
+					log.WithError(errDrain).Error("failed to drain Home executions after unsafe subscriber replacement")
+					s.cancelServiceRun()
+					return
+				}
+				client.SuppressTakeover()
+				registry = executionregistry.New()
+				releaseFlusher = home.NewReleaseFlusher(nil, nil)
+				registry.SetReleaseSink(releaseFlusher.MarkDirty)
+			}
+			log.WithError(errRun).Warn("home config subscription lifetime ended")
+			if !published.Load() && !waitForHomeSubscriberRetry(homeCtx, homeSubscriberPreAckRetryBackoff) {
+				return
+			}
+			previousClient = client
+			continue
+		}
+
 		drainBound := time.Duration(cancelBound.Load())
 		drainCtx, cancelDrain := context.WithTimeout(context.WithoutCancel(parentCtx), drainBound)
 		errDrain := registry.Drain(drainCtx)
@@ -577,13 +631,7 @@ func (s *Service) runHomeSubscriber(homeCtx context.Context, parentCtx context.C
 			}
 			return
 		}
-		if errRun != nil && homeCtx.Err() == nil {
-			log.WithError(errRun).Warn("home config subscription lifetime ended")
-		}
-		if !published.Load() && errRun != nil && !waitForHomeSubscriberRetry(homeCtx, homeSubscriberPreAckRetryBackoff) {
-			return
-		}
-		previousClient = client
+		return
 	}
 }