[v1.x] Scope experimental tasks to the session that created them (#2720)

tests/experimental/tasks/server/test_task_scope.py

@@ -0,0 +1,150 @@
+"""Unit tests for the task session-scope helpers.
+
+A session scope is an opaque marker assigned to each session by
+TaskSupport.configure_session(). Task IDs generated by run_task() embed it so
+the default task handlers can tell which session created a task. See
+test_task_visibility.py for the end-to-end behaviour these helpers produce.
+"""
+
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager
+
+import anyio
+import pytest
+
+from mcp.server import Server
+from mcp.server.experimental.task_scope import (
+    new_session_scope,
+    scoped_task_id,
+    session_scope_of,
+    task_in_session_scope,
+    task_listable_in_session_scope,
+)
+from mcp.server.experimental.task_support import TaskSupport
+from mcp.server.lowlevel import NotificationOptions
+from mcp.server.models import InitializationOptions
+from mcp.server.session import ServerSession
+from mcp.shared.message import SessionMessage
+
+
+def test_new_session_scope_is_unique() -> None:
+    assert new_session_scope() != new_session_scope()
+
+
+def test_scoped_task_id_round_trips_its_scope() -> None:
+    scope = new_session_scope()
+
+    task_id = scoped_task_id(scope)
+
+    assert session_scope_of(task_id) == scope
+
+
+def test_scoped_task_ids_are_unique_within_a_scope() -> None:
+    scope = new_session_scope()
+
+    assert scoped_task_id(scope) != scoped_task_id(scope)
+
+
+@pytest.mark.parametrize(
+    "task_id",
+    [
+        "plain-task-id",
+        "550e8400-e29b-41d4-a716-446655440000",  # bare uuid4
+        "",
+        # Right shape but the scope half is not 32 hex chars.
+        "not-a-scope:550e8400-e29b-41d4-a716-446655440000",
+        # Right scope half but the suffix is not a uuid4.
+        "0123456789abcdef0123456789abcdef:not-a-uuid",
+        # Uppercase hex is not produced by new_session_scope().
+        "0123456789ABCDEF0123456789ABCDEF:550e8400-e29b-41d4-a716-446655440000",
+    ],
+)
+def test_session_scope_of_returns_none_for_unscoped_ids(task_id: str) -> None:
+    assert session_scope_of(task_id) is None
+
+
+def test_a_scoped_task_is_usable_only_from_the_scope_that_created_it() -> None:
+    scope = new_session_scope()
+    task_id = scoped_task_id(scope)
+
+    assert task_in_session_scope(task_id, scope) is True
+    assert task_in_session_scope(task_id, new_session_scope()) is False
+    assert task_in_session_scope(task_id, None) is False
+
+
+def test_an_unscoped_task_is_usable_from_any_scope() -> None:
+    assert task_in_session_scope("plain-task-id", new_session_scope()) is True
+    assert task_in_session_scope("plain-task-id", None) is True
+
+
+def test_a_scoped_task_is_listable_only_in_the_scope_that_created_it() -> None:
+    scope = new_session_scope()
+    task_id = scoped_task_id(scope)
+
+    assert task_listable_in_session_scope(task_id, scope) is True
+    assert task_listable_in_session_scope(task_id, new_session_scope()) is False
+    assert task_listable_in_session_scope(task_id, None) is False
+
+
+def test_an_unscoped_task_is_never_listable() -> None:
+    assert task_listable_in_session_scope("plain-task-id", new_session_scope()) is False
+    assert task_listable_in_session_scope("plain-task-id", None) is False
+
+
+@asynccontextmanager
+async def _make_session() -> AsyncIterator[ServerSession]:
+    """Create a ServerSession suitable for inspecting configure_session()."""
+    server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](1)
+    client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage | Exception](1)
+    options = InitializationOptions(
+        server_name="test",
+        server_version="0",
+        capabilities=Server("test").get_capabilities(NotificationOptions(), {}),
+    )
+    async with (
+        server_to_client_receive,
+        client_to_server_send,
+        ServerSession(client_to_server_receive, server_to_client_send, options) as session,
+    ):
+        yield session
+
+
+@pytest.mark.anyio
+async def test_configure_session_assigns_a_scope() -> None:
+    support = TaskSupport.in_memory()
+    async with _make_session() as session:
+        assert session.experimental.task_session_scope is None
+
+        support.configure_session(session)
+
+        assert session.experimental.task_session_scope is not None
+
+
+@pytest.mark.anyio
+async def test_configure_session_assigns_distinct_scopes_per_session() -> None:
+    support = TaskSupport.in_memory()
+    async with _make_session() as first, _make_session() as second:
+        support.configure_session(first)
+        support.configure_session(second)
+
+        assert first.experimental.task_session_scope != second.experimental.task_session_scope
+
+
+@pytest.mark.anyio
+async def test_configure_session_is_idempotent() -> None:
+    support = TaskSupport.in_memory()
+    async with _make_session() as session:
+        support.configure_session(session)
+        scope = session.experimental.task_session_scope
+        support.configure_session(session)
+
+        assert session.experimental.task_session_scope == scope
+
+
+@pytest.mark.anyio
+async def test_configure_session_assigns_no_scope_to_stateless_sessions() -> None:
+    support = TaskSupport.in_memory()
+    async with _make_session() as session:
+        support.configure_session(session, stateless=True)
+
+        assert session.experimental.task_session_scope is None