Coverage for agentos/core/event_bus.py: 0%
174 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-07-10 07:44 +0800
1"""AgentOS Event Bus — async pub/sub with backpressure control.
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
11Design: ~350 lines, zero external deps beyond stdlib + asyncio.
12"""
14from __future__ import annotations
16import asyncio
17import logging
18import time
19from collections import defaultdict
20from collections.abc import Awaitable, Callable
21from dataclasses import dataclass, field
22from enum import Enum
23from typing import Any
24from uuid import uuid4
26logger = logging.getLogger(__name__)
29# ============================================================================
30# Core types
31# ============================================================================
34class EventPriority(int, Enum):
35 LOW = 0
36 NORMAL = 50
37 HIGH = 100
38 CRITICAL = 200
41@dataclass
42class Event:
43 """An event published on the bus."""
45 topic: str
46 payload: Any = None
47 event_id: str = field(default_factory=lambda: str(uuid4())[:12])
48 timestamp: float = field(default_factory=time.time)
49 priority: EventPriority = EventPriority.NORMAL
50 source: str = ""
51 correlation_id: str = ""
52 metadata: dict[str, Any] = field(default_factory=dict)
55@dataclass
56class DeadLetter:
57 """Event that failed processing and was sent to DLQ."""
59 event: Event
60 handler_name: str
61 error: str
62 failed_at: float = field(default_factory=time.time)
63 retry_count: int = 0
66# ============================================================================
67# Subscription
68# ============================================================================
71@dataclass
72class Subscription:
73 """A handler subscribed to a topic pattern."""
75 topic_pattern: str
76 handler: Callable[[Event], Awaitable[Any]]
77 handler_name: str
78 concurrency: int = 1 # Max concurrent executions
79 is_pattern: bool = False # True if topic_pattern contains wildcards
81 def matches(self, topic: str) -> bool:
82 if not self.is_pattern:
83 return self.topic_pattern == topic
84 # Support * wildcard: "agent.*" matches "agent.start", "agent.stop"
85 pattern = self.topic_pattern.replace("*", ".*")
86 import re
88 return bool(re.match(f"^{pattern}$", topic))
91# ============================================================================
92# Event Bus
93# ============================================================================
96class EventBus:
97 """Async pub/sub event bus with backpressure and DLQ."""
99 def __init__(
100 self,
101 max_queue_size: int = 10000,
102 dlq_enabled: bool = True,
103 dlq_max_size: int = 1000,
104 worker_count: int = 4,
105 ):
106 self._subscriptions: dict[str, list[Subscription]] = defaultdict(list)
107 self._dead_letters: list[DeadLetter] = []
108 self._dlq_enabled = dlq_enabled
109 self._dlq_max_size = dlq_max_size
110 self._queue: asyncio.Queue[tuple[Event, Subscription]] = asyncio.Queue(
111 maxsize=max_queue_size
112 )
113 self._workers: list[asyncio.Task] = []
114 self._worker_count = worker_count
115 self._running = False
116 self._lock = asyncio.Lock()
117 self._event_count: int = 0
119 # ----- subscription management -----
121 def subscribe(
122 self,
123 topic: str,
124 handler: Callable[[Event], Awaitable[Any]],
125 handler_name: str | None = None,
126 concurrency: int = 1,
127 ) -> Subscription:
128 """Subscribe a handler to a topic.
130 Topic can contain * wildcard: 'agent.*', 'order.*.created'
131 """
132 name = handler_name or handler.__name__
133 is_pattern = "*" in topic
134 sub = Subscription(
135 topic_pattern=topic,
136 handler=handler,
137 handler_name=name,
138 concurrency=concurrency,
139 is_pattern=is_pattern,
140 )
141 self._subscriptions[topic].append(sub)
142 logger.debug("Subscribed '%s' to topic '%s'", name, topic)
143 return sub
145 def unsubscribe(self, topic: str, handler_name: str) -> bool:
146 """Remove a subscription by handler name."""
147 subs = self._subscriptions.get(topic, [])
148 original_len = len(subs)
149 self._subscriptions[topic] = [s for s in subs if s.handler_name != handler_name]
150 removed = original_len != len(self._subscriptions[topic])
151 if removed:
152 logger.debug("Unsubscribed '%s' from topic '%s'", handler_name, topic)
153 return removed
155 def unsubscribe_all(self, handler_name: str) -> int:
156 """Remove all subscriptions for a handler across topics."""
157 count = 0
158 for topic in list(self._subscriptions.keys()):
159 if self.unsubscribe(topic, handler_name):
160 count += 1
161 return count
163 # ----- publishing -----
165 async def publish(self, event: Event) -> int:
166 """Publish an event to all matching subscribers.
168 Returns: number of subscribers the event was dispatched to.
169 """
170 matching: list[Subscription] = []
172 # Exact topic match first
173 if event.topic in self._subscriptions:
174 matching.extend(self._subscriptions[event.topic])
176 # Wildcard pattern matches
177 for topic, subs in self._subscriptions.items():
178 if "*" in topic:
179 for sub in subs:
180 if sub.matches(event.topic) and sub not in matching:
181 matching.append(sub)
183 # Sort by priority — higher priority processed first
184 matching.sort(key=lambda s: event.priority.value, reverse=True)
186 for sub in matching:
187 await self._queue.put((event, sub))
189 if matching:
190 self._event_count += 1
191 logger.debug(
192 "Published '%s' to %d subscribers [total events: %d]",
193 event.topic,
194 len(matching),
195 self._event_count,
196 )
198 return len(matching)
200 async def publish_nowait(self, event: Event) -> int:
201 """Non-blocking publish — drops event if queue is full."""
202 try:
203 return await asyncio.wait_for(self.publish(event), timeout=0.1)
204 except TimeoutError:
205 logger.warning("Event bus queue full — event '%s' dropped", event.topic)
206 return 0
208 def emit_sync(self, event: Event) -> None:
209 """Fire-and-forget from sync context."""
210 try:
211 loop = asyncio.get_running_loop()
212 loop.create_task(self.publish(event))
213 except RuntimeError:
214 logger.warning("No running event loop — event '%s' not published", event.topic)
216 # ----- processing -----
218 async def start(self) -> None:
219 """Start background workers."""
220 if self._running:
221 return
222 self._running = True
223 self._workers = [asyncio.create_task(self._worker(i)) for i in range(self._worker_count)]
224 logger.info("EventBus started with %d workers", self._worker_count)
226 async def stop(self, grace_period: float = 5.0) -> None:
227 """Gracefully stop all workers, draining remaining events."""
228 self._running = False
230 # Wait for queue to drain
231 try:
232 await asyncio.wait_for(self._queue.join(), timeout=grace_period)
233 except TimeoutError:
234 logger.warning(
235 "EventBus shutdown timeout — %d events remaining in queue",
236 self._queue.qsize(),
237 )
239 for w in self._workers:
240 w.cancel()
241 await asyncio.gather(*self._workers, return_exceptions=True)
242 logger.info("EventBus stopped — processed %d events", self._event_count)
244 async def _worker(self, worker_id: int) -> None:
245 """Background worker processing events from the queue."""
246 logger.debug("EventBus worker %d started", worker_id)
248 while self._running:
249 try:
250 event, sub = await asyncio.wait_for(self._queue.get(), timeout=0.5)
251 except TimeoutError:
252 continue
253 except asyncio.CancelledError:
254 break
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,
262 event.topic,
263 exc,
264 )
265 if self._dlq_enabled:
266 self._add_to_dlq(event, sub, str(exc))
267 finally:
268 self._queue.task_done()
270 def _add_to_dlq(self, event: Event, sub: Subscription, error: str) -> None:
271 if len(self._dead_letters) >= self._dlq_max_size:
272 self._dead_letters.pop(0) # Drop oldest
273 self._dead_letters.append(
274 DeadLetter(
275 event=event,
276 handler_name=sub.handler_name,
277 error=error,
278 )
279 )
281 # ----- DLQ operations -----
283 def get_dlq(self) -> list[DeadLetter]:
284 return list(self._dead_letters)
286 async def replay_dlq(self, max_events: int = 100) -> int:
287 """Replay dead letter events."""
288 to_replay = self._dead_letters[:max_events]
289 self._dead_letters = self._dead_letters[max_events:]
291 count = 0
292 for dl in to_replay:
293 dl.retry_count += 1
294 await self._queue.put(
295 (
296 dl.event,
297 Subscription(
298 topic_pattern=dl.event.topic,
299 handler=lambda e: None, # Original handler not preserved
300 handler_name=f"dlq_replay_{dl.handler_name}",
301 ),
302 )
303 )
304 count += 1
306 logger.info("Replayed %d dead letter events", count)
307 return count
309 def clear_dlq(self) -> int:
310 count = len(self._dead_letters)
311 self._dead_letters.clear()
312 logger.info("Cleared %d dead letter events", count)
313 return count
315 # ----- inspection -----
317 @property
318 def queue_size(self) -> int:
319 return self._queue.qsize()
321 @property
322 def total_events(self) -> int:
323 return self._event_count
325 @property
326 def subscription_count(self) -> int:
327 return sum(len(subs) for subs in self._subscriptions.values())
329 def list_topics(self) -> list[str]:
330 return sorted(self._subscriptions.keys())
333# ============================================================================
334# Helpers
335# ============================================================================
338def event(topic: str, **kwargs) -> Event:
339 """Convenience factory for creating events."""
340 return Event(topic=topic, **kwargs)
343# ============================================================================
344# Module-level instance
345# ============================================================================
347default_bus = EventBus()