၊၊||၊ Add streaming support

fasta2a/__init__.py

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

fasta2a/applications.py

@@ -1,10 +1,12 @@
 from __future__ import annotations as _annotations
 
+import json
 from collections.abc import AsyncIterator, Sequence
 from contextlib import asynccontextmanager
 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
@@ -21,6 +23,8 @@
     a2a_request_ta,
     a2a_response_ta,
     agent_card_ta,
+    stream_event_ta,
+    stream_message_request_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 = False,
         # 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,25 @@ 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':
+            # Parse the streaming request
+            stream_request = stream_message_request_ta.validate_json(data)
+
+            # Create an async generator wrapper that formats events as JSON-RPC responses
+            async def sse_generator():
+                request_id = stream_request.get('id')
+                async for event in self.task_manager.stream_message(stream_request):
+                    # Serialize event to ensure proper camelCase conversion
+                    event_dict = stream_event_ta.dump_python(event, mode='json', by_alias=True)
+
+                    # Wrap in JSON-RPC response
+                    jsonrpc_response = {'jsonrpc': '2.0', 'id': request_id, 'result': event_dict}
+
+                    # Convert to JSON string
+                    yield json.dumps(jsonrpc_response)
+
+            # Return SSE response
+            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__)
 
@@ -30,12 +31,32 @@ class Broker(ABC):
     @abstractmethod
     async def run_task(self, params: TaskSendParams) -> None:
         """Send a task to be executed by the worker."""
-        raise NotImplementedError('send_run_task is not implemented yet.')
+        ...
 
     @abstractmethod
     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

@@ -5,11 +5,22 @@
 import uuid
 from abc import ABC, abstractmethod
 from datetime import datetime
-from typing import Any, Generic
+from typing import TYPE_CHECKING, Any, Generic
 
 from typing_extensions import TypeVar
 
-from .schema import Artifact, Message, Task, TaskState, TaskStatus
+from .schema import (
+    Artifact,
+    Message,
+    Task,
+    TaskArtifactUpdateEvent,
+    TaskState,
+    TaskStatus,
+    TaskStatusUpdateEvent,
+)
+
+if TYPE_CHECKING:
+    from .broker import Broker
 
 ContextT = TypeVar('ContextT', default=Any)
 
@@ -129,3 +140,82 @@ async def update_context(self, context_id: str, context: ContextT) -> None:
     async def load_context(self, context_id: str) -> ContextT | None:
         """Retrieve the stored context given the `context_id`."""
         return self.contexts.get(context_id)
+
+
+class StreamingStorageWrapper(Storage[ContextT]):
+    """A storage wrapper that publishes streaming events when tasks are updated.
+
+    This wrapper intercepts update_task calls and publishes TaskStatusUpdateEvent
+    and TaskArtifactUpdateEvent to the broker, enabling SSE streaming without
+    modifying the underlying storage or worker implementations.
+    """
+
+    def __init__(self, storage: Storage[ContextT], broker: Broker):
+        self._storage = storage
+        self._broker = broker
+
+    async def load_task(self, task_id: str, history_length: int | None = None) -> Task | None:
+        return await self._storage.load_task(task_id, history_length)
+
+    async def submit_task(self, context_id: str, message: Message) -> Task:
+        return await self._storage.submit_task(context_id, message)
+
+    async def update_task(
+        self,
+        task_id: str,
+        state: TaskState,
+        new_artifacts: list[Artifact] | None = None,
+        new_messages: list[Message] | None = None,
+    ) -> Task:
+        """Update task and publish streaming events."""
+        # Update the underlying storage first
+        task = await self._storage.update_task(task_id, state, new_artifacts, new_messages)
+
+        # Determine if this is a final state
+        final = state in ('completed', 'failed', 'canceled')
+
+        # For non-final updates, publish status first
+        if not final:
+            status_event = TaskStatusUpdateEvent(
+                kind='status-update',
+                task_id=task_id,
+                context_id=task['context_id'],
+                status=task['status'],
+                final=False,
+            )
+            await self._broker.send_stream_event(task_id, status_event)
+
+        # Publish message events BEFORE final status (so subscriber receives 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:
+                artifact_event = TaskArtifactUpdateEvent(
+                    kind='artifact-update',
+                    task_id=task_id,
+                    context_id=task['context_id'],
+                    artifact=artifact,
+                )
+                await self._broker.send_stream_event(task_id, artifact_event)
+
+        # For final updates, publish status LAST (after messages and artifacts)
+        if final:
+            status_event = TaskStatusUpdateEvent(
+                kind='status-update',
+                task_id=task_id,
+                context_id=task['context_id'],
+                status=task['status'],
+                final=True,
+            )
+            await self._broker.send_stream_event(task_id, status_event)
+
+        return task
+
+    async def load_context(self, context_id: str) -> ContextT | None:
+        return await self._storage.load_context(context_id)
+
+    async def update_context(self, context_id: str, context: ContextT) -> None:
+        await self._storage.update_context(context_id, context)