Add ServerCallContext into task store operations

Author: pstephengoogle

src/a2a/server/request_handlers/default_request_handler.py

@@ -109,7 +109,7 @@ async def on_get_task(
         context: ServerCallContext | None = None,
     ) -> Task | None:
         """Default handler for 'tasks/get'."""
-        task: Task | None = await self.task_store.get(params.id)
+        task: Task | None = await self.task_store.get(params.id, context)
         if not task:
             raise ServerError(error=TaskNotFoundError())
 
@@ -141,7 +141,7 @@ async def on_cancel_task(
 
         Attempts to cancel the task managed by the `AgentExecutor`.
         """
-        task: Task | None = await self.task_store.get(params.id)
+        task: Task | None = await self.task_store.get(params.id, context)
         if not task:
             raise ServerError(error=TaskNotFoundError())
 
@@ -158,6 +158,7 @@ async def on_cancel_task(
             context_id=task.context_id,
             task_store=self.task_store,
             initial_message=None,
+            context=context,
         )
         result_aggregator = ResultAggregator(task_manager)
 
@@ -217,6 +218,7 @@ async def _setup_message_execution(
             context_id=params.message.context_id,
             task_store=self.task_store,
             initial_message=params.message,
+            context=context,
         )
         task: Task | None = await task_manager.get_task()
 
@@ -417,7 +419,7 @@ async def on_set_task_push_notification_config(
         if not self._push_config_store:
             raise ServerError(error=UnsupportedOperationError())
 
-        task: Task | None = await self.task_store.get(params.task_id)
+        task: Task | None = await self.task_store.get(params.task_id, context)
         if not task:
             raise ServerError(error=TaskNotFoundError())
 
@@ -440,7 +442,7 @@ async def on_get_task_push_notification_config(
         if not self._push_config_store:
             raise ServerError(error=UnsupportedOperationError())
 
-        task: Task | None = await self.task_store.get(params.id)
+        task: Task | None = await self.task_store.get(params.id, context)
         if not task:
             raise ServerError(error=TaskNotFoundError())
 
@@ -469,7 +471,7 @@ async def on_resubscribe_to_task(
         Allows a client to re-attach to a running streaming task's event stream.
         Requires the task and its queue to still be active.
         """
-        task: Task | None = await self.task_store.get(params.id)
+        task: Task | None = await self.task_store.get(params.id, context)
         if not task:
             raise ServerError(error=TaskNotFoundError())
 
@@ -485,6 +487,7 @@ async def on_resubscribe_to_task(
             context_id=task.context_id,
             task_store=self.task_store,
             initial_message=None,
+            context=context,
         )
 
         result_aggregator = ResultAggregator(task_manager)
@@ -509,7 +512,7 @@ async def on_list_task_push_notification_config(
         if not self._push_config_store:
             raise ServerError(error=UnsupportedOperationError())
 
-        task: Task | None = await self.task_store.get(params.id)
+        task: Task | None = await self.task_store.get(params.id, context)
         if not task:
             raise ServerError(error=TaskNotFoundError())
 
@@ -536,7 +539,7 @@ async def on_delete_task_push_notification_config(
         if not self._push_config_store:
             raise ServerError(error=UnsupportedOperationError())
 
-        task: Task | None = await self.task_store.get(params.id)
+        task: Task | None = await self.task_store.get(params.id, context)
         if not task:
             raise ServerError(error=TaskNotFoundError())
 

src/a2a/server/tasks/database_task_store.py

@@ -19,6 +19,7 @@
         "or 'pip install a2a-sdk[sql]'"
     ) from e
 
+from a2a.server.context import ServerCallContext
 from a2a.server.models import Base, TaskModel, create_task_model
 from a2a.server.tasks.task_store import TaskStore
 from a2a.types import Task  # Task is the Pydantic model
@@ -119,15 +120,19 @@ def _from_orm(self, task_model: TaskModel) -> Task:
         # Pydantic's model_validate will parse the nested dicts/lists from JSON
         return Task.model_validate(task_data_from_db)
 
-    async def save(self, task: Task) -> None:
+    async def save(
+        self, task: Task, context: ServerCallContext | None = None
+    ) -> None:
         """Saves or updates a task in the database."""
         await self._ensure_initialized()
         db_task = self._to_orm(task)
         async with self.async_session_maker.begin() as session:
             await session.merge(db_task)
             logger.debug('Task %s saved/updated successfully.', task.id)
 
-    async def get(self, task_id: str) -> Task | None:
+    async def get(
+        self, task_id: str, context: ServerCallContext | None = None
+    ) -> Task | None:
         """Retrieves a task from the database by ID."""
         await self._ensure_initialized()
         async with self.async_session_maker() as session:
@@ -142,7 +147,9 @@ async def get(self, task_id: str) -> Task | None:
             logger.debug('Task %s not found in store.', task_id)
             return None
 
-    async def delete(self, task_id: str) -> None:
+    async def delete(
+        self, task_id: str, context: ServerCallContext | None = None
+    ) -> None:
         """Deletes a task from the database by ID."""
         await self._ensure_initialized()
 

src/a2a/server/tasks/inmemory_task_store.py

@@ -1,6 +1,7 @@
 import asyncio
 import logging
 
+from a2a.server.context import ServerCallContext
 from a2a.server.tasks.task_store import TaskStore
 from a2a.types import Task
 
@@ -21,13 +22,17 @@ def __init__(self) -> None:
         self.tasks: dict[str, Task] = {}
         self.lock = asyncio.Lock()
 
-    async def save(self, task: Task) -> None:
+    async def save(
+        self, task: Task, context: ServerCallContext | None = None
+    ) -> None:
         """Saves or updates a task in the in-memory store."""
         async with self.lock:
             self.tasks[task.id] = task
             logger.debug('Task %s saved successfully.', task.id)
 
-    async def get(self, task_id: str) -> Task | None:
+    async def get(
+        self, task_id: str, context: ServerCallContext | None = None
+    ) -> Task | None:
         """Retrieves a task from the in-memory store by ID."""
         async with self.lock:
             logger.debug('Attempting to get task with id: %s', task_id)
@@ -38,7 +43,9 @@ async def get(self, task_id: str) -> Task | None:
                 logger.debug('Task %s not found in store.', task_id)
             return task
 
-    async def delete(self, task_id: str) -> None:
+    async def delete(
+        self, task_id: str, context: ServerCallContext | None = None
+    ) -> None:
         """Deletes a task from the in-memory store by ID."""
         async with self.lock:
             logger.debug('Attempting to delete task with id: %s', task_id)

src/a2a/server/tasks/task_manager.py

@@ -1,5 +1,6 @@
 import logging
 
+from a2a.server.context import ServerCallContext
 from a2a.server.events.event_queue import Event
 from a2a.server.tasks.task_store import TaskStore
 from a2a.types import (
@@ -31,6 +32,7 @@ def __init__(
         context_id: str | None,
         task_store: TaskStore,
         initial_message: Message | None,
+        context: ServerCallContext | None = None,
     ):
         """Initializes the TaskManager.
 
@@ -49,6 +51,7 @@ def __init__(
         self.task_store = task_store
         self._initial_message = initial_message
         self._current_task: Task | None = None
+        self._call_context: ServerCallContext | None = context
         logger.debug(
             'TaskManager initialized with task_id: %s, context_id: %s',
             task_id,
@@ -74,7 +77,9 @@ async def get_task(self) -> Task | None:
         logger.debug(
             'Attempting to get task from store with id: %s', self.task_id
         )
-        self._current_task = await self.task_store.get(self.task_id)
+        self._current_task = await self.task_store.get(
+            self.task_id, self._context
+        )
         if self._current_task:
             logger.debug('Task %s retrieved successfully.', self.task_id)
         else:
@@ -167,7 +172,7 @@ async def ensure_task(
             logger.debug(
                 'Attempting to retrieve existing task with id: %s', self.task_id
             )
-            task = await self.task_store.get(self.task_id)
+            task = await self.task_store.get(self.task_id, self._context)
 
         if not task:
             logger.info(
@@ -231,7 +236,7 @@ async def _save_task(self, task: Task) -> None:
             task: The `Task` object to save.
         """
         logger.debug('Saving task with id: %s', task.id)
-        await self.task_store.save(task)
+        await self.task_store.save(task, self._context)
         self._current_task = task
         if not self.task_id:
             logger.info('New task created with id: %s', task.id)

src/a2a/server/tasks/task_store.py

@@ -1,22 +1,34 @@
 from abc import ABC, abstractmethod
 
+from a2a.server.context import ServerCallContext
 from a2a.types import Task
 
-
 class TaskStore(ABC):
     """Agent Task Store interface.
 
     Defines the methods for persisting and retrieving `Task` objects.
     """
 
     @abstractmethod
-    async def save(self, task: Task) -> None:
+    async def save(
+        self,
+        task: Task,
+        context: ServerCallContext | None = None
+    ) -> None:
         """Saves or updates a task in the store."""
 
     @abstractmethod
-    async def get(self, task_id: str) -> Task | None:
+    async def get(
+        self,
+        task_id: str,
+        context: ServerCallContext | None = None
+    ) -> Task | None:
         """Retrieves a task from the store by ID."""
 
     @abstractmethod
-    async def delete(self, task_id: str) -> None:
+    async def delete(
+        self,
+        task_id: str,
+        context: ServerCallContext | None = None
+    ) -> None:
         """Deletes a task from the store by ID."""