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

Preserve MCP sessions after call rejections

Browse source 
Track whether MCP failures occur while connecting or running a tool call. A provider can map a per-call authorization denial to HTTP 403, so status alone is not a session verdict. Keep request-level permission and protocol failures local to the call while retiring connections for connection-level failures.
This commit is contained in:
yoni 2026-09-01 06:50:04 +00:00
parent 37138a2332
commit c7e9bfb580
3 changed files with 152 additions and 37 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -13,16 +13,19 @@ from agents.exceptions import UserError
from mcp.shared.exceptions import McpError
FailureKind = Literal["auth", "rate_limit", "server", "transport", "timeout", "protocol", "unknown"]
FailureKind = Literal[
"auth", "permission", "rate_limit", "server", "transport", "timeout", "protocol", "unknown"
]
_PRIORITY: dict[FailureKind, int] = {
"auth": 0,
"rate_limit": 1,
"server": 2,
"protocol": 3,
"timeout": 4,
"transport": 5,
"unknown": 6,
"permission": 1,
"rate_limit": 2,
"server": 3,
"protocol": 4,
"timeout": 5,
"transport": 6,
"unknown": 7,
}
_HTTP_ERROR_RE = re.compile(r"\bHTTP error\s+(\d{3})\b", re.IGNORECASE)
@ -40,7 +43,7 @@ class FailureInfo:
@property
def retryable(self) -> bool:
return self.kind != "auth"
return self.kind not in {"auth", "permission"}
def _retry_after(value: str | None) -> float | None:
@ -67,8 +70,10 @@ def _from_status(
request_method: str | None = None,
request_path: str | None = None,
) -> FailureInfo:
if status in (401, 403):
if status == 401:
kind: FailureKind = "auth"
elif status == 403:
kind = "permission"
elif status == 429:
kind = "rate_limit"
elif 500 <= status <= 599:

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

@ -29,12 +29,12 @@ lifetime, and ``cleanup()``. Three consequences:
agent loop.
Failure handling follows connection-pool discipline: discard on error, rebuild on
next use. The instant a call errors, the supervisor disposes that session on its own
task and never awaits on it again; the next call lazily builds a fresh one. A
classified transient failure is retried once on the rebuilt session and then, if it
keeps failing, temporarily quarantines the connection so the circuit breaker plus the
next call can recover it. Authentication failures and repeated transient exhaustion
permanently retire a connection.
next use. A failure while connecting or rebuilding describes the session. A
non-2xx response from a tool call describes that request, not the session. Permission
and protocol failures from a call return a failed tool output while the connection
stays usable. Other classified failures are retried on the rebuilt session and then,
if they keep failing, temporarily quarantine the connection. Authentication failures
and repeated transient exhaustion permanently retire a connection.
Security: the connection's :class:`~strix.tools.mcp.config.McpConnectionConfig`
holds a live bearer credential and is kept here in memory only, on the same
@ -52,7 +52,7 @@ import logging
import secrets
import time
import weakref
from typing import TYPE_CHECKING, Any, cast
from typing import TYPE_CHECKING, Any, Literal, cast
from strix.tools.mcp.config import DEFAULT_MAX_CONCURRENT_CALLS
from strix.tools.mcp.failures import FailureInfo, HttpStatusRecorder, classify
@ -72,6 +72,8 @@ if TYPE_CHECKING:
# sessions), and its return value becomes the caller's result.
Job = Callable[[MCPServer], Awaitable[Any]]
_Phase = Literal["connect", "call"]
logger = logging.getLogger(__name__)
@ -118,10 +120,11 @@ class McpConnectionUnavailableError(RuntimeError):
@dataclasses.dataclass
class _Outcome:
"""What running one job resolved to: a value, or the connection being dead."""
"""What running one job resolved to: a value, a call failure, or a dead connection."""
value: Any = None
dead: bool = False
call_failure: FailureInfo | None = None
@dataclasses.dataclass
@ -130,6 +133,7 @@ class _Request:
job: Job
future: asyncio.Future[_Outcome]
phase: _Phase
class SupervisedMcpSession:
@ -342,11 +346,14 @@ class SupervisedMcpSession:
async def list_tools(self) -> list[MCPTool]:
"""List the connection's tools, retrying transient session failures.
Raises :class:`McpConnectionUnavailableError` when the connection is dead.
Raises :class:`McpConnectionUnavailableError` when the connection is dead
and never returns a call failure.
"""
outcome = await self._run_job(lambda server: server.list_tools())
outcome = await self._run_job(lambda server: server.list_tools(), phase="connect")
if outcome.dead:
raise McpConnectionUnavailableError(self._unavailable_message())
if outcome.call_failure is not None:
raise RuntimeError("MCP list_tools returned a call failure")
return cast("list[MCPTool]", outcome.value)
async def dispatch(
@ -360,8 +367,9 @@ class SupervisedMcpSession:
"""Run one tool call with bounded retries for transient session failures.
Returns the tool output on success, or the standard failed-tool output
(``success: False``) with a "connection unavailable" message when the
connection is dead.
(``success: False``) when the provider rejects the call or the connection
is unavailable. A call rejection keeps the connection usable because the
provider rejected the request, not the session.
"""
from strix.tools.mcp.client import dispatch_mcp_call
@ -374,7 +382,11 @@ class SupervisedMcpSession:
result_transform=result_transform,
)
outcome = await self._run_job(job)
outcome = await self._run_job(job, phase="call")
if outcome.call_failure is not None:
from strix.tools.mcp.client import _errored_tool_output
return _errored_tool_output(self._call_rejected_message(outcome.call_failure))
if outcome.dead:
from strix.tools.mcp.client import _errored_tool_output
@ -383,13 +395,13 @@ class SupervisedMcpSession:
# -- job routing ----------------------------------------------------------
async def _run_job(self, job: Job) -> _Outcome:
async def _run_job(self, job: Job, *, phase: _Phase) -> _Outcome:
"""Route one job to the owning task (supervised) or run it inline (adopted)."""
if self._supervised:
return await self._submit(job)
return await self._execute(job)
return await self._submit(job, phase)
return await self._execute(job, phase)
async def _submit(self, job: Job) -> _Outcome:
async def _submit(self, job: Job, phase: _Phase) -> _Outcome:
"""Hand a job to the supervising task and await its result as a value."""
if self._dead or self._closing or self._task is None or self._task.done():
return _Outcome(dead=True)
@ -399,7 +411,7 @@ class SupervisedMcpSession:
if self._queue is None:
self._pending.discard(future)
return _Outcome(dead=True)
await self._queue.put(_Request(job=job, future=future))
await self._queue.put(_Request(job=job, future=future, phase=phase))
# The task may have ended between the guard above and the put; ``_fail_pending``
# would then never see this future, so resolve it here.
if self._task.done() and not future.done():
@ -458,7 +470,7 @@ class SupervisedMcpSession:
failure = self._recorder.take() if self._recorder is not None else None
failure = failure or FailureInfo("transport", reason="session cancelled")
self._last_failure = failure
if failure.kind == "auth":
if failure.kind in {"auth", "permission"}:
self._mark_dead(failure, attempt=1)
return
await self._quarantine(failure, attempt=1)
@ -467,7 +479,7 @@ class SupervisedMcpSession:
continue
if request is None: # shutdown sentinel
return
outcome = await self._execute(request.job)
outcome = await self._execute(request.job, request.phase)
if not request.future.done():
request.future.set_result(outcome)
self._pending.discard(request.future)
@ -476,7 +488,7 @@ class SupervisedMcpSession:
# -- run one job with bounded classified retries --------------------------
async def _execute(self, job: Job) -> _Outcome: # noqa: PLR0912
async def _execute(self, job: Job, phase: _Phase) -> _Outcome: # noqa: PLR0912
"""Run one job on a healthy session, disposing it the instant it errors.
Discard-on-error, rebuild-on-next-use is the whole discipline here, and it
@ -499,9 +511,11 @@ class SupervisedMcpSession:
The rebuild itself happens lazily at the top of the loop: once a failure has
set ``_server`` to None, the next iteration builds a fresh session (guarded by
:meth:`_reconnect`) and retries the call on it. A genuine shutdown
(``_closing``) and a real external cancellation still propagate; only the
transport's teardown cancellation is contained.
:meth:`_reconnect`) and retries the operation on it. A permission or protocol
failure from a call returns immediately after disposal because it describes
that request, not the session. A genuine shutdown (``_closing``) and a real
external cancellation still propagate; only the transport's teardown
cancellation is contained.
"""
if self._dead:
return _Outcome(dead=True)
@ -538,7 +552,7 @@ class SupervisedMcpSession:
failure = reconnect_failure or FailureInfo(
"transport", reason="reconnect failed"
)
outcome = await self._handle_failure(failure, attempt)
outcome = await self._handle_failure(failure, attempt, phase="connect")
if outcome is not None:
return outcome
continue
@ -576,16 +590,30 @@ class SupervisedMcpSession:
# Session is disposed and _server is None; _handle_failure may sleep for
# backoff safely, and the next loop iteration rebuilds and retries.
outcome = await self._handle_failure(failure, attempt)
outcome = await self._handle_failure(failure, attempt, phase=phase)
if outcome is not None:
return outcome
return _Outcome(dead=True)
async def _handle_failure(self, failure: FailureInfo, attempt: int) -> _Outcome | None:
async def _handle_failure(
self, failure: FailureInfo, attempt: int, *, phase: _Phase
) -> _Outcome | None:
self._last_failure = failure
if failure.kind == "auth":
self._mark_dead(failure, attempt=attempt)
return _Outcome(dead=True)
if failure.kind == "permission":
if phase == "call":
return _Outcome(call_failure=failure)
self._mark_dead(failure, attempt=attempt)
return _Outcome(dead=True)
if (
phase == "call"
and failure.kind == "protocol"
and failure.status is not None
and 400 <= failure.status <= 499
):
return _Outcome(call_failure=failure)
if attempt == _MAX_ATTEMPTS:
await self._quarantine(failure, attempt=attempt)
return _Outcome(dead=True)
@ -691,6 +719,24 @@ class SupervisedMcpSession:
# -- helpers --------------------------------------------------------------
def _call_rejected_message(self, failure: FailureInfo) -> str:
if failure.kind == "permission":
return (
f"MCP connection {self._name!r} rejected this call (status={failure.status}): "
"the provider denied this specific request, not the connection. The connection "
"is still available. Check the arguments — resource and project identifiers, "
"and required fields — and whether the configured credential is allowed to read "
"that resource, then retry."
)
if failure.kind == "protocol":
return (
f"MCP connection {self._name!r} rejected this call as invalid "
f"(status={failure.status}): the request itself was malformed, not the "
"connection. The connection is still available. Check the tool's required "
"arguments and value formats with describe_mcp, then retry."
)
raise AssertionError(f"Unexpected call failure kind: {failure.kind}")
async def _safe_cleanup(self) -> None:
"""Dispose the live session on this task, completing teardown even under a
firing cancel scope.

66
tests/test_mcp_resilience.py
View file

@ -47,13 +47,14 @@ def _http_error(status: int, *, retry_after: str | None = None) -> httpx.HTTPSta
("exc", "kind"),
[
(_http_error(401), "auth"),
(_http_error(403), "permission"),
(_http_error(429), "rate_limit"),
(_http_error(503), "server"),
(_http_error(404), "protocol"),
(httpx.ReadTimeout("timed out"), "timeout"),
(httpx.ConnectError("disconnected"), "transport"),
(McpError(ErrorData(code=-1, message="bad response")), "protocol"),
(UserError("Failed to call tool: HTTP error 403"), "auth"),
(UserError("Failed to call tool: HTTP error 403"), "permission"),
],
)
def test_classifies_failures(exc: BaseException, kind: str) -> None:
@ -71,6 +72,14 @@ def test_classifies_nested_exception_groups_by_specificity() -> None:
assert info.retryable is False
def test_classifies_permission_before_rate_limit() -> None:
error = ExceptionGroup("outer", [_http_error(429), _http_error(403)])
info = classify(error)
assert info.kind == "permission"
assert info.status == 403
assert info.retryable is False
@pytest.mark.parametrize("control_flow", [SystemExit, KeyboardInterrupt])
@pytest.mark.asyncio
async def test_control_flow_exceptions_propagate(
@ -150,6 +159,16 @@ def _sequence_server(name: str, error: BaseException | None = None) -> Any:
return server
def _list_tools_error_server(name: str, error: BaseException) -> Any:
server = FakeMCPServer(name, [_mcp_tool("read")])
async def list_tools(*_args: Any, **_kwargs: Any) -> Any:
raise error
server.list_tools = list_tools
return server
@pytest.mark.asyncio
async def test_rate_limit_retries_and_succeeds(
monkeypatch: pytest.MonkeyPatch,
@ -252,6 +271,51 @@ async def test_auth_failure_dies_without_retry(monkeypatch: pytest.MonkeyPatch)
await session.aclose()
@pytest.mark.parametrize(
("status", "name"),
[(403, "permission-call"), (400, "protocol-call")],
)
@pytest.mark.asyncio
async def test_call_http_rejection_preserves_session(
monkeypatch: pytest.MonkeyPatch,
status: int,
name: str,
) -> None:
first = _sequence_server(name, _http_error(status))
second = _sequence_server(name)
builds = iter([first, second])
monkeypatch.setattr(mcp_client, "_build_server", lambda _config: _built_server(next(builds)))
session = mcp_session.SupervisedMcpSession(_config(name))
assert await session.start()
result = await session.dispatch("read", {}, label=f"{name}_read")
assert result["success"] is False
assert "not the connection" in result["content"]
assert session.is_dead is False
assert session.is_unavailable is False
assert session._quarantine_count == 0
result = await session.dispatch("read", {}, label=f"{name}_read")
assert result == {"type": "text", "text": "routed:read"}
assert session.is_dead is False
await session.aclose()
@pytest.mark.asyncio
async def test_call_http_403_during_list_tools_dies() -> None:
server = _list_tools_error_server("connect-403", _http_error(403))
session = mcp_session.SupervisedMcpSession.adopt(
server,
name="connect-403",
config=_config("connect-403"),
)
with pytest.raises(mcp_session.McpConnectionUnavailableError):
await session.list_tools()
assert session.is_dead is True
await session.aclose()
@pytest.mark.asyncio
async def test_cancelled_call_uses_recorded_status(
monkeypatch: pytest.MonkeyPatch,