[v1.x] Scope experimental tasks to the session that created them (#2720)
Some checks failed
Main branch checks / checks (push) Has been cancelled
Some checks failed
Main branch checks / checks (push) Has been cancelled
This commit is contained in:
parent
ce267b6fc5
commit
62137874ff
4 changed files with 731 additions and 14 deletions
|
|
@ -506,13 +506,14 @@ async def test_default_task_handlers_via_enable_tasks() -> None:
|
|||
# Create a task directly in the store for testing
|
||||
task = await store.create_task(TaskMetadata(ttl=60000))
|
||||
|
||||
# Test list_tasks (default handler)
|
||||
# Test list_tasks (default handler). Tasks created directly in the
|
||||
# store have no session scope, so they are reachable by ID but not
|
||||
# included in tasks/list (see test_task_scope.py).
|
||||
list_result = await client_session.send_request(
|
||||
ClientRequest(ListTasksRequest()),
|
||||
ListTasksResult,
|
||||
)
|
||||
assert len(list_result.tasks) == 1
|
||||
assert list_result.tasks[0].taskId == task.taskId
|
||||
assert list_result.tasks == []
|
||||
|
||||
# Test get_task (default handler - found)
|
||||
get_result = await client_session.send_request(
|
||||
|
|
|
|||
150
tests/experimental/tasks/server/test_task_scope.py
Normal file
150
tests/experimental/tasks/server/test_task_scope.py
Normal file
|
|
@ -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
|
||||
321
tests/experimental/tasks/server/test_task_visibility.py
Normal file
321
tests/experimental/tasks/server/test_task_visibility.py
Normal file
|
|
@ -0,0 +1,321 @@
|
|||
"""End-to-end tests for which clients can see and control a task.
|
||||
|
||||
Every test runs a real server and one or more in-memory client sessions. A
|
||||
task started with run_task() belongs to the client session that started it:
|
||||
that session can poll it, list it, and cancel it, while every other session
|
||||
is told the task does not exist. Tasks whose IDs carry no session marker
|
||||
(explicitly chosen IDs, or tasks on stateless servers) are usable by any
|
||||
session that knows the ID, but are never listed.
|
||||
"""
|
||||
|
||||
from collections.abc import AsyncIterator, Awaitable, Callable
|
||||
from contextlib import AsyncExitStack
|
||||
from typing import Any
|
||||
|
||||
import anyio
|
||||
import pytest
|
||||
from anyio.abc import TaskGroup
|
||||
|
||||
from mcp.client.session import ClientSession
|
||||
from mcp.server import Server
|
||||
from mcp.server.experimental.task_context import ServerTaskContext
|
||||
from mcp.shared.exceptions import McpError
|
||||
from mcp.shared.experimental.tasks.in_memory_task_store import InMemoryTaskStore
|
||||
from mcp.shared.experimental.tasks.store import TaskStore
|
||||
from mcp.shared.message import SessionMessage
|
||||
from mcp.types import (
|
||||
TASK_REQUIRED,
|
||||
CallToolResult,
|
||||
CreateTaskResult,
|
||||
ListTasksResult,
|
||||
TextContent,
|
||||
Tool,
|
||||
ToolExecution,
|
||||
)
|
||||
|
||||
# The `connect` fixture: each call opens a new client session against the test server.
|
||||
Connect = Callable[..., Awaitable[ClientSession]]
|
||||
|
||||
# Enough tasks that the bundled in-memory store needs more than one page (of 10)
|
||||
# to list them, so listings that span store pages are exercised.
|
||||
MORE_TASKS_THAN_ONE_STORE_PAGE = 11
|
||||
|
||||
|
||||
def build_task_server(store: TaskStore | None = None) -> Server:
|
||||
"""Build a server exposing three task tools.
|
||||
|
||||
- "greet" finishes immediately and returns a greeting.
|
||||
- "long_running_job" keeps running until the server shuts down.
|
||||
- "nightly_export" is a singleton job: every invocation uses the
|
||||
explicitly chosen task ID "the-nightly-export".
|
||||
"""
|
||||
server = Server("task-visibility-test-server")
|
||||
server.experimental.enable_tasks(store=store)
|
||||
|
||||
@server.list_tools()
|
||||
async def list_tools() -> list[Tool]:
|
||||
return [
|
||||
Tool(
|
||||
name=name,
|
||||
description=name,
|
||||
inputSchema={"type": "object"},
|
||||
execution=ToolExecution(taskSupport=TASK_REQUIRED),
|
||||
)
|
||||
for name in ("greet", "long_running_job", "nightly_export")
|
||||
]
|
||||
|
||||
@server.call_tool()
|
||||
async def call_tool(name: str, arguments: dict[str, Any]) -> CallToolResult | CreateTaskResult:
|
||||
async def greet(task: ServerTaskContext) -> CallToolResult:
|
||||
return CallToolResult(content=[TextContent(type="text", text=f"Hello, {arguments['name']}!")])
|
||||
|
||||
async def long_running_job(task: ServerTaskContext) -> CallToolResult:
|
||||
await anyio.sleep_forever()
|
||||
raise AssertionError("unreachable") # pragma: no cover
|
||||
|
||||
async def nightly_export(task: ServerTaskContext) -> CallToolResult:
|
||||
return CallToolResult(content=[TextContent(type="text", text="exported")])
|
||||
|
||||
run_task = server.request_context.experimental.run_task
|
||||
if name == "nightly_export":
|
||||
return await run_task(nightly_export, task_id="the-nightly-export")
|
||||
return await run_task(greet if name == "greet" else long_running_job)
|
||||
|
||||
return server
|
||||
|
||||
|
||||
async def open_client(
|
||||
server: Server, task_group: TaskGroup, stack: AsyncExitStack, *, stateless: bool = False
|
||||
) -> ClientSession:
|
||||
"""Connect a new client session to `server` over in-memory streams."""
|
||||
server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](10)
|
||||
client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](10)
|
||||
|
||||
async def run_server() -> None:
|
||||
await server.run(
|
||||
client_to_server_receive,
|
||||
server_to_client_send,
|
||||
server.create_initialization_options(),
|
||||
stateless=stateless,
|
||||
)
|
||||
|
||||
task_group.start_soon(run_server)
|
||||
client = await stack.enter_async_context(ClientSession(server_to_client_receive, client_to_server_send))
|
||||
await client.initialize()
|
||||
return client
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def task_server() -> Server:
|
||||
return build_task_server()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def connect(task_server: Server) -> AsyncIterator[Connect]:
|
||||
"""A factory that opens a new client session against the test server on each call."""
|
||||
async with anyio.create_task_group() as task_group, AsyncExitStack() as stack:
|
||||
|
||||
async def _connect(*, stateless: bool = False) -> ClientSession:
|
||||
return await open_client(task_server, task_group, stack, stateless=stateless)
|
||||
|
||||
yield _connect
|
||||
task_group.cancel_scope.cancel()
|
||||
|
||||
|
||||
async def start_task(client: ClientSession, tool: str = "long_running_job", **arguments: Any) -> str:
|
||||
"""Start `tool` as a task and return the new task's ID."""
|
||||
result = await client.experimental.call_tool_as_task(tool, arguments)
|
||||
return result.task.taskId
|
||||
|
||||
|
||||
async def wait_until_finished(client: ClientSession, task_id: str) -> None:
|
||||
"""Poll the task until it reaches a terminal status."""
|
||||
with anyio.fail_after(5):
|
||||
async for _ in client.experimental.poll_task(task_id):
|
||||
pass
|
||||
|
||||
|
||||
async def listed_task_ids(client: ClientSession) -> list[str]:
|
||||
"""Return the IDs of every task the server lists for this client."""
|
||||
return [task.taskId for task in (await client.experimental.list_tasks()).tasks]
|
||||
|
||||
|
||||
# --- What the client that started a task can do with it ---
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_a_client_can_poll_its_own_task_to_completion_and_read_the_result(connect: Connect) -> None:
|
||||
client = await connect()
|
||||
task_id = await start_task(client, "greet", name="Ada")
|
||||
await wait_until_finished(client, task_id)
|
||||
|
||||
result = await client.experimental.get_task_result(task_id, CallToolResult)
|
||||
|
||||
assert result.content == [TextContent(type="text", text="Hello, Ada!")]
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_a_client_sees_its_own_task_when_listing_tasks(connect: Connect) -> None:
|
||||
client = await connect()
|
||||
task_id = await start_task(client)
|
||||
|
||||
listed = await listed_task_ids(client)
|
||||
|
||||
assert listed == [task_id]
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_a_client_can_cancel_its_own_task(connect: Connect) -> None:
|
||||
client = await connect()
|
||||
task_id = await start_task(client)
|
||||
|
||||
cancelled = await client.experimental.cancel_task(task_id)
|
||||
|
||||
assert cancelled.status == "cancelled"
|
||||
|
||||
|
||||
# --- What a client cannot do with a task started by another client ---
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_a_client_cannot_get_the_status_of_another_clients_task(connect: Connect) -> None:
|
||||
creator = await connect()
|
||||
other_client = await connect()
|
||||
task_id = await start_task(creator)
|
||||
|
||||
with pytest.raises(McpError, match="Task not found"):
|
||||
await other_client.experimental.get_task(task_id)
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_a_client_cannot_get_the_result_of_another_clients_task(connect: Connect) -> None:
|
||||
creator = await connect()
|
||||
other_client = await connect()
|
||||
task_id = await start_task(creator, "greet", name="Ada")
|
||||
await wait_until_finished(creator, task_id)
|
||||
|
||||
with pytest.raises(McpError, match="Task not found"):
|
||||
await other_client.experimental.get_task_result(task_id, CallToolResult)
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_a_client_cannot_cancel_another_clients_task(connect: Connect) -> None:
|
||||
creator = await connect()
|
||||
other_client = await connect()
|
||||
task_id = await start_task(creator)
|
||||
|
||||
with pytest.raises(McpError, match="Task not found"):
|
||||
await other_client.experimental.cancel_task(task_id)
|
||||
|
||||
# The task is unaffected.
|
||||
assert (await creator.experimental.get_task(task_id)).status == "working"
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_a_client_does_not_see_another_clients_task_when_listing_tasks(connect: Connect) -> None:
|
||||
creator = await connect()
|
||||
other_client = await connect()
|
||||
await start_task(creator)
|
||||
|
||||
listed = await listed_task_ids(other_client)
|
||||
|
||||
assert listed == []
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_each_client_lists_only_its_own_tasks(connect: Connect) -> None:
|
||||
first_client = await connect()
|
||||
second_client = await connect()
|
||||
first_task = await start_task(first_client)
|
||||
second_task = await start_task(second_client)
|
||||
|
||||
assert await listed_task_ids(first_client) == [first_task]
|
||||
assert await listed_task_ids(second_client) == [second_task]
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_listing_tasks_reveals_nothing_about_other_clients_tasks_however_many_there_are(
|
||||
connect: Connect,
|
||||
) -> None:
|
||||
"""The listing must not identify other clients' tasks through any field, including the pagination cursor."""
|
||||
creator = await connect()
|
||||
other_client = await connect()
|
||||
for _ in range(MORE_TASKS_THAN_ONE_STORE_PAGE):
|
||||
await start_task(creator)
|
||||
|
||||
listing = await other_client.experimental.list_tasks()
|
||||
|
||||
assert listing == ListTasksResult(tasks=[], nextCursor=None)
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_a_client_with_more_than_one_store_page_of_tasks_lists_all_of_them(connect: Connect) -> None:
|
||||
client = await connect()
|
||||
started = {await start_task(client) for _ in range(MORE_TASKS_THAN_ONE_STORE_PAGE)}
|
||||
|
||||
listing = await client.experimental.list_tasks()
|
||||
|
||||
assert {task.taskId for task in listing.tasks} == started
|
||||
assert listing.nextCursor is None
|
||||
|
||||
|
||||
# --- Tasks that do not belong to any client session ---
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
# Choosing the task ID instead of letting the SDK generate one is deprecated for
|
||||
# exactly the behaviour this test demonstrates: the task is not tied to the
|
||||
# session that created it.
|
||||
@pytest.mark.filterwarnings("ignore:Passing an explicit task_id")
|
||||
async def test_a_task_whose_id_was_chosen_by_the_server_is_accessible_to_every_client(connect: Connect) -> None:
|
||||
creator = await connect()
|
||||
other_client = await connect()
|
||||
await wait_until_finished(creator, await start_task(creator, "nightly_export"))
|
||||
|
||||
status = await other_client.experimental.get_task("the-nightly-export")
|
||||
|
||||
assert status.status == "completed"
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_a_stateless_server_serves_a_task_to_any_session_that_knows_its_id(connect: Connect) -> None:
|
||||
first_session = await connect(stateless=True)
|
||||
second_session = await connect(stateless=True)
|
||||
task_id = await start_task(first_session, "greet", name="Ada")
|
||||
await wait_until_finished(second_session, task_id)
|
||||
|
||||
result = await second_session.experimental.get_task_result(task_id, CallToolResult)
|
||||
|
||||
assert result.content == [TextContent(type="text", text="Hello, Ada!")]
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_a_stateless_server_lists_no_tasks(connect: Connect) -> None:
|
||||
session = await connect(stateless=True)
|
||||
await start_task(session)
|
||||
|
||||
listed = await listed_task_ids(session)
|
||||
|
||||
assert listed == []
|
||||
|
||||
|
||||
# --- The behaviour does not depend on the bundled in-memory store ---
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_clients_are_isolated_when_the_server_uses_a_custom_task_store() -> None:
|
||||
class CustomTaskStore(InMemoryTaskStore):
|
||||
"""A stand-in for a user-provided TaskStore implementation."""
|
||||
|
||||
server = build_task_server(store=CustomTaskStore())
|
||||
|
||||
async with anyio.create_task_group() as task_group, AsyncExitStack() as stack:
|
||||
creator = await open_client(server, task_group, stack)
|
||||
other_client = await open_client(server, task_group, stack)
|
||||
task_id = await start_task(creator)
|
||||
|
||||
with pytest.raises(McpError, match="Task not found"):
|
||||
await other_client.experimental.get_task(task_id)
|
||||
|
||||
assert (await creator.experimental.get_task(task_id)).status == "working"
|
||||
task_group.cancel_scope.cancel()
|
||||
|
|
@ -3,6 +3,7 @@
|
|||
import pytest
|
||||
|
||||
from mcp.server.experimental.request_context import Experimental
|
||||
from mcp.server.experimental.task_context import ServerTaskContext
|
||||
from mcp.shared.exceptions import McpError
|
||||
from mcp.types import (
|
||||
METHOD_NOT_FOUND,
|
||||
|
|
@ -11,6 +12,7 @@ from mcp.types import (
|
|||
TASK_REQUIRED,
|
||||
ClientCapabilities,
|
||||
ClientTasksCapability,
|
||||
Result,
|
||||
TaskMetadata,
|
||||
Tool,
|
||||
ToolExecution,
|
||||
|
|
@ -164,3 +166,30 @@ def test_can_use_tool_forbidden_without_task_support() -> None:
|
|||
def test_can_use_tool_none_without_task_support() -> None:
|
||||
exp = Experimental(_client_capabilities=ClientCapabilities())
|
||||
assert exp.can_use_tool(None) is True
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_run_task_with_an_explicit_task_id_emits_a_deprecation_warning() -> None:
|
||||
"""An explicitly provided task ID is not associated with the creating session, so passing one is deprecated."""
|
||||
exp = Experimental(task_metadata=TaskMetadata(ttl=60000))
|
||||
|
||||
async def work(task: ServerTaskContext) -> Result:
|
||||
raise AssertionError("unreachable") # pragma: no cover
|
||||
|
||||
with pytest.warns(DeprecationWarning, match="not associated with the session"):
|
||||
# Task support is not configured, so the call fails after the
|
||||
# deprecated argument has been reported.
|
||||
with pytest.raises(RuntimeError, match="Task support not enabled"):
|
||||
# The deliberate use of the deprecated overload is the point of this test.
|
||||
await exp.run_task(work, task_id="explicitly-chosen") # pyright: ignore[reportDeprecated]
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_run_task_without_a_task_id_does_not_warn() -> None:
|
||||
exp = Experimental(task_metadata=TaskMetadata(ttl=60000))
|
||||
|
||||
async def work(task: ServerTaskContext) -> Result:
|
||||
raise AssertionError("unreachable") # pragma: no cover
|
||||
|
||||
with pytest.raises(RuntimeError, match="Task support not enabled"):
|
||||
await exp.run_task(work)
|
||||
|
|
|
|||
Loading…
Reference in a new issue