94 lines
3.2 KiB
Python
94 lines
3.2 KiB
Python
"""In-process message bus for Agent OS agents and task events."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass, field
|
|
import time
|
|
import uuid
|
|
from typing import Any, Callable
|
|
|
|
|
|
MessageHandler = Callable[["AgentMessage"], None]
|
|
|
|
|
|
@dataclass
|
|
class AgentMessage:
|
|
sender: str
|
|
recipient: str
|
|
message_type: str
|
|
payload: dict[str, Any] = field(default_factory=dict)
|
|
message_id: str = field(default_factory=lambda: str(uuid.uuid4()))
|
|
created_at: float = field(default_factory=time.time)
|
|
|
|
def snapshot(self) -> dict[str, Any]:
|
|
return {
|
|
"message_id": self.message_id,
|
|
"sender": self.sender,
|
|
"recipient": self.recipient,
|
|
"message_type": self.message_type,
|
|
"payload": dict(self.payload),
|
|
"created_at": self.created_at,
|
|
}
|
|
|
|
|
|
@dataclass
|
|
class TaskEvent:
|
|
task_id: str
|
|
event_type: str
|
|
source: str
|
|
detail: dict[str, Any] = field(default_factory=dict)
|
|
created_at: float = field(default_factory=time.time)
|
|
|
|
def to_message(self) -> AgentMessage:
|
|
return AgentMessage(
|
|
sender=self.source,
|
|
recipient="task-events",
|
|
message_type=self.event_type,
|
|
payload={"task_id": self.task_id, "detail": self.detail, "created_at": self.created_at},
|
|
)
|
|
|
|
|
|
class MessageBus:
|
|
def __init__(self) -> None:
|
|
self._subscribers: dict[str, list[MessageHandler]] = {}
|
|
self._routes: dict[str, str] = {}
|
|
self._history: list[dict[str, Any]] = []
|
|
|
|
def subscribe(self, topic: str, handler: MessageHandler) -> None:
|
|
self._subscribers.setdefault(topic, []).append(handler)
|
|
|
|
def unsubscribe(self, topic: str, handler: MessageHandler) -> bool:
|
|
handlers = self._subscribers.get(topic, [])
|
|
if handler not in handlers:
|
|
return False
|
|
handlers.remove(handler)
|
|
return True
|
|
|
|
def route(self, message_type: str, topic: str) -> None:
|
|
self._routes[message_type] = topic
|
|
|
|
def publish(self, topic: str, message: AgentMessage) -> dict[str, Any]:
|
|
self._history.append({"topic": topic, **message.snapshot()})
|
|
delivered = 0
|
|
for handler in list(self._subscribers.get(topic, [])):
|
|
handler(message)
|
|
delivered += 1
|
|
return {"topic": topic, "message_id": message.message_id, "delivered": delivered}
|
|
|
|
def send(self, message: AgentMessage) -> dict[str, Any]:
|
|
topic = self._routes.get(message.message_type, message.recipient)
|
|
return self.publish(topic, message)
|
|
|
|
def broadcast(self, sender: str, message_type: str, payload: dict[str, Any] | None = None) -> dict[str, Any]:
|
|
message = AgentMessage(sender=sender, recipient="broadcast", message_type=message_type, payload=dict(payload or {}))
|
|
total = 0
|
|
for topic in list(self._subscribers):
|
|
total += self.publish(topic, message)["delivered"]
|
|
return {"message_id": message.message_id, "delivered": total}
|
|
|
|
def emit_task_event(self, event: TaskEvent) -> dict[str, Any]:
|
|
return self.send(event.to_message())
|
|
|
|
def history(self, limit: int = 100) -> list[dict[str, Any]]:
|
|
return self._history[-limit:]
|