pcode 0.1.0__py3-none-any.whl

pcode/mcp.py

@@ -0,0 +1,683 @@
+"""Explicit, conversation-scoped MCP activation; configuration alone does nothing."""
+
+import functools
+import json
+import logging
+import os
+import re
+import sys
+from collections.abc import Callable
+from contextlib import AsyncExitStack
+from dataclasses import replace
+from pathlib import Path
+from typing import Any, Literal
+from urllib.parse import urlsplit
+
+from pydantic import BaseModel, ConfigDict, ValidationError, model_validator
+
+from pcode.preferences import preferences_path
+
+_NAME = re.compile(r"[A-Za-z][A-Za-z0-9_-]{0,31}\Z")
+_ENV = re.compile(r"\$\{([A-Za-z_][A-Za-z0-9_]*)(?::-([^}]*))?\}")
+# Packages whose exceptions mean the OAuth flow itself failed.
+OAUTH_PACKAGES = ("mcp.client.auth", "fastmcp.client.auth", "pcode.mcp_oauth")
+STDIO_INIT_TIMEOUT = 30.0  # seconds from launch until the server answers its handshake
+
+
+class MCPConnectError(RuntimeError):
+    """An enabled server failed to connect: names it, with the HTTP status if any.
+
+    The message holds only fixed text, the server name, and the config path,
+    never anything the server sent, so it is safe to show as-is.
+    """
+
+    def __init__(self, server: str, message: str):
+        super().__init__(message)
+        self.server = server
+
+
+def mcp_transport(toolset: Any) -> Any:
+    """The client transport of the MCPToolset inside our prefix/defer wrappers."""
+    while (wrapped := getattr(toolset, "wrapped", None)) is not None:
+        toolset = wrapped
+    return getattr(getattr(toolset, "client", None), "transport", None)
+
+
+def _without_deferred_loading(toolset: Any) -> Any | None:
+    """Rebuild a toolset without its deferred-loading layer, or `None` if it has none.
+
+    Only the wrappers are rebuilt: the `MCPToolset` itself (with its client, its
+    OAuth object and whatever tokens it already holds) is the same instance, so
+    dropping deferral costs neither a reconnection nor a second sign-in.
+    """
+    from pydantic_ai.toolsets.deferred_loading import DeferredLoadingToolset
+
+    if isinstance(toolset, DeferredLoadingToolset):
+        return toolset.wrapped
+    wrapped = getattr(toolset, "wrapped", None)
+    if wrapped is None:
+        return None
+    inner = _without_deferred_loading(wrapped)
+    return None if inner is None else replace(toolset, wrapped=inner)
+
+
+def config_path() -> Path:
+    override = os.environ.get("PCODE_MCP_CONFIG", "").strip()
+    return Path(override).expanduser() if override else preferences_path().with_name("mcp.json")
+
+
+def configured_servers() -> dict[str, Any]:
+    """Read names without importing MCP, expanding secrets, or starting servers."""
+    path = config_path()
+    try:
+        data = json.loads(path.read_text())
+    except FileNotFoundError:
+        if not os.environ.get("PCODE_MCP_CONFIG", "").strip():
+            return {}
+        raise ValueError(f"MCP configuration not found: {path}") from None
+    except (OSError, ValueError):
+        raise ValueError(f"Cannot read MCP configuration: {path} (expected JSON)") from None
+    return _servers(data)
+
+
+def _servers(data: Any) -> dict[str, Any]:
+    if not isinstance(data, dict) or not isinstance(data.get("mcpServers"), dict):
+        raise ValueError("MCP configuration must contain an mcpServers object.")
+    servers = data["mcpServers"]
+    if any(not _NAME.fullmatch(name) for name in servers):
+        raise ValueError(
+            "MCP server names must start with a letter and contain only letters, digits, "
+            "underscores or hyphens (up to 32 characters)."
+        )
+    return servers
+
+
+def default_servers() -> list[str]:
+    """Names marked `"enabled": true`, without validating or expanding the rest."""
+    return sorted(
+        name
+        for name, raw in configured_servers().items()
+        if isinstance(raw, dict) and raw.get("enabled") is True
+    )
+
+
+def save_default(name: str, enabled: bool) -> None:
+    """Set or clear a server's `"enabled": true`, leaving the rest of the file as is.
+
+    Off is the default, so turning one off drops the key rather than writing false.
+    """
+    from filelock import FileLock
+
+    from pcode.preferences import write_json
+
+    # Write through a symlink (a dotfile-managed config) rather than replacing it.
+    path = config_path().resolve()
+    with FileLock(str(path) + ".lock", timeout=5):
+        try:
+            data = json.loads(path.read_text())
+        except (OSError, ValueError):
+            raise ValueError(f"Cannot read MCP configuration: {path} (expected JSON)") from None
+        server = _servers(data).get(name)
+        if not isinstance(server, dict):
+            raise ValueError(f"Unknown MCP server '{name}'. Use /mcp list.")
+        if enabled:
+            server["enabled"] = True
+        else:
+            server.pop("enabled", None)
+        write_json(path, data)
+
+
+def _expand(value: Any) -> Any:
+    if isinstance(value, str):
+
+        def replace(match: re.Match) -> str:
+            name, default = match.groups()
+            if name in os.environ:
+                return os.environ[name]
+            if default is not None:
+                return default
+            raise ValueError(f"Missing MCP environment variable: {name}")
+
+        return _ENV.sub(replace, value)
+    if isinstance(value, list):
+        return [_expand(item) for item in value]
+    if isinstance(value, dict):
+        return {key: _expand(item) for key, item in value.items()}
+    return value
+
+
+class ServerConfig(BaseModel):
+    model_config = ConfigDict(extra="forbid", strict=True)
+
+    command: str | None = None
+    args: list[str] | None = None
+    env: dict[str, str] | None = None
+    cwd: str | None = None
+    url: str | None = None
+    headers: dict[str, str] | None = None
+    auth: Literal["oauth"] | None = None
+    # A client created with the provider ahead of time, for services without
+    # dynamic client registration (Google's MCP servers, for one).
+    client_id: str | None = None
+    client_secret: str | None = None
+    # Enable at startup, /new, and resume instead of waiting for /mcp enable.
+    enabled: bool = False
+    # Off by default: a server's schemas otherwise sit in every request of the
+    # conversation, while tool search costs one call for the tools actually used.
+    direct: bool = False
+    # What the server is for, shown to the model beside its name in the list of
+    # enabled servers. Only worth setting when the name does not already say it.
+    description: str | None = None
+
+    @model_validator(mode="after")
+    def transport(self):
+        if bool(self.command) == bool(self.url):
+            raise ValueError("Specify exactly one of command or url.")
+        if self.client_secret is not None and self.client_id is None:
+            raise ValueError("client_secret needs client_id.")
+        if self.client_id is not None and (self.auth != "oauth" or not self.client_id.strip()):
+            raise ValueError('client_id needs auth: "oauth".')
+        if self.command:
+            if not self.command.strip() or any(
+                item is not None
+                for item in (self.headers, self.auth, self.client_id, self.client_secret)
+            ):
+                raise ValueError("Invalid stdio options.")
+            return self
+        if any(item is not None for item in (self.command, self.args, self.env, self.cwd)):
+            raise ValueError("Invalid HTTP options.")
+        parsed = urlsplit(self.url)
+        if parsed.scheme not in {"http", "https"} or not parsed.hostname:
+            raise ValueError("Expected an HTTP(S) URL.")
+        if self.auth and any(key.lower() == "authorization" for key in self.headers or {}):
+            raise ValueError("OAuth cannot be combined with an Authorization header.")
+        return self
+
+
+def _config(name: str, raw: Any) -> ServerConfig:
+    try:
+        return ServerConfig.model_validate(_expand(raw))
+    except ValidationError:
+        # Pydantic errors include input values: never print credentials from config.
+        raise ValueError(
+            f"Invalid MCP server '{name}'. Use command/args/env/cwd for stdio or url/headers "
+            'for HTTP (optional auth: "oauth" with client_id/client_secret, enabled: true, '
+            "direct: true, description); "
+            "other fields are not supported."
+        ) from None
+
+
+def _connect_failure(
+    server: str, config: ServerConfig, status: int | None, oauth_challenge: bool, error: Exception
+) -> str:
+    if status is None:
+        return f"MCP server '{server}' failed to connect ({type(error).__name__})."
+    from http import HTTPStatus
+
+    try:
+        reason = f" {HTTPStatus(status).phrase}"
+    except ValueError:
+        reason = ""
+    # A config edit reaches a server only when it is enabled again.
+    entry = f"its entry in {config_path()}"
+    then = f", then run `/mcp disable {server}` and `/mcp enable {server}`"
+    authorization = any(key.lower() == "authorization" for key in config.headers or {})
+    if status in (401, 403):
+        if config.auth == "oauth":
+            hint = f"Sign in again: `/mcp logout {server}`, then `/mcp enable {server}`."
+        elif oauth_challenge and not authorization:
+            hint = f'It offers OAuth sign-in: add "auth": "oauth" to {entry}{then}.'
+        elif config.headers:
+            hint = f"Check the credentials in the headers of {entry}{then}."
+        else:
+            hint = f"It needs credentials: add an Authorization header to {entry}{then}."
+    elif status == 404:
+        hint = f"Check the url in {entry}{then}."
+    elif status >= 500:
+        hint = "The server itself failed; try again later."
+    else:
+        hint = ""
+    return f"MCP server '{server}' failed to connect: HTTP {status}{reason}. {hint}".strip()
+
+
+@functools.cache
+def _recording_transport() -> type:
+    """A streamable HTTP transport that keeps its endpoint's last HTTP error.
+
+    The MCP SDK reduces any failed response but a 404 to a bare `MCPError`
+    ("Server returned an error response"), dropping the status. FastMCP
+    registers `_capture_session_id` as a response hook on every client it
+    builds, whatever the factory, so extending it sees each response.
+    """
+    import httpx2
+    from fastmcp.client.transports import StreamableHttpTransport
+
+    class RecordingTransport(StreamableHttpTransport):
+        failed_status: int | None = None
+        oauth_challenge = False
+
+        async def _capture_session_id(self, response: httpx2.Response) -> None:
+            await super()._capture_session_id(response)
+            request = response.request
+            # The OAuth flow's discovery, registration, and token requests travel
+            # the same client; only MCP messages to the endpoint itself count.
+            if request.method != "POST" or request.url != httpx2.URL(self.url):
+                return
+            failed = response.status_code >= 400
+            self.failed_status = response.status_code if failed else None
+            challenge = response.headers.get("www-authenticate", "") if failed else ""
+            self.oauth_challenge = "resource_metadata=" in challenge
+
+    return RecordingTransport
+
+
+@functools.cache
+def _named_toolset() -> type:
+    from dataclasses import dataclass
+
+    from pydantic_ai.toolsets import WrapperToolset
+
+    @dataclass
+    class NamedServer(WrapperToolset):
+        """Says which server failed to connect; the SDK's errors never do."""
+
+        server: str = ""
+        config: ServerConfig | None = None
+
+        async def __aenter__(self):
+            transport = mcp_transport(self.wrapped)
+            if hasattr(transport, "failed_status"):
+                transport.failed_status = None
+            try:
+                return await super().__aenter__()
+            except Exception as error:
+                status = getattr(transport, "failed_status", None)
+                # An OAuth failure has its own cause; the endpoint's 401 is just
+                # the challenge that started the flow.
+                if find_cause(error, _from_packages(OAUTH_PACKAGES)) is not None:
+                    status = None
+                challenge = getattr(transport, "oauth_challenge", False)
+                message = _connect_failure(self.server, self.config, status, challenge, error)
+                raise MCPConnectError(self.server, message) from error
+
+        async def call_tool(self, name, tool_args, ctx, tool):
+            from pydantic_ai.exceptions import ModelRetry
+
+            auth = getattr(mcp_transport(self.wrapped), "auth", None)
+            try:
+                return await super().call_tool(name, tool_args, ctx, tool)
+            except ModelRetry as error:
+                # A server that lists tools without sign-in can demand it on the
+                # first call. When that sign-in fails the session closes, and the
+                # model is told only "Connection closed", which it retries as a
+                # hiccup. Fail the turn with the sign-in error instead. Network
+                # errors during the flow stay retries.
+                failure = getattr(auth, "failure", None)
+                if failure is None or find_cause(failure, _from_packages(OAUTH_PACKAGES)) is None:
+                    raise
+                raise failure from error
+
+    return NamedServer
+
+
+def build_toolset(name: str, raw: Any, *, interactive: bool = True):
+    """Construct only the selected server. Connections are owned by each agent run.
+
+    A non-interactive OAuth server may use stored tokens and refresh them, but
+    fails with SignInRequired rather than opening a browser.
+    """
+    config = _config(name, raw)
+    from fastmcp.client.transports import StdioTransport
+    from pydantic_ai.mcp import MCPToolset
+
+    try:
+        if config.command:
+            transport = StdioTransport(
+                command=config.command,
+                args=config.args or [],
+                env=config.env,
+                cwd=config.cwd,
+                # FastMCP otherwise keeps subprocesses alive after toolset exit.
+                keep_alive=False,
+            )
+            # FastMCP's five-second handshake deadline counts the server's own
+            # startup, which a `uvx`/`npx` launch beside other servers can exceed.
+            toolset = MCPToolset(transport, id=name, init_timeout=STDIO_INIT_TIMEOUT)
+        else:
+            # FastMCP owns PKCE, browser sign-in, and refresh; pcode owns the socket
+            # and the credential file (see mcp_oauth).
+            from pcode.mcp_oauth import LoopbackOAuth
+
+            auth = (
+                LoopbackOAuth(
+                    interactive=interactive,
+                    client_id=config.client_id,
+                    client_secret=config.client_secret,
+                )
+                if config.auth == "oauth"
+                else None
+            )
+            from fastmcp.mcp_config import infer_transport_type_from_url
+
+            if infer_transport_type_from_url(config.url) == "sse":
+                server, http = config.url, {"headers": config.headers, "auth": auth}
+            else:
+                server = _recording_transport()(config.url, headers=config.headers, auth=auth)
+                http = {}
+            toolset = MCPToolset(
+                server,
+                id=name,
+                **http,
+                # The default five-second handshake deadline also covers OAuth.
+                # Give interactive sign-in a bounded five-minute window instead.
+                **({"init_timeout": 300} if auth is not None else {}),
+            )
+        toolset = _named_toolset()(toolset, server=name, config=config)
+        # Hidden until Pydantic AI's auto-injected ToolSearch reveals them, so a
+        # server's schemas cost one `search_tools` call instead of every request.
+        if not config.direct:
+            toolset = toolset.defer_loading()
+        return toolset.prefixed(f"mcp_{name}")
+    except (ValueError, TypeError):
+        raise ValueError(
+            f"Cannot configure MCP server '{name}'; check its transport options."
+        ) from None
+
+
+def deferred_schemas_rejected(error: BaseException) -> bool:
+    """Whether the provider rejected this request's hidden tool schemas.
+
+    Deferred loading is a request-shape choice, not history: a provider that
+    will not pair withheld schemas with its tool search rejects every request
+    the same way, so the session is stuck until the tools are sent in full.
+    Narrow on purpose -- a 400 naming the search surface is the provider saying
+    this pairing is impossible; any other 400 is a request pcode built wrong.
+    """
+    from pydantic_ai.exceptions import ModelHTTPError
+
+    return (
+        isinstance(error, ModelHTTPError)
+        and error.status_code == 400
+        and ("tool_search" in str(error.body) or "defer_loading" in str(error.body))
+    )
+
+
+def _from_packages(packages: tuple[str, ...]):
+    def raised_there(error: BaseException) -> bool:
+        module = type(error).__module__
+        return any(module == pkg or module.startswith(f"{pkg}.") for pkg in packages)
+
+    return raised_there
+
+
+def find_cause(error: BaseException, kind) -> BaseException | None:
+    """The first error in the chain that is a `kind`, or that `kind(error)` accepts."""
+    matches = kind if not isinstance(kind, type) else lambda current: isinstance(current, kind)
+    seen: set[int] = set()
+    pending = [error]
+    while pending:
+        current = pending.pop()
+        if current is None or id(current) in seen:
+            continue
+        seen.add(id(current))
+        if matches(current):
+            return current
+        pending.extend([current.__cause__, current.__context__])
+        pending.extend(getattr(current, "exceptions", ()))
+    return None
+
+
+def error_message(error: BaseException) -> str:
+    """Show redacted MCP causes, not model/provider troubleshooting advice."""
+    from pcode.diagnostics import error_details
+
+    detail = error_details(error)
+    messages = []
+    while detail:
+        message = " ".join(str(detail.get("message", "")).split())
+        summary = detail["type"] + (f": {message}" if message else "")
+        if summary not in messages:
+            messages.append(summary)
+        detail = detail.get("cause", detail.get("context", {}))
+    return (
+        "MCP connection or sign-in failed. "
+        + " Caused by: ".join(messages)[:2000]
+        + " Check the MCP server configuration, authentication, and connectivity."
+    )
+
+
+@functools.cache
+def _isolated() -> type:
+    """A per-server layer that turns a failed connection into an empty toolset.
+
+    pydantic-ai enters every run toolset in one exit stack, so one server whose
+    handshake fails would otherwise fail the whole turn and take every other
+    tool, built-in or MCP, down with it. The same holds for a connection that
+    drops later: listing its tools, or closing it when the run ends, raises
+    from inside the run, so both degrade the same way the handshake does.
+    """
+    from dataclasses import dataclass, field
+
+    from pydantic_ai.toolsets import WrapperToolset
+
+    @dataclass
+    class IsolatedServer(WrapperToolset):
+        server: str = ""
+        state: Any = None
+        # Per instance: exit only what this instance entered.
+        _entered: bool = field(default=False, init=False, compare=False, repr=False)
+
+        async def __aenter__(self):
+            try:
+                await self.wrapped.__aenter__()
+            except Exception as error:  # CancelledError is not an Exception.
+                self.state.connect_failed(self.server, error)
+            else:
+                self._entered = True
+                self.state.unavailable.pop(self.server, None)
+            return self
+
+        async def __aexit__(self, *args: Any) -> bool | None:
+            if not self._entered:
+                return None
+            self._entered = False
+            try:
+                return await self.wrapped.__aexit__(*args)
+            except Exception:
+                # The turn is already over; a server that went away meanwhile
+                # must not replace its result (or its real error) with this one.
+                logging.getLogger(__name__).debug(
+                    "MCP server %r failed to disconnect", self.server, exc_info=True
+                )
+                return None
+
+        # Keyed on the shared state, not `_entered`: pydantic-ai may hand these
+        # calls a per-step copy of this wrapper rather than the entered instance.
+        async def get_tools(self, ctx):
+            if self.server in self.state.unavailable:
+                return {}
+            try:
+                return await self.wrapped.get_tools(ctx)
+            except Exception as error:
+                self.state.connect_failed(self.server, error)
+                return {}
+
+        async def get_instructions(self, ctx):
+            if self.server in self.state.unavailable:
+                return None
+            return await super().get_instructions(ctx)
+
+    return IsolatedServer
+
+
+@functools.cache
+def _live() -> type:
+    """The enabled servers as of each model request, not as of the turn's start.
+
+    A run otherwise fixes its toolsets when it starts, so `/mcp enable` during a
+    long turn would wait for the next one. Instead each step (a model request,
+    or the tool calls answering one) compares the run's entered servers with
+    `MCPState.enabled`: a server enabled meanwhile is connected and listed from
+    the next request on, one disabled is disconnected, and the rest keep their
+    connections. One instance per run (`for_run`), so concurrent runs (a side
+    question, a worker) never share entered state. (`visit_and_replace` would
+    copy one without that state, but only durable execution calls it.)
+
+    The `<mcp-servers>` notice reads `MCPState` itself a moment later, in
+    `before_model_request`, so a change landing between the two can list a
+    server one request before (or after) its tools. The next request agrees.
+    """
+    from dataclasses import dataclass, field
+
+    from pydantic_ai.toolsets import CombinedToolset
+
+    @dataclass
+    class LiveServers(CombinedToolset):
+        toolsets: list = field(default_factory=list)
+        state: Any = None
+        # name -> (the enabled toolset, the entered wrapper around it)
+        _entered: dict = field(default_factory=dict, init=False, compare=False, repr=False)
+
+        async def for_run(self, ctx):
+            return LiveServers(state=self.state)
+
+        async def for_run_step(self, ctx):
+            await self._sync()
+            return self
+
+        async def __aenter__(self):
+            # The run's exit stack only exits what entered successfully, so a
+            # cancelled first sync must close the servers it already opened.
+            try:
+                await self._sync()
+            except BaseException:
+                await self.__aexit__(*sys.exc_info())
+                raise
+            return self
+
+        async def __aexit__(self, *args: Any) -> bool | None:
+            entered, self._entered, self.toolsets = self._entered, {}, []
+            # A stack, so one exit that is cancelled or fails still runs the rest.
+            async with AsyncExitStack() as stack:
+                for _, wrapper in entered.values():
+                    stack.push_async_exit(wrapper)
+            return None
+
+        async def _sync(self) -> None:
+            wanted = dict(self.state.enabled)
+            for name, (toolset, wrapper) in list(self._entered.items()):
+                # Identity, not name: undefer() swaps a server's toolset in place.
+                if wanted.get(name) is not toolset:
+                    del self._entered[name]
+                    await wrapper.__aexit__(None, None, None)
+            isolated = _isolated()
+            for name, toolset in wanted.items():
+                if name not in self._entered:
+                    wrapper = isolated(toolset, server=name, state=self.state)
+                    # Never raises for a failed connection: IsolatedServer
+                    # records it and lists no tools instead.
+                    await wrapper.__aenter__()
+                    self._entered[name] = (toolset, wrapper)
+            self.toolsets = [wrapper for _, wrapper in self._entered.values()]
+
+    return LiveServers
+
+
+class MCPState:
+    """Never persisted. Disabled servers have no toolsets, connections, or prompt cost."""
+
+    def __init__(self) -> None:
+        self.enabled: dict[str, Any] = {}
+        # Captured at enable, like the toolset, so a later config edit changes
+        # neither until the server is disabled and enabled again.
+        self.descriptions: dict[str, str] = {}
+        # Enabled servers whose latest connection attempt failed, with a message
+        # safe to show (fixed text and an exception type, nothing the server
+        # sent). Cleared by the next successful connection, retried every turn.
+        self.unavailable: dict[str, str] = {}
+        # Set by the runtime to surface a failure and save its frames.
+        self.on_connect_failure: Callable[[str, BaseException], None] = lambda name, error: None
+
+    def connect_failed(self, name: str, error: BaseException) -> None:
+        self.unavailable[name] = (
+            f"MCP server '{name}' failed to connect ({type(error).__name__}); "
+            "continuing without its tools this turn."
+        )
+        self.on_connect_failure(name, error)
+
+    async def enable(self, name: str, *, interactive: bool = True) -> None:
+        if name in self.enabled:
+            return
+        servers = configured_servers()
+        if name not in servers:
+            raise ValueError(f"Unknown MCP server '{name}'. Use /mcp list.")
+        # Read before building, so a bad entry fails here, not after enabling.
+        description = _config(name, servers[name]).description
+        toolset = build_toolset(name, servers[name], interactive=interactive)
+        # Only OAuth needs an enable-time connection. Entering the MCP toolset
+        # initializes the server and completes native auth without a model call.
+        # Publish it only after successful login AND connection cleanup, retaining
+        # the same OAuth object (and its loaded tokens) for subsequent turns.
+        if getattr(mcp_transport(toolset), "auth", None) is not None:
+            try:
+                async with toolset:
+                    pass
+            except Exception as error:
+                # FastMCP wraps connection failures; callers need the sign-in
+                # verdict itself to tell "run /mcp enable" from a real failure.
+                from pcode.mcp_oauth import SignInRequired
+
+                if (sign_in := find_cause(error, SignInRequired)) is not None:
+                    raise sign_in from error
+                raise
+        self.enabled[name] = toolset
+        if description:
+            self.descriptions[name] = description
+
+    async def forget(self, name: str) -> None:
+        """Drop stored OAuth credentials for a server, and its toolset if enabled."""
+        self.enabled.pop(name, None)
+        self.descriptions.pop(name, None)
+        self.unavailable.pop(name, None)
+        servers = configured_servers()
+        if name not in servers:
+            raise ValueError(f"Unknown MCP server '{name}'. Use /mcp list.")
+        config = _config(name, servers[name])
+        if config.auth != "oauth":
+            raise ValueError(f"MCP server '{name}' does not use OAuth.")
+        from fastmcp.client.auth.oauth import TokenStorageAdapter
+
+        from pcode.mcp_oauth import CredentialStore
+
+        await TokenStorageAdapter(CredentialStore(), config.url.rstrip("/")).clear()
+
+    def disable(self, name: str) -> None:
+        # No config read: disabling must work even if the config was removed or broken.
+        if name not in self.enabled:
+            raise ValueError(f"MCP server '{name}' is not enabled.")
+        del self.enabled[name]
+        self.descriptions.pop(name, None)
+        self.unavailable.pop(name, None)
+
+    def live(self):
+        """The enabled servers as one toolset that follows `enabled` while a run is
+        in progress, each isolated so one that cannot connect costs only its tools."""
+        return _live()(state=self)
+
+    def servers(self) -> dict[str, str | None]:
+        """Enabled server names, sorted, with the descriptions configured for them."""
+        return {name: self.descriptions.get(name) for name in sorted(self.enabled)}
+
+    def undefer(self) -> list[str]:
+        """Send every enabled server's schemas up front, for the rest of the session.
+
+        The recovery for a provider that rejects withheld schemas: the tools stay
+        available (at their full prompt cost) instead of the session failing every
+        request. Returns the servers that were still deferring.
+        """
+        undeferred = []
+        for name, toolset in list(self.enabled.items()):
+            if (direct := _without_deferred_loading(toolset)) is not None:
+                self.enabled[name] = direct
+                undeferred.append(name)
+        return undeferred