From 9e1248567487748d2ed7eb1bbad2c54d13e638f0 Mon Sep 17 00:00:00 2001 From: "Ioannis L." <44038245+blitzcrieg1@users.noreply.github.com> Date: Sun, 23 Aug 2026 20:16:53 +0300 Subject: [PATCH] feat(mcp): attach initialize handshake to schema fingerprint events Capture serverInfo.version and tools.listChanged from initialize so operators can separate release churn from digest-only rug-pull shapes. --- .../agentmetry/core/audit/ingest.py | 59 +++++++++++++------ .../agentmetry/core/diagnostics/mcp_schema.py | 51 +++++++++++++++- .../tests/test_mcp_audit_proxy.py | 48 +++++++++++++-- apps/orchestrator/tests/test_mcp_schema.py | 57 +++++++++++++++++- apps/orchestrator/tools/mcp_audit_proxy.py | 41 +++++++++++-- 5 files changed, 226 insertions(+), 30 deletions(-) diff --git a/apps/orchestrator/agentmetry/core/audit/ingest.py b/apps/orchestrator/agentmetry/core/audit/ingest.py index 772e60a..0b20c22 100644 --- a/apps/orchestrator/agentmetry/core/audit/ingest.py +++ b/apps/orchestrator/agentmetry/core/audit/ingest.py @@ -177,7 +177,7 @@ def _get_sink(): return _sink -def _schema_payload_fields(payload: dict[str, Any]) -> tuple[str, str, int, str]: +def _schema_payload_fields(payload: dict[str, Any]) -> tuple[str, str, int, str, str, bool | None]: tool = payload.get("tool") if isinstance(payload.get("tool"), dict) else {} server = str(tool.get("server") or payload.get("server") or "") fingerprint = str(payload.get("schema_fingerprint") or "") @@ -186,7 +186,11 @@ def _schema_payload_fields(payload: dict[str, Any]) -> tuple[str, str, int, str] except (TypeError, ValueError): tool_count = 0 source = str(payload.get("adapter") or "mcp_proxy") - return server, fingerprint, tool_count, source + server_version = str(payload.get("server_version") or "") + list_changed = payload.get("list_changed") + if list_changed is not None and not isinstance(list_changed, bool): + list_changed = None + return server, fingerprint, tool_count, source, server_version, list_changed def build_schema_canonical(payload: dict[str, Any], status: str) -> dict[str, Any]: @@ -199,13 +203,29 @@ def build_schema_canonical(payload: dict[str, Any], status: str) -> dict[str, An from agentmetry.core.audit.atlas import RUG_PULL from agentmetry.core.diagnostics.mcp_schema import server_id - server, fingerprint, tool_count, _source = _schema_payload_fields(payload) + server, fingerprint, tool_count, _source, server_version, list_changed = _schema_payload_fields(payload) outcome = "changed" if status == "changed" else "success" reason = ( "MCP tool schema changed; config may be unchanged (rug-pull candidate)" if status == "changed" else "MCP tool schema observed" ) + mcp_schema: dict[str, Any] = { + "server_id": server_id(server) if server else "", + "fingerprint": fingerprint, + "tool_count": tool_count, + "status": status, + # Only a schema that MOVED is the technique. `new` is the first + # sight of a server and `same` is a quiet reconnect; tagging either + # as a rug pull would put a Defense Evasion label on installing a + # tool. ATT&CK has no id for this at all, which is the clearest + # case in the product for ATLAS existing alongside it. + **({"atlas": dict(RUG_PULL)} if status == "changed" else {}), + } + if server_version: + mcp_schema["server_version"] = server_version + if list_changed is not None: + mcp_schema["list_changed"] = list_changed return { "schema_version": SCHEMA_VERSION, "event_id": str(uuid.uuid4()), @@ -223,18 +243,7 @@ def build_schema_canonical(payload: dict[str, Any], status: str) -> dict[str, An "actor": {"type": "system", "id": "agentmetry", "role": "recorder"}, "action": {"type": "mcp_schema", "outcome": outcome, "reason": reason}, "agent": {"name": "agentmetry", "skill_id": ""}, - "mcp_schema": { - "server_id": server_id(server) if server else "", - "fingerprint": fingerprint, - "tool_count": tool_count, - "status": status, - # Only a schema that MOVED is the technique. `new` is the first - # sight of a server and `same` is a quiet reconnect; tagging either - # as a rug pull would put a Defense Evasion label on installing a - # tool. ATT&CK has no id for this at all, which is the clearest - # case in the product for ATLAS existing alongside it. - **({"atlas": dict(RUG_PULL)} if status == "changed" else {}), - }, + "mcp_schema": mcp_schema, } @@ -260,20 +269,34 @@ async def _ingest_observed_schema(payload: dict[str, Any]) -> dict[str, Any]: record_observation, ) - server, fingerprint, tool_count, source = _schema_payload_fields(payload) + server, fingerprint, tool_count, source, server_version, list_changed = _schema_payload_fields(payload) status = classify_observation(server, fingerprint) canonical = build_schema_canonical(payload, status) if status == "same": # Only the timestamp moves, and nothing alerts on it, so there is # nothing to make durable first. - record_observation(server, fingerprint, tool_count, source=source) + record_observation( + server, + fingerprint, + tool_count, + source=source, + server_version=server_version, + list_changed=list_changed, + ) return canonical get_trail_db().insert(canonical) sink = _get_sink() if sink is None: raise RuntimeError("No audit sinks configured") await sink.emit(canonical) - record_observation(server, fingerprint, tool_count, source=source) + record_observation( + server, + fingerprint, + tool_count, + source=source, + server_version=server_version, + list_changed=list_changed, + ) return canonical diff --git a/apps/orchestrator/agentmetry/core/diagnostics/mcp_schema.py b/apps/orchestrator/agentmetry/core/diagnostics/mcp_schema.py index 11f8c82..dd09288 100644 --- a/apps/orchestrator/agentmetry/core/diagnostics/mcp_schema.py +++ b/apps/orchestrator/agentmetry/core/diagnostics/mcp_schema.py @@ -68,6 +68,33 @@ def server_id(name: str) -> str: return hashlib.sha256(name.encode("utf-8")).hexdigest()[:16] +def parse_initialize_result(result: Any) -> dict[str, Any]: + """Handshake fields that help separate releases from rug pulls. + + ``serverInfo.version`` is attacker-controlled but useful as a benign-churn + handle: digest moved + version moved often means a shipped release; digest + moved + version stable is the shape worth investigating first. + + ``capabilities.tools.listChanged`` records whether the server promised + ``notifications/tools/list_changed``; a gap between that promise and a + silent listing change is only visible if both are captured. + """ + out: dict[str, Any] = {} + if not isinstance(result, dict): + return out + info = result.get("serverInfo") + if isinstance(info, dict): + version = info.get("version") + if version is not None and str(version).strip(): + out["server_version"] = str(version).strip() + caps = result.get("capabilities") + if isinstance(caps, dict): + tools = caps.get("tools") + if isinstance(tools, dict) and "listChanged" in tools: + out["list_changed"] = bool(tools["listChanged"]) + return out + + class ToolsListBuffer: """Accumulate paginated `tools/list` pages until the cursor is exhausted. @@ -110,6 +137,8 @@ class SchemaRecord: observed_at: str previous: str = "" source: str = "" + server_version: str = "" + list_changed: bool | None = None @dataclass @@ -150,12 +179,15 @@ def load_store(path: Path | None = None) -> SchemaStore: for name, rec in block.items(): if not isinstance(rec, dict) or not rec.get("fingerprint"): continue + list_changed = rec.get("list_changed") servers[str(name)] = SchemaRecord( fingerprint=str(rec.get("fingerprint") or ""), tool_count=int(rec.get("tool_count") or 0), observed_at=str(rec.get("observed_at") or ""), previous=str(rec.get("previous") or ""), source=str(rec.get("source") or ""), + server_version=str(rec.get("server_version") or ""), + list_changed=list_changed if isinstance(list_changed, bool) else None, ) return SchemaStore(servers=servers) @@ -170,6 +202,8 @@ def _dump(store: SchemaStore) -> dict[str, Any]: "observed_at": rec.observed_at, "previous": rec.previous, "source": rec.source, + **({"server_version": rec.server_version} if rec.server_version else {}), + **({"list_changed": rec.list_changed} if rec.list_changed is not None else {}), } for name, rec in sorted(store.servers.items()) }, @@ -218,6 +252,8 @@ def record_observation( tool_count: int, *, source: str = "mcp_proxy", + server_version: str = "", + list_changed: bool | None = None, path: Path | None = None, now: str | None = None, ) -> str: @@ -235,6 +271,10 @@ def record_observation( if existing and existing.fingerprint == fp: existing.observed_at = stamp existing.tool_count = tool_count + if server_version: + existing.server_version = server_version + if list_changed is not None: + existing.list_changed = list_changed _write_store(store, target) return "same" previous = existing.fingerprint if existing else "" @@ -244,6 +284,8 @@ def record_observation( observed_at=stamp, previous=previous, source=source, + server_version=server_version, + list_changed=list_changed, ) _write_store(store, target) return "changed" if previous else "new" @@ -262,7 +304,14 @@ def schema_summary_lines(store: SchemaStore | None = None) -> list[str]: ] for name, rec in sorted(store.servers.items()): extra = f", was {rec.previous[:16]}" if rec.previous else "" + version = f", v={rec.server_version}" if rec.server_version else "" + notify = ( + ", listChanged" + if rec.list_changed + else (", no listChanged" if rec.list_changed is False else "") + ) lines.append( - f" schema {name}: {rec.fingerprint[:16]} ({rec.tool_count} tools{extra})" + f" schema {name}: {rec.fingerprint[:16]} " + f"({rec.tool_count} tools{version}{notify}{extra})" ) return lines diff --git a/apps/orchestrator/tests/test_mcp_audit_proxy.py b/apps/orchestrator/tests/test_mcp_audit_proxy.py index d6f3f7d..0dde3a8 100644 --- a/apps/orchestrator/tests/test_mcp_audit_proxy.py +++ b/apps/orchestrator/tests/test_mcp_audit_proxy.py @@ -62,6 +62,18 @@ def test_schema_payload_moves_when_the_description_does(): assert a["schema_fingerprint"] != b["schema_fingerprint"] +def test_schema_payload_carries_initialize_handshake(): + payload = proxy.build_schema_payload( + "postmark", + [{"name": "t", "description": "x"}], + "sess", + server_version="3.1.0", + list_changed=True, + ) + assert payload["server_version"] == "3.1.0" + assert payload["list_changed"] is True + + def test_call_payload_preserves_already_qualified_name(): msg = {"method": "tools/call", "params": {"name": "mcp__x.read", "arguments": {}}} payload = proxy.build_call_payload(msg, "vault_fs", "s") @@ -89,7 +101,7 @@ def test_correlation_env_override(monkeypatch): assert proxy._correlation_id() == "fixed-corr" -async def _drive_stdout(lines, pending, monkeypatch): +async def _drive_stdout(lines, pending, monkeypatch, handshake=None): """Run the real stdout relay over canned server output, capturing ingests.""" sent = [] monkeypatch.setattr(proxy, "post_ingest", lambda payload, **kw: sent.append(payload)) @@ -98,8 +110,9 @@ async def _drive_stdout(lines, pending, monkeypatch): reader.feed_data(json.dumps(line).encode() + b"\n") reader.feed_eof() buf = proxy.ToolsListBuffer() - await proxy._relay_stdout(reader, pending, "postmark", buf) - return sent + hs = handshake if handshake is not None else {} + await proxy._relay_stdout(reader, pending, "postmark", buf, hs) + return sent, hs @pytest.mark.asyncio @@ -116,7 +129,7 @@ async def test_a_failed_page_does_not_make_the_next_listing_look_poisoned(monkey tool_a = {"name": "a", "description": "A", "inputSchema": {"type": "object"}} tool_b = {"name": "b", "description": "B", "inputSchema": {"type": "object"}} pending = {str(i): {"kind": "list", "server": "postmark"} for i in range(1, 5)} - sent = await _drive_stdout( + sent, _ = await _drive_stdout( [ {"jsonrpc": "2.0", "id": 1, "result": {"tools": [tool_a], "nextCursor": "p2"}}, {"jsonrpc": "2.0", "id": 2, "error": {"code": -32000, "message": "dropped"}}, @@ -132,3 +145,30 @@ async def test_a_failed_page_does_not_make_the_next_listing_look_poisoned(monkey from agentmetry.core.diagnostics.mcp_schema import fingerprint_tools assert sent[0]["schema_fingerprint"] == fingerprint_tools([tool_a, tool_b]) + + +@pytest.mark.asyncio +async def test_initialize_handshake_is_attached_to_the_next_tools_list(monkeypatch): + tool = {"name": "a", "description": "A", "inputSchema": {"type": "object"}} + pending = { + "1": {"kind": "init", "server": "postmark"}, + "2": {"kind": "list", "server": "postmark"}, + } + sent, _ = await _drive_stdout( + [ + { + "jsonrpc": "2.0", + "id": 1, + "result": { + "serverInfo": {"name": "postmark", "version": "15.0.0"}, + "capabilities": {"tools": {"listChanged": True}}, + }, + }, + {"jsonrpc": "2.0", "id": 2, "result": {"tools": [tool]}}, + ], + pending, + monkeypatch, + ) + assert len(sent) == 1 + assert sent[0]["server_version"] == "15.0.0" + assert sent[0]["list_changed"] is True diff --git a/apps/orchestrator/tests/test_mcp_schema.py b/apps/orchestrator/tests/test_mcp_schema.py index 9d28f59..1996cd7 100644 --- a/apps/orchestrator/tests/test_mcp_schema.py +++ b/apps/orchestrator/tests/test_mcp_schema.py @@ -23,6 +23,7 @@ classify_observation, fingerprint_tools, load_store, + parse_initialize_result, record_observation, schema_summary_lines, server_id, @@ -79,6 +80,22 @@ def test_server_id_is_opaque_and_stable(): assert len(server_id("github")) == 16 +def test_parse_initialize_result_extracts_version_and_list_changed(): + parsed = parse_initialize_result( + { + "protocolVersion": "2024-11-05", + "serverInfo": {"name": "demo", "version": "1.2.3"}, + "capabilities": {"tools": {"listChanged": True}}, + } + ) + assert parsed == {"server_version": "1.2.3", "list_changed": True} + + +def test_parse_initialize_result_ignores_missing_fields(): + assert parse_initialize_result({}) == {} + assert parse_initialize_result(None) == {} + + @pytest.fixture() def schema_home(tmp_path, monkeypatch): monkeypatch.setattr(settings, "audit_export_path", tmp_path / "audit-forward.jsonl") @@ -116,12 +133,50 @@ def test_digest_moves_when_one_server_changes_and_not_when_order_does(schema_hom def test_summary_names_servers_locally_but_not_descriptions(schema_home): - record_observation("very-secret-internal-tool", fingerprint_tools([_tool()]), 1) + record_observation( + "very-secret-internal-tool", + fingerprint_tools([_tool()]), + 1, + server_version="1.0.0", + list_changed=True, + ) text = "\n".join(schema_summary_lines()) assert "very-secret-internal-tool" in text + assert "v=1.0.0" in text + assert "listChanged" in text assert "Send an email" not in text +@pytest.mark.asyncio +async def test_ingest_carries_initialize_handshake_fields(schema_home, monkeypatch): + monkeypatch.setattr(settings, "audit_export_enabled", True) + monkeypatch.setattr(settings, "audit_ingest_enabled", True) + monkeypatch.setattr(settings, "audit_sink", "file") + monkeypatch.setattr(settings, "audit_db_path", schema_home / "audit.db") + from agentmetry.core.audit.trail_db import reset_trail_db + + reset_trail_db() + reset_ingest_sink_cache() + fp = fingerprint_tools([_tool()]) + payload = { + "source_app": "mcp_proxy", + "adapter": "mcp_audit_proxy", + "event_type": "mcp_schema", + "schema_fingerprint": fp, + "schema_tool_count": 1, + "tool": {"server": "github"}, + "server_version": "2.4.1", + "list_changed": False, + } + event = await ingest_external_event(payload) + assert event["mcp_schema"]["server_version"] == "2.4.1" + assert event["mcp_schema"]["list_changed"] is False + store = load_store() + assert store.servers["github"].server_version == "2.4.1" + assert store.servers["github"].list_changed is False + reset_ingest_sink_cache() + + @pytest.mark.asyncio async def test_ingest_emits_on_change_not_on_reconnect(schema_home, monkeypatch): """A tools/list on every session start must not become a trail flood.""" diff --git a/apps/orchestrator/tools/mcp_audit_proxy.py b/apps/orchestrator/tools/mcp_audit_proxy.py index 8f9d30d..923ccc3 100644 --- a/apps/orchestrator/tools/mcp_audit_proxy.py +++ b/apps/orchestrator/tools/mcp_audit_proxy.py @@ -11,8 +11,9 @@ Correlation: all calls in one proxy process share a per-process session id (override with AGENTMETRY_CORRELATION_ID) — NOT the JSON-RPC request id, which collides across sessions. The JSON-RPC id is used only to match a response to -its request so a server error becomes a tool_failed event, and so a -paginated tools/list can be assembled before it is hashed. +its request so a server error becomes a tool_failed event, so a paginated +tools/list can be assembled before it is hashed, and so an initialize response +can be paired with the next completed tools/list. Redaction: tool arguments are hashed in-process (input_hash); plaintext args never cross the wire to the orchestrator. Tool descriptions are hashed the @@ -42,6 +43,7 @@ from agentmetry.core.diagnostics.mcp_schema import ( # noqa: E402 ToolsListBuffer, fingerprint_tools, + parse_initialize_result, ) # Per-process session id — ties every tool call in this MCP connection together. @@ -63,10 +65,15 @@ def _qualified(server_name: str, tool_name: str) -> str: def build_schema_payload( - server_name: str, tools: list[Any], correlation_id: str + server_name: str, + tools: list[Any], + correlation_id: str, + *, + server_version: str = "", + list_changed: bool | None = None, ) -> dict[str, Any]: """Hash-only ingest of a completed `tools/list`. Descriptions stay here.""" - return { + payload: dict[str, Any] = { "source_app": _source_app(), "adapter": "mcp_audit_proxy", "event_type": "mcp_schema", @@ -75,6 +82,11 @@ def build_schema_payload( "schema_tool_count": len(tools), "tool": {"server": server_name}, } + if server_version: + payload["server_version"] = server_version + if list_changed is not None: + payload["list_changed"] = list_changed + return payload def build_call_payload( @@ -139,6 +151,9 @@ async def _relay_stdin( continue rid = msg.get("id") method = msg.get("method") + if method == "initialize" and rid is not None: + pending[str(rid)] = {"kind": "init", "server": server_name} + continue if method == "tools/list" and rid is not None: pending[str(rid)] = {"kind": "list", "server": server_name} continue @@ -160,6 +175,7 @@ async def _relay_stdout( pending: dict[str, dict[str, str]], server_name: str, list_buf: ToolsListBuffer, + handshake: dict[str, Any], ) -> None: correlation = _correlation_id() while True: @@ -179,6 +195,11 @@ async def _relay_stdout( ctx = pending.pop(str(rid), None) if ctx is None: continue + if ctx.get("kind") == "init": + if not msg.get("error"): + handshake.clear() + handshake.update(parse_initialize_result(msg.get("result"))) + continue if ctx.get("kind") == "list": if msg.get("error"): # Drop the pages already accumulated. Keeping them would let a @@ -189,7 +210,14 @@ async def _relay_stdout( done = list_buf.add_page(msg.get("result")) if done is not None: post_ingest( - build_schema_payload(server_name, done, correlation), quiet=True + build_schema_payload( + server_name, + done, + correlation, + server_version=str(handshake.get("server_version") or ""), + list_changed=handshake.get("list_changed"), + ), + quiet=True, ) continue err_payload = build_error_payload(msg, ctx, correlation) @@ -211,9 +239,10 @@ async def run_proxy(command: list[str], server_name: str) -> int: pending: dict[str, dict[str, str]] = {} list_buf = ToolsListBuffer() + handshake: dict[str, Any] = {} stdin_task = asyncio.create_task(_relay_stdin(proc.stdin, server_name, pending)) stdout_task = asyncio.create_task( - _relay_stdout(proc.stdout, pending, server_name, list_buf) + _relay_stdout(proc.stdout, pending, server_name, list_buf, handshake) ) async def _stderr() -> None: