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

control_plane_api/app/observability/tracing.py

@@ -0,0 +1,382 @@
+"""
+OpenTelemetry tracing configuration and utilities.
+
+This module sets up the OpenTelemetry SDK with OTLP exporter support
+and provides helper functions for creating spans with organizational context.
+"""
+
+import structlog
+from typing import Optional, Dict, Any
+from contextlib import contextmanager
+
+from opentelemetry import trace
+from opentelemetry.sdk.trace import TracerProvider
+from opentelemetry.sdk.trace.export import BatchSpanProcessor
+from opentelemetry.sdk.resources import Resource, SERVICE_NAME, SERVICE_VERSION
+from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter as GRPCExporter
+from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter as HTTPExporter
+from opentelemetry.sdk.trace.sampling import (
+    ParentBasedTraceIdRatio,
+    ParentBased,
+    ALWAYS_ON,
+    ALWAYS_OFF,
+)
+
+from control_plane_api.app.config import settings
+
+logger = structlog.get_logger(__name__)
+
+# Global tracer instance
+_tracer: Optional[trace.Tracer] = None
+_tracer_provider: Optional[TracerProvider] = None
+_local_storage_processor = None  # LocalStorageSpanProcessor instance
+
+
+def _parse_resource_attributes() -> Dict[str, str]:
+    """Parse OTEL_RESOURCE_ATTRIBUTES from environment variable.
+
+    Format: "key1=value1,key2=value2"
+    """
+    attributes = {}
+    if settings.OTEL_RESOURCE_ATTRIBUTES:
+        for pair in settings.OTEL_RESOURCE_ATTRIBUTES.split(","):
+            if "=" in pair:
+                key, value = pair.split("=", 1)
+                attributes[key.strip()] = value.strip()
+    return attributes
+
+
+def _create_sampler():
+    """Create appropriate sampler based on configuration."""
+    sampler_name = settings.OTEL_TRACES_SAMPLER.lower()
+
+    if sampler_name == "always_off":
+        return ALWAYS_OFF
+    elif sampler_name == "always_on":
+        return ALWAYS_ON
+    elif sampler_name == "parentbased_always_on":
+        return ParentBased(root=ALWAYS_ON)
+    elif sampler_name == "parentbased_always_off":
+        return ParentBased(root=ALWAYS_OFF)
+    elif sampler_name == "parentbased_traceidratio":
+        ratio = settings.OTEL_TRACES_SAMPLER_ARG or 1.0
+        return ParentBased(root=ParentBasedTraceIdRatio(ratio))
+    elif sampler_name == "traceidratio":
+        ratio = settings.OTEL_TRACES_SAMPLER_ARG or 1.0
+        return ParentBasedTraceIdRatio(ratio)
+    else:
+        logger.warning(
+            "unknown_sampler_type",
+            sampler=sampler_name,
+            fallback="parentbased_always_on"
+        )
+        return ParentBased(root=ALWAYS_ON)
+
+
+def setup_telemetry() -> None:
+    """
+    Initialize OpenTelemetry SDK with OTLP exporter and/or local storage.
+
+    This should be called once at application startup.
+    If OTEL_ENABLED is False, tracing will be disabled (no-op tracer).
+    Local storage can work independently without external OTLP endpoint.
+    """
+    global _tracer, _tracer_provider
+
+    if not settings.OTEL_ENABLED:
+        logger.info("otel_disabled", reason="OTEL_ENABLED=false")
+        # Set no-op tracer
+        trace.set_tracer_provider(trace.NoOpTracerProvider())
+        _tracer = trace.get_tracer(__name__)
+        return
+
+    # Check if we have any tracing destination (external OTLP or local storage)
+    has_otlp_endpoint = bool(settings.OTEL_EXPORTER_OTLP_ENDPOINT)
+    has_local_storage = getattr(settings, 'OTEL_LOCAL_STORAGE_ENABLED', True)
+
+    if not has_otlp_endpoint and not has_local_storage:
+        logger.warning(
+            "otel_no_destination_configured",
+            message="Neither OTLP endpoint nor local storage enabled, tracing disabled"
+        )
+        # Set no-op tracer
+        trace.set_tracer_provider(trace.NoOpTracerProvider())
+        _tracer = trace.get_tracer(__name__)
+        return
+
+    try:
+        # Create resource with service information
+        resource_attributes = {
+            SERVICE_NAME: settings.OTEL_SERVICE_NAME,
+            SERVICE_VERSION: settings.api_version,
+            "deployment.environment": settings.environment,
+        }
+
+        # Add custom resource attributes from config
+        resource_attributes.update(_parse_resource_attributes())
+
+        resource = Resource.create(resource_attributes)
+
+        # Create sampler
+        sampler = _create_sampler()
+
+        # Create tracer provider
+        _tracer_provider = TracerProvider(resource=resource, sampler=sampler)
+
+        # Only add OTLP exporter if endpoint is configured
+        if has_otlp_endpoint:
+            # Create appropriate exporter based on protocol
+            if settings.OTEL_EXPORTER_OTLP_PROTOCOL == "http":
+                # HTTP/protobuf exporter
+                endpoint = settings.OTEL_EXPORTER_OTLP_ENDPOINT
+                if not endpoint.endswith("/v1/traces"):
+                    endpoint = f"{endpoint}/v1/traces"
+                exporter = HTTPExporter(endpoint=endpoint)
+                logger.info(
+                    "otel_exporter_configured",
+                    protocol="http",
+                    endpoint=endpoint
+                )
+            else:
+                # gRPC exporter (default)
+                exporter = GRPCExporter(
+                    endpoint=settings.OTEL_EXPORTER_OTLP_ENDPOINT,
+                    insecure=True  # TODO: Configure TLS for production
+                )
+                logger.info(
+                    "otel_exporter_configured",
+                    protocol="grpc",
+                    endpoint=settings.OTEL_EXPORTER_OTLP_ENDPOINT
+                )
+
+            # Add batch span processor with exporter
+            span_processor = BatchSpanProcessor(exporter)
+            _tracer_provider.add_span_processor(span_processor)
+        else:
+            logger.info(
+                "otel_exporter_skipped",
+                reason="No OTLP endpoint configured, using local storage only"
+            )
+
+        # Add local storage processor for observability UI
+        global _local_storage_processor
+        if has_local_storage:
+            try:
+                from control_plane_api.app.observability.local_span_processor import (
+                    setup_local_storage_processor,
+                )
+                _local_storage_processor = setup_local_storage_processor(
+                    _tracer_provider,
+                )
+            except Exception as e:
+                logger.warning(
+                    "local_storage_processor_setup_failed",
+                    error=str(e),
+                    message="Continuing without local trace storage"
+                )
+
+        # Set global tracer provider
+        trace.set_tracer_provider(_tracer_provider)
+
+        # Create tracer instance
+        _tracer = trace.get_tracer(
+            instrumenting_module_name="control_plane_api",
+            instrumenting_library_version=settings.api_version
+        )
+
+        logger.info(
+            "otel_initialized",
+            service_name=settings.OTEL_SERVICE_NAME,
+            otlp_endpoint=settings.OTEL_EXPORTER_OTLP_ENDPOINT or "none",
+            otlp_protocol=settings.OTEL_EXPORTER_OTLP_PROTOCOL if has_otlp_endpoint else "disabled",
+            local_storage=bool(_local_storage_processor),
+            sampler=settings.OTEL_TRACES_SAMPLER,
+            sampler_arg=settings.OTEL_TRACES_SAMPLER_ARG,
+        )
+
+    except Exception as e:
+        logger.error(
+            "otel_initialization_failed",
+            error=str(e),
+            exc_info=True
+        )
+        # Fall back to no-op tracer on error
+        trace.set_tracer_provider(trace.NoOpTracerProvider())
+        _tracer = trace.get_tracer(__name__)
+
+
+def get_tracer() -> trace.Tracer:
+    """Get the global tracer instance."""
+    global _tracer
+    if _tracer is None:
+        # Return default tracer if not initialized
+        return trace.get_tracer(__name__)
+    return _tracer
+
+
+@contextmanager
+def create_span_with_context(
+    name: str,
+    organization_id: Optional[str] = None,
+    user_id: Optional[str] = None,
+    user_email: Optional[str] = None,
+    user_name: Optional[str] = None,
+    user_avatar: Optional[str] = None,
+    execution_id: Optional[str] = None,
+    request_id: Optional[str] = None,
+    **attributes: Any,
+):
+    """
+    Create a span with standard organizational and user context.
+
+    Args:
+        name: Span name (e.g., "streaming.phase_1_connect")
+        organization_id: Organization slug (e.g., "kubiya-ai")
+        user_id: User UUID
+        user_email: User email address
+        user_name: User display name
+        user_avatar: User avatar URL
+        execution_id: Execution UUID (if applicable)
+        request_id: Request ID from middleware
+        **attributes: Additional span attributes
+
+    Usage:
+        with create_span_with_context(
+            "worker.registration",
+            organization_id="kubiya-ai",
+            user_id="user-123",
+            user_email="user@example.com",
+            user_name="John Doe",
+            worker_queue_id="queue-456"
+        ) as span:
+            # Do work
+            span.set_attribute("worker.id", worker_id)
+    """
+    tracer = get_tracer()
+
+    with tracer.start_as_current_span(name) as span:
+        # Add standard organizational context
+        if organization_id:
+            span.set_attribute("organization.id", organization_id)
+
+        if user_id:
+            span.set_attribute("user.id", user_id)
+
+        if user_email:
+            span.set_attribute("user.email", user_email)
+
+        if user_name:
+            span.set_attribute("user.name", user_name)
+
+        if user_avatar:
+            span.set_attribute("user.avatar", user_avatar)
+
+        if execution_id:
+            span.set_attribute("execution.id", execution_id)
+
+        if request_id:
+            span.set_attribute("request.id", request_id)
+
+        # Add custom attributes
+        for key, value in attributes.items():
+            if value is not None:
+                # Convert value to string for non-primitive types
+                if isinstance(value, (str, int, float, bool)):
+                    span.set_attribute(key, value)
+                else:
+                    span.set_attribute(key, str(value))
+
+        yield span
+
+
+def get_current_trace_id() -> Optional[str]:
+    """
+    Get the current trace ID as a hex string.
+
+    Returns:
+        Trace ID in format: "4bf92f3577b34da6a3ce929d0e0e4736"
+        None if no active span
+    """
+    span = trace.get_current_span()
+    if span and span.is_recording():
+        trace_id = span.get_span_context().trace_id
+        return format(trace_id, '032x')
+    return None
+
+
+def get_current_span_id() -> Optional[str]:
+    """
+    Get the current span ID as a hex string.
+
+    Returns:
+        16-character hex string of the 64-bit span ID, or None if no active span
+    """
+    span = trace.get_current_span()
+    if span and span.is_recording():
+        span_id = span.get_span_context().span_id
+        return format(span_id, '016x')
+    return None
+
+
+def add_span_event(name: str, attributes: Optional[Dict[str, Any]] = None) -> None:
+    """
+    Add an event (log) to the current span for better trace visibility.
+
+    This makes logs visible in the Jaeger UI as span events.
+
+    Args:
+        name: Event name (e.g., "Database query started", "LLM response received")
+        attributes: Optional key-value pairs with event details
+
+    Example:
+        add_span_event("Agent execution started", {"agent_id": "abc123", "prompt": "Deploy app"})
+        # ... do work ...
+        add_span_event("Agent execution completed", {"status": "success", "duration_ms": 1234})
+    """
+    span = trace.get_current_span()
+    if span and span.is_recording():
+        event_attrs = attributes or {}
+        span.add_event(name, event_attrs)
+
+
+def add_span_error(error: Exception, attributes: Optional[Dict[str, Any]] = None) -> None:
+    """
+    Record an error in the current span with full context.
+
+    Args:
+        error: The exception that occurred
+        attributes: Optional additional context about the error
+
+    Example:
+        try:
+            risky_operation()
+        except Exception as e:
+            add_span_error(e, {"operation": "database_write", "table": "users"})
+            raise
+    """
+    span = trace.get_current_span()
+    if span and span.is_recording():
+        from opentelemetry.trace import StatusCode
+
+        span.set_status(StatusCode.ERROR, str(error))
+        span.record_exception(error)
+
+        if attributes:
+            for key, value in attributes.items():
+                span.set_attribute(f"error.{key}", str(value))
+
+
+def shutdown_telemetry() -> None:
+    """
+    Shutdown the OpenTelemetry SDK and flush remaining spans.
+
+    This should be called during application shutdown.
+    """
+    global _tracer_provider
+
+    if _tracer_provider:
+        try:
+            _tracer_provider.shutdown()
+            logger.info("otel_shutdown_complete")
+        except Exception as e:
+            logger.error("otel_shutdown_failed", error=str(e), exc_info=True)

control_plane_api/app/policies/README.md

@@ -0,0 +1,149 @@
+# Policy Modules for Agent Control Plane
+
+This directory contains pre-built OPA (Open Policy Agent) policy modules written in Rego for common agent control plane operations.
+
+## Overview
+
+Policies are stored and evaluated in the OPA Watchdog Enforcer Service. These modules provide templates and examples for:
+
+- **Execution Policies**: Control when and how agents can execute
+- **Resource Policies**: Limit resource usage and access
+- **Time-based Policies**: Business hours and rate limiting
+- **Tool Enforcement Policies**: Real-time validation of tool executions (NEW)
+- **Tool Policies**: Control which tools agents can use
+- **Team Policies**: Collaboration and approval workflows
+
+### Tool Enforcement Policies (NEW)
+
+The `tool_enforcement/` directory contains policies for **real-time tool call enforcement** during agent execution. These policies are evaluated before each tool execution and can block or allow tools based on:
+
+- User roles and permissions
+- Tool risk level (critical/high/medium/low)
+- Environment (production/staging/dev)
+- Time of day (business hours)
+- Tool source (MCP/builtin/skill)
+- Command patterns (dangerous bash commands)
+
+See the [Tool Enforcement Guide](../../../docs/TOOL_ENFORCEMENT.md) for detailed documentation.
+
+## Policy Structure
+
+All policies follow this structure:
+
+```rego
+package <policy_name>
+
+# Default decision (deny by default)
+default allow = false
+
+# Rules that grant permissions
+allow {
+    # conditions
+}
+
+# Violations that explain why requests are denied
+violations[msg] {
+    # condition
+    msg := "Explanation"
+}
+```
+
+## Policy Priority and Inheritance
+
+Policies are associated with entities (agents, teams, environments) with automatic inheritance:
+
+1. **Environment** policies (priority: 300) - apply to all agents/teams in the environment
+2. **Team** policies (priority: 200) - apply to all agents in the team
+3. **Agent** policies (priority: 100) - apply to specific agent
+
+When the same policy is defined at multiple levels, **higher priority wins**.
+
+## Usage
+
+### 1. Create a Policy
+
+```bash
+curl -X POST https://your-control-plane.com/api/v1/policies \
+  -H "Authorization: Bearer $TOKEN" \
+  -H "Content-Type: application/json" \
+  -d '{
+    "name": "business-hours-only",
+    "policy_content": "<rego_policy>",
+    "description": "Allow executions only during business hours",
+    "enabled": true,
+    "tags": ["time", "compliance"]
+  }'
+```
+
+### 2. Associate with Entity
+
+```bash
+curl -X POST https://your-control-plane.com/api/v1/policies/associations \
+  -H "Authorization: Bearer $TOKEN" \
+  -H "Content-Type: application/json" \
+  -d '{
+    "policy_id": "<policy_uuid>",
+    "policy_name": "business-hours-only",
+    "entity_type": "environment",
+    "entity_id": "<environment_uuid>",
+    "enabled": true
+  }'
+```
+
+### 3. Evaluate Policy
+
+Policies are automatically evaluated during agent execution. You can also test manually:
+
+```bash
+curl -X POST https://your-control-plane.com/api/v1/policies/evaluate/agent/<agent_id> \
+  -H "Authorization: Bearer $TOKEN" \
+  -H "Content-Type: application/json" \
+  -d '{
+    "input_data": {
+      "action": "execute",
+      "user": "alice@company.com",
+      "time": "2025-01-15T14:30:00Z"
+    }
+  }'
+```
+
+## Available Policy Modules
+
+### Tool Enforcement Policies (`tool_enforcement/`)
+
+Real-time tool call validation policies:
+
+- **`role_based_tool_access.rego`** - Restrict command execution tools to admin/devops roles
+- **`production_safeguards.rego`** - Block critical/high-risk tools in production environment
+- **`bash_command_validation.rego`** - Block dangerous bash patterns (rm -rf, dd, mkfs, etc.)
+- **`business_hours_enforcement.rego`** - Restrict high-risk tools to business hours (9 AM - 6 PM UTC, Mon-Fri)
+- **`mcp_tool_allowlist.rego`** - Whitelist approved MCP tools
+
+All tool enforcement policies use package `kubiya.tool_enforcement` and evaluate against structured tool execution input. See [Tool Enforcement Guide](../../../docs/TOOL_ENFORCEMENT.md) for input schema and examples.
+
+### General Policy Modules
+
+See individual `.rego` files for detailed policies:
+
+- `business_hours.rego` - Restrict executions to business hours
+- `approved_users.rego` - Require user approval list
+- `rate_limiting.rego` - Limit execution frequency
+- `resource_limits.rego` - Enforce resource quotas
+- `tool_restrictions.rego` - Control tool usage
+- `approval_workflow.rego` - Require approvals for sensitive actions
+- `environment_restrictions.rego` - Restrict access to environments
+
+## Policy Development Tips
+
+1. **Test policies locally** using OPA CLI before deploying
+2. **Start with deny-by-default** for security
+3. **Use clear violation messages** for debugging
+4. **Document assumptions** in policy comments
+5. **Version policies** using tags or metadata
+6. **Test edge cases** thoroughly
+
+## References
+
+- [OPA Documentation](https://www.openpolicyagent.org/docs/latest/)
+- [Rego Language Reference](https://www.openpolicyagent.org/docs/latest/policy-reference/)
+- [OPA Watchdog Enforcer API](https://enforcer-psi.vercel.app/api/docs)

control_plane_api/app/policies/approved_users.rego

@@ -0,0 +1,62 @@
+# Approved Users Policy
+# Only allows executions from a whitelist of approved users
+
+package approved_users
+
+import future.keywords.if
+import future.keywords.in
+
+# Default deny
+default allow = false
+
+# List of approved users (emails)
+# This can be customized per organization or pulled from metadata
+approved_user_list := {
+    "admin@company.com",
+    "devops@company.com",
+    "alice@company.com",
+    "bob@company.com"
+}
+
+# Allow if user is in the approved list
+allow if {
+    input.user in approved_user_list
+}
+
+# Also allow if user email domain is approved
+allow if {
+    user_email := input.user
+    contains(user_email, "@")
+    domain := split(user_email, "@")[1]
+    domain in approved_domains
+}
+
+# Approved email domains
+approved_domains := {
+    "company.com",
+    "trusted-partner.com"
+}
+
+# Violations
+violations[msg] if {
+    not input.user in approved_user_list
+    user_email := input.user
+    contains(user_email, "@")
+    domain := split(user_email, "@")[1]
+    not domain in approved_domains
+    msg := sprintf("User '%v' is not in the approved user list or approved domain", [input.user])
+}
+
+violations[msg] if {
+    not input.user
+    msg := "No user specified in the request"
+}
+
+# Metadata
+metadata := {
+    "name": "approved-users-only",
+    "description": "Only allow executions from approved users or domains",
+    "version": "1.0.0",
+    "author": "Kubiya",
+    "tags": ["security", "access-control", "users"]
+}

control_plane_api/app/policies/business_hours.rego

@@ -0,0 +1,51 @@
+# Business Hours Policy
+# Restricts agent executions to business hours (9 AM - 5 PM, Monday-Friday)
+
+package business_hours
+
+import future.keywords.if
+import future.keywords.in
+
+# Default deny
+default allow = false
+
+# Allow executions during business hours
+allow if {
+    is_business_hours
+    not is_weekend
+}
+
+# Check if current time is within business hours (9 AM - 5 PM)
+is_business_hours if {
+    time_hour := time.parse_rfc3339_ns(input.time)
+    hour := time.clock(time_hour)[0]
+    hour >= 9
+    hour < 17
+}
+
+# Check if current day is a weekend
+is_weekend if {
+    time_day := time.parse_rfc3339_ns(input.time)
+    weekday := time.weekday(time_day)
+    weekday in ["Saturday", "Sunday"]
+}
+
+# Violations
+violations[msg] if {
+    not is_business_hours
+    msg := sprintf("Execution requested outside business hours (9 AM - 5 PM). Current time: %v", [input.time])
+}
+
+violations[msg] if {
+    is_weekend
+    msg := sprintf("Execution requested on weekend. Current day: %v", [time.weekday(time.parse_rfc3339_ns(input.time))])
+}
+
+# Metadata
+metadata := {
+    "name": "business-hours-only",
+    "description": "Restrict agent executions to business hours (9 AM - 5 PM, Monday-Friday)",
+    "version": "1.0.0",
+    "author": "Kubiya",
+    "tags": ["time", "compliance", "business-hours"]
+}