b60439152b
- harness/es.py: append-only event store + проекции (образец — кооп-портал) - harness/levels.py: уровни 1с/1м/1ч/1д, контекст ∝ масштабу - harness/scheduler.py: уровневые потоки, идемпотентные тики, absorbed-стабы - harness/isolated_runner.py + _llm_worker.py: LLM-тик в сабпроцессе (глобальный asyncio-рантайм Strands не переживает 2+ create_harness) - cli.py: init/task/tasks/status/events + requeue (написан самим минутным агентом) - 9router: отдельный ключ sk-simbiot, модель auto (glm-5.3 первым)
331 lines
13 KiB
Python
331 lines
13 KiB
Python
"""OpenTelemetry instrumentation for Model Context Protocol (MCP) tracing.
|
|
|
|
Enables distributed tracing across MCP client-server boundaries. The client
|
|
side is handled with `inject_trace_context`, which `MCPClient` uses to merge
|
|
the current OpenTelemetry context into the `_meta` field of outgoing tool
|
|
calls through the public `meta` parameter of the official `mcp` package.
|
|
|
|
The server side covers MCP servers hosted in the same process. It extracts
|
|
client-injected context from incoming messages so server-side spans join the
|
|
client's trace. This requires patching private internals of the `mcp` package
|
|
and those internals are only stable within the 1.x line, so the patches are
|
|
gated to `mcp` 1.x. The `mcp` 2.x line propagates trace context natively.
|
|
|
|
Based on: https://github.com/traceloop/openllmetry/tree/main/packages/opentelemetry-instrumentation-mcp
|
|
Related issue: https://github.com/modelcontextprotocol/modelcontextprotocol/issues/246
|
|
"""
|
|
|
|
import logging
|
|
import threading
|
|
from collections.abc import AsyncGenerator, Callable
|
|
from contextlib import _AsyncGeneratorContextManager, asynccontextmanager
|
|
from dataclasses import dataclass
|
|
from typing import Any
|
|
|
|
from mcp.shared.message import SessionMessage
|
|
from mcp.types import JSONRPCMessage, JSONRPCRequest
|
|
from opentelemetry import context, propagate
|
|
from wrapt import ObjectProxy, register_post_import_hook, wrap_function_wrapper
|
|
|
|
from ._compat import MCP_V2
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Module-level flag to ensure instrumentation is applied only once. The lock
|
|
# makes the check-and-set atomic so concurrent MCPClient construction cannot
|
|
# stack duplicate wrappers.
|
|
_instrumentation_applied = False
|
|
_instrumentation_lock = threading.Lock()
|
|
|
|
|
|
def inject_trace_context(meta: dict[str, Any] | None) -> dict[str, Any] | None:
|
|
"""Merge the current OpenTelemetry context into MCP request metadata.
|
|
|
|
Injects the active trace context (for example `traceparent` and
|
|
`tracestate`) into a copy of `meta` so it can be sent as the `_meta`
|
|
field of an MCP request. This enables server-side context extraction and
|
|
trace continuation across the client-server boundary.
|
|
|
|
Telemetry must never fail the tool call itself, so propagator errors are
|
|
logged and swallowed; the caller's metadata is still returned.
|
|
|
|
Args:
|
|
meta: Existing request metadata, or None. The input is not mutated.
|
|
|
|
Returns:
|
|
A new dict containing the caller's entries plus the injected trace
|
|
context, or None when there is no metadata to send.
|
|
"""
|
|
carrier: dict[str, Any] = dict(meta) if meta else {}
|
|
try:
|
|
propagate.get_global_textmap().inject(carrier)
|
|
except Exception:
|
|
logger.warning("failed to inject trace context into mcp request meta", exc_info=True)
|
|
return carrier or None
|
|
|
|
|
|
@dataclass(slots=True, frozen=True)
|
|
class ItemWithContext:
|
|
"""Wrapper for items that need to carry OpenTelemetry context.
|
|
|
|
Used to preserve tracing context across async boundaries in MCP sessions,
|
|
ensuring that distributed traces remain connected even when messages are
|
|
processed asynchronously.
|
|
|
|
Attributes:
|
|
item: The original item being wrapped
|
|
ctx: The OpenTelemetry context associated with the item
|
|
"""
|
|
|
|
item: Any
|
|
ctx: context.Context
|
|
|
|
|
|
def mcp_instrumentation() -> None:
|
|
"""Apply OpenTelemetry instrumentation patches for in-process MCP servers.
|
|
|
|
Instruments two areas of server-side MCP communication:
|
|
1. Transport-level: Extracts context from incoming messages
|
|
2. Session-level: Preserves context across async message processing
|
|
boundaries
|
|
|
|
Together the patches let an MCP server hosted in the same process join
|
|
the trace of the client that injected context into the request's `_meta`
|
|
field. Client-side injection does not require patching; `MCPClient`
|
|
injects context via `inject_trace_context`.
|
|
|
|
The patches wrap private internals of the `mcp` package that are only
|
|
stable within the 1.x line, so they are applied only when `mcp` 1.x is
|
|
installed. The `mcp` 2.x line propagates trace context natively.
|
|
|
|
This function is idempotent - multiple calls will not accumulate wrappers.
|
|
"""
|
|
global _instrumentation_applied
|
|
|
|
with _instrumentation_lock:
|
|
# Return early if instrumentation has already been applied
|
|
if _instrumentation_applied:
|
|
return
|
|
_instrumentation_applied = True
|
|
|
|
if MCP_V2:
|
|
return
|
|
|
|
def transport_wrapper() -> Callable[
|
|
[Callable[..., Any], Any, Any, Any], _AsyncGeneratorContextManager[tuple[Any, Any]]
|
|
]:
|
|
"""Create a wrapper for MCP transport connections.
|
|
|
|
Returns a context manager that wraps transport read/write streams
|
|
with context extraction capabilities. The wrapped reader will
|
|
automatically extract OpenTelemetry context from incoming messages.
|
|
|
|
Returns:
|
|
An async context manager that yields wrapped transport streams
|
|
"""
|
|
|
|
@asynccontextmanager
|
|
async def traced_method(
|
|
wrapped: Callable[..., Any], instance: Any, args: Any, kwargs: Any
|
|
) -> AsyncGenerator[tuple[Any, Any], None]:
|
|
async with wrapped(*args, **kwargs) as result:
|
|
try:
|
|
read_stream, write_stream = result
|
|
except ValueError:
|
|
read_stream, write_stream, _ = result
|
|
yield TransportContextExtractingReader(read_stream), write_stream
|
|
|
|
return traced_method
|
|
|
|
def session_init_wrapper() -> Callable[[Any, Any, tuple[Any, ...], dict[str, Any]], None]:
|
|
"""Create a wrapper for MCP session initialization.
|
|
|
|
Wraps session message streams to enable bidirectional context flow.
|
|
The reader extracts and activates context, while the writer preserves
|
|
context for async processing.
|
|
|
|
Returns:
|
|
A function that wraps session initialization
|
|
"""
|
|
|
|
def traced_method(
|
|
wrapped: Callable[..., Any], instance: Any, args: tuple[Any, ...], kwargs: dict[str, Any]
|
|
) -> None:
|
|
wrapped(*args, **kwargs)
|
|
reader = getattr(instance, "_incoming_message_stream_reader", None)
|
|
writer = getattr(instance, "_incoming_message_stream_writer", None)
|
|
if reader and writer:
|
|
instance._incoming_message_stream_reader = SessionContextAttachingReader(reader)
|
|
instance._incoming_message_stream_writer = SessionContextSavingWriter(writer)
|
|
|
|
return traced_method
|
|
|
|
# Apply patches
|
|
register_post_import_hook(
|
|
lambda _: wrap_function_wrapper(
|
|
"mcp.server.streamable_http", "StreamableHTTPServerTransport.connect", transport_wrapper()
|
|
),
|
|
"mcp.server.streamable_http",
|
|
)
|
|
|
|
register_post_import_hook(
|
|
lambda _: wrap_function_wrapper("mcp.server.session", "ServerSession.__init__", session_init_wrapper()),
|
|
"mcp.server.session",
|
|
)
|
|
|
|
# Mark instrumentation as applied
|
|
_instrumentation_applied = True
|
|
|
|
|
|
class TransportContextExtractingReader(ObjectProxy):
|
|
"""A proxy reader that extracts OpenTelemetry context from MCP messages.
|
|
|
|
Wraps an async message stream reader to automatically extract and activate
|
|
OpenTelemetry context from the _meta field of incoming MCP requests. This
|
|
enables server-side trace continuation from client-injected context.
|
|
|
|
The reader handles both SessionMessage and JSONRPCMessage formats, and
|
|
supports both dict and Pydantic model parameter structures.
|
|
"""
|
|
|
|
def __init__(self, wrapped: Any) -> None:
|
|
"""Initialize the context-extracting reader.
|
|
|
|
Args:
|
|
wrapped: The original async stream reader to wrap
|
|
"""
|
|
super().__init__(wrapped)
|
|
|
|
async def __aenter__(self) -> Any:
|
|
"""Enter the async context manager by delegating to the wrapped object."""
|
|
return await self.__wrapped__.__aenter__()
|
|
|
|
async def __aexit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> Any:
|
|
"""Exit the async context manager by delegating to the wrapped object."""
|
|
return await self.__wrapped__.__aexit__(exc_type, exc_value, traceback)
|
|
|
|
async def __aiter__(self) -> AsyncGenerator[Any, None]:
|
|
"""Iterate over messages, extracting and activating context as needed.
|
|
|
|
For each incoming message, checks if it contains tracing context in
|
|
the _meta field. If found, extracts and activates the context for
|
|
the duration of message processing, then properly detaches it.
|
|
|
|
Yields:
|
|
Messages from the wrapped stream, processed under the appropriate
|
|
OpenTelemetry context
|
|
"""
|
|
async for item in self.__wrapped__:
|
|
if isinstance(item, SessionMessage):
|
|
request = item.message.root # type: ignore[union-attr]
|
|
elif type(item) is JSONRPCMessage:
|
|
request = item.root # type: ignore[unreachable]
|
|
else:
|
|
yield item
|
|
continue
|
|
|
|
if isinstance(request, JSONRPCRequest) and request.params:
|
|
# Handle both dict and Pydantic model params
|
|
if hasattr(request.params, "get"):
|
|
# Dict-like access
|
|
meta = request.params.get("_meta")
|
|
elif hasattr(request.params, "_meta"):
|
|
# Direct attribute access for Pydantic models
|
|
meta = getattr(request.params, "_meta", None)
|
|
else:
|
|
meta = None
|
|
|
|
if meta:
|
|
extracted_context = propagate.extract(meta)
|
|
restore = context.attach(extracted_context)
|
|
try:
|
|
yield item
|
|
continue
|
|
finally:
|
|
context.detach(restore)
|
|
yield item
|
|
|
|
|
|
class SessionContextSavingWriter(ObjectProxy):
|
|
"""A proxy writer that preserves OpenTelemetry context with outgoing items.
|
|
|
|
Wraps an async message stream writer to capture the current OpenTelemetry
|
|
context and associate it with outgoing items. This enables context
|
|
preservation across async boundaries in MCP session processing.
|
|
"""
|
|
|
|
def __init__(self, wrapped: Any) -> None:
|
|
"""Initialize the context-saving writer.
|
|
|
|
Args:
|
|
wrapped: The original async stream writer to wrap
|
|
"""
|
|
super().__init__(wrapped)
|
|
|
|
async def __aenter__(self) -> Any:
|
|
"""Enter the async context manager by delegating to the wrapped object."""
|
|
return await self.__wrapped__.__aenter__()
|
|
|
|
async def __aexit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> Any:
|
|
"""Exit the async context manager by delegating to the wrapped object."""
|
|
return await self.__wrapped__.__aexit__(exc_type, exc_value, traceback)
|
|
|
|
async def send(self, item: Any) -> Any:
|
|
"""Send an item while preserving the current OpenTelemetry context.
|
|
|
|
Captures the current context and wraps the item with it, enabling
|
|
the receiving side to restore the appropriate tracing context.
|
|
|
|
Args:
|
|
item: The item to send through the stream
|
|
|
|
Returns:
|
|
Result of sending the wrapped item
|
|
"""
|
|
ctx = context.get_current()
|
|
return await self.__wrapped__.send(ItemWithContext(item, ctx))
|
|
|
|
|
|
class SessionContextAttachingReader(ObjectProxy):
|
|
"""A proxy reader that restores OpenTelemetry context from wrapped items.
|
|
|
|
Wraps an async message stream reader to detect ItemWithContext instances
|
|
and restore their associated OpenTelemetry context during processing.
|
|
This completes the context preservation cycle started by SessionContextSavingWriter.
|
|
"""
|
|
|
|
def __init__(self, wrapped: Any) -> None:
|
|
"""Initialize the context-attaching reader.
|
|
|
|
Args:
|
|
wrapped: The original async stream reader to wrap
|
|
"""
|
|
super().__init__(wrapped)
|
|
|
|
async def __aenter__(self) -> Any:
|
|
"""Enter the async context manager by delegating to the wrapped object."""
|
|
return await self.__wrapped__.__aenter__()
|
|
|
|
async def __aexit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> Any:
|
|
"""Exit the async context manager by delegating to the wrapped object."""
|
|
return await self.__wrapped__.__aexit__(exc_type, exc_value, traceback)
|
|
|
|
async def __aiter__(self) -> AsyncGenerator[Any, None]:
|
|
"""Iterate over items, restoring context for ItemWithContext instances.
|
|
|
|
For items wrapped with context, temporarily activates the associated
|
|
OpenTelemetry context during processing, then properly detaches it.
|
|
Regular items are yielded without context modification.
|
|
|
|
Yields:
|
|
Unwrapped items processed under their associated OpenTelemetry context
|
|
"""
|
|
async for item in self.__wrapped__:
|
|
if isinstance(item, ItemWithContext):
|
|
restore = context.attach(item.ctx)
|
|
try:
|
|
yield item.item
|
|
finally:
|
|
context.detach(restore)
|
|
else:
|
|
yield item
|