Coverage for agentos/core/event_bus.py: 0%

1"""AgentOS Event Bus — async pub/sub with backpressure control. 

2 

3Production-grade internal event system: 

4- Topic-based publish/subscribe 

5- Async handlers with configurable concurrency 

6- Dead letter queue (DLQ) for failed events 

7- Backpressure control with bounded queues 

8- Event replay & audit 

9- Wildcard subscriptions 

10 

11Design: ~350 lines, zero external deps beyond stdlib + asyncio. 

12""" 

13 

14from __future__ import annotations 

15 

16import asyncio 

17import logging 

18import time 

19from collections import defaultdict 

20from dataclasses import dataclass, field 

21from enum import Enum 

22from typing import Any, Awaitable, Callable, Dict, List, Optional, Set 

23from uuid import uuid4 

24 

25logger = logging.getLogger(__name__) 

26 

27 

28# ============================================================================ 

29# Core types 

30# ============================================================================ 

31 

32class EventPriority(int, Enum): 

33 LOW = 0 

34 NORMAL = 50 

35 HIGH = 100 

36 CRITICAL = 200 

37 

38 

39@dataclass 

40class Event: 

41 """An event published on the bus.""" 

42 

43 topic: str 

44 payload: Any = None 

45 event_id: str = field(default_factory=lambda: str(uuid4())[:12]) 

46 timestamp: float = field(default_factory=time.time) 

47 priority: EventPriority = EventPriority.NORMAL 

48 source: str = "" 

49 correlation_id: str = "" 

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

51 

52 

53@dataclass 

54class DeadLetter: 

55 """Event that failed processing and was sent to DLQ.""" 

56 

57 event: Event 

58 handler_name: str 

59 error: str 

60 failed_at: float = field(default_factory=time.time) 

61 retry_count: int = 0 

62 

63 

64# ============================================================================ 

65# Subscription 

66# ============================================================================ 

67 

68@dataclass 

69class Subscription: 

70 """A handler subscribed to a topic pattern.""" 

71 

72 topic_pattern: str 

73 handler: Callable[[Event], Awaitable[Any]] 

74 handler_name: str 

75 concurrency: int = 1 # Max concurrent executions 

76 is_pattern: bool = False # True if topic_pattern contains wildcards 

77 

78 def matches(self, topic: str) -> bool: 

79 if not self.is_pattern: 

80 return self.topic_pattern == topic 

81 # Support * wildcard: "agent.*" matches "agent.start", "agent.stop" 

82 pattern = self.topic_pattern.replace("*", ".*") 

83 import re 

84 return bool(re.match(f"^{pattern}$", topic)) 

85 

86 

87# ============================================================================ 

88# Event Bus 

89# ============================================================================ 

90 

91class EventBus: 

92 """Async pub/sub event bus with backpressure and DLQ.""" 

93 

94 def __init__( 

95 self, 

96 max_queue_size: int = 10000, 

97 dlq_enabled: bool = True, 

98 dlq_max_size: int = 1000, 

99 worker_count: int = 4, 

100 ): 

101 self._subscriptions: Dict[str, List[Subscription]] = defaultdict(list) 

102 self._dead_letters: List[DeadLetter] = [] 

103 self._dlq_enabled = dlq_enabled 

104 self._dlq_max_size = dlq_max_size 

105 self._queue: asyncio.Queue[tuple[Event, Subscription]] = asyncio.Queue( 

106 maxsize=max_queue_size 

107 ) 

108 self._workers: List[asyncio.Task] = [] 

109 self._worker_count = worker_count 

110 self._running = False 

111 self._lock = asyncio.Lock() 

112 self._event_count: int = 0 

113 

114 # ----- subscription management ----- 

115 

116 def subscribe( 

117 self, 

118 topic: str, 

119 handler: Callable[[Event], Awaitable[Any]], 

120 handler_name: Optional[str] = None, 

121 concurrency: int = 1, 

122 ) -> Subscription: 

123 """Subscribe a handler to a topic. 

124 

125 Topic can contain * wildcard: 'agent.*', 'order.*.created' 

126 """ 

127 name = handler_name or handler.__name__ 

128 is_pattern = "*" in topic 

129 sub = Subscription( 

130 topic_pattern=topic, 

131 handler=handler, 

132 handler_name=name, 

133 concurrency=concurrency, 

134 is_pattern=is_pattern, 

135 ) 

136 self._subscriptions[topic].append(sub) 

137 logger.debug("Subscribed '%s' to topic '%s'", name, topic) 

138 return sub 

139 

140 def unsubscribe(self, topic: str, handler_name: str) -> bool: 

141 """Remove a subscription by handler name.""" 

142 subs = self._subscriptions.get(topic, []) 

143 original_len = len(subs) 

144 self._subscriptions[topic] = [ 

145 s for s in subs if s.handler_name != handler_name 

146 ] 

147 removed = original_len != len(self._subscriptions[topic]) 

148 if removed: 

149 logger.debug("Unsubscribed '%s' from topic '%s'", handler_name, topic) 

150 return removed 

151 

152 def unsubscribe_all(self, handler_name: str) -> int: 

153 """Remove all subscriptions for a handler across topics.""" 

154 count = 0 

155 for topic in list(self._subscriptions.keys()): 

156 if self.unsubscribe(topic, handler_name): 

157 count += 1 

158 return count 

159 

160 # ----- publishing ----- 

161 

162 async def publish(self, event: Event) -> int: 

163 """Publish an event to all matching subscribers. 

164 

165 Returns: number of subscribers the event was dispatched to. 

166 """ 

167 matching: List[Subscription] = [] 

168 

169 # Exact topic match first 

170 if event.topic in self._subscriptions: 

171 matching.extend(self._subscriptions[event.topic]) 

172 

173 # Wildcard pattern matches 

174 for topic, subs in self._subscriptions.items(): 

175 if "*" in topic: 

176 for sub in subs: 

177 if sub.matches(event.topic) and sub not in matching: 

178 matching.append(sub) 

179 

180 # Sort by priority — higher priority processed first 

181 matching.sort(key=lambda s: event.priority.value, reverse=True) 

182 

183 for sub in matching: 

184 await self._queue.put((event, sub)) 

185 

186 if matching: 

187 self._event_count += 1 

188 logger.debug( 

189 "Published '%s' to %d subscribers [total events: %d]", 

190 event.topic, len(matching), self._event_count, 

191 ) 

192 

193 return len(matching) 

194 

195 async def publish_nowait(self, event: Event) -> int: 

196 """Non-blocking publish — drops event if queue is full.""" 

197 try: 

198 return await asyncio.wait_for(self.publish(event), timeout=0.1) 

199 except asyncio.TimeoutError: 

200 logger.warning("Event bus queue full — event '%s' dropped", event.topic) 

201 return 0 

202 

203 def emit_sync(self, event: Event) -> None: 

204 """Fire-and-forget from sync context.""" 

205 try: 

206 loop = asyncio.get_running_loop() 

207 loop.create_task(self.publish(event)) 

208 except RuntimeError: 

209 logger.warning("No running event loop — event '%s' not published", event.topic) 

210 

211 # ----- processing ----- 

212 

213 async def start(self) -> None: 

214 """Start background workers.""" 

215 if self._running: 

216 return 

217 self._running = True 

218 self._workers = [ 

219 asyncio.create_task(self._worker(i)) 

220 for i in range(self._worker_count) 

221 ] 

222 logger.info("EventBus started with %d workers", self._worker_count) 

223 

224 async def stop(self, grace_period: float = 5.0) -> None: 

225 """Gracefully stop all workers, draining remaining events.""" 

226 self._running = False 

227 

228 # Wait for queue to drain 

229 try: 

230 await asyncio.wait_for(self._queue.join(), timeout=grace_period) 

231 except asyncio.TimeoutError: 

232 logger.warning( 

233 "EventBus shutdown timeout — %d events remaining in queue", 

234 self._queue.qsize(), 

235 ) 

236 

237 for w in self._workers: 

238 w.cancel() 

239 await asyncio.gather(*self._workers, return_exceptions=True) 

240 logger.info("EventBus stopped — processed %d events", self._event_count) 

241 

242 async def _worker(self, worker_id: int) -> None: 

243 """Background worker processing events from the queue.""" 

244 logger.debug("EventBus worker %d started", worker_id) 

245 

246 while self._running: 

247 try: 

248 event, sub = await asyncio.wait_for( 

249 self._queue.get(), timeout=0.5 

250 ) 

251 except asyncio.TimeoutError: 

252 continue 

253 except asyncio.CancelledError: 

254 break 

255 

256 try: 

257 await sub.handler(event) 

258 except Exception as exc: 

259 logger.error( 

260 "Handler '%s' failed for event '%s': %s", 

261 sub.handler_name, event.topic, exc, 

262 ) 

263 if self._dlq_enabled: 

264 self._add_to_dlq(event, sub, str(exc)) 

265 finally: 

266 self._queue.task_done() 

267 

268 def _add_to_dlq(self, event: Event, sub: Subscription, error: str) -> None: 

269 if len(self._dead_letters) >= self._dlq_max_size: 

270 self._dead_letters.pop(0) # Drop oldest 

271 self._dead_letters.append(DeadLetter( 

272 event=event, 

273 handler_name=sub.handler_name, 

274 error=error, 

275 )) 

276 

277 # ----- DLQ operations ----- 

278 

279 def get_dlq(self) -> List[DeadLetter]: 

280 return list(self._dead_letters) 

281 

282 async def replay_dlq(self, max_events: int = 100) -> int: 

283 """Replay dead letter events.""" 

284 to_replay = self._dead_letters[:max_events] 

285 self._dead_letters = self._dead_letters[max_events:] 

286 

287 count = 0 

288 for dl in to_replay: 

289 dl.retry_count += 1 

290 await self._queue.put((dl.event, Subscription( 

291 topic_pattern=dl.event.topic, 

292 handler=lambda e: None, # Original handler not preserved 

293 handler_name=f"dlq_replay_{dl.handler_name}", 

294 ))) 

295 count += 1 

296 

297 logger.info("Replayed %d dead letter events", count) 

298 return count 

299 

300 def clear_dlq(self) -> int: 

301 count = len(self._dead_letters) 

302 self._dead_letters.clear() 

303 logger.info("Cleared %d dead letter events", count) 

304 return count 

305 

306 # ----- inspection ----- 

307 

308 @property 

309 def queue_size(self) -> int: 

310 return self._queue.qsize() 

311 

312 @property 

313 def total_events(self) -> int: 

314 return self._event_count 

315 

316 @property 

317 def subscription_count(self) -> int: 

318 return sum(len(subs) for subs in self._subscriptions.values()) 

319 

320 def list_topics(self) -> List[str]: 

321 return sorted(self._subscriptions.keys()) 

322 

323 

324# ============================================================================ 

325# Helpers 

326# ============================================================================ 

327 

328def event(topic: str, **kwargs) -> Event: 

329 """Convenience factory for creating events.""" 

330 return Event(topic=topic, **kwargs) 

331 

332 

333# ============================================================================ 

334# Module-level instance 

335# ============================================================================ 

336 

337default_bus = EventBus()