🎥 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

@@ -44,6 +44,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,
@@ -69,6 +70,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']
@@ -102,7 +104,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=AgentCapabilities(streaming=self.streaming, push_notifications=False),
             )
             if self.provider is not None:
                 agent_card['provider'] = self.provider
@@ -133,6 +135,11 @@ async def _agent_run_endpoint(self, request: Request) -> Response:
         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 +154,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

@@ -992,3 +992,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",
 ]