scriptoriumadmin/strix
1
0
Fork 0
mirror of https://github.com/usestrix/strix.git synced 2026-11-14 09:01:25 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 1

Fix/str 2014 mcp refresh bearer on 401 (#1717)

Browse source 
* fix(mcp): refresh an expired bearer once on 401 instead of dropping the connection

* fix(mcp): clear recorded HTTP failures after successful responses

* fix(mcp): only clear a recorded auth failure on a later success

* fix(mcp): preserve eager refresh callbacks and share failed refresh attempts
This commit is contained in:
alex s 2026-11-09 17:32:30 -05:00 • committed by GitHub
parent 215482eb22
commit 18fa4c8449
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 351 additions and 11 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

56
strix/tools/mcp/auth.py Normal file
View file

@ -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

29
strix/tools/mcp/client.py
View file

@ -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

10
strix/tools/mcp/failures.py
View file

@ -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(

15
strix/tools/mcp/registry.py
View file

@ -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)

20
strix/tools/mcp/session.py
View file

@ -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:

223
tests/test_mcp_auth.py Normal file
View file

@ -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()

9
tests/test_mcp_resilience.py
View file

@ -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()