diff --git a/tests/test_mcp_trace_propagation.py b/tests/test_mcp_trace_propagation.py new file mode 100644 index 000000000000..48e18bc2ad18 --- /dev/null +++ b/tests/test_mcp_trace_propagation.py @@ -0,0 +1,218 @@ +"""Tests for tools/mcp_trace_propagation.py — opt-in W3C trace propagation. + +Covers the three properties the feature promises: + +1. Gate — everything is a no-op unless ``mcp.trace_propagation`` is true. +2. Capture — the caller's context is read on the calling thread via the + standard propagation API (or a registered provider), validated against + the W3C grammar, and never raises. +3. Injection — the header exists on the shared client exactly for the + duration of one RPC, restores any pre-existing value, and no-ops for + stdio transports (client=None). + +Plus the design-rationale test: capture on a different thread than the one +owning the span yields nothing, which is why capture must happen before the +call crosses onto the MCP daemon loop. + +No live network calls; the OpenTelemetry API is faked through the provider +hook and a stub ``opentelemetry`` module so the suite passes with or +without the real SDK installed. +""" + +import sys +import threading +import types + +import pytest + +from tools import mcp_trace_propagation as mtp + + +VALID_TP = "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01" + + +@pytest.fixture(autouse=True) +def _clean_state(tmp_path, monkeypatch): + """Isolate HERMES_HOME, clear any registered provider, default gate off.""" + monkeypatch.setenv("HERMES_HOME", str(tmp_path)) + mtp.register_traceparent_provider(None) + yield + mtp.register_traceparent_provider(None) + + +def _enable(monkeypatch, enabled=True): + monkeypatch.setattr( + mtp, "is_enabled", lambda: enabled, + ) + + +def _fake_config(monkeypatch, cfg): + import hermes_cli.config as config_mod + + monkeypatch.setattr(config_mod, "load_config_readonly", lambda: cfg) + + +# --------------------------------------------------------------------------- +# The gate +# --------------------------------------------------------------------------- + +class TestGate: + def test_disabled_by_default(self, monkeypatch): + _fake_config(monkeypatch, {}) + assert mtp.is_enabled() is False + + def test_enabled_via_config(self, monkeypatch): + _fake_config(monkeypatch, {"mcp": {"trace_propagation": True}}) + assert mtp.is_enabled() is True + + def test_null_mcp_section_is_off(self, monkeypatch): + # A config carrying `mcp:` with no body loads as {"mcp": None}. + _fake_config(monkeypatch, {"mcp": None}) + assert mtp.is_enabled() is False + + def test_config_error_means_off(self, monkeypatch): + import hermes_cli.config as config_mod + + def boom(): + raise RuntimeError("unreadable config") + + monkeypatch.setattr(config_mod, "load_config_readonly", boom) + assert mtp.is_enabled() is False + + def test_disabled_gate_short_circuits_even_with_a_provider(self, monkeypatch): + _enable(monkeypatch, False) + mtp.register_traceparent_provider(lambda: VALID_TP) + assert mtp.current_traceparent() is None + + +# --------------------------------------------------------------------------- +# Capture — provider path +# --------------------------------------------------------------------------- + +class TestProviderCapture: + def test_valid_provider_output_is_used(self, monkeypatch): + _enable(monkeypatch) + mtp.register_traceparent_provider(lambda: VALID_TP) + assert mtp.current_traceparent() == VALID_TP + + @pytest.mark.parametrize( + "bad", + [ + None, + "", + "not-a-traceparent", + "01-" + "a" * 32 + "-" + "b" * 16 + "-01", # unknown version + "00-" + "A" * 32 + "-" + "b" * 16 + "-01", # uppercase hex + "00-" + "0" * 32 + "-" + "b" * 16 + "-01", # all-zero trace-id + "00-" + "a" * 32 + "-" + "0" * 16 + "-01", # all-zero parent-id + "00-" + "a" * 32 + "-" + "b" * 16 + "-01; evil: header", + 42, + ], + ) + def test_malformed_provider_output_is_discarded(self, monkeypatch, bad): + _enable(monkeypatch) + mtp.register_traceparent_provider(lambda: bad) + assert mtp.current_traceparent() is None + + def test_raising_provider_is_survived(self, monkeypatch): + _enable(monkeypatch) + + def boom(): + raise RuntimeError("plugin bug") + + mtp.register_traceparent_provider(boom) + assert mtp.current_traceparent() is None + + +# --------------------------------------------------------------------------- +# Capture — standard propagation API path (stubbed opentelemetry) +# --------------------------------------------------------------------------- + +class TestAmbientCapture: + def _stub_otel(self, monkeypatch, inject): + stub = types.ModuleType("opentelemetry") + stub.propagate = types.SimpleNamespace(inject=inject) + monkeypatch.setitem(sys.modules, "opentelemetry", stub) + + def test_ambient_context_is_captured(self, monkeypatch): + _enable(monkeypatch) + self._stub_otel( + monkeypatch, + lambda carrier: carrier.__setitem__("traceparent", VALID_TP), + ) + assert mtp.current_traceparent() == VALID_TP + + def test_no_active_span_yields_none(self, monkeypatch): + _enable(monkeypatch) + self._stub_otel(monkeypatch, lambda carrier: None) # injects nothing + assert mtp.current_traceparent() is None + + def test_missing_sdk_yields_none(self, monkeypatch): + _enable(monkeypatch) + monkeypatch.setitem(sys.modules, "opentelemetry", None) # ImportError + assert mtp.current_traceparent() is None + + def test_capture_on_the_wrong_thread_sees_nothing(self, monkeypatch): + """The design constraint: ambient context is thread-local. Capturing + on any thread but the span's own yields nothing — which is why the + tool handler captures BEFORE the call crosses to the MCP loop.""" + _enable(monkeypatch) + span_thread = threading.current_thread() + + def thread_local_inject(carrier): + if threading.current_thread() is span_thread: + carrier["traceparent"] = VALID_TP + + self._stub_otel(monkeypatch, thread_local_inject) + + assert mtp.current_traceparent() == VALID_TP # span's own thread + + seen_elsewhere = [] + other = threading.Thread( + target=lambda: seen_elsewhere.append(mtp.current_traceparent()) + ) + other.start() + other.join() + assert seen_elsewhere == [None] # the daemon thread would see this + + +# --------------------------------------------------------------------------- +# Injection +# --------------------------------------------------------------------------- + +class _FakeClient: + """Just the surface injected_headers touches: a headers mapping.""" + + def __init__(self, headers=None): + self.headers = dict(headers or {}) + + +class TestInjectedHeaders: + def test_header_exists_exactly_for_the_duration_of_the_block(self): + client = _FakeClient() + with mtp.injected_headers(client, VALID_TP): + assert client.headers["traceparent"] == VALID_TP + assert "traceparent" not in client.headers + + def test_header_is_removed_even_when_the_rpc_raises(self): + client = _FakeClient() + with pytest.raises(RuntimeError): + with mtp.injected_headers(client, VALID_TP): + raise RuntimeError("transport dropped") + assert "traceparent" not in client.headers + + def test_preexisting_header_is_restored_not_dropped(self): + client = _FakeClient({"traceparent": "00-" + "c" * 32 + "-" + "d" * 16 + "-00"}) + with mtp.injected_headers(client, VALID_TP): + assert client.headers["traceparent"] == VALID_TP + assert client.headers["traceparent"].startswith("00-cccc") + + def test_stdio_transport_is_a_noop(self): + with mtp.injected_headers(None, VALID_TP): + pass # nothing to assert — it must simply not raise + + def test_no_traceparent_is_a_noop(self): + client = _FakeClient({"x": "y"}) + with mtp.injected_headers(client, None): + assert client.headers == {"x": "y"} + assert client.headers == {"x": "y"} diff --git a/tools/mcp_tool.py b/tools/mcp_tool.py index 993d9a13c80f..5657e36346e9 100644 --- a/tools/mcp_tool.py +++ b/tools/mcp_tool.py @@ -111,6 +111,7 @@ from typing import Any, Coroutine, Dict, List, Optional, Set, Tuple from urllib.parse import urlparse +from tools.mcp_trace_propagation import current_traceparent, injected_headers from tools.registry import tool_error logger = logging.getLogger(__name__) @@ -1840,11 +1841,16 @@ class MCPServerTask: "_idle_timeout_seconds", "_max_lifetime_seconds", "_recycled_reason", "initialize_result", "_ping_unsupported", "_reconnect_retries", "_session_proven", "_was_parked", + "_http_client", ) def __init__(self, name: str): self.name = name self.session: Optional[Any] = None + # The transport's shared httpx client, exposed so tool calls can + # inject per-RPC trace-context headers (None for stdio, and while + # no HTTP transport is live). See tools/mcp_trace_propagation.py. + self._http_client: Optional[Any] = None self.tool_timeout: float = _DEFAULT_TOOL_TIMEOUT self._task: Optional[asyncio.Task] = None self._ready = asyncio.Event() @@ -2922,6 +2928,7 @@ async def _strip_auth_on_cross_origin_redirect(response): # http_client is provided, so we wrap in async-with. try: async with httpx.AsyncClient(**client_kwargs) as http_client: + self._http_client = http_client async with streamable_http_client(url, http_client=http_client) as ( read_stream, write_stream, _get_session_id, ): @@ -2949,6 +2956,11 @@ async def _strip_auth_on_cross_origin_redirect(response): # Streamable-HTTP transport TaskGroup dropped: reconnect # immediately instead of backoff/park (#66092). reason = self._reconnect_or_reraise_group(_eg) + finally: + # The async-with above owns the client; past this point it is + # closed, so drop the injection reference rather than letting + # header writes mutate a dead client across a reconnect gap. + self._http_client = None return reason else: # Deprecated API (mcp < 1.24.0): manages httpx client internally. @@ -4942,6 +4954,14 @@ def _handler(args: dict, **kwargs) -> str: ) return tool_error(f"MCP server '{server_name}' is not connected") + # Capture the caller's trace context HERE, on the agent thread. The + # coroutine below runs on the MCP daemon loop, where the agent's span + # is invisible (contextvars don't cross run_coroutine_threadsafe), so + # capturing any later is capturing nothing. One capture per tool + # invocation, injected around the RPC below (#52211). No-op (None) + # unless mcp.trace_propagation is enabled in config.yaml. + _traceparent = current_traceparent() + async def _call(): _mark_server_call_started(server) async with server._rpc_lock: @@ -4951,7 +4971,15 @@ async def _call(): # it and detect the gateway platform / session for routing. server._pending_call_context = contextvars.copy_context() try: - result = await server.session.call_tool(tool_name, arguments=args) + # Per-call injection on the shared client's headers — the + # POST is written by the transport's own task, which never + # sees this task's contextvars, but does read the client's + # default headers. The rpc lock serializes set/restore with + # the request that uses it. Deliberately NOT set at + # connection time: one connection serves many calls, each + # needing its own parent span (see PR #60466 review). + with injected_headers(server._http_client, _traceparent): + result = await server.session.call_tool(tool_name, arguments=args) finally: server._pending_call_context = None # The RPC round-trip completed — the session is demonstrably diff --git a/tools/mcp_trace_propagation.py b/tools/mcp_trace_propagation.py new file mode 100644 index 000000000000..b7f167283716 --- /dev/null +++ b/tools/mcp_trace_propagation.py @@ -0,0 +1,156 @@ +"""W3C trace-context propagation for MCP tool calls (opt-in). + +When an agent instrumented with OpenTelemetry calls a tool on an MCP server +that is itself instrumented, the two sides today produce disconnected trace +populations: the agent has a span for the tool call, the server has spans for +serving it, and nothing joins them. This module fixes that by injecting a +W3C ``traceparent`` header on the HTTP transport around each tool-call RPC, +so the server's spans become children of the agent's. + +Everything here is off by default and degrades to a no-op. Enable with:: + + mcp: + trace_propagation: true + +Two design constraints shape this module (they are the review feedback on +PR #60466, which attempted the same feature at connection-setup time): + +* **Capture happens on the agent thread, before the thread boundary.** MCP + RPCs run on a dedicated daemon event loop; ``trace.get_current_span()`` + over there observes nothing, because contextvars do not cross + ``run_coroutine_threadsafe``. The caller captures its own context and + hands the formatted header across. + +* **Injection happens per call, not per connection.** MCP tool calls reuse + one long-lived ``ClientSession``; a header fixed when the transport is + established would pin every later call on that connection to a single + stale span. The header is set on the shared httpx client immediately + before one RPC and removed immediately after, under the per-server RPC + lock that already serializes calls on a session. + +The trace context is read with the standard OpenTelemetry propagation API, +so any instrumentation that sets ambient context works unmodified. Tracing +plugins that keep their own span registry instead of attaching context (for +example hermes-otel, whose ``get_current_traceparent`` exists for exactly +this interop) can register a provider callback:: + + from tools.mcp_trace_propagation import register_traceparent_provider + register_traceparent_provider(get_current_traceparent) + +Provider output is validated against the W3C ``traceparent`` grammar before +use; anything else is discarded. No path in this module may raise into a +tool call — a broken telemetry setup must cost traces, never tool calls. +""" + +from __future__ import annotations + +import re +from contextlib import contextmanager +from typing import Any, Callable, Optional + +# version 00, 16-byte trace-id, 8-byte parent-id, 1-byte flags — lowercase hex +# (https://www.w3.org/TR/trace-context/#traceparent-header-field-values) +_TRACEPARENT_RE = re.compile(r"^00-[0-9a-f]{32}-[0-9a-f]{16}-[0-9a-f]{2}$") +_ALL_ZERO_TRACE_ID = "0" * 32 +_ALL_ZERO_PARENT_ID = "0" * 16 + +_provider: Optional[Callable[[], Optional[str]]] = None + + +def register_traceparent_provider( + provider: Optional[Callable[[], Optional[str]]], +) -> None: + """Register a callback that returns the current W3C traceparent, or None. + + For tracing integrations that track the active span themselves rather + than attaching it to ambient OpenTelemetry context. The callback is + consulted before the standard propagation API. Pass ``None`` to clear. + """ + global _provider + _provider = provider + + +def is_enabled() -> bool: + """The ``mcp.trace_propagation`` config gate (default: off).""" + try: + from hermes_cli.config import load_config_readonly + + section = load_config_readonly().get("mcp") or {} + return bool(section.get("trace_propagation", False)) + except Exception: + return False + + +def _valid(traceparent: Any) -> Optional[str]: + """Return ``traceparent`` if it is a well-formed W3C header value.""" + if not isinstance(traceparent, str) or not _TRACEPARENT_RE.match(traceparent): + return None + _version, trace_id, parent_id, _flags = traceparent.split("-") + if trace_id == _ALL_ZERO_TRACE_ID or parent_id == _ALL_ZERO_PARENT_ID: + return None # all-zero ids are explicitly invalid per the spec + return traceparent + + +def current_traceparent() -> Optional[str]: + """Capture the caller's trace context as a ``traceparent`` value. + + Must be called on the thread that owns the span — i.e. the agent thread + inside the tool handler, before the call is scheduled onto the MCP loop. + Returns ``None`` when the feature is disabled, no provider/SDK is + available, or there is no active span. Never raises. + """ + if not is_enabled(): + return None + + if _provider is not None: + try: + candidate = _valid(_provider()) + except Exception: + candidate = None + if candidate is not None: + return candidate + + try: + from opentelemetry import propagate + + carrier: dict = {} + propagate.inject(carrier) + return _valid(carrier.get("traceparent")) + except Exception: + return None + + +@contextmanager +def injected_headers(client: Any, traceparent: Optional[str]): + """Set ``traceparent`` on ``client``'s default headers for one RPC. + + ``client`` is the server's shared ``httpx.AsyncClient`` (None for stdio + transports, where there is nothing to inject into — the block is then a + no-op). Restores any pre-existing header value on exit so a client is + never left carrying a stale span, which is the precise failure mode of + connection-time injection. + + Client default headers, not contextvars, because the HTTP POST is + written by the transport's own writer task — context attached to the + calling task never reaches it, but the shared client's headers do. The + caller must hold the server's RPC lock (call sites already do), which + serializes header set/restore with the request that uses it. + """ + if client is None or not traceparent: + yield + return + + headers = client.headers + sentinel = object() + previous = headers.get("traceparent", sentinel) + headers["traceparent"] = traceparent + try: + yield + finally: + try: + if previous is sentinel: + headers.pop("traceparent", None) + else: + headers["traceparent"] = previous + except Exception: + pass diff --git a/website/docs/user-guide/features/mcp.md b/website/docs/user-guide/features/mcp.md index 75e22e569507..38da67fd537c 100644 --- a/website/docs/user-guide/features/mcp.md +++ b/website/docs/user-guide/features/mcp.md @@ -586,6 +586,43 @@ mcp- That makes MCP servers easier to reason about at the toolset level. +## Trace propagation (OpenTelemetry) + +Off by default. If your agent is instrumented with OpenTelemetry (via a +tracing plugin or your own SDK setup) and your HTTP MCP servers are too, +enabling trace propagation joins the two sides of every tool call into one +distributed trace — the server's spans become children of the agent's span, +so a trace viewer can answer "which MCP call made this turn slow": + +```yaml +mcp: + trace_propagation: true +``` + +With the gate on, Hermes captures the caller's active span as a W3C +[`traceparent`](https://www.w3.org/TR/trace-context/) header once per tool +invocation and injects it on the HTTP transport for exactly that one RPC. +There is nothing to configure per server, and stdio servers are unaffected +(there is no header to carry it in). + +Notes: + +- The trace context is read with the standard OpenTelemetry propagation + API. Tracing integrations that track spans themselves instead of setting + ambient context can register a callback: + + ```python + from tools.mcp_trace_propagation import register_traceparent_provider + register_traceparent_provider(get_current_traceparent) + ``` + +- No OpenTelemetry installation is required when the gate is off (and a + missing SDK with the gate on simply means no header is sent). +- Only [HTTP servers](#http-servers) on the current MCP SDK transport + (`mcp >= 1.24.0`) receive the header. +- This propagates *your* trace ids to *your* servers on requests the agent + already makes; nothing is reported anywhere else. + ## Security model ### Stdio env filtering