golf-mcp 0.2.16__py3-none-any.whl

golf/telemetry/instrumentation.py

@@ -0,0 +1,1402 @@
+"""Component-level OpenTelemetry instrumentation for Golf-built servers."""
+
+import asyncio
+import functools
+import os
+import sys
+import time
+import json
+from collections.abc import Callable
+from contextlib import asynccontextmanager
+from typing import Any, TypeVar
+from collections.abc import AsyncGenerator
+from collections import OrderedDict
+
+from opentelemetry import baggage, trace
+
+from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
+from opentelemetry.sdk.resources import Resource
+from opentelemetry.sdk.trace import TracerProvider
+from opentelemetry.sdk.trace.export import BatchSpanProcessor, ConsoleSpanExporter
+from opentelemetry.trace import Status, StatusCode
+
+from starlette.middleware.base import BaseHTTPMiddleware
+
+T = TypeVar("T")
+
+# Global tracer instance
+_tracer: trace.Tracer | None = None
+_provider: TracerProvider | None = None
+_detailed_tracing_enabled: bool = False
+
+
+def _safe_serialize(data: Any, max_length: int = 1000) -> str | None:
+    """Safely serialize data to string with length limit."""
+    try:
+        if isinstance(data, str):
+            serialized = data
+        else:
+            serialized = json.dumps(data, default=str, ensure_ascii=False)
+
+        if len(serialized) > max_length:
+            return serialized[:max_length] + "..." + f" (truncated from {len(serialized)} chars)"
+        return serialized
+    except (TypeError, ValueError):
+        # Fallback for non-serializable objects
+        try:
+            return str(data)[:max_length] + "..." if len(str(data)) > max_length else str(data)
+        except Exception:
+            return None
+
+
+def set_detailed_tracing(enabled: bool) -> None:
+    """Enable or disable detailed tracing with input/output capture."""
+    global _detailed_tracing_enabled
+    _detailed_tracing_enabled = enabled
+
+
+def init_telemetry(service_name: str = "golf-mcp-server") -> TracerProvider | None:
+    """Initialize OpenTelemetry with environment-based configuration.
+
+    Returns None if required environment variables are not set.
+    """
+    global _provider
+
+    # Check for required environment variables based on exporter type
+    exporter_type = os.environ.get("OTEL_TRACES_EXPORTER", "console").lower()
+
+    # For OTLP HTTP exporter, check if endpoint is configured
+    if exporter_type == "otlp_http":
+        endpoint = os.environ.get("OTEL_EXPORTER_OTLP_ENDPOINT")
+        if not endpoint:
+            print(
+                "[WARNING] OpenTelemetry tracing is disabled: "
+                "OTEL_EXPORTER_OTLP_ENDPOINT is not set for OTLP HTTP exporter"
+            )
+            return None
+
+    # Create resource with service information
+    resource_attributes = {
+        "service.name": os.environ.get("OTEL_SERVICE_NAME", service_name),
+        "service.version": os.environ.get("SERVICE_VERSION", "1.0.0"),
+        "service.instance.id": os.environ.get("SERVICE_INSTANCE_ID", "default"),
+    }
+
+    resource = Resource.create(resource_attributes)
+
+    # Create provider
+    provider = TracerProvider(resource=resource)
+
+    # Configure exporter based on type
+    try:
+        if exporter_type == "otlp_http":
+            endpoint = os.environ.get("OTEL_EXPORTER_OTLP_ENDPOINT", "http://localhost:4318/v1/traces")
+            headers = os.environ.get("OTEL_EXPORTER_OTLP_HEADERS", "")
+
+            # Parse headers if provided
+            header_dict = {}
+            if headers:
+                for header in headers.split(","):
+                    if "=" in header:
+                        key, value = header.split("=", 1)
+                        header_dict[key.strip()] = value.strip()
+
+            exporter = OTLPSpanExporter(endpoint=endpoint, headers=header_dict if header_dict else None)
+
+        else:
+            # Default to console exporter
+            exporter = ConsoleSpanExporter(out=sys.stderr)
+    except Exception:
+        import traceback
+
+        traceback.print_exc()
+        raise
+
+    # Add batch processor for better performance
+    try:
+        processor = BatchSpanProcessor(
+            exporter,
+            max_queue_size=2048,
+            schedule_delay_millis=1000,  # Export every 1 second instead of
+            # default 5 seconds
+            max_export_batch_size=512,
+            export_timeout_millis=5000,
+        )
+        provider.add_span_processor(processor)
+    except Exception:
+        import traceback
+
+        traceback.print_exc()
+        raise
+
+    # Set as global provider
+    try:
+        # Check if a provider is already set to avoid the warning
+        existing_provider = trace.get_tracer_provider()
+        if existing_provider is None or str(type(existing_provider).__name__) == "ProxyTracerProvider":
+            # Only set if no provider exists or it's the default proxy provider
+            trace.set_tracer_provider(provider)
+        _provider = provider
+    except Exception:
+        import traceback
+
+        traceback.print_exc()
+        raise
+
+    return provider
+
+
+def get_tracer() -> trace.Tracer:
+    """Get or create the global tracer instance."""
+    global _tracer, _provider
+
+    # If no provider is set, telemetry is disabled - return no-op tracer
+    if _provider is None:
+        return trace.get_tracer("golf.mcp.components.noop", "1.0.0")
+
+    if _tracer is None:
+        _tracer = trace.get_tracer("golf.mcp.components", "1.0.0")
+    return _tracer
+
+
+def instrument_tool(func: Callable[..., T], tool_name: str) -> Callable[..., T]:
+    """Instrument a tool function with OpenTelemetry tracing."""
+    global _provider
+
+    # If telemetry is disabled, return the original function
+    if _provider is None:
+        return func
+
+    tracer = get_tracer()
+
+    @functools.wraps(func)
+    async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
+        # Record metrics timing
+        import time
+
+        start_time = time.time()
+
+        # Create a more descriptive span name
+        span_name = f"mcp.tool.{tool_name}.execute"
+
+        # start_as_current_span automatically uses the current context and manages it
+        with tracer.start_as_current_span(span_name) as span:
+            # Add essential attributes only
+            span.set_attribute("mcp.component.type", "tool")
+            span.set_attribute("mcp.tool.name", tool_name)
+            span.set_attribute(
+                "mcp.tool.module",
+                func.__module__ if hasattr(func, "__module__") else "unknown",
+            )
+
+            # Add minimal execution context
+            if args or kwargs:
+                span.set_attribute("mcp.execution.has_params", True)
+
+            # Capture inputs if detailed tracing is enabled
+            if _detailed_tracing_enabled and (args or kwargs):
+                input_data = {"args": args, "kwargs": kwargs} if args or kwargs else None
+                if input_data:
+                    input_str = _safe_serialize(input_data)
+                    if input_str:
+                        span.set_attribute("mcp.tool.input", input_str)
+
+            # Extract Context parameter if present
+            ctx = kwargs.get("ctx")
+            if ctx:
+                # Only extract known MCP context attributes
+                ctx_attrs = [
+                    "request_id",
+                    "session_id",
+                    "client_id",
+                    "user_id",
+                    "tenant_id",
+                ]
+                for attr in ctx_attrs:
+                    if hasattr(ctx, attr):
+                        value = getattr(ctx, attr)
+                        if value is not None:
+                            span.set_attribute(f"mcp.context.{attr}", str(value))
+
+            # Also check baggage for session ID
+            session_id_from_baggage = baggage.get_baggage("mcp.session.id")
+            if session_id_from_baggage:
+                span.set_attribute("mcp.session.id", session_id_from_baggage)
+
+            # Add event for tool execution start
+            span.add_event("tool.execution.started", {"tool.name": tool_name})
+
+            try:
+                result = await func(*args, **kwargs)
+                span.set_status(Status(StatusCode.OK))
+
+                # Add event for successful completion
+                span.add_event("tool.execution.completed", {"tool.name": tool_name})
+
+                # Record metrics for successful execution
+                try:
+                    from golf.metrics import get_metrics_collector
+
+                    metrics_collector = get_metrics_collector()
+                    metrics_collector.increment_tool_execution(tool_name, "success")
+                    metrics_collector.record_tool_duration(tool_name, time.time() - start_time)
+                except ImportError:
+                    # Metrics not available, continue without metrics
+                    pass
+
+                # Capture result metadata
+                if result is not None:
+                    span.set_attribute("mcp.tool.result.type", type(result).__name__)
+
+                    if isinstance(result, list | dict) and hasattr(result, "__len__"):
+                        span.set_attribute("mcp.tool.result.size", len(result))
+                    elif isinstance(result, str):
+                        span.set_attribute("mcp.tool.result.length", len(result))
+
+                    # Capture full output if detailed tracing is enabled
+                    if _detailed_tracing_enabled:
+                        output_str = _safe_serialize(result)
+                        if output_str:
+                            span.set_attribute("mcp.tool.output", output_str)
+
+                return result
+            except Exception as e:
+                span.record_exception(e)
+                span.set_status(Status(StatusCode.ERROR, str(e)))
+
+                # Add event for error
+                span.add_event(
+                    "tool.execution.error",
+                    {
+                        "tool.name": tool_name,
+                        "error.type": type(e).__name__,
+                        "error.message": str(e),
+                    },
+                )
+
+                # Record metrics for failed execution
+                try:
+                    from golf.metrics import get_metrics_collector
+
+                    metrics_collector = get_metrics_collector()
+                    metrics_collector.increment_tool_execution(tool_name, "error")
+                    metrics_collector.increment_error("tool", type(e).__name__)
+                except ImportError:
+                    # Metrics not available, continue without metrics
+                    pass
+
+                raise
+
+    @functools.wraps(func)
+    def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
+        # Record metrics timing
+        import time
+
+        start_time = time.time()
+
+        # Create a more descriptive span name
+        span_name = f"mcp.tool.{tool_name}.execute"
+
+        # start_as_current_span automatically uses the current context and manages it
+        with tracer.start_as_current_span(span_name) as span:
+            # Add essential attributes only
+            span.set_attribute("mcp.component.type", "tool")
+            span.set_attribute("mcp.tool.name", tool_name)
+            span.set_attribute(
+                "mcp.tool.module",
+                func.__module__ if hasattr(func, "__module__") else "unknown",
+            )
+
+            # Add execution context
+            span.set_attribute("mcp.execution.args_count", len(args))
+            span.set_attribute("mcp.execution.kwargs_count", len(kwargs))
+
+            # Extract Context parameter if present
+            ctx = kwargs.get("ctx")
+            if ctx:
+                # Only extract known MCP context attributes
+                ctx_attrs = [
+                    "request_id",
+                    "session_id",
+                    "client_id",
+                    "user_id",
+                    "tenant_id",
+                ]
+                for attr in ctx_attrs:
+                    if hasattr(ctx, attr):
+                        value = getattr(ctx, attr)
+                        if value is not None:
+                            span.set_attribute(f"mcp.context.{attr}", str(value))
+
+            # Also check baggage for session ID
+            session_id_from_baggage = baggage.get_baggage("mcp.session.id")
+            if session_id_from_baggage:
+                span.set_attribute("mcp.session.id", session_id_from_baggage)
+
+            # Add event for tool execution start
+            span.add_event("tool.execution.started", {"tool.name": tool_name})
+
+            try:
+                result = func(*args, **kwargs)
+                span.set_status(Status(StatusCode.OK))
+
+                # Add event for successful completion
+                span.add_event("tool.execution.completed", {"tool.name": tool_name})
+
+                # Record metrics for successful execution
+                try:
+                    from golf.metrics import get_metrics_collector
+
+                    metrics_collector = get_metrics_collector()
+                    metrics_collector.increment_tool_execution(tool_name, "success")
+                    metrics_collector.record_tool_duration(tool_name, time.time() - start_time)
+                except ImportError:
+                    # Metrics not available, continue without metrics
+                    pass
+
+                # Capture result metadata
+                if result is not None:
+                    span.set_attribute("mcp.tool.result.type", type(result).__name__)
+
+                    if isinstance(result, list | dict) and hasattr(result, "__len__"):
+                        span.set_attribute("mcp.tool.result.size", len(result))
+                    elif isinstance(result, str):
+                        span.set_attribute("mcp.tool.result.length", len(result))
+
+                    # Capture full output if detailed tracing is enabled
+                    if _detailed_tracing_enabled:
+                        output_str = _safe_serialize(result)
+                        if output_str:
+                            span.set_attribute("mcp.tool.output", output_str)
+
+                return result
+            except Exception as e:
+                span.record_exception(e)
+                span.set_status(Status(StatusCode.ERROR, str(e)))
+
+                # Add event for error
+                span.add_event(
+                    "tool.execution.error",
+                    {
+                        "tool.name": tool_name,
+                        "error.type": type(e).__name__,
+                        "error.message": str(e),
+                    },
+                )
+
+                # Record metrics for failed execution
+                try:
+                    from golf.metrics import get_metrics_collector
+
+                    metrics_collector = get_metrics_collector()
+                    metrics_collector.increment_tool_execution(tool_name, "error")
+                    metrics_collector.increment_error("tool", type(e).__name__)
+                except ImportError:
+                    # Metrics not available, continue without metrics
+                    pass
+
+                raise
+
+    # Return appropriate wrapper based on function type
+    if asyncio.iscoroutinefunction(func):
+        return async_wrapper
+    else:
+        return sync_wrapper
+
+
+def instrument_resource(func: Callable[..., T], resource_uri: str) -> Callable[..., T]:
+    """Instrument a resource function with OpenTelemetry tracing."""
+    global _provider
+
+    # If telemetry is disabled, return the original function
+    if _provider is None:
+        return func
+
+    tracer = get_tracer()
+
+    # Determine if this is a template based on URI pattern
+    is_template = "{" in resource_uri
+
+    @functools.wraps(func)
+    async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
+        # Create a more descriptive span name
+        span_name = f"mcp.resource.{'template' if is_template else 'static'}.read"
+        with tracer.start_as_current_span(span_name) as span:
+            # Add essential attributes only
+            span.set_attribute("mcp.component.type", "resource")
+            span.set_attribute("mcp.resource.uri", resource_uri)
+            span.set_attribute("mcp.resource.is_template", is_template)
+            span.set_attribute(
+                "mcp.resource.module",
+                func.__module__ if hasattr(func, "__module__") else "unknown",
+            )
+
+            # Extract Context parameter if present
+            ctx = kwargs.get("ctx")
+            if ctx:
+                # Only extract known MCP context attributes
+                ctx_attrs = [
+                    "request_id",
+                    "session_id",
+                    "client_id",
+                    "user_id",
+                    "tenant_id",
+                ]
+                for attr in ctx_attrs:
+                    if hasattr(ctx, attr):
+                        value = getattr(ctx, attr)
+                        if value is not None:
+                            span.set_attribute(f"mcp.context.{attr}", str(value))
+
+            # Also check baggage for session ID
+            session_id_from_baggage = baggage.get_baggage("mcp.session.id")
+            if session_id_from_baggage:
+                span.set_attribute("mcp.session.id", session_id_from_baggage)
+
+            # Add event for resource read start
+            span.add_event("resource.read.started", {"resource.uri": resource_uri})
+
+            try:
+                result = await func(*args, **kwargs)
+                span.set_status(Status(StatusCode.OK))
+
+                # Add event for successful read
+                span.add_event("resource.read.completed", {"resource.uri": resource_uri})
+
+                # Add result metadata
+                if hasattr(result, "__len__"):
+                    span.set_attribute("mcp.resource.result.size", len(result))
+
+                # Determine content type if possible
+                if isinstance(result, str):
+                    span.set_attribute("mcp.resource.result.type", "text")
+                    span.set_attribute("mcp.resource.result.length", len(result))
+                elif isinstance(result, bytes):
+                    span.set_attribute("mcp.resource.result.type", "binary")
+                    span.set_attribute("mcp.resource.result.size_bytes", len(result))
+                elif isinstance(result, dict):
+                    span.set_attribute("mcp.resource.result.type", "object")
+                    span.set_attribute("mcp.resource.result.keys_count", len(result))
+                elif isinstance(result, list):
+                    span.set_attribute("mcp.resource.result.type", "array")
+                    span.set_attribute("mcp.resource.result.items_count", len(result))
+
+                return result
+            except Exception as e:
+                span.record_exception(e)
+                span.set_status(Status(StatusCode.ERROR, str(e)))
+
+                # Add event for error
+                span.add_event(
+                    "resource.read.error",
+                    {
+                        "resource.uri": resource_uri,
+                        "error.type": type(e).__name__,
+                        "error.message": str(e),
+                    },
+                )
+                raise
+
+    @functools.wraps(func)
+    def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
+        # Create a more descriptive span name
+        span_name = f"mcp.resource.{'template' if is_template else 'static'}.read"
+        with tracer.start_as_current_span(span_name) as span:
+            # Add essential attributes only
+            span.set_attribute("mcp.component.type", "resource")
+            span.set_attribute("mcp.resource.uri", resource_uri)
+            span.set_attribute("mcp.resource.is_template", is_template)
+            span.set_attribute(
+                "mcp.resource.module",
+                func.__module__ if hasattr(func, "__module__") else "unknown",
+            )
+
+            # Extract Context parameter if present
+            ctx = kwargs.get("ctx")
+            if ctx:
+                # Only extract known MCP context attributes
+                ctx_attrs = [
+                    "request_id",
+                    "session_id",
+                    "client_id",
+                    "user_id",
+                    "tenant_id",
+                ]
+                for attr in ctx_attrs:
+                    if hasattr(ctx, attr):
+                        value = getattr(ctx, attr)
+                        if value is not None:
+                            span.set_attribute(f"mcp.context.{attr}", str(value))
+
+            # Also check baggage for session ID
+            session_id_from_baggage = baggage.get_baggage("mcp.session.id")
+            if session_id_from_baggage:
+                span.set_attribute("mcp.session.id", session_id_from_baggage)
+
+            # Add event for resource read start
+            span.add_event("resource.read.started", {"resource.uri": resource_uri})
+
+            try:
+                result = func(*args, **kwargs)
+                span.set_status(Status(StatusCode.OK))
+
+                # Add event for successful read
+                span.add_event("resource.read.completed", {"resource.uri": resource_uri})
+
+                # Add result metadata
+                if hasattr(result, "__len__"):
+                    span.set_attribute("mcp.resource.result.size", len(result))
+
+                # Determine content type if possible
+                if isinstance(result, str):
+                    span.set_attribute("mcp.resource.result.type", "text")
+                    span.set_attribute("mcp.resource.result.length", len(result))
+                elif isinstance(result, bytes):
+                    span.set_attribute("mcp.resource.result.type", "binary")
+                    span.set_attribute("mcp.resource.result.size_bytes", len(result))
+                elif isinstance(result, dict):
+                    span.set_attribute("mcp.resource.result.type", "object")
+                    span.set_attribute("mcp.resource.result.keys_count", len(result))
+                elif isinstance(result, list):
+                    span.set_attribute("mcp.resource.result.type", "array")
+                    span.set_attribute("mcp.resource.result.items_count", len(result))
+
+                return result
+            except Exception as e:
+                span.record_exception(e)
+                span.set_status(Status(StatusCode.ERROR, str(e)))
+
+                # Add event for error
+                span.add_event(
+                    "resource.read.error",
+                    {
+                        "resource.uri": resource_uri,
+                        "error.type": type(e).__name__,
+                        "error.message": str(e),
+                    },
+                )
+                raise
+
+    if asyncio.iscoroutinefunction(func):
+        return async_wrapper
+    else:
+        return sync_wrapper
+
+
+def instrument_elicitation(func: Callable[..., T], elicitation_type: str = "elicit") -> Callable[..., T]:
+    """Instrument an elicitation function with OpenTelemetry tracing."""
+    global _provider
+
+    # If telemetry is disabled, return the original function
+    if _provider is None:
+        return func
+
+    tracer = get_tracer()
+
+    @functools.wraps(func)
+    async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
+        # If telemetry is disabled at runtime, call original function
+        global _provider
+        if _provider is None:
+            return await func(*args, **kwargs)
+
+        # Record metrics timing
+        start_time = time.time()
+
+        # Create a more descriptive span name
+        span_name = f"mcp.elicitation.{elicitation_type}.request"
+        with tracer.start_as_current_span(span_name) as span:
+            # Add essential attributes
+            span.set_attribute("mcp.component.type", "elicitation")
+            span.set_attribute("mcp.elicitation.type", elicitation_type)
+
+            # Capture elicitation parameters if detailed tracing is enabled
+            if _detailed_tracing_enabled:
+                # Extract message from first argument (common pattern)
+                if args:
+                    message = args[0] if isinstance(args[0], str) else None
+                    if message:
+                        span.set_attribute("mcp.elicitation.message", _safe_serialize(message, 500))
+
+                # Extract response_type from kwargs/args
+                response_type = kwargs.get("response_type") or (args[1] if len(args) > 1 else None)
+                if response_type is not None:
+                    if isinstance(response_type, list):
+                        span.set_attribute("mcp.elicitation.response_type", "choice")
+                        span.set_attribute("mcp.elicitation.choices", str(response_type))
+                    elif hasattr(response_type, "__name__"):
+                        span.set_attribute("mcp.elicitation.response_type", response_type.__name__)
+                    else:
+                        span.set_attribute("mcp.elicitation.response_type", str(type(response_type).__name__))
+
+            # Extract Context parameter if present
+            ctx = kwargs.get("ctx")
+            if ctx:
+                ctx_attrs = ["request_id", "session_id", "client_id", "user_id", "tenant_id"]
+                for attr in ctx_attrs:
+                    if hasattr(ctx, attr):
+                        value = getattr(ctx, attr)
+                        if value is not None:
+                            span.set_attribute(f"mcp.context.{attr}", str(value))
+
+            # Add event for elicitation start
+            span.add_event("elicitation.request.started")
+
+            try:
+                result = await func(*args, **kwargs)
+                span.set_status(Status(StatusCode.OK))
+
+                # Add event for successful completion
+                span.add_event("elicitation.request.completed")
+
+                # Capture result metadata
+                if result is not None and _detailed_tracing_enabled:
+                    if isinstance(result, str):
+                        span.set_attribute("mcp.elicitation.result.content", _safe_serialize(result, 500))
+                    elif isinstance(result, (list, dict)) and hasattr(result, "__len__"):
+                        span.set_attribute("mcp.elicitation.result.size", len(result))
+                        span.set_attribute("mcp.elicitation.result.content", _safe_serialize(result, 1000))
+
+                # Record metrics for successful elicitation
+                try:
+                    from golf.metrics import get_metrics_collector
+
+                    metrics_collector = get_metrics_collector()
+                    metrics_collector.increment_elicitation(elicitation_type, "success")
+                    metrics_collector.record_elicitation_duration(elicitation_type, time.time() - start_time)
+                except ImportError:
+                    pass
+
+                return result
+            except Exception as e:
+                span.record_exception(e)
+                span.set_status(Status(StatusCode.ERROR, str(e)))
+
+                # Add event for error
+                span.add_event(
+                    "elicitation.request.error",
+                    {
+                        "error.type": type(e).__name__,
+                        "error.message": str(e),
+                    },
+                )
+
+                # Record metrics for failed elicitation
+                try:
+                    from golf.metrics import get_metrics_collector
+
+                    metrics_collector = get_metrics_collector()
+                    metrics_collector.increment_elicitation(elicitation_type, "error")
+                    metrics_collector.increment_error("elicitation", type(e).__name__)
+                except ImportError:
+                    pass
+
+                raise
+
+    @functools.wraps(func)
+    def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
+        # If telemetry is disabled at runtime, call original function
+        global _provider
+        if _provider is None:
+            return func(*args, **kwargs)
+
+        # Record metrics timing
+        start_time = time.time()
+
+        # Create a more descriptive span name
+        span_name = f"mcp.elicitation.{elicitation_type}.request"
+        with tracer.start_as_current_span(span_name) as span:
+            # Add essential attributes
+            span.set_attribute("mcp.component.type", "elicitation")
+            span.set_attribute("mcp.elicitation.type", elicitation_type)
+
+            # Capture elicitation parameters if detailed tracing is enabled
+            if _detailed_tracing_enabled:
+                if args:
+                    message = args[0] if isinstance(args[0], str) else None
+                    if message:
+                        span.set_attribute("mcp.elicitation.message", _safe_serialize(message, 500))
+
+            # Add event for elicitation start
+            span.add_event("elicitation.request.started")
+
+            try:
+                result = func(*args, **kwargs)
+                span.set_status(Status(StatusCode.OK))
+
+                # Add event for successful completion
+                span.add_event("elicitation.request.completed")
+
+                # Record metrics for successful elicitation
+                try:
+                    from golf.metrics import get_metrics_collector
+
+                    metrics_collector = get_metrics_collector()
+                    metrics_collector.increment_elicitation(elicitation_type, "success")
+                    metrics_collector.record_elicitation_duration(elicitation_type, time.time() - start_time)
+                except ImportError:
+                    pass
+
+                return result
+            except Exception as e:
+                span.record_exception(e)
+                span.set_status(Status(StatusCode.ERROR, str(e)))
+
+                # Add event for error
+                span.add_event(
+                    "elicitation.request.error",
+                    {
+                        "error.type": type(e).__name__,
+                        "error.message": str(e),
+                    },
+                )
+
+                # Record metrics for failed elicitation
+                try:
+                    from golf.metrics import get_metrics_collector
+
+                    metrics_collector = get_metrics_collector()
+                    metrics_collector.increment_elicitation(elicitation_type, "error")
+                    metrics_collector.increment_error("elicitation", type(e).__name__)
+                except ImportError:
+                    pass
+
+                raise
+
+    if asyncio.iscoroutinefunction(func):
+        return async_wrapper
+    else:
+        return sync_wrapper
+
+
+def instrument_sampling(func: Callable[..., T], sampling_type: str = "sample") -> Callable[..., T]:
+    """Instrument a sampling function with OpenTelemetry tracing."""
+    global _provider
+
+    # If telemetry is disabled, return the original function
+    if _provider is None:
+        return func
+
+    tracer = get_tracer()
+
+    @functools.wraps(func)
+    async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
+        # If telemetry is disabled at runtime, call original function
+        global _provider
+        if _provider is None:
+            return await func(*args, **kwargs)
+
+        # Record metrics timing
+        start_time = time.time()
+
+        # Create a more descriptive span name
+        span_name = f"mcp.sampling.{sampling_type}.request"
+        with tracer.start_as_current_span(span_name) as span:
+            # Add essential attributes
+            span.set_attribute("mcp.component.type", "sampling")
+            span.set_attribute("mcp.sampling.type", sampling_type)
+
+            # Capture sampling parameters
+            messages = kwargs.get("messages") or (args[0] if args else None)
+            if messages and _detailed_tracing_enabled:
+                if isinstance(messages, str):
+                    span.set_attribute("mcp.sampling.messages.content", _safe_serialize(messages, 1000))
+                elif isinstance(messages, list):
+                    span.set_attribute("mcp.sampling.messages.type", "list")
+                    span.set_attribute("mcp.sampling.messages.count", len(messages))
+                    span.set_attribute("mcp.sampling.messages.content", _safe_serialize(messages, 1000))
+
+            # Capture other sampling parameters
+            system_prompt = kwargs.get("system_prompt")
+            if system_prompt and _detailed_tracing_enabled:
+                span.set_attribute("mcp.sampling.system_prompt.length", len(str(system_prompt)))
+                span.set_attribute("mcp.sampling.system_prompt.content", _safe_serialize(system_prompt, 500))
+
+            temperature = kwargs.get("temperature")
+            if temperature is not None:
+                span.set_attribute("mcp.sampling.temperature", temperature)
+
+            max_tokens = kwargs.get("max_tokens")
+            if max_tokens is not None:
+                span.set_attribute("mcp.sampling.max_tokens", max_tokens)
+
+            model_preferences = kwargs.get("model_preferences")
+            if model_preferences:
+                if isinstance(model_preferences, str):
+                    span.set_attribute("mcp.sampling.model_preferences", model_preferences)
+                elif isinstance(model_preferences, list):
+                    span.set_attribute("mcp.sampling.model_preferences", ",".join(model_preferences))
+
+            # Extract Context parameter if present
+            ctx = kwargs.get("ctx")
+            if ctx:
+                ctx_attrs = ["request_id", "session_id", "client_id", "user_id", "tenant_id"]
+                for attr in ctx_attrs:
+                    if hasattr(ctx, attr):
+                        value = getattr(ctx, attr)
+                        if value is not None:
+                            span.set_attribute(f"mcp.context.{attr}", str(value))
+
+            # Add event for sampling start
+            span.add_event("sampling.request.started")
+
+            try:
+                result = await func(*args, **kwargs)
+                span.set_status(Status(StatusCode.OK))
+
+                # Add event for successful completion
+                span.add_event("sampling.request.completed")
+
+                # Capture result metadata
+                if result is not None and _detailed_tracing_enabled and isinstance(result, str):
+                    span.set_attribute("mcp.sampling.result.content", _safe_serialize(result, 1000))
+
+                # Record metrics for successful sampling
+                try:
+                    from golf.metrics import get_metrics_collector
+
+                    metrics_collector = get_metrics_collector()
+                    metrics_collector.increment_sampling(sampling_type, "success")
+                    metrics_collector.record_sampling_duration(sampling_type, time.time() - start_time)
+                    if isinstance(result, str):
+                        metrics_collector.record_sampling_tokens(sampling_type, len(result.split()))
+                except ImportError:
+                    pass
+
+                return result
+            except Exception as e:
+                span.record_exception(e)
+                span.set_status(Status(StatusCode.ERROR, str(e)))
+
+                # Add event for error
+                span.add_event(
+                    "sampling.request.error",
+                    {
+                        "error.type": type(e).__name__,
+                        "error.message": str(e),
+                    },
+                )
+
+                # Record metrics for failed sampling
+                try:
+                    from golf.metrics import get_metrics_collector
+
+                    metrics_collector = get_metrics_collector()
+                    metrics_collector.increment_sampling(sampling_type, "error")
+                    metrics_collector.increment_error("sampling", type(e).__name__)
+                except ImportError:
+                    pass
+
+                raise
+
+    @functools.wraps(func)
+    def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
+        # If telemetry is disabled at runtime, call original function
+        global _provider
+        if _provider is None:
+            return func(*args, **kwargs)
+
+        # Record metrics timing
+        start_time = time.time()
+
+        # Create a more descriptive span name
+        span_name = f"mcp.sampling.{sampling_type}.request"
+        with tracer.start_as_current_span(span_name) as span:
+            # Add essential attributes
+            span.set_attribute("mcp.component.type", "sampling")
+            span.set_attribute("mcp.sampling.type", sampling_type)
+
+            # Add event for sampling start
+            span.add_event("sampling.request.started")
+
+            try:
+                result = func(*args, **kwargs)
+                span.set_status(Status(StatusCode.OK))
+
+                # Add event for successful completion
+                span.add_event("sampling.request.completed")
+
+                # Record metrics for successful sampling
+                try:
+                    from golf.metrics import get_metrics_collector
+
+                    metrics_collector = get_metrics_collector()
+                    metrics_collector.increment_sampling(sampling_type, "success")
+                    metrics_collector.record_sampling_duration(sampling_type, time.time() - start_time)
+                except ImportError:
+                    pass
+
+                return result
+            except Exception as e:
+                span.record_exception(e)
+                span.set_status(Status(StatusCode.ERROR, str(e)))
+
+                # Add event for error
+                span.add_event(
+                    "sampling.request.error",
+                    {
+                        "error.type": type(e).__name__,
+                        "error.message": str(e),
+                    },
+                )
+                raise
+
+    if asyncio.iscoroutinefunction(func):
+        return async_wrapper
+    else:
+        return sync_wrapper
+
+
+def instrument_prompt(func: Callable[..., T], prompt_name: str) -> Callable[..., T]:
+    """Instrument a prompt function with OpenTelemetry tracing."""
+    global _provider
+
+    # If telemetry is disabled, return the original function
+    if _provider is None:
+        return func
+
+    tracer = get_tracer()
+
+    @functools.wraps(func)
+    async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
+        # Create a more descriptive span name
+        span_name = f"mcp.prompt.{prompt_name}.generate"
+        with tracer.start_as_current_span(span_name) as span:
+            # Add essential attributes only
+            span.set_attribute("mcp.component.type", "prompt")
+            span.set_attribute("mcp.prompt.name", prompt_name)
+            span.set_attribute(
+                "mcp.prompt.module",
+                func.__module__ if hasattr(func, "__module__") else "unknown",
+            )
+
+            # Extract Context parameter if present
+            ctx = kwargs.get("ctx")
+            if ctx:
+                # Only extract known MCP context attributes
+                ctx_attrs = [
+                    "request_id",
+                    "session_id",
+                    "client_id",
+                    "user_id",
+                    "tenant_id",
+                ]
+                for attr in ctx_attrs:
+                    if hasattr(ctx, attr):
+                        value = getattr(ctx, attr)
+                        if value is not None:
+                            span.set_attribute(f"mcp.context.{attr}", str(value))
+
+            # Also check baggage for session ID
+            session_id_from_baggage = baggage.get_baggage("mcp.session.id")
+            if session_id_from_baggage:
+                span.set_attribute("mcp.session.id", session_id_from_baggage)
+
+            # Add event for prompt generation start
+            span.add_event("prompt.generation.started", {"prompt.name": prompt_name})
+
+            try:
+                result = await func(*args, **kwargs)
+                span.set_status(Status(StatusCode.OK))
+
+                # Add event for successful generation
+                span.add_event("prompt.generation.completed", {"prompt.name": prompt_name})
+
+                # Add message count and type information
+                if isinstance(result, list):
+                    span.set_attribute("mcp.prompt.result.message_count", len(result))
+                    span.set_attribute("mcp.prompt.result.type", "message_list")
+
+                    # Analyze message types if they have role attributes
+                    roles = []
+                    for msg in result:
+                        if hasattr(msg, "role"):
+                            roles.append(msg.role)
+                        elif isinstance(msg, dict) and "role" in msg:
+                            roles.append(msg["role"])
+
+                    if roles:
+                        unique_roles = list(set(roles))
+                        span.set_attribute("mcp.prompt.result.roles", ",".join(unique_roles))
+                        span.set_attribute(
+                            "mcp.prompt.result.role_counts",
+                            str({role: roles.count(role) for role in unique_roles}),
+                        )
+                elif isinstance(result, str):
+                    span.set_attribute("mcp.prompt.result.type", "string")
+                    span.set_attribute("mcp.prompt.result.length", len(result))
+                else:
+                    span.set_attribute("mcp.prompt.result.type", type(result).__name__)
+
+                return result
+            except Exception as e:
+                span.record_exception(e)
+                span.set_status(Status(StatusCode.ERROR, str(e)))
+
+                # Add event for error
+                span.add_event(
+                    "prompt.generation.error",
+                    {
+                        "prompt.name": prompt_name,
+                        "error.type": type(e).__name__,
+                        "error.message": str(e),
+                    },
+                )
+                raise
+
+    @functools.wraps(func)
+    def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
+        # Create a more descriptive span name
+        span_name = f"mcp.prompt.{prompt_name}.generate"
+        with tracer.start_as_current_span(span_name) as span:
+            # Add essential attributes only
+            span.set_attribute("mcp.component.type", "prompt")
+            span.set_attribute("mcp.prompt.name", prompt_name)
+            span.set_attribute(
+                "mcp.prompt.module",
+                func.__module__ if hasattr(func, "__module__") else "unknown",
+            )
+
+            # Extract Context parameter if present
+            ctx = kwargs.get("ctx")
+            if ctx:
+                # Only extract known MCP context attributes
+                ctx_attrs = [
+                    "request_id",
+                    "session_id",
+                    "client_id",
+                    "user_id",
+                    "tenant_id",
+                ]
+                for attr in ctx_attrs:
+                    if hasattr(ctx, attr):
+                        value = getattr(ctx, attr)
+                        if value is not None:
+                            span.set_attribute(f"mcp.context.{attr}", str(value))
+
+            # Also check baggage for session ID
+            session_id_from_baggage = baggage.get_baggage("mcp.session.id")
+            if session_id_from_baggage:
+                span.set_attribute("mcp.session.id", session_id_from_baggage)
+
+            # Add event for prompt generation start
+            span.add_event("prompt.generation.started", {"prompt.name": prompt_name})
+
+            try:
+                result = func(*args, **kwargs)
+                span.set_status(Status(StatusCode.OK))
+
+                # Add event for successful generation
+                span.add_event("prompt.generation.completed", {"prompt.name": prompt_name})
+
+                # Add message count and type information
+                if isinstance(result, list):
+                    span.set_attribute("mcp.prompt.result.message_count", len(result))
+                    span.set_attribute("mcp.prompt.result.type", "message_list")
+
+                    # Analyze message types if they have role attributes
+                    roles = []
+                    for msg in result:
+                        if hasattr(msg, "role"):
+                            roles.append(msg.role)
+                        elif isinstance(msg, dict) and "role" in msg:
+                            roles.append(msg["role"])
+
+                    if roles:
+                        unique_roles = list(set(roles))
+                        span.set_attribute("mcp.prompt.result.roles", ",".join(unique_roles))
+                        span.set_attribute(
+                            "mcp.prompt.result.role_counts",
+                            str({role: roles.count(role) for role in unique_roles}),
+                        )
+                elif isinstance(result, str):
+                    span.set_attribute("mcp.prompt.result.type", "string")
+                    span.set_attribute("mcp.prompt.result.length", len(result))
+                else:
+                    span.set_attribute("mcp.prompt.result.type", type(result).__name__)
+
+                return result
+            except Exception as e:
+                span.record_exception(e)
+                span.set_status(Status(StatusCode.ERROR, str(e)))
+
+                # Add event for error
+                span.add_event(
+                    "prompt.generation.error",
+                    {
+                        "prompt.name": prompt_name,
+                        "error.type": type(e).__name__,
+                        "error.message": str(e),
+                    },
+                )
+                raise
+
+    if asyncio.iscoroutinefunction(func):
+        return async_wrapper
+    else:
+        return sync_wrapper
+
+
+# Add the BoundedSessionTracker class before SessionTracingMiddleware
+class BoundedSessionTracker:
+    """Memory-safe session tracker with automatic expiration."""
+
+    def __init__(self, max_sessions: int = 1000, session_ttl: int = 3600) -> None:
+        self.max_sessions = max_sessions
+        self.session_ttl = session_ttl
+        self.sessions: OrderedDict[str, float] = OrderedDict()
+        self.last_cleanup = time.time()
+
+    def track_session(self, session_id: str) -> bool:
+        """Track a session, returns True if it's new."""
+        current_time = time.time()
+
+        # Periodic cleanup (every 5 minutes)
+        if current_time - self.last_cleanup > 300:
+            self._cleanup_expired(current_time)
+            self.last_cleanup = current_time
+
+        # Check if session exists and is still valid
+        if session_id in self.sessions:
+            # Move to end (mark as recently used)
+            self.sessions.move_to_end(session_id)
+            return False
+
+        # New session
+        self.sessions[session_id] = current_time
+
+        # Enforce max size
+        while len(self.sessions) > self.max_sessions:
+            self.sessions.popitem(last=False)  # Remove oldest
+
+        return True
+
+    def _cleanup_expired(self, current_time: float) -> None:
+        """Remove expired sessions."""
+        expired = [sid for sid, timestamp in self.sessions.items() if current_time - timestamp > self.session_ttl]
+        for sid in expired:
+            del self.sessions[sid]
+
+    def get_active_session_count(self) -> int:
+        return len(self.sessions)
+
+
+class SessionTracingMiddleware(BaseHTTPMiddleware):
+    def __init__(self, app: Any) -> None:
+        super().__init__(app)
+        # Use memory-safe session tracker instead of unbounded collections
+        self.session_tracker = BoundedSessionTracker(max_sessions=1000, session_ttl=3600)
+
+    async def dispatch(self, request: Any, call_next: Callable[..., Any]) -> Any:
+        # Record HTTP request timing
+        import time
+
+        start_time = time.time()
+
+        # Extract session ID from query params or headers
+        session_id = request.query_params.get("session_id")
+        if not session_id:
+            # Check headers as fallback
+            session_id = request.headers.get("x-session-id")
+
+        # Track session metrics using memory-safe tracker
+        if session_id:
+            is_new_session = self.session_tracker.track_session(session_id)
+
+            if is_new_session:
+                try:
+                    from golf.metrics import get_metrics_collector
+
+                    metrics_collector = get_metrics_collector()
+                    metrics_collector.increment_session()
+                except ImportError:
+                    pass
+            else:
+                # Record session duration for existing sessions
+                try:
+                    from golf.metrics import get_metrics_collector
+
+                    metrics_collector = get_metrics_collector()
+                    # Use a default duration since we don't track exact start
+                    # times anymore
+                    # This is less precise but memory-safe
+                    metrics_collector.record_session_duration(300.0)  # 5 min default
+                except ImportError:
+                    pass
+
+        # Create a descriptive span name based on the request
+        method = request.method
+        path = request.url.path
+
+        # Determine the operation type from the path
+        operation_type = "unknown"
+        if "/mcp" in path:
+            operation_type = "mcp.request"
+        elif "/sse" in path:
+            operation_type = "sse.stream"
+        elif "/auth" in path:
+            operation_type = "auth"
+
+        span_name = f"{operation_type}.{method.lower()}"
+
+        tracer = get_tracer()
+        with tracer.start_as_current_span(span_name) as span:
+            # Add essential HTTP attributes
+            span.set_attribute("http.method", method)
+            span.set_attribute("http.target", path)
+            span.set_attribute("http.host", request.url.hostname or "unknown")
+
+            # Add session tracking
+            if session_id:
+                span.set_attribute("mcp.session.id", session_id)
+                span.set_attribute(
+                    "mcp.session.active_count",
+                    self.session_tracker.get_active_session_count(),
+                )
+                # Add to baggage for propagation
+                ctx = baggage.set_baggage("mcp.session.id", session_id)
+                from opentelemetry import context
+
+                token = context.attach(ctx)
+            else:
+                token = None
+
+            # Add request size if available
+            content_length = request.headers.get("content-length")
+            if content_length:
+                span.set_attribute("http.request.size", int(content_length))
+
+            # Add event for request start
+            span.add_event("http.request.started", {"method": method, "path": path})
+
+            try:
+                response = await call_next(request)
+
+                # Add response attributes
+                span.set_attribute("http.status_code", response.status_code)
+
+                # Set span status based on HTTP status
+                if response.status_code >= 400:
+                    span.set_status(Status(StatusCode.ERROR, f"HTTP {response.status_code}"))
+                else:
+                    span.set_status(Status(StatusCode.OK))
+
+                # Add event for request completion
+                span.add_event(
+                    "http.request.completed",
+                    {
+                        "method": method,
+                        "path": path,
+                        "status_code": response.status_code,
+                    },
+                )
+
+                # Record HTTP request metrics
+                try:
+                    from golf.metrics import get_metrics_collector
+
+                    metrics_collector = get_metrics_collector()
+
+                    # Clean up path for metrics (remove query params, normalize)
+                    clean_path = path.split("?")[0]  # Remove query parameters
+                    if clean_path.startswith("/"):
+                        clean_path = clean_path[1:] or "root"  # Remove leading slash, handle root
+
+                    metrics_collector.increment_http_request(method, response.status_code, clean_path)
+                    metrics_collector.record_http_duration(method, clean_path, time.time() - start_time)
+                except ImportError:
+                    # Metrics not available, continue without metrics
+                    pass
+
+                return response
+            except Exception as e:
+                span.record_exception(e)
+                span.set_status(Status(StatusCode.ERROR, str(e)))
+
+                # Add event for error
+                span.add_event(
+                    "http.request.error",
+                    {
+                        "method": method,
+                        "path": path,
+                        "error.type": type(e).__name__,
+                        "error.message": str(e),
+                    },
+                )
+
+                # Record HTTP error metrics
+                try:
+                    from golf.metrics import get_metrics_collector
+
+                    metrics_collector = get_metrics_collector()
+
+                    # Clean up path for metrics
+                    clean_path = path.split("?")[0]
+                    if clean_path.startswith("/"):
+                        clean_path = clean_path[1:] or "root"
+
+                    metrics_collector.increment_http_request(method, 500, clean_path)  # Assume 500 for exceptions
+                    metrics_collector.increment_error("http", type(e).__name__)
+                except ImportError:
+                    pass
+
+                raise
+            finally:
+                if token:
+                    context.detach(token)
+
+
+@asynccontextmanager
+async def telemetry_lifespan(mcp_instance: Any) -> AsyncGenerator[None, None]:
+    """Simplified lifespan for telemetry initialization and cleanup."""
+    global _provider
+
+    # Initialize telemetry with the server name
+    provider = init_telemetry(service_name=mcp_instance.name)
+
+    # If provider is None, telemetry is disabled
+    if provider is None:
+        # Just yield without any telemetry setup
+        yield
+        return
+
+    # Try to add session tracking middleware if possible
+    try:
+        # Try to add middleware to FastMCP app if it has Starlette app
+        if hasattr(mcp_instance, "app") or hasattr(mcp_instance, "_app"):
+            app = getattr(mcp_instance, "app", getattr(mcp_instance, "_app", None))
+            if app and hasattr(app, "add_middleware"):
+                app.add_middleware(SessionTracingMiddleware)
+
+        # Also try to instrument FastMCP's internal handlers
+        if hasattr(mcp_instance, "_tool_manager") and hasattr(mcp_instance._tool_manager, "tools"):
+            # The tools should already be instrumented when they were registered
+            pass
+
+        # Try to patch FastMCP's request handling to ensure context propagation
+        if hasattr(mcp_instance, "handle_request"):
+            original_handle_request = mcp_instance.handle_request
+
+            async def traced_handle_request(*args: Any, **kwargs: Any) -> Any:
+                tracer = get_tracer()
+                with tracer.start_as_current_span("mcp.handle_request") as span:
+                    span.set_attribute("mcp.request.handler", "handle_request")
+                    return await original_handle_request(*args, **kwargs)
+
+            mcp_instance.handle_request = traced_handle_request
+
+    except Exception:
+        # Silently continue if middleware setup fails
+        import traceback
+
+        traceback.print_exc()
+
+    try:
+        # Yield control back to FastMCP
+        yield
+    finally:
+        # Cleanup - shutdown the provider
+        if _provider and hasattr(_provider, "shutdown"):
+            _provider.force_flush()
+            _provider.shutdown()
+            _provider = None