Add streaming support

.gitignore

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

fasta2a/__init__.py

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

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
@@ -18,9 +19,12 @@
     AgentCard,
     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
@@ -42,6 +46,7 @@ def __init__(
         provider: AgentProvider | None = None,
         skills: list[Skill] | None = None,
         docs_url: str | None = '/docs',
+        streaming: bool = True,
         # Starlette
         debug: bool = False,
         routes: Sequence[Route] | None = None,
@@ -67,6 +72,7 @@ def __init__(
         self.provider = provider
         self.skills = skills 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']
@@ -100,7 +106,7 @@ async def _agent_card_endpoint(self, request: Request) -> Response:
                 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
+                    streaming=self.streaming, push_notifications=False, state_transition_history=False
                 ),
             )
             if self.provider is not None:
@@ -131,6 +137,20 @@ async def _agent_run_endpoint(self, request: Request) -> Response:
 
         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

@@ -808,3 +808,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

fasta2a/worker.py

@@ -10,11 +10,12 @@
 from opentelemetry.trace import get_tracer, use_span
 from typing_extensions import assert_never
 
+from .schema import 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,6 +57,66 @@ async def _handle_task_operation(self, task_operation: TaskOperation) -> None:
         except Exception:
             await self.storage.update_task(task_operation['params']['id'], state='failed')
 
+    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.send_stream_event(
+                task_id,
+                TaskStatusUpdateEvent(
+                    kind='status-update',
+                    task_id=task_id,
+                    context_id=task['context_id'],
+                    status=task['status'],
+                    final=False,
+                ),
+            )
+
+        # Publish message events before final status so subscribers receive them
+        if new_messages:
+            for message in new_messages:
+                await self.broker.send_stream_event(task_id, message)
+
+        # Publish artifact events
+        if new_artifacts:
+            for artifact in new_artifacts:
+                await self.broker.send_stream_event(
+                    task_id,
+                    TaskArtifactUpdateEvent(
+                        kind='artifact-update',
+                        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.send_stream_event(
+                task_id,
+                TaskStatusUpdateEvent(
+                    kind='status-update',
+                    task_id=task_id,
+                    context_id=task['context_id'],
+                    status=task['status'],
+                    final=True,
+                ),
+            )
+
     @abstractmethod
     async def run_task(self, params: TaskSendParams) -> None: ...