fix(mcp): scope error teardown to the erroring session; serialize ref…

internal/agent/tools/mcp/init.go

@@ -601,8 +601,9 @@ func getOrRenewClient(ctx context.Context, cfg *config.ConfigStore, name string)
 
 	state, _ := states.Get(name)
 	// StateError closes the dead session and clears its tools, prompts, and
-	// resources from the registry.
-	updateState(name, StateError, maybeTimeoutErr(pingErr, timeout), nil, state.Counts)
+	// resources from the registry. Report the failure against the exact
+	// session that failed the ping so only that session is torn down.
+	updateState(name, StateError, maybeTimeoutErr(pingErr, timeout), sess, state.Counts)
 
 	// Capture the generation so a reconcile teardown that lands mid-renewal
 	// invalidates this rebuild instead of letting it clobber the newer one.
@@ -744,23 +745,37 @@ func updateState(name string, state State, err error, client *ClientSession, cou
 		info.Config = config.MCPConfig{}
 		info.PendingConfig = nil
 	case StateError:
-		// A session that has errored is dead to us. Atomically remove it and
-		// close it so the child process and its stdio pipes are released — the
-		// bare map delete this used to do leaked both. Clearing the tool
-		// registry keeps the agent from advertising tools it can no longer
-		// call: without it, crush_info / the `/mcp` menu and the tool list
-		// handed to the LLM diverge, so a server still reads "connected, N
-		// tools" while every call fails with "tool not found".
-		if old, ok := sessions.Take(name); ok {
-			closeSession(name, old)
+		// A session that has errored is dead to us: close it so the child
+		// process and its stdio pipes are released, and clear its registry
+		// entries so the agent stops advertising capabilities it can no
+		// longer call (without that, crush_info / the `/mcp` menu and the
+		// tool list handed to the LLM diverge). Crucially, close exactly the
+		// session that errored (the client argument): if the registry
+		// already holds a DIFFERENT session — a newer, healthy one another
+		// path installed — leave it and its registrations alone. Closing
+		// "whatever is in the map" here let a stale error transition (e.g. a
+		// refresh that raced a renewal) tear down the healthy replacement.
+		switch {
+		case client != nil:
+			if cur, ok := sessions.Get(name); ok && cur == client {
+				sessions.Del(name)
+				allTools.Del(name)
+				allPrompts.Del(name)
+				allResources.Del(name)
+			}
+			closeSession(name, client)
+		default:
+			// No specific session errored (e.g. connect itself failed);
+			// anything still registered under this name is unusable.
+			if old, ok := sessions.Take(name); ok {
+				closeSession(name, old)
+			}
+			allTools.Del(name)
+			allPrompts.Del(name)
+			allResources.Del(name)
 		}
-		// Drop every registry entry for the dead server. Leaving prompts or
-		// resources behind lets a disconnected server keep advertising
-		// capabilities the agent can no longer fulfil, the same divergence the
-		// tool clear prevents.
-		allTools.Del(name)
-		allPrompts.Del(name)
-		allResources.Del(name)
+		// Never publish a dead session on the state.
+		info.Client = nil
 	}
 	states.Set(name, info)
 

internal/agent/tools/mcp/lifecycle_test.go

@@ -123,6 +123,83 @@ func TestUpdateState_ErrorClosesSessionAndClearsTools(t *testing.T) {
 	require.Equal(t, StateError, info.State)
 }
 
+// TestUpdateState_ErrorFromStaleSessionPreservesHealthyReplacement pins the
+// teardown scoping: a StateError reported against a session that is NO LONGER
+// the registered one (a renewal already replaced it) must not tear down the
+// healthy replacement or its registrations. Before the fix updateState closed
+// whatever session was in the map, so a stale error transition — e.g. a
+// refresh whose list call timed out after another path had already renewed —
+// killed the fresh session and wiped its tools.
+func TestUpdateState_ErrorFromStaleSessionPreservesHealthyReplacement(t *testing.T) {
+	const name = "test-stale-error"
+	t.Cleanup(func() {
+		sessions.Del(name)
+		allTools.Del(name)
+		allPrompts.Del(name)
+		allResources.Del(name)
+		states.Del(name)
+	})
+
+	stale, staleCtx := liveSession(t, "old_tool")
+	fresh, freshCtx := liveSession(t, "new_tool")
+
+	// The registry holds the fresh session and its registrations.
+	sessions.Set(name, fresh)
+	allTools.Set(name, []*Tool{{Name: "new_tool"}})
+	allPrompts.Set(name, []*Prompt{{Name: "new_prompt"}})
+
+	// A stale error arrives for the OLD session.
+	updateState(name, StateError, errors.New("ping timeout"), stale, Counts{})
+
+	// The fresh session must still be registered and open.
+	got, ok := sessions.Get(name)
+	require.True(t, ok, "healthy replacement session was removed")
+	require.Same(t, fresh, got)
+	require.NoError(t, freshCtx.Err(), "healthy replacement session was closed")
+	_, ok = allTools.Get(name)
+	require.True(t, ok, "healthy replacement's tools were cleared")
+	_, ok = allPrompts.Get(name)
+	require.True(t, ok, "healthy replacement's prompts were cleared")
+
+	// The stale session must have been closed.
+	require.ErrorIs(t, staleCtx.Err(), context.Canceled, "stale session must still be closed")
+}
+
+// TestUpdateState_ErrorFromCurrentSessionClearsEverything pins the complement:
+// when the erroring session IS the registered one, the teardown must behave
+// exactly as before the scoping — session removed and closed, every registry
+// entry cleared, and the published state must not carry the dead session.
+func TestUpdateState_ErrorFromCurrentSessionClearsEverything(t *testing.T) {
+	const name = "test-current-error"
+	t.Cleanup(func() {
+		sessions.Del(name)
+		allTools.Del(name)
+		allPrompts.Del(name)
+		allResources.Del(name)
+		states.Del(name)
+	})
+
+	sess, sessCtx := liveSession(t, "do_thing")
+	sessions.Set(name, sess)
+	allTools.Set(name, []*Tool{{Name: "do_thing"}})
+	allPrompts.Set(name, []*Prompt{{Name: "a_prompt"}})
+
+	updateState(name, StateError, errors.New("pipe broke"), sess, Counts{})
+
+	_, ok := sessions.Get(name)
+	require.False(t, ok, "errored current session must be removed")
+	require.ErrorIs(t, sessCtx.Err(), context.Canceled, "errored current session must be closed")
+	_, ok = allTools.Get(name)
+	require.False(t, ok, "errored current session's tools must be cleared")
+	_, ok = allPrompts.Get(name)
+	require.False(t, ok, "errored current session's prompts must be cleared")
+
+	info, ok := GetState(name)
+	require.True(t, ok)
+	require.Equal(t, StateError, info.State)
+	require.Nil(t, info.Client, "a dead session must never be published on the state")
+}
+
 // TestUpdateState_ConfigBookkeeping pins the config snapshot reconcile relies
 // on: StateConnected records the config now in effect and clears any pending
 // attempt, StateStarting records the config the in-flight attempt is using,

internal/agent/tools/mcp/prompts.go

@@ -48,6 +48,13 @@ func GetPromptMessages(ctx context.Context, cfg *config.ConfigStore, clientName,
 // RefreshPrompts gets the updated list of prompts from the MCP and updates the
 // global state.
 func RefreshPrompts(ctx context.Context, name string) {
+	// Serialize with session renewal so the registered session can't be
+	// swapped between the Get and the state update below — a stale error
+	// transition would otherwise tear down the healthy replacement.
+	mu := renewLock(name)
+	mu.Lock()
+	defer mu.Unlock()
+
 	session, ok := sessions.Get(name)
 	if !ok {
 		slog.Warn("Refresh prompts: no session", "name", name)
@@ -56,7 +63,7 @@ func RefreshPrompts(ctx context.Context, name string) {
 
 	prompts, err := getPrompts(ctx, session)
 	if err != nil {
-		updateState(name, StateError, err, nil, Counts{})
+		updateState(name, StateError, err, session, Counts{})
 		return
 	}
 

internal/agent/tools/mcp/resources.go

@@ -58,6 +58,13 @@ func ReadResource(ctx context.Context, cfg *config.ConfigStore, name, uri string
 // RefreshResources gets the updated list of resources from the MCP and updates the
 // global state.
 func RefreshResources(ctx context.Context, name string) {
+	// Serialize with session renewal so the registered session can't be
+	// swapped between the Get and the state update below — a stale error
+	// transition would otherwise tear down the healthy replacement.
+	mu := renewLock(name)
+	mu.Lock()
+	defer mu.Unlock()
+
 	session, ok := sessions.Get(name)
 	if !ok {
 		slog.Warn("Refresh resources: no session", "name", name)
@@ -66,7 +73,7 @@ func RefreshResources(ctx context.Context, name string) {
 
 	resources, err := getResources(ctx, session)
 	if err != nil {
-		updateState(name, StateError, err, nil, Counts{})
+		updateState(name, StateError, err, session, Counts{})
 		return
 	}
 

internal/agent/tools/mcp/tools.go

@@ -111,6 +111,13 @@ func RunTool(ctx context.Context, cfg *config.ConfigStore, name, toolName string
 // RefreshTools gets the updated list of tools from the MCP and updates the
 // global state.
 func RefreshTools(ctx context.Context, cfg *config.ConfigStore, name string) {
+	// Serialize with session renewal so the registered session can't be
+	// swapped between the Get and the state update below — a stale error
+	// transition would otherwise tear down the healthy replacement.
+	mu := renewLock(name)
+	mu.Lock()
+	defer mu.Unlock()
+
 	session, ok := sessions.Get(name)
 	if !ok {
 		slog.Warn("Refresh tools: no session", "name", name)
@@ -119,7 +126,7 @@ func RefreshTools(ctx context.Context, cfg *config.ConfigStore, name string) {
 
 	tools, err := getTools(ctx, session)
 	if err != nil {
-		updateState(name, StateError, err, nil, Counts{})
+		updateState(name, StateError, err, session, Counts{})
 		return
 	}