Max
98 lines
4.4 KiB
Python
98 lines
4.4 KiB
Python
"""Open a `subscriptions/listen` stream, watch one URI and the tool list, then close it."""
|
|
|
|
import anyio
|
|
import mcp_types as types
|
|
|
|
from mcp.client import Client
|
|
from stories._harness import Target, run_client
|
|
|
|
SUBSCRIPTION_ID = "io.modelcontextprotocol/subscriptionId"
|
|
|
|
|
|
async def main(target: Target, *, mode: str = "auto") -> None:
|
|
# Stream frames arrive as ordinary server notifications; `message_handler`
|
|
# is constructor-only on `Client`, so the list it fills exists first.
|
|
received: list[types.ServerNotification] = []
|
|
arrival = anyio.Event()
|
|
|
|
async def on_message(message: object) -> None:
|
|
nonlocal arrival
|
|
if isinstance(
|
|
message,
|
|
types.SubscriptionsAcknowledgedNotification
|
|
| types.ResourceUpdatedNotification
|
|
| types.ToolListChangedNotification,
|
|
):
|
|
received.append(message)
|
|
arrival.set()
|
|
arrival = anyio.Event()
|
|
|
|
async def wait_for(count: int) -> None:
|
|
with anyio.fail_after(10):
|
|
while len(received) < count:
|
|
await arrival.wait()
|
|
|
|
async with Client(target, mode=mode, message_handler=on_message) as client:
|
|
before = await client.list_tools()
|
|
assert "search" not in {tool.name for tool in before.tools}
|
|
|
|
async with anyio.create_task_group() as tg:
|
|
# There is no client-side listen API yet, so the story drops to the
|
|
# `client.session` escape hatch. The request parks for the stream's
|
|
# lifetime, so it runs as a task; cancelling it releases the local
|
|
# awaiting scope. In-memory that also ends the server's stream; over
|
|
# HTTP today nothing aborts the POST, so the server-side stream ends
|
|
# when the connection closes (the `Client` exit right below).
|
|
async def listen() -> None:
|
|
request = types.SubscriptionsListenRequest(
|
|
params=types.SubscriptionsListenRequestParams(
|
|
notifications=types.SubscriptionFilter(
|
|
tools_list_changed=True, resource_subscriptions=["note://todo"]
|
|
)
|
|
)
|
|
)
|
|
await client.session.send_request(request, types.SubscriptionsListenResult)
|
|
|
|
tg.start_soon(listen)
|
|
|
|
# ── the ack is the first frame: it echoes the honored filter, tagged ──
|
|
await wait_for(1)
|
|
ack = received[0]
|
|
assert isinstance(ack, types.SubscriptionsAcknowledgedNotification), ack
|
|
assert ack.params.notifications.tools_list_changed is True
|
|
assert ack.params.notifications.resource_subscriptions == ["note://todo"]
|
|
assert ack.params.meta is not None and SUBSCRIPTION_ID in ack.params.meta
|
|
|
|
# ── exact-URI filtering: an unsubscribed note edit stays silent ──
|
|
await client.call_tool("edit_note", {"name": "journal", "text": "day two"})
|
|
# ── the subscribed URI delivers, carrying the same subscription id ──
|
|
await client.call_tool("edit_note", {"name": "todo", "text": "water plants"})
|
|
await wait_for(2)
|
|
updated = received[1]
|
|
assert isinstance(updated, types.ResourceUpdatedNotification), updated
|
|
assert updated.params.uri == "note://todo"
|
|
assert updated.params.meta is not None
|
|
assert updated.params.meta[SUBSCRIPTION_ID] == ack.params.meta[SUBSCRIPTION_ID]
|
|
assert len(received) == 2, "the journal edit must not have been delivered"
|
|
|
|
# ── a runtime tool registration announces itself ──
|
|
await client.call_tool("enable_search", {})
|
|
await wait_for(3)
|
|
assert isinstance(received[2], types.ToolListChangedNotification), received[2]
|
|
|
|
# The client is done listening: cancel the parked request and let
|
|
# the connection teardown below end the stream server-side.
|
|
tg.cancel_scope.cancel()
|
|
|
|
# list_changed told us to re-fetch - the new tool is callable, and the
|
|
# session outlives the closed stream.
|
|
tools = await client.list_tools()
|
|
assert "search" in {tool.name for tool in tools.tools}
|
|
result = await client.call_tool("search", {"query": "water"})
|
|
content = result.content[0]
|
|
assert isinstance(content, types.TextContent) and content.text == "todo", result
|
|
|
|
|
|
if __name__ == "__main__":
|
|
run_client(main)
|