From c380922c8dc6220676476e0f475dd9c3a3524c14 Mon Sep 17 00:00:00 2001 From: Rico Schrage Date: Sat, 16 May 2026 01:02:47 +0200 Subject: [PATCH 01/10] Adding session construct, adding gathering for conv. transactional messages, some minor conv APIs. --- mango/__init__.py | 6 + mango/agent/conversation.py | 161 ++++++++++++ mango/agent/core.py | 226 +++++++++++++++- mango/agent/decorators.py | 269 +++++++++++++++++++ mango/agent/role.py | 281 +++++++++++++++++++- mango/express/health.py | 150 +++++++++++ mango/express/topology.py | 42 ++- tests/unit_tests/role/conversation_test.py | 248 ++++++++++++++++++ tests/unit_tests/role/decorators_test.py | 286 +++++++++++++++++++++ tests/unit_tests/role/edge_health_test.py | 228 ++++++++++++++++ tests/unit_tests/role/emit_event_test.py | 88 +++++++ tests/unit_tests/role/gather_test.py | 281 ++++++++++++++++++++ 12 files changed, 2248 insertions(+), 18 deletions(-) create mode 100644 mango/agent/conversation.py create mode 100644 mango/agent/decorators.py create mode 100644 mango/express/health.py create mode 100644 tests/unit_tests/role/conversation_test.py create mode 100644 tests/unit_tests/role/decorators_test.py create mode 100644 tests/unit_tests/role/edge_health_test.py create mode 100644 tests/unit_tests/role/emit_event_test.py create mode 100644 tests/unit_tests/role/gather_test.py diff --git a/mango/__init__.py b/mango/__init__.py index 94df304..433f884 100644 --- a/mango/__init__.py +++ b/mango/__init__.py @@ -15,6 +15,12 @@ RoleContext, WaitingMessagePreprocessor, ) +from .agent.decorators import ( + on_message, + on_event, + periodic, +) +from .express.health import EdgeHealth, TopologyHealth from .container.factory import ( create_tcp as create_tcp_container, create_mqtt as create_mqtt_container, diff --git a/mango/agent/conversation.py b/mango/agent/conversation.py new file mode 100644 index 0000000..a1b40e1 --- /dev/null +++ b/mango/agent/conversation.py @@ -0,0 +1,161 @@ +"""Multi-hop conversation primitive. + +A *conversation* is a logically-grouped sequence of messages threaded +by a shared id. Mango's existing ``tracking_id`` mechanism solves the +single request/response case (see :meth:`AgentDelegates.send_tracked_message` +and :meth:`AgentDelegates.reply_to`), but multi-hop protocols — gossip, +auctions, holonic ADMM coordination — need the same id to route many +messages over time *without* being consumed on first reply. + +:class:`Conversation` is that abstraction: an async context manager +that holds a conversation id, a clock-aware timeout, mutable state, +and a receive queue. Anyone in the conversation (initiator or joiner) +can ``async for msg, meta in conversation`` and ``await +conversation.send(addr, payload)``. + +Two entry points on :class:`mango.RoleContext`: + +* :meth:`RoleContext.open_conversation` — initiator: generates a new id + and starts a fresh conversation. +* :meth:`RoleContext.join_conversation` — joiner: re-uses the id from + an inbound message's ``meta`` so the responder participates in the + same exchange. + +Both return a :class:`Conversation`. The timeout is enforced via the +agent's scheduler clock so simulation and real-time modes behave the +same way. + +Example — gossip-style initiator:: + + async with self.context.open_conversation( + timeout=10.0, + state={"target": -5.0, "delta": 0.0, "lambda": 0.01}, + ) as conv: + await conv.send(neighbours[0], GossipStep(payload=...)) + async for msg, meta in conv: + update(conv.state, msg) + if conv.state["delta"] >= conv.state["target"]: + conv.converge() # exits the loop on next iteration + continue + next_hop = pick(...) + await conv.send(next_hop, GossipStep(payload=...)) + +Example — joiner side:: + + @on_message(GossipStep) + async def on_step(self, msg, meta): + async with self.context.join_conversation(meta) as conv: + ... # same async for loop as above +""" + +from __future__ import annotations + +import asyncio +from typing import Any + +# Metadata key used to carry the conversation id alongside ``tracking_id``. +# Distinct from ``tracking_id`` because the latter is consumed by the +# single-shot reply machinery; a conversation may legitimately carry +# both (e.g. when a participant uses ``reply_to`` inside a session). +CONVERSATION_ID_KEY = "conversation_id" + + +class Conversation: + """Active conversation handle owned by one agent. + + The state dict is purely user-controlled — mango does not interpret + it. ``converge()`` ends the async iteration after the current + message; ``cancel()`` ends it immediately with no further yields. + Both are idempotent. + """ + + def __init__( + self, + *, + owner, + conversation_id: str, + state: dict[str, Any] | None = None, + timeout: float | None = None, + ) -> None: + self._owner = owner # RoleContext + self._conversation_id = conversation_id + self.state: dict[str, Any] = dict(state) if state else {} + self._timeout = timeout + self._queue: asyncio.Queue = asyncio.Queue() + self._converged: bool = False + self._cancelled: bool = False + self._timeout_handle = None # asyncio.Task | None + + # -- introspection ------------------------------------------------ + @property + def conversation_id(self) -> str: + return self._conversation_id + + @property + def closed(self) -> bool: + return self._converged or self._cancelled + + # -- public control ----------------------------------------------- + def converge(self) -> None: + """Mark the conversation as converged. The async iterator + will terminate after delivering any already-queued message.""" + self._converged = True + # Wake the receiver if it's idle waiting on the queue. + self._queue.put_nowait((_SENTINEL_END, None)) + + def cancel(self) -> None: + """Drop the conversation immediately — discards any queued + messages and terminates the async iterator on the next pull.""" + self._cancelled = True + self._queue.put_nowait((_SENTINEL_END, None)) + + # -- send helpers ------------------------------------------------- + async def send(self, receiver_addr, content: Any, **kwargs) -> bool: + """Send *content* tagged with this conversation's id.""" + meta_extra = {CONVERSATION_ID_KEY: self._conversation_id} + meta_extra.update(kwargs) + return await self._owner.send_message( + content, + receiver_addr=receiver_addr, + **meta_extra, + ) + + async def broadcast(self, receivers, content: Any, **kwargs) -> None: + """Send *content* to every receiver in *receivers*.""" + for addr in receivers: + await self.send(addr, content, **kwargs) + + # -- internal hooks (used by RoleContext) ------------------------- + def _on_inbound(self, content: Any, meta: dict) -> None: + if self.closed: + return + self._queue.put_nowait((content, meta)) + + def _fire_timeout(self) -> None: + if self.closed: + return + # Treat timeout as cancellation — partial state remains + # accessible to the caller after exiting the context. + self._cancelled = True + self._queue.put_nowait((_SENTINEL_END, None)) + + # -- iteration ---------------------------------------------------- + def __aiter__(self): + return self + + async def __anext__(self): + if self._cancelled: + raise StopAsyncIteration + content, meta = await self._queue.get() + if content is _SENTINEL_END: + raise StopAsyncIteration + if self._converged: + # Deliver this final message, then end on the next pull by + # putting a sentinel back into the queue. + self._queue.put_nowait((_SENTINEL_END, None)) + return content, meta + + +# Sentinel placed in the queue to end the async iterator. A module- +# level singleton so identity checks are cheap and unambiguous. +_SENTINEL_END = object() diff --git a/mango/agent/core.py b/mango/agent/core.py index cce36b0..a00e543 100644 --- a/mango/agent/core.py +++ b/mango/agent/core.py @@ -5,18 +5,24 @@ connections to other agents. """ +from __future__ import annotations + import asyncio import logging import uuid from abc import ABC from dataclasses import dataclass, field from enum import Enum -from typing import Any +from typing import TYPE_CHECKING, Any from ..messages.message import AgentAddress from ..util.clock import Clock from ..util.scheduling import ScheduledProcessTask, ScheduledTask, Scheduler +if TYPE_CHECKING: + from mango.agent.conversation import Conversation + from mango.express.health import TopologyHealth + logger = logging.getLogger(__name__) @@ -70,7 +76,7 @@ class TopologyNeighbor: def __init__( self, agent: Any, - description: "AgentDescription", + description: AgentDescription, characteristic: str = "", ) -> None: self._agent = agent @@ -78,7 +84,7 @@ def __init__( self.characteristic = characteristic @property - def address(self) -> "AgentAddress": + def address(self) -> AgentAddress: """The neighbor's current :class:`AgentAddress` (resolved lazily).""" return self._agent.addr @@ -98,6 +104,12 @@ def __init__(self) -> None: self._tid_to_characteristic: dict[str, str] = {} self._tid_to_connectors: dict[str, list[tuple[str, TopologyNeighbor]]] = {} self._marked_connector_for: list[str] = [] + # Optional :class:`TopologyHealth` instance per topology — see + # :mod:`mango.express.health`. Populated by ``Topology.inject`` + # for every topology that was configured with ``edge_health``. + # The same instance is shared across every agent in the + # topology (it stores scores for all directed edges). + self._tid_to_health: dict[str, TopologyHealth] = {} def neighbors( self, @@ -107,7 +119,7 @@ def neighbors( has_characteristic: str | None = None, include_connectors: tuple[str, ...] | list[str] = (), match_func: Any = None, - ) -> list["AgentAddress"]: + ) -> list[AgentAddress]: """Return addresses of neighbors in topology *tid* with edge *state*. :param state: only return neighbors reachable via edges in this state @@ -149,7 +161,7 @@ def connectors( *, include_connectors: tuple[str, ...] | list[str] = (), match_func: Any = None, - ) -> list["AgentAddress"]: + ) -> list[AgentAddress]: """Return addresses of connector agents for topology *tid*.""" result: list[AgentAddress] = [] for conn_type, n in self._tid_to_connectors.get(tid, []): @@ -164,6 +176,69 @@ def connection_types(self, tid: str = "default") -> list[str]: """Return the connection type labels for connectors in topology *tid*.""" return [ct for ct, _ in self._tid_to_connectors.get(tid, [])] + # ------------------------------------------------------------------ + # Edge-health queries (see mango.express.health) + # ------------------------------------------------------------------ + def has_health(self, tid: str = "default") -> bool: + """True when topology *tid* has continuous link-health enabled.""" + return tid in self._tid_to_health + + def health_runtime(self, tid: str = "default") -> TopologyHealth | None: + """Return the :class:`TopologyHealth` runtime for *tid*, or None.""" + return self._tid_to_health.get(tid) + + +class _GatherCollector: + """Aggregates multi-reply tracked responses for :meth:`Agent.open_gather`. + + The collector exposes a single :meth:`wait` coroutine that resolves + when either *expected* replies have arrived or :meth:`finish` has + been called explicitly (used by :meth:`RoleContext.gather` to + implement the timeout / quorum policy). Replies are stored under + the responding agent's :class:`AgentAddress` so the caller can + match each response to its source. + """ + + def __init__( + self, + *, + expected: int, + reply_type: type | tuple[type, ...] | None = None, + ) -> None: + self._expected = max(0, int(expected)) + self._reply_type = reply_type + self.responses: dict[AgentAddress, Any] = {} + # Lazily created so the collector can be constructed before a + # running event loop is required (Agent.open_gather runs in + # whichever context the caller is in). + self._done: asyncio.Event | None = None + + def _event(self) -> asyncio.Event: + if self._done is None: + self._done = asyncio.Event() + return self._done + + def on_reply(self, content: Any, meta: dict) -> None: + if self._reply_type is not None and not isinstance(content, self._reply_type): + return + sender_id = meta.get("sender_id") + sender_addr = meta.get("sender_addr") + addr = AgentAddress(protocol_addr=sender_addr, aid=sender_id) + # First reply per sender wins — late duplicates (e.g. retries) + # are dropped so the caller sees a stable mapping. + if addr in self.responses: + return + self.responses[addr] = content + if self._expected and len(self.responses) >= self._expected: + self._event().set() + + def finish(self) -> None: + """Force the collector to resolve (used on timeout / quorum hit).""" + self._event().set() + + async def wait(self) -> None: + await self._event().wait() + class AgentContext: def __init__(self, container) -> None: @@ -220,6 +295,17 @@ def __init__(self) -> None: self._description: AgentDescription = AgentDescription() self._forwarding_rules: list[ForwardingRule] = [] self._transaction_handlers: dict[str, tuple] = {} + # Multi-shot reply collectors keyed by tracking_id. Used by + # ``RoleContext.gather`` (and any caller of + # :meth:`Agent.open_gather`) to aggregate replies from many + # receivers under a single id without consuming the entry on + # the first reply. See :meth:`_handle_tracked_reply`. + self._gather_collectors: dict[str, _GatherCollector] = {} + # Open conversations keyed by conversation_id (see + # :mod:`mango.agent.conversation`). Messages whose meta + # carries a matching id are routed to the conversation's + # async-iterator queue. + self._conversations: dict[str, Conversation] = {} self._behavior_message_subs: list[tuple] = [] self._behavior_global_event_handlers: list[tuple] = [] self._behavior_agent_event_handlers: list[tuple] = [] @@ -410,18 +496,125 @@ def _handle_tracked_reply(self, content: Any, meta: dict) -> bool: """Internal: check if *meta* contains a tracked reply; call handler. Returns ``True`` if the message was handled as a tracked reply. + Two flavours are supported: + + * Single-shot ``_transaction_handlers`` entries (created by + :meth:`send_tracked_message`). Popped on first matching + reply. + * Multi-shot ``_gather_collectors`` entries (created by + :meth:`open_gather`). Not popped — every matching reply + is delivered to the collector until the caller closes it. """ tracking_id = meta.get("tracking_id") - if ( - tracking_id - and meta.get("reply") - and tracking_id in self._transaction_handlers - ): + if not tracking_id or not meta.get("reply"): + return False + if tracking_id in self._transaction_handlers: (handler,) = self._transaction_handlers.pop(tracking_id) handler(content, meta) return True + collector = self._gather_collectors.get(tracking_id) + if collector is not None: + collector.on_reply(content, meta) + return True return False + def _route_to_conversation(self, content: Any, meta: dict) -> None: + """Deliver *content*/*meta* to any open conversation whose id + matches ``meta[CONVERSATION_ID_KEY]``. No-op when the message + has no conversation id or no matching conversation exists. + """ + from mango.agent.conversation import CONVERSATION_ID_KEY + + conv_id = meta.get(CONVERSATION_ID_KEY) + if not conv_id: + return + conv = self._conversations.get(conv_id) + if conv is None: + return + conv._on_inbound(content, meta) + + def open_conversation(self, conv: Conversation) -> None: + """Register *conv* so inbound messages with its id are routed + to it. Used internally by ``RoleContext.open_conversation``; + end users should not need to call this directly.""" + if conv.conversation_id in self._conversations: + raise ValueError( + f"conversation {conv.conversation_id!r} already open on {self.aid}" + ) + self._conversations[conv.conversation_id] = conv + + def close_conversation(self, conv: Conversation) -> None: + """Unregister *conv*. Called when the context manager exits.""" + self._conversations.pop(conv.conversation_id, None) + + def _nudge_topology_health(self, meta: dict) -> None: + """Multiplicatively recover edge scores on every received message. + + Called once per inbox dequeue (before role-level dispatch). No-op + unless this agent participates in at least one topology that + was created with ``edge_health=...``. All clock reads go + through the agent's scheduler so the decay model is identical + under real-time and simulation clocks. + """ + svc = self.service_of_type(TopologyService, None) + if svc is None or not svc._tid_to_health: + return + sender_id = meta.get("sender_id") + sender_addr = meta.get("sender_addr") + if not sender_id: + return + peer_addr = AgentAddress(protocol_addr=sender_addr, aid=sender_id) + scheduler = getattr(self, "scheduler", None) + if scheduler is None or scheduler.clock is None: + return + now = scheduler.clock.time + my_addr = self.addr + for tid, health in svc._tid_to_health.items(): + # Only nudge if the sender is actually a neighbour of this + # agent in that topology — keeps unrelated traffic + # (cross-topology messages, broadcast lists) from inflating + # scores for non-neighbours. + if any( + n.address == peer_addr + for ns in svc._tid_to_state_to_neighbors.get(tid, {}).values() + for n in ns + ): + health.nudge(my_addr, peer_addr, now) + + def open_gather( + self, + tracking_id: str, + *, + expected: int, + reply_type: type | tuple[type, ...] | None = None, + ) -> _GatherCollector: + """Register a multi-reply collector for *tracking_id*. + + Returns an opaque collector handle; the caller awaits + :meth:`_GatherCollector.wait` and must call + :meth:`close_gather` (or use :meth:`RoleContext.gather`, which + does both) to release the registration. + + :param expected: number of replies that satisfies the collector + without waiting for the timeout. Used by ``gather`` to + return as soon as every recipient has answered. + :param reply_type: when set, only replies whose ``content`` is + an instance of *reply_type* are accepted; everything else + is dropped silently. Filters out tracking-id collisions + with unrelated reply traffic. + """ + if tracking_id in self._gather_collectors: + raise ValueError( + f"gather collector already open for tracking_id={tracking_id!r}" + ) + collector = _GatherCollector(expected=expected, reply_type=reply_type) + self._gather_collectors[tracking_id] = collector + return collector + + def close_gather(self, tracking_id: str) -> None: + """Release a previously opened gather collector.""" + self._gather_collectors.pop(tracking_id, None) + @property def current_timestamp(self) -> float: """ @@ -838,6 +1031,19 @@ async def _check_inbox(self): if not forwarded: # Check tracked reply handlers + # Update topology link-health (if any topology this + # agent belongs to has ``edge_health`` enabled). + # Done before user-defined dispatch so the role's + # ``@on_message`` handler sees a freshly-nudged + # neighbour score should it read one. + self._nudge_topology_health(meta) + # Route to any open conversation that matches the + # message's conversation_id. Doesn't suppress + # normal handler dispatch — a role can still react + # to the message through ``@on_message`` AND join + # the conversation; the routing simply ensures the + # async-iterator pulls the message too. + self._route_to_conversation(content, meta) if not self._handle_tracked_reply(content, meta): for _cond, _handler, _proc in self._behavior_message_subs: if _cond(content, meta): diff --git a/mango/agent/decorators.py b/mango/agent/decorators.py new file mode 100644 index 0000000..97af404 --- /dev/null +++ b/mango/agent/decorators.py @@ -0,0 +1,269 @@ +"""Declarative dispatch for :class:`mango.agent.role.Role`. + +Three decorators that move the mechanical parts of ``setup()`` out of +hand-written boilerplate and into class-level metadata: + +* :func:`on_message` — subscribe to a message type with an optional filter. +* :func:`on_event` — subscribe to a co-located event type. +* :func:`periodic` — schedule a periodic task at role-attach time. + +A role that uses these decorators does not need to implement +``setup()`` at all unless it has other initialisation logic. When +``setup()`` is implemented, the decorator wiring runs first (so the +explicit ``setup`` body can override or extend it). + +The decorators are pure annotations: they stash configuration on the +decorated method via the attribute name :data:`_MANGO_DISPATCH_META`. +The :class:`mango.agent.role.Role` base class collects them at class +creation time via ``__init_subclass__`` and replays the registrations +through :meth:`RoleContext.subscribe_message` / +:meth:`RoleContext.subscribe_event` / +:meth:`RoleContext.schedule_periodic_task` when the role is bound. + +Async coroutine handlers for ``on_message`` are scheduled as instant +tasks automatically — the underlying ``handle_message`` callback is +synchronous, so an async handler would otherwise raise +``RuntimeWarning: coroutine was never awaited``. This removes the +repeated ``def _wrap(coro_fn): ...`` shim every existing role declares. + +Example:: + + class Echo(Role): + @on_message(str, where=lambda self, c, m: c.startswith("ping")) + async def on_ping(self, content, meta): + await self.context.send_message("pong", receiver_addr=sender_addr(meta)) + + @on_event(MyEvent) + def on_event(self, event, src): + ... + + @periodic(every=1.0) + async def heartbeat(self): + ... +""" + +from __future__ import annotations + +import asyncio +from collections.abc import Callable +from dataclasses import dataclass, field +from typing import Any + +#: Attribute name on a decorated method that carries the dispatch metadata. +_MANGO_DISPATCH_META = "_mango_dispatch_meta" + + +@dataclass +class _Dispatch: + """Per-method dispatch metadata collected by the decorators. + + A single method may carry multiple subscriptions (e.g. one + ``@on_message`` and one ``@on_event``) by stacking decorators — + each adds an entry to the corresponding list. + """ + + message_subs: list[dict[str, Any]] = field(default_factory=list) + event_subs: list[dict[str, Any]] = field(default_factory=list) + periodic: list[dict[str, Any]] = field(default_factory=list) + + +def _get_or_create_meta(method: Callable) -> _Dispatch: + meta = getattr(method, _MANGO_DISPATCH_META, None) + if meta is None: + meta = _Dispatch() + try: + setattr(method, _MANGO_DISPATCH_META, meta) + except (AttributeError, TypeError): + # Bound methods and some descriptors are read-only — the + # decorators below are intended for regular functions on + # class bodies, where setattr always succeeds. + raise TypeError( + "mango dispatch decorators must be applied to plain " + "functions defined in a class body, not " + f"{type(method).__name__}" + ) from None + return meta + + +def on_message( + message_type: type, + *, + where: Callable[..., bool] | None = None, + priority: int = 0, +) -> Callable: + """Subscribe the decorated method to *message_type*. + + The handler is called as ``handler(self, content, meta)``. Async + handlers are scheduled as instant tasks automatically. + + :param message_type: only deliver messages where + ``isinstance(content, message_type)`` is true. + :param where: optional extra filter. Called as + ``where(self, content, meta) -> bool`` — receives ``self`` + so the filter can read role state (e.g. a sector tag) without + capturing it in a closure at class-define time. If ``None`` + any message of *message_type* is accepted. + :param priority: forwarded to :meth:`RoleContext.subscribe_message` + (lower runs first). + """ + + def decorator(method: Callable) -> Callable: + meta = _get_or_create_meta(method) + meta.message_subs.append( + { + "message_type": message_type, + "where": where, + "priority": priority, + } + ) + return method + + return decorator + + +def on_event(event_type: type) -> Callable: + """Subscribe the decorated method to a co-located event type. + + The handler is called as ``handler(self, event, source)`` and runs + synchronously inside :meth:`RoleContext.emit_event`. + """ + + def decorator(method: Callable) -> Callable: + meta = _get_or_create_meta(method) + meta.event_subs.append({"event_type": event_type}) + return method + + return decorator + + +def periodic( + every: float | str, + *, + only_if: Callable[..., bool] | None = None, +) -> Callable: + """Schedule the decorated coroutine method as a periodic task. + + :param every: period in seconds, or a string key looked up on the + role instance at attach time (e.g. ``"poll_period_s"`` reads + ``role.poll_period_s`` — useful for sector-dependent periods + configured per role instance). + :param only_if: optional predicate ``only_if(self) -> bool`` evaluated + at every firing. When false the task body is skipped — the + scheduler still runs but the handler returns early. This + replaces the "if not leader: return" guard that every periodic + coroutine in scare repeats by hand. + """ + + def decorator(method: Callable) -> Callable: + if not asyncio.iscoroutinefunction(method): + raise TypeError( + f"@periodic must decorate an async coroutine function; " + f"{method.__qualname__} is not a coroutine" + ) + meta = _get_or_create_meta(method) + meta.periodic.append({"every": every, "only_if": only_if}) + return method + + return decorator + + +def collect_dispatch(cls: type) -> dict[str, _Dispatch]: + """Walk ``cls`` (and bases) and collect every dispatch-decorated method. + + Returns a dict keyed by method name; the same method name can carry + multiple subscriptions because :class:`_Dispatch` is a list of + each kind. Subclass methods *replace* base-class entries with the + same name (standard MRO), but unrelated method names from a base + class are still collected — so a base role that declares + ``@periodic`` decorated methods is inherited by subclasses + transparently. + """ + out: dict[str, _Dispatch] = {} + for klass in reversed(cls.__mro__): + for name, attr in klass.__dict__.items(): + meta = getattr(attr, _MANGO_DISPATCH_META, None) + if isinstance(meta, _Dispatch): + out[name] = meta + return out + + +def _bind_async(method: Callable, role: Any) -> Callable: + """Return a sync callback that schedules *method* as an instant task. + + Used to bridge async ``@on_message`` handlers to the synchronous + ``subscribe_message`` callback contract. Mirrors the + ``def _wrap(coro_fn): ...`` shim every legacy role declares. + """ + + def _sync(content: Any, meta: dict) -> None: + coro = method(role, content, meta) + role.context.schedule_instant_task(coro) + + return _sync + + +def _bind_sync(method: Callable, role: Any) -> Callable: + """Return a sync callback that calls *method* directly (for sync handlers).""" + + def _sync(content: Any, meta: dict) -> None: + method(role, content, meta) + + return _sync + + +def apply_dispatch(role: Any) -> None: + """Apply every collected dispatch entry on ``role`` to the bound context. + + Called by :class:`Role._bind` after the role context is wired but + before user-defined ``setup()``. Splitting this out keeps the + decoration mechanism orthogonal from the Role base class — tests + can call ``apply_dispatch(role)`` against a mock context. + """ + dispatch = collect_dispatch(type(role)) + ctx = role.context + for name, meta in dispatch.items(): + method = getattr(type(role), name) # unbound function + for sub in meta.message_subs: + mtype = sub["message_type"] + where = sub["where"] + priority = sub["priority"] + if where is None: + + def condition(content, _meta, _mtype=mtype): + return isinstance(content, _mtype) + else: + + def condition(content, _meta, _mtype=mtype, _w=where, _r=role): + return isinstance(content, _mtype) and _w(_r, content, _meta) + + if asyncio.iscoroutinefunction(method): + callback = _bind_async(method, role) + else: + callback = _bind_sync(method, role) + ctx.subscribe_message(role, callback, condition, priority=priority) + for sub in meta.event_subs: + etype = sub["event_type"] + handler = getattr(role, name) + ctx.subscribe_event(role, etype, handler) + for sub in meta.periodic: + every = sub["every"] + only_if = sub["only_if"] + if isinstance(every, str): + delay = getattr(role, every) + else: + delay = float(every) + coro_method = getattr(role, name) + if only_if is None: + schedule = coro_method + else: + + def _gated(_role=role, _coro=coro_method, _gate=only_if): + async def _run(): + if not _gate(_role): + return + await _coro() + + return _run() + + schedule = _gated + ctx.schedule_periodic_task(schedule, delay=delay) diff --git a/mango/agent/role.py b/mango/agent/role.py index bb8d656..0e0d979 100644 --- a/mango/agent/role.py +++ b/mango/agent/role.py @@ -189,6 +189,12 @@ def __init__(self, scheduler): self._role_event_type_to_handler = {} self._scheduler = scheduler self._data = DataContainer() + # Back-reference to the owning :class:`RoleAgent`. Set by + # ``RoleAgent.__init__`` so role-level helpers like + # :meth:`RoleContext.gather` can reach the agent-level + # gather/transaction machinery without an indirection through + # the agent context. + self._agent: Agent | None = None def get_or_create_model(self, cls): """Creates or return (when already created) a central role model. @@ -352,8 +358,22 @@ def subscribe_send(self, role: Role, method: Callable): else: self._send_msg_subs[role] = [method] - def emit_event(self, event: Any, event_source: Any = None): - subs = self._role_event_type_to_handler[type(event)] + def emit_event(self, event: Any, event_source: Any = None, *, strict: bool = False): + """Dispatch *event* to every subscribed handler on this agent. + + :param strict: when True, raise :class:`KeyError` if no role is + subscribed to ``type(event)``. Default is False — events + without subscribers are silently dropped, matching the + fire-and-forget semantics callers expect from a + notification API and removing the ``try/except KeyError`` + guard pattern that otherwise wraps every ``emit_event`` + call site. + """ + subs = self._role_event_type_to_handler.get(type(event)) + if subs is None: + if strict: + raise KeyError(type(event)) + return for _, method in subs: method(event, event_source) @@ -491,15 +511,18 @@ async def send_message( **kwargs, ) - def emit_event(self, event: Any, event_source: Any = None): + def emit_event(self, event: Any, event_source: Any = None, *, strict: bool = False): """Emit an custom event to other roles. :param event: the event :type event: Any :param event_source: emitter of the event (mostly the emitting role), defaults to None :type event_source: Any, optional + :param strict: when True, raise :class:`KeyError` if no role + is subscribed to ``type(event)``. Default False — see + :meth:`RoleHandler.emit_event` for the rationale. """ - self._role_handler.emit_event(event, event_source) + self._role_handler.emit_event(event, event_source, strict=strict) def subscribe_event(self, role: Role, event_type: Any, handler_method: Callable): """Subscribe to specific event types. The listener will be evaluated based @@ -518,6 +541,199 @@ def deactivate(self, role) -> None: def activate(self, role) -> None: self._role_handler.activate(role) + # ------------------------------------------------------------------ + # Topology link-health queries (see mango.express.health) + # ------------------------------------------------------------------ + + def neighbour_score(self, neighbour_addr, *, tid: str = "default") -> float | None: + """Return the current edge health for one neighbour, or ``None`` + when the topology does not have ``edge_health`` enabled.""" + agent = self._role_handler._agent + if agent is None: + return None + from mango.agent.core import TopologyService + + svc = agent.service_of_type(TopologyService, None) + if svc is None: + return None + health = svc.health_runtime(tid) + if health is None: + return None + now = agent.scheduler.clock.time if agent.scheduler else 0.0 + return health.score(agent.addr, neighbour_addr, now) + + def live_neighbours( + self, + tid: str = "default", + *, + threshold: float | None = None, + ): + """Return the subset of ``topology_neighbors(tid=tid)`` whose + edge health is at or above *threshold*. + + Falls back to the full neighbour list when the topology has no + ``edge_health`` configured — callers can use this method + unconditionally without having to branch on whether tracking is + enabled. + """ + agent = self._role_handler._agent + if agent is None: + return [] + from mango.agent.core import State, TopologyService + + svc = agent.service_of_type(TopologyService, None) + if svc is None: + return [] + all_neighbours = svc.neighbors(state=State.NORMAL, tid=tid) + health = svc.health_runtime(tid) + if health is None: + return all_neighbours + now = agent.scheduler.clock.time if agent.scheduler else 0.0 + my_addr = agent.addr + return [ + n + for n in all_neighbours + if health.is_live(my_addr, n, now, threshold=threshold) + ] + + def open_conversation( + self, + *, + state: dict | None = None, + timeout: float | None = None, + ): + """Open a fresh multi-hop conversation as the initiator. + + Returns an async context manager whose body yields a + :class:`~mango.agent.conversation.Conversation`. A new id is + generated; pass it (or use :meth:`Conversation.send`) so other + agents can route their messages back. + + The optional *timeout* is enforced via the agent's scheduler + clock so behaviour is identical under :class:`AsyncioClock` + (real time) and :class:`ExternalClock` (simulation). + """ + import uuid as _uuid + + from mango.agent.conversation import Conversation + + return _ConversationContext( + role_context=self, + conv=Conversation( + owner=self, + conversation_id=str(_uuid.uuid4()), + state=state, + timeout=timeout, + ), + ) + + def join_conversation( + self, + meta: dict, + *, + state: dict | None = None, + timeout: float | None = None, + ): + """Join an existing conversation as a non-initiator. + + Reads the conversation id from ``meta`` (the same dict passed + to ``@on_message`` handlers) and registers a local handle so + subsequent messages tagged with that id route here too. Use + when a role wants to process a multi-message exchange from the + responder side — typical for gossip forwarders. + """ + from mango.agent.conversation import CONVERSATION_ID_KEY, Conversation + + conv_id = meta.get(CONVERSATION_ID_KEY) + if not conv_id: + raise ValueError( + "join_conversation requires a meta dict carrying " + f"{CONVERSATION_ID_KEY!r}" + ) + return _ConversationContext( + role_context=self, + conv=Conversation( + owner=self, + conversation_id=conv_id, + state=state, + timeout=timeout, + ), + ) + + async def gather( + self, + content: Any, + receivers, + *, + reply_type: type | tuple[type, ...] | None = None, + timeout: float = 5.0, + min_fraction: float = 1.0, + ) -> dict["AgentAddress", Any]: + """Send *content* to every address in *receivers* and collect replies. + + Returns a dict mapping the responding agent's + :class:`AgentAddress` to the reply content. Senders are + expected to use :meth:`AgentDelegates.reply_to` (or otherwise + echo ``tracking_id`` with ``reply=True``) — every legacy + request/response pair in mango already follows that convention, + so existing responder roles work unchanged. + + :param receivers: iterable of :class:`AgentAddress` targets. + :param reply_type: optional class/tuple — replies not matching + this type are silently dropped (filters out tracking-id + collisions with unrelated traffic). + :param timeout: hard wallclock cap; if no quorum is reached by + then, returns whatever replies have arrived so far. + :param min_fraction: between 0 and 1. When :math:`\\geq` this + fraction of receivers have replied, the call returns + without waiting further. Defaults to 1.0 — wait for all, + time out otherwise. + """ + import uuid as _uuid + + agent = self._role_handler._agent + if agent is None: + raise RuntimeError( + "RoleContext.gather requires a RoleAgent — the role's " + "context is not bound to an agent yet." + ) + receivers = list(receivers) + n = len(receivers) + if n == 0: + return {} + expected = max(1, int(round(min_fraction * n))) + + tracking_id = str(_uuid.uuid4()) + collector = agent.open_gather( + tracking_id, expected=expected, reply_type=reply_type + ) + # Timeout must follow mango's simulation clock — not wall time + # — so ``gather`` behaves identically under ``AsyncioClock`` + # (real-time) and ``ExternalClock`` (simulation). ``wait_for`` + # uses ``loop.time()`` and would block real seconds in sim mode. + try: + for addr in receivers: + await self.send_message( + content, + receiver_addr=addr, + tracking_id=tracking_id, + ) + clock = agent.scheduler.clock + done_fut = asyncio.ensure_future(collector.wait()) + timeout_fut = asyncio.ensure_future(clock.sleep(timeout)) + try: + await asyncio.wait( + {done_fut, timeout_fut}, + return_when=asyncio.FIRST_COMPLETED, + ) + finally: + for fut in (done_fut, timeout_fut): + if not fut.done(): + fut.cancel() + finally: + agent.close_gather(tracking_id) + return dict(collector.responses) + def on_start(self): self._role_handler.on_start() @@ -539,6 +755,7 @@ def __init__(self): """ super().__init__() self._role_handler = RoleHandler(None) + self._role_handler._agent = self self._role_context = RoleContext(self._role_handler, self.aid, self.inbox) def on_start(self): @@ -612,6 +829,12 @@ def _bind(self, context: RoleContext) -> None: :param context: the role context """ self._context = context + # Apply class-level @on_message / @on_event / @periodic + # decorators before user-defined ``setup`` runs, so the explicit + # setup body can override or extend the declarative wiring. + from mango.agent.decorators import apply_dispatch + + apply_dispatch(self) @property def context(self) -> RoleContext: @@ -666,3 +889,53 @@ def on_agent_event(self, event: Any) -> None: :param event: the event object """ + + +class _ConversationContext: + """Async context manager returned by ``RoleContext.open_conversation`` / + ``RoleContext.join_conversation``. + + Owns the lifecycle of one :class:`~mango.agent.conversation.Conversation`: + registers it with the agent on ``__aenter__``, schedules an + optional clock-aware timeout, and unregisters on ``__aexit__``. + Splitting the context manager out keeps ``Conversation`` itself + free of mango-agent references and easy to unit-test. + """ + + def __init__(self, *, role_context: "RoleContext", conv) -> None: + self._role_context = role_context + self._conv = conv + self._timeout_future = None + + async def __aenter__(self): + agent = self._role_context._role_handler._agent + if agent is None: + raise RuntimeError( + "Conversation requires a bound RoleAgent — the role's " + "context is not attached yet." + ) + agent.open_conversation(self._conv) + if self._conv._timeout is not None: + # Use the agent's scheduler clock so timeouts respect + # simulation time when running under an ExternalClock. + clock = agent.scheduler.clock if agent.scheduler else None + if clock is not None: + self._timeout_future = asyncio.ensure_future( + self._fire_after(clock, self._conv._timeout) + ) + return self._conv + + async def __aexit__(self, exc_type, exc, tb): + if self._timeout_future is not None and not self._timeout_future.done(): + self._timeout_future.cancel() + agent = self._role_context._role_handler._agent + if agent is not None: + agent.close_conversation(self._conv) + return False + + async def _fire_after(self, clock, delay: float) -> None: + try: + await clock.sleep(delay) + except asyncio.CancelledError: + return + self._conv._fire_timeout() diff --git a/mango/express/health.py b/mango/express/health.py new file mode 100644 index 0000000..3ac13ca --- /dev/null +++ b/mango/express/health.py @@ -0,0 +1,150 @@ +"""Per-topology link-health tracking. + +Multi-agent gossip protocols typically need a notion of "is this +neighbour still alive?" so a sender can route around silent peers. +Without framework support every role re-invents the same machinery — +a per-neighbour decay timer, a multiplicative recovery rule on every +received message, and a filter that drops neighbours whose score is +below some threshold. + +This module makes that machinery first-class on a :class:`Topology` +via the :class:`EdgeHealth` config and the :class:`TopologyHealth` +runtime. When a topology is built with ``edge_health=EdgeHealth(...)``, +every agent in the topology gets an auto-installed receive hook that +multiplicatively recovers the corresponding edge score on every +incoming message. Roles can then ask the topology for the +``live_neighbours()`` set — neighbours whose current score is at or +above the configured threshold — without touching the bookkeeping. + +The decay follows the clock attached to the agent's scheduler, so +behaviour is identical under real-time (:class:`AsyncioClock`) and +simulation (:class:`ExternalClock`). + +Example:: + + with create_topology(tid="groups", edge_health=EdgeHealth( + decay_per_s=0.125, recovery_rate=0.6, liveness_threshold=0.5, + )) as topo: + ... + + # In a role: + live = self.context.live_neighbours(tid="groups") # filtered list + weights = self.context.neighbour_scores(tid="groups") # K per neighbour +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from mango.agent.core import AgentAddress + + +@dataclass(frozen=True) +class EdgeHealth: + """Configuration for continuous link-health tracking on a topology. + + The score for a directed (owner → neighbour) edge starts at + ``initial`` and: + + * decays linearly with elapsed clock time at ``decay_per_s``; + * recovers multiplicatively toward 1.0 by ``score += (1-score) * + recovery_rate`` on every incoming message from the neighbour. + + Defaults match the scare baseline: a poll period of ~8 s on a + sensitive sector implies one missed beat per second of silence is + ``≈ 1/8`` of the full score, while every received message rebuilds + 60 % of the missing fraction. + + :param decay_per_s: score decay per second of silence. + :param recovery_rate: fraction of (1-score) recovered per received + message. Must be in ``(0, 1]``. + :param initial: starting score for a never-contacted neighbour + (default 1.0 — optimistic bootstrap). + :param liveness_threshold: default cutoff for ``live_neighbours``. + Callers can override per-call. + """ + + decay_per_s: float = 0.125 + recovery_rate: float = 0.6 + initial: float = 1.0 + liveness_threshold: float = 0.5 + + def __post_init__(self) -> None: + if not (0.0 < self.recovery_rate <= 1.0): + raise ValueError("recovery_rate must be in (0, 1]") + if self.decay_per_s < 0.0: + raise ValueError("decay_per_s must be non-negative") + + +class TopologyHealth: + """Per-topology runtime state — one instance per :class:`Topology`. + + Stores a score per directed (owner, neighbour) pair keyed by their + string addresses. The owner side is included because the same + agent can be a member of multiple topologies, each with its own + health view of the same neighbour. + + All clock reads go through the ``clock_fn`` callable supplied at + construction time — ``lambda: scheduler.clock.time``. This keeps + decay rate-aware of simulation time without coupling the runtime + to a specific clock class. + """ + + def __init__(self, params: EdgeHealth) -> None: + self.params = params + # (owner_addr_str, neighbour_addr_str) -> (score, last_t) + self._state: dict[tuple[str, str], tuple[float, float]] = {} + + def _key( + self, owner_addr: AgentAddress, neighbour_addr: AgentAddress + ) -> tuple[str, str]: + return (str(owner_addr), str(neighbour_addr)) + + def nudge( + self, + owner_addr: AgentAddress, + neighbour_addr: AgentAddress, + now: float, + ) -> None: + """Apply the multiplicative recovery on (owner → neighbour). + + Called by the message-receive hook on every inbound message + from *neighbour_addr*. The decay is first applied up to + *now*, then the recovery rule lifts the score. + """ + key = self._key(owner_addr, neighbour_addr) + score = self._decayed_score(key, now) + score = min(1.0, score + (1.0 - score) * self.params.recovery_rate) + self._state[key] = (score, now) + + def score( + self, + owner_addr: AgentAddress, + neighbour_addr: AgentAddress, + now: float, + ) -> float: + """Return the current decayed score for (owner → neighbour).""" + key = self._key(owner_addr, neighbour_addr) + decayed = self._decayed_score(key, now) + # Memoise the decay snapshot so a stale entry doesn't compound + # over many reads. + self._state[key] = (decayed, now) + return decayed + + def is_live( + self, + owner_addr: AgentAddress, + neighbour_addr: AgentAddress, + now: float, + *, + threshold: float | None = None, + ) -> bool: + cutoff = threshold if threshold is not None else self.params.liveness_threshold + return self.score(owner_addr, neighbour_addr, now) >= cutoff + + def _decayed_score(self, key: tuple[str, str], now: float) -> float: + score, last_t = self._state.get(key, (self.params.initial, now)) + elapsed = max(0.0, now - last_t) + return max(0.0, score - elapsed * self.params.decay_per_s) diff --git a/mango/express/topology.py b/mango/express/topology.py index a55ac11..7826cf5 100644 --- a/mango/express/topology.py +++ b/mango/express/topology.py @@ -9,7 +9,7 @@ from collections.abc import Callable, Iterator from contextlib import contextmanager -from typing import Any +from typing import TYPE_CHECKING, Any import networkx as nx @@ -22,6 +22,9 @@ TopologyService, ) +if TYPE_CHECKING: + from mango.express.health import EdgeHealth + AGENT_NODE_KEY = "node" STATE_EDGE_KEY = "state" @@ -59,11 +62,27 @@ class Topology: this to distinguish neighbor sets (default ``"default"``) """ - def __init__(self, graph: nx.Graph, *, tid: str = "default") -> None: + def __init__( + self, + graph: nx.Graph, + *, + tid: str = "default", + edge_health: EdgeHealth | None = None, + ) -> None: self._graph: nx.Graph = graph self._tid: str = tid self._connectors: list[tuple[str, TopologyNeighbor]] = [] self._connections: list[tuple[str, Topology]] = [] + # Optional per-topology link-health tracking. When set, a + # :class:`TopologyHealth` runtime is shared across every agent + # in this topology and an auto-nudge subscription is installed + # at inject time. See :mod:`mango.express.health`. + from mango.express.health import TopologyHealth + + self._edge_health: EdgeHealth | None = edge_health + self._health: TopologyHealth | None = ( + TopologyHealth(edge_health) if edge_health is not None else None + ) for node in self._graph.nodes: self._graph.nodes[node][AGENT_NODE_KEY] = AgentNode() @@ -226,6 +245,13 @@ def _build_and_inject(self, *, update_connected: bool = True) -> None: svc._tid_to_characteristic[self._tid] = agent_node._characteristic_for( agent ) + # Bind the shared :class:`TopologyHealth` runtime so the + # agent's receive hook can nudge per-edge scores. None + # when the topology was built without ``edge_health``. + if self._health is not None: + svc._tid_to_health[self._tid] = self._health + else: + svc._tid_to_health.pop(self._tid, None) # Transfer any marks from mark_as_connector for conn_type in svc._marked_connector_for: @@ -319,12 +345,20 @@ def custom_topology(graph: nx.Graph) -> Topology: @contextmanager def create_topology( - *, directed: bool = False, tid: str = "default" + *, + directed: bool = False, + tid: str = "default", + edge_health: EdgeHealth | None = None, ) -> Iterator[Topology]: """Context manager that builds a topology and injects neighborhoods on exit. :param directed: use a directed graph (default ``False``) :param tid: topology identifier (default ``"default"``) + :param edge_health: opt-in continuous link-health tracking. When + set, every agent in the topology gets an auto-nudge hook that + recovers per-neighbour scores on incoming messages, and the + :meth:`mango.RoleContext.live_neighbours` query becomes + available for this ``tid``. See :mod:`mango.express.health`. :yield: an empty :class:`Topology` to populate Example:: @@ -338,7 +372,7 @@ def create_topology( topology.add_edge(n1, n3) """ graph = nx.DiGraph() if directed else nx.Graph() - topology = Topology(graph, tid=tid) + topology = Topology(graph, tid=tid, edge_health=edge_health) yield topology topology.inject() diff --git a/tests/unit_tests/role/conversation_test.py b/tests/unit_tests/role/conversation_test.py new file mode 100644 index 0000000..8093ce9 --- /dev/null +++ b/tests/unit_tests/role/conversation_test.py @@ -0,0 +1,248 @@ +"""Tests for :class:`mango.agent.conversation.Conversation`. + +Three flavours of coverage: + +* Real-time (TCP) — initiator opens a conversation, joiner receives + the message and joins it, both sides exchange multiple messages + under one id. +* Convergence / cancellation — the iterator terminates correctly when + the caller signals end. +* Simulation-time timeout — exercises the conversation's clock-aware + timeout under an :class:`ExternalClock` so we know the timeout + advances with simulation time, not wall time. +""" + +from __future__ import annotations + +import asyncio +from dataclasses import dataclass + +import pytest + +from mango import ( + Agent, + Role, + RoleAgent, + activate, + create_tcp_container, + on_message, + sender_addr, +) +from mango.express.api import run_with_simulation +from mango.simulation.world import step_simulation + + +@dataclass +class _Step: + """Multi-hop payload — carries a counter the participants increment.""" + + counter: int + + +# --------------------------------------------------------------------------- +# End-to-end multi-hop conversation on TCP. +# --------------------------------------------------------------------------- + + +class _Joiner(Role): + """Joiner-side: reads incoming step, replies with the next counter. + Stays in the conversation so the initiator can keep volleying.""" + + def __init__(self) -> None: + super().__init__() + self.processed: list[int] = [] + + @on_message(_Step) + async def on_step(self, content: _Step, meta: dict) -> None: + self.processed.append(content.counter) + async with self.context.join_conversation(meta) as conv: + await conv.send(sender_addr(meta), _Step(counter=content.counter + 1)) + + +class _Initiator(Role): + """Opens a conversation, sends N messages and reads N replies on + the same conversation id.""" + + def __init__(self, peer_addr, *, rounds: int, timeout: float = 5.0) -> None: + super().__init__() + self.peer_addr = peer_addr + self.rounds = rounds + self.timeout = timeout + self.seen: list[int] = [] + + async def run(self) -> None: + async with self.context.open_conversation(timeout=self.timeout) as conv: + await conv.send(self.peer_addr, _Step(counter=0)) + async for content, meta in conv: + self.seen.append(content.counter) + if len(self.seen) >= self.rounds: + conv.converge() + continue + await conv.send(self.peer_addr, _Step(counter=content.counter + 1)) + + +@pytest.mark.asyncio +async def test_conversation_round_trip(): + """Three back-and-forth volleys carry the same conversation id — + every reply routes to the initiator's async iterator, the joiner's + fresh ``join_conversation`` succeeds every time.""" + container = create_tcp_container(addr=("127.0.0.1", 5580)) + joiner_agent = container.register(RoleAgent()) + joiner = _Joiner() + joiner_agent.add_role(joiner) + + initiator_agent = container.register(RoleAgent()) + initiator = _Initiator(joiner_agent.addr, rounds=3) + initiator_agent.add_role(initiator) + + async with activate([container]): + await initiator.run() + + # Initiator received exactly ``rounds`` replies with strictly + # increasing counters. + assert initiator.seen == [1, 3, 5] + # Joiner processed every initiator-sent step. + assert joiner.processed == [0, 2, 4] + + +# --------------------------------------------------------------------------- +# Convergence and cancellation +# --------------------------------------------------------------------------- + + +class _BroadcastRecv(Role): + @on_message(_Step) + async def on_step(self, content: _Step, meta: dict) -> None: + async with self.context.join_conversation(meta) as conv: + # Echo once, then converge. + await conv.send(sender_addr(meta), _Step(counter=content.counter * 10)) + + +@pytest.mark.asyncio +async def test_conversation_converge_delivers_final_message(): + """``converge()`` lets the iterator deliver the message in flight + before terminating. The next pull then exits cleanly.""" + + class CollectingInitiator(Role): + def __init__(self, peer): + super().__init__() + self.peer = peer + self.seen = [] + self.iterations = 0 + + async def run(self): + async with self.context.open_conversation(timeout=2.0) as conv: + await conv.send(self.peer, _Step(counter=5)) + async for content, _meta in conv: + self.iterations += 1 + self.seen.append(content.counter) + conv.converge() # end after this one + # On exit `closed` should be set. + assert conv.closed + + container = create_tcp_container(addr=("127.0.0.1", 5581)) + recv = container.register(RoleAgent()) + recv.add_role(_BroadcastRecv()) + init_agent = container.register(RoleAgent()) + init = CollectingInitiator(recv.addr) + init_agent.add_role(init) + + async with activate([container]): + await init.run() + + assert init.seen == [50] + assert init.iterations == 1 + + +@pytest.mark.asyncio +async def test_conversation_cancel_drops_remaining_messages(): + """``cancel()`` ends iteration on the next pull without delivering + whatever was already in flight.""" + + container = create_tcp_container(addr=("127.0.0.1", 5582)) + init_agent = container.register(RoleAgent()) + + class Solo(Role): + async def run(self): + async with self.context.open_conversation(timeout=2.0) as conv: + conv.cancel() + pulls = 0 + async for _, _meta in conv: # pragma: no cover - shouldn't loop + pulls += 1 + self.pulls = pulls + + solo = Solo() + init_agent.add_role(solo) + + async with activate([container]): + await solo.run() + + assert solo.pulls == 0 + + +# --------------------------------------------------------------------------- +# Simulation-clock timeout — the load-bearing requirement. +# --------------------------------------------------------------------------- + + +class SimAgent(Agent): + """A bare :class:`Agent` for testing — we only need the scheduler + clock; the message inbox is unused.""" + + def handle_message(self, content, meta): + pass + + +@pytest.mark.asyncio +async def test_conversation_timeout_follows_simulation_clock(): + """An open conversation under ``run_with_simulation`` must time out + when *simulation* time crosses the deadline — not wall time. + + Strategy: open a conversation with ``timeout=5.0`` simulated + seconds and advance the clock by one 5-second step. The + timeout future on the conversation's context manager fires when + ``clock.sleep(5.0)`` resolves, which under ``ExternalClock`` only + happens after the world advances past that instant. + """ + from mango import RoleAgent + from mango.agent.role import Role + + class TimedRole(Role): + def __init__(self): + super().__init__() + self.iteration_done_at = None + self.before_advance = None + + async def watch(self, world): + async with self.context.open_conversation(timeout=5.0) as conv: + self.before_advance = world.clock.time + + # Schedule the world advance as a background task so + # the async iterator can proceed. ``step_simulation`` + # advances the clock by 5 s, which makes the + # conversation's ``clock.sleep(5.0)`` future fire. + async def _advance(): + # Give the conversation context manager a moment + # to register its sleep future before we tick. + await asyncio.sleep(0) + await step_simulation(world, step_size_s=5.0) + + advancer = asyncio.create_task(_advance()) + # Drain the conversation — when the timeout fires the + # iterator exits. + async for _, _meta in conv: # pragma: no cover — empty stream + pass + await advancer + self.iteration_done_at = world.clock.time + + agent = RoleAgent() + role = TimedRole() + agent.add_role(role) + + async with run_with_simulation(agent) as world: + await role.watch(world) + + # The conversation exited after simulation time advanced 5 s, not + # 5 wall seconds — confirms the timeout used clock.sleep. + assert role.before_advance == pytest.approx(0.0) + assert role.iteration_done_at == pytest.approx(5.0) diff --git a/tests/unit_tests/role/decorators_test.py b/tests/unit_tests/role/decorators_test.py new file mode 100644 index 0000000..ce6ae07 --- /dev/null +++ b/tests/unit_tests/role/decorators_test.py @@ -0,0 +1,286 @@ +"""Unit tests for :mod:`mango.agent.decorators`. + +These tests do not spin up a container — they exercise +``apply_dispatch`` directly against a stub context that records every +subscribe / schedule call. That keeps the tests focused on the +decoration semantics (collection order, async-wrap, predicate +threading, ``only_if`` gating) without depending on the rest of the +agent runtime. +""" + +from __future__ import annotations + +import asyncio +from dataclasses import dataclass, field +from typing import Any + +import pytest + +from mango import Role, on_event, on_message, periodic +from mango.agent.decorators import apply_dispatch, collect_dispatch + + +@dataclass +class _StubContext: + """Minimal recording context. Mirrors only the surface + ``apply_dispatch`` touches.""" + + message_subs: list[dict] = field(default_factory=list) + event_subs: list[dict] = field(default_factory=list) + periodic_subs: list[dict] = field(default_factory=list) + + def subscribe_message(self, role, method, condition, priority=0): + self.message_subs.append( + { + "role": role, + "method": method, + "condition": condition, + "priority": priority, + } + ) + + def subscribe_event(self, role, event_type, method): + self.event_subs.append( + {"role": role, "event_type": event_type, "method": method} + ) + + def schedule_periodic_task(self, coro_func, delay): + self.periodic_subs.append({"coro": coro_func, "delay": delay}) + + def schedule_instant_task(self, coro): + # The async @on_message wrap calls this when a matching message + # arrives. Capture and immediately run synchronously for the + # tests; the coroutine itself shouldn't actually do anything + # except mutate role state. + asyncio.get_event_loop().run_until_complete(coro) + + +class _StubRole: + """A non-:class:`Role` carrier so we can probe the dispatch + machinery without inheriting the full Role lifecycle.""" + + def __init__(self): + self.context = _StubContext() + self.calls: list[tuple[str, Any]] = [] + + +class TestOnMessageDecorator: + def test_collects_single_subscription(self): + class R(_StubRole): + @on_message(str) + def handle(self, content, meta): + self.calls.append(("handle", content)) + + role = R() + apply_dispatch(role) + assert len(role.context.message_subs) == 1 + entry = role.context.message_subs[0] + assert entry["priority"] == 0 + assert entry["condition"]("hello", {}) is True + assert entry["condition"](42, {}) is False + + def test_where_predicate_receives_self(self): + """The ``where`` predicate is bound to the role instance, so it + can inspect role state (e.g. a sector tag) at filter time.""" + + class R(_StubRole): + sector = "electricity" + + @on_message(dict, where=lambda self, c, m: c.get("sector") == self.sector) + def handle(self, content, meta): + pass + + role = R() + apply_dispatch(role) + cond = role.context.message_subs[0]["condition"] + assert cond({"sector": "electricity"}, {}) is True + assert cond({"sector": "gas"}, {}) is False + assert cond("not a dict", {}) is False # type guard still holds + + def test_priority_is_forwarded(self): + class R(_StubRole): + @on_message(int, priority=-5) + def first(self, content, meta): + pass + + @on_message(int, priority=5) + def last(self, content, meta): + pass + + role = R() + apply_dispatch(role) + priorities = sorted(s["priority"] for s in role.context.message_subs) + assert priorities == [-5, 5] + + def test_async_handler_is_scheduled_as_task(self): + class R(_StubRole): + @on_message(str) + async def handle(self, content, meta): + self.calls.append(("async", content)) + + role = R() + + # Replace schedule_instant_task with an awaitable trap so we + # can verify the wrap behaviour without spinning up a real + # event loop. + scheduled = [] + + def fake_instant(coro): + scheduled.append(coro) + coro.close() # don't actually run; just confirm it's a coro + + role.context.schedule_instant_task = fake_instant + apply_dispatch(role) + + # Synthesize a matching delivery. + method = role.context.message_subs[0]["method"] + method("hello", {}) + assert len(scheduled) == 1 + + +class TestOnEventDecorator: + def test_collects_event_subscription(self): + @dataclass + class MyEvent: + value: int + + class R(_StubRole): + @on_event(MyEvent) + def on_event(self, event, src): + self.calls.append(("event", event)) + + role = R() + apply_dispatch(role) + assert len(role.context.event_subs) == 1 + assert role.context.event_subs[0]["event_type"] is MyEvent + + def test_multiple_event_types(self): + @dataclass + class A: ... + + @dataclass + class B: ... + + class R(_StubRole): + @on_event(A) + @on_event(B) + def on_either(self, event, src): + pass + + role = R() + apply_dispatch(role) + types = {s["event_type"] for s in role.context.event_subs} + assert types == {A, B} + + +class TestPeriodicDecorator: + def test_collects_periodic_with_literal_delay(self): + class R(_StubRole): + @periodic(every=2.5) + async def tick(self): + self.calls.append(("tick",)) + + role = R() + apply_dispatch(role) + assert len(role.context.periodic_subs) == 1 + assert role.context.periodic_subs[0]["delay"] == pytest.approx(2.5) + + def test_periodic_reads_role_attribute_when_string(self): + class R(_StubRole): + poll_period_s = 0.75 + + @periodic(every="poll_period_s") + async def tick(self): + pass + + role = R() + apply_dispatch(role) + assert role.context.periodic_subs[0]["delay"] == pytest.approx(0.75) + + def test_periodic_rejects_sync_function(self): + with pytest.raises(TypeError, match="async coroutine"): + + class R(_StubRole): + @periodic(every=1.0) + def tick(self): # noqa: B903 — sync intentionally for the test + pass + + def test_only_if_gates_the_body(self): + runs: list[bool] = [] + + class R(_StubRole): + allowed = False + + @periodic(every=0.1, only_if=lambda self: self.allowed) + async def tick(self): + runs.append(True) + + role = R() + apply_dispatch(role) + wrapped = role.context.periodic_subs[0]["coro"] + + loop = asyncio.new_event_loop() + try: + # Gate closed → no run. + loop.run_until_complete(wrapped()) + assert runs == [] + # Gate open → runs. + role.allowed = True + loop.run_until_complete(wrapped()) + assert runs == [True] + finally: + loop.close() + + +class TestCollectDispatch: + def test_collects_from_base_class(self): + class Base(_StubRole): + @on_message(str) + def base_handler(self, content, meta): + pass + + class Sub(Base): + @on_message(int) + def sub_handler(self, content, meta): + pass + + meta = collect_dispatch(Sub) + assert set(meta.keys()) == {"base_handler", "sub_handler"} + + def test_subclass_overrides_keep_metadata(self): + class Base(_StubRole): + @on_message(str) + def handler(self, content, meta): + pass + + class Sub(Base): + @on_message(int) + def handler(self, content, meta): # noqa: D401 — override + pass + + meta = collect_dispatch(Sub) + # Only the override is kept; the base entry is masked. + assert "handler" in meta + assert len(meta["handler"].message_subs) == 1 + assert meta["handler"].message_subs[0]["message_type"] is int + + +class TestRoleIntegration: + """End-to-end via a real :class:`Role` — verifies that ``_bind`` + triggers ``apply_dispatch`` so a role using only decorators (no + ``setup`` body) gets fully wired.""" + + def test_role_without_setup_still_subscribes(self): + captured = [] + + class MyRole(Role): + @on_message(str) + def handler(self, content, meta): + captured.append(content) + + # Stub context that records subscribe_message calls. + ctx = _StubContext() + role = MyRole() + role._context = ctx + apply_dispatch(role) + assert len(ctx.message_subs) == 1 diff --git a/tests/unit_tests/role/edge_health_test.py b/tests/unit_tests/role/edge_health_test.py new file mode 100644 index 0000000..6a7cd4c --- /dev/null +++ b/tests/unit_tests/role/edge_health_test.py @@ -0,0 +1,228 @@ +"""Tests for :class:`EdgeHealth` and the auto-nudge integration with topology. + +Three layers of coverage: + +1. :class:`TopologyHealth` unit tests — pure-state machine, no agents. +2. End-to-end test on TCP containers — confirms auto-nudge fires on + every received message and ``live_neighbours`` filters correctly. +3. Simulation-clock test — confirms decay uses :class:`ExternalClock` + so behaviour is identical under simulation time. +""" + +from __future__ import annotations + +import asyncio +from dataclasses import dataclass + +import pytest + +from mango import ( + EdgeHealth, + Role, + RoleAgent, + activate, + create_tcp_container, + create_topology, + on_message, + sender_addr, +) +from mango.express.health import TopologyHealth + +# --------------------------------------------------------------------------- +# Pure unit tests on the runtime — no agents, no clock. +# --------------------------------------------------------------------------- + + +class _A: + """Tiny stand-in for an :class:`AgentAddress`; only its ``str`` + representation is used by :class:`TopologyHealth`.""" + + def __init__(self, key: str) -> None: + self.key = key + + def __str__(self) -> str: # noqa: D401 + return self.key + + +class TestTopologyHealthRuntime: + def test_initial_score_at_bootstrap(self): + h = TopologyHealth(EdgeHealth(initial=0.7)) + # First read uses the configured initial value. + score = h.score(_A("me"), _A("peer"), now=0.0) + assert score == pytest.approx(0.7) + + def test_score_decays_linearly_with_silence(self): + h = TopologyHealth(EdgeHealth(decay_per_s=0.5, initial=1.0)) + # Establish the entry. + _ = h.score(_A("me"), _A("peer"), now=0.0) + # Two seconds of silence → score = 1.0 - 2 * 0.5 = 0.0 + assert h.score(_A("me"), _A("peer"), now=2.0) == pytest.approx(0.0) + + def test_nudge_recovers_multiplicatively(self): + h = TopologyHealth(EdgeHealth(decay_per_s=0.5, recovery_rate=0.5, initial=1.0)) + # Start at 1, decay to 0 by t=2. + h.score(_A("me"), _A("peer"), now=0.0) + assert h.score(_A("me"), _A("peer"), now=2.0) == pytest.approx(0.0) + # Nudge at t=2 → score = 0 + (1-0) * 0.5 = 0.5 + h.nudge(_A("me"), _A("peer"), now=2.0) + assert h.score(_A("me"), _A("peer"), now=2.0) == pytest.approx(0.5) + # Another nudge → score = 0.5 + (1-0.5) * 0.5 = 0.75 + h.nudge(_A("me"), _A("peer"), now=2.0) + assert h.score(_A("me"), _A("peer"), now=2.0) == pytest.approx(0.75) + + def test_is_live_uses_threshold(self): + h = TopologyHealth( + EdgeHealth(decay_per_s=0.5, initial=1.0, liveness_threshold=0.6) + ) + # At t=0 score is 1.0 → live. + assert h.is_live(_A("me"), _A("peer"), now=0.0) is True + # At t=1 score is 0.5 → not live by configured threshold (0.6), + # but live by custom threshold (0.4). + assert h.is_live(_A("me"), _A("peer"), now=1.0) is False + assert h.is_live(_A("me"), _A("peer"), now=1.0, threshold=0.4) is True + + def test_per_owner_scores_are_independent(self): + """The same neighbour observed by two different owners has two + independent scores — the runtime is a single shared instance + across all agents in the topology, but the (owner, neighbour) + key isolates each agent's view.""" + h = TopologyHealth(EdgeHealth(decay_per_s=0.0, initial=1.0)) + h.nudge(_A("alice"), _A("bob"), now=0.0) + # Alice's view of Bob recovered, Carol's didn't. + assert h.score(_A("alice"), _A("bob"), now=0.0) > 0.0 + # Carol hasn't been pinged — still at the initial. + assert h.score(_A("carol"), _A("bob"), now=0.0) == pytest.approx(1.0) + + def test_recovery_rate_validation(self): + with pytest.raises(ValueError): + EdgeHealth(recovery_rate=0.0) + with pytest.raises(ValueError): + EdgeHealth(recovery_rate=-0.1) + with pytest.raises(ValueError): + EdgeHealth(recovery_rate=1.1) + with pytest.raises(ValueError): + EdgeHealth(decay_per_s=-1.0) + + +# --------------------------------------------------------------------------- +# End-to-end on TCP containers — auto-nudge + live_neighbours. +# --------------------------------------------------------------------------- + + +@dataclass +class _Ping: + counter: int + + +class _PingHandler(Role): + """Receives pings and replies — keeps the test wiring symmetric.""" + + def __init__(self) -> None: + super().__init__() + self.received: list[int] = [] + + @on_message(_Ping) + async def on_ping(self, content: _Ping, meta: dict) -> None: + self.received.append(content.counter) + # Reply so the sender's edge to us also gets nudged. + await self.context.send_message("ack", receiver_addr=sender_addr(meta)) + + +@pytest.mark.asyncio +async def test_auto_nudge_on_received_message(): + """A message received from a topology neighbour must recover that + neighbour's edge score on the receiving agent's TopologyService.""" + container = create_tcp_container(addr=("127.0.0.1", 5571)) + alice = container.register(RoleAgent()) + bob = container.register(RoleAgent()) + alice.add_role(_PingHandler()) + bob.add_role(_PingHandler()) + + with create_topology( + tid="pair", + edge_health=EdgeHealth( + decay_per_s=0.0, # disabled — isolate the nudge effect + recovery_rate=0.5, + initial=0.0, # start cold so nudge is observable + liveness_threshold=0.4, + ), + ) as t: + n0 = t.add_node(alice) + n1 = t.add_node(bob) + t.add_edge(n0, n1) + + async with activate([container]): + # Before any message: Alice's view of Bob is at the initial 0. + alice_ctx = alice._role_context + assert alice_ctx.neighbour_score(bob.addr, tid="pair") == pytest.approx(0.0) + # Send a ping; Alice should observe Bob's ack arrive and nudge. + await alice.send_message(_Ping(counter=1), receiver_addr=bob.addr) + # Let the ping → ack roundtrip complete. + await asyncio.sleep(0.2) + score = alice_ctx.neighbour_score(bob.addr, tid="pair") + # One ack from Bob → recovery_rate=0.5 of the missing fraction + # against a starting score of 0 → exactly 0.5. + assert score == pytest.approx(0.5) + + +@pytest.mark.asyncio +async def test_live_neighbours_filters_silent_peers(): + """A neighbour whose score has decayed below the threshold is + excluded from ``live_neighbours``.""" + container = create_tcp_container(addr=("127.0.0.1", 5572)) + me = container.register(RoleAgent()) + chatty = container.register(RoleAgent()) + silent = container.register(RoleAgent()) + me.add_role(_PingHandler()) + chatty.add_role(_PingHandler()) + silent.add_role(_PingHandler()) + + with create_topology( + tid="trio", + edge_health=EdgeHealth( + decay_per_s=0.0, + recovery_rate=0.9, + initial=0.0, + liveness_threshold=0.5, + ), + ) as t: + n_me = t.add_node(me) + n_chatty = t.add_node(chatty) + n_silent = t.add_node(silent) + t.add_edge(n_me, n_chatty) + t.add_edge(n_me, n_silent) + + async with activate([container]): + # Ping only chatty. + await me.send_message(_Ping(counter=1), receiver_addr=chatty.addr) + await asyncio.sleep(0.2) + live = me._role_context.live_neighbours(tid="trio") + # Chatty replied, silent didn't. Only chatty's score crossed + # the threshold. + assert chatty.addr in live + assert silent.addr not in live + + +@pytest.mark.asyncio +async def test_live_neighbours_falls_back_when_no_health(): + """Without ``edge_health`` configured, ``live_neighbours`` returns + the full unfiltered neighbour list — callers can use the method + unconditionally.""" + container = create_tcp_container(addr=("127.0.0.1", 5573)) + a = container.register(RoleAgent()) + b = container.register(RoleAgent()) + a.add_role(_PingHandler()) + b.add_role(_PingHandler()) + + with create_topology(tid="bare") as t: + n0 = t.add_node(a) + n1 = t.add_node(b) + t.add_edge(n0, n1) + + async with activate([container]): + live = a._role_context.live_neighbours(tid="bare") + # No nudges have happened — b is still in the list because + # health tracking was never enabled. + assert b.addr in live + # neighbour_score returns None for a no-health topology. + assert a._role_context.neighbour_score(b.addr, tid="bare") is None diff --git a/tests/unit_tests/role/emit_event_test.py b/tests/unit_tests/role/emit_event_test.py new file mode 100644 index 0000000..541cd03 --- /dev/null +++ b/tests/unit_tests/role/emit_event_test.py @@ -0,0 +1,88 @@ +"""Tests for the relaxed :meth:`RoleContext.emit_event` semantics. + +Before: emitting an event with no subscribed listener raised +``KeyError``, forcing every call site to wrap the emit in +``try/except KeyError``. Now: missing listeners are silently dropped +by default; the legacy strict behaviour is opt-in via ``strict=True``. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import pytest + +from mango import RoleAgent +from mango.agent.role import RoleContext, RoleHandler + + +@dataclass +class _SomeEvent: + value: int + + +def _make_handler() -> RoleHandler: + """Build a bare RoleHandler suitable for emit_event tests. No + real container is needed because emit_event is purely role-local.""" + h = RoleHandler(scheduler=None) + return h + + +class TestEmitEventNoListener: + def test_default_is_silent(self): + """Emitting an event nobody subscribed to is a no-op.""" + handler = _make_handler() + # Must not raise. + handler.emit_event(_SomeEvent(value=1)) + + def test_strict_raises_keyerror(self): + """Opting back into the legacy strict behaviour.""" + handler = _make_handler() + with pytest.raises(KeyError): + handler.emit_event(_SomeEvent(value=1), strict=True) + + def test_listener_receives_event_silent_mode(self): + """The relaxed default still delivers events to actual listeners.""" + handler = _make_handler() + received: list[_SomeEvent] = [] + # Subscribe with the legacy API (role is irrelevant for this test). + handler.subscribe_event( + role=object(), + event_type=_SomeEvent, + method=lambda ev, src: received.append(ev), + ) + handler.emit_event(_SomeEvent(value=42)) + assert received and received[0].value == 42 + + def test_listener_receives_event_strict_mode(self): + """Strict mode is purely about the no-listener case — when a + listener exists, the event is delivered normally.""" + handler = _make_handler() + received: list[_SomeEvent] = [] + handler.subscribe_event( + role=object(), + event_type=_SomeEvent, + method=lambda ev, src: received.append(ev), + ) + handler.emit_event(_SomeEvent(value=7), strict=True) + assert received and received[0].value == 7 + + +class TestRoleContextEmitEvent: + """The same semantics must reach callers via :class:`RoleContext`, + which is the public surface most roles use.""" + + def test_role_context_silent_no_listener(self): + # Use a RoleAgent so the role-context wiring is realistic; we + # don't need a container because emit_event is role-local. + agent = RoleAgent() + # The context is reachable via the private handle on the agent. + ctx: RoleContext = agent._role_context + # No listener subscribed — must not raise. + ctx.emit_event(_SomeEvent(value=1)) + + def test_role_context_strict_raises(self): + agent = RoleAgent() + ctx: RoleContext = agent._role_context + with pytest.raises(KeyError): + ctx.emit_event(_SomeEvent(value=1), strict=True) diff --git a/tests/unit_tests/role/gather_test.py b/tests/unit_tests/role/gather_test.py new file mode 100644 index 0000000..9334d37 --- /dev/null +++ b/tests/unit_tests/role/gather_test.py @@ -0,0 +1,281 @@ +"""End-to-end tests for :meth:`RoleContext.gather`. + +``gather`` is the simplification target for the most common scare / +mes pattern: send a request to N agents, collect their replies under +one id, return either when everyone answered or after a timeout. +These tests run against real TCP containers so the message-passing +infrastructure (tracking_id threading, reply matching) is exercised +end-to-end, not stubbed. +""" + +from __future__ import annotations + +import asyncio +from dataclasses import dataclass +from typing import Any + +import pytest + +from mango import ( + Role, + RoleAgent, + activate, + create_tcp_container, + on_message, +) + + +@dataclass +class _Ask: + """Request payload — carried by the gather caller.""" + + topic: str + + +@dataclass +class _Reply: + """Per-responder reply payload.""" + + value: float + + +class _Responder(Role): + """Replies to every :class:`_Ask` it receives with its configured value. + + Uses :meth:`AgentDelegates.reply_to` so the ``tracking_id`` threads + back to the caller automatically — that is the contract the gather + machinery relies on. + """ + + def __init__(self, value: float): + super().__init__() + self.value = value + self.received: int = 0 + + @on_message(_Ask) + async def on_ask(self, content: _Ask, meta: dict) -> None: + self.received += 1 + await self.context.reply_to(_Reply(value=self.value), received_meta=meta) + + +class _SilentResponder(Role): + """Never replies — used to verify the timeout branch.""" + + @on_message(_Ask) + async def on_ask(self, content: _Ask, meta: dict) -> None: # noqa: ARG002 + return None + + +class _Caller(Role): + """Owns the ``gather`` call. Stored result is asserted by the test.""" + + def __init__(self, receivers, *, timeout: float, min_fraction: float = 1.0): + super().__init__() + self.receivers = receivers + self.timeout = timeout + self.min_fraction = min_fraction + self.responses: dict[Any, _Reply] | None = None + self.elapsed: float = 0.0 + + async def run(self) -> None: + start = asyncio.get_event_loop().time() + self.responses = await self.context.gather( + _Ask(topic="x"), + receivers=self.receivers, + reply_type=_Reply, + timeout=self.timeout, + min_fraction=self.min_fraction, + ) + self.elapsed = asyncio.get_event_loop().time() - start + + +@pytest.mark.asyncio +async def test_gather_collects_all_replies(): + container = create_tcp_container(addr=("127.0.0.1", 5556)) + responder_agents = [] + responder_addrs = [] + for v in (1.0, 2.0, 3.0): + a = container.register(RoleAgent()) + a.add_role(_Responder(value=v)) + responder_agents.append(a) + responder_addrs.append(a.addr) + + caller = container.register(RoleAgent()) + caller_role = _Caller(responder_addrs, timeout=2.0) + caller.add_role(caller_role) + + async with activate([container]): + await caller_role.run() + + assert caller_role.responses is not None + assert len(caller_role.responses) == 3 + values = sorted(r.value for r in caller_role.responses.values()) + assert values == [1.0, 2.0, 3.0] + # All three responded, so we never hit the timeout — gather should + # have returned essentially immediately. + assert caller_role.elapsed < 2.0 + + +@pytest.mark.asyncio +async def test_gather_returns_on_timeout_with_partial_results(): + """With one silent responder and ``min_fraction=1.0``, gather waits + until the timeout and returns whatever did arrive.""" + container = create_tcp_container(addr=("127.0.0.1", 5557)) + addrs = [] + for v in (1.0, 2.0): + a = container.register(RoleAgent()) + a.add_role(_Responder(value=v)) + addrs.append(a.addr) + silent = container.register(RoleAgent()) + silent.add_role(_SilentResponder()) + addrs.append(silent.addr) + + caller = container.register(RoleAgent()) + caller_role = _Caller(addrs, timeout=0.5) + caller.add_role(caller_role) + + async with activate([container]): + await caller_role.run() + + assert caller_role.responses is not None + # Two responders answered; the silent one did not. + assert len(caller_role.responses) == 2 + # The timeout should have been hit (we waited ≥ timeout-but-less than 2x). + assert 0.4 <= caller_role.elapsed < 1.5 + + +@pytest.mark.asyncio +async def test_gather_returns_early_on_quorum(): + """``min_fraction=0.5`` lets gather return as soon as half have replied — + no need to wait for the silent member.""" + container = create_tcp_container(addr=("127.0.0.1", 5558)) + addrs = [] + for v in (1.0, 2.0): + a = container.register(RoleAgent()) + a.add_role(_Responder(value=v)) + addrs.append(a.addr) + silent = container.register(RoleAgent()) + silent.add_role(_SilentResponder()) + addrs.append(silent.addr) + + caller = container.register(RoleAgent()) + caller_role = _Caller(addrs, timeout=5.0, min_fraction=0.5) + caller.add_role(caller_role) + + async with activate([container]): + await caller_role.run() + + # Quorum is 2 (round(0.5 * 3) = 2). Should return as soon as the + # two real responders answered — well under the 5 s timeout. + assert caller_role.responses is not None + assert len(caller_role.responses) >= 2 + assert caller_role.elapsed < 1.0 + + +@pytest.mark.asyncio +async def test_gather_empty_receivers_returns_empty_dict(): + """A degenerate ``gather`` with no receivers must not block.""" + container = create_tcp_container(addr=("127.0.0.1", 5559)) + caller = container.register(RoleAgent()) + caller_role = _Caller([], timeout=2.0) + caller.add_role(caller_role) + + async with activate([container]): + await caller_role.run() + + assert caller_role.responses == {} + assert caller_role.elapsed < 0.5 + + +@pytest.mark.asyncio +async def test_gather_timeout_uses_simulation_clock(): + """Under :class:`ExternalClock` the timeout must advance with the + simulation, not with wall time. Strategy: open a gather to a + silent receiver with ``timeout=10`` simulated seconds, advance the + world clock by one 10-second step in a background task, and + confirm the call returns when the *clock* crosses the deadline.""" + from mango import RoleAgent + from mango.express.api import run_with_simulation + from mango.simulation.world import step_simulation + + class SilentInSim(Role): + @on_message(_Ask) + async def on_ask(self, content: _Ask, meta: dict) -> None: + return None + + class SimCaller(Role): + def __init__(self, peer): + super().__init__() + self.peer = peer + self.elapsed_sim_time = None + + async def run(self, world): + start = world.clock.time + advancer = asyncio.create_task( + _advance_after_sleep(world, step_size_s=10.0) + ) + responses = await self.context.gather( + _Ask(topic="x"), + receivers=[self.peer], + reply_type=_Reply, + timeout=10.0, + min_fraction=1.0, + ) + await advancer + self.elapsed_sim_time = world.clock.time - start + assert responses == {} + + async def _advance_after_sleep(world, *, step_size_s): + # Yield once so the gather() call has time to register its + # ``clock.sleep`` future before we tick. + await asyncio.sleep(0) + await step_simulation(world, step_size_s=step_size_s) + + silent_agent = RoleAgent() + silent_agent.add_role(SilentInSim()) + caller_agent = RoleAgent() + + async with run_with_simulation(silent_agent, caller_agent) as world: + caller = SimCaller(silent_agent.addr) + caller_agent.add_role(caller) + await caller.run(world) + + # 10 simulated seconds elapsed — confirms gather used clock.sleep, + # not wall-clock asyncio.wait_for. + assert caller.elapsed_sim_time == pytest.approx(10.0) + + +@pytest.mark.asyncio +async def test_gather_filters_by_reply_type(): + """A reply of an unrelated type with the same tracking_id is dropped.""" + + class _NoiseResponder(Role): + """Sends back a string instead of a ``_Reply`` — ``reply_type`` + filter must reject it without raising.""" + + @on_message(_Ask) + async def on_ask(self, content: _Ask, meta: dict) -> None: # noqa: ARG002 + await self.context.reply_to("not-a-Reply", received_meta=meta) + + container = create_tcp_container(addr=("127.0.0.1", 5560)) + addrs = [] + a = container.register(RoleAgent()) + a.add_role(_Responder(value=1.0)) + addrs.append(a.addr) + noisy = container.register(RoleAgent()) + noisy.add_role(_NoiseResponder()) + addrs.append(noisy.addr) + + caller = container.register(RoleAgent()) + caller_role = _Caller(addrs, timeout=0.6) + caller.add_role(caller_role) + + async with activate([container]): + await caller_role.run() + + assert caller_role.responses is not None + # Only the well-typed reply survives. + assert len(caller_role.responses) == 1 + (reply,) = caller_role.responses.values() + assert isinstance(reply, _Reply) + assert reply.value == 1.0 From d9af8a47934a81dc7bfcff40e0a2b88fa44eb4a8 Mon Sep 17 00:00:00 2001 From: Rico Schrage Date: Sun, 17 May 2026 14:54:11 +0200 Subject: [PATCH 02/10] Refactoring of conversation and simplifying message dispatch and event handling. --- mango/agent/conversation.py | 141 +++++++++++-------------- mango/agent/core.py | 44 +++----- mango/agent/role.py | 205 ++++++++++++++++++++---------------- 3 files changed, 193 insertions(+), 197 deletions(-) diff --git a/mango/agent/conversation.py b/mango/agent/conversation.py index a1b40e1..8024ea4 100644 --- a/mango/agent/conversation.py +++ b/mango/agent/conversation.py @@ -1,51 +1,41 @@ """Multi-hop conversation primitive. -A *conversation* is a logically-grouped sequence of messages threaded -by a shared id. Mango's existing ``tracking_id`` mechanism solves the -single request/response case (see :meth:`AgentDelegates.send_tracked_message` -and :meth:`AgentDelegates.reply_to`), but multi-hop protocols — gossip, -auctions, holonic ADMM coordination — need the same id to route many -messages over time *without* being consumed on first reply. +A *conversation* groups a sequence of messages that share a single id +so participants can volley back and forth without each reply being +consumed on receipt (the way ``tracking_id`` is in +:meth:`AgentDelegates.send_tracked_message`). Use it for protocols +that span many hops — gossip, auctions, holonic ADMM coordination. -:class:`Conversation` is that abstraction: an async context manager -that holds a conversation id, a clock-aware timeout, mutable state, -and a receive queue. Anyone in the conversation (initiator or joiner) -can ``async for msg, meta in conversation`` and ``await -conversation.send(addr, payload)``. - -Two entry points on :class:`mango.RoleContext`: - -* :meth:`RoleContext.open_conversation` — initiator: generates a new id - and starts a fresh conversation. -* :meth:`RoleContext.join_conversation` — joiner: re-uses the id from - an inbound message's ``meta`` so the responder participates in the - same exchange. - -Both return a :class:`Conversation`. The timeout is enforced via the -agent's scheduler clock so simulation and real-time modes behave the -same way. - -Example — gossip-style initiator:: +The handle is an async context manager that yields a +:class:`Conversation` carrying the id, a user-controlled state dict, +and a receive queue:: async with self.context.open_conversation( timeout=10.0, - state={"target": -5.0, "delta": 0.0, "lambda": 0.01}, + state={"target": -5.0, "delta": 0.0}, ) as conv: - await conv.send(neighbours[0], GossipStep(payload=...)) + await conv.send(neighbour, GossipStep(...)) async for msg, meta in conv: update(conv.state, msg) if conv.state["delta"] >= conv.state["target"]: - conv.converge() # exits the loop on next iteration + conv.converge() continue - next_hop = pick(...) - await conv.send(next_hop, GossipStep(payload=...)) + await conv.send(pick_next_hop(), GossipStep(...)) -Example — joiner side:: +Responders join the existing exchange via the inbound ``meta``:: @on_message(GossipStep) async def on_step(self, msg, meta): async with self.context.join_conversation(meta) as conv: - ... # same async for loop as above + ... + +Two control methods end the iteration: + +* :meth:`Conversation.converge` — graceful: any messages already on + the queue still deliver, then the iterator exits. +* :meth:`Conversation.cancel` — abrupt: queued messages are discarded + and the iterator exits on the next pull. Also called by the context + manager when its clock-aware timeout fires. """ from __future__ import annotations @@ -53,40 +43,38 @@ async def on_step(self, msg, meta): import asyncio from typing import Any -# Metadata key used to carry the conversation id alongside ``tracking_id``. -# Distinct from ``tracking_id`` because the latter is consumed by the -# single-shot reply machinery; a conversation may legitimately carry -# both (e.g. when a participant uses ``reply_to`` inside a session). +# Carried alongside ``tracking_id`` because the latter is consumed by +# the single-shot reply machinery; a conversation message may legitimately +# bear both (e.g. when a participant calls ``reply_to`` inside a session). CONVERSATION_ID_KEY = "conversation_id" class Conversation: - """Active conversation handle owned by one agent. + """A live, agent-owned conversation handle. - The state dict is purely user-controlled — mango does not interpret - it. ``converge()`` ends the async iteration after the current - message; ``cancel()`` ends it immediately with no further yields. - Both are idempotent. + The ``state`` dict is user-controlled — mango never reads it. """ + # Queue sentinel that signals end-of-iteration. A class-level + # singleton so identity checks are unambiguous. + _END = object() + def __init__( self, *, - owner, + owner, # RoleContext conversation_id: str, state: dict[str, Any] | None = None, timeout: float | None = None, ) -> None: - self._owner = owner # RoleContext + self._owner = owner self._conversation_id = conversation_id self.state: dict[str, Any] = dict(state) if state else {} self._timeout = timeout self._queue: asyncio.Queue = asyncio.Queue() - self._converged: bool = False - self._cancelled: bool = False - self._timeout_handle = None # asyncio.Task | None + self._converged = False + self._cancelled = False - # -- introspection ------------------------------------------------ @property def conversation_id(self) -> str: return self._conversation_id @@ -95,67 +83,62 @@ def conversation_id(self) -> str: def closed(self) -> bool: return self._converged or self._cancelled - # -- public control ----------------------------------------------- def converge(self) -> None: - """Mark the conversation as converged. The async iterator - will terminate after delivering any already-queued message.""" + """Signal graceful end-of-iteration. + + Messages already on the queue still deliver; the iterator exits + once the queue drains. Subsequent inbound messages are dropped. + Idempotent. + """ + if self.closed: + return self._converged = True - # Wake the receiver if it's idle waiting on the queue. - self._queue.put_nowait((_SENTINEL_END, None)) + # Wake an idle ``__anext__`` if the queue was empty; otherwise + # the sentinel simply lands behind anything still pending and + # ends iteration after it drains. + self._queue.put_nowait((self._END, None)) def cancel(self) -> None: - """Drop the conversation immediately — discards any queued - messages and terminates the async iterator on the next pull.""" + """Signal abrupt end-of-iteration. + + Drops any queued messages and ends the iterator on the next + pull. Also called by the context manager when its timeout + fires. Idempotent. + """ + if self._cancelled: + return self._cancelled = True - self._queue.put_nowait((_SENTINEL_END, None)) + self._queue.put_nowait((self._END, None)) - # -- send helpers ------------------------------------------------- async def send(self, receiver_addr, content: Any, **kwargs) -> bool: """Send *content* tagged with this conversation's id.""" - meta_extra = {CONVERSATION_ID_KEY: self._conversation_id} - meta_extra.update(kwargs) return await self._owner.send_message( content, receiver_addr=receiver_addr, - **meta_extra, + **{CONVERSATION_ID_KEY: self._conversation_id, **kwargs}, ) async def broadcast(self, receivers, content: Any, **kwargs) -> None: - """Send *content* to every receiver in *receivers*.""" + """Send *content* to every address in *receivers*.""" for addr in receivers: await self.send(addr, content, **kwargs) - # -- internal hooks (used by RoleContext) ------------------------- + # -- agent-facing hook -------------------------------------------- def _on_inbound(self, content: Any, meta: dict) -> None: + """Called by the agent's router when a matching id arrives.""" if self.closed: return self._queue.put_nowait((content, meta)) - def _fire_timeout(self) -> None: - if self.closed: - return - # Treat timeout as cancellation — partial state remains - # accessible to the caller after exiting the context. - self._cancelled = True - self._queue.put_nowait((_SENTINEL_END, None)) - # -- iteration ---------------------------------------------------- def __aiter__(self): return self async def __anext__(self): + # cancel() is abrupt: anything still on the queue is dropped. if self._cancelled: raise StopAsyncIteration content, meta = await self._queue.get() - if content is _SENTINEL_END: + if content is self._END: raise StopAsyncIteration - if self._converged: - # Deliver this final message, then end on the next pull by - # putting a sentinel back into the queue. - self._queue.put_nowait((_SENTINEL_END, None)) return content, meta - - -# Sentinel placed in the queue to end the async iterator. A module- -# level singleton so identity checks are cheap and unambiguous. -_SENTINEL_END = object() diff --git a/mango/agent/core.py b/mango/agent/core.py index a00e543..f5dd16a 100644 --- a/mango/agent/core.py +++ b/mango/agent/core.py @@ -191,12 +191,12 @@ def health_runtime(self, tid: str = "default") -> TopologyHealth | None: class _GatherCollector: """Aggregates multi-reply tracked responses for :meth:`Agent.open_gather`. - The collector exposes a single :meth:`wait` coroutine that resolves - when either *expected* replies have arrived or :meth:`finish` has - been called explicitly (used by :meth:`RoleContext.gather` to - implement the timeout / quorum policy). Replies are stored under - the responding agent's :class:`AgentAddress` so the caller can - match each response to its source. + Replies are stored under the responding agent's + :class:`AgentAddress` so the caller can match each response to its + source. :meth:`wait` resolves once *expected* distinct replies + have arrived, or when :meth:`finish` is called explicitly (which + :meth:`RoleContext.gather` uses to honour the timeout / quorum + policy). """ def __init__( @@ -208,36 +208,28 @@ def __init__( self._expected = max(0, int(expected)) self._reply_type = reply_type self.responses: dict[AgentAddress, Any] = {} - # Lazily created so the collector can be constructed before a - # running event loop is required (Agent.open_gather runs in - # whichever context the caller is in). - self._done: asyncio.Event | None = None - - def _event(self) -> asyncio.Event: - if self._done is None: - self._done = asyncio.Event() - return self._done + self._done = asyncio.Event() def on_reply(self, content: Any, meta: dict) -> None: if self._reply_type is not None and not isinstance(content, self._reply_type): return - sender_id = meta.get("sender_id") - sender_addr = meta.get("sender_addr") - addr = AgentAddress(protocol_addr=sender_addr, aid=sender_id) + addr = AgentAddress( + protocol_addr=meta.get("sender_addr"), aid=meta.get("sender_id") + ) # First reply per sender wins — late duplicates (e.g. retries) # are dropped so the caller sees a stable mapping. if addr in self.responses: return self.responses[addr] = content if self._expected and len(self.responses) >= self._expected: - self._event().set() + self._done.set() def finish(self) -> None: """Force the collector to resolve (used on timeout / quorum hit).""" - self._event().set() + self._done.set() async def wait(self) -> None: - await self._event().wait() + await self._done.wait() class AgentContext: @@ -1037,12 +1029,10 @@ async def _check_inbox(self): # ``@on_message`` handler sees a freshly-nudged # neighbour score should it read one. self._nudge_topology_health(meta) - # Route to any open conversation that matches the - # message's conversation_id. Doesn't suppress - # normal handler dispatch — a role can still react - # to the message through ``@on_message`` AND join - # the conversation; the routing simply ensures the - # async-iterator pulls the message too. + # Push to any matching open conversation in addition + # to (not instead of) normal dispatch — a role can + # react via ``@on_message`` and pull from the + # conversation iterator on the same message. self._route_to_conversation(content, meta) if not self._handle_tracked_reply(content, meta): for _cond, _handler, _proc in self._behavior_message_subs: diff --git a/mango/agent/role.py b/mango/agent/role.py index 0e0d979..e2db76f 100644 --- a/mango/agent/role.py +++ b/mango/agent/role.py @@ -596,6 +596,22 @@ def live_neighbours( if health.is_live(my_addr, n, now, threshold=threshold) ] + def _bound_agent(self, operation: str): + """Return the owning :class:`RoleAgent`, or raise a clear error. + + Conversation- and gather-style helpers need access to the + agent's scheduler clock and message-routing tables; this + wrapper produces a uniform error when a role is used before + its agent attaches. + """ + agent = self._role_handler._agent + if agent is None: + raise RuntimeError( + f"{operation} requires a bound RoleAgent — the role's " + "context is not attached yet." + ) + return agent + def open_conversation( self, *, @@ -606,25 +622,17 @@ def open_conversation( Returns an async context manager whose body yields a :class:`~mango.agent.conversation.Conversation`. A new id is - generated; pass it (or use :meth:`Conversation.send`) so other - agents can route their messages back. + generated; other agents reply by echoing it (every call to + :meth:`Conversation.send` carries it automatically). The optional *timeout* is enforced via the agent's scheduler clock so behaviour is identical under :class:`AsyncioClock` (real time) and :class:`ExternalClock` (simulation). """ - import uuid as _uuid + import uuid - from mango.agent.conversation import Conversation - - return _ConversationContext( - role_context=self, - conv=Conversation( - owner=self, - conversation_id=str(_uuid.uuid4()), - state=state, - timeout=timeout, - ), + return self._make_conversation( + conversation_id=str(uuid.uuid4()), state=state, timeout=timeout ) def join_conversation( @@ -636,13 +644,13 @@ def join_conversation( ): """Join an existing conversation as a non-initiator. - Reads the conversation id from ``meta`` (the same dict passed - to ``@on_message`` handlers) and registers a local handle so - subsequent messages tagged with that id route here too. Use - when a role wants to process a multi-message exchange from the - responder side — typical for gossip forwarders. + Reads the id from ``meta`` (the dict passed to ``@on_message`` + handlers) and registers a local handle so subsequent messages + tagged with that id route here too. Typical use: a gossip + forwarder processes a multi-message exchange from the + responder side. """ - from mango.agent.conversation import CONVERSATION_ID_KEY, Conversation + from mango.agent.conversation import CONVERSATION_ID_KEY conv_id = meta.get(CONVERSATION_ID_KEY) if not conv_id: @@ -650,11 +658,24 @@ def join_conversation( "join_conversation requires a meta dict carrying " f"{CONVERSATION_ID_KEY!r}" ) + return self._make_conversation( + conversation_id=conv_id, state=state, timeout=timeout + ) + + def _make_conversation( + self, + *, + conversation_id: str, + state: dict | None, + timeout: float | None, + ) -> "_ConversationContext": + from mango.agent.conversation import Conversation + return _ConversationContext( role_context=self, conv=Conversation( owner=self, - conversation_id=conv_id, + conversation_id=conversation_id, state=state, timeout=timeout, ), @@ -671,65 +692,50 @@ async def gather( ) -> dict["AgentAddress", Any]: """Send *content* to every address in *receivers* and collect replies. - Returns a dict mapping the responding agent's - :class:`AgentAddress` to the reply content. Senders are + Returns a dict mapping each responding agent's + :class:`AgentAddress` to its reply content. Responders are expected to use :meth:`AgentDelegates.reply_to` (or otherwise - echo ``tracking_id`` with ``reply=True``) — every legacy - request/response pair in mango already follows that convention, - so existing responder roles work unchanged. + echo ``tracking_id`` with ``reply=True``) — every existing + request/response pair in mango follows that convention, so + responder roles work unchanged. :param receivers: iterable of :class:`AgentAddress` targets. - :param reply_type: optional class/tuple — replies not matching - this type are silently dropped (filters out tracking-id - collisions with unrelated traffic). - :param timeout: hard wallclock cap; if no quorum is reached by - then, returns whatever replies have arrived so far. - :param min_fraction: between 0 and 1. When :math:`\\geq` this - fraction of receivers have replied, the call returns - without waiting further. Defaults to 1.0 — wait for all, - time out otherwise. - """ - import uuid as _uuid - - agent = self._role_handler._agent - if agent is None: - raise RuntimeError( - "RoleContext.gather requires a RoleAgent — the role's " - "context is not bound to an agent yet." - ) + :param reply_type: optional class/tuple — replies of any other + type are silently dropped, filtering out tracking-id + collisions with unrelated traffic. + :param timeout: cap on how long to wait, measured on the + agent's scheduler clock. When it elapses, returns + whatever replies have arrived so far. + :param min_fraction: between 0 and 1. Returns as soon as + ``ceil(min_fraction * len(receivers))`` replies have + arrived. Defaults to 1.0 — wait for all, time out + otherwise. + """ + import uuid + + agent = self._bound_agent("RoleContext.gather") receivers = list(receivers) - n = len(receivers) - if n == 0: + if not receivers: return {} - expected = max(1, int(round(min_fraction * n))) + expected = max(1, int(round(min_fraction * len(receivers)))) - tracking_id = str(_uuid.uuid4()) + tracking_id = str(uuid.uuid4()) collector = agent.open_gather( tracking_id, expected=expected, reply_type=reply_type ) - # Timeout must follow mango's simulation clock — not wall time - # — so ``gather`` behaves identically under ``AsyncioClock`` - # (real-time) and ``ExternalClock`` (simulation). ``wait_for`` - # uses ``loop.time()`` and would block real seconds in sim mode. try: for addr in receivers: await self.send_message( - content, - receiver_addr=addr, - tracking_id=tracking_id, - ) - clock = agent.scheduler.clock - done_fut = asyncio.ensure_future(collector.wait()) - timeout_fut = asyncio.ensure_future(clock.sleep(timeout)) - try: - await asyncio.wait( - {done_fut, timeout_fut}, - return_when=asyncio.FIRST_COMPLETED, + content, receiver_addr=addr, tracking_id=tracking_id ) - finally: - for fut in (done_fut, timeout_fut): - if not fut.done(): - fut.cancel() + # Timeout follows the agent's scheduler clock so behaviour + # is identical under AsyncioClock (real time) and + # ExternalClock (simulation). asyncio.wait_for would block + # real seconds in sim mode. + await _race_first_completed( + collector.wait(), + agent.scheduler.clock.sleep(timeout), + ) finally: agent.close_gather(tracking_id) return dict(collector.responses) @@ -891,51 +897,68 @@ def on_agent_event(self, event: Any) -> None: """ +async def _race_first_completed(*awaitables) -> None: + """Run *awaitables* concurrently; return when the first finishes, + cancelling the rest. Used by gather and the conversation timeout + to race a "we're done" signal against a clock-aware timer. + """ + tasks = [asyncio.ensure_future(a) for a in awaitables] + try: + await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED) + finally: + for task in tasks: + if not task.done(): + task.cancel() + + class _ConversationContext: - """Async context manager returned by ``RoleContext.open_conversation`` / - ``RoleContext.join_conversation``. + """Async context manager returned by ``RoleContext.open_conversation`` + and ``RoleContext.join_conversation``. Owns the lifecycle of one :class:`~mango.agent.conversation.Conversation`: - registers it with the agent on ``__aenter__``, schedules an - optional clock-aware timeout, and unregisters on ``__aexit__``. - Splitting the context manager out keeps ``Conversation`` itself - free of mango-agent references and easy to unit-test. + + * ``__aenter__`` registers the conversation with the agent so + inbound messages route to it, and schedules the optional + clock-aware timeout. + * ``__aexit__`` unregisters the conversation and cancels the + timeout future. + + Keeping the context manager separate from :class:`Conversation` + itself leaves the data class free of mango-agent references and + easy to unit-test. """ def __init__(self, *, role_context: "RoleContext", conv) -> None: self._role_context = role_context self._conv = conv - self._timeout_future = None + self._timeout_task: asyncio.Task | None = None async def __aenter__(self): - agent = self._role_context._role_handler._agent - if agent is None: - raise RuntimeError( - "Conversation requires a bound RoleAgent — the role's " - "context is not attached yet." - ) + agent = self._role_context._bound_agent("Conversation") agent.open_conversation(self._conv) - if self._conv._timeout is not None: - # Use the agent's scheduler clock so timeouts respect - # simulation time when running under an ExternalClock. - clock = agent.scheduler.clock if agent.scheduler else None - if clock is not None: - self._timeout_future = asyncio.ensure_future( - self._fire_after(clock, self._conv._timeout) - ) + # Schedule the timeout on the agent's scheduler clock so it + # respects simulation time under an ExternalClock. + timeout = self._conv._timeout + clock = agent.scheduler.clock if agent.scheduler else None + if timeout is not None and clock is not None: + self._timeout_task = asyncio.ensure_future( + self._cancel_after(clock, timeout) + ) return self._conv async def __aexit__(self, exc_type, exc, tb): - if self._timeout_future is not None and not self._timeout_future.done(): - self._timeout_future.cancel() + if self._timeout_task is not None and not self._timeout_task.done(): + self._timeout_task.cancel() agent = self._role_context._role_handler._agent if agent is not None: agent.close_conversation(self._conv) return False - async def _fire_after(self, clock, delay: float) -> None: + async def _cancel_after(self, clock, delay: float) -> None: try: await clock.sleep(delay) except asyncio.CancelledError: return - self._conv._fire_timeout() + # Timeout == abrupt close: anything still queued is dropped, + # but conv.state remains readable after the context exits. + self._conv.cancel() From 73b7aa0026e3f5a2967525940b64fcdd6f40da42 Mon Sep 17 00:00:00 2001 From: Rico Schrage Date: Wed, 17 Jun 2026 19:28:15 +0200 Subject: [PATCH 03/10] Fixing awaiting tasks while susp. --- mango/util/scheduling.py | 29 +++++++++++++++++++++++++---- 1 file changed, 25 insertions(+), 4 deletions(-) diff --git a/mango/util/scheduling.py b/mango/util/scheduling.py index 1022c63..1dfc608 100644 --- a/mango/util/scheduling.py +++ b/mango/util/scheduling.py @@ -58,8 +58,13 @@ class Suspendable: Wraps a coroutine, intercepting __await__ to add the functionality of suspending. """ - def __init__(self, coro, ext_contr_event=None, kill_event=None): + def __init__(self, coro, ext_contr_event=None, kill_event=None, notify_task=None): self._coro = coro + # Optional owning ScheduledTask; when set, the wrapper marks it + # "sleeping" whenever the coroutine parks on a pending future so that + # simulation termination detection treats request/reply drivers (which + # await replies rather than the clock) as idle instead of deadlocking. + self._notify_task = notify_task self._kill_event = kill_event if ext_contr_event is not None: @@ -94,11 +99,24 @@ def __await__(self): return err.value else: send = iter_send + # Parked on a pending future (awaiting a reply/message rather than + # progressing) -> mark the owning task sleeping for the duration of + # the suspension. Idempotent, so it nests safely with clock sleeps. + parked = ( + self._notify_task is not None + and isinstance(signal, asyncio.Future) + and not signal.done() + ) + if parked: + self._notify_task.notify_sleeping() try: # pass signal via yielding it message = yield signal except BaseException as err: send, message = iter_throw, err + finally: + if parked: + self._notify_task.notify_running() def suspend(self): """ @@ -158,11 +176,14 @@ def __init__(self, clock: Clock = None, observable=True, on_stop=None) -> None: self._is_done = asyncio.Future() def notify_sleeping(self): - if self._is_observable: + # Idempotent so it composes with Suspendable's per-await notifications + # (a clock-sleep already marks the task sleeping before the wrapper + # also sees the yielded sleep future). + if self._is_observable and not self._is_sleeping.done(): self._is_sleeping.set_result(True) def notify_running(self): - if self._is_observable: + if self._is_observable and self._is_sleeping.done(): self._is_sleeping = asyncio.Future() @abstractmethod @@ -521,7 +542,7 @@ def schedule_task(self, task: ScheduledTask, src=None) -> asyncio.Task: """ l_task = None if self.suspendable: - coro = Suspendable(task.run()) + coro = Suspendable(task.run(), notify_task=task) l_task = asyncio.ensure_future(coro) else: coro = task.run() From 995b1a0633963da787d4c1939e4de82af6842d1d Mon Sep 17 00:00:00 2001 From: Rico Schrage Date: Tue, 7 Jul 2026 17:28:40 +0200 Subject: [PATCH 04/10] Fixing smells. Auto install. --- mango/agent/conversation.py | 8 ++ mango/agent/core.py | 77 ++++++++++++++++--- mango/agent/decorators.py | 22 ++++-- mango/agent/role.py | 60 +++++++++++++-- mango/container/mp.py | 3 +- mango/express/topology.py | 17 ++-- mango/simulation/communication.py | 4 +- mango/simulation/world.py | 4 + .../simulation/test_visualization.py | 4 +- 9 files changed, 162 insertions(+), 37 deletions(-) diff --git a/mango/agent/conversation.py b/mango/agent/conversation.py index 8024ea4..73e3c33 100644 --- a/mango/agent/conversation.py +++ b/mango/agent/conversation.py @@ -108,6 +108,14 @@ def cancel(self) -> None: if self._cancelled: return self._cancelled = True + # Drain anything already queued so a consumer blocked inside + # ``queue.get()`` receives the sentinel next, not a stale message — + # honouring the "drops queued messages" contract even mid-await. + while not self._queue.empty(): + try: + self._queue.get_nowait() + except asyncio.QueueEmpty: + break self._queue.put_nowait((self._END, None)) async def send(self, receiver_addr, content: Any, **kwargs) -> bool: diff --git a/mango/agent/core.py b/mango/agent/core.py index 25215da..f87f3a5 100644 --- a/mango/agent/core.py +++ b/mango/agent/core.py @@ -188,6 +188,20 @@ def health_runtime(self, tid: str = "default") -> TopologyHealth | None: return self._tid_to_health.get(tid) +def _addr_from_meta(meta: dict) -> AgentAddress: + """Build an :class:`AgentAddress` from a received ``meta``. + + The JSON codec decodes a ``(host, port)`` protocol address back into a + *list*, which is unhashable and never compares equal to the tuple form + kept internally. Normalise it to a tuple so the result is usable as a + dict key and equality-comparable against topology neighbour addresses. + """ + protocol_addr = meta.get("sender_addr") + if isinstance(protocol_addr, list): + protocol_addr = tuple(protocol_addr) + return AgentAddress(protocol_addr=protocol_addr, aid=meta.get("sender_id")) + + class _GatherCollector: """Aggregates multi-reply tracked responses for :meth:`Agent.open_gather`. @@ -213,9 +227,7 @@ def __init__( def on_reply(self, content: Any, meta: dict) -> None: if self._reply_type is not None and not isinstance(content, self._reply_type): return - addr = AgentAddress( - protocol_addr=meta.get("sender_addr"), aid=meta.get("sender_id") - ) + addr = _addr_from_meta(meta) # First reply per sender wins — late duplicates (e.g. retries) # are dropped so the caller sees a stable mapping. if addr in self.responses: @@ -298,6 +310,11 @@ def __init__(self) -> None: # carries a matching id are routed to the conversation's # async-iterator queue. self._conversations: dict[str, Conversation] = {} + # Reference count per open conversation id so multiple + # ``join_conversation`` contexts on the same id (e.g. an + # ``@on_message`` handler that re-fires while an earlier join is + # still open) share one handle instead of colliding. + self._conversation_refs: dict[str, int] = {} self._behavior_message_subs: list[tuple] = [] self._behavior_global_event_handlers: list[tuple] = [] self._behavior_agent_event_handlers: list[tuple] = [] @@ -526,18 +543,51 @@ def _route_to_conversation(self, content: Any, meta: dict) -> None: conv._on_inbound(content, meta) def open_conversation(self, conv: Conversation) -> None: - """Register *conv* so inbound messages with its id are routed - to it. Used internally by ``RoleContext.open_conversation``; - end users should not need to call this directly.""" + """Register *conv* as the initiator of a fresh conversation. + + Used internally by ``RoleContext.open_conversation``; end users + should not need to call this directly. Raises if the id is already + open — the initiator generates a unique id, so a collision here is a + genuine programming error. + """ if conv.conversation_id in self._conversations: raise ValueError( f"conversation {conv.conversation_id!r} already open on {self.aid}" ) self._conversations[conv.conversation_id] = conv + self._conversation_refs[conv.conversation_id] = 1 + + def join_conversation(self, conv: Conversation) -> Conversation: + """Register a *join* on a possibly-already-open conversation. + + Returns the handle to actually use: the existing one when the id is + already open (incrementing its reference count), otherwise *conv*. + This lets an ``@on_message`` handler that re-fires while an earlier + join is still open share a single handle instead of raising. + """ + cid = conv.conversation_id + existing = self._conversations.get(cid) + if existing is not None: + self._conversation_refs[cid] += 1 + return existing + self._conversations[cid] = conv + self._conversation_refs[cid] = 1 + return conv def close_conversation(self, conv: Conversation) -> None: - """Unregister *conv*. Called when the context manager exits.""" - self._conversations.pop(conv.conversation_id, None) + """Release one reference to *conv*; unregister at zero. + + Called when a conversation context manager exits. + """ + cid = conv.conversation_id + refs = self._conversation_refs.get(cid) + if refs is None: + return + if refs <= 1: + self._conversations.pop(cid, None) + self._conversation_refs.pop(cid, None) + else: + self._conversation_refs[cid] = refs - 1 def _nudge_topology_health(self, meta: dict) -> None: """Multiplicatively recover edge scores on every received message. @@ -552,10 +602,9 @@ def _nudge_topology_health(self, meta: dict) -> None: if svc is None or not svc._tid_to_health: return sender_id = meta.get("sender_id") - sender_addr = meta.get("sender_addr") if not sender_id: return - peer_addr = AgentAddress(protocol_addr=sender_addr, aid=sender_id) + peer_addr = _addr_from_meta(meta) scheduler = getattr(self, "scheduler", None) if scheduler is None or scheduler.clock is None: return @@ -1119,6 +1168,14 @@ async def shutdown(self): and deregister from the container""" await self.on_stop() + # Backstop: cancel any conversations still open (e.g. one opened + # with timeout=None whose peer never replied) so their iterators + # unblock and the registry does not leak past agent lifetime. + for conv in list(self._conversations.values()): + conv.cancel() + self._conversations.clear() + self._conversation_refs.clear() + if not self._stopped.done(): self._stopped.set_result(True) self.context.deregister(self.aid) diff --git a/mango/agent/decorators.py b/mango/agent/decorators.py index 97af404..15ac5ad 100644 --- a/mango/agent/decorators.py +++ b/mango/agent/decorators.py @@ -9,8 +9,10 @@ A role that uses these decorators does not need to implement ``setup()`` at all unless it has other initialisation logic. When -``setup()`` is implemented, the decorator wiring runs first (so the -explicit ``setup`` body can override or extend it). +``setup()`` is implemented, the decorator wiring runs first; the +explicit ``setup`` body can then *add* further subscriptions. There is +no unsubscribe, so ``setup`` extends the declarative wiring — it cannot +remove or replace it. The decorators are pure annotations: they stash configuration on the decorated method via the attribute name :data:`_MANGO_DISPATCH_META`. @@ -44,7 +46,7 @@ async def heartbeat(self): from __future__ import annotations -import asyncio +import inspect from collections.abc import Callable from dataclasses import dataclass, field from typing import Any @@ -125,10 +127,18 @@ def on_event(event_type: type) -> Callable: """Subscribe the decorated method to a co-located event type. The handler is called as ``handler(self, event, source)`` and runs - synchronously inside :meth:`RoleContext.emit_event`. + synchronously inside :meth:`RoleContext.emit_event`. Async handlers + are rejected: ``emit_event`` does not await, so an ``async def`` would + silently never run. """ def decorator(method: Callable) -> Callable: + if inspect.iscoroutinefunction(method): + raise TypeError( + f"@on_event must decorate a synchronous method; " + f"{method.__qualname__} is a coroutine (emit_event does not " + f"await handlers)" + ) meta = _get_or_create_meta(method) meta.event_subs.append({"event_type": event_type}) return method @@ -155,7 +165,7 @@ def periodic( """ def decorator(method: Callable) -> Callable: - if not asyncio.iscoroutinefunction(method): + if not inspect.iscoroutinefunction(method): raise TypeError( f"@periodic must decorate an async coroutine function; " f"{method.__qualname__} is not a coroutine" @@ -236,7 +246,7 @@ def condition(content, _meta, _mtype=mtype): def condition(content, _meta, _mtype=mtype, _w=where, _r=role): return isinstance(content, _mtype) and _w(_r, content, _meta) - if asyncio.iscoroutinefunction(method): + if inspect.iscoroutinefunction(method): callback = _bind_async(method, role) else: callback = _bind_sync(method, role) diff --git a/mango/agent/role.py b/mango/agent/role.py index 7342465..3fb8710 100644 --- a/mango/agent/role.py +++ b/mango/agent/role.py @@ -562,6 +562,31 @@ def neighbour_score(self, neighbour_addr, *, tid: str = "default") -> float | No now = agent.scheduler.clock.time if agent.scheduler else 0.0 return health.score(agent.addr, neighbour_addr, now) + def neighbour_scores(self, tid: str = "default") -> dict: + """Return ``{neighbour_addr: score}`` for every neighbour in *tid*. + + Empty when the agent is unbound or the topology has no + ``edge_health`` configured. Complements :meth:`neighbour_score` + (single neighbour) and :meth:`live_neighbours` (threshold filter). + """ + agent = self._role_handler._agent + if agent is None: + return {} + from mango.agent.core import State, TopologyService + + svc = agent.service_of_type(TopologyService, None) + if svc is None: + return {} + health = svc.health_runtime(tid) + if health is None: + return {} + now = agent.scheduler.clock.time if agent.scheduler else 0.0 + my_addr = agent.addr + return { + n: health.score(my_addr, n, now) + for n in svc.neighbors(state=State.NORMAL, tid=tid) + } + def live_neighbours( self, tid: str = "default", @@ -632,7 +657,10 @@ def open_conversation( import uuid return self._make_conversation( - conversation_id=str(uuid.uuid4()), state=state, timeout=timeout + conversation_id=str(uuid.uuid4()), + state=state, + timeout=timeout, + is_join=False, ) def join_conversation( @@ -659,7 +687,7 @@ def join_conversation( f"{CONVERSATION_ID_KEY!r}" ) return self._make_conversation( - conversation_id=conv_id, state=state, timeout=timeout + conversation_id=conv_id, state=state, timeout=timeout, is_join=True ) def _make_conversation( @@ -668,6 +696,7 @@ def _make_conversation( conversation_id: str, state: dict | None, timeout: float | None, + is_join: bool, ) -> "_ConversationContext": from mango.agent.conversation import Conversation @@ -679,6 +708,7 @@ def _make_conversation( state=state, timeout=timeout, ), + is_join=is_join, ) async def gather( @@ -711,13 +741,14 @@ async def gather( arrived. Defaults to 1.0 — wait for all, time out otherwise. """ + import math import uuid agent = self._bound_agent("RoleContext.gather") receivers = list(receivers) if not receivers: return {} - expected = max(1, int(round(min_fraction * len(receivers)))) + expected = max(1, math.ceil(min_fraction * len(receivers))) tracking_id = str(uuid.uuid4()) collector = agent.open_gather( @@ -837,7 +868,7 @@ def _bind(self, context: RoleContext) -> None: self._context = context # Apply class-level @on_message / @on_event / @periodic # decorators before user-defined ``setup`` runs, so the explicit - # setup body can override or extend the declarative wiring. + # setup body can add to (not remove from) the declarative wiring. from mango.agent.decorators import apply_dispatch apply_dispatch(self) @@ -928,19 +959,32 @@ class _ConversationContext: easy to unit-test. """ - def __init__(self, *, role_context: "RoleContext", conv) -> None: + def __init__( + self, *, role_context: "RoleContext", conv, is_join: bool = False + ) -> None: self._role_context = role_context self._conv = conv + self._is_join = is_join self._timeout_task: asyncio.Task | None = None async def __aenter__(self): agent = self._role_context._bound_agent("Conversation") - agent.open_conversation(self._conv) + reused = False + if self._is_join: + # Reuse an already-open handle (shared refcount) so a re-firing + # handler does not collide on the same id. + registered = agent.join_conversation(self._conv) + reused = registered is not self._conv + self._conv = registered + else: + agent.open_conversation(self._conv) # Schedule the timeout on the agent's scheduler clock so it - # respects simulation time under an ExternalClock. + # respects simulation time under an ExternalClock. A join that + # merely reused an existing handle does not arm its own timeout — + # the original owner's timeout governs the shared conversation. timeout = self._conv._timeout clock = agent.scheduler.clock if agent.scheduler else None - if timeout is not None and clock is not None: + if not reused and timeout is not None and clock is not None: self._timeout_task = asyncio.ensure_future( self._cancel_after(clock, timeout) ) diff --git a/mango/container/mp.py b/mango/container/mp.py index 84247c1..39203fc 100644 --- a/mango/container/mp.py +++ b/mango/container/mp.py @@ -1,4 +1,5 @@ import asyncio +import inspect import logging import os import warnings @@ -119,7 +120,7 @@ async def start_agent_loop(): ) message_pipe.close() event_pipe.close() - if asyncio.iscoroutinefunction(agent_creator): + if inspect.iscoroutinefunction(agent_creator): await agent_creator(container) else: agent_creator(container) diff --git a/mango/express/topology.py b/mango/express/topology.py index fd900bd..4864afd 100644 --- a/mango/express/topology.py +++ b/mango/express/topology.py @@ -112,7 +112,8 @@ def add_node(self, *agents: Agent) -> int: :param agents: zero or more agents to place at this node :return: integer node ID """ - node_id = max(self._graph.nodes, default=-1) + 1 + int_labels = [n for n in self._graph.nodes if isinstance(n, int)] + node_id = (max(int_labels) + 1) if int_labels else 0 self._graph.add_node(node_id, **{AGENT_NODE_KEY: AgentNode(list(agents))}) return node_id @@ -257,9 +258,7 @@ def _build_and_inject(self, *, update_connected: bool = True) -> None: self_neighbor = TopologyNeighbor( agent=agent, description=agent.description ) - existing = { - (ct, n.description.uid) for ct, n in self._connectors - } + existing = {(ct, n.description.uid) for ct, n in self._connectors} if (conn_type, agent.description.uid) not in existing: self._connectors.append((conn_type, self_neighbor)) @@ -501,10 +500,14 @@ def connect_topologies( :param directed: if ``True`` only ``topology_one → topology_two`` (default ``False``) """ - topology_one._connections.append((connection_type, topology_two)) + link_one = (connection_type, topology_two) + if link_one not in topology_one._connections: + topology_one._connections.append(link_one) topology_one._build_and_inject() if not directed: - topology_two._connections.append((connection_type, topology_one)) + link_two = (connection_type, topology_one) + if link_two not in topology_two._connections: + topology_two._connections.append(link_two) topology_two._build_and_inject() @@ -613,14 +616,12 @@ def topology_to_aid_graph(topology: Topology) -> nx.Graph: # Use with create_distribution_based_com_sim for topology-aware delays """ g: nx.Graph = nx.Graph() - aid_to_node: dict[str, int] = {} # Add one vertex per agent for node_id in topology.graph.nodes: agent_node: AgentNode = topology.graph.nodes[node_id][AGENT_NODE_KEY] for agent in agent_node.agents: g.add_node(agent.aid, agent=agent) - aid_to_node[agent.aid] = node_id # Same-node agents are fully connected (NORMAL) for node_id in topology.graph.nodes: diff --git a/mango/simulation/communication.py b/mango/simulation/communication.py index 29bf9e9..3f1dbbc 100644 --- a/mango/simulation/communication.py +++ b/mango/simulation/communication.py @@ -84,12 +84,14 @@ def __init__( loss_percent: float = 0.0, default_delay_s: float = 0.0, delay_s_directed_edge_dict: dict[tuple[str | None, str], float] | None = None, + rng: random.Random | None = None, ): self.loss_percent = loss_percent self.default_delay_s = default_delay_s self.delay_s_directed_edge_dict: dict[tuple[str | None, str], float] = ( delay_s_directed_edge_dict or {} ) + self._rng = rng or random def calculate_communication( self, @@ -100,7 +102,7 @@ def calculate_communication( for msg in messages: key = (msg.sender_id, msg.receiver_id) delay_s = self.delay_s_directed_edge_dict.get(key, self.default_delay_s) - reached = random.random() >= self.loss_percent + reached = self._rng.random() >= self.loss_percent results.append(PackageResult(reached=reached, delay_s=delay_s)) return CommunicationSimulationResult(package_results=results) diff --git a/mango/simulation/world.py b/mango/simulation/world.py index adeb1e0..23e4b55 100644 --- a/mango/simulation/world.py +++ b/mango/simulation/world.py @@ -188,6 +188,10 @@ def register(self, agent: Agent, suggested_aid: str | None = None) -> Agent: raise ValueError("Agent is already registered to a container") self._agents[aid] = agent agent._do_register(self, aid) + + install = getattr(self.environment, "install", None) + if callable(install): + install(agent, agent_id=aid) logger.debug("Registered agent '%s' with world", aid) if self.running: agent._do_start() diff --git a/tests/unit_tests/simulation/test_visualization.py b/tests/unit_tests/simulation/test_visualization.py index ecb032b..4aa1a8d 100644 --- a/tests/unit_tests/simulation/test_visualization.py +++ b/tests/unit_tests/simulation/test_visualization.py @@ -245,9 +245,7 @@ async def test_plot_recordings_colormap_hidden_axes_and_write_to(tmp_path): await step_simulation(world, step_size_s=1.0) out = tmp_path / "grid.png" - fig = plot_recordings( - world, figsize=(12, 8), colormap="viridis", write_to=str(out) - ) + fig = plot_recordings(world, figsize=(12, 8), colormap="viridis", write_to=str(out)) assert fig is not None assert out.exists() From 87d1ed050bde1c1ead319de72be7e20c41041fd1 Mon Sep 17 00:00:00 2001 From: Rico Schrage Date: Tue, 7 Jul 2026 17:36:27 +0200 Subject: [PATCH 05/10] Matplotlib as test requirement. --- pyproject.toml | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 03372d4..54fd793 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -39,7 +39,8 @@ test = [ "pytest", "pytest-cov", "pytest-asyncio", - "pre-commit" + "pre-commit", + "matplotlib" ] [project.urls] From 0b3c771259bd7d9b0cc7648a89fcb88ce2647e2c Mon Sep 17 00:00:00 2001 From: Rico Schrage Date: Tue, 7 Jul 2026 17:59:44 +0200 Subject: [PATCH 06/10] Doc gap. --- docs/source/topology.rst | 287 ++++++++++++++++++++++++++++++++++++++- 1 file changed, 281 insertions(+), 6 deletions(-) diff --git a/docs/source/topology.rst b/docs/source/topology.rst index 2d40c45..3ce81ea 100644 --- a/docs/source/topology.rst +++ b/docs/source/topology.rst @@ -15,6 +15,28 @@ Under the hood, topologies are backed by `networkx `_ ``Graph`` objects, so you can use all of networkx's graph construction helpers. +.. grid:: 1 2 2 3 + :gutter: 3 + + .. grid-item-card:: Build & assign + :shadow: sm + + Construct a graph, then place agents on its nodes — one by one, in + round-robin, or by a predicate. + + .. grid-item-card:: Query neighbours + :shadow: sm + + From any agent or role, list direct neighbours and filter them by link + state, characteristic, or a custom predicate. + + .. grid-item-card:: Connect topologies + :shadow: sm + + Bridge two independent topologies through *connector* agents for + hierarchical and multi-overlay designs. + + Building a topology from scratch ================================= @@ -65,13 +87,53 @@ agents as nodes, then wire them up with edges: objects. By default only *normal* (active) links are returned; you can filter by link state using the ``state`` parameter (see `Link states`_ below). +The ``with`` block is a convenience: :func:`~mango.create_topology` calls +:meth:`~mango.Topology.inject` for you when the block exits, pushing the +neighbourhood of every node into each agent's +:class:`~mango.TopologyService`. Build the graph first; **inject before you +start sending** (register the agents in a container after, as above). + + +Ready-made graph shapes +======================= + +For common structures you do not need to wire edges by hand. Each constructor +returns a :class:`~mango.Topology` you can populate with :func:`~mango.per_node` +(see `Assigning agents to nodes`_): + +.. list-table:: + :widths: 30 70 + :header-rows: 1 + + * - Constructor + - Shape + * - :func:`~mango.complete_topology` ``(n)`` + - Fully connected — every node linked to every other. + * - :func:`~mango.star_topology` ``(n)`` + - One hub (node ``0``) connected to ``n - 1`` leaves. + * - :func:`~mango.cycle_topology` ``(n)`` + - A ring — each node linked to its two neighbours. + * - :func:`~mango.graph_topology` ``(graph)`` + - Any existing ``networkx`` graph (alias: :func:`~mango.custom_topology`). + +.. code-block:: python + + from mango import star_topology, cycle_topology, graph_topology + import networkx as nx -Using a pre-built graph -======================== + hub_and_spokes = star_topology(5) # node 0 is the hub + ring = cycle_topology(6) + from_graph = graph_topology(nx.wheel_graph(7)) -If you already have a networkx graph, use :func:`~mango.per_node` to iterate -over the topology and attach one agent per node. Convenience constructors -like :func:`~mango.complete_topology` build common graph shapes for you: + +Assigning agents to nodes +========================= + +A freshly-built graph has empty nodes. There are three ways to place agents, +each ending in an automatic :meth:`~mango.Topology.inject`: + +**One agent per node** — iterate the nodes with :func:`~mango.per_node` and +call :meth:`AgentNode.add` on each: .. testcode:: @@ -110,6 +172,127 @@ like :func:`~mango.complete_topology` build common graph shapes for you: 1 1 +.. warning:: + + :func:`~mango.per_node` injects **after** the loop finishes. Breaking out + of the loop early skips the injection — iterate to the end (or call + :meth:`~mango.Topology.inject` yourself). + +**Round-robin** — spread a list of agents across the nodes with +:func:`~mango.auto_assign` (wraps around when there are more agents than +nodes): + +.. code-block:: python + + from mango import complete_topology, auto_assign + + topology = complete_topology(3) + auto_assign(topology, my_agents) # agent i → node i % 3 + +**By predicate** — :func:`~mango.assign_agents` places every agent for which +``condition(agent, node)`` is true. The node is passed as the second argument; +a matching agent is added to *every* node whose predicate returns true, so the +predicate usually tests the agent's type or attributes: + +.. code-block:: python + + from mango import assign_agents + + assign_agents( + lambda agent, node: isinstance(agent, SensorAgent), + topology, + all_agents, + ) + +.. note:: + + More than one agent may live on a single node. Agents sharing a node are + automatically each other's neighbours (a ``NORMAL`` link), so a node acts + like a fully-connected local cluster. + + +Node characteristics +==================== + +A *characteristic* is a short string label attached to an agent within a node +(for example ``"leader"`` or ``"aggregator"``). Neighbours can then filter by +role in the graph instead of by identity. + +Assign a characteristic while building the topology with +:meth:`~mango.Topology.set_characteristic`: + +.. code-block:: python + + with create_topology() as topology: + hub = topology.add_node(hub_agent) + leaf = topology.add_node(leaf_agent) + topology.add_edge(hub, leaf) + topology.set_characteristic(hub, hub_agent, "leader") + +At runtime, filter neighbours by characteristic and read your own label: + +.. code-block:: python + + # only neighbours tagged "leader" + leaders = agent.neighbors(has_characteristic="leader") + + # your own role in the graph + from mango import topology_characteristic + my_role = topology_characteristic(agent) # "" if none was set + + +Querying neighbours from agents and roles +========================================= + +:meth:`~mango.Agent.neighbors` is the direct method on an agent. When you work +with the :doc:`role system `, use the matching free functions — +they accept **either an agent or a role** and read the same +:class:`~mango.TopologyService`: + +.. list-table:: + :widths: 40 60 + :header-rows: 1 + + * - Function + - Returns + * - :func:`~mango.topology_neighbors` ``(agent_or_role)`` + - neighbour :class:`~mango.AgentAddress` list (with the same filters as + :meth:`~mango.Agent.neighbors`). + * - :func:`~mango.topology_node_id` ``(agent_or_role)`` + - the integer node ID the agent occupies. + * - :func:`~mango.topology_characteristic` ``(agent_or_role)`` + - the agent's characteristic label (``""`` if unset). + * - :func:`~mango.topology_connectors` ``(agent_or_role)`` + - connector addresses reachable across a topology link. + * - :func:`~mango.topology_connection_types` ``(agent_or_role)`` + - the connection-type labels available to the agent. + +All of them take a ``tid`` keyword so an agent that belongs to several +topologies can pick which neighbour set it means: + +.. code-block:: python + + from mango import Role, topology_neighbors + + class GossipRole(Role): + async def broadcast(self, payload): + for addr in topology_neighbors(self, tid="overlay"): + await self.context.send_message(payload, addr) + +The full filter set is shared by :meth:`~mango.Agent.neighbors` and +:func:`~mango.topology_neighbors`: + +.. code-block:: python + + topology_neighbors( + agent_or_role, + state=State.NORMAL, # by link state + tid="default", # which topology + has_characteristic="leader", # by neighbour label + include_connectors=("uplink",), # also reach connectors (see below) + match_func=lambda desc: ..., # arbitrary predicate on AgentDescription + ) + Link states ============ @@ -139,7 +322,99 @@ Use the ``state`` argument to query neighbours in a specific state: active_neighbours = agent.neighbors(state=State.NORMAL) inactive_links = agent.neighbors(state=State.INACTIVE) +To change the graph *after* it has been injected — flip a link to +``BROKEN`` when a peer goes silent, add or remove nodes — wrap the edits in +:func:`~mango.modify_topology`. It re-injects the updated neighbourhoods when +the block exits: + +.. code-block:: python + + from mango import modify_topology, State + + with modify_topology(topology) as t: + t.set_edge_state(0, 1, State.BROKEN) + t.remove_node(2) + # every affected agent now sees the updated neighbour set + + +Connecting multiple topologies +============================== + +Large systems are often built from several independent topologies — one per +region, per voltage level, per organisation — that must still exchange a few +messages across the boundary. Rather than merging them into one graph, mango +links them through *connector* agents. + +A connector is an agent nominated to represent its topology to the outside. +Nominate connectors either while building the topology with +:meth:`~mango.Topology.set_as_connector`, or ahead of time with +:func:`~mango.mark_as_connector` (the mark is picked up on the next inject): + +.. code-block:: python + + from mango import ( + complete_topology, per_node, connect_topologies, topology_connectors, + ) + + region_a = complete_topology(3, tid="a") + region_b = complete_topology(3, tid="b") + for node in per_node(region_a): + node.add(RegionAgent()) + for node in per_node(region_b): + node.add(RegionAgent()) + + # nominate one bridge agent on each side + region_a.set_as_connector(region_a.agents[0], connector_type="uplink") + region_b.set_as_connector(region_b.agents[0], connector_type="uplink") + + # link the two topologies via matching connector types + connect_topologies(region_a, region_b, connection_type="uplink") + +After :func:`~mango.connect_topologies`, each side's connector can see the +other side's connectors — without any node in ``region_a`` becoming a graph +neighbour of a node in ``region_b``: + +.. code-block:: python + + bridge = region_a.agents[0] + for addr in topology_connectors(bridge, tid="a"): + await bridge.send_message("cross-region hello", addr) + +Connectors are intentionally kept separate from ordinary neighbours so +intra-topology algorithms are unaffected. When you *do* want a single call to +reach both, pass ``include_connectors`` to +:meth:`~mango.Agent.neighbors` / :func:`~mango.topology_neighbors`: + +.. code-block:: python + + everyone = topology_neighbors(bridge, tid="a", include_connectors=("uplink",)) + +Pass ``directed=True`` to :func:`~mango.connect_topologies` for a one-way link +(``region_a`` reaches ``region_b`` but not vice versa). + + +Exporting to an agent-level graph +================================= + +A topology node may hold several agents, so the topology graph is not always +one-agent-per-node. :func:`~mango.topology_to_aid_graph` expands it into a +flat ``networkx`` graph whose nodes are agent IDs and whose edges carry the +link :class:`~mango.State`. This is the natural input for distance-based +communication delays in a simulation: + +.. code-block:: python + + from mango import topology_to_aid_graph + from mango.simulation import create_distribution_based_com_sim + + aid_graph = topology_to_aid_graph(topology) + com_sim = create_distribution_based_com_sim( + aid_graph, default_delay_per_edge_ms=20.0 + ) + .. seealso:: :doc:`simulation` — use :func:`~mango.run_with_simulation` to run - topology-based agents in a simulation world. + topology-based agents in a simulation world, and feed + :func:`~mango.topology_to_aid_graph` into the communication simulation for + delays that grow with graph distance. From 3f7f27c654cfac5038918697d79aae090dc4c998 Mon Sep 17 00:00:00 2001 From: Rico Schrage Date: Thu, 9 Jul 2026 00:24:48 +0200 Subject: [PATCH 07/10] Refactor conversation such that its usable in Agents as well. --- docs/source/index.rst | 10 +- docs/source/role-api.rst | 114 +++ docs/source/topology.rst | 104 +++ docs/source/transactions.rst | 229 +++++ mango/__init__.py | 1 + mango/agent/conversation.py | 156 +++- mango/agent/core.py | 279 ++++-- mango/agent/role.py | 216 +---- mango/container/core.py | 7 +- mango/express/api.py | 2 +- mango/simulation/__init__.py | 19 +- mango/simulation/container.py | 196 +++++ mango/simulation/environment.py | 2 +- mango/simulation/recording.py | 235 +++++ mango/simulation/visualization.py | 2 +- mango/simulation/world.py | 814 ++++++------------ tests/unit_tests/role/conversation_test.py | 333 ++++++- tests/unit_tests/role/gather_test.py | 35 + .../unit_tests/simulation/test_simulation.py | 12 +- 19 files changed, 1900 insertions(+), 866 deletions(-) create mode 100644 docs/source/transactions.rst create mode 100644 mango/simulation/container.py create mode 100644 mango/simulation/recording.py diff --git a/docs/source/index.rst b/docs/source/index.rst index e49718e..9b3fc4d 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -69,7 +69,14 @@ Features **Role system** ^^^ - Compose agent behaviour from small, reusable ``Role`` classes with shared state and event subscriptions. + Compose agent behaviour from small, reusable ``Role`` classes with shared state and event subscriptions — wired declaratively with ``@on_message``, ``@on_event``, and ``@periodic``. + + .. grid-item-card:: + :shadow: sm + + **Transactional messaging** + ^^^ + Multi-reply ``gather`` with quorum and timeout, and multi-hop conversations for gossip, auctions, and negotiation — clock-aware in real time and simulation. .. grid-item-card:: :shadow: sm @@ -227,6 +234,7 @@ Where to go next agents-container message exchange + transactions role-api scheduling topology diff --git a/docs/source/role-api.rst b/docs/source/role-api.rst index 2ca670d..db8b490 100644 --- a/docs/source/role-api.rst +++ b/docs/source/role-api.rst @@ -509,6 +509,120 @@ the event — useful when multiple roles can emit the same event type. have run. Avoid long-running or awaiting logic inside event handlers. +---- + +Declarative dispatch with decorators +==================================== + +The subscription and scheduling calls above (:meth:`~mango.RoleContext.subscribe_message`, +:meth:`~mango.RoleContext.subscribe_event`, periodic tasks) can all be written +declaratively by decorating the handler method directly. A role that uses the +decorators often needs no ``setup`` at all — the wiring is read from the class +at bind time: + +.. list-table:: + :widths: 26 74 + :header-rows: 1 + + * - Decorator + - Replaces + * - :func:`~mango.on_message` + - a :meth:`~mango.RoleContext.subscribe_message` call. + * - :func:`~mango.on_event` + - a :meth:`~mango.RoleContext.subscribe_event` call. + * - :func:`~mango.periodic` + - a periodic ``schedule_periodic_task`` call. + +.. code-block:: python + + from mango import Role, on_message, on_event, periodic, sender_addr + + class Worker(Role): + @on_message(Task) + async def on_task(self, content, meta): + await self.context.reply_to(Result(...), meta) + + @on_event(ConfigChanged) + def on_config(self, event, source): + self.config = event.config + + @periodic(every=1.0) + async def heartbeat(self): + await self.context.send_message(Beat(), self.leader) + +The equivalent hand-written ``setup`` would register three callbacks; the +decorated form keeps each handler next to its trigger and removes the +boilerplate. + +on_message +---------- + +:func:`~mango.on_message` delivers messages where +``isinstance(content, message_type)`` is true. The handler is called as +``handler(self, content, meta)``. **Async handlers are supported directly** — +they are scheduled as instant tasks automatically (no manual coroutine shim). + +Two keyword options refine the subscription: + +* ``where`` — an extra predicate ``where(self, content, meta) -> bool``. It + receives ``self``, so it can read role state instead of capturing it in a + class-time closure. +* ``priority`` — dispatch order when several handlers match (lower runs first, + default ``0``), mirroring :meth:`~mango.RoleContext.subscribe_message`. + +.. code-block:: python + + class Router(Role): + @on_message(Packet, where=lambda self, c, m: c.ttl > 0, priority=0) + async def forward(self, content, meta): + ... + +.. note:: + + Each async ``@on_message`` invocation runs as an independent task, so + handlers for different messages may run concurrently and out of order. If a + handler mutates role state and must not interleave, register it the + imperative way with a + :class:`~mango.WaitingMessagePreprocessor` (see `Message preprocessors`_). + +on_event +-------- + +:func:`~mango.on_event` subscribes to a co-located inter-role event (see +`Inter-role events`_). The handler is ``handler(self, event, source)`` +and runs **synchronously** inside :meth:`~mango.RoleContext.emit_event`, so it +must be a plain (non-async) method — decorating an ``async def`` raises a +``TypeError`` at class-definition time rather than silently never running. + +periodic +-------- + +:func:`~mango.periodic` schedules an async method to run on a fixed period. +``every`` is either a number of seconds or the *name* of an instance attribute +read at attach time (handy for per-instance periods). The optional ``only_if`` +predicate gates each firing: + +.. code-block:: python + + class Poller(Role): + poll_period_s = 0.5 + + @periodic(every="poll_period_s", only_if=lambda self: self.is_leader) + async def poll(self): + await self.context.gather(Ping(), self.peers) + +When ``only_if(self)`` is false the scheduled task still fires but returns +early, replacing the ``if not leader: return`` guard by hand. + +.. tip:: + + Decorators and ``setup`` compose: decorator wiring is applied **before** + ``setup`` runs, so ``setup`` can add further subscriptions. There is no + unsubscribe — ``setup`` extends the declarative wiring, it cannot remove it. + Stacking decorators on one method (e.g. two ``@on_message``) is supported, + and decorated handlers on a base ``Role`` are inherited by subclasses. + + ---- Deactivating and activating roles diff --git a/docs/source/topology.rst b/docs/source/topology.rst index 3ce81ea..382dc83 100644 --- a/docs/source/topology.rst +++ b/docs/source/topology.rst @@ -294,6 +294,33 @@ The full filter set is shared by :meth:`~mango.Agent.neighbors` and ) +Broadcasting to neighbours +========================== + +Sending to every neighbour is common enough to have its own helper. +:func:`~mango.broadcast_to_neighbors` resolves the neighbour set (with the same +filters as :func:`~mango.topology_neighbors`) and awaits a ``send_message`` to +each, returning the addresses it reached: + +.. code-block:: python + + from mango import Role, broadcast_to_neighbors + + class GossipRole(Role): + async def on_ready(self): + await broadcast_to_neighbors(self, Rumor("hello")) + + async def relay(self, rumor): + # only push to leaders on the "overlay" topology + await broadcast_to_neighbors( + self, rumor, tid="overlay", has_characteristic="leader" + ) + +Any extra keyword arguments are forwarded to ``send_message`` (e.g. +``tracking_id``), and ``include_connectors`` extends the broadcast across a +topology link. + + Link states ============ @@ -393,6 +420,83 @@ Pass ``directed=True`` to :func:`~mango.connect_topologies` for a one-way link (``region_a`` reaches ``region_b`` but not vice versa). +Tracking link health +==================== + +Gossip and peer-to-peer protocols often need to know whether a neighbour is +still responsive, so a sender can route around silent peers. Enable +*continuous link-health tracking* by passing an :class:`~mango.EdgeHealth` +config to :func:`~mango.create_topology` (or a :class:`~mango.Topology` +directly). Every agent in the topology then gets an automatic hook that: + +* **decays** each neighbour's score over time (``decay_per_s`` per second of + silence), and +* **recovers** it on every message received from that neighbour + (``score += (1 - score) * recovery_rate``). + +.. code-block:: python + + from mango import create_topology, EdgeHealth + + with create_topology( + tid="grid", + edge_health=EdgeHealth( + decay_per_s=0.125, # lose ~1/8 of the score per silent second + recovery_rate=0.6, # rebuild 60 % of the gap on each message + liveness_threshold=0.5, # "live" at or above this score + ), + ) as topology: + ... + +.. list-table:: + :widths: 26 14 60 + :header-rows: 1 + + * - Field + - Default + - Meaning + * - ``decay_per_s`` + - ``0.125`` + - Score lost per second of silence. + * - ``recovery_rate`` + - ``0.6`` + - Fraction of the missing score rebuilt per received message ``(0, 1]``. + * - ``initial`` + - ``1.0`` + - Starting score for a never-contacted neighbour (optimistic). + * - ``liveness_threshold`` + - ``0.5`` + - Default cutoff for :meth:`~mango.RoleContext.live_neighbours`. + +From a role, query health through the context. These are no-ops-with-fallback +when the topology has **no** ``edge_health`` configured, so a role can call them +unconditionally: + +.. code-block:: python + + class GossipRole(Role): + async def push(self, payload): + # neighbours at or above the liveness threshold (or all of them, + # if this topology has no health tracking) + for addr in self.context.live_neighbours(tid="grid"): + await self.context.send_message(payload, addr) + + def report(self): + scores = self.context.neighbour_scores(tid="grid") # {addr: score} + one = self.context.neighbour_score(some_addr, tid="grid") + +* :meth:`~mango.RoleContext.live_neighbours` — the neighbours at or above the + threshold (override per call with ``threshold=``). +* :meth:`~mango.RoleContext.neighbour_scores` — a ``{AgentAddress: score}`` map + for all neighbours. +* :meth:`~mango.RoleContext.neighbour_score` — the score for one neighbour, or + ``None`` when tracking is off. + +All scores are read on the agent's scheduler clock, so decay behaves +identically in real time and under a simulation +:class:`~mango.util.clock.ExternalClock`. + + Exporting to an agent-level graph ================================= diff --git a/docs/source/transactions.rst b/docs/source/transactions.rst new file mode 100644 index 0000000..446e973 --- /dev/null +++ b/docs/source/transactions.rst @@ -0,0 +1,229 @@ +======================= +Transactional messaging +======================= + +Many multi-agent protocols are not fire-and-forget: an agent asks a question +and needs the answer, polls a group of peers and waits for a quorum, or runs a +back-and-forth exchange that spans many hops. mango provides three primitives +for this, each a step up in scope: + +.. list-table:: + :widths: 25 40 35 + :header-rows: 1 + + * - Primitive + - Use it when… + - Waits for + * - :meth:`~mango.Agent.send_tracked_message` + - you send one request and want the one matching reply. + - a single reply (a callback fires). + * - :meth:`~mango.Agent.gather` + - you ask *many* peers the same thing in one round. + - a quorum of replies, or a timeout. + * - :meth:`~mango.Agent.open_conversation` + - the exchange runs over many messages / many hops. + - as long as you keep the conversation open. + +All three primitives are available on plain agents (``self.gather(...)``) +and inside roles (``self.context.gather(...)``) alike. +The single tracked request/reply pair (``send_tracked_message`` / ``reply_to``) +is covered on the :doc:`message exchange ` page. This page +covers the two multi-message primitives. +Both are clock-aware: their timeouts run on the agent's scheduler clock, so +they behave identically under real time (:class:`~mango.util.clock.AsyncioClock`) +and simulation time (:class:`~mango.util.clock.ExternalClock`). + + +Collecting replies with gather +============================== + +:meth:`~mango.Agent.gather` sends one message to many receivers and +returns their replies as a dict keyed by the responding agent's +:class:`~mango.AgentAddress`. It is the right tool for one-round request +scatter/gather — a price poll, a capability query, a distributed sum: + +.. code-block:: python + + from mango import Role, sender_addr + + class Coordinator(Role): + async def poll_prices(self, sellers): + replies = await self.context.gather( + PriceRequest(), + sellers, # iterable of AgentAddress + reply_type=PriceOffer, # ignore any non-PriceOffer traffic + timeout=2.0, # scheduler-clock seconds + ) + # replies: {AgentAddress: PriceOffer} + best = min(replies.values(), key=lambda offer: offer.price) + return best + +Responders need no special API — they just reply, echoing the ``tracking_id``. +:meth:`~mango.Agent.reply_to` does that automatically: + +.. code-block:: python + + class Seller(Role): + @on_message(PriceRequest) + async def on_request(self, content, meta): + await self.context.reply_to(PriceOffer(price=self.price), meta) + +**Quorum and partial results.** By default ``gather`` waits for *all* +receivers or times out. Lower ``min_fraction`` to return as soon as a fraction +has answered — useful when stragglers should not hold up progress: + +.. code-block:: python + + # return as soon as ⌈0.5 · N⌉ replies are in (majority quorum) + replies = await self.context.gather(Ping(), peers, min_fraction=0.5) + +On timeout, ``gather`` returns whatever arrived so far rather than raising — so +always be prepared for fewer entries than receivers. The first reply from each +sender wins; late duplicates are dropped so the mapping is stable. + +.. list-table:: + :widths: 25 20 55 + :header-rows: 1 + + * - Parameter + - Default + - Meaning + * - ``reply_type`` + - ``None`` + - Only accept replies of this type/tuple; others are dropped. + * - ``timeout`` + - ``5.0`` + - Scheduler-clock seconds before returning partial results. + * - ``min_fraction`` + - ``1.0`` + - Return early once ``⌈min_fraction · len(receivers)⌉`` replied. + + +Multi-hop conversations +======================= + +A *conversation* groups a sequence of messages under one shared id so +participants can volley back and forth without every reply being consumed on +receipt. Reach for it when a single tracked reply or one ``gather`` round is +not enough — gossip, auctions, multi-round negotiation, iterative distributed +optimisation. + +The initiator opens a conversation, sends into it, and iterates the replies as +they arrive: + +.. code-block:: python + + from mango import Role, sender_addr + + class Negotiator(Role): + async def run(self, peer): + async with self.context.open_conversation( + timeout=10.0, + state={"round": 0}, + ) as conv: + await conv.send(peer, Offer(price=100)) + async for content, meta in conv: + conv.state["round"] += 1 + if self.acceptable(content) or conv.state["round"] >= 5: + conv.converge() # graceful stop + continue + await conv.send(peer, Offer(price=content.price - 5)) + +The handle yielded by ``open_conversation`` is a +:class:`~mango.Conversation`: + +.. list-table:: + :widths: 32 68 + :header-rows: 1 + + * - Member + - Description + * - ``await conv.send(addr, content, **kwargs)`` + - Send tagged with this conversation's id; returns the send result. + * - ``await conv.broadcast(addrs, content)`` + - Send the same content to several receivers; returns a dict of + per-receiver send results keyed by address. + * - ``conv.state`` + - A user-owned ``dict`` for carrying protocol state; mango never reads it. + * - ``async for content, meta in conv`` + - Iterate inbound messages of this conversation until it closes. + * - ``conv.converge()`` + - **Graceful** stop: already-queued messages still deliver, then the + iterator ends. + * - ``conv.cancel()`` + - **Abrupt** stop: queued messages are dropped; the iterator ends on the + next pull. + * - ``conv.closed`` + - ``True`` once converged or cancelled. + +.. note:: + + A conversation message is delivered to matching ``@on_message`` handlers + *and* pushed to the conversation iterator — the iterator receives it in + addition to, not instead of, normal dispatch. Replying with + :meth:`~mango.Agent.reply_to` inside a conversation keeps the id, so the + reply routes back into the initiator's iterator automatically. + +Responding side +--------------- + +A responder joins the existing exchange from inside an ``@on_message`` handler +using the inbound ``meta`` (which carries the conversation id). The simplest +responders reply once and let the ``with`` block close the handle: + +.. code-block:: python + + class Peer(Role): + @on_message(Offer) + async def on_offer(self, content, meta): + async with self.context.join_conversation(meta) as conv: + await conv.send(sender_addr(meta), Offer(price=content.price + 3)) + +You may also keep the joined conversation open and iterate it for follow-ups. +Every join owns an independent handle with its own ``state`` and ``timeout``; +when several handles for the same id are open on one agent (e.g. a re-firing +handler while an earlier join is still iterating), each inbound message is +delivered to *all* of them. + +``conv.state`` lives on the handle, so it does not persist across separate +``join_conversation`` blocks — a responder that accumulates state over many +handler invocations should keep it on the role (or agent) instance instead. + +Ending an exchange +------------------ + +Two control methods end iteration, plus the timeout: + +* :meth:`~mango.agent.conversation.Conversation.converge` — graceful: drain + what is already queued, then stop. Use it when the protocol reached a + result. +* :meth:`~mango.agent.conversation.Conversation.cancel` — abrupt: drop the + queue and stop. Use it to abandon. +* **Timeout** — pass ``timeout=`` to ``open_conversation`` / + ``join_conversation``. When it elapses (measured on the scheduler clock) the + conversation is cancelled automatically. ``conv.state`` remains readable + after the ``with`` block exits, so you can inspect the outcome. + +Leaving the ``async with`` block also cancels the handle — an iterator that +escaped the block terminates instead of waiting forever. + +.. warning:: + + ``open_conversation`` defaults to ``timeout=None`` (no timeout) — unlike + ``gather``, conversations are long-lived by design. An un-timed + conversation whose peer never replies will block its ``async for`` + forever — always set a ``timeout`` for protocols that can stall. (Open + conversations are cleaned up on agent shutdown as a backstop.) + +.. note:: + + ``conv.send`` also stamps the id onto content that has a + ``conversation_id`` attribute (e.g. ACL messages created with + :func:`~mango.create_acl`), so the id survives transports that send such + content without the surrounding meta. + +.. seealso:: + + :doc:`simulation` — conversation and ``gather`` timeouts advance with + simulation time under :func:`~mango.run_with_simulation`, so the same + protocol code runs unchanged in real time and in a discrete-event world. diff --git a/mango/__init__.py b/mango/__init__.py index f972a21..c7b08ab 100644 --- a/mango/__init__.py +++ b/mango/__init__.py @@ -1,4 +1,5 @@ from .messages.message import create_acl, Performatives +from .agent.conversation import Conversation from .agent.core import ( Agent, AgentAddress, diff --git a/mango/agent/conversation.py b/mango/agent/conversation.py index 73e3c33..13adeb5 100644 --- a/mango/agent/conversation.py +++ b/mango/agent/conversation.py @@ -10,7 +10,7 @@ :class:`Conversation` carrying the id, a user-controlled state dict, and a receive queue:: - async with self.context.open_conversation( + async with self.open_conversation( timeout=10.0, state={"target": -5.0, "delta": 0.0}, ) as conv: @@ -22,6 +22,9 @@ continue await conv.send(pick_next_hop(), GossipStep(...)) +(Inside a role, use ``self.context.open_conversation`` — the API is +available on plain agents and role contexts alike.) + Responders join the existing exchange via the inbound ``meta``:: @on_message(GossipStep) @@ -29,23 +32,42 @@ async def on_step(self, msg, meta): async with self.context.join_conversation(meta) as conv: ... +Every join owns an independent handle with its own ``state`` and +``timeout``; when several handles for the same id are open on one +agent (e.g. an ``@on_message`` handler that re-fires while an earlier +join is still open), each inbound message is delivered to *all* of +them. Note that a conversation message is also dispatched to matching +``@on_message`` handlers as usual — the iterator receives it *in +addition to*, not instead of, normal dispatch. + Two control methods end the iteration: * :meth:`Conversation.converge` — graceful: any messages already on the queue still deliver, then the iterator exits. * :meth:`Conversation.cancel` — abrupt: queued messages are discarded - and the iterator exits on the next pull. Also called by the context - manager when its clock-aware timeout fires. + and the iterator exits on the next pull. Also called when the + clock-aware timeout fires and when the context manager exits. """ from __future__ import annotations import asyncio -from typing import Any +import logging +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from collections.abc import Iterable + + from mango.agent.core import AgentDelegates + from mango.messages.message import AgentAddress + from mango.util.clock import Clock + +logger = logging.getLogger(__name__) # Carried alongside ``tracking_id`` because the latter is consumed by # the single-shot reply machinery; a conversation message may legitimately -# bear both (e.g. when a participant calls ``reply_to`` inside a session). +# bear both (e.g. when a participant calls ``reply_to`` inside a session — +# ``reply_to`` echoes the conversation id automatically). CONVERSATION_ID_KEY = "conversation_id" @@ -56,13 +78,15 @@ class Conversation: """ # Queue sentinel that signals end-of-iteration. A class-level - # singleton so identity checks are unambiguous. + # singleton so identity checks are unambiguous. Once enqueued it is + # re-enqueued by every consumer that reads it, so it acts as a latch: + # any number of concurrent or later iterators terminate. _END = object() def __init__( self, *, - owner, # RoleContext + owner: AgentDelegates, conversation_id: str, state: dict[str, Any] | None = None, timeout: float | None = None, @@ -79,6 +103,11 @@ def __init__( def conversation_id(self) -> str: return self._conversation_id + @property + def timeout(self) -> float | None: + """Scheduler-clock timeout this handle was opened with, if any.""" + return self._timeout + @property def closed(self) -> bool: return self._converged or self._cancelled @@ -93,17 +122,14 @@ def converge(self) -> None: if self.closed: return self._converged = True - # Wake an idle ``__anext__`` if the queue was empty; otherwise - # the sentinel simply lands behind anything still pending and - # ends iteration after it drains. self._queue.put_nowait((self._END, None)) def cancel(self) -> None: """Signal abrupt end-of-iteration. Drops any queued messages and ends the iterator on the next - pull. Also called by the context manager when its timeout - fires. Idempotent. + pull. Also called when the timeout fires and when the + conversation's context manager exits. Idempotent. """ if self._cancelled: return @@ -118,23 +144,49 @@ def cancel(self) -> None: break self._queue.put_nowait((self._END, None)) - async def send(self, receiver_addr, content: Any, **kwargs) -> bool: - """Send *content* tagged with this conversation's id.""" + async def send(self, receiver_addr: AgentAddress, content: Any, **kwargs) -> bool: + """Send *content* tagged with this conversation's id. + + The tag always wins over a ``conversation_id`` kwarg. When + *content* itself has a ``conversation_id`` attribute (e.g. an + ACL message), it is stamped too — container transports send such + content without the surrounding meta, so the id must travel + inside the message to survive a cross-container hop. + """ + if hasattr(content, CONVERSATION_ID_KEY): + try: + setattr(content, CONVERSATION_ID_KEY, self._conversation_id) + except AttributeError: + logger.warning( + "Conversation %s: content of type %s has an immutable " + "conversation_id attribute; the id travels only in meta " + "and will not survive transports that drop it.", + self._conversation_id, + type(content).__name__, + ) return await self._owner.send_message( content, receiver_addr=receiver_addr, - **{CONVERSATION_ID_KEY: self._conversation_id, **kwargs}, + **{**kwargs, CONVERSATION_ID_KEY: self._conversation_id}, ) - async def broadcast(self, receivers, content: Any, **kwargs) -> None: - """Send *content* to every address in *receivers*.""" - for addr in receivers: - await self.send(addr, content, **kwargs) + async def broadcast( + self, receivers: Iterable[AgentAddress], content: Any, **kwargs + ) -> dict[AgentAddress, bool]: + """Send *content* to every address in *receivers*. + + :return: send result per receiver, keyed by address. + """ + return {addr: await self.send(addr, content, **kwargs) for addr in receivers} # -- agent-facing hook -------------------------------------------- def _on_inbound(self, content: Any, meta: dict) -> None: """Called by the agent's router when a matching id arrives.""" if self.closed: + logger.debug( + "Dropping inbound message for closed conversation %s", + self._conversation_id, + ) return self._queue.put_nowait((content, meta)) @@ -148,5 +200,71 @@ async def __anext__(self): raise StopAsyncIteration content, meta = await self._queue.get() if content is self._END: + # Re-enqueue the sentinel so every other consumer — blocked + # concurrently or arriving later — terminates as well. + self._queue.put_nowait((self._END, None)) raise StopAsyncIteration return content, meta + + +class _ConversationContext: + """Async context manager returned by ``open_conversation`` and + ``join_conversation`` (see :class:`mango.agent.core.AgentDelegates`). + + Owns the lifecycle of one :class:`Conversation`: + + * ``__aenter__`` registers the handle with the owning agent so + inbound messages route to it, and arms the optional clock-aware + timeout. + * ``__aexit__`` disarms the timeout, unregisters the handle, and + cancels it — an iterator that escaped the block terminates instead + of waiting on a queue that can never be fed again. ``conv.state`` + stays readable after exit. + """ + + def __init__(self, *, owner: AgentDelegates, conv: Conversation) -> None: + self._owner = owner + self._conv = conv + self._timeout_task: asyncio.Task | None = None + + async def __aenter__(self) -> Conversation: + agent = self._owner._bound_agent("Conversation") + agent._register_conversation(self._conv) + # The timeout runs on the agent's scheduler clock so it respects + # simulation time under an ExternalClock. + timeout = self._conv.timeout + if timeout is not None: + clock = agent.scheduler.clock if agent.scheduler else None + if clock is None: + logger.warning( + "Conversation %s: timeout=%s requested but the agent has " + "no scheduler clock — the timeout will not be enforced.", + self._conv.conversation_id, + timeout, + ) + else: + self._timeout_task = asyncio.ensure_future( + self._cancel_after(clock, timeout) + ) + return self._conv + + async def __aexit__(self, exc_type, exc, tb): + if self._timeout_task is not None and not self._timeout_task.done(): + self._timeout_task.cancel() + try: + agent = self._owner._bound_agent("Conversation") + except RuntimeError: + agent = None + if agent is not None: + agent._unregister_conversation(self._conv) + self._conv.cancel() + return False + + async def _cancel_after(self, clock: Clock, delay: float) -> None: + try: + await clock.sleep(delay) + except asyncio.CancelledError: + return + # Timeout == abrupt close: anything still queued is dropped, + # but conv.state remains readable after the context exits. + self._conv.cancel() diff --git a/mango/agent/core.py b/mango/agent/core.py index f87f3a5..cf9ccb4 100644 --- a/mango/agent/core.py +++ b/mango/agent/core.py @@ -9,6 +9,7 @@ import asyncio import logging +import math import uuid from abc import ABC from dataclasses import dataclass, field @@ -18,9 +19,11 @@ from ..messages.message import AgentAddress from ..util.clock import Clock from ..util.scheduling import ScheduledProcessTask, ScheduledTask, Scheduler +from .conversation import CONVERSATION_ID_KEY, Conversation, _ConversationContext if TYPE_CHECKING: - from mango.agent.conversation import Conversation + from collections.abc import Iterable + from mango.express.health import TopologyHealth logger = logging.getLogger(__name__) @@ -202,14 +205,28 @@ def _addr_from_meta(meta: dict) -> AgentAddress: return AgentAddress(protocol_addr=protocol_addr, aid=meta.get("sender_id")) +async def _race_first_completed(*awaitables) -> None: + """Run *awaitables* concurrently; return when the first finishes, + cancelling the rest. Used by :meth:`AgentDelegates.gather` to race + the collector's "we're done" signal against a clock-aware timer. + """ + tasks = [asyncio.ensure_future(a) for a in awaitables] + try: + await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED) + finally: + for task in tasks: + if not task.done(): + task.cancel() + + class _GatherCollector: - """Aggregates multi-reply tracked responses for :meth:`Agent.open_gather`. + """Aggregates multi-reply tracked responses for :meth:`AgentDelegates.open_gather`. Replies are stored under the responding agent's :class:`AgentAddress` so the caller can match each response to its source. :meth:`wait` resolves once *expected* distinct replies have arrived, or when :meth:`finish` is called explicitly (which - :meth:`RoleContext.gather` uses to honour the timeout / quorum + :meth:`AgentDelegates.gather` uses to honour the timeout / quorum policy). """ @@ -300,21 +317,17 @@ def __init__(self) -> None: self._forwarding_rules: list[ForwardingRule] = [] self._transaction_handlers: dict[str, tuple] = {} # Multi-shot reply collectors keyed by tracking_id. Used by - # ``RoleContext.gather`` (and any caller of - # :meth:`Agent.open_gather`) to aggregate replies from many - # receivers under a single id without consuming the entry on - # the first reply. See :meth:`_handle_tracked_reply`. + # :meth:`gather` (and any caller of :meth:`open_gather`) to + # aggregate replies from many receivers under a single id + # without consuming the entry on the first reply. See + # :meth:`_handle_tracked_reply`. self._gather_collectors: dict[str, _GatherCollector] = {} - # Open conversations keyed by conversation_id (see - # :mod:`mango.agent.conversation`). Messages whose meta - # carries a matching id are routed to the conversation's - # async-iterator queue. - self._conversations: dict[str, Conversation] = {} - # Reference count per open conversation id so multiple - # ``join_conversation`` contexts on the same id (e.g. an - # ``@on_message`` handler that re-fires while an earlier join is - # still open) share one handle instead of colliding. - self._conversation_refs: dict[str, int] = {} + # Open conversation handles keyed by conversation_id (see + # :mod:`mango.agent.conversation`). Several handles may be open + # for the same id (e.g. an ``@on_message`` handler that re-fires + # while an earlier join is still open) — every matching inbound + # message is fanned out to all of them. + self._conversations: dict[str, list[Conversation]] = {} self._behavior_message_subs: list[tuple] = [] self._behavior_global_event_handlers: list[tuple] = [] self._behavior_agent_event_handlers: list[tuple] = [] @@ -443,6 +456,17 @@ def delete_forwarding_rule( ) ] + def _bound_agent(self, operation: str) -> Agent: + """Return the agent whose registries and scheduler serve *operation*. + + Identity on an agent itself; delegating contexts (e.g. + ``RoleContext``) override this to return their owning agent — + message routing consults the *agent's* registries, so tracked + handlers, gather collectors and conversations must be + registered there, not on the delegate. + """ + return self + # ------------------------------------------------------------------ # Tracked / reply-to messaging # ------------------------------------------------------------------ @@ -468,7 +492,8 @@ async def send_tracked_message( """ tracking_id = str(uuid.uuid4()) if response_handler is not None: - self._transaction_handlers[tracking_id] = (response_handler,) + agent = self._bound_agent("send_tracked_message") + agent._transaction_handlers[tracking_id] = (response_handler,) return await self.send_message( content, receiver_addr=receiver_addr, @@ -485,7 +510,8 @@ async def reply_to( """Convenience helper to reply to a received message. Extracts the sender address from *received_meta* and sends *content* - back, preserving any ``tracking_id`` for transaction matching. + back, preserving any ``tracking_id`` for transaction matching and any + ``conversation_id`` so the reply routes into an open conversation. :param content: reply content :param received_meta: the ``meta`` dict from the received message @@ -494,10 +520,13 @@ async def reply_to( sender_id = received_meta.get("sender_id") sender_addr = received_meta.get("sender_addr") tracking_id = received_meta.get("tracking_id") + conversation_id = received_meta.get(CONVERSATION_ID_KEY) reply_addr = AgentAddress(protocol_addr=sender_addr, aid=sender_id) extra: dict = {"reply": True} if tracking_id: extra["tracking_id"] = tracking_id + if conversation_id: + extra[CONVERSATION_ID_KEY] = conversation_id extra.update(kwargs) return await self.send_message(content, receiver_addr=reply_addr, **extra) @@ -528,66 +557,35 @@ def _handle_tracked_reply(self, content: Any, meta: dict) -> bool: return False def _route_to_conversation(self, content: Any, meta: dict) -> None: - """Deliver *content*/*meta* to any open conversation whose id - matches ``meta[CONVERSATION_ID_KEY]``. No-op when the message - has no conversation id or no matching conversation exists. + """Deliver *content*/*meta* to every open conversation handle whose + id matches ``meta[CONVERSATION_ID_KEY]``. No-op when the message + has no conversation id or no matching handle exists. """ - from mango.agent.conversation import CONVERSATION_ID_KEY - conv_id = meta.get(CONVERSATION_ID_KEY) if not conv_id: return - conv = self._conversations.get(conv_id) - if conv is None: - return - conv._on_inbound(content, meta) + for conv in list(self._conversations.get(conv_id, ())): + conv._on_inbound(content, meta) - def open_conversation(self, conv: Conversation) -> None: - """Register *conv* as the initiator of a fresh conversation. + def _register_conversation(self, conv: Conversation) -> None: + """Register a conversation handle for inbound routing. - Used internally by ``RoleContext.open_conversation``; end users - should not need to call this directly. Raises if the id is already - open — the initiator generates a unique id, so a collision here is a - genuine programming error. + Called by the conversation context manager on entry; several + handles may share one id — messages fan out to all of them. """ - if conv.conversation_id in self._conversations: - raise ValueError( - f"conversation {conv.conversation_id!r} already open on {self.aid}" - ) - self._conversations[conv.conversation_id] = conv - self._conversation_refs[conv.conversation_id] = 1 - - def join_conversation(self, conv: Conversation) -> Conversation: - """Register a *join* on a possibly-already-open conversation. + self._conversations.setdefault(conv.conversation_id, []).append(conv) - Returns the handle to actually use: the existing one when the id is - already open (incrementing its reference count), otherwise *conv*. - This lets an ``@on_message`` handler that re-fires while an earlier - join is still open share a single handle instead of raising. - """ - cid = conv.conversation_id - existing = self._conversations.get(cid) - if existing is not None: - self._conversation_refs[cid] += 1 - return existing - self._conversations[cid] = conv - self._conversation_refs[cid] = 1 - return conv - - def close_conversation(self, conv: Conversation) -> None: - """Release one reference to *conv*; unregister at zero. - - Called when a conversation context manager exits. - """ - cid = conv.conversation_id - refs = self._conversation_refs.get(cid) - if refs is None: + def _unregister_conversation(self, conv: Conversation) -> None: + """Remove one previously registered handle. Idempotent.""" + handles = self._conversations.get(conv.conversation_id) + if handles is None: return - if refs <= 1: - self._conversations.pop(cid, None) - self._conversation_refs.pop(cid, None) - else: - self._conversation_refs[cid] = refs - 1 + try: + handles.remove(conv) + except ValueError: + pass + if not handles: + self._conversations.pop(conv.conversation_id, None) def _nudge_topology_health(self, meta: dict) -> None: """Multiplicatively recover edge scores on every received message. @@ -633,7 +631,7 @@ def open_gather( Returns an opaque collector handle; the caller awaits :meth:`_GatherCollector.wait` and must call - :meth:`close_gather` (or use :meth:`RoleContext.gather`, which + :meth:`close_gather` (or use :meth:`gather`, which does both) to release the registration. :param expected: number of replies that satisfies the collector @@ -644,17 +642,144 @@ def open_gather( is dropped silently. Filters out tracking-id collisions with unrelated reply traffic. """ - if tracking_id in self._gather_collectors: + collectors = self._bound_agent("open_gather")._gather_collectors + if tracking_id in collectors: raise ValueError( f"gather collector already open for tracking_id={tracking_id!r}" ) collector = _GatherCollector(expected=expected, reply_type=reply_type) - self._gather_collectors[tracking_id] = collector + collectors[tracking_id] = collector return collector def close_gather(self, tracking_id: str) -> None: """Release a previously opened gather collector.""" - self._gather_collectors.pop(tracking_id, None) + self._bound_agent("close_gather")._gather_collectors.pop(tracking_id, None) + + async def gather( + self, + content: Any, + receivers: Iterable[AgentAddress], + *, + reply_type: type | tuple[type, ...] | None = None, + timeout: float = 5.0, + min_fraction: float = 1.0, + ) -> dict[AgentAddress, Any]: + """Send *content* to every address in *receivers* and collect replies. + + Returns a dict mapping each responding agent's + :class:`AgentAddress` to its reply content. Responders are + expected to use :meth:`reply_to` (or otherwise echo + ``tracking_id`` with ``reply=True``) — every existing + request/response pair in mango follows that convention, so + responders work unchanged. + + :param receivers: iterable of :class:`AgentAddress` targets. + :param reply_type: optional class/tuple — replies of any other + type are silently dropped, filtering out tracking-id + collisions with unrelated traffic. + :param timeout: cap on how long to wait, measured on the + agent's scheduler clock. When it elapses, returns + whatever replies have arrived so far. + :param min_fraction: between 0 and 1. Returns as soon as + ``ceil(min_fraction * len(receivers))`` replies have + arrived. Defaults to 1.0 — wait for all, time out + otherwise. + """ + agent = self._bound_agent("gather") + receivers = list(receivers) + if not receivers: + return {} + expected = max(1, math.ceil(min_fraction * len(receivers))) + + tracking_id = str(uuid.uuid4()) + collector = agent.open_gather( + tracking_id, expected=expected, reply_type=reply_type + ) + try: + for addr in receivers: + await self.send_message( + content, receiver_addr=addr, tracking_id=tracking_id + ) + # Timeout follows the agent's scheduler clock so behaviour + # is identical under AsyncioClock (real time) and + # ExternalClock (simulation). asyncio.wait_for would block + # real seconds in sim mode. + await _race_first_completed( + collector.wait(), + agent.scheduler.clock.sleep(timeout), + ) + finally: + agent.close_gather(tracking_id) + return dict(collector.responses) + + def open_conversation( + self, + *, + state: dict | None = None, + timeout: float | None = None, + ) -> _ConversationContext: + """Open a fresh multi-hop conversation as the initiator. + + Returns an async context manager whose body yields a + :class:`~mango.agent.conversation.Conversation`. A new id is + generated; other agents reply by echoing it (every call to + :meth:`Conversation.send` and every :meth:`reply_to` carries it + automatically). + + The optional *timeout* is enforced via the agent's scheduler + clock so behaviour is identical under :class:`AsyncioClock` + (real time) and :class:`ExternalClock` (simulation). Unlike + :meth:`gather`, it defaults to ``None`` — conversations are + long-lived by design, so no arbitrary deadline is imposed. + """ + return self._make_conversation( + conversation_id=str(uuid.uuid4()), state=state, timeout=timeout + ) + + def join_conversation( + self, + meta: dict, + *, + state: dict | None = None, + timeout: float | None = None, + ) -> _ConversationContext: + """Join an existing conversation as a non-initiator. + + Reads the id from *meta* (the dict passed to message handlers) + and registers a local handle so subsequent messages tagged with + that id route here too. Typical use: a gossip forwarder + processes a multi-message exchange from the responder side. + + Every join owns an independent handle with its own *state* and + *timeout*; when several handles for the same id are open on one + agent, each inbound message is delivered to all of them. + """ + conv_id = meta.get(CONVERSATION_ID_KEY) + if not conv_id: + raise ValueError( + "join_conversation requires a meta dict carrying " + f"{CONVERSATION_ID_KEY!r}" + ) + return self._make_conversation( + conversation_id=conv_id, state=state, timeout=timeout + ) + + def _make_conversation( + self, + *, + conversation_id: str, + state: dict | None, + timeout: float | None, + ) -> _ConversationContext: + return _ConversationContext( + owner=self, + conv=Conversation( + owner=self, + conversation_id=conversation_id, + state=state, + timeout=timeout, + ), + ) @property def current_timestamp(self) -> float: @@ -1171,10 +1296,10 @@ async def shutdown(self): # Backstop: cancel any conversations still open (e.g. one opened # with timeout=None whose peer never replied) so their iterators # unblock and the registry does not leak past agent lifetime. - for conv in list(self._conversations.values()): - conv.cancel() + for handles in list(self._conversations.values()): + for conv in list(handles): + conv.cancel() self._conversations.clear() - self._conversation_refs.clear() if not self._stopped.done(): self._stopped.set_result(True) diff --git a/mango/agent/role.py b/mango/agent/role.py index 3fb8710..1cfd4fa 100644 --- a/mango/agent/role.py +++ b/mango/agent/role.py @@ -626,7 +626,7 @@ def _bound_agent(self, operation: str): Conversation- and gather-style helpers need access to the agent's scheduler clock and message-routing tables; this - wrapper produces a uniform error when a role is used before + override produces a uniform error when a role is used before its agent attaches. """ agent = self._role_handler._agent @@ -637,140 +637,6 @@ def _bound_agent(self, operation: str): ) return agent - def open_conversation( - self, - *, - state: dict | None = None, - timeout: float | None = None, - ): - """Open a fresh multi-hop conversation as the initiator. - - Returns an async context manager whose body yields a - :class:`~mango.agent.conversation.Conversation`. A new id is - generated; other agents reply by echoing it (every call to - :meth:`Conversation.send` carries it automatically). - - The optional *timeout* is enforced via the agent's scheduler - clock so behaviour is identical under :class:`AsyncioClock` - (real time) and :class:`ExternalClock` (simulation). - """ - import uuid - - return self._make_conversation( - conversation_id=str(uuid.uuid4()), - state=state, - timeout=timeout, - is_join=False, - ) - - def join_conversation( - self, - meta: dict, - *, - state: dict | None = None, - timeout: float | None = None, - ): - """Join an existing conversation as a non-initiator. - - Reads the id from ``meta`` (the dict passed to ``@on_message`` - handlers) and registers a local handle so subsequent messages - tagged with that id route here too. Typical use: a gossip - forwarder processes a multi-message exchange from the - responder side. - """ - from mango.agent.conversation import CONVERSATION_ID_KEY - - conv_id = meta.get(CONVERSATION_ID_KEY) - if not conv_id: - raise ValueError( - "join_conversation requires a meta dict carrying " - f"{CONVERSATION_ID_KEY!r}" - ) - return self._make_conversation( - conversation_id=conv_id, state=state, timeout=timeout, is_join=True - ) - - def _make_conversation( - self, - *, - conversation_id: str, - state: dict | None, - timeout: float | None, - is_join: bool, - ) -> "_ConversationContext": - from mango.agent.conversation import Conversation - - return _ConversationContext( - role_context=self, - conv=Conversation( - owner=self, - conversation_id=conversation_id, - state=state, - timeout=timeout, - ), - is_join=is_join, - ) - - async def gather( - self, - content: Any, - receivers, - *, - reply_type: type | tuple[type, ...] | None = None, - timeout: float = 5.0, - min_fraction: float = 1.0, - ) -> dict["AgentAddress", Any]: - """Send *content* to every address in *receivers* and collect replies. - - Returns a dict mapping each responding agent's - :class:`AgentAddress` to its reply content. Responders are - expected to use :meth:`AgentDelegates.reply_to` (or otherwise - echo ``tracking_id`` with ``reply=True``) — every existing - request/response pair in mango follows that convention, so - responder roles work unchanged. - - :param receivers: iterable of :class:`AgentAddress` targets. - :param reply_type: optional class/tuple — replies of any other - type are silently dropped, filtering out tracking-id - collisions with unrelated traffic. - :param timeout: cap on how long to wait, measured on the - agent's scheduler clock. When it elapses, returns - whatever replies have arrived so far. - :param min_fraction: between 0 and 1. Returns as soon as - ``ceil(min_fraction * len(receivers))`` replies have - arrived. Defaults to 1.0 — wait for all, time out - otherwise. - """ - import math - import uuid - - agent = self._bound_agent("RoleContext.gather") - receivers = list(receivers) - if not receivers: - return {} - expected = max(1, math.ceil(min_fraction * len(receivers))) - - tracking_id = str(uuid.uuid4()) - collector = agent.open_gather( - tracking_id, expected=expected, reply_type=reply_type - ) - try: - for addr in receivers: - await self.send_message( - content, receiver_addr=addr, tracking_id=tracking_id - ) - # Timeout follows the agent's scheduler clock so behaviour - # is identical under AsyncioClock (real time) and - # ExternalClock (simulation). asyncio.wait_for would block - # real seconds in sim mode. - await _race_first_completed( - collector.wait(), - agent.scheduler.clock.sleep(timeout), - ) - finally: - agent.close_gather(tracking_id) - return dict(collector.responses) - def on_start(self): self._role_handler.on_start() @@ -926,83 +792,3 @@ def on_agent_event(self, event: Any) -> None: :param event: the event object """ - - -async def _race_first_completed(*awaitables) -> None: - """Run *awaitables* concurrently; return when the first finishes, - cancelling the rest. Used by gather and the conversation timeout - to race a "we're done" signal against a clock-aware timer. - """ - tasks = [asyncio.ensure_future(a) for a in awaitables] - try: - await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED) - finally: - for task in tasks: - if not task.done(): - task.cancel() - - -class _ConversationContext: - """Async context manager returned by ``RoleContext.open_conversation`` - and ``RoleContext.join_conversation``. - - Owns the lifecycle of one :class:`~mango.agent.conversation.Conversation`: - - * ``__aenter__`` registers the conversation with the agent so - inbound messages route to it, and schedules the optional - clock-aware timeout. - * ``__aexit__`` unregisters the conversation and cancels the - timeout future. - - Keeping the context manager separate from :class:`Conversation` - itself leaves the data class free of mango-agent references and - easy to unit-test. - """ - - def __init__( - self, *, role_context: "RoleContext", conv, is_join: bool = False - ) -> None: - self._role_context = role_context - self._conv = conv - self._is_join = is_join - self._timeout_task: asyncio.Task | None = None - - async def __aenter__(self): - agent = self._role_context._bound_agent("Conversation") - reused = False - if self._is_join: - # Reuse an already-open handle (shared refcount) so a re-firing - # handler does not collide on the same id. - registered = agent.join_conversation(self._conv) - reused = registered is not self._conv - self._conv = registered - else: - agent.open_conversation(self._conv) - # Schedule the timeout on the agent's scheduler clock so it - # respects simulation time under an ExternalClock. A join that - # merely reused an existing handle does not arm its own timeout — - # the original owner's timeout governs the shared conversation. - timeout = self._conv._timeout - clock = agent.scheduler.clock if agent.scheduler else None - if not reused and timeout is not None and clock is not None: - self._timeout_task = asyncio.ensure_future( - self._cancel_after(clock, timeout) - ) - return self._conv - - async def __aexit__(self, exc_type, exc, tb): - if self._timeout_task is not None and not self._timeout_task.done(): - self._timeout_task.cancel() - agent = self._role_context._role_handler._agent - if agent is not None: - agent.close_conversation(self._conv) - return False - - async def _cancel_after(self, clock, delay: float) -> None: - try: - await clock.sleep(delay) - except asyncio.CancelledError: - return - # Timeout == abrupt close: anything still queued is dropped, - # but conv.state remains readable after the context exits. - self._conv.cancel() diff --git a/mango/container/core.py b/mango/container/core.py index 6fd00e5..863e1be 100644 --- a/mango/container/core.py +++ b/mango/container/core.py @@ -112,13 +112,14 @@ def _reserve_aid(self, suggested_aid=None): self._aid_counter += 1 return aid - def register(self, agent: Agent, suggested_aid: str = None): + def register(self, agent: Agent, suggested_aid: str = None, **kwargs): """ Register *agent* and return the agent id :param agent: The agent instance :param suggested_aid: (Optional) suggested aid, if the aid is already taken, a generated aid is used. Using the generated aid-style ("agentX") is not allowed. + :param kwargs: additional keyword arguments passed to :meth:`on_register` :return The agent ID """ @@ -127,6 +128,7 @@ def register(self, agent: Agent, suggested_aid: str = None): raise ValueError("Agent is already registered to a container") self._agents[aid] = agent agent._do_register(self, aid) + self.on_register(agent, aid, **kwargs) logger.debug("Successfully registered agent;%s", aid) if self.running: agent._do_start() @@ -135,6 +137,9 @@ def register(self, agent: Agent, suggested_aid: str = None): agent.on_ready() return agent + def on_register(self, agent: Agent, aid: str, **kwargs) -> None: + """Hook called after *agent* is registered, before it is started.""" + def _get_aid(self, agent): for aid, a in self._agents.items(): if id(a) == id(agent): diff --git a/mango/express/api.py b/mango/express/api.py index 60cef52..cdccbea 100644 --- a/mango/express/api.py +++ b/mango/express/api.py @@ -373,7 +373,7 @@ def _matches_agent(agent): return True return False - for agent in list(world._agents.values()): + for agent in list(world.agents.values()): if not _matches_agent(agent): continue diff --git a/mango/simulation/__init__.py b/mango/simulation/__init__.py index 759c74d..00443da 100644 --- a/mango/simulation/__init__.py +++ b/mango/simulation/__init__.py @@ -25,21 +25,26 @@ plot_world, show_communication_data, ) -from .world import ( - AgentsRecording, - DISCRETE_EVENT, +from .container import ( MessageTransaction, - SimulationResult, - SimulationWorld, + SimulationContainer, +) +from .recording import ( + AgentsRecording, WorldRecording, collect_agent_data, collect_data, - create_world, - discrete_step_until, position_history, record_agent, record_agent_having, record_position, record_world, +) +from .world import ( + DISCRETE_EVENT, + SimulationResult, + SimulationWorld, + create_world, + discrete_step_until, step_simulation, ) diff --git a/mango/simulation/container.py b/mango/simulation/container.py new file mode 100644 index 0000000..75e4119 --- /dev/null +++ b/mango/simulation/container.py @@ -0,0 +1,196 @@ +""" +SimulationContainer – the container implementation backing +:class:`~mango.simulation.world.SimulationWorld`. + +Implements the mango container contract for a clock-driven simulation: +``send_message`` queues messages with a delivery time computed by the +communication simulation; the world delivers them while stepping. +""" + +import bisect +import logging +from collections.abc import Callable +from dataclasses import dataclass +from typing import Any + +from ..agent.core import Agent, AgentAddress +from ..container.core import Container +from ..messages.codecs import JSON +from ..util.clock import ExternalClock +from .communication import CommunicationSimulation, MessagePackage + +logger = logging.getLogger(__name__) + +SIMULATION_ADDR = "simulation" + + +@dataclass +class MessageTransaction: + """Records a message that was delivered during the simulation.""" + + sender_id: str | None + receiver_id: str + sent_time: float + arriving_time: float + content: Any + + +class SimulationContainer(Container): + """Local container for agents living in a :class:`SimulationWorld`. + + Unlike networked containers there is no explicit ``start()`` phase: + the container is running from construction on, and messages are held + in a pending queue until the world steps the clock past their + delivery time. + """ + + def __init__( + self, + clock: ExternalClock, + communication_sim: CommunicationSimulation, + on_agent_registered: Callable[..., None] | None = None, + ): + super().__init__( + addr=SIMULATION_ADDR, + name=SIMULATION_ADDR, + codec=JSON(), + clock=clock, + ) + self.communication_sim = communication_sim + self._on_agent_registered = on_agent_registered + self.running = True + + # Pending message queue: sorted list of (delivery_time, seq, sent_time, content, meta) + # seq is a monotonically increasing tie-breaker so bisect.insort never + # needs to compare content or meta (which may not be orderable). + self._pending_messages: list[tuple[float, int, float, Any, dict]] = [] + self._msg_seq: int = 0 + + self.recorded_messages: list[MessageTransaction] = [] + + def on_register(self, agent: Agent, aid: str, **kwargs) -> None: + if self._on_agent_registered is not None: + self._on_agent_registered(agent, aid, **kwargs) + + def deregister(self, aid: str) -> None: + self._agents.pop(aid, None) + + async def send_message( + self, + content: Any, + receiver_addr: AgentAddress, + sender_id: str | None = None, + **kwargs, + ) -> bool: + """Send a message, applying communication simulation. + + Messages are queued with a delivery time determined by the + communication simulation. They are delivered during the next + simulation step. + """ + meta: dict[str, Any] = { + "sender_id": sender_id, + "sender_addr": self.addr, + "receiver_id": receiver_addr.aid, + "receiver_addr": self.addr, + "network_protocol": "simulation", + } + meta.update(kwargs) + + sent_time = self.clock.time + package = MessagePackage( + sender_id=sender_id, + receiver_id=receiver_addr.aid, + sent_time=sent_time, + content=(content, meta), + ) + result = self.communication_sim.calculate_communication( + current_time=sent_time, + messages=[package], + ).package_results[0] + + if not result.reached: + logger.debug( + "Message from %s to %s dropped (loss simulation)", + sender_id, + receiver_addr.aid, + ) + return False + + delivery_time = sent_time + result.delay_s + seq = self._msg_seq + self._msg_seq += 1 + bisect.insort( + self._pending_messages, + (delivery_time, seq, sent_time, content, meta), + ) + return True + + async def _deliver_messages_due(self, up_to_time: float) -> int: + """Deliver all pending messages with delivery_time <= up_to_time. + + Returns the number of messages delivered. + """ + delivered = 0 + remaining: list[tuple[float, int, float, Any, dict]] = [] + + for delivery_time, seq, sent_time, content, meta in self._pending_messages: + if delivery_time <= up_to_time: + receiver_id = meta.get("receiver_id") + agent = self._agents.get(receiver_id) + if agent is not None: + await agent.inbox.put((0, content, meta)) + self.recorded_messages.append( + MessageTransaction( + sender_id=meta.get("sender_id"), + receiver_id=receiver_id, + sent_time=sent_time, + arriving_time=delivery_time, + content=content, + ) + ) + delivered += 1 + else: + logger.warning( + "Unknown receiver '%s'; dropping message", receiver_id + ) + else: + remaining.append((delivery_time, seq, sent_time, content, meta)) + + self._pending_messages = remaining + return delivered + + def _determine_next_step_size(self) -> float | None: + """Return seconds to the next scheduled event, or None if none.""" + candidates: list[float] = [] + + if self._pending_messages: + next_msg_arrival = self._pending_messages[0][0] + candidates.append(max(0.0, next_msg_arrival - self.clock.time)) + + next_task = self.clock.get_next_activity() + if next_task is not None: + candidates.append(max(0.0, next_task - self.clock.time)) + + if not candidates: + return None + return min(candidates) + + async def as_agent_process(self, agent_creator, mirror_container_creator=None): + raise NotImplementedError( + "Agent subprocesses are not supported in a simulation container" + ) + + def as_agent_process_lazy(self, agent_creator, mirror_container_creator=None): + raise NotImplementedError( + "Agent subprocesses are not supported in a simulation container" + ) + + async def shutdown(self) -> None: + """Shut down all agents.""" + self.running = False + for agent in list(self._agents.values()): + try: + await agent.shutdown() + except Exception: + logger.exception("Error shutting down agent '%s'", agent.aid) diff --git a/mango/simulation/environment.py b/mango/simulation/environment.py index 12abdb6..c9556b4 100644 --- a/mango/simulation/environment.py +++ b/mango/simulation/environment.py @@ -131,7 +131,7 @@ def agents_within(self, center, radius: float, agents: list) -> list: Example:: - nearby = space.agents_within(my_agent, 5.0, world._agents.values()) + nearby = space.agents_within(my_agent, 5.0, world.agents.values()) """ result = [] for agent in agents: diff --git a/mango/simulation/recording.py b/mango/simulation/recording.py new file mode 100644 index 0000000..c952394 --- /dev/null +++ b/mango/simulation/recording.py @@ -0,0 +1,235 @@ +""" +Recording utilities for :class:`~mango.simulation.world.SimulationWorld`. + +Recorders register a collector on the world that is invoked after every +simulation step; the collected values are stored in +:class:`WorldRecording` / :class:`AgentsRecording` instances accessible +via ``world.data_collections`` and ``world.data_agent_collections``. +""" + +from collections.abc import Callable +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any + +from ..agent.core import Agent + +if TYPE_CHECKING: + from .world import SimulationWorld + + +@dataclass +class WorldRecording: + """Time-series recording of world-level data.""" + + timeseries: list[Any] = field(default_factory=list) + time: list[float] = field(default_factory=list) + + +@dataclass +class AgentsRecording: + """Per-agent time-series recording. + + ``timeseries`` maps each agent AID to a list of recorded values. + ``agent_time`` maps each agent AID to the elapsed simulation seconds at + which each of that agent's values was recorded; it stays aligned with + ``timeseries`` even when agents are recorded sparsely (e.g. registered + mid-simulation or gated by a filter). ``time`` holds every step's + timestamp as a shared axis for agents recorded on every step. + """ + + timeseries: dict[str, list[Any]] = field(default_factory=dict) + agent_time: dict[str, list[float]] = field(default_factory=dict) + time: list[float] = field(default_factory=list) + + +def collect_data( + world: "SimulationWorld", + key: str, + collector: Callable[["SimulationWorld", WorldRecording], None], +) -> None: + """Register a world-level data collector. + + *collector* is called after every simulation step with the world and the + :class:`WorldRecording` identified by *key*. + + Example:: + + collect_data(world, "total_msgs", lambda w, rec: ( + rec.timeseries.append(len(w.recorded_messages)), + rec.time.append(w.clock.time), + )) + """ + recording = world.new_world_recording(key) + + def _run(w: "SimulationWorld") -> None: + collector(w, recording) + + world.add_data_collector(_run) + + +def collect_agent_data( + world: "SimulationWorld", + key: str, + collector: Callable[["SimulationWorld", Agent, AgentsRecording], None], +) -> None: + """Register an agent-level data collector. + + *collector* is called for every agent after every simulation step with + the world, the agent, and the :class:`AgentsRecording` for *key*. A + shared ``time`` entry is appended once per step. + + Example:: + + collect_agent_data(world, "state", lambda w, a, rec: ( + rec.timeseries.setdefault(a.aid, []).append(a.some_state), + )) + """ + recording = world.new_agents_recording(key) + + def _run(w: "SimulationWorld") -> None: + for agent in w.agents.values(): + collector(w, agent, recording) + recording.time.append(w.clock.time) + + world.add_data_collector(_run) + + +def record_world( + world: "SimulationWorld", + key: str, + recorder: Callable[[], Any], +) -> None: + """Record a world-level scalar after every step. + + *recorder* is a zero-argument callable whose return value is appended to + the recording's ``timeseries``. + + Example:: + + record_world(world, "agent_count", lambda: len(world.agents)) + """ + recording = world.new_world_recording(key) + + def _run(w: "SimulationWorld") -> None: + recording.timeseries.append(recorder()) + recording.time.append(w.clock.time) + + world.add_data_collector(_run) + + +def record_agent( + world: "SimulationWorld", + key: str, + recorder: Callable[[Agent], Any], + filter_fn: Callable[[Agent], bool] | None = None, +) -> None: + """Record a per-agent scalar after every step. + + *recorder* receives each agent and returns the value to store. An + optional *filter_fn* restricts recording to a subset of agents — pass + an ``isinstance``-based predicate to record only agents of a particular + type:: + + record_agent(world, "soc", lambda a: a.soc_kwh, + filter_fn=lambda a: isinstance(a, EVAgent)) + + :param world: the simulation world + :param key: recording key + :param recorder: ``(agent) -> value`` callable + :param filter_fn: optional ``(agent) -> bool`` predicate; ``None`` records + all registered agents + """ + recording = world.new_agents_recording(key) + + def _run(w: "SimulationWorld") -> None: + for agent in w.agents.values(): + if filter_fn is None or filter_fn(agent): + recording.timeseries.setdefault(agent.aid, []).append(recorder(agent)) + recording.agent_time.setdefault(agent.aid, []).append(w.clock.time) + recording.time.append(w.clock.time) + + world.add_data_collector(_run) + + +def record_agent_having( + world: "SimulationWorld", + key: str, + role_type: type, + recorder: Callable[[Agent], Any], +) -> None: + """Record a per-agent scalar for agents that carry a specific role type. + + Only agents that have at least one role that is an instance of + *role_type* are included in the recording. *recorder* receives the + agent and returns the value to store. + + :param world: the simulation world + :param key: recording key + :param role_type: only record agents that have a role of this type + :param recorder: ``(agent) -> value`` callable + + Example:: + + record_agent_having(world, "energy", EnergyRole, lambda a: a.roles[0].energy) + """ + recording = world.new_agents_recording(key) + + def _run(w: "SimulationWorld") -> None: + for agent in w.agents.values(): + if hasattr(agent, "roles") and any( + isinstance(r, role_type) for r in agent.roles + ): + recording.timeseries.setdefault(agent.aid, []).append(recorder(agent)) + recording.agent_time.setdefault(agent.aid, []).append(w.clock.time) + recording.time.append(w.clock.time) + + world.add_data_collector(_run) + + +def record_position( + world: "SimulationWorld", + key: str = "positions", + filter_fn: Callable[[Agent], bool] | None = None, +) -> None: + """Record the spatial position of every agent after each step. + + Only agents that have a position in the world's space are recorded. + An optional *filter_fn* restricts recording to a subset of agents. + + :param world: the simulation world + :param key: recording key (default ``"positions"``) + :param filter_fn: ``(agent) -> bool`` predicate; ``None`` means all agents + + Example:: + + record_position(world) + history = position_history(world) + # history.timeseries["agent0"] -> list of Position2D + """ + recording = world.new_agents_recording(key) + + def _run(w: "SimulationWorld") -> None: + space = w.environment.space + for agent in w.agents.values(): + if space.has_position(agent): + if filter_fn is None or filter_fn(agent): + recording.timeseries.setdefault(agent.aid, []).append( + space.location(agent) + ) + recording.agent_time.setdefault(agent.aid, []).append(w.clock.time) + recording.time.append(w.clock.time) + + world.add_data_collector(_run) + + +def position_history( + world: "SimulationWorld", + key: str = "positions", +) -> AgentsRecording: + """Return the :class:`AgentsRecording` populated by :func:`record_position`. + + :param world: the simulation world + :param key: recording key (default ``"positions"``) + :return: the recording + """ + return world.data_agent_collections.get(key, AgentsRecording()) diff --git a/mango/simulation/visualization.py b/mango/simulation/visualization.py index 8d83785..1c11a7e 100644 --- a/mango/simulation/visualization.py +++ b/mango/simulation/visualization.py @@ -341,7 +341,7 @@ def show_communication_data( def _agent_label(world: SimulationWorld, aid: str) -> str: """Return display label for an agent: name if set, else AID.""" - agent = world._agents.get(aid) + agent = world.agents.get(aid) if agent is not None and agent.name: return f"{agent.name} ({aid})" return aid diff --git a/mango/simulation/world.py b/mango/simulation/world.py index 23e4b55..c887ec1 100644 --- a/mango/simulation/world.py +++ b/mango/simulation/world.py @@ -1,10 +1,15 @@ """ -SimulationWorld – a self-contained simulation container for mango. +SimulationWorld – a self-contained simulation world for mango. Mirrors the ``World`` type from Mango.jl. Agents registered in a SimulationWorld share an :class:`~mango.util.clock.ExternalClock` and can be stepped forward in discrete or fixed-size time increments. +The world is a facade: agent registration and message transport are +handled by a :class:`~mango.simulation.container.SimulationContainer` +subcomponent, while the world drives the simulation loop, the +environment, and the recording infrastructure. + Typical usage:: async def run(): @@ -12,8 +17,8 @@ async def run(): agent = world.register(MyAgent()) async with world: - await step_simulation(world, step_size_s=1.0) - await step_simulation(world, step_size_s=1.0) + await world.step(step_size_s=1.0) + await world.step(step_size_s=1.0) asyncio.run(run()) @@ -23,74 +28,70 @@ async def run(): world = create_world(start_time=0.0) agent = world.register(MyAgent()) async with world: - await discrete_step_until(world, max_advance_time_s=60.0) + await world.step_until(max_advance_time_s=60.0) asyncio.run(run()) + +The module-level :func:`step_simulation` and :func:`discrete_step_until` +are thin aliases for :meth:`SimulationWorld.step` and +:meth:`SimulationWorld.step_until`. """ import asyncio -import bisect import logging from collections.abc import Callable -from dataclasses import dataclass, field +from dataclasses import dataclass from typing import Any from mango.agent.core import Agent, AgentAddress -from mango.messages.codecs import JSON +from mango.messages.codecs import Codec from mango.util.clock import ExternalClock -from .communication import ( - CommunicationSimulation, - MessagePackage, - SimpleCommunicationSimulation, -) +from .communication import CommunicationSimulation, SimpleCommunicationSimulation +from .container import MessageTransaction, SimulationContainer from .environment import DefaultEnvironment, Environment, WorldObserver +from .recording import ( + AgentsRecording, + WorldRecording, + collect_agent_data, + collect_data, + position_history, + record_agent, + record_agent_having, + record_position, + record_world, +) logger = logging.getLogger(__name__) DISCRETE_EVENT: float = -1.0 AGENT_PREFIX: str = "agent" - -@dataclass -class WorldRecording: - """Time-series recording of world-level data.""" - - timeseries: list[Any] = field(default_factory=list) - time: list[float] = field(default_factory=list) - - -@dataclass -class AgentsRecording: - """Per-agent time-series recording. - - ``timeseries`` maps each agent AID to a list of recorded values. - ``agent_time`` maps each agent AID to the elapsed simulation seconds at - which each of that agent's values was recorded; it stays aligned with - ``timeseries`` even when agents are recorded sparsely (e.g. registered - mid-simulation or gated by a filter). ``time`` holds every step's - timestamp as a shared axis for agents recorded on every step. - """ - - timeseries: dict[str, list[Any]] = field(default_factory=dict) - agent_time: dict[str, list[float]] = field(default_factory=dict) - time: list[float] = field(default_factory=list) - - -@dataclass -class MessageTransaction: - """Records a message that was delivered during the simulation.""" - - sender_id: str | None - receiver_id: str - sent_time: float - arriving_time: float - content: Any +__all__ = [ + "AGENT_PREFIX", + "AgentsRecording", + "DISCRETE_EVENT", + "MessageTransaction", + "SimulationContainer", + "SimulationResult", + "SimulationWorld", + "WorldRecording", + "collect_agent_data", + "collect_data", + "create_world", + "discrete_step_until", + "position_history", + "record_agent", + "record_agent_having", + "record_position", + "record_world", + "step_simulation", +] @dataclass class SimulationResult: - """Return value of :func:`step_simulation`.""" + """Return value of :meth:`SimulationWorld.step`.""" time_elapsed_s: float step_size_s: float @@ -104,7 +105,7 @@ def __init__(self, world: "SimulationWorld"): self._world = world def dispatch_global_event(self, clock, event: Any) -> None: - for agent in self._world._agents.values(): + for agent in self._world.agents.values(): for _cond, _handler in agent._behavior_global_event_handlers: if _cond(event): _handler(agent, event) @@ -118,14 +119,16 @@ def dispatch_global_event(self, clock, event: Any) -> None: class SimulationWorld: - """A local, clock-driven simulation container. + """A local, clock-driven simulation world. Do not instantiate directly; use :func:`create_world` instead. - The world acts as the *container* that agents register against. It - satisfies the minimal container interface expected by mango's - :func:`~mango.util.termination_detection.tasks_complete_or_sleeping` - helper (``inbox``, ``_agents``). + The world is a facade over a + :class:`~mango.simulation.container.SimulationContainer` (which agents + register against and which handles message transport), the + :class:`~mango.simulation.environment.Environment`, and the recording + infrastructure. Stepping the simulation is done via :meth:`step` and + :meth:`step_until`. """ def __init__( @@ -134,88 +137,94 @@ def __init__( communication_sim: CommunicationSimulation, environment: Environment | None = None, ): - self.clock: ExternalClock = clock - self.communication_sim: CommunicationSimulation = communication_sim self.environment: Environment = environment or DefaultEnvironment() - - # Container-like state - self._agents: dict[str, Agent] = {} - self._aid_counter: int = 0 - self.addr: str = "simulation" - self.running: bool = True - self.ready: bool = False - self.inbox: asyncio.Queue | None = ( - None # not used but expected by termination util + self._container = SimulationContainer( + clock=clock, + communication_sim=communication_sim, + on_agent_registered=self._install_in_environment, ) - # codec is needed by agents that call container.codec; provide a default - self.codec = JSON() - - # Pending message queue: sorted list of (delivery_time, seq, sent_time, content, meta) - # seq is a monotonically increasing tie-breaker so bisect.insort never - # needs to compare content or meta (which may not be orderable). - self._pending_messages: list[tuple[float, int, float, Any, dict]] = [] - self._msg_seq: int = 0 - # Recording infrastructure self.data_collections: dict[str, WorldRecording] = {} self.data_agent_collections: dict[str, AgentsRecording] = {} self._data_collectors: list[Callable] = [] - # Transaction log - self.recorded_messages: list[MessageTransaction] = [] - - # Internal state self._initialized: bool = False # Wire environment to agent dispatcher - observer = _AgentDispatchObserver(self) - self.environment.add_observer(observer) + self.environment.add_observer(_AgentDispatchObserver(self)) + + # ------------------------------------------------------------------ + # Container facade + # ------------------------------------------------------------------ + + @property + def container(self) -> SimulationContainer: + return self._container + + @property + def clock(self) -> ExternalClock: + return self._container.clock + + @property + def communication_sim(self) -> CommunicationSimulation: + return self._container.communication_sim + + @property + def addr(self) -> str: + return self._container.addr @property def name(self) -> str: - return self.addr + return self._container.name + + @property + def codec(self) -> Codec: + return self._container.codec + + @property + def agents(self) -> dict[str, Agent]: + return self._container._agents + + @property + def running(self) -> bool: + return self._container.running + + @running.setter + def running(self, value: bool) -> None: + self._container.running = value + + @property + def ready(self) -> bool: + return self._container.ready + + @property + def recorded_messages(self) -> list[MessageTransaction]: + return self._container.recorded_messages - def register(self, agent: Agent, suggested_aid: str | None = None) -> Agent: + @recorded_messages.setter + def recorded_messages(self, value: list[MessageTransaction]) -> None: + self._container.recorded_messages = value + + def register( + self, agent: Agent, suggested_aid: str | None = None, **install_kwargs + ) -> Agent: """Register *agent* with the world and return it. :param agent: agent instance to register :param suggested_aid: optional preferred agent ID + :param install_kwargs: forwarded to the environment's ``install`` hook :return: the registered agent (same object) """ - aid = self._reserve_aid(suggested_aid) - if agent.context._container: - raise ValueError("Agent is already registered to a container") - self._agents[aid] = agent - agent._do_register(self, aid) - - install = getattr(self.environment, "install", None) - if callable(install): - install(agent, agent_id=aid) - logger.debug("Registered agent '%s' with world", aid) - if self.running: - agent._do_start() - if self.ready: - agent.on_ready() - return agent + return self._container.register( + agent, suggested_aid=suggested_aid, **install_kwargs + ) def deregister(self, aid: str) -> None: - self._agents.pop(aid, None) + self._container.deregister(aid) def is_aid_available(self, aid: str) -> bool: - pattern_clash = ( - aid.startswith(AGENT_PREFIX) and aid[len(AGENT_PREFIX) :].isnumeric() - ) - return aid not in self._agents and not pattern_clash - - def _reserve_aid(self, suggested_aid: str | None = None) -> str: - if suggested_aid is not None and self.is_aid_available(suggested_aid): - return suggested_aid - while True: - aid = f"{AGENT_PREFIX}{self._aid_counter}" - self._aid_counter += 1 - if aid not in self._agents: - return aid + return self._container.is_aid_available(aid) async def send_message( self, @@ -228,119 +237,31 @@ async def send_message( Messages are queued with a delivery time determined by the communication simulation. They are delivered during the next call - to :func:`step_simulation`. + to :meth:`step`. """ - meta: dict[str, Any] = { - "sender_id": sender_id, - "sender_addr": self.addr, - "receiver_id": receiver_addr.aid, - "receiver_addr": self.addr, - "network_protocol": "simulation", - } - meta.update(kwargs) - - sent_time = self.clock.time - package = MessagePackage( - sender_id=sender_id, - receiver_id=receiver_addr.aid, - sent_time=sent_time, - content=(content, meta), - ) - result = self.communication_sim.calculate_communication( - current_time=sent_time, - messages=[package], - ).package_results[0] - - if not result.reached: - logger.debug( - "Message from %s to %s dropped (loss simulation)", - sender_id, - receiver_addr.aid, - ) - return False - - delivery_time = sent_time + result.delay_s - # Keep list sorted by delivery_time; seq breaks ties without comparing - # content or meta (which may not be orderable). - seq = self._msg_seq - self._msg_seq += 1 - bisect.insort( - self._pending_messages, - (delivery_time, seq, sent_time, content, meta), + return await self._container.send_message( + content, receiver_addr=receiver_addr, sender_id=sender_id, **kwargs ) - return True - def _start_agents(self) -> None: - """Ensure all agents have been started (internal use).""" - # Agents are started when registered, but we call on_ready here - self.ready = True - for agent in self._agents.values(): - agent.on_ready() + async def shutdown(self) -> None: + """Shut down all agents.""" + await self._container.shutdown() - def _initialize_if_needed(self) -> None: - if not self._initialized: - self._start_agents() # call on_ready on all agents - self.environment.initialize(list(self._agents.values()), self.clock) - self._initialized = True - self._do_recordings() + def __getitem__(self, key: str | int) -> Agent: + if isinstance(key, int): + return list(self.agents.values())[key] + return self.agents[key] - async def _deliver_messages_due(self, up_to_time: float) -> int: - """Deliver all pending messages with delivery_time <= up_to_time. + def _install_in_environment(self, agent: Agent, aid: str, **install_kwargs) -> None: + install = getattr(self.environment, "install", None) + if callable(install): + install(agent, agent_id=aid, **install_kwargs) - Returns the number of messages delivered. - """ - delivered = 0 - remaining: list[tuple[float, int, float, Any, dict]] = [] - - for delivery_time, seq, sent_time, content, meta in self._pending_messages: - if delivery_time <= up_to_time: - receiver_id = meta.get("receiver_id") - agent = self._agents.get(receiver_id) - if agent is not None: - await agent.inbox.put((0, content, meta)) - self.recorded_messages.append( - MessageTransaction( - sender_id=meta.get("sender_id"), - receiver_id=receiver_id, - sent_time=sent_time, - arriving_time=delivery_time, - content=content, - ) - ) - delivered += 1 - else: - logger.warning( - "Unknown receiver '%s'; dropping message", receiver_id - ) - else: - remaining.append((delivery_time, seq, sent_time, content, meta)) - - self._pending_messages = remaining - return delivered - - def _determine_next_step_size(self) -> float | None: - """Return seconds to the next scheduled event, or None if none.""" - candidates: list[float] = [] - - # Next message arrival - if self._pending_messages: - next_msg_arrival = self._pending_messages[0][0] - candidates.append(max(0.0, next_msg_arrival - self.clock.time)) - - # Next task wakeup - next_task = self.clock.get_next_activity() - if next_task is not None: - candidates.append(max(0.0, next_task - self.clock.time)) - - if not candidates: - return None - return min(candidates) + # ------------------------------------------------------------------ + # Recording + # ------------------------------------------------------------------ - def _do_recordings(self) -> None: - for collector in self._data_collectors: - collector(self) - - def _new_world_recording(self, key: str) -> WorldRecording: + def new_world_recording(self, key: str) -> WorldRecording: if key in self.data_collections: raise ValueError( f"A recording is already registered for key '{key}'; " @@ -350,7 +271,7 @@ def _new_world_recording(self, key: str) -> WorldRecording: self.data_collections[key] = recording return recording - def _new_agents_recording(self, key: str) -> AgentsRecording: + def new_agents_recording(self, key: str) -> AgentsRecording: if key in self.data_agent_collections: raise ValueError( f"A recording is already registered for key '{key}'; " @@ -360,30 +281,150 @@ def _new_agents_recording(self, key: str) -> AgentsRecording: self.data_agent_collections[key] = recording return recording + def add_data_collector( + self, collector: Callable[["SimulationWorld"], None] + ) -> None: + self._data_collectors.append(collector) + + def _do_recordings(self) -> None: + for collector in self._data_collectors: + collector(self) + + # ------------------------------------------------------------------ + # Simulation loop + # ------------------------------------------------------------------ + + def _initialize_if_needed(self) -> None: + if not self._initialized: + self._container.on_ready() + self.environment.initialize(list(self.agents.values()), self.clock) + self._initialized = True + self._do_recordings() + + async def _wait_for_agents(self) -> None: + """Wait until all agent tasks are complete or sleeping.""" + from mango.util.termination_detection import tasks_complete_or_sleeping + + await tasks_complete_or_sleeping(self._container) + + async def step( + self, + step_size_s: float = DISCRETE_EVENT, + max_advance_time_s: float = -1.0, + ) -> SimulationResult | None: + """Advance the simulation by *step_size_s* seconds. + + When *step_size_s* is ``DISCRETE_EVENT`` (the default), the step size + is determined automatically as the time until the next scheduled + event (message arrival or agent task wakeup). + + :param step_size_s: step size in seconds, or ``DISCRETE_EVENT`` + :param max_advance_time_s: abort if the determined discrete step would + exceed this value; ``-1`` means no limit + :return: :class:`SimulationResult`, or ``None`` if there is nothing to do + + Example:: + + result = await world.step(step_size_s=1.0) + result = await world.step() # discrete-event + """ + self._initialize_if_needed() + + # Allow freshly scheduled tasks (e.g. from on_ready / on_start) to reach + # their first sleeping or done point. This is necessary so that periodic + # tasks register their first wakeup with the ExternalClock before we + # determine the discrete step size. + await self._wait_for_agents() + + actual_step = step_size_s + + if actual_step == DISCRETE_EVENT: + actual_step = self._container._determine_next_step_size() + if actual_step is None: + return None + if max_advance_time_s >= 0 and actual_step > max_advance_time_s: + return None + + new_time = self.clock.time + actual_step + + # ---- call on_step hooks ---- + for agent in list(self.agents.values()): + agent.on_step(self.environment, self.clock, actual_step) + if hasattr(agent, "roles"): + for role in agent.roles: + role.on_step(self.environment, self.clock, actual_step) + + # ---- step environment ---- + self.environment.step(self.clock, actual_step) + + # ---- advance clock (wakes up sleeping tasks) ---- + self.clock.set_time(new_time) + + # ---- convergence loop: process messages and tasks ---- + total_delivered = 0 + state_changed = True + while state_changed: + # Wait for agents to settle + await self._wait_for_agents() + + # Deliver messages whose delivery_time has now arrived + delivered = await self._container._deliver_messages_due(new_time) + total_delivered += delivered + state_changed = delivered > 0 + + # Yield control so delivered messages are processed + if state_changed: + await asyncio.sleep(0) + + self._do_recordings() + + return SimulationResult( + time_elapsed_s=actual_step, + step_size_s=actual_step, + messages_delivered=total_delivered, + ) + + async def step_until(self, max_advance_time_s: float) -> list[SimulationResult]: + """Run a discrete-event simulation until *max_advance_time_s* has elapsed. + + The simulation stops when the world clock has advanced by + *max_advance_time_s* from its current position, or when there are no + more events to process. + + :param max_advance_time_s: maximum total time to simulate in seconds + :return: list of :class:`SimulationResult` from each step + + Example:: + + results = await world.step_until(max_advance_time_s=3600.0) + """ + self._initialize_if_needed() + start_time = self.clock.time + max_time = start_time + max_advance_time_s + results: list[SimulationResult] = [] + + while self.clock.time < max_time: + remaining = max_time - self.clock.time + result = await self.step( + step_size_s=DISCRETE_EVENT, + max_advance_time_s=remaining, + ) + if result is None: + break + results.append(result) + + return results + async def __aenter__(self) -> "SimulationWorld": self._initialize_if_needed() # Allow all freshly scheduled tasks (e.g. from on_ready) to start and # reach their first sleeping/done point before the caller proceeds. - await _wait_for_agents(self) + await self._wait_for_agents() return self async def __aexit__(self, *_) -> None: await self.shutdown() - async def shutdown(self) -> None: - """Shut down all agents.""" - self.running = False - for agent in list(self._agents.values()): - try: - await agent.shutdown() - except Exception: - logger.exception("Error shutting down agent '%s'", agent.aid) - - def __getitem__(self, key: str | int) -> Agent: - if isinstance(key, int): - return list(self._agents.values())[key] - return self._agents[key] - def create_world( start_time: float = 0.0, @@ -412,89 +453,14 @@ def create_world( return SimulationWorld(clock=clock, communication_sim=sim, environment=environment) -async def _wait_for_agents(world: SimulationWorld) -> None: - """Wait until all agent tasks are complete or sleeping.""" - from mango.util.termination_detection import tasks_complete_or_sleeping - - await tasks_complete_or_sleeping(world) - - async def step_simulation( world: SimulationWorld, step_size_s: float = DISCRETE_EVENT, max_advance_time_s: float = -1.0, ) -> SimulationResult | None: - """Advance the simulation by *step_size_s* seconds. - - When *step_size_s* is ``DISCRETE_EVENT`` (the default), the step size is - determined automatically as the time until the next scheduled event - (message arrival or agent task wakeup). - - :param world: the simulation world to step - :param step_size_s: step size in seconds, or ``DISCRETE_EVENT`` - :param max_advance_time_s: abort if the determined discrete step would - exceed this value; ``-1`` means no limit - :return: :class:`SimulationResult`, or ``None`` if there is nothing to do - - Example:: - - result = await step_simulation(world, step_size_s=1.0) - result = await step_simulation(world) # discrete-event - """ - world._initialize_if_needed() - - # Allow freshly scheduled tasks (e.g. from on_ready / on_start) to reach - # their first sleeping or done point. This is necessary so that periodic - # tasks register their first wakeup with the ExternalClock before we - # determine the discrete step size. - await _wait_for_agents(world) - - actual_step = step_size_s - - if actual_step == DISCRETE_EVENT: - actual_step = world._determine_next_step_size() - if actual_step is None: - return None - if max_advance_time_s >= 0 and actual_step > max_advance_time_s: - return None - - new_time = world.clock.time + actual_step - - # ---- call on_step hooks ---- - for agent in list(world._agents.values()): - agent.on_step(world.environment, world.clock, actual_step) - if hasattr(agent, "roles"): - for role in agent.roles: - role.on_step(world.environment, world.clock, actual_step) - - # ---- step environment ---- - world.environment.step(world.clock, actual_step) - - # ---- advance clock (wakes up sleeping tasks) ---- - world.clock.set_time(new_time) - - # ---- convergence loop: process messages and tasks ---- - total_delivered = 0 - state_changed = True - while state_changed: - # Wait for agents to settle - await _wait_for_agents(world) - - # Deliver messages whose delivery_time has now arrived - delivered = await world._deliver_messages_due(new_time) - total_delivered += delivered - state_changed = delivered > 0 - - # Yield control so delivered messages are processed - if state_changed: - await asyncio.sleep(0) - - world._do_recordings() - - return SimulationResult( - time_elapsed_s=actual_step, - step_size_s=actual_step, - messages_delivered=total_delivered, + """Alias for :meth:`SimulationWorld.step` (which is preferred).""" + return await world.step( + step_size_s=step_size_s, max_advance_time_s=max_advance_time_s ) @@ -502,227 +468,5 @@ async def discrete_step_until( world: SimulationWorld, max_advance_time_s: float, ) -> list[SimulationResult]: - """Run a discrete-event simulation until *max_advance_time_s* has elapsed. - - The simulation stops when the world clock has advanced by - *max_advance_time_s* from its current position, or when there are no - more events to process. - - :param world: the simulation world - :param max_advance_time_s: maximum total time to simulate in seconds - :return: list of :class:`SimulationResult` from each step - - Example:: - - results = await discrete_step_until(world, max_advance_time_s=3600.0) - """ - world._initialize_if_needed() - start_time = world.clock.time - max_time = start_time + max_advance_time_s - results: list[SimulationResult] = [] - - while world.clock.time < max_time: - remaining = max_time - world.clock.time - result = await step_simulation( - world, - step_size_s=DISCRETE_EVENT, - max_advance_time_s=remaining, - ) - if result is None: - break - results.append(result) - - return results - - -def collect_data( - world: SimulationWorld, - key: str, - collector: Callable[["SimulationWorld", WorldRecording], None], -) -> None: - """Register a world-level data collector. - - *collector* is called after every simulation step with the world and the - :class:`WorldRecording` identified by *key*. - - Example:: - - collect_data(world, "total_msgs", lambda w, rec: ( - rec.timeseries.append(len(w.recorded_messages)), - rec.time.append(w.clock.time), - )) - """ - recording = world._new_world_recording(key) - - def _run(w: SimulationWorld) -> None: - collector(w, recording) - - world._data_collectors.append(_run) - - -def collect_agent_data( - world: SimulationWorld, - key: str, - collector: Callable[["SimulationWorld", Agent, AgentsRecording], None], -) -> None: - """Register an agent-level data collector. - - *collector* is called for every agent after every simulation step with - the world, the agent, and the :class:`AgentsRecording` for *key*. A - shared ``time`` entry is appended once per step. - - Example:: - - collect_agent_data(world, "state", lambda w, a, rec: ( - rec.timeseries.setdefault(a.aid, []).append(a.some_state), - )) - """ - recording = world._new_agents_recording(key) - - def _run(w: SimulationWorld) -> None: - for agent in w._agents.values(): - collector(w, agent, recording) - recording.time.append(w.clock.time) - - world._data_collectors.append(_run) - - -def record_world( - world: SimulationWorld, - key: str, - recorder: Callable[[], Any], -) -> None: - """Record a world-level scalar after every step. - - *recorder* is a zero-argument callable whose return value is appended to - the recording's ``timeseries``. - - Example:: - - record_world(world, "agent_count", lambda: len(world._agents)) - """ - recording = world._new_world_recording(key) - - def _run(w: SimulationWorld) -> None: - recording.timeseries.append(recorder()) - recording.time.append(w.clock.time) - - world._data_collectors.append(_run) - - -def record_agent( - world: SimulationWorld, - key: str, - recorder: Callable[[Agent], Any], - filter_fn: Callable[[Agent], bool] | None = None, -) -> None: - """Record a per-agent scalar after every step. - - *recorder* receives each agent and returns the value to store. An - optional *filter_fn* restricts recording to a subset of agents — pass - an ``isinstance``-based predicate to record only agents of a particular - type:: - - record_agent(world, "soc", lambda a: a.soc_kwh, - filter_fn=lambda a: isinstance(a, EVAgent)) - - :param world: the simulation world - :param key: recording key - :param recorder: ``(agent) -> value`` callable - :param filter_fn: optional ``(agent) -> bool`` predicate; ``None`` records - all registered agents - """ - recording = world._new_agents_recording(key) - - def _run(w: SimulationWorld) -> None: - for agent in w._agents.values(): - if filter_fn is None or filter_fn(agent): - recording.timeseries.setdefault(agent.aid, []).append(recorder(agent)) - recording.agent_time.setdefault(agent.aid, []).append(w.clock.time) - recording.time.append(w.clock.time) - - world._data_collectors.append(_run) - - -def record_agent_having( - world: SimulationWorld, - key: str, - role_type: type, - recorder: Callable[[Agent], Any], -) -> None: - """Record a per-agent scalar for agents that carry a specific role type. - - Only agents that have at least one role that is an instance of - *role_type* are included in the recording. *recorder* receives the - agent and returns the value to store. - - :param world: the simulation world - :param key: recording key - :param role_type: only record agents that have a role of this type - :param recorder: ``(agent) -> value`` callable - - Example:: - - record_agent_having(world, "energy", EnergyRole, lambda a: a.roles[0].energy) - """ - recording = world._new_agents_recording(key) - - def _run(w: SimulationWorld) -> None: - for agent in w._agents.values(): - if hasattr(agent, "roles") and any( - isinstance(r, role_type) for r in agent.roles - ): - recording.timeseries.setdefault(agent.aid, []).append(recorder(agent)) - recording.agent_time.setdefault(agent.aid, []).append(w.clock.time) - recording.time.append(w.clock.time) - - world._data_collectors.append(_run) - - -def record_position( - world: SimulationWorld, - key: str = "positions", - filter_fn: Callable[[Agent], bool] | None = None, -) -> None: - """Record the spatial position of every agent after each step. - - Only agents that have a position in the world's space are recorded. - An optional *filter_fn* restricts recording to a subset of agents. - - :param world: the simulation world - :param key: recording key (default ``"positions"``) - :param filter_fn: ``(agent) -> bool`` predicate; ``None`` means all agents - - Example:: - - record_position(world) - history = position_history(world) - # history.timeseries["agent0"] -> list of Position2D - """ - recording = world._new_agents_recording(key) - - def _run(w: SimulationWorld) -> None: - space = w.environment.space - for agent in w._agents.values(): - if space.has_position(agent): - if filter_fn is None or filter_fn(agent): - recording.timeseries.setdefault(agent.aid, []).append( - space.location(agent) - ) - recording.agent_time.setdefault(agent.aid, []).append(w.clock.time) - recording.time.append(w.clock.time) - - world._data_collectors.append(_run) - - -def position_history( - world: SimulationWorld, - key: str = "positions", -) -> AgentsRecording: - """Return the :class:`AgentsRecording` populated by :func:`record_position`. - - :param world: the simulation world - :param key: recording key (default ``"positions"``) - :return: the recording - """ - return world.data_agent_collections.get(key, AgentsRecording()) + """Alias for :meth:`SimulationWorld.step_until` (which is preferred).""" + return await world.step_until(max_advance_time_s) diff --git a/tests/unit_tests/role/conversation_test.py b/tests/unit_tests/role/conversation_test.py index 8093ce9..4784c0e 100644 --- a/tests/unit_tests/role/conversation_test.py +++ b/tests/unit_tests/role/conversation_test.py @@ -1,10 +1,13 @@ """Tests for :class:`mango.agent.conversation.Conversation`. -Three flavours of coverage: +Coverage flavours: +* Handle-level — termination latch (multi-consumer wakeup, re-iteration), + send tagging, broadcast results. * Real-time (TCP) — initiator opens a conversation, joiner receives the message and joins it, both sides exchange multiple messages - under one id. + under one id; also across two containers and with plain (role-less) + agents. * Convergence / cancellation — the iterator terminates correctly when the caller signals end. * Simulation-time timeout — exercises the conversation's clock-aware @@ -21,14 +24,18 @@ from mango import ( Agent, + Conversation, Role, RoleAgent, activate, + create_acl, create_tcp_container, + json_serializable, on_message, sender_addr, ) from mango.express.api import run_with_simulation +from mango.messages.codecs import JSON from mango.simulation.world import step_simulation @@ -180,6 +187,328 @@ async def run(self): assert solo.pulls == 0 +# --------------------------------------------------------------------------- +# Handle-level behaviour — no container needed. +# --------------------------------------------------------------------------- + + +class _RecordingOwner: + """Stub owner capturing what ``Conversation.send`` forwards.""" + + def __init__(self, results: dict | None = None): + self.sent = [] + self._results = results or {} + + async def send_message(self, content, receiver_addr, **kwargs): + self.sent.append((receiver_addr, content, kwargs)) + return self._results.get(receiver_addr, True) + + +def _handle(owner=None, **kwargs) -> Conversation: + return Conversation(owner=owner, conversation_id="cid-1", **kwargs) + + +@pytest.mark.asyncio +async def test_converge_wakes_all_concurrent_iterators(): + """The end sentinel is a latch: every consumer blocked in ``__anext__`` + terminates, not just the one that happens to read it first.""" + conv = _handle() + + async def consume(): + async for _ in conv: # pragma: no cover - stream stays empty + pass + return True + + consumers = [asyncio.create_task(consume()) for _ in range(3)] + await asyncio.sleep(0) # let all consumers block in queue.get() + conv.converge() + results = await asyncio.wait_for(asyncio.gather(*consumers), timeout=1.0) + assert results == [True, True, True] + + +@pytest.mark.asyncio +async def test_iteration_after_converge_terminates_immediately(): + """A second drain loop on an already-converged handle must exit + instead of awaiting an empty queue forever.""" + conv = _handle() + conv._on_inbound("x", {}) + conv.converge() + + assert [c async for c, _ in conv] == ["x"] + assert [c async for c, _ in conv] == [] + + +@pytest.mark.asyncio +async def test_cancel_wakes_all_concurrent_iterators(): + conv = _handle() + + async def consume(): + return [c async for c, _ in conv] + + consumers = [asyncio.create_task(consume()) for _ in range(2)] + await asyncio.sleep(0) # both block in queue.get() + conv.cancel() + results = await asyncio.wait_for(asyncio.gather(*consumers), timeout=1.0) + assert results == [[], []] + + +@pytest.mark.asyncio +async def test_send_id_wins_over_kwargs_and_stamps_content(): + """The conversation's own id beats a stray ``conversation_id`` kwarg, + and content carrying a ``conversation_id`` attribute (ACL-style) is + stamped so the id survives split-content transports.""" + + @dataclass + class _AclLike: + conversation_id: str | None = None + + owner = _RecordingOwner() + conv = _handle(owner) + payload = _AclLike() + + assert await conv.send("addr-1", payload, conversation_id="spoofed") + + _, sent_content, kwargs = owner.sent[0] + assert kwargs["conversation_id"] == "cid-1" + assert sent_content.conversation_id == "cid-1" + + +@pytest.mark.asyncio +async def test_broadcast_returns_result_per_receiver(): + owner = _RecordingOwner(results={"bad": False}) + conv = _handle(owner) + + results = await conv.broadcast(["good", "bad"], _Step(counter=0)) + + assert results == {"good": True, "bad": False} + assert all(kw["conversation_id"] == "cid-1" for _, _, kw in owner.sent) + + +@pytest.mark.asyncio +async def test_handle_escaping_its_context_terminates(): + """Leaving the ``async with`` cancels the handle, so an iterator that + escaped the block ends instead of hanging on a dead queue.""" + container = create_tcp_container(addr=("127.0.0.1", 5583)) + agent = container.register(RoleAgent()) + role = Role() + agent.add_role(role) + + async with activate([container]): + async with role.context.open_conversation() as conv: + pass + assert conv.closed + with pytest.raises(StopAsyncIteration): + await conv.__anext__() + + +# --------------------------------------------------------------------------- +# Fan-out: several handles on the same id all receive every message. +# --------------------------------------------------------------------------- + + +class _DoubleJoiner(Agent): + """Joins the same conversation twice concurrently; both handles must + independently receive the follow-up message.""" + + def __init__(self): + super().__init__() + self.ready = asyncio.Event() + self.got: tuple | None = None + + def handle_message(self, content, meta): + if content == "start": + self.schedule_instant_task(self._double_join(dict(meta))) + + async def _double_join(self, meta): + async with self.join_conversation(meta) as first: + async with self.join_conversation(meta) as second: + self.ready.set() + got_first = await first.__anext__() + got_second = await second.__anext__() + self.got = (got_first[0], got_second[0]) + + +@pytest.mark.asyncio +async def test_overlapping_joins_each_receive_every_message(): + container = create_tcp_container(addr=("127.0.0.1", 5584)) + joiner = container.register(_DoubleJoiner()) + initiator = container.register(RoleAgent()) + role = Role() + initiator.add_role(role) + + async with activate([container]): + async with role.context.open_conversation(timeout=5.0) as conv: + await conv.send(joiner.addr, "start") + await asyncio.wait_for(joiner.ready.wait(), timeout=2.0) + await conv.send(joiner.addr, "payload") + while joiner.got is None: + await asyncio.sleep(0.01) + + assert joiner.got == ("payload", "payload") + + +# --------------------------------------------------------------------------- +# Plain (role-less) agents: conversations, reply_to threading, gather. +# --------------------------------------------------------------------------- + + +class _PlainResponder(Agent): + """Replies via ``reply_to`` — the conversation id must thread back so + the initiator's iterator receives the reply.""" + + def handle_message(self, content, meta): + if isinstance(content, _Step): + self.schedule_instant_task( + self.reply_to(_Step(counter=content.counter + 1), meta) + ) + + +class _PlainInitiator(Agent): + def __init__(self, rounds: int): + super().__init__() + self.rounds = rounds + self.seen: list[int] = [] + + def handle_message(self, content, meta): + pass + + async def run(self, peer) -> None: + async with self.open_conversation(timeout=5.0) as conv: + await conv.send(peer, _Step(counter=0)) + async for content, _meta in conv: + self.seen.append(content.counter) + if len(self.seen) >= self.rounds: + conv.converge() + continue + await conv.send(peer, _Step(counter=content.counter + 1)) + + +@pytest.mark.asyncio +async def test_plain_agent_conversation_with_reply_to(): + container = create_tcp_container(addr=("127.0.0.1", 5585)) + responder = container.register(_PlainResponder()) + initiator = container.register(_PlainInitiator(rounds=2)) + + async with activate([container]): + await initiator.run(responder.addr) + + assert initiator.seen == [1, 3] + + +class _PongAgent(Agent): + def handle_message(self, content, meta): + if content == "ping": + self.schedule_instant_task(self.reply_to("pong", meta)) + + +class _QuietAgent(Agent): + def handle_message(self, content, meta): + pass + + +@pytest.mark.asyncio +async def test_plain_agent_gather(): + container = create_tcp_container(addr=("127.0.0.1", 5586)) + responders = [container.register(_PongAgent()) for _ in range(3)] + caller = container.register(_QuietAgent()) + + async with activate([container]): + replies = await caller.gather("ping", [r.addr for r in responders], timeout=2.0) + + assert len(replies) == 3 + assert set(replies.values()) == {"pong"} + assert set(replies.keys()) == {r.addr for r in responders} + + +# --------------------------------------------------------------------------- +# Cross-container: the id must survive the codec, for plain and ACL content. +# --------------------------------------------------------------------------- + + +@json_serializable +@dataclass +class _WireStep: + counter: int + + +def _wire_codec() -> JSON: + codec = JSON() + codec.add_serializer(*_WireStep.__serializer__()) + return codec + + +class _WireJoiner(Role): + def __init__(self) -> None: + super().__init__() + self.processed: list[int] = [] + + @on_message(_WireStep) + async def on_step(self, content: _WireStep, meta: dict) -> None: + self.processed.append(content.counter) + async with self.context.join_conversation(meta) as conv: + await conv.send(sender_addr(meta), _WireStep(counter=content.counter + 1)) + + +class _WireInitiator(Role): + def __init__(self, peer_addr, *, rounds: int, acl: bool = False) -> None: + super().__init__() + self.peer_addr = peer_addr + self.rounds = rounds + self.acl = acl + self.seen: list[int] = [] + + def _payload(self, counter: int): + step = _WireStep(counter=counter) + if not self.acl: + return step + return create_acl( + step, receiver_addr=self.peer_addr, sender_addr=self.context.addr + ) + + async def run(self) -> None: + async with self.context.open_conversation(timeout=5.0) as conv: + await conv.send(self.peer_addr, self._payload(0)) + async for content, _meta in conv: + self.seen.append(content.counter) + if len(self.seen) >= self.rounds: + conv.converge() + continue + await conv.send(self.peer_addr, self._payload(content.counter + 1)) + + +async def _run_two_container_conversation(port_a: int, port_b: int, *, acl: bool): + container_a = create_tcp_container(addr=("127.0.0.1", port_a), codec=_wire_codec()) + container_b = create_tcp_container(addr=("127.0.0.1", port_b), codec=_wire_codec()) + + joiner_agent = container_b.register(RoleAgent()) + joiner = _WireJoiner() + joiner_agent.add_role(joiner) + + initiator_agent = container_a.register(RoleAgent()) + initiator = _WireInitiator(joiner_agent.addr, rounds=2, acl=acl) + initiator_agent.add_role(initiator) + + async with activate([container_a, container_b]): + await asyncio.wait_for(initiator.run(), timeout=5.0) + + assert initiator.seen == [1, 3] + assert joiner.processed == [0, 2] + + +@pytest.mark.asyncio +async def test_conversation_across_two_containers(): + """The conversation id survives JSON-codec serialization over TCP.""" + await _run_two_container_conversation(5587, 5588, acl=False) + + +@pytest.mark.asyncio +async def test_conversation_across_two_containers_with_acl_content(): + """ACL content is sent without the kwargs meta, so the id must travel + in the ACL message's own ``conversation_id`` field (stamped by + ``conv.send``).""" + await _run_two_container_conversation(5589, 5590, acl=True) + + # --------------------------------------------------------------------------- # Simulation-clock timeout — the load-bearing requirement. # --------------------------------------------------------------------------- diff --git a/tests/unit_tests/role/gather_test.py b/tests/unit_tests/role/gather_test.py index 9334d37..a006770 100644 --- a/tests/unit_tests/role/gather_test.py +++ b/tests/unit_tests/role/gather_test.py @@ -279,3 +279,38 @@ async def on_ask(self, content: _Ask, meta: dict) -> None: # noqa: ARG002 (reply,) = caller_role.responses.values() assert isinstance(reply, _Reply) assert reply.value == 1.0 + + +@pytest.mark.asyncio +async def test_role_send_tracked_message_response_handler_fires(): + """A role's ``send_tracked_message`` must register its response handler + on the *agent* (whose inbox matches replies), not on the RoleContext's + own inherited registry — a handler registered there never fires.""" + + class _TrackedCaller(Role): + def __init__(self, peer): + super().__init__() + self.peer = peer + self.reply = None + + async def run(self): + got_reply = asyncio.get_event_loop().create_future() + await self.context.send_tracked_message( + _Ask(topic="x"), + self.peer, + response_handler=lambda content, meta: got_reply.set_result(content), + ) + self.reply = await asyncio.wait_for(got_reply, timeout=2.0) + + container = create_tcp_container(addr=("127.0.0.1", 5591)) + responder = container.register(RoleAgent()) + responder.add_role(_Responder(value=7.0)) + + caller = container.register(RoleAgent()) + caller_role = _TrackedCaller(responder.addr) + caller.add_role(caller_role) + + async with activate([container]): + await caller_role.run() + + assert caller_role.reply == _Reply(value=7.0) diff --git a/tests/unit_tests/simulation/test_simulation.py b/tests/unit_tests/simulation/test_simulation.py index 2ead8b8..371199f 100644 --- a/tests/unit_tests/simulation/test_simulation.py +++ b/tests/unit_tests/simulation/test_simulation.py @@ -504,7 +504,7 @@ async def test_deregister(self): world = create_world() a = world.register(SimpleAgent()) world.deregister(a.aid) - assert a.aid not in world._agents + assert a.aid not in world.agents await world.shutdown() def test_deregister_nonexistent_is_noop(self): @@ -672,16 +672,16 @@ async def test_message_unknown_receiver_dropped(): world.register(SimpleAgent()) # just to have an agent # Manually inject a message with unknown receiver async with world: - await world._deliver_messages_due.__func__(world, 999.0) if False else None # noqa: this branch is tested via direct queue # Insert directly into pending with bogus receiver import bisect - world._pending_messages.clear() + container = world.container + container._pending_messages.clear() bisect.insort( - world._pending_messages, + container._pending_messages, (0.0, 0, 0.0, "msg", {"receiver_id": "ghost", "sender_id": None}), ) - delivered = await world._deliver_messages_due(10.0) + delivered = await container._deliver_messages_due(10.0) assert delivered == 0 # dropped, not counted @@ -725,7 +725,7 @@ async def test_shutdown_sets_running_false(): async def test_record_world(): world = create_world() world.register(SimpleAgent()) - record_world(world, "agent_count", lambda: len(world._agents)) + record_world(world, "agent_count", lambda: len(world.agents)) async with world: await step_simulation(world, step_size_s=1.0) await step_simulation(world, step_size_s=1.0) From fde9c636d3199c08f1f447c86418eb27bd0a7c2e Mon Sep 17 00:00:00 2001 From: Rico Schrage Date: Sat, 11 Jul 2026 07:02:49 +0200 Subject: [PATCH 08/10] Setter for com sim. --- mango/simulation/world.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/mango/simulation/world.py b/mango/simulation/world.py index c887ec1..cdaadec 100644 --- a/mango/simulation/world.py +++ b/mango/simulation/world.py @@ -170,6 +170,10 @@ def clock(self) -> ExternalClock: def communication_sim(self) -> CommunicationSimulation: return self._container.communication_sim + @communication_sim.setter + def communication_sim(self, value: CommunicationSimulation) -> None: + self._container.communication_sim = value + @property def addr(self) -> str: return self._container.addr From b9c74faf720f574d1ca421b747419c8606048ebc Mon Sep 17 00:00:00 2001 From: Rico Schrage Date: Tue, 21 Jul 2026 18:45:54 +0200 Subject: [PATCH 09/10] Introducing dead lock prevention mechanism for arbitrary tasks. Fixing deadlock for specific types of async tasks. --- mango/agent/conversation.py | 8 +- mango/agent/core.py | 18 ++-- mango/simulation/world.py | 4 +- mango/util/scheduling.py | 64 +++++++++------ .../unit_tests/simulation/test_simulation.py | 82 +++++++++++++++++++ 5 files changed, 143 insertions(+), 33 deletions(-) diff --git a/mango/agent/conversation.py b/mango/agent/conversation.py index 13adeb5..cc488cd 100644 --- a/mango/agent/conversation.py +++ b/mango/agent/conversation.py @@ -55,6 +55,8 @@ async def on_step(self, msg, meta): import logging from typing import TYPE_CHECKING, Any +from mango.util.scheduling import sleeping_wait + if TYPE_CHECKING: from collections.abc import Iterable @@ -198,7 +200,11 @@ async def __anext__(self): # cancel() is abrupt: anything still on the queue is dropped. if self._cancelled: raise StopAsyncIteration - content, meta = await self._queue.get() + # Only an inbound message (or converge/cancel/timeout) feeds the + # queue, so a parked iterator counts as sleeping for stepped + # simulations — the message can only arrive in a later step. + with sleeping_wait(): + content, meta = await self._queue.get() if content is self._END: # Re-enqueue the sentinel so every other consumer — blocked # concurrently or arriving later — terminates as well. diff --git a/mango/agent/core.py b/mango/agent/core.py index cf9ccb4..f1b50fa 100644 --- a/mango/agent/core.py +++ b/mango/agent/core.py @@ -18,7 +18,12 @@ from ..messages.message import AgentAddress from ..util.clock import Clock -from ..util.scheduling import ScheduledProcessTask, ScheduledTask, Scheduler +from ..util.scheduling import ( + ScheduledProcessTask, + ScheduledTask, + Scheduler, + sleeping_wait, +) from .conversation import CONVERSATION_ID_KEY, Conversation, _ConversationContext if TYPE_CHECKING: @@ -704,10 +709,13 @@ async def gather( # is identical under AsyncioClock (real time) and # ExternalClock (simulation). asyncio.wait_for would block # real seconds in sim mode. - await _race_first_completed( - collector.wait(), - agent.scheduler.clock.sleep(timeout), - ) + # Both racers resolve externally (reply message / clock step), + # so the wait counts as sleeping for stepped simulations. + with sleeping_wait(): + await _race_first_completed( + collector.wait(), + agent.scheduler.clock.sleep(timeout), + ) finally: agent.close_gather(tracking_id) return dict(collector.responses) diff --git a/mango/simulation/world.py b/mango/simulation/world.py index cdaadec..191c7c7 100644 --- a/mango/simulation/world.py +++ b/mango/simulation/world.py @@ -46,6 +46,7 @@ async def run(): from mango.agent.core import Agent, AgentAddress from mango.messages.codecs import Codec from mango.util.clock import ExternalClock +from mango.util.termination_detection import tasks_complete_or_sleeping from .communication import CommunicationSimulation, SimpleCommunicationSimulation from .container import MessageTransaction, SimulationContainer @@ -306,9 +307,6 @@ def _initialize_if_needed(self) -> None: self._do_recordings() async def _wait_for_agents(self) -> None: - """Wait until all agent tasks are complete or sleeping.""" - from mango.util.termination_detection import tasks_complete_or_sleeping - await tasks_complete_or_sleeping(self._container) async def step( diff --git a/mango/util/scheduling.py b/mango/util/scheduling.py index 1dfc608..34ae054 100644 --- a/mango/util/scheduling.py +++ b/mango/util/scheduling.py @@ -7,6 +7,8 @@ import logging from abc import abstractmethod from asyncio import Future +from contextlib import contextmanager +from contextvars import ContextVar from dataclasses import dataclass from datetime import datetime, timezone from multiprocessing import Manager @@ -53,18 +55,46 @@ def suspend_task(self): self.run_task_event.clear() +_current_scheduled_task: ContextVar["ScheduledTask | None"] = ContextVar( + "_current_scheduled_task", default=None +) + + +@contextmanager +def sleeping_wait(): + """Mark the current scheduled task as sleeping for the duration of a wait + that only an external event can resolve (an incoming message, a clock + step). Stepped-simulation termination detection treats sleeping tasks as + idle; without this scope it would wait for the task to finish and + deadlock, as the resolving event can only arrive in a later step. + No-op outside a scheduled task. Only autonomous waits (wall-clock timers, + executors, I/O) must not be wrapped — those have to finish within a step. + """ + task = _current_scheduled_task.get() + if task is None: + yield + return + was_sleeping = task._is_observable and task._is_sleeping.done() + task.notify_sleeping() + try: + yield + finally: + if not was_sleeping: + task.notify_running() + + +async def _run_as_current_task(task: "ScheduledTask"): + _current_scheduled_task.set(task) + return await task.run() + + class Suspendable: """ Wraps a coroutine, intercepting __await__ to add the functionality of suspending. """ - def __init__(self, coro, ext_contr_event=None, kill_event=None, notify_task=None): + def __init__(self, coro, ext_contr_event=None, kill_event=None): self._coro = coro - # Optional owning ScheduledTask; when set, the wrapper marks it - # "sleeping" whenever the coroutine parks on a pending future so that - # simulation termination detection treats request/reply drivers (which - # await replies rather than the clock) as idle instead of deadlocking. - self._notify_task = notify_task self._kill_event = kill_event if ext_contr_event is not None: @@ -99,24 +129,11 @@ def __await__(self): return err.value else: send = iter_send - # Parked on a pending future (awaiting a reply/message rather than - # progressing) -> mark the owning task sleeping for the duration of - # the suspension. Idempotent, so it nests safely with clock sleeps. - parked = ( - self._notify_task is not None - and isinstance(signal, asyncio.Future) - and not signal.done() - ) - if parked: - self._notify_task.notify_sleeping() try: # pass signal via yielding it message = yield signal except BaseException as err: send, message = iter_throw, err - finally: - if parked: - self._notify_task.notify_running() def suspend(self): """ @@ -176,9 +193,8 @@ def __init__(self, clock: Clock = None, observable=True, on_stop=None) -> None: self._is_done = asyncio.Future() def notify_sleeping(self): - # Idempotent so it composes with Suspendable's per-await notifications - # (a clock-sleep already marks the task sleeping before the wrapper - # also sees the yielded sleep future). + # Idempotent so explicit notifications compose with a surrounding + # sleeping_wait() scope. if self._is_observable and not self._is_sleeping.done(): self._is_sleeping.set_result(True) @@ -542,10 +558,10 @@ def schedule_task(self, task: ScheduledTask, src=None) -> asyncio.Task: """ l_task = None if self.suspendable: - coro = Suspendable(task.run(), notify_task=task) + coro = Suspendable(_run_as_current_task(task)) l_task = asyncio.ensure_future(coro) else: - coro = task.run() + coro = _run_as_current_task(task) l_task = asyncio.create_task(coro) l_task.add_done_callback(task.on_stop) l_task.add_done_callback(_raise_exceptions) diff --git a/tests/unit_tests/simulation/test_simulation.py b/tests/unit_tests/simulation/test_simulation.py index 371199f..713b0ad 100644 --- a/tests/unit_tests/simulation/test_simulation.py +++ b/tests/unit_tests/simulation/test_simulation.py @@ -1,3 +1,4 @@ +import asyncio from typing import Any from unittest.mock import MagicMock @@ -1250,3 +1251,84 @@ async def _tick(self): assert world.clock.time <= 3.0 assert len(results) >= 1 assert agent.ticks >= 1 + + +@pytest.mark.asyncio +async def test_gather_from_scheduled_task_does_not_deadlock_step(): + # Regression: a gather awaiting replies inside a scheduled task must count + # as sleeping for termination detection — otherwise world.step() hangs, + # since the replies can only be delivered after detection returns. + world = create_world() + + class GatherAgent(Agent): + def __init__(self): + super().__init__() + self.peer = None + self.responses = None + + def handle_message(self, content, meta): + pass + + def on_ready(self): + self.schedule_instant_task(self.run_gather()) + + async def run_gather(self): + self.responses = await self.gather( + "ping", receivers=[self.peer], timeout=30.0 + ) + + gatherer = world.register(GatherAgent()) + responder = world.register(ReplyAgent()) + gatherer.peer = responder.addr + + async def drive(): + # world entry/exit also run termination detection, so the timeout + # must span the whole world lifecycle to catch a deadlock anywhere + async with world: + for _ in range(4): + await step_simulation(world, step_size_s=1.0) + if gatherer.responses is not None: + break + + await asyncio.wait_for(drive(), timeout=10) + assert gatherer.responses == {responder.addr: "pong"} + + +@pytest.mark.asyncio +async def test_conversation_from_scheduled_task_does_not_deadlock_step(): + # Regression: a conversation iterator parked on its queue inside a + # scheduled task must count as sleeping for termination detection. + world = create_world() + + class InitiatorAgent(Agent): + def __init__(self): + super().__init__() + self.peer = None + self.received = None + + def handle_message(self, content, meta): + pass + + def on_ready(self): + self.schedule_instant_task(self.run_conversation()) + + async def run_conversation(self): + async with self.open_conversation(timeout=30.0) as conv: + await conv.send(self.peer, "ping") + async for content, _ in conv: + self.received = content + conv.converge() + + initiator = world.register(InitiatorAgent()) + responder = world.register(ReplyAgent()) + initiator.peer = responder.addr + + async def drive(): + async with world: + for _ in range(4): + await step_simulation(world, step_size_s=1.0) + if initiator.received is not None: + break + + await asyncio.wait_for(drive(), timeout=10) + assert initiator.received == "pong" From ecda32e7156805dfd6e0905aa9fab239842584fd Mon Sep 17 00:00:00 2001 From: Rico Schrage Date: Tue, 28 Jul 2026 18:56:59 +0200 Subject: [PATCH 10/10] Readme. --- readme.md | 38 ++++++++++++++++++++++++-------------- 1 file changed, 24 insertions(+), 14 deletions(-) diff --git a/readme.md b/readme.md index 6542cfc..6a4efdf 100644 --- a/readme.md +++ b/readme.md @@ -75,7 +75,7 @@ async def main(): # All agents on a single host share one container. container = create_tcp_container(addr=("127.0.0.1", 5555)) - sender = container.register(ReportingAgent()) + sender = container.register(ReportingAgent()) receiver = container.register(ReportingAgent()) async with activate(container): @@ -149,7 +149,7 @@ class SensorAgent(Agent): def __init__(self): super().__init__() - self.monitor_addr = None # set after both agents are registered + self.monitor_addr = None # set after both agents are registered self.readings_sent = 0 def on_ready(self): @@ -173,7 +173,7 @@ class MonitorAgent(Agent): async def run(): - world = create_world(start_time=0.0) + world = create_world(start_time=0.0) sensor = world.register(SensorAgent()) monitor = world.register(MonitorAgent()) @@ -181,8 +181,12 @@ async def run(): sensor.monitor_addr = monitor.addr # Record the running count of received messages after every step. - record_agent(world, "received", lambda a: a.received, - filter_fn=lambda a: isinstance(a, MonitorAgent)) + record_agent( + world, + "received", + lambda a: a.received, + filter_fn=lambda a: isinstance(a, MonitorAgent), + ) async with world: await discrete_step_until(world, max_advance_time_s=60.0) @@ -212,7 +216,9 @@ from mango import step_simulation async with world: for _ in range(10): result = await step_simulation(world, step_size_s=1.0) - print(f"t = {world.clock.time:.1f} s messages delivered: {result.messages_delivered}") + print( + f"t = {world.clock.time:.1f} s messages delivered: {result.messages_delivered}" + ) ``` ### Communication modelling @@ -225,8 +231,8 @@ from mango import create_world, SimpleCommunicationSimulation world = create_world( start_time=0.0, communication_sim=SimpleCommunicationSimulation( - default_delay_s=0.1, # baseline one-way latency - loss_percent=0.01, # independent loss probability per message + default_delay_s=0.1, # baseline one-way latency + loss_percent=0.01, # independent loss probability per message delay_s_directed_edge_dict={ ("agent0", "agent1"): 0.5, # directed per-link override }, @@ -256,11 +262,15 @@ Agents can be embedded in a continuous 2-D area. Agents without a pre-assigned p ```python from mango import ( - create_world, Area2D, DefaultEnvironment, - record_position, position_history, discrete_step_until, + create_world, + Area2D, + DefaultEnvironment, + record_position, + position_history, + discrete_step_until, ) -env = DefaultEnvironment(space=Area2D(width=100.0, height=100.0)) +env = DefaultEnvironment(space=Area2D(width=100.0, height=100.0)) world = create_world(start_time=0.0, environment=env) # ... register agents ... @@ -295,9 +305,9 @@ record_agent(world, "energy_kwh", lambda a: a.energy_kwh) async with world: await discrete_step_until(world, max_advance_time_s=3600.0) -world.data_collections["msg_count"].timeseries # list of scalars -world.data_agent_collections["energy_kwh"].timeseries # dict[aid -> list] -world.data_agent_collections["energy_kwh"].time # shared time axis +world.data_collections["msg_count"].timeseries # list of scalars +world.data_agent_collections["energy_kwh"].timeseries # dict[aid -> list] +world.data_agent_collections["energy_kwh"].time # shared time axis ``` All exchanged messages are also logged automatically in `world.recorded_messages` as `MessageTransaction` objects, each carrying sender, receiver, send time, and arrival time.