Support session-binding in TaskStore

Author: LucaButBoring

examples/shared/in_memory_task_store.py

@@ -41,7 +41,9 @@ def __init__(self) -> None:
         self._tasks: dict[str, StoredTask] = {}
         self._cleanup_tasks: dict[str, asyncio.Task[None]] = {}
 
-    async def create_task(self, task: TaskMetadata, request_id: RequestId, request: Request[Any, Any]) -> None:
+    async def create_task(
+        self, task: TaskMetadata, request_id: RequestId, request: Request[Any, Any], session_id: str | None = None
+    ) -> None:
         """Create a new task with the given metadata and original request."""
         task_id = task.taskId
 
@@ -61,7 +63,7 @@ async def create_task(self, task: TaskMetadata, request_id: RequestId, request:
         if task.keepAlive is not None:
             self._schedule_cleanup(task_id, task.keepAlive / 1000.0)
 
-    async def get_task(self, task_id: str) -> Task | None:
+    async def get_task(self, task_id: str, session_id: str | None = None) -> Task | None:
         """Get the current status of a task."""
         stored = self._tasks.get(task_id)
         if stored is None:
@@ -70,7 +72,7 @@ async def get_task(self, task_id: str) -> Task | None:
         # Return a copy to prevent external modification
         return Task(**stored.task.model_dump())
 
-    async def store_task_result(self, task_id: str, result: Result) -> None:
+    async def store_task_result(self, task_id: str, result: Result, session_id: str | None = None) -> None:
         """Store the result of a completed task."""
         stored = self._tasks.get(task_id)
         if stored is None:
@@ -84,7 +86,7 @@ async def store_task_result(self, task_id: str, result: Result) -> None:
             self._cancel_cleanup(task_id)
             self._schedule_cleanup(task_id, stored.task.keepAlive / 1000.0)
 
-    async def get_task_result(self, task_id: str) -> Result:
+    async def get_task_result(self, task_id: str, session_id: str | None = None) -> Result:
         """Retrieve the stored result of a task."""
         stored = self._tasks.get(task_id)
         if stored is None:
@@ -95,7 +97,9 @@ async def get_task_result(self, task_id: str) -> Result:
 
         return stored.result
 
-    async def update_task_status(self, task_id: str, status: TaskStatus, error: str | None = None) -> None:
+    async def update_task_status(
+        self, task_id: str, status: TaskStatus, error: str | None = None, session_id: str | None = None
+    ) -> None:
         """Update a task's status."""
         stored = self._tasks.get(task_id)
         if stored is None:
@@ -110,7 +114,7 @@ async def update_task_status(self, task_id: str, status: TaskStatus, error: str
             self._cancel_cleanup(task_id)
             self._schedule_cleanup(task_id, stored.task.keepAlive / 1000.0)
 
-    async def list_tasks(self, cursor: str | None = None) -> dict[str, Any]:
+    async def list_tasks(self, cursor: str | None = None, session_id: str | None = None) -> dict[str, Any]:
         """
         List tasks, optionally starting from a pagination cursor.
 
@@ -134,7 +138,7 @@ async def list_tasks(self, cursor: str | None = None) -> dict[str, Any]:
 
         return {"tasks": tasks, "nextCursor": next_cursor}
 
-    async def delete_task(self, task_id: str) -> None:
+    async def delete_task(self, task_id: str, session_id: str | None = None) -> None:
         """Delete a task from storage."""
         if task_id not in self._tasks:
             raise ValueError(f"Task with ID {task_id} not found")

examples/snippets/clients/streamable_task_get_result.py

@@ -0,0 +1,34 @@
+"""
+Run from the repository root:
+    uv run examples/snippets/clients/task_based_tool_client.py
+
+Prerequisites:
+    The task_based_tool server must be running on http://localhost:8000
+    Start it with:
+        cd examples/snippets && uv run server task_based_tool streamable-http
+"""
+
+import asyncio
+
+from mcp import ClientSession
+from mcp.client.streamable_http import MCP_SESSION_ID, streamablehttp_client
+from mcp.types import CallToolResult
+
+
+async def main():
+    async with streamablehttp_client(
+        "http://localhost:3000/mcp",
+        headers={MCP_SESSION_ID: "5771f709-66f5-4176-9f32-ce91e3117df2"},
+        terminate_on_close=False,
+    ) as (
+        read_stream,
+        write_stream,
+        _,
+    ):
+        async with ClientSession(read_stream, write_stream) as session:
+            result = await session.get_task_result("736054ac-5f10-409e-a06a-526761ea827a", CallToolResult)
+            print(result)
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

examples/snippets/clients/streamable_tasks.py

@@ -0,0 +1,122 @@
+"""
+Run from the repository root:
+    uv run examples/snippets/clients/task_based_tool_client.py
+
+Prerequisites:
+    The task_based_tool server must be running on http://localhost:8000
+    Start it with:
+        cd examples/snippets && uv run server task_based_tool streamable-http
+"""
+
+import asyncio
+
+from mcp import ClientSession, types
+from mcp.client.streamable_http import streamablehttp_client
+from mcp.shared.context import RequestContext
+from mcp.shared.request import TaskHandlerOptions
+
+
+async def elicitation_handler(
+    context: RequestContext[ClientSession, None], params: types.ElicitRequestParams
+) -> types.ElicitResult | types.ErrorData:
+    """
+    Handle elicitation requests from the server.
+
+    This handler collects user feedback with a predefined schema including:
+    - rating (1-5, required)
+    - comments (optional text up to 500 chars)
+    - recommend (boolean, required)
+    """
+    print(f"\n🎯 Elicitation request received: {params.message}")
+    print(f"Schema: {params.requestedSchema}")
+    await asyncio.sleep(5)
+
+    # In a real application, you would collect this data from the user
+    # For this example, we'll return mock data
+    feedback_data: dict[str, str | int | float | bool | None] = {
+        "rating": 5,
+        "comments": "The task execution was excellent and fast!",
+        "recommend": True,
+    }
+
+    print(f"📝 Returning feedback: {feedback_data}")
+
+    return types.ElicitResult(action="accept", content=feedback_data)
+
+
+async def main():
+    """
+    Demonstrate task-based execution with begin_call_tool.
+
+    This example shows how to:
+    1. Start a long-running tool call with begin_call_tool()
+    2. Get task status updates through callbacks
+    3. Wait for the final result with polling
+    4. Handle elicitation requests from the server
+    """
+    # Connect to the task-based tool example server via streamable HTTP
+    async with streamablehttp_client("http://localhost:3000/mcp", terminate_on_close=False) as (
+        read_stream,
+        write_stream,
+        _,
+    ):
+        async with ClientSession(read_stream, write_stream, elicitation_callback=elicitation_handler) as session:
+            # Initialize the connection
+            await session.initialize()
+
+            print("Starting task-based tool execution...")
+
+            # Track callback invocations
+            task_created = False
+            status_updates: list[str] = []
+
+            async def on_task_created() -> None:
+                """Called when the task is first created."""
+                nonlocal task_created
+                task_created = True
+                print("✓ Task created on server")
+
+            async def on_task_status(task_result: types.GetTaskResult) -> None:
+                """Called whenever the task status is polled."""
+                status_updates.append(task_result.status)
+                print(f"  Status ({task_result.taskId}): {task_result.status}")
+
+            # Begin the tool call (returns immediately with a PendingRequest)
+            print("\nCalling begin_call_tool...")
+            # pending_request = session.begin_call_tool(
+            #     "collect-user-info",
+            #     arguments={"infoType": "feedback"},
+            # )
+            pending_request = session.begin_call_tool(
+                "delay",
+                arguments={},
+            )
+
+            print("Tool call initiated! Now waiting for result with task polling...\n")
+
+            # Wait for the result with task callbacks
+            result = await pending_request.result(
+                TaskHandlerOptions(on_task_created=on_task_created, on_task_status=on_task_status)
+            )
+
+            # Display the result
+            print("\n✓ Tool execution completed!")
+            if result.content:
+                content_block = result.content[0]
+                if isinstance(content_block, types.TextContent):
+                    print(f"Result: {content_block.text}")
+                else:
+                    print(f"Result: {content_block}")
+            else:
+                print("Result: No content")
+
+            # Show callback statistics
+            print("\nTask callbacks:")
+            print(f"  - Task created callback: {'Yes' if task_created else 'No'}")
+            print(f"  - Status updates received: {len(status_updates)}")
+            if status_updates:
+                print(f"  - Final status: {status_updates[-1]}")
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

src/mcp/client/session.py

@@ -128,6 +128,7 @@ def __init__(
         message_handler: MessageHandlerFnT | None = None,
         client_info: types.Implementation | None = None,
         task_store: TaskStore | None = None,
+        session_id: str | None = None,
     ) -> None:
         super().__init__(
             read_stream,
@@ -136,6 +137,7 @@ def __init__(
             types.ServerNotification,
             read_timeout_seconds=read_timeout_seconds,
             task_store=task_store,
+            session_id=session_id,
         )
         self._client_info = client_info or DEFAULT_CLIENT_INFO
         self._sampling_callback = sampling_callback or _default_sampling_callback
@@ -605,6 +607,29 @@ async def delete_task(self, task_id: str) -> types.EmptyResult:
         )
 
     async def _received_request(self, responder: RequestResponder[types.ServerRequest, types.ClientResult]) -> None:
+        # Handle task creation if task metadata is present
+        if responder.request_meta and responder.request_meta.task and self._task_store:
+            task_meta = responder.request_meta.task
+            # Create the task in the task store
+            await self._task_store.create_task(
+                task_meta,
+                responder.request_id,
+                responder.request.root,
+                session_id=self._session_id,  # type: ignore[arg-type]
+            )
+            # Send task created notification with related task metadata
+            notification_params = types.TaskCreatedNotificationParams(
+                _meta=types.NotificationParams.Meta(
+                    **{types.RELATED_TASK_META_KEY: types.RelatedTaskMetadata(taskId=task_meta.taskId)}
+                )
+            )
+            await self.send_notification(
+                types.ClientNotification(
+                    types.TaskCreatedNotification(method="notifications/tasks/created", params=notification_params)
+                ),
+                related_request_id=responder.request_id,
+            )
+
         ctx = RequestContext[ClientSession, Any](
             request_id=responder.request_id,
             meta=responder.request_meta,
@@ -638,7 +663,7 @@ async def _received_request(self, responder: RequestResponder[types.ServerReques
             case types.GetTaskRequest(params=params):
                 # Handle get task requests if task store is available
                 if self._task_store:
-                    task = await self._task_store.get_task(params.taskId)
+                    task = await self._task_store.get_task(params.taskId, session_id=self._session_id)
                     if task is None:
                         with responder:
                             await responder.respond(
@@ -666,7 +691,7 @@ async def _received_request(self, responder: RequestResponder[types.ServerReques
             case types.GetTaskPayloadRequest(params=params):
                 # Handle get task result requests if task store is available
                 if self._task_store:
-                    task = await self._task_store.get_task(params.taskId)
+                    task = await self._task_store.get_task(params.taskId, session_id=self._session_id)
                     if task is None:
                         with responder:
                             await responder.respond(
@@ -683,7 +708,7 @@ async def _received_request(self, responder: RequestResponder[types.ServerReques
                                 )
                             )
                     else:
-                        result = await self._task_store.get_task_result(params.taskId)
+                        result = await self._task_store.get_task_result(params.taskId, session_id=self._session_id)
                         # Add related-task metadata
                         result_dict = result.model_dump(by_alias=True, mode="json", exclude_none=True)
                         if "_meta" not in result_dict:
@@ -701,7 +726,9 @@ async def _received_request(self, responder: RequestResponder[types.ServerReques
                 # Handle list tasks requests if task store is available
                 if self._task_store:
                     try:
-                        result = await self._task_store.list_tasks(params.cursor if params else None)
+                        result = await self._task_store.list_tasks(
+                            params.cursor if params else None, session_id=self._session_id
+                        )
                         with responder:
                             await responder.respond(
                                 types.ClientResult(

src/mcp/server/lowlevel/server.py

@@ -664,6 +664,8 @@ async def run(
         # the initialization lifecycle, but can do so with any available node
         # rather than requiring initialization for each connection.
         stateless: bool = False,
+        # Optional session identifier to pass to ServerSession for multi-session task stores
+        session_id: str | None = None,
     ):
         async with AsyncExitStack() as stack:
             lifespan_context = await stack.enter_async_context(self.lifespan(self))
@@ -674,6 +676,7 @@ async def run(
                     initialization_options,
                     stateless=stateless,
                     task_store=self.task_store,
+                    session_id=session_id,
                 )
             )