Fix tests

🏠 Remote-Dev: homespace

Author: localden

tests/experimental/tasks/server/test_run_task_flow.py

@@ -31,6 +31,7 @@
     GetTaskResult,
     ListToolsResult,
     PaginatedRequestParams,
+    TaskMetadata,
     TextContent,
 )
 
@@ -365,3 +366,26 @@ async def work(task: ServerTaskContext) -> CallToolResult:
                     break
 
         assert status.status_message == "Manually failed"
+
+
+@pytest.mark.anyio
+async def test_run_task_without_session_id_raises() -> None:
+    """Test that run_task raises when session has no session_id."""
+    task_support = TaskSupport.in_memory()
+
+    mock_session = Mock()
+    mock_session.session_id = None
+
+    experimental = Experimental(
+        task_metadata=TaskMetadata(ttl=60000),
+        _client_capabilities=None,
+        _session=mock_session,
+        _task_support=task_support,
+    )
+
+    async def work(task: ServerTaskContext) -> CallToolResult:
+        raise NotImplementedError
+
+    async with task_support.run():
+        with pytest.raises(RuntimeError, match="Session ID is required"):
+            await experimental.run_task(work)

tests/experimental/tasks/server/test_server.py

@@ -399,6 +399,53 @@ async def run_server() -> None:
             tg.cancel_scope.cancel()
 
 
+@pytest.mark.anyio
+async def test_default_task_handlers_require_session_id() -> None:
+    """Test that default task handlers reject requests when session has no session_id."""
+    server = Server("test-no-session-id")
+    server.experimental.enable_tasks()
+
+    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 message_handler(
+        message: RequestResponder[ServerRequest, ClientResult] | ServerNotification | Exception,
+    ) -> None: ...  # pragma: no branch
+
+    async def run_server() -> None:
+        async with ServerSession(
+            client_to_server_receive,
+            server_to_client_send,
+            InitializationOptions(
+                server_name="test-server",
+                server_version="1.0.0",
+                capabilities=server.get_capabilities(
+                    notification_options=NotificationOptions(),
+                    experimental_capabilities={},
+                ),
+            ),
+        ) as server_session:
+            # session_id is None (no session_id passed)
+            async for message in server_session.incoming_messages:
+                await server._handle_message(message, server_session, {}, False)
+
+    async with anyio.create_task_group() as tg:
+        tg.start_soon(run_server)
+
+        async with ClientSession(
+            server_to_client_receive,
+            client_to_server_send,
+            message_handler=message_handler,
+        ) as client_session:
+            await client_session.initialize()
+
+            # All default task handlers should fail with "Session ID is required"
+            with pytest.raises(MCPError, match="Session ID is required"):
+                await client_session.experimental.list_tasks()
+
+            tg.cancel_scope.cancel()
+
+
 @pytest.mark.anyio
 async def test_build_elicit_form_request() -> None:
     """Test that _build_elicit_form_request builds a proper elicitation request."""

tests/experimental/tasks/server/test_server_task_context.py

@@ -725,3 +725,19 @@ async def test_create_message_as_task_raises_without_handler() -> None:
         )
 
     store.cleanup()
+
+
+@pytest.mark.anyio
+async def test_server_task_context_requires_session_id() -> None:
+    """Test that ServerTaskContext raises when session has no session_id."""
+    store = InMemoryTaskStore()
+    queue = InMemoryTaskMessageQueue()
+    task = await store.create_task(TaskMetadata(ttl=60000), session_id="test-session")
+
+    mock_session = Mock()
+    mock_session.session_id = None
+
+    with pytest.raises(RuntimeError, match="Session ID is required for task operations"):
+        ServerTaskContext(task=task, store=store, session=mock_session, queue=queue)
+
+    store.cleanup()