fix(mcp): harden MCP tool registry lifecycle and address PR #155 review findings

- Publish connected MCP tools before raising on partial connect failure, so
  servers that connected stay callable when a sibling fails.
- Rebuild the published MCP registry atomically on disconnect/refresh/reconnect:
  a failed refresh no longer drops live tools, and disconnecting a server falls
  its shadowed tool names back to other still-connected servers instead of
  orphaning them.
- Convert refresh inventory timeouts/errors to MCPRuntimeError at the boundary.
- Validate /mcp verb arity instead of silently using only the first operand.
- Convert MCP config errors to typer.BadParameter in CLI parsing.
- Stage suggestion prefill unconditionally so an empty prefill clears stale text.
- Document telemetry slice budget, in-place mcpServers normalization, and recall
  windowing params; strengthen suggest/config/emoji assertions; drop a duplicate
  ExitPlanMode description test; add MCP lifecycle failure-path tests.

Author: elkaix

src/pythinker_code/cli/__init__.py

@@ -270,6 +270,8 @@ def _load_mcp_configs_from_cli_inputs(
             file_configs.append(project_mcp_file)
 
     configs: list[Any] = []
+    from pythinker_code.exception import MCPConfigError
+
     from .mcp import prepare_mcp_config_dict
 
     for conf in file_configs:
@@ -285,12 +287,19 @@ def _load_mcp_configs_from_cli_inputs(
                 f"Cannot read MCP config file {conf}: {e}",
                 param_hint="--mcp-config-file",
             ) from e
+        except MCPConfigError as e:
+            raise typer.BadParameter(
+                f"Invalid MCP config in file {conf}: {e}",
+                param_hint="--mcp-config-file",
+            ) from e
 
     for conf in raw_mcp_config:
         try:
             configs.append(prepare_mcp_config_dict(json.loads(conf)))
         except json.JSONDecodeError as e:
             raise typer.BadParameter(f"Invalid JSON: {e}", param_hint="--mcp-config") from e
+        except MCPConfigError as e:
+            raise typer.BadParameter(f"Invalid MCP config: {e}", param_hint="--mcp-config") from e
 
     for path in _yaml_files_with_misplaced_mcp_servers():
         from pythinker_code.utils.logging import logger

src/pythinker_code/soul/toolset.py

@@ -624,6 +624,23 @@ def _publish_connected_mcp_tools(self, runtime: Runtime) -> None:
 
                 runtime.mcp_tools[mcp_tool_runtime_key(server_name, tool.name)] = tool
 
+    def _rebuild_published_mcp_tools(self, runtime: Runtime) -> None:
+        """Atomically rebuild the published MCP tool registry from connected servers.
+
+        Drop every currently-published MCP tool (non-MCP tools are preserved), then
+        republish all *connected* servers in configured order via
+        :meth:`_publish_connected_mcp_tools`. This keeps last-wins collision order
+        deterministic and lets a disconnect/refresh re-claim a tool name another
+        still-connected server provides, instead of orphaning it. The method runs
+        synchronously (no ``await`` between the drop and the republish), so the two
+        registries are never observed half-rebuilt.
+        """
+        stale = [name for name, tool in self._tool_dict.items() if isinstance(tool, MCPTool)]
+        for name in stale:
+            del self._tool_dict[name]
+        runtime.mcp_tools.clear()
+        self._publish_connected_mcp_tools(runtime)
+
     def hide(self, tool_name: str) -> bool:
         """Hide a tool from the LLM tool list. Returns True if the tool exists."""
         if tool_name in self._tool_dict:
@@ -1300,10 +1317,12 @@ async def _connect():
             results = await asyncio.gather(*tasks) if tasks else []
             failed_servers = {name: error for name, error in results if error is not None}
 
+            # Publish before raising so servers that DID connect become callable in
+            # this session even when another server fails the aggregate connect.
+            self._publish_connected_mcp_tools(runtime)
             if failed_servers:
                 _toast_mcp("mcp connection failed")
                 raise MCPRuntimeError(f"Failed to connect MCP servers: {failed_servers}")
-            self._publish_connected_mcp_tools(runtime)
             if unauthorized_servers:
                 _toast_mcp("mcp authorization needed")
             else:
@@ -1353,19 +1372,6 @@ async def wait_for_mcp_tools(self) -> None:
             if self._mcp_loading_task is task and task.done():
                 self._mcp_loading_task = None
 
-    def _unregister_mcp_server_tools(self, server_name: str, runtime: Runtime) -> None:
-        info = self._mcp_servers.get(server_name)
-        if info is None:
-            return
-        from pythinker_code.utils.mcp_names import mcp_tool_runtime_key
-
-        for tool in info.tools:
-            registered = self._tool_dict.get(tool.name)
-            if registered is tool:
-                del self._tool_dict[tool.name]
-            runtime.mcp_tools.pop(mcp_tool_runtime_key(server_name, tool.name), None)
-        info.tools = []
-
     async def _inventory_mcp_server(
         self, server_name: str, server_info: MCPServerInfo, runtime: Runtime
     ) -> None:
@@ -1427,16 +1433,6 @@ async def _connect_mcp_server(
             server_info.error = _classify_mcp_connect_error(e, server_name)
             return server_name, e
 
-    def _publish_mcp_server_tools(self, server_name: str, runtime: Runtime) -> None:
-        info = self._mcp_servers.get(server_name)
-        if info is None:
-            return
-        self._register_mcp_tools(server_name, info.tools)
-        from pythinker_code.utils.mcp_names import mcp_tool_runtime_key
-
-        for tool in info.tools:
-            runtime.mcp_tools[mcp_tool_runtime_key(server_name, tool.name)] = tool
-
     def _ensure_mcp_idle(self) -> None:
         if self._mcp_loading_task is not None and not self._mcp_loading_task.done():
             raise MCPRuntimeError("MCP servers are still loading")
@@ -1463,7 +1459,10 @@ async def disconnect_mcp_server(self, server_name: str, runtime: Runtime) -> Non
         if info is None:
             raise MCPRuntimeError(f"Unknown MCP server: {server_name}")
         await self._stop_mcp_session_holder(info)
-        self._unregister_mcp_server_tools(server_name, runtime)
+        # Drop this server's inventory, then rebuild the published registry so any
+        # tool name it was shadowing falls back to another still-connected server.
+        info.tools = []
+        self._rebuild_published_mcp_tools(runtime)
         prior_error = info.error
         close_error: str | None = None
         try:
@@ -1504,12 +1503,21 @@ async def refresh_mcp_server(self, server_name: str, runtime: Runtime) -> None:
             raise MCPRuntimeError(
                 f"MCP server '{server_name}' is not connected (status={info.status})"
             )
-        self._unregister_mcp_server_tools(server_name, runtime)
-        await asyncio.wait_for(
-            self._inventory_mcp_server(server_name, info, runtime),
-            timeout=runtime.config.mcp.client.startup_timeout_ms / 1000,
-        )
-        self._publish_mcp_server_tools(server_name, runtime)
+        # Inventory first: on failure ``info.tools`` keeps its last-known-good value
+        # and the live registry is untouched, so a failed refresh never drops tools.
+        # Convert raw timeout/inventory errors to MCPRuntimeError so callers (e.g. the
+        # /mcp slash handler) receive a single typed boundary error.
+        try:
+            await asyncio.wait_for(
+                self._inventory_mcp_server(server_name, info, runtime),
+                timeout=runtime.config.mcp.client.startup_timeout_ms / 1000,
+            )
+        except TimeoutError as exc:
+            raise MCPRuntimeError(f"Refresh of MCP server '{server_name}' timed out") from exc
+        except Exception as exc:
+            raise MCPRuntimeError(f"Failed to refresh MCP server '{server_name}': {exc}") from exc
+        # Inventory succeeded; swap the old tool set for the new one atomically.
+        self._rebuild_published_mcp_tools(runtime)
 
     async def reconnect_mcp_server(self, server_name: str, runtime: Runtime) -> None:
         """Close and reconnect one MCP server from its stored config."""
@@ -1532,7 +1540,7 @@ async def reconnect_mcp_server(self, server_name: str, runtime: Runtime) -> None
             raise MCPRuntimeError(
                 info.error or f"Failed to reconnect MCP server '{server_name}': {error}"
             )
-        self._publish_mcp_server_tools(server_name, runtime)
+        self._rebuild_published_mcp_tools(runtime)
 
     async def cleanup(self) -> None:
         """Cleanup any resources held by the toolset."""

src/pythinker_code/telemetry/names.py

@@ -37,5 +37,7 @@ def _bounded_mcp_label(server: str, tool: str) -> str:
     if len(label) <= _TELEMETRY_TOOL_NAME_MAX:
         return label
     digest = hashlib.sha256(label.encode("utf-8")).hexdigest()[:8]
+    # Slice budget: "mcp__" (5) + server[:16] + "__" (2) + tool[:24] = 47, then
+    # head[:55] + "_" + 8-char digest = 56 — always within _TELEMETRY_TOOL_NAME_MAX (64).
     head = f"mcp__{safe_server[:16]}__{safe_tool[:24]}"
     return f"{head[: _TELEMETRY_TOOL_NAME_MAX - 9]}_{digest}"

src/pythinker_code/tools/recall/__init__.py

@@ -103,7 +103,15 @@ def _render_transcript(
 
     Internal (``_``-prefixed) roles are skipped. Each message's text is sanitized;
     a block that trips the secret/injection scanner becomes ``[redacted]`` rather
-    than leaking or silently vanishing. Stops once the char budget is reached.
+    than leaking or silently vanishing.
+
+    Windowing operates on the renderable (post-filter) message stream — i.e. after
+    internal-role and empty-segment messages are dropped:
+
+    - ``message_offset``: skip this many renderable messages before emitting any.
+    - ``max_messages``: emit at most this many renderable messages (``None`` = no limit).
+
+    Stops once the char budget or ``max_messages`` is reached.
     """
     out: list[str] = []
     used = 0

src/pythinker_code/ui/shell/slash.py

@@ -2258,6 +2258,13 @@ async def mcp(app: Shell, args: str):
             if soul is None:
                 return
             server_name = parts[1] if len(parts) > 1 else None
+            # Reject extra operands instead of silently using only the first.
+            if verb in {"reconnect", "disconnect", "retry"} and len(parts) != 2:
+                console.print(f"[{_get_tok_mcp().warning}]Usage: /mcp {verb} <server>[/]")
+                return
+            if verb == "refresh" and len(parts) > 2:
+                console.print(f"[{_get_tok_mcp().warning}]Usage: /mcp refresh [server][/]")
+                return
             if verb == "retry":
                 verb = "reconnect"
             try:

src/pythinker_code/ui/shell/visualize/_interactive.py

@@ -517,8 +517,10 @@ def dispatch_wire_message(self, msg: WireMessage) -> None:
 
     def display_suggestion(self, event: Suggestion) -> None:
         super().display_suggestion(event)
-        if event.prefill.strip():
-            self._prompt_session.stage_suggestion_prefill(event.prefill)
+        # Stage unconditionally: an empty prefill clears any prior staged value
+        # (stage_suggestion_prefill stores ``None`` for blank input), so a later
+        # suggestion without a prefill cannot leave stale Esc+s text behind.
+        self._prompt_session.stage_suggestion_prefill(event.prefill)
 
     # -- Running prompt rendering --------------------------------------------
 

src/pythinker_code/utils/mcp_names.py

@@ -36,7 +36,11 @@ def mcp_tool_runtime_key(server_name: str, tool_name: str) -> str:
 
 
 def normalize_mcp_servers_in_config(config: dict[str, Any]) -> dict[str, Any]:
-    """Re-key ``mcpServers`` entries to normalized names; fail on collisions."""
+    """Re-key ``mcpServers`` entries to normalized names; fail on collisions.
+
+    Mutates ``config`` in place (reassigns ``config["mcpServers"]``) and returns
+    the same dict for chaining.
+    """
     servers = config.get("mcpServers")
     if not isinstance(servers, dict):
         return config

tests/core/test_exit_plan_mode_verification.py

@@ -1,4 +1,3 @@
-from pythinker_code.tools.plan import ExitPlanMode
 from pythinker_code.tools.plan.__init__ import _plan_lacks_verification_section
 
 
@@ -7,6 +6,7 @@ def test_plan_lacks_verification_section_detects_missing_heading() -> None:
     assert not _plan_lacks_verification_section("## Plan\n## Verification\nmake test\n")
 
 
-def test_exit_plan_mode_description_requires_verification_section():
-    tool = ExitPlanMode()
-    assert "Verification section" in tool.base.description
+# The ExitPlanMode description assertion lives in
+# tests/tools/test_tool_descriptions.py::test_exit_plan_mode_description_requires_verification_section
+# (it checks both the "Verification section" heading and the "smallest command,
+# test, or check" guidance). This file is scoped to the _plan_lacks_verification_section helper.

tests/core/test_mcp_lifecycle.py

@@ -201,3 +201,106 @@ class _FakeClient:
 
     with pytest.raises(MCPRuntimeError, match="command not found"):
         await toolset.reconnect_mcp_server("alpha", runtime)
+
+
+@pytest.mark.asyncio
+async def test_refresh_failure_preserves_live_tools(monkeypatch: pytest.MonkeyPatch) -> None:
+    """A failed refresh must not drop the server's still-working tools."""
+    toolset = PythinkerToolset()
+    runtime = _runtime()
+    live_tool = _fake_mcp_tool("alpha", "LiveTool")
+    info = MCPServerInfo(
+        status="connected",
+        client=cast(Any, SimpleNamespace()),
+        tools=[live_tool],
+        resources=[],
+        prompts=[],
+    )
+    toolset._mcp_servers["alpha"] = info
+    toolset.add(live_tool)
+    runtime.mcp_tools["mcp__alpha__LiveTool"] = live_tool
+
+    async def _failing_inventory(_server: str, _info: MCPServerInfo, _runtime: Any) -> None:
+        raise RuntimeError("list_tools blew up")
+
+    monkeypatch.setattr(toolset, "_inventory_mcp_server", _failing_inventory)
+
+    with pytest.raises(MCPRuntimeError, match="Failed to refresh"):
+        await toolset.refresh_mcp_server("alpha", runtime)
+
+    assert toolset.find("LiveTool") is live_tool
+    assert runtime.mcp_tools["mcp__alpha__LiveTool"] is live_tool
+    assert info.tools == [live_tool]
+
+
+@pytest.mark.asyncio
+async def test_disconnect_reclaims_shadowed_tool_from_other_server() -> None:
+    """Disconnecting the winning server must fall the tool name back, not orphan it."""
+    toolset = PythinkerToolset()
+    runtime = _runtime()
+    shared_alpha = _fake_mcp_tool("alpha", "Shared")
+    shared_beta = _fake_mcp_tool("beta", "Shared")
+    toolset._mcp_servers["alpha"] = MCPServerInfo(
+        status="connected",
+        client=cast(Any, SimpleNamespace(close=AsyncMock())),
+        tools=[shared_alpha],
+        resources=[],
+        prompts=[],
+        server_config={"command": "echo"},
+    )
+    toolset._mcp_servers["beta"] = MCPServerInfo(
+        status="connected",
+        client=cast(Any, SimpleNamespace(close=AsyncMock())),
+        tools=[shared_beta],
+        resources=[],
+        prompts=[],
+        server_config={"command": "echo"},
+    )
+    # Configured order publishes beta last, so it wins the shared name.
+    toolset._publish_connected_mcp_tools(runtime)
+    assert toolset.find("Shared") is shared_beta
+
+    await toolset.disconnect_mcp_server("beta", runtime)
+
+    assert toolset.find("Shared") is shared_alpha
+    assert runtime.mcp_tools["mcp__alpha__Shared"] is shared_alpha
+    assert "mcp__beta__Shared" not in runtime.mcp_tools
+    assert toolset._mcp_servers["beta"].status == "failed"
+
+
+@pytest.mark.asyncio
+async def test_partial_connect_publishes_connected_servers(monkeypatch: pytest.MonkeyPatch) -> None:
+    """A server that connects must be published even when a sibling fails the aggregate."""
+    from fastmcp.mcp_config import MCPConfig
+
+    toolset = PythinkerToolset()
+    runtime = _runtime()
+    good_tool = _fake_mcp_tool("good", "GoodTool")
+
+    async def _connect(
+        server_name: str, server_info: MCPServerInfo, _runtime: Any
+    ) -> tuple[str, Exception | None]:
+        if server_name == "good":
+            server_info.status = "connected"
+            server_info.tools = [good_tool]
+            return server_name, None
+        server_info.status = "failed"
+        server_info.error = "boom"
+        return server_name, RuntimeError("boom")
+
+    monkeypatch.setattr(toolset, "_connect_mcp_server", _connect)
+    monkeypatch.setattr(
+        "pythinker_code.soul.toolset._configure_mcp_client_handlers",
+        lambda *args, **kwargs: None,
+    )
+    monkeypatch.setattr("fastmcp.Client", lambda *args, **kwargs: SimpleNamespace())
+
+    config = MCPConfig.model_validate(
+        {"mcpServers": {"good": {"command": "echo"}, "bad": {"command": "echo"}}}
+    )
+
+    with pytest.raises(MCPRuntimeError, match="Failed to connect"):
+        await toolset.load_mcp_tools([config], runtime, in_background=False)
+
+    assert toolset.find("GoodTool") is good_tool
+    assert runtime.mcp_tools["mcp__good__GoodTool"] is good_tool

tests/core/test_toolset.py

@@ -309,7 +309,7 @@ def fire_and_forget_trigger(*args: Any, **kwargs: Any) -> None:
         )
     )
     assert isinstance(result, asyncio.Task)
-    await result
+    _ = await result  # drive the task to completion; the awaited value is unused
 
     assert post_inputs == [{"payload": {"value": "original"}}]