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

@@ -5,6 +5,7 @@
 from pathlib import Path
 from typing import Any
 
+from sse_starlette import EventSourceResponse
 from starlette.applications import Starlette
 from starlette.middleware import Middleware
 from starlette.requests import Request
@@ -16,11 +17,15 @@
 from .schema import (
     AgentCapabilities,
     AgentCard,
+    AgentExtension,
     AgentProvider,
     Skill,
+    StreamMessageResponse,
     a2a_request_ta,
     a2a_response_ta,
     agent_card_ta,
+    stream_message_request_ta,
+    stream_message_response_ta,
 )
 from .storage import Storage
 from .task_manager import TaskManager
@@ -41,7 +46,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,
@@ -66,7 +73,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']
@@ -90,6 +99,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.',
@@ -99,9 +113,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=False, push_notifications=False, state_transition_history=False
-                ),
+                capabilities=capabilities,
             )
             if self.provider is not None:
                 agent_card['provider'] = self.provider
@@ -129,8 +141,30 @@ 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
+
         if a2a_request['method'] == 'message/send':
             jsonrpc_response = await self.task_manager.send_message(a2a_request)
+        elif a2a_request['method'] == 'message/stream':
+            stream_request = stream_message_request_ta.validate_json(data)
+
+            async def sse_generator():
+                request_id = stream_request.get('id')
+                async for event in self.task_manager.stream_message(stream_request):
+                    jsonrpc_response = StreamMessageResponse(
+                        jsonrpc='2.0',
+                        id=request_id,
+                        result=event,
+                    )
+                    yield stream_message_response_ta.dump_json(jsonrpc_response, by_alias=True).decode()
+
+            return EventSourceResponse(sse_generator())
         elif a2a_request['method'] == 'tasks/get':
             jsonrpc_response = await self.task_manager.get_task(a2a_request)
         elif a2a_request['method'] == 'tasks/cancel':

fasta2a/broker.py

@@ -7,11 +7,12 @@
 from typing import Annotated, Any, Generic, Literal, TypeVar
 
 import anyio
+from anyio.streams.memory import MemoryObjectSendStream
 from opentelemetry.trace import Span, get_current_span, get_tracer
 from pydantic import Discriminator
 from typing_extensions import Self, TypedDict
 
-from .schema import TaskIdParams, TaskSendParams
+from .schema import StreamEvent, TaskIdParams, TaskSendParams
 
 tracer = get_tracer(__name__)
 
@@ -37,6 +38,26 @@ async def cancel_task(self, params: TaskIdParams) -> None:
         """Cancel a task."""
         raise NotImplementedError('send_cancel_task is not implemented yet.')
 
+    @abstractmethod
+    async def send_stream_event(self, task_id: str, event: StreamEvent) -> None:
+        """Send a streaming event from worker to subscribers.
+
+        This is used by workers to publish status updates, messages, and artifacts
+        during task execution. Events are forwarded to all active subscribers of
+        the given task_id.
+        """
+        ...
+
+    @abstractmethod
+    def subscribe_to_stream(self, task_id: str) -> AsyncIterator[StreamEvent]:
+        """Subscribe to streaming events for a specific task.
+
+        Returns an async iterator that yields events published by workers for the
+        given task_id. The iterator completes when a TaskStatusUpdateEvent with
+        final=True is received or the subscription is cancelled.
+        """
+        ...
+
     @abstractmethod
     async def __aenter__(self) -> Self: ...
 
@@ -73,6 +94,10 @@ class _TaskOperation(TypedDict, Generic[OperationT, ParamsT]):
 class InMemoryBroker(Broker):
     """A broker that schedules tasks in memory."""
 
+    def __init__(self) -> None:
+        self._event_subscribers: dict[str, list[MemoryObjectSendStream[StreamEvent]]] = {}
+        self._subscriber_lock: anyio.Lock | None = None
+
     async def __aenter__(self):
         self.aexit_stack = AsyncExitStack()
         await self.aexit_stack.__aenter__()
@@ -81,6 +106,8 @@ async def __aenter__(self):
         await self.aexit_stack.enter_async_context(self._read_stream)
         await self.aexit_stack.enter_async_context(self._write_stream)
 
+        self._subscriber_lock = anyio.Lock()
+
         return self
 
     async def __aexit__(self, exc_type: Any, exc_value: Any, traceback: Any):
@@ -96,3 +123,65 @@ async def receive_task_operations(self) -> AsyncIterator[TaskOperation]:
         """Receive task operations from the broker."""
         async for task_operation in self._read_stream:
             yield task_operation
+
+    async def send_stream_event(self, task_id: str, event: StreamEvent) -> None:
+        """Send a streaming event from worker to subscribers."""
+        assert self._subscriber_lock is not None, 'Broker not initialized'
+
+        async with self._subscriber_lock:
+            subscribers = self._event_subscribers.get(task_id, [])
+            if not subscribers:
+                return
+
+            # Send event to all subscribers, removing closed streams
+            active_subscribers: list[MemoryObjectSendStream[StreamEvent]] = []
+            for stream in subscribers:
+                try:
+                    await stream.send(event)
+                    active_subscribers.append(stream)
+                except (anyio.ClosedResourceError, anyio.BrokenResourceError):
+                    # Subscriber disconnected, remove from list
+                    pass
+
+            # Update subscriber list with only active ones
+            if active_subscribers:
+                self._event_subscribers[task_id] = active_subscribers
+            elif task_id in self._event_subscribers:
+                # No active subscribers left, clean up
+                del self._event_subscribers[task_id]
+
+    async def subscribe_to_stream(self, task_id: str) -> AsyncIterator[StreamEvent]:
+        """Subscribe to streaming events for a specific task."""
+        assert self._subscriber_lock is not None, 'Broker not initialized'
+
+        # Create a new stream for this subscriber
+        send_stream, receive_stream = anyio.create_memory_object_stream[StreamEvent](max_buffer_size=100)
+
+        # Register the subscriber
+        async with self._subscriber_lock:
+            if task_id not in self._event_subscribers:
+                self._event_subscribers[task_id] = []
+            self._event_subscribers[task_id].append(send_stream)
+
+        try:
+            async with receive_stream:
+                async for event in receive_stream:
+                    yield event
+
+                    # Check if this is a final status update
+                    if isinstance(event, dict) and event.get('kind') == 'status-update' and event.get('final', False):
+                        break
+        finally:
+            # Clean up subscription on exit
+            async with self._subscriber_lock:
+                if task_id in self._event_subscribers:
+                    try:
+                        self._event_subscribers[task_id].remove(send_stream)
+                        if not self._event_subscribers[task_id]:
+                            del self._event_subscribers[task_id]
+                    except ValueError:
+                        # Already removed
+                        pass
+
+            # Close the send stream
+            await send_stream.aclose()

fasta2a/schema.py

@@ -93,6 +93,15 @@ class AgentCapabilities(TypedDict):
     state_transition_history: NotRequired[bool]
     """Whether the agent exposes status change history for tasks."""
 
+    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):
@@ -808,3 +817,9 @@ class JSONRPCResponse(JSONRPCMessage, Generic[ResultT, ErrorT]):
 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/task_manager.py

@@ -60,7 +60,9 @@
 
 from __future__ import annotations as _annotations
 
+import asyncio
 import uuid
+from collections.abc import AsyncGenerator
 from contextlib import AsyncExitStack
 from dataclasses import dataclass, field
 from typing import Any
@@ -78,8 +80,8 @@
     SendMessageResponse,
     SetTaskPushNotificationRequest,
     SetTaskPushNotificationResponse,
+    StreamEvent,
     StreamMessageRequest,
-    StreamMessageResponse,
     TaskNotFoundError,
     TaskSendParams,
 )
@@ -156,9 +158,44 @@ async def cancel_task(self, request: CancelTaskRequest) -> CancelTaskResponse:
             )
         return CancelTaskResponse(jsonrpc='2.0', id=request['id'], result=task)
 
-    async def stream_message(self, request: StreamMessageRequest) -> StreamMessageResponse:
-        """Stream messages using Server-Sent Events."""
-        raise NotImplementedError('message/stream method is not implemented yet.')
+    async def stream_message(self, request: StreamMessageRequest) -> AsyncGenerator[StreamEvent, None]:
+        """Handle a streaming message request.
+
+        This method:
+        1. Creates and submits a new task
+        2. Yields the initial task object
+        3. Subscribes to the broker's event stream
+        4. Starts task execution asynchronously
+        5. Streams all events until completion
+        """
+        # Extract parameters
+        params = request['params']
+        message = params['message']
+        context_id = message.get('context_id', str(uuid.uuid4()))
+
+        # Create and submit the task
+        task = await self.storage.submit_task(context_id, message)
+
+        # Yield the initial task
+        yield task
+
+        # Prepare broker params
+        broker_params: TaskSendParams = {'id': task['id'], 'context_id': context_id, 'message': message}
+        config = params.get('configuration', {})
+        history_length = config.get('history_length')
+        if history_length is not None:
+            broker_params['history_length'] = history_length
+
+        metadata = params.get('metadata')
+        if metadata is not None:
+            broker_params['metadata'] = metadata
+
+        # Start task execution in background
+        asyncio.create_task(self.broker.run_task(broker_params))
+
+        # Stream events from broker
+        async for event in self.broker.subscribe_to_stream(task['id']):
+            yield event
 
     async def set_task_push_notification(
         self, request: SetTaskPushNotificationRequest