From 45b28aced1361492e9ca9aa1273ca933ffe65e2b Mon Sep 17 00:00:00 2001 From: Lucas Faria Date: Tue, 4 Aug 2026 19:49:15 -0300 Subject: [PATCH 1/4] mcp: add version-matrix test scaffolding for mcp 1.x/2.x The 2026-07-28 spec ships as `mcp` 2.x, a breaking rewrite of the same PyPI package that can't coexist with 1.x in one venv. Make the posthog.mcp suite valid on both SDKs so v2 support can be built and validated incrementally. What: - `posthog/mcp/_mcp_version.py`: `installed_mcp_generation() -> 1|2|None` probe (importlib.metadata; never raises), used by both runtime and tests. - `posthog/test/mcp/_helpers.py`: `requires_mcp_v1` / `requires_mcp_v2` skipif markers keyed off the probe. - Marked every v1-internals test (`request_handlers` shape, fastmcp import, stateless-token flows) with the v1 marker, preferring module-level `pytestmark` and guarding crash-prone module imports so v2 collection is clean. jlowin `fastmcp`'s server layer raises a rewritten ImportError under mcp 2.x, so its module guards the `from fastmcp import FastMCP` by hand (importorskip mis-handles that rewrite). - `posthog/test/mcp/test_mcp_version.py`: generation-probe tests that run in both envs and anchor the marker mutual-exclusivity invariant. - Pinned the `test` extra to `mcp>=1.28.1,<2` (lock refreshed to match) so a no-upper-bound resolve can't silently flip CI to 2.0. - CI `tests-mcp-v2` job (3.12): sync test extra, install `mcp>=2,<3` over it, run the mcp subset. `scripts/validate-mcp-matrix.sh` does the same locally across two throwaway venvs and prints a PASS/FAIL matrix. How tested: - v1 env (mcp 1.29): `pytest posthog/test/mcp` -> 151 passed, 1 skipped. - v2 env (mcp 2.0): `pytest posthog/test/mcp` -> 111 passed, 11 skipped, no collection errors; v2 probe tests pass, v1-only tests skip cleanly. - `ruff@0.11.12 check .` / `format --check .` clean; `mypy ... | mypy-baseline filter` -> no issues (190 files). Generated-By: PostHog Code Task-Id: c17ecd78-bb05-4b3c-b621-458abba9c6aa --- .github/workflows/ci.yml | 38 +++++++++++++++++++ posthog/mcp/_mcp_version.py | 34 +++++++++++++++++ posthog/test/mcp/_helpers.py | 21 +++++++++++ posthog/test/mcp/test_fastmcp.py | 17 ++++++--- posthog/test/mcp/test_fastmcp_v2.py | 18 ++++++++- posthog/test/mcp/test_features_m4.py | 20 +++++++--- posthog/test/mcp/test_lowlevel.py | 8 ++++ posthog/test/mcp/test_mcp_version.py | 52 ++++++++++++++++++++++++++ posthog/test/mcp/test_review_fixes.py | 18 ++++++--- posthog/test/mcp/test_session_token.py | 3 ++ pyproject.toml | 7 +++- scripts/validate-mcp-matrix.sh | 48 ++++++++++++++++++++++++ uv.lock | 2 +- 13 files changed, 265 insertions(+), 21 deletions(-) create mode 100644 posthog/mcp/_mcp_version.py create mode 100644 posthog/test/mcp/test_mcp_version.py create mode 100755 scripts/validate-mcp-matrix.sh diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index b00a7b3c..c3232c8c 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -151,6 +151,44 @@ jobs: run: | pytest --verbose --timeout=30 + tests-mcp-v2: + name: MCP SDK 2.x compatibility + runs-on: ubuntu-latest + + steps: + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + with: + fetch-depth: 1 + + - name: Set up Python 3.12 + uses: actions/setup-python@5fda3b95a4ea91299a34e894583c3862153e4b97 # v7.0.0 + with: + python-version: '3.12' + + - name: Install uv + uses: astral-sh/setup-uv@c771a70e6277c0a99b617c7a806ffedaca235ff9 # v9.0.0 + with: + version: "0.11.32" + enable-cache: true + + - name: Install test dependencies + shell: bash + run: | + UV_PROJECT_ENVIRONMENT=$pythonLocation uv sync --extra test + + # The `test` extra pins mcp<2; install the 2.x SDK over it so the same + # posthog/test/mcp suite runs against the 2026-07-28 SDK. jlowin's + # `fastmcp` (needs mcp 1.x) is uninstalled by this resolution, so its + # tests skip via `importorskip("fastmcp")`. + - name: Upgrade to MCP SDK 2.x + shell: bash + run: | + UV_PROJECT_ENVIRONMENT=$pythonLocation uv pip install "mcp>=2,<3" + + - name: Run MCP tests against SDK 2.x + run: | + pytest posthog/test/mcp --verbose --timeout=30 + mutation-tests: name: Targeted mutation tests runs-on: ubuntu-latest diff --git a/posthog/mcp/_mcp_version.py b/posthog/mcp/_mcp_version.py new file mode 100644 index 00000000..37a08c5b --- /dev/null +++ b/posthog/mcp/_mcp_version.py @@ -0,0 +1,34 @@ +# 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 + +"""Which generation of the official ``mcp`` SDK is installed. + +The 2026-07-28 spec ships as ``mcp`` 2.x, a breaking rewrite of the same PyPI +package (``mcp.server.fastmcp`` gone, low-level handlers now keyed by method +string, an official middleware protocol). The two generations can't coexist in +one venv, so ``instrument()`` and the test suite both branch on this probe rather +than on a try/except over a module import that means different things per version. +""" + +from __future__ import annotations + +from typing import Literal, Optional + + +def installed_mcp_generation() -> Optional[Literal[1, 2]]: + """The major version of the installed ``mcp`` SDK as a generation number: + ``1`` for ``mcp>=1,<2`` (the handshake-era SDK), ``2`` for ``mcp>=2,<3`` (the + 2026-07-28 SDK), or ``None`` when ``mcp`` isn't installed or its version is + unreadable/outside the supported range. Never raises.""" + try: + from importlib.metadata import version + + major = int(version("mcp").split(".")[0]) + except Exception: # noqa: BLE001 - a probe must never break its callers + return None + if major == 1: + return 1 + if major == 2: + return 2 + return None diff --git a/posthog/test/mcp/_helpers.py b/posthog/test/mcp/_helpers.py index 21afc62a..1e69f725 100644 --- a/posthog/test/mcp/_helpers.py +++ b/posthog/test/mcp/_helpers.py @@ -7,6 +7,27 @@ import asyncio import concurrent.futures +import pytest + +from posthog.mcp._mcp_version import installed_mcp_generation + +# The 2026-07-28 SDK (`mcp` 2.x) is a breaking rewrite of the same PyPI package, +# so the two generations can't share a venv. The suite runs unchanged in both: +# tests that reach into a generation's private server internals (fastmcp import, +# `request_handlers` shape, session-token flows) carry the matching marker and +# skip cleanly in the other env. Prefer marking a whole module via +# `pytestmark = requires_mcp_v1` when every test in it is generation-specific. +_GENERATION = installed_mcp_generation() + +requires_mcp_v1 = pytest.mark.skipif( + _GENERATION != 1, + reason=f"requires mcp 1.x internals; installed generation is {_GENERATION}", +) +requires_mcp_v2 = pytest.mark.skipif( + _GENERATION != 2, + reason=f"requires mcp 2.x internals; installed generation is {_GENERATION}", +) + class FakeClient: """Records capture() calls instead of sending them.""" diff --git a/posthog/test/mcp/test_fastmcp.py b/posthog/test/mcp/test_fastmcp.py index 1784835f..e7eb7a27 100644 --- a/posthog/test/mcp/test_fastmcp.py +++ b/posthog/test/mcp/test_fastmcp.py @@ -4,17 +4,24 @@ import pytest -import mcp.types as mcp_types -from mcp.server.fastmcp import FastMCP +# `mcp.server.fastmcp` was removed in mcp 2.x, so importing it crashes collection +# there; skip the whole module on any non-v1 generation before that import runs. +pytest.importorskip("mcp.server.fastmcp") -from posthog.mcp import instrument -from posthog.mcp.types import MCPAnalyticsOptions, UserIdentity -from posthog.test.mcp._helpers import ( +import mcp.types as mcp_types # noqa: E402 +from mcp.server.fastmcp import FastMCP # noqa: E402 + +from posthog.mcp import instrument # noqa: E402 +from posthog.mcp.types import MCPAnalyticsOptions, UserIdentity # noqa: E402 +from posthog.test.mcp._helpers import ( # noqa: E402 FakeClient, events_named as _events, flush_background as _flush, + requires_mcp_v1, ) +pytestmark = requires_mcp_v1 + def make_server(): server = FastMCP("test-server") diff --git a/posthog/test/mcp/test_fastmcp_v2.py b/posthog/test/mcp/test_fastmcp_v2.py index e40213d3..8daf39de 100644 --- a/posthog/test/mcp/test_fastmcp_v2.py +++ b/posthog/test/mcp/test_fastmcp_v2.py @@ -3,10 +3,24 @@ import pytest -pytest.importorskip("fastmcp") +# jlowin's `fastmcp` top-level package still imports under mcp 2.x, but its server +# layer (where `FastMCP` lives) needs mcp 1.x internals and raises a *rewritten* +# ImportError there -- which `importorskip` treats as a real error, not a skip. So +# guard the import by hand and skip the whole module cleanly when the server layer +# is unavailable. These tests also drive the v1 low-level handler shape +# (`_mcp_server.request_handlers` keyed by request type), so they're v1-only. +try: + from fastmcp import FastMCP +except ImportError: + pytest.skip( + "fastmcp server support unavailable (needs mcp 1.x)", allow_module_level=True + ) + +from posthog.test.mcp._helpers import requires_mcp_v1 # noqa: E402 + +pytestmark = requires_mcp_v1 import mcp.types as mcp_types # noqa: E402 -from fastmcp import FastMCP # noqa: E402 from posthog.mcp import instrument # noqa: E402 from posthog.mcp.types import MCPAnalyticsOptions # noqa: E402 diff --git a/posthog/test/mcp/test_features_m4.py b/posthog/test/mcp/test_features_m4.py index 0baa56c5..c504874f 100644 --- a/posthog/test/mcp/test_features_m4.py +++ b/posthog/test/mcp/test_features_m4.py @@ -1,17 +1,25 @@ """Tests for M4 parity features: get_more_tools (missing capability) + conversation_id.""" -import mcp.types as mcp_types -from mcp.server.fastmcp import FastMCP -from mcp.server.lowlevel import Server +import pytest -from posthog.mcp import PostHogMCP, get_more_tools_result, instrument -from posthog.mcp.types import MCPAnalyticsOptions -from posthog.test.mcp._helpers import ( +# `mcp.server.fastmcp` was removed in mcp 2.x; skip the module there before import. +pytest.importorskip("mcp.server.fastmcp") + +import mcp.types as mcp_types # noqa: E402 +from mcp.server.fastmcp import FastMCP # noqa: E402 +from mcp.server.lowlevel import Server # noqa: E402 + +from posthog.mcp import PostHogMCP, get_more_tools_result, instrument # noqa: E402 +from posthog.mcp.types import MCPAnalyticsOptions # noqa: E402 +from posthog.test.mcp._helpers import ( # noqa: E402 FakeClient, events_named as _events, flush_background as _flush, + requires_mcp_v1, ) +pytestmark = requires_mcp_v1 + def make_fastmcp(): server = FastMCP("m4-fastmcp") diff --git a/posthog/test/mcp/test_lowlevel.py b/posthog/test/mcp/test_lowlevel.py index 79bc063f..b26f884b 100644 --- a/posthog/test/mcp/test_lowlevel.py +++ b/posthog/test/mcp/test_lowlevel.py @@ -8,8 +8,16 @@ FakeClient, events_named as _events, flush_background as _flush, + requires_mcp_v1, ) +# The v1 low-level server keeps handlers in a public `request_handlers` dict keyed +# by request TYPE; mcp 2.x renamed it `_request_handlers` keyed by method string. +# These tests drive the v1 shape directly, so they're v1-only. `mcp.server.lowlevel` +# and `mcp.types` still import in v2, so a module marker (not `importorskip`) is the +# right guard. +pytestmark = requires_mcp_v1 + def make_server(): server = Server("test-lowlevel") diff --git a/posthog/test/mcp/test_mcp_version.py b/posthog/test/mcp/test_mcp_version.py new file mode 100644 index 00000000..b04b0a03 --- /dev/null +++ b/posthog/test/mcp/test_mcp_version.py @@ -0,0 +1,52 @@ +"""Generation-probe tests. These run in BOTH the mcp-v1 and mcp-v2 envs and are +the anchor that the same suite is valid on either SDK: the probe must agree with +the installed `mcp`, and the two generation markers must be mutually exclusive.""" + +from importlib.metadata import version + +import pytest + +from posthog.mcp._mcp_version import installed_mcp_generation +from posthog.test.mcp._helpers import requires_mcp_v1, requires_mcp_v2 + + +def test_probe_matches_installed_mcp(): + try: + major = int(version("mcp").split(".")[0]) + except Exception: # noqa: BLE001 - mcp genuinely absent + major = None + + generation = installed_mcp_generation() + if major in (1, 2): + assert generation == major + else: + assert generation is None + + +def test_probe_never_raises_without_mcp(monkeypatch): + def _missing(_name): + from importlib.metadata import PackageNotFoundError + + raise PackageNotFoundError("mcp") + + monkeypatch.setattr("importlib.metadata.version", _missing) + assert installed_mcp_generation() is None + + +@requires_mcp_v1 +def test_v1_marker_runs_only_on_v1(): + assert installed_mcp_generation() == 1 + + +@requires_mcp_v2 +def test_v2_marker_runs_only_on_v2(): + assert installed_mcp_generation() == 2 + + +def test_generation_markers_are_mutually_exclusive(): + # Exactly one of the two skip conditions is False for a supported install, so + # a generation-specific test is never collected-and-run in the wrong env. + generation = installed_mcp_generation() + if generation is None: + pytest.skip("no supported mcp installed") + assert (generation == 1) != (generation == 2) diff --git a/posthog/test/mcp/test_review_fixes.py b/posthog/test/mcp/test_review_fixes.py index f5f97572..65ec445f 100644 --- a/posthog/test/mcp/test_review_fixes.py +++ b/posthog/test/mcp/test_review_fixes.py @@ -15,18 +15,24 @@ import pytest -import mcp.types as mcp_types -from mcp.server.fastmcp import FastMCP -from mcp.server.lowlevel import Server +# `mcp.server.fastmcp` was removed in mcp 2.x; skip the module there before import. +pytest.importorskip("mcp.server.fastmcp") -from posthog.mcp import PostHogMCP, instrument -from posthog.mcp.types import MCPAnalyticsOptions, UserIdentity -from posthog.test.mcp._helpers import ( +import mcp.types as mcp_types # noqa: E402 +from mcp.server.fastmcp import FastMCP # noqa: E402 +from mcp.server.lowlevel import Server # noqa: E402 + +from posthog.mcp import PostHogMCP, instrument # noqa: E402 +from posthog.mcp.types import MCPAnalyticsOptions, UserIdentity # noqa: E402 +from posthog.test.mcp._helpers import ( # noqa: E402 FakeClient, events_named as _events, flush_background as _flush, + requires_mcp_v1, ) +pytestmark = requires_mcp_v1 + def _call_request(name, arguments): return mcp_types.CallToolRequest( diff --git a/posthog/test/mcp/test_session_token.py b/posthog/test/mcp/test_session_token.py index cf536e80..b58c8a73 100644 --- a/posthog/test/mcp/test_session_token.py +++ b/posthog/test/mcp/test_session_token.py @@ -21,6 +21,7 @@ read_mcp_session_header, ) from posthog.mcp.types import MCPAnalyticsOptions +from posthog.test.mcp._helpers import requires_mcp_v1 # --- codec ------------------------------------------------------------------- @@ -414,6 +415,7 @@ async def send(message): # --- end-to-end against a real stateless FastMCP transport ------------------- +@requires_mcp_v1 def test_middleware_end_to_end_with_stateless_fastmcp(): """Cross-version sanity: mount a real stateless FastMCP streamable-HTTP app, add the middleware, and confirm (1) the `initialize` response carries a @@ -483,6 +485,7 @@ def rpc(method, params=None, id=1, extra=None): assert resp2.status_code == 200, resp2.text +@requires_mcp_v1 def test_instrument_autowires_stateless_mint_no_manual_middleware(): """instrument() alone (no app.add_middleware) makes the FastMCP streamable-HTTP app mint the session token -- the zero-config path. mcp.run() uses the same diff --git a/pyproject.toml b/pyproject.toml index 230f32d9..306eea3e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -88,7 +88,12 @@ test = [ "pydantic>=2.12.0", "parameterized>=0.8.1", "claude-agent-sdk", - "mcp>=1.28.1", + # Pinned below 2.0: `mcp` 2.x is a breaking rewrite (2026-07-28 spec) that + # can't coexist with 1.x in one venv, and a lock refresh with no upper bound + # would silently flip this env to 2.0. The v2 SDK is exercised by the separate + # `tests-mcp-v2` CI job (and `scripts/validate-mcp-matrix.sh`), which installs + # `mcp>=2,<3` over this env. + "mcp>=1.28.1,<2", "fastmcp>=2.0", "opentelemetry-sdk>=1.20.0", "opentelemetry-exporter-otlp-proto-http>=1.20.0", diff --git a/scripts/validate-mcp-matrix.sh b/scripts/validate-mcp-matrix.sh new file mode 100755 index 00000000..d7425643 --- /dev/null +++ b/scripts/validate-mcp-matrix.sh @@ -0,0 +1,48 @@ +#!/usr/bin/env bash +# Validate the posthog.mcp suite against both generations of the `mcp` SDK. +# +# The 2026-07-28 spec ships as `mcp` 2.x, a breaking rewrite of the same PyPI +# package that can't coexist with 1.x in one venv. This builds two throwaway uv +# venvs (v1 and v2), runs the mcp test subset in each, and prints a PASS/FAIL +# matrix. Non-interactive; exits non-zero if either env fails. +# +# Usage: scripts/validate-mcp-matrix.sh +set -uo pipefail + +REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +cd "$REPO_ROOT" + +WORKDIR="$(mktemp -d)" +trap 'rm -rf "$WORKDIR"' EXIT + +V1_ENV="$WORKDIR/v1" +V2_ENV="$WORKDIR/v2" +PYTHON_VERSION="3.12" + +run_env() { + # $1 = label, $2 = venv path, $3 = mcp spec to install over the base env + local label="$1" env_path="$2" mcp_spec="$3" + echo "=== [$label] building venv ($mcp_spec) ===" + uv venv "$env_path" --python "$PYTHON_VERSION" >/dev/null || return 1 + UV_PROJECT_ENVIRONMENT="$env_path" uv sync --extra test >/dev/null || return 1 + UV_PROJECT_ENVIRONMENT="$env_path" uv pip install --python "$env_path/bin/python" "$mcp_spec" >/dev/null || return 1 + echo "--- [$label] installed mcp: $("$env_path/bin/python" -c 'from importlib.metadata import version; print(version("mcp"))')" + echo "=== [$label] running posthog/test/mcp ===" + "$env_path/bin/python" -m pytest posthog/test/mcp --timeout=30 -q +} + +run_env "mcp-v1" "$V1_ENV" "mcp>=1.28.1,<2" +V1_STATUS=$? + +run_env "mcp-v2" "$V2_ENV" "mcp>=2,<3" +V2_STATUS=$? + +result() { [ "$1" -eq 0 ] && echo "PASS" || echo "FAIL"; } + +echo "" +echo "================ MCP version matrix ================" +printf " %-10s %s\n" "mcp 1.x" "$(result "$V1_STATUS")" +printf " %-10s %s\n" "mcp 2.x" "$(result "$V2_STATUS")" +echo "====================================================" + +[ "$V1_STATUS" -eq 0 ] && [ "$V2_STATUS" -eq 0 ] diff --git a/uv.lock b/uv.lock index ed760aa1..c832f1c3 100644 --- a/uv.lock +++ b/uv.lock @@ -2597,7 +2597,7 @@ requires-dist = [ { name = "langgraph", marker = "extra == 'test'", specifier = ">=1.0" }, { name = "langgraph-checkpoint", marker = "extra == 'test'", specifier = ">=4.1.1" }, { name = "lxml", marker = "extra == 'dev'" }, - { name = "mcp", marker = "extra == 'test'", specifier = ">=1.28.1" }, + { name = "mcp", marker = "extra == 'test'", specifier = ">=1.28.1,<2" }, { name = "mypy", marker = "extra == 'dev'" }, { name = "mypy-baseline", marker = "extra == 'dev'" }, { name = "openai-agents", marker = "extra == 'test'", specifier = ">=0.18" }, From e275ecbc297f40653d45f423f4aa1a6e6ca8f218 Mon Sep 17 00:00:00 2001 From: Lucas Faria Date: Tue, 4 Aug 2026 19:58:19 -0300 Subject: [PATCH 2/4] mcp: add mcp 2.x (2026-07-28) middleware adapter MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The v1 adapters import `mcp.server.fastmcp` at module level, so on mcp 2.x `instrument()` raised a bare ModuleNotFoundError before it could dispatch. Make compat detection lazy/per-generation and add a capture-only adapter that hooks the official 2.x `ServerMiddleware` seam. What: - `_compatibility.py`: dropped module-level v1 imports; every predicate imports lazily and returns False on ImportError. `is_fastmcp` no longer crashes on 2.x; added `is_mcpserver_v2` (mcp.server.mcpserver.MCPServer) and made `is_low_level_server` generation-agnostic. - `_instrument_v2.py`: attaches one `(ctx, call_next)` middleware to `server.middleware` (the same public list `MCPServer` and low-level `Server` expose — no private-attr patching). Captures `tools/call` (reusing `record_tool_call`), `tools/list` (reusing `record_tools_list`, read-only — no response mutation), and `server/discover` (reusing the lazy `_maybe_emit_initialize`). Client name/version + protocol come from the per-request `_meta` envelope (`io.modelcontextprotocol/clientInfo`, `.../protocolVersion`) on 2026-07-28 sessions, or from `initialize` params on a legacy-negotiated session. Identify flows via `prepare_request`. - MRTR: an `input_required` result is NOT an error; it stamps the new `$mcp_result_type` property (added to `constants.py`, threaded through `_capture.py`/`_posthog_events.py`/`record_tool_call`). Full round-trip stitching is out of scope. The middleware sees results as wire dicts (`{"isError": ...}` / `{"resultType": "input_required", ...}`), so error detection handles both dict and model shapes. - `__init__.py`: `instrument()` dispatches on `installed_mcp_generation()`; `_canonical_server` now unwraps `_lowlevel_server` too (v2's wrapper attr). On an attach failure the no-op fallback logs an actionable message naming the detected generation and both supported ranges. `_warn_if_unsupported_mcp_ version` updated to advertise mcp>=1.26,<2 and mcp>=2,<3. - No context injection, no get_more_tools, no stateless minting on v2 (SEP-2567 removed the Mcp-Session-Id header). v1 paths are behavior-preserved. How tested: - New `test_instrument_v2.py` (7 tests) drives a real `MCPServer` over the SDK in-memory transport through modern (server/discover) and legacy (initialize) handshakes: tool call w/ envelope identity, isError -> $exception, tools/list names, initialize-once, identify attribution, MRTR result_type-not-error. - v2 env: `pytest posthog/test/mcp` -> 118 passed, 11 skipped. - v1 env: 151 passed, 2 skipped (byte-for-byte v1 behavior; only the version- warning assertion string updated). Full-suite `--collect-only`: 1981 tests, no import errors. - `ruff@0.11.12 check .`/`format --check .` clean; `mypy | mypy-baseline filter` -> no issues (192 files). Generated-By: PostHog Code Task-Id: c17ecd78-bb05-4b3c-b621-458abba9c6aa --- posthog/mcp/__init__.py | 108 ++++++--- posthog/mcp/_capture.py | 1 + posthog/mcp/_compatibility.py | 41 +++- posthog/mcp/_instrument_v2.py | 317 +++++++++++++++++++++++++ posthog/mcp/_instrumentation.py | 2 + posthog/mcp/_posthog_events.py | 2 + posthog/mcp/constants.py | 5 + posthog/test/mcp/test_instrument_v2.py | 180 ++++++++++++++ posthog/test/mcp/test_review_fixes.py | 2 +- 9 files changed, 622 insertions(+), 36 deletions(-) create mode 100644 posthog/mcp/_instrument_v2.py create mode 100644 posthog/test/mcp/test_instrument_v2.py diff --git a/posthog/mcp/__init__.py b/posthog/mcp/__init__.py index dbab60cd..e2cc15c5 100644 --- a/posthog/mcp/__init__.py +++ b/posthog/mcp/__init__.py @@ -157,10 +157,11 @@ 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).""" + """PostHog MCP analytics supports two generations of the ``mcp`` SDK: 1.x + (``mcp>=1.26,<2``, the ``request_handlers`` seam) and 2.x (``mcp>=2,<3``, the + 2026-07-28 ``ServerMiddleware`` seam). Since ``mcp`` is an unpinned peer + dependency, advise at runtime when the installed version is outside both + supported ranges rather than failing hard (a near-neighbor may still work).""" try: from importlib.metadata import version @@ -168,22 +169,75 @@ 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}. " - "Instrumentation hooks private SDK internals and may behave unexpectedly." + f"Warning: PostHog MCP analytics supports mcp>=1.26,<2 and mcp>=2,<3; found {installed}. " + "Instrumentation may behave unexpectedly on this version." ) 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, else the server + itself. v1 FastMCP (official and jlowin's) exposes ``_mcp_server``; the mcp 2.x + ``MCPServer`` exposes ``_lowlevel_server``. 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 +def _instrument_generation_1( + server: Any, + data: MCPAnalyticsData, + is_fastmcp: Any, + is_fastmcp_v2: Any, + is_low_level_server: Any, +) -> None: + """Dispatch for the mcp 1.x SDK: the ``request_handlers`` monkey-patch seam + plus zero-config stateless minting (an ASGI wrap that is a no-op for stdio / + low-level servers).""" + from ._instrument_fastmcp import instrument_fastmcp + from ._instrument_lowlevel import instrument_fastmcp_v2, instrument_low_level + + if is_fastmcp(server): + instrument_fastmcp(server, data) + elif is_fastmcp_v2(server): + instrument_fastmcp_v2(server, data) + elif is_low_level_server(server): + 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." + ) + + # Zero-config stateless minting: wrap the server's ASGI-app factories so a + # stateless/multi-pod deployment keeps one $session_id + the client harness + # across pods with no extra setup. No-op for stdio / low-level servers. + autowire_stateless_mint(server) + + +def _instrument_generation_2(server: Any, data: MCPAnalyticsData) -> None: + """Dispatch for the mcp 2.x SDK (2026-07-28): attach the analytics + ``ServerMiddleware`` to the ``MCPServer`` or low-level ``Server``. A capture-only + adapter — no context injection, no tool-list mutation, no stateless minting + (SEP-2567 removed the ``Mcp-Session-Id`` header, so there's nothing to mint).""" + from ._compatibility import is_low_level_server, is_mcpserver_v2 + from ._instrument_v2 import instrument_low_level_v2, instrument_mcpserver_v2 + + if is_mcpserver_v2(server): + instrument_mcpserver_v2(server, data) + elif is_low_level_server(server): + instrument_low_level_v2(server, data) + else: + raise TypeError( + f"Unsupported server type for mcp 2.x: {type(server)!r}. Pass an " + "mcp.server.mcpserver.MCPServer or a low-level mcp.server.lowlevel.Server." + ) + + def instrument( server: Any, posthog_client: Optional[Client] = None, @@ -222,10 +276,10 @@ def instrument( "(PostHogMCP for custom dispatchers works without it.)" ) _warn_if_unsupported_mcp_version() + from ._mcp_version import installed_mcp_generation 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 + generation = installed_mcp_generation() key = _canonical_server(server) try: @@ -241,24 +295,20 @@ def instrument( data = MCPAnalyticsData(options=opts, sink=sink, session_id=new_session_id()) set_server_tracking_data(key, data) - if is_fastmcp(server): - instrument_fastmcp(server, data) - elif is_fastmcp_v2(server): - instrument_fastmcp_v2(server, data) - elif is_low_level_server(server): - instrument_low_level(server, data) + if generation == 2: + _instrument_generation_2(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." + _instrument_generation_1( + server, data, is_fastmcp, is_fastmcp_v2, is_low_level_server ) - # Zero-config stateless minting: wrap the server's ASGI-app factories so a - # stateless/multi-pod deployment keeps one $session_id + the client harness - # across pods with no extra setup. No-op for stdio / low-level servers. - autowire_stateless_mint(server) - return McpAnalytics(key) except Exception as error: # noqa: BLE001 - log(f"Warning: failed to instrument server - {error}") + # Degrade to a no-op so the host app keeps working, but make the failure + # actionable: name what happened, the detected generation, and the + # versions we support — never a bare ModuleNotFoundError or silent no-op. + log( + f"Warning: failed to instrument server (mcp generation {generation}, " + f"supported: 1.x as mcp>=1.26,<2 and 2.x as mcp>=2,<3) - {error}" + ) return _NoopAnalytics() diff --git a/posthog/mcp/_capture.py b/posthog/mcp/_capture.py index 26cf67fd..bda3c19a 100644 --- a/posthog/mcp/_capture.py +++ b/posthog/mcp/_capture.py @@ -64,6 +64,7 @@ def capture_event( "is_error": event_input.get("is_error"), "error": event_input.get("error"), "conversation_id": event_input.get("conversation_id"), + "result_type": event_input.get("result_type"), "properties": event_input.get("properties"), } diff --git a/posthog/mcp/_compatibility.py b/posthog/mcp/_compatibility.py index 37b017df..0fb701a9 100644 --- a/posthog/mcp/_compatibility.py +++ b/posthog/mcp/_compatibility.py @@ -2,24 +2,35 @@ # 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 check imports lazily and swallows import failures: the ``mcp`` SDK ships +two mutually incompatible generations under one package name (1.x removed +``mcp.server.fastmcp`` in 2.x, which added ``mcp.server.mcpserver.MCPServer``), +so a module-level import of either would crash ``instrument()`` on the other +generation. Each predicate returns ``False`` cleanly when its target class isn't +importable, letting ``instrument()`` dispatch on whichever generation is present. +""" 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 v1 SDK's high-level server (``mcp.server.fastmcp.FastMCP``). Returns + False on mcp 2.x, where that module was removed.""" + try: + from mcp.server.fastmcp import FastMCP + except ImportError: + return False return isinstance(server, FastMCP) 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.""" + from the official SDK. Its server layer needs mcp 1.x internals, so the import + raises under mcp 2.x — returns False there (and when ``fastmcp`` is absent).""" try: from fastmcp import FastMCP as FastMCPv2 except ImportError: @@ -27,5 +38,23 @@ def is_fastmcp_v2(server: Any) -> bool: return isinstance(server, FastMCPv2) +def is_mcpserver_v2(server: Any) -> bool: + """The mcp 2.x high-level server (``mcp.server.mcpserver.MCPServer``), which + replaced ``FastMCP``. Returns False on mcp 1.x, where that module is absent.""" + try: + from mcp.server.mcpserver import MCPServer + except ImportError: + return False + return isinstance(server, MCPServer) + + def is_low_level_server(server: Any) -> bool: + """A raw ``mcp.server.lowlevel.Server``. Present in both generations, but its + handler seam differs (v1: public ``request_handlers`` keyed by request type; + v2: private ``_request_handlers`` keyed by method string), so callers must + branch on generation before wrapping it.""" + try: + from mcp.server.lowlevel import Server as LowLevelServer + except ImportError: + return False return isinstance(server, LowLevelServer) diff --git a/posthog/mcp/_instrument_v2.py b/posthog/mcp/_instrument_v2.py new file mode 100644 index 00000000..dee0a1c7 --- /dev/null +++ b/posthog/mcp/_instrument_v2.py @@ -0,0 +1,317 @@ +# 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 2.x (2026-07-28 spec) adapter. + +The v2 SDK dropped the ``request_handlers`` monkey-patch seam the v1 adapters use +and added an official context-tier middleware protocol (``ServerMiddleware``): a +``(ctx, call_next)`` callable appended to ``server.middleware`` that wraps every +inbound request. We attach one such middleware instead of patching internals. + +Unlike v1, the middleware observes results as already-serialized wire dicts +(``{"content": [...], "isError": ...}`` for ``tools/call``, ``{"tools": [...]}`` +for ``tools/list``) or a raised ``MCPError`` on failure. Per-request client +identity and protocol version ride the ``_meta`` envelope +(``io.modelcontextprotocol/clientInfo`` etc.) on 2026-07-28 sessions, exposed on +``ctx.meta``; a client that negotiated an older protocol against this SDK sends +no envelope, so client info is recovered from the ``initialize`` params instead. + +This adapter is capture-only: it does not inject a ``context`` parameter, mutate +the tool list (no ``get_more_tools``), or stitch MRTR round-trips. An +``input_required`` interim result is stamped with ``$mcp_result_type`` and is +never treated as an error. +""" + +from __future__ import annotations + +import time +from typing import Any, Dict, Optional, Tuple + +from ._instrumentation import ( + build_tool_call_request, + prepare_request, + record_tool_call, + record_tools_list, +) +from ._internal import MCPAnalyticsData + +_MIDDLEWARE_FLAG = "__posthog_mcp_v2_middleware__" + + +def instrument_mcpserver_v2(server: Any, data: MCPAnalyticsData) -> None: + """Instrument the mcp 2.x high-level ``MCPServer`` by attaching the analytics + middleware to its (low-level-backed) ``middleware`` list.""" + low_level = getattr(server, "_lowlevel_server", None) + 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 + ) + _attach_middleware(server, data) + + +def instrument_low_level_v2(server: Any, data: MCPAnalyticsData) -> None: + """Instrument a raw mcp 2.x ``mcp.server.lowlevel.Server`` via its public + ``middleware`` list (the same seam ``MCPServer`` exposes).""" + data.server_name = getattr(server, "name", None) + data.server_version = getattr(server, "version", None) + _attach_middleware(server, data) + + +def _attach_middleware(server: Any, data: MCPAnalyticsData) -> None: + """Append the analytics ``ServerMiddleware`` to ``server.middleware``. + + Both ``MCPServer`` (via a property) and the low-level ``Server`` expose the + same ``middleware`` list; appending is the official attach mechanism (no + private-attr patching). Idempotent: a flagged middleware is never added twice. + """ + middleware = getattr(server, "middleware", None) + if middleware is None or not hasattr(middleware, "append"): + raise TypeError( + "mcp 2.x server exposes no `middleware` list to attach analytics to; " + f"got {type(server)!r}. Supported: mcp>=2,<3 MCPServer / lowlevel Server." + ) + if any(getattr(mw, _MIDDLEWARE_FLAG, False) for mw in middleware): + return + middleware.append(_build_middleware(data)) + + +def _build_middleware(data: MCPAnalyticsData) -> Any: + async def analytics_middleware(ctx: Any, call_next: Any) -> Any: + # Observe only the methods we capture; everything else passes straight + # through so we never alter dispatch for other requests or notifications. + method = getattr(ctx, "method", None) + if method == "tools/call": + return await _on_tool_call(data, ctx, call_next) + if method == "tools/list": + return await _on_tools_list(data, ctx, call_next) + if method == "server/discover": + # A discover carries client info in its envelope but no tool payload; + # emit the lazy $mcp_initialize (and identify) off it, then continue. + await _prepare(data, ctx, {"method": "server/discover", "params": {}}) + return await call_next(ctx) + return await call_next(ctx) + + setattr(analytics_middleware, _MIDDLEWARE_FLAG, True) + return analytics_middleware + + +def _client_info(ctx: Any) -> Tuple[Optional[str], Optional[str]]: + """Client name/version from the 2026-07-28 ``_meta`` envelope, or from an + ``initialize`` request's params when an older-protocol session sends no + envelope. Best-effort — returns ``(None, None)`` when neither is present.""" + meta = getattr(ctx, "meta", None) + if isinstance(meta, dict): + try: + from mcp_types import CLIENT_INFO_META_KEY + + info = meta.get(CLIENT_INFO_META_KEY) + except ImportError: + info = None + if isinstance(info, dict): + name = info.get("name") + version = info.get("version") + if name or version: + return name, version + # Legacy handshake on the v2 SDK: `initialize` params carry clientInfo. + params = getattr(ctx, "params", None) + if isinstance(params, dict): + info = params.get("clientInfo") + if isinstance(info, dict): + return info.get("name"), info.get("version") + return None, None + + +def _protocol_version(ctx: Any) -> Optional[str]: + version = getattr(ctx, "protocol_version", None) + if isinstance(version, str) and version: + return version + meta = getattr(ctx, "meta", None) + if isinstance(meta, dict): + try: + from mcp_types import PROTOCOL_VERSION_META_KEY + + value = meta.get(PROTOCOL_VERSION_META_KEY) + except ImportError: + value = None + if isinstance(value, str) and value: + return value + return None + + +async def _prepare( + data: MCPAnalyticsData, ctx: Any, request: Dict[str, Any] +) -> Tuple[str, Optional[str], Optional[str], Optional[str]]: + """Resolve session (emitting identify + lazy initialize) for this request. + + v2 traffic has no self-encoded token and no ``Mcp-Session-Id`` header by + construction (SEP-2567 removed it), so ``mcp_session_id``/``token`` are always + None here; sessions come out derived (Stage C) when identified, generated + otherwise. Returns the session id plus the resolved client/protocol info so + the caller stamps the same values on the captured event.""" + client_name, client_version = _client_info(ctx) + protocol_version = _protocol_version(ctx) + extra = {"session_id": None} + session_id = await prepare_request( + data, + mcp_session_id=None, + client_name=client_name, + client_version=client_version, + protocol_version=protocol_version, + request=request, + extra=extra, + ) + return session_id, client_name, client_version, protocol_version + + +def _result_type(result: Any) -> Optional[str]: + """The 2026-07-28 ``resultType`` of a wire result, when non-default. Returns + None for a plain "complete" result so we only stamp interesting values.""" + value: Any = None + if isinstance(result, dict): + value = result.get("resultType") + else: + value = getattr(result, "result_type", None) + if isinstance(value, str) and value and value != "complete": + return value + return None + + +def _is_error_result(result: Any) -> bool: + """A v2 tool result signals a tool error via ``isError``/``is_error`` — but an + ``input_required`` interim result (MRTR) is NOT an error, just a continuation.""" + if _result_type(result) == "input_required": + return False + if isinstance(result, dict): + return result.get("isError") is True + return ( + getattr(result, "is_error", None) is True + or getattr(result, "isError", None) is True + ) + + +async def _on_tool_call(data: MCPAnalyticsData, ctx: Any, call_next: Any) -> Any: + params = ctx.params if isinstance(ctx.params, dict) else {} + # The tools/call wire always carries a string name; default defensively so a + # malformed request captures an event rather than raising into dispatch. + raw_name = params.get("name") + name: str = raw_name if isinstance(raw_name, str) else "" + arguments = params.get("arguments") or {} + request = build_tool_call_request(name, arguments) + + session_id, client_name, client_version, protocol_version = await _prepare( + data, ctx, request + ) + + start = time.monotonic() + try: + result = await call_next(ctx) + except Exception as error: + # The middleware chain surfaces a request-side failure as a raised + # MCPError; capture it before re-raising so the failed call isn't dropped. + 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, + ) + raise + duration_ms = (time.monotonic() - start) * 1000 + + await record_tool_call( + data, + session_id, + name=name, + arguments=arguments, + result=_ForceErrorFlag(result) if _is_error_result(result) else result, + duration_ms=duration_ms, + client_name=client_name, + client_version=client_version, + protocol_version=protocol_version, + result_type=_result_type(result), + ) + return result + + +class _ForceErrorFlag: + """Adapts a v2 wire result so ``record_tool_call``'s ``isError`` detection + (which reads ``.isError``/``["isError"]``) fires for the value we already + classified as a tool error — including a model-shaped result whose flag is + ``is_error``. Wraps, never mutates, the result the tool actually returned.""" + + isError = True + + def __init__(self, wrapped: Any) -> None: + self._wrapped = wrapped + + def __getattr__(self, item: str) -> Any: + return getattr(self._wrapped, item) + + +async def _on_tools_list(data: MCPAnalyticsData, ctx: Any, call_next: Any) -> Any: + request = {"method": "tools/list", "params": {}} + session_id, client_name, client_version, protocol_version = await _prepare( + data, ctx, request + ) + + start = time.monotonic() + try: + result = await call_next(ctx) + 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, + ) + raise + duration_ms = (time.monotonic() - start) * 1000 + + names = _listed_tool_names(result) + empty = len(names) == 0 + await record_tools_list( + data, + session_id, + names=names, + request=request, + response=result if isinstance(result, dict) else None, + 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, + ) + return result + + +def _listed_tool_names(result: Any) -> list: + """Tool names out of a v2 ``tools/list`` result (a wire dict ``{"tools": + [{"name": ...}, ...]}``, or a model with a ``.tools`` list). Read-only — the + v2 adapter never mutates the response.""" + tools: Any = None + if isinstance(result, dict): + tools = result.get("tools") + else: + tools = getattr(result, "tools", None) + if not isinstance(tools, list): + return [] + names = [] + for tool in tools: + name = ( + tool.get("name") if isinstance(tool, dict) else getattr(tool, "name", None) + ) + if isinstance(name, str): + names.append(name) + return names diff --git a/posthog/mcp/_instrumentation.py b/posthog/mcp/_instrumentation.py index ff5380ab..700a1ab4 100644 --- a/posthog/mcp/_instrumentation.py +++ b/posthog/mcp/_instrumentation.py @@ -286,6 +286,7 @@ async def record_tool_call( client_version: Optional[str] = None, protocol_version: Optional[str] = None, conversation_id: Optional[str] = None, + result_type: Optional[str] = None, extra: Optional[Dict[str, Any]] = None, ) -> None: # Analytics must never change what the tool returns or raises: any failure @@ -304,6 +305,7 @@ async def record_tool_call( "client_version": client_version, "protocol_version": protocol_version, "conversation_id": conversation_id, + "result_type": result_type, "is_error": False, } set_event_intent(event, await resolve_tool_call_intent(data, request, extra)) diff --git a/posthog/mcp/_posthog_events.py b/posthog/mcp/_posthog_events.py index fd6f7363..619459c1 100644 --- a/posthog/mcp/_posthog_events.py +++ b/posthog/mcp/_posthog_events.py @@ -139,6 +139,8 @@ def _add_common_properties(event: Event, properties: Dict[str, Any]) -> None: properties[_P.INTENT_SOURCE] = event["user_intent_source"] if event.get("is_error") is not None: properties[_P.IS_ERROR] = event["is_error"] + if event.get("result_type") and _is_tool_call(event): + properties[_P.RESULT_TYPE] = event["result_type"] if event.get("parameters") is not None: properties[_P.PARAMETERS] = event["parameters"] if event.get("response") is not None: diff --git a/posthog/mcp/constants.py b/posthog/mcp/constants.py index eb2de16e..73310064 100644 --- a/posthog/mcp/constants.py +++ b/posthog/mcp/constants.py @@ -67,6 +67,11 @@ class PostHogMCPAnalyticsProperty: PARAMETERS = "$mcp_parameters" RESOURCE_NAME = "$mcp_resource_name" RESPONSE = "$mcp_response" + # The 2026-07-28 `resultType` of a tool result (e.g. "input_required" for an + # MRTR interim result). Present only on mcp 2.x traffic that returns a + # non-"complete" result, so downstream can segment MRTR calls without full + # request stitching (which is out of scope here). + RESULT_TYPE = "$mcp_result_type" SERVER_NAME = "$mcp_server_name" SERVER_VERSION = "$mcp_server_version" SESSION_ID = "$session_id" diff --git a/posthog/test/mcp/test_instrument_v2.py b/posthog/test/mcp/test_instrument_v2.py new file mode 100644 index 00000000..fa1794b2 --- /dev/null +++ b/posthog/test/mcp/test_instrument_v2.py @@ -0,0 +1,180 @@ +"""End-to-end tests for the mcp 2.x (2026-07-28 spec) middleware adapter. + +Drives a real `MCPServer` over the SDK's in-memory transport so the adapter is +exercised through the official `ServerMiddleware` seam exactly as production +traffic hits it: `server/discover`, `tools/list`, and `tools/call` all flow +through `instrument()`'s attached middleware. Runs only in the mcp-v2 env.""" + +import pytest + +from posthog.test.mcp._helpers import ( + FakeClient, + events_named as _events, + requires_mcp_v2, +) + +pytestmark = requires_mcp_v2 + +# `mcp.server.mcpserver` is the mcp 2.x high-level server module; it doesn't exist +# in mcp 1.x, so importing it would crash collection there. Skip the module before +# that import runs (the marker above additionally documents intent). +pytest.importorskip("mcp.server.mcpserver") + +from mcp.server.mcpserver import MCPServer # noqa: E402 +from mcp.client._memory import InMemoryTransport # noqa: E402 +from mcp.client.session import ClientSession # noqa: E402 + +from mcp_types import InputRequiredResult # noqa: E402 + +from posthog.mcp import instrument # noqa: E402 +from posthog.mcp.types import MCPAnalyticsOptions, UserIdentity # noqa: E402 + + +def make_server(name="v2-probe"): + server = MCPServer(name) + + @server.tool() + def add(a: int, b: int) -> int: + return a + b + + @server.tool() + def boom() -> int: + raise ValueError("kaboom") + + return server + + +async def drive( + server, *, modern=True, calls=(("add", {"a": 2, "b": 3}),), list_tools=True +): + """Connect an in-memory client, run a modern (server/discover) or legacy + (initialize) handshake, then list tools and make the given tool calls. + Returns the collected client-side results in call order.""" + results = [] + async with InMemoryTransport(server) as (read, write): + async with ClientSession(read, write) as session: + if modern: + await session.discover() + else: + await session.initialize() + if list_tools: + await session.list_tools() + for name, args in calls: + results.append(await session.call_tool(name, args)) + return results + + +async def test_tool_call_captured_with_envelope_identity(): + server = make_server() + client = FakeClient() + instrument(server, client) + + await drive(server, calls=(("add", {"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_is_error"] is False + # Client identity + protocol version come from the per-request _meta envelope. + assert props["$mcp_client_name"] == "mcp" + assert props["$mcp_protocol_version"] == "2026-07-28" + + +async def test_tool_error_result_captured_as_error(): + server = make_server() + client = FakeClient() + instrument(server, client) + + await drive(server, calls=(("boom", {}),)) + + calls = _events(client, "$mcp_tool_call") + assert calls and calls[0]["properties"]["$mcp_is_error"] is True + exceptions = _events(client, "$exception") + assert exceptions, "an isError tool result should also emit $exception" + + +async def test_tools_list_captured_with_names(): + server = make_server() + client = FakeClient() + instrument(server, client) + + await drive(server, calls=()) + + listed = _events(client, "$mcp_tools_list") + assert listed + names = listed[0]["properties"]["$mcp_listed_tool_names"] + assert "add" in names and "boom" in names + + +async def test_initialize_emitted_once_per_session(): + server = make_server() + client = FakeClient() + instrument(server, client) + + await drive(server, calls=(("add", {"a": 1, "b": 1}), ("add", {"a": 2, "b": 2}))) + + assert len(_events(client, "$mcp_initialize")) == 1 + assert len(_events(client, "$mcp_tool_call")) == 2 + + +async def test_identify_flows_through_v2_adapter(): + server = make_server() + client = FakeClient() + instrument( + server, + client, + MCPAnalyticsOptions( + identify=UserIdentity(distinct_id="user-42", properties={"plan": "pro"}) + ), + ) + + await drive(server, calls=(("add", {"a": 1, "b": 1}),)) + + identifies = _events(client, "$identify") + assert identifies and identifies[0]["distinct_id"] == "user-42" + # The tool call is attributed to the identified user, not the anonymous session. + calls = _events(client, "$mcp_tool_call") + assert calls and calls[0]["distinct_id"] == "user-42" + + +async def test_mrtr_input_required_stamps_result_type_not_error(): + # An MRTR interim result (`resultType: "input_required"`) is a continuation, + # not a failure: it must NOT be flagged as an error, and it stamps + # `$mcp_result_type` so downstream can segment MRTR calls. Full round-trip + # stitching is out of scope. + server = MCPServer("mrtr") + + @server.tool() + def needs_input() -> InputRequiredResult: + return InputRequiredResult(input_requests={}, request_state="s1") + + client = FakeClient() + instrument(server, client) + + async with InMemoryTransport(server) as (read, write): + async with ClientSession(read, write) as session: + await session.discover() + await session.call_tool("needs_input", {}, allow_input_required=True) + + calls = _events(client, "$mcp_tool_call") + assert calls + props = calls[0]["properties"] + assert props["$mcp_is_error"] is False + assert props["$mcp_result_type"] == "input_required" + + +async def test_legacy_negotiation_on_v2_sdk_still_captures(): + # A client that runs the classic `initialize` handshake against the v2 SDK + # negotiates an older protocol: the _meta envelope is absent, so client info + # must be recovered from the initialize params instead. The adapter must + # still capture the tool call (client name may be backfilled or absent, but + # the event must fire and not be an error). + server = make_server() + client = FakeClient() + instrument(server, client) + + await drive(server, modern=False, calls=(("add", {"a": 4, "b": 5}),)) + + calls = _events(client, "$mcp_tool_call") + assert calls and calls[0]["properties"]["$mcp_is_error"] is False diff --git a/posthog/test/mcp/test_review_fixes.py b/posthog/test/mcp/test_review_fixes.py index 65ec445f..fc5e0289 100644 --- a/posthog/test/mcp/test_review_fixes.py +++ b/posthog/test/mcp/test_review_fixes.py @@ -340,4 +340,4 @@ async def _lt(): server, FakeClient(), MCPAnalyticsOptions(logger=lambda m: logs.append(m)) ) - assert any("tested against mcp>=1.26" in m for m in logs) + assert any("supports mcp>=1.26,<2 and mcp>=2,<3" in m for m in logs) From 08011b679d3793677857f2353f4fc7263e84eb70 Mon Sep 17 00:00:00 2001 From: Lucas Faria Date: Tue, 4 Aug 2026 20:08:33 -0300 Subject: [PATCH 3/4] mcp: derive sessions for stateless traffic + session-id provenance MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit SEP-2567 removed the Mcp-Session-Id header on 2026-07-28, so stateless / per- request servers have no transport session to correlate a user's tool calls into one $session_id — each request would mint its own. Per SEP-2567's telemetry guidance, derive the session from the authenticated principal + client instead. What: - `_derived_sessions.py`: a module-level (process-shared, so per-request server instances correlate) `DerivedSessionRegistry` mapping `(distinct_id, client_name, client_version)` -> a rolling `ses_` UUIDv7. Thread-safe (lock), LRU-bounded (10k), idle-evicts entries past 2x the inactivity timeout, rolls a session after the timeout. Fork-reset via `os.register_at_fork`, mirroring the background-loop reset. - `session.resolve_session_id`: new precedence token > mcp > sticky-mcp > derived > generated. Derived is taken only when a `distinct_id` is present — deriving anonymously would merge unrelated users under one session. - `_internal.py`: split identity resolution out of `handle_identify` into `resolve_identity(data, request, extra)` (callback invoked at most once, side-effect-free). `prepare_request` now resolves identity FIRST and threads it into both `resolve_session_id` (for the derived key) and `handle_identify` (dedup still keyed by the resolved session id). - Provenance: `$mcp_session_id_source` (token|mcp|derived|generated) added to `constants.py` and stamped on every $mcp_* event (and $identify). Threaded through `record_tool_call`/`record_tools_list`/`record_missing_capability`/ `_maybe_emit_initialize` and the v1 (fastmcp, lowlevel) + v2 adapters; the source is snapshotted at resolution time (shared per-server state) rather than re-read at capture. Additive. - The v2 adapter passes no token / no mcp header by construction, so its sessions are `derived` when identified, `generated` otherwise. How tested: - `test_derived_sessions.py` (parameterized): same key within gap -> one session; gap expiry -> new; different distinct_id/client_name/client_version -> different; no distinct_id -> generated; LRU bound; idle eviction; concurrency (8 threads, one key -> one session); fork reset; full precedence table (token/mcp/derived/generated); derived requires distinct_id. - Provenance parameterized tests assert `$mcp_session_id_source` on initialize/ tools_list/tool_call/identify in both the v1 (test_lowlevel) and v2 (test_instrument_v2) adapters: identified -> derived, anonymous -> generated. - v1 env: 167 passed, 2 skipped. v2 env: 134 passed, 13 skipped. Full-suite `--collect-only`: 1997 tests, no import errors. - `ruff@0.11.12 check .`/`format --check .` clean; `mypy | mypy-baseline filter` -> no issues (194 files). Generated-By: PostHog Code Task-Id: c17ecd78-bb05-4b3c-b621-458abba9c6aa --- posthog/mcp/_capture.py | 1 + posthog/mcp/_derived_sessions.py | 134 +++++++++++++++++++ posthog/mcp/_instrument_fastmcp.py | 9 +- posthog/mcp/_instrument_lowlevel.py | 9 +- posthog/mcp/_instrument_v2.py | 36 ++++-- posthog/mcp/_instrumentation.py | 57 +++++++-- posthog/mcp/_internal.py | 52 ++++++-- posthog/mcp/_posthog_events.py | 3 + posthog/mcp/constants.py | 5 + posthog/mcp/session.py | 25 ++++ posthog/test/mcp/test_derived_sessions.py | 149 ++++++++++++++++++++++ posthog/test/mcp/test_instrument_v2.py | 30 +++++ posthog/test/mcp/test_lowlevel.py | 36 ++++++ 13 files changed, 507 insertions(+), 39 deletions(-) create mode 100644 posthog/mcp/_derived_sessions.py create mode 100644 posthog/test/mcp/test_derived_sessions.py diff --git a/posthog/mcp/_capture.py b/posthog/mcp/_capture.py index bda3c19a..88d1ed1c 100644 --- a/posthog/mcp/_capture.py +++ b/posthog/mcp/_capture.py @@ -39,6 +39,7 @@ def capture_event( full_event: Dict[str, Any] = { "id": event_input.get("id") or "", "session_id": session_id, + "session_id_source": event_input.get("session_id_source"), "event_type": event_input.get("event_type") or MCPAnalyticsEventType.CUSTOM, "event_name": event_input.get("event_name"), "timestamp": timestamp, diff --git a/posthog/mcp/_derived_sessions.py b/posthog/mcp/_derived_sessions.py new file mode 100644 index 00000000..2c65650d --- /dev/null +++ b/posthog/mcp/_derived_sessions.py @@ -0,0 +1,134 @@ +# 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 + +"""Derived sessions for stateless MCP traffic. + +The 2026-07-28 spec (SEP-2567) removed the ``Mcp-Session-Id`` header, so a +stateless / per-request server has no transport session to correlate a user's +tool calls into one ``$session_id``. SEP-2567's telemetry guidance is to derive +the session from "the authenticated principal ... or a request-level correlation +ID". This registry does exactly that: it maps +``(distinct_id, client_name, client_version)`` to a stable ``ses_`` id that rolls +over after the same inactivity timeout as the in-memory fallback. + +The registry is module-level (not per-server) so that per-request server +instances created within one process share it — a stateless deployment spins up a +fresh server per request, and each would otherwise mint its own session. It never +derives without a ``distinct_id``: an anonymous key would merge unrelated users +under one session, so callers fall back to a generated session instead. +""" + +from __future__ import annotations + +import os +import threading +from collections import OrderedDict +from datetime import datetime, timezone +from typing import Optional + +from .constants import INACTIVITY_TIMEOUT_IN_MINUTES +from ._ids import new_prefixed_id + +# Bound the registry so a long-lived process serving many distinct users can't +# grow it without limit; the LRU evicts the least-recently-resolved key. +_MAX_ENTRIES = 10_000 + +# An entry idle beyond twice the inactivity timeout can never be reused (a reuse +# within one timeout is required, and past one timeout the session rolls anyway), +# so it is eligible for opportunistic eviction on the next resolve. +_IDLE_EVICTION_SECONDS = 2 * INACTIVITY_TIMEOUT_IN_MINUTES * 60 + +_DerivedKey = tuple[str, str, str] + + +class DerivedSessionRegistry: + """Thread-safe, bounded map from an identity+client key to a rolling session + id. Standalone (not tied to ``MCPAnalyticsData``) so a process's per-request + server instances share one registry.""" + + def __init__(self) -> None: + # OrderedDict as an LRU: move_to_end on reuse, popitem(last=False) to evict. + self._entries: "OrderedDict[_DerivedKey, tuple[str, datetime]]" = OrderedDict() + self._lock = threading.Lock() + + def resolve( + self, + distinct_id: str, + client_name: str, + client_version: str, + *, + now: Optional[datetime] = None, + ) -> str: + """The stable session id for this key, minting a new one on first sight or + after the inactivity timeout. ``now`` is injectable for deterministic tests.""" + now = now or datetime.now(timezone.utc) + key = (distinct_id, client_name or "", client_version or "") + timeout_seconds = INACTIVITY_TIMEOUT_IN_MINUTES * 60 + with self._lock: + self._evict_idle_locked(now) + existing = self._entries.get(key) + if existing is not None: + session_id, last_activity = existing + if (now - last_activity).total_seconds() <= timeout_seconds: + self._entries[key] = (session_id, now) + self._entries.move_to_end(key) + return session_id + session_id = new_prefixed_id("ses") + self._entries[key] = (session_id, now) + self._entries.move_to_end(key) + self._enforce_bound_locked() + return session_id + + def _evict_idle_locked(self, now: datetime) -> None: + # Entries are time-ordered by last activity only loosely (LRU by access), + # so scan for any idle beyond the eviction horizon. Bounded by _MAX_ENTRIES. + stale = [ + key + for key, (_sid, last_activity) in self._entries.items() + if (now - last_activity).total_seconds() > _IDLE_EVICTION_SECONDS + ] + for key in stale: + del self._entries[key] + + def _enforce_bound_locked(self) -> None: + while len(self._entries) > _MAX_ENTRIES: + self._entries.popitem(last=False) + + def size(self) -> int: + with self._lock: + return len(self._entries) + + def clear(self) -> None: + with self._lock: + self._entries.clear() + + +_DERIVED_REGISTRY = DerivedSessionRegistry() + + +def derive_session_id( + distinct_id: str, + client_name: Optional[str], + client_version: Optional[str], + *, + now: Optional[datetime] = None, +) -> str: + """Resolve the process-shared derived session id for an identified request.""" + return _DERIVED_REGISTRY.resolve( + distinct_id, client_name or "", client_version or "", now=now + ) + + +def _reset_derived_registry_after_fork() -> None: + """Drop registry state inherited by a forked child. + + The lock may have been held by a thread that does not survive ``fork()``, so + replace the whole registry rather than acquiring the inherited lock. Mirrors + the background-loop fork reset in ``_instrumentation.py``.""" + global _DERIVED_REGISTRY + _DERIVED_REGISTRY = DerivedSessionRegistry() + + +if hasattr(os, "register_at_fork"): + os.register_at_fork(after_in_child=_reset_derived_registry_after_fork) diff --git a/posthog/mcp/_instrument_fastmcp.py b/posthog/mcp/_instrument_fastmcp.py index 6fd871f2..bf57ff98 100644 --- a/posthog/mcp/_instrument_fastmcp.py +++ b/posthog/mcp/_instrument_fastmcp.py @@ -101,7 +101,7 @@ async def wrapped( request = build_tool_call_request(name, arguments) extra: Dict[str, Any] = {"session_id": mcp_session_id} - session_id = await prepare_request( + session_id, session_id_source = await prepare_request( data, mcp_session_id=mcp_session_id, client_name=client_name, @@ -123,6 +123,7 @@ async def wrapped( client_name=client_name, client_version=client_version, protocol_version=protocol_version, + session_id_source=session_id_source, extra=extra, ) return [ @@ -169,6 +170,7 @@ async def wrapped( client_version=client_version, protocol_version=protocol_version, conversation_id=None if minted else conversation_id, + session_id_source=session_id_source, extra=extra, ) raise @@ -195,6 +197,7 @@ async def wrapped( client_version=client_version, protocol_version=protocol_version, conversation_id=delivered_conversation_id, + session_id_source=session_id_source, extra=extra, ) return result @@ -233,7 +236,7 @@ async def list_handler(req: Any) -> Any: extra: Dict[str, Any] = {"session_id": mcp_session_id} # 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( + session_id, session_id_source = await prepare_request( data, mcp_session_id=mcp_session_id, client_name=client_name, @@ -259,6 +262,7 @@ async def list_handler(req: Any) -> Any: client_name=client_name, client_version=client_version, protocol_version=protocol_version, + session_id_source=session_id_source, extra=extra, ) raise @@ -312,6 +316,7 @@ async def list_handler(req: Any) -> Any: client_name=client_name, client_version=client_version, protocol_version=protocol_version, + session_id_source=session_id_source, extra=extra, ) diff --git a/posthog/mcp/_instrument_lowlevel.py b/posthog/mcp/_instrument_lowlevel.py index cbe6c500..ee73d9e4 100644 --- a/posthog/mcp/_instrument_lowlevel.py +++ b/posthog/mcp/_instrument_lowlevel.py @@ -104,7 +104,7 @@ async def handler(req: Any) -> Any: request = build_tool_call_request(name, arguments) extra = {"session_id": mcp_session_id} - session_id = await prepare_request( + session_id, session_id_source = await prepare_request( data, mcp_session_id=mcp_session_id, client_name=client_name, @@ -126,6 +126,7 @@ async def handler(req: Any) -> Any: client_name=client_name, client_version=client_version, protocol_version=protocol_version, + session_id_source=session_id_source, extra=extra, ) return mcp_types.ServerResult( @@ -176,6 +177,7 @@ async def handler(req: Any) -> Any: client_version=client_version, protocol_version=protocol_version, conversation_id=None if minted else conversation_id, + session_id_source=session_id_source, extra=extra, ) raise @@ -211,6 +213,7 @@ async def handler(req: Any) -> Any: client_version=client_version, protocol_version=protocol_version, conversation_id=delivered_conversation_id, + session_id_source=session_id_source, extra=extra, ) return result @@ -245,7 +248,7 @@ async def handler(req: Any) -> Any: extra = {"session_id": mcp_session_id} # 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( + session_id, session_id_source = await prepare_request( data, mcp_session_id=mcp_session_id, client_name=client_name, @@ -271,6 +274,7 @@ async def handler(req: Any) -> Any: client_name=client_name, client_version=client_version, protocol_version=protocol_version, + session_id_source=session_id_source, extra=extra, ) raise @@ -330,6 +334,7 @@ async def handler(req: Any) -> Any: client_name=client_name, client_version=client_version, protocol_version=protocol_version, + session_id_source=session_id_source, extra=extra, ) diff --git a/posthog/mcp/_instrument_v2.py b/posthog/mcp/_instrument_v2.py index dee0a1c7..f5a68ea2 100644 --- a/posthog/mcp/_instrument_v2.py +++ b/posthog/mcp/_instrument_v2.py @@ -141,18 +141,18 @@ def _protocol_version(ctx: Any) -> Optional[str]: async def _prepare( data: MCPAnalyticsData, ctx: Any, request: Dict[str, Any] -) -> Tuple[str, Optional[str], Optional[str], Optional[str]]: +) -> Tuple[str, str, Optional[str], Optional[str], Optional[str]]: """Resolve session (emitting identify + lazy initialize) for this request. v2 traffic has no self-encoded token and no ``Mcp-Session-Id`` header by construction (SEP-2567 removed it), so ``mcp_session_id``/``token`` are always - None here; sessions come out derived (Stage C) when identified, generated - otherwise. Returns the session id plus the resolved client/protocol info so - the caller stamps the same values on the captured event.""" + None here; sessions come out ``derived`` (Stage C) when identified, + ``generated`` otherwise. Returns the session id, its provenance source, and the + resolved client/protocol info so the caller stamps them on the captured event.""" client_name, client_version = _client_info(ctx) protocol_version = _protocol_version(ctx) extra = {"session_id": None} - session_id = await prepare_request( + session_id, session_id_source = await prepare_request( data, mcp_session_id=None, client_name=client_name, @@ -161,7 +161,7 @@ async def _prepare( request=request, extra=extra, ) - return session_id, client_name, client_version, protocol_version + return session_id, session_id_source, client_name, client_version, protocol_version def _result_type(result: Any) -> Optional[str]: @@ -199,9 +199,13 @@ async def _on_tool_call(data: MCPAnalyticsData, ctx: Any, call_next: Any) -> Any arguments = params.get("arguments") or {} request = build_tool_call_request(name, arguments) - session_id, client_name, client_version, protocol_version = await _prepare( - data, ctx, request - ) + ( + session_id, + session_id_source, + client_name, + client_version, + protocol_version, + ) = await _prepare(data, ctx, request) start = time.monotonic() try: @@ -219,6 +223,7 @@ async def _on_tool_call(data: MCPAnalyticsData, ctx: Any, call_next: Any) -> Any client_name=client_name, client_version=client_version, protocol_version=protocol_version, + session_id_source=session_id_source, ) raise duration_ms = (time.monotonic() - start) * 1000 @@ -234,6 +239,7 @@ async def _on_tool_call(data: MCPAnalyticsData, ctx: Any, call_next: Any) -> Any client_version=client_version, protocol_version=protocol_version, result_type=_result_type(result), + session_id_source=session_id_source, ) return result @@ -255,9 +261,13 @@ def __getattr__(self, item: str) -> Any: async def _on_tools_list(data: MCPAnalyticsData, ctx: Any, call_next: Any) -> Any: request = {"method": "tools/list", "params": {}} - session_id, client_name, client_version, protocol_version = await _prepare( - data, ctx, request - ) + ( + session_id, + session_id_source, + client_name, + client_version, + protocol_version, + ) = await _prepare(data, ctx, request) start = time.monotonic() try: @@ -274,6 +284,7 @@ async def _on_tools_list(data: MCPAnalyticsData, ctx: Any, call_next: Any) -> An client_name=client_name, client_version=client_version, protocol_version=protocol_version, + session_id_source=session_id_source, ) raise duration_ms = (time.monotonic() - start) * 1000 @@ -292,6 +303,7 @@ async def _on_tools_list(data: MCPAnalyticsData, ctx: Any, call_next: Any) -> An client_name=client_name, client_version=client_version, protocol_version=protocol_version, + session_id_source=session_id_source, ) return result diff --git a/posthog/mcp/_instrumentation.py b/posthog/mcp/_instrumentation.py index 700a1ab4..8d081fff 100644 --- a/posthog/mcp/_instrumentation.py +++ b/posthog/mcp/_instrumentation.py @@ -20,7 +20,12 @@ from ._event_types import MCPAnalyticsEventType from ._exceptions import capture_exception from ._intent import resolve_tool_call_intent, set_event_intent -from ._internal import MCPAnalyticsData, handle_identify, resolve_event_properties +from ._internal import ( + MCPAnalyticsData, + handle_identify, + resolve_event_properties, + resolve_identity, +) from .logger import log from ._sanitization import build_captured_mcp_parameters from .session import resolve_session_id @@ -185,6 +190,7 @@ async def _maybe_emit_initialize( client_version: Optional[str], extra: Optional[Dict[str, Any]], protocol_version: Optional[str] = None, + session_id_source: Optional[str] = None, ) -> None: """Lazily emit ``$mcp_initialize`` once per session. The Python MCP SDK handles ``InitializeRequest`` inside the session layer (not ``request_handlers``), so we @@ -195,6 +201,7 @@ async def _maybe_emit_initialize( event: Dict[str, Any] = { "event_type": MCPAnalyticsEventType.MCP_INITIALIZE, "session_id": session_id, + "session_id_source": session_id_source, "client_name": client_name, "client_version": client_version, "protocol_version": protocol_version, @@ -250,27 +257,51 @@ async def prepare_request( extra: Optional[Dict[str, Any]], token: Optional[SessionTokenPayload] = None, protocol_version: 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. +) -> tuple[str, str]: + """Resolve identity, then the session id, run identify, then lazily emit + initialize. Returns ``(session_id, session_id_source)`` to stamp on this + request's event; ``session_id_source`` is one of ``token`` | ``mcp`` | + ``derived`` | ``generated``. ``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. - Identify runs *before* initialize so the resolved identity is already in the cache - when ``capture_event`` builds the initialize event — otherwise the first + Identity is resolved FIRST (its callback invoked at most once), then passed to + both session resolution — so an identified stateless request can derive a stable + session from its ``distinct_id`` — and to ``handle_identify``. Identify runs + *before* initialize so the resolved identity is already in the cache when + ``capture_event`` builds the initialize event — otherwise the first ``$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) - identify_event = await handle_identify(data, session_id, request, extra) + identity = await resolve_identity(data, request, extra) + session_id = await resolve_session_id( + data, + mcp_session_id, + token=token, + identity=identity, + client_name=client_name, + client_version=client_version, + ) + # Snapshot the provenance under no further mutation: `data.session_source` is + # shared per-server state, so read it right after resolution and thread the + # value onto every event this request emits rather than re-reading at capture. + session_id_source = data.session_source + identify_event = await handle_identify(data, session_id, identity, request, extra) if identify_event: + identify_event["session_id_source"] = session_id_source fire_and_forget(capture_event(data, identify_event), data) await _maybe_emit_initialize( - data, session_id, client_name, client_version, extra, protocol_version + data, + session_id, + client_name, + client_version, + extra, + protocol_version, + session_id_source, ) - return session_id + return session_id, session_id_source async def record_tool_call( @@ -287,6 +318,7 @@ async def record_tool_call( protocol_version: Optional[str] = None, conversation_id: Optional[str] = None, result_type: Optional[str] = None, + session_id_source: Optional[str] = None, extra: Optional[Dict[str, Any]] = None, ) -> None: # Analytics must never change what the tool returns or raises: any failure @@ -296,6 +328,7 @@ async def record_tool_call( event: Dict[str, Any] = { "event_type": MCPAnalyticsEventType.MCP_TOOLS_CALL, "session_id": session_id, + "session_id_source": session_id_source, "resource_name": name, "tool_description": data.tool_descriptions.get(name), "tool_category": data.tool_categories.get(name), @@ -387,6 +420,7 @@ async def record_missing_capability( client_name: Optional[str] = None, client_version: Optional[str] = None, protocol_version: Optional[str] = None, + session_id_source: Optional[str] = None, extra: Optional[Dict[str, Any]] = None, ) -> None: """Record a ``get_more_tools`` call as ``$mcp_missing_capability``, with the @@ -396,6 +430,7 @@ async def record_missing_capability( event: Dict[str, Any] = { "event_type": MCPAnalyticsEventType.MCP_MISSING_CAPABILITY, "session_id": session_id, + "session_id_source": session_id_source, "resource_name": tool_name, "parameters": build_captured_mcp_parameters(request), "client_name": client_name, @@ -424,12 +459,14 @@ async def record_tools_list( client_name: Optional[str] = None, client_version: Optional[str] = None, protocol_version: Optional[str] = None, + session_id_source: Optional[str] = None, extra: Optional[Dict[str, Any]] = None, ) -> None: try: event: Dict[str, Any] = { "event_type": MCPAnalyticsEventType.MCP_TOOLS_LIST, "session_id": session_id, + "session_id_source": session_id_source, "listed_tool_names": names, "parameters": build_captured_mcp_parameters(request), "response": _wrap_response(response) if response is not None else None, diff --git a/posthog/mcp/_internal.py b/posthog/mcp/_internal.py index 5c529541..9cb900cf 100644 --- a/posthog/mcp/_internal.py +++ b/posthog/mcp/_internal.py @@ -132,33 +132,59 @@ async def _maybe_await(value: Any) -> Any: return value -async def handle_identify( +async def resolve_identity( data: MCPAnalyticsData, - session_id: str, request: Dict[str, Any], extra: Optional[Dict[str, Any]] = None, -) -> Optional[Dict[str, Any]]: - """Resolve the optional ``identify`` callback, dedupe against the identity - cache, and return an ``$identify`` event to emit only when the identity has - materially changed (otherwise ``None``).""" +) -> Optional[UserIdentity]: + """Invoke the customer's ``identify`` callback (or read the static identity) + and return the resolved :class:`UserIdentity`, or ``None`` when no identify is + configured / it returns nothing / it raises. + + Runs *before* session resolution so the resolved ``distinct_id`` is available + to derive a session key (see :func:`.session.resolve_session_id`). It performs + no caching or ``$identify`` decision — that stays in :func:`handle_identify`, + which is keyed by the resolved session id. Kept side-effect-free so it can be + called at most once per request and its result threaded to both consumers.""" if not data.options.identify: return None - try: identify = data.options.identify if isinstance(identify, UserIdentity): - identity_result: Optional[UserIdentity] = identify - else: - identity_result = await _maybe_await(identify(request, extra)) + return identify + result = await _maybe_await(identify(request, extra)) + return result or None + except Exception as error: # noqa: BLE001 + log(f"Error: identify function threw while resolving identity - {error}") + return None + - if not identity_result: +async def handle_identify( + data: MCPAnalyticsData, + session_id: str, + identity: Optional[UserIdentity], + request: Dict[str, Any], + extra: Optional[Dict[str, Any]] = None, +) -> Optional[Dict[str, Any]]: + """Given the already-resolved ``identity`` (from :func:`resolve_identity`), + dedupe against the per-session identity cache and return an ``$identify`` event + to emit only when the identity has materially changed (otherwise ``None``). + + Split from identity resolution so the callback is invoked at most once per + request while the dedupe cache stays keyed by ``session_id`` (which is resolved + before this runs).""" + if not data.options.identify: + return None + if not identity: + if data.options.identify: log( f"Warning: Supplied identify function returned null for session {session_id}" ) - return None + return None + try: previous = data.identified_sessions.get(session_id) - merged = merge_identities(previous, identity_result) + merged = merge_identities(previous, identity) has_changed = not (previous and are_identities_equal(previous, merged)) data.identified_sessions.set(session_id, merged) diff --git a/posthog/mcp/_posthog_events.py b/posthog/mcp/_posthog_events.py index 619459c1..a8ac6a83 100644 --- a/posthog/mcp/_posthog_events.py +++ b/posthog/mcp/_posthog_events.py @@ -81,6 +81,9 @@ def _add_session_id(event: Event, properties: Dict[str, Any]) -> None: session_id = event.get("session_id") if isinstance(session_id, str) and len(session_id) > 0: properties[_P.SESSION_ID] = session_id + source = event.get("session_id_source") + if isinstance(source, str) and source: + properties[_P.SESSION_ID_SOURCE] = source def _add_conversation_id(event: Event, properties: Dict[str, Any]) -> None: diff --git a/posthog/mcp/constants.py b/posthog/mcp/constants.py index 73310064..570e3ab4 100644 --- a/posthog/mcp/constants.py +++ b/posthog/mcp/constants.py @@ -75,6 +75,11 @@ class PostHogMCPAnalyticsProperty: SERVER_NAME = "$mcp_server_name" SERVER_VERSION = "$mcp_server_version" SESSION_ID = "$session_id" + # How the $session_id was resolved: "token" (self-encoded session token), + # "mcp" (transport Mcp-Session-Id), "derived" (from the identified principal + + # client, for stateless 2026-07-28 traffic), or "generated" (fresh). Lets + # downstream tell a correlated session from a per-request generated one. + SESSION_ID_SOURCE = "$mcp_session_id_source" SOURCE = "$mcp_source" TOOL_CATEGORY = "$mcp_tool_category" TOOL_DESCRIPTION = "$mcp_tool_description" diff --git a/posthog/mcp/session.py b/posthog/mcp/session.py index 902342ec..d2ddc5fb 100644 --- a/posthog/mcp/session.py +++ b/posthog/mcp/session.py @@ -12,9 +12,11 @@ from typing import Optional from .constants import INACTIVITY_TIMEOUT_IN_MINUTES +from ._derived_sessions import derive_session_id from ._ids import deterministic_prefixed_id, new_prefixed_id from ._internal import MCPAnalyticsData from .session_token import SessionTokenPayload +from .types import UserIdentity __all__ = ["derive_session_id_from_mcp_session"] @@ -34,10 +36,15 @@ async def resolve_session_id( mcp_session_id: Optional[str], *, token: Optional[SessionTokenPayload] = None, + identity: Optional[UserIdentity] = None, + client_name: Optional[str] = None, + client_version: 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. + Precedence: ``token`` > ``mcp_session_id`` > sticky-mcp > derived > generated. + ``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. @@ -47,6 +54,13 @@ async def resolve_session_id( for a request that didn't replay the token would merge unrelated clients under one ``$session_id``. A compliant client replays the header on every request, so a genuine token session never needs the fallback. + + ``identity`` enables the *derived* source: on stateless traffic with no token + and no MCP session header (SEP-2567 dropped the header on 2026-07-28), an + identified request derives a stable session from ``(distinct_id, client_name, + client_version)`` via the process-shared registry. Only taken when a + ``distinct_id`` is present — deriving from an anonymous key would merge + unrelated users under one ``$session_id``, so those fall through to generated. """ async with data.session_lock: now = datetime.now(timezone.utc) @@ -73,6 +87,17 @@ async def resolve_session_id( data.last_activity = now return data.session_id + # Derived: stateless traffic (no token, no MCP header) that is identified. + # The registry is process-shared, so per-request server instances correlate + # into one session per (distinct_id, client). Never derived anonymously. + if identity is not None and identity.distinct_id: + data.session_id = derive_session_id( + identity.distinct_id, client_name, client_version, now=now + ) + data.session_source = "derived" + data.last_activity = now + return data.session_id + # Memory fallback (single-owner transports like stdio). A leftover token # session must NOT leak to a credential-less request, so anything that # isn't already a generated session starts fresh; generated sessions diff --git a/posthog/test/mcp/test_derived_sessions.py b/posthog/test/mcp/test_derived_sessions.py new file mode 100644 index 00000000..d625b34b --- /dev/null +++ b/posthog/test/mcp/test_derived_sessions.py @@ -0,0 +1,149 @@ +"""Tests for derived sessions: stateless MCP traffic (mcp 2.x / SEP-2567 has no +`Mcp-Session-Id` header) gets a stable `$session_id` derived from the identified +principal + client, so an identified user's calls correlate into one session +instead of fragmenting into one-per-request. Pure-posthog, runs in both envs.""" + +from datetime import datetime, timedelta, timezone + +from parameterized import parameterized + +from posthog.mcp import _derived_sessions as ds +from posthog.mcp._derived_sessions import ( + DerivedSessionRegistry, + _MAX_ENTRIES, +) +from posthog.mcp.constants import INACTIVITY_TIMEOUT_IN_MINUTES + + +def _key(distinct_id="u1", client_name="cli", client_version="1.0"): + return (distinct_id, client_name, client_version) + + +def test_same_key_within_gap_reuses_session(): + reg = DerivedSessionRegistry() + now = datetime.now(timezone.utc) + first = reg.resolve(*_key(), now=now) + second = reg.resolve(*_key(), now=now + timedelta(minutes=1)) + assert first == second + assert first.startswith("ses_") + + +def test_gap_expiry_rolls_session(): + reg = DerivedSessionRegistry() + now = datetime.now(timezone.utc) + first = reg.resolve(*_key(), now=now) + later = now + timedelta(minutes=INACTIVITY_TIMEOUT_IN_MINUTES + 1) + second = reg.resolve(*_key(), now=later) + assert first != second + + +@parameterized.expand( + [ + ("different_distinct_id", _key(distinct_id="u2")), + ("different_client_name", _key(client_name="other")), + ("different_client_version", _key(client_version="2.0")), + ] +) +def test_different_key_gets_different_session(_name, other_key): + reg = DerivedSessionRegistry() + now = datetime.now(timezone.utc) + base = reg.resolve(*_key(), now=now) + other = reg.resolve(*other_key, now=now) + assert base != other + + +def test_lru_bound_respected(): + reg = DerivedSessionRegistry() + now = datetime.now(timezone.utc) + for i in range(_MAX_ENTRIES + 100): + reg.resolve(f"u{i}", "cli", "1.0", now=now) + assert reg.size() <= _MAX_ENTRIES + + +def test_idle_entries_evicted_beyond_two_timeouts(): + reg = DerivedSessionRegistry() + start = datetime.now(timezone.utc) + reg.resolve("idle", "cli", "1.0", now=start) + # A much later resolve for a different key triggers eviction of the idle one. + far = start + timedelta(minutes=2 * INACTIVITY_TIMEOUT_IN_MINUTES + 1) + reg.resolve("fresh", "cli", "1.0", now=far) + assert reg.size() == 1 + + +def test_concurrent_resolves_of_one_key_yield_one_session(): + import threading + + reg = DerivedSessionRegistry() + now = datetime.now(timezone.utc) + results = [] + barrier = threading.Barrier(8) + + def worker(): + barrier.wait() + results.append(reg.resolve(*_key(), now=now)) + + threads = [threading.Thread(target=worker) for _ in range(8)] + for t in threads: + t.start() + for t in threads: + t.join() + + assert len(set(results)) == 1 + + +def test_fork_reset_clears_registry(): + ds._DERIVED_REGISTRY.resolve("u", "cli", "1.0") + assert ds._DERIVED_REGISTRY.size() >= 1 + ds._reset_derived_registry_after_fork() + assert ds._DERIVED_REGISTRY.size() == 0 + + +@parameterized.expand( + [ + # (has_token, mcp_session_id, has_identity) -> expected source, in precedence order. + ("token_wins", True, "sess-1", True, "token"), + ("mcp_over_derived", False, "sess-1", True, "mcp"), + ("derived_when_identified", False, None, True, "derived"), + ("generated_when_anonymous", False, None, False, "generated"), + ] +) +async def test_resolve_session_id_precedence( + _name, has_token, mcp_session_id, has_identity, expected_source +): + from posthog.mcp._internal import MCPAnalyticsData + from posthog.mcp.session import resolve_session_id + from posthog.mcp.session_token import SessionTokenPayload + from posthog.mcp.types import MCPAnalyticsOptions, UserIdentity + + data = MCPAnalyticsData(options=MCPAnalyticsOptions()) + token = ( + SessionTokenPayload( + session_id="ses_from_token", client_name="c", client_version="1" + ) + if has_token + else None + ) + identity = UserIdentity(distinct_id="u1") if has_identity else None + + await resolve_session_id( + data, + mcp_session_id, + token=token, + identity=identity, + client_name="cli", + client_version="1.0", + ) + assert data.session_source == expected_source + + +async def test_derived_requires_distinct_id(): + from posthog.mcp._internal import MCPAnalyticsData + from posthog.mcp.session import resolve_session_id + from posthog.mcp.types import MCPAnalyticsOptions, UserIdentity + + data = MCPAnalyticsData(options=MCPAnalyticsOptions()) + # An identity with an empty distinct_id must NOT derive (would merge anon users). + await resolve_session_id( + data, None, identity=UserIdentity(distinct_id=""), client_name="cli" + ) + assert data.session_source == "generated" diff --git a/posthog/test/mcp/test_instrument_v2.py b/posthog/test/mcp/test_instrument_v2.py index fa1794b2..3b5d5ee1 100644 --- a/posthog/test/mcp/test_instrument_v2.py +++ b/posthog/test/mcp/test_instrument_v2.py @@ -164,6 +164,36 @@ def needs_input() -> InputRequiredResult: assert props["$mcp_result_type"] == "input_required" +@pytest.mark.parametrize( + "identify, expected_source", + [ + (UserIdentity(distinct_id="u1"), "derived"), + (None, "generated"), + ], +) +async def test_session_id_source_on_every_event(identify, expected_source): + # v2 traffic has no token and no Mcp-Session-Id header, so an identified + # request derives a stable session and an anonymous one gets a generated one. + # Every $mcp_* event (and $identify) carries the provenance. + server = make_server() + client = FakeClient() + instrument( + server, + client, + MCPAnalyticsOptions(identify=identify) if identify else None, + ) + + await drive(server, calls=(("add", {"a": 1, "b": 1}),)) + + names = ["$mcp_initialize", "$mcp_tools_list", "$mcp_tool_call"] + if identify: + names.append("$identify") + for name in names: + events = _events(client, name) + assert events, f"expected a {name} event" + assert events[0]["properties"]["$mcp_session_id_source"] == expected_source + + async def test_legacy_negotiation_on_v2_sdk_still_captures(): # A client that runs the classic `initialize` handshake against the v2 SDK # negotiates an older protocol: the _meta envelope is absent, so client info diff --git a/posthog/test/mcp/test_lowlevel.py b/posthog/test/mcp/test_lowlevel.py index b26f884b..4b01fd99 100644 --- a/posthog/test/mcp/test_lowlevel.py +++ b/posthog/test/mcp/test_lowlevel.py @@ -1,9 +1,11 @@ """End-to-end tests for the low-level mcp.server.Server adapter (Milestone 3).""" import mcp.types as mcp_types +import pytest from mcp.server.lowlevel import Server from posthog.mcp import instrument +from posthog.mcp.types import MCPAnalyticsOptions, UserIdentity from posthog.test.mcp._helpers import ( FakeClient, events_named as _events, @@ -133,3 +135,37 @@ async def test_initialize_emitted_once(): assert len(_events(client, "$mcp_initialize")) == 1 assert len(_events(client, "$mcp_tool_call")) == 2 + + +@pytest.mark.parametrize( + "identify, expected_source", + [ + (UserIdentity(distinct_id="u1"), "derived"), + (None, "generated"), + ], +) +async def test_session_id_source_on_every_v1_event(identify, expected_source): + # A low-level (stdio-style) call with no Mcp-Session-Id header and no token: + # identified -> derived, anonymous -> generated. Provenance rides every event. + server = make_server() + client = FakeClient() + instrument( + server, client, MCPAnalyticsOptions(identify=identify) if identify else None + ) + + await server.request_handlers[mcp_types.ListToolsRequest]( + mcp_types.ListToolsRequest(method="tools/list") + ) + handler = server.request_handlers[mcp_types.CallToolRequest] + await handler( + _call_request("echo", {"msg": "hi", "context": "listing then calling"}) + ) + await _flush() + + names = ["$mcp_initialize", "$mcp_tools_list", "$mcp_tool_call"] + if identify: + names.append("$identify") + for name in names: + events = _events(client, name) + assert events, f"expected a {name} event" + assert events[0]["properties"]["$mcp_session_id_source"] == expected_source From 217175c740712671e3677744140409357de1203c Mon Sep 17 00:00:00 2001 From: Lucas Faria Date: Tue, 4 Aug 2026 20:15:28 -0300 Subject: [PATCH 4/4] mcp: fix v2 error-result capture and legacy-era session correlation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two review fixes on the v2 adapter. (1) Drop the _ForceErrorFlag wrapper: for dict-shaped error results (the common v2 wire case) it masked dict-ness, so _to_jsonable returned the raw wrapper into the event's $mcp_response — unserializable in production. is_tool_result_error now reads the 2.x models' snake_case is_error directly and results pass through unwrapped. (2) Thread ctx.session_id into session resolution: legacy-era clients on the v2 SDK still carry a transport session id, which now resolves with "mcp" provenance instead of falling through to derived/generated. Tested: scripts/validate-mcp-matrix.sh — mcp 1.x PASS, mcp 2.x PASS (135 passed / 13 skipped); ruff 0.11.12 format+check clean. New regression tests: JSON-serializability of captured error responses, and stable "mcp" session provenance for a stubbed legacy-era ctx. Generated-By: PostHog Code Task-Id: c17ecd78-bb05-4b3c-b621-458abba9c6aa --- posthog/mcp/_instrument_v2.py | 52 ++++++++------------------ posthog/mcp/_instrumentation.py | 8 +++- posthog/test/mcp/test_instrument_v2.py | 38 +++++++++++++++++++ 3 files changed, 60 insertions(+), 38 deletions(-) diff --git a/posthog/mcp/_instrument_v2.py b/posthog/mcp/_instrument_v2.py index f5a68ea2..de4485bb 100644 --- a/posthog/mcp/_instrument_v2.py +++ b/posthog/mcp/_instrument_v2.py @@ -144,17 +144,22 @@ async def _prepare( ) -> Tuple[str, str, Optional[str], Optional[str], Optional[str]]: """Resolve session (emitting identify + lazy initialize) for this request. - v2 traffic has no self-encoded token and no ``Mcp-Session-Id`` header by - construction (SEP-2567 removed it), so ``mcp_session_id``/``token`` are always - None here; sessions come out ``derived`` (Stage C) when identified, - ``generated`` otherwise. Returns the session id, its provenance source, and the - resolved client/protocol info so the caller stamps them on the captured event.""" + 2026-07-28 traffic has no self-encoded token and no ``Mcp-Session-Id`` header + (SEP-2567 removed it), so ``ctx.session_id`` is None and sessions come out + ``derived`` (Stage C) when identified, ``generated`` otherwise. A legacy-era + client on this SDK still gets a transport session id, which ``ctx.session_id`` + surfaces — pass it through so mixed fleets keep MCP-session correlation. + Returns the session id, its provenance source, and the resolved + client/protocol info so the caller stamps them on the captured event.""" client_name, client_version = _client_info(ctx) protocol_version = _protocol_version(ctx) - extra = {"session_id": None} + mcp_session_id = getattr(ctx, "session_id", None) + if not isinstance(mcp_session_id, str) or not mcp_session_id: + mcp_session_id = None + extra = {"session_id": mcp_session_id} session_id, session_id_source = await prepare_request( data, - mcp_session_id=None, + mcp_session_id=mcp_session_id, client_name=client_name, client_version=client_version, protocol_version=protocol_version, @@ -177,19 +182,6 @@ def _result_type(result: Any) -> Optional[str]: return None -def _is_error_result(result: Any) -> bool: - """A v2 tool result signals a tool error via ``isError``/``is_error`` — but an - ``input_required`` interim result (MRTR) is NOT an error, just a continuation.""" - if _result_type(result) == "input_required": - return False - if isinstance(result, dict): - return result.get("isError") is True - return ( - getattr(result, "is_error", None) is True - or getattr(result, "isError", None) is True - ) - - async def _on_tool_call(data: MCPAnalyticsData, ctx: Any, call_next: Any) -> Any: params = ctx.params if isinstance(ctx.params, dict) else {} # The tools/call wire always carries a string name; default defensively so a @@ -228,12 +220,15 @@ async def _on_tool_call(data: MCPAnalyticsData, ctx: Any, call_next: Any) -> Any raise duration_ms = (time.monotonic() - start) * 1000 + # The result passes through unwrapped: record_tool_call's error detection + # handles both the wire dict's `isError` and a 2.x model's `is_error`, and an + # MRTR `input_required` result carries neither, so it lands as a non-error. await record_tool_call( data, session_id, name=name, arguments=arguments, - result=_ForceErrorFlag(result) if _is_error_result(result) else result, + result=result, duration_ms=duration_ms, client_name=client_name, client_version=client_version, @@ -244,21 +239,6 @@ async def _on_tool_call(data: MCPAnalyticsData, ctx: Any, call_next: Any) -> Any return result -class _ForceErrorFlag: - """Adapts a v2 wire result so ``record_tool_call``'s ``isError`` detection - (which reads ``.isError``/``["isError"]``) fires for the value we already - classified as a tool error — including a model-shaped result whose flag is - ``is_error``. Wraps, never mutates, the result the tool actually returned.""" - - isError = True - - def __init__(self, wrapped: Any) -> None: - self._wrapped = wrapped - - def __getattr__(self, item: str) -> Any: - return getattr(self._wrapped, item) - - async def _on_tools_list(data: MCPAnalyticsData, ctx: Any, call_next: Any) -> Any: request = {"method": "tools/list", "params": {}} ( diff --git a/posthog/mcp/_instrumentation.py b/posthog/mcp/_instrumentation.py index 8d081fff..2bd1b2f6 100644 --- a/posthog/mcp/_instrumentation.py +++ b/posthog/mcp/_instrumentation.py @@ -146,10 +146,14 @@ 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. + mcp 2.x pydantic models expose the same flag as snake_case ``is_error``.""" if isinstance(result, dict): return result.get("isError") is True - return getattr(result, "isError", None) is True + return ( + getattr(result, "isError", None) is True + or getattr(result, "is_error", None) is True + ) def build_tool_call_request( diff --git a/posthog/test/mcp/test_instrument_v2.py b/posthog/test/mcp/test_instrument_v2.py index 3b5d5ee1..5a96f587 100644 --- a/posthog/test/mcp/test_instrument_v2.py +++ b/posthog/test/mcp/test_instrument_v2.py @@ -92,6 +92,44 @@ async def test_tool_error_result_captured_as_error(): assert calls and calls[0]["properties"]["$mcp_is_error"] is True exceptions = _events(client, "$exception") assert exceptions, "an isError tool result should also emit $exception" + # The captured response must be the tool's own JSON-shaped result, not an + # adapter wrapper: FakeClient never JSON-encodes, so guard serializability + # here or a wrapper leaks a repr into production payloads unnoticed. + import json + + response = calls[0]["properties"].get("$mcp_response") + json.dumps(response) + + +async def test_legacy_transport_session_id_resolves_mcp_source(): + # A legacy-era client on the v2 SDK still gets a transport session id, which + # the middleware ctx surfaces as `ctx.session_id`. It must resolve with "mcp" + # provenance (deterministic per MCP session), not fall through to generated. + from posthog.mcp._instrument_v2 import _prepare + from posthog.mcp._internal import MCPAnalyticsData + from posthog.mcp._sink import McpEventSink + from posthog.mcp.session import new_session_id + + class StubCtx: + method = "tools/call" + params = {"name": "add", "arguments": {}} + meta = None + protocol_version = "2025-06-18" + session_id = "legacy-transport-session" + + client = FakeClient() + data = MCPAnalyticsData( + options=MCPAnalyticsOptions(), + sink=McpEventSink(client), + session_id=new_session_id(), + ) + + request = {"method": "tools/call", "params": {"name": "add", "arguments": {}}} + session_a, source, *_ = await _prepare(data, StubCtx(), request) + session_b, _, *_ = await _prepare(data, StubCtx(), request) + + assert source == "mcp" + assert session_a == session_b, "same MCP session must map to one $session_id" async def test_tools_list_captured_with_names():