Make ServerCallContext mandatory in CopyingTaskStore

Author: sokoliva

src/a2a/server/tasks/copying_task_store.py

@@ -25,15 +25,15 @@ def __init__(self, underlying_store: TaskStore):
         self._store = underlying_store
 
     async def save(
-        self, task: Task, context: ServerCallContext | None = None
+        self, task: Task, context: ServerCallContext
     ) -> None:
         """Saves a copy of the task to the underlying store."""
         task_copy = Task()
         task_copy.CopyFrom(task)
         await self._store.save(task_copy, context)
 
     async def get(
-        self, task_id: str, context: ServerCallContext | None = None
+        self, task_id: str, context: ServerCallContext
     ) -> Task | None:
         """Retrieves a task from the underlying store and returns a copy."""
         task = await self._store.get(task_id, context)
@@ -46,7 +46,7 @@ async def get(
     async def list(
         self,
         params: ListTasksRequest,
-        context: ServerCallContext | None = None,
+        context: ServerCallContext,
     ) -> ListTasksResponse:
         """Retrieves a list of tasks from the underlying store and returns a copy."""
         response = await self._store.list(params, context)
@@ -55,7 +55,7 @@ async def list(
         return response_copy
 
     async def delete(
-        self, task_id: str, context: ServerCallContext | None = None
+        self, task_id: str, context: ServerCallContext
     ) -> None:
         """Deletes a task from the underlying store."""
         await self._store.delete(task_id, context)

src/a2a/server/tasks/inmemory_task_store.py

@@ -208,27 +208,27 @@ def __init__(
         )
 
     async def save(
-        self, task: Task, context: ServerCallContext | None = None
+        self, task: Task, context: ServerCallContext
     ) -> None:
         """Saves or updates a task in the store."""
         await self._store.save(task, context)
 
     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."""
         return await self._store.get(task_id, context)
 
     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."""
         return await self._store.list(params, context)
 
     async def delete(
-        self, task_id: str, context: ServerCallContext | None = None
+        self, task_id: str, context: ServerCallContext
     ) -> None:
         """Deletes a task from the store by ID."""
         await self._store.delete(task_id, context)

tests/server/tasks/test_inmemory_task_store.py

@@ -344,10 +344,10 @@ async def test_inmemory_task_store_copying_behavior(use_copying: bool):
     original_task = Task(
         id='test_task', status=TaskStatus(state=TaskState.TASK_STATE_WORKING)
     )
-    await store.save(original_task)
+    await store.save(original_task, TEST_CONTEXT)
 
     # Retrieve it
-    retrieved_task = await store.get('test_task')
+    retrieved_task = await store.get('test_task', TEST_CONTEXT)
     assert retrieved_task is not None
 
     if use_copying:
@@ -359,7 +359,7 @@ async def test_inmemory_task_store_copying_behavior(use_copying: bool):
     retrieved_task.status.state = TaskState.TASK_STATE_COMPLETED
 
     # Retrieve it again, it should NOT be modified in the store if use_copying=True
-    retrieved_task_2 = await store.get('test_task')
+    retrieved_task_2 = await store.get('test_task', TEST_CONTEXT)
     assert retrieved_task_2 is not None
 
     if use_copying: