Coverage for agentos/tools/event_bus.py: 91%

1""" 

2Event Bus / Pub-Sub for AgentOS. 

3 

4EventBus — in-process pub/sub with topic wildcards, async dispatch, and replay. 

5supports: exact topic match, `*` single-level, `**` multi-level wildcards. 

6""" 

7 

8import fnmatch 

9import threading 

10import time 

11from collections import defaultdict 

12from collections.abc import Callable 

13from dataclasses import dataclass, field 

14from typing import Any 

15 

16# ============================================================================ 

17# Event 

18# ============================================================================ 

19 

20 

21@dataclass 

22class Event: 

23 """A published event with topic, payload, and metadata.""" 

24 

25 topic: str 

26 data: Any = None 

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

28 source: str = "" 

29 

30 

31Subscriber = Callable[[Event], None] 

32UnsubscribeHandle = Callable[[], None] 

33 

34 

35# ============================================================================ 

36# EventBus 

37# ============================================================================ 

38 

39 

40class EventBus: 

41 """Thread-safe in-process pub/sub event bus. 

42 

43 Topics use dot-separated paths: 'agent.tool.call', 'system.shutdown'. 

44 Wildcards: 'agent.*.call' (single level), 'agent.**' (multi-level). 

45 """ 

46 

47 def __init__(self): 

48 self._subscribers: dict[str, list[Subscriber]] = defaultdict(list) 

49 self._lock = threading.RLock() 

50 self._history: list[Event] = [] 

51 self._max_history: int = 1000 

52 self._total_published: int = 0 

53 self._total_delivered: int = 0 

54 self._running: bool = True 

55 

56 def subscribe(self, topic: str, callback: Subscriber) -> UnsubscribeHandle: 

57 """Subscribe to a topic (supports wildcards). Returns unsubscribe handle.""" 

58 with self._lock: 

59 self._subscribers[topic].append(callback) 

60 # Return a closure that removes this specific callback 

61 sub_list = self._subscribers[topic] 

62 

63 def unsubscribe(): 

64 with self._lock: 

65 if callback in sub_list: 

66 sub_list.remove(callback) 

67 

68 return unsubscribe 

69 

70 def unsubscribe(self, topic: str, callback: Subscriber) -> bool: 

71 """Remove a specific subscriber. Returns True if found and removed.""" 

72 with self._lock: 

73 subs = self._subscribers.get(topic) 

74 if subs and callback in subs: 

75 subs.remove(callback) 

76 return True 

77 return False 

78 

79 def publish(self, topic: str, data: Any = None, source: str = "") -> None: 

80 """Publish an event to all matching subscribers synchronously.""" 

81 event = Event(topic=topic, data=data, source=source) 

82 with self._lock: 

83 self._history.append(event) 

84 if len(self._history) > self._max_history: 

85 self._history = self._history[-self._max_history :] 

86 self._total_published += 1 

87 

88 if not self._running: 

89 return 

90 

91 for sub_topic, callbacks in list(self._subscribers.items()): 

92 if self._topic_match(sub_topic, topic): 

93 for cb in callbacks[:]: # copy for safe iteration 

94 try: 

95 cb(event) 

96 self._total_delivered += 1 

97 except Exception: 

98 pass 

99 

100 def _topic_match(self, pattern: str, topic: str) -> bool: 

101 """Match a topic against a pattern with wildcard support. 

102 - `*` matches a single level (dot-delimited). 

103 - `**` matches zero or more levels. 

104 """ 

105 if "**" in pattern or "*" in pattern: 

106 return self._wildcard_match(pattern, topic) 

107 return pattern == topic 

108 

109 def _wildcard_match(self, pattern: str, topic: str) -> bool: 

110 """fnmatch-style glob matching on dot-delimited topic paths.""" 

111 return fnmatch.fnmatch(topic, pattern) 

112 

113 def get_history(self, limit: int = 100) -> list[Event]: 

114 """Get recent published events.""" 

115 with self._lock: 

116 return list(self._history[-limit:]) 

117 

118 def clear_history(self) -> None: 

119 with self._lock: 

120 self._history.clear() 

121 

122 def stop(self) -> None: 

123 """Stop delivering events (still records history).""" 

124 with self._lock: 

125 self._running = False 

126 

127 def start(self) -> None: 

128 with self._lock: 

129 self._running = True 

130 

131 def subscriber_count(self, topic: str | None = None) -> int: 

132 """Count subscribers. If topic given, counts for that pattern only.""" 

133 with self._lock: 

134 if topic: 

135 return len(self._subscribers.get(topic, [])) 

136 return sum(len(v) for v in self._subscribers.values()) 

137 

138 @property 

139 def stats(self) -> dict[str, Any]: 

140 with self._lock: 

141 return { 

142 "total_published": self._total_published, 

143 "total_delivered": self._total_delivered, 

144 "subscriber_count": self.subscriber_count(), 

145 "history_size": len(self._history), 

146 "topics": list(self._subscribers.keys()), 

147 } 

148 

149 

150# ============================================================================ 

151# TopicFilter 

152# ============================================================================ 

153 

154 

155class TopicFilter: 

156 """Pre-compiled topic filter chain for high-throughput event routing.""" 

157 

158 def __init__(self): 

159 self._filters: dict[str, Callable[[Event], bool]] = {} 

160 self._lock = threading.Lock() 

161 

162 def add(self, name: str, predicate: Callable[[Event], bool]) -> None: 

163 with self._lock: 

164 self._filters[name] = predicate 

165 

166 def remove(self, name: str) -> bool: 

167 with self._lock: 

168 return self._filters.pop(name, None) is not None 

169 

170 def evaluate(self, event: Event) -> list[str]: 

171 """Return names of all matching filters for this event.""" 

172 matches = [] 

173 with self._lock: 

174 for name, pred in self._filters.items(): 

175 try: 

176 if pred(event): 

177 matches.append(name) 

178 except Exception: 

179 pass 

180 return matches 

181 

182 

183# ============================================================================ 

184# Global singleton 

185# ============================================================================ 

186 

187_default_bus: EventBus | None = None 

188_default_lock = threading.Lock() 

189 

190 

191def get_event_bus() -> EventBus: 

192 """Get or create the global default EventBus.""" 

193 global _default_bus 

194 if _default_bus is None: 

195 with _default_lock: 

196 if _default_bus is None: 

197 _default_bus = EventBus() 

198 return _default_bus