refactor: simplify MCP spans — remove enrichment logic and protocol version caching

- Always create nested CLIENT spans for tools/call instead of enriching
  the parent execute_tool span
- Remove _ACTIVE_TOOL_EXECUTION_SPAN contextvar (no longer needed)
- Remove enrich_span_with_mcp_attributes() helper
- Remove _otel_error_type preservation in FunctionTool.invoke()
- Remove _mcp_protocol_version instance variable; protocol version is
  only set on the initialize span where it is available

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
This commit is contained in:
Tao Chen
2026-06-04 15:05:46 -07:00
Unverified
parent df32fc02f0
commit f666dd516f
2 changed files with 7 additions and 30 deletions
+7 -11
View File
@@ -285,7 +285,6 @@ class MCPTool:
self._supports_logging: bool | None = None
self._ping_available: bool = True
self._pending_reload_tasks: set[asyncio.Task[None]] = set()
self._mcp_protocol_version: str | None = None
def __str__(self) -> str:
return f"MCPTool(name={self.name}, description={self.description})"
@@ -295,10 +294,7 @@ class MCPTool:
Subclasses override to add transport-specific attributes (server address, port, etc.).
"""
attrs: dict[str, Any] = {}
if self._mcp_protocol_version:
attrs[OtelAttr.MCP_PROTOCOL_VERSION] = self._mcp_protocol_version
return attrs
return {}
def _parse_prompt_result_from_mcp(
self,
@@ -812,9 +808,9 @@ class MCPTool:
try:
with create_mcp_client_span("initialize", attributes=self._mcp_base_span_attributes()) as init_span:
initialize_result = await session.initialize()
self._mcp_protocol_version = getattr(initialize_result, "protocolVersion", None)
if self._mcp_protocol_version:
init_span.set_attribute(OtelAttr.MCP_PROTOCOL_VERSION, self._mcp_protocol_version)
protocol_version = getattr(initialize_result, "protocolVersion", None)
if protocol_version:
init_span.set_attribute(OtelAttr.MCP_PROTOCOL_VERSION, protocol_version)
self._set_server_capabilities(getattr(initialize_result, "capabilities", None))
except (Exception, asyncio.CancelledError) as ex:
if await self._close_and_check_cancelled(ex):
@@ -835,9 +831,9 @@ class MCPTool:
# If the session is not initialized, we need to reinitialize it
with create_mcp_client_span("initialize", attributes=self._mcp_base_span_attributes()) as init_span:
initialize_result = await self.session.initialize()
self._mcp_protocol_version = getattr(initialize_result, "protocolVersion", None)
if self._mcp_protocol_version:
init_span.set_attribute(OtelAttr.MCP_PROTOCOL_VERSION, self._mcp_protocol_version)
protocol_version = getattr(initialize_result, "protocolVersion", None)
if protocol_version:
init_span.set_attribute(OtelAttr.MCP_PROTOCOL_VERSION, protocol_version)
self._set_server_capabilities(getattr(initialize_result, "capabilities", None))
elif self._server_capabilities is None:
self._set_server_capabilities(getattr(self.session, "_server_capabilities", None))
@@ -28,7 +28,6 @@ def _make_connected_mcp_tool(
*,
supports_tools: bool = True,
supports_prompts: bool = True,
protocol_version: str | None = "2025-06-18",
) -> MCPTool:
"""Create an MCPTool with a mocked session, ready for testing."""
tool = MCPTool(name=name)
@@ -38,7 +37,6 @@ def _make_connected_mcp_tool(
tool._supports_prompts = supports_prompts
tool.load_tools_flag = True
tool.load_prompts_flag = True
tool._mcp_protocol_version = protocol_version
return tool
@@ -129,7 +127,6 @@ async def test_mcp_initialize_span(span_exporter: InMemorySpanExporter):
from agent_framework.observability import OtelAttr
with create_mcp_client_span("initialize", attributes=self_._mcp_base_span_attributes()) as init_span:
self_._mcp_protocol_version = "2025-06-18"
init_span.set_attribute(OtelAttr.MCP_PROTOCOL_VERSION, "2025-06-18")
self_.session = mock_session_cls
@@ -167,7 +164,6 @@ async def test_mcp_tools_list_span(span_exporter: InMemorySpanExporter):
span = list_spans[0]
assert span.kind == SpanKind.CLIENT
assert span.attributes[OtelAttr.MCP_METHOD_NAME] == "tools/list"
assert span.attributes.get(OtelAttr.MCP_PROTOCOL_VERSION) == "2025-06-18"
# endregion
@@ -215,7 +211,6 @@ async def test_mcp_tools_call_creates_client_span_when_no_parent(span_exporter:
assert span.name == "tools/call get-weather"
assert span.attributes[OtelAttr.MCP_METHOD_NAME] == "tools/call"
assert span.attributes[OtelAttr.TOOL_NAME] == "get-weather"
assert span.attributes.get(OtelAttr.MCP_PROTOCOL_VERSION) == "2025-06-18"
async def test_mcp_tools_call_tool_error_sets_error_type(span_exporter: InMemorySpanExporter):
@@ -315,20 +310,6 @@ def test_mcp_websocket_tool_default_port():
assert attrs[OtelAttr.PORT] == 443
def test_mcp_protocol_version_in_attributes():
"""Protocol version should be included in span attributes."""
tool = _make_connected_mcp_tool(protocol_version="2025-06-18")
attrs = tool._mcp_base_span_attributes()
assert attrs[OtelAttr.MCP_PROTOCOL_VERSION] == "2025-06-18"
def test_mcp_protocol_version_omitted_when_none():
"""Protocol version should be omitted from span attributes when unknown."""
tool = _make_connected_mcp_tool(protocol_version=None)
attrs = tool._mcp_base_span_attributes()
assert OtelAttr.MCP_PROTOCOL_VERSION not in attrs
# endregion