mcpx: recover resumed HTTP sessions and clean removed servers (#1416)

pkgs/mcpx/internal/daemon/registry.go

@@ -50,7 +50,15 @@ type Registry struct {
 	paths Paths
 	hash  string
 
-	mu sync.RWMutex
+	// reloadMu serializes configuration swaps with asynchronous cleanup from
+	// earlier swaps. A cleanup must not remove sessions created by a later
+	// reload that re-adds the same server name.
+	reloadMu sync.Mutex
+	// pendingSessionCleanup protects persisted records until the removed
+	// pool has drained and can terminate them with its last live config.
+	// Access is serialized by reloadMu.
+	pendingSessionCleanup map[string]int
+	mu                    sync.RWMutex
 	// pools may be shared between names; views never are. A view carries the
 	// per-name presentation -- namespace, tool subset, prelude, profiles --
 	// that must not follow the shared process.

pkgs/mcpx/internal/daemon/registry_handoff_test.go

@@ -19,11 +19,12 @@ import (
 // on initialize and counts the DELETEs that end sessions.
 type orphanUpstream struct {
 	*httptest.Server
-	mu      sync.Mutex
-	live    map[string]bool
-	issued  int
-	deletes int
-	deleted []string
+	mu            sync.Mutex
+	live          map[string]bool
+	issued        int
+	deletes       int
+	deleted       []string
+	deleteMarkers []string
 	// blockDelete, when non-nil, holds a DELETE in the handler until it is
 	// closed (or the request goes away); deleteStarted is signalled when a
 	// blocked DELETE arrives. They simulate a blackholed upstream (#291).
@@ -60,6 +61,7 @@ func (u *orphanUpstream) serve(w http.ResponseWriter, r *http.Request) {
 	case http.MethodDelete:
 		u.deletes++
 		u.deleted = append(u.deleted, sid)
+		u.deleteMarkers = append(u.deleteMarkers, r.Header.Get("X-Reload-Marker"))
 		if u.live[sid] {
 			delete(u.live, sid)
 			w.WriteHeader(http.StatusOK)
@@ -106,6 +108,12 @@ func (u *orphanUpstream) counts() (deletes int, deleted []string) {
 	return u.deletes, append([]string(nil), u.deleted...)
 }
 
+func (u *orphanUpstream) deletedMarkers() []string {
+	u.mu.Lock()
+	defer u.mu.Unlock()
+	return append([]string(nil), u.deleteMarkers...)
+}
+
 // detachRemote starts a session on a registry whose only server is the
 // upstream, then detaches it for a successor.
 func detachRemote(t *testing.T, url string) (*Registry, *RegistryHandoff) {

pkgs/mcpx/internal/daemon/reload.go

@@ -27,6 +27,8 @@ func (r *Registry) Reload(cfg *config.Config) (added, removed []string, err erro
 		return nil, nil, err
 	}
 
+	r.reloadMu.Lock()
+	defer r.reloadMu.Unlock()
 	r.mu.Lock()
 	defer r.mu.Unlock()
 
@@ -84,26 +86,74 @@ func (r *Registry) Reload(cfg *config.Config) (added, removed []string, err erro
 		keep[p] = true
 	}
 	var orphans []*pool.Pool
+	removedPools := map[*pool.Pool]bool{}
 	for name, p := range r.pools {
 		if _, still := pools[name]; !still {
 			removed = append(removed, name)
+			removedPools[p] = true
 		}
 		if !keep[p] {
 			orphans = append(orphans, p)
 		}
 	}
 	sort.Strings(removed)
-
 	r.cfg = cfg
 	r.pools, r.views, r.order = pools, views, order
 	r.hash = HashConfig(mustJSON(cfg.MCPServers))
+	// Mark orphans closed before the replacement map is visible. A later
+	// reload must not resume a record whose old client is still draining and
+	// will DELETE that upstream session.
+	for _, p := range orphans {
+		p.CloseDrained()
+	}
+	if r.pendingSessionCleanup == nil {
+		r.pendingSessionCleanup = map[string]int{}
+	}
+	for p := range removedPools {
+		r.pendingSessionCleanup[p.Name()]++
+	}
 
 	// Drained, not closed: a call running on an orphan finishes on it, as it
 	// does on a lazy restart. Close would cut one that ran past its grace.
 	go func() {
 		for _, p := range orphans {
-			p.CloseDrained()
+			p.WaitDrained()
+			if removedPools[p] {
+				r.reloadMu.Lock()
+				r.mu.RLock()
+				_, reconfigured := r.pools[p.Name()]
+				r.mu.RUnlock()
+				if !reconfigured {
+					p.CleanupRemovedSessionRecords()
+				}
+				if r.pendingSessionCleanup[p.Name()] <= 1 {
+					delete(r.pendingSessionCleanup, p.Name())
+				} else {
+					r.pendingSessionCleanup[p.Name()]--
+				}
+				r.reloadMu.Unlock()
+			}
+		}
+		r.reloadMu.Lock()
+		r.mu.RLock()
+		configured := make(map[string]bool, len(r.pools))
+		for name := range r.pools {
+			configured[name] = true
+		}
+		// A later reload can have no orphan pools while an earlier removed
+		// pool is still draining. Keep those records out of the generic
+		// discard path until the owner has attempted DELETE with its last
+		// live config.
+		for name := range r.pendingSessionCleanup {
+			configured[name] = true
+		}
+		var sessions pool.SessionStore
+		if r.hooks != nil {
+			sessions = r.hooks.Sessions
 		}
+		r.mu.RUnlock()
+		pool.DiscardUnconfiguredSessionRecords(sessions, configured)
+		r.reloadMu.Unlock()
 	}()
 	return added, removed, nil
 }