Address PR review findings

Author: Kludex

src/mcp/server/lowlevel/server.py

@@ -63,7 +63,7 @@ async def main():
 from mcp.server.lowlevel.experimental import ExperimentalHandlers
 from mcp.server.models import InitializationOptions
 from mcp.server.session import ServerSession
-from mcp.server.streamable_http import EventStore
+from mcp.server.streamable_http import MCP_SESSION_ID_HEADER, EventStore
 from mcp.server.streamable_http_manager import StreamableHTTPASGIApp, StreamableHTTPSessionManager
 from mcp.server.transport_security import TransportSecuritySettings
 from mcp.shared._otel import build_server_span_attributes, extract_trace_context, otel_span
@@ -73,7 +73,6 @@ async def main():
 from mcp.shared.session import RequestResponder
 
 logger = logging.getLogger(__name__)
-MCP_SESSION_ID_HEADER = "mcp-session-id"
 
 LifespanResultT = TypeVar("LifespanResultT", default=Any)
 

src/mcp/shared/_otel.py

@@ -4,7 +4,7 @@
 
 from collections.abc import Iterator
 from contextlib import contextmanager
-from typing import Any
+from typing import Any, cast
 
 from opentelemetry.context import Context
 from opentelemetry.propagate import extract, inject
@@ -51,7 +51,16 @@ def build_client_span_attributes(
         "jsonrpc.request.id": request_id,
     }
 
-    if params is not None and (resource_uri := params.get("uri")) is not None:
+    resource_uri = None
+    if params is not None:
+        resource_uri = params.get("uri")
+        if resource_uri is None:
+            ref = params.get("ref")
+            if isinstance(ref, dict):
+                typed_ref = cast(dict[str, Any], ref)
+                resource_uri = typed_ref.get("uri")
+
+    if resource_uri is not None:
         attributes["mcp.resource.uri"] = resource_uri
 
     return attributes
@@ -75,6 +84,9 @@ def build_server_span_attributes(
     }
 
     resource_uri = getattr(params, "uri", None)
+    if resource_uri is None:
+        ref = getattr(params, "ref", None)
+        resource_uri = getattr(ref, "uri", None)
     if resource_uri is not None:
         attributes["mcp.resource.uri"] = str(resource_uri)
 

tests/shared/test_otel.py

@@ -78,3 +78,49 @@ def test_resource() -> str:
     assert server_span["attributes"]["rpc.method"] == "resources/read"
     assert server_span["attributes"]["mcp.method.name"] == "resources/read"
     assert server_span["attributes"]["mcp.resource.uri"] == "test://resource"
+
+    # Server span should be in the same trace as the client span (context propagation).
+    assert server_span["context"]["trace_id"] == client_span["context"]["trace_id"]
+
+
+@pytest.mark.filterwarnings("ignore::RuntimeWarning")
+async def test_completion_spans_include_resource_template_uri(capfire: CaptureLogfire):
+    """Verify completion spans include the referenced resource template URI."""
+    server = MCPServer("test")
+
+    @server.completion()
+    async def handle_completion(
+        ref: types.ResourceTemplateReference | types.PromptReference,
+        argument: types.CompletionArgument,
+        context: types.CompletionContext | None,
+    ) -> types.Completion:
+        assert isinstance(ref, types.ResourceTemplateReference)
+        assert argument.name == "path"
+        assert argument.value == "rea"
+        assert context is None
+        return types.Completion(values=["README.md"])
+
+    async with Client(server) as client:
+        result = await client.complete(
+            ref=types.ResourceTemplateReference(type="ref/resource", uri="repo://files/{path}"),
+            argument={"name": "path", "value": "rea"},
+        )
+
+    assert result.completion.values == ["README.md"]
+
+    spans = capfire.exporter.exported_spans_as_dict()
+
+    client_span = next(s for s in spans if s["name"] == "MCP send completion/complete")
+    server_span = next(s for s in spans if s["name"] == "MCP handle completion/complete")
+
+    assert client_span["attributes"]["rpc.system"] == "mcp"
+    assert client_span["attributes"]["rpc.method"] == "completion/complete"
+    assert client_span["attributes"]["mcp.method.name"] == "completion/complete"
+    assert client_span["attributes"]["mcp.resource.uri"] == "repo://files/{path}"
+
+    assert server_span["attributes"]["rpc.system"] == "mcp"
+    assert server_span["attributes"]["rpc.service"] == "test"
+    assert server_span["attributes"]["rpc.method"] == "completion/complete"
+    assert server_span["attributes"]["mcp.method.name"] == "completion/complete"
+    assert server_span["attributes"]["mcp.resource.uri"] == "repo://files/{path}"
+    assert server_span["context"]["trace_id"] == client_span["context"]["trace_id"]

tests/shared/test_streamable_http.py

@@ -23,6 +23,7 @@
 import requests
 import uvicorn
 from httpx_sse import ServerSentEvent
+from logfire.testing import CaptureLogfire
 from starlette.applications import Starlette
 from starlette.requests import Request
 from starlette.routing import Mount
@@ -1081,6 +1082,32 @@ async def test_streamable_http_client_resource_read(initialized_client_session:
     assert response.contents[0].text == "Read test-resource"
 
 
+@pytest.mark.anyio
+@pytest.mark.filterwarnings("ignore::RuntimeWarning")
+async def test_streamable_http_server_span_includes_session_id(capfire: CaptureLogfire):
+    """Verify streamable HTTP server spans include the negotiated MCP session ID."""
+    app = _create_server()
+    mcp_app = app.streamable_http_app(host="testserver")
+
+    async with (
+        mcp_app.router.lifespan_context(mcp_app),
+        httpx.ASGITransport(mcp_app) as transport,
+        httpx.AsyncClient(transport=transport, base_url="http://testserver") as http_client,
+        streamable_http_client("http://testserver/mcp", http_client=http_client) as (read_stream, write_stream),
+        ClientSession(read_stream, write_stream) as session,
+    ):
+        await session.initialize()
+        response = await session.read_resource(uri="foobar://test-resource")
+
+    assert response.contents[0].uri == "foobar://test-resource"
+
+    spans = capfire.exporter.exported_spans_as_dict()
+    server_span = next(s for s in spans if s["name"] == "MCP handle resources/read")
+
+    assert server_span["attributes"]["mcp.session.id"]
+    assert server_span["attributes"]["mcp.resource.uri"] == "foobar://test-resource"
+
+
 @pytest.mark.anyio
 async def test_streamable_http_client_tool_invocation(initialized_client_session: ClientSession):
     """Test client tool invocation."""