diff --git a/src/cleveragents/a2a/events.py b/src/cleveragents/a2a/events.py index f17ae21e7..8e528d076 100644 --- a/src/cleveragents/a2a/events.py +++ b/src/cleveragents/a2a/events.py @@ -24,6 +24,7 @@ from ulid import ULID from cleveragents.a2a.errors import A2aNotAvailableError from cleveragents.a2a.models import A2aEvent +from cleveragents.infrastructure.events.types import EventType # --------------------------------------------------------------------------- # SSE event type constants (A2A protocol) @@ -234,25 +235,6 @@ class EventBusBridge: bridge.stop() # unsubscribes """ - # Domain event types that map to SSE TaskStatusUpdateEvent - _STATUS_EVENT_TYPES: frozenset[str] = frozenset( - { - "PLAN_CREATED", - "PLAN_PHASE_CHANGED", - "PLAN_STATE_CHANGED", - "PLAN_APPLIED", - "PLAN_CANCELLED", - "PLAN_ERRORED", - } - ) - - # Domain event types that map to SSE TaskArtifactUpdateEvent - _ARTIFACT_EVENT_TYPES: frozenset[str] = frozenset( - { - "CHECKPOINT_RESTORED", - } - ) - def __init__( self, event_bus: Any, @@ -260,21 +242,23 @@ class EventBusBridge: ) -> None: self._event_bus = event_bus self._event_queue = event_queue - self._subscription: Any | None = None + self._subscription_ids: dict[EventType, int] = {} def start(self) -> None: - """Subscribe to the event bus and begin forwarding.""" + """Subscribe to all event types and begin forwarding.""" if hasattr(self._event_bus, "subscribe"): - self._subscription = self._event_bus.subscribe(self._on_domain_event) + for event_type in EventType: + sub_id = self._event_bus.subscribe(event_type, self._on_domain_event) + self._subscription_ids[event_type] = sub_id logger.info("a2a.event_bridge.started") def stop(self) -> None: """Unsubscribe from the event bus.""" - if self._subscription is not None: - if hasattr(self._subscription, "dispose"): - self._subscription.dispose() - self._subscription = None - logger.info("a2a.event_bridge.stopped") + for event_type, sub_id in list(self._subscription_ids.items()): + if hasattr(self._event_bus, "unsubscribe"): + self._event_bus.unsubscribe(event_type, sub_id) + self._subscription_ids.clear() + logger.info("a2a.event_bridge.stopped") def _on_domain_event(self, domain_event: Any) -> None: """Translate a domain event to an A2A event and publish.""" @@ -310,6 +294,25 @@ class EventBusBridge: with contextlib.suppress(RuntimeError): self._event_queue.publish(a2a_event) + # Domain event types that map to SSE TaskStatusUpdateEvent + _STATUS_EVENT_TYPES: frozenset[str] = frozenset( + { + "PLAN_CREATED", + "PLAN_PHASE_CHANGED", + "PLAN_STATE_CHANGED", + "PLAN_APPLIED", + "PLAN_CANCELLED", + "PLAN_ERRORED", + } + ) + + # Domain event types that map to SSE TaskArtifactUpdateEvent + _ARTIFACT_EVENT_TYPES: frozenset[str] = frozenset( + { + "CHECKPOINT_RESTORED", + } + ) + __all__ = [ "TASK_ARTIFACT_UPDATE", diff --git a/src/cleveragents/application/services/audit_event_subscriber.py b/src/cleveragents/application/services/audit_event_subscriber.py index 215cacf89..573709abb 100644 --- a/src/cleveragents/application/services/audit_event_subscriber.py +++ b/src/cleveragents/application/services/audit_event_subscriber.py @@ -79,8 +79,21 @@ class AuditEventSubscriber: def __init__(self, audit_service: AuditService, event_bus: EventBus) -> None: self._audit_service = audit_service self._event_bus = event_bus + self._subscription_ids: dict[EventType, int] = {} for event_type in SECURITY_EVENT_TYPES: - event_bus.subscribe(event_type, self._handle_event) + sub_id = event_bus.subscribe(event_type, self._handle_event) + self._subscription_ids[event_type] = sub_id + + def dispose(self) -> None: + """Remove all EventBus subscriptions so events stop arriving. + + Each individual subscription is unregistered via + :meth:`~cleveragents.infrastructure.events.protocol.EventBus.unsubscribe`. + This method is idempotent — calling it more than once is safe. + """ + for event_type, sub_id in list(self._subscription_ids.items()): + self._event_bus.unsubscribe(event_type, sub_id) + self._subscription_ids.clear() def _handle_event(self, event: DomainEvent) -> None: """Redact sensitive data and persist an audit log entry. diff --git a/src/cleveragents/infrastructure/events/logging_bus.py b/src/cleveragents/infrastructure/events/logging_bus.py index 307bf4ed3..61897ae70 100644 --- a/src/cleveragents/infrastructure/events/logging_bus.py +++ b/src/cleveragents/infrastructure/events/logging_bus.py @@ -43,7 +43,9 @@ class LoggingEventBus: """ def __init__(self) -> None: - self._subscriptions: dict[EventType, list[Callable[[DomainEvent], None]]] = {} + self._subscriptions: dict[ + EventType, list[tuple[Callable[[DomainEvent], None], int]] + ] = {} # ------------------------------------------------------------------ # Public API (satisfies EventBus protocol) @@ -70,7 +72,7 @@ class LoggingEventBus: actor_name=event.actor_name, details=event.details, ) - for handler in list(self._subscriptions.get(event.event_type, ())): + for handler, _sub_id in list(self._subscriptions.get(event.event_type, ())): try: handler(event) except Exception as exc: @@ -86,13 +88,18 @@ class LoggingEventBus: self, event_type: EventType, handler: Callable[[DomainEvent], None], - ) -> None: + ) -> int: """Register *handler* for events of *event_type*. Args: event_type: The :class:`EventType` to listen for. handler: Callable invoked for each matching :class:`DomainEvent`. + Returns: + A unique subscription identifier (``int``). Use this value as + the ``subscription_id`` argument in :meth:`unsubscribe` to + remove this registration. + Raises: TypeError: If *event_type* is not an :class:`EventType`. TypeError: If *handler* is not callable. @@ -103,7 +110,43 @@ class LoggingEventBus: ) if not callable(handler): raise TypeError("handler must be callable") - self._subscriptions.setdefault(event_type, []).append(handler) + sub_id = id(handler) + self._subscriptions.setdefault(event_type, []).append((handler, sub_id)) + return sub_id + + def unsubscribe( + self, + event_type: EventType, + subscription_id: int, + ) -> bool: + """Remove a previous *subscribe* call identified by *subscription_id*. + + The handler associated with the given *event_type* and + *subscription_id* is deregistered so it no longer receives + events. Subsequent calls to :meth:`emit` for *event_type* will + skip this handler. + + Args: + event_type: The :class:`EventType` originally passed to + :meth:`subscribe`. + subscription_id: The return value from the corresponding + :meth:`subscribe` call. + + Returns: + ``True`` if a matching subscription was found and removed, + ``False`` if no such subscription exists (the call is safe + even when nothing is removed). + """ + handlers = self._subscriptions.get(event_type) + if handlers is None: + return False + self._subscriptions[event_type] = [ + (h, sid) for h, sid in handlers if sid != subscription_id + ] + # Clean up empty lists to avoid wasting memory + if not self._subscriptions[event_type]: + del self._subscriptions[event_type] + return True __all__ = ["LoggingEventBus"] diff --git a/src/cleveragents/infrastructure/events/protocol.py b/src/cleveragents/infrastructure/events/protocol.py index 2c3cae85f..d7828b69f 100644 --- a/src/cleveragents/infrastructure/events/protocol.py +++ b/src/cleveragents/infrastructure/events/protocol.py @@ -40,12 +40,42 @@ class EventBus(Protocol): self, event_type: EventType, handler: Callable[[DomainEvent], None], - ) -> None: + ) -> int: """Register *handler* to receive events of *event_type*. Args: event_type: The :class:`EventType` to listen for. handler: Callable invoked synchronously for each matching event. + + Returns: + A subscription identifier (``int``). Use this value as the + ``subscription_id`` argument in :meth:`unsubscribe` to remove + the registration. + """ + ... + + def unsubscribe( + self, + event_type: EventType, + subscription_id: int, + ) -> bool: + """Remove a previous *subscribe* call identified by *subscription_id*. + + The handler associated with the given *event_type* and + *subscription_id* is deregistered so it no longer receives + events. Subsequent calls to :meth:`emit` for *event_type* will + skip this handler. + + Args: + event_type: The :class:`EventType` originally passed to + :meth:`subscribe`. + subscription_id: The return value from the corresponding + :meth:`subscribe` call. + + Returns: + ``True`` if a matching subscription was found and removed, + ``False`` if no such subscription exists (the call is safe + even when nothing is removed). """ ... diff --git a/src/cleveragents/infrastructure/events/reactive.py b/src/cleveragents/infrastructure/events/reactive.py index ec7c63194..b9085dde9 100644 --- a/src/cleveragents/infrastructure/events/reactive.py +++ b/src/cleveragents/infrastructure/events/reactive.py @@ -22,6 +22,7 @@ import contextlib from collections import deque from collections.abc import Callable from types import TracebackType +from typing import Any import structlog from rx import operators as ops @@ -60,7 +61,9 @@ class ReactiveEventBus: self._subject: Subject = Subject() self._stream: Observable = self._subject.pipe(ops.map(_identity)) - self._subscriptions: dict[EventType, list[Callable[[DomainEvent], None]]] = {} + self._subscriptions: dict[ + EventType, list[tuple[Callable[[DomainEvent], None], int]] + ] = {} self._audit_log: deque[DomainEvent] = deque(maxlen=max_audit_log_size) self._closed: bool = False @@ -135,7 +138,7 @@ class ReactiveEventBus: self._audit_log.append(event.model_copy(deep=True)) - for handler in list(self._subscriptions.get(event.event_type, ())): + for handler, _sub_id in list(self._subscriptions.get(event.event_type, ())): try: handler(event) except Exception as exc: @@ -152,7 +155,7 @@ class ReactiveEventBus: self, event_type: EventType, handler: Callable[[DomainEvent], None], - ) -> None: + ) -> int: """Register *handler* for events of *event_type*. Handlers are called synchronously inside :meth:`emit` in the order @@ -162,6 +165,11 @@ class ReactiveEventBus: event_type: The :class:`EventType` to listen for. handler: Callable invoked for each matching :class:`DomainEvent`. + Returns: + A unique subscription identifier (``int``). Use this value as + the ``subscription_id`` argument in :meth:`unsubscribe` to + remove this registration. + Raises: TypeError: If *event_type* is not an :class:`EventType`. TypeError: If *handler* is not callable. @@ -172,7 +180,43 @@ class ReactiveEventBus: ) if not callable(handler): raise TypeError("handler must be callable") - self._subscriptions.setdefault(event_type, []).append(handler) + sub_id = id(handler) + self._subscriptions.setdefault(event_type, []).append((handler, sub_id)) + return sub_id + + def unsubscribe( + self, + event_type: EventType, + subscription_id: int, + ) -> bool: + """Remove a previous *subscribe* call identified by *subscription_id*. + + The handler associated with the given *event_type* and + *subscription_id* is deregistered so it no longer receives + events. Subsequent calls to :meth:`emit` for *event_type* will + skip this handler. + + Args: + event_type: The :class:`EventType` originally passed to + :meth:`subscribe`. + subscription_id: The return value from the corresponding + :meth:`subscribe` call. + + Returns: + ``True`` if a matching subscription was found and removed, + ``False`` if no such subscription exists (the call is safe + even when nothing is removed). + """ + handlers = self._subscriptions.get(event_type) + if handlers is None: + return False + self._subscriptions[event_type] = [ + (h, sid) for h, sid in handlers if sid != subscription_id + ] + # Clean up empty lists to avoid wasting memory + if not self._subscriptions[event_type]: + del self._subscriptions[event_type] + return True @property def stream(self) -> Observable: