feat: Support MCP Authentication

adapters/claude/src/nemo_fabric_adapters/claude/adapter.py

@@ -12,6 +12,7 @@
 import os
 import shutil
 import subprocess
+from collections.abc import Mapping
 from dataclasses import asdict
 from dataclasses import dataclass
 from dataclasses import is_dataclass
@@ -31,6 +32,7 @@
 from claude_agent_sdk import HookMatcher
 from claude_agent_sdk._errors import MessageParseError
 from nemo_fabric_adapters.common import lifecycle
+from nemo_fabric_adapters.common import mcp_auth
 from nemo_fabric_adapters.common import relay_artifacts
 from nemo_fabric_adapters.common import relay_gateway
 from nemo_fabric_adapters.common import relay_hooks
@@ -59,8 +61,11 @@
     "ANTHROPIC_SERVICE_ACCOUNT_ID",
     "ANTHROPIC_WORKSPACE_ID",
     "APPDATA",
+    "BROWSER",
     "CLAUDE_CONFIG_DIR",
     "COMSPEC",
+    "DBUS_SESSION_BUS_ADDRESS",
+    "DISPLAY",
     "HOME",
     "HTTP_PROXY",
     "HTTPS_PROXY",
@@ -80,9 +85,12 @@
     "TMPDIR",
     "USER",
     "USERPROFILE",
+    "WAYLAND_DISPLAY",
+    "XAUTHORITY",
     "XDG_CACHE_HOME",
     "XDG_CONFIG_HOME",
     "XDG_DATA_HOME",
+    "XDG_RUNTIME_DIR",
     "http_proxy",
     "https_proxy",
     "no_proxy",
@@ -291,16 +299,9 @@ def _model_environment(
 
 
 def _mcp_servers(payload: dict[str, Any]) -> dict[str, Any]:
-    native = (
-        _mapping(common_utils.capability_plan(payload), name="capability_plan").get(
-            "native"
-        )
-        or {}
-    )
-    servers = _mapping(native, name="capability_plan.native").get("mcp_servers") or {}
+    servers = _native_mcp_server_specs(payload)
     result: dict[str, Any] = {}
-    for name, raw in sorted(_mapping(servers, name="native MCP servers").items()):
-        server = _mapping(raw, name=f"MCP server {name}")
+    for name, server in sorted(servers.items()):
         transport = server.get("transport")
         url = server.get("url")
         if not isinstance(url, str) or not url:
@@ -324,9 +325,163 @@ def _mcp_servers(payload: dict[str, Any]) -> dict[str, Any]:
                 "claude_invalid_configuration",
                 f"unsupported MCP transport: {transport}",
             )
+        if headers := server.get("custom_headers"):
+            try:
+                result[name]["headers"] = mcp_auth.normalize_custom_headers(
+                    name, headers
+                )
+            except mcp_auth.McpAuthConfigError as error:
+                raise AdapterConfigError(
+                    "claude_invalid_configuration", str(error)
+                ) from error
+        if authentication := server.get("authentication"):
+            if transport == "stdio":
+                raise AdapterConfigError(
+                    "claude_invalid_configuration",
+                    f"MCP server {name} authentication is not supported for stdio; "
+                    "provide credentials through env",
+                )
+            _mcp_oauth_config(name, authentication)
     return result
 
 
+def _native_mcp_server_specs(payload: dict[str, Any]) -> dict[str, dict[str, Any]]:
+    native = (
+        _mapping(common_utils.capability_plan(payload), name="capability_plan").get(
+            "native"
+        )
+        or {}
+    )
+    servers = _mapping(native, name="capability_plan.native").get("mcp_servers") or {}
+    return {
+        name: _mapping(raw, name=f"MCP server {name}")
+        for name, raw in _mapping(servers, name="native MCP servers").items()
+    }
+
+
+def _mcp_oauth_config(name: str, value: Any) -> mcp_auth.McpOAuth2Config:
+    try:
+        if isinstance(value, Mapping) and value.get("type") == "service_account":
+            raise mcp_auth.McpAuthConfigError(
+                f"MCP server {name!r} service_account authentication is not supported by Claude"
+            )
+        return mcp_auth.parse_oauth2_config(name, value)
+    except mcp_auth.McpAuthConfigError as error:
+        raise AdapterConfigError("claude_invalid_configuration", str(error)) from error
+
+
+def _authenticated_mcp_servers(payload: dict[str, Any]) -> dict[str, dict[str, Any]]:
+    servers = _native_mcp_server_specs(payload)
+    mapped = _mcp_servers(payload)
+    return {
+        name: mapped[name]
+        for name, server in sorted(servers.items())
+        if server.get("authentication")
+    }
+
+
+def _mcp_authentication(payload: dict[str, Any], name: str) -> mcp_auth.McpOAuth2Config:
+    server = _native_mcp_server_specs(payload)[name]
+    return _mcp_oauth_config(name, server.get("authentication"))
+
+
+async def _self_authenticate_http_mcp_server(
+    server_name: str,
+    server_url: str,
+    config: mcp_auth.McpOAuth2Config,
+    *,
+    timeout: float,
+) -> str:
+    """Run an OAuth authorization_code flow in-process and return the access token.
+
+    Uses mcp_auth.create_mcp_oauth_provider and an httpx request to trigger the
+    401→OAuth dance. On success, the token is stored in the provider context and
+    returned as a string. Raises ClaudeAdapterError on failure; re-raises
+    TimeoutError so the caller can map it to claude_mcp_authentication_timed_out.
+    """
+    import httpx
+
+    try:
+        provider = mcp_auth.create_mcp_oauth_provider(
+            server_name,
+            server_url,
+            config,
+            client_name=config.client_name or "nemo-fabric",
+        )
+    except mcp_auth.McpAuthConfigError as error:
+        raise ClaudeAdapterError(
+            "claude_mcp_authentication_failed",
+            f"MCP server {server_name!r} OAuth configuration is invalid: {error}",
+            metadata={"server": server_name},
+        ) from error
+
+    try:
+        async with asyncio.timeout(timeout):
+            async with httpx.AsyncClient(auth=provider) as http_client:
+                await http_client.get(server_url)
+    except TimeoutError:
+        raise
+    except mcp_auth.McpAuthConfigError as error:
+        raise ClaudeAdapterError(
+            "claude_mcp_authentication_failed",
+            f"Claude Code could not authenticate MCP server {server_name!r}: {error}",
+            metadata={"server": server_name},
+        ) from error
+    except asyncio.CancelledError:
+        raise
+    except Exception as error:
+        raise ClaudeAdapterError(
+            "claude_mcp_authentication_failed",
+            f"Claude Code could not authenticate MCP server {server_name!r}",
+            metadata={"server": server_name},
+        ) from error
+
+    context = getattr(provider, "context", None)
+    tokens = getattr(context, "current_tokens", None) if context is not None else None
+    access_token: str | None = getattr(tokens, "access_token", None) if tokens is not None else None
+    if not access_token:
+        raise ClaudeAdapterError(
+            "claude_mcp_authentication_failed",
+            f"MCP server {server_name!r} OAuth flow did not return an access token",
+            metadata={"server": server_name},
+        )
+    return access_token
+
+
+def _update_staged_mcp_server_header(
+    config_path: Path,
+    server_name: str,
+    token: str,
+) -> None:
+    """Inject an Authorization: Bearer header into the staged MCP config file.
+
+    Merges the header into any existing headers dict without clobbering custom_headers.
+    Writes atomically via a .tmp sibling so Claude never sees a partial file.
+    The token is written as a literal value — it is a short-lived runtime credential,
+    not a static secret, and the file is already owner-only (0o600).
+    """
+    tmp_path = config_path.parent / (config_path.name + ".tmp")
+    try:
+        config = json.loads(config_path.read_text(encoding="utf-8"))
+        server_entry = config.setdefault("mcpServers", {}).setdefault(server_name, {})
+        server_entry.setdefault("headers", {})["Authorization"] = f"Bearer {token}"
+        descriptor = os.open(tmp_path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600)
+        try:
+            with os.fdopen(descriptor, "w", encoding="utf-8") as stream:
+                json.dump(config, stream, indent=2, sort_keys=True)
+                stream.write("\n")
+        except BaseException:
+            tmp_path.unlink(missing_ok=True)
+            raise
+        os.replace(tmp_path, config_path)
+    except OSError as error:
+        raise ClaudeAdapterError(
+            "claude_mcp_configuration_failed",
+            f"MCP server {server_name!r} authorization header could not be staged",
+            metadata={"server": server_name},
+        ) from error
+
+
 def _stage_mcp_config(payload: dict[str, Any]) -> ClaudeMcpSettings | None:
     # Dictionary-valued ClaudeAgentOptions.mcp_servers are JSON-serialized by
     # claude-agent-sdk into the literal `--mcp-config` command-line argument,
@@ -362,9 +517,11 @@ def _stage_mcp_config(payload: dict[str, Any]) -> ClaudeMcpSettings | None:
                     "claude_invalid_configuration",
                     f"MCP server {server_name} env values must be strings",
                 )
-            projection_key = sha256(
-                f"{fabric_runtime_id}\0{server_name}\0{variable_name}".encode()
-            ).hexdigest().upper()
+            projection_key = (
+                sha256(f"{fabric_runtime_id}\0{server_name}\0{variable_name}".encode())
+                .hexdigest()
+                .upper()
+            )
             projected_name = f"NEMO_FABRIC_CLAUDE_MCP_{projection_key}"
             projected_environment[variable_name] = f"${{{projected_name}}}"
             environment[projected_name] = value
@@ -403,6 +560,7 @@ def _cleanup_mcp_config(config_path: Path | None) -> None:
         return
     try:
         config_path.unlink(missing_ok=True)
+        (config_path.parent / (config_path.name + ".tmp")).unlink(missing_ok=True)
         config_path.parent.rmdir()
     except OSError:
         LOGGER.exception("Claude MCP runtime configuration could not be removed")
@@ -930,6 +1088,7 @@ def __init__(self) -> None:
         self._fabric_runtime_id: str | None = None
         self._claude_session_id: str | None = None
         self._client: ClaudeSDKClient | None = None
+        self._mcp_authentication_checked = False
         self._relay: ClaudeRelaySettings | None = None
         self._gateway_process: subprocess.Popen[Any] | None = None
         self._mcp_config_path: Path | None = None
@@ -947,8 +1106,10 @@ async def start(self, payload: dict[str, Any]) -> None:
             self._relay = relay
             self._gateway_process = _start_relay_gateway(payload, relay)
             options = build_options(payload, relay=relay)
+
             if isinstance(options.mcp_servers, Path):
                 self._mcp_config_path = options.mcp_servers
+
             client = ClaudeSDKClient(options)
             await client.connect()
         except ClaudeAdapterError as error:
@@ -993,8 +1154,20 @@ async def invoke(self, invocation: dict[str, Any]) -> dict[str, Any]:
         try:
             prompt = request_prompt(payload)
             invocation_timeout = timeout_seconds(payload)
+            await self._authenticate_mcp_servers(
+                payload,
+                client,
+                invocation_timeout,
+            )
         except ClaudeAdapterError as error:
             output = adapter_failure(error)
+        except TimeoutError:
+            output = _failure(
+                "claude_mcp_authentication_timed_out",
+                "Claude MCP authentication timed out",
+            )
+        except ClaudeSDKError as error:
+            output = sdk_failure(error)
         else:
             relay = self._relay
             atif_before = (
@@ -1037,6 +1210,72 @@ async def invoke(self, invocation: dict[str, Any]) -> dict[str, Any]:
             output = _relay_output(output, self._relay)
         return output
 
+    async def _authenticate_mcp_servers(
+        self,
+        payload: dict[str, Any],
+        client: ClaudeSDKClient,
+        invocation_timeout: float,
+    ) -> None:
+        if self._mcp_authentication_checked:
+            return
+
+        servers = _authenticated_mcp_servers(payload)
+        if servers and self._mcp_config_path is None:
+            raise ClaudeAdapterError(
+                "claude_mcp_configuration_failed",
+                "MCP staged configuration is unavailable for OAuth authentication",
+            )
+
+        for name, server in servers.items():
+            authentication = _mcp_authentication(payload, name)
+            auth_timeout = min(invocation_timeout, authentication.authorization_timeout_seconds)
+            token = await _self_authenticate_http_mcp_server(
+                name,
+                server["url"],
+                authentication,
+                timeout=auth_timeout,
+            )
+            assert self._mcp_config_path is not None
+            _update_staged_mcp_server_header(self._mcp_config_path, name, token)
+            await client.reconnect_mcp_server(name)
+            status = await self._mcp_server_status(client, name)
+            if status != "connected":
+                raise ClaudeAdapterError(
+                    "claude_mcp_unavailable",
+                    f"Claude MCP server {name!r} is unavailable after OAuth authentication",
+                    metadata={"server": name, "status": status},
+                )
+
+        self._mcp_authentication_checked = True
+
+    async def _mcp_server_status(
+        self,
+        client: ClaudeSDKClient,
+        name: str,
+    ) -> str:
+        loop = asyncio.get_running_loop()
+        deadline = loop.time() + 5.0
+        while True:
+            response = await client.get_mcp_status()
+            match = next(
+                (
+                    server
+                    for server in response.get("mcpServers", [])
+                    if server.get("name") == name
+                ),
+                None,
+            )
+            if match is None:
+                raise ClaudeAdapterError(
+                    "claude_mcp_unavailable",
+                    f"Claude MCP server {name!r} was not loaded",
+                    metadata={"server": name, "status": "missing"},
+                )
+            status = str(match.get("status") or "unknown")
+            if status != "pending" or loop.time() >= deadline:
+                return status
+            await asyncio.sleep(0.25)
+
     async def _run_query(
         self,
         payload: dict[str, Any],
@@ -1105,6 +1344,7 @@ async def stop(self) -> None:
         self._start_payload = None
         self._fabric_runtime_id = None
         self._claude_session_id = None
+        self._mcp_authentication_checked = False
         self._unusable = True
 
         disconnect_error: BaseException | None = None