Add missing server to client task tests

Author: LucaButBoring

tests/client/test_server_to_client_tasks.py

@@ -0,0 +1,261 @@
+"""Tests for server-to-client task requests.
+
+These tests verify that the client can properly handle task-related requests
+initiated by the server, which is part of bidirectional task-based execution.
+"""
+
+from typing import Any
+
+import pytest
+
+import mcp.types as types
+from examples.shared.in_memory_task_store import InMemoryTaskStore
+from mcp.server import Server
+from mcp.shared.memory import create_connected_server_and_client_session
+
+# Mark all tests in this module to ignore memory stream cleanup warnings
+pytestmark = pytest.mark.filterwarnings(
+    "ignore:Exception ignored.*MemoryObject.*Stream:pytest.PytestUnraisableExceptionWarning"
+)
+
+
+@pytest.mark.anyio
+async def test_server_queries_client_get_task_success():
+    """Test server using session.get_task() to query client's task."""
+    server_task_store = InMemoryTaskStore()
+    client_task_store = InMemoryTaskStore()
+    server = Server("test", task_store=server_task_store)
+
+    # Create a task in the CLIENT's store
+    task_id = "client-task-123"
+    task_meta = types.TaskMetadata(taskId=task_id, keepAlive=60000)
+    request = types.ClientRequest(types.PingRequest())
+    await client_task_store.create_task(task_meta, "req-1", request.root)
+
+    # Register list_tools to advertise our tool
+    @server.list_tools()
+    async def list_tools() -> list[types.Tool]:
+        return [
+            types.Tool(
+                name="query_client_task",
+                description="Query a task from the client",
+                inputSchema={"type": "object", "properties": {}, "required": []},
+            )
+        ]
+
+    # Create a tool that will use the server session to query the client
+    @server.call_tool()
+    async def query_client_task(name: str, arguments: dict[str, Any]) -> list[types.TextContent]:
+        # Server session queries client's task via request_context
+        session = server.request_context.session
+        result = await session.get_task(task_id)
+        return [types.TextContent(type="text", text=f"Task {result.taskId} status: {result.status}")]
+
+    async with create_connected_server_and_client_session(
+        server,
+        # Client needs task store to handle the GetTaskRequest
+    ) as client_session:
+        # Set the client's task store
+        client_session._task_store = client_task_store
+
+        await client_session.initialize()
+
+        # Call the tool which will trigger server->client task query
+        result = await client_session.call_tool("query_client_task", {})
+
+        assert len(result.content) == 1
+        assert isinstance(result.content[0], types.TextContent)
+        assert "client-task-123" in result.content[0].text
+        assert "submitted" in result.content[0].text
+
+
+@pytest.mark.anyio
+async def test_server_queries_client_get_task_not_found():
+    """Test server querying client for non-existent task."""
+    server_task_store = InMemoryTaskStore()
+    client_task_store = InMemoryTaskStore()
+    server = Server("test", task_store=server_task_store)
+
+    @server.list_tools()
+    async def list_tools() -> list[types.Tool]:
+        return [
+            types.Tool(
+                name="query_nonexistent_task",
+                description="Query nonexistent task",
+                inputSchema={"type": "object", "properties": {}, "required": []},
+            )
+        ]
+
+    @server.call_tool()
+    async def query_nonexistent_task(name: str, arguments: dict[str, Any]) -> list[types.TextContent]:
+        try:
+            session = server.request_context.session
+            await session.get_task("nonexistent-task")
+            return [types.TextContent(type="text", text="Should have failed")]
+        except Exception as e:
+            return [types.TextContent(type="text", text=f"Error: {str(e)}")]
+
+    async with create_connected_server_and_client_session(server) as client_session:
+        client_session._task_store = client_task_store
+        await client_session.initialize()
+
+        result = await client_session.call_tool("query_nonexistent_task", {})
+
+        assert len(result.content) == 1
+        assert isinstance(result.content[0], types.TextContent)
+        assert "Error" in result.content[0].text
+        assert "Task not found" in result.content[0].text or "INVALID_PARAMS" in result.content[0].text
+
+
+@pytest.mark.anyio
+async def test_server_queries_client_without_task_store():
+    """Test server querying client without task store configured."""
+    server = Server("test", task_store=InMemoryTaskStore())
+
+    @server.list_tools()
+    async def list_tools() -> list[types.Tool]:
+        return [
+            types.Tool(
+                name="query_without_store",
+                description="Query without task store",
+                inputSchema={"type": "object", "properties": {}, "required": []},
+            )
+        ]
+
+    @server.call_tool()
+    async def query_without_store(name: str, arguments: dict[str, Any]) -> list[types.TextContent]:
+        try:
+            session = server.request_context.session
+            await session.get_task("some-task")
+            return [types.TextContent(type="text", text="Should have failed")]
+        except Exception as e:
+            return [types.TextContent(type="text", text=f"Error: {str(e)}")]
+
+    async with create_connected_server_and_client_session(server) as client_session:
+        # No task store on client
+        await client_session.initialize()
+
+        result = await client_session.call_tool("query_without_store", {})
+
+        assert len(result.content) == 1
+        assert isinstance(result.content[0], types.TextContent)
+        assert "Error" in result.content[0].text
+        assert "Task store not configured" in result.content[0].text or "INVALID_REQUEST" in result.content[0].text
+
+
+@pytest.mark.anyio
+async def test_server_queries_client_get_task_result_success():
+    """Test server querying client for completed task result."""
+    server_task_store = InMemoryTaskStore()
+    client_task_store = InMemoryTaskStore()
+    server = Server("test", task_store=server_task_store)
+
+    # Create a completed task with result in CLIENT's store
+    # This is a server-to-client request (like sampling), so result is ClientResult
+    task_id = "client-completed-task"
+    task_meta = types.TaskMetadata(taskId=task_id)
+    request = types.ServerRequest(types.PingRequest())
+    await client_task_store.create_task(task_meta, "req-1", request.root)
+    result = types.ClientResult(types.EmptyResult())
+    await client_task_store.store_task_result(task_id, result.root)
+    await client_task_store.update_task_status(task_id, "completed")
+
+    @server.list_tools()
+    async def list_tools() -> list[types.Tool]:
+        return [
+            types.Tool(
+                name="query_client_result",
+                description="Query client task result",
+                inputSchema={"type": "object", "properties": {}, "required": []},
+            )
+        ]
+
+    @server.call_tool()
+    async def query_client_result(name: str, arguments: dict[str, Any]) -> list[types.TextContent]:
+        session = server.request_context.session
+        result = await session.get_task_result(task_id, types.ClientResult)
+        return [types.TextContent(type="text", text=f"Got result: {type(result.root).__name__}")]  # type: ignore[attr-defined]
+
+    async with create_connected_server_and_client_session(server) as client_session:
+        client_session._task_store = client_task_store
+        await client_session.initialize()
+
+        result = await client_session.call_tool("query_client_result", {})
+
+        assert len(result.content) == 1
+        assert isinstance(result.content[0], types.TextContent)
+        assert "EmptyResult" in result.content[0].text
+
+
+@pytest.mark.anyio
+async def test_server_queries_client_list_tasks():
+    """Test server querying client's task list."""
+    server_task_store = InMemoryTaskStore()
+    client_task_store = InMemoryTaskStore()
+    server = Server("test", task_store=server_task_store)
+
+    # Create multiple tasks in CLIENT's store
+    for i in range(3):
+        task_id = f"client-task-{i}"
+        task_meta = types.TaskMetadata(taskId=task_id, keepAlive=60000)
+        request = types.ClientRequest(types.PingRequest())
+        await client_task_store.create_task(task_meta, f"req-{i}", request.root)
+
+    @server.list_tools()
+    async def list_tools() -> list[types.Tool]:
+        return [
+            types.Tool(
+                name="list_client_tasks",
+                description="List client tasks",
+                inputSchema={"type": "object", "properties": {}, "required": []},
+            )
+        ]
+
+    @server.call_tool()
+    async def list_client_tasks(name: str, arguments: dict[str, Any]) -> list[types.TextContent]:
+        session = server.request_context.session
+        result = await session.list_tasks()
+        return [types.TextContent(type="text", text=f"Found {len(result.tasks)} tasks")]
+
+    async with create_connected_server_and_client_session(server) as client_session:
+        client_session._task_store = client_task_store
+        await client_session.initialize()
+
+        result = await client_session.call_tool("list_client_tasks", {})
+
+        assert len(result.content) == 1
+        assert isinstance(result.content[0], types.TextContent)
+        assert "Found 3 tasks" in result.content[0].text
+
+
+@pytest.mark.anyio
+async def test_server_list_tasks_empty():
+    """Test server querying client with empty task list."""
+    server = Server("test", task_store=InMemoryTaskStore())
+    client_task_store = InMemoryTaskStore()
+
+    @server.list_tools()
+    async def list_tools() -> list[types.Tool]:
+        return [
+            types.Tool(
+                name="list_empty_tasks",
+                description="List empty tasks",
+                inputSchema={"type": "object", "properties": {}, "required": []},
+            )
+        ]
+
+    @server.call_tool()
+    async def list_empty_tasks(name: str, arguments: dict[str, Any]) -> list[types.TextContent]:
+        session = server.request_context.session
+        result = await session.list_tasks()
+        return [types.TextContent(type="text", text=f"Found {len(result.tasks)} tasks")]
+
+    async with create_connected_server_and_client_session(server) as client_session:
+        client_session._task_store = client_task_store
+        await client_session.initialize()
+
+        result = await client_session.call_tool("list_empty_tasks", {})
+
+        assert len(result.content) == 1
+        assert isinstance(result.content[0], types.TextContent)
+        assert "Found 0 tasks" in result.content[0].text