From f0a61d900ce98acad3247caad242dd6eafd8a539 Mon Sep 17 00:00:00 2001 From: Weishi Zeng Date: Tue, 16 Jun 2026 16:30:51 -0700 Subject: [PATCH 1/3] agent_discovery: cache resolved agent location to skip re-discovery on repeat sends send_message resolved its target via a full mngr discovery (find_all_agents) on every call. Cache each agent's resolved AgentMatch by name and feed it straight to mngr's pre-resolved agents_to_message API, so repeat messages to the same agent skip discovery and go to the known host. On a stale cache hit (the send reaches no agent -- destroyed, recreated, or moved hosts) the entry is dropped and the agent is re-resolved, preserving correctness. A name resolving to multiple agents is never cached, preserving the existing "message all matches" behavior. STOPPED agents are still auto-started (is_start_desired=True) regardless of cache state. The resolver/sender are injected as MutableModel callables so the cache and stale-fallback orchestration is unit-tested without monkeypatching the mngr API. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../changelog/wz-message-cache-agent-match.md | 1 + .../imbue/system_interface/agent_discovery.py | 105 ++++++++++++++-- .../system_interface/agent_discovery_test.py | 119 +++++++++++++++++- 3 files changed, 211 insertions(+), 14 deletions(-) create mode 100644 apps/system_interface/changelog/wz-message-cache-agent-match.md diff --git a/apps/system_interface/changelog/wz-message-cache-agent-match.md b/apps/system_interface/changelog/wz-message-cache-agent-match.md new file mode 100644 index 000000000..f0a176277 --- /dev/null +++ b/apps/system_interface/changelog/wz-message-cache-agent-match.md @@ -0,0 +1 @@ +- `agent_discovery.send_message` now caches each agent's resolved location (its `AgentMatch`) keyed by agent name, so repeat messages to the same agent skip mngr's full agent discovery and send straight to the known host. This leverages mngr's pre-resolved `agents_to_message` API: the first message to an agent resolves via `find_all_agents` as before, and subsequent messages reuse the cached location. If a cached send reaches no agent (the agent was destroyed, recreated, or moved hosts), the entry is dropped and the agent is re-resolved, so correctness is preserved. STOPPED agents are still auto-started regardless of cache state. diff --git a/apps/system_interface/imbue/system_interface/agent_discovery.py b/apps/system_interface/imbue/system_interface/agent_discovery.py index c1c7fcee1..c74f0e3c2 100644 --- a/apps/system_interface/imbue/system_interface/agent_discovery.py +++ b/apps/system_interface/imbue/system_interface/agent_discovery.py @@ -3,18 +3,25 @@ from __future__ import annotations import os +import threading +from collections.abc import Callable +from collections.abc import Sequence from pathlib import Path from loguru import logger as _loguru_logger from pydantic import Field +from pydantic import PrivateAttr from imbue.concurrency_group.concurrency_group import ConcurrencyGroup from imbue.imbue_common.frozen_model import FrozenModel +from imbue.imbue_common.mutable_model import MutableModel +from imbue.mngr.api.find import AgentMatch from imbue.mngr.api.find import find_all_agents from imbue.mngr.api.find import find_one_agent from imbue.mngr.api.find import resolve_to_started_host_and_running_agent from imbue.mngr.api.list import ErrorBehavior from imbue.mngr.api.list import list_agents +from imbue.mngr.api.message import MessageResult from imbue.mngr.api.message import send_message_to_agents from imbue.mngr.config.data_types import MngrContext from imbue.mngr.config.loader import load_config @@ -167,30 +174,102 @@ def discover_agents( return agents -def send_message(agent_name: str, message: str) -> bool: - """Send a message to an agent. Returns True on success. +class _AgentMatchCache(MutableModel): + """Thread-safe cache of resolved agent locations, keyed by agent name.""" - STOPPED agents are automatically started before the message is sent - (`is_start_desired=True`), so messaging is possible regardless of agent state. - """ - mngr_ctx, cg = _get_mngr_context() - try: - matches = find_all_agents( + _lock: threading.Lock = PrivateAttr(default_factory=threading.Lock) + _matches: dict[str, AgentMatch] = PrivateAttr(default_factory=dict) + + def get(self, agent_name: str) -> AgentMatch | None: + with self._lock: + return self._matches.get(agent_name) + + def put(self, agent_name: str, match: AgentMatch) -> None: + with self._lock: + self._matches[agent_name] = match + + def invalidate(self, agent_name: str) -> None: + with self._lock: + self._matches.pop(agent_name, None) + + +_AGENT_MATCH_CACHE = _AgentMatchCache() + + +class _AgentResolver(MutableModel): + """Resolves an agent name to its mngr locations within a fixed context.""" + + mngr_ctx: MngrContext + + def __call__(self, agent_name: str) -> Sequence[AgentMatch]: + return find_all_agents( addresses=(AgentAddress(agent=AgentName(agent_name)),), filter_all=False, target_state=None, - mngr_ctx=mngr_ctx, + mngr_ctx=self.mngr_ctx, ) - result = send_message_to_agents( - mngr_ctx=mngr_ctx, - message_content=message, + + +class _AgentSender(MutableModel): + """Sends a fixed message to pre-resolved agents, auto-starting STOPPED ones.""" + + mngr_ctx: MngrContext + message: str + + def __call__(self, matches: Sequence[AgentMatch]) -> MessageResult: + return send_message_to_agents( + mngr_ctx=self.mngr_ctx, + message_content=self.message, agents_to_message=matches, error_behavior=ErrorBehavior.CONTINUE, is_start_desired=True, ) + + +def _send_to_cached_agent( + agent_name: str, + cache: _AgentMatchCache, + resolve: Callable[[str], Sequence[AgentMatch]], + send_to: Callable[[Sequence[AgentMatch]], MessageResult], +) -> bool: + """Route a message to ``agent_name``, resolving its location only when needed. + + On a cache hit the message goes straight to the known location -- no mngr + discovery. If that send reaches no agent (the cached location is stale: the + agent was destroyed, recreated, or moved hosts), the entry is dropped and the + agent is re-resolved from scratch. A uniquely-resolved location is cached for + reuse; a name that resolves to several agents is never cached, so every such + send still reaches all of them. + """ + cached = cache.get(agent_name) + if cached is not None: + if send_to((cached,)).successful_agents: + return True + cache.invalidate(agent_name) + matches = tuple(resolve(agent_name)) + if len(matches) == 1: + cache.put(agent_name, matches[0]) + return bool(send_to(matches).successful_agents) + + +def send_message(agent_name: str, message: str) -> bool: + """Send a message to an agent. Returns True on success. + + The agent's location is resolved via mngr discovery once and cached, so + repeat messages to the same agent skip discovery and go straight to its host + (this is what mngr's pre-resolved ``agents_to_message`` API enables). STOPPED + agents are auto-started (`is_start_desired=True`) regardless of cache state. + """ + mngr_ctx, cg = _get_mngr_context() + try: + return _send_to_cached_agent( + agent_name, + _AGENT_MATCH_CACHE, + _AgentResolver(mngr_ctx=mngr_ctx), + _AgentSender(mngr_ctx=mngr_ctx, message=message), + ) finally: cg.__exit__(None, None, None) - return len(result.successful_agents) > 0 def start_agent(agent_name: str) -> None: diff --git a/apps/system_interface/imbue/system_interface/agent_discovery_test.py b/apps/system_interface/imbue/system_interface/agent_discovery_test.py index ce6813d6e..29be7dba0 100644 --- a/apps/system_interface/imbue/system_interface/agent_discovery_test.py +++ b/apps/system_interface/imbue/system_interface/agent_discovery_test.py @@ -1,9 +1,21 @@ """Tests for agent_discovery module.""" +from collections.abc import Sequence from pathlib import Path import pytest - +from pydantic import Field + +from imbue.imbue_common.mutable_model import MutableModel +from imbue.mngr.api.find import AgentMatch +from imbue.mngr.api.message import MessageResult +from imbue.mngr.primitives import AgentId +from imbue.mngr.primitives import AgentName +from imbue.mngr.primitives import HostId +from imbue.mngr.primitives import HostName +from imbue.mngr.primitives import ProviderInstanceName +from imbue.system_interface.agent_discovery import _AgentMatchCache +from imbue.system_interface.agent_discovery import _send_to_cached_agent from imbue.system_interface.agent_discovery import read_claude_config_dir_from_env_file @@ -94,3 +106,108 @@ def test_falls_back_to_home_claude_when_nothing_else_exists(monkeypatch: pytest. result = read_claude_config_dir_from_env_file(agent_state_dir) assert result == Path.home() / ".claude" + + +def _make_match(name: str, host: str = "host-a") -> AgentMatch: + return AgentMatch( + agent_id=AgentId.generate(), + agent_name=AgentName(name), + host_id=HostId.generate(), + host_name=HostName(host), + provider_name=ProviderInstanceName("local"), + ) + + +class _RecordingResolver(MutableModel): + """Stand-in for find_all_agents: returns a fixed result and records calls.""" + + result: tuple[AgentMatch, ...] + calls: list[str] = Field(default_factory=list) + + def __call__(self, agent_name: str) -> Sequence[AgentMatch]: + self.calls.append(agent_name) + return self.result + + +class _RecordingSender(MutableModel): + """Stand-in for send_message_to_agents: reaches only the live agents given.""" + + live_agent_ids: frozenset[AgentId] + calls: list[tuple[AgentMatch, ...]] = Field(default_factory=list) + + def __call__(self, matches: Sequence[AgentMatch]) -> MessageResult: + self.calls.append(tuple(matches)) + reached = [str(m.agent_name) for m in matches if m.agent_id in self.live_agent_ids] + return MessageResult(successful_agents=reached) + + +def test_cache_miss_resolves_caches_and_sends() -> None: + cache = _AgentMatchCache() + match = _make_match("alpha") + resolver = _RecordingResolver(result=(match,)) + sender = _RecordingSender(live_agent_ids=frozenset({match.agent_id})) + + assert _send_to_cached_agent("alpha", cache, resolver, sender) is True + assert resolver.calls == ["alpha"] + assert sender.calls == [(match,)] + assert cache.get("alpha") == match + + +def test_cache_hit_skips_resolution() -> None: + cache = _AgentMatchCache() + match = _make_match("alpha") + cache.put("alpha", match) + resolver = _RecordingResolver(result=(_make_match("alpha"),)) + sender = _RecordingSender(live_agent_ids=frozenset({match.agent_id})) + + assert _send_to_cached_agent("alpha", cache, resolver, sender) is True + assert resolver.calls == [] + assert sender.calls == [(match,)] + + +def test_stale_cache_hit_reresolves_and_updates() -> None: + cache = _AgentMatchCache() + stale = _make_match("alpha") + fresh = _make_match("alpha") + cache.put("alpha", stale) + resolver = _RecordingResolver(result=(fresh,)) + sender = _RecordingSender(live_agent_ids=frozenset({fresh.agent_id})) + + assert _send_to_cached_agent("alpha", cache, resolver, sender) is True + assert resolver.calls == ["alpha"] + assert sender.calls == [(stale,), (fresh,)] + assert cache.get("alpha") == fresh + + +def test_multiple_matches_are_not_cached() -> None: + cache = _AgentMatchCache() + first = _make_match("alpha", host="host-a") + second = _make_match("alpha", host="host-b") + resolver = _RecordingResolver(result=(first, second)) + sender = _RecordingSender(live_agent_ids=frozenset({first.agent_id, second.agent_id})) + + assert _send_to_cached_agent("alpha", cache, resolver, sender) is True + assert sender.calls == [(first, second)] + assert cache.get("alpha") is None + + +def test_no_match_returns_false_and_caches_nothing() -> None: + cache = _AgentMatchCache() + resolver = _RecordingResolver(result=()) + sender = _RecordingSender(live_agent_ids=frozenset()) + + assert _send_to_cached_agent("ghost", cache, resolver, sender) is False + assert resolver.calls == ["ghost"] + assert cache.get("ghost") is None + + +def test_agent_match_cache_put_get_invalidate() -> None: + cache = _AgentMatchCache() + match = _make_match("alpha") + + assert cache.get("alpha") is None + cache.put("alpha", match) + assert cache.get("alpha") == match + cache.invalidate("alpha") + assert cache.get("alpha") is None + cache.invalidate("alpha") From 0d6afa17766cc23198fe2fe35efe17ae8fdc987b Mon Sep 17 00:00:00 2001 From: Weishi Zeng Date: Tue, 16 Jun 2026 19:14:35 -0700 Subject: [PATCH 2/3] agent_discovery: simplify the message cache to match codebase idioms Reworks the caching from the previous commit to follow existing precedents instead of bespoke machinery (per review feedback): - Cache is now a plain module-level dict (_AGENT_LOCATION_CACHE), like mngr's own get_provider_instance _instance_cache, instead of an _AgentMatchCache class. No lock: dict get/set/pop are individually atomic and the compound race is benign for a best-effort cache (mngr's cache is unlocked for the same reason). - Drops the _AgentResolver/_AgentSender MutableModel __call__ wrappers, which only existed to bind mngr_ctx without a closure/partial. The resolve/send seam now follows welcome_resend.py's idiom: typed Callable aliases injected as defaulted parameters on one orchestration helper (_send_message_to_agent), defaulting to plain module-level functions (_resolve_agent/_send_to_agents). - send_message stays a thin wrapper -- _send_message_to_agent(name, msg, mngr_ctx) -- with no global passed in and no wrapper objects constructed. Tests use plain inline def fakes (inline functions are allowed in test files; the ratchet excludes them) instead of MutableModel stubs, plus a real-but-cheap MngrContext fixture (empty MNGR_PROJECT_CONFIG_DIR so no config files load under mngr's pytest guard). Behavior unchanged; 11 unit tests pass, ty/ruff/ratchets clean. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../imbue/system_interface/agent_discovery.py | 117 +++++-------- .../system_interface/agent_discovery_test.py | 156 ++++++++++-------- 2 files changed, 132 insertions(+), 141 deletions(-) diff --git a/apps/system_interface/imbue/system_interface/agent_discovery.py b/apps/system_interface/imbue/system_interface/agent_discovery.py index c74f0e3c2..43cf903c8 100644 --- a/apps/system_interface/imbue/system_interface/agent_discovery.py +++ b/apps/system_interface/imbue/system_interface/agent_discovery.py @@ -3,18 +3,15 @@ from __future__ import annotations import os -import threading from collections.abc import Callable from collections.abc import Sequence from pathlib import Path from loguru import logger as _loguru_logger from pydantic import Field -from pydantic import PrivateAttr from imbue.concurrency_group.concurrency_group import ConcurrencyGroup from imbue.imbue_common.frozen_model import FrozenModel -from imbue.imbue_common.mutable_model import MutableModel from imbue.mngr.api.find import AgentMatch from imbue.mngr.api.find import find_all_agents from imbue.mngr.api.find import find_one_agent @@ -174,100 +171,76 @@ def discover_agents( return agents -class _AgentMatchCache(MutableModel): - """Thread-safe cache of resolved agent locations, keyed by agent name.""" +# Agent name -> its last resolved location. A best-effort cache: concurrent sends +# race benignly on this dict (each get/set/pop is atomic; the worst case is an +# extra re-resolve), so it needs no lock. +_AGENT_LOCATION_CACHE: dict[str, AgentMatch] = {} - _lock: threading.Lock = PrivateAttr(default_factory=threading.Lock) - _matches: dict[str, AgentMatch] = PrivateAttr(default_factory=dict) +ResolveAgentFn = Callable[[str, MngrContext], Sequence[AgentMatch]] +SendToAgentsFn = Callable[[Sequence[AgentMatch], str, MngrContext], MessageResult] - def get(self, agent_name: str) -> AgentMatch | None: - with self._lock: - return self._matches.get(agent_name) - def put(self, agent_name: str, match: AgentMatch) -> None: - with self._lock: - self._matches[agent_name] = match +def _resolve_agent(agent_name: str, mngr_ctx: MngrContext) -> Sequence[AgentMatch]: + """Resolve an agent name to its matching locations via mngr discovery.""" + return find_all_agents( + addresses=(AgentAddress(agent=AgentName(agent_name)),), + filter_all=False, + target_state=None, + mngr_ctx=mngr_ctx, + ) - def invalidate(self, agent_name: str) -> None: - with self._lock: - self._matches.pop(agent_name, None) +def _send_to_agents(matches: Sequence[AgentMatch], message: str, mngr_ctx: MngrContext) -> MessageResult: + """Send a message to a pre-resolved set of agents, auto-starting STOPPED ones.""" + return send_message_to_agents( + mngr_ctx=mngr_ctx, + message_content=message, + agents_to_message=matches, + error_behavior=ErrorBehavior.CONTINUE, + is_start_desired=True, + ) -_AGENT_MATCH_CACHE = _AgentMatchCache() - -class _AgentResolver(MutableModel): - """Resolves an agent name to its mngr locations within a fixed context.""" - - mngr_ctx: MngrContext - - def __call__(self, agent_name: str) -> Sequence[AgentMatch]: - return find_all_agents( - addresses=(AgentAddress(agent=AgentName(agent_name)),), - filter_all=False, - target_state=None, - mngr_ctx=self.mngr_ctx, - ) - - -class _AgentSender(MutableModel): - """Sends a fixed message to pre-resolved agents, auto-starting STOPPED ones.""" - - mngr_ctx: MngrContext - message: str - - def __call__(self, matches: Sequence[AgentMatch]) -> MessageResult: - return send_message_to_agents( - mngr_ctx=self.mngr_ctx, - message_content=self.message, - agents_to_message=matches, - error_behavior=ErrorBehavior.CONTINUE, - is_start_desired=True, - ) - - -def _send_to_cached_agent( +def _send_message_to_agent( agent_name: str, - cache: _AgentMatchCache, - resolve: Callable[[str], Sequence[AgentMatch]], - send_to: Callable[[Sequence[AgentMatch]], MessageResult], + message: str, + mngr_ctx: MngrContext, + *, + resolve: ResolveAgentFn = _resolve_agent, + send: SendToAgentsFn = _send_to_agents, ) -> bool: - """Route a message to ``agent_name``, resolving its location only when needed. + """Send to ``agent_name``, reusing its cached location and re-resolving if stale. On a cache hit the message goes straight to the known location -- no mngr discovery. If that send reaches no agent (the cached location is stale: the agent was destroyed, recreated, or moved hosts), the entry is dropped and the - agent is re-resolved from scratch. A uniquely-resolved location is cached for - reuse; a name that resolves to several agents is never cached, so every such - send still reaches all of them. + agent is re-resolved. A uniquely-resolved location is cached for reuse; a name + matching several agents is never cached, so every such send reaches all of them. + + ``resolve``/``send`` default to the real mngr calls and are injected in tests. """ - cached = cache.get(agent_name) + cached = _AGENT_LOCATION_CACHE.get(agent_name) if cached is not None: - if send_to((cached,)).successful_agents: + if send((cached,), message, mngr_ctx).successful_agents: return True - cache.invalidate(agent_name) - matches = tuple(resolve(agent_name)) + _AGENT_LOCATION_CACHE.pop(agent_name, None) + matches = tuple(resolve(agent_name, mngr_ctx)) if len(matches) == 1: - cache.put(agent_name, matches[0]) - return bool(send_to(matches).successful_agents) + _AGENT_LOCATION_CACHE[agent_name] = matches[0] + return bool(send(matches, message, mngr_ctx).successful_agents) def send_message(agent_name: str, message: str) -> bool: """Send a message to an agent. Returns True on success. - The agent's location is resolved via mngr discovery once and cached, so - repeat messages to the same agent skip discovery and go straight to its host - (this is what mngr's pre-resolved ``agents_to_message`` API enables). STOPPED - agents are auto-started (`is_start_desired=True`) regardless of cache state. + The agent's location is resolved via mngr discovery and cached, so repeat + messages to the same agent skip discovery and go straight to its host (this is + what mngr's pre-resolved ``agents_to_message`` API enables). STOPPED agents are + auto-started (`is_start_desired=True`) regardless of cache state. """ mngr_ctx, cg = _get_mngr_context() try: - return _send_to_cached_agent( - agent_name, - _AGENT_MATCH_CACHE, - _AgentResolver(mngr_ctx=mngr_ctx), - _AgentSender(mngr_ctx=mngr_ctx, message=message), - ) + return _send_message_to_agent(agent_name, message, mngr_ctx) finally: cg.__exit__(None, None, None) diff --git a/apps/system_interface/imbue/system_interface/agent_discovery_test.py b/apps/system_interface/imbue/system_interface/agent_discovery_test.py index 29be7dba0..29267df66 100644 --- a/apps/system_interface/imbue/system_interface/agent_discovery_test.py +++ b/apps/system_interface/imbue/system_interface/agent_discovery_test.py @@ -1,21 +1,22 @@ """Tests for agent_discovery module.""" +from collections.abc import Iterator from collections.abc import Sequence from pathlib import Path import pytest -from pydantic import Field -from imbue.imbue_common.mutable_model import MutableModel from imbue.mngr.api.find import AgentMatch from imbue.mngr.api.message import MessageResult +from imbue.mngr.config.data_types import MngrContext from imbue.mngr.primitives import AgentId from imbue.mngr.primitives import AgentName from imbue.mngr.primitives import HostId from imbue.mngr.primitives import HostName from imbue.mngr.primitives import ProviderInstanceName -from imbue.system_interface.agent_discovery import _AgentMatchCache -from imbue.system_interface.agent_discovery import _send_to_cached_agent +from imbue.system_interface.agent_discovery import _AGENT_LOCATION_CACHE +from imbue.system_interface.agent_discovery import _get_mngr_context +from imbue.system_interface.agent_discovery import _send_message_to_agent from imbue.system_interface.agent_discovery import read_claude_config_dir_from_env_file @@ -118,96 +119,113 @@ def _make_match(name: str, host: str = "host-a") -> AgentMatch: ) -class _RecordingResolver(MutableModel): - """Stand-in for find_all_agents: returns a fixed result and records calls.""" +@pytest.fixture +def mngr_ctx(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Iterator[MngrContext]: + """A real MngrContext rooted at empty tmp dirs (no project config files to load).""" + config_dir = tmp_path / "cfg" + config_dir.mkdir() + monkeypatch.setenv("MNGR_PROJECT_CONFIG_DIR", str(config_dir)) + monkeypatch.setenv("MNGR_HOST_DIR", str(tmp_path / "host")) + ctx, cg = _get_mngr_context() + try: + yield ctx + finally: + cg.__exit__(None, None, None) - result: tuple[AgentMatch, ...] - calls: list[str] = Field(default_factory=list) - def __call__(self, agent_name: str) -> Sequence[AgentMatch]: - self.calls.append(agent_name) - return self.result +@pytest.fixture(autouse=True) +def _clear_agent_location_cache() -> Iterator[None]: + _AGENT_LOCATION_CACHE.clear() + yield + _AGENT_LOCATION_CACHE.clear() -class _RecordingSender(MutableModel): - """Stand-in for send_message_to_agents: reaches only the live agents given.""" +def test_cache_miss_resolves_caches_and_sends(mngr_ctx: MngrContext) -> None: + match = _make_match("alpha") + resolve_calls: list[str] = [] + send_calls: list[tuple[AgentMatch, ...]] = [] - live_agent_ids: frozenset[AgentId] - calls: list[tuple[AgentMatch, ...]] = Field(default_factory=list) + def _resolve(name: str, ctx: MngrContext) -> Sequence[AgentMatch]: + resolve_calls.append(name) + return (match,) - def __call__(self, matches: Sequence[AgentMatch]) -> MessageResult: - self.calls.append(tuple(matches)) - reached = [str(m.agent_name) for m in matches if m.agent_id in self.live_agent_ids] - return MessageResult(successful_agents=reached) + def _send(matches: Sequence[AgentMatch], message: str, ctx: MngrContext) -> MessageResult: + send_calls.append(tuple(matches)) + return MessageResult(successful_agents=[str(m.agent_name) for m in matches]) + assert _send_message_to_agent("alpha", "hi", mngr_ctx, resolve=_resolve, send=_send) is True + assert resolve_calls == ["alpha"] + assert send_calls == [(match,)] + assert _AGENT_LOCATION_CACHE["alpha"] == match -def test_cache_miss_resolves_caches_and_sends() -> None: - cache = _AgentMatchCache() - match = _make_match("alpha") - resolver = _RecordingResolver(result=(match,)) - sender = _RecordingSender(live_agent_ids=frozenset({match.agent_id})) - assert _send_to_cached_agent("alpha", cache, resolver, sender) is True - assert resolver.calls == ["alpha"] - assert sender.calls == [(match,)] - assert cache.get("alpha") == match +def test_cache_hit_skips_resolution(mngr_ctx: MngrContext) -> None: + match = _make_match("alpha") + _AGENT_LOCATION_CACHE["alpha"] = match + resolve_calls: list[str] = [] + send_calls: list[tuple[AgentMatch, ...]] = [] + def _resolve(name: str, ctx: MngrContext) -> Sequence[AgentMatch]: + resolve_calls.append(name) + return (_make_match("alpha"),) -def test_cache_hit_skips_resolution() -> None: - cache = _AgentMatchCache() - match = _make_match("alpha") - cache.put("alpha", match) - resolver = _RecordingResolver(result=(_make_match("alpha"),)) - sender = _RecordingSender(live_agent_ids=frozenset({match.agent_id})) + def _send(matches: Sequence[AgentMatch], message: str, ctx: MngrContext) -> MessageResult: + send_calls.append(tuple(matches)) + return MessageResult(successful_agents=[str(m.agent_name) for m in matches if m.agent_id == match.agent_id]) - assert _send_to_cached_agent("alpha", cache, resolver, sender) is True - assert resolver.calls == [] - assert sender.calls == [(match,)] + assert _send_message_to_agent("alpha", "hi", mngr_ctx, resolve=_resolve, send=_send) is True + assert resolve_calls == [] + assert send_calls == [(match,)] -def test_stale_cache_hit_reresolves_and_updates() -> None: - cache = _AgentMatchCache() +def test_stale_cache_hit_reresolves_and_updates(mngr_ctx: MngrContext) -> None: stale = _make_match("alpha") fresh = _make_match("alpha") - cache.put("alpha", stale) - resolver = _RecordingResolver(result=(fresh,)) - sender = _RecordingSender(live_agent_ids=frozenset({fresh.agent_id})) + _AGENT_LOCATION_CACHE["alpha"] = stale + resolve_calls: list[str] = [] + send_calls: list[tuple[AgentMatch, ...]] = [] - assert _send_to_cached_agent("alpha", cache, resolver, sender) is True - assert resolver.calls == ["alpha"] - assert sender.calls == [(stale,), (fresh,)] - assert cache.get("alpha") == fresh + def _resolve(name: str, ctx: MngrContext) -> Sequence[AgentMatch]: + resolve_calls.append(name) + return (fresh,) + def _send(matches: Sequence[AgentMatch], message: str, ctx: MngrContext) -> MessageResult: + send_calls.append(tuple(matches)) + return MessageResult(successful_agents=[str(m.agent_name) for m in matches if m.agent_id == fresh.agent_id]) -def test_multiple_matches_are_not_cached() -> None: - cache = _AgentMatchCache() + assert _send_message_to_agent("alpha", "hi", mngr_ctx, resolve=_resolve, send=_send) is True + assert resolve_calls == ["alpha"] + assert send_calls == [(stale,), (fresh,)] + assert _AGENT_LOCATION_CACHE["alpha"] == fresh + + +def test_multiple_matches_are_not_cached(mngr_ctx: MngrContext) -> None: first = _make_match("alpha", host="host-a") second = _make_match("alpha", host="host-b") - resolver = _RecordingResolver(result=(first, second)) - sender = _RecordingSender(live_agent_ids=frozenset({first.agent_id, second.agent_id})) + send_calls: list[tuple[AgentMatch, ...]] = [] - assert _send_to_cached_agent("alpha", cache, resolver, sender) is True - assert sender.calls == [(first, second)] - assert cache.get("alpha") is None + def _resolve(name: str, ctx: MngrContext) -> Sequence[AgentMatch]: + return (first, second) + def _send(matches: Sequence[AgentMatch], message: str, ctx: MngrContext) -> MessageResult: + send_calls.append(tuple(matches)) + return MessageResult(successful_agents=[str(m.agent_name) for m in matches]) -def test_no_match_returns_false_and_caches_nothing() -> None: - cache = _AgentMatchCache() - resolver = _RecordingResolver(result=()) - sender = _RecordingSender(live_agent_ids=frozenset()) + assert _send_message_to_agent("alpha", "hi", mngr_ctx, resolve=_resolve, send=_send) is True + assert send_calls == [(first, second)] + assert "alpha" not in _AGENT_LOCATION_CACHE - assert _send_to_cached_agent("ghost", cache, resolver, sender) is False - assert resolver.calls == ["ghost"] - assert cache.get("ghost") is None +def test_no_match_returns_false_and_caches_nothing(mngr_ctx: MngrContext) -> None: + resolve_calls: list[str] = [] -def test_agent_match_cache_put_get_invalidate() -> None: - cache = _AgentMatchCache() - match = _make_match("alpha") + def _resolve(name: str, ctx: MngrContext) -> Sequence[AgentMatch]: + resolve_calls.append(name) + return () + + def _send(matches: Sequence[AgentMatch], message: str, ctx: MngrContext) -> MessageResult: + return MessageResult(successful_agents=[]) - assert cache.get("alpha") is None - cache.put("alpha", match) - assert cache.get("alpha") == match - cache.invalidate("alpha") - assert cache.get("alpha") is None - cache.invalidate("alpha") + assert _send_message_to_agent("ghost", "hi", mngr_ctx, resolve=_resolve, send=_send) is False + assert resolve_calls == ["ghost"] + assert "ghost" not in _AGENT_LOCATION_CACHE From c6218983b93dc2abfe27a7952285ffd848021938 Mon Sep 17 00:00:00 2001 From: Weishi Zeng Date: Tue, 16 Jun 2026 19:46:40 -0700 Subject: [PATCH 3/3] agent_discovery: drop test-harness reference from _send_message_to_agent docstring The docstring described the prod code's current behavior plus a note that resolve/send "are injected in tests" -- prod docstrings shouldn't narrate the test setup. The defaulted Callable params are self-evident from the signature. Co-Authored-By: Claude Opus 4.8 (1M context) --- apps/system_interface/imbue/system_interface/agent_discovery.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/apps/system_interface/imbue/system_interface/agent_discovery.py b/apps/system_interface/imbue/system_interface/agent_discovery.py index 43cf903c8..25497e03c 100644 --- a/apps/system_interface/imbue/system_interface/agent_discovery.py +++ b/apps/system_interface/imbue/system_interface/agent_discovery.py @@ -216,8 +216,6 @@ def _send_message_to_agent( agent was destroyed, recreated, or moved hosts), the entry is dropped and the agent is re-resolved. A uniquely-resolved location is cached for reuse; a name matching several agents is never cached, so every such send reaches all of them. - - ``resolve``/``send`` default to the real mngr calls and are injected in tests. """ cached = _AGENT_LOCATION_CACHE.get(agent_name) if cached is not None: