File size: 1,257 Bytes
2f203f5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 | """OpenTelemetry parent spans for each MCP tool invocation (nests httpx child spans)."""
from __future__ import annotations
import functools
import inspect
from collections.abc import Callable
from typing import Any, cast
from opentelemetry import trace
_tracer = trace.get_tracer("data360.mcp.tools")
def instrument_mcp_tool(fn: Callable[..., Any], *, tool_name: str) -> Callable[..., Any]:
"""Wrap a tool function so each call runs under ``mcp.tool.<name>`` span."""
if inspect.iscoroutinefunction(fn):
fn_async = cast("Callable[..., Any]", fn)
@functools.wraps(fn_async)
async def _async_impl(*args: Any, **kwargs: Any) -> Any:
with _tracer.start_as_current_span(
f"mcp.tool.{tool_name}",
attributes={"mcp.tool.name": tool_name},
):
return await fn_async(*args, **kwargs)
return _async_impl
fn_sync = cast("Callable[..., Any]", fn)
@functools.wraps(fn_sync)
def _sync_impl(*args: Any, **kwargs: Any) -> Any:
with _tracer.start_as_current_span(
f"mcp.tool.{tool_name}",
attributes={"mcp.tool.name": tool_name},
):
return fn_sync(*args, **kwargs)
return _sync_impl
|