Address feedback

🏠 Remote-Dev: homespace

Author: localden

src/mcp/client/_memory.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import uuid
 from collections.abc import AsyncIterator
 from contextlib import AbstractAsyncContextManager, asynccontextmanager
 from types import TracebackType
@@ -50,12 +51,14 @@ async def _connect(self) -> AsyncIterator[TransportStreams]:
 
             async with anyio.create_task_group() as tg:
                 # Start server in background
+                memory_session_id = uuid.uuid4().hex
                 tg.start_soon(
                     lambda: actual_server.run(
                         server_read,
                         server_write,
                         actual_server.create_initialization_options(),
                         raise_exceptions=self._raise_exceptions,
+                        session_id=memory_session_id,
                     )
                 )
 

src/mcp/server/experimental/request_context.py

@@ -188,6 +188,8 @@ async def work(task: ServerTaskContext) -> CallToolResult:
         task_group = support.task_group
 
         session_id = self._session.session_id
+        if session_id is None:
+            raise RuntimeError("Session ID is required for task operations but session has no ID.")
         task = await support.store.create_task(self.task_metadata, task_id, session_id=session_id)
 
         task_ctx = ServerTaskContext(

src/mcp/server/experimental/task_context.py

@@ -88,7 +88,11 @@ def __init__(
             queue: The message queue for elicitation/sampling
             handler: The result handler for response routing (required for elicit/create_message)
         """
-        self._ctx = TaskContext(task=task, store=store)
+        session_id = session.session_id
+        if session_id is None:
+            raise RuntimeError("Session ID is required for task operations but session has no ID.")
+        self._session_id = session_id
+        self._ctx = TaskContext(task=task, store=store, session_id=session_id)
         self._session = session
         self._queue = queue
         self._handler = handler
@@ -210,7 +214,7 @@ async def elicit(
             raise RuntimeError("handler is required for elicit(). Pass handler= to ServerTaskContext.")
 
         # Update status to input_required
-        await self._store.update_task(self.task_id, status=TASK_STATUS_INPUT_REQUIRED)
+        await self._store.update_task(self.task_id, status=TASK_STATUS_INPUT_REQUIRED, session_id=self._session_id)
 
         # Build the request using session's helper
         request = self._session._build_elicit_form_request(  # pyright: ignore[reportPrivateUsage]
@@ -234,12 +238,12 @@ async def elicit(
         try:
             # Wait for response (routed back via TaskResultHandler)
             response_data = await resolver.wait()
-            await self._store.update_task(self.task_id, status=TASK_STATUS_WORKING)
+            await self._store.update_task(self.task_id, status=TASK_STATUS_WORKING, session_id=self._session_id)
             return ElicitResult.model_validate(response_data)
         except anyio.get_cancelled_exc_class():
             # This path is tested in test_elicit_restores_status_on_cancellation
             # which verifies status is restored to "working" after cancellation.
-            await self._store.update_task(self.task_id, status=TASK_STATUS_WORKING)
+            await self._store.update_task(self.task_id, status=TASK_STATUS_WORKING, session_id=self._session_id)
             raise
 
     async def elicit_url(
@@ -279,7 +283,7 @@ async def elicit_url(
             raise RuntimeError("handler is required for elicit_url(). Pass handler= to ServerTaskContext.")
 
         # Update status to input_required
-        await self._store.update_task(self.task_id, status=TASK_STATUS_INPUT_REQUIRED)
+        await self._store.update_task(self.task_id, status=TASK_STATUS_INPUT_REQUIRED, session_id=self._session_id)
 
         # Build the request using session's helper
         request = self._session._build_elicit_url_request(  # pyright: ignore[reportPrivateUsage]
@@ -304,10 +308,10 @@ async def elicit_url(
         try:
             # Wait for response (routed back via TaskResultHandler)
             response_data = await resolver.wait()
-            await self._store.update_task(self.task_id, status=TASK_STATUS_WORKING)
+            await self._store.update_task(self.task_id, status=TASK_STATUS_WORKING, session_id=self._session_id)
             return ElicitResult.model_validate(response_data)
         except anyio.get_cancelled_exc_class():  # pragma: no cover
-            await self._store.update_task(self.task_id, status=TASK_STATUS_WORKING)
+            await self._store.update_task(self.task_id, status=TASK_STATUS_WORKING, session_id=self._session_id)
             raise
 
     async def create_message(
@@ -362,7 +366,7 @@ async def create_message(
             raise RuntimeError("handler is required for create_message(). Pass handler= to ServerTaskContext.")
 
         # Update status to input_required
-        await self._store.update_task(self.task_id, status=TASK_STATUS_INPUT_REQUIRED)
+        await self._store.update_task(self.task_id, status=TASK_STATUS_INPUT_REQUIRED, session_id=self._session_id)
 
         # Build the request using session's helper
         request = self._session._build_create_message_request(  # pyright: ignore[reportPrivateUsage]
@@ -394,12 +398,12 @@ async def create_message(
         try:
             # Wait for response (routed back via TaskResultHandler)
             response_data = await resolver.wait()
-            await self._store.update_task(self.task_id, status=TASK_STATUS_WORKING)
+            await self._store.update_task(self.task_id, status=TASK_STATUS_WORKING, session_id=self._session_id)
             return CreateMessageResult.model_validate(response_data)
         except anyio.get_cancelled_exc_class():
             # This path is tested in test_create_message_restores_status_on_cancellation
             # which verifies status is restored to "working" after cancellation.
-            await self._store.update_task(self.task_id, status=TASK_STATUS_WORKING)
+            await self._store.update_task(self.task_id, status=TASK_STATUS_WORKING, session_id=self._session_id)
             raise
 
     async def elicit_as_task(
@@ -435,7 +439,7 @@ async def elicit_as_task(
             raise RuntimeError("handler is required for elicit_as_task()")
 
         # Update status to input_required
-        await self._store.update_task(self.task_id, status=TASK_STATUS_INPUT_REQUIRED)
+        await self._store.update_task(self.task_id, status=TASK_STATUS_INPUT_REQUIRED, session_id=self._session_id)
 
         request = self._session._build_elicit_form_request(  # pyright: ignore[reportPrivateUsage]
             message=message,
@@ -472,11 +476,11 @@ async def elicit_as_task(
                 ElicitResult,
             )
 
-            await self._store.update_task(self.task_id, status=TASK_STATUS_WORKING)
+            await self._store.update_task(self.task_id, status=TASK_STATUS_WORKING, session_id=self._session_id)
             return result
 
         except anyio.get_cancelled_exc_class():  # pragma: no cover
-            await self._store.update_task(self.task_id, status=TASK_STATUS_WORKING)
+            await self._store.update_task(self.task_id, status=TASK_STATUS_WORKING, session_id=self._session_id)
             raise
 
     async def create_message_as_task(
@@ -531,7 +535,7 @@ async def create_message_as_task(
             raise RuntimeError("handler is required for create_message_as_task()")
 
         # Update status to input_required
-        await self._store.update_task(self.task_id, status=TASK_STATUS_INPUT_REQUIRED)
+        await self._store.update_task(self.task_id, status=TASK_STATUS_INPUT_REQUIRED, session_id=self._session_id)
 
         # Build request WITH task field for task-augmented sampling
         request = self._session._build_create_message_request(  # pyright: ignore[reportPrivateUsage]
@@ -577,9 +581,9 @@ async def create_message_as_task(
                 CreateMessageResult,
             )
 
-            await self._store.update_task(self.task_id, status=TASK_STATUS_WORKING)
+            await self._store.update_task(self.task_id, status=TASK_STATUS_WORKING, session_id=self._session_id)
             return result
 
         except anyio.get_cancelled_exc_class():  # pragma: no cover
-            await self._store.update_task(self.task_id, status=TASK_STATUS_WORKING)
+            await self._store.update_task(self.task_id, status=TASK_STATUS_WORKING, session_id=self._session_id)
             raise

src/mcp/server/experimental/task_result_handler.py

@@ -80,7 +80,7 @@ async def handle(
         request: GetTaskPayloadRequest,
         session: ServerSession,
         request_id: RequestId,
-        session_id: str | None = None,
+        session_id: str,
     ) -> GetTaskPayloadResult:
         """Handle a tasks/result request.
 
@@ -95,7 +95,7 @@ async def handle(
             request: The GetTaskPayloadRequest
             session: The server session for sending messages
             request_id: The request ID for relatedRequestId routing
-            session_id: Optional session identifier for access control.
+            session_id: Session identifier for access control.
 
         Returns:
             GetTaskPayloadResult with the task's final payload

src/mcp/server/lowlevel/experimental.py

@@ -147,13 +147,22 @@ def enable_tasks(
         if on_cancel_task is not None:
             self._add_request_handler("tasks/cancel", on_cancel_task)
 
+        def _require_session_id(ctx: ServerRequestContext[LifespanResultT]) -> str:
+            session_id = ctx.session.session_id
+            if session_id is None:
+                raise MCPError(
+                    code=INVALID_PARAMS,
+                    message="Session ID is required for task operations.",
+                )
+            return session_id
+
         # Fill in defaults for any not provided
         if not self._has_handler("tasks/get"):
 
             async def _default_get_task(
                 ctx: ServerRequestContext[LifespanResultT], params: GetTaskRequestParams
             ) -> GetTaskResult:
-                session_id = ctx.session.session_id
+                session_id = _require_session_id(ctx)
                 task = await task_support.store.get_task(params.task_id, session_id=session_id)
                 if task is None:
                     raise MCPError(code=INVALID_PARAMS, message=f"Task not found: {params.task_id}")
@@ -175,7 +184,7 @@ async def _default_get_task_result(
                 ctx: ServerRequestContext[LifespanResultT], params: GetTaskPayloadRequestParams
             ) -> GetTaskPayloadResult:
                 assert ctx.request_id is not None
-                session_id = ctx.session.session_id
+                session_id = _require_session_id(ctx)
                 req = GetTaskPayloadRequest(params=params)
                 result = await task_support.handler.handle(req, ctx.session, ctx.request_id, session_id=session_id)
                 return result
@@ -188,7 +197,7 @@ async def _default_list_tasks(
                 ctx: ServerRequestContext[LifespanResultT], params: PaginatedRequestParams | None
             ) -> ListTasksResult:
                 cursor = params.cursor if params else None
-                session_id = ctx.session.session_id
+                session_id = _require_session_id(ctx)
                 tasks, next_cursor = await task_support.store.list_tasks(cursor, session_id=session_id)
                 return ListTasksResult(tasks=tasks, next_cursor=next_cursor)
 
@@ -199,7 +208,7 @@ async def _default_list_tasks(
             async def _default_cancel_task(
                 ctx: ServerRequestContext[LifespanResultT], params: CancelTaskRequestParams
             ) -> CancelTaskResult:
-                session_id = ctx.session.session_id
+                session_id = _require_session_id(ctx)
                 result = await cancel_task(task_support.store, params.task_id, session_id=session_id)
                 return result
 

src/mcp/server/lowlevel/server.py

@@ -374,8 +374,6 @@ 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 for task isolation. When provided (e.g.,
-        # from the transport's mcp_session_id), tasks are bound to this ID.
         session_id: str | None = None,
     ):
         async with AsyncExitStack() as stack:

src/mcp/server/session.py

@@ -83,6 +83,7 @@ def __init__(
         write_stream: MemoryObjectSendStream[SessionMessage],
         init_options: InitializationOptions,
         stateless: bool = False,
+        *,
         session_id: str | None = None,
     ) -> None:
         super().__init__(read_stream, write_stream)

src/mcp/shared/experimental/tasks/context.py

@@ -21,19 +21,20 @@ class TaskContext:
     use ServerTaskContext from mcp.server.experimental.
 
     Example (distributed worker):
-        async def worker_job(task_id: str):
+        async def worker_job(task_id: str, session_id: str):
             store = RedisTaskStore(redis_url)
-            task = await store.get_task(task_id)
-            ctx = TaskContext(task=task, store=store)
+            task = await store.get_task(task_id, session_id=session_id)
+            ctx = TaskContext(task=task, store=store, session_id=session_id)
 
             await ctx.update_status("Working...")
             result = await do_work()
             await ctx.complete(result)
     """
 
-    def __init__(self, task: Task, store: TaskStore):
+    def __init__(self, task: Task, store: TaskStore, *, session_id: str):
         self._task = task
         self._store = store
+        self._session_id = session_id
         self._cancelled = False
 
     @property
@@ -68,6 +69,7 @@ async def update_status(self, message: str) -> None:
         self._task = await self._store.update_task(
             self.task_id,
             status_message=message,
+            session_id=self._session_id,
         )
 
     async def complete(self, result: Result) -> None:
@@ -76,10 +78,11 @@ async def complete(self, result: Result) -> None:
         Args:
             result: The task result
         """
-        await self._store.store_result(self.task_id, result)
+        await self._store.store_result(self.task_id, result, session_id=self._session_id)
         self._task = await self._store.update_task(
             self.task_id,
             status=TASK_STATUS_COMPLETED,
+            session_id=self._session_id,
         )
 
     async def fail(self, error: str) -> None:
@@ -92,4 +95,5 @@ async def fail(self, error: str) -> None:
             self.task_id,
             status=TASK_STATUS_FAILED,
             status_message=error,
+            session_id=self._session_id,
         )

src/mcp/shared/experimental/tasks/helpers.py

@@ -50,7 +50,8 @@ def is_terminal(status: TaskStatus) -> bool:
 async def cancel_task(
     store: TaskStore,
     task_id: str,
-    session_id: str | None = None,
+    *,
+    session_id: str,
 ) -> CancelTaskResult:
     """Cancel a task with spec-compliant validation.
 
@@ -63,7 +64,7 @@ async def cancel_task(
     Args:
         store: The task store
         task_id: The task identifier to cancel
-        session_id: Optional session identifier for access control.
+        session_id: Session identifier for access control.
 
     Returns:
         CancelTaskResult with the cancelled task state
@@ -75,7 +76,7 @@ async def cancel_task(
 
     Example:
         async def handle_cancel(ctx, params: CancelTaskRequestParams) -> CancelTaskResult:
-            return await cancel_task(store, params.task_id)
+            return await cancel_task(store, params.task_id, session_id=ctx.session.session_id)
     """
     task = await store.get_task(task_id, session_id=session_id)
     if task is None:
@@ -124,6 +125,8 @@ def create_task_state(
 async def task_execution(
     task_id: str,
     store: TaskStore,
+    *,
+    session_id: str,
 ) -> AsyncIterator[TaskContext]:
     """Context manager for safe task execution (pure, no server dependencies).
 
@@ -136,6 +139,7 @@ async def task_execution(
     Args:
         task_id: The task identifier to execute
         store: The task store (must be accessible by the worker)
+        session_id: Session identifier for access control.
 
     Yields:
         TaskContext for updating status and completing/failing the task
@@ -144,18 +148,18 @@ async def task_execution(
         ValueError: If the task is not found in the store
 
     Example (distributed worker):
-        async def worker_process(task_id: str):
+        async def worker_process(task_id: str, session_id: str):
             store = RedisTaskStore(redis_url)
-            async with task_execution(task_id, store) as ctx:
+            async with task_execution(task_id, store, session_id=session_id) as ctx:
                 await ctx.update_status("Working...")
                 result = await do_work()
                 await ctx.complete(result)
     """
-    task = await store.get_task(task_id)
+    task = await store.get_task(task_id, session_id=session_id)
     if task is None:
         raise ValueError(f"Task {task_id} not found")
 
-    ctx = TaskContext(task, store)
+    ctx = TaskContext(task, store, session_id=session_id)
     try:
         yield ctx
     except Exception as e: