make ServerCallCOntext mandatory

Author: sokoliva

src/a2a/contrib/tasks/vertex_task_store.py

@@ -44,9 +44,7 @@ def __init__(
         self._client = client
         self._agent_engine_resource_id = agent_engine_resource_id
 
-    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."""
         compat_task = to_compat_task(task)
         previous_task = await self._get_stored_task(compat_task.id)
@@ -206,7 +204,7 @@ async def _get_stored_task(
         return a2a_task
 
     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."""
         a2a_task = await self._get_stored_task(task_id)
@@ -217,13 +215,11 @@ async def get(
     async def list(
         self,
         params: ListTasksRequest,
-        context: ServerCallContext | None = None,
+        context: ServerCallContext,
     ) -> ListTasksResponse:
         """Retrieves a list of tasks from the store."""
         raise NotImplementedError
 
-    async def delete(
-        self, task_id: str, context: ServerCallContext | None = None
-    ) -> None:
+    async def delete(self, task_id: str, context: ServerCallContext) -> None:
         """The backend doesn't support deleting tasks, so this is not implemented."""
         raise NotImplementedError

src/a2a/server/agent_execution/context.py

@@ -140,7 +140,7 @@ def configuration(self) -> SendMessageConfiguration | None:
         return self._params.configuration if self._params else None
 
     @property
-    def call_context(self) -> ServerCallContext | None:
+    def call_context(self) -> ServerCallContext:
         """The server call context associated with this request."""
         return self._call_context
 
@@ -157,22 +157,17 @@ def add_activated_extension(self, uri: str) -> None:
         This causes the extension to be indicated back to the client in the
         response.
         """
-        if self._call_context:
-            self._call_context.activated_extensions.add(uri)
+        self._call_context.activated_extensions.add(uri)
 
     @property
     def tenant(self) -> str:
         """The tenant associated with this request."""
-        return self._call_context.tenant if self._call_context else ''
+        return self._call_context.tenant
 
     @property
     def requested_extensions(self) -> set[str]:
         """Extensions that the client requested to activate."""
-        return (
-            self._call_context.requested_extensions
-            if self._call_context
-            else set()
-        )
+        return self._call_context.requested_extensions
 
     def _check_or_generate_task_id(self) -> None:
         """Ensures a task ID is present, generating one if necessary."""

src/a2a/server/request_handlers/default_request_handler.py

@@ -290,7 +290,7 @@ async def _setup_message_execution(
             await self._push_config_store.set_info(
                 task_id,
                 params.configuration.task_push_notification_config,
-                context or ServerCallContext(),
+                context,
             )
 
         queue = await self._queue_manager.create_or_tap(task_id)
@@ -504,7 +504,7 @@ async def on_create_task_push_notification_config(
         await self._push_config_store.set_info(
             task_id,
             params,
-            context or ServerCallContext(),
+            context,
         )
 
         return params
@@ -529,10 +529,7 @@ async def on_get_task_push_notification_config(
             raise TaskNotFoundError
 
         push_notification_configs: list[TaskPushNotificationConfig] = (
-            await self._push_config_store.get_info(
-                task_id, context or ServerCallContext()
-            )
-            or []
+            await self._push_config_store.get_info(task_id, context) or []
         )
 
         for config in push_notification_configs:
@@ -603,7 +600,7 @@ async def on_list_task_push_notification_configs(
             raise TaskNotFoundError
 
         push_notification_config_list = await self._push_config_store.get_info(
-            task_id, context or ServerCallContext()
+            task_id, context
         )
 
         return ListTaskPushNotificationConfigsResponse(
@@ -629,6 +626,4 @@ async def on_delete_task_push_notification_config(
         if not task:
             raise TaskNotFoundError
 
-        await self._push_config_store.delete_info(
-            task_id, context or ServerCallContext(), config_id
-        )
+        await self._push_config_store.delete_info(task_id, context, config_id)

tests/contrib/tasks/test_vertex_task_store.py

@@ -62,6 +62,7 @@ def backend_type(request) -> str:
 
 
 from a2a.contrib.tasks.vertex_task_store import VertexTaskStore
+from a2a.server.context import ServerCallContext
 from a2a.types.a2a_pb2 import (
     Artifact,
     Part,
@@ -140,9 +141,11 @@ async def test_save_task(vertex_store: VertexTaskStore) -> None:
     task_to_save = Task()
     task_to_save.CopyFrom(MINIMAL_TASK_OBJ)
     task_to_save.id = 'save-test-task-2'
-    await vertex_store.save(task_to_save)
+    await vertex_store.save(task_to_save, ServerCallContext())
 
-    retrieved_task = await vertex_store.get(task_to_save.id)
+    retrieved_task = await vertex_store.get(
+        task_to_save.id, ServerCallContext()
+    )
     assert retrieved_task is not None
     assert retrieved_task.id == task_to_save.id
 
@@ -156,9 +159,11 @@ async def test_get_task(vertex_store: VertexTaskStore) -> None:
     task_to_save = Task()
     task_to_save.CopyFrom(MINIMAL_TASK_OBJ)
     task_to_save.id = task_id
-    await vertex_store.save(task_to_save)
+    await vertex_store.save(task_to_save, ServerCallContext())
 
-    retrieved_task = await vertex_store.get(task_to_save.id)
+    retrieved_task = await vertex_store.get(
+        task_to_save.id, ServerCallContext()
+    )
     assert retrieved_task is not None
     assert retrieved_task.id == task_to_save.id
     assert retrieved_task.context_id == task_to_save.context_id
@@ -170,7 +175,9 @@ async def test_get_nonexistent_task(
     vertex_store: VertexTaskStore,
 ) -> None:
     """Test retrieving a nonexistent task."""
-    retrieved_task = await vertex_store.get('nonexistent-task-id')
+    retrieved_task = await vertex_store.get(
+        'nonexistent-task-id', ServerCallContext()
+    )
     assert retrieved_task is None
 
 
@@ -196,8 +203,8 @@ async def test_save_and_get_detailed_task(
     test_task.metadata['key1'] = 'value1'
     test_task.metadata['key2'] = 123
 
-    await vertex_store.save(test_task)
-    retrieved_task = await vertex_store.get(test_task.id)
+    await vertex_store.save(test_task, ServerCallContext())
+    retrieved_task = await vertex_store.get(test_task.id, ServerCallContext())
 
     assert retrieved_task is not None
     assert retrieved_task.id == test_task.id
@@ -221,9 +228,11 @@ async def test_update_task_status_and_metadata(
         artifacts=[],
         history=[],
     )
-    await vertex_store.save(original_task)
+    await vertex_store.save(original_task, ServerCallContext())
 
-    retrieved_before_update = await vertex_store.get(task_id)
+    retrieved_before_update = await vertex_store.get(
+        task_id, ServerCallContext()
+    )
     assert retrieved_before_update is not None
     assert (
         retrieved_before_update.status.state == TaskState.TASK_STATE_SUBMITTED
@@ -236,9 +245,11 @@ async def test_update_task_status_and_metadata(
     updated_task.status.timestamp.FromJsonString('2023-01-02T11:00:00Z')
     updated_task.metadata.update({'update_key': 'update_value'})
 
-    await vertex_store.save(updated_task)
+    await vertex_store.save(updated_task, ServerCallContext())
 
-    retrieved_after_update = await vertex_store.get(task_id)
+    retrieved_after_update = await vertex_store.get(
+        task_id, ServerCallContext()
+    )
     assert retrieved_after_update is not None
     assert retrieved_after_update.status.state == TaskState.TASK_STATE_COMPLETED
     assert retrieved_after_update.metadata == {'update_key': 'update_value'}
@@ -260,9 +271,11 @@ async def test_update_task_add_artifact(vertex_store: VertexTaskStore) -> None:
         ],
         history=[],
     )
-    await vertex_store.save(original_task)
+    await vertex_store.save(original_task, ServerCallContext())
 
-    retrieved_before_update = await vertex_store.get(task_id)
+    retrieved_before_update = await vertex_store.get(
+        task_id, ServerCallContext()
+    )
     assert retrieved_before_update is not None
     assert (
         retrieved_before_update.status.state == TaskState.TASK_STATE_SUBMITTED
@@ -281,9 +294,11 @@ async def test_update_task_add_artifact(vertex_store: VertexTaskStore) -> None:
         )
     )
 
-    await vertex_store.save(updated_task)
+    await vertex_store.save(updated_task, ServerCallContext())
 
-    retrieved_after_update = await vertex_store.get(task_id)
+    retrieved_after_update = await vertex_store.get(
+        task_id, ServerCallContext()
+    )
     assert retrieved_after_update is not None
     assert retrieved_after_update.status.state == TaskState.TASK_STATE_WORKING
 
@@ -321,9 +336,11 @@ async def test_update_task_update_artifact(
         ],
         history=[],
     )
-    await vertex_store.save(original_task)
+    await vertex_store.save(original_task, ServerCallContext())
 
-    retrieved_before_update = await vertex_store.get(task_id)
+    retrieved_before_update = await vertex_store.get(
+        task_id, ServerCallContext()
+    )
     assert retrieved_before_update is not None
     assert (
         retrieved_before_update.status.state == TaskState.TASK_STATE_SUBMITTED
@@ -337,9 +354,11 @@ async def test_update_task_update_artifact(
 
     updated_task.artifacts[0].parts[0].text = 'ahoy'
 
-    await vertex_store.save(updated_task)
+    await vertex_store.save(updated_task, ServerCallContext())
 
-    retrieved_after_update = await vertex_store.get(task_id)
+    retrieved_after_update = await vertex_store.get(
+        task_id, ServerCallContext()
+    )
     assert retrieved_after_update is not None
     assert retrieved_after_update.status.state == TaskState.TASK_STATE_WORKING
 
@@ -377,9 +396,11 @@ async def test_update_task_delete_artifact(
         ],
         history=[],
     )
-    await vertex_store.save(original_task)
+    await vertex_store.save(original_task, ServerCallContext())
 
-    retrieved_before_update = await vertex_store.get(task_id)
+    retrieved_before_update = await vertex_store.get(
+        task_id, ServerCallContext()
+    )
     assert retrieved_before_update is not None
     assert (
         retrieved_before_update.status.state == TaskState.TASK_STATE_SUBMITTED
@@ -393,9 +414,11 @@ async def test_update_task_delete_artifact(
 
     del updated_task.artifacts[1]
 
-    await vertex_store.save(updated_task)
+    await vertex_store.save(updated_task, ServerCallContext())
 
-    retrieved_after_update = await vertex_store.get(task_id)
+    retrieved_after_update = await vertex_store.get(
+        task_id, ServerCallContext()
+    )
     assert retrieved_after_update is not None
     assert retrieved_after_update.status.state == TaskState.TASK_STATE_WORKING
 
@@ -426,8 +449,10 @@ async def test_metadata_field_mapping(
         context_id='session-meta-1',
         status=TaskStatus(state=TaskState.TASK_STATE_SUBMITTED),
     )
-    await vertex_store.save(task_no_metadata)
-    retrieved_no_metadata = await vertex_store.get('task-metadata-test-1')
+    await vertex_store.save(task_no_metadata, ServerCallContext())
+    retrieved_no_metadata = await vertex_store.get(
+        'task-metadata-test-1', ServerCallContext()
+    )
     assert retrieved_no_metadata is not None
     assert retrieved_no_metadata.metadata == {}
 
@@ -439,8 +464,10 @@ async def test_metadata_field_mapping(
         status=TaskStatus(state=TaskState.TASK_STATE_SUBMITTED),
         metadata=simple_metadata,
     )
-    await vertex_store.save(task_simple_metadata)
-    retrieved_simple = await vertex_store.get('task-metadata-test-2')
+    await vertex_store.save(task_simple_metadata, ServerCallContext())
+    retrieved_simple = await vertex_store.get(
+        'task-metadata-test-2', ServerCallContext()
+    )
     assert retrieved_simple is not None
     assert retrieved_simple.metadata == simple_metadata
 
@@ -463,8 +490,10 @@ async def test_metadata_field_mapping(
         status=TaskStatus(state=TaskState.TASK_STATE_SUBMITTED),
         metadata=complex_metadata,
     )
-    await vertex_store.save(task_complex_metadata)
-    retrieved_complex = await vertex_store.get('task-metadata-test-3')
+    await vertex_store.save(task_complex_metadata, ServerCallContext())
+    retrieved_complex = await vertex_store.get(
+        'task-metadata-test-3', ServerCallContext()
+    )
     assert retrieved_complex is not None
     assert retrieved_complex.metadata == complex_metadata
 
@@ -474,16 +503,18 @@ async def test_metadata_field_mapping(
         context_id='session-meta-4',
         status=TaskStatus(state=TaskState.TASK_STATE_SUBMITTED),
     )
-    await vertex_store.save(task_update_metadata)
+    await vertex_store.save(task_update_metadata, ServerCallContext())
 
     # Update metadata
     task_update_metadata.metadata.Clear()
     task_update_metadata.metadata.update(
         {'updated': True, 'timestamp': '2024-01-01'}
     )
-    await vertex_store.save(task_update_metadata)
+    await vertex_store.save(task_update_metadata, ServerCallContext())
 
-    retrieved_updated = await vertex_store.get('task-metadata-test-4')
+    retrieved_updated = await vertex_store.get(
+        'task-metadata-test-4', ServerCallContext()
+    )
     assert retrieved_updated is not None
     assert retrieved_updated.metadata == {
         'updated': True,
@@ -492,8 +523,10 @@ async def test_metadata_field_mapping(
 
     # Test 5: Update metadata from dict to None
     task_update_metadata.metadata.Clear()
-    await vertex_store.save(task_update_metadata)
+    await vertex_store.save(task_update_metadata, ServerCallContext())
 
-    retrieved_none = await vertex_store.get('task-metadata-test-4')
+    retrieved_none = await vertex_store.get(
+        'task-metadata-test-4', ServerCallContext()
+    )
     assert retrieved_none is not None
     assert retrieved_none.metadata == {}