diff --git a/.sampo/changesets/bold-king-mielikki.md b/.sampo/changesets/bold-king-mielikki.md new file mode 100644 index 000000000..0d37204ed --- /dev/null +++ b/.sampo/changesets/bold-king-mielikki.md @@ -0,0 +1,6 @@ +--- +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 a8dcd1020..a032629f7 100644 --- a/posthog/mcp/README.md +++ b/posthog/mcp/README.md @@ -21,6 +21,75 @@ 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`. 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 +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()`. +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 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/_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/_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..05ee0f29f --- /dev/null +++ b/posthog/mcp/_model_parameters.py @@ -0,0 +1,128 @@ +"""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 not can_inject_model_parameter(schema): + log( + f'WARN: Tool "{tool_name}" has an application-owned llm_model or ' + "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"] = {} + 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..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 @@ -25,11 +26,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 +60,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 +73,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 +106,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 +140,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 +235,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 +265,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 ----------------------------------------------------- @@ -259,12 +279,17 @@ 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.""" - if is_context_enabled(context): - description = get_context_description(context) - prepared = [self._inject_context(tool, description) for tool in tools] - else: - prepared = list(tools) + 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: + 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,23 +297,49 @@ 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.""" + ``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() if isinstance(raw_context, str) and raw_context.strip() 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( + _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 +407,50 @@ def _inject_context(self, tool: Any, description: Optional[str]) -> Any: pass return tool + def _inject_models(self, tools: List[Any]) -> List[Any]: + if not is_capture_model_enabled(self._capture_model): + self._model_parameter_injected = {} + 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 + 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, ownership: Dict[str, bool]) -> 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: + prepared = copy.copy(tool) + if prepared is tool: + ownership[name] = False + return tool + if hasattr(prepared, "input_schema"): + prepared.input_schema = new_schema + else: + prepared.inputSchema = new_schema + return prepared + except Exception: # noqa: BLE001 - read-only descriptors fail closed + ownership[name] = False + return tool + def _apply_intent( event: Dict[str, Any], intent: Optional[str], source: Optional[str] @@ -367,13 +462,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..0891f0e95 --- /dev/null +++ b/posthog/test/mcp/test_model_parameters.py @@ -0,0 +1,104 @@ +from types import SimpleNamespace + +import pytest +from jsonschema import Draft202012Validator + +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") == (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( + "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..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 ( @@ -11,8 +15,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 +215,117 @@ 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"] + + +@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", + "inputSchema": {"type": "object", "properties": {"q": {"type": "string"}}}, + } + ] + + prepared_tools = client.prepare_tool_list(tools) + 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"} 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", + llm_model=call.llm_model, + llm_model_source=call.llm_model_source, + ) + await _flush() + + props = _events(captured, "$mcp_tool_call")[0]["properties"] + 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(): + 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_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_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_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"] 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