diff --git a/strix/tools/mcp/auth.py b/strix/tools/mcp/auth.py new file mode 100644 index 000000000..320efad6e --- /dev/null +++ b/strix/tools/mcp/auth.py @@ -0,0 +1,56 @@ +"""In-memory bearer refresh for HTTP MCP connections.""" + +from __future__ import annotations + +import asyncio +import logging +from typing import TYPE_CHECKING + +import httpx + + +if TYPE_CHECKING: + from collections.abc import AsyncGenerator, Awaitable, Callable + + +logger = logging.getLogger(__name__) + + +class RefreshingBearerAuth(httpx.Auth): + def __init__( + self, token: str | None, refresh_token: Callable[[], Awaitable[str | None]] + ) -> None: + self._token = token + self._refresh_token = refresh_token + self._lock = asyncio.Lock() + self._refresh_generation = 0 + + def _authorize(self, request: httpx.Request, token: str | None) -> None: + if token: + request.headers["Authorization"] = f"Bearer {token}" + else: + request.headers.pop("Authorization", None) + + async def async_auth_flow( + self, request: httpx.Request + ) -> AsyncGenerator[httpx.Request, httpx.Response]: + sent_token = self._token + sent_generation = self._refresh_generation + self._authorize(request, sent_token) + response = yield request + if response.status_code != 401: + return + async with self._lock: + if self._token == sent_token and self._refresh_generation == sent_generation: + try: + token = await self._refresh_token() + except Exception as exc: # noqa: BLE001 + logger.warning("MCP bearer refresh failed (%s)", type(exc).__name__) + token = None + self._refresh_generation += 1 + if token and token != sent_token: + self._token = token + retry_token = self._token + if retry_token and retry_token != sent_token: + self._authorize(request, retry_token) + yield request diff --git a/strix/tools/mcp/client.py b/strix/tools/mcp/client.py index b1f3f902c..86500969b 100644 --- a/strix/tools/mcp/client.py +++ b/strix/tools/mcp/client.py @@ -39,10 +39,11 @@ from strix.tools.mcp.session import McpConnectionUnavailableError, SupervisedMcp if TYPE_CHECKING: - from collections.abc import Callable + from collections.abc import Awaitable, Callable, Mapping import httpx + from strix.tools.mcp.auth import RefreshingBearerAuth from strix.tools.mcp.config import McpConnectionConfig from strix.tools.mcp.registry import McpConnectionRequest, McpRegistry @@ -121,7 +122,9 @@ class _QuietMCPServerStdio(MCPServerStdio): return _quiet_stdio_streams(self.params) -def _build_server(config: McpConnectionConfig) -> BuiltMcpServer: +def _build_server( + config: McpConnectionConfig, bearer_auth: RefreshingBearerAuth | None = None +) -> BuiltMcpServer: """Construct (but do not connect) the SDK server for one connection. The returned tuple carries the server and, for HTTP connections, a recorder @@ -161,13 +164,13 @@ def _build_server(config: McpConnectionConfig) -> BuiltMcpServer: timeout: httpx.Timeout | None = None, auth: httpx.Auth | None = None, ) -> httpx.AsyncClient: - client = create_mcp_http_client(headers=headers, timeout=timeout, auth=auth) + client = create_mcp_http_client(headers=headers, timeout=timeout, auth=bearer_auth or auth) client.event_hooks.setdefault("response", []).append(recorder) return client http_params: MCPServerStreamableHttpParams = { "url": cast("str", config.url), - "headers": _auth_headers(config), + "headers": {} if bearer_auth is not None else _auth_headers(config), "timeout": config.http_timeout_seconds, "sse_read_timeout": config.sse_read_timeout_seconds, "httpx_client_factory": httpx_client_factory, @@ -300,9 +303,12 @@ async def connect_mcp_servers( configs: list[McpConnectionConfig], *, max_concurrency: int = 6, + token_refreshers: Mapping[str, Callable[[], Awaitable[str | None]]] | None = None, ) -> list[ConnectedMcpServer]: """Connect MCP configs concurrently under a fixed bound. + ``token_refreshers`` optionally supplies a bearer refresher by connection name. + Each connection becomes a :class:`~strix.tools.mcp.session.SupervisedMcpSession` that owns ``connect()``, the held-open session, and ``cleanup()`` on one dedicated task, so a later background failure in one session is contained to @@ -325,7 +331,10 @@ async def connect_mcp_servers( async def connect_one(config: McpConnectionConfig) -> ConnectedMcpServer | None: async with semaphore: - session = SupervisedMcpSession(config) + session = SupervisedMcpSession( + config, + refresh_token=token_refreshers.get(config.name) if token_refreshers else None, + ) sessions.append(session) try: if not await session.start(): @@ -381,7 +390,14 @@ async def attach_mcp_requests( ends). """ request_by_name = {request.config.name: request for request in requests} - connections = await connect_mcp_servers([request.config for request in requests]) + connections = await connect_mcp_servers( + [request.config for request in requests], + token_refreshers={ + request.config.name: request.refresh_token + for request in requests + if request.refresh_token is not None + }, + ) for connection in connections: request = request_by_name[connection.name] registry.add( @@ -391,5 +407,6 @@ async def attach_mcp_requests( purpose=request.purpose or connection.notes, provider=request.provider, result_transform=request.result_transform, + refresh_token=request.refresh_token, ) return connections diff --git a/strix/tools/mcp/failures.py b/strix/tools/mcp/failures.py index a5b82ca87..3710ce1a2 100644 --- a/strix/tools/mcp/failures.py +++ b/strix/tools/mcp/failures.py @@ -129,12 +129,20 @@ def classify(exc: BaseException) -> FailureInfo: class HttpStatusRecorder: - """Capture the last non-success response from one HTTP connection.""" + """Capture the last non-success response from one HTTP connection. + + A later success clears a recorded 401, so a bearer refreshed and retried + in the auth flow does not leave a stale auth failure behind. + """ def __init__(self) -> None: self._failure: FailureInfo | None = None async def __call__(self, response: httpx.Response) -> None: + if 200 <= response.status_code < 300: + if self._failure is not None and self._failure.kind == "auth": + self._failure = None + return if not 200 <= response.status_code < 300: request = response.request self._failure = _from_status( diff --git a/strix/tools/mcp/registry.py b/strix/tools/mcp/registry.py index 6e87c958b..914afe5a5 100644 --- a/strix/tools/mcp/registry.py +++ b/strix/tools/mcp/registry.py @@ -32,7 +32,7 @@ from strix.tools.mcp.session import McpConnectionUnavailableError, SupervisedMcp if TYPE_CHECKING: - from collections.abc import Callable + from collections.abc import Awaitable, Callable from agents.mcp import MCPServer @@ -96,6 +96,9 @@ class McpConnectionEntry: tool_count: int = 0 result_transform: ResultTransform | None = None provider: str | None = None + refresh_token: Callable[[], Awaitable[str | None]] | None = dataclasses.field( + default=None, repr=False + ) state: McpConnectionState = "configured" _catalog: list[Any] | None = dataclasses.field(default=None, repr=False) _connect_task: asyncio.Task[SupervisedMcpSession] | None = dataclasses.field( @@ -180,7 +183,7 @@ class McpConnectionEntry: f"MCP connection {self.name!r} has no connection configuration." ) self._set_state("connecting") - session = SupervisedMcpSession(config) + session = SupervisedMcpSession(config, refresh_token=self.refresh_token) try: started = await session.start() except BaseException: @@ -298,12 +301,17 @@ class McpConnectionRequest: at the single dispatch point (strix-pro's sanitizer; empty for the command-line path). ``purpose`` is the human label ``list_mcps`` reports as the connection's description; when unset it falls back to ``config.notes``. + ``refresh_token`` optionally fetches a new bearer asynchronously after an HTTP + 401; returning no token or the current token leaves the failure unchanged. """ config: McpConnectionConfig provider: str | None = None result_transform: ResultTransform | None = None purpose: str | None = None + refresh_token: Callable[[], Awaitable[str | None]] | None = dataclasses.field( + default=None, repr=False + ) class McpCallInfo(NamedTuple): @@ -337,6 +345,7 @@ class McpRegistry: purpose=request.purpose or request.config.notes, result_transform=request.result_transform, provider=request.provider, + refresh_token=request.refresh_token, ) entry.set_status_sink(self._status_sink) self._entries[entry.name] = entry @@ -353,6 +362,7 @@ class McpRegistry: tool_count: int = 0, result_transform: ResultTransform | None = None, provider: str | None = None, + refresh_token: Callable[[], Awaitable[str | None]] | None = None, ) -> McpConnectionEntry: """Register one connection under ``name`` (last write wins). @@ -374,6 +384,7 @@ class McpRegistry: tool_count=tool_count, result_transform=result_transform, provider=provider, + refresh_token=refresh_token, state="connected", ) entry.set_status_sink(self._status_sink) diff --git a/strix/tools/mcp/session.py b/strix/tools/mcp/session.py index 9cf854086..faf0c674a 100644 --- a/strix/tools/mcp/session.py +++ b/strix/tools/mcp/session.py @@ -54,6 +54,7 @@ import time import weakref from typing import TYPE_CHECKING, Any, Literal, cast +from strix.tools.mcp.auth import RefreshingBearerAuth from strix.tools.mcp.config import DEFAULT_MAX_CONCURRENT_CALLS from strix.tools.mcp.failures import FailureInfo, HttpStatusRecorder, classify @@ -157,9 +158,19 @@ class SupervisedMcpSession: :attr:`name`, :attr:`server`, :attr:`config`, :attr:`is_dead`. """ - def __init__(self, config: McpConnectionConfig) -> None: + def __init__( + self, + config: McpConnectionConfig, + *, + refresh_token: Callable[[], Awaitable[str | None]] | None = None, + ) -> None: self._name = config.name self._config: McpConnectionConfig | None = config + self._bearer_auth = ( + RefreshingBearerAuth(config.auth.token if config.auth else None, refresh_token) + if refresh_token is not None and config.transport == "http" + else None + ) self._server: MCPServer | None = None self._supervised = True self._task: asyncio.Task[None] | None = None @@ -193,6 +204,7 @@ class SupervisedMcpSession: self = cls.__new__(cls) self._name = name self._config = config + self._bearer_auth = None self._server = server self._supervised = False self._task = None @@ -701,7 +713,11 @@ class SupervisedMcpSession: if self._config is None: raise RuntimeError(f"MCP connection {self._name!r} has no config to connect") - built = _build_server(self._config) + built = ( + _build_server(self._config, self._bearer_auth) + if self._bearer_auth is not None + else _build_server(self._config) + ) server = built.server self._recorder = built.recorder try: diff --git a/tests/test_mcp_auth.py b/tests/test_mcp_auth.py new file mode 100644 index 000000000..e8691c246 --- /dev/null +++ b/tests/test_mcp_auth.py @@ -0,0 +1,223 @@ +"""HTTP bearer refresh regressions without network access.""" + +from __future__ import annotations + +import asyncio +from typing import Any, cast +from unittest.mock import AsyncMock + +import httpx +import pytest +from mcp.types import CallToolResult, TextContent, Tool + +from strix.tools.mcp import BearerAuth, McpConnectionConfig, McpConnectionRequest, McpRegistry +from strix.tools.mcp import client as mcp_client +from strix.tools.mcp import session as mcp_session +from strix.tools.mcp.auth import RefreshingBearerAuth + + +def _config() -> McpConnectionConfig: + return McpConnectionConfig( + name="refresh", + url="https://example.test/mcp", + auth=BearerAuth(token="old"), # noqa: S106 # nosec B106 + ) + + +def _client(auth: RefreshingBearerAuth | None, handler: Any) -> httpx.AsyncClient: + built = mcp_client._build_server(_config(), auth) + params = cast("Any", built.server).params + assert params["headers"] == ({} if auth else {"Authorization": "Bearer old"}) + client = cast("httpx.AsyncClient", params["httpx_client_factory"](headers=params["headers"])) + client._transport = httpx.MockTransport(handler) + return client + + +@pytest.mark.asyncio +async def test_refresh_retries_once_and_persists_across_server_builds() -> None: + refresh = AsyncMock(return_value="new") + session = mcp_session.SupervisedMcpSession(_config(), refresh_token=refresh) + seen: list[str] = [] + + def handler(request: httpx.Request) -> httpx.Response: + seen.append(request.headers["Authorization"]) + return httpx.Response(200 if seen[-1] == "Bearer new" else 401) + + for _ in range(2): + async with _client(session._bearer_auth, handler) as client: + response = await client.post("https://example.test/mcp", content=b"payload") + assert response.status_code == 200 + assert seen == ["Bearer old", "Bearer new", "Bearer new"] + refresh.assert_awaited_once() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("replacement", ["new", None, "old", "raises"]) +async def test_concurrent_unauthorized_requests_share_refresh(replacement: str | None) -> None: + refresh = AsyncMock(return_value=replacement) + if replacement == "raises": + refresh.side_effect = ValueError("refresh unavailable") + auth = RefreshingBearerAuth("old", refresh) + both_sent = asyncio.Event() + old_requests = 0 + + async def handler(request: httpx.Request) -> httpx.Response: + nonlocal old_requests + if request.headers["Authorization"] == "Bearer old": + old_requests += 1 + if old_requests == 5: + both_sent.set() + await both_sent.wait() + return httpx.Response(401) + return httpx.Response(200) + + async with _client(auth, handler) as client: + responses = await asyncio.wait_for( + asyncio.gather(*(client.get("https://example.test/mcp") for _ in range(5))), + timeout=2, + ) + assert [response.status_code for response in responses] == [ + 200 if replacement == "new" else 401 + ] * 5 + assert old_requests == 5 + refresh.assert_awaited_once() + + +@pytest.mark.parametrize("replacement", [None, "", "old", "new", "raises", "disabled"]) +@pytest.mark.asyncio +async def test_unrecovered_401_still_retires_connection( + replacement: str | None, caplog: pytest.LogCaptureFixture +) -> None: + refresh = AsyncMock(return_value=replacement) + if replacement == "raises": + refresh.side_effect = ValueError("private-token-value") + auth = None if replacement == "disabled" else RefreshingBearerAuth("old", refresh) + seen: list[str] = [] + + def handler(request: httpx.Request) -> httpx.Response: + seen.append(request.headers["Authorization"]) + return httpx.Response(401) + + async with _client(auth, handler) as client: + server = AsyncMock() + + async def call_tool(*_args: Any, **_kwargs: Any) -> None: + response = await client.get("https://example.test/mcp") + response.raise_for_status() + + server.call_tool = call_tool + session = mcp_session.SupervisedMcpSession.adopt(server, name="refresh") + result = await session.dispatch("read", {}, label="refresh_read") + assert result["success"] is False + assert session.is_dead + await session.aclose() + assert seen == (["Bearer old", "Bearer new"] if replacement == "new" else ["Bearer old"]) + assert refresh.await_count == (0 if replacement == "disabled" else 1) + assert "private-token-value" not in caplog.text + if replacement == "raises": + assert "ValueError" in caplog.text + + +@pytest.mark.asyncio +async def test_non_401_does_not_refresh() -> None: + refresh = AsyncMock(return_value="new") + auth = RefreshingBearerAuth("old", refresh) + async with _client(auth, lambda _: httpx.Response(403)) as client: + assert (await client.get("https://example.test/mcp")).status_code == 403 + refresh.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_registry_passes_refresher_to_session(monkeypatch: pytest.MonkeyPatch) -> None: + refresh = AsyncMock(return_value="new") + registry = McpRegistry() + entry = registry.register(McpConnectionRequest(_config(), refresh_token=refresh)) + server = AsyncMock() + captured: list[RefreshingBearerAuth] = [] + + def build(_config: McpConnectionConfig, auth: RefreshingBearerAuth) -> Any: + captured.append(auth) + return mcp_client.BuiltMcpServer(server, None) + + monkeypatch.setattr(mcp_client, "_build_server", build) + session = await entry.ensure_connected() + assert entry.refresh_token is refresh + assert captured == [session._bearer_auth] + await session._open() + assert captured == [session._bearer_auth, session._bearer_auth] + await session.aclose() + + +@pytest.mark.asyncio +async def test_refreshed_token_can_refresh_again_after_expiry() -> None: + refresh = AsyncMock(side_effect=["new", "newer"]) + auth = RefreshingBearerAuth("old", refresh) + valid = "new" + + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response(200 if request.headers["Authorization"] == f"Bearer {valid}" else 401) + + async with _client(auth, handler) as client: + assert (await client.get("https://example.test/mcp")).status_code == 200 + valid = "newer" + assert (await client.get("https://example.test/mcp")).status_code == 200 + assert refresh.await_count == 2 + + +@pytest.mark.parametrize("with_refresher", [True, False]) +@pytest.mark.asyncio +async def test_eager_attachment_preserves_bearer_refresh( + monkeypatch: pytest.MonkeyPatch, with_refresher: bool +) -> None: + refresh = AsyncMock(return_value="new") + original_build = mcp_client._build_server + valid = "old" + + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response(200 if request.headers["Authorization"] == f"Bearer {valid}" else 401) + + def build(config: McpConnectionConfig, auth: RefreshingBearerAuth | None = None) -> Any: + built = original_build(config, auth) + params = cast("Any", built.server).params + client = params["httpx_client_factory"](headers=params["headers"]) + client._transport = httpx.MockTransport(handler) + server = AsyncMock() + server.use_structured_content = False + server.list_tools.return_value = [Tool(name="read", inputSchema={"type": "object"})] + server.cleanup.side_effect = client.aclose + + async def call_tool(*_args: Any, **_kwargs: Any) -> CallToolResult: + response = await client.post("https://example.test/mcp") + response.raise_for_status() + return CallToolResult(content=[TextContent(type="text", text="ok")]) + + server.call_tool = call_tool + return mcp_client.BuiltMcpServer(server, built.recorder) + + monkeypatch.setattr(mcp_client, "_build_server", build) + registry = McpRegistry() + try: + connections = await mcp_client.attach_mcp_requests( + [McpConnectionRequest(_config(), refresh_token=refresh if with_refresher else None)], + registry, + ) + assert len(connections) == 1 + entry = registry.get("refresh") + assert entry is not None + assert entry.refresh_token is (refresh if with_refresher else None) + session = await entry.ensure_connected() + assert await session.dispatch("read", {}, label="refresh_read") == { + "type": "text", + "text": "ok", + } + valid = "new" + result = await session.dispatch("read", {}, label="refresh_read") + if with_refresher: + assert result == {"type": "text", "text": "ok"} + assert not session.is_dead + refresh.assert_awaited_once() + else: + assert result["success"] is False + assert session.is_dead + finally: + await registry.close() diff --git a/tests/test_mcp_resilience.py b/tests/test_mcp_resilience.py index 209acf1c5..fc0e0cff2 100644 --- a/tests/test_mcp_resilience.py +++ b/tests/test_mcp_resilience.py @@ -111,6 +111,15 @@ async def test_retry_after_parses_seconds_and_http_date() -> None: assert 0 <= retry_after <= 20 +@pytest.mark.asyncio +async def test_successful_response_clears_previous_auth_failure() -> None: + recorder = HttpStatusRecorder() + await recorder(_http_error(401).response) + request = httpx.Request("POST", "https://provider.example/tools") + await recorder(httpx.Response(200, request=request)) + assert recorder.take() is None + + @pytest.mark.asyncio async def test_recorder_only_keeps_non_sensitive_request_metadata() -> None: recorder = HttpStatusRecorder()