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