feat: Enforce ServerCallContext in all TaskStore operations

Author: sokoliva

src/a2a/server/owner_resolver.py

@@ -4,13 +4,10 @@
 
 
 # Definition
-OwnerResolver = Callable[[ServerCallContext | None], str]
+OwnerResolver = Callable[[ServerCallContext], str]
 
 
 # Example Default Implementation
-def resolve_user_scope(context: ServerCallContext | None) -> str:
+def resolve_user_scope(context: ServerCallContext) -> str:
     """Resolves the owner scope based on the user in the context."""
-    if not context:
-        return 'unknown'
-    # Example: Basic user name. Adapt as needed for your user model.
     return context.user.user_name

src/a2a/server/tasks/database_task_store.py

@@ -169,9 +169,7 @@ def _from_orm(self, task_model: TaskModel) -> Task:
         # Legacy conversion
         return compat_task_model_to_core(task_model)
 
-    async def save(
-        self, task: Task, context: ServerCallContext | None = None
-    ) -> None:
+    async def save(self, task: Task, context: ServerCallContext) -> None:
         """Saves or updates a task in the database for the resolved owner."""
         await self._ensure_initialized()
         owner = self.owner_resolver(context)
@@ -185,7 +183,7 @@ async def save(
             )
 
     async def get(
-        self, task_id: str, context: ServerCallContext | None = None
+        self, task_id: str, context: ServerCallContext
     ) -> Task | None:
         """Retrieves a task from the database by ID, for the given owner."""
         await self._ensure_initialized()
@@ -216,7 +214,7 @@ async def get(
     async def list(
         self,
         params: a2a_pb2.ListTasksRequest,
-        context: ServerCallContext | None = None,
+        context: ServerCallContext,
     ) -> a2a_pb2.ListTasksResponse:
         """Retrieves tasks from the database based on provided parameters, for the given owner."""
         await self._ensure_initialized()
@@ -315,9 +313,7 @@ async def list(
                 page_size=page_size,
             )
 
-    async def delete(
-        self, task_id: str, context: ServerCallContext | None = None
-    ) -> None:
+    async def delete(self, task_id: str, context: ServerCallContext) -> None:
         """Deletes a task from the database by ID, for the given owner."""
         await self._ensure_initialized()
         owner = self.owner_resolver(context)

src/a2a/server/tasks/inmemory_task_store.py

@@ -34,9 +34,7 @@ def __init__(
     def _get_owner_tasks(self, owner: str) -> dict[str, Task]:
         return self.tasks.get(owner, {})
 
-    async def save(
-        self, task: Task, context: ServerCallContext | None = None
-    ) -> None:
+    async def save(self, task: Task, context: ServerCallContext) -> None:
         """Saves or updates a task in the in-memory store for the resolved owner."""
         owner = self.owner_resolver(context)
         if owner not in self.tasks:
@@ -49,7 +47,7 @@ async def save(
             )
 
     async def get(
-        self, task_id: str, context: ServerCallContext | None = None
+        self, task_id: str, context: ServerCallContext
     ) -> Task | None:
         """Retrieves a task from the in-memory store by ID, for the given owner."""
         owner = self.owner_resolver(context)
@@ -76,7 +74,7 @@ async def get(
     async def list(
         self,
         params: a2a_pb2.ListTasksRequest,
-        context: ServerCallContext | None = None,
+        context: ServerCallContext,
     ) -> a2a_pb2.ListTasksResponse:
         """Retrieves a list of tasks from the store, for the given owner."""
         owner = self.owner_resolver(context)
@@ -155,9 +153,7 @@ async def list(
             page_size=page_size,
         )
 
-    async def delete(
-        self, task_id: str, context: ServerCallContext | None = None
-    ) -> None:
+    async def delete(self, task_id: str, context: ServerCallContext) -> None:
         """Deletes a task from the in-memory store by ID, for the given owner."""
         owner = self.owner_resolver(context)
         async with self.lock:

src/a2a/server/tasks/task_store.py

@@ -11,27 +11,23 @@ class TaskStore(ABC):
     """
 
     @abstractmethod
-    async def save(
-        self, task: Task, context: ServerCallContext | None = None
-    ) -> None:
+    async def save(self, task: Task, context: ServerCallContext) -> None:
         """Saves or updates a task in the store."""
 
     @abstractmethod
     async def get(
-        self, task_id: str, context: ServerCallContext | None = None
+        self, task_id: str, context: ServerCallContext
     ) -> Task | None:
         """Retrieves a task from the store by ID."""
 
     @abstractmethod
     async def list(
         self,
         params: ListTasksRequest,
-        context: ServerCallContext | None = None,
+        context: ServerCallContext,
     ) -> ListTasksResponse:
         """Retrieves a list of tasks from the store."""
 
     @abstractmethod
-    async def delete(
-        self, task_id: str, context: ServerCallContext | None = None
-    ) -> None:
+    async def delete(self, task_id: str, context: ServerCallContext) -> None:
         """Deletes a task from the store by ID."""