Add more edge case tests for StC task flows

Author: LucaButBoring

tests/client/test_server_to_client_tasks.py

@@ -228,6 +228,89 @@ async def list_client_tasks(name: str, arguments: dict[str, Any]) -> list[types.
         assert "Found 3 tasks" in result.content[0].text
 
 
+@pytest.mark.anyio
+async def test_server_queries_client_get_task_payload_not_found():
+    """Test server querying client for non-existent task payload."""
+    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_payload",
+                description="Query nonexistent task payload",
+                inputSchema={"type": "object", "properties": {}, "required": []},
+            )
+        ]
+
+    @server.call_tool()
+    async def query_nonexistent_payload(name: str, arguments: dict[str, Any]) -> list[types.TextContent]:
+        try:
+            session = server.request_context.session
+            await session.get_task_result("nonexistent-task", types.ClientResult)
+            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_payload", {})
+
+        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_get_task_payload_not_completed():
+    """Test server querying client for task payload when task is not completed."""
+    server_task_store = InMemoryTaskStore()
+    client_task_store = InMemoryTaskStore()
+    server = Server("test", task_store=server_task_store)
+
+    # Create a task that is NOT completed
+    task_id = "client-pending-task"
+    task_meta = types.TaskMetadata(taskId=task_id, keepAlive=60000)
+    request = types.ServerRequest(types.PingRequest())
+    await client_task_store.create_task(task_meta, "req-1", request.root)
+    await client_task_store.update_task_status(task_id, "submitted")
+
+    @server.list_tools()
+    async def list_tools() -> list[types.Tool]:
+        return [
+            types.Tool(
+                name="query_incomplete_payload",
+                description="Query incomplete task payload",
+                inputSchema={"type": "object", "properties": {}, "required": []},
+            )
+        ]
+
+    @server.call_tool()
+    async def query_incomplete_payload(name: str, arguments: dict[str, Any]) -> list[types.TextContent]:
+        try:
+            session = server.request_context.session
+            await session.get_task_result(task_id, types.ClientResult)
+            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_incomplete_payload", {})
+
+        assert len(result.content) == 1
+        assert isinstance(result.content[0], types.TextContent)
+        assert "Error" in result.content[0].text
+        assert "not 'completed'" in result.content[0].text or "INVALID_PARAMS" in result.content[0].text
+
+
 @pytest.mark.anyio
 async def test_server_list_tasks_empty():
     """Test server querying client with empty task list."""