1"""

2 A2A Streaming — real-time task status updates via SSE for A2A protocol.

3

4 Provides push-based task lifecycle notifications so agents don't poll.

5"""

6

7 from __future__ import annotations

8

9 import asyncio

10 import json

11 import time

12 from dataclasses import dataclass, field

13 from enum import Enum

14 from typing import Any, AsyncIterator, Callable, Dict, Optional

15

16 from agentos.protocols.a2a import A2ATask, TaskState

17

18

19 class A2AStreamEvent(str, Enum):

20 """A2A-specific streaming event types."""

21

22 TASK_CREATED = "task.created"

23 TASK_STARTED = "task.started"

24 TASK_PROGRESS = "task.progress"

25 TASK_COMPLETED = "task.completed"

26 TASK_FAILED = "task.failed"

27 TASK_CANCELLED = "task.cancelled"

28 ARTIFACT_ADDED = "artifact.added"

29 HEARTBEAT = "heartbeat"

30

31

32@dataclass

33 class TaskProgress:

34 """Progress update within a running task."""

35

36 percent: float = 0.0

37 message: str = ""

38 step: str = ""

39 metadata: dict[str, Any] = field(default_factory=dict)

40

41

42 class A2AStreamSession:

43 """Manages a streaming connection for a single task.

44

45 Agents subscribe to receive push updates as the task progresses.

46 """

47

48 def __init__(self, task: A2ATask):

49 self.task_id = task.task_id

50 self._subscribers: list[asyncio.Queue[dict]] = []

51 self._closed = False

52 self._heartbeat_task: Optional[asyncio.Task] = None

53

54 async def start(self, heartbeat_s: float = 30.0):

55 """Start heartbeat loop."""

56 async def _pulse():

57 while not self._closed:

58 await asyncio.sleep(heartbeat_s)

59 if not self._closed:

60 await self._broadcast({

61 "event": A2AStreamEvent.HEARTBEAT,

62 "task_id": self.task_id,

63 "timestamp": time.time(),

64 })

65 self._heartbeat_task = asyncio.create_task(_pulse())

66

67 def subscribe(self) -> asyncio.Queue[dict]:

68 """Register a new subscriber. Returns a queue of SSE events."""

69 q: asyncio.Queue[dict] = asyncio.Queue(maxsize=64)

70 self._subscribers.append(q)

71 return q

72

73 def unsubscribe(self, sub: asyncio.Queue):

74 """Remove a subscriber."""

75 try:

76 self._subscribers.remove(sub)

77 except ValueError:

78 pass

79

80 async def emit(self, event: A2AStreamEvent, data: dict | None = None):

81 """Push an event to all subscribers."""

82 payload = {

83 "event": event.value,

84 "task_id": self.task_id,

85 "timestamp": time.time(),

86 }

87 if data:

88 payload["data"] = data

89 await self._broadcast(payload)

90

91 async def _broadcast(self, payload: dict):

92 dead: list[asyncio.Queue] = []

93 for q in self._subscribers:

94 try:

95 q.put_nowait(payload)

96 except asyncio.QueueFull:

97 dead.append(q)

98 for q in dead:

99 self.unsubscribe(q)

100

101 async def close(self):

102 """Shut down the stream."""

103 self._closed = True

104 if self._heartbeat_task:

105 self._heartbeat_task.cancel()

106 # Close all subscriber queues

107 for q in self._subscribers:

108 try:

109 q.put_nowait(None) # Sentinel

110 except asyncio.QueueFull:

111 pass

112 self._subscribers.clear()

113

114 async def iter_events(self, subscriber: asyncio.Queue) -> AsyncIterator[dict]:

115 """Async iterator yielding SSE-compatible event dicts."""

116 while True:

117 event = await subscriber.get()

118 if event is None:

119 break

120 yield event

121

122 def to_sse(self, event: dict) -> str:

123 """Format a single event dict into SSE wire format."""

124 lines: list[str] = [f"event: {event['event']}"]

125 for key in ("task_id", "timestamp"):

126 if key in event:

127 lines.append(f"id: {key}={event[key]}")

128 data_str = json.dumps(event.get("data", {}), ensure_ascii=False)

129 for line in data_str.split("\n"):

130 lines.append(f"data: {line}")

131 return "\n".join(lines) + "\n\n"

132

133

134 class StreamingAggregator:

135 """流式结果聚合器 — 合规测试套件要求。"""

136

137 def __init__(self):

138 self._chunks: list[str] = []

139

140 def collect(self, chunk: str) -> None:

141 self._chunks.append(chunk)

142

143 def aggregated(self) -> str:

144 return "".join(self._chunks)

145

146

147 class A2AStreamManager:

148 """Global manager for A2A task streaming sessions.

149

150 Tracks all active task streams and dispatches events on state transitions.

151 """

152

153 def __init__(self):

154 self._sessions: dict[str, A2AStreamSession] = {}

155 self._on_state_change: Optional[Callable] = None

156

157 def on_state_change(self, callback: Callable[[A2ATask, TaskState, TaskState], Any]):

158 """Register a hook called on every state transition (old_state, new_state)."""

159 self._on_state_change = callback

160

161 def create_session(self, task: A2ATask) -> A2AStreamSession:

162 """Create a streaming session for a new task."""

163 session = A2AStreamSession(task)

164 self._sessions[task.task_id] = session

165 return session

166

167 def get_session(self, task_id: str) -> Optional[A2AStreamSession]:

168 return self._sessions.get(task_id)

169

170 async def notify_state_change(self, task: A2ATask, old_state: TaskState):

171 """Called when a task transitions state."""

172 session = self._sessions.get(task.task_id)

173 if not session:

174 return

175

176 event_map = {

177 TaskState.SUBMITTED: A2AStreamEvent.TASK_CREATED,

178 TaskState.WORKING: A2AStreamEvent.TASK_STARTED,

179 TaskState.COMPLETED: A2AStreamEvent.TASK_COMPLETED,

180 TaskState.FAILED: A2AStreamEvent.TASK_FAILED,

181 TaskState.CANCELLED: A2AStreamEvent.TASK_CANCELLED,

182 }

183 event = event_map.get(task.state, A2AStreamEvent.TASK_PROGRESS)

184 await session.emit(event, {

185 "previous_state": old_state.value,

186 "current_state": task.state.value,

187 "error": task.error,

188 })

189

190 if task.is_terminal():

191 await session.close()

192 del self._sessions[task.task_id]

193

194 async def notify_artifact(self, task_id: str, artifact_name: str):

195 """Called when an artifact is added to a task."""

196 session = self._sessions.get(task_id)

197 if session:

198 await session.emit(A2AStreamEvent.ARTIFACT_ADDED, {

199 "artifact_name": artifact_name,

200 })

201

202 async def notify_progress(

203 self,

204 task_id: str,

205 progress: TaskProgress,

206 ):

207 """Push a progress update to subscribers."""

208 session = self._sessions.get(task_id)

209 if session:

210 await session.emit(A2AStreamEvent.TASK_PROGRESS, {

211 "percent": progress.percent,

212 "message": progress.message,

213 "step": progress.step,

214 "metadata": progress.metadata,

215 })

216

217 async def shutdown(self):

218 """Gracefully close all sessions."""

219 for sid in list(self._sessions.keys()):

220 session = self._sessions[sid]

221 await session.close()

222 self._sessions.clear()