kubiya-control-plane-api 0.9.15__py3-none-any.whl

control_plane_api/app/observability/metrics.py

@@ -0,0 +1,303 @@
+"""
+Prometheus metrics for Kubiya Control Plane.
+
+This module provides comprehensive metrics for monitoring the AI orchestration layer:
+- HTTP request metrics (latency, error rates, request counts)
+- Execution metrics (active count, failures, duration, wait time)
+- Worker and queue metrics
+- LLM/AI metrics (requests, latency, tokens)
+- Tool execution metrics
+- Business metrics (agents, jobs, organizations)
+
+Metrics naming convention follows Prometheus best practices:
+- Prefix: kubiya_control_plane_ for HTTP metrics
+- Prefix: kubiya_ for business metrics
+- Suffix: _total for counters, _seconds for durations, _count for gauges
+"""
+
+import re
+import os
+import structlog
+from typing import Tuple, Optional
+from functools import lru_cache
+
+from prometheus_client import (
+    Counter,
+    Histogram,
+    Gauge,
+    REGISTRY,
+    generate_latest,
+    CONTENT_TYPE_LATEST,
+    CollectorRegistry,
+)
+
+logger = structlog.get_logger(__name__)
+
+PROMETHEUS_MULTIPROC_DIR = os.environ.get('PROMETHEUS_MULTIPROC_DIR')
+
+HTTP_REQUESTS_TOTAL = Counter(
+    'kubiya_control_plane_http_requests_total',
+    'Total number of HTTP requests received by the Control Plane API',
+    ['method', 'endpoint', 'status_code'],
+)
+
+HTTP_REQUEST_DURATION_SECONDS = Histogram(
+    'kubiya_control_plane_http_request_duration_seconds',
+    'HTTP request duration in seconds (latency)',
+    ['method', 'endpoint'],
+    buckets=(0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0, 30.0, 60.0, 120.0),
+)
+
+
+ACTIVE_TASKS_COUNT = Gauge(
+    'kubiya_active_tasks_count',
+    'Number of currently active tasks/executions by type and status',
+    ['execution_type', 'status'],
+)
+
+TASK_FAILURE_TOTAL = Gauge(
+    'kubiya_task_failure_total',
+    'Total number of failed task/executions by type',
+    ['execution_type'],
+)
+
+
+# Execution duration - how long executions take
+EXECUTION_DURATION_SECONDS = Gauge(
+    'kubiya_execution_duration_seconds',
+    'Average execution duration in seconds by type and status (from completed executions)',
+    ['execution_type', 'status'],
+)
+
+# Worker queue depth - pending executions per queue
+WORKER_QUEUE_DEPTH = Gauge(
+    'kubiya_worker_queue_depth',
+    'Number of pending executions waiting in queue',
+    ['queue_id'],
+)
+
+# LLM requests total
+LLM_REQUESTS_TOTAL = Gauge(
+    'kubiya_llm_requests_total',
+    'Total number of LLM API requests by model and status',
+    ['model', 'status'],
+)
+
+# LLM latency
+LLM_LATENCY_SECONDS = Gauge(
+    'kubiya_llm_latency_seconds',
+    'Average LLM request latency in seconds by model',
+    ['model'],
+)
+
+# LLM tokens total
+LLM_TOKENS_TOTAL = Gauge(
+    'kubiya_llm_tokens_total',
+    'Total tokens used in LLM requests by model and type',
+    ['model', 'token_type'],
+)
+
+
+# Active streaming connections
+STREAMING_CONNECTIONS_ACTIVE = Gauge(
+    'kubiya_streaming_connections_active',
+    'Number of active SSE streaming connections',
+)
+
+# Tool execution duration
+TOOL_EXECUTION_DURATION_SECONDS = Gauge(
+    'kubiya_tool_execution_duration_seconds',
+    'Average tool execution duration in seconds by tool name',
+    ['tool_name'],
+)
+
+# Tool executions total
+TOOL_EXECUTIONS_TOTAL = Gauge(
+    'kubiya_tool_executions_total',
+    'Total number of tool executions by tool name and status',
+    ['tool_name', 'status'],
+)
+
+# Execution wait time (time from created to running)
+EXECUTION_WAIT_TIME_SECONDS = Gauge(
+    'kubiya_execution_wait_time_seconds',
+    'Average wait time from creation to running in seconds by type',
+    ['execution_type'],
+)
+
+# Webhook requests total
+WEBHOOK_REQUESTS_TOTAL = Gauge(
+    'kubiya_webhook_requests_total',
+    'Total number of webhook trigger requests by status',
+    ['status'],
+)
+
+
+# Executions by organization
+EXECUTIONS_BY_ORG_TOTAL = Gauge(
+    'kubiya_executions_by_org_total',
+    'Total number of executions by organization',
+    ['organization_id'],
+)
+
+# Active agents
+AGENTS_ACTIVE = Gauge(
+    'kubiya_agents_active',
+    'Number of active agents by organization',
+    ['organization_id'],
+)
+
+# Scheduled jobs total
+SCHEDULED_JOBS_TOTAL = Gauge(
+    'kubiya_scheduled_jobs_total',
+    'Total number of scheduled jobs by status',
+    ['status'],
+)
+
+# Conversation turns total
+CONVERSATION_TURNS_TOTAL = Gauge(
+    'kubiya_conversation_turns_total',
+    'Total conversation turns by execution type',
+    ['execution_type'],
+)
+
+
+UUID_PATTERN = re.compile(
+    r'[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}',
+    re.IGNORECASE
+)
+NUMERIC_ID_PATTERN = re.compile(r'/\d+(?=/|$)')
+
+
+@lru_cache(maxsize=1000)
+def normalize_endpoint(path: str) -> str:
+    """Normalize endpoint paths to reduce cardinality in metrics."""
+    normalized = UUID_PATTERN.sub('{id}', path)
+    normalized = NUMERIC_ID_PATTERN.sub('/{id}', normalized)
+    return normalized
+
+
+def get_metrics_response() -> Tuple[bytes, str]:
+    """Generate Prometheus metrics response."""
+    try:
+        if PROMETHEUS_MULTIPROC_DIR:
+            from prometheus_client.multiprocess import MultiProcessCollector
+            registry = CollectorRegistry()
+            MultiProcessCollector(registry)
+            output = generate_latest(registry)
+        else:
+            output = generate_latest(REGISTRY)
+        return output, CONTENT_TYPE_LATEST
+    except Exception as e:
+        logger.error("metrics_generation_failed", error=str(e))
+        return b"# Error generating metrics\n", CONTENT_TYPE_LATEST
+
+
+def record_http_request(
+    method: str,
+    endpoint: str,
+    status_code: int,
+    duration_seconds: float,
+    trace_id: Optional[str] = None,
+):
+    """Record an HTTP request in Prometheus metrics."""
+    HTTP_REQUESTS_TOTAL.labels(
+        method=method,
+        endpoint=endpoint,
+        status_code=str(status_code),
+    ).inc()
+    
+    HTTP_REQUEST_DURATION_SECONDS.labels(
+        method=method,
+        endpoint=endpoint,
+    ).observe(duration_seconds)
+    
+    if status_code >= 500:
+        logger.error(
+            "prometheus_5xx_recorded",
+            method=method,
+            endpoint=endpoint,
+            status_code=status_code,
+            duration_seconds=round(duration_seconds, 4),
+            trace_id=trace_id or "unknown",
+        )
+
+
+def update_active_tasks(execution_type: str, status: str, count: int):
+    """Update the active tasks gauge."""
+    ACTIVE_TASKS_COUNT.labels(execution_type=execution_type, status=status).set(count)
+
+
+def update_task_failures(execution_type: str, count: int):
+    """Update the task failures gauge."""
+    TASK_FAILURE_TOTAL.labels(execution_type=execution_type).set(count)
+
+
+def update_execution_duration(execution_type: str, status: str, avg_seconds: float):
+    """Update average execution duration."""
+    EXECUTION_DURATION_SECONDS.labels(execution_type=execution_type, status=status).set(avg_seconds)
+
+
+def update_worker_queue_depth(queue_id: str, count: int):
+    """Update worker queue depth."""
+    WORKER_QUEUE_DEPTH.labels(queue_id=queue_id).set(count)
+
+
+def update_llm_requests(model: str, status: str, count: int):
+    """Update LLM requests total."""
+    LLM_REQUESTS_TOTAL.labels(model=model, status=status).set(count)
+
+
+def update_llm_latency(model: str, avg_seconds: float):
+    """Update average LLM latency."""
+    LLM_LATENCY_SECONDS.labels(model=model).set(avg_seconds)
+
+
+def update_llm_tokens(model: str, token_type: str, count: int):
+    """Update LLM tokens total."""
+    LLM_TOKENS_TOTAL.labels(model=model, token_type=token_type).set(count)
+
+
+def update_streaming_connections(count: int):
+    """Update active streaming connections."""
+    STREAMING_CONNECTIONS_ACTIVE.set(count)
+
+
+def update_tool_execution_duration(tool_name: str, avg_seconds: float):
+    """Update average tool execution duration."""
+    TOOL_EXECUTION_DURATION_SECONDS.labels(tool_name=tool_name).set(avg_seconds)
+
+
+def update_tool_executions(tool_name: str, status: str, count: int):
+    """Update tool executions total."""
+    TOOL_EXECUTIONS_TOTAL.labels(tool_name=tool_name, status=status).set(count)
+
+
+def update_execution_wait_time(execution_type: str, avg_seconds: float):
+    """Update average execution wait time."""
+    EXECUTION_WAIT_TIME_SECONDS.labels(execution_type=execution_type).set(avg_seconds)
+
+
+def update_webhook_requests(status: str, count: int):
+    """Update webhook requests total."""
+    WEBHOOK_REQUESTS_TOTAL.labels(status=status).set(count)
+
+
+def update_executions_by_org(organization_id: str, count: int):
+    """Update executions by organization."""
+    EXECUTIONS_BY_ORG_TOTAL.labels(organization_id=organization_id).set(count)
+
+
+def update_agents_active(organization_id: str, count: int):
+    """Update active agents count."""
+    AGENTS_ACTIVE.labels(organization_id=organization_id).set(count)
+
+
+def update_scheduled_jobs(status: str, count: int):
+    """Update scheduled jobs total."""
+    SCHEDULED_JOBS_TOTAL.labels(status=status).set(count)
+
+
+def update_conversation_turns(execution_type: str, count: int):
+    """Update conversation turns total."""
+    CONVERSATION_TURNS_TOTAL.labels(execution_type=execution_type).set(count)

control_plane_api/app/observability/middleware.py

@@ -0,0 +1,246 @@
+"""
+OpenTelemetry middleware for FastAPI.
+
+This module provides middleware to:
+- Add trace ID to response headers (X-Trace-ID)
+- Enrich spans with organizational and user context from request.state
+- Set span status based on HTTP status codes
+"""
+
+import structlog
+from starlette.middleware.base import BaseHTTPMiddleware
+from starlette.requests import Request
+from starlette.responses import Response
+from opentelemetry import trace
+from opentelemetry.trace import Status, StatusCode
+
+logger = structlog.get_logger(__name__)
+
+
+class TraceContextMiddleware(BaseHTTPMiddleware):
+    """
+    Middleware to enrich spans with organizational context and add trace ID to responses.
+
+    This middleware should be added AFTER the OpenTelemetry FastAPI instrumentation,
+    so that it can enrich the automatically created spans.
+    """
+
+    # Paths to exclude from tracing (health checks, metrics, etc.)
+    EXCLUDED_PATHS = {
+        "/api/health",
+        "/health",
+        "/health/live",
+        "/health/ready",
+        "/health/detailed",
+        "/health/event-bus",
+        "/health/temporal-credentials",
+        "/ready",
+        "/metrics",
+        "/favicon.ico"
+    }
+
+    async def dispatch(self, request: Request, call_next):
+        """
+        Process request and enrich span with organizational context.
+
+        Args:
+            request: Incoming HTTP request
+            call_next: Next middleware in chain
+
+        Returns:
+            HTTP response with X-Trace-ID header
+        """
+        # Skip tracing for health checks and other excluded paths
+        if request.url.path in self.EXCLUDED_PATHS:
+            return await call_next(request)
+
+        # Get current span (created by FastAPI instrumentation)
+        span = trace.get_current_span()
+
+        # Add span event for request received
+        if span and span.is_recording():
+            from control_plane_api.app.observability import add_span_event
+            add_span_event(
+                f"HTTP request received: {request.method} {request.url.path}",
+                {
+                    "http.method": request.method,
+                    "http.path": request.url.path,
+                    "http.query": request.url.query if request.url.query else "",
+                    "client.host": request.client.host if request.client else "unknown",
+                }
+            )
+
+        # Capture request body for non-GET requests (for debugging)
+        request_body = None
+        if request.method in ("POST", "PUT", "PATCH"):
+            try:
+                body_bytes = await request.body()
+                if body_bytes and len(body_bytes) < 10000:  # Only capture if < 10KB
+                    request_body = body_bytes.decode('utf-8')
+                    # Re-create request with body for downstream handlers
+                    async def receive():
+                        return {"type": "http.request", "body": body_bytes}
+                    request._receive = receive
+            except Exception as e:
+                logger.warning("failed_to_capture_request_body", error=str(e))
+
+        # Enrich span with organizational context from request.state
+        # These are set by the auth dependency (get_current_organization)
+        if span and span.is_recording():
+            try:
+                # Debug: Check if organization exists
+                has_org = hasattr(request.state, "organization")
+                logger.debug(
+                    "trace_context_check",
+                    has_organization=has_org,
+                    path=request.url.path
+                )
+
+                # Add organization context
+                if has_org:
+                    org = request.state.organization
+                    if isinstance(org, dict):
+                        span.set_attribute("organization.id", org.get("id", ""))
+                        span.set_attribute("organization.name", org.get("name", ""))
+
+                        from control_plane_api.app.observability import add_span_event
+                        add_span_event(
+                            "Organizational context added to span",
+                            {
+                                "organization.id": org.get("id", ""),
+                                "organization.name": org.get("name", ""),
+                                "user.email": org.get("user_email", ""),
+                            }
+                        )
+
+                        logger.info(
+                            "span_enriched_with_org",
+                            org_id=org.get("id"),
+                            path=request.url.path
+                        )
+
+                        # Add user context
+                        if org.get("user_id"):
+                            span.set_attribute("user.id", org["user_id"])
+                        if org.get("user_email"):
+                            span.set_attribute("user.email", org["user_email"])
+                        if org.get("user_name"):
+                            span.set_attribute("user.name", org["user_name"])
+                        if org.get("user_avatar"):
+                            span.set_attribute("user.avatar", org["user_avatar"])
+
+                # Add request ID
+                if hasattr(request.state, "request_id"):
+                    span.set_attribute("request.id", request.state.request_id)
+
+                # Add request path and method
+                span.set_attribute("http.route", request.url.path)
+                span.set_attribute("http.method", request.method)
+
+                # Add query parameters
+                if request.url.query:
+                    span.set_attribute("http.query", request.url.query)
+
+                # Add request body for debugging (sanitize sensitive data)
+                if request_body:
+                    # Sanitize passwords, tokens, etc.
+                    sanitized_body = request_body
+                    for sensitive_key in ["password", "token", "secret", "api_key", "apiKey"]:
+                        if sensitive_key in sanitized_body.lower():
+                            sanitized_body = sanitized_body[:100] + "...[REDACTED]"
+                            break
+                    span.set_attribute("http.request.body", sanitized_body[:500])  # Max 500 chars
+
+            except Exception as e:
+                logger.warning(
+                    "span_enrichment_failed",
+                    error=str(e),
+                    exc_info=True
+                )
+
+        # Get trace ID before processing request for logging correlation
+        trace_id = None
+        span_id = None
+        if span and span.is_recording():
+            trace_id = format(span.get_span_context().trace_id, '032x')
+            span_id = format(span.get_span_context().span_id, '016x')
+
+        # Process request
+        response = await call_next(request)
+
+        # Add trace ID to response headers and capture response
+        if span and span.is_recording():
+            try:
+                response.headers["X-Trace-ID"] = trace_id
+                response.headers["X-Span-ID"] = span_id
+
+                # Set span status based on HTTP status code
+                from control_plane_api.app.observability import add_span_event
+                if response.status_code >= 500:
+                    span.set_status(Status(StatusCode.ERROR, f"HTTP {response.status_code}"))
+                    span.set_attribute("error", True)
+                    add_span_event(
+                        f"HTTP response: Server error {response.status_code}",
+                        {
+                            "http.status_code": response.status_code,
+                            "status": "error",
+                        }
+                    )
+                elif response.status_code >= 400:
+                    # Client errors are not span errors (they're expected)
+                    span.set_attribute("http.client_error", True)
+                    span.set_status(Status(StatusCode.OK))
+                    add_span_event(
+                        f"HTTP response: Client error {response.status_code}",
+                        {
+                            "http.status_code": response.status_code,
+                            "status": "client_error",
+                        }
+                    )
+                else:
+                    span.set_status(Status(StatusCode.OK))
+                    add_span_event(
+                        f"HTTP response: Success {response.status_code}",
+                        {
+                            "http.status_code": response.status_code,
+                            "status": "success",
+                        }
+                    )
+
+                # Add HTTP status code attribute
+                span.set_attribute("http.status_code", response.status_code)
+
+                # NOW enrich span with user context (AFTER route handler ran, so organization is set)
+                # This is the correct place because auth dependency sets request.state.organization
+                if hasattr(request.state, "organization"):
+                    org = request.state.organization
+                    if isinstance(org, dict):
+                        span.set_attribute("organization.id", org.get("id", ""))
+                        if org.get("user_id"):
+                            span.set_attribute("user.id", org["user_id"])
+                        if org.get("user_email"):
+                            span.set_attribute("user.email", org["user_email"])
+                        if org.get("user_name"):
+                            span.set_attribute("user.name", org["user_name"])
+                        if org.get("user_avatar"):
+                            span.set_attribute("user.avatar", org["user_avatar"])
+
+                # Log request completion with trace correlation
+                logger.info(
+                    "http_request_completed",
+                    method=request.method,
+                    path=request.url.path,
+                    status_code=response.status_code,
+                    trace_id=trace_id,
+                    span_id=span_id,
+                    organization_id=getattr(request.state, "organization", {}).get("id") if hasattr(request.state, "organization") else None
+                )
+
+            except Exception as e:
+                logger.warning(
+                    "trace_id_header_failed",
+                    error=str(e),
+                    exc_info=True
+                )
+
+        return response

control_plane_api/app/observability/optional.py

@@ -0,0 +1,115 @@
+"""
+Optional OpenTelemetry imports for environments with size constraints.
+
+This module provides safe imports for OpenTelemetry that gracefully degrade
+when the packages are not available (e.g., in Vercel serverless deployments
+where package size matters).
+"""
+
+import os
+from typing import Optional, Any
+from contextlib import nullcontext
+
+# Check if we should disable tracing entirely
+TRACING_ENABLED = os.getenv("TRACING_ENABLED", "true").lower() in ("true", "1", "yes")
+
+# Try to import OpenTelemetry, but fall back gracefully
+try:
+    if not TRACING_ENABLED:
+        raise ImportError("Tracing disabled via TRACING_ENABLED env var")
+
+    from opentelemetry import trace as _trace
+    from opentelemetry.trace import Status, StatusCode, Span, Tracer
+
+    HAS_OPENTELEMETRY = True
+    trace = _trace
+
+except ImportError:
+    HAS_OPENTELEMETRY = False
+
+    # Create no-op implementations
+    class NoOpSpan:
+        """No-op span that does nothing."""
+        def set_attribute(self, key: str, value: Any) -> None:
+            pass
+
+        def set_attributes(self, attributes: dict) -> None:
+            pass
+
+        def set_status(self, status: Any) -> None:
+            pass
+
+        def record_exception(self, exception: Exception) -> None:
+            pass
+
+        def is_recording(self) -> bool:
+            return False
+
+        def get_span_context(self) -> Any:
+            return None
+
+        def __enter__(self):
+            return self
+
+        def __exit__(self, *args):
+            pass
+
+    class NoOpTracer:
+        """No-op tracer that does nothing."""
+        def start_as_current_span(self, name: str, *args, **kwargs):
+            return nullcontext(NoOpSpan())
+
+        def start_span(self, name: str, *args, **kwargs):
+            return NoOpSpan()
+
+    class NoOpTracerProvider:
+        """No-op tracer provider."""
+        def get_tracer(self, *args, **kwargs) -> NoOpTracer:
+            return NoOpTracer()
+
+    class NoOpTrace:
+        """No-op trace module replacement."""
+        @staticmethod
+        def get_tracer(name: str, version: str = "") -> NoOpTracer:
+            return NoOpTracer()
+
+        @staticmethod
+        def get_current_span() -> NoOpSpan:
+            return NoOpSpan()
+
+        @staticmethod
+        def get_tracer_provider() -> NoOpTracerProvider:
+            return NoOpTracerProvider()
+
+        @staticmethod
+        def set_tracer_provider(provider: Any) -> None:
+            pass
+
+    # Status and StatusCode placeholders
+    class Status:
+        """No-op Status."""
+        def __init__(self, status_code: Any, description: str = ""):
+            pass
+
+    class StatusCode:
+        """No-op StatusCode."""
+        OK = "OK"
+        ERROR = "ERROR"
+        UNSET = "UNSET"
+
+    # Type hints
+    Span = NoOpSpan
+    Tracer = NoOpTracer
+
+    # Create the no-op trace module
+    trace = NoOpTrace()
+
+
+def get_tracer(name: str, version: str = "") -> Tracer:
+    """Get a tracer, or a no-op tracer if OpenTelemetry is not available."""
+    return trace.get_tracer(name, version)
+
+
+def get_current_span() -> Span:
+    """Get the current span, or a no-op span if OpenTelemetry is not available."""
+    return trace.get_current_span()