From 3afedd1a32f4d1b064c1dd17c309f710518ccf2e Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Fri, 13 Mar 2026 23:10:18 +0000 Subject: [PATCH] Add OTel instrumentation for MCP tool and prompt calls following OTel 1.40.0 conventions Co-authored-by: TaoChenOSU <12570346+TaoChenOSU@users.noreply.github.com> --- python/packages/core/agent_framework/_mcp.py | 206 ++++++++++++------ .../core/agent_framework/observability.py | 39 ++++ python/packages/core/tests/core/test_mcp.py | 190 ++++++++++++++-- 3 files changed, 343 insertions(+), 92 deletions(-) diff --git a/python/packages/core/agent_framework/_mcp.py b/python/packages/core/agent_framework/_mcp.py index 81227c7e73..c990b9dd61 100644 --- a/python/packages/core/agent_framework/_mcp.py +++ b/python/packages/core/agent_framework/_mcp.py @@ -25,7 +25,7 @@ from mcp.client.websocket import websocket_client from mcp.shared.context import RequestContext from mcp.shared.exceptions import McpError from mcp.shared.session import RequestResponder -from opentelemetry import propagate +from opentelemetry import propagate, trace from ._tools import FunctionTool from ._types import ( @@ -33,6 +33,7 @@ from ._types import ( Message, ) from .exceptions import ToolException, ToolExecutionException +from .observability import OtelAttr, get_mcp_call_span if sys.version_info >= (3, 11): from typing import Self # pragma: no cover @@ -487,6 +488,7 @@ class MCPTool: self.is_connected: bool = False self._tools_loaded: bool = False self._prompts_loaded: bool = False + self._mcp_protocol_version: str | None = None def __str__(self) -> str: return f"MCPTool(name={self.name}, description={self.description})" @@ -590,7 +592,8 @@ class MCPTool: inner_exception=ex, ) from ex try: - await session.initialize() + init_result = await session.initialize() + self._mcp_protocol_version = init_result.protocolVersion except Exception as ex: await self._safe_close_exit_stack() # Provide context about initialization failure @@ -605,7 +608,8 @@ class MCPTool: self.session = session elif self.session._request_id == 0: # type: ignore[reportPrivateUsage] # If the session is not initialized, we need to reinitialize it - await self.session.initialize() + init_result = await self.session.initialize() + self._mcp_protocol_version = init_result.protocolVersion logger.debug("Connected to MCP server: %s", self.session) self.is_connected = True if self.load_tools_flag: @@ -927,50 +931,83 @@ class MCPTool: } } - # Inject OpenTelemetry trace context into MCP _meta for distributed tracing. - otel_meta = _inject_otel_into_mcp_meta() - parser = self.parse_tool_results or _parse_tool_result_from_mcp - # Try the operation, reconnecting once if the connection is closed - for attempt in range(2): - try: - result = await self.session.call_tool(tool_name, arguments=filtered_kwargs, meta=otel_meta) # type: ignore - if result.isError: - parsed = parser(result) - text = ( - "\n".join(c.text for c in parsed if c.type == "text" and c.text) - if isinstance(parsed, list) - else str(parsed) - ) - raise ToolExecutionException(text or str(parsed)) - return parser(result) - except ToolExecutionException: - raise - except ClosedResourceError as cl_ex: - if attempt == 0: - # First attempt failed, try reconnecting - logger.info("MCP connection closed unexpectedly. Reconnecting...") - try: - await self.connect(reset=True) - continue # Retry the operation - except Exception as reconn_ex: + span_attributes: dict[str, Any] = { + OtelAttr.MCP_METHOD_NAME: "tools/call", + OtelAttr.TOOL_NAME: tool_name, + OtelAttr.OPERATION: OtelAttr.TOOL_EXECUTION_OPERATION, + OtelAttr.JSONRPC_PROTOCOL_VERSION: "2.0", + } + if self._mcp_protocol_version: + span_attributes[OtelAttr.MCP_PROTOCOL_VERSION] = self._mcp_protocol_version + + with get_mcp_call_span(span_attributes) as span: + # Try the operation, reconnecting once if the connection is closed + for attempt in range(2): + try: + # Capture the JSON-RPC request ID before the call is made. + # The MCP SDK stores the next request ID in the private `_request_id` + # attribute; no public API is available. We use getattr with a default + # so this degrades gracefully if the attribute is renamed in a future + # version of the library. + request_id = getattr(self.session, "_request_id", None) # type: ignore[union-attr] + if request_id is not None: + span.set_attribute(OtelAttr.JSONRPC_REQUEST_ID, str(request_id)) + + # Inject OpenTelemetry trace context into MCP _meta for distributed tracing. + otel_meta = _inject_otel_into_mcp_meta() + + result = await self.session.call_tool(tool_name, arguments=filtered_kwargs, meta=otel_meta) # type: ignore + if result.isError: + parsed = parser(result) + text = ( + "\n".join(c.text for c in parsed if c.type == "text" and c.text) + if isinstance(parsed, list) + else str(parsed) + ) + error_msg = text or str(parsed) + span.set_attribute(OtelAttr.ERROR_TYPE, "ToolError") + span.set_status(trace.StatusCode.ERROR, error_msg) + raise ToolExecutionException(error_msg) + return parser(result) + except ToolExecutionException: + raise + except ClosedResourceError as cl_ex: + if attempt == 0: + # First attempt failed, try reconnecting + logger.info("MCP connection closed unexpectedly. Reconnecting...") + try: + await self.connect(reset=True) + continue # Retry the operation + except Exception as reconn_ex: + error_type = type(reconn_ex).__name__ + span.set_attribute(OtelAttr.ERROR_TYPE, error_type) + span.set_status(trace.StatusCode.ERROR, str(reconn_ex)) + raise ToolExecutionException( + "Failed to reconnect to MCP server.", + inner_exception=reconn_ex, + ) from reconn_ex + else: + # Second attempt also failed, give up + logger.error(f"MCP connection closed unexpectedly after reconnection: {cl_ex}") + span.set_attribute(OtelAttr.ERROR_TYPE, type(cl_ex).__name__) + span.set_status(trace.StatusCode.ERROR, str(cl_ex)) raise ToolExecutionException( - "Failed to reconnect to MCP server.", - inner_exception=reconn_ex, - ) from reconn_ex - else: - # Second attempt also failed, give up - logger.error(f"MCP connection closed unexpectedly after reconnection: {cl_ex}") - raise ToolExecutionException( - f"Failed to call tool '{tool_name}' - connection lost.", - inner_exception=cl_ex, - ) from cl_ex - except McpError as mcp_exc: - raise ToolExecutionException(mcp_exc.error.message, inner_exception=mcp_exc) from mcp_exc - except Exception as ex: - raise ToolExecutionException(f"Failed to call tool '{tool_name}'.", inner_exception=ex) from ex - raise ToolExecutionException(f"Failed to call tool '{tool_name}' after retries.") + f"Failed to call tool '{tool_name}' - connection lost.", + inner_exception=cl_ex, + ) from cl_ex + except McpError as mcp_exc: + span.set_attribute(OtelAttr.ERROR_TYPE, str(mcp_exc.error.code)) + span.set_status(trace.StatusCode.ERROR, mcp_exc.error.message) + raise ToolExecutionException(mcp_exc.error.message, inner_exception=mcp_exc) from mcp_exc + except Exception as ex: + span.set_attribute(OtelAttr.ERROR_TYPE, type(ex).__name__) + span.set_status(trace.StatusCode.ERROR, str(ex)) + raise ToolExecutionException(f"Failed to call tool '{tool_name}'.", inner_exception=ex) from ex + span.set_attribute(OtelAttr.ERROR_TYPE, "RetryExhausted") + span.set_status(trace.StatusCode.ERROR, f"Failed to call tool '{tool_name}' after retries.") + raise ToolExecutionException(f"Failed to call tool '{tool_name}' after retries.") async def get_prompt(self, prompt_name: str, **kwargs: Any) -> str: """Call a prompt with the given arguments. @@ -995,35 +1032,60 @@ class MCPTool: parser = self.parse_prompt_results or _parse_prompt_result_from_mcp - # Try the operation, reconnecting once if the connection is closed - for attempt in range(2): - try: - prompt_result = await self.session.get_prompt(prompt_name, arguments=kwargs) # type: ignore - return parser(prompt_result) - except ClosedResourceError as cl_ex: - if attempt == 0: - # First attempt failed, try reconnecting - logger.info("MCP connection closed unexpectedly. Reconnecting...") - try: - await self.connect(reset=True) - continue # Retry the operation - except Exception as reconn_ex: + span_attributes: dict[str, Any] = { + OtelAttr.MCP_METHOD_NAME: "prompts/get", + OtelAttr.PROMPT_NAME: prompt_name, + OtelAttr.JSONRPC_PROTOCOL_VERSION: "2.0", + } + if self._mcp_protocol_version: + span_attributes[OtelAttr.MCP_PROTOCOL_VERSION] = self._mcp_protocol_version + + with get_mcp_call_span(span_attributes) as span: + # Try the operation, reconnecting once if the connection is closed + for attempt in range(2): + try: + # Capture the JSON-RPC request ID before the call is made. + # See call_tool for rationale on using getattr with a default. + request_id = getattr(self.session, "_request_id", None) # type: ignore[union-attr] + if request_id is not None: + span.set_attribute(OtelAttr.JSONRPC_REQUEST_ID, str(request_id)) + + prompt_result = await self.session.get_prompt(prompt_name, arguments=kwargs) # type: ignore + return parser(prompt_result) + except ClosedResourceError as cl_ex: + if attempt == 0: + # First attempt failed, try reconnecting + logger.info("MCP connection closed unexpectedly. Reconnecting...") + try: + await self.connect(reset=True) + continue # Retry the operation + except Exception as reconn_ex: + span.set_attribute(OtelAttr.ERROR_TYPE, type(reconn_ex).__name__) + span.set_status(trace.StatusCode.ERROR, str(reconn_ex)) + raise ToolExecutionException( + "Failed to reconnect to MCP server.", + inner_exception=reconn_ex, + ) from reconn_ex + else: + # Second attempt also failed, give up + logger.error(f"MCP connection closed unexpectedly after reconnection: {cl_ex}") + span.set_attribute(OtelAttr.ERROR_TYPE, type(cl_ex).__name__) + span.set_status(trace.StatusCode.ERROR, str(cl_ex)) raise ToolExecutionException( - "Failed to reconnect to MCP server.", - inner_exception=reconn_ex, - ) from reconn_ex - else: - # Second attempt also failed, give up - logger.error(f"MCP connection closed unexpectedly after reconnection: {cl_ex}") - raise ToolExecutionException( - f"Failed to call prompt '{prompt_name}' - connection lost.", - inner_exception=cl_ex, - ) from cl_ex - except McpError as mcp_exc: - raise ToolExecutionException(mcp_exc.error.message, inner_exception=mcp_exc) from mcp_exc - except Exception as ex: - raise ToolExecutionException(f"Failed to call prompt '{prompt_name}'.", inner_exception=ex) from ex - raise ToolExecutionException(f"Failed to get prompt '{prompt_name}' after retries.") + f"Failed to call prompt '{prompt_name}' - connection lost.", + inner_exception=cl_ex, + ) from cl_ex + except McpError as mcp_exc: + span.set_attribute(OtelAttr.ERROR_TYPE, str(mcp_exc.error.code)) + span.set_status(trace.StatusCode.ERROR, mcp_exc.error.message) + raise ToolExecutionException(mcp_exc.error.message, inner_exception=mcp_exc) from mcp_exc + except Exception as ex: + span.set_attribute(OtelAttr.ERROR_TYPE, type(ex).__name__) + span.set_status(trace.StatusCode.ERROR, str(ex)) + raise ToolExecutionException(f"Failed to call prompt '{prompt_name}'.", inner_exception=ex) from ex + span.set_attribute(OtelAttr.ERROR_TYPE, "RetryExhausted") + span.set_status(trace.StatusCode.ERROR, f"Failed to get prompt '{prompt_name}' after retries.") + raise ToolExecutionException(f"Failed to get prompt '{prompt_name}' after retries.") async def __aenter__(self) -> Self: """Enter the async context manager. diff --git a/python/packages/core/agent_framework/observability.py b/python/packages/core/agent_framework/observability.py index bcc4e1365d..7e6d2b1ea9 100644 --- a/python/packages/core/agent_framework/observability.py +++ b/python/packages/core/agent_framework/observability.py @@ -79,6 +79,7 @@ __all__ = [ "create_metric_views", "create_resource", "enable_instrumentation", + "get_mcp_call_span", "get_meter", "get_tracer", ] @@ -272,6 +273,15 @@ class OtelAttr(str, Enum): AGENT_CREATE_OPERATION = "create_agent" AGENT_INVOKE_OPERATION = "invoke_agent" + # MCP-specific attributes + # https://github.com/open-telemetry/semantic-conventions/blob/main/docs/gen-ai/mcp.md + MCP_METHOD_NAME = "mcp.method.name" + MCP_PROTOCOL_VERSION = "mcp.protocol.version" + MCP_SESSION_ID = "mcp.session.id" + JSONRPC_REQUEST_ID = "jsonrpc.request.id" + JSONRPC_PROTOCOL_VERSION = "jsonrpc.protocol.version" + PROMPT_NAME = "gen_ai.prompt.name" + # Agent Framework specific attributes MEASUREMENT_FUNCTION_TAG_NAME = "agent_framework.function.name" MEASUREMENT_FUNCTION_INVOCATION_DURATION = "agent_framework.function.invocation.duration" @@ -1683,6 +1693,35 @@ def get_function_span( ) +def get_mcp_call_span( + attributes: dict[str, Any], +) -> _AgnosticContextManager[trace.Span]: + """Start a CLIENT span for an MCP call (tool or prompt). + + Creates a span following the OTel 1.40.0 semantic conventions for MCP: + https://github.com/open-telemetry/semantic-conventions/blob/main/docs/gen-ai/mcp.md + + Args: + attributes: The span attributes. Must contain ``mcp.method.name``. + When the call is tool-related, ``gen_ai.tool.name`` is used in the span name. + When the call is prompt-related, ``gen_ai.prompt.name`` is used instead. + + Returns: + A context manager that starts the span as the current span. + """ + method_name = attributes.get(OtelAttr.MCP_METHOD_NAME, "mcp") + target = attributes.get(OtelAttr.TOOL_NAME) or attributes.get(OtelAttr.PROMPT_NAME) + span_name = f"{method_name} {target}" if target else method_name + return get_tracer().start_as_current_span( + name=span_name, + kind=trace.SpanKind.CLIENT, + attributes=attributes, + set_status_on_exception=False, + end_on_exit=True, + record_exception=False, + ) + + @contextlib.contextmanager def _get_span( attributes: dict[str, Any], diff --git a/python/packages/core/tests/core/test_mcp.py b/python/packages/core/tests/core/test_mcp.py index 70aff972fe..323e4398a9 100644 --- a/python/packages/core/tests/core/test_mcp.py +++ b/python/packages/core/tests/core/test_mcp.py @@ -2872,14 +2872,11 @@ async def test_mcp_tool_filters_framework_kwargs(): @pytest.mark.parametrize( - "use_span,expect_traceparent", - [ - (True, True), - (False, False), - ], + "use_parent_span", + [True, False], ) -async def test_mcp_tool_call_tool_otel_meta(use_span, expect_traceparent, span_exporter): - """call_tool propagates OTel trace context via meta only when a span is active.""" +async def test_mcp_tool_call_tool_otel_meta(use_parent_span, span_exporter): + """call_tool always propagates OTel trace context via meta because it creates its own CLIENT span.""" from opentelemetry import trace class TestServer(MCPTool): @@ -2911,25 +2908,178 @@ async def test_mcp_tool_call_tool_otel_meta(use_span, expect_traceparent, span_e async with server: await server.load_tools() - if use_span: + if use_parent_span: tracer = trace.get_tracer("test") with tracer.start_as_current_span("test_span"): await server.functions[0].invoke(param="test_value") else: - # Use an invalid span to ensure no trace context is injected; - # call server.call_tool directly to bypass FunctionTool.invoke's own span. - with trace.use_span(trace.NonRecordingSpan(trace.INVALID_SPAN_CONTEXT)): - await server.call_tool("test_tool", param="test_value") + # call_tool creates its own MCP CLIENT span, so traceparent is always injected. + await server.call_tool("test_tool", param="test_value") meta = server.session.call_tool.call_args.kwargs.get("meta") - if expect_traceparent: - # When a valid span is active, we expect some propagation fields to be injected, - # but we do not assume any specific header name to keep this test propagator-agnostic. - assert meta is not None - assert isinstance(meta, dict) - assert len(meta) > 0 - else: - assert meta is None + # call_tool always creates an OTel span, so traceparent is always propagated. + assert meta is not None + assert isinstance(meta, dict) + assert len(meta) > 0 + + +# endregion + + +# region: OTel MCP spans + + +@pytest.fixture +def mcp_test_server_class(): + """Factory for a minimal MCPTool subclass usable in span tests.""" + + class _TestServer(MCPTool): + def __init__(self, **kwargs): + super().__init__(name="test_server", **kwargs) + self._mock_session = Mock(spec=ClientSession) + # MCP SDK stores _request_id as int; call_tool converts it to str for the span. + self._mock_session._request_id = 42 + self._mock_session.call_tool = AsyncMock( + return_value=types.CallToolResult( + content=[types.TextContent(type="text", text="ok")] + ) + ) + self._mock_session.get_prompt = AsyncMock( + return_value=types.GetPromptResult( + messages=[ + types.PromptMessage( + role="user", + content=types.TextContent(type="text", text="hello"), + ) + ] + ) + ) + self._mock_session.list_tools = AsyncMock( + return_value=types.ListToolsResult(tools=[]) + ) + + async def connect(self): + self.session = self._mock_session + self._tools_loaded = True + self._prompts_loaded = True + self.is_connected = True + + def get_mcp_client(self): + return None + + return _TestServer + + +async def test_call_tool_creates_mcp_client_span(span_exporter, mcp_test_server_class): + """call_tool emits a CLIENT span following OTel MCP conventions.""" + from opentelemetry.trace import SpanKind + + server = mcp_test_server_class() + async with server: + await server.call_tool("my_tool", param="value") + + spans = span_exporter.get_finished_spans() + assert len(spans) == 1 + span = spans[0] + assert span.name == "tools/call my_tool" + assert span.kind == SpanKind.CLIENT + assert span.attributes["mcp.method.name"] == "tools/call" + assert span.attributes["gen_ai.tool.name"] == "my_tool" + assert span.attributes["gen_ai.operation.name"] == "execute_tool" + assert span.attributes["jsonrpc.protocol.version"] == "2.0" + assert span.attributes["jsonrpc.request.id"] == "42" + + +async def test_call_tool_span_sets_error_on_mcp_error(span_exporter, mcp_test_server_class): + """call_tool sets ERROR status and error.type when McpError is raised.""" + from mcp.shared.exceptions import McpError + from opentelemetry.trace import StatusCode + + server = mcp_test_server_class() + async with server: + server._mock_session.call_tool = AsyncMock( + side_effect=McpError( + types.ErrorData(code=-32600, message="bad request") + ) + ) + with pytest.raises(Exception): + await server.call_tool("my_tool") + + spans = span_exporter.get_finished_spans() + assert len(spans) == 1 + span = spans[0] + assert span.status.status_code == StatusCode.ERROR + assert span.attributes["error.type"] == "-32600" + + +async def test_call_tool_span_sets_error_on_tool_error_result(span_exporter, mcp_test_server_class): + """call_tool sets ERROR status when the tool result itself is an error.""" + from opentelemetry.trace import StatusCode + + server = mcp_test_server_class() + async with server: + server._mock_session.call_tool = AsyncMock( + return_value=types.CallToolResult( + isError=True, + content=[types.TextContent(type="text", text="tool error")], + ) + ) + with pytest.raises(Exception): + await server.call_tool("my_tool") + + spans = span_exporter.get_finished_spans() + assert len(spans) == 1 + span = spans[0] + assert span.status.status_code == StatusCode.ERROR + assert span.attributes["error.type"] == "ToolError" + + +async def test_call_tool_span_includes_protocol_version(span_exporter, mcp_test_server_class): + """call_tool includes mcp.protocol.version when set.""" + server = mcp_test_server_class() + async with server: + server._mcp_protocol_version = "2025-06-18" + await server.call_tool("my_tool") + + spans = span_exporter.get_finished_spans() + assert len(spans) == 1 + assert spans[0].attributes["mcp.protocol.version"] == "2025-06-18" + + +async def test_get_prompt_creates_mcp_client_span(span_exporter, mcp_test_server_class): + """get_prompt emits a CLIENT span following OTel MCP conventions.""" + from opentelemetry.trace import SpanKind + + server = mcp_test_server_class() + async with server: + await server.get_prompt("my_prompt", arg="value") + + spans = span_exporter.get_finished_spans() + assert len(spans) == 1 + span = spans[0] + assert span.name == "prompts/get my_prompt" + assert span.kind == SpanKind.CLIENT + assert span.attributes["mcp.method.name"] == "prompts/get" + assert span.attributes["gen_ai.prompt.name"] == "my_prompt" + assert span.attributes["jsonrpc.protocol.version"] == "2.0" + assert span.attributes["jsonrpc.request.id"] == "42" + + +async def test_get_prompt_span_sets_error_on_exception(span_exporter, mcp_test_server_class): + """get_prompt sets ERROR status and error.type on exception.""" + from opentelemetry.trace import StatusCode + + server = mcp_test_server_class() + async with server: + server._mock_session.get_prompt = AsyncMock(side_effect=RuntimeError("fail")) + with pytest.raises(Exception): + await server.get_prompt("my_prompt") + + spans = span_exporter.get_finished_spans() + assert len(spans) == 1 + span = spans[0] + assert span.status.status_code == StatusCode.ERROR + assert span.attributes["error.type"] == "RuntimeError" # endregion