From a52a78af675392d37e2712a7def2dffb9f9e08b8 Mon Sep 17 00:00:00 2001 From: Georgis Andonis Date: Thu, 20 Aug 2026 14:47:22 +0300 Subject: [PATCH 01/16] =?UTF-8?q?test(mcp):=20dual-major=20test=20setup=20?= =?UTF-8?q?=E2=80=94=20run=20the=20MCP=20suite=20against=20SDK=20v1=20and?= =?UTF-8?q?=20v2?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The suite now runs twice in CI: the existing `tests` matrix stays on mcp>=1.26,<2, and a new `tests-mcp-v2` job swaps in mcp>=2,<3 (spec 2026-07-28) and runs posthog/test/mcp. A conftest splits collection by installed major, since each major's seams fail at import on the other. New coverage, red until the SDK changes land: - test_v2_mcpserver / test_v2_lowlevel: the v2 adapters driven directly, mirroring the v1 files (capture, errors, intent injection/stripping, identify, report_missing, idempotency, late registration) - test_v2_wire_dual_era: raw JSON-RPC over the real streamable-http app in both protocol eras — stateless topology, envelope identity, no session header on 2026-07-28, conversation-anchored sessions across two fresh instances, and the self-encoded token surviving pods on the legacy era - test_conversation_session: the cross-SDK session derivation contract (byte-for-byte vectors against posthog-js), the minted-shape guard, and v1 anchoring flows - test_no_crash: instrument() degrades to a logged no-op instead of crashing, on both majors Generated-By: PostHog Code Task-Id: ebafcb71-b03b-443d-b40c-d527ed4a04f4 --- .github/workflows/ci.yml | 38 ++ posthog/test/mcp/_helpers.py | 5 + posthog/test/mcp/_helpers_v2.py | 89 +++++ posthog/test/mcp/conftest.py | 29 ++ posthog/test/mcp/test_conversation_session.py | 211 +++++++++++ posthog/test/mcp/test_features_m4.py | 22 +- posthog/test/mcp/test_no_crash.py | 88 +++++ posthog/test/mcp/test_review_fixes.py | 7 +- posthog/test/mcp/test_session_token.py | 2 + posthog/test/mcp/test_units.py | 21 +- posthog/test/mcp/test_v2_lowlevel.py | 249 +++++++++++++ posthog/test/mcp/test_v2_mcpserver.py | 285 +++++++++++++++ posthog/test/mcp/test_v2_wire_dual_era.py | 341 ++++++++++++++++++ 13 files changed, 1373 insertions(+), 14 deletions(-) create mode 100644 posthog/test/mcp/_helpers_v2.py create mode 100644 posthog/test/mcp/conftest.py create mode 100644 posthog/test/mcp/test_conversation_session.py create mode 100644 posthog/test/mcp/test_no_crash.py create mode 100644 posthog/test/mcp/test_v2_lowlevel.py create mode 100644 posthog/test/mcp/test_v2_mcpserver.py create mode 100644 posthog/test/mcp/test_v2_wire_dual_era.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 49eeccddc..99b0cb0d7 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -151,6 +151,44 @@ jobs: run: | pytest --verbose --timeout=30 + tests-mcp-v2: + # The MCP suite again, against MCP Python SDK v2 (spec 2026-07-28). The + # `tests` matrix covers mcp 1.x on every Python version; this lane swaps + # in mcp>=2 (and drops jlowin fastmcp, which pins mcp<2) and runs only + # posthog/test/mcp — conftest.py there splits collection by major. + name: MCP SDK v2 tests (Python ${{ matrix.python-version }}) + runs-on: ubuntu-latest + strategy: + matrix: + python-version: ['3.10', '3.14'] + + steps: + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + with: + fetch-depth: 1 + + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@5fda3b95a4ea91299a34e894583c3862153e4b97 # v7.0.0 + with: + python-version: ${{ matrix.python-version }} + + - name: Install uv + uses: astral-sh/setup-uv@c771a70e6277c0a99b617c7a806ffedaca235ff9 # v9.0.0 + with: + version: "0.11.32" + enable-cache: true + + - name: Install test dependencies with MCP SDK v2 + shell: bash + run: | + UV_PROJECT_ENVIRONMENT=$pythonLocation uv sync --extra test + uv pip uninstall --python $pythonLocation fastmcp + uv pip install --python $pythonLocation 'mcp>=2,<3' + + - name: Run MCP tests against SDK v2 + run: | + pytest posthog/test/mcp --verbose --timeout=30 + mutation-tests: name: Targeted mutation tests runs-on: ubuntu-latest diff --git a/posthog/test/mcp/_helpers.py b/posthog/test/mcp/_helpers.py index 21afc62a9..b4e019524 100644 --- a/posthog/test/mcp/_helpers.py +++ b/posthog/test/mcp/_helpers.py @@ -6,6 +6,11 @@ import asyncio import concurrent.futures +import importlib.metadata + +MCP_MAJOR = int(importlib.metadata.version("mcp").split(".")[0]) +"""Installed MCP SDK major. The suite runs under both 1.x and 2.x in CI; use +this (and ``conftest.collect_ignore``) to scope tests coupled to one major.""" class FakeClient: diff --git a/posthog/test/mcp/_helpers_v2.py b/posthog/test/mcp/_helpers_v2.py new file mode 100644 index 000000000..57d7d67b5 --- /dev/null +++ b/posthog/test/mcp/_helpers_v2.py @@ -0,0 +1,89 @@ +"""Fixtures for the MCP Python SDK v2 (``mcp>=2``) tests. + +Only imported by the ``test_v2_*`` files, which ``conftest.py`` excludes from +collection under mcp 1.x — so this module may assume v2 symbols exist. + +v2 low-level handlers receive ``(ctx, params)`` where ``ctx`` is a +``ServerRequestContext``. The adapters only read a handful of attributes from +it (``protocol_version``, ``session.client_params``, ``request.headers``), all +defensively, so a ``SimpleNamespace`` with the same shape drives the wrapped +handlers without standing up a session or transport. +""" + +from types import SimpleNamespace +from typing import Any, Dict, Optional + + +def fake_ctx( + protocol_version: str = "2026-07-28", + client_name: Optional[str] = "test-client", + client_version: Optional[str] = "9.9.9", + method: str = "tools/call", + headers: Optional[Dict[str, str]] = None, +) -> Any: + """A ``ServerRequestContext``-shaped stand-in, as the v2 runner would build. + + ``client_params`` mirrors v2's snake_case ``InitializeRequestParams`` + (synthesized from the per-request envelope on the modern era, or from the + handshake on the legacy era). + """ + client_params = None + if client_name is not None: + client_params = SimpleNamespace( + client_info=SimpleNamespace(name=client_name, version=client_version), + protocol_version=protocol_version, + ) + session = SimpleNamespace(client_params=client_params) + request = SimpleNamespace(headers=headers) if headers is not None else None + return SimpleNamespace( + session=session, + protocol_version=protocol_version, + method=method, + params=None, + request_id=1, + meta=None, + request=request, + ) + + +# --- wire-level helpers (test_v2_wire_dual_era) -------------------------------- + +MODERN_PROTOCOL_VERSION = "2026-07-28" +LEGACY_PROTOCOL_VERSION = "2025-11-25" + + +def modern_meta( + client_name: str = "wire-client", client_version: str = "1.2.3" +) -> Dict[str, Any]: + """The 2026-07-28 per-request ``_meta`` envelope (protocol version and + client capabilities are required; client info is a SHOULD).""" + return { + "io.modelcontextprotocol/protocolVersion": MODERN_PROTOCOL_VERSION, + "io.modelcontextprotocol/clientCapabilities": {}, + "io.modelcontextprotocol/clientInfo": { + "name": client_name, + "version": client_version, + }, + } + + +def modern_headers(method: str, tool_name: Optional[str] = None) -> Dict[str, str]: + """Required modern-era HTTP headers: the protocol version must match the + envelope, ``Mcp-Method`` the body method, and ``Mcp-Name`` the tool name + on name-bearing methods (SEP-2243).""" + headers = { + "content-type": "application/json", + "accept": "application/json, text/event-stream", + "mcp-protocol-version": MODERN_PROTOCOL_VERSION, + "mcp-method": method, + } + if tool_name is not None: + headers["mcp-name"] = tool_name + return headers + + +def legacy_headers() -> Dict[str, str]: + return { + "content-type": "application/json", + "accept": "application/json, text/event-stream", + } diff --git a/posthog/test/mcp/conftest.py b/posthog/test/mcp/conftest.py new file mode 100644 index 000000000..10083fffc --- /dev/null +++ b/posthog/test/mcp/conftest.py @@ -0,0 +1,29 @@ +"""Split the MCP test suite by installed MCP SDK major. + +The suite runs twice in CI: once against ``mcp>=1.26,<2`` (the ``tests`` job) +and once against ``mcp>=2,<3`` (the ``tests-mcp-v2`` job). Files coupled to one +major's seams import symbols the other major doesn't ship, so they are excluded +from *collection* (a skip marker can't help — the failure is at import time). +Version-agnostic files (units, truncation, session tokens, PostHogMCP, ids) +collect under both majors. +""" + +from posthog.test.mcp._helpers import MCP_MAJOR + +_V1_ONLY = [ + # module-level `from mcp.server.fastmcp import ...` / v1 request_handlers seams + "test_fastmcp.py", + "test_fastmcp_v2.py", + "test_features_m4.py", + "test_lowlevel.py", + "test_review_fixes.py", +] + +_V2_ONLY = [ + # module-level `from mcp.server.mcpserver import ...` / v2 handler seams + "test_v2_mcpserver.py", + "test_v2_lowlevel.py", + "test_v2_wire_dual_era.py", +] + +collect_ignore = _V2_ONLY if MCP_MAJOR < 2 else _V1_ONLY diff --git a/posthog/test/mcp/test_conversation_session.py b/posthog/test/mcp/test_conversation_session.py new file mode 100644 index 000000000..bbc3d195f --- /dev/null +++ b/posthog/test/mcp/test_conversation_session.py @@ -0,0 +1,211 @@ +"""Conversation-anchored sessions (posthog-js ADR-0004) — runs under both MCP majors. + +The 2026-07-28 revision removed protocol-level sessions, so the only thing that +can carry a session across stateless pods is the agent-echoed +``conversation_id`` handle. ``$session_id`` is derived from it deterministically +and unsalted so two pods that never met still agree — and the derivation is a +cross-SDK contract: posthog-js's ``deriveSessionIdFromConversation`` and this +package's ``derive_session_id_from_conversation`` must match byte for byte. +""" + +import pytest + +from posthog.mcp import derive_session_id_from_conversation +from posthog.mcp._conversation_id import resolve_conversation_id +from posthog.mcp.session import resolve_session_id +from posthog.mcp._internal import MCPAnalyticsData +from posthog.mcp.types import MCPAnalyticsOptions +from posthog.test.mcp._helpers import ( + MCP_MAJOR, + FakeClient, + events_named as _events, + flush_background as _flush, +) + +# A handle shaped exactly like the ids we mint (lowercase uuidv7). +MINTED_SHAPE_HANDLE = "0198d3a7-1111-7222-8333-444455556666" + + +# --- the cross-SDK derivation contract ---------------------------------------- + + +def test_derivation_matches_the_typescript_sdk_byte_for_byte(): + # Vectors computed from posthog-js/packages/mcp `deriveSessionIdFromConversation` + # (src/extensions/ids.ts). If this test fails, the same conversation splits + # into two sessions depending on which SDK served the call — do NOT update + # the expectations without changing both SDKs in lockstep. + assert ( + derive_session_id_from_conversation("conv-123") + == "ses_19c018eaeb9263330c016d3a3a41474b" + ) + assert ( + derive_session_id_from_conversation("0198d3a7-1111-7222-8333-444455556666") + == "ses_57a5f3768678e803a4af9566ca8a661b" + ) + assert ( + derive_session_id_from_conversation("a") + == "ses_8601ec8c0eec655f4ec03fd0b1129ba7" + ) + + +def test_derivation_is_deterministic_and_distinct(): + assert derive_session_id_from_conversation( + "h1" + ) == derive_session_id_from_conversation("h1") + assert derive_session_id_from_conversation( + "h1" + ) != derive_session_id_from_conversation("h2") + + +# --- the minted-shape guard ----------------------------------------------------- + + +def test_echo_of_a_mintable_handle_is_accepted(): + cid, minted = resolve_conversation_id( + True, {"conversation_id": MINTED_SHAPE_HANDLE}, "t", "get_more_tools" + ) + assert minted is False + assert cid == MINTED_SHAPE_HANDLE + + +def test_uppercased_echo_is_lowercased_before_hashing(): + # Some hosts normalise uuids to uppercase; the hash behind $session_id is + # case-sensitive, so the echo must be folded back or it lands in a + # different session than the call that minted it. + cid, minted = resolve_conversation_id( + True, {"conversation_id": MINTED_SHAPE_HANDLE.upper()}, "t", "get_more_tools" + ) + assert minted is False + assert cid == MINTED_SHAPE_HANDLE + + +def test_invented_handle_is_not_anchored(): + # Two unrelated users both sending "conv-1" must NOT share a session, so a + # value we could not have minted is replaced with a fresh handle. + cid, minted = resolve_conversation_id( + True, {"conversation_id": "conv-1"}, "t", "get_more_tools" + ) + assert minted is True + assert cid != "conv-1" + + +# --- session resolution --------------------------------------------------------- + + +async def test_conversation_wins_session_resolution_without_touching_state(): + data = MCPAnalyticsData( + options=MCPAnalyticsOptions(), sink=None, session_id="ses_memory" + ) + before = data.session_id + + resolved = await resolve_session_id( + data, "transport-session", conversation_id=MINTED_SHAPE_HANDLE + ) + + assert resolved == derive_session_id_from_conversation(MINTED_SHAPE_HANDLE) + # per-request, never sticky: the shared state a concurrent chat reads is untouched + assert data.session_id == before + + +async def test_without_conversation_resolution_is_unchanged(): + data = MCPAnalyticsData( + options=MCPAnalyticsOptions(), sink=None, session_id="ses_memory" + ) + resolved = await resolve_session_id(data, "transport-session") + assert resolved != derive_session_id_from_conversation(MINTED_SHAPE_HANDLE) + assert resolved.startswith("ses_") + + +# --- end-to-end on the v1 high-level server (anchoring is not era-gated) -------- + + +@pytest.mark.skipif(MCP_MAJOR != 1, reason="v1 FastMCP server") +async def test_v1_calls_sharing_a_handle_land_in_one_session(): + from mcp.server.fastmcp import FastMCP + from posthog.mcp import instrument + + server = FastMCP("conv-v1") + + @server.tool() + def echo(msg: str) -> str: + return msg + + client = FakeClient() + instrument(server, client, MCPAnalyticsOptions(enable_conversation_id=True)) + + await server._tool_manager.call_tool( + "echo", {"msg": "a", "conversation_id": MINTED_SHAPE_HANDLE, "context": "first"} + ) + await server._tool_manager.call_tool( + "echo", + {"msg": "b", "conversation_id": MINTED_SHAPE_HANDLE, "context": "second"}, + ) + await _flush() + + calls = _events(client, "$mcp_tool_call") + expected = derive_session_id_from_conversation(MINTED_SHAPE_HANDLE) + assert [c["properties"]["$session_id"] for c in calls] == [expected, expected] + assert [c["properties"]["$mcp_conversation_id"] for c in calls] == [ + MINTED_SHAPE_HANDLE, + MINTED_SHAPE_HANDLE, + ] + + +@pytest.mark.skipif(MCP_MAJOR != 1, reason="v1 FastMCP server") +async def test_v1_minted_then_echoed_reuses_one_session(): + from mcp.server.fastmcp import FastMCP + from posthog.mcp import instrument + + server = FastMCP("conv-v1-mint") + + @server.tool() + def echo(msg: str) -> str: + return msg + + client = FakeClient() + instrument(server, client, MCPAnalyticsOptions(enable_conversation_id=True)) + + # First call omits the handle: the SDK mints one and prompts the agent for it. + # convert_result=True is the production shape (FastMCP's call path), giving + # the (content, structured) tuple the prompt-back can ride. + await server._tool_manager.call_tool( + "echo", {"msg": "a", "context": "first"}, convert_result=True + ) + await _flush() + first = _events(client, "$mcp_tool_call")[0]["properties"] + minted = first["$mcp_conversation_id"] + assert minted + + # The agent echoes it back. + await server._tool_manager.call_tool( + "echo", {"msg": "b", "conversation_id": minted, "context": "second"} + ) + await _flush() + + calls = _events(client, "$mcp_tool_call") + expected = derive_session_id_from_conversation(minted) + assert [c["properties"]["$session_id"] for c in calls] == [expected, expected] + + +@pytest.mark.skipif(MCP_MAJOR != 1, reason="v1 FastMCP server") +async def test_v1_feature_off_keeps_transport_sessions(): + from mcp.server.fastmcp import FastMCP + from posthog.mcp import instrument + + server = FastMCP("conv-v1-off") + + @server.tool() + def echo(msg: str) -> str: + return msg + + client = FakeClient() + instrument(server, client) # enable_conversation_id defaults off + + await server._tool_manager.call_tool("echo", {"msg": "a", "context": "x"}) + await _flush() + + props = _events(client, "$mcp_tool_call")[0]["properties"] + assert "$mcp_conversation_id" not in props + assert props["$session_id"] != derive_session_id_from_conversation( + MINTED_SHAPE_HANDLE + ) diff --git a/posthog/test/mcp/test_features_m4.py b/posthog/test/mcp/test_features_m4.py index 0baa56c5f..a434dbfca 100644 --- a/posthog/test/mcp/test_features_m4.py +++ b/posthog/test/mcp/test_features_m4.py @@ -144,24 +144,27 @@ async def test_conversation_id_reused_when_supplied(): client = FakeClient() instrument(server, client, MCPAnalyticsOptions(enable_conversation_id=True)) + # A handle we could have minted (lowercase uuidv7) — an invented value would + # be replaced with a fresh mint, matching posthog-js. + handle = "0198d3a7-1111-7222-8333-444455556666" call_handler = server.request_handlers[mcp_types.CallToolRequest] await call_handler( - _call_request( - "echo", {"msg": "hi", "conversation_id": "conv-123", "context": "x"} - ) + _call_request("echo", {"msg": "hi", "conversation_id": handle, "context": "x"}) ) await _flush() calls = _events(client, "$mcp_tool_call") - assert calls[0]["properties"]["$mcp_conversation_id"] == "conv-123" + assert calls[0]["properties"]["$mcp_conversation_id"] == handle # the injected conversation_id is stripped from captured params (surfaces only as $mcp_conversation_id) args = calls[0]["properties"]["$mcp_parameters"]["request"]["params"]["arguments"] assert "conversation_id" not in args -async def test_conversation_id_not_stamped_when_prompt_back_undeliverable(): - # A tool that errors -> the minted prompt-back can't be delivered, so we must NOT - # record an orphan $mcp_conversation_id the agent never received. +async def test_minted_conversation_id_rides_errored_results(): + # A tool that errors on the FIRST call of a conversation is exactly when the + # agent needs the handle — the low-level decorator converts the raise into an + # isError result whose content still carries the prompt-back, so the minted + # id IS stamped (parity with posthog-js: errored results included on purpose). server = Server("conv-err") @server.list_tools() @@ -189,7 +192,10 @@ async def _ct(name, arguments): assert out.root.isError is True calls = _events(client, "$mcp_tool_call") assert calls and calls[0]["properties"]["$mcp_is_error"] is True - assert "$mcp_conversation_id" not in calls[0]["properties"] + conv_id = calls[0]["properties"].get("$mcp_conversation_id") + assert conv_id + texts = [c.text for c in out.root.content if getattr(c, "type", None) == "text"] + assert any(f"conversation_id={conv_id}" in t for t in texts) async def test_event_properties_applied_to_all_event_types(): diff --git a/posthog/test/mcp/test_no_crash.py b/posthog/test/mcp/test_no_crash.py new file mode 100644 index 000000000..752ef3b13 --- /dev/null +++ b/posthog/test/mcp/test_no_crash.py @@ -0,0 +1,88 @@ +"""``instrument()`` must never crash the host app — runs under both MCP majors. + +The compatibility layer used to import ``mcp.server.fastmcp`` at module scope, +which raises ``ImportError`` on mcp>=2 and propagated straight out of +``instrument()`` into the host (the Python twin of posthog-js#4449, but a crash +instead of a silent no-op). These tests pin the graceful-degradation contract +on whichever major is installed. +""" + +import pytest + +from posthog.mcp import instrument +from posthog.test.mcp._helpers import MCP_MAJOR, FakeClient + + +async def test_unsupported_server_returns_noop_handle(): + handle = instrument(object(), FakeClient()) + # graceful no-op: capture and flush do nothing and do not raise + await handle.capture("anything") + await handle.flush() + + +async def test_unsupported_server_logs_instead_of_raising(): + from posthog.mcp import set_logger + from posthog.mcp.types import MCPAnalyticsOptions + + lines = [] + try: + instrument(object(), FakeClient(), MCPAnalyticsOptions(logger=lines.append)) + finally: + set_logger(None) + assert any("failed to instrument" in line.lower() for line in lines) + + +def test_supported_server_type_detected_on_installed_major(): + """The high-level server class of the installed major must be detected — + the exact regression of posthog-js#4449 was the compatibility gate + rejecting every server of the newer major.""" + from posthog.mcp import _compatibility as compat + + if MCP_MAJOR >= 2: + from mcp.server.mcpserver import MCPServer + + assert compat.is_mcpserver(MCPServer("probe")) is True + assert compat.is_fastmcp(MCPServer("probe")) is False + else: + from mcp.server.fastmcp import FastMCP + + assert compat.is_fastmcp(FastMCP("probe")) is True + assert compat.is_mcpserver(FastMCP("probe")) is False + + +def test_low_level_server_detected_on_installed_major(): + from mcp.server.lowlevel import Server + + from posthog.mcp import _compatibility as compat + + assert compat.is_low_level_server(Server("probe")) is True + assert compat.is_low_level_server(object()) is False + + +@pytest.mark.parametrize( + ("installed", "should_warn"), + [ + ("1.25.0", True), + ("1.26.0", False), + ("1.29.0", False), + ("2.0.0", False), + ("2.9.9", False), + ("3.0.0", True), + ], +) +def test_version_advisory_fires_only_outside_supported_range( + monkeypatch, installed, should_warn +): + import posthog.mcp as mcp_pkg + from posthog.mcp import logger as mcp_logger + + lines = [] + monkeypatch.setattr(mcp_logger, "_active_logger", lines.append) + monkeypatch.setattr( + "importlib.metadata.version", lambda name: installed if name == "mcp" else "0" + ) + + mcp_pkg._warn_if_unsupported_mcp_version() + + warned = any("tested against" in line for line in lines) + assert warned is should_warn diff --git a/posthog/test/mcp/test_review_fixes.py b/posthog/test/mcp/test_review_fixes.py index f5f975728..0b329a050 100644 --- a/posthog/test/mcp/test_review_fixes.py +++ b/posthog/test/mcp/test_review_fixes.py @@ -77,9 +77,12 @@ def summarize(text: str, context: str) -> str: client = FakeClient() instrument(server, client, MCPAnalyticsOptions(enable_conversation_id=True)) + # A handle we could have minted (lowercase uuidv7) — an invented value would + # be replaced with a fresh mint, matching posthog-js. + handle = "0198d3a7-1111-7222-8333-444455556666" result = await server._tool_manager.call_tool( "summarize", - {"text": "hi", "context": "my own context", "conversation_id": "conv-xyz"}, + {"text": "hi", "context": "my own context", "conversation_id": handle}, convert_result=True, ) await _flush() @@ -90,7 +93,7 @@ def summarize(text: str, context: str) -> str: assert any("ctx=my own context" in t for t in text_blocks) calls = _events(client, "$mcp_tool_call") - assert calls and calls[0]["properties"]["$mcp_conversation_id"] == "conv-xyz" + assert calls and calls[0]["properties"]["$mcp_conversation_id"] == handle # --- E: tools/list response + duration, and failure capture ------------------- diff --git a/posthog/test/mcp/test_session_token.py b/posthog/test/mcp/test_session_token.py index cf536e808..82ac19e8c 100644 --- a/posthog/test/mcp/test_session_token.py +++ b/posthog/test/mcp/test_session_token.py @@ -421,6 +421,7 @@ def test_middleware_end_to_end_with_stateless_fastmcp(): on a fresh stateless transport that replays the token is accepted (not rejected) -- i.e. the same session survives across pods.""" pytest.importorskip("starlette.testclient") + pytest.importorskip("mcp.server.fastmcp") # v1-only server; skipped under mcp>=2 from mcp.server.fastmcp import FastMCP from mcp.server.transport_security import TransportSecuritySettings from starlette.testclient import TestClient @@ -488,6 +489,7 @@ def test_instrument_autowires_stateless_mint_no_manual_middleware(): app mint the session token -- the zero-config path. mcp.run() uses the same factory internally, so it's covered too.""" pytest.importorskip("starlette.testclient") + pytest.importorskip("mcp.server.fastmcp") # v1-only server; skipped under mcp>=2 from mcp.server.fastmcp import FastMCP from mcp.server.transport_security import TransportSecuritySettings from starlette.testclient import TestClient diff --git a/posthog/test/mcp/test_units.py b/posthog/test/mcp/test_units.py index 94524705b..3182a38fd 100644 --- a/posthog/test/mcp/test_units.py +++ b/posthog/test/mcp/test_units.py @@ -136,10 +136,21 @@ def test_resolve_conversation_id_skips_missing_capability_tool(): ) -def test_resolve_conversation_id_uses_supplied(): +def test_resolve_conversation_id_uses_supplied_when_mintable_shape(): + # Only an echo of a handle we could have minted (a uuidv7) is accepted — + # the handle becomes $session_id, so an invented value ("conv-1") must not + # anchor two unrelated callers to one session (parity with posthog-js). + handle = "0198d3a7-1111-7222-8333-444455556666" assert resolve_conversation_id( + True, {"conversation_id": handle}, "t", "get_more_tools" + ) == (handle, False) + + +def test_resolve_conversation_id_replaces_invented_values(): + cid, minted = resolve_conversation_id( True, {"conversation_id": "conv-1"}, "t", "get_more_tools" - ) == ("conv-1", False) + ) + assert minted is True and cid != "conv-1" def test_resolve_conversation_id_mints_when_absent(): @@ -149,7 +160,9 @@ def test_resolve_conversation_id_mints_when_absent(): def test_can_inject_prompt_back(): assert can_inject_prompt_back({"content": []}) is True - assert can_inject_prompt_back({"content": [], "isError": True}) is False + # Errored results carry the prompt-back on purpose: a first-call failure is + # exactly when the agent needs the handle (parity with posthog-js). + assert can_inject_prompt_back({"content": [], "isError": True}) is True assert can_inject_prompt_back({"content": "not a list"}) is False assert can_inject_prompt_back("not a dict") is False @@ -160,7 +173,7 @@ def test_inject_prompt_back_appends_block(): def test_inject_prompt_back_noop_when_not_injectable(): - result = {"isError": True, "content": []} + result = {"content": "not a list"} assert inject_prompt_back(result, "conv-9") is result diff --git a/posthog/test/mcp/test_v2_lowlevel.py b/posthog/test/mcp/test_v2_lowlevel.py new file mode 100644 index 000000000..baec19986 --- /dev/null +++ b/posthog/test/mcp/test_v2_lowlevel.py @@ -0,0 +1,249 @@ +"""End-to-end tests for the MCP Python SDK v2 low-level ``Server`` adapter. + +v2 replaced the public ``request_handlers`` dict (keyed by request class) with +constructor-injected handlers stored in ``_request_handlers`` (keyed by method +string) behind ``add_request_handler``/``get_request_handler``. Handlers take +``(ctx, params)`` and raise through to JSON-RPC errors instead of auto- +converting to ``is_error`` results. +""" + +import pytest + +import mcp.types as mcp_types +from mcp.server.lowlevel import Server + +from posthog.mcp import instrument +from posthog.mcp.types import MCPAnalyticsOptions +from posthog.test.mcp._helpers import ( + FakeClient, + events_named as _events, + flush_background as _flush, +) +from posthog.test.mcp._helpers_v2 import fake_ctx + + +def make_server(): + async def on_call_tool(ctx, params): + if params.name == "boom": + raise ValueError("explode") + if params.name == "soft-fail": + return mcp_types.CallToolResult( + content=[mcp_types.TextContent(type="text", text="failed politely")], + is_error=True, + ) + args = params.arguments or {} + return mcp_types.CallToolResult( + content=[ + mcp_types.TextContent( + type="text", text=str(args.get("a", 0) + args.get("b", 0)) + ) + ] + ) + + async def on_list_tools(ctx, params): + return mcp_types.ListToolsResult( + tools=[ + mcp_types.Tool( + name="add", + description="Add two numbers", + input_schema={ + "type": "object", + "properties": { + "a": {"type": "integer"}, + "b": {"type": "integer"}, + }, + "required": ["a", "b"], + }, + ) + ] + ) + + return Server( + "test-low-v2", + version="1.2.3", + on_call_tool=on_call_tool, + on_list_tools=on_list_tools, + ) + + +async def _call_tool(server, name, arguments, ctx=None): + entry = server.get_request_handler("tools/call") + params = mcp_types.CallToolRequestParams(name=name, arguments=arguments) + return await entry.handler(ctx or fake_ctx(), params) + + +async def _list_tools(server, ctx=None): + entry = server.get_request_handler("tools/list") + return await entry.handler(ctx or fake_ctx(method="tools/list"), None) + + +# --- tools/list -------------------------------------------------------------- + + +async def test_list_tools_injects_optional_context_and_captures(): + server = make_server() + client = FakeClient() + instrument(server, client) + + result = await _list_tools(server) + await _flush() + + add_tool = next(t for t in result.tools if t.name == "add") + # optional on the raw low-level path: this schema is also the call's + # validation schema, so `context` must not become required + assert "context" in add_tool.input_schema["properties"] + assert "context" not in add_tool.input_schema.get("required", []) + + listed = _events(client, "$mcp_tools_list") + assert listed + assert listed[0]["properties"]["$mcp_listed_tool_names"] == ["add"] + assert listed[0]["properties"]["$mcp_server_name"] == "test-low-v2" + + +# --- tools/call -------------------------------------------------------------- + + +async def test_tool_call_captured_with_client_identity(): + server = make_server() + client = FakeClient() + instrument(server, client) + + result = await _call_tool( + server, "add", {"a": 2, "b": 3, "context": "adding for a report"} + ) + await _flush() + + assert result.content[0].text == "5" + + calls = _events(client, "$mcp_tool_call") + assert len(calls) == 1 + props = calls[0]["properties"] + assert props["$mcp_tool_name"] == "add" + assert props["$mcp_intent"] == "adding for a report" + assert props["$mcp_is_error"] is False + assert props["$mcp_client_name"] == "test-client" + assert props["$mcp_protocol_version"] == "2026-07-28" + + +async def test_context_left_in_arguments_for_raw_handlers(): + """On the raw low-level path injected keys are NOT stripped — the schema + advertises them as optional and a ``(name, arguments)`` handler ignores + extra keys.""" + seen = {} + + async def on_call_tool(ctx, params): + seen["arguments"] = dict(params.arguments or {}) + return mcp_types.CallToolResult( + content=[mcp_types.TextContent(type="text", text="ok")] + ) + + server = Server("raw", on_call_tool=on_call_tool) + client = FakeClient() + instrument(server, client) + + await _call_tool( + server, "anything", {"x": 1, "context": "raw handlers see everything"} + ) + await _flush() + + assert seen["arguments"] == {"x": 1, "context": "raw handlers see everything"} + + +async def test_raised_error_is_captured_and_reraised(): + server = make_server() + client = FakeClient() + instrument(server, client) + + with pytest.raises(ValueError): + await _call_tool(server, "boom", {"context": "attempting the risky operation"}) + await _flush() + + calls = _events(client, "$mcp_tool_call") + assert calls and calls[0]["properties"]["$mcp_is_error"] is True + exceptions = _events(client, "$exception") + assert exceptions + assert exceptions[0]["properties"]["$exception_list"][-1]["value"] == "explode" + + +async def test_is_error_result_is_captured(): + server = make_server() + client = FakeClient() + instrument(server, client) + + result = await _call_tool( + server, "soft-fail", {"context": "expecting a polite failure"} + ) + await _flush() + + assert result.is_error is True + calls = _events(client, "$mcp_tool_call") + assert calls and calls[0]["properties"]["$mcp_is_error"] is True + + +async def test_late_registration_is_wrapped(): + """Handlers registered via ``add_request_handler`` *after* ``instrument()`` + must still be wrapped (the JS #4449 lesson: adapters that hand over a bare + server and register handlers afterwards).""" + server = Server("late-reg") # no handlers yet + client = FakeClient() + instrument(server, client) + + async def late_call_tool(ctx, params): + return mcp_types.CallToolResult( + content=[mcp_types.TextContent(type="text", text="late ok")] + ) + + server.add_request_handler( + "tools/call", mcp_types.CallToolRequestParams, late_call_tool + ) + + result = await _call_tool(server, "anything", {"context": "late registration"}) + await _flush() + + assert result.content[0].text == "late ok" + calls = _events(client, "$mcp_tool_call") + assert len(calls) == 1 + assert calls[0]["properties"]["$mcp_tool_name"] == "anything" + + +async def test_initialize_and_session_reuse_across_calls(): + server = make_server() + client = FakeClient() + instrument(server, client) + + await _call_tool(server, "add", {"a": 1, "b": 1, "context": "first"}) + await _call_tool(server, "add", {"a": 2, "b": 2, "context": "second"}) + await _flush() + + assert len(_events(client, "$mcp_initialize")) == 1 + calls = _events(client, "$mcp_tool_call") + session_ids = {c["properties"]["$session_id"] for c in calls} + assert len(session_ids) == 1 + + +async def test_instrument_is_idempotent(): + server = make_server() + client = FakeClient() + instrument(server, client) + wrapped = server.get_request_handler("tools/call").handler + instrument(server, client) + assert server.get_request_handler("tools/call").handler is wrapped + + +async def test_report_missing_appends_virtual_tool(): + server = make_server() + client = FakeClient() + instrument(server, client, MCPAnalyticsOptions(report_missing=True)) + + result = await _list_tools(server) + names = [t.name for t in result.tools] + assert "get_more_tools" in names + + call_result = await _call_tool( + server, "get_more_tools", {"context": "need an email tool"} + ) + await _flush() + + assert call_result.is_error is False + missing = _events(client, "$mcp_missing_capability") + assert missing and missing[0]["properties"]["$mcp_intent"] == "need an email tool" diff --git a/posthog/test/mcp/test_v2_mcpserver.py b/posthog/test/mcp/test_v2_mcpserver.py new file mode 100644 index 000000000..37306fdae --- /dev/null +++ b/posthog/test/mcp/test_v2_mcpserver.py @@ -0,0 +1,285 @@ +"""End-to-end tests for the MCP Python SDK v2 high-level server adapter +(``mcp.server.mcpserver.MCPServer``, the renamed FastMCP). + +Mirrors ``test_fastmcp.py``; drives the wrapped seams directly — the low-level +``_request_handlers`` entries for tools/list and tools/call — with a fake +``ServerRequestContext``, exactly as the v2 runner would invoke them. +""" + +import mcp.types as mcp_types +from mcp.server.mcpserver import MCPServer + +from posthog.mcp import instrument +from posthog.mcp.types import MCPAnalyticsOptions, UserIdentity +from posthog.test.mcp._helpers import ( + FakeClient, + events_named as _events, + flush_background as _flush, +) +from posthog.test.mcp._helpers_v2 import fake_ctx + + +def make_server(): + server = MCPServer("test-server-v2") + + @server.tool() + def add(a: int, b: int) -> int: + return a + b + + @server.tool() + def boom() -> str: + raise ValueError("explode") + + return server + + +async def _list_tools(server, ctx=None): + entry = server._lowlevel_server.get_request_handler("tools/list") + return await entry.handler(ctx or fake_ctx(method="tools/list"), None) + + +async def _call_tool(server, name, arguments, ctx=None): + entry = server._lowlevel_server.get_request_handler("tools/call") + params = mcp_types.CallToolRequestParams(name=name, arguments=arguments) + return await entry.handler(ctx or fake_ctx(), params) + + +# --- tools/list -------------------------------------------------------------- + + +async def test_list_tools_injects_context_and_captures(): + server = make_server() + client = FakeClient() + instrument(server, client) + + result = await _list_tools(server) + await _flush() + + add_tool = next(t for t in result.tools if t.name == "add") + assert "context" in add_tool.input_schema["properties"] + assert "context" in add_tool.input_schema["required"] + + listed = _events(client, "$mcp_tools_list") + assert listed + assert set(listed[0]["properties"]["$mcp_listed_tool_names"]) == {"add", "boom"} + assert listed[0]["properties"]["$mcp_client_name"] == "test-client" + assert listed[0]["properties"]["$mcp_protocol_version"] == "2026-07-28" + + +async def test_context_injection_can_be_disabled(): + server = make_server() + client = FakeClient() + instrument(server, client, MCPAnalyticsOptions(context=False)) + + result = await _list_tools(server) + add_tool = next(t for t in result.tools if t.name == "add") + assert "context" not in add_tool.input_schema.get("properties", {}) + + +# --- tools/call -------------------------------------------------------------- + + +async def test_tool_call_captures_intent_and_strips_context(): + server = make_server() + client = FakeClient() + instrument(server, client) + + received = {} + original_add = server._tool_manager.get_tool("add").fn + + def spy_add(a: int, b: int) -> int: + received["args"] = {"a": a, "b": b} + return original_add(a, b) + + server._tool_manager.get_tool("add").fn = spy_add + + result = await _call_tool( + server, + "add", + {"a": 2, "b": 3, "context": "summing two numbers for the user's report"}, + ) + await _flush() + + # the tool executed cleanly and the injected `context` never reached it + assert result.is_error is False + assert received["args"] == {"a": 2, "b": 3} + + calls = _events(client, "$mcp_tool_call") + assert len(calls) == 1 + props = calls[0]["properties"] + assert props["$mcp_tool_name"] == "add" + assert props["$mcp_intent"] == "summing two numbers for the user's report" + assert props["$mcp_intent_source"] == "context_parameter" + assert props["$mcp_is_error"] is False + assert props["$mcp_client_name"] == "test-client" + assert props["$mcp_client_version"] == "9.9.9" + assert props["$mcp_protocol_version"] == "2026-07-28" + assert "$mcp_duration_ms" in props + # context is stripped from captured parameters too + assert "context" not in props["$mcp_parameters"]["request"]["params"]["arguments"] + + +async def test_tool_owning_context_keeps_it(): + server = make_server() + + received = {} + + @server.tool() + def search(query: str, context: str) -> str: + received["context"] = context + return f"{query}!" + + client = FakeClient() + instrument(server, client) + + listed = await _list_tools(server) + search_tool = next(t for t in listed.tools if t.name == "search") + # the tool's own `context` parameter is not clobbered by injection + assert search_tool.input_schema["properties"]["context"]["type"] == "string" + + await _call_tool(server, "search", {"query": "q", "context": "the real argument"}) + await _flush() + + assert received["context"] == "the real argument" + + +async def test_tool_call_error_is_captured_and_converted(): + server = make_server() + client = FakeClient() + instrument(server, client) + + # MCPServer converts the raise into CallToolResult(is_error=True) *outside* + # the wrapped ToolManager seam, so the client still gets the error result... + result = await _call_tool( + server, "boom", {"context": "attempting the risky operation"} + ) + await _flush() + assert result.is_error is True + + # ...and the wrapper saw the raise: error event + $exception sibling. + calls = _events(client, "$mcp_tool_call") + assert calls and calls[0]["properties"]["$mcp_is_error"] is True + exceptions = _events(client, "$exception") + assert exceptions + exception_list = exceptions[0]["properties"]["$exception_list"] + assert exception_list[-1]["value"] == "explode" + + +async def test_initialize_emitted_once_per_session(): + server = make_server() + client = FakeClient() + instrument(server, client) + + await _call_tool( + server, "add", {"a": 1, "b": 1, "context": "first call to warm up"} + ) + await _call_tool( + server, "add", {"a": 2, "b": 2, "context": "second call for the total"} + ) + await _flush() + + assert len(_events(client, "$mcp_initialize")) == 1 + assert len(_events(client, "$mcp_tool_call")) == 2 + + +async def test_identify_sets_distinct_id_and_groups(): + server = make_server() + client = FakeClient() + instrument( + server, + client, + MCPAnalyticsOptions( + identify=lambda request, extra: UserIdentity( + distinct_id="user_42", + properties={"plan": "pro"}, + groups={"organization": "org_7"}, + ) + ), + ) + + await _call_tool( + server, "add", {"a": 1, "b": 2, "context": "checking identity flows through"} + ) + await _flush() + + calls = _events(client, "$mcp_tool_call") + assert calls[0]["distinct_id"] == "user_42" + assert calls[0]["properties"]["$groups"] == {"organization": "org_7"} + assert _events(client, "$identify") + + +async def test_report_missing_advertises_and_captures(): + server = make_server() + client = FakeClient() + instrument(server, client, MCPAnalyticsOptions(report_missing=True)) + + listed = await _list_tools(server) + names = [t.name for t in listed.tools] + assert "get_more_tools" in names + + result = await _call_tool( + server, "get_more_tools", {"context": "need a tool to send emails"} + ) + await _flush() + + assert result.is_error is False + missing = _events(client, "$mcp_missing_capability") + assert missing + assert missing[0]["properties"]["$mcp_intent"] == "need a tool to send emails" + + +async def test_instrument_is_idempotent(): + server = make_server() + client = FakeClient() + instrument(server, client) + wrapped_call = server._tool_manager.call_tool + instrument(server, client) + assert server._tool_manager.call_tool is wrapped_call # not double-wrapped + + +async def test_tool_error_reraise_preserved_for_mcp_layer(): + """The wrapper re-raises: MCPServer's own conversion still yields the same + error text an uninstrumented server produces (analytics never mutates the + tool path).""" + bare = make_server() + bare_entry = bare._lowlevel_server.get_request_handler("tools/call") + bare_result = await bare_entry.handler( + fake_ctx(), mcp_types.CallToolRequestParams(name="boom", arguments={}) + ) + + instrumented = make_server() + instrument(instrumented, FakeClient()) + result = await _call_tool(instrumented, "boom", {"context": "same failure"}) + await _flush() + + assert result.is_error is True + assert result.content[0].text == bare_result.content[0].text + + +async def test_late_registered_tool_is_covered(): + server = make_server() + client = FakeClient() + instrument(server, client) + + @server.tool() + def late(x: int) -> int: + return x * 2 + + result = await _call_tool(server, "late", {"x": 21, "context": "doubling a number"}) + await _flush() + + assert result.is_error is False + calls = _events(client, "$mcp_tool_call") + assert calls and calls[0]["properties"]["$mcp_tool_name"] == "late" + + +async def test_anonymous_events_do_not_create_person_profiles(): + server = make_server() + client = FakeClient() + instrument(server, client) + + await _call_tool(server, "add", {"a": 1, "b": 1, "context": "anonymous call"}) + await _flush() + + calls = _events(client, "$mcp_tool_call") + assert calls[0]["properties"]["$process_person_profile"] is False diff --git a/posthog/test/mcp/test_v2_wire_dual_era.py b/posthog/test/mcp/test_v2_wire_dual_era.py new file mode 100644 index 000000000..a6439df38 --- /dev/null +++ b/posthog/test/mcp/test_v2_wire_dual_era.py @@ -0,0 +1,341 @@ +"""Wire-level dual-era gate for MCP Python SDK v2 (analog of the posthog-js +``harness/dual-era`` matrix, scaled to pytest). + +Drives a real ``streamable_http_app`` with raw JSON-RPC over httpx's +ASGITransport — deliberately not the SDK ``Client``, which negotiates the +legacy era and would never exercise 2026-07-28. Two lanes: + +* **legacy** (2025-11-25): initialize handshake, then tools/list + tools/call. +* **modern** (2026-07-28): no initialize; every request carries the reserved + ``_meta`` envelope and the ``MCP-Protocol-Version``/``Mcp-Method``/``Mcp-Name`` + headers. + +The stateless topology (``stateless_http=True``, fresh transport per request) +is the one the new spec is built around and the one that broke the JS SDK in +posthog-js#4449 — so it is the only topology tested here. +""" + +import json +import re +from contextlib import asynccontextmanager + +import httpx +from mcp.server.mcpserver import MCPServer + +from posthog.mcp import ( + decode_session_id, + derive_session_id_from_conversation, + instrument, +) +from posthog.mcp.types import MCPAnalyticsOptions +from posthog.test.mcp._helpers import ( + FakeClient, + events_named as _events, + flush_background as _flush, +) +from posthog.test.mcp._helpers_v2 import ( + LEGACY_PROTOCOL_VERSION, + MODERN_PROTOCOL_VERSION, + legacy_headers, + modern_headers, + modern_meta, +) + + +def make_server(): + server = MCPServer("wire-server") + + @server.tool() + def add(a: int, b: int) -> int: + return a + b + + @server.tool() + def boom() -> str: + raise ValueError("wire explode") + + return server + + +@asynccontextmanager +async def wire(server): + """The instrumented server as a live HTTP surface, without binding a port.""" + app = server.streamable_http_app(json_response=True, stateless_http=True) + async with server.session_manager.run(): + transport = httpx.ASGITransport(app=app) + async with httpx.AsyncClient( + transport=transport, base_url="http://127.0.0.1:8000" + ) as http: + yield http + + +def rpc(method, params, request_id=1): + return {"jsonrpc": "2.0", "id": request_id, "method": method, "params": params} + + +async def modern_call(http, name, arguments, request_id=1): + body = rpc( + "tools/call", + {"name": name, "arguments": arguments, "_meta": modern_meta()}, + request_id, + ) + return await http.post( + "/mcp", json=body, headers=modern_headers("tools/call", name) + ) + + +# --- modern era (2026-07-28) --------------------------------------------------- + + +async def test_modern_tool_call_captured_with_envelope_identity(): + server = make_server() + client = FakeClient() + instrument(server, client) + + async with wire(server) as http: + response = await modern_call( + http, "add", {"a": 2, "b": 3, "context": "adding on the wire"} + ) + await _flush() + + assert response.status_code == 200 + payload = response.json() + assert payload["result"]["isError"] is False + assert payload["result"]["content"][0]["text"] == "5" + # 2026-07-28 removed protocol sessions: the header must not come back + assert response.headers.get("mcp-session-id") is None + + calls = _events(client, "$mcp_tool_call") + assert len(calls) == 1 + props = calls[0]["properties"] + assert props["$mcp_tool_name"] == "add" + assert props["$mcp_intent"] == "adding on the wire" + assert props["$mcp_client_name"] == "wire-client" + assert props["$mcp_client_version"] == "1.2.3" + assert props["$mcp_protocol_version"] == MODERN_PROTOCOL_VERSION + + +async def test_modern_tools_list_advertises_injected_params(): + server = make_server() + client = FakeClient() + instrument(server, client, MCPAnalyticsOptions(enable_conversation_id=True)) + + async with wire(server) as http: + body = rpc("tools/list", {"_meta": modern_meta()}) + response = await http.post( + "/mcp", json=body, headers=modern_headers("tools/list") + ) + await _flush() + + assert response.status_code == 200 + tools = {t["name"]: t for t in response.json()["result"]["tools"]} + assert "context" in tools["add"]["inputSchema"]["properties"] + assert "conversation_id" in tools["add"]["inputSchema"]["properties"] + + listed = _events(client, "$mcp_tools_list") + assert listed and set(listed[0]["properties"]["$mcp_listed_tool_names"]) == { + "add", + "boom", + } + + +async def test_modern_error_is_captured_and_returned(): + server = make_server() + client = FakeClient() + instrument(server, client) + + async with wire(server) as http: + response = await modern_call(http, "boom", {"context": "expected to fail"}) + await _flush() + + assert response.status_code == 200 + assert response.json()["result"]["isError"] is True + + calls = _events(client, "$mcp_tool_call") + assert calls and calls[0]["properties"]["$mcp_is_error"] is True + assert _events(client, "$exception") + + +async def test_modern_conversation_anchors_session_across_instances(): + """The cross-pod contract (ADR-0004): two fresh server instances that never + shared state agree on ``$session_id`` through the agent-echoed handle alone. + On 2026-07-28 there is no session header, so this is the only correlation.""" + client = FakeClient() + options = MCPAnalyticsOptions(enable_conversation_id=True) + + # Pod A: the agent sends no handle, so the SDK mints one and prompts back. + server_a = make_server() + instrument(server_a, client, options) + async with wire(server_a) as http: + response = await modern_call( + http, "add", {"a": 1, "b": 1, "context": "first call"} + ) + await _flush() + + content = response.json()["result"]["content"] + prompt_back = next( + block["text"] + for block in content + if "conversation_id=" in block.get("text", "") + ) + minted = re.search(r"conversation_id=([0-9a-f-]+)", prompt_back).group(1) + + # Pod B: a different process; the agent echoes the handle. + server_b = make_server() + instrument(server_b, client, options) + async with wire(server_b) as http: + await modern_call( + http, + "add", + {"a": 2, "b": 2, "context": "second call", "conversation_id": minted}, + ) + await _flush() + + calls = _events(client, "$mcp_tool_call") + assert len(calls) == 2 + expected = derive_session_id_from_conversation(minted) + assert [c["properties"]["$session_id"] for c in calls] == [expected, expected] + assert [c["properties"]["$mcp_conversation_id"] for c in calls] == [minted, minted] + + +async def test_modern_result_shape_survives_instrumentation(): + """Alive check: the instrumented result matches the bare server's, minus + analytics-owned additions.""" + bare = make_server() + async with wire(bare) as http: + bare_response = await modern_call(http, "add", {"a": 4, "b": 5}) + + instrumented = make_server() + instrument(instrumented, FakeClient()) + async with wire(instrumented) as http: + response = await modern_call(http, "add", {"a": 4, "b": 5, "context": "alive"}) + await _flush() + + bare_result = bare_response.json()["result"] + result = response.json()["result"] + assert result["content"] == bare_result["content"] + assert result["isError"] is bare_result["isError"] + assert result.get("structuredContent") == bare_result.get("structuredContent") + + +# --- legacy era (2025-11-25) on the v2 SDK ------------------------------------- + + +async def test_legacy_handshake_lane_still_works(): + server = make_server() + client = FakeClient() + instrument(server, client) + + async with wire(server) as http: + init = rpc( + "initialize", + { + "protocolVersion": LEGACY_PROTOCOL_VERSION, + "capabilities": {}, + "clientInfo": {"name": "legacy-probe", "version": "0.1"}, + }, + ) + response = await http.post("/mcp", json=init, headers=legacy_headers()) + assert response.status_code == 200 + session_header = response.headers.get("mcp-session-id") + + await http.post( + "/mcp", + json={"jsonrpc": "2.0", "method": "notifications/initialized"}, + headers=legacy_headers(), + ) + + headers = legacy_headers() + if session_header: + headers["mcp-session-id"] = session_header + call = rpc( + "tools/call", + {"name": "add", "arguments": {"a": 5, "b": 6, "context": "legacy lane"}}, + 2, + ) + response = await http.post("/mcp", json=call, headers=headers) + await _flush() + + assert response.status_code == 200 + assert json.loads(response.content)["result"]["content"][0]["text"] == "11" + + calls = _events(client, "$mcp_tool_call") + assert len(calls) == 1 + assert calls[0]["properties"]["$mcp_tool_name"] == "add" + assert calls[0]["properties"]["$mcp_is_error"] is False + + +async def test_legacy_stateless_token_survives_across_instances(): + """The multi-pod legacy story on the v2 SDK: instrument() auto-wires the + self-encoded ``Mcp-Session-Id`` token onto the stateless app, so client + identity and one ``$session_id`` survive a fresh server instance when the + client replays the header.""" + client = FakeClient() + + # Pod A mints the token on initialize. + server_a = make_server() + instrument(server_a, client) + async with wire(server_a) as http: + init = rpc( + "initialize", + { + "protocolVersion": LEGACY_PROTOCOL_VERSION, + "capabilities": {}, + "clientInfo": {"name": "legacy-probe", "version": "0.1"}, + }, + ) + response = await http.post("/mcp", json=init, headers=legacy_headers()) + token_header = response.headers.get("mcp-session-id") + token = decode_session_id(token_header) + assert token is not None + assert token.client_name == "legacy-probe" + assert token.protocol_version == LEGACY_PROTOCOL_VERSION + + # Pod B — a fresh process that never saw the handshake — gets the replay. + # A compliant legacy client replays both the session header and the + # negotiated MCP-Protocol-Version on every subsequent request. + server_b = make_server() + instrument(server_b, client) + async with wire(server_b) as http: + headers = { + **legacy_headers(), + "mcp-session-id": token_header, + "mcp-protocol-version": LEGACY_PROTOCOL_VERSION, + } + call = rpc( + "tools/call", + {"name": "add", "arguments": {"a": 1, "b": 2, "context": "pod B"}}, + 2, + ) + response = await http.post("/mcp", json=call, headers=headers) + await _flush() + + assert response.status_code == 200 + calls = _events(client, "$mcp_tool_call") + assert len(calls) == 1 + props = calls[0]["properties"] + # identity recovered from the token, not the (nonexistent) handshake + assert props["$mcp_client_name"] == "legacy-probe" + assert props["$mcp_protocol_version"] == LEGACY_PROTOCOL_VERSION + assert props["$session_id"] == token.session_id + + +async def test_modern_rejects_initialize_but_analytics_stays_out_of_it(): + """The SDK itself refuses initialize on a modern-locked flow — analytics + must not change that error contract.""" + server = make_server() + instrument(server, FakeClient()) + + async with wire(server) as http: + # initialize carrying a modern envelope is a protocol error in v2 + body = rpc( + "initialize", + { + "protocolVersion": MODERN_PROTOCOL_VERSION, + "capabilities": {}, + "_meta": modern_meta(), + }, + ) + response = await http.post("/mcp", json=body, headers=legacy_headers()) + + # whatever the SDK answers (error payload), the app must not 500 + assert response.status_code < 500 From 004c745dacdf38004caad76f8a189e62770f9cb7 Mon Sep 17 00:00:00 2001 From: Georgis Andonis Date: Thu, 20 Aug 2026 14:47:39 +0300 Subject: [PATCH 02/16] feat(mcp): support MCP Python SDK v2 and the 2026-07-28 spec revision MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit instrument() now works on mcp 1.x and 2.x: - New _instrument_v2 adapters: the high-level MCPServer (renamed FastMCP) wraps ToolManager.call_tool — the seam every dispatch routes through — and the low-level Server wraps the string-keyed handler registry through the public add_request_handler/get_request_handler API, late registrations included (the posthog-js#4449 lesson). Strip-vs-leave policy per entry point matches the 1.x pair: the high-level path strips injected context/conversation_id (v2 validates against the function signature), the raw low-level path leaves them optional in the schema. - _compatibility probes are import-tolerant and shape-based (posthog-js ADR-0005): the old module-level `from mcp.server.fastmcp import ...` raised ImportError straight out of instrument() on mcp>=2, crashing the host app. Unsupported servers now degrade to a logged no-op handle. - Conversation-anchored sessions (posthog-js ADR-0004): with enable_conversation_id, $session_id is derived deterministically from the agent-echoed conversation_id — the only correlation that survives 2026-07-28's per-request instances. Byte-for-byte parity with @posthog/mcp (new export derive_session_id_from_conversation), the minted-shape (uuidv7) guard so invented handles can't merge unrelated callers, lowercased echoes, and the prompt-back riding errored results so a first-call failure keeps the conversation together. Applies to the 1.x adapters too — anchoring is not era-gated. - Dual-shape attribute reads (isError/is_error, inputSchema/input_schema, clientInfo/client_info); captured payloads dump by_alias so both majors emit the camelCase wire shape. - Version advisory widened to mcp>=1.26,<3. No new runtime dependencies; mcp stays a lazily-imported peer dependency. jlowin's fastmcp keeps the 1.x seams (it pins mcp<2). Generated-By: PostHog Code Task-Id: ebafcb71-b03b-443d-b40c-d527ed4a04f4 --- .sampo/changesets/mcp-sdk-v2-support.md | 5 + examples/mcp_analytics_demo.py | 3 + posthog/mcp/__init__.py | 93 +++- posthog/mcp/_compatibility.py | 48 +- posthog/mcp/_conversation_id.py | 29 +- posthog/mcp/_instrument_fastmcp.py | 27 +- posthog/mcp/_instrument_lowlevel.py | 21 +- posthog/mcp/_instrument_v2.py | 611 ++++++++++++++++++++++++ posthog/mcp/_instrumentation.py | 26 +- posthog/mcp/session.py | 35 +- posthog/mcp/version.py | 2 +- references/public_api_snapshot.txt | 2 + 12 files changed, 839 insertions(+), 63 deletions(-) create mode 100644 .sampo/changesets/mcp-sdk-v2-support.md create mode 100644 posthog/mcp/_instrument_v2.py diff --git a/.sampo/changesets/mcp-sdk-v2-support.md b/.sampo/changesets/mcp-sdk-v2-support.md new file mode 100644 index 000000000..07736cbca --- /dev/null +++ b/.sampo/changesets/mcp-sdk-v2-support.md @@ -0,0 +1,5 @@ +--- +posthog: minor +--- + +feat(mcp): support the MCP Python SDK 2.x and the 2026-07-28 spec revision. `instrument()` now wraps `mcp.server.mcpserver.MCPServer` (the renamed FastMCP) and the v2 low-level `Server` (constructor-injected handlers, string-keyed registry, late `add_request_handler` registrations included), capturing tool calls, tools/list, errors, intent, client identity, and `$mcp_protocol_version` on both protocol eras — SDK 1.x paths are unchanged. Conversation-anchored sessions land as the cross-pod correlation for the stateless era (parity with `@posthog/mcp`): with `enable_conversation_id`, `$session_id` is derived deterministically from the agent-echoed `conversation_id` (new export `derive_session_id_from_conversation`), only handles the SDK could have minted (uuidv7) anchor, and the prompt-back now rides errored results too so a first-call failure keeps the conversation together. `instrument()` also no longer crashes on an unsupported or unrecognized MCP SDK — it degrades to a logged no-op. diff --git a/examples/mcp_analytics_demo.py b/examples/mcp_analytics_demo.py index 5ce20981e..23f30b6cc 100644 --- a/examples/mcp_analytics_demo.py +++ b/examples/mcp_analytics_demo.py @@ -19,6 +19,9 @@ import os import mcp.types as mcp_types + +# MCP SDK 1.x. On mcp>=2 the class moved: `from mcp.server.mcpserver import +# MCPServer` — instrument() works the same on both. from mcp.server.fastmcp import FastMCP from posthog import Posthog diff --git a/posthog/mcp/__init__.py b/posthog/mcp/__init__.py index dbab60cdb..66645cc1e 100644 --- a/posthog/mcp/__init__.py +++ b/posthog/mcp/__init__.py @@ -4,15 +4,23 @@ """PostHog MCP analytics SDK — product analytics for Model Context Protocol servers. -Wrap a Python MCP server (``FastMCP`` or low-level ``mcp.server.Server``) so every -tool call, agent intent, and failure is captured to PostHog as a ``$mcp_*`` event:: +Wrap a Python MCP server so every tool call, agent intent, and failure is +captured to PostHog as a ``$mcp_*`` event. Works with the MCP Python SDK 1.x +*and* 2.x (the 2026-07-28 spec revision) — the high-level server class moved +between majors, but ``instrument()`` is the same:: from posthog import Posthog from posthog.mcp import instrument - from mcp.server.fastmcp import FastMCP + + # MCP SDK 2.x (spec 2026-07-28) + from mcp.server.mcpserver import MCPServer + server = MCPServer("my-server") + + # MCP SDK 1.x + # from mcp.server.fastmcp import FastMCP + # server = FastMCP("my-server") posthog = Posthog("phc_...", host="https://us.i.posthog.com") - server = FastMCP("my-server") analytics = instrument(server, posthog) Install is just ``pip install posthog``. ``instrument()`` needs the MCP SDK at runtime, @@ -43,7 +51,11 @@ ) from .logger import log, set_logger from .posthog_mcp import PostHogMCP -from .session import derive_session_id_from_mcp_session, new_session_id +from .session import ( + derive_session_id_from_conversation, + derive_session_id_from_mcp_session, + new_session_id, +) from .session_token import ( MCP_SESSION_HEADER, SessionTokenPayload, @@ -77,6 +89,10 @@ "PreparedToolCall", "get_more_tools_result", "derive_session_id_from_mcp_session", + # Conversation-anchored sessions: the cross-SDK derivation contract with + # posthog-js (the 2026-07-28 revision has no protocol sessions, so the + # agent-echoed conversation_id is the only cross-pod session carrier). + "derive_session_id_from_conversation", # Self-encoded session tokens for stateless / multi-pod servers. Minted onto # the `Mcp-Session-Id` response header by PostHogMcpStatelessSessionMiddleware # and decoded on every request; codec is exported for custom HTTP layers. @@ -157,10 +173,12 @@ def _resolve_client(posthog_client: Optional[Client]) -> Optional[Client]: def _warn_if_unsupported_mcp_version() -> None: - """The adapters hook private MCP SDK seams (``_tool_manager``, ``_mcp_server``, - ``request_handlers``) tested against ``mcp>=1.26,<2``. Since ``mcp`` is a peer - dependency we don't pin, advise at runtime when the installed version is outside - that range rather than failing hard (older/newer may still mostly work).""" + """The adapters hook private MCP SDK seams (``_tool_manager``, ``_mcp_server`` + / ``_lowlevel_server``, the request-handler registries) tested against + ``mcp>=1.26,<3`` — both the 1.x line and the 2.x line (spec 2026-07-28). + Since ``mcp`` is a peer dependency we don't pin, advise at runtime when the + installed version is outside that range rather than failing hard (older/newer + may still mostly work).""" try: from importlib.metadata import version @@ -168,19 +186,22 @@ def _warn_if_unsupported_mcp_version() -> None: major, minor = (int(p) for p in installed.split(".")[:2]) except Exception: # noqa: BLE001 - never let a version probe break instrument() return - if (major, minor) < (1, 26) or major >= 2: + if (major, minor) < (1, 26) or major >= 3: log( - f"Warning: PostHog MCP analytics is tested against mcp>=1.26,<2; found {installed}. " + f"Warning: PostHog MCP analytics is tested against mcp>=1.26,<3; found {installed}. " "Instrumentation hooks private SDK internals and may behave unexpectedly." ) def _canonical_server(server: Any) -> Any: - """The underlying low-level server for high-level wrappers (official FastMCP and - jlowin's fastmcp 2.0 both expose ``_mcp_server``), else the server itself. Used as - the tracking key so instrumenting a wrapper and its underlying server resolve to - one state instead of two divergent ones (matching the TS SDK).""" - low_level = getattr(server, "_mcp_server", None) + """The underlying low-level server for high-level wrappers (SDK 1.x FastMCP and + jlowin's fastmcp expose ``_mcp_server``; SDK 2.x MCPServer renamed it + ``_lowlevel_server``), else the server itself. Used as the tracking key so + instrumenting a wrapper and its underlying server resolve to one state instead + of two divergent ones (matching the TS SDK).""" + low_level = getattr(server, "_mcp_server", None) or getattr( + server, "_lowlevel_server", None + ) return low_level if low_level is not None else server @@ -197,8 +218,9 @@ def instrument( state instead of double-wrapping. Degrades to a no-op handle on any failure so the host application keeps working. - :param server: A ``FastMCP`` server (official ``mcp.server.fastmcp`` or jlowin's - ``fastmcp`` 2.0) or a low-level ``mcp.server.Server``. + :param server: A high-level server — SDK 1.x ``mcp.server.fastmcp.FastMCP``, + SDK 2.x ``mcp.server.mcpserver.MCPServer``, or jlowin's ``fastmcp.FastMCP`` + — or a low-level ``mcp.server.lowlevel.Server`` (either SDK major). :param posthog_client: A posthog ``Client`` you construct and own (call ``shutdown()`` on exit to flush). Falls back to the global client. :param options: Optional :class:`MCPAnalyticsOptions`. @@ -222,13 +244,20 @@ def instrument( "(PostHogMCP for custom dispatchers works without it.)" ) _warn_if_unsupported_mcp_version() - from ._compatibility import is_fastmcp, is_fastmcp_v2, is_low_level_server - from ._instrument_fastmcp import instrument_fastmcp - from ._instrument_lowlevel import instrument_fastmcp_v2, instrument_low_level key = _canonical_server(server) try: + # Imported inside the try: the adapters touch major-specific modules, and + # an import error must degrade to the no-op handle, not crash the host. + from ._compatibility import ( + is_fastmcp, + is_fastmcp_v2, + is_low_level_server, + is_mcpserver, + uses_v2_handler_registry, + ) + client = _resolve_client(posthog_client) if client is None: log("Warning: no PostHog client available; MCP events will not be sent.") @@ -242,15 +271,31 @@ def instrument( set_server_tracking_data(key, data) if is_fastmcp(server): + from ._instrument_fastmcp import instrument_fastmcp + instrument_fastmcp(server, data) + elif is_mcpserver(server): + from ._instrument_v2 import instrument_mcpserver_v2 + + instrument_mcpserver_v2(server, data) elif is_fastmcp_v2(server): + from ._instrument_lowlevel import instrument_fastmcp_v2 + instrument_fastmcp_v2(server, data) elif is_low_level_server(server): - instrument_low_level(server, data) + if uses_v2_handler_registry(server): + from ._instrument_v2 import instrument_lowlevel_v2 + + instrument_lowlevel_v2(server, data) + else: + from ._instrument_lowlevel import instrument_low_level + + instrument_low_level(server, data) else: raise TypeError( - f"Unsupported server type: {type(server)!r}. Pass a FastMCP (official or jlowin's " - "fastmcp 2.0) or a low-level mcp.server.Server." + f"Unsupported server type: {type(server)!r}. Pass a high-level server " + "(mcp.server.fastmcp.FastMCP on SDK 1.x, mcp.server.mcpserver.MCPServer " + "on SDK 2.x, or jlowin's fastmcp.FastMCP) or a low-level mcp.server.Server." ) # Zero-config stateless minting: wrap the server's ASGI-app factories so a diff --git a/posthog/mcp/_compatibility.py b/posthog/mcp/_compatibility.py index 37b017dfa..cd4d16b27 100644 --- a/posthog/mcp/_compatibility.py +++ b/posthog/mcp/_compatibility.py @@ -2,21 +2,40 @@ # Copyright (c) 2025 MCPcat # Licensed under the MIT License: https://github.com/MCPCat/mcpcat-typescript-sdk/blob/main/LICENSE -"""Detect which kind of MCP server was passed to ``instrument()``.""" +"""Detect which kind of MCP server was passed to ``instrument()``. + +Every probe is import-tolerant: the classes live at different paths per MCP SDK +major (``mcp.server.fastmcp.FastMCP`` on 1.x, ``mcp.server.mcpserver.MCPServer`` +on 2.x), so a probe whose class doesn't exist on the installed major answers +``False`` instead of raising — an unconditional import here is exactly what +would crash ``instrument()`` on the other major. +""" from __future__ import annotations from typing import Any -from mcp.server.fastmcp import FastMCP -from mcp.server.lowlevel import Server as LowLevelServer - def is_fastmcp(server: Any) -> bool: - """The official SDK's high-level server (``mcp.server.fastmcp.FastMCP``).""" + """The MCP SDK 1.x high-level server (``mcp.server.fastmcp.FastMCP``). + The module was renamed in 2.x, so this is False whenever mcp>=2 is installed.""" + try: + from mcp.server.fastmcp import FastMCP + except ImportError: + return False return isinstance(server, FastMCP) +def is_mcpserver(server: Any) -> bool: + """The MCP SDK 2.x high-level server (``mcp.server.mcpserver.MCPServer``, + the renamed FastMCP). False whenever mcp<2 is installed.""" + try: + from mcp.server.mcpserver import MCPServer + except ImportError: + return False + return isinstance(server, MCPServer) + + def is_fastmcp_v2(server: Any) -> bool: """jlowin's standalone FastMCP 2.0 (``fastmcp.FastMCP``), a separate package from the official SDK. Returns False if ``fastmcp`` isn't installed.""" @@ -28,4 +47,23 @@ def is_fastmcp_v2(server: Any) -> bool: def is_low_level_server(server: Any) -> bool: + """The low-level ``mcp.server.lowlevel.Server`` — the import path is the + same on both majors; use :func:`uses_v2_handler_registry` to tell which + handler seam it carries.""" + try: + from mcp.server.lowlevel import Server as LowLevelServer + except ImportError: + return False return isinstance(server, LowLevelServer) + + +def uses_v2_handler_registry(server: Any) -> bool: + """Which major's handler seam a low-level server carries, decided by shape + rather than package version (a la posthog-js ADR-0005): 1.x exposes the + public ``request_handlers`` dict keyed by request class; 2.x replaced it + with ``add_request_handler``/``get_request_handler`` keyed by method string.""" + if hasattr(server, "add_request_handler") and hasattr( + server, "get_request_handler" + ): + return True + return not hasattr(server, "request_handlers") diff --git a/posthog/mcp/_conversation_id.py b/posthog/mcp/_conversation_id.py index e4714d4f5..0187b01a3 100644 --- a/posthog/mcp/_conversation_id.py +++ b/posthog/mcp/_conversation_id.py @@ -10,6 +10,7 @@ from __future__ import annotations import copy +import re from typing import Any, Dict, Optional, Tuple from .constants import DEFAULT_CONVERSATION_ID_DESCRIPTION @@ -18,6 +19,16 @@ CONVERSATION_ID_PARAM_NAME = "conversation_id" +# The shape of every id we mint: a uuidv7. Used to tell an echo of our own +# handle from a value the agent made up. The shape check matters because the +# handle becomes ``$session_id`` — without it, two unrelated users both sending +# "conv-1" would share a session (byte-parity with posthog-js's +# MINTED_CONVERSATION_ID). +_MINTED_CONVERSATION_ID = re.compile( + r"^[0-9a-f]{8}-[0-9a-f]{4}-7[0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}$", + re.IGNORECASE, +) + def add_conversation_id_to_schema( input_schema: Optional[Dict[str, Any]], tool_name: str = "unknown" @@ -71,20 +82,28 @@ def resolve_conversation_id( missing_capability_tool_name: str, ) -> Tuple[Optional[str], bool]: """Return ``(conversation_id, minted)``. Disabled or get_more_tools → ``(None, False)``; - agent supplied → ``(value, False)``; agent omitted → ``(new uuid, True)``.""" + agent echoed a handle we could have minted → ``(value, False)``; anything + else (omitted, or a value the agent made up) → ``(new uuid, True)``. + + Lowercased on the way in: the shape test is case-insensitive but the hash + behind ``$session_id`` is not, so an uppercased echo (some hosts normalise + uuids) would land in a different session than the call that minted it.""" if not enabled or tool_name == missing_capability_tool_name: return None, False supplied = extract_conversation_id(args) - if supplied: - return supplied, False + if supplied and _MINTED_CONVERSATION_ID.match(supplied): + return supplied.lower(), False return _uuid7(), True def can_inject_prompt_back(result: Any) -> bool: + """Whether the prompt-back can ride this result's ``content`` — the only + requirement is a list to append to. Errored results included on purpose: a + tool that fails on the first call of a conversation is exactly when the + agent needs the handle, or the retry starts a fresh conversation and the + failure and its fix land in different sessions.""" if not isinstance(result, dict): return False - if result.get("isError") is True: - return False return isinstance(result.get("content"), list) diff --git a/posthog/mcp/_instrument_fastmcp.py b/posthog/mcp/_instrument_fastmcp.py index 6fd871f2d..5480586e7 100644 --- a/posthog/mcp/_instrument_fastmcp.py +++ b/posthog/mcp/_instrument_fastmcp.py @@ -101,6 +101,15 @@ async def wrapped( request = build_tool_call_request(name, arguments) extra: Dict[str, Any] = {"session_id": mcp_session_id} + # Resolve the conversation handle before the session: when the agent + # carries (or is about to receive) one, it anchors $session_id for every + # event of this request (ADR-0004) — the only correlation that survives + # the 2026-07-28 revision's per-request server instances. + missing_name = resolve_missing_capability_tool_name(data.options) + conversation_id, minted = resolve_conversation_id( + data.options.enable_conversation_id, arguments, name, missing_name + ) + session_id = await prepare_request( data, mcp_session_id=mcp_session_id, @@ -110,9 +119,9 @@ async def wrapped( request=request, extra=extra, token=token, + conversation_id=conversation_id, ) - missing_name = resolve_missing_capability_tool_name(data.options) if data.options.report_missing and name == missing_name: await record_missing_capability( data, @@ -129,10 +138,6 @@ async def wrapped( mcp_types.TextContent(type="text", text=get_more_tools_result_text()) ] - conversation_id, minted = resolve_conversation_id( - data.options.enable_conversation_id, arguments, name, missing_name - ) - # Strip each injected key independently. A tool can declare its own # `context` (kept) while `conversation_id` is still SDK-injected (stripped), # so coupling both to context-ownership leaked conversation_id into the tool. @@ -328,8 +333,10 @@ def _inject_prompt_back(result: Any, conversation_id: str) -> Any: """Append the conversation_id prompt-back to a tool result so the agent echoes it on later calls. Handles every shape ToolManager.call_tool can return: a ``(content_list, structured)`` tuple (the convert_result=True production path), - a bare content list, or a ``{content: [...]}`` dict. Returns the result unchanged - (so the caller can detect non-delivery) for shapes we can't append to.""" + a bare content list, or a ``{content: [...]}`` dict — errored dicts included on + purpose (a first-call failure is exactly when the agent needs the handle). + Returns the result unchanged (so the caller can detect non-delivery) for + shapes we can't append to.""" block = mcp_types.TextContent( type="text", text=build_prompt_back(conversation_id)["text"] ) @@ -337,11 +344,7 @@ def _inject_prompt_back(result: Any, conversation_id: str) -> Any: return ([*result[0], block], result[1]) if isinstance(result, list): return [*result, block] - if ( - isinstance(result, dict) - and isinstance(result.get("content"), list) - and not result.get("isError") - ): + if isinstance(result, dict) and isinstance(result.get("content"), list): return {**result, "content": [*result["content"], block]} return result diff --git a/posthog/mcp/_instrument_lowlevel.py b/posthog/mcp/_instrument_lowlevel.py index cbe6c5002..fc443f27d 100644 --- a/posthog/mcp/_instrument_lowlevel.py +++ b/posthog/mcp/_instrument_lowlevel.py @@ -104,6 +104,13 @@ async def handler(req: Any) -> Any: request = build_tool_call_request(name, arguments) extra = {"session_id": mcp_session_id} + # Resolve the conversation handle before the session: when present it + # anchors $session_id for every event of this request (ADR-0004). + missing_name = resolve_missing_capability_tool_name(data.options) + conversation_id, minted = resolve_conversation_id( + data.options.enable_conversation_id, arguments, name, missing_name + ) + session_id = await prepare_request( data, mcp_session_id=mcp_session_id, @@ -113,9 +120,9 @@ async def handler(req: Any) -> Any: request=request, extra=extra, token=token, + conversation_id=conversation_id, ) - missing_name = resolve_missing_capability_tool_name(data.options) if data.options.report_missing and name == missing_name: await record_missing_capability( data, @@ -139,10 +146,6 @@ async def handler(req: Any) -> Any: ) ) - conversation_id, minted = resolve_conversation_id( - data.options.enable_conversation_id, arguments, name, missing_name - ) - # On raw low-level servers `context`/`conversation_id` are injected as # *optional* schema properties and left in place (a (name, arguments) # handler ignores extra keys). FastMCP 2.0 validates against the function @@ -186,12 +189,14 @@ async def handler(req: Any) -> Any: call_result = getattr(result, "root", result) # Inject the prompt-back before capture; only stamp a minted conversation_id - # when it was actually delivered (not on isError / non-list results), so we - # don't record an orphan id the agent never received. + # when it was actually delivered (non-list results can't carry it), so we + # don't record an orphan id the agent never received. Errored results carry + # it on purpose: a first-call failure is exactly when the agent needs the + # handle, or the retry starts a fresh conversation. delivered_conversation_id = conversation_id if minted and conversation_id: content = getattr(call_result, "content", None) - if not getattr(call_result, "isError", False) and isinstance(content, list): + if isinstance(content, list): content.append( mcp_types.TextContent( type="text", text=build_prompt_back(conversation_id)["text"] diff --git a/posthog/mcp/_instrument_v2.py b/posthog/mcp/_instrument_v2.py new file mode 100644 index 000000000..99af12bd9 --- /dev/null +++ b/posthog/mcp/_instrument_v2.py @@ -0,0 +1,611 @@ +# Portions of this package are derived from MCPCat/mcpcat-typescript-sdk +# Copyright (c) 2025 MCPcat +# Licensed under the MIT License: https://github.com/MCPCat/mcpcat-typescript-sdk/blob/main/LICENSE + +"""MCP Python SDK 2.x adapters (spec revision 2026-07-28). + +Two entry points, mirroring the 1.x pair: + +* :func:`instrument_mcpserver_v2` — the high-level ``mcp.server.mcpserver.MCPServer`` + (the renamed FastMCP). Tool calls are wrapped at ``ToolManager.call_tool``, the + one seam every dispatch routes through (late-registered tools covered for + free); tools/list at the underlying low-level registry. +* :func:`instrument_lowlevel_v2` — the low-level ``mcp.server.lowlevel.Server``. + v2 replaced the public ``request_handlers`` dict (keyed by request class) with + ``add_request_handler``/``get_request_handler`` keyed by method string, and + handlers changed shape to ``(ctx, params)``; both existing registrations and + later ``add_request_handler`` calls are wrapped (the posthog-js#4449 lesson: + adapters that hand over a bare server register handlers *after* instrument()). + +Era note: a single v2 server serves both protocol eras request by request — the +legacy 2025-11-25 handshake and the stateless 2026-07-28 envelope. Nothing here +branches on era: ``ctx.protocol_version`` is captured as-is, client identity +comes from ``ctx.session.client_params`` (synthesized from the per-request +envelope on the modern era), and on 2026-07-28 — which removed protocol +sessions — cross-pod correlation comes from ``enable_conversation_id``. + +v2 models expose snake_case attributes (``is_error``, ``input_schema``, +``client_info``); the wire JSON keeps the camelCase aliases. +""" + +from __future__ import annotations + +import time +from typing import Any, Dict, Optional, Tuple + +import mcp.types as mcp_types + +from ._context_parameters import ( + add_context_parameter_to_schema, + get_context_description, + is_context_enabled, +) +from ._conversation_id import ( + add_conversation_id_to_schema, + build_prompt_back, + resolve_conversation_id, +) +from ._instrumentation import ( + _to_jsonable, + build_tool_call_request, + prepare_request, + read_tool_category, + record_missing_capability, + record_tool_call, + record_tools_list, + resolve_session_and_client, +) +from ._internal import MCPAnalyticsData +from .logger import log +from .tools import ( + GET_MORE_TOOLS_NAME as _GET_MORE_TOOLS_NAME, + build_report_missing_descriptor, + get_more_tools_result_text, + resolve_missing_capability_tool_name, +) + +_WRAPPED_FLAG = "__posthog_mcp_wrapped__" + +# The two methods we instrument, and how the tools/list wrapper treats the +# injected `context` parameter per entry point (see _wrap_v2_list_tools). +_CALL_METHOD = "tools/call" +_LIST_METHOD = "tools/list" + + +def instrument_mcpserver_v2(server: Any, data: MCPAnalyticsData) -> None: + """Instrument a v2 ``MCPServer``. Injected ``context``/``conversation_id`` + are STRIPPED before dispatch (v2 validates tool arguments against the + function signature and rejects unexpected keys), unless the tool's own + schema declares the parameter — then it's a real argument the agent's value + belongs to.""" + low_level = getattr(server, "_lowlevel_server", None) + if low_level is None: + log("Warning: MCPServer has no _lowlevel_server; cannot instrument.") + return + data.server_name = getattr(server, "name", None) or getattr(low_level, "name", None) + data.server_version = getattr(server, "version", None) or getattr( + low_level, "version", None + ) + _wrap_tool_manager_call_v2(server, data) + _wrap_v2_list_tools(low_level, data, context_required=True, high_level=server) + _patch_add_request_handler(low_level, data, wrap_call=False, high_level=server) + + +def instrument_lowlevel_v2(server: Any, data: MCPAnalyticsData) -> None: + """Instrument a raw v2 low-level ``Server``. ``context`` is injected as an + *optional* schema property and NOT stripped — the schema doubles as the + call's validation surface, and a typical ``(ctx, params)`` handler ignores + extra argument keys.""" + data.server_name = getattr(server, "name", None) + data.server_version = getattr(server, "version", None) + _wrap_v2_call_tool(server, data) + _wrap_v2_list_tools(server, data, context_required=False) + _patch_add_request_handler(server, data, wrap_call=True) + + +# --- registry plumbing --------------------------------------------------------- + + +def _replace_handler(server: Any, method: str, wrapped: Any, params_type: Any) -> None: + """Re-register through the public API so the SDK keeps owning validation.""" + server.add_request_handler(method, params_type, wrapped) + + +def _patch_add_request_handler( + server: Any, data: MCPAnalyticsData, *, wrap_call: bool, high_level: Any = None +) -> None: + """Wrap ``add_request_handler`` so handlers registered *after* instrument() + for the instrumented methods get wrapped too. Registrations for other + methods pass through untouched.""" + original_add = server.add_request_handler + if getattr(original_add, _WRAPPED_FLAG, False): + return + + def add_request_handler(method: str, params_type: Any, handler: Any) -> None: + original_add(method, params_type, handler) + if getattr(handler, _WRAPPED_FLAG, False): + return + if method == _CALL_METHOD and wrap_call: + _wrap_v2_call_tool(server, data) + elif method == _LIST_METHOD: + _wrap_v2_list_tools( + server, + data, + context_required=high_level is not None, + high_level=high_level, + ) + + setattr(add_request_handler, _WRAPPED_FLAG, True) + server.add_request_handler = add_request_handler + + +# --- ctx readers ----------------------------------------------------------------- + + +def _ctx_client_info(ctx: Any) -> Tuple[Optional[str], Optional[str]]: + try: + client_params = ctx.session.client_params + info = getattr(client_params, "client_info", None) + if info is not None: + return getattr(info, "name", None), getattr(info, "version", None) + except Exception: # noqa: BLE001 + pass + return None, None + + +def _ctx_protocol_version(ctx: Any) -> Optional[str]: + version = getattr(ctx, "protocol_version", None) + if isinstance(version, str) and version: + return version + try: + return ctx.session.client_params.protocol_version + except Exception: # noqa: BLE001 + return None + + +def _ctx_mcp_session_id(ctx: Any) -> Optional[str]: + """Best-effort transport session id (the ``Mcp-Session-Id`` header on the + legacy era — 2026-07-28 removed it). ``ctx.request`` carries the transport's + HTTP request when there is one; ``None`` on stdio.""" + try: + headers = getattr(getattr(ctx, "request", None), "headers", None) + if headers is not None: + return headers.get("mcp-session-id") + except Exception: # noqa: BLE001 + pass + return None + + +def _resolve_ctx( + ctx: Any, +) -> Tuple[Optional[Any], Optional[str], Optional[str], Optional[str], Optional[str]]: + """(token, client_name, client_version, protocol_version, mcp_session_id) + for a request, with token-carried identity backfilled for stateless pods.""" + client_name, client_version = _ctx_client_info(ctx) + protocol_version = _ctx_protocol_version(ctx) + mcp_session_id = _ctx_mcp_session_id(ctx) + token, client_name, client_version, protocol_version = resolve_session_and_client( + mcp_session_id, client_name, client_version, protocol_version + ) + return token, client_name, client_version, protocol_version, mcp_session_id + + +def _params_to_request(method: str, params: Any) -> Dict[str, Any]: + params_dict: Any = {} + if params is not None and hasattr(params, "model_dump"): + try: + params_dict = params.model_dump(mode="json", by_alias=True) + except Exception: # noqa: BLE001 + params_dict = {} + return {"method": method, "params": params_dict} + + +# --- tool ownership -------------------------------------------------------------- + + +def _tool_owns_param_v2(high_level: Any, name: str, param: str) -> bool: + """Whether the tool's own JSON schema declares ``param`` — then it's a real + tool argument we must neither inject over nor strip. Read from the tool's + declared parameters rather than the function signature so a tool taking the + SDK's ``Context`` object under a ``context`` name isn't mistaken for owning + our string parameter.""" + try: + tool = high_level._tool_manager.get_tool(name) + properties = (getattr(tool, "parameters", None) or {}).get("properties", {}) + return param in properties + except Exception: # noqa: BLE001 + return False + + +# --- high-level: ToolManager.call_tool seam -------------------------------------- + + +def _wrap_tool_manager_call_v2(server: Any, data: MCPAnalyticsData) -> None: + tool_manager = getattr(server, "_tool_manager", None) + if tool_manager is None: + log("Warning: MCPServer has no _tool_manager; tool calls will not be captured.") + return + + original = tool_manager.call_tool + if getattr(original, _WRAPPED_FLAG, False): + return + + async def wrapped( + name: str, + arguments: Dict[str, Any], + context: Any = None, + convert_result: bool = False, + ) -> Any: + ctx = getattr(context, "_request_context", None) + token, client_name, client_version, protocol_version, mcp_session_id = ( + _resolve_ctx(ctx) + ) + request = build_tool_call_request(name, arguments) + extra: Dict[str, Any] = {"session_id": mcp_session_id, "ctx": ctx} + + missing_name = resolve_missing_capability_tool_name(data.options) + conversation_id, minted = resolve_conversation_id( + data.options.enable_conversation_id, arguments, name, missing_name + ) + + session_id = await prepare_request( + data, + mcp_session_id=mcp_session_id, + client_name=client_name, + client_version=client_version, + protocol_version=protocol_version, + request=request, + extra=extra, + token=token, + conversation_id=conversation_id, + ) + + if data.options.report_missing and name == missing_name: + await record_missing_capability( + data, + session_id, + tool_name=missing_name, + context=(arguments or {}).get("context"), + arguments=arguments, + client_name=client_name, + client_version=client_version, + protocol_version=protocol_version, + extra=extra, + ) + return mcp_types.CallToolResult( + content=[ + mcp_types.TextContent( + type="text", text=get_more_tools_result_text() + ) + ] + ) + + # v2 validates against the function signature and rejects unexpected + # keys, so injected parameters are stripped before dispatch — but never + # one the tool's own schema declares (that's a real argument). + call_arguments = arguments + if isinstance(arguments, dict): + strip_keys = set() + if not _tool_owns_param_v2(server, name, "context"): + strip_keys.add("context") + if data.options.enable_conversation_id and not _tool_owns_param_v2( + server, name, "conversation_id" + ): + strip_keys.add("conversation_id") + if strip_keys: + call_arguments = { + k: v for k, v in arguments.items() if k not in strip_keys + } + + start = time.monotonic() + try: + result = await original( + name, call_arguments, context=context, convert_result=convert_result + ) + except Exception as error: + # The raise is converted to CallToolResult(is_error=True) one layer + # up (MCPServer._handle_call_tool), so the prompt-back never rides + # it — a minted (undelivered) conversation_id is not stamped. + await record_tool_call( + data, + session_id, + name=name, + arguments=arguments, + error=error, + duration_ms=(time.monotonic() - start) * 1000, + client_name=client_name, + client_version=client_version, + protocol_version=protocol_version, + conversation_id=None if minted else conversation_id, + extra=extra, + ) + raise + duration_ms = (time.monotonic() - start) * 1000 + + delivered_conversation_id = conversation_id + if minted and conversation_id: + if not _append_prompt_back(result, conversation_id): + delivered_conversation_id = None + + await record_tool_call( + data, + session_id, + name=name, + arguments=arguments, + result=result, + duration_ms=duration_ms, + client_name=client_name, + client_version=client_version, + protocol_version=protocol_version, + conversation_id=delivered_conversation_id, + extra=extra, + ) + return result + + setattr(wrapped, _WRAPPED_FLAG, True) + tool_manager.call_tool = wrapped + + +def _append_prompt_back(result: Any, conversation_id: str) -> bool: + """Append the conversation prompt-back to a result's ``content`` list (model + or dict shape). Errored results included on purpose — a first-call failure + is exactly when the agent needs the handle. Returns False for shapes with no + content list to ride (e.g. MRTR ``input_required`` results).""" + block = mcp_types.TextContent( + type="text", text=build_prompt_back(conversation_id)["text"] + ) + content = ( + result.get("content") + if isinstance(result, dict) + else getattr(result, "content", None) + ) + if isinstance(content, list): + content.append(block) + return True + return False + + +# --- low-level: tools/call ------------------------------------------------------ + + +def _wrap_v2_call_tool(server: Any, data: MCPAnalyticsData) -> None: + entry = server.get_request_handler(_CALL_METHOD) + if entry is None or getattr(entry.handler, _WRAPPED_FLAG, False): + return + original = entry.handler + + async def handler(ctx: Any, params: Any) -> Any: + name = params.name + arguments = dict(params.arguments or {}) + token, client_name, client_version, protocol_version, mcp_session_id = ( + _resolve_ctx(ctx) + ) + request = build_tool_call_request(name, arguments) + extra: Dict[str, Any] = {"session_id": mcp_session_id, "ctx": ctx} + + missing_name = resolve_missing_capability_tool_name(data.options) + conversation_id, minted = resolve_conversation_id( + data.options.enable_conversation_id, arguments, name, missing_name + ) + + session_id = await prepare_request( + data, + mcp_session_id=mcp_session_id, + client_name=client_name, + client_version=client_version, + protocol_version=protocol_version, + request=request, + extra=extra, + token=token, + conversation_id=conversation_id, + ) + + if data.options.report_missing and name == missing_name: + await record_missing_capability( + data, + session_id, + tool_name=missing_name, + context=arguments.get("context"), + arguments=arguments, + client_name=client_name, + client_version=client_version, + protocol_version=protocol_version, + extra=extra, + ) + return mcp_types.CallToolResult( + content=[ + mcp_types.TextContent( + type="text", text=get_more_tools_result_text() + ) + ] + ) + + start = time.monotonic() + try: + result = await original(ctx, params) + except Exception as error: + # v2 low-level handlers raise through to JSON-RPC errors (no + # auto-conversion) — capture before re-raising. The prompt-back was + # never delivered, so a minted conversation_id is not stamped. + await record_tool_call( + data, + session_id, + name=name, + arguments=arguments, + error=error, + duration_ms=(time.monotonic() - start) * 1000, + client_name=client_name, + client_version=client_version, + protocol_version=protocol_version, + conversation_id=None if minted else conversation_id, + extra=extra, + ) + raise + duration_ms = (time.monotonic() - start) * 1000 + + delivered_conversation_id = conversation_id + if minted and conversation_id: + if not _append_prompt_back(result, conversation_id): + delivered_conversation_id = None + + await record_tool_call( + data, + session_id, + name=name, + arguments=arguments, + result=result, + duration_ms=duration_ms, + client_name=client_name, + client_version=client_version, + protocol_version=protocol_version, + conversation_id=delivered_conversation_id, + extra=extra, + ) + return result + + setattr(handler, _WRAPPED_FLAG, True) + _replace_handler(server, _CALL_METHOD, handler, entry.params_type) + + +# --- tools/list ------------------------------------------------------------------- + + +def _wrap_v2_list_tools( + server: Any, + data: MCPAnalyticsData, + *, + context_required: bool, + high_level: Any = None, +) -> None: + entry = server.get_request_handler(_LIST_METHOD) + if entry is None or getattr(entry.handler, _WRAPPED_FLAG, False): + return + original = entry.handler + + async def handler(ctx: Any, params: Any) -> Any: + token, client_name, client_version, protocol_version, mcp_session_id = ( + _resolve_ctx(ctx) + ) + request = _params_to_request(_LIST_METHOD, params) + extra: Dict[str, Any] = {"session_id": mcp_session_id, "ctx": ctx} + # Resolve session, emit $mcp_initialize (once per session) and identify + # here too — a client may list tools without ever calling one. + session_id = await prepare_request( + data, + mcp_session_id=mcp_session_id, + client_name=client_name, + client_version=client_version, + protocol_version=protocol_version, + request=request, + extra=extra, + token=token, + ) + + start = time.monotonic() + try: + result = await original(ctx, params) + except Exception as error: + await record_tools_list( + data, + session_id, + names=[], + request=request, + duration_ms=(time.monotonic() - start) * 1000, + is_error=True, + error=error, + client_name=client_name, + client_version=client_version, + protocol_version=protocol_version, + extra=extra, + ) + raise + duration_ms = (time.monotonic() - start) * 1000 + + tools = list(getattr(result, "tools", []) or []) + # Zero advertised tools is treated as an errored tools/list (parity with + # the TS SDK), checked before we append our own virtual tool. + empty = len(tools) == 0 + + names = [] + for tool in tools: + names.append(tool.name) + if getattr(tool, "description", None): + data.tool_descriptions[tool.name] = tool.description + category = read_tool_category(tool) + if category: + data.tool_categories[tool.name] = category + + context_enabled = is_context_enabled(data.options.context) + description = get_context_description(data.options.context) + for tool in tools: + if tool.name == _GET_MORE_TOOLS_NAME: + continue + schema = getattr(tool, "input_schema", None) + owns_context = ( + _tool_owns_param_v2(high_level, tool.name, "context") + if high_level is not None + else _schema_has_param(schema, "context") + ) + # required follows the entry point: the raw low-level path validates + # the call against this same schema (optional); the high-level path + # strips before dispatch (required-advisory). + if context_enabled and not owns_context: + schema = add_context_parameter_to_schema( + schema, tool.name, description, required=context_required + ) + if data.options.enable_conversation_id and not _schema_has_param( + schema, "conversation_id" + ): + schema = add_conversation_id_to_schema(schema, tool.name) + if schema is not getattr(tool, "input_schema", None): + try: + tool.input_schema = schema + except Exception: # noqa: BLE001 - some schema attrs may be read-only + log(f"WARN: could not set input_schema on tool {tool.name}") + + if data.options.report_missing: + missing_name = resolve_missing_capability_tool_name(data.options) + if not any(t.name == missing_name for t in tools): + _append_get_more_tools_v2(result, missing_name) + names.append(missing_name) + + await record_tools_list( + data, + session_id, + names=names, + request=request, + response=_to_jsonable(result), + duration_ms=duration_ms, + is_error=empty, + error="tools/list returned no tools" if empty else None, + client_name=client_name, + client_version=client_version, + protocol_version=protocol_version, + extra=extra, + ) + + return result + + setattr(handler, _WRAPPED_FLAG, True) + _replace_handler(server, _LIST_METHOD, handler, entry.params_type) + + +def _append_get_more_tools_v2(result: Any, name: str) -> None: + descriptor = build_report_missing_descriptor(name) + tool = mcp_types.Tool( + name=descriptor["name"], + description=descriptor["description"], + input_schema=descriptor["inputSchema"], + annotations=descriptor["annotations"], + ) + tools_list = getattr(result, "tools", None) + if isinstance(tools_list, list): + tools_list.append(tool) + + +def _schema_has_param(schema: Any, name: str) -> bool: + return ( + isinstance(schema, dict) + and isinstance(schema.get("properties"), dict) + and name in schema["properties"] + ) diff --git a/posthog/mcp/_instrumentation.py b/posthog/mcp/_instrumentation.py index ff5380abd..9ab365ca1 100644 --- a/posthog/mcp/_instrumentation.py +++ b/posthog/mcp/_instrumentation.py @@ -141,10 +141,15 @@ def drain_pending_sync(owner: Any, timeout: Optional[float] = None) -> None: def is_tool_result_error(result: Any) -> bool: - """MCP tool results signal errors via ``isError: true`` rather than raising.""" + """MCP tool results signal errors via ``isError: true`` rather than raising. + The attribute is ``isError`` on MCP SDK 1.x models and ``is_error`` on 2.x + (wire JSON unchanged); check both shapes.""" if isinstance(result, dict): - return result.get("isError") is True - return getattr(result, "isError", None) is True + return result.get("isError") is True or result.get("is_error") is True + return ( + getattr(result, "isError", None) is True + or getattr(result, "is_error", None) is True + ) def build_tool_call_request( @@ -158,8 +163,11 @@ def build_tool_call_request( def _to_jsonable(obj: Any) -> Any: if hasattr(obj, "model_dump"): + # by_alias so captured payloads keep the camelCase wire shape on both MCP + # SDK majors (2.x renamed model attributes to snake_case but kept the + # aliases); 1.x field names are already the wire names, so this is a no-op. try: - return obj.model_dump(mode="json") + return obj.model_dump(mode="json", by_alias=True) except Exception: # noqa: BLE001 return str(obj) if isinstance(obj, (list, tuple)): @@ -250,10 +258,16 @@ async def prepare_request( extra: Optional[Dict[str, Any]], token: Optional[SessionTokenPayload] = None, protocol_version: Optional[str] = None, + conversation_id: Optional[str] = None, ) -> str: """Resolve the session id, run identify, then lazily emit initialize. Returns the session id to stamp on the event for this request. + ``conversation_id`` is the agent's handle for this request (echoed or freshly + minted); when present it anchors the session (ADR-0004) so every event of the + request — identify, initialize, and the call itself — lands in the + conversation's session rather than this instance's. + ``token`` is the decoded self-encoded session token (see ``session_token.py``); when present it takes precedence over ``mcp_session_id`` and carries the client identity across stateless pods. @@ -263,7 +277,9 @@ async def prepare_request( ``$mcp_initialize`` is anonymous even when identify resolves on the same request. (Still not byte-parity with the TS SDK, which wraps the real initialize handler; the Python SDK handles initialize in the session layer, not ``request_handlers``.)""" - session_id = await resolve_session_id(data, mcp_session_id, token=token) + session_id = await resolve_session_id( + data, mcp_session_id, token=token, conversation_id=conversation_id + ) identify_event = await handle_identify(data, session_id, request, extra) if identify_event: fire_and_forget(capture_event(data, identify_event), data) diff --git a/posthog/mcp/session.py b/posthog/mcp/session.py index 902342ec5..a7edbc8d4 100644 --- a/posthog/mcp/session.py +++ b/posthog/mcp/session.py @@ -16,7 +16,7 @@ from ._internal import MCPAnalyticsData from .session_token import SessionTokenPayload -__all__ = ["derive_session_id_from_mcp_session"] +__all__ = ["derive_session_id_from_mcp_session", "derive_session_id_from_conversation"] def new_session_id() -> str: @@ -29,18 +29,44 @@ def derive_session_id_from_mcp_session(mcp_session_id: str) -> str: return deterministic_prefixed_id("ses", mcp_session_id) +def derive_session_id_from_conversation(conversation_id: str) -> str: + """Derive the SDK session id from the agent's conversation handle. + + Deterministic and unsalted on purpose: two pods that never met must agree on + the session, and the 2026-07-28 protocol revision leaves them no shared + state to agree through. Hashed rather than used verbatim so an MCP session + can never collide with a Session Replay id. + + This is the cross-SDK contract: posthog-js's ``deriveSessionIdFromConversation`` + produces the same value byte for byte, or the same conversation splits into + two sessions depending on which SDK served the call. + """ + return deterministic_prefixed_id("ses", conversation_id) + + async def resolve_session_id( data: MCPAnalyticsData, mcp_session_id: Optional[str], *, token: Optional[SessionTokenPayload] = None, + conversation_id: Optional[str] = None, ) -> str: """Resolve the session id for a request. Mutates per-server state under a lock so concurrent async requests can't race on session rotation. + Priority mirrors posthog-js ``getSessionId``: the agent's ``conversation_id`` + handle first (the only id that survives the 2026-07-28 revision's + per-request server instances), then our self-encoded session token, then the + transport's MCP session id, then this instance's own memory. + + ``conversation_id`` is resolved *per request and never stored*: the handle + belongs to one chat, and persisting it on shared ``data`` (or advancing + ``last_activity``) would leak one chat's session onto a concurrent chat's + request. + ``token`` is our self-encoded session token (see :mod:`.session_token`), - decoded from the replayed ``Mcp-Session-Id`` header. It is the only source - that survives a stateless / multi-pod deployment, so it takes precedence. + decoded from the replayed ``Mcp-Session-Id`` header. It is the only other + source that survives a stateless / multi-pod deployment. The token session is resolved *per request*, never sticky: ``data`` is shared by every client hitting this server instance, so reusing a stored token session @@ -48,6 +74,9 @@ async def resolve_session_id( one ``$session_id``. A compliant client replays the header on every request, so a genuine token session never needs the fallback. """ + if conversation_id: + return derive_session_id_from_conversation(conversation_id) + async with data.session_lock: now = datetime.now(timezone.utc) diff --git a/posthog/mcp/version.py b/posthog/mcp/version.py index c91ca8052..862043880 100644 --- a/posthog/mcp/version.py +++ b/posthog/mcp/version.py @@ -4,4 +4,4 @@ # Version of the PostHog MCP analytics SDK surface. Stamped as ``sdk_version`` # on every captured event. The package itself ships inside posthog. -__version__ = "0.2.0" +__version__ = "0.3.0" diff --git a/references/public_api_snapshot.txt b/references/public_api_snapshot.txt index b2adc71e5..3c311a1c3 100644 --- a/references/public_api_snapshot.txt +++ b/references/public_api_snapshot.txt @@ -350,6 +350,7 @@ alias posthog.mcp.asgi.log -> posthog.mcp.logger.log alias posthog.mcp.asgi.new_session_id -> posthog.mcp.session.new_session_id alias posthog.mcp.asgi.read_mcp_session_header -> posthog.mcp.session_token.read_mcp_session_header alias posthog.mcp.decode_session_id -> posthog.mcp.session_token.decode_session_id +alias posthog.mcp.derive_session_id_from_conversation -> posthog.mcp.session.derive_session_id_from_conversation alias posthog.mcp.derive_session_id_from_mcp_session -> posthog.mcp.session.derive_session_id_from_mcp_session alias posthog.mcp.encode_session_id -> posthog.mcp.session_token.encode_session_id alias posthog.mcp.get_mcp_session -> posthog.mcp.asgi.get_mcp_session @@ -1128,6 +1129,7 @@ function posthog.mcp.asgi.autowire_stateless_mint(server: Any) -> None function posthog.mcp.asgi.get_mcp_session(request_or_scope: Any) -> Optional[SessionTokenPayload] function posthog.mcp.instrument(server: Any, posthog_client: Optional[Client] = None, options: Optional[MCPAnalyticsOptions] = None) -> McpAnalytics function posthog.mcp.logger.set_logger(logger: Optional[LoggerFn]) -> None +function posthog.mcp.session.derive_session_id_from_conversation(conversation_id: str) -> str function posthog.mcp.session.derive_session_id_from_mcp_session(mcp_session_id: str) -> str function posthog.mcp.session_token.decode_session_id(value: Any) -> Optional[SessionTokenPayload] function posthog.mcp.session_token.encode_session_id(payload: SessionTokenPayload) -> str From 9fda7cca546910f958a6a1b729964501c372d7d4 Mon Sep 17 00:00:00 2001 From: Georgis Andonis Date: Thu, 20 Aug 2026 14:50:20 +0300 Subject: [PATCH 03/16] chore(mcp): refresh public API snapshot for the 0.3.0 sdk-surface bump Generated-By: PostHog Code Task-Id: ebafcb71-b03b-443d-b40c-d527ed4a04f4 --- references/public_api_snapshot.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/references/public_api_snapshot.txt b/references/public_api_snapshot.txt index 3c311a1c3..21b8b1fc2 100644 --- a/references/public_api_snapshot.txt +++ b/references/public_api_snapshot.txt @@ -753,7 +753,7 @@ attribute posthog.mcp.types.PreparedToolCall.is_missing_capability: bool = False attribute posthog.mcp.types.UserIdentity.distinct_id: str attribute posthog.mcp.types.UserIdentity.groups: Optional[Dict[str, str]] = None attribute posthog.mcp.types.UserIdentity.properties: Optional[JsonRecord] = None -attribute posthog.mcp.version.__version__ = '0.2.0' +attribute posthog.mcp.version.__version__ = '0.3.0' attribute posthog.metrics = None attribute posthog.metrics_capture.DEFAULT_HISTOGRAM_BOUNDS = [0, 5, 10, 25, 50, 75, 100, 250, 500, 750, 1000, 2500, 5000, 7500, 10000] attribute posthog.metrics_capture.MetricAttributeValue = Union[str, int, float, bool] From 021823f05afac9c5c0df66cd66c7f5eb9b1f2de2 Mon Sep 17 00:00:00 2001 From: Georgis Andonis Date: Thu, 20 Aug 2026 15:13:19 +0300 Subject: [PATCH 04/16] fix(mcp): capture only a scalar projection of extra on $identify events MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The v2 adapters hand the raw request ctx to identify/event_properties/ intent_fallback callbacks via `extra` so hosts can read headers — but handle_identify embedded the whole dict into the captured $identify parameters, where the sanitizer leaves opaque objects untouched and truncation stringifies them: whatever the context repr carries (headers, transport state) would ship to PostHog without key-based redaction. Callbacks keep the full extra; captured parameters now carry only JSON-safe scalars (e.g. session_id). Generated-By: PostHog Code Task-Id: ebafcb71-b03b-443d-b40c-d527ed4a04f4 --- posthog/mcp/_internal.py | 20 ++++++++++++++++- posthog/test/mcp/test_v2_mcpserver.py | 31 +++++++++++++++++++++++++++ 2 files changed, 50 insertions(+), 1 deletion(-) diff --git a/posthog/mcp/_internal.py b/posthog/mcp/_internal.py index 5c5295413..5f77fee7c 100644 --- a/posthog/mcp/_internal.py +++ b/posthog/mcp/_internal.py @@ -132,6 +132,24 @@ async def _maybe_await(value: Any) -> Any: return value +def _captured_extra(extra: Optional[Dict[str, Any]]) -> Optional[Dict[str, Any]]: + """Project ``extra`` down to JSON-safe scalars before it is captured. + + Callbacks (``identify``, ``event_properties``, ``intent_fallback``) receive + the full dict — on MCP SDK v2 that includes the raw request ``ctx`` so hosts + can read headers. Captured event parameters must not: the sanitizer leaves + opaque objects untouched and truncation stringifies them, which would ship + whatever the object's repr carries (headers, auth material, transport state) + to PostHog without key-based redaction.""" + if extra is None: + return None + return { + key: value + for key, value in extra.items() + if value is None or isinstance(value, (str, int, float, bool)) + } + + async def handle_identify( data: MCPAnalyticsData, session_id: str, @@ -170,7 +188,7 @@ async def handle_identify( "session_id": session_id, "resource_name": _get_request_resource_name(request), "event_type": MCPAnalyticsEventType.IDENTIFY, - "parameters": {"request": request, "extra": extra}, + "parameters": {"request": request, "extra": _captured_extra(extra)}, "timestamp": datetime.now(timezone.utc), } except Exception as error: # noqa: BLE001 diff --git a/posthog/test/mcp/test_v2_mcpserver.py b/posthog/test/mcp/test_v2_mcpserver.py index 37306fdae..f591cc17e 100644 --- a/posthog/test/mcp/test_v2_mcpserver.py +++ b/posthog/test/mcp/test_v2_mcpserver.py @@ -208,6 +208,37 @@ async def test_identify_sets_distinct_id_and_groups(): assert _events(client, "$identify") +async def test_identify_callback_sees_ctx_but_capture_does_not(): + """The v2 adapters hand the raw request ``ctx`` to callbacks via ``extra`` so + hosts can read headers — but the captured $identify parameters must carry + only a scalar projection, or the stringified context (headers, transport + state) would ship to PostHog without key-based redaction.""" + server = make_server() + client = FakeClient() + seen = {} + + def identify(request, extra): + seen["extra"] = extra + return UserIdentity(distinct_id="user_9") + + instrument(server, client, MCPAnalyticsOptions(identify=identify)) + + await _call_tool( + server, "add", {"a": 1, "b": 1, "context": "identity capture check"} + ) + await _flush() + + # the callback got the live context object... + assert seen["extra"]["ctx"] is not None + + # ...but the captured event only carries JSON-safe scalars + identify_events = _events(client, "$identify") + assert identify_events + captured_extra = identify_events[0]["properties"]["$mcp_parameters"]["extra"] + assert "ctx" not in captured_extra + assert set(captured_extra) <= {"session_id"} + + async def test_report_missing_advertises_and_captures(): server = make_server() client = FakeClient() From fddfaf1a2cd9c76638ff0ef05cd53fa42b9928fc Mon Sep 17 00:00:00 2001 From: Georgis Andonis Date: Thu, 20 Aug 2026 15:26:09 +0300 Subject: [PATCH 05/16] =?UTF-8?q?ci(mcp):=20name=20the=20MCP=20gate=20per?= =?UTF-8?q?=20SDK=20major=20=E2=80=94=20MCP=20SDK=20v1/v2=20(Python=20X.Y)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Reshape tests-mcp-v2 into a tests-mcp matrix over {v1, v2} x {3.10, 3.14}, so the checks list shows an explicit per-major signal instead of the v1 side being buried inside the whole-repo tests matrix. The v1 leg uses the lockfile's mcp 1.x as-is; the v2 leg swaps in mcp>=2,<3 and drops jlowin fastmcp (pins mcp<2). The main tests matrix is unchanged. Generated-By: PostHog Code Task-Id: ebafcb71-b03b-443d-b40c-d527ed4a04f4 --- .github/workflows/ci.yml | 23 +++++++++++++++-------- posthog/test/mcp/conftest.py | 5 +++-- 2 files changed, 18 insertions(+), 10 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 99b0cb0d7..8232b58f9 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -151,16 +151,18 @@ jobs: run: | pytest --verbose --timeout=30 - tests-mcp-v2: - # The MCP suite again, against MCP Python SDK v2 (spec 2026-07-28). The - # `tests` matrix covers mcp 1.x on every Python version; this lane swaps - # in mcp>=2 (and drops jlowin fastmcp, which pins mcp<2) and runs only - # posthog/test/mcp — conftest.py there splits collection by major. - name: MCP SDK v2 tests (Python ${{ matrix.python-version }}) + tests-mcp: + # The MCP suite as a named gate per MCP Python SDK major. The v1 leg uses + # the lockfile's mcp 1.x (also exercised incidentally by the `tests` + # matrix — this leg exists as an explicit, named signal); the v2 leg + # (spec 2026-07-28) swaps in mcp>=2 and drops jlowin fastmcp, which pins + # mcp<2. posthog/test/mcp/conftest.py splits collection by major. + name: MCP SDK ${{ matrix.mcp-major }} (Python ${{ matrix.python-version }}) runs-on: ubuntu-latest strategy: matrix: python-version: ['3.10', '3.14'] + mcp-major: ['v1', 'v2'] steps: - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 @@ -178,14 +180,19 @@ jobs: version: "0.11.32" enable-cache: true - - name: Install test dependencies with MCP SDK v2 + - name: Install test dependencies shell: bash run: | UV_PROJECT_ENVIRONMENT=$pythonLocation uv sync --extra test + + - name: Swap in MCP SDK v2 + if: matrix.mcp-major == 'v2' + shell: bash + run: | uv pip uninstall --python $pythonLocation fastmcp uv pip install --python $pythonLocation 'mcp>=2,<3' - - name: Run MCP tests against SDK v2 + - name: Run MCP tests against SDK ${{ matrix.mcp-major }} run: | pytest posthog/test/mcp --verbose --timeout=30 diff --git a/posthog/test/mcp/conftest.py b/posthog/test/mcp/conftest.py index 10083fffc..65c0e887d 100644 --- a/posthog/test/mcp/conftest.py +++ b/posthog/test/mcp/conftest.py @@ -1,7 +1,8 @@ """Split the MCP test suite by installed MCP SDK major. -The suite runs twice in CI: once against ``mcp>=1.26,<2`` (the ``tests`` job) -and once against ``mcp>=2,<3`` (the ``tests-mcp-v2`` job). Files coupled to one +The suite runs against both majors in CI — the ``tests-mcp`` matrix has a +``mcp>=1.26,<2`` leg and a ``mcp>=2,<3`` leg (the main ``tests`` matrix also +exercises the v1 side incidentally). Files coupled to one major's seams import symbols the other major doesn't ship, so they are excluded from *collection* (a skip marker can't help — the failure is at import time). Version-agnostic files (units, truncation, session tokens, PostHogMCP, ids) From 94ac11310d2fa7f3f09e53d3c2442f56ab3e853c Mon Sep 17 00:00:00 2001 From: Georgis Andonis Date: Thu, 20 Aug 2026 17:13:57 +0300 Subject: [PATCH 06/16] fix(mcp): deliver the conversation handle via structuredContent too MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Manual testing against Claude Code found the conversation feature inert for any tool that declares an output schema: clients that read structuredContent never render the content blocks, so the [SERVER]: Reuse conversation_id text block was invisible and the agent had nothing to echo back. Every call minted a fresh handle, and each landed in its own session. Ports the second delivery channel from @posthog/mcp (ADR-0004, posthog-js #4430/#4431), which measured the same 0% echo rate before fixing it: - declare an optional `_mcp_instructions` key on the tool's advertised outputSchema at tools/list (never `required`) - mirror {conversation_id} into the result's structuredContent on *every* response, not just the minting one, so an agent that dropped the handle can read it back The declaration is what makes the write safe — clients validate structuredContent against the advertised schema, so an undeclared key fails the customer's whole tool result under additionalProperties: false. Only tools we declared on are ever written to, and an instance that never served a tools/list fails closed. Composed schemas (oneOf/allOf/anyOf/$ref) and tools owning the key are skipped; the text block still carries them. Wired into all four adapters (v1 FastMCP, v1 low-level, jlowin fastmcp, v2) with shape-tolerant reads: the (content, structured) tuple from FastMCP 1.x's convert_result path, CallToolResult models (structuredContent on 1.x, structured_content on 2.x), and plain dicts. Verified on the live playground server: outputSchema declares the key and structuredContent carries the handle alongside the tool's own payload. Generated-By: PostHog Code Task-Id: ebafcb71-b03b-443d-b40c-d527ed4a04f4 --- .sampo/changesets/mcp-sdk-v2-support.md | 2 +- posthog/mcp/_instrument_fastmcp.py | 37 +++- posthog/mcp/_instrument_lowlevel.py | 50 +++-- posthog/mcp/_instrument_v2.py | 45 ++++- posthog/mcp/_internal.py | 6 + posthog/mcp/_output_instructions.py | 185 ++++++++++++++++++ posthog/test/mcp/test_conversation_session.py | 42 ++++ posthog/test/mcp/test_output_instructions.py | 171 ++++++++++++++++ posthog/test/mcp/test_v2_mcpserver.py | 97 +++++++++ 9 files changed, 608 insertions(+), 27 deletions(-) create mode 100644 posthog/mcp/_output_instructions.py create mode 100644 posthog/test/mcp/test_output_instructions.py diff --git a/.sampo/changesets/mcp-sdk-v2-support.md b/.sampo/changesets/mcp-sdk-v2-support.md index 07736cbca..20a9fcf4e 100644 --- a/.sampo/changesets/mcp-sdk-v2-support.md +++ b/.sampo/changesets/mcp-sdk-v2-support.md @@ -2,4 +2,4 @@ posthog: minor --- -feat(mcp): support the MCP Python SDK 2.x and the 2026-07-28 spec revision. `instrument()` now wraps `mcp.server.mcpserver.MCPServer` (the renamed FastMCP) and the v2 low-level `Server` (constructor-injected handlers, string-keyed registry, late `add_request_handler` registrations included), capturing tool calls, tools/list, errors, intent, client identity, and `$mcp_protocol_version` on both protocol eras — SDK 1.x paths are unchanged. Conversation-anchored sessions land as the cross-pod correlation for the stateless era (parity with `@posthog/mcp`): with `enable_conversation_id`, `$session_id` is derived deterministically from the agent-echoed `conversation_id` (new export `derive_session_id_from_conversation`), only handles the SDK could have minted (uuidv7) anchor, and the prompt-back now rides errored results too so a first-call failure keeps the conversation together. `instrument()` also no longer crashes on an unsupported or unrecognized MCP SDK — it degrades to a logged no-op. +feat(mcp): support the MCP Python SDK 2.x and the 2026-07-28 spec revision. `instrument()` now wraps `mcp.server.mcpserver.MCPServer` (the renamed FastMCP) and the v2 low-level `Server` (constructor-injected handlers, string-keyed registry, late `add_request_handler` registrations included), capturing tool calls, tools/list, errors, intent, client identity, and `$mcp_protocol_version` on both protocol eras — SDK 1.x paths are unchanged. Conversation-anchored sessions land as the cross-pod correlation for the stateless era (parity with `@posthog/mcp`): with `enable_conversation_id`, `$session_id` is derived deterministically from the agent-echoed `conversation_id` (new export `derive_session_id_from_conversation`), only handles the SDK could have minted (uuidv7) anchor, and the prompt-back now rides errored results too so a first-call failure keeps the conversation together. The handle is delivered over both channels a tool result has — a `content` text block on the minting response, and an `_mcp_instructions` key declared on the tool's `outputSchema` and mirrored into `structuredContent` on every response. The second channel is what makes the feature work for tools with structured output at all: clients that read `structuredContent` never render `content`, so the agent had no handle to echo (0% echo rate measured against Claude Code before the mirror). `instrument()` also no longer crashes on an unsupported or unrecognized MCP SDK — it degrades to a logged no-op. diff --git a/posthog/mcp/_instrument_fastmcp.py b/posthog/mcp/_instrument_fastmcp.py index 5480586e7..a3ac50e72 100644 --- a/posthog/mcp/_instrument_fastmcp.py +++ b/posthog/mcp/_instrument_fastmcp.py @@ -50,6 +50,10 @@ resolve_session_and_client, ) from ._internal import MCPAnalyticsData +from ._output_instructions import ( + add_instructions_to_output_schema, + mirror_instructions_into_structured_content, +) from .logger import log from .tools import ( GET_MORE_TOOLS_NAME as _GET_MORE_TOOLS_NAME, @@ -178,16 +182,26 @@ async def wrapped( ) raise - # Inject the prompt-back first, then capture the delivered result. Only stamp - # a minted conversation_id when it was actually appended to what the agent got. + # Deliver the handle first, then capture the result the agent actually got. + # Two channels: mirrored into structuredContent on every response (for + # tools whose output schema we declared the key on — clients that read + # structuredContent never see the text block), and the prompt-back text + # block on the minting response only. delivered_conversation_id = conversation_id - if minted and conversation_id: - injected = _inject_prompt_back(result, conversation_id) - if injected is result: - delivered_conversation_id = ( - None # not injectable (e.g. tuple/scalar result) + if conversation_id: + delivered = False + if data.tool_output_instructions.get(name): + result, delivered = mirror_instructions_into_structured_content( + result, conversation_id ) - result = injected + if minted: + injected = _inject_prompt_back(result, conversation_id) + if injected is not result: + delivered = True + result = injected + # Only a minted handle can be lost — one the agent supplied, it has. + if not delivered: + delivered_conversation_id = None await record_tool_call( data, @@ -298,6 +312,13 @@ async def list_handler(req: Any) -> Any: tool.inputSchema = schema except Exception: # noqa: BLE001 - some schema attrs may be read-only log(f"WARN: could not set inputSchema on tool {tool.name}") + # Declare the structuredContent channel and remember the answer: + # clients that read structuredContent never see the content text + # block, and only a declared key may be written back on a call. + if data.options.enable_conversation_id: + data.tool_output_instructions[tool.name] = ( + add_instructions_to_output_schema(tool) + ) if data.options.report_missing: missing_name = resolve_missing_capability_tool_name(data.options) diff --git a/posthog/mcp/_instrument_lowlevel.py b/posthog/mcp/_instrument_lowlevel.py index fc443f27d..8d5d09790 100644 --- a/posthog/mcp/_instrument_lowlevel.py +++ b/posthog/mcp/_instrument_lowlevel.py @@ -44,6 +44,10 @@ resolve_session_and_client, ) from ._internal import MCPAnalyticsData +from ._output_instructions import ( + add_instructions_to_output_schema, + mirror_instructions_into_structured_content, +) from .logger import log from .tools import ( GET_MORE_TOOLS_NAME as _GET_MORE_TOOLS_NAME, @@ -188,22 +192,33 @@ async def handler(req: Any) -> Any: # CallToolResult(isError=True); record_tool_call detects that from the result. call_result = getattr(result, "root", result) - # Inject the prompt-back before capture; only stamp a minted conversation_id - # when it was actually delivered (non-list results can't carry it), so we - # don't record an orphan id the agent never received. Errored results carry - # it on purpose: a first-call failure is exactly when the agent needs the - # handle, or the retry starts a fresh conversation. + # Deliver the handle before capture, over both channels a result has: + # mirrored into structuredContent on every response (for tools whose + # output schema we declared the key on — clients that read + # structuredContent never see the text block), and the prompt-back text + # block on the minting response only. Only stamp a minted conversation_id + # when it actually reached the agent, so we don't record an orphan id. + # Errored results carry it on purpose: a first-call failure is exactly + # when the agent needs the handle, or the retry starts a fresh conversation. delivered_conversation_id = conversation_id - if minted and conversation_id: - content = getattr(call_result, "content", None) - if isinstance(content, list): - content.append( - mcp_types.TextContent( - type="text", text=build_prompt_back(conversation_id)["text"] - ) + if conversation_id: + delivered = False + if data.tool_output_instructions.get(name): + _, delivered = mirror_instructions_into_structured_content( + call_result, conversation_id ) - else: - delivered_conversation_id = None + if minted: + content = getattr(call_result, "content", None) + if isinstance(content, list): + content.append( + mcp_types.TextContent( + type="text", text=build_prompt_back(conversation_id)["text"] + ) + ) + delivered = True + # Only a minted handle can be lost — one the agent supplied, it has. + if not delivered: + delivered_conversation_id = None await record_tool_call( data, @@ -316,6 +331,13 @@ async def handler(req: Any) -> Any: tool.inputSchema = schema except Exception: # noqa: BLE001 log(f"WARN: could not set inputSchema on tool {tool.name}") + # Declare the structuredContent channel and remember the answer: + # clients that read structuredContent never see the content text + # block, and only a declared key may be written back on a call. + if data.options.enable_conversation_id: + data.tool_output_instructions[tool.name] = ( + add_instructions_to_output_schema(tool) + ) if data.options.report_missing: missing_name = resolve_missing_capability_tool_name(data.options) diff --git a/posthog/mcp/_instrument_v2.py b/posthog/mcp/_instrument_v2.py index 99af12bd9..44e1d6050 100644 --- a/posthog/mcp/_instrument_v2.py +++ b/posthog/mcp/_instrument_v2.py @@ -56,6 +56,10 @@ resolve_session_and_client, ) from ._internal import MCPAnalyticsData +from ._output_instructions import ( + add_instructions_to_output_schema, + mirror_instructions_into_structured_content, +) from .logger import log from .tools import ( GET_MORE_TOOLS_NAME as _GET_MORE_TOOLS_NAME, @@ -323,8 +327,12 @@ async def wrapped( duration_ms = (time.monotonic() - start) * 1000 delivered_conversation_id = conversation_id - if minted and conversation_id: - if not _append_prompt_back(result, conversation_id): + if conversation_id: + result, delivered = _deliver_conversation_id( + data, result, name, conversation_id, minted + ) + # Only a minted handle can be lost this way — one the agent supplied, it has. + if minted and not delivered: delivered_conversation_id = None await record_tool_call( @@ -365,6 +373,24 @@ def _append_prompt_back(result: Any, conversation_id: str) -> bool: return False +def _deliver_conversation_id( + data: MCPAnalyticsData, result: Any, name: str, conversation_id: str, minted: bool +) -> Tuple[Any, bool]: + """Hand the conversation handle back over both channels a result has: + mirrored into ``structuredContent`` on every response (for tools whose + output schema we declared the key on), and as a ``content`` text block on + the minting response only. Returns ``(result, delivered)`` — a minted handle + the agent never received must not be stamped on the event.""" + delivered = False + if data.tool_output_instructions.get(name): + result, delivered = mirror_instructions_into_structured_content( + result, conversation_id + ) + if minted and _append_prompt_back(result, conversation_id): + delivered = True + return result, delivered + + # --- low-level: tools/call ------------------------------------------------------ @@ -444,8 +470,12 @@ async def handler(ctx: Any, params: Any) -> Any: duration_ms = (time.monotonic() - start) * 1000 delivered_conversation_id = conversation_id - if minted and conversation_id: - if not _append_prompt_back(result, conversation_id): + if conversation_id: + result, delivered = _deliver_conversation_id( + data, result, name, conversation_id, minted + ) + # Only a minted handle can be lost this way — one the agent supplied, it has. + if minted and not delivered: delivered_conversation_id = None await record_tool_call( @@ -562,6 +592,13 @@ async def handler(ctx: Any, params: Any) -> Any: tool.input_schema = schema except Exception: # noqa: BLE001 - some schema attrs may be read-only log(f"WARN: could not set input_schema on tool {tool.name}") + # Declare the structuredContent channel and remember the answer: + # clients that read structuredContent never see the content text + # block, and only a declared key may be written back on a call. + if data.options.enable_conversation_id: + data.tool_output_instructions[tool.name] = ( + add_instructions_to_output_schema(tool) + ) if data.options.report_missing: missing_name = resolve_missing_capability_tool_name(data.options) diff --git a/posthog/mcp/_internal.py b/posthog/mcp/_internal.py index 5f77fee7c..1a29dbd47 100644 --- a/posthog/mcp/_internal.py +++ b/posthog/mcp/_internal.py @@ -66,6 +66,12 @@ class MCPAnalyticsData: identified_sessions: IdentityCache = field(default_factory=IdentityCache) tool_categories: Dict[str, str] = field(default_factory=dict) tool_descriptions: Dict[str, str] = field(default_factory=dict) + # Which tools got `_mcp_instructions` declared on their advertised output + # schema at tools/list. Only those may be mirrored into on a call — writing + # an undeclared key fails the customer's whole result under + # `additionalProperties: false`. Absent means "never served a listing for + # this tool", which fails closed. + tool_output_instructions: Dict[str, bool] = field(default_factory=dict) # Bounded FIFO of sessions we've emitted $mcp_initialize for, so a long-lived # server can't accumulate one entry per session forever. initialized_sessions: "OrderedDict[str, None]" = field(default_factory=OrderedDict) diff --git a/posthog/mcp/_output_instructions.py b/posthog/mcp/_output_instructions.py new file mode 100644 index 000000000..d33d3d581 --- /dev/null +++ b/posthog/mcp/_output_instructions.py @@ -0,0 +1,185 @@ +# Portions of this package are derived from MCPCat/mcpcat-typescript-sdk +# Copyright (c) 2025 MCPcat +# Licensed under the MIT License: https://github.com/MCPCat/mcpcat-typescript-sdk/blob/main/LICENSE + +"""Second delivery channel for the conversation handle: ``structuredContent``. + +Two halves that must stay in this order: + + 1. declare ``_mcp_instructions`` on the tool's advertised output schema + (at ``tools/list``) + 2. write it into the result's ``structuredContent`` (on every response) + +Needed because clients that read ``structuredContent`` — which they do whenever +a tool declares an output schema — never see the ``content`` text block that +carries the handle. Measured against Claude Code, the echo rate was 100% for +schema-less tools and 0% for schema-declaring ones before this mirror existed +(posthog-js ADR-0004). + +The declaration is what makes the write safe: MCP clients validate +``structuredContent`` against the schema from ``tools/list``, and generated +schemas commonly carry ``additionalProperties: false``, so an undeclared key is +not ignored — it fails the entire tool result. Only tools that got the +declaration are ever written to, and an instance that never served a listing +fails closed. +""" + +from __future__ import annotations + +import copy +from typing import Any, Dict, Optional, Tuple + +from .logger import log + +MCP_INSTRUCTIONS_KEY = "_mcp_instructions" + +_INSTRUCTIONS_FIELD_DESCRIPTION = "Server-issued metadata for this conversation." +_CONVERSATION_ID_FIELD_DESCRIPTION = "The server-issued conversation identifier." + +# `outputSchema` on MCP SDK 1.x models, `output_schema` on 2.x (same wire field). +_OUTPUT_SCHEMA_ATTRS = ("outputSchema", "output_schema") +_STRUCTURED_CONTENT_ATTRS = ("structuredContent", "structured_content") + + +def _read_attr(obj: Any, names: Tuple[str, ...]) -> Tuple[Optional[str], Any]: + """The first of ``names`` this object actually carries, and its value.""" + for name in names: + if hasattr(obj, name): + return name, getattr(obj, name) + return None, None + + +def can_declare_output_instructions(output_schema: Any) -> bool: + """True when :data:`MCP_INSTRUCTIONS_KEY` can safely be declared on this + tool's advertised output schema. + + Requires a plain-object JSON Schema we can extend. A tool with no output + schema has nothing to mirror into and keeps working through ``content``; a + composed schema (``oneOf``/``allOf``/``anyOf``/``$ref``) has no single + ``properties`` bag to add to. Both stay content-only, matching the policy on + the input side. + """ + if not isinstance(output_schema, dict): + return False + if ( + output_schema.get("$ref") + or output_schema.get("oneOf") + or output_schema.get("allOf") + or output_schema.get("anyOf") + ): + return False + properties = output_schema.get("properties") + # A malformed `properties` is harmless until we try to declare into it, so + # refuse rather than raise inside the tools/list wrapper and fail the listing. + if properties is not None and not isinstance(properties, dict): + return False + return not properties or MCP_INSTRUCTIONS_KEY not in properties + + +def add_instructions_to_output_schema(tool: Any) -> bool: + """Declare an optional :data:`MCP_INSTRUCTIONS_KEY` on ``tool``'s output + schema, in place. Returns whether the declaration was made — the caller + records that answer as ownership, and only declared tools are ever mirrored + into. + + The property is never added to ``required``: a result without it must stay + valid, since every tool result predating this change lacks it. + """ + attr, original = _read_attr(tool, _OUTPUT_SCHEMA_ATTRS) + if attr is None or not original: + # No output schema means the client reads `content`, where the handle + # already rides. Nothing to declare, and nothing broken. + return False + + name = getattr(tool, "name", "unknown") + if not can_declare_output_instructions(original): + properties = original.get("properties") if isinstance(original, dict) else None + if isinstance(properties, dict) and MCP_INSTRUCTIONS_KEY in properties: + log( + f"WARN: Tool \"{name}\" already declares '{MCP_INSTRUCTIONS_KEY}' in its " + "output schema. Leaving it alone." + ) + else: + log( + f'WARN: Tool "{name}" has a complex output schema (oneOf/allOf/anyOf/$ref). ' + f"Skipping '{MCP_INSTRUCTIONS_KEY}' declaration; its session handle stays " + "content-only." + ) + return False + + # Deep copy: the server may reuse or freeze the schema object it handed us. + schema = copy.deepcopy(original) + if not isinstance(schema.get("properties"), dict): + schema["properties"] = {} + schema["properties"][MCP_INSTRUCTIONS_KEY] = { + "type": "object", + "description": _INSTRUCTIONS_FIELD_DESCRIPTION, + "properties": { + "conversation_id": { + "type": "string", + "description": _CONVERSATION_ID_FIELD_DESCRIPTION, + } + }, + } + try: + setattr(tool, attr, schema) + except Exception: # noqa: BLE001 - some schema attrs may be read-only + log(f"WARN: could not set {attr} on tool {name}") + return False + return True + + +def build_conversation_instructions(conversation_id: str) -> Dict[str, Any]: + """The payload mirrored into ``structuredContent``.""" + return {"conversation_id": conversation_id} + + +def mirror_instructions_into_structured_content( + result: Any, conversation_id: str +) -> Tuple[Any, bool]: + """Write :data:`MCP_INSTRUCTIONS_KEY` into a result's ``structuredContent``. + + Returns ``(result, delivered)``. Unlike the text block this rides *every* + response rather than only the one that minted the handle, so an agent that + dropped it can read it back. + + Handles every shape a tool result takes across the adapters: the + ``(content, structured)`` tuple from FastMCP 1.x's ``convert_result`` path, + a ``CallToolResult`` model (``structuredContent`` on 1.x, + ``structured_content`` on 2.x), and a plain dict. Leaves the result untouched + when there is no plain-object structured content to extend, or when the tool + already produced its own key — customer data wins. + """ + payload = build_conversation_instructions(conversation_id) + + # FastMCP 1.x convert_result path: (content_list, structured) + if isinstance(result, tuple) and len(result) == 2: + structured = result[1] + if not isinstance(structured, dict) or MCP_INSTRUCTIONS_KEY in structured: + return result, False + return (result[0], {**structured, MCP_INSTRUCTIONS_KEY: payload}), True + + if isinstance(result, dict): + for key in _STRUCTURED_CONTENT_ATTRS: + structured = result.get(key) + if isinstance(structured, dict) and MCP_INSTRUCTIONS_KEY not in structured: + return { + **result, + key: {**structured, MCP_INSTRUCTIONS_KEY: payload}, + }, True + return result, False + + # CallToolResult model (or the ServerResult wrapper around one). + target = getattr(result, "root", result) + attr, structured = _read_attr(target, _STRUCTURED_CONTENT_ATTRS) + if ( + attr is None + or not isinstance(structured, dict) + or MCP_INSTRUCTIONS_KEY in structured + ): + return result, False + try: + setattr(target, attr, {**structured, MCP_INSTRUCTIONS_KEY: payload}) + except Exception: # noqa: BLE001 - never let delivery break the tool path + return result, False + return result, True diff --git a/posthog/test/mcp/test_conversation_session.py b/posthog/test/mcp/test_conversation_session.py index bbc3d195f..b926770ca 100644 --- a/posthog/test/mcp/test_conversation_session.py +++ b/posthog/test/mcp/test_conversation_session.py @@ -187,6 +187,48 @@ def echo(msg: str) -> str: assert [c["properties"]["$session_id"] for c in calls] == [expected, expected] +@pytest.mark.skipif(MCP_MAJOR != 1, reason="v1 FastMCP server") +async def test_v1_structured_output_tool_gets_the_conversation_handle(): + """Clients that read ``structuredContent`` never render ``content``, so a + tool declaring an output schema must get the handle mirrored into the + structured half or the agent can never echo it back (ADR-0004).""" + from typing import Any + + import mcp.types as mcp_types + from mcp.server.fastmcp import FastMCP + + from posthog.mcp import instrument + from posthog.mcp._output_instructions import MCP_INSTRUCTIONS_KEY + + server = FastMCP("structured-v1") + + @server.tool() + def totals(event: str) -> dict[str, Any]: + return {"event": event, "total": 7} + + client = FakeClient() + instrument(server, client, MCPAnalyticsOptions(enable_conversation_id=True)) + + handler = server._mcp_server.request_handlers[mcp_types.ListToolsRequest] + listed = await handler(mcp_types.ListToolsRequest(method="tools/list")) + tool = next(t for t in listed.root.tools if t.name == "totals") + assert MCP_INSTRUCTIONS_KEY in tool.outputSchema["properties"] + + result = await server._tool_manager.call_tool( + "totals", {"event": "pageview", "context": "structured"}, convert_result=True + ) + await _flush() + + structured = result[1] + handle = structured[MCP_INSTRUCTIONS_KEY]["conversation_id"] + assert handle + assert structured["total"] == 7 + assert ( + _events(client, "$mcp_tool_call")[0]["properties"]["$mcp_conversation_id"] + == handle + ) + + @pytest.mark.skipif(MCP_MAJOR != 1, reason="v1 FastMCP server") async def test_v1_feature_off_keeps_transport_sessions(): from mcp.server.fastmcp import FastMCP diff --git a/posthog/test/mcp/test_output_instructions.py b/posthog/test/mcp/test_output_instructions.py new file mode 100644 index 000000000..312ffe996 --- /dev/null +++ b/posthog/test/mcp/test_output_instructions.py @@ -0,0 +1,171 @@ +"""The ``structuredContent`` delivery channel for the conversation handle. + +Clients that read ``structuredContent`` — which they do whenever a tool declares +an output schema — never see the ``content`` text block carrying the handle, so +the echo rate on schema-declaring tools is 0% without this mirror +(posthog-js ADR-0004). Runs under both MCP SDK majors. +""" + +from types import SimpleNamespace + +from posthog.mcp._output_instructions import ( + MCP_INSTRUCTIONS_KEY, + add_instructions_to_output_schema, + can_declare_output_instructions, + mirror_instructions_into_structured_content, +) + + +def _tool(output_schema, name="demo"): + return SimpleNamespace(name=name, outputSchema=output_schema) + + +# --- declaration (tools/list half) --------------------------------------------- + + +def test_declares_key_on_a_plain_object_schema(): + tool = _tool({"type": "object", "properties": {"total": {"type": "integer"}}}) + + assert add_instructions_to_output_schema(tool) is True + + props = tool.outputSchema["properties"] + assert ( + props[MCP_INSTRUCTIONS_KEY]["properties"]["conversation_id"]["type"] == "string" + ) + # never required — every result predating this change lacks the key + assert MCP_INSTRUCTIONS_KEY not in tool.outputSchema.get("required", []) + # the customer's own properties survive + assert "total" in props + + +def test_declaration_does_not_mutate_the_original_schema(): + original = {"type": "object", "properties": {"total": {"type": "integer"}}} + tool = _tool(original) + + add_instructions_to_output_schema(tool) + + assert MCP_INSTRUCTIONS_KEY not in original["properties"] + + +def test_no_output_schema_is_not_declared(): + tool = _tool(None) + assert add_instructions_to_output_schema(tool) is False + assert tool.outputSchema is None + + +def test_composed_schemas_are_skipped(): + for schema in ( + {"oneOf": [{"type": "object"}]}, + {"allOf": [{"type": "object"}]}, + {"anyOf": [{"type": "object"}]}, + {"$ref": "#/defs/Thing"}, + ): + assert can_declare_output_instructions(schema) is False + tool = _tool(schema) + assert add_instructions_to_output_schema(tool) is False + + +def test_tool_owning_the_key_is_left_alone(): + schema = { + "type": "object", + "properties": {MCP_INSTRUCTIONS_KEY: {"type": "string"}}, + } + tool = _tool(schema) + + assert add_instructions_to_output_schema(tool) is False + # customer's declaration untouched + assert tool.outputSchema["properties"][MCP_INSTRUCTIONS_KEY] == {"type": "string"} + + +def test_malformed_properties_are_refused_rather_than_raising(): + assert ( + can_declare_output_instructions({"type": "object", "properties": True}) is False + ) + assert ( + can_declare_output_instructions({"type": "object", "properties": []}) is False + ) + assert ( + add_instructions_to_output_schema(_tool({"type": "object", "properties": True})) + is False + ) + + +def test_snake_case_output_schema_attr_is_supported(): + """MCP SDK 2.x models expose ``output_schema``; 1.x exposes ``outputSchema``.""" + tool = SimpleNamespace( + name="v2tool", output_schema={"type": "object", "properties": {}} + ) + + assert add_instructions_to_output_schema(tool) is True + assert MCP_INSTRUCTIONS_KEY in tool.output_schema["properties"] + + +# --- mirroring (tools/call half) ------------------------------------------------ + + +def test_mirrors_into_a_fastmcp_v1_tuple_result(): + result = ([{"type": "text", "text": "{}"}], {"total": 7}) + + mirrored, delivered = mirror_instructions_into_structured_content(result, "conv-1") + + assert delivered is True + assert mirrored[1][MCP_INSTRUCTIONS_KEY] == {"conversation_id": "conv-1"} + assert mirrored[1]["total"] == 7 + assert mirrored[0] is result[0] # content untouched + + +def test_mirrors_into_a_model_result_both_attr_shapes(): + for attr in ("structuredContent", "structured_content"): + result = SimpleNamespace(**{attr: {"total": 7}}) + + _, delivered = mirror_instructions_into_structured_content(result, "conv-2") + + assert delivered is True + assert getattr(result, attr)[MCP_INSTRUCTIONS_KEY] == { + "conversation_id": "conv-2" + } + + +def test_mirrors_through_a_serverresult_wrapper(): + inner = SimpleNamespace(structuredContent={"total": 1}) + result = SimpleNamespace(root=inner) + + _, delivered = mirror_instructions_into_structured_content(result, "conv-3") + + assert delivered is True + assert inner.structuredContent[MCP_INSTRUCTIONS_KEY] == { + "conversation_id": "conv-3" + } + + +def test_mirrors_into_a_dict_result_without_mutating_it(): + result = {"structuredContent": {"total": 7}} + + mirrored, delivered = mirror_instructions_into_structured_content(result, "conv-4") + + assert delivered is True + assert mirrored["structuredContent"][MCP_INSTRUCTIONS_KEY] == { + "conversation_id": "conv-4" + } + assert MCP_INSTRUCTIONS_KEY not in result["structuredContent"] + + +def test_no_structured_content_is_not_delivered(): + for result in ( + SimpleNamespace(content=[]), # content-only tool + {"content": []}, + ([{"type": "text"}], None), # tuple with no structured half + "not a result", + None, + ): + _, delivered = mirror_instructions_into_structured_content(result, "conv-5") + assert delivered is False + + +def test_customer_key_wins_over_the_mirror(): + result = SimpleNamespace(structuredContent={MCP_INSTRUCTIONS_KEY: "mine"}) + + _, delivered = mirror_instructions_into_structured_content(result, "conv-6") + + assert delivered is False + assert result.structuredContent[MCP_INSTRUCTIONS_KEY] == "mine" diff --git a/posthog/test/mcp/test_v2_mcpserver.py b/posthog/test/mcp/test_v2_mcpserver.py index f591cc17e..1c75fd935 100644 --- a/posthog/test/mcp/test_v2_mcpserver.py +++ b/posthog/test/mcp/test_v2_mcpserver.py @@ -6,6 +6,8 @@ ``ServerRequestContext``, exactly as the v2 runner would invoke them. """ +from typing import Any + import mcp.types as mcp_types from mcp.server.mcpserver import MCPServer @@ -239,6 +241,101 @@ def identify(request, extra): assert set(captured_extra) <= {"session_id"} +async def test_structured_output_tool_gets_the_conversation_handle(): + """A tool that declares an output schema is served to clients that read + ``structuredContent`` and never render ``content`` — so the handle has to + ride the structured half too, or the agent can never echo it back.""" + from posthog.mcp._output_instructions import MCP_INSTRUCTIONS_KEY + + server = MCPServer("structured-v2") + + @server.tool() + def totals(event: str) -> dict[str, Any]: + return {"event": event, "total": 7} + + client = FakeClient() + instrument(server, client, MCPAnalyticsOptions(enable_conversation_id=True)) + + listed = await _list_tools(server) + tool = next(t for t in listed.tools if t.name == "totals") + # 1. declared on the advertised output schema, and never required + assert MCP_INSTRUCTIONS_KEY in tool.output_schema["properties"] + assert MCP_INSTRUCTIONS_KEY not in tool.output_schema.get("required", []) + + result = await _call_tool( + server, "totals", {"event": "pageview", "context": "structured"} + ) + await _flush() + + # 2. mirrored into the result the client actually reads + handle = result.structured_content[MCP_INSTRUCTIONS_KEY]["conversation_id"] + assert handle + # the tool's own payload is intact alongside it + assert result.structured_content["total"] == 7 + # and it matches what analytics recorded + assert ( + _events(client, "$mcp_tool_call")[0]["properties"]["$mcp_conversation_id"] + == handle + ) + + +async def test_mirror_rides_every_response_not_just_the_minting_one(): + """Unlike the text block, the structured mirror repeats — an agent that + dropped the handle can read it back on any later response.""" + from posthog.mcp._output_instructions import MCP_INSTRUCTIONS_KEY + + server = MCPServer("structured-v2-repeat") + + @server.tool() + def totals(event: str) -> dict[str, Any]: + return {"event": event, "total": 7} + + client = FakeClient() + instrument(server, client, MCPAnalyticsOptions(enable_conversation_id=True)) + await _list_tools(server) + + first = await _call_tool(server, "totals", {"event": "a", "context": "first"}) + handle = first.structured_content[MCP_INSTRUCTIONS_KEY]["conversation_id"] + + # the agent echoes it; the mirror must still be present on this response + second = await _call_tool( + server, "totals", {"event": "b", "context": "second", "conversation_id": handle} + ) + await _flush() + + assert second.structured_content[MCP_INSTRUCTIONS_KEY]["conversation_id"] == handle + sessions = { + c["properties"]["$session_id"] for c in _events(client, "$mcp_tool_call") + } + assert len(sessions) == 1 + + +async def test_mirror_is_skipped_when_no_listing_declared_the_key(): + """Fail closed: writing an undeclared key fails the customer's whole result + under ``additionalProperties: false``, so an instance that never served a + tools/list must not mirror.""" + from posthog.mcp._output_instructions import MCP_INSTRUCTIONS_KEY + + server = MCPServer("structured-v2-nolist") + + @server.tool() + def totals(event: str) -> dict[str, Any]: + return {"event": event, "total": 7} + + client = FakeClient() + instrument(server, client, MCPAnalyticsOptions(enable_conversation_id=True)) + + # no _list_tools() call — ownership is unknown + result = await _call_tool(server, "totals", {"event": "a", "context": "no listing"}) + await _flush() + + assert MCP_INSTRUCTIONS_KEY not in (result.structured_content or {}) + # the text-block channel still carried it, so the handle is still recorded + assert _events(client, "$mcp_tool_call")[0]["properties"].get( + "$mcp_conversation_id" + ) + + async def test_report_missing_advertises_and_captures(): server = make_server() client = FakeClient() From e9c490dcb8561fd87b9f1179392f0417142a1e0b Mon Sep 17 00:00:00 2001 From: Georgis Andonis Date: Thu, 20 Aug 2026 17:53:27 +0300 Subject: [PATCH 07/16] feat(mcp): uniform ctx in callbacks + exported get_request_headers MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adding v2 support left the callback contract divergent: the v2 adapters passed the request context as extra["ctx"] while the v1 adapters passed only session_id, so an identify() written against one major silently got nothing on the other — and on v1 there was no way to reach headers from extra at all. Header-based identification is the main use of identify, and the failure is invisible: no headers, no distinct id, every event anonymous. - all four adapters now pass the SDK's own per-request context as extra["ctx"], unchanged and identically shaped across majors - new exported get_request_headers(extra) flattens Starlette Headers, plain mappings, or any iterable of pairs into a lowercase-keyed dict; returns None on stdio and never raises - the context object still never reaches an event: the scalar projection added earlier for the $identify capture strips it Ports the intent of @posthog/mcp's getRequestHeaders (ADR-0006: hand hosts the raw context rather than synthesising a fake uniform shape, and give them one helper for the thing they actually need). Generated-By: PostHog Code Task-Id: ebafcb71-b03b-443d-b40c-d527ed4a04f4 --- .sampo/changesets/mcp-sdk-v2-support.md | 2 +- posthog/mcp/__init__.py | 5 ++ posthog/mcp/_instrument_fastmcp.py | 14 ++++- posthog/mcp/_instrument_lowlevel.py | 12 +++- posthog/mcp/request_headers.py | 69 +++++++++++++++++++++++ posthog/test/mcp/test_fastmcp.py | 32 +++++++++++ posthog/test/mcp/test_request_headers.py | 72 ++++++++++++++++++++++++ posthog/test/mcp/test_v2_lowlevel.py | 22 ++++++++ references/public_api_snapshot.txt | 3 + 9 files changed, 226 insertions(+), 5 deletions(-) create mode 100644 posthog/mcp/request_headers.py create mode 100644 posthog/test/mcp/test_request_headers.py diff --git a/.sampo/changesets/mcp-sdk-v2-support.md b/.sampo/changesets/mcp-sdk-v2-support.md index 20a9fcf4e..be30a386e 100644 --- a/.sampo/changesets/mcp-sdk-v2-support.md +++ b/.sampo/changesets/mcp-sdk-v2-support.md @@ -2,4 +2,4 @@ posthog: minor --- -feat(mcp): support the MCP Python SDK 2.x and the 2026-07-28 spec revision. `instrument()` now wraps `mcp.server.mcpserver.MCPServer` (the renamed FastMCP) and the v2 low-level `Server` (constructor-injected handlers, string-keyed registry, late `add_request_handler` registrations included), capturing tool calls, tools/list, errors, intent, client identity, and `$mcp_protocol_version` on both protocol eras — SDK 1.x paths are unchanged. Conversation-anchored sessions land as the cross-pod correlation for the stateless era (parity with `@posthog/mcp`): with `enable_conversation_id`, `$session_id` is derived deterministically from the agent-echoed `conversation_id` (new export `derive_session_id_from_conversation`), only handles the SDK could have minted (uuidv7) anchor, and the prompt-back now rides errored results too so a first-call failure keeps the conversation together. The handle is delivered over both channels a tool result has — a `content` text block on the minting response, and an `_mcp_instructions` key declared on the tool's `outputSchema` and mirrored into `structuredContent` on every response. The second channel is what makes the feature work for tools with structured output at all: clients that read `structuredContent` never render `content`, so the agent had no handle to echo (0% echo rate measured against Claude Code before the mirror). `instrument()` also no longer crashes on an unsupported or unrecognized MCP SDK — it degrades to a logged no-op. +feat(mcp): support the MCP Python SDK 2.x and the 2026-07-28 spec revision. `instrument()` now wraps `mcp.server.mcpserver.MCPServer` (the renamed FastMCP) and the v2 low-level `Server` (constructor-injected handlers, string-keyed registry, late `add_request_handler` registrations included), capturing tool calls, tools/list, errors, intent, client identity, and `$mcp_protocol_version` on both protocol eras — SDK 1.x paths are unchanged. Conversation-anchored sessions land as the cross-pod correlation for the stateless era (parity with `@posthog/mcp`): with `enable_conversation_id`, `$session_id` is derived deterministically from the agent-echoed `conversation_id` (new export `derive_session_id_from_conversation`), only handles the SDK could have minted (uuidv7) anchor, and the prompt-back now rides errored results too so a first-call failure keeps the conversation together. The handle is delivered over both channels a tool result has — a `content` text block on the minting response, and an `_mcp_instructions` key declared on the tool's `outputSchema` and mirrored into `structuredContent` on every response. The second channel is what makes the feature work for tools with structured output at all: clients that read `structuredContent` never render `content`, so the agent had no handle to echo (0% echo rate measured against Claude Code before the mirror). Host callbacks (`identify`, `intent_fallback`, `event_properties`, `before_send`) now receive the SDK's own per-request context as `extra["ctx"]` on both SDK majors, and a new exported `get_request_headers(extra)` reads HTTP headers off it — the shape differs between majors, and a hand-rolled read that works on one silently returns nothing on the other, which sends every event out anonymous. `instrument()` also no longer crashes on an unsupported or unrecognized MCP SDK — it degrades to a logged no-op. diff --git a/posthog/mcp/__init__.py b/posthog/mcp/__init__.py index 66645cc1e..8814427b6 100644 --- a/posthog/mcp/__init__.py +++ b/posthog/mcp/__init__.py @@ -51,6 +51,7 @@ ) from .logger import log, set_logger from .posthog_mcp import PostHogMCP +from .request_headers import get_request_headers from .session import ( derive_session_id_from_conversation, derive_session_id_from_mcp_session, @@ -88,6 +89,10 @@ "CaptureEventData", "PreparedToolCall", "get_more_tools_result", + # Read HTTP headers inside identify/intent_fallback/event_properties/ + # before_send callbacks on either SDK major: the raw per-request context + # arrives as extra["ctx"] and its shape differs between them. + "get_request_headers", "derive_session_id_from_mcp_session", # Conversation-anchored sessions: the cross-SDK derivation contract with # posthog-js (the 2026-07-28 revision has no protocol sessions, so the diff --git a/posthog/mcp/_instrument_fastmcp.py b/posthog/mcp/_instrument_fastmcp.py index a3ac50e72..2ac25c861 100644 --- a/posthog/mcp/_instrument_fastmcp.py +++ b/posthog/mcp/_instrument_fastmcp.py @@ -103,7 +103,14 @@ async def wrapped( ) ) request = build_tool_call_request(name, arguments) - extra: Dict[str, Any] = {"session_id": mcp_session_id} + # `ctx` is the SDK's own per-request context, handed to host callbacks + # unchanged and identically on both SDK majors (read headers off it with + # the exported `get_request_headers`). Never captured — the event + # pipeline keeps only a scalar projection of `extra`. + extra: Dict[str, Any] = { + "session_id": mcp_session_id, + "ctx": getattr(context, "request_context", None), + } # Resolve the conversation handle before the session: when the agent # carries (or is about to receive) one, it anchors $session_id for every @@ -249,7 +256,10 @@ async def list_handler(req: Any) -> Any: ) ) request = request_to_dict(req) - extra: Dict[str, Any] = {"session_id": mcp_session_id} + extra: Dict[str, Any] = { + "session_id": mcp_session_id, + "ctx": _low_level_request_context(server), + } # Resolve session, emit $mcp_initialize (once per session) and identify here # too — a client may list tools without ever calling one. session_id = await prepare_request( diff --git a/posthog/mcp/_instrument_lowlevel.py b/posthog/mcp/_instrument_lowlevel.py index 8d5d09790..c9d44e0b4 100644 --- a/posthog/mcp/_instrument_lowlevel.py +++ b/posthog/mcp/_instrument_lowlevel.py @@ -106,7 +106,11 @@ async def handler(req: Any) -> Any: ) ) request = build_tool_call_request(name, arguments) - extra = {"session_id": mcp_session_id} + # `ctx` is the SDK's own per-request context, handed to host callbacks + # unchanged and identically on both SDK majors (read headers off it with + # the exported `get_request_headers`). Never captured — the event + # pipeline keeps only a scalar projection of `extra`. + extra = {"session_id": mcp_session_id, "ctx": _request_context(server)} # Resolve the conversation handle before the session: when present it # anchors $session_id for every event of this request (ADR-0004). @@ -262,7 +266,11 @@ async def handler(req: Any) -> Any: ) ) request = request_to_dict(req) - extra = {"session_id": mcp_session_id} + # `ctx` is the SDK's own per-request context, handed to host callbacks + # unchanged and identically on both SDK majors (read headers off it with + # the exported `get_request_headers`). Never captured — the event + # pipeline keeps only a scalar projection of `extra`. + extra = {"session_id": mcp_session_id, "ctx": _request_context(server)} # Resolve session, emit $mcp_initialize (once per session) and identify here # too — a client may list tools without ever calling one. session_id = await prepare_request( diff --git a/posthog/mcp/request_headers.py b/posthog/mcp/request_headers.py new file mode 100644 index 000000000..2b7edcac4 --- /dev/null +++ b/posthog/mcp/request_headers.py @@ -0,0 +1,69 @@ +# Portions of this package are derived from MCPCat/mcpcat-typescript-sdk +# Copyright (c) 2025 MCPcat +# Licensed under the MIT License: https://github.com/MCPCat/mcpcat-typescript-sdk/blob/main/LICENSE + +"""Read HTTP request headers inside a host callback, on either MCP SDK major. + +``identify``, ``intent_fallback``, ``event_properties`` and ``before_send`` +receive the SDK's own per-request context under ``extra["ctx"]``, unchanged. We +deliberately do not synthesise a uniform shape for it: the two majors expose +different objects, and a fabricated one is a convincing partial lie about a +shape the SDK actually changed. Headers are the one thing nearly every callback +wants, so they get a helper instead:: + + from posthog.mcp import get_request_headers + + def identify(request, extra): + headers = get_request_headers(extra) or {} + token = headers.get("authorization") + ... + +Returns ``None`` when the request did not arrive over HTTP — stdio and +in-memory transports carry no headers at all. +""" + +from __future__ import annotations + +from typing import Any, Dict, Optional + +__all__ = ["get_request_headers"] + +RequestHeaderBag = Dict[str, str] + + +def get_request_headers(extra: Any) -> Optional[RequestHeaderBag]: + """The request's HTTP headers as a plain dict with lowercase keys, or ``None``. + + Accepts the ``extra`` dict handed to a callback, or the raw per-request + context itself, so it works whichever one a host happens to hold. + """ + ctx = extra + if isinstance(extra, dict): + ctx = extra.get("ctx") + if ctx is None: + return None + + # Both majors reach the transport's request the same way from their own + # context object (`ServerRequestContext` on 2.x, `RequestContext` on 1.x); + # `request` is None on stdio. + source = getattr(getattr(ctx, "request", None), "headers", None) + if source is None: + return None + return _to_header_bag(source) + + +def _to_header_bag(source: Any) -> Optional[RequestHeaderBag]: + """Flatten a Starlette ``Headers``, a mapping, or anything iterable of pairs + into a lowercase-keyed dict. Never raises: a header read must not take a + tool call down with it.""" + try: + # Starlette's Headers and dict both expose .items(); Headers already + # lowercases, a plain dict may not, so normalise either way. + items = source.items() if hasattr(source, "items") else source + bag: RequestHeaderBag = {} + for key, value in items: + if isinstance(key, str) and isinstance(value, str): + bag[key.lower()] = value + return bag + except Exception: # noqa: BLE001 - best effort, never break the tool path + return None diff --git a/posthog/test/mcp/test_fastmcp.py b/posthog/test/mcp/test_fastmcp.py index 1784835ff..087863f0c 100644 --- a/posthog/test/mcp/test_fastmcp.py +++ b/posthog/test/mcp/test_fastmcp.py @@ -205,3 +205,35 @@ async def test_unsupported_server_returns_noop_handle(): # graceful no-op: capture and flush do nothing and do not raise await handle.capture("anything") await handle.flush() + + +async def test_callbacks_can_read_headers_through_the_helper(): + """Same callback body as the v2 lane's equivalent test: `extra["ctx"]` is the + SDK's own per-request context on both majors, read via `get_request_headers`.""" + from types import SimpleNamespace + + from posthog.mcp import get_request_headers + + server = make_server() + client = FakeClient() + seen = {} + + def identify(request, extra): + seen["headers"] = get_request_headers(extra) + return None + + instrument(server, client, MCPAnalyticsOptions(identify=identify)) + + # FastMCP hands the tool a Context whose .request_context carries the request. + context = SimpleNamespace( + request_context=SimpleNamespace( + request=SimpleNamespace(headers={"Authorization": "Bearer t0ken"}), + session=SimpleNamespace(client_params=None), + ) + ) + await server._tool_manager.call_tool( + "add", {"a": 1, "b": 1, "context": "header read"}, context=context + ) + await _flush() + + assert seen["headers"] == {"authorization": "Bearer t0ken"} diff --git a/posthog/test/mcp/test_request_headers.py b/posthog/test/mcp/test_request_headers.py new file mode 100644 index 000000000..0c6b752f1 --- /dev/null +++ b/posthog/test/mcp/test_request_headers.py @@ -0,0 +1,72 @@ +"""``get_request_headers`` — reading HTTP headers inside a host callback. + +Host callbacks receive the SDK's own per-request context under ``extra["ctx"]``, +unchanged. Its shape differs between MCP SDK majors, so header reads go through +this helper instead of a hand-rolled path that silently returns ``None`` on the +other major (the failure mode is invisible: ``identify`` returns nothing and +every event goes out anonymous). Runs under both majors. +""" + +from types import SimpleNamespace + +from posthog.mcp import get_request_headers + + +def _ctx(headers): + return SimpleNamespace(request=SimpleNamespace(headers=headers)) + + +def test_reads_headers_from_a_callback_extra(): + extra = {"session_id": "abc", "ctx": _ctx({"Authorization": "Bearer t0ken"})} + + assert get_request_headers(extra) == {"authorization": "Bearer t0ken"} + + +def test_keys_are_lowercased(): + extra = { + "ctx": _ctx({"X-Anthropic-Client": "claude-code", "USER-AGENT": "probe/1"}) + } + + assert get_request_headers(extra) == { + "x-anthropic-client": "claude-code", + "user-agent": "probe/1", + } + + +def test_accepts_a_raw_context_too(): + """A host holding the context itself shouldn't have to wrap it in a dict.""" + assert get_request_headers(_ctx({"a": "b"})) == {"a": "b"} + + +def test_starlette_style_headers_are_supported(): + class Headers: + """Starlette's Headers: already-lowercased, .items() of pairs.""" + + def items(self): + return [("content-type", "application/json"), ("mcp-session-id", "s1")] + + assert get_request_headers({"ctx": _ctx(Headers())}) == { + "content-type": "application/json", + "mcp-session-id": "s1", + } + + +def test_stdio_and_missing_context_return_none(): + assert get_request_headers({"ctx": _ctx(None)}) is None # HTTP-less transport + assert get_request_headers({"ctx": SimpleNamespace(request=None)}) is None + assert get_request_headers({"session_id": "abc"}) is None # no ctx at all + assert get_request_headers(None) is None + + +def test_never_raises_on_a_hostile_header_object(): + class Exploding: + def items(self): + raise RuntimeError("nope") + + assert get_request_headers({"ctx": _ctx(Exploding())}) is None + + +def test_non_string_header_values_are_skipped(): + extra = {"ctx": _ctx({"good": "yes", "bad": 42, 7: "alsobad"})} + + assert get_request_headers(extra) == {"good": "yes"} diff --git a/posthog/test/mcp/test_v2_lowlevel.py b/posthog/test/mcp/test_v2_lowlevel.py index baec19986..788bfd207 100644 --- a/posthog/test/mcp/test_v2_lowlevel.py +++ b/posthog/test/mcp/test_v2_lowlevel.py @@ -247,3 +247,25 @@ async def test_report_missing_appends_virtual_tool(): assert call_result.is_error is False missing = _events(client, "$mcp_missing_capability") assert missing and missing[0]["properties"]["$mcp_intent"] == "need an email tool" + + +async def test_callbacks_can_read_headers_through_the_helper(): + """The same `identify` body must work on both SDK majors: `extra["ctx"]` is + the SDK's own context and `get_request_headers` normalises the read.""" + from posthog.mcp import get_request_headers + + server = make_server() + client = FakeClient() + seen = {} + + def identify(request, extra): + seen["headers"] = get_request_headers(extra) + return None + + instrument(server, client, MCPAnalyticsOptions(identify=identify)) + + ctx = fake_ctx(headers={"Authorization": "Bearer t0ken", "User-Agent": "probe/1"}) + await _call_tool(server, "add", {"a": 1, "b": 1, "context": "header read"}, ctx=ctx) + await _flush() + + assert seen["headers"] == {"authorization": "Bearer t0ken", "user-agent": "probe/1"} diff --git a/references/public_api_snapshot.txt b/references/public_api_snapshot.txt index 21b8b1fc2..10ad20201 100644 --- a/references/public_api_snapshot.txt +++ b/references/public_api_snapshot.txt @@ -355,6 +355,7 @@ alias posthog.mcp.derive_session_id_from_mcp_session -> posthog.mcp.session.deri alias posthog.mcp.encode_session_id -> posthog.mcp.session_token.encode_session_id alias posthog.mcp.get_mcp_session -> posthog.mcp.asgi.get_mcp_session alias posthog.mcp.get_more_tools_result -> posthog.mcp.tools.get_more_tools_result +alias posthog.mcp.get_request_headers -> posthog.mcp.request_headers.get_request_headers alias posthog.mcp.set_logger -> posthog.mcp.logger.set_logger alias posthog.metrics_capture.VERSION -> posthog.version.VERSION alias posthog.metrics_capture.remove_trailing_slash -> posthog.utils.remove_trailing_slash @@ -1129,6 +1130,7 @@ function posthog.mcp.asgi.autowire_stateless_mint(server: Any) -> None function posthog.mcp.asgi.get_mcp_session(request_or_scope: Any) -> Optional[SessionTokenPayload] function posthog.mcp.instrument(server: Any, posthog_client: Optional[Client] = None, options: Optional[MCPAnalyticsOptions] = None) -> McpAnalytics function posthog.mcp.logger.set_logger(logger: Optional[LoggerFn]) -> None +function posthog.mcp.request_headers.get_request_headers(extra: Any) -> Optional[RequestHeaderBag] function posthog.mcp.session.derive_session_id_from_conversation(conversation_id: str) -> str function posthog.mcp.session.derive_session_id_from_mcp_session(mcp_session_id: str) -> str function posthog.mcp.session_token.decode_session_id(value: Any) -> Optional[SessionTokenPayload] @@ -1423,6 +1425,7 @@ module posthog.mcp.asgi module posthog.mcp.constants module posthog.mcp.logger module posthog.mcp.posthog_mcp +module posthog.mcp.request_headers module posthog.mcp.session module posthog.mcp.session_token module posthog.mcp.tools From 5ed208661b187b88bb0c813b8c5c5a482383849f Mon Sep 17 00:00:00 2001 From: Georgis Andonis Date: Thu, 20 Aug 2026 18:49:38 +0300 Subject: [PATCH 08/16] =?UTF-8?q?fix(mcp):=20review=20findings=20=E2=80=94?= =?UTF-8?q?=20never=20mutate=20the=20caller's=20result,=20guard=20ctx=20re?= =?UTF-8?q?ad?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Four defects found by the qa-swarm delegation pass (fable technical lens + opus security lens), all reproduced before fixing: HIGH — strict output schemas broke after a tool-cache rebuild. mcp 1.x validates a call's structuredContent against its *cached* tool definition, and rebuilds that cache from the internal `req=None` listing pass whenever an unlisted tool name is called (our own advertised `get_more_tools` does it). We skipped injecting on that pass, so the cache lost the `_mcp_instructions` declaration while the mirror kept writing the key: every later call to a tool with `additionalProperties: false` came back `isError: "Additional properties are not allowed"`. The internal pass now gets the same injections; only capture is skipped. HIGH — cross-caller handle leak. The mirror and the prompt-back mutated the result object in place, so a tool returning a shared/cached CallToolResult pinned one conversation's handle onto it and served it to every later caller — and since the handle outranks the transport session, that collapses unrelated clients into one session, where identity merging can write one user's person properties onto another's profile. Both channels now copy (model_copy) instead of mutating; the low-level path rewraps and returns the copy. MEDIUM — a re-listing disabled the mirror. add_instructions_to_output_schema could not tell its own prior declaration from a customer's, so servers that return persistent Tool objects flipped ownership to False on the second tools/list, silently switching off the feature for the schema-reading clients it exists for (and blaming the customer in the log). It now recognises our declaration by its description sentinel. MEDIUM — `getattr(context, "request_context", None)` was unguarded: FastMCP's property *raises* outside a request, and getattr only swallows AttributeError, so the public `FastMCP.call_tool()` entry point started raising ValueError from analytics. Read it through a guarded helper, like every other context reader here. Also de-footguns the get_request_headers docstring: its example resolved an Authorization header to a user id rather than handing back the raw token, with an explicit warning that person properties and custom event properties are not redacted. Regression tests for all four, each verified to fail without its fix. Generated-By: PostHog Code Task-Id: ebafcb71-b03b-443d-b40c-d527ed4a04f4 --- posthog/mcp/_instrument_fastmcp.py | 83 ++++++++++++++------ posthog/mcp/_instrument_lowlevel.py | 23 ++++-- posthog/mcp/_instrument_v2.py | 38 ++++++--- posthog/mcp/_output_instructions.py | 38 ++++++++- posthog/mcp/request_headers.py | 10 ++- posthog/test/mcp/test_fastmcp.py | 14 ++++ posthog/test/mcp/test_output_instructions.py | 53 +++++++++++++ posthog/test/mcp/test_review_fixes.py | 68 ++++++++++++++++ 8 files changed, 280 insertions(+), 47 deletions(-) diff --git a/posthog/mcp/_instrument_fastmcp.py b/posthog/mcp/_instrument_fastmcp.py index 2ac25c861..338a6d3dc 100644 --- a/posthog/mcp/_instrument_fastmcp.py +++ b/posthog/mcp/_instrument_fastmcp.py @@ -109,7 +109,7 @@ async def wrapped( # pipeline keeps only a scalar projection of `extra`. extra: Dict[str, Any] = { "session_id": mcp_session_id, - "ctx": getattr(context, "request_context", None), + "ctx": _tool_call_request_context(context), } # Resolve the conversation handle before the session: when the agent @@ -232,6 +232,38 @@ async def wrapped( # --- tools/list seam --------------------------------------------------------- +def _inject_tool_schemas(server: Any, data: MCPAnalyticsData, tools: list) -> None: + """Advertise the analytics parameters on a listing's tools, in place. + + Runs on both the client-facing listing and the SDK's internal cache- + population pass, so the schema the SDK validates against always matches the + one we advertised — see the note in ``list_handler``. + """ + context_enabled = is_context_enabled(data.options.context) + description = get_context_description(data.options.context) + for tool in tools: + if tool.name == _GET_MORE_TOOLS_NAME: + continue + owns_context = _tool_owns_context(server, tool.name) + schema = getattr(tool, "inputSchema", None) + if context_enabled and not owns_context: + schema = add_context_parameter_to_schema(schema, tool.name, description) + if data.options.enable_conversation_id: + schema = add_conversation_id_to_schema(schema, tool.name) + if schema is not getattr(tool, "inputSchema", None): + try: + tool.inputSchema = schema + except Exception: # noqa: BLE001 - some schema attrs may be read-only + log(f"WARN: could not set inputSchema on tool {tool.name}") + # Declare the structuredContent channel and remember the answer: + # clients that read structuredContent never see the content text + # block, and only a declared key may be written back on a call. + if data.options.enable_conversation_id: + data.tool_output_instructions[tool.name] = ( + add_instructions_to_output_schema(tool) + ) + + def _wrap_list_tools_handler(server: Any, data: MCPAnalyticsData) -> None: low_level = getattr(server, "_mcp_server", None) if low_level is None: @@ -243,9 +275,16 @@ def _wrap_list_tools_handler(server: Any, data: MCPAnalyticsData) -> None: async def list_handler(req: Any) -> Any: # The low-level server calls the handler with None to populate its tool - # cache; don't capture or inject on that internal pass. + # cache. Skip analytics on that internal pass — but still inject the + # schemas. That cache is the surface the SDK validates calls against + # (`jsonschema.validate(arguments, inputSchema)` and + # `(structuredContent, outputSchema)`), and it is rebuilt from scratch + # whenever an unlisted tool name is called. If it lacks the keys we + # advertise and write, the SDK rejects the customer's own tool result. if req is None: - return await original(req) + result = await original(req) + _inject_tool_schemas(server, data, extract_tools(result)) + return result client_name, client_version = _low_level_client_info(server) protocol_version = _low_level_protocol_version(server) @@ -306,29 +345,7 @@ async def list_handler(req: Any) -> Any: if category: data.tool_categories[tool.name] = category - context_enabled = is_context_enabled(data.options.context) - description = get_context_description(data.options.context) - for tool in tools: - if tool.name == _GET_MORE_TOOLS_NAME: - continue - owns_context = _tool_owns_context(server, tool.name) - schema = getattr(tool, "inputSchema", None) - if context_enabled and not owns_context: - schema = add_context_parameter_to_schema(schema, tool.name, description) - if data.options.enable_conversation_id: - schema = add_conversation_id_to_schema(schema, tool.name) - if schema is not getattr(tool, "inputSchema", None): - try: - tool.inputSchema = schema - except Exception: # noqa: BLE001 - some schema attrs may be read-only - log(f"WARN: could not set inputSchema on tool {tool.name}") - # Declare the structuredContent channel and remember the answer: - # clients that read structuredContent never see the content text - # block, and only a declared key may be written back on a call. - if data.options.enable_conversation_id: - data.tool_output_instructions[tool.name] = ( - add_instructions_to_output_schema(tool) - ) + _inject_tool_schemas(server, data, tools) if data.options.report_missing: missing_name = resolve_missing_capability_tool_name(data.options) @@ -400,6 +417,20 @@ def _tool_owns_context(server: Any, name: str) -> bool: return _tool_owns_param(server, name, "context") +def _tool_call_request_context(context: Any) -> Any: + """The request context behind a FastMCP ``Context``, or ``None``. + + ``Context.request_context`` is a property that *raises* ``ValueError`` when + the call happens outside a request — which the public ``FastMCP.call_tool()`` + entry point does. A bare ``getattr(..., None)`` only swallows + ``AttributeError``, so it would let that escape into the customer's tool call. + """ + try: + return context.request_context + except (LookupError, ValueError, AttributeError): + return None + + def _low_level_request_context(server: Any) -> Any: """The underlying low-level server's request_context, set during a request. The tools/list handler runs on ``server._mcp_server``, so client info / session id diff --git a/posthog/mcp/_instrument_lowlevel.py b/posthog/mcp/_instrument_lowlevel.py index c9d44e0b4..404cdc8c3 100644 --- a/posthog/mcp/_instrument_lowlevel.py +++ b/posthog/mcp/_instrument_lowlevel.py @@ -208,21 +208,34 @@ async def handler(req: Any) -> Any: if conversation_id: delivered = False if data.tool_output_instructions.get(name): - _, delivered = mirror_instructions_into_structured_content( + call_result, delivered = mirror_instructions_into_structured_content( call_result, conversation_id ) if minted: content = getattr(call_result, "content", None) if isinstance(content, list): - content.append( - mcp_types.TextContent( - type="text", text=build_prompt_back(conversation_id)["text"] - ) + block = mcp_types.TextContent( + type="text", text=build_prompt_back(conversation_id)["text"] ) + # Copy rather than append in place — a shared or cached result + # object would accumulate a block per conversation and leak + # earlier callers' handles to later ones. + copy_model = getattr(call_result, "model_copy", None) + if callable(copy_model): + call_result = copy_model(update={"content": [*content, block]}) + else: + content.append(block) delivered = True # Only a minted handle can be lost — one the agent supplied, it has. if not delivered: delivered_conversation_id = None + # Hand back whatever copy we made, rewrapped as the SDK expects. + if call_result is not getattr(result, "root", result): + result = ( + mcp_types.ServerResult(call_result) + if hasattr(result, "root") + else call_result + ) await record_tool_call( data, diff --git a/posthog/mcp/_instrument_v2.py b/posthog/mcp/_instrument_v2.py index 44e1d6050..dd56b82d6 100644 --- a/posthog/mcp/_instrument_v2.py +++ b/posthog/mcp/_instrument_v2.py @@ -354,23 +354,34 @@ async def wrapped( tool_manager.call_tool = wrapped -def _append_prompt_back(result: Any, conversation_id: str) -> bool: +def _append_prompt_back(result: Any, conversation_id: str) -> Tuple[Any, bool]: """Append the conversation prompt-back to a result's ``content`` list (model or dict shape). Errored results included on purpose — a first-call failure is exactly when the agent needs the handle. Returns False for shapes with no - content list to ride (e.g. MRTR ``input_required`` results).""" + content list to ride (e.g. MRTR ``input_required`` results). Returns + ``(result, delivered)`` — the result may be a copy, never a mutation.""" block = mcp_types.TextContent( type="text", text=build_prompt_back(conversation_id)["text"] ) - content = ( - result.get("content") - if isinstance(result, dict) - else getattr(result, "content", None) - ) - if isinstance(content, list): - content.append(block) - return True - return False + if isinstance(result, dict): + content = result.get("content") + if isinstance(content, list): + return {**result, "content": [*content, block]}, True + return result, False + content = getattr(result, "content", None) + if not isinstance(content, list): + return result, False + # Copy rather than append in place: a shared or cached result object would + # otherwise accumulate a block per conversation and hand each caller the + # previous callers' handles. + copy_model = getattr(result, "model_copy", None) + if callable(copy_model): + try: + return copy_model(update={"content": [*content, block]}), True + except Exception: # noqa: BLE001 - never let delivery break the tool path + return result, False + content.append(block) + return result, True def _deliver_conversation_id( @@ -386,8 +397,9 @@ def _deliver_conversation_id( result, delivered = mirror_instructions_into_structured_content( result, conversation_id ) - if minted and _append_prompt_back(result, conversation_id): - delivered = True + if minted: + result, appended = _append_prompt_back(result, conversation_id) + delivered = delivered or appended return result, delivered diff --git a/posthog/mcp/_output_instructions.py b/posthog/mcp/_output_instructions.py index d33d3d581..e1fb084ae 100644 --- a/posthog/mcp/_output_instructions.py +++ b/posthog/mcp/_output_instructions.py @@ -49,6 +49,15 @@ def _read_attr(obj: Any, names: Tuple[str, ...]) -> Tuple[Optional[str], Any]: return None, None +def _is_our_declaration(declaration: Any) -> bool: + """Whether an existing :data:`MCP_INSTRUCTIONS_KEY` property is one we wrote + on a previous listing, told apart from a customer's by our description.""" + return ( + isinstance(declaration, dict) + and declaration.get("description") == _INSTRUCTIONS_FIELD_DESCRIPTION + ) + + def can_declare_output_instructions(output_schema: Any) -> bool: """True when :data:`MCP_INSTRUCTIONS_KEY` can safely be declared on this tool's advertised output schema. @@ -95,6 +104,13 @@ def add_instructions_to_output_schema(tool: Any) -> bool: if not can_declare_output_instructions(original): properties = original.get("properties") if isinstance(original, dict) else None if isinstance(properties, dict) and MCP_INSTRUCTIONS_KEY in properties: + # Our own declaration from an earlier listing: servers that hand back + # persistent Tool objects (a module-level list, or two servers sharing + # tools) hit this on every re-list. Report it as declared — reading it + # as customer-owned would silently switch the mirror off for exactly + # the clients it exists for, and blame the customer in the log. + if _is_our_declaration(properties[MCP_INSTRUCTIONS_KEY]): + return True log( f"WARN: Tool \"{name}\" already declares '{MCP_INSTRUCTIONS_KEY}' in its " "output schema. Leaving it alone." @@ -178,8 +194,28 @@ def mirror_instructions_into_structured_content( or MCP_INSTRUCTIONS_KEY in structured ): return result, False + updated = {**structured, MCP_INSTRUCTIONS_KEY: payload} + # Copy rather than mutate: a tool is free to return a shared or cached + # result object, and pinning one conversation's handle onto it would serve + # that handle to every later caller (and, through the handle, collapse + # unrelated clients into one session). + copy_model = getattr(target, "model_copy", None) + if callable(copy_model): + try: + new_target = copy_model(update={attr: updated}) + except Exception: # noqa: BLE001 - never let delivery break the tool path + return result, False + if target is result: + return new_target, True + rewrap = getattr(result, "model_copy", None) + if callable(rewrap): + try: + return rewrap(update={"root": new_target}), True + except Exception: # noqa: BLE001 + return result, False + return result, False try: - setattr(target, attr, {**structured, MCP_INSTRUCTIONS_KEY: payload}) + setattr(target, attr, updated) except Exception: # noqa: BLE001 - never let delivery break the tool path return result, False return result, True diff --git a/posthog/mcp/request_headers.py b/posthog/mcp/request_headers.py index 2b7edcac4..800efa1bd 100644 --- a/posthog/mcp/request_headers.py +++ b/posthog/mcp/request_headers.py @@ -15,8 +15,14 @@ def identify(request, extra): headers = get_request_headers(extra) or {} - token = headers.get("authorization") - ... + return UserIdentity(distinct_id=user_id_for(headers.get("authorization"))) + +**Never return a raw header value** (an ``Authorization`` bearer token, a +cookie, an API key) as a ``UserIdentity`` property or from ``event_properties``. +Resolve it to an id or a role first. Person properties and custom event +properties are the one part of the payload the SDK does not redact, and they +persist on the person profile — a token written there outlives event retention +and is visible in the PostHog UI. Returns ``None`` when the request did not arrive over HTTP — stdio and in-memory transports carry no headers at all. diff --git a/posthog/test/mcp/test_fastmcp.py b/posthog/test/mcp/test_fastmcp.py index 087863f0c..b5e902c18 100644 --- a/posthog/test/mcp/test_fastmcp.py +++ b/posthog/test/mcp/test_fastmcp.py @@ -237,3 +237,17 @@ def identify(request, extra): await _flush() assert seen["headers"] == {"authorization": "Bearer t0ken"} + + +async def test_public_call_tool_entrypoint_still_works_outside_a_request(): + """`FastMCP.call_tool()` is public API for in-process invocation, where there + is no request context. `Context.request_context` *raises* there, so reading + it unguarded would push an analytics error into the customer's tool path.""" + server = make_server() + instrument(server, FakeClient()) + + result = await server.call_tool("add", {"a": 2, "b": 3, "context": "no request"}) + await _flush() + + text_blocks = [c.text for c in result[0] if getattr(c, "type", None) == "text"] + assert "5" in text_blocks diff --git a/posthog/test/mcp/test_output_instructions.py b/posthog/test/mcp/test_output_instructions.py index 312ffe996..0fd569606 100644 --- a/posthog/test/mcp/test_output_instructions.py +++ b/posthog/test/mcp/test_output_instructions.py @@ -169,3 +169,56 @@ def test_customer_key_wins_over_the_mirror(): assert delivered is False assert result.structuredContent[MCP_INSTRUCTIONS_KEY] == "mine" + + +def test_our_own_declaration_is_recognised_on_a_relisting(): + """Servers that hand back persistent Tool objects re-list the *same* object. + Reading our own prior declaration as customer-owned would flip ownership to + False and silently switch the mirror off for schema-reading clients.""" + tool = _tool({"type": "object", "properties": {"total": {"type": "integer"}}}) + + assert add_instructions_to_output_schema(tool) is True + # second listing of the very same object — still ours, still declared + assert add_instructions_to_output_schema(tool) is True + assert MCP_INSTRUCTIONS_KEY in tool.outputSchema["properties"] + + +def test_a_customer_key_is_still_not_claimed_as_ours(): + tool = _tool( + { + "type": "object", + "properties": {MCP_INSTRUCTIONS_KEY: {"type": "string"}}, + } + ) + + assert add_instructions_to_output_schema(tool) is False + + +def test_mirror_does_not_mutate_a_shared_model_result(): + """A tool may return a cached/shared result object. Pinning one + conversation's handle onto it would serve that handle to every later caller + — and, through the handle, collapse unrelated clients into one session.""" + import mcp.types as mcp_types + + # `structuredContent` on SDK 1.x, `structured_content` on 2.x (same wire field). + def structured(result): + for attr in ("structuredContent", "structured_content"): + if hasattr(result, attr): + return getattr(result, attr) + raise AssertionError("no structured content attribute") + + shared = mcp_types.CallToolResult( + content=[mcp_types.TextContent(type="text", text="{}")], + structuredContent={"total": 7}, + ) + + first, delivered = mirror_instructions_into_structured_content(shared, "conv-A") + + assert delivered is True + assert first is not shared # a copy, not the caller's object + assert structured(first)[MCP_INSTRUCTIONS_KEY] == {"conversation_id": "conv-A"} + # the shared object the customer owns is untouched, so the next caller is clean + assert MCP_INSTRUCTIONS_KEY not in structured(shared) + + second, _ = mirror_instructions_into_structured_content(shared, "conv-B") + assert structured(second)[MCP_INSTRUCTIONS_KEY] == {"conversation_id": "conv-B"} diff --git a/posthog/test/mcp/test_review_fixes.py b/posthog/test/mcp/test_review_fixes.py index 0b329a050..468c3af79 100644 --- a/posthog/test/mcp/test_review_fixes.py +++ b/posthog/test/mcp/test_review_fixes.py @@ -338,3 +338,71 @@ async def _lt(): ) assert any("tested against mcp>=1.26" in m for m in logs) + + +# --- F: the SDK's tool cache must carry the keys we advertise ----------------- + + +async def test_strict_output_schema_survives_a_tool_cache_rebuild(): + """mcp 1.x validates a call's structuredContent against its *cached* tool + definition. That cache is rebuilt from the internal `req=None` listing pass + whenever an unlisted tool name is called — including our own advertised + `get_more_tools`. If the rebuild dropped our `_mcp_instructions` + declaration while the mirror kept writing the key, the SDK would reject the + customer's own successful result under `additionalProperties: false`.""" + from typing import Any + + from pydantic import BaseModel, ConfigDict + + from posthog.mcp._output_instructions import MCP_INSTRUCTIONS_KEY + + class Totals(BaseModel): + model_config = ConfigDict(extra="forbid") # -> additionalProperties: false + total: int + + server = FastMCP("strict-output") + + @server.tool() + def totals(event: str) -> Totals: + return Totals(total=7) + + client = FakeClient() + instrument(server, client, MCPAnalyticsOptions(enable_conversation_id=True)) + + # FastMCP's lowlevel call path reads the request contextvar; set one so the + # handler can run outside a live session. + from types import SimpleNamespace + + from mcp.server.lowlevel.server import request_ctx + from mcp.shared.context import RequestContext + + request_ctx.set( + RequestContext( + request_id=1, + meta=None, + session=SimpleNamespace(client_params=None), + lifespan_context=None, + request=None, + ) + ) + + call = server._mcp_server.request_handlers[mcp_types.CallToolRequest] + listed = await server._mcp_server.request_handlers[mcp_types.ListToolsRequest]( + mcp_types.ListToolsRequest(method="tools/list") + ) + tool = next(t for t in listed.root.tools if t.name == "totals") + assert MCP_INSTRUCTIONS_KEY in tool.outputSchema["properties"] + + first = await call(_call_request("totals", {"event": "a", "context": "first"})) + assert first.root.isError is False + + # Force a cache rebuild the way a real client does: call a name the cache + # doesn't know yet. + await call(_call_request("not-a-real-tool", {"context": "typo"})) + + # The customer's tool must still work — analytics never breaks the result. + after: Any = await call( + _call_request("totals", {"event": "b", "context": "second"}) + ) + assert after.root.isError is False, after.root.content[0].text + assert after.root.structuredContent[MCP_INSTRUCTIONS_KEY]["conversation_id"] From 07dd1c8b5bea26bcb319cb8bd4c53e4746795fc4 Mon Sep 17 00:00:00 2001 From: Georgis Andonis Date: Thu, 20 Aug 2026 19:00:12 +0300 Subject: [PATCH 09/16] refactor: apply simplify pass MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Reuse existing helpers in the v2 adapter instead of hand-rolled copies: - _ctx_mcp_session_id now goes through get_request_headers + read_mcp_session_header, which already handle case-insensitive keys, list-valued headers and whitespace — strictly better than the raw ctx.request.headers read it replaces - shared schema_has_param() in _context_parameters (was duplicated between the lowlevel and v2 adapters) - shared params_to_request_dict() in _instrumentation (request_to_dict now delegates to it) - one tool-manager lookup per call instead of two when checking ownership of both injected parameters The three adapters stay separate on purpose — they hook different SDK seams with different semantics; only genuinely identical logic is shared. Generated-By: PostHog Code Task-Id: ebafcb71-b03b-443d-b40c-d527ed4a04f4 --- posthog/mcp/_context_parameters.py | 11 +++++ posthog/mcp/_instrument_v2.py | 66 +++++++++++++----------------- posthog/mcp/_instrumentation.py | 13 +++++- 3 files changed, 52 insertions(+), 38 deletions(-) diff --git a/posthog/mcp/_context_parameters.py b/posthog/mcp/_context_parameters.py index 6aed132fd..326e8eb25 100644 --- a/posthog/mcp/_context_parameters.py +++ b/posthog/mcp/_context_parameters.py @@ -20,6 +20,17 @@ def is_context_enabled(context: Union[bool, MCPAnalyticsContextOptions, None]) - return context is not False +def schema_has_param(schema: Any, name: str) -> bool: + """Whether a (already-serialized) JSON Schema dict declares a top-level + property named ``name``. Shared by the lowlevel and v2 adapters, which both + need to tell an injected parameter apart from one the tool already owns.""" + return ( + isinstance(schema, dict) + and isinstance(schema.get("properties"), dict) + and name in schema["properties"] + ) + + def get_context_description( context: Union[bool, MCPAnalyticsContextOptions, None], ) -> Optional[str]: diff --git a/posthog/mcp/_instrument_v2.py b/posthog/mcp/_instrument_v2.py index dd56b82d6..f701045f6 100644 --- a/posthog/mcp/_instrument_v2.py +++ b/posthog/mcp/_instrument_v2.py @@ -39,6 +39,7 @@ add_context_parameter_to_schema, get_context_description, is_context_enabled, + schema_has_param, ) from ._conversation_id import ( add_conversation_id_to_schema, @@ -48,6 +49,7 @@ from ._instrumentation import ( _to_jsonable, build_tool_call_request, + params_to_request_dict, prepare_request, read_tool_category, record_missing_capability, @@ -61,6 +63,8 @@ mirror_instructions_into_structured_content, ) from .logger import log +from .request_headers import get_request_headers +from .session_token import read_mcp_session_header from .tools import ( GET_MORE_TOOLS_NAME as _GET_MORE_TOOLS_NAME, build_report_missing_descriptor, @@ -170,14 +174,12 @@ def _ctx_protocol_version(ctx: Any) -> Optional[str]: def _ctx_mcp_session_id(ctx: Any) -> Optional[str]: """Best-effort transport session id (the ``Mcp-Session-Id`` header on the legacy era — 2026-07-28 removed it). ``ctx.request`` carries the transport's - HTTP request when there is one; ``None`` on stdio.""" - try: - headers = getattr(getattr(ctx, "request", None), "headers", None) - if headers is not None: - return headers.get("mcp-session-id") - except Exception: # noqa: BLE001 - pass - return None + HTTP request when there is one; ``None`` on stdio. + + Reuses the same header-bag normalisation and case-insensitive lookup the + public ``get_request_headers``/``read_mcp_session_header`` helpers already + provide, instead of a second hand-rolled ``ctx.request.headers`` read.""" + return read_mcp_session_header(get_request_headers(ctx)) def _resolve_ctx( @@ -194,17 +196,18 @@ def _resolve_ctx( return token, client_name, client_version, protocol_version, mcp_session_id -def _params_to_request(method: str, params: Any) -> Dict[str, Any]: - params_dict: Any = {} - if params is not None and hasattr(params, "model_dump"): - try: - params_dict = params.model_dump(mode="json", by_alias=True) - except Exception: # noqa: BLE001 - params_dict = {} - return {"method": method, "params": params_dict} +# --- tool ownership -------------------------------------------------------------- -# --- tool ownership -------------------------------------------------------------- +def _tool_own_properties_v2(high_level: Any, name: str) -> Dict[str, Any]: + """The tool's own declared JSON Schema ``properties``, read once per call + site so checking ownership of both ``context`` and ``conversation_id`` + doesn't look the tool up from the manager twice.""" + try: + tool = high_level._tool_manager.get_tool(name) + return (getattr(tool, "parameters", None) or {}).get("properties", {}) + except Exception: # noqa: BLE001 + return {} def _tool_owns_param_v2(high_level: Any, name: str, param: str) -> bool: @@ -213,12 +216,7 @@ def _tool_owns_param_v2(high_level: Any, name: str, param: str) -> bool: declared parameters rather than the function signature so a tool taking the SDK's ``Context`` object under a ``context`` name isn't mistaken for owning our string parameter.""" - try: - tool = high_level._tool_manager.get_tool(name) - properties = (getattr(tool, "parameters", None) or {}).get("properties", {}) - return param in properties - except Exception: # noqa: BLE001 - return False + return param in _tool_own_properties_v2(high_level, name) # --- high-level: ToolManager.call_tool seam -------------------------------------- @@ -289,11 +287,13 @@ async def wrapped( # one the tool's own schema declares (that's a real argument). call_arguments = arguments if isinstance(arguments, dict): + own_properties = _tool_own_properties_v2(server, name) strip_keys = set() - if not _tool_owns_param_v2(server, name, "context"): + if "context" not in own_properties: strip_keys.add("context") - if data.options.enable_conversation_id and not _tool_owns_param_v2( - server, name, "conversation_id" + if ( + data.options.enable_conversation_id + and "conversation_id" not in own_properties ): strip_keys.add("conversation_id") if strip_keys: @@ -528,7 +528,7 @@ async def handler(ctx: Any, params: Any) -> Any: token, client_name, client_version, protocol_version, mcp_session_id = ( _resolve_ctx(ctx) ) - request = _params_to_request(_LIST_METHOD, params) + request = params_to_request_dict(_LIST_METHOD, params, by_alias=True) extra: Dict[str, Any] = {"session_id": mcp_session_id, "ctx": ctx} # Resolve session, emit $mcp_initialize (once per session) and identify # here too — a client may list tools without ever calling one. @@ -586,7 +586,7 @@ async def handler(ctx: Any, params: Any) -> Any: owns_context = ( _tool_owns_param_v2(high_level, tool.name, "context") if high_level is not None - else _schema_has_param(schema, "context") + else schema_has_param(schema, "context") ) # required follows the entry point: the raw low-level path validates # the call against this same schema (optional); the high-level path @@ -595,7 +595,7 @@ async def handler(ctx: Any, params: Any) -> Any: schema = add_context_parameter_to_schema( schema, tool.name, description, required=context_required ) - if data.options.enable_conversation_id and not _schema_has_param( + if data.options.enable_conversation_id and not schema_has_param( schema, "conversation_id" ): schema = add_conversation_id_to_schema(schema, tool.name) @@ -650,11 +650,3 @@ def _append_get_more_tools_v2(result: Any, name: str) -> None: tools_list = getattr(result, "tools", None) if isinstance(tools_list, list): tools_list.append(tool) - - -def _schema_has_param(schema: Any, name: str) -> bool: - return ( - isinstance(schema, dict) - and isinstance(schema.get("properties"), dict) - and name in schema["properties"] - ) diff --git a/posthog/mcp/_instrumentation.py b/posthog/mcp/_instrumentation.py index 9ab365ca1..4ce2d4a59 100644 --- a/posthog/mcp/_instrumentation.py +++ b/posthog/mcp/_instrumentation.py @@ -382,10 +382,21 @@ def request_to_dict(req: Any) -> Dict[str, Any]: """Shape a request object into the JSON-RPC-ish dict the sanitizer expects.""" method = getattr(req, "method", None) or "tools/list" params = getattr(req, "params", None) + return params_to_request_dict(method, params) + + +def params_to_request_dict( + method: str, params: Any, *, by_alias: bool = False +) -> Dict[str, Any]: + """Shape a bare ``(method, params)`` pair into the same JSON-RPC-ish dict + ``request_to_dict`` builds from a request object. v2's request handlers + receive ``params`` directly rather than a ``req`` wrapper, so there's no + object to hand ``request_to_dict``; ``by_alias`` lets v2 keep the wire's + camelCase aliases (its models expose snake_case attributes).""" params_dict: Any = {} if params is not None and hasattr(params, "model_dump"): try: - params_dict = params.model_dump(mode="json") + params_dict = params.model_dump(mode="json", by_alias=by_alias) except Exception: # noqa: BLE001 params_dict = {} return {"method": method, "params": params_dict} From e335876d9826a30c85b441519db4dbb02015ac22 Mon Sep 17 00:00:00 2001 From: Georgis Andonis Date: Thu, 20 Aug 2026 19:14:06 +0300 Subject: [PATCH 10/16] fix(mcp): finish the tool-cache fix on the low-level adapter MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Round 2 of the review loop caught that the previous commit's cache-rebuild fix stopped at the FastMCP adapter. The raw low-level v1 adapter has the identical exposure on the *input* side, and it is worse there: this adapter advertises `context`/`conversation_id` without stripping them, relying on the advertised schema doubling as the call's validation schema. So a cache rebuilt from the un-injected internal listing pass rejects exactly the arguments the SDK told the agent to send — "Input validation error: Additional properties are not allowed ('context' was unexpected)" — for any tool with `additionalProperties: false`. One unknown tool name (or a fresh stateless instance) triggers it, and the conversation prompt-back tells the agent to keep sending the argument, so it persists until the next client-facing tools/list. The internal pass now gets the same injections here too, via a shared `_inject_tool_schemas` helper mirroring the FastMCP one. Regression test verified to fail without the fix. Also from round 2: - `_tool_own_properties_v2` fails closed on a malformed schema again — the simplify pass had narrowed the guard so a non-dict `properties` would raise (or answer by substring) in the tool-call hot path - finish the `schema_has_param` dedup the simplify pass started: the low-level adapter now uses the shared helper instead of its own copy Generated-By: PostHog Code Task-Id: ebafcb71-b03b-443d-b40c-d527ed4a04f4 --- posthog/mcp/_instrument_lowlevel.py | 92 +++++++++++++++------------ posthog/mcp/_instrument_v2.py | 6 +- posthog/test/mcp/test_review_fixes.py | 43 +++++++++++++ 3 files changed, 101 insertions(+), 40 deletions(-) diff --git a/posthog/mcp/_instrument_lowlevel.py b/posthog/mcp/_instrument_lowlevel.py index 404cdc8c3..6b9baad52 100644 --- a/posthog/mcp/_instrument_lowlevel.py +++ b/posthog/mcp/_instrument_lowlevel.py @@ -24,6 +24,7 @@ add_context_parameter_to_schema, get_context_description, is_context_enabled, + schema_has_param, ) from ._conversation_id import ( add_conversation_id_to_schema, @@ -256,6 +257,45 @@ async def handler(req: Any) -> Any: handlers[mcp_types.CallToolRequest] = handler +def _inject_tool_schemas( + data: MCPAnalyticsData, tools: list, *, context_required: bool +) -> None: + """Advertise the analytics parameters on a listing's tools, in place. + + Runs on both the client-facing listing and the SDK's internal cache- + population pass, so the schema the SDK validates against always matches the + one we advertised — see the note in ``handler``. + """ + context_enabled = is_context_enabled(data.options.context) + description = get_context_description(data.options.context) + for tool in tools: + if tool.name == _GET_MORE_TOOLS_NAME: + continue + schema = getattr(tool, "inputSchema", None) + # required follows the path: raw low-level validates the call against + # this same schema (optional), FastMCP 2.0 strips it first (required-advisory). + if context_enabled and not schema_has_param(schema, "context"): + schema = add_context_parameter_to_schema( + schema, tool.name, description, required=context_required + ) + if data.options.enable_conversation_id and not schema_has_param( + schema, "conversation_id" + ): + schema = add_conversation_id_to_schema(schema, tool.name) + if schema is not getattr(tool, "inputSchema", None): + try: + tool.inputSchema = schema + except Exception: # noqa: BLE001 + log(f"WARN: could not set inputSchema on tool {tool.name}") + # Declare the structuredContent channel and remember the answer: + # clients that read structuredContent never see the content text + # block, and only a declared key may be written back on a call. + if data.options.enable_conversation_id: + data.tool_output_instructions[tool.name] = ( + add_instructions_to_output_schema(tool) + ) + + def _wrap_list_tools( server: Any, data: MCPAnalyticsData, *, context_required: bool ) -> None: @@ -265,10 +305,19 @@ def _wrap_list_tools( return async def handler(req: Any) -> Any: - # The server calls the handler with None to populate its tool cache; - # don't capture or inject on that internal pass. + # The server calls the handler with None to populate its tool cache. + # Skip analytics there — but still inject, because that cache is the + # schema the SDK validates calls against. This adapter advertises + # `context`/`conversation_id` without stripping them, so a cache built + # from un-injected schemas rejects the very arguments we told the agent + # to send ("Additional properties are not allowed") on any tool with + # `additionalProperties: false`. if req is None: - return await original(req) + result = await original(req) + _inject_tool_schemas( + data, extract_tools(result), context_required=context_required + ) + return result client_name, client_version = _client_info(server) protocol_version = _protocol_version(server) @@ -331,34 +380,7 @@ async def handler(req: Any) -> Any: # TS SDK) — captured before we append our own get_more_tools virtual tool. empty = len(tools) == 0 - context_enabled = is_context_enabled(data.options.context) - description = get_context_description(data.options.context) - for tool in tools: - if tool.name == _GET_MORE_TOOLS_NAME: - continue - schema = getattr(tool, "inputSchema", None) - # required follows the path: raw low-level validates the call against - # this same schema (optional), FastMCP 2.0 strips it first (required-advisory). - if context_enabled and not _schema_has_param(schema, "context"): - schema = add_context_parameter_to_schema( - schema, tool.name, description, required=context_required - ) - if data.options.enable_conversation_id and not _schema_has_param( - schema, "conversation_id" - ): - schema = add_conversation_id_to_schema(schema, tool.name) - if schema is not getattr(tool, "inputSchema", None): - try: - tool.inputSchema = schema - except Exception: # noqa: BLE001 - log(f"WARN: could not set inputSchema on tool {tool.name}") - # Declare the structuredContent channel and remember the answer: - # clients that read structuredContent never see the content text - # block, and only a declared key may be written back on a call. - if data.options.enable_conversation_id: - data.tool_output_instructions[tool.name] = ( - add_instructions_to_output_schema(tool) - ) + _inject_tool_schemas(data, tools, context_required=context_required) if data.options.report_missing: missing_name = resolve_missing_capability_tool_name(data.options) @@ -387,14 +409,6 @@ async def handler(req: Any) -> Any: handlers[mcp_types.ListToolsRequest] = handler -def _schema_has_param(schema: Any, name: str) -> bool: - return ( - isinstance(schema, dict) - and isinstance(schema.get("properties"), dict) - and name in schema["properties"] - ) - - async def _tool_owned_injected_keys(high_level: Any, name: str) -> set: """Which of (``context``, ``conversation_id``) the jlowin FastMCP tool declares itself, read from its function signature. These are real tool arguments we must diff --git a/posthog/mcp/_instrument_v2.py b/posthog/mcp/_instrument_v2.py index f701045f6..10ca6b043 100644 --- a/posthog/mcp/_instrument_v2.py +++ b/posthog/mcp/_instrument_v2.py @@ -205,9 +205,13 @@ def _tool_own_properties_v2(high_level: Any, name: str) -> Dict[str, Any]: doesn't look the tool up from the manager twice.""" try: tool = high_level._tool_manager.get_tool(name) - return (getattr(tool, "parameters", None) or {}).get("properties", {}) + properties = (getattr(tool, "parameters", None) or {}).get("properties") except Exception: # noqa: BLE001 return {} + # Fail closed on a malformed schema: the caller does `param in ` in the + # tool-call hot path, where a None or a string would raise or answer by + # substring. + return properties if isinstance(properties, dict) else {} def _tool_owns_param_v2(high_level: Any, name: str, param: str) -> bool: diff --git a/posthog/test/mcp/test_review_fixes.py b/posthog/test/mcp/test_review_fixes.py index 468c3af79..ddcf3edd0 100644 --- a/posthog/test/mcp/test_review_fixes.py +++ b/posthog/test/mcp/test_review_fixes.py @@ -406,3 +406,46 @@ def totals(event: str) -> Totals: ) assert after.root.isError is False, after.root.content[0].text assert after.root.structuredContent[MCP_INSTRUCTIONS_KEY]["conversation_id"] + + +async def test_lowlevel_strict_input_schema_survives_a_tool_cache_rebuild(): + """Same cache-rebuild trap as the FastMCP case, on the *input* side. The raw + lowlevel adapter advertises `context`/`conversation_id` without stripping + them, relying on the advertised schema doubling as the validation schema — + so if the SDK's cache is rebuilt from an un-injected listing, it rejects the + very arguments we told the agent to send.""" + server = Server("strict-input") + + @server.list_tools() + async def _lt(): + return [ + mcp_types.Tool( + name="echo", + description="Echo", + inputSchema={ + "type": "object", + "properties": {"msg": {"type": "string"}}, + "required": ["msg"], + "additionalProperties": False, + }, + ) + ] + + @server.call_tool() + async def _ct(name, arguments): + return [mcp_types.TextContent(type="text", text=str(arguments.get("msg")))] + + client = FakeClient() + instrument(server, client, MCPAnalyticsOptions(enable_conversation_id=True)) + + call = server.request_handlers[mcp_types.CallToolRequest] + await server.request_handlers[mcp_types.ListToolsRequest](_list_request()) + + first = await call(_call_request("echo", {"msg": "a", "context": "first"})) + assert first.root.isError is False + + # Force a cache rebuild the way a real client does: an unknown tool name. + await call(_call_request("not-a-real-tool", {"context": "typo"})) + + after = await call(_call_request("echo", {"msg": "b", "context": "second"})) + assert after.root.isError is False, after.root.content[0].text From 45f56307e1f34149db65ac94d809b30585813d25 Mon Sep 17 00:00:00 2001 From: Georgis Andonis Date: Thu, 20 Aug 2026 19:25:00 +0300 Subject: [PATCH 11/16] fix(mcp): don't require a parameter we strip before the SDK validates MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Round 3 of the review loop. On the jlowin-FastMCP path the adapter strips the injected `context`/`conversation_id` before dispatch, because that SDK validates tool arguments against the function signature. But the listing wrapper marked `context` *required* in the advertised schema — which is also the schema mcp 1.x's low-level server validates the (already stripped) arguments against. Under `FastMCP(strict_input_validation=True)` every call therefore failed with "Input validation error: 'context' is a required property". Latent on main, where an un-injected cache rebuild intermittently cured it; the previous commit's internal-pass injection made it permanent. Advertise `context` as optional on this path instead — requiring a parameter the validator can never see is self-contradictory, and losing the "required" nudge only softens intent capture, where the alternative is a hard failure of the customer's tool call. Regression test added; the default `strict_input_validation=False` and both MCP SDK 2.x paths were unaffected either way. Generated-By: PostHog Code Task-Id: ebafcb71-b03b-443d-b40c-d527ed4a04f4 --- posthog/mcp/_instrument_lowlevel.py | 7 +++++- posthog/test/mcp/test_fastmcp_v2.py | 36 +++++++++++++++++++++++++++++ 2 files changed, 42 insertions(+), 1 deletion(-) diff --git a/posthog/mcp/_instrument_lowlevel.py b/posthog/mcp/_instrument_lowlevel.py index 6b9baad52..58a90bb86 100644 --- a/posthog/mcp/_instrument_lowlevel.py +++ b/posthog/mcp/_instrument_lowlevel.py @@ -84,7 +84,12 @@ def instrument_fastmcp_v2(server: Any, data: MCPAnalyticsData) -> None: low_level, "version", None ) _wrap_call_tool(low_level, data, strip_injected=True, high_level=server) - _wrap_list_tools(low_level, data, context_required=True) + # `context` is advertised but NOT marked required here. This adapter strips + # the injected parameters before the SDK's own input validation runs, so a + # schema that requires `context` contradicts the arguments the SDK actually + # sees: under `FastMCP(strict_input_validation=True)` every call fails with + # "'context' is a required property". + _wrap_list_tools(low_level, data, context_required=False) def _wrap_call_tool( diff --git a/posthog/test/mcp/test_fastmcp_v2.py b/posthog/test/mcp/test_fastmcp_v2.py index e40213d34..9ebc11459 100644 --- a/posthog/test/mcp/test_fastmcp_v2.py +++ b/posthog/test/mcp/test_fastmcp_v2.py @@ -115,3 +115,39 @@ async def test_jlowin_report_missing_advertises_get_more_tools(): await _flush() assert out.root.isError is False assert _events(client, "$mcp_missing_capability") + + +async def test_strict_input_validation_still_accepts_calls(): + """This adapter strips the injected parameters before the SDK validates, so + the advertised schema must not mark them required — otherwise every call + under `strict_input_validation=True` fails on a parameter the SDK never + sees.""" + import mcp.types as t + + server = FastMCP("strict", strict_input_validation=True) + + @server.tool() + def echo(msg: str) -> str: + return msg + + client = FakeClient() + instrument(server, client, MCPAnalyticsOptions(enable_conversation_id=True)) + + low = server._mcp_server + await low.request_handlers[t.ListToolsRequest]( + t.ListToolsRequest(method="tools/list") + ) + result = await low.request_handlers[t.CallToolRequest]( + t.CallToolRequest( + method="tools/call", + params=t.CallToolRequestParams( + name="echo", arguments={"msg": "hi", "context": "strict validation"} + ), + ) + ) + await _flush() + + assert result.root.isError is False, result.root.content[0].text + assert _events(client, "$mcp_tool_call")[0]["properties"]["$mcp_intent"] == ( + "strict validation" + ) From f8f4c1a78a946d8711a3cd231a819ff570e30f1b Mon Sep 17 00:00:00 2001 From: Georgis Andonis Date: Fri, 21 Aug 2026 11:39:14 +0300 Subject: [PATCH 12/16] docs(mcp): reframe the changeset around cross-SDK parity The changeset claimed "SDK 1.x paths are unchanged", which stopped being true as the PR grew: conversation-anchored sessions, the structuredContent delivery channel, the uuidv7 handle guard, the prompt-back on errored results and the callback context all reach v1 servers, and three separately reproduced bugs on v1 paths were fixed. Since no mainstream client speaks the 2026-07-28 era yet, that half is what current users actually get. Restructured into v2 support / cross-SDK parity / fixes affecting existing 1.x users, and calls out the two deliberate behavioural changes. Generated-By: PostHog Code Task-Id: ebafcb71-b03b-443d-b40c-d527ed4a04f4 --- .sampo/changesets/mcp-sdk-v2-support.md | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/.sampo/changesets/mcp-sdk-v2-support.md b/.sampo/changesets/mcp-sdk-v2-support.md index be30a386e..24fdc4d1b 100644 --- a/.sampo/changesets/mcp-sdk-v2-support.md +++ b/.sampo/changesets/mcp-sdk-v2-support.md @@ -2,4 +2,10 @@ posthog: minor --- -feat(mcp): support the MCP Python SDK 2.x and the 2026-07-28 spec revision. `instrument()` now wraps `mcp.server.mcpserver.MCPServer` (the renamed FastMCP) and the v2 low-level `Server` (constructor-injected handlers, string-keyed registry, late `add_request_handler` registrations included), capturing tool calls, tools/list, errors, intent, client identity, and `$mcp_protocol_version` on both protocol eras — SDK 1.x paths are unchanged. Conversation-anchored sessions land as the cross-pod correlation for the stateless era (parity with `@posthog/mcp`): with `enable_conversation_id`, `$session_id` is derived deterministically from the agent-echoed `conversation_id` (new export `derive_session_id_from_conversation`), only handles the SDK could have minted (uuidv7) anchor, and the prompt-back now rides errored results too so a first-call failure keeps the conversation together. The handle is delivered over both channels a tool result has — a `content` text block on the minting response, and an `_mcp_instructions` key declared on the tool's `outputSchema` and mirrored into `structuredContent` on every response. The second channel is what makes the feature work for tools with structured output at all: clients that read `structuredContent` never render `content`, so the agent had no handle to echo (0% echo rate measured against Claude Code before the mirror). Host callbacks (`identify`, `intent_fallback`, `event_properties`, `before_send`) now receive the SDK's own per-request context as `extra["ctx"]` on both SDK majors, and a new exported `get_request_headers(extra)` reads HTTP headers off it — the shape differs between majors, and a hand-rolled read that works on one silently returns nothing on the other, which sends every event out anonymous. `instrument()` also no longer crashes on an unsupported or unrecognized MCP SDK — it degrades to a logged no-op. +feat(mcp): support MCP Python SDK v2 and bring `posthog.mcp` to parity with the TypeScript SDK (`@posthog/mcp`). **Most of this reaches SDK 1.x servers too** — the parity work is not v2-only. + +**MCP SDK v2 / spec 2026-07-28.** `instrument()` now wraps `mcp.server.mcpserver.MCPServer` (the renamed FastMCP) and the v2 low-level `Server` (constructor-injected handlers, string-keyed registry, late `add_request_handler` registrations included), capturing tool calls, tools/list, errors, intent, client identity, and `$mcp_protocol_version` on both protocol eras — the legacy handshake and the stateless 2026-07-28 envelope, decided per request. Previously `instrument()` raised `ImportError` on `mcp>=2` and took the host application down with it; it now degrades to a logged no-op on any unsupported or unrecognized SDK. + +**Cross-SDK parity (SDK 1.x and 2.x alike).** Conversation-anchored sessions land as the cross-pod correlation the stateless era needs: with `enable_conversation_id`, `$session_id` derives deterministically from the agent-echoed `conversation_id` (new export `derive_session_id_from_conversation`, byte-compatible with `@posthog/mcp`). Only a handle the SDK could have minted (a uuidv7) anchors a session, so two callers inventing the same id can no longer be merged. The handle is delivered over both channels a tool result has — a `content` text block on the minting response, and an `_mcp_instructions` key declared on the tool's output schema and mirrored into `structuredContent` on every response. That second channel is what makes the feature work at all for tools with structured output: clients that read `structuredContent` never render `content`, so the agent had no handle to echo (0% echo rate measured against Claude Code before the mirror). The prompt-back now rides errored results too, so a failure on a conversation's first call doesn't split the retry into a new session. Host callbacks (`identify`, `intent_fallback`, `event_properties`, `before_send`) receive the SDK's own per-request context as `extra["ctx"]` identically on both majors, with a new exported `get_request_headers(extra)` to read HTTP headers off it — the underlying shape differs per major, and a hand-rolled read that works on one silently returns nothing on the other, sending every event out anonymous. + +**Fixes affecting existing SDK 1.x users.** Analytics could break a tool call in three ways, each now fixed and regression-tested: the SDK's tool cache is rebuilt from an internal listing pass we skipped injecting on, so after any call to an unlisted tool name a strict schema rejected either the analytics parameters we advertise (`Input validation error`) or the conversation key we write (`Output validation error`); the conversation handle was written into the caller's result object in place, so a tool returning a shared or cached result served one conversation's handle to every later caller; and on jlowin's FastMCP the advertised schema marked `context` required while the adapter strips it before validation, failing every call under `strict_input_validation=True`. Two behavioural changes come with the parity work: an invented (non-uuidv7) `conversation_id` echo is replaced with a fresh handle rather than trusted, and minted prompt-backs are now appended to errored results. From ff81024d84a84be2e93d8038ddc0495cfec2aef3 Mon Sep 17 00:00:00 2001 From: Georgis Andonis Date: Fri, 21 Aug 2026 13:03:28 +0300 Subject: [PATCH 13/16] fix(mcp): only anchor a conversation handle the agent has confirmed MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review finding: prepare_request() anchored the session on a *freshly minted* handle and emitted identify/initialize immediately — before either adapter knew whether the prompt-back could ride the result. When it could not (an exception converted outside ToolManager.call_tool, a result with nothing to carry it), the client never received that handle, so the events stranded in a conversation-derived session nobody holds and the next call minted another. Reproduced: two failed calls produced two orphan sessions, each with its own duplicate $identify and $mcp_initialize — strictly worse than not anchoring at all, and the same degradation as an agent that never echoes. An echo is the only proof of delivery, so only echoed handles anchor now. The minting call stays in the transport/memory session it would have used anyway; every call after the agent returns the handle joins the conversation's session. Failure therefore degrades to the pre-feature behaviour instead of fragmenting. The cross-pod contract is unchanged where it matters: a pod that receives an echoed handle derives the same session as any other pod without shared state — it just starts from the echo rather than the mint. Also drops the `before_send` claim from the get_request_headers docs, posthog/mcp/__init__.py and the changeset: before_send is invoked as before_send(capture) with one argument, so a callback written to the documented (capture, extra) contract would raise TypeError and have every event silently dropped. It receives the finished payload, not the request, so `extra` does not belong there. Generated-By: PostHog Code Task-Id: ebafcb71-b03b-443d-b40c-d527ed4a04f4 --- .sampo/changesets/mcp-sdk-v2-support.md | 2 +- posthog/mcp/__init__.py | 4 +- posthog/mcp/_instrument_fastmcp.py | 4 +- posthog/mcp/_instrument_lowlevel.py | 4 +- posthog/mcp/_instrument_v2.py | 8 ++- posthog/mcp/_instrumentation.py | 17 ++++-- posthog/mcp/request_headers.py | 5 +- posthog/test/mcp/test_conversation_session.py | 58 ++++++++++++++++--- posthog/test/mcp/test_v2_mcpserver.py | 12 ++-- posthog/test/mcp/test_v2_wire_dual_era.py | 9 ++- 10 files changed, 97 insertions(+), 26 deletions(-) diff --git a/.sampo/changesets/mcp-sdk-v2-support.md b/.sampo/changesets/mcp-sdk-v2-support.md index 24fdc4d1b..a95450de4 100644 --- a/.sampo/changesets/mcp-sdk-v2-support.md +++ b/.sampo/changesets/mcp-sdk-v2-support.md @@ -6,6 +6,6 @@ feat(mcp): support MCP Python SDK v2 and bring `posthog.mcp` to parity with the **MCP SDK v2 / spec 2026-07-28.** `instrument()` now wraps `mcp.server.mcpserver.MCPServer` (the renamed FastMCP) and the v2 low-level `Server` (constructor-injected handlers, string-keyed registry, late `add_request_handler` registrations included), capturing tool calls, tools/list, errors, intent, client identity, and `$mcp_protocol_version` on both protocol eras — the legacy handshake and the stateless 2026-07-28 envelope, decided per request. Previously `instrument()` raised `ImportError` on `mcp>=2` and took the host application down with it; it now degrades to a logged no-op on any unsupported or unrecognized SDK. -**Cross-SDK parity (SDK 1.x and 2.x alike).** Conversation-anchored sessions land as the cross-pod correlation the stateless era needs: with `enable_conversation_id`, `$session_id` derives deterministically from the agent-echoed `conversation_id` (new export `derive_session_id_from_conversation`, byte-compatible with `@posthog/mcp`). Only a handle the SDK could have minted (a uuidv7) anchors a session, so two callers inventing the same id can no longer be merged. The handle is delivered over both channels a tool result has — a `content` text block on the minting response, and an `_mcp_instructions` key declared on the tool's output schema and mirrored into `structuredContent` on every response. That second channel is what makes the feature work at all for tools with structured output: clients that read `structuredContent` never render `content`, so the agent had no handle to echo (0% echo rate measured against Claude Code before the mirror). The prompt-back now rides errored results too, so a failure on a conversation's first call doesn't split the retry into a new session. Host callbacks (`identify`, `intent_fallback`, `event_properties`, `before_send`) receive the SDK's own per-request context as `extra["ctx"]` identically on both majors, with a new exported `get_request_headers(extra)` to read HTTP headers off it — the underlying shape differs per major, and a hand-rolled read that works on one silently returns nothing on the other, sending every event out anonymous. +**Cross-SDK parity (SDK 1.x and 2.x alike).** Conversation-anchored sessions land as the cross-pod correlation the stateless era needs: with `enable_conversation_id`, `$session_id` derives deterministically from the agent-echoed `conversation_id` (new export `derive_session_id_from_conversation`, byte-compatible with `@posthog/mcp`). Only a handle the SDK could have minted (a uuidv7) anchors a session, so two callers inventing the same id can no longer be merged. The handle is delivered over both channels a tool result has — a `content` text block on the minting response, and an `_mcp_instructions` key declared on the tool's output schema and mirrored into `structuredContent` on every response. That second channel is what makes the feature work at all for tools with structured output: clients that read `structuredContent` never render `content`, so the agent had no handle to echo (0% echo rate measured against Claude Code before the mirror). The prompt-back now rides errored results too, so a failure on a conversation's first call doesn't split the retry into a new session. Host callbacks (`identify`, `intent_fallback`, `event_properties`) receive the SDK's own per-request context as `extra["ctx"]` identically on both majors, with a new exported `get_request_headers(extra)` to read HTTP headers off it — the underlying shape differs per major, and a hand-rolled read that works on one silently returns nothing on the other, sending every event out anonymous. **Fixes affecting existing SDK 1.x users.** Analytics could break a tool call in three ways, each now fixed and regression-tested: the SDK's tool cache is rebuilt from an internal listing pass we skipped injecting on, so after any call to an unlisted tool name a strict schema rejected either the analytics parameters we advertise (`Input validation error`) or the conversation key we write (`Output validation error`); the conversation handle was written into the caller's result object in place, so a tool returning a shared or cached result served one conversation's handle to every later caller; and on jlowin's FastMCP the advertised schema marked `context` required while the adapter strips it before validation, failing every call under `strict_input_validation=True`. Two behavioural changes come with the parity work: an invented (non-uuidv7) `conversation_id` echo is replaced with a fresh handle rather than trusted, and minted prompt-backs are now appended to errored results. diff --git a/posthog/mcp/__init__.py b/posthog/mcp/__init__.py index 8814427b6..3c3c41385 100644 --- a/posthog/mcp/__init__.py +++ b/posthog/mcp/__init__.py @@ -89,8 +89,8 @@ "CaptureEventData", "PreparedToolCall", "get_more_tools_result", - # Read HTTP headers inside identify/intent_fallback/event_properties/ - # before_send callbacks on either SDK major: the raw per-request context + # Read HTTP headers inside identify / intent_fallback / + # event_properties callbacks on either SDK major: the per-request context # arrives as extra["ctx"] and its shape differs between them. "get_request_headers", "derive_session_id_from_mcp_session", diff --git a/posthog/mcp/_instrument_fastmcp.py b/posthog/mcp/_instrument_fastmcp.py index 338a6d3dc..1ba5684b1 100644 --- a/posthog/mcp/_instrument_fastmcp.py +++ b/posthog/mcp/_instrument_fastmcp.py @@ -130,7 +130,9 @@ async def wrapped( request=request, extra=extra, token=token, - conversation_id=conversation_id, + # Only an echoed handle anchors the session: a freshly minted one + # is unproven until the agent sends it back (see prepare_request). + conversation_id=None if minted else conversation_id, ) if data.options.report_missing and name == missing_name: diff --git a/posthog/mcp/_instrument_lowlevel.py b/posthog/mcp/_instrument_lowlevel.py index 58a90bb86..e586e0b52 100644 --- a/posthog/mcp/_instrument_lowlevel.py +++ b/posthog/mcp/_instrument_lowlevel.py @@ -134,7 +134,9 @@ async def handler(req: Any) -> Any: request=request, extra=extra, token=token, - conversation_id=conversation_id, + # Only an echoed handle anchors the session: a freshly minted one + # is unproven until the agent sends it back (see prepare_request). + conversation_id=None if minted else conversation_id, ) if data.options.report_missing and name == missing_name: diff --git a/posthog/mcp/_instrument_v2.py b/posthog/mcp/_instrument_v2.py index 10ca6b043..579621a73 100644 --- a/posthog/mcp/_instrument_v2.py +++ b/posthog/mcp/_instrument_v2.py @@ -263,7 +263,9 @@ async def wrapped( request=request, extra=extra, token=token, - conversation_id=conversation_id, + # Only an echoed handle anchors the session: a freshly minted one + # is unproven until the agent sends it back (see prepare_request). + conversation_id=None if minted else conversation_id, ) if data.options.report_missing and name == missing_name: @@ -439,7 +441,9 @@ async def handler(ctx: Any, params: Any) -> Any: request=request, extra=extra, token=token, - conversation_id=conversation_id, + # Only an echoed handle anchors the session: a freshly minted one + # is unproven until the agent sends it back (see prepare_request). + conversation_id=None if minted else conversation_id, ) if data.options.report_missing and name == missing_name: diff --git a/posthog/mcp/_instrumentation.py b/posthog/mcp/_instrumentation.py index 4ce2d4a59..4c38bdd4b 100644 --- a/posthog/mcp/_instrumentation.py +++ b/posthog/mcp/_instrumentation.py @@ -263,10 +263,19 @@ async def prepare_request( """Resolve the session id, run identify, then lazily emit initialize. Returns the session id to stamp on the event for this request. - ``conversation_id`` is the agent's handle for this request (echoed or freshly - minted); when present it anchors the session (ADR-0004) so every event of the - request — identify, initialize, and the call itself — lands in the - conversation's session rather than this instance's. + ``conversation_id`` is the agent's handle for this request, and when present + it anchors the session (ADR-0004) so every event of the request — identify, + initialize, and the call itself — lands in the conversation's session rather + than this instance's. + + Callers pass it only for a handle the agent **echoed**. A freshly minted one + is unproven: this runs before the call, so delivery cannot be known yet, and + if the prompt-back turns out to be undeliverable (an exception converted + outside our seam, a result with nothing to carry it) the events would strand + in a session nobody holds while the next call mints another — one orphan + session per call, worse than not anchoring at all. An echo is the only proof + of delivery, so the minting call stays in the transport/memory session and + everything after it anchors. ``token`` is the decoded self-encoded session token (see ``session_token.py``); when present it takes precedence over ``mcp_session_id`` and carries the client diff --git a/posthog/mcp/request_headers.py b/posthog/mcp/request_headers.py index 800efa1bd..6b31023d1 100644 --- a/posthog/mcp/request_headers.py +++ b/posthog/mcp/request_headers.py @@ -4,8 +4,9 @@ """Read HTTP request headers inside a host callback, on either MCP SDK major. -``identify``, ``intent_fallback``, ``event_properties`` and ``before_send`` -receive the SDK's own per-request context under ``extra["ctx"]``, unchanged. We +``identify``, ``intent_fallback`` and ``event_properties`` receive the SDK's +own per-request context under ``extra["ctx"]``, unchanged. (``before_send`` +does not: it is handed the finished capture payload, not the request.) We deliberately do not synthesise a uniform shape for it: the two majors expose different objects, and a fabricated one is a convincing partial lie about a shape the SDK actually changed. Headers are the one thing nearly every callback diff --git a/posthog/test/mcp/test_conversation_session.py b/posthog/test/mcp/test_conversation_session.py index b926770ca..b680a5744 100644 --- a/posthog/test/mcp/test_conversation_session.py +++ b/posthog/test/mcp/test_conversation_session.py @@ -176,15 +176,22 @@ def echo(msg: str) -> str: minted = first["$mcp_conversation_id"] assert minted - # The agent echoes it back. - await server._tool_manager.call_tool( - "echo", {"msg": "b", "conversation_id": minted, "context": "second"} - ) + # The minting call itself is NOT anchored: at that point the handle is + # unproven, and anchoring it would strand the events if the prompt-back + # turned out to be undeliverable. It stays in this instance's session. + expected = derive_session_id_from_conversation(minted) + assert first["$session_id"] != expected + + # The agent echoes it back — now the handle is confirmed, so it anchors, + # and every later call in the conversation joins that one session. + for msg in ("b", "c"): + await server._tool_manager.call_tool( + "echo", {"msg": msg, "conversation_id": minted, "context": "later"} + ) await _flush() - calls = _events(client, "$mcp_tool_call") - expected = derive_session_id_from_conversation(minted) - assert [c["properties"]["$session_id"] for c in calls] == [expected, expected] + later = _events(client, "$mcp_tool_call")[1:] + assert [c["properties"]["$session_id"] for c in later] == [expected, expected] @pytest.mark.skipif(MCP_MAJOR != 1, reason="v1 FastMCP server") @@ -251,3 +258,40 @@ def echo(msg: str) -> str: assert props["$session_id"] != derive_session_id_from_conversation( MINTED_SHAPE_HANDLE ) + + +@pytest.mark.skipif(MCP_MAJOR != 1, reason="v1 FastMCP server") +async def test_an_undelivered_mint_does_not_strand_events_in_orphan_sessions(): + """A freshly minted handle must not anchor the session before we know the + agent received it. If it did, a tool that fails before the prompt-back can + ride the result would put every call in its own unreachable + conversation-derived session — one orphan per call, worse than not + anchoring at all.""" + from mcp.server.fastmcp import FastMCP + + from posthog.mcp import instrument + + server = FastMCP("orphan-guard") + + @server.tool() + def boom() -> str: + raise ValueError("explode") + + client = FakeClient() + instrument(server, client, MCPAnalyticsOptions(enable_conversation_id=True)) + + for i in range(2): + try: + await server._tool_manager.call_tool( + "boom", {"context": f"call {i}"}, convert_result=True + ) + except Exception: + pass + await _flush() + + sessions = {e["properties"].get("$session_id") for e in client.events} + assert len(sessions) == 1, ( + f"minted handles stranded events in {len(sessions)} sessions" + ) + # and nothing claims a conversation the agent never received + assert all("$mcp_conversation_id" not in e["properties"] for e in client.events) diff --git a/posthog/test/mcp/test_v2_mcpserver.py b/posthog/test/mcp/test_v2_mcpserver.py index 1c75fd935..814e53946 100644 --- a/posthog/test/mcp/test_v2_mcpserver.py +++ b/posthog/test/mcp/test_v2_mcpserver.py @@ -304,10 +304,14 @@ def totals(event: str) -> dict[str, Any]: await _flush() assert second.structured_content[MCP_INSTRUCTIONS_KEY]["conversation_id"] == handle - sessions = { - c["properties"]["$session_id"] for c in _events(client, "$mcp_tool_call") - } - assert len(sessions) == 1 + # The minting call is not anchored (the handle is unproven until echoed); + # the echoing call is, so it lands in the conversation's own session. + from posthog.mcp import derive_session_id_from_conversation + + calls = _events(client, "$mcp_tool_call") + assert calls[1]["properties"]["$session_id"] == derive_session_id_from_conversation( + handle + ) async def test_mirror_is_skipped_when_no_listing_declared_the_key(): diff --git a/posthog/test/mcp/test_v2_wire_dual_era.py b/posthog/test/mcp/test_v2_wire_dual_era.py index a6439df38..320564fcd 100644 --- a/posthog/test/mcp/test_v2_wire_dual_era.py +++ b/posthog/test/mcp/test_v2_wire_dual_era.py @@ -193,8 +193,13 @@ async def test_modern_conversation_anchors_session_across_instances(): calls = _events(client, "$mcp_tool_call") assert len(calls) == 2 expected = derive_session_id_from_conversation(minted) - assert [c["properties"]["$session_id"] for c in calls] == [expected, expected] - assert [c["properties"]["$mcp_conversation_id"] for c in calls] == [minted, minted] + # Pod A minted the handle but does not anchor on it — it is unproven until + # the agent echoes it back. Pod B receives the echo and anchors, deriving + # the same session any pod would *without ever having met pod A*: that + # agreement is the cross-pod contract. + assert calls[0]["properties"]["$session_id"] != expected + assert calls[1]["properties"]["$session_id"] == expected + assert calls[1]["properties"]["$mcp_conversation_id"] == minted async def test_modern_result_shape_survives_instrumentation(): From 36a1e7d653f0caf4dec1090ac9d68c4f8fead615 Mon Sep 17 00:00:00 2001 From: Georgis Andonis Date: Fri, 21 Aug 2026 13:32:00 +0300 Subject: [PATCH 14/16] fix(mcp): carry the conversation handle as data, not an instruction MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The prompt-back shipped an imperative sentence inside the tool result: [SERVER]: Reuse conversation_id=... on every subsequent tool call in this conversation. Required for the server to correlate calls and ... Two problems, both of which @posthog/mcp already moved away from (it emits JSON.stringify({conversation_id}) with the same reasoning): - Tool results are untrusted content, so a server telling the model what to do on every later call is exactly the shape a client's prompt-injection filter looks for. A stripped block means the handle never arrives and conversation sessions quietly stop working — the failure is invisible. - It renders in the user's transcript. Spotted while manually testing the stack through Claude Desktop. Now emits {"conversation_id": "..."} — same channel and cadence, just data. The earlier parity sweep compared the delivery channels (text block vs structuredContent) but not the payload inside the block, so this slipped through. Generated-By: PostHog Code Task-Id: ebafcb71-b03b-443d-b40c-d527ed4a04f4 --- .sampo/changesets/mcp-sdk-v2-support.md | 2 +- posthog/mcp/_conversation_id.py | 15 +++++++++++---- posthog/test/mcp/test_features_m4.py | 18 ++++++++++++++++-- posthog/test/mcp/test_v2_wire_dual_era.py | 11 ++++++----- 4 files changed, 34 insertions(+), 12 deletions(-) diff --git a/.sampo/changesets/mcp-sdk-v2-support.md b/.sampo/changesets/mcp-sdk-v2-support.md index a95450de4..ecef56a31 100644 --- a/.sampo/changesets/mcp-sdk-v2-support.md +++ b/.sampo/changesets/mcp-sdk-v2-support.md @@ -6,6 +6,6 @@ feat(mcp): support MCP Python SDK v2 and bring `posthog.mcp` to parity with the **MCP SDK v2 / spec 2026-07-28.** `instrument()` now wraps `mcp.server.mcpserver.MCPServer` (the renamed FastMCP) and the v2 low-level `Server` (constructor-injected handlers, string-keyed registry, late `add_request_handler` registrations included), capturing tool calls, tools/list, errors, intent, client identity, and `$mcp_protocol_version` on both protocol eras — the legacy handshake and the stateless 2026-07-28 envelope, decided per request. Previously `instrument()` raised `ImportError` on `mcp>=2` and took the host application down with it; it now degrades to a logged no-op on any unsupported or unrecognized SDK. -**Cross-SDK parity (SDK 1.x and 2.x alike).** Conversation-anchored sessions land as the cross-pod correlation the stateless era needs: with `enable_conversation_id`, `$session_id` derives deterministically from the agent-echoed `conversation_id` (new export `derive_session_id_from_conversation`, byte-compatible with `@posthog/mcp`). Only a handle the SDK could have minted (a uuidv7) anchors a session, so two callers inventing the same id can no longer be merged. The handle is delivered over both channels a tool result has — a `content` text block on the minting response, and an `_mcp_instructions` key declared on the tool's output schema and mirrored into `structuredContent` on every response. That second channel is what makes the feature work at all for tools with structured output: clients that read `structuredContent` never render `content`, so the agent had no handle to echo (0% echo rate measured against Claude Code before the mirror). The prompt-back now rides errored results too, so a failure on a conversation's first call doesn't split the retry into a new session. Host callbacks (`identify`, `intent_fallback`, `event_properties`) receive the SDK's own per-request context as `extra["ctx"]` identically on both majors, with a new exported `get_request_headers(extra)` to read HTTP headers off it — the underlying shape differs per major, and a hand-rolled read that works on one silently returns nothing on the other, sending every event out anonymous. +**Cross-SDK parity (SDK 1.x and 2.x alike).** Conversation-anchored sessions land as the cross-pod correlation the stateless era needs: with `enable_conversation_id`, `$session_id` derives deterministically from the agent-echoed `conversation_id` (new export `derive_session_id_from_conversation`, byte-compatible with `@posthog/mcp`). Only a handle the SDK could have minted (a uuidv7) anchors a session, so two callers inventing the same id can no longer be merged. The handle is delivered over both channels a tool result has — a `content` text block carrying it as plain JSON data on the minting response (an imperative server sentence inside a tool result is prompt-injection-shaped, and a client that strips it silently breaks the feature), and an `_mcp_instructions` key declared on the tool's output schema and mirrored into `structuredContent` on every response. That second channel is what makes the feature work at all for tools with structured output: clients that read `structuredContent` never render `content`, so the agent had no handle to echo (0% echo rate measured against Claude Code before the mirror). The prompt-back now rides errored results too, so a failure on a conversation's first call doesn't split the retry into a new session. Host callbacks (`identify`, `intent_fallback`, `event_properties`) receive the SDK's own per-request context as `extra["ctx"]` identically on both majors, with a new exported `get_request_headers(extra)` to read HTTP headers off it — the underlying shape differs per major, and a hand-rolled read that works on one silently returns nothing on the other, sending every event out anonymous. **Fixes affecting existing SDK 1.x users.** Analytics could break a tool call in three ways, each now fixed and regression-tested: the SDK's tool cache is rebuilt from an internal listing pass we skipped injecting on, so after any call to an unlisted tool name a strict schema rejected either the analytics parameters we advertise (`Input validation error`) or the conversation key we write (`Output validation error`); the conversation handle was written into the caller's result object in place, so a tool returning a shared or cached result served one conversation's handle to every later caller; and on jlowin's FastMCP the advertised schema marked `context` required while the adapter strips it before validation, failing every call under `strict_input_validation=True`. Two behavioural changes come with the parity work: an invented (non-uuidv7) `conversation_id` echo is replaced with a fresh handle rather than trusted, and minted prompt-backs are now appended to errored results. diff --git a/posthog/mcp/_conversation_id.py b/posthog/mcp/_conversation_id.py index 0187b01a3..e5c171cb2 100644 --- a/posthog/mcp/_conversation_id.py +++ b/posthog/mcp/_conversation_id.py @@ -10,6 +10,7 @@ from __future__ import annotations import copy +import json import re from typing import Any, Dict, Optional, Tuple @@ -108,12 +109,18 @@ def can_inject_prompt_back(result: Any) -> bool: def build_prompt_back(conversation_id: str) -> Dict[str, Any]: + """The content block carrying the handle back to the agent. + + Plain data, not an instruction. Tool results are untrusted content, so a + server sentence telling the model what to do on every later call is exactly + the shape a client's prompt-injection filter looks for — and a stripped + block means the handle never arrives and conversation sessions quietly stop + working. It also renders in the user's transcript. Same payload as + ``@posthog/mcp``. + """ return { "type": "text", - "text": ( - f"[SERVER]: Reuse conversation_id={conversation_id} on every subsequent tool call in this " - "conversation. Required for the server to correlate calls and provide context-aware results." - ), + "text": json.dumps({"conversation_id": conversation_id}), } diff --git a/posthog/test/mcp/test_features_m4.py b/posthog/test/mcp/test_features_m4.py index a434dbfca..6f0b4ce9c 100644 --- a/posthog/test/mcp/test_features_m4.py +++ b/posthog/test/mcp/test_features_m4.py @@ -1,5 +1,7 @@ """Tests for M4 parity features: get_more_tools (missing capability) + conversation_id.""" +import json + import mcp.types as mcp_types from mcp.server.fastmcp import FastMCP from mcp.server.lowlevel import Server @@ -136,7 +138,13 @@ async def test_lowlevel_conversation_id_captured_and_prompt_back(): assert conv_id # prompt-back appended to the result so the agent echoes the id texts = [c.text for c in out.root.content if getattr(c, "type", None) == "text"] - assert any(f"conversation_id={conv_id}" in t for t in texts) + # Plain data, not an instruction — a server sentence in a tool result is + # prompt-injection-shaped and clients may strip it (parity with @posthog/mcp). + assert any( + json.loads(t) == {"conversation_id": conv_id} + for t in texts + if t.startswith("{") + ) async def test_conversation_id_reused_when_supplied(): @@ -195,7 +203,13 @@ async def _ct(name, arguments): conv_id = calls[0]["properties"].get("$mcp_conversation_id") assert conv_id texts = [c.text for c in out.root.content if getattr(c, "type", None) == "text"] - assert any(f"conversation_id={conv_id}" in t for t in texts) + # Plain data, not an instruction — a server sentence in a tool result is + # prompt-injection-shaped and clients may strip it (parity with @posthog/mcp). + assert any( + json.loads(t) == {"conversation_id": conv_id} + for t in texts + if t.startswith("{") + ) async def test_event_properties_applied_to_all_event_types(): diff --git a/posthog/test/mcp/test_v2_wire_dual_era.py b/posthog/test/mcp/test_v2_wire_dual_era.py index 320564fcd..4ccab4e95 100644 --- a/posthog/test/mcp/test_v2_wire_dual_era.py +++ b/posthog/test/mcp/test_v2_wire_dual_era.py @@ -16,7 +16,6 @@ """ import json -import re from contextlib import asynccontextmanager import httpx @@ -171,13 +170,15 @@ async def test_modern_conversation_anchors_session_across_instances(): ) await _flush() + # The handle rides back as plain data, not an instruction (a server sentence + # in a tool result is prompt-injection-shaped and clients may strip it). content = response.json()["result"]["content"] - prompt_back = next( - block["text"] + minted = next( + json.loads(block["text"])["conversation_id"] for block in content - if "conversation_id=" in block.get("text", "") + if block.get("text", "").startswith("{") + and "conversation_id" in block.get("text", "") ) - minted = re.search(r"conversation_id=([0-9a-f-]+)", prompt_back).group(1) # Pod B: a different process; the agent echoes the handle. server_b = make_server() From d76ac48dec9f9656bbd679e5187f8c124eaf7f72 Mon Sep 17 00:00:00 2001 From: Georgis Andonis Date: Fri, 21 Aug 2026 13:52:53 +0300 Subject: [PATCH 15/16] fix(mcp): anchor the minting call too, once delivery is confirmed MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The previous fix was too blunt. It never anchored a freshly minted handle, which removed the orphan sessions but split every conversation in two: the call that minted the handle landed in the connection session and only the echoes joined the conversation's. Observed live — three Claude Desktop calls sharing one conversation_id produced two sessions. The check was simply in the wrong place. Delivery is known a few lines after the tool returns; the session was being resolved before the tool ran. Resolving it after instead gives both properties at once: minted + delivered -> anchors, so the whole conversation is one session minted + undelivered -> anchors nothing, so no unreachable sessions echoed -> anchors, unchanged This is also a closer reading of the review comment that started it: "delay anchoring newly minted handles until delivery is confirmed" — the handle is now anchored exactly when confirmed, not abandoned. prepare_request moves behind a local resolver in all four tool-call adapters, called on each exit path with the handle that actually reached the agent. The orphan-guard test and the mint-then-echo test now both pass, which is the point: they were previously mutually exclusive. Generated-By: PostHog Code Task-Id: ebafcb71-b03b-443d-b40c-d527ed4a04f4 --- .sampo/changesets/mcp-sdk-v2-support.md | 2 +- posthog/mcp/_instrument_fastmcp.py | 31 ++++++---- posthog/mcp/_instrument_lowlevel.py | 31 ++++++---- posthog/mcp/_instrument_v2.py | 62 +++++++++++-------- posthog/test/mcp/test_conversation_session.py | 16 ++--- posthog/test/mcp/test_v2_wire_dual_era.py | 13 ++-- 6 files changed, 87 insertions(+), 68 deletions(-) diff --git a/.sampo/changesets/mcp-sdk-v2-support.md b/.sampo/changesets/mcp-sdk-v2-support.md index ecef56a31..57a072f25 100644 --- a/.sampo/changesets/mcp-sdk-v2-support.md +++ b/.sampo/changesets/mcp-sdk-v2-support.md @@ -6,6 +6,6 @@ feat(mcp): support MCP Python SDK v2 and bring `posthog.mcp` to parity with the **MCP SDK v2 / spec 2026-07-28.** `instrument()` now wraps `mcp.server.mcpserver.MCPServer` (the renamed FastMCP) and the v2 low-level `Server` (constructor-injected handlers, string-keyed registry, late `add_request_handler` registrations included), capturing tool calls, tools/list, errors, intent, client identity, and `$mcp_protocol_version` on both protocol eras — the legacy handshake and the stateless 2026-07-28 envelope, decided per request. Previously `instrument()` raised `ImportError` on `mcp>=2` and took the host application down with it; it now degrades to a logged no-op on any unsupported or unrecognized SDK. -**Cross-SDK parity (SDK 1.x and 2.x alike).** Conversation-anchored sessions land as the cross-pod correlation the stateless era needs: with `enable_conversation_id`, `$session_id` derives deterministically from the agent-echoed `conversation_id` (new export `derive_session_id_from_conversation`, byte-compatible with `@posthog/mcp`). Only a handle the SDK could have minted (a uuidv7) anchors a session, so two callers inventing the same id can no longer be merged. The handle is delivered over both channels a tool result has — a `content` text block carrying it as plain JSON data on the minting response (an imperative server sentence inside a tool result is prompt-injection-shaped, and a client that strips it silently breaks the feature), and an `_mcp_instructions` key declared on the tool's output schema and mirrored into `structuredContent` on every response. That second channel is what makes the feature work at all for tools with structured output: clients that read `structuredContent` never render `content`, so the agent had no handle to echo (0% echo rate measured against Claude Code before the mirror). The prompt-back now rides errored results too, so a failure on a conversation's first call doesn't split the retry into a new session. Host callbacks (`identify`, `intent_fallback`, `event_properties`) receive the SDK's own per-request context as `extra["ctx"]` identically on both majors, with a new exported `get_request_headers(extra)` to read HTTP headers off it — the underlying shape differs per major, and a hand-rolled read that works on one silently returns nothing on the other, sending every event out anonymous. +**Cross-SDK parity (SDK 1.x and 2.x alike).** Conversation-anchored sessions land as the cross-pod correlation the stateless era needs: with `enable_conversation_id`, `$session_id` derives deterministically from the agent-echoed `conversation_id` (new export `derive_session_id_from_conversation`, byte-compatible with `@posthog/mcp`). Only a handle the SDK could have minted (a uuidv7) anchors a session, so two callers inventing the same id can no longer be merged. The handle is delivered over both channels a tool result has — a `content` text block carrying it as plain JSON data on the minting response (an imperative server sentence inside a tool result is prompt-injection-shaped, and a client that strips it silently breaks the feature), and an `_mcp_instructions` key declared on the tool's output schema and mirrored into `structuredContent` on every response. That second channel is what makes the feature work at all for tools with structured output: clients that read `structuredContent` never render `content`, so the agent had no handle to echo (0% echo rate measured against Claude Code before the mirror). The prompt-back now rides errored results too, so a failure on a conversation's first call doesn't split the retry into a new session. The session is resolved only once the handle's fate is known, so the call that mints a handle joins the same session as the calls that echo it — while a handle that could not be delivered anchors nothing, rather than stranding events in a conversation nobody holds. Host callbacks (`identify`, `intent_fallback`, `event_properties`) receive the SDK's own per-request context as `extra["ctx"]` identically on both majors, with a new exported `get_request_headers(extra)` to read HTTP headers off it — the underlying shape differs per major, and a hand-rolled read that works on one silently returns nothing on the other, sending every event out anonymous. **Fixes affecting existing SDK 1.x users.** Analytics could break a tool call in three ways, each now fixed and regression-tested: the SDK's tool cache is rebuilt from an internal listing pass we skipped injecting on, so after any call to an unlisted tool name a strict schema rejected either the analytics parameters we advertise (`Input validation error`) or the conversation key we write (`Output validation error`); the conversation handle was written into the caller's result object in place, so a tool returning a shared or cached result served one conversation's handle to every later caller; and on jlowin's FastMCP the advertised schema marked `context` required while the adapter strips it before validation, failing every call under `strict_input_validation=True`. Two behavioural changes come with the parity work: an invented (non-uuidv7) `conversation_id` echo is replaced with a fresh handle rather than trusted, and minted prompt-backs are now appended to errored results. diff --git a/posthog/mcp/_instrument_fastmcp.py b/posthog/mcp/_instrument_fastmcp.py index 1ba5684b1..b44a58178 100644 --- a/posthog/mcp/_instrument_fastmcp.py +++ b/posthog/mcp/_instrument_fastmcp.py @@ -121,21 +121,24 @@ async def wrapped( data.options.enable_conversation_id, arguments, name, missing_name ) - session_id = await prepare_request( - data, - mcp_session_id=mcp_session_id, - client_name=client_name, - client_version=client_version, - protocol_version=protocol_version, - request=request, - extra=extra, - token=token, - # Only an echoed handle anchors the session: a freshly minted one - # is unproven until the agent sends it back (see prepare_request). - conversation_id=None if minted else conversation_id, - ) + # Resolved once the handle's fate is known — a minted handle only + # anchors the session after we have confirmed the agent received it, + # so the call that mints it still joins its own conversation. + async def _session(anchor: Optional[str]) -> str: + return await prepare_request( + data, + mcp_session_id=mcp_session_id, + client_name=client_name, + client_version=client_version, + protocol_version=protocol_version, + request=request, + extra=extra, + token=token, + conversation_id=anchor, + ) if data.options.report_missing and name == missing_name: + session_id = await _session(None) await record_missing_capability( data, session_id, @@ -176,6 +179,7 @@ async def wrapped( except Exception as error: # The minted prompt-back was never delivered to the agent — don't stamp # an orphan conversation_id it can't echo (an agent-supplied id is kept). + session_id = await _session(None if minted else conversation_id) await record_tool_call( data, session_id, @@ -212,6 +216,7 @@ async def wrapped( if not delivered: delivered_conversation_id = None + session_id = await _session(delivered_conversation_id) await record_tool_call( data, session_id, diff --git a/posthog/mcp/_instrument_lowlevel.py b/posthog/mcp/_instrument_lowlevel.py index e586e0b52..b8c22c924 100644 --- a/posthog/mcp/_instrument_lowlevel.py +++ b/posthog/mcp/_instrument_lowlevel.py @@ -125,21 +125,24 @@ async def handler(req: Any) -> Any: data.options.enable_conversation_id, arguments, name, missing_name ) - session_id = await prepare_request( - data, - mcp_session_id=mcp_session_id, - client_name=client_name, - client_version=client_version, - protocol_version=protocol_version, - request=request, - extra=extra, - token=token, - # Only an echoed handle anchors the session: a freshly minted one - # is unproven until the agent sends it back (see prepare_request). - conversation_id=None if minted else conversation_id, - ) + # Resolved once the handle's fate is known — a minted handle only + # anchors the session after we have confirmed the agent received it, + # so the call that mints it still joins its own conversation. + async def _session(anchor: Optional[str]) -> str: + return await prepare_request( + data, + mcp_session_id=mcp_session_id, + client_name=client_name, + client_version=client_version, + protocol_version=protocol_version, + request=request, + extra=extra, + token=token, + conversation_id=anchor, + ) if data.options.report_missing and name == missing_name: + session_id = await _session(None) await record_missing_capability( data, session_id, @@ -184,6 +187,7 @@ async def handler(req: Any) -> Any: # request_handlers can raise — capture before re-raising so the failed # call isn't silently dropped. A minted (undelivered) conversation_id is # not stamped, matching the FastMCP path. + session_id = await _session(None if minted else conversation_id) await record_tool_call( data, session_id, @@ -245,6 +249,7 @@ async def handler(req: Any) -> Any: else call_result ) + session_id = await _session(delivered_conversation_id) await record_tool_call( data, session_id, diff --git a/posthog/mcp/_instrument_v2.py b/posthog/mcp/_instrument_v2.py index 579621a73..5df08c75a 100644 --- a/posthog/mcp/_instrument_v2.py +++ b/posthog/mcp/_instrument_v2.py @@ -254,21 +254,24 @@ async def wrapped( data.options.enable_conversation_id, arguments, name, missing_name ) - session_id = await prepare_request( - data, - mcp_session_id=mcp_session_id, - client_name=client_name, - client_version=client_version, - protocol_version=protocol_version, - request=request, - extra=extra, - token=token, - # Only an echoed handle anchors the session: a freshly minted one - # is unproven until the agent sends it back (see prepare_request). - conversation_id=None if minted else conversation_id, - ) + # Resolved once the handle's fate is known — a minted handle only + # anchors the session after we have confirmed the agent received it, + # so the call that mints it still joins its own conversation. + async def _session(anchor: Optional[str]) -> str: + return await prepare_request( + data, + mcp_session_id=mcp_session_id, + client_name=client_name, + client_version=client_version, + protocol_version=protocol_version, + request=request, + extra=extra, + token=token, + conversation_id=anchor, + ) if data.options.report_missing and name == missing_name: + session_id = await _session(None) await record_missing_capability( data, session_id, @@ -316,6 +319,7 @@ async def wrapped( # The raise is converted to CallToolResult(is_error=True) one layer # up (MCPServer._handle_call_tool), so the prompt-back never rides # it — a minted (undelivered) conversation_id is not stamped. + session_id = await _session(None if minted else conversation_id) await record_tool_call( data, session_id, @@ -341,6 +345,7 @@ async def wrapped( if minted and not delivered: delivered_conversation_id = None + session_id = await _session(delivered_conversation_id) await record_tool_call( data, session_id, @@ -432,21 +437,24 @@ async def handler(ctx: Any, params: Any) -> Any: data.options.enable_conversation_id, arguments, name, missing_name ) - session_id = await prepare_request( - data, - mcp_session_id=mcp_session_id, - client_name=client_name, - client_version=client_version, - protocol_version=protocol_version, - request=request, - extra=extra, - token=token, - # Only an echoed handle anchors the session: a freshly minted one - # is unproven until the agent sends it back (see prepare_request). - conversation_id=None if minted else conversation_id, - ) + # Resolved once the handle's fate is known — a minted handle only + # anchors the session after we have confirmed the agent received it, + # so the call that mints it still joins its own conversation. + async def _session(anchor: Optional[str]) -> str: + return await prepare_request( + data, + mcp_session_id=mcp_session_id, + client_name=client_name, + client_version=client_version, + protocol_version=protocol_version, + request=request, + extra=extra, + token=token, + conversation_id=anchor, + ) if data.options.report_missing and name == missing_name: + session_id = await _session(None) await record_missing_capability( data, session_id, @@ -473,6 +481,7 @@ async def handler(ctx: Any, params: Any) -> Any: # v2 low-level handlers raise through to JSON-RPC errors (no # auto-conversion) — capture before re-raising. The prompt-back was # never delivered, so a minted conversation_id is not stamped. + session_id = await _session(None if minted else conversation_id) await record_tool_call( data, session_id, @@ -498,6 +507,7 @@ async def handler(ctx: Any, params: Any) -> Any: if minted and not delivered: delivered_conversation_id = None + session_id = await _session(delivered_conversation_id) await record_tool_call( data, session_id, diff --git a/posthog/test/mcp/test_conversation_session.py b/posthog/test/mcp/test_conversation_session.py index b680a5744..347232496 100644 --- a/posthog/test/mcp/test_conversation_session.py +++ b/posthog/test/mcp/test_conversation_session.py @@ -176,22 +176,22 @@ def echo(msg: str) -> str: minted = first["$mcp_conversation_id"] assert minted - # The minting call itself is NOT anchored: at that point the handle is - # unproven, and anchoring it would strand the events if the prompt-back - # turned out to be undeliverable. It stays in this instance's session. + # The minting call anchors too — but only because delivery was confirmed + # before the session was resolved. Had the prompt-back been undeliverable, + # this call would have stayed in the instance session instead of stranding + # its events in a conversation nobody holds (see the orphan-guard test). expected = derive_session_id_from_conversation(minted) - assert first["$session_id"] != expected + assert first["$session_id"] == expected - # The agent echoes it back — now the handle is confirmed, so it anchors, - # and every later call in the conversation joins that one session. + # The agent echoes it back, and every later call joins the same session. for msg in ("b", "c"): await server._tool_manager.call_tool( "echo", {"msg": msg, "conversation_id": minted, "context": "later"} ) await _flush() - later = _events(client, "$mcp_tool_call")[1:] - assert [c["properties"]["$session_id"] for c in later] == [expected, expected] + calls = _events(client, "$mcp_tool_call") + assert [c["properties"]["$session_id"] for c in calls] == [expected] * 3 @pytest.mark.skipif(MCP_MAJOR != 1, reason="v1 FastMCP server") diff --git a/posthog/test/mcp/test_v2_wire_dual_era.py b/posthog/test/mcp/test_v2_wire_dual_era.py index 4ccab4e95..a29334ae3 100644 --- a/posthog/test/mcp/test_v2_wire_dual_era.py +++ b/posthog/test/mcp/test_v2_wire_dual_era.py @@ -194,13 +194,12 @@ async def test_modern_conversation_anchors_session_across_instances(): calls = _events(client, "$mcp_tool_call") assert len(calls) == 2 expected = derive_session_id_from_conversation(minted) - # Pod A minted the handle but does not anchor on it — it is unproven until - # the agent echoes it back. Pod B receives the echo and anchors, deriving - # the same session any pod would *without ever having met pod A*: that - # agreement is the cross-pod contract. - assert calls[0]["properties"]["$session_id"] != expected - assert calls[1]["properties"]["$session_id"] == expected - assert calls[1]["properties"]["$mcp_conversation_id"] == minted + # Pod A anchors on the handle it minted (delivery was confirmed before the + # session was resolved), and pod B derives the same session from the echo + # *without ever having met pod A*. That agreement, across two processes + # sharing no state, is the cross-pod contract. + assert [c["properties"]["$session_id"] for c in calls] == [expected, expected] + assert [c["properties"]["$mcp_conversation_id"] for c in calls] == [minted, minted] async def test_modern_result_shape_survives_instrumentation(): From 80c9c35ea8a175b47324b9bd7d596fad6075fd63 Mon Sep 17 00:00:00 2001 From: Georgis Andonis Date: Fri, 21 Aug 2026 14:03:14 +0300 Subject: [PATCH 16/16] fix(mcp): settle the shared session before the tool body runs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review finding, and a regression from the previous commit: moving prepare_request() to after the call meant data.session_id still held the *previous* request's value while the tool executed. McpAnalytics.capture() reads that field, so a custom event emitted from inside a tool body was attributed to the previous caller's session — and through the identity cache, their person. prime_session() now settles the transport/memory session before dispatch; the conversation anchor is still finalised afterwards, once delivery is known. An in-tool event predates that decision, so it belongs to the transport session, which is the correct answer available at that moment. Regression test uses two callers with distinct Mcp-Session-Id headers and asserts B's tool body sees B's session — verified to fail without the fix (the first version of this test passed either way and was worthless). Generated-By: PostHog Code Task-Id: ebafcb71-b03b-443d-b40c-d527ed4a04f4 --- posthog/mcp/_instrument_fastmcp.py | 5 ++ posthog/mcp/_instrument_lowlevel.py | 5 ++ posthog/mcp/_instrument_v2.py | 9 +++ posthog/mcp/_instrumentation.py | 18 ++++++ posthog/test/mcp/test_conversation_session.py | 57 +++++++++++++++++++ 5 files changed, 94 insertions(+) diff --git a/posthog/mcp/_instrument_fastmcp.py b/posthog/mcp/_instrument_fastmcp.py index b44a58178..b0ac4db59 100644 --- a/posthog/mcp/_instrument_fastmcp.py +++ b/posthog/mcp/_instrument_fastmcp.py @@ -42,6 +42,7 @@ build_tool_call_request, extract_tools, prepare_request, + prime_session, read_tool_category, record_missing_capability, record_tool_call, @@ -171,6 +172,10 @@ async def _session(anchor: Optional[str]) -> str: k: v for k, v in arguments.items() if k not in strip_keys } + # Settle the shared session before the tool body runs, so an in-tool + # `analytics.capture()` is attributed to this caller and not the last one. + await prime_session(data, mcp_session_id=mcp_session_id, token=token) + start = time.monotonic() try: result = await original( diff --git a/posthog/mcp/_instrument_lowlevel.py b/posthog/mcp/_instrument_lowlevel.py index b8c22c924..bfeb390b9 100644 --- a/posthog/mcp/_instrument_lowlevel.py +++ b/posthog/mcp/_instrument_lowlevel.py @@ -37,6 +37,7 @@ build_tool_call_request, extract_tools, prepare_request, + prime_session, read_tool_category, record_missing_capability, record_tool_call, @@ -178,6 +179,10 @@ async def _session(anchor: Optional[str]) -> str: if key not in owned: req.params.arguments.pop(key, None) + # Settle the shared session before the tool body runs, so an in-tool + # `analytics.capture()` is attributed to this caller and not the last one. + await prime_session(data, mcp_session_id=mcp_session_id, token=token) + start = time.monotonic() try: result = await original(req) diff --git a/posthog/mcp/_instrument_v2.py b/posthog/mcp/_instrument_v2.py index 5df08c75a..fc3cac941 100644 --- a/posthog/mcp/_instrument_v2.py +++ b/posthog/mcp/_instrument_v2.py @@ -51,6 +51,7 @@ build_tool_call_request, params_to_request_dict, prepare_request, + prime_session, read_tool_category, record_missing_capability, record_tool_call, @@ -310,6 +311,10 @@ async def _session(anchor: Optional[str]) -> str: k: v for k, v in arguments.items() if k not in strip_keys } + # Settle the shared session before the tool body runs, so an in-tool + # `analytics.capture()` is attributed to this caller and not the last one. + await prime_session(data, mcp_session_id=mcp_session_id, token=token) + start = time.monotonic() try: result = await original( @@ -474,6 +479,10 @@ async def _session(anchor: Optional[str]) -> str: ] ) + # Settle the shared session before the tool body runs, so an in-tool + # `analytics.capture()` is attributed to this caller and not the last one. + await prime_session(data, mcp_session_id=mcp_session_id, token=token) + start = time.monotonic() try: result = await original(ctx, params) diff --git a/posthog/mcp/_instrumentation.py b/posthog/mcp/_instrumentation.py index 4c38bdd4b..6f6a363ff 100644 --- a/posthog/mcp/_instrumentation.py +++ b/posthog/mcp/_instrumentation.py @@ -248,6 +248,24 @@ def resolve_session_and_client( return token, client_name, client_version, protocol_version +async def prime_session( + data: MCPAnalyticsData, + *, + mcp_session_id: Optional[str], + token: Optional[SessionTokenPayload] = None, +) -> None: + """Point the shared per-server session at *this* request before the tool body runs. + + ``McpAnalytics.capture()`` reads ``data.session_id`` for custom in-tool + events. The conversation anchor can only be resolved after the call (we + don't know until then whether the agent received the handle), so without + this the tool body would read whatever the *previous* request left behind + and attribute a custom event to the wrong caller. Emits nothing — it only + settles the transport/memory session an in-tool event should belong to. + """ + await resolve_session_id(data, mcp_session_id, token=token) + + async def prepare_request( data: MCPAnalyticsData, *, diff --git a/posthog/test/mcp/test_conversation_session.py b/posthog/test/mcp/test_conversation_session.py index 347232496..037888a6f 100644 --- a/posthog/test/mcp/test_conversation_session.py +++ b/posthog/test/mcp/test_conversation_session.py @@ -295,3 +295,60 @@ def boom() -> str: ) # and nothing claims a conversation the agent never received assert all("$mcp_conversation_id" not in e["properties"] for e in client.events) + + +@pytest.mark.skipif(MCP_MAJOR != 1, reason="v1 FastMCP server") +async def test_in_tool_events_are_not_attributed_to_the_previous_caller(): + """A custom event captured *inside* a tool body reads the shared + ``data.session_id``. The conversation anchor can only be resolved after the + call, so unless that field is settled first, caller B's in-tool event is + attributed to caller A's session — and, through the identity cache, to + caller A's person.""" + from types import SimpleNamespace + + from mcp.server.fastmcp import FastMCP + + from posthog.mcp import instrument + from posthog.mcp._internal import get_server_tracking_data + from posthog.mcp.session import derive_session_id_from_mcp_session + + def caller(session_header): + return SimpleNamespace( + request_context=SimpleNamespace( + request=SimpleNamespace(headers={"mcp-session-id": session_header}), + session=SimpleNamespace(client_params=None), + ) + ) + + server = FastMCP("in-tool") + seen = {} + + @server.tool() + def echo(msg: str) -> str: + return msg + + client = FakeClient() + handle = instrument(server, client, MCPAnalyticsOptions()) + + @server.tool() + def emits(msg: str) -> str: + data = get_server_tracking_data(handle._key) + seen["session_during_body"] = data.session_id if data else None + return msg + + # Caller A runs first and leaves its session behind on the shared state. + await server._tool_manager.call_tool( + "echo", {"msg": "a", "context": "caller A"}, context=caller("session-A") + ) + # Caller B's tool body must see *its own* session, not A's. + await server._tool_manager.call_tool( + "emits", {"msg": "b", "context": "caller B"}, context=caller("session-B") + ) + await _flush() + + assert seen["session_during_body"] == derive_session_id_from_mcp_session( + "session-B" + ) + assert seen["session_during_body"] != derive_session_id_from_mcp_session( + "session-A" + )