Serve subscriptions/listen with a pluggable event bus (SEP-2575) (#3035)

examples/stories/subscriptions/client.py

@@ -0,0 +1,97 @@
+"""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)