fix(mcp): retain parked startup tasks for clean shutdown

tools/mcp_tool.py

@@ -3021,7 +3021,7 @@ class MCPServerTask:
         self._register_discovered_tools_if_needed()
 
     def _register_discovered_tools_if_needed(self) -> None:
-        """Re-register tools after a post-ready reconnect if needed.
+        """Re-register tools after an owned server reconnects if needed.
 
         Initial registration is performed by ``_discover_and_register_server``
         after ``start()`` completes. During a later reconnect, outage handling
@@ -3029,6 +3029,9 @@ class MCPServerTask:
         A managed server can still be identified by its entry in ``_servers``;
         publish its freshly discovered tools before transport readiness is
         restored so a successful revival cannot come back with zero tools.
+        A server retained after a recoverable initial failure is likewise
+        registry-owned before its first successful session, so ownership also
+        authorizes its first publication.
         """
         if self._registered_tool_names:
             return
@@ -3039,6 +3042,12 @@ class MCPServerTask:
         self._registered_tool_names = _register_server_tools(
             self.name, self, self._config
         )
+        # A retained initial-failure server that just published tools has
+        # recovered: drop its stale connect error so status surfaces stop
+        # reporting it as failed.
+        with _lock:
+            if _servers.get(self.name) is self:
+                _server_connect_errors.pop(self.name, None)
 
     async def run(self, config: dict):
         """Long-lived coroutine: connect, discover tools, wait, disconnect.
@@ -3523,6 +3532,12 @@ class MCPServerTask:
 _servers: Dict[str, MCPServerTask] = {}
 _server_connecting: set[str] = set()
 _server_connect_errors: Dict[str, str] = {}
+# Discovery installs a task-local claim before calling ``_connect_server`` so
+# it can retain a recoverable parked task without making standalone probe calls
+# publish failed servers into module-global ownership.
+_connect_server_claim: contextvars.ContextVar[
+    Optional[Callable[[MCPServerTask], None]]
+] = contextvars.ContextVar("mcp_connect_server_claim", default=None)
 
 # Connection-retry cooldown (per-server isolation against restart storms).
 #
@@ -3632,9 +3647,8 @@ def _signal_reconnect(server: Any) -> bool:
     The tool handlers run on caller threads, while the server task and its
     ``_reconnect_event`` live on the background MCP loop. Setting an
     asyncio.Event from another thread must go through
-    ``loop.call_soon_threadsafe``; only fall back to a direct ``.set()``
-    when the loop isn't running (e.g. unit tests that drive the handler
-    synchronously).
+    ``loop.call_soon_threadsafe``; non-async adapters and tests without a
+    running loop can use a direct ``.set()``.
 
     Returns True if a reconnect signal was delivered, False if the server
     has no reconnect machinery (nothing to revive).
@@ -3643,7 +3657,11 @@ def _signal_reconnect(server: Any) -> bool:
     if event is None:
         return False
     loop = _mcp_loop
-    if loop is not None and loop.is_running():
+    if (
+        isinstance(event, asyncio.Event)
+        and loop is not None
+        and loop.is_running()
+    ):
         loop.call_soon_threadsafe(event.set)
     else:
         event.set()
@@ -4616,7 +4634,26 @@ async def _connect_server(name: str, config: dict) -> MCPServerTask:
         Exception: on connection or initialization failure.
     """
     server = MCPServerTask(name)
-    await server.start(config)
+    claim = _connect_server_claim.get()
+    claim_token = None
+    if claim is not None:
+        claim(server)
+        # ``start()`` creates the long-lived run task by copying this context.
+        # The ownership callback is only for this connection attempt; do not
+        # retain its discovery closure for the server's lifetime.
+        claim_token = _connect_server_claim.set(None)
+    try:
+        await server.start(config)
+    except BaseException:
+        # Discovery owns claimed tasks and decides whether a failed start is a
+        # live recoverable park or a terminal failure. Standalone probes have
+        # no revival owner, so they must reap their failed task locally.
+        if claim is None:
+            await server.shutdown()
+        raise
+    finally:
+        if claim_token is not None:
+            _connect_server_claim.reset(claim_token)
     return server
 
 
@@ -5788,10 +5825,42 @@ async def _discover_and_register_server(name: str, config: dict) -> List[str]:
     Returns list of registered tool names.
     """
     connect_timeout = config.get("connect_timeout", _DEFAULT_CONNECT_TIMEOUT)
-    server = await asyncio.wait_for(
-        _connect_server(name, config),
-        timeout=connect_timeout,
-    )
+    server: Optional[MCPServerTask] = None
+
+    def _claim_server(created: MCPServerTask) -> None:
+        nonlocal server
+        server = created
+
+    claim_token = _connect_server_claim.set(_claim_server)
+    try:
+        server = await asyncio.wait_for(
+            _connect_server(name, config),
+            timeout=connect_timeout,
+        )
+    except BaseException:
+        task = server._task if server is not None else None
+        task_cancelling = (
+            task.cancelling()
+            if task is not None and hasattr(task, "cancelling")
+            else 0
+        )
+        recoverable_park = (
+            server is not None
+            and server._error is not None
+            and task is not None
+            and not task.done()
+            and not task_cancelling
+        )
+        if recoverable_park:
+            with _lock:
+                _servers[name] = server
+        elif server is not None:
+            await server.shutdown()
+        raise
+    finally:
+        _connect_server_claim.reset(claim_token)
+
+    assert server is not None
     with _lock:
         _server_connecting.discard(name)
         _server_connect_errors.pop(name, None)
@@ -5932,7 +6001,11 @@ def register_mcp_servers(servers: Dict[str, dict]) -> List[str]:
 
     # Log a summary so ACP callers get visibility into what was registered.
     with _lock:
-        connected = [n for n in new_servers if n in _servers]
+        connected = [
+            n
+            for n in new_servers
+            if n in _servers and n not in _server_connect_errors
+        ]
         new_tool_count = sum(
             len(getattr(_servers[n], "_registered_tool_names", []))
             for n in connected
@@ -6005,7 +6078,11 @@ def discover_mcp_tools() -> List[str]:
             return tool_names
 
         with _lock:
-            connected_server_names = [name for name in new_server_names if name in _servers]
+            connected_server_names = [
+                name
+                for name in new_server_names
+                if name in _servers and name not in _server_connect_errors
+            ]
             new_tool_count = sum(
                 len(getattr(_servers[name], "_registered_tool_names", []))
                 for name in connected_server_names