Merge branch 'main' into doc/helper-utils

Author: holtskinner

.github/CODEOWNERS

@@ -4,4 +4,5 @@
 # For syntax help see:
 # https://help.github.com/en/github/creating-cloning-and-archiving-repositories/about-code-owners#codeowners-syntax
 
-*   @a2aproject/google-a2a-eng
+*                   @a2aproject/google-a2a-eng
+src/a2a/types.py    @a2a-bot

CHANGELOG.md

@@ -1,5 +1,27 @@
 # Changelog
 
+## [0.3.4](https://github.com/a2aproject/a2a-python/compare/v0.3.3...v0.3.4) (2025-09-02)
+
+
+### Features
+
+* Add `ServerCallContext` into task store operations ([#443](https://github.com/a2aproject/a2a-python/issues/443)) ([e3e5c4b](https://github.com/a2aproject/a2a-python/commit/e3e5c4b7dcb5106e943b9aeb8e761ed23cc166a2))
+* Add extensions support to `TaskUpdater.add_artifact` ([#436](https://github.com/a2aproject/a2a-python/issues/436)) ([598d8a1](https://github.com/a2aproject/a2a-python/commit/598d8a10e61be83bcb7bc9377365f7c42bc6af41))
+
+
+### Bug Fixes
+
+* convert auth_required state in proto utils ([#444](https://github.com/a2aproject/a2a-python/issues/444)) ([ac12f05](https://github.com/a2aproject/a2a-python/commit/ac12f0527d923800192c47dc1bd2e7eed262dfe6))
+* handle concurrent task completion during cancellation ([#449](https://github.com/a2aproject/a2a-python/issues/449)) ([f4c9c18](https://github.com/a2aproject/a2a-python/commit/f4c9c18cfef3ccab1ac7bb30cc7f8293cf3e3ef6))
+* Remove logger error from init on `rest_adapter` and `jsonrpc_app` ([#439](https://github.com/a2aproject/a2a-python/issues/439)) ([9193208](https://github.com/a2aproject/a2a-python/commit/9193208aabac2655a197732ff826e3c2d76f11b5))
+* resolve streaming endpoint deadlock by pre-consuming request body ([#426](https://github.com/a2aproject/a2a-python/issues/426)) ([4186731](https://github.com/a2aproject/a2a-python/commit/4186731df60f7adfcd25f19078d055aca26612a3))
+* Sync jsonrpc and rest implementation of authenticated agent card ([#441](https://github.com/a2aproject/a2a-python/issues/441)) ([9da9ecc](https://github.com/a2aproject/a2a-python/commit/9da9ecc96856a2474d75f986a1f45488c36f53e3))
+
+
+### Performance Improvements
+
+* Improve performance and code style for `proto_utils.py` ([#452](https://github.com/a2aproject/a2a-python/issues/452)) ([1e4b574](https://github.com/a2aproject/a2a-python/commit/1e4b57457386875b64362113356c615bc87315e3))
+
 ## [0.3.3](https://github.com/a2aproject/a2a-python/compare/v0.3.2...v0.3.3) (2025-08-22)
 
 

src/a2a/grpc/a2a_pb2.py

@@ -84,14 +84,16 @@ class TaskStatus(_message.Message):
     def __init__(self, state: _Optional[_Union[TaskState, str]] = ..., update: _Optional[_Union[Message, _Mapping]] = ..., timestamp: _Optional[_Union[datetime.datetime, _timestamp_pb2.Timestamp, _Mapping]] = ...) -> None: ...
 
 class Part(_message.Message):
-    __slots__ = ("text", "file", "data")
+    __slots__ = ("text", "file", "data", "metadata")
     TEXT_FIELD_NUMBER: _ClassVar[int]
     FILE_FIELD_NUMBER: _ClassVar[int]
     DATA_FIELD_NUMBER: _ClassVar[int]
+    METADATA_FIELD_NUMBER: _ClassVar[int]
     text: str
     file: FilePart
     data: DataPart
-    def __init__(self, text: _Optional[str] = ..., file: _Optional[_Union[FilePart, _Mapping]] = ..., data: _Optional[_Union[DataPart, _Mapping]] = ...) -> None: ...
+    metadata: _struct_pb2.Struct
+    def __init__(self, text: _Optional[str] = ..., file: _Optional[_Union[FilePart, _Mapping]] = ..., data: _Optional[_Union[DataPart, _Mapping]] = ..., metadata: _Optional[_Union[_struct_pb2.Struct, _Mapping]] = ...) -> None: ...
 
 class FilePart(_message.Message):
     __slots__ = ("file_with_uri", "file_with_bytes", "mime_type", "name")

src/a2a/grpc/a2a_pb2.pyi

@@ -133,7 +133,7 @@ async def consume_all(self) -> AsyncGenerator[Event]:
                 # continue polling until there is a final event
                 continue
             except asyncio.TimeoutError:  # pyright: ignore [reportUnusedExcept]
-                # This class was made an alias of build-in TimeoutError after 3.11
+                # This class was made an alias of built-in TimeoutError after 3.11
                 continue
             except (QueueClosed, asyncio.QueueEmpty):
                 # Confirm that the queue is closed, e.g. we aren't on

src/a2a/server/events/event_consumer.py

@@ -67,6 +67,7 @@ class DefaultRequestHandler(RequestHandler):
     """
 
     _running_agents: dict[str, asyncio.Task]
+    _background_tasks: set[asyncio.Task]
 
     def __init__(  # noqa: PLR0913
         self,
@@ -102,14 +103,17 @@ def __init__(  # noqa: PLR0913
         # TODO: Likely want an interface for managing this, like AgentExecutionManager.
         self._running_agents = {}
         self._running_agents_lock = asyncio.Lock()
+        # Tracks background tasks (e.g., deferred cleanups) to avoid orphaning
+        # asyncio tasks and to surface unexpected exceptions.
+        self._background_tasks = set()
 
     async def on_get_task(
         self,
         params: TaskQueryParams,
         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 +145,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 +162,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)
 
@@ -180,14 +185,21 @@ async def on_cancel_task(
 
         consumer = EventConsumer(queue)
         result = await result_aggregator.consume_all(consumer)
-        if isinstance(result, Task):
-            return result
+        if not isinstance(result, Task):
+            raise ServerError(
+                error=InternalError(
+                    message='Agent did not return valid response for cancel'
+                )
+            )
 
-        raise ServerError(
-            error=InternalError(
-                message='Agent did not return valid response for cancel'
+        if result.status.state != TaskState.canceled:
+            raise ServerError(
+                error=TaskNotCancelableError(
+                    message=f'Task cannot be canceled - current state: {result.status.state}'
+                )
             )
-        )
+
+        return result
 
     async def _run_event_stream(
         self, request: RequestContext, queue: EventQueue
@@ -217,6 +229,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()
 
@@ -346,10 +359,11 @@ async def push_notification_callback() -> None:
             raise
         finally:
             if interrupted_or_non_blocking:
-                # TODO: Track this disconnected cleanup task.
-                asyncio.create_task(  # noqa: RUF006
+                cleanup_task = asyncio.create_task(
                     self._cleanup_producer(producer_task, task_id)
                 )
+                cleanup_task.set_name(f'cleanup_producer:{task_id}')
+                self._track_background_task(cleanup_task)
             else:
                 await self._cleanup_producer(producer_task, task_id)
 
@@ -385,7 +399,11 @@ async def on_message_send_stream(
                 )
                 yield event
         finally:
-            await self._cleanup_producer(producer_task, task_id)
+            cleanup_task = asyncio.create_task(
+                self._cleanup_producer(producer_task, task_id)
+            )
+            cleanup_task.set_name(f'cleanup_producer:{task_id}')
+            self._track_background_task(cleanup_task)
 
     async def _register_producer(
         self, task_id: str, producer_task: asyncio.Task
@@ -394,6 +412,29 @@ async def _register_producer(
         async with self._running_agents_lock:
             self._running_agents[task_id] = producer_task
 
+    def _track_background_task(self, task: asyncio.Task) -> None:
+        """Tracks a background task and logs exceptions on completion.
+
+        This avoids unreferenced tasks (and associated lint warnings) while
+        ensuring any exceptions are surfaced in logs.
+        """
+        self._background_tasks.add(task)
+
+        def _on_done(completed: asyncio.Task) -> None:
+            try:
+                # Retrieve result to raise exceptions, if any
+                completed.result()
+            except asyncio.CancelledError:
+                name = completed.get_name()
+                logger.debug('Background task %s cancelled', name)
+            except Exception:
+                name = completed.get_name()
+                logger.exception('Background task %s failed', name)
+            finally:
+                self._background_tasks.discard(completed)
+
+        task.add_done_callback(_on_done)
+
     async def _cleanup_producer(
         self,
         producer_task: asyncio.Task,
@@ -417,7 +458,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 +481,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 +510,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 +526,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 +551,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 +578,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/request_handlers/default_request_handler.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/database_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/inmemory_task_store.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.
 
@@ -40,6 +42,7 @@ def __init__(
             task_store: The `TaskStore` instance for persistence.
             initial_message: The `Message` that initiated the task, if any.
                              Used when creating a new task object.
+            context: The `ServerCallContext` that this task is produced under.
         """
         if task_id is not None and not (isinstance(task_id, str) and task_id):
             raise ValueError('Task ID must be a non-empty string')
@@ -49,6 +52,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 +78,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._call_context
+        )
         if self._current_task:
             logger.debug('Task %s retrieved successfully.', self.task_id)
         else:
@@ -167,7 +173,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._call_context)
 
         if not task:
             logger.info(
@@ -231,7 +237,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._call_context)
         self._current_task = task
         if not self.task_id:
             logger.info('New task created with id: %s', task.id)