diff --git a/posthog/mcp/_instrument_fastmcp.py b/posthog/mcp/_instrument_fastmcp.py index b0ac4db5..e4d8a84d 100644 --- a/posthog/mcp/_instrument_fastmcp.py +++ b/posthog/mcp/_instrument_fastmcp.py @@ -26,41 +26,22 @@ import mcp.types as mcp_types -from ._context_parameters import ( - add_context_parameter_to_schema, - get_context_description, - is_context_enabled, -) -from ._conversation_id import ( - add_conversation_id_to_schema, - build_prompt_back, - resolve_conversation_id, -) +from ._conversation_id import build_prompt_back from ._instrumentation import ( _to_jsonable, append_get_more_tools, - build_tool_call_request, + collect_listed_tools, extract_tools, - prepare_request, - prime_session, - read_tool_category, - record_missing_capability, - record_tool_call, - record_tools_list, + mutate_tool_schema, request_to_dict, resolve_session_and_client, + start_tool_call_lifecycle, + start_tools_list_lifecycle, ) from ._internal import MCPAnalyticsData -from ._output_instructions import ( - add_instructions_to_output_schema, - mirror_instructions_into_structured_content, -) +from ._output_instructions import mirror_instructions_into_structured_content from .logger import log -from .tools import ( - GET_MORE_TOOLS_NAME as _GET_MORE_TOOLS_NAME, - get_more_tools_result_text, - resolve_missing_capability_tool_name, -) +from .tools import get_more_tools_result_text, resolve_missing_capability_tool_name _WRAPPED_FLAG = "__posthog_mcp_wrapped__" @@ -103,54 +84,25 @@ async def wrapped( mcp_session_id, client_name, client_version, protocol_version ) ) - request = build_tool_call_request(name, arguments) - # `ctx` is the SDK's own per-request context, handed to host callbacks - # unchanged and identically on both SDK majors (read headers off it with - # the exported `get_request_headers`). Never captured — the event - # pipeline keeps only a scalar projection of `extra`. - extra: Dict[str, Any] = { - "session_id": mcp_session_id, - "ctx": _tool_call_request_context(context), - } - - # Resolve the conversation handle before the session: when the agent - # carries (or is about to receive) one, it anchors $session_id for every - # event of this request (ADR-0004) — the only correlation that survives - # the 2026-07-28 revision's per-request server instances. - missing_name = resolve_missing_capability_tool_name(data.options) - conversation_id, minted = resolve_conversation_id( - data.options.enable_conversation_id, arguments, name, missing_name + # Context lookup and dispatch stay adapter-specific; the lifecycle owns + # only the common session/capture policy. + lifecycle = start_tool_call_lifecycle( + data, + name=name, + arguments=arguments, + mcp_session_id=mcp_session_id, + token=token, + client_name=client_name, + client_version=client_version, + protocol_version=protocol_version, + extra={ + "session_id": mcp_session_id, + "ctx": _tool_call_request_context(context), + }, ) - # Resolved once the handle's fate is known — a minted handle only - # anchors the session after we have confirmed the agent received it, - # so the call that mints it still joins its own conversation. - async def _session(anchor: Optional[str]) -> str: - return await prepare_request( - data, - mcp_session_id=mcp_session_id, - client_name=client_name, - client_version=client_version, - protocol_version=protocol_version, - request=request, - extra=extra, - token=token, - conversation_id=anchor, - ) - - if data.options.report_missing and name == missing_name: - session_id = await _session(None) - await record_missing_capability( - data, - session_id, - tool_name=missing_name, - context=(arguments or {}).get("context"), - arguments=arguments, - client_name=client_name, - client_version=client_version, - protocol_version=protocol_version, - extra=extra, - ) + if lifecycle.is_missing_capability: + await lifecycle.record_missing_capability() return [ mcp_types.TextContent(type="text", text=get_more_tools_result_text()) ] @@ -174,7 +126,7 @@ async def _session(anchor: Optional[str]) -> str: # Settle the shared session before the tool body runs, so an in-tool # `analytics.capture()` is attributed to this caller and not the last one. - await prime_session(data, mcp_session_id=mcp_session_id, token=token) + await lifecycle.prime_session() start = time.monotonic() try: @@ -184,20 +136,7 @@ async def _session(anchor: Optional[str]) -> str: except Exception as error: # The minted prompt-back was never delivered to the agent — don't stamp # an orphan conversation_id it can't echo (an agent-supplied id is kept). - session_id = await _session(None if minted else conversation_id) - await record_tool_call( - data, - session_id, - name=name, - arguments=arguments, - error=error, - duration_ms=(time.monotonic() - start) * 1000, - client_name=client_name, - client_version=client_version, - protocol_version=protocol_version, - conversation_id=None if minted else conversation_id, - extra=extra, - ) + await lifecycle.record_error(error, (time.monotonic() - start) * 1000) raise # Deliver the handle first, then capture the result the agent actually got. @@ -205,35 +144,22 @@ async def _session(anchor: Optional[str]) -> str: # tools whose output schema we declared the key on — clients that read # structuredContent never see the text block), and the prompt-back text # block on the minting response only. - delivered_conversation_id = conversation_id - if conversation_id: - delivered = False + delivered = False + if lifecycle.conversation_id: if data.tool_output_instructions.get(name): result, delivered = mirror_instructions_into_structured_content( - result, conversation_id + result, lifecycle.conversation_id ) - if minted: - injected = _inject_prompt_back(result, conversation_id) + if lifecycle.minted_conversation_id: + injected = _inject_prompt_back(result, lifecycle.conversation_id) if injected is not result: delivered = True result = injected - # Only a minted handle can be lost — one the agent supplied, it has. - if not delivered: - delivered_conversation_id = None - session_id = await _session(delivered_conversation_id) - await record_tool_call( - data, - session_id, - name=name, - arguments=arguments, - result=result, - duration_ms=(time.monotonic() - start) * 1000, - client_name=client_name, - client_version=client_version, - protocol_version=protocol_version, - conversation_id=delivered_conversation_id, - extra=extra, + await lifecycle.record_result( + result, + (time.monotonic() - start) * 1000, + conversation_id_delivered=delivered, ) return result @@ -251,29 +177,14 @@ def _inject_tool_schemas(server: Any, data: MCPAnalyticsData, tools: list) -> No population pass, so the schema the SDK validates against always matches the one we advertised — see the note in ``list_handler``. """ - context_enabled = is_context_enabled(data.options.context) - description = get_context_description(data.options.context) for tool in tools: - if tool.name == _GET_MORE_TOOLS_NAME: - continue - owns_context = _tool_owns_context(server, tool.name) - schema = getattr(tool, "inputSchema", None) - if context_enabled and not owns_context: - schema = add_context_parameter_to_schema(schema, tool.name, description) - if data.options.enable_conversation_id: - schema = add_conversation_id_to_schema(schema, tool.name) - if schema is not getattr(tool, "inputSchema", None): - try: - tool.inputSchema = schema - except Exception: # noqa: BLE001 - some schema attrs may be read-only - log(f"WARN: could not set inputSchema on tool {tool.name}") - # Declare the structuredContent channel and remember the answer: - # clients that read structuredContent never see the content text - # block, and only a declared key may be written back on a call. - if data.options.enable_conversation_id: - data.tool_output_instructions[tool.name] = ( - add_instructions_to_output_schema(tool) - ) + mutate_tool_schema( + data, + tool, + schema_attribute="inputSchema", + owns_context=_tool_owns_context(server, tool.name), + context_required=True, + ) def _wrap_list_tools_handler(server: Any, data: MCPAnalyticsData) -> None: @@ -313,49 +224,27 @@ async def list_handler(req: Any) -> Any: } # Resolve session, emit $mcp_initialize (once per session) and identify here # too — a client may list tools without ever calling one. - session_id = await prepare_request( + lifecycle = await start_tools_list_lifecycle( data, + request=request, + extra=extra, mcp_session_id=mcp_session_id, + token=token, client_name=client_name, client_version=client_version, protocol_version=protocol_version, - request=request, - extra=extra, - token=token, ) start = time.monotonic() try: result = await original(req) except Exception as error: - await record_tools_list( - data, - session_id, - names=[], - request=request, - duration_ms=(time.monotonic() - start) * 1000, - is_error=True, - error=error, - client_name=client_name, - client_version=client_version, - protocol_version=protocol_version, - extra=extra, - ) + await lifecycle.record_error(error, (time.monotonic() - start) * 1000) raise duration_ms = (time.monotonic() - start) * 1000 tools = extract_tools(result) - # Zero advertised tools is treated as an errored tools/list (parity with the - # TS SDK), checked before we append our own get_more_tools virtual tool. - empty = len(tools) == 0 - - names = [] - for tool in tools: - names.append(tool.name) - if getattr(tool, "description", None): - data.tool_descriptions[tool.name] = tool.description - category = read_tool_category(tool) - if category: - data.tool_categories[tool.name] = category + # Empty is computed before adding the virtual missing-capability tool. + names, empty = collect_listed_tools(data, tools) _inject_tool_schemas(server, data, tools) @@ -365,19 +254,11 @@ async def list_handler(req: Any) -> Any: append_get_more_tools(result, missing_name) names.append(missing_name) - await record_tools_list( - data, - session_id, + await lifecycle.record_result( names=names, - request=request, response=_to_jsonable(result), duration_ms=duration_ms, - is_error=empty, - error="tools/list returned no tools" if empty else None, - client_name=client_name, - client_version=client_version, - protocol_version=protocol_version, - extra=extra, + is_empty=empty, ) return result diff --git a/posthog/mcp/_instrument_lowlevel.py b/posthog/mcp/_instrument_lowlevel.py index bfeb390b..cc143255 100644 --- a/posthog/mcp/_instrument_lowlevel.py +++ b/posthog/mcp/_instrument_lowlevel.py @@ -20,42 +20,23 @@ import mcp.types as mcp_types -from ._context_parameters import ( - add_context_parameter_to_schema, - get_context_description, - is_context_enabled, - schema_has_param, -) -from ._conversation_id import ( - add_conversation_id_to_schema, - build_prompt_back, - resolve_conversation_id, -) +from ._context_parameters import schema_has_param +from ._conversation_id import build_prompt_back from ._instrumentation import ( _to_jsonable, append_get_more_tools, - build_tool_call_request, + collect_listed_tools, extract_tools, - prepare_request, - prime_session, - read_tool_category, - record_missing_capability, - record_tool_call, - record_tools_list, + mutate_tool_schema, request_to_dict, resolve_session_and_client, + start_tool_call_lifecycle, + start_tools_list_lifecycle, ) from ._internal import MCPAnalyticsData -from ._output_instructions import ( - add_instructions_to_output_schema, - mirror_instructions_into_structured_content, -) +from ._output_instructions import mirror_instructions_into_structured_content from .logger import log -from .tools import ( - GET_MORE_TOOLS_NAME as _GET_MORE_TOOLS_NAME, - get_more_tools_result_text, - resolve_missing_capability_tool_name, -) +from .tools import get_more_tools_result_text, resolve_missing_capability_tool_name _WRAPPED_FLAG = "__posthog_mcp_wrapped__" @@ -112,49 +93,20 @@ async def handler(req: Any) -> Any: mcp_session_id, client_name, client_version, protocol_version ) ) - request = build_tool_call_request(name, arguments) - # `ctx` is the SDK's own per-request context, handed to host callbacks - # unchanged and identically on both SDK majors (read headers off it with - # the exported `get_request_headers`). Never captured — the event - # pipeline keeps only a scalar projection of `extra`. - extra = {"session_id": mcp_session_id, "ctx": _request_context(server)} - - # Resolve the conversation handle before the session: when present it - # anchors $session_id for every event of this request (ADR-0004). - missing_name = resolve_missing_capability_tool_name(data.options) - conversation_id, minted = resolve_conversation_id( - data.options.enable_conversation_id, arguments, name, missing_name + lifecycle = start_tool_call_lifecycle( + data, + name=name, + arguments=arguments, + mcp_session_id=mcp_session_id, + token=token, + client_name=client_name, + client_version=client_version, + protocol_version=protocol_version, + extra={"session_id": mcp_session_id, "ctx": _request_context(server)}, ) - # Resolved once the handle's fate is known — a minted handle only - # anchors the session after we have confirmed the agent received it, - # so the call that mints it still joins its own conversation. - async def _session(anchor: Optional[str]) -> str: - return await prepare_request( - data, - mcp_session_id=mcp_session_id, - client_name=client_name, - client_version=client_version, - protocol_version=protocol_version, - request=request, - extra=extra, - token=token, - conversation_id=anchor, - ) - - if data.options.report_missing and name == missing_name: - session_id = await _session(None) - await record_missing_capability( - data, - session_id, - tool_name=missing_name, - context=arguments.get("context"), - arguments=arguments, - client_name=client_name, - client_version=client_version, - protocol_version=protocol_version, - extra=extra, - ) + if lifecycle.is_missing_capability: + await lifecycle.record_missing_capability() return mcp_types.ServerResult( mcp_types.CallToolResult( content=[ @@ -181,7 +133,7 @@ async def _session(anchor: Optional[str]) -> str: # Settle the shared session before the tool body runs, so an in-tool # `analytics.capture()` is attributed to this caller and not the last one. - await prime_session(data, mcp_session_id=mcp_session_id, token=token) + await lifecycle.prime_session() start = time.monotonic() try: @@ -192,20 +144,7 @@ async def _session(anchor: Optional[str]) -> str: # request_handlers can raise — capture before re-raising so the failed # call isn't silently dropped. A minted (undelivered) conversation_id is # not stamped, matching the FastMCP path. - session_id = await _session(None if minted else conversation_id) - await record_tool_call( - data, - session_id, - name=name, - arguments=arguments, - error=error, - duration_ms=(time.monotonic() - start) * 1000, - client_name=client_name, - client_version=client_version, - protocol_version=protocol_version, - conversation_id=None if minted else conversation_id, - extra=extra, - ) + await lifecycle.record_error(error, (time.monotonic() - start) * 1000) raise duration_ms = (time.monotonic() - start) * 1000 @@ -221,18 +160,18 @@ async def _session(anchor: Optional[str]) -> str: # when it actually reached the agent, so we don't record an orphan id. # Errored results carry it on purpose: a first-call failure is exactly # when the agent needs the handle, or the retry starts a fresh conversation. - delivered_conversation_id = conversation_id - if conversation_id: - delivered = False + delivered = False + if lifecycle.conversation_id: if data.tool_output_instructions.get(name): call_result, delivered = mirror_instructions_into_structured_content( - call_result, conversation_id + call_result, lifecycle.conversation_id ) - if minted: + if lifecycle.minted_conversation_id: content = getattr(call_result, "content", None) if isinstance(content, list): block = mcp_types.TextContent( - type="text", text=build_prompt_back(conversation_id)["text"] + type="text", + text=build_prompt_back(lifecycle.conversation_id)["text"], ) # Copy rather than append in place — a shared or cached result # object would accumulate a block per conversation and leak @@ -243,9 +182,6 @@ async def _session(anchor: Optional[str]) -> str: else: content.append(block) delivered = True - # Only a minted handle can be lost — one the agent supplied, it has. - if not delivered: - delivered_conversation_id = None # Hand back whatever copy we made, rewrapped as the SDK expects. if call_result is not getattr(result, "root", result): result = ( @@ -254,19 +190,10 @@ async def _session(anchor: Optional[str]) -> str: else call_result ) - session_id = await _session(delivered_conversation_id) - await record_tool_call( - data, - session_id, - name=name, - arguments=arguments, - result=call_result, - duration_ms=duration_ms, - client_name=client_name, - client_version=client_version, - protocol_version=protocol_version, - conversation_id=delivered_conversation_id, - extra=extra, + await lifecycle.record_result( + call_result, + duration_ms, + conversation_id_delivered=delivered, ) return result @@ -283,34 +210,15 @@ def _inject_tool_schemas( population pass, so the schema the SDK validates against always matches the one we advertised — see the note in ``handler``. """ - context_enabled = is_context_enabled(data.options.context) - description = get_context_description(data.options.context) for tool in tools: - if tool.name == _GET_MORE_TOOLS_NAME: - continue schema = getattr(tool, "inputSchema", None) - # required follows the path: raw low-level validates the call against - # this same schema (optional), FastMCP 2.0 strips it first (required-advisory). - if context_enabled and not schema_has_param(schema, "context"): - schema = add_context_parameter_to_schema( - schema, tool.name, description, required=context_required - ) - if data.options.enable_conversation_id and not schema_has_param( - schema, "conversation_id" - ): - schema = add_conversation_id_to_schema(schema, tool.name) - if schema is not getattr(tool, "inputSchema", None): - try: - tool.inputSchema = schema - except Exception: # noqa: BLE001 - log(f"WARN: could not set inputSchema on tool {tool.name}") - # Declare the structuredContent channel and remember the answer: - # clients that read structuredContent never see the content text - # block, and only a declared key may be written back on a call. - if data.options.enable_conversation_id: - data.tool_output_instructions[tool.name] = ( - add_instructions_to_output_schema(tool) - ) + mutate_tool_schema( + data, + tool, + schema_attribute="inputSchema", + owns_context=schema_has_param(schema, "context"), + context_required=context_required, + ) def _wrap_list_tools( @@ -352,50 +260,29 @@ async def handler(req: Any) -> Any: extra = {"session_id": mcp_session_id, "ctx": _request_context(server)} # Resolve session, emit $mcp_initialize (once per session) and identify here # too — a client may list tools without ever calling one. - session_id = await prepare_request( + lifecycle = await start_tools_list_lifecycle( data, + request=request, + extra=extra, mcp_session_id=mcp_session_id, + token=token, client_name=client_name, client_version=client_version, protocol_version=protocol_version, - request=request, - extra=extra, - token=token, ) start = time.monotonic() try: result = await original(req) except Exception as error: - await record_tools_list( - data, - session_id, - names=[], - request=request, - duration_ms=(time.monotonic() - start) * 1000, - is_error=True, - error=error, - client_name=client_name, - client_version=client_version, - protocol_version=protocol_version, - extra=extra, - ) + await lifecycle.record_error(error, (time.monotonic() - start) * 1000) raise duration_ms = (time.monotonic() - start) * 1000 tools = extract_tools(result) - names = [] - for tool in tools: - names.append(tool.name) - if getattr(tool, "description", None): - data.tool_descriptions[tool.name] = tool.description - category = read_tool_category(tool) - if category: - data.tool_categories[tool.name] = category - - # Zero advertised tools is treated as an errored tools/list (parity with the - # TS SDK) — captured before we append our own get_more_tools virtual tool. - empty = len(tools) == 0 + # Zero advertised tools is treated as an errored tools/list before the + # virtual missing-capability tool is appended. + names, empty = collect_listed_tools(data, tools) _inject_tool_schemas(data, tools, context_required=context_required) @@ -405,19 +292,11 @@ async def handler(req: Any) -> Any: append_get_more_tools(result, missing_name) names.append(missing_name) - await record_tools_list( - data, - session_id, + await lifecycle.record_result( names=names, - request=request, response=_to_jsonable(result), duration_ms=duration_ms, - is_error=empty, - error="tools/list returned no tools" if empty else None, - client_name=client_name, - client_version=client_version, - protocol_version=protocol_version, - extra=extra, + is_empty=empty, ) return result diff --git a/posthog/mcp/_instrument_v2.py b/posthog/mcp/_instrument_v2.py index f97e5043..437015be 100644 --- a/posthog/mcp/_instrument_v2.py +++ b/posthog/mcp/_instrument_v2.py @@ -35,39 +35,23 @@ import mcp.types as mcp_types -from ._context_parameters import ( - add_context_parameter_to_schema, - get_context_description, - is_context_enabled, - schema_has_param, -) -from ._conversation_id import ( - add_conversation_id_to_schema, - build_prompt_back, - resolve_conversation_id, -) +from ._context_parameters import schema_has_param +from ._conversation_id import build_prompt_back from ._instrumentation import ( _to_jsonable, - build_tool_call_request, + collect_listed_tools, + mutate_tool_schema, params_to_request_dict, - prepare_request, - prime_session, - read_tool_category, - record_missing_capability, - record_tool_call, - record_tools_list, resolve_session_and_client, + start_tool_call_lifecycle, + start_tools_list_lifecycle, ) from ._internal import MCPAnalyticsData -from ._output_instructions import ( - add_instructions_to_output_schema, - mirror_instructions_into_structured_content, -) +from ._output_instructions import mirror_instructions_into_structured_content from .logger import log from .request_headers import get_request_headers from .session_token import read_mcp_session_header from .tools import ( - GET_MORE_TOOLS_NAME as _GET_MORE_TOOLS_NAME, build_report_missing_descriptor, get_more_tools_result_text, resolve_missing_capability_tool_name, @@ -261,43 +245,20 @@ async def wrapped( token, client_name, client_version, protocol_version, mcp_session_id = ( _resolve_ctx(ctx) ) - request = build_tool_call_request(name, arguments) - extra: Dict[str, Any] = {"session_id": mcp_session_id, "ctx": ctx} - - missing_name = resolve_missing_capability_tool_name(data.options) - conversation_id, minted = resolve_conversation_id( - data.options.enable_conversation_id, arguments, name, missing_name + lifecycle = start_tool_call_lifecycle( + data, + name=name, + arguments=arguments, + mcp_session_id=mcp_session_id, + token=token, + client_name=client_name, + client_version=client_version, + protocol_version=protocol_version, + extra={"session_id": mcp_session_id, "ctx": ctx}, ) - # Resolved once the handle's fate is known — a minted handle only - # anchors the session after we have confirmed the agent received it, - # so the call that mints it still joins its own conversation. - async def _session(anchor: Optional[str]) -> str: - return await prepare_request( - data, - mcp_session_id=mcp_session_id, - client_name=client_name, - client_version=client_version, - protocol_version=protocol_version, - request=request, - extra=extra, - token=token, - conversation_id=anchor, - ) - - if data.options.report_missing and name == missing_name: - session_id = await _session(None) - await record_missing_capability( - data, - session_id, - tool_name=missing_name, - context=(arguments or {}).get("context"), - arguments=arguments, - client_name=client_name, - client_version=client_version, - protocol_version=protocol_version, - extra=extra, - ) + if lifecycle.is_missing_capability: + await lifecycle.record_missing_capability() return mcp_types.CallToolResult( content=[ mcp_types.TextContent( @@ -327,7 +288,7 @@ async def _session(anchor: Optional[str]) -> str: # Settle the shared session before the tool body runs, so an in-tool # `analytics.capture()` is attributed to this caller and not the last one. - await prime_session(data, mcp_session_id=mcp_session_id, token=token) + await lifecycle.prime_session() start = time.monotonic() try: @@ -335,48 +296,24 @@ async def _session(anchor: Optional[str]) -> str: name, call_arguments, context=context, convert_result=convert_result ) except Exception as error: - # The raise is converted to CallToolResult(is_error=True) one layer - # up (MCPServer._handle_call_tool), so the prompt-back never rides - # it — a minted (undelivered) conversation_id is not stamped. - session_id = await _session(None if minted else conversation_id) - await record_tool_call( - data, - session_id, - name=name, - arguments=arguments, - error=error, - duration_ms=(time.monotonic() - start) * 1000, - client_name=client_name, - client_version=client_version, - protocol_version=protocol_version, - conversation_id=None if minted else conversation_id, - extra=extra, - ) + # The outer MCPServer layer converts the raise after this seam; no + # freshly minted handle could have been delivered yet. + await lifecycle.record_error(error, (time.monotonic() - start) * 1000) raise duration_ms = (time.monotonic() - start) * 1000 - delivered_conversation_id = conversation_id - if conversation_id: + delivered = False + if lifecycle.conversation_id: result, delivered = _deliver_conversation_id( - data, result, name, conversation_id, minted + data, + result, + name, + lifecycle.conversation_id, + lifecycle.minted_conversation_id, ) - # Only a minted handle can be lost this way — one the agent supplied, it has. - if minted and not delivered: - delivered_conversation_id = None - session_id = await _session(delivered_conversation_id) - await record_tool_call( - data, - session_id, - name=name, - arguments=arguments, - result=result, - duration_ms=duration_ms, - client_name=client_name, - client_version=client_version, - protocol_version=protocol_version, - conversation_id=delivered_conversation_id, - extra=extra, + await lifecycle.record_result( + result, duration_ms, conversation_id_delivered=delivered ) return result @@ -448,43 +385,20 @@ async def handler(ctx: Any, params: Any) -> Any: token, client_name, client_version, protocol_version, mcp_session_id = ( _resolve_ctx(ctx) ) - request = build_tool_call_request(name, arguments) - extra: Dict[str, Any] = {"session_id": mcp_session_id, "ctx": ctx} - - missing_name = resolve_missing_capability_tool_name(data.options) - conversation_id, minted = resolve_conversation_id( - data.options.enable_conversation_id, arguments, name, missing_name + lifecycle = start_tool_call_lifecycle( + data, + name=name, + arguments=arguments, + mcp_session_id=mcp_session_id, + token=token, + client_name=client_name, + client_version=client_version, + protocol_version=protocol_version, + extra={"session_id": mcp_session_id, "ctx": ctx}, ) - # Resolved once the handle's fate is known — a minted handle only - # anchors the session after we have confirmed the agent received it, - # so the call that mints it still joins its own conversation. - async def _session(anchor: Optional[str]) -> str: - return await prepare_request( - data, - mcp_session_id=mcp_session_id, - client_name=client_name, - client_version=client_version, - protocol_version=protocol_version, - request=request, - extra=extra, - token=token, - conversation_id=anchor, - ) - - if data.options.report_missing and name == missing_name: - session_id = await _session(None) - await record_missing_capability( - data, - session_id, - tool_name=missing_name, - context=arguments.get("context"), - arguments=arguments, - client_name=client_name, - client_version=client_version, - protocol_version=protocol_version, - extra=extra, - ) + if lifecycle.is_missing_capability: + await lifecycle.record_missing_capability() return mcp_types.CallToolResult( content=[ mcp_types.TextContent( @@ -495,54 +409,30 @@ async def _session(anchor: Optional[str]) -> str: # Settle the shared session before the tool body runs, so an in-tool # `analytics.capture()` is attributed to this caller and not the last one. - await prime_session(data, mcp_session_id=mcp_session_id, token=token) + await lifecycle.prime_session() start = time.monotonic() try: result = await original(ctx, params) except Exception as error: - # v2 low-level handlers raise through to JSON-RPC errors (no - # auto-conversion) — capture before re-raising. The prompt-back was - # never delivered, so a minted conversation_id is not stamped. - session_id = await _session(None if minted else conversation_id) - await record_tool_call( - data, - session_id, - name=name, - arguments=arguments, - error=error, - duration_ms=(time.monotonic() - start) * 1000, - client_name=client_name, - client_version=client_version, - protocol_version=protocol_version, - conversation_id=None if minted else conversation_id, - extra=extra, - ) + # Raw v2 handlers raise through to JSON-RPC errors; no freshly minted + # handle could have been delivered yet. + await lifecycle.record_error(error, (time.monotonic() - start) * 1000) raise duration_ms = (time.monotonic() - start) * 1000 - delivered_conversation_id = conversation_id - if conversation_id: + delivered = False + if lifecycle.conversation_id: result, delivered = _deliver_conversation_id( - data, result, name, conversation_id, minted + data, + result, + name, + lifecycle.conversation_id, + lifecycle.minted_conversation_id, ) - # Only a minted handle can be lost this way — one the agent supplied, it has. - if minted and not delivered: - delivered_conversation_id = None - session_id = await _session(delivered_conversation_id) - await record_tool_call( - data, - session_id, - name=name, - arguments=arguments, - result=result, - duration_ms=duration_ms, - client_name=client_name, - client_version=client_version, - protocol_version=protocol_version, - conversation_id=delivered_conversation_id, - extra=extra, + await lifecycle.record_result( + result, duration_ms, conversation_id_delivered=delivered ) return result @@ -573,85 +463,43 @@ async def handler(ctx: Any, params: Any) -> Any: extra: Dict[str, Any] = {"session_id": mcp_session_id, "ctx": ctx} # Resolve session, emit $mcp_initialize (once per session) and identify # here too — a client may list tools without ever calling one. - session_id = await prepare_request( + lifecycle = await start_tools_list_lifecycle( data, + request=request, + extra=extra, mcp_session_id=mcp_session_id, + token=token, client_name=client_name, client_version=client_version, protocol_version=protocol_version, - request=request, - extra=extra, - token=token, ) start = time.monotonic() try: result = await original(ctx, params) except Exception as error: - await record_tools_list( - data, - session_id, - names=[], - request=request, - duration_ms=(time.monotonic() - start) * 1000, - is_error=True, - error=error, - client_name=client_name, - client_version=client_version, - protocol_version=protocol_version, - extra=extra, - ) + await lifecycle.record_error(error, (time.monotonic() - start) * 1000) raise duration_ms = (time.monotonic() - start) * 1000 tools = list(getattr(result, "tools", []) or []) - # Zero advertised tools is treated as an errored tools/list (parity with - # the TS SDK), checked before we append our own virtual tool. - empty = len(tools) == 0 + # Empty is computed before adding the virtual missing-capability tool. + names, empty = collect_listed_tools(data, tools) - names = [] - for tool in tools: - names.append(tool.name) - if getattr(tool, "description", None): - data.tool_descriptions[tool.name] = tool.description - category = read_tool_category(tool) - if category: - data.tool_categories[tool.name] = category - - context_enabled = is_context_enabled(data.options.context) - description = get_context_description(data.options.context) for tool in tools: - if tool.name == _GET_MORE_TOOLS_NAME: - continue schema = getattr(tool, "input_schema", None) owns_context = ( _tool_owns_param_v2(high_level, tool.name, "context") if high_level is not None else schema_has_param(schema, "context") ) - # required follows the entry point: the raw low-level path validates - # the call against this same schema (optional); the high-level path - # strips before dispatch (required-advisory). - if context_enabled and not owns_context: - schema = add_context_parameter_to_schema( - schema, tool.name, description, required=context_required - ) - if data.options.enable_conversation_id and not schema_has_param( - schema, "conversation_id" - ): - schema = add_conversation_id_to_schema(schema, tool.name) - if schema is not getattr(tool, "input_schema", None): - try: - tool.input_schema = schema - except Exception: # noqa: BLE001 - some schema attrs may be read-only - log(f"WARN: could not set input_schema on tool {tool.name}") - # Declare the structuredContent channel and remember the answer: - # clients that read structuredContent never see the content text - # block, and only a declared key may be written back on a call. - if data.options.enable_conversation_id: - data.tool_output_instructions[tool.name] = ( - add_instructions_to_output_schema(tool) - ) + mutate_tool_schema( + data, + tool, + schema_attribute="input_schema", + owns_context=owns_context, + context_required=context_required, + ) if data.options.report_missing: missing_name = resolve_missing_capability_tool_name(data.options) @@ -659,19 +507,11 @@ async def handler(ctx: Any, params: Any) -> Any: _append_get_more_tools_v2(result, missing_name) names.append(missing_name) - await record_tools_list( - data, - session_id, + await lifecycle.record_result( names=names, - request=request, response=_to_jsonable(result), duration_ms=duration_ms, - is_error=empty, - error="tools/list returned no tools" if empty else None, - client_name=client_name, - client_version=client_version, - protocol_version=protocol_version, - extra=extra, + is_empty=empty, ) return result diff --git a/posthog/mcp/_instrumentation.py b/posthog/mcp/_instrumentation.py index be274c6c..b9c052b0 100644 --- a/posthog/mcp/_instrumentation.py +++ b/posthog/mcp/_instrumentation.py @@ -13,20 +13,30 @@ import concurrent.futures import os import threading +from dataclasses import dataclass from datetime import datetime, timezone from typing import Any, Dict, List, Optional from ._capture import capture_event +from ._context_parameters import ( + add_context_parameter_to_schema, + get_context_description, + is_context_enabled, + schema_has_param, +) +from ._conversation_id import add_conversation_id_to_schema, resolve_conversation_id from ._event_types import MCPAnalyticsEventType from ._exceptions import capture_exception from ._intent import resolve_tool_call_intent, set_event_intent from ._internal import MCPAnalyticsData, handle_identify, resolve_event_properties +from ._output_instructions import add_instructions_to_output_schema from .logger import log, warn from .request_headers import get_request from ._sanitization import build_captured_mcp_parameters from ._transport_identity import stamp_transport_identity from .session import resolve_session_id, resolve_session_id_with_source from .session_token import SessionTokenPayload, decode_session_id +from .tools import GET_MORE_TOOLS_NAME, resolve_missing_capability_tool_name # Keep strong refs to in-flight capture tasks/futures and their lifecycle owners so # they aren't GC'd mid-flight and lifecycle drains can select only their own work. @@ -369,6 +379,140 @@ async def prepare_request( return session_id +@dataclass(frozen=True) +class ToolCallLifecycle: + """Common analytics policy for one tool call. + + Adapters still own request-context lookup, argument stripping, dispatch, result + shape, and conversation-id delivery. This object only keeps the shared session, + missing-capability, and capture ordering in one place. + """ + + data: MCPAnalyticsData + name: str + arguments: Optional[Dict[str, Any]] + request: Dict[str, Any] + extra: Dict[str, Any] + mcp_session_id: Optional[str] + token: Optional[SessionTokenPayload] + client_name: Optional[str] + client_version: Optional[str] + protocol_version: Optional[str] + missing_name: str + conversation_id: Optional[str] + minted_conversation_id: bool + + @property + def is_missing_capability(self) -> bool: + return self.data.options.report_missing and self.name == self.missing_name + + async def prepare_session(self, conversation_id: Optional[str]) -> str: + return await prepare_request( + self.data, + mcp_session_id=self.mcp_session_id, + client_name=self.client_name, + client_version=self.client_version, + protocol_version=self.protocol_version, + request=self.request, + extra=self.extra, + token=self.token, + conversation_id=conversation_id, + ) + + async def prime_session(self) -> None: + await prime_session( + self.data, mcp_session_id=self.mcp_session_id, token=self.token + ) + + async def record_missing_capability(self) -> None: + session_id = await self.prepare_session(None) + await record_missing_capability( + self.data, + session_id, + tool_name=self.missing_name, + context=(self.arguments or {}).get("context"), + arguments=self.arguments, + client_name=self.client_name, + client_version=self.client_version, + protocol_version=self.protocol_version, + extra=self.extra, + ) + + async def record_error(self, error: Any, duration_ms: float) -> None: + # A freshly minted handle cannot anchor or be captured when dispatch + # raised: no adapter had an opportunity to deliver it to the agent. + conversation_id = None if self.minted_conversation_id else self.conversation_id + session_id = await self.prepare_session(conversation_id) + await record_tool_call( + self.data, + session_id, + name=self.name, + arguments=self.arguments, + error=error, + duration_ms=duration_ms, + client_name=self.client_name, + client_version=self.client_version, + protocol_version=self.protocol_version, + conversation_id=conversation_id, + extra=self.extra, + ) + + async def record_result( + self, result: Any, duration_ms: float, *, conversation_id_delivered: bool + ) -> None: + conversation_id = self.conversation_id + if self.minted_conversation_id and not conversation_id_delivered: + conversation_id = None + session_id = await self.prepare_session(conversation_id) + await record_tool_call( + self.data, + session_id, + name=self.name, + arguments=self.arguments, + result=result, + duration_ms=duration_ms, + client_name=self.client_name, + client_version=self.client_version, + protocol_version=self.protocol_version, + conversation_id=conversation_id, + extra=self.extra, + ) + + +def start_tool_call_lifecycle( + data: MCPAnalyticsData, + *, + name: str, + arguments: Optional[Dict[str, Any]], + mcp_session_id: Optional[str], + token: Optional[SessionTokenPayload], + client_name: Optional[str], + client_version: Optional[str], + protocol_version: Optional[str], + extra: Dict[str, Any], +) -> ToolCallLifecycle: + """Resolve adapter-independent policy for a tool call without dispatching it.""" + missing_name = resolve_missing_capability_tool_name(data.options) + conversation_id, minted = resolve_conversation_id( + data.options.enable_conversation_id, arguments, name, missing_name + ) + return ToolCallLifecycle( + data=data, + name=name, + arguments=arguments, + request=build_tool_call_request(name, arguments), + extra=extra, + mcp_session_id=mcp_session_id, + token=token, + client_name=client_name, + client_version=client_version, + protocol_version=protocol_version, + missing_name=missing_name, + conversation_id=conversation_id, + minted_conversation_id=minted, + ) + + async def record_tool_call( data: MCPAnalyticsData, session_id: str, @@ -459,6 +603,59 @@ def read_tool_category(tool: Any) -> Optional[str]: return None +def collect_listed_tools(data: MCPAnalyticsData, tools: list) -> tuple[List[str], bool]: + """Cache common tool metadata and return the pre-injection listing summary.""" + names = [] + for tool in tools: + names.append(tool.name) + if getattr(tool, "description", None): + data.tool_descriptions[tool.name] = tool.description + category = read_tool_category(tool) + if category: + data.tool_categories[tool.name] = category + return names, not tools + + +def mutate_tool_schema( + data: MCPAnalyticsData, + tool: Any, + *, + schema_attribute: str, + owns_context: bool, + context_required: bool, +) -> None: + """Apply the common analytics schema pipeline and write it back in place. + + The adapter explicitly supplies its SDK model's schema attribute and its own + 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: + 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" + ): + schema = add_conversation_id_to_schema(schema, tool.name) + if schema is not original_schema: + try: + setattr(tool, schema_attribute, schema) + except Exception: # noqa: BLE001 - some schema attrs may be read-only + log(f"WARN: could not set {schema_attribute} on tool {tool.name}") + if data.options.enable_conversation_id: + data.tool_output_instructions[tool.name] = add_instructions_to_output_schema( + tool + ) + + def request_to_dict(req: Any) -> Dict[str, Any]: """Shape a request object into the JSON-RPC-ish dict the sanitizer expects.""" method = getattr(req, "method", None) or "tools/list" @@ -483,6 +680,89 @@ def params_to_request_dict( return {"method": method, "params": params_dict} +@dataclass(frozen=True) +class ToolsListLifecycle: + """Common capture lifecycle for one client-facing tools/list dispatch.""" + + data: MCPAnalyticsData + session_id: str + request: Dict[str, Any] + extra: Dict[str, Any] + client_name: Optional[str] + client_version: Optional[str] + protocol_version: Optional[str] + + async def record_error(self, error: Any, duration_ms: float) -> None: + await record_tools_list( + self.data, + self.session_id, + names=[], + request=self.request, + duration_ms=duration_ms, + is_error=True, + error=error, + client_name=self.client_name, + client_version=self.client_version, + protocol_version=self.protocol_version, + extra=self.extra, + ) + + async def record_result( + self, + *, + names: List[str], + response: Any, + duration_ms: float, + is_empty: bool, + ) -> None: + await record_tools_list( + self.data, + self.session_id, + names=names, + request=self.request, + response=response, + duration_ms=duration_ms, + is_error=is_empty, + error="tools/list returned no tools" if is_empty else None, + client_name=self.client_name, + client_version=self.client_version, + protocol_version=self.protocol_version, + extra=self.extra, + ) + + +async def start_tools_list_lifecycle( + data: MCPAnalyticsData, + *, + request: Dict[str, Any], + extra: Dict[str, Any], + mcp_session_id: Optional[str], + token: Optional[SessionTokenPayload], + client_name: Optional[str], + client_version: Optional[str], + protocol_version: Optional[str], +) -> ToolsListLifecycle: + session_id = await prepare_request( + data, + mcp_session_id=mcp_session_id, + client_name=client_name, + client_version=client_version, + protocol_version=protocol_version, + request=request, + extra=extra, + token=token, + ) + return ToolsListLifecycle( + data=data, + session_id=session_id, + request=request, + extra=extra, + client_name=client_name, + client_version=client_version, + protocol_version=protocol_version, + ) + + async def record_missing_capability( data: MCPAnalyticsData, session_id: str, diff --git a/posthog/test/mcp/test_fastmcp.py b/posthog/test/mcp/test_fastmcp.py index b5e902c1..f93d069b 100644 --- a/posthog/test/mcp/test_fastmcp.py +++ b/posthog/test/mcp/test_fastmcp.py @@ -21,6 +21,7 @@ def make_server(): @server.tool() def add(a: int, b: int) -> int: + """Add two numbers.""" return a + b @server.tool() @@ -74,6 +75,9 @@ async def test_tool_call_captures_intent_and_strips_context(): client = FakeClient() instrument(server, client) + # Prime the listing metadata used by the shared call lifecycle. + await _list_tools(server) + received = {} original_add = server._tool_manager.get_tool("add").fn @@ -95,6 +99,7 @@ def spy_add(a: int, b: int) -> int: assert len(calls) == 1 props = calls[0]["properties"] assert props["$mcp_tool_name"] == "add" + assert props["$mcp_tool_description"] == "Add two numbers." assert props["$mcp_intent"] == "summing two numbers for the user's report" assert props["$mcp_intent_source"] == "context_parameter" assert props["$mcp_is_error"] is False diff --git a/posthog/test/mcp/test_lowlevel.py b/posthog/test/mcp/test_lowlevel.py index 79bc063f..eb50ff1e 100644 --- a/posthog/test/mcp/test_lowlevel.py +++ b/posthog/test/mcp/test_lowlevel.py @@ -85,6 +85,8 @@ async def test_tool_call_success_captures_intent(): assert len(calls) == 1 props = calls[0]["properties"] assert props["$mcp_tool_name"] == "echo" + # tools/list metadata and tools/call capture share the same lifecycle policy. + assert props["$mcp_tool_description"] == "Echo back a message" assert props["$mcp_intent"] == "echoing a message for the test" assert props["$mcp_is_error"] is False # context is stripped from captured parameters diff --git a/posthog/test/mcp/test_units.py b/posthog/test/mcp/test_units.py index 3182a38f..04c8cb02 100644 --- a/posthog/test/mcp/test_units.py +++ b/posthog/test/mcp/test_units.py @@ -3,6 +3,7 @@ end-to-end adapter tests by exercising edge branches directly.""" from datetime import datetime, timedelta, timezone +from types import SimpleNamespace from posthog.mcp._conversation_id import ( add_conversation_id_to_schema, @@ -11,6 +12,7 @@ inject_prompt_back, resolve_conversation_id, ) +from posthog.mcp._instrumentation import mutate_tool_schema from posthog.mcp._intent import _get_context_argument, resolve_tool_call_intent from posthog.mcp._internal import ( IdentityCache, @@ -100,6 +102,24 @@ def test_add_conversation_id_skips_when_already_present(): assert add_conversation_id_to_schema(schema, "t") is schema +def test_schema_pipeline_does_not_warn_for_owned_conversation_id(monkeypatch): + warnings = [] + monkeypatch.setattr("posthog.mcp._conversation_id.log", warnings.append) + schema = {"type": "object", "properties": {"conversation_id": {"type": "string"}}} + tool = SimpleNamespace(name="t", input_schema=schema) + + mutate_tool_schema( + _data(context=False, enable_conversation_id=True), + tool, + schema_attribute="input_schema", + owns_context=False, + context_required=False, + ) + + assert tool.input_schema is schema + assert warnings == [] + + def test_add_conversation_id_skips_complex_schema(): schema = {"oneOf": [{"type": "object"}]} assert add_conversation_id_to_schema(schema, "t") is schema diff --git a/posthog/test/mcp/test_v2_lowlevel.py b/posthog/test/mcp/test_v2_lowlevel.py index 788bfd20..eb543f86 100644 --- a/posthog/test/mcp/test_v2_lowlevel.py +++ b/posthog/test/mcp/test_v2_lowlevel.py @@ -108,6 +108,8 @@ async def test_tool_call_captured_with_client_identity(): client = FakeClient() instrument(server, client) + # Prime the listing metadata used by the shared call lifecycle. + await _list_tools(server) result = await _call_tool( server, "add", {"a": 2, "b": 3, "context": "adding for a report"} ) @@ -119,6 +121,7 @@ async def test_tool_call_captured_with_client_identity(): assert len(calls) == 1 props = calls[0]["properties"] assert props["$mcp_tool_name"] == "add" + assert props["$mcp_tool_description"] == "Add two numbers" assert props["$mcp_intent"] == "adding for a report" assert props["$mcp_is_error"] is False assert props["$mcp_client_name"] == "test-client" diff --git a/posthog/test/mcp/test_v2_mcpserver.py b/posthog/test/mcp/test_v2_mcpserver.py index 814e5394..fb938fb2 100644 --- a/posthog/test/mcp/test_v2_mcpserver.py +++ b/posthog/test/mcp/test_v2_mcpserver.py @@ -26,6 +26,7 @@ def make_server(): @server.tool() def add(a: int, b: int) -> int: + """Add two numbers.""" return a + b @server.tool() @@ -86,6 +87,9 @@ async def test_tool_call_captures_intent_and_strips_context(): client = FakeClient() instrument(server, client) + # Prime the listing metadata used by the shared call lifecycle. + await _list_tools(server) + received = {} original_add = server._tool_manager.get_tool("add").fn @@ -110,6 +114,7 @@ def spy_add(a: int, b: int) -> int: assert len(calls) == 1 props = calls[0]["properties"] assert props["$mcp_tool_name"] == "add" + assert props["$mcp_tool_description"] == "Add two numbers." assert props["$mcp_intent"] == "summing two numbers for the user's report" assert props["$mcp_intent_source"] == "context_parameter" assert props["$mcp_is_error"] is False @@ -360,6 +365,20 @@ async def test_report_missing_advertises_and_captures(): assert missing[0]["properties"]["$mcp_intent"] == "need a tool to send emails" +async def test_report_missing_accepts_omitted_arguments(): + server = make_server() + client = FakeClient() + instrument(server, client, MCPAnalyticsOptions(report_missing=True)) + + result = await _call_tool(server, "get_more_tools", None) + await _flush() + + assert result.is_error is False + missing = _events(client, "$mcp_missing_capability") + assert missing + assert "$mcp_intent" not in missing[0]["properties"] + + async def test_instrument_is_idempotent(): server = make_server() client = FakeClient()