feat: extensions

.gitignore

@@ -1,2 +1,3 @@
 .cache
 __pycache__
+.coverage*

fasta2a/__init__.py

@@ -1,7 +1,15 @@
 from .applications import FastA2A
 from .broker import Broker
-from .schema import Skill
+from .schema import AgentExtension, Skill, StreamEvent
 from .storage import Storage
 from .worker import Worker
 
-__all__ = ['FastA2A', 'Skill', 'Storage', 'Broker', 'Worker']
+__all__ = [
+    'AgentExtension',
+    'Broker',
+    'FastA2A',
+    'Skill',
+    'Storage',
+    'StreamEvent',
+    'Worker',
+]

fasta2a/applications.py

@@ -17,6 +17,7 @@
     A2AResponse,
     AgentCapabilities,
     AgentCard,
+    AgentExtension,
     AgentInterface,
     AgentProvider,
     Skill,
@@ -43,7 +44,9 @@ def __init__(
         description: str | None = None,
         provider: AgentProvider | None = None,
         skills: list[Skill] | None = None,
+        extensions: list[AgentExtension] | None = None,
         docs_url: str | None = '/docs',
+        streaming: bool = True,
         # Starlette
         debug: bool = False,
         routes: Sequence[Route] | None = None,
@@ -68,7 +71,9 @@ def __init__(
         self.description = description
         self.provider = provider
         self.skills = skills or []
+        self.extensions = extensions or []
         self.docs_url = docs_url
+        self.streaming = streaming
         # NOTE: For now, I don't think there's any reason to support any other input/output modes.
         self.default_input_modes = ['application/json']
         self.default_output_modes = ['application/json']
@@ -92,6 +97,11 @@ async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
 
     async def _agent_card_endpoint(self, request: Request) -> Response:
         if self._agent_card_json_schema is None:
+            capabilities = AgentCapabilities(
+                streaming=self.streaming, push_notifications=False, state_transition_history=False
+            )
+            if self.extensions:
+                capabilities['extensions'] = self.extensions
             agent_card = AgentCard(
                 name=self.name,
                 description=self.description or 'An AI agent exposed as an A2A agent.',
@@ -102,7 +112,7 @@ async def _agent_card_endpoint(self, request: Request) -> Response:
                 skills=self.skills,
                 default_input_modes=self.default_input_modes,
                 default_output_modes=self.default_output_modes,
-                capabilities=AgentCapabilities(streaming=True, push_notifications=False),
+                capabilities=capabilities,
             )
             if self.provider is not None:
                 agent_card['provider'] = self.provider
@@ -130,9 +140,22 @@ async def _agent_run_endpoint(self, request: Request) -> Response:
         data = await request.body()
         a2a_request = a2a_request_ta.validate_json(data)
 
+        # Parse activated extensions from the A2A-Extensions header
+        extensions_header = request.headers.get('a2a-extensions', '')
+        activated_extensions: list[str] = (
+            [uri.strip() for uri in extensions_header.split(',') if uri.strip()] if extensions_header else []
+        )
+        # Stash on the request state so workers / handlers can inspect them
+        request.state.activated_extensions = activated_extensions
+
         jsonrpc_response: A2AResponse
         if a2a_request['method'] == 'message/send':
             jsonrpc_response = await self.task_manager.send_message(a2a_request)
+        elif a2a_request['method'] == 'message/stream':
+            return StreamingResponse(
+                self.task_manager.stream_message(a2a_request),
+                media_type='text/event-stream',
+            )
         elif a2a_request['method'] == 'tasks/get':
             jsonrpc_response = await self.task_manager.get_task(a2a_request)
         elif a2a_request['method'] == 'tasks/cancel':
@@ -147,11 +170,6 @@ async def _agent_run_endpoint(self, request: Request) -> Response:
             jsonrpc_response = await self.task_manager.delete_task_push_notification_config(a2a_request)
         elif a2a_request['method'] == 'tasks/list':
             jsonrpc_response = await self.task_manager.list_tasks(a2a_request)
-        elif a2a_request['method'] == 'message/stream':
-            return StreamingResponse(
-                self.task_manager.stream_message(a2a_request),
-                media_type='text/event-stream',
-            )
         elif a2a_request['method'] == 'tasks/resubscribe':
             return StreamingResponse(
                 self.task_manager.resubscribe_task(a2a_request),

fasta2a/schema.py

@@ -90,6 +90,15 @@ class AgentCapabilities(TypedDict):
     extensions: NotRequired[list[AgentExtension]]
     """A list of protocol extensions supported by the agent."""
 
+    extensions: NotRequired[list[AgentExtension]]
+    """A2A extensions supported by this agent.
+
+    Each extension is declared as an ``AgentExtension`` object with a
+    unique ``uri``, optional ``description``, ``required`` flag, and
+    ``params`` configuration.  Clients activate extensions by sending
+    the selected URIs in the ``A2A-Extensions`` HTTP header.
+    """
+
 
 @pydantic.with_config({'alias_generator': to_camel})
 class HttpSecurityScheme(TypedDict):
@@ -992,3 +1001,9 @@ class StreamResponse(TypedDict):
 send_message_response_ta: TypeAdapter[SendMessageResponse] = TypeAdapter(SendMessageResponse)
 stream_message_request_ta: TypeAdapter[StreamMessageRequest] = TypeAdapter(StreamMessageRequest)
 stream_message_response_ta: TypeAdapter[StreamMessageResponse] = TypeAdapter(StreamMessageResponse)
+
+# Type for streaming events (used by broker and task manager)
+StreamEvent = Union[Task, Message, TaskStatusUpdateEvent, TaskArtifactUpdateEvent]
+"""A streaming event that can be sent during message/stream requests."""
+
+stream_event_ta: TypeAdapter[StreamEvent] = TypeAdapter(StreamEvent)

fasta2a/storage.py

@@ -9,7 +9,13 @@
 
 from typing_extensions import TypeVar
 
-from .schema import Artifact, Message, Task, TaskState, TaskStatus
+from .schema import (
+    Artifact,
+    Message,
+    Task,
+    TaskState,
+    TaskStatus,
+)
 
 ContextT = TypeVar('ContextT', default=Any)
 

fasta2a/worker.py

@@ -10,11 +10,12 @@
 from opentelemetry.trace import get_tracer, use_span
 from typing_extensions import assert_never
 
+from .schema import StreamResponse, TaskArtifactUpdateEvent, TaskStatusUpdateEvent
 from .storage import ContextT, Storage
 
 if TYPE_CHECKING:
     from .broker import Broker, TaskOperation
-    from .schema import Artifact, Message, TaskIdParams, TaskSendParams
+    from .schema import Artifact, Message, TaskIdParams, TaskSendParams, TaskState
 
 tracer = get_tracer(__name__)
 
@@ -56,7 +57,7 @@ async def _handle_task_operation(self, task_operation: TaskOperation) -> None:
         except Exception:
             task_id = task_operation['params']['id']
             task = await self.storage.update_task(task_id, state='failed')
-            from .schema import StreamResponse, TaskStatus, TaskStatusUpdateEvent
+            from .schema import TaskStatus
 
             await self.broker.event_bus.emit(
                 task_id,
@@ -70,6 +71,68 @@ async def _handle_task_operation(self, task_operation: TaskOperation) -> None:
             )
             await self.broker.event_bus.close(task_id)
 
+    async def update_task(
+        self,
+        task_id: str,
+        state: TaskState,
+        new_artifacts: list[Artifact] | None = None,
+        new_messages: list[Message] | None = None,
+    ) -> None:
+        """Update a task's state in storage and publish streaming events to the broker.
+
+        This is the primary method workers should use to update task state. It handles
+        both persisting the update and notifying any stream subscribers.
+        """
+        task = await self.storage.update_task(task_id, state, new_artifacts, new_messages)
+
+        final = state in ('completed', 'failed', 'canceled')
+
+        # For non-final updates, publish status first
+        if not final:
+            await self.broker.event_bus.emit(
+                task_id,
+                StreamResponse(
+                    status_update=TaskStatusUpdateEvent(
+                        task_id=task_id,
+                        context_id=task['context_id'],
+                        status=task['status'],
+                    ),
+                ),
+            )
+
+        # Publish message events before final status so subscribers receive them
+        if new_messages:
+            for message in new_messages:
+                await self.broker.event_bus.emit(task_id, StreamResponse(message=message))
+
+        # Publish artifact events
+        if new_artifacts:
+            for artifact in new_artifacts:
+                await self.broker.event_bus.emit(
+                    task_id,
+                    StreamResponse(
+                        artifact_update=TaskArtifactUpdateEvent(
+                            task_id=task_id,
+                            context_id=task['context_id'],
+                            artifact=artifact,
+                        ),
+                    ),
+                )
+
+        # For final updates, publish status last (after messages and artifacts)
+        if final:
+            await self.broker.event_bus.emit(
+                task_id,
+                StreamResponse(
+                    status_update=TaskStatusUpdateEvent(
+                        task_id=task_id,
+                        context_id=task['context_id'],
+                        status=task['status'],
+                    ),
+                ),
+            )
+            await self.broker.event_bus.close(task_id)
+
     @abstractmethod
     async def run_task(self, params: TaskSendParams) -> None: ...
 

pyproject.toml

@@ -48,6 +48,7 @@ dependencies = [
     "pydantic>=2.10",
     "opentelemetry-api>=1.28.0",
     "eval_type_backport>=0.2.2; python_version <= '3.9'",
+    "sse-starlette>=2.0.0",
 ]
 
 [project.optional-dependencies]
@@ -58,8 +59,11 @@ dev = [
     "asgi-lifespan",
     "coverage",
     "httpx",
+    "httpx-sse",
     "inline-snapshot",
     "pytest",
+    "pytest-asyncio",
+    "pytest-mock",
     "ruff",
     "pyright",
 ]