diff --git a/PR.md b/PR.md new file mode 100644 index 00000000..d843eb44 --- /dev/null +++ b/PR.md @@ -0,0 +1,113 @@ +# PR Title + +feat(memory): expose expert memory as an MCP server for external agents + +--- + +## Summary + +Adds a memory **MCP server** (Streamable HTTP at `/mcp/memory`) so external +agents (coding agents, bots, other AI tools) can directly **read / write / +update Octop expert memory**, aligned 1:1 with the in-process +`MemoryService` capabilities. Every write stamps a `source` marker that is +traceable on recall. + +## Why + +Octop experts accumulate rich memory (facts, conversations, decisions), but +today only the Octop dashboard / in-process agent can access it. External +agents that need to reuse that expertise (e.g. a coding agent asking a +business expert's accumulated knowledge) have no way in. This PR exposes the +same memory surface over the standard MCP protocol so any MCP-capable agent +can join the loop. + +## What + +- **New module** `src/octop/infra/agents/memory_mcp.py` — FastMCP server + bound to one expert per connection, plus token auth and header routing. +- **Mount** in `api/app.py` (`build_app`) at `/mcp/memory`, with + `streamable_http` task groups wired into the FastAPI lifespan. +- **Tests** `tests/unit/agents/test_memory_mcp.py` (13 tests). + +### Tools + +| Tool | Purpose | Backing API | +|------|---------|-------------| +| `memory_recall(query, limit=5)` | Recall memories (full pipeline: tokenize → FTS → rerank → dedupe); returns structured snippets + rendered markdown | `recall_for_prompt` | +| `memory_save(content, source, topic?)` | Persist a structured fact directly into the atom/tree (durable, no extraction) | `Memory.store` | +| `memory_capture(content, source, session_id?)` | Write an **L0 raw event** (goes through extraction); visible immediately via `memory_search_raw` | `Memory.add_raw` | +| `memory_search_raw(query, limit=10)` | FTS-search L0 raw events (capture visible before extraction) | `Memory.search_raw` | +| `memory_update(atom_id, new_content, source)` | Deprecate old atom + persist the new fact | `deprecate_atom` + `store` | + +### Expert binding & auth + +- **One connection binds one expert**: endpoint is a single `/mcp/memory`; + the expert is selected at connect time via the `X-Octop-Agent-Id` header — + callers never pass an agent id per tool call (they don't know the id list). +- **Auth**: independent token via `OCTOP_MEMORY_MCP_TOKEN` (fail-closed when + unset). Authorization via `Authorization: Bearer` or `X-Octop-Memory-Token`. + +### raw vs atom (for callers) + +- `memory_capture` → **L0 raw event** (evidence layer), distilled later by + the extraction pipeline (`extract → candidate → promote → atom`). Use it to + record raw conversations/events; the record is visible immediately via + `memory_search_raw` and recallable via `memory_recall` once promoted. +- `memory_save` → **atom/tree directly** (durable, no extraction). Use it + when the fact is already known. + +## Implementation notes + +- Lives in `infra/agents/` with no api-layer dependency: opens the agent + `Memory` instance via `open_memory_kwargs` + `Memory(...)` (workspace + resolved from the agent registry). +- DNS rebinding protection disabled (`TransportSecuritySettings`) because + Octop runs behind a reverse proxy (Host is the public domain, not localhost). +- `streamable_http_path` collapsed to `/` so the endpoint is exactly + `/mcp/memory` (the SDK default `/mcp` would yield `/mcp/memory/mcp`). +- One `FastMCP` per expert, routed by an ASGI dispatcher on the + `X-Octop-Agent-Id` header; missing/unknown agent → 404. + +## Usage example + +```bash +export OCTOP_MEMORY_MCP_TOKEN="" +``` + +```json +{ + "mcpServers": { + "octop-memory": { + "type": "streamable_http", + "url": "http:///mcp/memory/", + "headers": { + "Authorization": "Bearer ", + "X-Octop-Agent-Id": "" + } + } + } +} +``` + +```text +memory_recall(query="what are the key project decisions?") +memory_save(content="the release window is every Tuesday", source="coding-agent", topic="release") +memory_capture(content="user reported: the report panel banner is not rendering", source="review-bot", session_id="review-2026-08-20") +memory_search_raw(query="report panel banner") +memory_update(atom_id="atom_xxx", new_content="updated fact", source="coding-agent") +``` + +## Testing + +- `tests/unit/agents/test_memory_mcp.py` — 13 tests: tool registration, + recall pipeline, capture (raw) semantics, search_raw, update, token + middleware (401 / accept), header routing, 404 unknown agent, unified mount. +- Verified locally by booting the server and exercising the MCP endpoints: + health, 401 without token, `initialize` (binds expert via header), + `tools/list` (5 tools), `tools/call memory_recall`. + +## Checklist + +- [x] No internal/hard-coded environment-specific values in the diff +- [x] `make lint` clean (ruff) +- [x] Unit tests pass diff --git a/src/octop/api/app.py b/src/octop/api/app.py index d30f3599..5f5477f1 100644 --- a/src/octop/api/app.py +++ b/src/octop/api/app.py @@ -248,6 +248,23 @@ async def acme_http01_challenge(token: str) -> PlainTextResponse: ], ) + # 专家记忆 MCP server（对外暴露，独立 token 鉴权，未配置 OCTOP_MEMORY_MCP_TOKEN 时不挂载） + from octop.infra.agents.memory_mcp import mount_memory_mcp + + memory_mcp_managers = mount_memory_mcp(app, server) + if memory_mcp_managers: + from contextlib import AsyncExitStack, asynccontextmanager + + @asynccontextmanager + async def _memory_mcp_lifespan(application: FastAPI): + # streamable_http_app 的 task group 依赖 lifespan，挂载后须手动并入 + async with AsyncExitStack() as stack: + for mgr in memory_mcp_managers: + await stack.enter_async_context(mgr.run()) + yield + + app.router.lifespan_context = _memory_mcp_lifespan + if enable_api_docs: @app.get("/api/docs", include_in_schema=False) diff --git a/src/octop/infra/agents/memory_mcp.py b/src/octop/infra/agents/memory_mcp.py new file mode 100644 index 00000000..931208ec --- /dev/null +++ b/src/octop/infra/agents/memory_mcp.py @@ -0,0 +1,359 @@ +"""Expose Octop expert memory as an MCP server for external agents. + +External agents (coding agents, bots) can read/write Octop expert memory over +MCP (Streamable HTTP), aligned with the in-process ``MemoryService`` +capabilities. Every write stamps a ``source`` marker that can be traced back +on recall. + +Expert binding: the endpoint is a single ``/mcp/memory`` mount; the expert is +selected at connect time via the ``X-Octop-Agent-Id`` header (one connection +binds one expert — the caller never passes an agent id per tool call). + +raw vs atom (aligned with ``MemoryService``): + +* ``memory_capture`` -> ``add_raw``: writes an **L0 raw event**, which goes + through the extraction pipeline (extract -> candidate -> promote -> atom). + Use it to record raw conversations / events. The record is visible + immediately via ``memory_search_raw``; ``memory_recall`` returns it only + after extraction promotes it to an atom. +* ``memory_save`` -> ``store``: persists a structured fact directly into the + canonical atom/tree (durable, no extraction). Use it when you already know + the exact fact to remember. + +Auth: independent token via ``OCTOP_MEMORY_MCP_TOKEN`` (fail-closed when +unset), enforced by the ASGI middleware in ``mount_memory_mcp``. +""" + +from __future__ import annotations + +import logging +import os +from typing import Any + +from mcp.server.fastmcp import FastMCP +from mcp.server.transport_security import TransportSecuritySettings + +from octop.infra.agents.memory_backend import open_memory_kwargs +from octop.infra.server import OctopServer + +logger = logging.getLogger(__name__) + + +def _open_memory(server: OctopServer, agent_id: str) -> Any: + """Open the agent's ``Memory`` instance (sqlite by default, postgres opt-in). + + Mirrors ``api.common.memory_client._open_memory_for_agent`` but stays in + ``infra/`` (no api dependency). Workspace is resolved from the agent + registry, falling back to the Octop default layout. + """ + from harness_memory.core import Memory # noqa: PLC0415 + + runtime = getattr(server, "app_runtime", None) + registry = getattr(runtime, "agent_registry", None) if runtime is not None else None + if registry is not None and hasattr(registry, "resolve_workspace_dir"): + workspace = registry.resolve_workspace_dir(agent_id) + else: + services = server.services + assert services is not None, "server.services required when agent_registry unavailable" + paths = getattr(server, "paths", None) or services.paths + workspace = paths.ensure_agent_workspace(agent_id) + + services = server.services + assert services is not None, "server.services required for memory backend" + row = services.agent_repo.get(agent_id) + cfg: dict[str, Any] = {} + if row is not None and row.config_json: + import json # noqa: PLC0415 + + try: + parsed = json.loads(row.config_json) + if isinstance(parsed, dict): + cfg = parsed + except json.JSONDecodeError: + cfg = {} + + ns, backend, backend_config = open_memory_kwargs( + agent_id=agent_id, + cfg=cfg, + octop_config=services.config, + workspace_dir=workspace, + ) + return Memory(namespace=ns, backend=backend, backend_config=backend_config) + + +def build_memory_mcp(server: OctopServer, agent_id: str) -> FastMCP: + """Build an MCP server bound to one expert (``agent_id`` captured in closure).""" + mcp = FastMCP( + f"octop-memory-{agent_id}", + # Octop runs behind a reverse proxy (Host is the public domain, forwarded + # by nginx), not a localhost dev scenario — the mcp SDK's localhost + # DNS-rebinding protection does not apply and would reject the Host + # with 421 unless the domain is allow-listed. + transport_security=TransportSecuritySettings(enable_dns_rebinding_protection=False), + ) + # Collapse the streamable-HTTP path to "/" so the endpoint is exactly + # /mcp/memory (the default "/mcp" would make it /mcp/memory/mcp). + mcp.settings.streamable_http_path = "/" + + def _memory() -> Any: + return _open_memory(server, agent_id) + + @mcp.tool() + def memory_recall(query: str, limit: int = 5) -> dict[str, Any]: + """Recall memories from this expert (aligned with the in-process recall_inject). + + Runs the full recall pipeline (tokenization -> FTS -> rerank -> dedupe) + and returns structured snippets plus a rendered markdown block ready to + inject into a system prompt. + + Args: + query: free-form question / keywords (pass the whole sentence; the + pipeline tokenizes CJK into n-grams internally). + limit: max number of snippets to return. + """ + from harness_memory.pipeline.recall import recall_for_prompt # noqa: PLC0415 + + memory = _memory() + result = recall_for_prompt(memory, query, limit=limit) + return { + "memories": [ + { + "source_id": s.source_id, + "timestamp": s.timestamp_iso, + "layer": s.layer, + "text": s.text, + } + for s in result.snippets + ], + "count": len(result.snippets), + "rendered": result.rendered, + } + + @mcp.tool() + def memory_save( + content: str, + source: str, + topic: str | None = None, + ) -> dict[str, Any]: + """Persist a structured fact directly (atom/tree, durable, no extraction). + + Use this when you already know the exact fact to remember — it is + immediately recallable via ``memory_recall``. The source marker is + stored in ``metadata.source``. + + Args: + content: the fact to remember. + source: who/what recorded it (e.g. "coding-agent"), for traceability. + topic: optional topic label. + """ + memory = _memory() + node = memory.store(content, topic=topic, metadata={"source": source}) + return {"node_id": node.id, "content": node.content, "source": source} + + @mcp.tool() + def memory_capture(content: str, source: str, session_id: str | None = None) -> dict[str, Any]: + """Record a raw event to L0 (goes through extraction: extract -> candidate -> atom). + + Use this to record raw conversations / events that the extraction + pipeline will later distill into atoms. The record is NOT immediately + recallable via ``memory_recall`` (that reads atoms); query it right + away with ``memory_search_raw``. The source marker is stored in + ``payload.source``. + + Example:: + + memory_capture( + content="user reported: the report panel banner is not rendering", + source="review-bot", + session_id="review-2026-08-20", + ) + # -> {"event_id": "...", "recorded": true, ...} + # later: memory_recall(query="report panel banner not rendering") + + Args: + content: the raw conversation / event text. + source: who/what recorded it, for traceability. + session_id: optional stable session id (e.g. caller name) so the + extraction pipeline can group events by session. + """ + memory = _memory() + raw = memory.add_raw( + content, + event_type="manual", + host="mcp-external", + session_id=session_id, + payload={"source": source}, + ) + return { + "event_id": raw.id, + "source": source, + "recorded": True, + "note": ( + "raw (L0) event recorded; visible now via memory_search_raw, " + "recallable via memory_recall after the extraction pipeline " + "promotes it to an atom" + ), + } + + @mcp.tool() + def memory_search_raw(query: str, limit: int = 10) -> dict[str, Any]: + """FTS-search L0 raw events of this expert (capture visible immediately). + + Unlike ``memory_recall`` (which reads atoms), this searches the raw + event layer, so records written by ``memory_capture`` are visible right + away, before extraction promotes them. + + Args: + query: keywords to match against raw event content. + limit: max number of events to return. + """ + memory = _memory() + events = memory.search_raw(query, limit=limit) + return { + "events": [ + { + "event_id": e.id, + "timestamp": e.timestamp.isoformat(), + "session_id": e.session_id, + "user": e.user, + "source": (e.payload or {}).get("source") if e.payload else None, + "content": e.content, + } + for e in events + ], + "count": len(events), + } + + @mcp.tool() + def memory_update( + atom_id: str, + new_content: str, + source: str, + note: str = "mcp update", + ) -> dict[str, Any]: + """Update a memory: deprecate the old atom and persist the new fact. + + Args: + atom_id: id of the atom to supersede. + new_content: the replacement fact. + source: who/what updated it, for traceability. + note: deprecation note. + """ + memory = _memory() + deprecated = memory.deprecate_atom(atom_id, actor="user", note=note) + node = memory.store(new_content, metadata={"source": source, "supersedes": atom_id}) + return { + "deprecated": deprecated, + "deprecated_atom_id": atom_id, + "new_node_id": node.id, + "source": source, + } + + return mcp + + +def _memory_mcp_token() -> str | None: + """Read the MCP auth token (empty string treated as unconfigured).""" + return (os.environ.get("OCTOP_MEMORY_MCP_TOKEN") or "").strip() or None + + +class _TokenAuthMiddleware: + """ASGI middleware enforcing ``Authorization: Bearer`` or ``X-Octop-Memory-Token``.""" + + def __init__(self, app: Any, token: str) -> None: + self._app = app + self._token = token + + async def __call__(self, scope: dict[str, Any], receive: Any, send: Any) -> None: + if scope.get("type") != "http": + await self._app(scope, receive, send) + return + + headers = { + k.decode("latin-1").lower(): v.decode("latin-1") for k, v in scope.get("headers", []) + } + auth = headers.get("authorization", "") + provided = auth[7:].strip() if auth.startswith("Bearer ") else "" + if not provided: + provided = headers.get("x-octop-memory-token", "").strip() + + if provided != self._token: + body = b'{"error":"unauthorized"}' + await send( + { + "type": "http.response.start", + "status": 401, + "headers": [ + (b"content-type", b"application/json"), + (b"content-length", str(len(body)).encode()), + ], + } + ) + await send({"type": "http.response.body", "body": body}) + return + + await self._app(scope, receive, send) + + +class _AgentRouter: + """ASGI dispatcher routing to the per-expert MCP app by ``X-Octop-Agent-Id`` header.""" + + def __init__(self, mcp_apps: dict[str, Any]) -> None: + self._mcp_apps = mcp_apps + + async def __call__(self, scope: dict[str, Any], receive: Any, send: Any) -> None: + if scope.get("type") != "http": + return # lifespan is wired into the host FastAPI manually; http only here + + headers = { + k.decode("latin-1").lower(): v.decode("latin-1") for k, v in scope.get("headers", []) + } + agent_id = headers.get("x-octop-agent-id", "").strip() + target = self._mcp_apps.get(agent_id) + if target is None: + body = b'{"error":"missing or unknown agent_id (X-Octop-Agent-Id)"}' + await send( + { + "type": "http.response.start", + "status": 404, + "headers": [ + (b"content-type", b"application/json"), + (b"content-length", str(len(body)).encode()), + ], + } + ) + await send({"type": "http.response.body", "body": body}) + return + await target(scope, receive, send) + + +def mount_memory_mcp(app: Any, server: OctopServer) -> list[Any]: + """Mount the memory MCP endpoint at ``/mcp/memory``; the expert is selected + per connection via the ``X-Octop-Agent-Id`` header (one connection binds one + expert; the URL stays uniform and does not leak expert ids). + + Does not mount when ``OCTOP_MEMORY_MCP_TOKEN`` is unset (fail-closed). + Returns the session managers that must be initialized in the host FastAPI + lifespan (``streamable_http_app`` task groups depend on it). + """ + token = _memory_mcp_token() + if token is None: + return [] + + services = server.services + assert services is not None, "server.services required for memory MCP mount" + + managers: list[Any] = [] + mcp_apps: dict[str, Any] = {} + rows = services.agent_repo.list_all(include_disabled=False) + for row in rows: + agent_id = row.agent_id + mcp = build_memory_mcp(server, agent_id) + mcp_apps[agent_id] = mcp.streamable_http_app() + managers.append(mcp._session_manager) + + app.mount("/mcp/memory", _TokenAuthMiddleware(_AgentRouter(mcp_apps), token)) + return managers + + +__all__ = ["build_memory_mcp", "mount_memory_mcp"] diff --git a/tests/unit/agents/test_memory_mcp.py b/tests/unit/agents/test_memory_mcp.py new file mode 100644 index 00000000..1be2924c --- /dev/null +++ b/tests/unit/agents/test_memory_mcp.py @@ -0,0 +1,229 @@ +"""Unit tests for the expert memory MCP server (infra/agents/memory_mcp).""" + +from __future__ import annotations + +from unittest import mock + +import pytest + +from octop.infra.agents import memory_mcp as mm + + +@pytest.fixture +def fake_memory(monkeypatch): + mem = mock.MagicMock() + mem.recall.return_value = [] + node = mock.MagicMock() + node.id = "node1" + node.content = "remember X" + mem.store.return_value = node + mem.add_raw.return_value = mock.MagicMock(id="evt1") + mem.deprecate_atom.return_value = True + monkeypatch.setattr(mm, "_open_memory", lambda server, agent_id: mem) + return mem + + +def _tools(mcp): + return mcp._tool_manager._tools + + +def test_build_binds_agent_id(monkeypatch): + """Tools capture agent_id in the closure; callers never pass it.""" + captured = {} + + def fake_open(server, agent_id): + captured["agent_id"] = agent_id + mem = mock.MagicMock() + mem.store.return_value = mock.MagicMock(id="n1", content="x") + return mem + + monkeypatch.setattr(mm, "_open_memory", fake_open) + mcp = mm.build_memory_mcp(mock.MagicMock(), agent_id="EXPERT42") + _tools(mcp)["memory_save"].fn(content="x", source="s") + assert captured["agent_id"] == "EXPERT42" + + +def test_build_registers_five_tools(fake_memory): + mcp = mm.build_memory_mcp(mock.MagicMock(), "A1") + assert set(_tools(mcp)) == { + "memory_recall", + "memory_save", + "memory_capture", + "memory_update", + "memory_search_raw", + } + + +def test_memory_recall_uses_full_pipeline(fake_memory, monkeypatch): + """memory_recall runs the full recall pipeline (recall_for_prompt).""" + import harness_memory.pipeline.recall as _recall + + class _Snippet: + source_id = "atom-1" + timestamp_iso = "2026-08-19T00:00:00+00:00" + layer = "atom" + text = "billing-migration is the local clone" + + fake_result = mock.MagicMock() + fake_result.snippets = [_Snippet()] + fake_result.rendered = "markdown" + monkeypatch.setattr(_recall, "recall_for_prompt", lambda m, q, limit: fake_result) + + mcp = mm.build_memory_mcp(mock.MagicMock(), "A1") + result = _tools(mcp)["memory_recall"].fn(query="billing-migration", limit=3) + assert result["count"] == 1 + assert result["memories"][0]["text"] == "billing-migration is the local clone" + assert result["rendered"] == "markdown" + + +def test_memory_save_goes_store(fake_memory): + mcp = mm.build_memory_mcp(mock.MagicMock(), "A1") + result = _tools(mcp)["memory_save"].fn(content="remember X", source="coding-agent") + kwargs = fake_memory.store.call_args.kwargs + assert kwargs["topic"] is None + assert kwargs["metadata"] == {"source": "coding-agent"} + assert result["source"] == "coding-agent" + + +def test_memory_capture_goes_add_raw(fake_memory): + mcp = mm.build_memory_mcp(mock.MagicMock(), "A1") + result = _tools(mcp)["memory_capture"].fn( + content="raw conversation", source="review-bot", session_id="review-1" + ) + kwargs = fake_memory.add_raw.call_args.kwargs + assert kwargs["event_type"] == "manual" + assert kwargs["host"] == "mcp-external" + assert kwargs["session_id"] == "review-1" + assert kwargs["payload"] == {"source": "review-bot"} + assert result["recorded"] is True + assert "raw (L0)" in result["note"] + + +def test_memory_search_raw_queries_l0(fake_memory): + class _Evt: + id = "evt1" + timestamp = __import__("datetime").datetime(2026, 8, 19) + session_id = "review-1" + user = "u1" + payload = {"source": "review-bot"} + content = "report panel banner hidden" + + fake_memory.search_raw.return_value = [_Evt()] + mcp = mm.build_memory_mcp(mock.MagicMock(), "A1") + result = _tools(mcp)["memory_search_raw"].fn(query="report panel banner", limit=5) + fake_memory.search_raw.assert_called_once_with("report panel banner", limit=5) + assert result["count"] == 1 + assert result["events"][0]["event_id"] == "evt1" + assert result["events"][0]["source"] == "review-bot" + + +def test_memory_update_deprecates_and_saves(fake_memory): + mcp = mm.build_memory_mcp(mock.MagicMock(), "A1") + result = _tools(mcp)["memory_update"].fn( + atom_id="atom1", new_content="new fact", source="review-bot" + ) + fake_memory.deprecate_atom.assert_called_once_with("atom1", actor="user", note="mcp update") + assert fake_memory.store.call_args.kwargs["metadata"] == { + "source": "review-bot", + "supersedes": "atom1", + } + assert result["deprecated"] is True + + +def _asgi_scope(headers: list[tuple[bytes, bytes]] | None = None) -> dict: + return {"type": "http", "headers": headers or []} + + +@pytest.mark.asyncio +async def test_token_middleware_rejects_bad_token(): + inner_called = False + + async def _inner(scope, receive, send): + nonlocal inner_called + inner_called = True + + mw = mm._TokenAuthMiddleware(_inner, "secret") + sent = [] + scope = _asgi_scope([(b"authorization", b"Bearer wrong")]) + + async def _send(msg): + sent.append(msg) + + await mw(scope, lambda: {}, _send) + assert inner_called is False + assert sent[0]["status"] == 401 + + +@pytest.mark.asyncio +async def test_token_middleware_accepts_bearer(): + inner_called = False + + async def _inner(scope, receive, send): + nonlocal inner_called + inner_called = True + + mw = mm._TokenAuthMiddleware(_inner, "secret") + scope = _asgi_scope([(b"authorization", b"Bearer secret")]) + await mw(scope, lambda: {}, lambda msg: None) + assert inner_called is True + + +def test_mount_fail_closed_without_token(monkeypatch): + monkeypatch.delenv("OCTOP_MEMORY_MCP_TOKEN", raising=False) + app = mock.MagicMock() + assert mm.mount_memory_mcp(app, mock.MagicMock()) == [] + app.mount.assert_not_called() + + +def test_mount_unified_path_with_header_router(monkeypatch): + from types import SimpleNamespace + + monkeypatch.setenv("OCTOP_MEMORY_MCP_TOKEN", "secret") + app = mock.MagicMock() + server = SimpleNamespace( + services=SimpleNamespace( + agent_repo=mock.MagicMock( + list_all=lambda include_disabled: [ + SimpleNamespace(agent_id="A1"), + SimpleNamespace(agent_id="A2"), + ] + ) + ) + ) + managers = mm.mount_memory_mcp(app, server) + assert len(managers) == 2 + # unified path mounted exactly once + app.mount.assert_called_once() + assert app.mount.call_args.args[0] == "/mcp/memory" + + +@pytest.mark.asyncio +async def test_agent_router_routes_by_header(): + """_AgentRouter routes to the right app by X-Octop-Agent-Id header.""" + called = {} + + class _FakeApp: + def __init__(self, aid): + self._aid = aid + + async def __call__(self, scope, receive, send): + called["agent"] = self._aid + + router = mm._AgentRouter({"A1": _FakeApp("A1"), "A2": _FakeApp("A2")}) + scope = {"type": "http", "headers": [(b"x-octop-agent-id", b"A2")]} + await router(scope, lambda: {}, lambda msg: None) + assert called["agent"] == "A2" + + +@pytest.mark.asyncio +async def test_agent_router_404_unknown_agent(): + """Unknown agent_id returns 404.""" + router = mm._AgentRouter({"A1": mock.MagicMock()}) + scope = {"type": "http", "headers": [(b"x-octop-agent-id", b"NOPE")]} + sent = [] + + async def _send(msg): + sent.append(msg) + + await router(scope, lambda: {}, _send) + assert sent[0]["status"] == 404