feat(mcp): opt-in W3C trace-context propagation for MCP tool calls

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"}

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