feat: add OpenTelemetry tracing for tool calls

Author: priyank766

README.md

@@ -121,6 +121,7 @@ kubeflow-mcp serve \
   --instruction-tier full \       # full | compact | minimal
   --transport stdio \             # stdio | http | sse
   --auth-token SECRET \           # bearer token for HTTP auth (dev/staging)
+  --otel-endpoint URL \           # OTLP HTTP endpoint (optional tracing)
   --log-level INFO \              # DEBUG | INFO | WARNING | ERROR
   --log-format console \          # console | json (auto-detected if omitted)
   --no-banner                     # suppress startup banner
@@ -165,6 +166,23 @@ kubeflow-mcp agent \
 
 </details>
 
+## Observability
+
+OpenTelemetry tracing is optional and can be enabled without changing tool code.
+
+- Install optional dependencies: `pip install ".[otel]"`
+- Enable tracing with CLI flag or env var:
+
+```bash
+kubeflow-mcp serve --otel-endpoint http://localhost:4318/v1/traces
+# or
+export KUBEFLOW_MCP_OTEL_ENDPOINT=http://localhost:4318/v1/traces
+kubeflow-mcp serve
+```
+
+Each tool invocation emits a span with attributes:
+`tool.name`, `tool.success`, `tool.duration_ms`, `kubeflow.persona`, and `correlation_id`.
+
 ## Development
 
 ```bash

kubeflow_mcp/cli.py

@@ -94,6 +94,12 @@ def cli() -> None:
     "Falls back to KUBEFLOW_MCP_AUTH_TOKEN env var, config file. "
     "Ignored for stdio transport.",
 )
+@click.option(
+    "--otel-endpoint",
+    default=None,
+    help="OpenTelemetry OTLP HTTP endpoint for tracing. "
+    "Falls back to KUBEFLOW_MCP_OTEL_ENDPOINT env var, config file.",
+)
 def serve(
     clients: str | None,
     persona: str | None,
@@ -104,6 +110,7 @@ def serve(
     instruction_tier: str | None,
     no_banner: bool,
     auth_token: str | None,
+    otel_endpoint: str | None,
 ) -> None:
     """Start the MCP server.
 
@@ -116,6 +123,7 @@ def serve(
     from kubeflow_mcp.core.logging import setup_logging
     from kubeflow_mcp.core.resilience import configure_circuit_breaker
     from kubeflow_mcp.core.server import configure_resilience, create_server
+    from kubeflow_mcp.core.telemetry import setup_tracing
 
     cfg = load_config()
 
@@ -128,8 +136,11 @@ def serve(
 
     if auth_token:
         cfg.auth.auth_token = auth_token
+    if otel_endpoint:
+        cfg.observability.otel_endpoint = otel_endpoint
 
     logger = setup_logging(level=log_level, format=log_format)
+    tracing_enabled = setup_tracing(endpoint=cfg.observability.otel_endpoint)
     logger.info(
         "Starting kubeflow-mcp",
         extra={
@@ -138,6 +149,7 @@ def serve(
             "transport": transport,
             "mode": mode,
             "instruction_tier": instruction_tier,
+            "tracing_enabled": tracing_enabled,
         },
     )
 

kubeflow_mcp/cli_test.py

@@ -97,6 +97,7 @@ def _make_serve_mocks(config=None):
     mock_load_config = MagicMock(return_value=config)
     mock_build_auth_provider = MagicMock(return_value=None)
     mock_configure_circuit_breaker = MagicMock()
+    mock_setup_tracing = MagicMock(return_value=False)
 
     fake_server_mod = MagicMock()
     fake_server_mod.create_server = mock_create_server
@@ -109,13 +110,16 @@ def _make_serve_mocks(config=None):
     fake_auth_mod.build_auth_provider = mock_build_auth_provider
     fake_resilience_mod = MagicMock()
     fake_resilience_mod.configure_circuit_breaker = mock_configure_circuit_breaker
+    fake_telemetry_mod = MagicMock()
+    fake_telemetry_mod.setup_tracing = mock_setup_tracing
 
     modules_patch = {
         "kubeflow_mcp.core.server": fake_server_mod,
         "kubeflow_mcp.core.logging": fake_logging_mod,
         "kubeflow_mcp.core.config": fake_config_mod,
         "kubeflow_mcp.core.auth": fake_auth_mod,
         "kubeflow_mcp.core.resilience": fake_resilience_mod,
+        "kubeflow_mcp.core.telemetry": fake_telemetry_mod,
     }
     return mock_server, mock_create_server, modules_patch
 
@@ -401,3 +405,40 @@ def test_serve_default_shows_banner():
 
     _, kwargs = mock_server.run.call_args
     assert kwargs.get("show_banner") is True
+
+
+def test_serve_calls_setup_tracing_with_config_endpoint():
+    from kubeflow_mcp.core.config import ObservabilityConfig
+
+    config = _make_default_config()
+    config.observability = ObservabilityConfig(otel_endpoint="http://otel-collector:4318/v1/traces")
+    mock_server, _, modules_patch = _make_serve_mocks(config=config)
+    fake_telemetry_mod = modules_patch["kubeflow_mcp.core.telemetry"]
+
+    with patch.dict(sys.modules, modules_patch):
+        runner = CliRunner()
+        runner.invoke(cli, ["serve"])
+
+    fake_telemetry_mod.setup_tracing.assert_called_once_with(
+        endpoint="http://otel-collector:4318/v1/traces"
+    )
+
+
+def test_serve_otel_endpoint_cli_overrides_config():
+    from kubeflow_mcp.core.config import ObservabilityConfig
+
+    config = _make_default_config()
+    config.observability = ObservabilityConfig(otel_endpoint="http://old-endpoint:4318/v1/traces")
+    mock_server, _, modules_patch = _make_serve_mocks(config=config)
+    fake_telemetry_mod = modules_patch["kubeflow_mcp.core.telemetry"]
+
+    with patch.dict(sys.modules, modules_patch):
+        runner = CliRunner()
+        runner.invoke(
+            cli,
+            ["serve", "--otel-endpoint", "http://new-endpoint:4318/v1/traces"],
+        )
+
+    fake_telemetry_mod.setup_tracing.assert_called_once_with(
+        endpoint="http://new-endpoint:4318/v1/traces"
+    )

kubeflow_mcp/core/config.py

@@ -50,6 +50,9 @@
             level: INFO
             format: json
 
+        observability:
+            otel_endpoint: http://localhost:4318/v1/traces
+
 Namespace restrictions are enforced via ``~/.kf-mcp-policy.yaml``
 (``policy.namespaces``), not through server config.
 """
@@ -135,6 +138,12 @@ class LoggingConfig(BaseModel):
     format: str | None = Field(default=None)
 
 
+class ObservabilityConfig(BaseModel):
+    """Observability configuration."""
+
+    otel_endpoint: str | None = Field(default=None)
+
+
 class Config(BaseModel):
     """Root configuration."""
 
@@ -144,6 +153,7 @@ class Config(BaseModel):
     trainer: TrainerConfig = Field(default_factory=TrainerConfig)
     optimizer: OptimizerConfig = Field(default_factory=OptimizerConfig)
     logging: LoggingConfig = Field(default_factory=LoggingConfig)
+    observability: ObservabilityConfig = Field(default_factory=ObservabilityConfig)
 
 
 def _find_config_file() -> Path | None:
@@ -229,6 +239,14 @@ def load_config(config_path: Path | None = None) -> Config:
         format=os.getenv("LOG_FORMAT", logging_file.get("format")),
     )
 
+    observability_file = file_config.get("observability", {})
+    observability = ObservabilityConfig(
+        otel_endpoint=os.getenv(
+            "KUBEFLOW_MCP_OTEL_ENDPOINT",
+            observability_file.get("otel_endpoint"),
+        )
+    )
+
     # Build client-specific configs
     trainer_file = file_config.get("trainer", {})
     trainer = TrainerConfig(
@@ -283,6 +301,7 @@ def load_config(config_path: Path | None = None) -> Config:
         trainer=trainer,
         optimizer=optimizer,
         logging=logging_config,
+        observability=observability,
     )
 
 

kubeflow_mcp/core/logging.py

@@ -49,7 +49,7 @@ def format(self, record: logging.LogRecord) -> str:
         if ctx is not None:
             log_dict["context"] = ctx
 
-        extra_keys = {"audit", "tool", "parameters", "success", "duration_ms"}
+        extra_keys = {"audit", "tool", "parameters", "success", "duration_ms", "tracing_enabled"}
         for key in extra_keys:
             if hasattr(record, key):
                 log_dict[key] = getattr(record, key)

kubeflow_mcp/core/server.py

@@ -41,10 +41,16 @@
     HEALTH_TOOLS,
 )
 from kubeflow_mcp.core.logging import with_correlation_id
-from kubeflow_mcp.core.policy import apply_policy_filters, get_allowed_tools, is_read_only
+from kubeflow_mcp.core.policy import (
+    apply_policy_filters,
+    get_allowed_tools,
+    get_effective_persona,
+    is_read_only,
+)
 from kubeflow_mcp.core.resilience import RateLimiter, get_breaker
 from kubeflow_mcp.core.resources import register_resources
 from kubeflow_mcp.core.security import mask_sensitive_data
+from kubeflow_mcp.core.telemetry import get_tracer
 
 logger = logging.getLogger(__name__)
 
@@ -88,66 +94,84 @@ def _audit_wrap(tool_func):
     @functools.wraps(tool_func)
     def wrapper(**kwargs):
         tool_name = tool_func.__name__
-
-        if _rate_limiter is not None and not _rate_limiter.acquire():
-            logger.warning("rate_limited", extra={"tool": tool_name})
-            return {
-                "error": "Rate limit exceeded. Retry after a brief pause.",
-                "error_code": ErrorCode.RATE_LIMITED,
-            }
-
-        breaker = get_breaker(tool_name)
-        if not breaker.can_execute():
-            logger.warning("circuit_open", extra={"tool": tool_name})
-            return {
-                "error": f"Circuit breaker open for '{tool_name}' — K8s API may be degraded. Retries automatically after recovery timeout.",
-                "error_code": ErrorCode.CIRCUIT_OPEN,
-            }
-
         cid = with_correlation_id()
-        masked = mask_sensitive_data(kwargs) if kwargs else {}
+        tracer = get_tracer("kubeflow_mcp.tools")
+        persona = get_effective_persona()
         start = time.monotonic()
-        try:
-            result = tool_func(**kwargs)
-            duration_ms = int((time.monotonic() - start) * 1000)
-            is_success = (
-                "error_code" not in result and "error" not in result
-                if isinstance(result, dict)
-                else True
-            )
-            if is_success:
-                breaker.record_success()
-            elif is_infrastructure_error(result):
-                breaker.record_failure()
 
-            logger.info(
-                "tool_call",
-                extra={
-                    "audit": True,
-                    "correlation_id": cid,
-                    "tool": tool_name,
-                    "parameters": masked,
-                    "success": is_success,
-                    "duration_ms": duration_ms,
-                },
-            )
-            return _inject_meta(result, tool_name)
-        except Exception:
-            duration_ms = int((time.monotonic() - start) * 1000)
-            breaker.record_failure()
-            logger.error(
-                "tool_call_failed",
-                extra={
-                    "audit": True,
-                    "correlation_id": cid,
-                    "tool": tool_name,
-                    "parameters": masked,
-                    "success": False,
-                    "duration_ms": duration_ms,
-                },
-                exc_info=True,
-            )
-            raise
+        with tracer.start_as_current_span("tool_call") as span:
+            span.set_attribute("tool.name", tool_name)
+            span.set_attribute("kubeflow.persona", persona)
+            span.set_attribute("correlation_id", cid)
+
+            if _rate_limiter is not None and not _rate_limiter.acquire():
+                duration_ms = int((time.monotonic() - start) * 1000)
+                span.set_attribute("tool.success", False)
+                span.set_attribute("tool.duration_ms", duration_ms)
+                logger.warning("rate_limited", extra={"tool": tool_name})
+                return {
+                    "error": "Rate limit exceeded. Retry after a brief pause.",
+                    "error_code": ErrorCode.RATE_LIMITED,
+                }
+
+            breaker = get_breaker(tool_name)
+            if not breaker.can_execute():
+                duration_ms = int((time.monotonic() - start) * 1000)
+                span.set_attribute("tool.success", False)
+                span.set_attribute("tool.duration_ms", duration_ms)
+                logger.warning("circuit_open", extra={"tool": tool_name})
+                return {
+                    "error": f"Circuit breaker open for '{tool_name}' — K8s API may be degraded. Retries automatically after recovery timeout.",
+                    "error_code": ErrorCode.CIRCUIT_OPEN,
+                }
+
+            masked = mask_sensitive_data(kwargs) if kwargs else {}
+            try:
+                result = tool_func(**kwargs)
+                duration_ms = int((time.monotonic() - start) * 1000)
+                is_success = (
+                    "error_code" not in result and "error" not in result
+                    if isinstance(result, dict)
+                    else True
+                )
+                span.set_attribute("tool.success", is_success)
+                span.set_attribute("tool.duration_ms", duration_ms)
+                if is_success:
+                    breaker.record_success()
+                elif is_infrastructure_error(result):
+                    breaker.record_failure()
+
+                logger.info(
+                    "tool_call",
+                    extra={
+                        "audit": True,
+                        "correlation_id": cid,
+                        "tool": tool_name,
+                        "parameters": masked,
+                        "success": is_success,
+                        "duration_ms": duration_ms,
+                    },
+                )
+                return _inject_meta(result, tool_name)
+            except Exception as exc:
+                duration_ms = int((time.monotonic() - start) * 1000)
+                breaker.record_failure()
+                span.set_attribute("tool.success", False)
+                span.set_attribute("tool.duration_ms", duration_ms)
+                span.record_exception(exc)
+                logger.error(
+                    "tool_call_failed",
+                    extra={
+                        "audit": True,
+                        "correlation_id": cid,
+                        "tool": tool_name,
+                        "parameters": masked,
+                        "success": False,
+                        "duration_ms": duration_ms,
+                    },
+                    exc_info=True,
+                )
+                raise
 
     return wrapper