Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,7 @@ dependencies = [
"pyjwt[crypto]>=2.10.1",
"typing-extensions>=4.13.0",
"typing-inspection>=0.4.1",
"opentelemetry-api>=1.23.0",
]

[project.optional-dependencies]
Expand Down Expand Up @@ -71,6 +72,7 @@ dev = [
"coverage[toml]>=7.10.7,<=7.13",
"pillow>=12.0",
"strict-no-cover",
"opentelemetry-sdk>=1.24.0",
]
docs = [
"mkdocs>=1.6.1",
Expand Down
31 changes: 27 additions & 4 deletions src/mcp/shared/session.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@

import anyio
from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
from opentelemetry.propagate import inject
from pydantic import BaseModel, TypeAdapter
from typing_extensions import Self

Expand Down Expand Up @@ -263,6 +264,9 @@ async def send_request(
# Store the callback for this request
self._progress_callbacks[request_id] = progress_callback

# Propagate opentelemetry trace context
self._inject_otel_context(request_data)

try:
jsonrpc_request = JSONRPCRequest(jsonrpc="2.0", id=request_id, **request_data)
await self._write_stream.send(SessionMessage(message=jsonrpc_request, metadata=metadata))
Expand Down Expand Up @@ -295,18 +299,37 @@ async def send_notification(
related_request_id: RequestId | None = None,
) -> None:
"""Emits a notification, which is a one-way message that does not expect a response."""

request_data = notification.model_dump(by_alias=True, mode="json", exclude_none=True)
# Propagate opentelemetry trace context
self._inject_otel_context(request_data)
jsonrpc_notification = JSONRPCNotification(jsonrpc="2.0", **request_data)

# Some transport implementations may need to set the related_request_id
# to attribute to the notifications to the request that triggered them.
jsonrpc_notification = JSONRPCNotification(
jsonrpc="2.0",
**notification.model_dump(by_alias=True, mode="json", exclude_none=True),
)
session_message = SessionMessage(
message=jsonrpc_notification,
metadata=ServerMessageMetadata(related_request_id=related_request_id) if related_request_id else None,
)
await self._write_stream.send(session_message)

def _inject_otel_context(self, request: dict[str, Any]) -> None:
"""Propagate OpenTelemetry context in `_meta`.

See
- SEP414 https://github.com/modelcontextprotocol/modelcontextprotocol/pull/414
- OpenTelemetry semantic conventions
https://github.com/open-telemetry/semantic-conventions/blob/v1.39.0/docs/gen-ai/mcp.md
"""

carrier: dict[str, str] = {}
inject(carrier)
if not carrier:
return

meta: dict[str, Any] = request.setdefault("params", {}).setdefault("_meta", {})
meta.update(carrier)

async def _send_response(self, request_id: RequestId, response: SendResultT | ErrorData) -> None:
if isinstance(response, ErrorData):
jsonrpc_error = JSONRPCError(jsonrpc="2.0", id=request_id, error=response)
Expand Down
Loading
Loading