From aeea8a3898660bdc5ce821a694d3cfc0d6390089 Mon Sep 17 00:00:00 2001 From: Lucas Faria Date: Mon, 7 Sep 2026 14:22:07 -0300 Subject: [PATCH 1/2] feat(mcp): capture model identifiers Capture model identifiers from recognized client metadata, with an SDK-owned self-report fallback for other clients. Preserve source provenance and fail closed when the application owns the field.\n\nVerify both MCP Python SDK 1.x and 2.x, including a 2026-07-28 wire-level call. --- .sampo/changesets/bold-king-mielikki.md | 5 + posthog/mcp/README.md | 62 ++++++++++ posthog/mcp/__init__.py | 4 + posthog/mcp/_capture.py | 2 + posthog/mcp/_instrument_fastmcp.py | 21 +++- posthog/mcp/_instrument_lowlevel.py | 12 +- posthog/mcp/_instrument_v2.py | 36 +++++- posthog/mcp/_instrumentation.py | 86 +++++++++++-- posthog/mcp/_internal.py | 3 + posthog/mcp/_model_parameters.py | 142 ++++++++++++++++++++++ posthog/mcp/_posthog_events.py | 4 + posthog/mcp/_sanitization.py | 19 ++- posthog/mcp/_truncation.py | 1 + posthog/mcp/constants.py | 9 ++ posthog/mcp/posthog_mcp.py | 126 +++++++++++++++++-- posthog/mcp/types.py | 20 ++- posthog/test/mcp/_helpers_v2.py | 9 +- posthog/test/mcp/test_fastmcp.py | 42 +++++++ posthog/test/mcp/test_model_parameters.py | 95 +++++++++++++++ posthog/test/mcp/test_pipeline.py | 25 ++++ posthog/test/mcp/test_posthog_mcp.py | 80 +++++++++++- posthog/test/mcp/test_truncation.py | 9 +- posthog/test/mcp/test_v2_wire_dual_era.py | 27 +++- references/public_api_snapshot.txt | 22 +++- 24 files changed, 817 insertions(+), 44 deletions(-) create mode 100644 .sampo/changesets/bold-king-mielikki.md create mode 100644 posthog/mcp/_model_parameters.py create mode 100644 posthog/test/mcp/test_model_parameters.py diff --git a/.sampo/changesets/bold-king-mielikki.md b/.sampo/changesets/bold-king-mielikki.md new file mode 100644 index 000000000..2783644a0 --- /dev/null +++ b/.sampo/changesets/bold-king-mielikki.md @@ -0,0 +1,5 @@ +--- +pypi/posthog: minor +--- + +Capture MCP model identifiers from client metadata or an SDK-owned self-report field. diff --git a/posthog/mcp/README.md b/posthog/mcp/README.md index a8dcd1020..fb6c12e31 100644 --- a/posthog/mcp/README.md +++ b/posthog/mcp/README.md @@ -21,6 +21,68 @@ Request headers use the same identity and package version, so SDK Health can com Because `$lib` is a client-level identity, `instrument()` relabels every event sent by the client passed to it. Use a client dedicated to MCP analytics if the application also captures unrelated events. +## Capture the calling model + +Model capture is off by default. Enable it for an instrumented MCP Python SDK 1.x or +2.x server: + +```python +from posthog.mcp import MCPAnalyticsOptions, instrument + +analytics = instrument( + server, + posthog, + MCPAnalyticsOptions(capture_model=True), +) +``` + +The SDK records the best model identifier visible to the server as +`$mcp_llm_model`. Recognized client metadata wins and sets +`$mcp_llm_model_source` to `client_metadata`. Otherwise, the SDK adds a required +`llm_model` string to each compatible tool schema and records the agent's answer +with source `self_reported`. + +The recognized metadata path is Codex's `x-codex-turn-metadata.model` field in +request `_meta`. Other clients, including Claude Code, use the self-report path +until they expose a stable model field. Missing, blank, and `unknown` values are +not recorded. + +MCP does not standardize or attest model identity. Both sources are unverified. +Use them to compare tool behavior across models, not for billing or access +control. + +Model self-reporting only runs when PostHog can prove it owns the injected field. +If a tool already declares `llm_model`, or uses a root `$ref`, `oneOf`, `allOf`, or +`anyOf` schema, PostHog leaves the schema and argument untouched. Client metadata +can still be captured in those cases. + +For a custom dispatcher, use the same option on `PostHogMCP` and pass request +metadata through explicitly: + +```python +from posthog.mcp import PostHogMCP + +posthog = PostHogMCP("phc_...", capture_model=True) +tools = posthog.prepare_tool_list(server_tools) +original_tool = next(tool for tool in server_tools if tool["name"] == tool_name) +call = posthog.prepare_tool_call( + tool_name, + raw_args, + request_meta=request.get("params", {}).get("_meta"), + original_tool=original_tool, +) +result = dispatch(tool_name, call.args) +posthog.capture_tool_call( + tool_name, + llm_model=call.llm_model, + llm_model_source=call.llm_model_source, +) +``` + +Passing `original_tool` keeps ownership accurate when `tools/list` and +`tools/call` reach different server replicas. A persistent single-process +dispatcher can omit it after calling `prepare_tool_list()`. + ## Stateless / multi-pod servers A stateless MCP server issues no session id, so `$session_id` fragments across pods diff --git a/posthog/mcp/__init__.py b/posthog/mcp/__init__.py index 3d9a58223..633459300 100644 --- a/posthog/mcp/__init__.py +++ b/posthog/mcp/__init__.py @@ -73,6 +73,8 @@ from .types import ( CaptureEventData, MCPAnalyticsContextOptions, + MCPAnalyticsModelOptions, + MCPAnalyticsModelSource, MCPAnalyticsOptions, PreparedToolCall, UserIdentity, @@ -85,6 +87,8 @@ "PostHogMCP", "MCPAnalyticsOptions", "MCPAnalyticsContextOptions", + "MCPAnalyticsModelOptions", + "MCPAnalyticsModelSource", "UserIdentity", "CaptureEventData", "PreparedToolCall", diff --git a/posthog/mcp/_capture.py b/posthog/mcp/_capture.py index 4544ce5f3..424501527 100644 --- a/posthog/mcp/_capture.py +++ b/posthog/mcp/_capture.py @@ -63,6 +63,8 @@ def capture_event( "response": event_input.get("response"), "user_intent": event_input.get("user_intent"), "user_intent_source": event_input.get("user_intent_source"), + "llm_model": event_input.get("llm_model"), + "llm_model_source": event_input.get("llm_model_source"), "is_error": event_input.get("is_error"), "error": event_input.get("error"), "error_type": event_input.get("error_type"), diff --git a/posthog/mcp/_instrument_fastmcp.py b/posthog/mcp/_instrument_fastmcp.py index e4d8a84d7..f61f3c361 100644 --- a/posthog/mcp/_instrument_fastmcp.py +++ b/posthog/mcp/_instrument_fastmcp.py @@ -39,6 +39,11 @@ start_tools_list_lifecycle, ) from ._internal import MCPAnalyticsData +from ._model_parameters import ( + can_inject_model_parameter, + is_capture_model_enabled, + request_meta_from_context, +) from ._output_instructions import mirror_instructions_into_structured_content from .logger import log from .tools import get_more_tools_result_text, resolve_missing_capability_tool_name @@ -90,6 +95,8 @@ async def wrapped( data, name=name, arguments=arguments, + request_meta=request_meta_from_context(_tool_call_request_context(context)), + allow_self_reported_model=_analytics_owns_model(server, data, name), mcp_session_id=mcp_session_id, token=token, client_name=client_name, @@ -119,6 +126,8 @@ async def wrapped( server, name, "conversation_id" ): strip_keys.add("conversation_id") + if _analytics_owns_model(server, data, name): + strip_keys.add("llm_model") if strip_keys: call_arguments = { k: v for k, v in arguments.items() if k not in strip_keys @@ -251,7 +260,7 @@ async def list_handler(req: Any) -> Any: if data.options.report_missing: missing_name = resolve_missing_capability_tool_name(data.options) if not any(t.name == missing_name for t in tools): - append_get_more_tools(result, missing_name) + append_get_more_tools(result, missing_name, data) names.append(missing_name) await lifecycle.record_result( @@ -310,6 +319,16 @@ def _tool_owns_context(server: Any, name: str) -> bool: return _tool_owns_param(server, name, "context") +def _analytics_owns_model(server: Any, data: MCPAnalyticsData, name: str) -> bool: + if not is_capture_model_enabled(data.options.capture_model): + return False + try: + tool = server._tool_manager.get_tool(name) + return can_inject_model_parameter(getattr(tool, "parameters", None)) + except Exception: # noqa: BLE001 - model analytics must never break dispatch + return data.tool_model_parameter_injected.get(name, False) + + def _tool_call_request_context(context: Any) -> Any: """The request context behind a FastMCP ``Context``, or ``None``. diff --git a/posthog/mcp/_instrument_lowlevel.py b/posthog/mcp/_instrument_lowlevel.py index cc1432554..d16c698e0 100644 --- a/posthog/mcp/_instrument_lowlevel.py +++ b/posthog/mcp/_instrument_lowlevel.py @@ -34,6 +34,7 @@ start_tools_list_lifecycle, ) from ._internal import MCPAnalyticsData +from ._model_parameters import request_meta_from_context from ._output_instructions import mirror_instructions_into_structured_content from .logger import log from .tools import get_more_tools_result_text, resolve_missing_capability_tool_name @@ -97,6 +98,10 @@ async def handler(req: Any) -> Any: data, name=name, arguments=arguments, + request_meta=request_meta_from_context(_request_context(server)), + allow_self_reported_model=data.tool_model_parameter_injected.get( + name, False + ), mcp_session_id=mcp_session_id, token=token, client_name=client_name, @@ -127,7 +132,10 @@ async def handler(req: Any) -> Any: # tools/list and across stateless per-request server instances. if strip_injected and req.params.arguments: owned = await _tool_owned_injected_keys(high_level, name) - for key in ("context", "conversation_id"): + injected_keys = ["context", "conversation_id"] + if data.tool_model_parameter_injected.get(name, False): + injected_keys.append("llm_model") + for key in injected_keys: if key not in owned: req.params.arguments.pop(key, None) @@ -289,7 +297,7 @@ async def handler(req: Any) -> Any: if data.options.report_missing: missing_name = resolve_missing_capability_tool_name(data.options) if not any(t.name == missing_name for t in tools): - append_get_more_tools(result, missing_name) + append_get_more_tools(result, missing_name, data) names.append(missing_name) await lifecycle.record_result( diff --git a/posthog/mcp/_instrument_v2.py b/posthog/mcp/_instrument_v2.py index 437015bee..c740ae273 100644 --- a/posthog/mcp/_instrument_v2.py +++ b/posthog/mcp/_instrument_v2.py @@ -47,6 +47,11 @@ start_tools_list_lifecycle, ) from ._internal import MCPAnalyticsData +from ._model_parameters import ( + can_inject_model_parameter, + is_capture_model_enabled, + request_meta_from_context, +) from ._output_instructions import mirror_instructions_into_structured_content from .logger import log from .request_headers import get_request_headers @@ -222,6 +227,18 @@ def _tool_owns_param_v2(high_level: Any, name: str, param: str) -> bool: return param in _tool_own_properties_v2(high_level, name) +def _analytics_owns_model_v2( + high_level: Any, data: MCPAnalyticsData, name: str +) -> bool: + if not is_capture_model_enabled(data.options.capture_model): + return False + try: + tool = high_level._tool_manager.get_tool(name) + return can_inject_model_parameter(getattr(tool, "parameters", None)) + except Exception: # noqa: BLE001 - model analytics must never break dispatch + return data.tool_model_parameter_injected.get(name, False) + + # --- high-level: ToolManager.call_tool seam -------------------------------------- @@ -249,6 +266,8 @@ async def wrapped( data, name=name, arguments=arguments, + request_meta=request_meta_from_context(ctx), + allow_self_reported_model=_analytics_owns_model_v2(server, data, name), mcp_session_id=mcp_session_id, token=token, client_name=client_name, @@ -281,6 +300,8 @@ async def wrapped( and "conversation_id" not in own_properties ): strip_keys.add("conversation_id") + if _analytics_owns_model_v2(server, data, name): + strip_keys.add("llm_model") if strip_keys: call_arguments = { k: v for k, v in arguments.items() if k not in strip_keys @@ -389,6 +410,10 @@ async def handler(ctx: Any, params: Any) -> Any: data, name=name, arguments=arguments, + request_meta=request_meta_from_context(ctx), + allow_self_reported_model=data.tool_model_parameter_injected.get( + name, False + ), mcp_session_id=mcp_session_id, token=token, client_name=client_name, @@ -504,7 +529,7 @@ async def handler(ctx: Any, params: Any) -> Any: if data.options.report_missing: missing_name = resolve_missing_capability_tool_name(data.options) if not any(t.name == missing_name for t in tools): - _append_get_more_tools_v2(result, missing_name) + _append_get_more_tools_v2(result, missing_name, data) names.append(missing_name) await lifecycle.record_result( @@ -520,7 +545,7 @@ async def handler(ctx: Any, params: Any) -> Any: _replace_handler(server, _LIST_METHOD, handler, entry.params_type) -def _append_get_more_tools_v2(result: Any, name: str) -> None: +def _append_get_more_tools_v2(result: Any, name: str, data: MCPAnalyticsData) -> None: descriptor = build_report_missing_descriptor(name) tool = mcp_types.Tool( name=descriptor["name"], @@ -528,6 +553,13 @@ def _append_get_more_tools_v2(result: Any, name: str) -> None: input_schema=descriptor["inputSchema"], annotations=descriptor["annotations"], ) + mutate_tool_schema( + data, + tool, + schema_attribute="input_schema", + owns_context=True, + context_required=True, + ) tools_list = getattr(result, "tools", None) if isinstance(tools_list, list): tools_list.append(tool) diff --git a/posthog/mcp/_instrumentation.py b/posthog/mcp/_instrumentation.py index b9c052b0f..328fc597f 100644 --- a/posthog/mcp/_instrumentation.py +++ b/posthog/mcp/_instrumentation.py @@ -29,6 +29,12 @@ 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 ._model_parameters import ( + add_model_parameter_to_schema, + get_model_description, + is_capture_model_enabled, + resolve_model, +) from ._output_instructions import add_instructions_to_output_schema from .logger import log, warn from .request_headers import get_request @@ -391,6 +397,8 @@ class ToolCallLifecycle: data: MCPAnalyticsData name: str arguments: Optional[Dict[str, Any]] + request_meta: Optional[Dict[str, Any]] + allow_self_reported_model: bool request: Dict[str, Any] extra: Dict[str, Any] mcp_session_id: Optional[str] @@ -432,6 +440,8 @@ async def record_missing_capability(self) -> None: tool_name=self.missing_name, context=(self.arguments or {}).get("context"), arguments=self.arguments, + request_meta=self.request_meta, + allow_self_reported_model=True, client_name=self.client_name, client_version=self.client_version, protocol_version=self.protocol_version, @@ -448,6 +458,8 @@ async def record_error(self, error: Any, duration_ms: float) -> None: session_id, name=self.name, arguments=self.arguments, + request_meta=self.request_meta, + allow_self_reported_model=self.allow_self_reported_model, error=error, duration_ms=duration_ms, client_name=self.client_name, @@ -469,6 +481,8 @@ async def record_result( session_id, name=self.name, arguments=self.arguments, + request_meta=self.request_meta, + allow_self_reported_model=self.allow_self_reported_model, result=result, duration_ms=duration_ms, client_name=self.client_name, @@ -484,6 +498,8 @@ def start_tool_call_lifecycle( *, name: str, arguments: Optional[Dict[str, Any]], + request_meta: Optional[Dict[str, Any]], + allow_self_reported_model: bool, mcp_session_id: Optional[str], token: Optional[SessionTokenPayload], client_name: Optional[str], @@ -500,6 +516,8 @@ def start_tool_call_lifecycle( data=data, name=name, arguments=arguments, + request_meta=request_meta, + allow_self_reported_model=allow_self_reported_model, request=build_tool_call_request(name, arguments), extra=extra, mcp_session_id=mcp_session_id, @@ -519,6 +537,8 @@ async def record_tool_call( *, name: str, arguments: Optional[Dict[str, Any]], + request_meta: Optional[Dict[str, Any]] = None, + allow_self_reported_model: bool = False, result: Any = None, error: Any = None, duration_ms: Optional[float] = None, @@ -538,7 +558,9 @@ async def record_tool_call( "resource_name": name, "tool_description": data.tool_descriptions.get(name), "tool_category": data.tool_categories.get(name), - "parameters": build_captured_mcp_parameters(request), + "parameters": build_captured_mcp_parameters( + request, strip_llm_model=allow_self_reported_model + ), "duration": duration_ms, "client_name": client_name, "client_version": client_version, @@ -547,6 +569,15 @@ async def record_tool_call( "is_error": False, } set_event_intent(event, await resolve_tool_call_intent(data, request, extra)) + if is_capture_model_enabled(data.options.capture_model): + model, source = resolve_model( + request_meta, + arguments, + allow_self_reported=allow_self_reported_model, + ) + if model: + event["llm_model"] = model + event["llm_model_source"] = source if error is not None: event["is_error"] = True @@ -574,7 +605,7 @@ def extract_tools(result: Any) -> list: return list(getattr(root, "tools", []) or []) -def append_get_more_tools(result: Any, name: str) -> None: +def append_get_more_tools(result: Any, name: str, data: MCPAnalyticsData) -> None: """Append the get_more_tools virtual tool to the real ListToolsResult.tools list.""" import mcp.types as mcp_types @@ -590,6 +621,13 @@ def append_get_more_tools(result: Any, name: str) -> None: root = getattr(result, "root", result) tools_list = getattr(root, "tools", None) if isinstance(tools_list, list): + mutate_tool_schema( + data, + tool, + schema_attribute="inputSchema", + owns_context=True, + context_required=True, + ) tools_list.append(tool) @@ -630,19 +668,38 @@ def mutate_tool_schema( ownership decision. Those are the parts that differ across MCP generations; context/conversation mutation and output-channel bookkeeping do not. """ - if tool.name == GET_MORE_TOOLS_NAME: - return schema = getattr(tool, schema_attribute, None) original_schema = schema - if is_context_enabled(data.options.context) and not owns_context: + if ( + tool.name != GET_MORE_TOOLS_NAME + and is_context_enabled(data.options.context) + and not owns_context + ): schema = add_context_parameter_to_schema( schema, tool.name, get_context_description(data.options.context), required=context_required, ) - if data.options.enable_conversation_id and not schema_has_param( - schema, "conversation_id" + if is_capture_model_enabled(data.options.capture_model): + model_was_injected = data.tool_model_parameter_injected.get(tool.name, False) + app_owns_model = ( + schema_has_param(schema, "llm_model") and not model_was_injected + ) + if not app_owns_model and not schema_has_param(schema, "llm_model"): + schema = add_model_parameter_to_schema( + schema, + tool.name, + get_model_description(data.options.capture_model), + required=context_required, + ) + data.tool_model_parameter_injected[tool.name] = ( + not app_owns_model and schema_has_param(schema, "llm_model") + ) + if ( + tool.name != GET_MORE_TOOLS_NAME + and data.options.enable_conversation_id + and not schema_has_param(schema, "conversation_id") ): schema = add_conversation_id_to_schema(schema, tool.name) if schema is not original_schema: @@ -770,6 +827,8 @@ async def record_missing_capability( tool_name: str, context: Optional[str], arguments: Optional[Dict[str, Any]], + request_meta: Optional[Dict[str, Any]] = None, + allow_self_reported_model: bool = False, client_name: Optional[str] = None, client_version: Optional[str] = None, protocol_version: Optional[str] = None, @@ -783,7 +842,9 @@ async def record_missing_capability( "event_type": MCPAnalyticsEventType.MCP_MISSING_CAPABILITY, "session_id": session_id, "resource_name": tool_name, - "parameters": build_captured_mcp_parameters(request), + "parameters": build_captured_mcp_parameters( + request, strip_llm_model=allow_self_reported_model + ), "client_name": client_name, "client_version": client_version, "protocol_version": protocol_version, @@ -791,6 +852,15 @@ async def record_missing_capability( if isinstance(context, str) and context.strip(): event["user_intent"] = context.strip() event["user_intent_source"] = "context_parameter" + if is_capture_model_enabled(data.options.capture_model): + model, source = resolve_model( + request_meta, + arguments, + allow_self_reported=allow_self_reported_model, + ) + if model: + event["llm_model"] = model + event["llm_model_source"] = source await _apply_event_properties(data, event, request, extra) stamp_transport_identity(event, extra) fire_and_forget(capture_event(data, event), data) diff --git a/posthog/mcp/_internal.py b/posthog/mcp/_internal.py index a91a6935f..34e547535 100644 --- a/posthog/mcp/_internal.py +++ b/posthog/mcp/_internal.py @@ -70,6 +70,9 @@ class MCPAnalyticsData: identified_sessions: IdentityCache = field(default_factory=IdentityCache) tool_categories: Dict[str, str] = field(default_factory=dict) tool_descriptions: Dict[str, str] = field(default_factory=dict) + # True only when PostHog added llm_model to this tool's advertised schema. + # Missing/False fails closed so an application-owned field is never read or stripped. + tool_model_parameter_injected: Dict[str, bool] = field(default_factory=dict) # Which tools got `_mcp_instructions` declared on their advertised output # schema at tools/list. Only those may be mirrored into on a call — writing # an undeclared key fails the customer's whole result under diff --git a/posthog/mcp/_model_parameters.py b/posthog/mcp/_model_parameters.py new file mode 100644 index 000000000..b3a0c3942 --- /dev/null +++ b/posthog/mcp/_model_parameters.py @@ -0,0 +1,142 @@ +"""Capture the calling model from MCP request metadata or an injected argument. + +MCP does not standardize or attest model identity. Codex currently exposes its +model in request ``_meta``; other clients can use the SDK-injected +``llm_model`` argument as an explicitly lower-confidence fallback. +""" + +from __future__ import annotations + +import copy +from typing import Any, Dict, Optional + +from .constants import DEFAULT_MODEL_PARAMETER_DESCRIPTION +from .logger import log +from .types import MCPAnalyticsModelOptions, MCPAnalyticsModelSource + +_CODEX_TURN_METADATA_KEY = "x-codex-turn-metadata" + + +def is_capture_model_enabled( + capture_model: object, +) -> bool: + return capture_model is True or isinstance(capture_model, MCPAnalyticsModelOptions) + + +def get_model_description(capture_model: object) -> Optional[str]: + if isinstance(capture_model, MCPAnalyticsModelOptions): + return capture_model.description + return None + + +def add_model_parameter_to_schema( + input_schema: Optional[Dict[str, Any]], + tool_name: str = "unknown", + description_override: Optional[str] = None, + required: bool = True, +) -> Optional[Dict[str, Any]]: + """Return a copied schema with an SDK-owned ``llm_model`` field. + + Existing application fields and complex schemas fail closed: the SDK must + never overwrite or later strip a value that belongs to the tool itself. + """ + schema = input_schema + if ( + schema + and isinstance(schema.get("properties"), dict) + and "llm_model" in schema["properties"] + ): + log( + f"WARN: Tool \"{tool_name}\" already has 'llm_model' parameter. " + "Skipping model injection." + ) + return schema + + if schema and any(schema.get(key) for key in ("$ref", "oneOf", "allOf", "anyOf")): + log( + f'WARN: Tool "{tool_name}" has complex schema ' + "($ref/oneOf/allOf/anyOf). Skipping model injection." + ) + return schema + + if not schema: + schema = {"type": "object", "properties": {}, "required": []} + + schema = copy.deepcopy(schema) + if not isinstance(schema.get("properties"), dict): + schema["properties"] = {} + if schema.get("additionalProperties") is False: + schema.pop("additionalProperties", None) + + schema["properties"]["llm_model"] = { + "type": "string", + "description": description_override or DEFAULT_MODEL_PARAMETER_DESCRIPTION, + } + if required: + required_list = schema.get("required") + if isinstance(required_list, list): + if "llm_model" not in required_list: + required_list.append("llm_model") + else: + schema["required"] = ["llm_model"] + return schema + + +def can_inject_model_parameter(input_schema: Any) -> bool: + if not isinstance(input_schema, dict): + return True + properties = input_schema.get("properties") + if isinstance(properties, dict) and "llm_model" in properties: + return False + return not any(input_schema.get(key) for key in ("$ref", "oneOf", "allOf", "anyOf")) + + +def request_meta_from_context(context: Any) -> Optional[Dict[str, Any]]: + """Read request ``_meta`` from either MCP Python SDK generation.""" + meta = getattr(context, "meta", None) + if isinstance(meta, dict): + return meta + model_dump = getattr(meta, "model_dump", None) + if callable(model_dump): + try: + dumped = model_dump(mode="python", by_alias=True) + return dumped if isinstance(dumped, dict) else None + except Exception: # noqa: BLE001 - metadata must never break a tool call + return None + dict_method = getattr(meta, "dict", None) + if callable(dict_method): + try: + dumped = dict_method(by_alias=True) + return dumped if isinstance(dumped, dict) else None + except Exception: # noqa: BLE001 - Pydantic v1 compatibility + return None + return None + + +def resolve_model( + request_meta: Optional[Dict[str, Any]], + arguments: Optional[Dict[str, Any]], + *, + allow_self_reported: bool, +) -> tuple[Optional[str], Optional[MCPAnalyticsModelSource]]: + """Resolve the strongest model identity visible to the MCP server.""" + codex_metadata = (request_meta or {}).get(_CODEX_TURN_METADATA_KEY) + if isinstance(codex_metadata, dict) and "model" in codex_metadata: + model = normalize_model(codex_metadata.get("model")) + if model: + return model, "client_metadata" + + if allow_self_reported: + model = normalize_model((arguments or {}).get("llm_model")) + if model: + return model, "self_reported" + return None, None + + +def normalize_model(model: Any) -> Optional[str]: + if not isinstance(model, str): + return None + normalized = model.strip() + if not normalized or normalized.lower() == "unknown": + return None + return normalized diff --git a/posthog/mcp/_posthog_events.py b/posthog/mcp/_posthog_events.py index c9d9d0b85..689746e84 100644 --- a/posthog/mcp/_posthog_events.py +++ b/posthog/mcp/_posthog_events.py @@ -143,6 +143,10 @@ def _add_common_properties(event: Event, properties: Dict[str, Any]) -> None: properties[_P.INTENT] = event["user_intent"] if event.get("user_intent_source"): properties[_P.INTENT_SOURCE] = event["user_intent_source"] + if event.get("llm_model"): + properties[_P.LLM_MODEL] = event["llm_model"] + if event.get("llm_model_source"): + properties[_P.LLM_MODEL_SOURCE] = event["llm_model_source"] if event.get("is_error") is not None: properties[_P.IS_ERROR] = event["is_error"] if event.get("is_error"): diff --git a/posthog/mcp/_sanitization.py b/posthog/mcp/_sanitization.py index fd547f418..363d4c7ad 100644 --- a/posthog/mcp/_sanitization.py +++ b/posthog/mcp/_sanitization.py @@ -110,6 +110,9 @@ def sanitize_event(event: Dict[str, Any]) -> Dict[str, Any]: if result.get("user_intent") is not None: result["user_intent"] = sanitize_captured_value(result["user_intent"]) + if result.get("llm_model") is not None: + result["llm_model"] = sanitize_captured_value(result["llm_model"]) + # An exception message is free text a server wrote, and it reaches PostHog # on the $exception sibling and — since it is also surfaced as # $mcp_error_message — on the primary event, so run it through the same @@ -205,7 +208,9 @@ def _sanitize_resource_block(block: Dict[str, Any]) -> Any: return sanitize_captured_value(block) -def build_captured_mcp_parameters(request: Any) -> Dict[str, Any]: +def build_captured_mcp_parameters( + request: Any, *, strip_llm_model: bool = False +) -> Dict[str, Any]: """Build the sanitized ``$mcp_parameters`` payload from a request, stripping the injected ``context`` argument before logging.""" if not _is_record(request): @@ -217,32 +222,34 @@ def build_captured_mcp_parameters(request: Any) -> Dict[str, Any]: captured_request[key] = sanitize_captured_value(request[key]) if "params" in request: - captured_request["params"] = _build_captured_mcp_params(request["params"]) + captured_request["params"] = _build_captured_mcp_params( + request["params"], strip_llm_model=strip_llm_model + ) return {"request": captured_request} -def _build_captured_mcp_params(params: Any) -> Any: +def _build_captured_mcp_params(params: Any, *, strip_llm_model: bool) -> Any: if not _is_record(params): return sanitize_captured_value(params) captured: Dict[str, Any] = {} for key, value in params.items(): captured[key] = ( - _build_captured_mcp_arguments(value) + _build_captured_mcp_arguments(value, strip_llm_model=strip_llm_model) if key == "arguments" else sanitize_captured_value(value) ) return captured -def _build_captured_mcp_arguments(arguments: Any) -> Any: +def _build_captured_mcp_arguments(arguments: Any, *, strip_llm_model: bool) -> Any: if not _is_record(arguments): return sanitize_captured_value(arguments) captured: Dict[str, Any] = {} for key, value in arguments.items(): - if key in _INJECTED_ARGUMENT_NAMES: + if key in _INJECTED_ARGUMENT_NAMES or (strip_llm_model and key == "llm_model"): continue captured[key] = sanitize_captured_value(value) return captured diff --git a/posthog/mcp/_truncation.py b/posthog/mcp/_truncation.py index 2d19c6564..ccf6401c7 100644 --- a/posthog/mcp/_truncation.py +++ b/posthog/mcp/_truncation.py @@ -36,6 +36,7 @@ _METADATA_FIELDS = ( ("user_intent", _MAX_USER_INTENT_LENGTH), + ("llm_model", _MAX_METADATA_LENGTH), ("resource_name", _MAX_RESOURCE_NAME_LENGTH), ("server_name", _MAX_METADATA_LENGTH), ("server_version", _MAX_METADATA_LENGTH), diff --git a/posthog/mcp/constants.py b/posthog/mcp/constants.py index f999eae11..af0e64a77 100644 --- a/posthog/mcp/constants.py +++ b/posthog/mcp/constants.py @@ -32,6 +32,13 @@ "the first call — never invent one, and do not issue parallel tool calls until you have it." ) +DEFAULT_MODEL_PARAMETER_DESCRIPTION = ( + "The exact model identifier you (the assistant) are running as, taken from your " + 'system prompt or environment (e.g. "claude-opus-4-8", "gpt-5.2"). Used for ' + 'analytics only. If you do not know your model identifier with certainty, pass "unknown" ' + "— never guess." +) + POSTHOG_MCP_ANALYTICS_SOURCE = "posthog_mcp_analytics" POSTHOG_MCP_LIB_NAME = "posthog-python-mcp" @@ -68,6 +75,8 @@ class PostHogMCPAnalyticsProperty: IS_ERROR = "$mcp_is_error" INTENT = "$mcp_intent" INTENT_SOURCE = "$mcp_intent_source" + LLM_MODEL = "$mcp_llm_model" + LLM_MODEL_SOURCE = "$mcp_llm_model_source" LISTED_TOOL_NAMES = "$mcp_listed_tool_names" PARAMETERS = "$mcp_parameters" RESOURCE_NAME = "$mcp_resource_name" diff --git a/posthog/mcp/posthog_mcp.py b/posthog/mcp/posthog_mcp.py index 926fcf6f1..b4b50220b 100644 --- a/posthog/mcp/posthog_mcp.py +++ b/posthog/mcp/posthog_mcp.py @@ -25,11 +25,21 @@ from ._exceptions import capture_exception from ._instrumentation import drain_pending_sync, fire_and_forget from ._lib_identity import apply_mcp_lib_identity +from ._model_parameters import ( + add_model_parameter_to_schema, + can_inject_model_parameter, + get_model_description, + is_capture_model_enabled, + normalize_model, + resolve_model, +) from ._sink import McpCaptureOptions, McpEventSink from .tools import build_report_missing_descriptor from .types import ( JsonRecord, MCPAnalyticsContextOptions, + MCPAnalyticsModelOptions, + MCPAnalyticsModelSource, PreparedToolCall, ) @@ -49,6 +59,7 @@ def __init__( api_key: str, missing_capability_tool_name: Optional[str] = None, mcp_exception_autocapture: bool = True, + capture_model: Union[bool, MCPAnalyticsModelOptions] = False, **kwargs: Any, ) -> None: super().__init__(api_key, **kwargs) @@ -61,6 +72,8 @@ def __init__( # from the inherited Client.enable_exception_autocapture (global uncaught-error # hook); this mirrors instrument()'s enable_exception_autocapture, default on. self._mcp_exception_autocapture = mcp_exception_autocapture + self._capture_model = capture_model + self._model_parameter_injected: Dict[str, bool] = {} # --- lifecycle ----------------------------------------------------------- @@ -92,6 +105,8 @@ def capture_tool_call( error_type: Optional[str] = None, category: Optional[str] = None, tool_description: Optional[str] = None, + llm_model: Optional[str] = None, + llm_model_source: Optional[MCPAnalyticsModelSource] = None, protocol_version: Optional[str] = None, distinct_id: Optional[str] = None, session_id: Optional[str] = None, @@ -124,6 +139,7 @@ def capture_tool_call( event["is_error"] = is_error event["error_type"] = error_type _apply_intent(event, intent, intent_source) + _apply_model(event, llm_model, llm_model_source) if is_error: event["error"] = capture_exception( error if error is not None else f"Tool {tool_name} returned an error" @@ -218,6 +234,8 @@ def capture_missing_capability( self, *, context: Optional[str] = None, + llm_model: Optional[str] = None, + llm_model_source: Optional[MCPAnalyticsModelSource] = None, parameters: Any = None, protocol_version: Optional[str] = None, distinct_id: Optional[str] = None, @@ -246,6 +264,7 @@ def capture_missing_capability( event["protocol_version"] = protocol_version event["parameters"] = parameters _apply_intent(event, context, "context_parameter") + _apply_model(event, llm_model, llm_model_source) self._emit(event) # --- prepare helpers ----------------------------------------------------- @@ -260,11 +279,15 @@ def prepare_tool_list( intent (captured as ``$mcp_intent``), and optionally append the ``get_more_tools`` virtual tool (``report_missing=True``). Returns a new list; dict tools are copied, tool objects are mutated in place.""" - if is_context_enabled(context): - description = get_context_description(context) - prepared = [self._inject_context(tool, description) for tool in tools] - else: - prepared = list(tools) + prepared = [] + context_description = get_context_description(context) + for tool in tools: + current = ( + self._inject_context(tool, context_description) + if is_context_enabled(context) + else tool + ) + prepared.append(current) if report_missing and not any( _tool_name(t) == self._missing_capability_tool_name for t in prepared @@ -272,10 +295,16 @@ def prepare_tool_list( prepared.append( build_report_missing_descriptor(self._missing_capability_tool_name) ) + prepared = self._inject_models(prepared) return prepared def prepare_tool_call( - self, name: str, args: Optional[JsonRecord] = None + self, + name: str, + args: Optional[JsonRecord] = None, + *, + request_meta: Optional[JsonRecord] = None, + original_tool: Any = None, ) -> PreparedToolCall: """Pull the agent's intent off the injected ``context`` argument, strip ``context`` from the arguments, and flag the ``get_more_tools`` virtual tool.""" @@ -285,10 +314,26 @@ def prepare_tool_call( if isinstance(raw_context, str) and raw_context.strip() else None ) + analytics_owns_model = False + if is_capture_model_enabled(self._capture_model): + if original_tool is not None: + analytics_owns_model = can_inject_model_parameter( + _tool_schema(original_tool) + ) + else: + analytics_owns_model = self._model_parameter_injected.get(name, False) + llm_model, llm_model_source = resolve_model( + request_meta, args, allow_self_reported=analytics_owns_model + ) + prepared_args = _strip_context(args) + if analytics_owns_model: + prepared_args = _strip_model(prepared_args) return PreparedToolCall( - args=_strip_context(args), + args=prepared_args, intent=intent, intent_source="context_parameter" if intent else None, + llm_model=llm_model, + llm_model_source=llm_model_source, is_missing_capability=name == self._missing_capability_tool_name, ) @@ -356,6 +401,45 @@ def _inject_context(self, tool: Any, description: Optional[str]) -> Any: pass return tool + def _inject_models(self, tools: List[Any]) -> List[Any]: + self._model_parameter_injected.clear() + if not is_capture_model_enabled(self._capture_model): + return tools + + ownership: Dict[str, bool] = {} + for tool in tools: + name = _tool_name(tool) + if name is None: + continue + can_inject = can_inject_model_parameter(_tool_schema(tool)) + ownership[name] = ownership.get(name, True) and can_inject + self._model_parameter_injected.update(ownership) + + return [ + self._inject_model(tool) + if ownership.get(_tool_name(tool) or "", True) + else tool + for tool in tools + ] + + def _inject_model(self, tool: Any) -> Any: + name = _tool_name(tool) or "unknown" + + schema = _tool_schema(tool) + new_schema = add_model_parameter_to_schema( + schema, name, get_model_description(self._capture_model) + ) + if isinstance(tool, dict): + return {**tool, "inputSchema": new_schema} + try: + if hasattr(tool, "input_schema"): + tool.input_schema = new_schema + else: + tool.inputSchema = new_schema + except Exception: # noqa: BLE001 - read-only descriptors fail closed + self._model_parameter_injected[name] = False + return tool + def _apply_intent( event: Dict[str, Any], intent: Optional[str], source: Optional[str] @@ -367,13 +451,41 @@ def _apply_intent( event["user_intent_source"] = source or "context_parameter" +def _apply_model( + event: Dict[str, Any], + model: Optional[str], + source: Optional[MCPAnalyticsModelSource], +) -> None: + normalized = normalize_model(model) + if not normalized: + return + event["llm_model"] = normalized + event["llm_model_source"] = source or "self_reported" + + def _strip_context(args: Optional[JsonRecord]) -> Optional[JsonRecord]: if not args or "context" not in args: return args return {k: v for k, v in args.items() if k != "context"} +def _strip_model(args: Optional[JsonRecord]) -> Optional[JsonRecord]: + if not args or "llm_model" not in args: + return args + return {k: v for k, v in args.items() if k != "llm_model"} + + def _tool_name(tool: Any) -> Optional[str]: if isinstance(tool, dict): return tool.get("name") return getattr(tool, "name", None) + + +def _tool_schema(tool: Any) -> Optional[Dict[str, Any]]: + if isinstance(tool, dict): + schema = tool.get("inputSchema") + else: + schema = getattr(tool, "input_schema", None) + if schema is None: + schema = getattr(tool, "inputSchema", None) + return schema if isinstance(schema, dict) else None diff --git a/posthog/mcp/types.py b/posthog/mcp/types.py index 47dd945a1..e5f64b434 100644 --- a/posthog/mcp/types.py +++ b/posthog/mcp/types.py @@ -17,13 +17,15 @@ from dataclasses import dataclass, field from datetime import datetime -from typing import Any, Awaitable, Callable, Dict, Optional, TypedDict, Union +from typing import Any, Awaitable, Callable, Dict, Literal, Optional, TypedDict, Union from .logger import LoggerFn __all__ = [ "MCPAnalyticsOptions", "MCPAnalyticsContextOptions", + "MCPAnalyticsModelOptions", + "MCPAnalyticsModelSource", "UserIdentity", "CaptureEventData", "PreparedToolCall", @@ -35,6 +37,7 @@ ErrorProperties = Dict[str, Any] MCPAnalyticsIntentSource = str # "context_parameter" | "inferred" +MCPAnalyticsModelSource = Literal["client_metadata", "self_reported"] # Internal MCP event as it flows through the SDK before capture. Modeled as a # plain dict (constructed and read with ``.get()`` throughout) to mirror the TS @@ -43,7 +46,8 @@ # duration, error, error_type, event_name, event_type, groups, id, identify_actor_data, # identify_actor_given_id, is_error, listed_tool_names, parameters, properties, # resource_name, response, server_name, server_version, session_id, timestamp, -# tool_category, tool_description, user_intent, user_intent_source. +# tool_category, tool_description, user_intent, user_intent_source, llm_model, +# llm_model_source. Event = Dict[str, Any] McpEvent = Dict[str, Any] @@ -81,6 +85,13 @@ class MCPAnalyticsContextOptions: description: Optional[str] = None +@dataclass +class MCPAnalyticsModelOptions: + """Configure the model field injected into tool input schemas.""" + + description: Optional[str] = None + + # request is a JSON-RPC-shaped dict; extra carries session_id / headers. IdentifyFn = Callable[ ..., Any @@ -100,6 +111,9 @@ class MCPAnalyticsOptions: enable_exception_autocapture: bool = True # Inject a required `context` parameter on every tool to capture user intent. context: Union[bool, MCPAnalyticsContextOptions] = True + # Capture the model from recognized client metadata, falling back to an + # SDK-injected llm_model argument. Off by default. + capture_model: Union[bool, MCPAnalyticsModelOptions] = False # Identify the calling user — a callable (request, extra) -> UserIdentity|None # (sync or async), or a static UserIdentity. identify: Optional[Union[IdentifyFn, UserIdentity]] = None @@ -128,6 +142,8 @@ class PreparedToolCall: args: Optional[JsonRecord] = None intent: Optional[str] = None intent_source: Optional[str] = None + llm_model: Optional[str] = None + llm_model_source: Optional[MCPAnalyticsModelSource] = None is_missing_capability: bool = False diff --git a/posthog/test/mcp/_helpers_v2.py b/posthog/test/mcp/_helpers_v2.py index 57d7d67b5..c6941611a 100644 --- a/posthog/test/mcp/_helpers_v2.py +++ b/posthog/test/mcp/_helpers_v2.py @@ -53,11 +53,13 @@ def fake_ctx( def modern_meta( - client_name: str = "wire-client", client_version: str = "1.2.3" + client_name: str = "wire-client", + client_version: str = "1.2.3", + model: Optional[str] = None, ) -> Dict[str, Any]: """The 2026-07-28 per-request ``_meta`` envelope (protocol version and client capabilities are required; client info is a SHOULD).""" - return { + meta: Dict[str, Any] = { "io.modelcontextprotocol/protocolVersion": MODERN_PROTOCOL_VERSION, "io.modelcontextprotocol/clientCapabilities": {}, "io.modelcontextprotocol/clientInfo": { @@ -65,6 +67,9 @@ def modern_meta( "version": client_version, }, } + if model is not None: + meta["x-codex-turn-metadata"] = {"model": model} + return meta def modern_headers(method: str, tool_name: Optional[str] = None) -> Dict[str, str]: diff --git a/posthog/test/mcp/test_fastmcp.py b/posthog/test/mcp/test_fastmcp.py index f93d069b7..f98bca160 100644 --- a/posthog/test/mcp/test_fastmcp.py +++ b/posthog/test/mcp/test_fastmcp.py @@ -1,6 +1,7 @@ """End-to-end tests for the FastMCP adapter (Milestone 2).""" import asyncio +from types import SimpleNamespace import pytest @@ -67,6 +68,23 @@ async def test_context_injection_can_be_disabled(): assert "context" not in add_tool.inputSchema.get("properties", {}) +async def test_list_tools_injects_model_into_real_and_virtual_tools(): + server = make_server() + client = FakeClient() + instrument( + server, + client, + MCPAnalyticsOptions(capture_model=True, report_missing=True), + ) + + result = await _list_tools(server) + tools = {tool.name: tool for tool in result.root.tools} + + for name in ("add", "boom", "get_more_tools"): + assert "llm_model" in tools[name].inputSchema["properties"] + assert "llm_model" in tools[name].inputSchema["required"] + + # --- tools/call -------------------------------------------------------------- @@ -108,6 +126,30 @@ def spy_add(a: int, b: int) -> int: assert "context" not in props["$mcp_parameters"]["request"]["params"]["arguments"] +async def test_tool_call_captures_client_model_without_prior_listing(): + server = make_server() + client = FakeClient() + instrument(server, client, MCPAnalyticsOptions(capture_model=True)) + + context = SimpleNamespace( + request_context=SimpleNamespace( + meta={"x-codex-turn-metadata": {"model": "gpt-5.6-sol"}} + ) + ) + result = await server._tool_manager.call_tool( + "add", + {"a": 2, "b": 3, "llm_model": "claude-opus-4-8"}, + context=context, + ) + await _flush() + + assert result == 5 + props = _events(client, "$mcp_tool_call")[0]["properties"] + assert props["$mcp_llm_model"] == "gpt-5.6-sol" + assert props["$mcp_llm_model_source"] == "client_metadata" + assert "llm_model" not in props["$mcp_parameters"]["request"]["params"]["arguments"] + + async def test_analytics_flush_drains_its_own_captures(): async def slow_before_send(event): await asyncio.sleep(0.05) diff --git a/posthog/test/mcp/test_model_parameters.py b/posthog/test/mcp/test_model_parameters.py new file mode 100644 index 000000000..d752e4fe2 --- /dev/null +++ b/posthog/test/mcp/test_model_parameters.py @@ -0,0 +1,95 @@ +from types import SimpleNamespace + +import pytest + +from posthog.mcp._model_parameters import ( + add_model_parameter_to_schema, + request_meta_from_context, + resolve_model, +) +from posthog.mcp.constants import DEFAULT_MODEL_PARAMETER_DESCRIPTION + + +@pytest.mark.parametrize( + "schema", + [ + None, + {}, + { + "type": "object", + "properties": {"query": {"type": "string"}}, + "required": ["query"], + "additionalProperties": False, + }, + ], +) +def test_add_model_parameter_to_schema(schema): + result = add_model_parameter_to_schema(schema, "search") + + assert result["properties"]["llm_model"] == { + "type": "string", + "description": DEFAULT_MODEL_PARAMETER_DESCRIPTION, + } + assert "llm_model" in result["required"] + assert result.get("additionalProperties") is not False + + +@pytest.mark.parametrize( + "schema", + [ + {"$ref": "#/$defs/Input"}, + {"oneOf": [{"type": "object", "properties": {}}]}, + { + "type": "object", + "properties": {"llm_model": {"type": "number"}}, + }, + ], +) +def test_add_model_parameter_to_schema_preserves_unsafe_or_owned_schema(schema): + assert add_model_parameter_to_schema(schema, "search") is schema + + +@pytest.mark.parametrize( + ("metadata", "argument", "allow_self_reported", "expected"), + [ + ( + {"x-codex-turn-metadata": {"model": " gpt-5.6-sol "}}, + "claude-opus-4-8", + True, + ("gpt-5.6-sol", "client_metadata"), + ), + ( + {"x-codex-turn-metadata": {"model": "unknown"}}, + "claude-opus-4-8", + True, + ("claude-opus-4-8", "self_reported"), + ), + ( + {"x-codex-turn-metadata": "gpt-5.6-sol"}, + "claude-opus-4-8", + True, + ("claude-opus-4-8", "self_reported"), + ), + (None, "unknown", True, (None, None)), + (None, "claude-opus-4-8", False, (None, None)), + ], +) +def test_resolve_model(metadata, argument, allow_self_reported, expected): + assert ( + resolve_model( + metadata, + {"llm_model": argument}, + allow_self_reported=allow_self_reported, + ) + == expected + ) + + +def test_request_meta_from_context_supports_pydantic_style_metadata(): + meta = SimpleNamespace( + model_dump=lambda **kwargs: {"x-codex-turn-metadata": {"model": "gpt-5.6-sol"}} + ) + + assert request_meta_from_context(SimpleNamespace(meta=meta)) == { + "x-codex-turn-metadata": {"model": "gpt-5.6-sol"} + } diff --git a/posthog/test/mcp/test_pipeline.py b/posthog/test/mcp/test_pipeline.py index f57c1572b..1c98d5af3 100644 --- a/posthog/test/mcp/test_pipeline.py +++ b/posthog/test/mcp/test_pipeline.py @@ -2,6 +2,8 @@ from datetime import datetime, timezone +import pytest + from posthog.mcp.constants import ( POSTHOG_MCP_ANALYTICS_SOURCE, PostHogMCPAnalyticsEvent, @@ -167,6 +169,8 @@ def test_build_tool_call_event_properties(): "protocol_version": "2025-06-18", "user_intent": "find churn cohort", "user_intent_source": "context_parameter", + "llm_model": "gpt-5.6-sol", + "llm_model_source": "client_metadata", "is_error": False, "timestamp": datetime.now(timezone.utc), } @@ -180,6 +184,8 @@ def test_build_tool_call_event_properties(): assert props[PostHogMCPAnalyticsProperty.PROTOCOL_VERSION] == "2025-06-18" assert props[PostHogMCPAnalyticsProperty.INTENT] == "find churn cohort" assert props[PostHogMCPAnalyticsProperty.INTENT_SOURCE] == "context_parameter" + assert props[PostHogMCPAnalyticsProperty.LLM_MODEL] == "gpt-5.6-sol" + assert props[PostHogMCPAnalyticsProperty.LLM_MODEL_SOURCE] == "client_metadata" assert props[PostHogMCPAnalyticsProperty.SESSION_ID] == "ses_abc" # anonymous (no identity) => person processing disabled assert props["$process_person_profile"] is False @@ -327,6 +333,25 @@ def test_build_captured_mcp_parameters_strips_context(): assert captured["request"]["method"] == "tools/call" +@pytest.mark.parametrize( + ("strip_llm_model", "expected_model"), + [(True, None), (False, "application-owned-model")], +) +def test_build_captured_mcp_parameters_only_strips_sdk_owned_model( + strip_llm_model, expected_model +): + request = { + "method": "tools/call", + "params": { + "name": "route", + "arguments": {"llm_model": "application-owned-model"}, + }, + } + + captured = build_captured_mcp_parameters(request, strip_llm_model=strip_llm_model) + assert captured["request"]["params"]["arguments"].get("llm_model") == expected_model + + async def test_process_mcp_event_basic(): event = { "event_type": MCPAnalyticsEventType.MCP_TOOLS_CALL, diff --git a/posthog/test/mcp/test_posthog_mcp.py b/posthog/test/mcp/test_posthog_mcp.py index cf1456e95..d27836384 100644 --- a/posthog/test/mcp/test_posthog_mcp.py +++ b/posthog/test/mcp/test_posthog_mcp.py @@ -11,8 +11,8 @@ from posthog.version import VERSION -def make_client(): - client = PostHogMCP("phc_test", host="https://us.i.posthog.com") +def make_client(**kwargs): + client = PostHogMCP("phc_test", host="https://us.i.posthog.com", **kwargs) captured = [] # Intercept the inherited Client.capture so nothing is sent over the network. client.capture = lambda event, **kwargs: captured.append({"event": event, **kwargs}) @@ -211,3 +211,79 @@ def test_prepare_tool_list_can_be_disabled(): tools = [{"name": "search", "inputSchema": {"type": "object", "properties": {}}}] prepared = client.prepare_tool_list(tools, context=False) assert "context" not in prepared[0]["inputSchema"]["properties"] + + +async def test_prepare_and_capture_model(): + client, captured = make_client(capture_model=True) + tools = [ + { + "name": "search", + "inputSchema": {"type": "object", "properties": {"q": {"type": "string"}}}, + } + ] + + prepared_tools = client.prepare_tool_list(tools) + assert ( + prepared_tools[0]["inputSchema"]["properties"]["llm_model"]["type"] == "string" + ) + assert "llm_model" in prepared_tools[0]["inputSchema"]["required"] + + call = client.prepare_tool_call( + "search", + {"q": "docs", "llm_model": "claude-opus-4-8"}, + request_meta={"x-codex-turn-metadata": {"model": "gpt-5.6-sol"}}, + ) + assert call.args == {"q": "docs"} + assert call.llm_model == "gpt-5.6-sol" + assert call.llm_model_source == "client_metadata" + + client.capture_tool_call( + "search", + llm_model=call.llm_model, + llm_model_source=call.llm_model_source, + ) + await _flush() + + props = _events(captured, "$mcp_tool_call")[0]["properties"] + assert props["$mcp_llm_model"] == "gpt-5.6-sol" + assert props["$mcp_llm_model_source"] == "client_metadata" + + +def test_prepare_tool_call_preserves_application_owned_model_argument(): + client, _ = make_client(capture_model=True) + tool = { + "name": "route", + "inputSchema": { + "type": "object", + "properties": {"llm_model": {"type": "string"}}, + "required": ["llm_model"], + }, + } + client.prepare_tool_list([tool]) + + call = client.prepare_tool_call( + "route", {"llm_model": "application-owned"}, original_tool=tool + ) + assert call.args == {"llm_model": "application-owned"} + assert call.llm_model is None + + +def test_prepare_tool_list_fails_closed_for_duplicate_tool_names(): + client, _ = make_client(capture_model=True) + tools = [ + {"name": "route", "inputSchema": {"type": "object", "properties": {}}}, + { + "name": "route", + "inputSchema": { + "type": "object", + "properties": {"llm_model": {"type": "string"}}, + }, + }, + ] + + prepared = client.prepare_tool_list(tools) + assert "llm_model" not in prepared[0]["inputSchema"]["properties"] + assert prepared[1]["inputSchema"]["properties"]["llm_model"] == {"type": "string"} + call = client.prepare_tool_call("route", {"llm_model": "application-owned"}) + assert call.args == {"llm_model": "application-owned"} + assert call.llm_model is None diff --git a/posthog/test/mcp/test_truncation.py b/posthog/test/mcp/test_truncation.py index af8a572c5..c05af44da 100644 --- a/posthog/test/mcp/test_truncation.py +++ b/posthog/test/mcp/test_truncation.py @@ -95,9 +95,16 @@ def __str__(self): def test_truncate_event_caps_metadata_fields(): - out = truncate_event({"user_intent": "i" * 5000, "resource_name": "r" * 500}) + out = truncate_event( + { + "user_intent": "i" * 5000, + "resource_name": "r" * 500, + "llm_model": "m" * 500, + } + ) assert len(out["user_intent"]) == 2048 + 3 assert len(out["resource_name"]) == 256 + 3 + assert len(out["llm_model"]) == 256 + 3 def test_truncate_event_caps_exception_value_and_frames(): diff --git a/posthog/test/mcp/test_v2_wire_dual_era.py b/posthog/test/mcp/test_v2_wire_dual_era.py index 8372bcbc6..4b7c11b39 100644 --- a/posthog/test/mcp/test_v2_wire_dual_era.py +++ b/posthog/test/mcp/test_v2_wire_dual_era.py @@ -71,10 +71,10 @@ def rpc(method, params, request_id=1): return {"jsonrpc": "2.0", "id": request_id, "method": method, "params": params} -async def modern_call(http, name, arguments, request_id=1): +async def modern_call(http, name, arguments, request_id=1, model=None): body = rpc( "tools/call", - {"name": name, "arguments": arguments, "_meta": modern_meta()}, + {"name": name, "arguments": arguments, "_meta": modern_meta(model=model)}, request_id, ) return await http.post( @@ -88,11 +88,19 @@ async def modern_call(http, name, arguments, request_id=1): async def test_modern_tool_call_captured_with_envelope_identity(): server = make_server() client = FakeClient() - instrument(server, client) + instrument(server, client, MCPAnalyticsOptions(capture_model=True)) async with wire(server) as http: response = await modern_call( - http, "add", {"a": 2, "b": 3, "context": "adding on the wire"} + http, + "add", + { + "a": 2, + "b": 3, + "context": "adding on the wire", + "llm_model": "claude-opus-4-8", + }, + model="gpt-5.6-sol", ) await _flush() @@ -111,12 +119,19 @@ async def test_modern_tool_call_captured_with_envelope_identity(): assert props["$mcp_client_name"] == "wire-client" assert props["$mcp_client_version"] == "1.2.3" assert props["$mcp_protocol_version"] == MODERN_PROTOCOL_VERSION + assert props["$mcp_llm_model"] == "gpt-5.6-sol" + assert props["$mcp_llm_model_source"] == "client_metadata" + assert "llm_model" not in props["$mcp_parameters"]["request"]["params"]["arguments"] async def test_modern_tools_list_advertises_injected_params(): server = make_server() client = FakeClient() - instrument(server, client, MCPAnalyticsOptions(enable_conversation_id=True)) + instrument( + server, + client, + MCPAnalyticsOptions(enable_conversation_id=True, capture_model=True), + ) async with wire(server) as http: body = rpc("tools/list", {"_meta": modern_meta()}) @@ -129,6 +144,8 @@ async def test_modern_tools_list_advertises_injected_params(): tools = {t["name"]: t for t in response.json()["result"]["tools"]} assert "context" in tools["add"]["inputSchema"]["properties"] assert "conversation_id" in tools["add"]["inputSchema"]["properties"] + assert "llm_model" in tools["add"]["inputSchema"]["properties"] + assert "llm_model" in tools["add"]["inputSchema"]["required"] listed = _events(client, "$mcp_tools_list") assert listed and set(listed[0]["properties"]["$mcp_listed_tool_names"]) == { diff --git a/references/public_api_snapshot.txt b/references/public_api_snapshot.txt index 06d7394ee..ffdb5373e 100644 --- a/references/public_api_snapshot.txt +++ b/references/public_api_snapshot.txt @@ -333,6 +333,8 @@ alias posthog.integrations.django.Client -> posthog.client.Client alias posthog.integrations.django.contexts -> posthog.contexts alias posthog.mcp.CaptureEventData -> posthog.mcp.types.CaptureEventData alias posthog.mcp.MCPAnalyticsContextOptions -> posthog.mcp.types.MCPAnalyticsContextOptions +alias posthog.mcp.MCPAnalyticsModelOptions -> posthog.mcp.types.MCPAnalyticsModelOptions +alias posthog.mcp.MCPAnalyticsModelSource -> posthog.mcp.types.MCPAnalyticsModelSource alias posthog.mcp.MCPAnalyticsOptions -> posthog.mcp.types.MCPAnalyticsOptions alias posthog.mcp.MCP_SESSION_HEADER -> posthog.mcp.session_token.MCP_SESSION_HEADER alias posthog.mcp.POSTHOG_MCP_ANALYTICS_SOURCE -> posthog.mcp.constants.POSTHOG_MCP_ANALYTICS_SOURCE @@ -759,6 +761,8 @@ attribute posthog.mcp.constants.PostHogMCPAnalyticsProperty.INTENT = '$mcp_inten attribute posthog.mcp.constants.PostHogMCPAnalyticsProperty.INTENT_SOURCE = '$mcp_intent_source' attribute posthog.mcp.constants.PostHogMCPAnalyticsProperty.IS_ERROR = '$mcp_is_error' attribute posthog.mcp.constants.PostHogMCPAnalyticsProperty.LISTED_TOOL_NAMES = '$mcp_listed_tool_names' +attribute posthog.mcp.constants.PostHogMCPAnalyticsProperty.LLM_MODEL = '$mcp_llm_model' +attribute posthog.mcp.constants.PostHogMCPAnalyticsProperty.LLM_MODEL_SOURCE = '$mcp_llm_model_source' attribute posthog.mcp.constants.PostHogMCPAnalyticsProperty.PARAMETERS = '$mcp_parameters' attribute posthog.mcp.constants.PostHogMCPAnalyticsProperty.PROTOCOL_VERSION = '$mcp_protocol_version' attribute posthog.mcp.constants.PostHogMCPAnalyticsProperty.RESOURCE_NAME = '$mcp_resource_name' @@ -779,7 +783,10 @@ attribute posthog.mcp.session_token.SessionTokenPayload.session_id: str attribute posthog.mcp.types.CaptureEventData.event: str attribute posthog.mcp.types.CaptureEventData.properties: Optional[JsonRecord] = None attribute posthog.mcp.types.MCPAnalyticsContextOptions.description: Optional[str] = None +attribute posthog.mcp.types.MCPAnalyticsModelOptions.description: Optional[str] = None +attribute posthog.mcp.types.MCPAnalyticsModelSource = Literal['client_metadata', 'self_reported'] attribute posthog.mcp.types.MCPAnalyticsOptions.before_send: Optional[BeforeSendFn] = None +attribute posthog.mcp.types.MCPAnalyticsOptions.capture_model: Union[bool, MCPAnalyticsModelOptions] = False attribute posthog.mcp.types.MCPAnalyticsOptions.context: Union[bool, MCPAnalyticsContextOptions] = True attribute posthog.mcp.types.MCPAnalyticsOptions.enable_conversation_id: bool = False attribute posthog.mcp.types.MCPAnalyticsOptions.enable_exception_autocapture: bool = True @@ -793,6 +800,8 @@ attribute posthog.mcp.types.PreparedToolCall.args: Optional[JsonRecord] = None attribute posthog.mcp.types.PreparedToolCall.intent: Optional[str] = None attribute posthog.mcp.types.PreparedToolCall.intent_source: Optional[str] = None attribute posthog.mcp.types.PreparedToolCall.is_missing_capability: bool = False +attribute posthog.mcp.types.PreparedToolCall.llm_model: Optional[str] = None +attribute posthog.mcp.types.PreparedToolCall.llm_model_source: Optional[MCPAnalyticsModelSource] = None attribute posthog.mcp.types.UserIdentity.distinct_id: str attribute posthog.mcp.types.UserIdentity.groups: Optional[Dict[str, str]] = None attribute posthog.mcp.types.UserIdentity.properties: Optional[JsonRecord] = None @@ -977,12 +986,13 @@ class posthog.mcp.McpAnalytics(key: Any) class posthog.mcp.asgi.PostHogMcpStatelessSessionMiddleware(app: Any) class posthog.mcp.constants.PostHogMCPAnalyticsEvent class posthog.mcp.constants.PostHogMCPAnalyticsProperty -class posthog.mcp.posthog_mcp.PostHogMCP(api_key: str, missing_capability_tool_name: Optional[str] = None, mcp_exception_autocapture: bool = True, **kwargs: Any) +class posthog.mcp.posthog_mcp.PostHogMCP(api_key: str, missing_capability_tool_name: Optional[str] = None, mcp_exception_autocapture: bool = True, capture_model: Union[bool, MCPAnalyticsModelOptions] = False, **kwargs: Any) class posthog.mcp.session_token.SessionTokenPayload(session_id: str, client_name: Optional[str] = None, client_version: Optional[str] = None, protocol_version: Optional[str] = None) class posthog.mcp.types.CaptureEventData(event: str, properties: Optional[JsonRecord] = None) class posthog.mcp.types.MCPAnalyticsContextOptions(description: Optional[str] = None) -class posthog.mcp.types.MCPAnalyticsOptions(logger: Optional[LoggerFn] = None, report_missing: bool = False, missing_capability_tool_name: Optional[str] = None, enable_conversation_id: bool = False, enable_exception_autocapture: bool = True, context: Union[bool, MCPAnalyticsContextOptions] = True, identify: Optional[Union[IdentifyFn, UserIdentity]] = None, intent_fallback: Optional[IntentFallbackFn] = None, before_send: Optional[BeforeSendFn] = None, event_properties: Optional[EventPropertiesFn] = None) -class posthog.mcp.types.PreparedToolCall(args: Optional[JsonRecord] = None, intent: Optional[str] = None, intent_source: Optional[str] = None, is_missing_capability: bool = False) +class posthog.mcp.types.MCPAnalyticsModelOptions(description: Optional[str] = None) +class posthog.mcp.types.MCPAnalyticsOptions(logger: Optional[LoggerFn] = None, report_missing: bool = False, missing_capability_tool_name: Optional[str] = None, enable_conversation_id: bool = False, enable_exception_autocapture: bool = True, context: Union[bool, MCPAnalyticsContextOptions] = True, capture_model: Union[bool, MCPAnalyticsModelOptions] = False, identify: Optional[Union[IdentifyFn, UserIdentity]] = None, intent_fallback: Optional[IntentFallbackFn] = None, before_send: Optional[BeforeSendFn] = None, event_properties: Optional[EventPropertiesFn] = None) +class posthog.mcp.types.PreparedToolCall(args: Optional[JsonRecord] = None, intent: Optional[str] = None, intent_source: Optional[str] = None, llm_model: Optional[str] = None, llm_model_source: Optional[MCPAnalyticsModelSource] = None, is_missing_capability: bool = False) class posthog.mcp.types.UserIdentity(distinct_id: str, properties: Optional[JsonRecord] = None, groups: Optional[Dict[str, str]] = None) class posthog.metrics_capture.PostHogMetrics(client, config: Optional[dict] = None) class posthog.poller.Poller(interval, execute, *args, **kwargs) @@ -1387,11 +1397,11 @@ method posthog.integrations.django.PosthogContextMiddleware.process_exception(re method posthog.mcp.McpAnalytics.capture(event: str, properties: Optional[dict] = None) -> None method posthog.mcp.McpAnalytics.flush() -> None method posthog.mcp.posthog_mcp.PostHogMCP.capture_initialize(*, client_name: Optional[str] = None, client_version: Optional[str] = None, protocol_version: Optional[str] = None, parameters: Any = None, response: Any = None, duration_ms: Optional[float] = None, distinct_id: Optional[str] = None, session_id: Optional[str] = None, client_user_agent: Optional[str] = None, vendor_client: Optional[str] = None, set_properties: Optional[JsonRecord] = None, groups: Optional[Dict[str, str]] = None, properties: Optional[JsonRecord] = None, timestamp: Optional[datetime] = None) -> None -method posthog.mcp.posthog_mcp.PostHogMCP.capture_missing_capability(*, context: Optional[str] = None, parameters: Any = None, protocol_version: Optional[str] = None, distinct_id: Optional[str] = None, session_id: Optional[str] = None, client_user_agent: Optional[str] = None, vendor_client: Optional[str] = None, set_properties: Optional[JsonRecord] = None, groups: Optional[Dict[str, str]] = None, properties: Optional[JsonRecord] = None, timestamp: Optional[datetime] = None) -> None -method posthog.mcp.posthog_mcp.PostHogMCP.capture_tool_call(tool_name: str, *, intent: Optional[str] = None, intent_source: Optional[str] = None, parameters: Any = None, response: Any = None, duration_ms: Optional[float] = None, is_error: bool = False, error: Any = None, error_type: Optional[str] = None, category: Optional[str] = None, tool_description: Optional[str] = None, protocol_version: Optional[str] = None, distinct_id: Optional[str] = None, session_id: Optional[str] = None, client_user_agent: Optional[str] = None, vendor_client: Optional[str] = None, set_properties: Optional[JsonRecord] = None, groups: Optional[Dict[str, str]] = None, properties: Optional[JsonRecord] = None, timestamp: Optional[datetime] = None) -> None +method posthog.mcp.posthog_mcp.PostHogMCP.capture_missing_capability(*, context: Optional[str] = None, llm_model: Optional[str] = None, llm_model_source: Optional[MCPAnalyticsModelSource] = None, parameters: Any = None, protocol_version: Optional[str] = None, distinct_id: Optional[str] = None, session_id: Optional[str] = None, client_user_agent: Optional[str] = None, vendor_client: Optional[str] = None, set_properties: Optional[JsonRecord] = None, groups: Optional[Dict[str, str]] = None, properties: Optional[JsonRecord] = None, timestamp: Optional[datetime] = None) -> None +method posthog.mcp.posthog_mcp.PostHogMCP.capture_tool_call(tool_name: str, *, intent: Optional[str] = None, intent_source: Optional[str] = None, parameters: Any = None, response: Any = None, duration_ms: Optional[float] = None, is_error: bool = False, error: Any = None, error_type: Optional[str] = None, category: Optional[str] = None, tool_description: Optional[str] = None, llm_model: Optional[str] = None, llm_model_source: Optional[MCPAnalyticsModelSource] = None, protocol_version: Optional[str] = None, distinct_id: Optional[str] = None, session_id: Optional[str] = None, client_user_agent: Optional[str] = None, vendor_client: Optional[str] = None, set_properties: Optional[JsonRecord] = None, groups: Optional[Dict[str, str]] = None, properties: Optional[JsonRecord] = None, timestamp: Optional[datetime] = None) -> None method posthog.mcp.posthog_mcp.PostHogMCP.capture_tools_list(*, tool_names: Optional[List[str]] = None, parameters: Any = None, response: Any = None, duration_ms: Optional[float] = None, is_error: bool = False, error: Any = None, error_type: Optional[str] = None, protocol_version: Optional[str] = None, distinct_id: Optional[str] = None, session_id: Optional[str] = None, client_user_agent: Optional[str] = None, vendor_client: Optional[str] = None, set_properties: Optional[JsonRecord] = None, groups: Optional[Dict[str, str]] = None, properties: Optional[JsonRecord] = None, timestamp: Optional[datetime] = None) -> None method posthog.mcp.posthog_mcp.PostHogMCP.flush(timeout_seconds: Optional[float] = 10) -> None -method posthog.mcp.posthog_mcp.PostHogMCP.prepare_tool_call(name: str, args: Optional[JsonRecord] = None) -> PreparedToolCall +method posthog.mcp.posthog_mcp.PostHogMCP.prepare_tool_call(name: str, args: Optional[JsonRecord] = None, *, request_meta: Optional[JsonRecord] = None, original_tool: Any = None) -> PreparedToolCall method posthog.mcp.posthog_mcp.PostHogMCP.prepare_tool_list(tools: List[Any], context: Union[bool, MCPAnalyticsContextOptions] = True, report_missing: bool = False) -> List[Any] method posthog.mcp.posthog_mcp.PostHogMCP.shutdown() -> None method posthog.metrics_capture.PostHogMetrics.count(name: str, value: float = 1, unit: Optional[str] = None, attributes: Optional[dict] = None) -> None From 11c48e9dbf850bd881f0a4125e4a50270a2464c4 Mon Sep 17 00:00:00 2001 From: Lucas Faria Date: Tue, 8 Sep 2026 15:39:38 -0300 Subject: [PATCH 2/2] fix(mcp): preserve model capture opt-in and tool ownership Gate custom-dispatcher model resolution behind capture_model and copy object tools before model injection so repeated listings preserve ownership. Publish the completed ownership map without an intermediate empty state. Keep additionalProperties constraints in all analytics schema injectors. Document adapter-specific requiredness rather than requiring fields that standalone FastMCP strips before input validation. Share the eligibility predicate and update prepare helper documentation and the changeset. Verification: - MCP v1 suite: 267 passed. - MCP v2 suite: 245 passed, 13 expected skips. - Ruff check and format check passed repository-wide. - Repository mypy/baseline check passed (229 source files). - Public API snapshot and git diff --check passed. - Regression cases failed before the fix for disabled capture, object ownership and strict validation through a low-level tool-cache rebuild. Full non-MCP test suite not rerun locally; CI covers the broader matrix. --- .sampo/changesets/bold-king-mielikki.md | 1 + posthog/mcp/README.md | 13 +++-- posthog/mcp/_context_parameters.py | 5 +- posthog/mcp/_conversation_id.py | 2 - posthog/mcp/_model_parameters.py | 20 ++----- posthog/mcp/posthog_mcp.py | 41 +++++++++------ posthog/test/mcp/test_model_parameters.py | 11 +++- posthog/test/mcp/test_posthog_mcp.py | 64 +++++++++++++++++++---- posthog/test/mcp/test_review_fixes.py | 15 +++++- posthog/test/mcp/test_units.py | 4 +- 10 files changed, 119 insertions(+), 57 deletions(-) diff --git a/.sampo/changesets/bold-king-mielikki.md b/.sampo/changesets/bold-king-mielikki.md index 2783644a0..0d37204ed 100644 --- a/.sampo/changesets/bold-king-mielikki.md +++ b/.sampo/changesets/bold-king-mielikki.md @@ -3,3 +3,4 @@ pypi/posthog: minor --- Capture MCP model identifiers from client metadata or an SDK-owned self-report field. +Model capture remains opt-in, preserves application-owned fields across repeated tool listings, and keeps strict input validation when adding analytics fields. diff --git a/posthog/mcp/README.md b/posthog/mcp/README.md index fb6c12e31..a032629f7 100644 --- a/posthog/mcp/README.md +++ b/posthog/mcp/README.md @@ -38,9 +38,13 @@ analytics = instrument( The SDK records the best model identifier visible to the server as `$mcp_llm_model`. Recognized client metadata wins and sets -`$mcp_llm_model_source` to `client_metadata`. Otherwise, the SDK adds a required -`llm_model` string to each compatible tool schema and records the agent's answer -with source `self_reported`. +`$mcp_llm_model_source` to `client_metadata`. The SDK also adds an `llm_model` +string to each compatible tool schema as a fallback, recorded with source +`self_reported`. It is required for custom dispatchers and the official high-level +MCP SDK adapters. Raw low-level servers and standalone `fastmcp.FastMCP` advertise +it as optional; the standalone adapter strips it before input validation, so +requiring it would reject calls. Existing schema strictness is preserved when +analytics fields are added. The recognized metadata path is Codex's `x-codex-turn-metadata.model` field in request `_meta`. Other clients, including Claude Code, use the self-report path @@ -82,6 +86,9 @@ posthog.capture_tool_call( Passing `original_tool` keeps ownership accurate when `tools/list` and `tools/call` reach different server replicas. A persistent single-process dispatcher can omit it after calling `prepare_tool_list()`. +Model injection copies tool objects instead of changing their original schemas. +Always advertise the returned list and pass the original application tool to +`prepare_tool_call()`. Repeatedly preparing the original list preserves ownership. ## Stateless / multi-pod servers diff --git a/posthog/mcp/_context_parameters.py b/posthog/mcp/_context_parameters.py index 326e8eb25..015663d6b 100644 --- a/posthog/mcp/_context_parameters.py +++ b/posthog/mcp/_context_parameters.py @@ -79,10 +79,7 @@ def add_context_parameter_to_schema( if not isinstance(schema.get("properties"), dict): schema["properties"] = {} - # additionalProperties: false would reject the injected context — remove it - # (the SDK adds this when converting Pydantic models to JSON Schema). - if schema.get("additionalProperties") is False: - schema.pop("additionalProperties", None) + # The declared context property is allowed even under additionalProperties: false. schema["properties"]["context"] = { "type": "string", diff --git a/posthog/mcp/_conversation_id.py b/posthog/mcp/_conversation_id.py index e5c171cb2..1f44f3c63 100644 --- a/posthog/mcp/_conversation_id.py +++ b/posthog/mcp/_conversation_id.py @@ -57,8 +57,6 @@ def add_conversation_id_to_schema( schema = copy.deepcopy(schema) if not isinstance(schema.get("properties"), dict): schema["properties"] = {} - if schema.get("additionalProperties") is False: - schema.pop("additionalProperties", None) schema["properties"][CONVERSATION_ID_PARAM_NAME] = { "type": "string", "description": DEFAULT_CONVERSATION_ID_DESCRIPTION, diff --git a/posthog/mcp/_model_parameters.py b/posthog/mcp/_model_parameters.py index b3a0c3942..05ee0f29f 100644 --- a/posthog/mcp/_model_parameters.py +++ b/posthog/mcp/_model_parameters.py @@ -41,21 +41,10 @@ def add_model_parameter_to_schema( never overwrite or later strip a value that belongs to the tool itself. """ schema = input_schema - if ( - schema - and isinstance(schema.get("properties"), dict) - and "llm_model" in schema["properties"] - ): + if not can_inject_model_parameter(schema): log( - f"WARN: Tool \"{tool_name}\" already has 'llm_model' parameter. " - "Skipping model injection." - ) - return schema - - if schema and any(schema.get(key) for key in ("$ref", "oneOf", "allOf", "anyOf")): - log( - f'WARN: Tool "{tool_name}" has complex schema ' - "($ref/oneOf/allOf/anyOf). Skipping model injection." + f'WARN: Tool "{tool_name}" has an application-owned llm_model or ' + "complex schema ($ref/oneOf/allOf/anyOf). Skipping model injection." ) return schema @@ -65,9 +54,6 @@ def add_model_parameter_to_schema( schema = copy.deepcopy(schema) if not isinstance(schema.get("properties"), dict): schema["properties"] = {} - if schema.get("additionalProperties") is False: - schema.pop("additionalProperties", None) - schema["properties"]["llm_model"] = { "type": "string", "description": description_override or DEFAULT_MODEL_PARAMETER_DESCRIPTION, diff --git a/posthog/mcp/posthog_mcp.py b/posthog/mcp/posthog_mcp.py index b4b50220b..8aac5a8b8 100644 --- a/posthog/mcp/posthog_mcp.py +++ b/posthog/mcp/posthog_mcp.py @@ -11,6 +11,7 @@ from __future__ import annotations +import copy from datetime import datetime, timezone from typing import Any, Dict, List, Optional, Union @@ -278,7 +279,8 @@ def prepare_tool_list( """Inject the ``context`` argument into every tool so agents state their intent (captured as ``$mcp_intent``), and optionally append the ``get_more_tools`` virtual tool (``report_missing=True``). Returns a new - list; dict tools are copied, tool objects are mutated in place.""" + list; dict tools are copied, context injection mutates tool objects in + place, and model injection copies them to preserve field ownership.""" prepared = [] context_description = get_context_description(context) for tool in tools: @@ -307,7 +309,9 @@ def prepare_tool_call( original_tool: Any = None, ) -> PreparedToolCall: """Pull the agent's intent off the injected ``context`` argument, strip - ``context`` from the arguments, and flag the ``get_more_tools`` virtual tool.""" + ``context`` from the arguments, and flag the ``get_more_tools`` virtual tool. + When model capture is enabled, resolve its value and source and strip + the SDK-owned ``llm_model`` argument before dispatch.""" raw_context = (args or {}).get("context") intent = ( raw_context.strip() @@ -315,6 +319,8 @@ def prepare_tool_call( else None ) analytics_owns_model = False + llm_model: Optional[str] = None + llm_model_source: Optional[MCPAnalyticsModelSource] = None if is_capture_model_enabled(self._capture_model): if original_tool is not None: analytics_owns_model = can_inject_model_parameter( @@ -322,9 +328,9 @@ def prepare_tool_call( ) else: analytics_owns_model = self._model_parameter_injected.get(name, False) - llm_model, llm_model_source = resolve_model( - request_meta, args, allow_self_reported=analytics_owns_model - ) + llm_model, llm_model_source = resolve_model( + request_meta, args, allow_self_reported=analytics_owns_model + ) prepared_args = _strip_context(args) if analytics_owns_model: prepared_args = _strip_model(prepared_args) @@ -402,8 +408,8 @@ def _inject_context(self, tool: Any, description: Optional[str]) -> Any: return tool def _inject_models(self, tools: List[Any]) -> List[Any]: - self._model_parameter_injected.clear() if not is_capture_model_enabled(self._capture_model): + self._model_parameter_injected = {} return tools ownership: Dict[str, bool] = {} @@ -413,16 +419,16 @@ def _inject_models(self, tools: List[Any]) -> List[Any]: continue can_inject = can_inject_model_parameter(_tool_schema(tool)) ownership[name] = ownership.get(name, True) and can_inject - self._model_parameter_injected.update(ownership) - - return [ - self._inject_model(tool) + prepared = [ + self._inject_model(tool, ownership) if ownership.get(_tool_name(tool) or "", True) else tool for tool in tools ] + self._model_parameter_injected = ownership + return prepared - def _inject_model(self, tool: Any) -> Any: + def _inject_model(self, tool: Any, ownership: Dict[str, bool]) -> Any: name = _tool_name(tool) or "unknown" schema = _tool_schema(tool) @@ -432,12 +438,17 @@ def _inject_model(self, tool: Any) -> Any: if isinstance(tool, dict): return {**tool, "inputSchema": new_schema} try: - if hasattr(tool, "input_schema"): - tool.input_schema = new_schema + prepared = copy.copy(tool) + if prepared is tool: + ownership[name] = False + return tool + if hasattr(prepared, "input_schema"): + prepared.input_schema = new_schema else: - tool.inputSchema = new_schema + prepared.inputSchema = new_schema + return prepared except Exception: # noqa: BLE001 - read-only descriptors fail closed - self._model_parameter_injected[name] = False + ownership[name] = False return tool diff --git a/posthog/test/mcp/test_model_parameters.py b/posthog/test/mcp/test_model_parameters.py index d752e4fe2..0891f0e95 100644 --- a/posthog/test/mcp/test_model_parameters.py +++ b/posthog/test/mcp/test_model_parameters.py @@ -1,6 +1,7 @@ from types import SimpleNamespace import pytest +from jsonschema import Draft202012Validator from posthog.mcp._model_parameters import ( add_model_parameter_to_schema, @@ -31,7 +32,15 @@ def test_add_model_parameter_to_schema(schema): "description": DEFAULT_MODEL_PARAMETER_DESCRIPTION, } assert "llm_model" in result["required"] - assert result.get("additionalProperties") is not False + assert result.get("additionalProperties") == (schema or {}).get( + "additionalProperties" + ) + if schema and schema.get("additionalProperties") is False: + validator = Draft202012Validator(result) + assert validator.is_valid({"query": "docs", "llm_model": "example-model"}) + assert not validator.is_valid( + {"query": "docs", "llm_model": "example-model", "undeclared": True} + ) @pytest.mark.parametrize( diff --git a/posthog/test/mcp/test_posthog_mcp.py b/posthog/test/mcp/test_posthog_mcp.py index d27836384..7b9b89fc5 100644 --- a/posthog/test/mcp/test_posthog_mcp.py +++ b/posthog/test/mcp/test_posthog_mcp.py @@ -1,7 +1,11 @@ """Tests for the PostHogMCP custom-dispatcher client (Milestone 3).""" +from types import SimpleNamespace from unittest import mock +import pytest +from mcp.types import Tool + from posthog.capture_mode import CaptureMode from posthog.mcp import PostHogMCP from posthog.test.mcp._helpers import ( @@ -213,8 +217,12 @@ def test_prepare_tool_list_can_be_disabled(): assert "context" not in prepared[0]["inputSchema"]["properties"] -async def test_prepare_and_capture_model(): - client, captured = make_client(capture_model=True) +@pytest.mark.parametrize( + "options", [{"capture_model": True}, {"capture_model": False}, {}] +) +async def test_prepare_and_capture_model(options: dict[str, bool]) -> None: + client, captured = make_client(**options) + enabled = options.get("capture_model", False) tools = [ { "name": "search", @@ -223,19 +231,20 @@ async def test_prepare_and_capture_model(): ] prepared_tools = client.prepare_tool_list(tools) - assert ( - prepared_tools[0]["inputSchema"]["properties"]["llm_model"]["type"] == "string" - ) - assert "llm_model" in prepared_tools[0]["inputSchema"]["required"] + schema = prepared_tools[0]["inputSchema"] + assert ("llm_model" in schema["properties"]) == enabled + assert ("llm_model" in schema.get("required", [])) == enabled call = client.prepare_tool_call( "search", {"q": "docs", "llm_model": "claude-opus-4-8"}, request_meta={"x-codex-turn-metadata": {"model": "gpt-5.6-sol"}}, ) - assert call.args == {"q": "docs"} - assert call.llm_model == "gpt-5.6-sol" - assert call.llm_model_source == "client_metadata" + assert call.args == ( + {"q": "docs"} if enabled else {"q": "docs", "llm_model": "claude-opus-4-8"} + ) + assert call.llm_model == ("gpt-5.6-sol" if enabled else None) + assert call.llm_model_source == ("client_metadata" if enabled else None) client.capture_tool_call( "search", @@ -245,8 +254,41 @@ async def test_prepare_and_capture_model(): await _flush() props = _events(captured, "$mcp_tool_call")[0]["properties"] - assert props["$mcp_llm_model"] == "gpt-5.6-sol" - assert props["$mcp_llm_model_source"] == "client_metadata" + assert props.get("$mcp_llm_model") == ("gpt-5.6-sol" if enabled else None) + assert props.get("$mcp_llm_model_source") == ( + "client_metadata" if enabled else None + ) + + +@pytest.mark.parametrize("sdk_tool", [False, True]) +@pytest.mark.parametrize("pass_original_tool", [False, True]) +def test_prepare_model_preserves_object_tool_ownership( + sdk_tool: bool, pass_original_tool: bool +) -> None: + client, _ = make_client(capture_model=True) + schema = {"type": "object", "properties": {"q": {"type": "string"}}} + tool = ( + Tool(name="search", inputSchema=schema) + if sdk_tool + else SimpleNamespace(name="search", input_schema=schema) + ) + schema_attribute = ( + "input_schema" if hasattr(tool, "input_schema") else "inputSchema" + ) + for _ in range(2): + prepared = client.prepare_tool_list([tool], context=False) + assert "llm_model" in getattr(prepared[0], schema_attribute)["properties"] + call = client.prepare_tool_call( + "search", + {"q": "docs", "llm_model": "example-model"}, + original_tool=tool if pass_original_tool else None, + ) + assert call.args == {"q": "docs"} + assert (call.llm_model, call.llm_model_source) == ( + "example-model", + "self_reported", + ) + assert "llm_model" not in getattr(tool, schema_attribute)["properties"] def test_prepare_tool_call_preserves_application_owned_model_argument(): diff --git a/posthog/test/mcp/test_review_fixes.py b/posthog/test/mcp/test_review_fixes.py index ddcf3edd0..d8be0b4e3 100644 --- a/posthog/test/mcp/test_review_fixes.py +++ b/posthog/test/mcp/test_review_fixes.py @@ -436,12 +436,20 @@ async def _ct(name, arguments): return [mcp_types.TextContent(type="text", text=str(arguments.get("msg")))] client = FakeClient() - instrument(server, client, MCPAnalyticsOptions(enable_conversation_id=True)) + instrument( + server, + client, + MCPAnalyticsOptions(enable_conversation_id=True, capture_model=True), + ) call = server.request_handlers[mcp_types.CallToolRequest] await server.request_handlers[mcp_types.ListToolsRequest](_list_request()) - first = await call(_call_request("echo", {"msg": "a", "context": "first"})) + first = await call( + _call_request( + "echo", {"msg": "a", "context": "first", "llm_model": "example-model"} + ) + ) assert first.root.isError is False # Force a cache rebuild the way a real client does: an unknown tool name. @@ -449,3 +457,6 @@ async def _ct(name, arguments): after = await call(_call_request("echo", {"msg": "b", "context": "second"})) assert after.root.isError is False, after.root.content[0].text + + invalid = await call(_call_request("echo", {"msg": "c", "undeclared": True})) + assert invalid.root.isError is True diff --git a/posthog/test/mcp/test_units.py b/posthog/test/mcp/test_units.py index 04c8cb02c..998510df1 100644 --- a/posthog/test/mcp/test_units.py +++ b/posthog/test/mcp/test_units.py @@ -125,11 +125,11 @@ def test_add_conversation_id_skips_complex_schema(): assert add_conversation_id_to_schema(schema, "t") is schema -def test_add_conversation_id_strips_additional_properties_false(): +def test_add_conversation_id_preserves_additional_properties_false(): out = add_conversation_id_to_schema( {"type": "object", "properties": {}, "additionalProperties": False}, "t" ) - assert "additionalProperties" not in out + assert out["additionalProperties"] is False assert "conversation_id" in out["properties"]