mirror of
https://github.com/turnstonelabs/turnstone.git
synced 2026-08-12 23:12:23 -06:00
fix(mcp): spawn static list_changed refreshes off the receive loop
The static-path notification handler awaited its catalog refresh inline in the SDK's receive loop, but the refresh issues a request on the same session — a request whose response only that (now parked) loop could route. The refresh never completed, and every user's calls on the shared per-node session stalled behind it, unbounded, until the health loop's ping timeout tore the transport down — which was also the only way a pushed catalog change ever landed. Port of the pool-path protocol (#836) onto the static primitives: - Refreshes are debounce-gated, coalesced per (server, kind), and spawned as tracked tasks; the runner serializes on the per-name connect lock so a refresh, a connect's discovery wiring, and the manual/periodic _refresh_server pass can never publish out of order (the remove -> re-add race is closed by lock identity, the static twin of the pool's entry-identity check). - The coalesce marker is cleared at lock-acquire so a change the in-flight list missed spawns exactly one successor; the finally discard is gated on non-acquisition so it never clobbers that successor's marker. - The debounce stamp survives a failed refresh (throttle over lost window) and every teardown/eviction path now pops it via the paired _drop_static_session_and_stamp, so a reconnected transport's first notification refreshes immediately. - All three static list calls are bounded by _CONNECT_TIMEOUT and discard their result if the state entry was replaced mid-flight; the resource pair rides one gather (mirrors the pool sibling). - Failure logging is (Exception, BaseExceptionGroup) type-name-only: an escaping group reaches _spawn_background's exc_info log, which serializes the chained httpx request carrying the configured bearer for auth_type=static servers; the recorded operator error string is type-name-only for the same reason. Non-list-changed notifications no longer clear the server's error pill (that pop was accidental — only a completed refresh proves anything). Includes a live end-to-end repro (FastMCP subprocess pushing tools/list_changed through a real receive loop): pre-fix the triggering call itself deadlocks (verified against main), post-fix it completes with the catalog landing on the original session, no teardown. Closes #839
This commit is contained in:
@@ -160,6 +160,20 @@ Earlier stable lines (`stable/1.6`, `stable/1.5`) are frozen.
|
||||
|
||||
### Fixed
|
||||
|
||||
- **Static MCP servers: a pushed catalog change no longer wedges the shared
|
||||
session (#839).** The static-path `*/list_changed` handler awaited its
|
||||
catalog refresh inline in the SDK's receive loop, but the refresh's own
|
||||
request can only be answered by that (now parked) loop — the refresh never
|
||||
completed, and every user's in-flight calls on the shared per-node session
|
||||
stalled behind it, unbounded, until the health loop's ping timeout tore the
|
||||
transport down (which was also the only way the changed catalog ever
|
||||
landed). Push refreshes now run as spawned tasks — debounced, coalesced per
|
||||
(server, kind), bounded by the connect timeout, and serialized on the
|
||||
per-server connect lock — and the manual/periodic refresh publishes under
|
||||
that same lock, so a slower publisher can no longer land a staler catalog
|
||||
over a fresher one. Every teardown path now also clears the notification
|
||||
debounce stamp, so a reconnected server's first push refreshes immediately.
|
||||
|
||||
- **OpenAI Responses streaming: truncated and refused responses no longer
|
||||
vanish.** A response that hit `max_output_tokens` terminates the stream
|
||||
with `response.incomplete`, which the stream consumer did not handle —
|
||||
|
||||
@@ -12,6 +12,8 @@ from contextlib import AsyncExitStack, suppress
|
||||
from typing import Any
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import anyio
|
||||
import mcp.types as mcp_types
|
||||
import pytest
|
||||
|
||||
from tests.conftest import _seed_static_state
|
||||
@@ -2877,6 +2879,485 @@ class TestNotificationDebounce:
|
||||
assert now - last_b >= mgr._NOTIFICATION_DEBOUNCE
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Static-path list_changed refresh — spawned runner protocol (#839)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _run_on_loop(loop: asyncio.AbstractEventLoop, coro: Any) -> Any:
|
||||
"""Submit *coro* to *loop*, wait for the result with a 5s timeout.
|
||||
|
||||
Mirrors ``test_mcp_user_pool.py``'s helper of the same name.
|
||||
"""
|
||||
fut = asyncio.run_coroutine_threadsafe(coro, loop)
|
||||
return fut.result(timeout=5)
|
||||
|
||||
|
||||
def _drain_background(mgr: MCPClientManager, loop: asyncio.AbstractEventLoop) -> None:
|
||||
"""Deterministically await the manager's tracked background tasks.
|
||||
|
||||
Replaces fixed sleeps for synchronizing with spawned refreshes:
|
||||
exact, and immune to slow-runner flake.
|
||||
"""
|
||||
|
||||
async def _drain() -> None:
|
||||
tasks = [t for t in list(mgr._background_tasks) if not t.done()]
|
||||
if tasks:
|
||||
await asyncio.gather(*tasks, return_exceptions=True)
|
||||
|
||||
_run_on_loop(loop, _drain())
|
||||
|
||||
|
||||
class TestStaticNotificationRefresh:
|
||||
"""The static ``*/list_changed`` protocol: spawn / serialize / coalesce.
|
||||
|
||||
Static twin of ``test_mcp_user_pool.py``'s notification-runner suite.
|
||||
The SDK awaits message handlers inline in its receive loop, so a
|
||||
handler that awaits a request on the SAME session self-deadlocks the
|
||||
loop (#839) — and static sessions are shared per node, so every
|
||||
user's in-flight calls on that server stall with it. The refresh must
|
||||
be spawned, serialized on the per-name connect lock, coalesced per
|
||||
(server, kind), and keep bearer-carrying exception chains out of the
|
||||
logs on failure.
|
||||
"""
|
||||
|
||||
def test_list_changed_handler_spawns_refresh_off_receive_loop(self, running_loop_mgr) -> None:
|
||||
"""The notification handler must SPAWN the refresh, not await it:
|
||||
an in-handler request on the same session can never receive its
|
||||
response (only the parked receive loop could route it), so
|
||||
push-driven static refreshes structurally never completed."""
|
||||
mgr, loop, _thread = running_loop_mgr
|
||||
refreshed: list[str] = []
|
||||
|
||||
async def _fake_refresh(name: str) -> tuple[list[str], list[str]]:
|
||||
refreshed.append(name)
|
||||
return [], []
|
||||
|
||||
mgr._refresh_server_tools = _fake_refresh # type: ignore[method-assign]
|
||||
handler = mgr._make_static_notification_handler("srv")
|
||||
|
||||
async def _fire() -> bool:
|
||||
mgr._static_connect_lock_for("srv")
|
||||
_seed_static_state(mgr, "srv", session=MagicMock())
|
||||
note = mcp_types.ServerNotification(
|
||||
mcp_types.ToolListChangedNotification(method="notifications/tools/list_changed")
|
||||
)
|
||||
await handler(note)
|
||||
# The handler returned WITHOUT running the refresh inline —
|
||||
# it must complete even though the refresh hasn't run yet.
|
||||
return len(refreshed) == 0
|
||||
|
||||
returned_before_refresh = _run_on_loop(loop, _fire())
|
||||
assert returned_before_refresh, "handler must not await the refresh inline"
|
||||
_drain_background(mgr, loop)
|
||||
assert refreshed == ["srv"]
|
||||
# Debounce stamp consumed at schedule time (storm dedupe); the
|
||||
# runner released its coalesce marker at lock acquire.
|
||||
assert "srv" in mgr._last_notification_refresh
|
||||
assert not mgr._static_refresh_pending
|
||||
|
||||
def test_handler_ignores_non_list_changed_messages(self, running_loop_mgr) -> None:
|
||||
"""Non-ServerNotification messages (RequestResponder / Exception)
|
||||
must schedule nothing — no stamp, no marker, no task."""
|
||||
mgr, loop, _thread = running_loop_mgr
|
||||
handler = mgr._make_static_notification_handler("srv")
|
||||
|
||||
async def _fire() -> int:
|
||||
before = len(mgr._background_tasks)
|
||||
await handler(MagicMock())
|
||||
return len(mgr._background_tasks) - before
|
||||
|
||||
assert _run_on_loop(loop, _fire()) == 0
|
||||
assert "srv" not in mgr._last_notification_refresh
|
||||
assert not mgr._static_refresh_pending
|
||||
|
||||
def test_notification_coalesces_when_refresh_already_queued(self, running_loop_mgr) -> None:
|
||||
"""While a runner is queued for a server+kind (coalesce marker
|
||||
set), further notifications spawn NOTHING — the parked runner's
|
||||
fresh list observes their change when it acquires the lock. This
|
||||
bounds the connect lock's waiter queue at one parked runner per
|
||||
server+kind: without it a notifying-but-slow server accretes
|
||||
waiters (admitted 1/5s, drained 1/30s) that starve the reconnect
|
||||
drivers sharing the lock, while ``_background_tasks`` grows
|
||||
without bound."""
|
||||
mgr, loop, _thread = running_loop_mgr
|
||||
refreshed: list[str] = []
|
||||
|
||||
async def _fake_refresh(name: str) -> tuple[list[str], list[str]]:
|
||||
refreshed.append(name)
|
||||
return [], []
|
||||
|
||||
mgr._refresh_server_tools = _fake_refresh # type: ignore[method-assign]
|
||||
handler = mgr._make_static_notification_handler("srv")
|
||||
|
||||
async def _fire_twice() -> int:
|
||||
lock = mgr._static_connect_lock_for("srv")
|
||||
_seed_static_state(mgr, "srv", session=MagicMock())
|
||||
await lock.acquire() # park the first runner
|
||||
note = mcp_types.ServerNotification(
|
||||
mcp_types.ToolListChangedNotification(method="notifications/tools/list_changed")
|
||||
)
|
||||
before = len(mgr._background_tasks)
|
||||
await handler(note)
|
||||
# Age the stamp past the debounce window: the MARKER (not
|
||||
# the stamp) must do the suppression for the second fire.
|
||||
mgr._last_notification_refresh["srv"] = time.monotonic() - 10.0
|
||||
await handler(note)
|
||||
spawned = len(mgr._background_tasks) - before
|
||||
lock.release()
|
||||
return spawned
|
||||
|
||||
spawned = _run_on_loop(loop, _fire_twice())
|
||||
_drain_background(mgr, loop)
|
||||
assert spawned == 1, "second notification must coalesce into the parked runner"
|
||||
assert refreshed == ["srv"]
|
||||
# The runner released its own marker at lock acquire.
|
||||
assert not mgr._static_refresh_pending
|
||||
|
||||
def test_notification_refresh_failure_keeps_debounce_stamp(self, running_loop_mgr) -> None:
|
||||
"""A failed spawned refresh must KEEP the debounce stamp: popping
|
||||
it re-arms the handler on every notification, so a fast-failing
|
||||
server spawns unthrottled refresh tasks at its notification rate.
|
||||
The bounded cost — a change announced in the remainder of the
|
||||
failed window waits for the server's next ``list_changed`` or the
|
||||
next reconnect — is the lesser failure (every teardown pops the
|
||||
stamp, so a reconnect's first notification refreshes immediately).
|
||||
The recorded operator error is type-name-only: the full exception
|
||||
chain can carry the configured bearer for ``auth_type=static``."""
|
||||
mgr, loop, _thread = running_loop_mgr
|
||||
|
||||
async def _seed() -> None:
|
||||
mgr._static_connect_lock_for("srv")
|
||||
_seed_static_state(mgr, "srv", session=MagicMock())
|
||||
|
||||
_run_on_loop(loop, _seed())
|
||||
mgr._last_notification_refresh["srv"] = 123.0
|
||||
|
||||
async def _boom(_name: str) -> tuple[list[str], list[str]]:
|
||||
raise TimeoutError("slow server")
|
||||
|
||||
_run_on_loop(loop, mgr._run_static_notification_refresh("srv", "tools", _boom))
|
||||
assert mgr._last_notification_refresh["srv"] == 123.0
|
||||
assert mgr._last_error["srv"] == "Refresh failed: TimeoutError"
|
||||
|
||||
async def _ok(_name: str) -> tuple[list[str], list[str]]:
|
||||
return [], []
|
||||
|
||||
mgr._last_notification_refresh["srv"] = 456.0
|
||||
_run_on_loop(loop, mgr._run_static_notification_refresh("srv", "tools", _ok))
|
||||
assert mgr._last_notification_refresh["srv"] == 456.0
|
||||
assert "srv" not in mgr._last_error # success clears the error
|
||||
|
||||
def test_notification_refresh_catches_exception_group(self, running_loop_mgr) -> None:
|
||||
"""A wedged anyio transport surfaces session-op failures as
|
||||
``BaseExceptionGroup`` — which ``except Exception`` misses. An
|
||||
escaping group reaches ``_spawn_background``'s failure log, whose
|
||||
``exc_info`` serializes the chained httpx request carrying the
|
||||
CONFIGURED bearer for ``auth_type=static`` servers. The runner
|
||||
must swallow the group (and keep the stamp) like any other
|
||||
refresh failure. The member below is a BaseException so the group
|
||||
does NOT collapse to ``ExceptionGroup`` (which ``Exception``
|
||||
would catch) — the anyio stray-cancel shape."""
|
||||
mgr, loop, _thread = running_loop_mgr
|
||||
|
||||
async def _seed() -> None:
|
||||
mgr._static_connect_lock_for("srv")
|
||||
_seed_static_state(mgr, "srv", session=MagicMock())
|
||||
|
||||
_run_on_loop(loop, _seed())
|
||||
mgr._last_notification_refresh["srv"] = 123.0
|
||||
|
||||
async def _wedge(_name: str) -> tuple[list[str], list[str]]:
|
||||
raise BaseExceptionGroup("wedged transport", [asyncio.CancelledError()])
|
||||
|
||||
_run_on_loop(loop, mgr._run_static_notification_refresh("srv", "tools", _wedge))
|
||||
assert mgr._last_notification_refresh["srv"] == 123.0
|
||||
assert mgr._last_error["srv"] == "Refresh failed: BaseExceptionGroup"
|
||||
|
||||
def test_notification_refresh_serializes_on_connect_lock(self, running_loop_mgr) -> None:
|
||||
"""The runner must take the per-name connect lock before
|
||||
refreshing: an unserialized refresh races an in-flight
|
||||
``_connect_one_locked``'s discovery wiring (which would overwrite
|
||||
the refresh's newer catalog with its older snapshot) and sibling
|
||||
same-server refreshes (the slower list call publishing the older
|
||||
catalog last)."""
|
||||
mgr, loop, _thread = running_loop_mgr
|
||||
refreshed: list[str] = []
|
||||
|
||||
async def _rec(name: str) -> tuple[list[str], list[str]]:
|
||||
refreshed.append(name)
|
||||
return [], []
|
||||
|
||||
async def _scenario() -> bool:
|
||||
lock = mgr._static_connect_lock_for("srv")
|
||||
_seed_static_state(mgr, "srv", session=MagicMock())
|
||||
await lock.acquire()
|
||||
task = asyncio.ensure_future(mgr._run_static_notification_refresh("srv", "tools", _rec))
|
||||
for _ in range(3):
|
||||
await asyncio.sleep(0)
|
||||
held_back = len(refreshed) == 0
|
||||
lock.release()
|
||||
await task
|
||||
return held_back
|
||||
|
||||
held_back = _run_on_loop(loop, _scenario())
|
||||
assert held_back, "refresh must wait for the connect lock"
|
||||
assert refreshed == ["srv"]
|
||||
|
||||
def test_notification_refresh_discards_when_server_removed_while_waiting(
|
||||
self, running_loop_mgr
|
||||
) -> None:
|
||||
"""``remove_server_sync`` retires the lock object after teardown;
|
||||
a runner that parked on the OLD lock must not touch state now
|
||||
owned by a NEW-lock holder (remove → re-add): the notification
|
||||
belonged to the old transport and the re-add publishes its own
|
||||
discovery under the new lock."""
|
||||
mgr, loop, _thread = running_loop_mgr
|
||||
refreshed: list[str] = []
|
||||
|
||||
async def _rec(name: str) -> tuple[list[str], list[str]]:
|
||||
refreshed.append(name)
|
||||
return [], []
|
||||
|
||||
async def _scenario() -> None:
|
||||
old_lock = mgr._static_connect_lock_for("srv")
|
||||
_seed_static_state(mgr, "srv", session=MagicMock())
|
||||
await old_lock.acquire()
|
||||
task = asyncio.ensure_future(mgr._run_static_notification_refresh("srv", "tools", _rec))
|
||||
for _ in range(3):
|
||||
await asyncio.sleep(0)
|
||||
# Simulate remove → re-add while the runner is parked: the
|
||||
# lock object is retired and a fresh one minted.
|
||||
mgr._static_connect_locks.pop("srv", None)
|
||||
mgr._static_connect_lock_for("srv")
|
||||
old_lock.release()
|
||||
await task
|
||||
|
||||
_run_on_loop(loop, _scenario())
|
||||
assert refreshed == []
|
||||
|
||||
def test_notification_refresh_skips_when_session_evicted_while_parked(
|
||||
self, running_loop_mgr
|
||||
) -> None:
|
||||
"""A session evicted while the runner was parked returns quietly:
|
||||
the reconnect's rediscovery republishes, and every teardown pops
|
||||
the debounce stamp, so the reconnected transport's first
|
||||
notification refreshes immediately."""
|
||||
mgr, loop, _thread = running_loop_mgr
|
||||
refreshed: list[str] = []
|
||||
|
||||
async def _rec(name: str) -> tuple[list[str], list[str]]:
|
||||
refreshed.append(name)
|
||||
return [], []
|
||||
|
||||
async def _scenario() -> None:
|
||||
lock = mgr._static_connect_lock_for("srv")
|
||||
state = _seed_static_state(mgr, "srv", session=MagicMock())
|
||||
await lock.acquire()
|
||||
task = asyncio.ensure_future(mgr._run_static_notification_refresh("srv", "tools", _rec))
|
||||
for _ in range(3):
|
||||
await asyncio.sleep(0)
|
||||
state.session = None # evicted while parked
|
||||
lock.release()
|
||||
await task
|
||||
|
||||
_run_on_loop(loop, _scenario())
|
||||
assert refreshed == []
|
||||
|
||||
def test_finally_does_not_clobber_successor_marker(self, running_loop_mgr) -> None:
|
||||
"""After the at-acquire discard, a marker present at the runner's
|
||||
exit belongs to the SUCCESSOR spawned during its in-flight list
|
||||
call — the finally must not discard it, or the handler mints
|
||||
runners past the one-parked-runner bound."""
|
||||
mgr, loop, _thread = running_loop_mgr
|
||||
marker = ("srv", "tools")
|
||||
|
||||
async def _refresh_readding(_name: str) -> tuple[list[str], list[str]]:
|
||||
# Our own marker was discarded at lock-acquire; a successor
|
||||
# spawned mid-refresh re-adds the same (server, kind) marker.
|
||||
assert marker not in mgr._static_refresh_pending
|
||||
mgr._static_refresh_pending.add(marker)
|
||||
return [], []
|
||||
|
||||
async def _scenario() -> None:
|
||||
mgr._static_connect_lock_for("srv")
|
||||
_seed_static_state(mgr, "srv", session=MagicMock())
|
||||
mgr._static_refresh_pending.add(marker) # our spawn's marker
|
||||
await mgr._run_static_notification_refresh("srv", "tools", _refresh_readding)
|
||||
|
||||
_run_on_loop(loop, _scenario())
|
||||
assert marker in mgr._static_refresh_pending, "successor's marker must survive our exit"
|
||||
mgr._static_refresh_pending.discard(marker)
|
||||
|
||||
def test_cancel_while_parked_releases_marker(self, running_loop_mgr) -> None:
|
||||
"""A runner cancelled while PARKED on the connect lock never
|
||||
reached the at-acquire discard — the finally must release its
|
||||
marker, or the server+kind never refreshes again."""
|
||||
mgr, loop, _thread = running_loop_mgr
|
||||
marker = ("srv", "tools")
|
||||
|
||||
async def _rec(_name: str) -> tuple[list[str], list[str]]:
|
||||
return [], []
|
||||
|
||||
async def _scenario() -> None:
|
||||
lock = mgr._static_connect_lock_for("srv")
|
||||
_seed_static_state(mgr, "srv", session=MagicMock())
|
||||
await lock.acquire()
|
||||
mgr._static_refresh_pending.add(marker)
|
||||
task = asyncio.ensure_future(mgr._run_static_notification_refresh("srv", "tools", _rec))
|
||||
for _ in range(3):
|
||||
await asyncio.sleep(0)
|
||||
task.cancel()
|
||||
with contextlib.suppress(asyncio.CancelledError):
|
||||
await task
|
||||
lock.release()
|
||||
|
||||
_run_on_loop(loop, _scenario())
|
||||
assert marker not in mgr._static_refresh_pending
|
||||
|
||||
def test_runner_releases_marker_when_lock_already_retired(self, running_loop_mgr) -> None:
|
||||
"""Spawn → remove completes (lock retired) → runner starts: it
|
||||
must release its marker and must NOT mint a fresh lock for the
|
||||
removed server."""
|
||||
mgr, loop, _thread = running_loop_mgr
|
||||
marker = ("gone", "tools")
|
||||
mgr._static_refresh_pending.add(marker)
|
||||
|
||||
async def _rec(_name: str) -> tuple[list[str], list[str]]:
|
||||
raise AssertionError("must not refresh a removed server")
|
||||
|
||||
_run_on_loop(loop, mgr._run_static_notification_refresh("gone", "tools", _rec))
|
||||
assert marker not in mgr._static_refresh_pending
|
||||
assert "gone" not in mgr._static_connect_locks
|
||||
|
||||
def test_teardown_static_session_pops_debounce_stamp(self, running_loop_mgr) -> None:
|
||||
"""Every teardown path must pop the debounce stamp — the
|
||||
keep-stamp-on-failure design leans on it: a reconnect's first
|
||||
``list_changed`` refreshes immediately, so a change announced in
|
||||
a failed window converges at the next reconnect."""
|
||||
mgr, loop, _thread = running_loop_mgr
|
||||
|
||||
async def _scenario() -> None:
|
||||
_seed_static_state(mgr, "srv", session=MagicMock())
|
||||
mgr._last_notification_refresh["srv"] = 123.0
|
||||
await mgr._teardown_static_session("srv")
|
||||
|
||||
_run_on_loop(loop, _scenario())
|
||||
assert "srv" not in mgr._last_notification_refresh
|
||||
|
||||
def test_owner_death_pops_debounce_stamp(self, running_loop_mgr) -> None:
|
||||
"""The unrequested-collapse path (owner done-callback) is a
|
||||
teardown too: the stamp must not outlive the transport."""
|
||||
mgr, loop, _thread = running_loop_mgr
|
||||
|
||||
async def _scenario() -> None:
|
||||
state = _seed_static_state(mgr, "srv", session=MagicMock())
|
||||
|
||||
async def _noop() -> None:
|
||||
pass
|
||||
|
||||
task = asyncio.ensure_future(_noop())
|
||||
await task
|
||||
state.owner_task = task
|
||||
mgr._last_notification_refresh["srv"] = 123.0
|
||||
mgr._on_static_owner_death("srv", task)
|
||||
|
||||
_run_on_loop(loop, _scenario())
|
||||
assert "srv" not in mgr._last_notification_refresh
|
||||
|
||||
def test_dead_transport_eviction_pops_debounce_stamp(self) -> None:
|
||||
"""The dispatch-observed transport-failure eviction is a teardown
|
||||
too — the stamp must not survive the session it throttled."""
|
||||
mgr = MCPClientManager({})
|
||||
_seed_static_state(mgr, "srv", session=MagicMock())
|
||||
mgr._last_notification_refresh["srv"] = 123.0
|
||||
mgr._record_and_evict_on_dead_transport("srv", anyio.ClosedResourceError())
|
||||
assert mgr._static_servers["srv"].session is None
|
||||
assert "srv" not in mgr._last_notification_refresh
|
||||
|
||||
def test_refresh_server_tools_bounded_by_timeout(self) -> None:
|
||||
"""A wedged server's list call must not hang a spawned refresh
|
||||
(and the connect lock it holds) forever — #839's unbounded
|
||||
``list_tools`` is what turned the receive-loop park into a
|
||||
permanent wedge on main."""
|
||||
mgr = MCPClientManager({})
|
||||
mgr._CONNECT_TIMEOUT = 0.05 # instance override of the class constant
|
||||
|
||||
async def _hang() -> Any:
|
||||
await asyncio.sleep(30)
|
||||
|
||||
session = MagicMock()
|
||||
session.list_tools = _hang
|
||||
_seed_static_state(mgr, "srv", session=session)
|
||||
|
||||
async def _run() -> None:
|
||||
await mgr._refresh_server_tools("srv")
|
||||
|
||||
with pytest.raises(TimeoutError):
|
||||
asyncio.run(_run())
|
||||
|
||||
def test_refresh_server_tools_discards_result_when_state_replaced(self) -> None:
|
||||
"""A state entry replaced (remove + re-add) while ``list_tools``
|
||||
was in flight owns the catalog now — publishing the stale result
|
||||
would clobber the new transport's discovery."""
|
||||
mgr = MCPClientManager({})
|
||||
session = MagicMock()
|
||||
old_state = _seed_static_state(mgr, "srv", session=session, tools=[])
|
||||
|
||||
async def _list_tools() -> Any:
|
||||
# Replace the state entry mid-flight (remove + re-add).
|
||||
mgr._static_servers.pop("srv")
|
||||
_seed_static_state(mgr, "srv", tools=[])
|
||||
result = MagicMock()
|
||||
result.tools = [_fake_mcp_tool("late")]
|
||||
return result
|
||||
|
||||
session.list_tools = _list_tools
|
||||
|
||||
async def _run() -> tuple[list[str], list[str]]:
|
||||
return await mgr._refresh_server_tools("srv")
|
||||
|
||||
added, removed = asyncio.run(_run())
|
||||
assert (added, removed) == ([], [])
|
||||
assert mgr._static_servers["srv"].tools == [] # new entry untouched
|
||||
assert old_state.tools == [] # stale result not published
|
||||
|
||||
def test_refresh_server_serializes_on_connect_lock(self, running_loop_mgr) -> None:
|
||||
"""The manual/periodic refresh is a catalog publisher too: it
|
||||
must queue behind the same per-name lock as connect wiring and
|
||||
the notification runner, or its older snapshot can land over a
|
||||
notification refresh's newer one."""
|
||||
mgr, loop, _thread = running_loop_mgr
|
||||
refreshed: list[str] = []
|
||||
|
||||
async def _rec_tools(name: str) -> tuple[list[str], list[str]]:
|
||||
refreshed.append(name)
|
||||
return [], []
|
||||
|
||||
async def _rec_none(_name: str) -> None:
|
||||
return None
|
||||
|
||||
mgr._refresh_server_tools = _rec_tools # type: ignore[method-assign]
|
||||
mgr._refresh_server_resources = _rec_none # type: ignore[method-assign]
|
||||
mgr._refresh_server_prompts = _rec_none # type: ignore[method-assign]
|
||||
|
||||
async def _scenario() -> bool:
|
||||
lock = mgr._static_connect_lock_for("srv")
|
||||
_seed_static_state(mgr, "srv", session=MagicMock())
|
||||
await lock.acquire()
|
||||
task = asyncio.ensure_future(mgr._refresh_server("srv"))
|
||||
for _ in range(3):
|
||||
await asyncio.sleep(0)
|
||||
held_back = len(refreshed) == 0
|
||||
lock.release()
|
||||
await task
|
||||
return held_back
|
||||
|
||||
held_back = _run_on_loop(loop, _scenario())
|
||||
assert held_back, "manual refresh must wait for the connect lock"
|
||||
assert refreshed == ["srv"]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# reconnect_sync — operator-driven full reconnect
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@@ -0,0 +1,190 @@
|
||||
"""Live push-refresh smoke test: a real ``tools/list_changed`` lands, no wedge.
|
||||
|
||||
End-to-end regression for #839: the static-path notification handler used to
|
||||
await its catalog refresh inline in the SDK's receive loop, but the refresh
|
||||
issues a request on the SAME session — a request whose response only that
|
||||
(now parked) receive loop could route. The refresh never completed, the
|
||||
receive loop wedged permanently, and every call on the shared per-node
|
||||
session stalled behind it; the only "recovery" was the health loop's ping
|
||||
timeout tearing the transport down.
|
||||
|
||||
A real streamable-http MCP server (FastMCP, subprocess) registers an extra
|
||||
tool at runtime inside a tool call and pushes ``notifications/tools/
|
||||
list_changed`` on the live session — so the notification and the call result
|
||||
are multiplexed on the real SDK receive loop, exactly the production shape.
|
||||
Pass criteria are discriminating on purpose:
|
||||
|
||||
* the TRIGGERING ``call_tool`` returns promptly — on the pre-fix code its
|
||||
result could never route past the parked receive loop, so this call is
|
||||
itself the deadlock repro;
|
||||
* the pushed catalog change lands on the SAME session object (health loop
|
||||
slowed to keep teardown/reconnect out of the picture) — the push did the
|
||||
work, not a rebuild;
|
||||
* the newly pushed tool DISPATCHES — the merged tool map was rebuilt, and
|
||||
the receive loop is still routing responses afterwards.
|
||||
|
||||
Self-contained (spawns its own server; no LLM backend, no network beyond
|
||||
127.0.0.1) — deliberately NOT marked ``live``, mirroring
|
||||
``test_mcp_live_flaky_server.py``. Wall clock ~5s.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import signal
|
||||
import socket
|
||||
import subprocess
|
||||
import sys
|
||||
import textwrap
|
||||
import time
|
||||
from typing import TYPE_CHECKING
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
||||
from turnstone.core.mcp_client import MCPClientManager
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from pathlib import Path
|
||||
|
||||
SERVER_SRC = textwrap.dedent(
|
||||
'''
|
||||
"""Streamable-http MCP server that grows a tool at runtime (#839 repro)."""
|
||||
import sys
|
||||
|
||||
from mcp.server.fastmcp import Context, FastMCP
|
||||
|
||||
port = int(sys.argv[1])
|
||||
mcp = FastMCP("push-victim", host="127.0.0.1", port=port)
|
||||
|
||||
|
||||
@mcp.tool()
|
||||
def ping_me(x: int) -> int:
|
||||
"""Return x + 1."""
|
||||
return x + 1
|
||||
|
||||
|
||||
def extra_tool(y: int) -> int:
|
||||
"""Return y * 2."""
|
||||
return y * 2
|
||||
|
||||
|
||||
@mcp.tool()
|
||||
async def register_extra(ctx: Context) -> str:
|
||||
"""Register extra_tool, then push tools/list_changed on this session."""
|
||||
mcp.add_tool(extra_tool)
|
||||
await ctx.session.send_tool_list_changed()
|
||||
return "registered"
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
mcp.run(transport="streamable-http")
|
||||
'''
|
||||
).lstrip()
|
||||
|
||||
|
||||
def _free_port() -> int:
|
||||
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
|
||||
s.bind(("127.0.0.1", 0))
|
||||
return int(s.getsockname()[1])
|
||||
|
||||
|
||||
def _wait_tcp_ready(port: int, timeout: float) -> bool:
|
||||
deadline = time.monotonic() + timeout
|
||||
while time.monotonic() < deadline:
|
||||
try:
|
||||
with socket.create_connection(("127.0.0.1", port), timeout=0.3):
|
||||
return True
|
||||
except OSError:
|
||||
time.sleep(0.05)
|
||||
return False
|
||||
|
||||
|
||||
def _wait_session_live(mgr: MCPClientManager, name: str, timeout: float) -> bool:
|
||||
deadline = time.monotonic() + timeout
|
||||
while time.monotonic() < deadline:
|
||||
state = mgr._static_servers.get(name)
|
||||
if state is not None and state.session is not None:
|
||||
return True
|
||||
time.sleep(0.05)
|
||||
return False
|
||||
|
||||
|
||||
def _wait_tool_visible(mgr: MCPClientManager, server: str, tool: str, timeout: float) -> bool:
|
||||
"""Poll the per-server catalog for *tool* — the push refresh landing."""
|
||||
deadline = time.monotonic() + timeout
|
||||
while time.monotonic() < deadline:
|
||||
state = mgr._static_servers.get(server)
|
||||
if state is not None and any(t["function"]["name"] == tool for t in state.tools):
|
||||
return True
|
||||
time.sleep(0.05)
|
||||
return False
|
||||
|
||||
|
||||
class TestPushRefreshNoDeadlock:
|
||||
def test_list_changed_push_refreshes_without_teardown(
|
||||
self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
# The subprocess runs sys.executable, so importability HERE is a
|
||||
# faithful proxy for the server side. Environment gaps skip, not fail.
|
||||
pytest.importorskip("mcp.server.fastmcp")
|
||||
script = tmp_path / "push_srv.py"
|
||||
script.write_text(SERVER_SRC)
|
||||
port = _free_port()
|
||||
|
||||
# Bound connect/discovery/refresh phases for a unit-test budget, but
|
||||
# SLOW the health loop right down: pre-fix, its ping-timeout teardown
|
||||
# was the accidental recovery path, and this test must prove the push
|
||||
# itself does the work on the ORIGINAL session.
|
||||
monkeypatch.setattr(MCPClientManager, "_CONNECT_TIMEOUT", 5)
|
||||
monkeypatch.setattr(MCPClientManager, "_TCP_PROBE_TIMEOUT", 1)
|
||||
|
||||
proc: subprocess.Popen[bytes] | None = None
|
||||
mgr: MCPClientManager | None = None
|
||||
try:
|
||||
proc = subprocess.Popen(
|
||||
[sys.executable, str(script), str(port)],
|
||||
stdout=subprocess.DEVNULL,
|
||||
stderr=subprocess.DEVNULL,
|
||||
)
|
||||
if not _wait_tcp_ready(port, 10.0):
|
||||
pytest.skip("push-refresh server subprocess did not come up")
|
||||
with patch(
|
||||
"turnstone.core.mcp_client.load_config",
|
||||
return_value={"static_health_check_seconds": 30},
|
||||
):
|
||||
mgr = MCPClientManager(
|
||||
{"push": {"type": "http", "url": f"http://127.0.0.1:{port}/mcp"}}
|
||||
)
|
||||
mgr.start()
|
||||
assert _wait_session_live(mgr, "push", 8.0), "initial connect failed"
|
||||
state = mgr._static_servers["push"]
|
||||
session_before = state.session
|
||||
assert not any(t["function"]["name"] == "mcp__push__extra_tool" for t in state.tools), (
|
||||
"extra_tool must not exist before the push"
|
||||
)
|
||||
|
||||
# THE repro: the server pushes list_changed while this call is in
|
||||
# flight, so its result and the notification share the receive
|
||||
# loop. Pre-fix, the inline-await handler parked that loop and
|
||||
# this call never returned.
|
||||
out = mgr.call_tool_sync("mcp__push__register_extra", {}, timeout=10)
|
||||
assert "registered" in out
|
||||
|
||||
# The push-driven refresh completes on its own — no teardown.
|
||||
assert _wait_tool_visible(mgr, "push", "mcp__push__extra_tool", 8.0), (
|
||||
"pushed tools/list_changed never refreshed the catalog"
|
||||
)
|
||||
assert mgr._static_servers["push"].session is session_before, (
|
||||
"catalog arrived via teardown/reconnect, not via the push refresh"
|
||||
)
|
||||
|
||||
# The merged map rebuilt AND the receive loop still routes:
|
||||
# the brand-new tool dispatches end-to-end.
|
||||
out = mgr.call_tool_sync("mcp__push__extra_tool", {"y": 21}, timeout=10)
|
||||
assert "42" in out
|
||||
finally:
|
||||
if mgr is not None:
|
||||
mgr.shutdown()
|
||||
if proc is not None:
|
||||
proc.send_signal(signal.SIGKILL)
|
||||
proc.wait(timeout=5)
|
||||
+247
-65
@@ -148,14 +148,15 @@ _MAX_PROMPTS_PER_SERVER = 1000
|
||||
# thundering herd of connects on every session start.
|
||||
_PRIME_MAX_CONCURRENCY = 4
|
||||
|
||||
# One row per pool ``*/list_changed`` kind: notification type → kind label.
|
||||
# The pool notification handler drives all three catalog kinds through this
|
||||
# table (plus its kind → bound-method map, which stays mypy-checked attribute
|
||||
# access rather than a name-string table) so the debounce/coalesce/log/spawn
|
||||
# protocol exists exactly once — a protocol change cannot silently diverge
|
||||
# between kinds. Exact-type lookup is safe: the SDK's discriminated-union
|
||||
# parser instantiates these concrete classes, never subclasses.
|
||||
_POOL_LIST_CHANGED_KINDS: dict[type, str] = {
|
||||
# One row per ``*/list_changed`` kind: notification type → kind label.
|
||||
# The static AND pool notification handlers drive all three catalog kinds
|
||||
# through this table (plus their kind → bound-method maps, which stay
|
||||
# mypy-checked attribute access rather than name-string tables) so the
|
||||
# debounce/coalesce/log/spawn protocol exists exactly once per path — a
|
||||
# protocol change cannot silently diverge between kinds. Exact-type lookup
|
||||
# is safe: the SDK's discriminated-union parser instantiates these concrete
|
||||
# classes, never subclasses.
|
||||
_LIST_CHANGED_KINDS: dict[type, str] = {
|
||||
mcp_types.ToolListChangedNotification: "tools",
|
||||
mcp_types.ResourceListChangedNotification: "resources",
|
||||
mcp_types.PromptListChangedNotification: "prompts",
|
||||
@@ -718,6 +719,14 @@ class MCPClientManager:
|
||||
|
||||
# Notification debounce (per-server)
|
||||
self._last_notification_refresh: dict[str, float] = {}
|
||||
# Coalescing markers for spawned static ``list_changed`` refreshes,
|
||||
# keyed ``(server, kind)``: set at spawn, cleared the moment the
|
||||
# runner acquires the per-name connect lock (before its list call).
|
||||
# While set, further notifications for the server+kind are dropped —
|
||||
# the parked runner's fresh list will observe their change — bounding
|
||||
# the connect lock's waiter queue at ONE parked runner per
|
||||
# server+kind. Mirrors ``_pool_refresh_pending``.
|
||||
self._static_refresh_pending: set[tuple[str, str]] = set()
|
||||
|
||||
# Last refresh outcome (Phase 9 — admin status indicator). Per-
|
||||
# server tuple of ``(unix_ts, outcome)`` where outcome is one of
|
||||
@@ -1323,6 +1332,21 @@ class MCPClientManager:
|
||||
)
|
||||
return disarmed
|
||||
|
||||
def _drop_static_session_and_stamp(self, name: str, state: StaticServerState) -> None:
|
||||
"""Null the session AND pop its notification debounce stamp (paired).
|
||||
|
||||
The static twin of :meth:`_drop_session_and_stamp` (pool). The
|
||||
stamp must not outlive the transport: the keep-stamp-on-failure
|
||||
design leans on every teardown popping it, so a reconnected
|
||||
transport's first ``list_changed`` refreshes immediately instead
|
||||
of being debounced against a pre-collapse stamp. An eviction site
|
||||
that nulled the session directly would silently re-open that
|
||||
stale-stamp hole. Safe when the session is already gone: both
|
||||
halves are idempotent.
|
||||
"""
|
||||
state.session = None
|
||||
self._last_notification_refresh.pop(name, None)
|
||||
|
||||
async def _teardown_static_session(self, name: str) -> None:
|
||||
"""Tear down a static server's session/transport (the ONE canonical order).
|
||||
|
||||
@@ -1357,7 +1381,7 @@ class MCPClientManager:
|
||||
state = self._static_servers.get(name)
|
||||
if state is None:
|
||||
return
|
||||
state.session = None
|
||||
self._drop_static_session_and_stamp(name, state)
|
||||
owner = state.owner_task
|
||||
close_requested = state.close_requested
|
||||
state.owner_task = None
|
||||
@@ -1526,10 +1550,16 @@ class MCPClientManager:
|
||||
def _make_static_notification_handler(self, name: str) -> Any:
|
||||
"""Build the per-server notification handler for a static session.
|
||||
|
||||
Dispatches tool, resource, and prompt list-change notifications to the
|
||||
appropriate refresh method (body unchanged from the pre-owner-task
|
||||
closure in ``_connect_one_locked``; extracted because the session cm
|
||||
is now entered by the transport owner).
|
||||
The handler itself NEVER awaits a request on the session: the SDK
|
||||
awaits message handlers inline in its receive loop, so an in-handler
|
||||
request on the same session can never receive its response — the
|
||||
response can only be routed by the receive loop that is parked
|
||||
awaiting this handler, and while it is parked EVERY in-flight and
|
||||
subsequent call on this shared per-node session stalls with it.
|
||||
List-change refreshes are debounced, coalesced per (server, kind),
|
||||
and SPAWNED as tracked tasks
|
||||
(:meth:`_run_static_notification_refresh`), mirroring the pool
|
||||
handler (:meth:`_make_pool_notification_handler`).
|
||||
"""
|
||||
|
||||
async def _on_notification(
|
||||
@@ -1537,8 +1567,9 @@ class MCPClientManager:
|
||||
) -> None:
|
||||
if not isinstance(msg, mcp_types.ServerNotification):
|
||||
return
|
||||
root = msg.root
|
||||
|
||||
kind = _LIST_CHANGED_KINDS.get(type(msg.root))
|
||||
if kind is None:
|
||||
return
|
||||
# Debounce: skip if we refreshed this server very recently
|
||||
now = time.monotonic()
|
||||
last = self._last_notification_refresh.get(name, 0.0)
|
||||
@@ -1549,27 +1580,134 @@ class MCPClientManager:
|
||||
now - last,
|
||||
)
|
||||
return
|
||||
|
||||
marker = (name, kind)
|
||||
if marker in self._static_refresh_pending:
|
||||
# A runner for this server+kind is queued but has not yet
|
||||
# issued its list call — it will observe this change when
|
||||
# it runs. Skipping bounds the connect lock's waiter queue
|
||||
# at one parked runner per server+kind, so a notifying-
|
||||
# but-slow server cannot accrete waiters that starve the
|
||||
# reconnect drivers sharing that lock.
|
||||
log.debug(
|
||||
"Coalescing static %s notification from '%s' (refresh queued)",
|
||||
kind,
|
||||
name,
|
||||
)
|
||||
return
|
||||
try:
|
||||
if isinstance(root, mcp_types.ToolListChangedNotification):
|
||||
log.info("Received tools/list_changed from '%s'", name)
|
||||
self._last_notification_refresh[name] = now
|
||||
await self._refresh_server_tools(name)
|
||||
elif isinstance(root, mcp_types.ResourceListChangedNotification):
|
||||
log.info("Received resources/list_changed from '%s'", name)
|
||||
self._last_notification_refresh[name] = now
|
||||
await self._refresh_server_resources(name)
|
||||
elif isinstance(root, mcp_types.PromptListChangedNotification):
|
||||
log.info("Received prompts/list_changed from '%s'", name)
|
||||
self._last_notification_refresh[name] = now
|
||||
await self._refresh_server_prompts(name)
|
||||
self._last_error.pop(name, None)
|
||||
# SPAWNED, never awaited — see the factory docstring: an
|
||||
# inline await here wedges the receive loop permanently.
|
||||
# The refresh runs as its own tracked task.
|
||||
#
|
||||
# Bound at dispatch time (not a module-level name table)
|
||||
# so instance-level overrides keep working and mypy
|
||||
# checks the attribute references.
|
||||
refreshers: dict[str, Callable[[str], Awaitable[Any]]] = {
|
||||
"tools": self._refresh_server_tools,
|
||||
"resources": self._refresh_server_resources,
|
||||
"prompts": self._refresh_server_prompts,
|
||||
}
|
||||
log.info("Received %s/list_changed from '%s'", kind, name)
|
||||
self._last_notification_refresh[name] = now
|
||||
self._static_refresh_pending.add(marker)
|
||||
self._spawn_background(
|
||||
self._run_static_notification_refresh(name, kind, refreshers[kind]),
|
||||
f"static {kind} refresh for '{name}'",
|
||||
)
|
||||
except Exception as exc:
|
||||
log.warning("Refresh after notification failed for '%s'", name, exc_info=True)
|
||||
self._set_error(name, f"Refresh failed: {exc}")
|
||||
# Scheduling failed (loop shutting down) — release the
|
||||
# coalesce marker or this server+kind never refreshes
|
||||
# again. Structured fields only — ``exc_info=True`` would
|
||||
# serialize the chained ``httpx.Request`` whose headers
|
||||
# carry the configured bearer for ``auth_type=static``.
|
||||
self._static_refresh_pending.discard(marker)
|
||||
log.warning(
|
||||
"Static refresh scheduling failed server=%s exc=%s",
|
||||
name,
|
||||
type(exc).__name__,
|
||||
)
|
||||
|
||||
return _on_notification
|
||||
|
||||
async def _run_static_notification_refresh(
|
||||
self,
|
||||
name: str,
|
||||
kind: str,
|
||||
refresh: Callable[[str], Awaitable[Any]],
|
||||
) -> None:
|
||||
"""Body of a spawned static ``list_changed`` refresh — never on the receive loop.
|
||||
|
||||
The static-path port of :meth:`_run_notification_refresh` — see that
|
||||
docstring for the full protocol rationale (serialize-then-list so the
|
||||
last publish is always the freshest; coalesce marker cleared at
|
||||
lock-ACQUIRE so a change the in-flight list missed spawns exactly one
|
||||
successor; debounce stamp SURVIVES failure and is popped by every
|
||||
teardown; ``finally`` discard gated on non-acquisition so it never
|
||||
clobbers a successor's marker; bounded wedged-notifier residual).
|
||||
The static primitives differ:
|
||||
|
||||
* Serialization is on the per-name CONNECT lock — the static path's
|
||||
one writer lock, shared with ``_connect_one_locked``'s discovery
|
||||
wiring, ``_refresh_server``, and the teardown protocol.
|
||||
* The remove → re-add race is closed by LOCK IDENTITY (the pool uses
|
||||
entry identity): ``remove_server_sync`` retires the lock object
|
||||
after teardown, so a runner that parked on the OLD lock must not
|
||||
touch state now owned by a NEW-lock holder.
|
||||
* ``(Exception, BaseExceptionGroup)`` is caught type-name-only for
|
||||
the same hygiene reason as the pool runner, but the credential at
|
||||
stake is the CONFIGURED bearer: for ``auth_type=static`` servers
|
||||
the chained ``httpx.Request`` headers carry it, and an escaping
|
||||
exception would reach ``_spawn_background``'s ``exc_info`` log.
|
||||
"""
|
||||
marker = (name, kind)
|
||||
lock = self._static_connect_locks.get(name)
|
||||
if lock is None:
|
||||
# Removed (lock retired) between spawn and run — nothing to
|
||||
# refresh; the marker is still ours to release. ``.get()``, not
|
||||
# the get-or-create helper: minting a fresh lock here would
|
||||
# resurrect an entry for a server that no longer exists.
|
||||
self._static_refresh_pending.discard(marker)
|
||||
return
|
||||
acquired = False
|
||||
try:
|
||||
async with lock:
|
||||
self._static_refresh_pending.discard(marker)
|
||||
acquired = True
|
||||
if self._static_connect_locks.get(name) is not lock:
|
||||
# Removed (and possibly re-added) while we were parked:
|
||||
# the notification belonged to the old transport, and a
|
||||
# re-add publishes its own discovery under the NEW lock.
|
||||
return
|
||||
state = self._static_servers.get(name)
|
||||
if state is None or state.session is None:
|
||||
# Torn down / evicted while we were parked. The
|
||||
# reconnect's rediscovery republishes, and every
|
||||
# teardown pops the debounce stamp, so the reconnected
|
||||
# transport's first notification refreshes immediately.
|
||||
return
|
||||
await refresh(name)
|
||||
self._last_error.pop(name, None)
|
||||
except (Exception, BaseExceptionGroup) as exc:
|
||||
# Structured fields only — ``exc_info`` would serialize the
|
||||
# chained ``httpx.Request`` whose headers carry the configured
|
||||
# bearer for ``auth_type=static`` servers.
|
||||
log.warning(
|
||||
"Static %s refresh after notification failed for '%s' exc=%s",
|
||||
kind,
|
||||
name,
|
||||
type(exc).__name__,
|
||||
)
|
||||
self._set_error(name, f"Refresh failed: {type(exc).__name__}")
|
||||
finally:
|
||||
if not acquired:
|
||||
# Cancelled (or failed) while PARKED — the marker still in
|
||||
# the set is OURS; release it or this server+kind never
|
||||
# refreshes again. After the at-acquire discard, a marker
|
||||
# present at exit belongs to the SUCCESSOR spawned during
|
||||
# our in-flight list — discarding it would let the handler
|
||||
# mint runners past the one-parked-runner bound.
|
||||
self._static_refresh_pending.discard(marker)
|
||||
|
||||
async def _static_transport_owner(
|
||||
self,
|
||||
name: str,
|
||||
@@ -1705,7 +1843,7 @@ class MCPClientManager:
|
||||
state.owner_task = None
|
||||
state.close_requested = None
|
||||
if state.session is not None:
|
||||
state.session = None
|
||||
self._drop_static_session_and_stamp(name, state)
|
||||
log.info("MCP server '%s' transport terminated; session evicted for reconnect", name)
|
||||
|
||||
async def _connect_one_locked(self, name: str, cfg: dict[str, Any]) -> None:
|
||||
@@ -1914,7 +2052,7 @@ class MCPClientManager:
|
||||
) -> None:
|
||||
if not isinstance(msg, mcp_types.ServerNotification):
|
||||
return
|
||||
kind = _POOL_LIST_CHANGED_KINDS.get(type(msg.root))
|
||||
kind = _LIST_CHANGED_KINDS.get(type(msg.root))
|
||||
if kind is None:
|
||||
return
|
||||
now = time.monotonic()
|
||||
@@ -3726,7 +3864,16 @@ class MCPClientManager:
|
||||
|
||||
old_names = {t["function"]["name"] for t in state.tools}
|
||||
|
||||
result = await session.list_tools()
|
||||
# ``asyncio.timeout`` mandatory — a wedged server must not hang a
|
||||
# spawned refresh (and the connect lock it holds) forever; the pool
|
||||
# sibling (:meth:`_refresh_pool_server_tools`) already complies.
|
||||
async with asyncio.timeout(self._CONNECT_TIMEOUT):
|
||||
result = await session.list_tools()
|
||||
if self._static_servers.get(name) is not state:
|
||||
# The state entry was replaced (remove + re-add) while
|
||||
# list_tools was in flight — this result belongs to the old
|
||||
# transport; publishing it would clobber the new discovery.
|
||||
return [], []
|
||||
capped = _cap_server_tools(name, result.tools)
|
||||
server_tools = [_mcp_to_openai(name, tool) for tool in capped]
|
||||
new_names = {t["function"]["name"] for t in server_tools}
|
||||
@@ -3970,31 +4117,45 @@ class MCPClientManager:
|
||||
/ ``state.prompts`` are bounded to whichever sub-refresh
|
||||
succeeded — the documented trade-off vs leaving orphan tasks
|
||||
running after the error is observed.
|
||||
|
||||
Serialized on the per-name connect lock: every publisher of a
|
||||
static per-server catalog — a connect's discovery wiring, a
|
||||
spawned notification refresh
|
||||
(:meth:`_run_static_notification_refresh`), and this
|
||||
manual/periodic pass — writes under the same lock, so each list
|
||||
call is issued only after the previous publisher finished and
|
||||
the last publish is always the freshest. An unserialized pass
|
||||
could land its older snapshot over a notification refresh's
|
||||
newer one, with no convergence until the next change. No caller
|
||||
holds the lock coming in: ``_refresh_all`` takes it per-branch
|
||||
(never nested), and the two post-reconnect schedule sites spawn
|
||||
this as its own task.
|
||||
"""
|
||||
results = await asyncio.gather(
|
||||
self._refresh_server_tools(name),
|
||||
self._refresh_server_resources(name),
|
||||
self._refresh_server_prompts(name),
|
||||
return_exceptions=True,
|
||||
)
|
||||
first_exc: BaseException | None = next(
|
||||
(r for r in results if isinstance(r, BaseException)), None
|
||||
)
|
||||
if first_exc is not None:
|
||||
self._last_refresh[name] = (
|
||||
time.time(),
|
||||
f"error:{type(first_exc).__name__}",
|
||||
async with self._static_connect_lock_for(name):
|
||||
results = await asyncio.gather(
|
||||
self._refresh_server_tools(name),
|
||||
self._refresh_server_resources(name),
|
||||
self._refresh_server_prompts(name),
|
||||
return_exceptions=True,
|
||||
)
|
||||
raise first_exc
|
||||
tool_diff = results[0]
|
||||
# ``return_exceptions=True`` widens the static type; on the all-
|
||||
# success path each entry is the awaited result. We narrow the
|
||||
# tool-diff entry to the documented ``(added, removed)`` shape.
|
||||
assert isinstance(tool_diff, tuple)
|
||||
added, removed = tool_diff
|
||||
self._last_error.pop(name, None)
|
||||
self._last_refresh[name] = (time.time(), "ok")
|
||||
return added, removed
|
||||
first_exc: BaseException | None = next(
|
||||
(r for r in results if isinstance(r, BaseException)), None
|
||||
)
|
||||
if first_exc is not None:
|
||||
self._last_refresh[name] = (
|
||||
time.time(),
|
||||
f"error:{type(first_exc).__name__}",
|
||||
)
|
||||
raise first_exc
|
||||
tool_diff = results[0]
|
||||
# ``return_exceptions=True`` widens the static type; on the all-
|
||||
# success path each entry is the awaited result. We narrow the
|
||||
# tool-diff entry to the documented ``(added, removed)`` shape.
|
||||
assert isinstance(tool_diff, tuple)
|
||||
added, removed = tool_diff
|
||||
self._last_error.pop(name, None)
|
||||
self._last_refresh[name] = (time.time(), "ok")
|
||||
return added, removed
|
||||
|
||||
async def _refresh_all(
|
||||
self, server_name: str | None = None
|
||||
@@ -4054,7 +4215,7 @@ class MCPClientManager:
|
||||
if _is_dead_transport(exc):
|
||||
dead_state = self._static_servers.get(name)
|
||||
if dead_state is not None:
|
||||
dead_state.session = None
|
||||
self._drop_static_session_and_stamp(name, dead_state)
|
||||
# Overwrite unconditionally with the freshest observed
|
||||
# outcome. Two cases produce the write:
|
||||
# (1) Reconnect branch: ``_connect_one`` raised before
|
||||
@@ -4164,8 +4325,22 @@ class MCPClientManager:
|
||||
# turn the second list_resource_templates() call into AttributeError.
|
||||
session = state.session
|
||||
|
||||
# ``asyncio.timeout`` mandatory — a wedged server must not hang a
|
||||
# spawned refresh (and the connect lock it holds) forever.
|
||||
async with asyncio.timeout(self._CONNECT_TIMEOUT):
|
||||
# 1-RTT (gather) instead of 2 sequential RTTs — both calls
|
||||
# share the same timeout budget and target disjoint catalogs
|
||||
# (resources vs. templates), so ordering is irrelevant.
|
||||
# Mirrors :meth:`_refresh_pool_server_resources`.
|
||||
res_result, tmpl_result = await asyncio.gather(
|
||||
session.list_resources(),
|
||||
session.list_resource_templates(),
|
||||
)
|
||||
if self._static_servers.get(name) is not state:
|
||||
# Entry replaced (remove + re-add) mid-flight — stale result.
|
||||
return
|
||||
|
||||
server_resources: list[dict[str, Any]] = []
|
||||
res_result = await session.list_resources()
|
||||
for r in res_result.resources:
|
||||
server_resources.append(
|
||||
{
|
||||
@@ -4176,7 +4351,6 @@ class MCPClientManager:
|
||||
"server": name,
|
||||
}
|
||||
)
|
||||
tmpl_result = await session.list_resource_templates()
|
||||
for t in tmpl_result.resourceTemplates:
|
||||
server_resources.append(
|
||||
{
|
||||
@@ -4222,8 +4396,15 @@ class MCPClientManager:
|
||||
# multi-await tomorrow; consistent capture-once idiom.
|
||||
session = state.session
|
||||
|
||||
# ``asyncio.timeout`` mandatory — a wedged server must not hang a
|
||||
# spawned refresh (and the connect lock it holds) forever.
|
||||
async with asyncio.timeout(self._CONNECT_TIMEOUT):
|
||||
prompt_result = await session.list_prompts()
|
||||
if self._static_servers.get(name) is not state:
|
||||
# Entry replaced (remove + re-add) mid-flight — stale result.
|
||||
return
|
||||
|
||||
server_prompts: list[dict[str, Any]] = []
|
||||
prompt_result = await session.list_prompts()
|
||||
for p in prompt_result.prompts:
|
||||
server_prompts.append(
|
||||
{
|
||||
@@ -4648,7 +4829,7 @@ class MCPClientManager:
|
||||
async def _close_all_owners() -> None:
|
||||
owners: list[asyncio.Task[None]] = []
|
||||
for srv_name, srv_state in list(self._static_servers.items()):
|
||||
srv_state.session = None
|
||||
self._drop_static_session_and_stamp(srv_name, srv_state)
|
||||
owner = srv_state.owner_task
|
||||
close_requested = srv_state.close_requested
|
||||
srv_state.owner_task = None
|
||||
@@ -4727,6 +4908,7 @@ class MCPClientManager:
|
||||
self._circuit_open_until.clear()
|
||||
self._circuit_trip_count.clear()
|
||||
self._last_notification_refresh.clear()
|
||||
self._static_refresh_pending.clear()
|
||||
self._last_pool_notification_refresh.clear()
|
||||
self._pool_refresh_pending.clear()
|
||||
# Pool state already cleared above when the loop was alive; this
|
||||
@@ -4872,7 +5054,7 @@ class MCPClientManager:
|
||||
fail_state.tools = []
|
||||
fail_state.resources = []
|
||||
fail_state.prompts = []
|
||||
fail_state.session = None
|
||||
self._drop_static_session_and_stamp(name, fail_state)
|
||||
self._rebuild_tools()
|
||||
self._rebuild_resources()
|
||||
self._rebuild_prompts()
|
||||
@@ -5801,7 +5983,7 @@ class MCPClientManager:
|
||||
busy = evict is not None and evict.in_flight > 0
|
||||
if dead and evict is not None and not busy and evict.session is session:
|
||||
self._cb_record_failure(name)
|
||||
evict.session = None
|
||||
self._drop_static_session_and_stamp(name, evict)
|
||||
asap = time.monotonic()
|
||||
self._static_reconnect_next[name] = asap # reconnect asap
|
||||
self._static_reconnect_attempt.pop(name, None)
|
||||
@@ -5915,7 +6097,7 @@ class MCPClientManager:
|
||||
if dead:
|
||||
evict = self._static_servers.get(server_name)
|
||||
if evict is not None:
|
||||
evict.session = None
|
||||
self._drop_static_session_and_stamp(server_name, evict)
|
||||
|
||||
async def _static_session_op(self, server_name: str, op: Coroutine[Any, Any, Any]) -> Any:
|
||||
"""Await a static session op on the mcp-loop, pinned against eviction.
|
||||
|
||||
Reference in New Issue
Block a user