agent_discovery: cache resolved agent location to skip re-discovery on repeat sends

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.

apps/system_interface/imbue/system_interface/agent_discovery.py

@@ -3,18 +3,22 @@
 from __future__ import annotations
 
 import os
+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 imbue.concurrency_group.concurrency_group import ConcurrencyGroup
 from imbue.imbue_common.frozen_model import FrozenModel
+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 +171,76 @@ def discover_agents(
     return agents
 
 
+# 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] = {}
+
+ResolveAgentFn = Callable[[str, MngrContext], Sequence[AgentMatch]]
+SendToAgentsFn = Callable[[Sequence[AgentMatch], str, MngrContext], MessageResult]
+
+
+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 _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,
+    )
+
+
+def _send_message_to_agent(
+    agent_name: str,
+    message: str,
+    mngr_ctx: MngrContext,
+    *,
+    resolve: ResolveAgentFn = _resolve_agent,
+    send: SendToAgentsFn = _send_to_agents,
+) -> bool:
+    """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. A uniquely-resolved location is cached for reuse; a name
+    matching several agents is never cached, so every such send reaches all of them.
+    """
+    cached = _AGENT_LOCATION_CACHE.get(agent_name)
+    if cached is not None:
+        if send((cached,), message, mngr_ctx).successful_agents:
+            return True
+        _AGENT_LOCATION_CACHE.pop(agent_name, None)
+    matches = tuple(resolve(agent_name, mngr_ctx))
+    if len(matches) == 1:
+        _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.
 
-    STOPPED agents are automatically started before the message is sent
-    (`is_start_desired=True`), so messaging is possible regardless of agent 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:
-        matches = find_all_agents(
-            addresses=(AgentAddress(agent=AgentName(agent_name)),),
-            filter_all=False,
-            target_state=None,
-            mngr_ctx=mngr_ctx,
-        )
-        result = send_message_to_agents(
-            mngr_ctx=mngr_ctx,
-            message_content=message,
-            agents_to_message=matches,
-            error_behavior=ErrorBehavior.CONTINUE,
-            is_start_desired=True,
-        )
+        return _send_message_to_agent(agent_name, message, mngr_ctx)
     finally:
         cg.__exit__(None, None, None)
-    return len(result.successful_agents) > 0
 
 
 def start_agent(agent_name: str) -> None:

apps/system_interface/imbue/system_interface/agent_discovery_test.py

@@ -1,9 +1,22 @@
 """Tests for agent_discovery module."""
 
+from collections.abc import Iterator
+from collections.abc import Sequence
 from pathlib import Path
 
 import pytest
 
+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 _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
 
 
@@ -94,3 +107,125 @@ 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"),
+    )
+
+
+@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)
+
+
+@pytest.fixture(autouse=True)
+def _clear_agent_location_cache() -> Iterator[None]:
+    _AGENT_LOCATION_CACHE.clear()
+    yield
+    _AGENT_LOCATION_CACHE.clear()
+
+
+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, ...]] = []
+
+    def _resolve(name: str, ctx: MngrContext) -> Sequence[AgentMatch]:
+        resolve_calls.append(name)
+        return (match,)
+
+    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_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 _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_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(mngr_ctx: MngrContext) -> None:
+    stale = _make_match("alpha")
+    fresh = _make_match("alpha")
+    _AGENT_LOCATION_CACHE["alpha"] = stale
+    resolve_calls: list[str] = []
+    send_calls: list[tuple[AgentMatch, ...]] = []
+
+    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])
+
+    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")
+    send_calls: list[tuple[AgentMatch, ...]] = []
+
+    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])
+
+    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
+
+
+def test_no_match_returns_false_and_caches_nothing(mngr_ctx: MngrContext) -> None:
+    resolve_calls: list[str] = []
+
+    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 _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