feat: Support MCP Authentication

adapter-contract/src/nemo_fabric_adapter_contract/models.py

@@ -218,6 +218,8 @@ class AgentMcpServerConfig(AgentContractBlock):
     url: str
     args: list[str] = _empty_list()
     env: dict[str, str] = _empty_dict()
+    authentication: dict[str, JsonValue] | None = _optional()
+    custom_headers: dict[str, str] = _empty_dict()
     allowed_tools: list[str] | None = _optional()
     blocked_tools: list[str] = _empty_list()
 

adapters/codex/src/nemo_fabric_adapters/codex/adapter.py

@@ -12,6 +12,7 @@
 import math
 import os
 import subprocess
+from collections.abc import Mapping
 from dataclasses import asdict, dataclass, is_dataclass
 from enum import Enum
 from pathlib import Path
@@ -26,14 +27,21 @@
     TransportClosedError,
     is_retryable_error,
 )
-from openai_codex.generated.v2_all import SkillsExtraRootsSetResponse
+from openai_codex.generated.v2_all import (
+    ListMcpServerStatusResponse,
+    McpAuthStatus,
+    McpServerOauthLoginCompletedNotification,
+    McpServerOauthLoginResponse,
+    SkillsExtraRootsSetResponse,
+)
 from openai_codex.types import Personality, ReasoningEffort, TurnStatus
 
 import nemo_fabric_adapters.common.relay_gateway as relay_gateway
 import nemo_fabric_adapters.common.relay_hooks as relay_hooks
 import nemo_fabric_adapters.common.relay_artifacts as relay_artifacts
 import nemo_fabric_adapters.common.utils as common_utils
 from nemo_fabric_adapters.common import lifecycle
+from nemo_fabric_adapters.common import mcp_auth
 
 
 DEFAULT_TIMEOUT_SECONDS = 1800.0
@@ -52,6 +60,7 @@
     "CODEX_HOME",
     "CODEX_SQLITE_HOME",
     "COMSPEC",
+    "DBUS_SESSION_BUS_ADDRESS",
     "HOME",
     "HTTP_PROXY",
     "HTTPS_PROXY",
@@ -74,6 +83,7 @@
     "XDG_CACHE_HOME",
     "XDG_CONFIG_HOME",
     "XDG_DATA_HOME",
+    "XDG_RUNTIME_DIR",
     "http_proxy",
     "https_proxy",
     "no_proxy",
@@ -198,9 +208,77 @@ def _native_mcp_servers(payload: dict[str, Any]) -> dict[str, dict[str, Any]]:
                 "codex_invalid_configuration",
                 f"unsupported Codex MCP transport: {transport}",
             )
+        if headers := server.get("custom_headers"):
+            try:
+                normalized_headers = mcp_auth.normalize_custom_headers(name, headers)
+            except mcp_auth.McpAuthConfigError as error:
+                raise AdapterConfigError(
+                    "codex_invalid_configuration", str(error)
+                ) from error
+            result[name]["http_headers"] = normalized_headers
+        if authentication := server.get("authentication"):
+            oauth = _mcp_oauth_config(name, authentication)
+            if oauth.client_secret_env:
+                raise AdapterConfigError(
+                    "codex_invalid_configuration",
+                    f"MCP server {name} authentication.client_secret_env is not supported by Codex",
+                )
+            if oauth.client_id:
+                raise AdapterConfigError(
+                    "codex_invalid_configuration",
+                    f"MCP server {name} authentication.client_id is not supported by Codex",
+                )
+            if oauth.client_name:
+                raise AdapterConfigError(
+                    "codex_invalid_configuration",
+                    f"MCP server {name} authentication.client_name is not supported by Codex",
+                )
+            if oauth.token_endpoint_auth_method:
+                raise AdapterConfigError(
+                    "codex_invalid_configuration",
+                    f"MCP server {name} authentication.token_endpoint_auth_method is not supported by Codex",
+                )
+            result[name]["auth"] = "oauth"
+            if oauth.scopes:
+                result[name]["scopes"] = list(oauth.scopes)
     return result
 
 
+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 Codex"
+            )
+        return mcp_auth.parse_oauth2_config(name, value)
+    except mcp_auth.McpAuthConfigError as error:
+        raise AdapterConfigError("codex_invalid_configuration", str(error)) from error
+
+
+def _mcp_oauth_callback_url(payload: dict[str, Any]) -> str | None:
+    servers = _mapping(
+        _native_capabilities(payload).get("mcp_servers"), name="native MCP servers"
+    )
+    values = {
+        oauth.redirect_uri
+        for name, raw in servers.items()
+        if (
+            (
+                authentication := _mapping(raw, name=f"MCP server {name}").get(
+                    "authentication"
+                )
+            )
+            and (oauth := _mcp_oauth_config(name, authentication)).redirect_uri
+        )
+    }
+    if len(values) > 1:
+        raise AdapterConfigError(
+            "codex_invalid_configuration",
+            "Codex supports only one MCP OAuth callback URL per adapter process",
+        )
+    return next(iter(values)) if values else None
+
+
 def _native_skill_paths(payload: dict[str, Any]) -> list[Path]:
     values = _native_capabilities(payload).get("skill_paths", [])
     if not isinstance(values, list) or any(
@@ -260,6 +338,151 @@ async def _register_skill_roots(codex: AsyncCodex, skill_paths: list[Path]) -> N
     )
 
 
+def _mcp_oauth_servers(
+    payload: dict[str, Any],
+) -> dict[str, mcp_auth.McpOAuth2Config]:
+    servers = _mapping(
+        _native_capabilities(payload).get("mcp_servers"),
+        name="native MCP servers",
+    )
+    result: dict[str, mcp_auth.McpOAuth2Config] = {}
+    for name, raw in servers.items():
+        server = _mapping(raw, name=f"MCP server {name}")
+        authentication = server.get("authentication")
+        if not authentication:
+            continue
+        result[name] = _mcp_oauth_config(name, authentication)
+    return result
+
+
+def _codex_protocol_client(codex: AsyncCodex) -> Any:
+    client = getattr(codex, "_client", None)
+    if not callable(getattr(client, "request", None)) or not callable(
+        getattr(client, "next_notification", None)
+    ):
+        raise AdapterConfigError(
+            "codex_invalid_configuration",
+            "Codex SDK does not expose the required MCP OAuth requests",
+        )
+    return client
+
+
+async def _mcp_auth_statuses(
+    client: Any, *, thread_id: str
+) -> dict[str, McpAuthStatus]:
+    statuses: dict[str, McpAuthStatus] = {}
+    cursor: str | None = None
+    while True:
+        params: dict[str, Any] = {
+            "detail": "toolsAndAuthOnly",
+            "threadId": thread_id,
+        }
+        if cursor is not None:
+            params["cursor"] = cursor
+        response = await client.request(
+            "mcpServerStatus/list",
+            params,
+            response_model=ListMcpServerStatusResponse,
+        )
+        statuses.update({server.name: server.auth_status for server in response.data})
+        cursor = response.next_cursor
+        if cursor is None:
+            return statuses
+
+
+async def _login_mcp_server(
+    client: Any,
+    *,
+    name: str,
+    scopes: list[str] | None,
+    thread_id: str,
+    timeout: float,
+) -> None:
+    params: dict[str, Any] = {
+        "name": name,
+        "threadId": thread_id,
+        "timeoutSecs": timeout,
+    }
+    if scopes:
+        params["scopes"] = scopes
+    response = await client.request(
+        "mcpServer/oauth/login",
+        params,
+        response_model=McpServerOauthLoginResponse,
+    )
+    opened = await mcp_auth.open_authorization_url(response.authorization_url)
+    if not opened:
+        raise AdapterConfigError(
+            "codex_mcp_authentication_failed",
+            f"Codex could not open a browser to authenticate MCP server {name!r}",
+        )
+
+    try:
+        async with asyncio.timeout(timeout):
+            while True:
+                notification = await client.next_notification()
+                completed = notification.payload
+                if (
+                    notification.method == "mcpServer/oauthLogin/completed"
+                    and isinstance(completed, McpServerOauthLoginCompletedNotification)
+                    and completed.name == name
+                    and completed.thread_id in {None, thread_id}
+                ):
+                    if completed.success:
+                        return
+                    raise AdapterConfigError(
+                        "codex_mcp_authentication_failed",
+                        f"Codex MCP OAuth login failed for server {name!r}",
+                    )
+    except TimeoutError as error:
+        raise AdapterConfigError(
+            "codex_mcp_authentication_failed",
+            f"Codex MCP OAuth login timed out for server {name!r}",
+        ) from error
+
+
+async def _authenticate_mcp_servers(
+    codex: AsyncCodex,
+    thread: Any,
+    payload: dict[str, Any],
+    invocation_timeout_seconds: float,
+) -> None:
+    oauth_servers = _mcp_oauth_servers(payload)
+    if not oauth_servers:
+        return
+
+    client = _codex_protocol_client(codex)
+    thread_id = str(thread.id)
+    try:
+        statuses = await _mcp_auth_statuses(client, thread_id=thread_id)
+        for name, oauth in oauth_servers.items():
+            status = statuses.get(name)
+            if status in {McpAuthStatus.o_auth, McpAuthStatus.bearer_token}:
+                continue
+            if status != McpAuthStatus.not_logged_in:
+                raise AdapterConfigError(
+                    "codex_mcp_authentication_failed",
+                    f"Codex MCP server {name!r} does not support the configured OAuth login",
+                )
+            await _login_mcp_server(
+                client,
+                name=name,
+                scopes=list(oauth.scopes) or None,
+                thread_id=thread_id,
+                timeout=min(
+                    invocation_timeout_seconds,
+                    oauth.authorization_timeout_seconds,
+                ),
+            )
+    except CodexAdapterError:
+        raise
+    except (CodexError, RuntimeError, OSError) as error:
+        raise AdapterConfigError(
+            "codex_mcp_authentication_failed",
+            "Codex MCP OAuth login could not be completed",
+        ) from error
+
+
 def resolve_cwd(payload: dict[str, Any]) -> Path:
     environment = _mapping(
         common_utils.environment_payload(payload), name="runtime environment"
@@ -620,6 +843,8 @@ def thread_config(
     mcp_servers = _native_mcp_servers(payload)
     if mcp_servers:
         config["mcp_servers"] = mcp_servers
+    if callback_url := _mcp_oauth_callback_url(payload):
+        config["mcp_oauth_callback_url"] = callback_url
     overrides = _mapping(
         _settings(payload).get("config_overrides"),
         name="harness.settings.config_overrides",
@@ -993,6 +1218,7 @@ def __init__(self) -> None:
         self._thread: Any = None
         self._relay: CodexRelaySettings | None = None
         self._gateway_process: subprocess.Popen[Any] | None = None
+        self._mcp_authentication_checked = False
         self._unusable = False
 
     async def start(self, payload: dict[str, Any]) -> None:
@@ -1071,9 +1297,18 @@ async def invoke(self, invocation: dict[str, Any]) -> dict[str, Any]:
 
         try:
             request_prompt(payload)
-            timeout_seconds(payload)
+            invocation_timeout_seconds = timeout_seconds(payload)
             _reasoning_effort(payload)
             _output_schema(payload)
+            if not self._mcp_authentication_checked:
+                await _authenticate_mcp_servers(
+                    self._client,
+                    self._thread,
+                    self._start_payload,
+                    invocation_timeout_seconds,
+                )
+                self._mcp_authentication_checked = True
+
             relay = self._relay
             atif_before = (
                 relay_artifacts.snapshot_atif_files(relay.plugin_config)
@@ -1123,6 +1358,7 @@ async def stop(self) -> None:
         self._start_payload = None
         self._thread = None
         self._fabric_runtime_id = None
+        self._mcp_authentication_checked = False
         self._unusable = True
 
         close_error: BaseException | None = None

adapters/common/pyproject.toml

@@ -25,6 +25,13 @@ license-files = ["LICENSE"]
 readme = "pypi.md"
 requires-python = ">=3.11"
 
+[project.optional-dependencies]
+mcp-oauth = [
+  "httpx>=0.27,<1",
+  # Hermes Agent 0.19 pins this version in its MCP extra.
+  "mcp>=1.26,<1.29",
+]
+
 [project.urls]
 Repository = "https://github.com/NVIDIA/NeMo-Fabric"
 Homepage = "https://github.com/NVIDIA/NeMo-Fabric"