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

control_plane_api/app/services/trace_retention.py

@@ -0,0 +1,354 @@
+"""
+Trace Retention Service.
+
+Handles automatic cleanup of old traces based on configurable retention period.
+Also provides storage statistics per organization.
+
+Features:
+- Configurable retention period (default: 30 days)
+- Per-organization storage stats
+- Batch deletion for performance
+- Scheduled cleanup job support
+"""
+
+import structlog
+from datetime import datetime, timezone, timedelta
+from typing import Optional, Dict, Any
+
+from sqlalchemy.orm import Session
+from sqlalchemy import func, delete
+
+from control_plane_api.app.config import settings
+from control_plane_api.app.database import get_session_local
+from control_plane_api.app.models.trace import Trace, Span
+
+logger = structlog.get_logger()
+
+# Configuration
+RETENTION_DAYS = getattr(settings, 'OTEL_LOCAL_STORAGE_RETENTION_DAYS', 30)
+BATCH_SIZE = 1000  # Number of traces to delete per batch
+
+
+class TraceRetentionService:
+    """Service for managing trace retention and cleanup."""
+
+    def __init__(self, retention_days: int = None):
+        """
+        Initialize the retention service.
+
+        Args:
+            retention_days: Number of days to retain traces (default from config)
+        """
+        self.retention_days = retention_days or RETENTION_DAYS
+        self._stats = {
+            "last_cleanup": None,
+            "traces_deleted": 0,
+            "spans_deleted": 0,
+            "errors": 0,
+        }
+
+    def get_cutoff_date(self) -> datetime:
+        """Get the cutoff date for trace retention."""
+        return datetime.now(timezone.utc) - timedelta(days=self.retention_days)
+
+    async def cleanup_old_traces(
+        self,
+        organization_id: Optional[str] = None,
+        batch_size: int = BATCH_SIZE,
+    ) -> Dict[str, int]:
+        """
+        Delete traces older than the retention period.
+
+        Args:
+            organization_id: Optional org to limit cleanup to
+            batch_size: Number of traces to delete per batch
+
+        Returns:
+            Dict with deletion statistics
+        """
+        cutoff_date = self.get_cutoff_date()
+        total_traces_deleted = 0
+        total_spans_deleted = 0
+
+        logger.info(
+            "trace_cleanup_starting",
+            retention_days=self.retention_days,
+            cutoff_date=cutoff_date.isoformat(),
+            organization_id=organization_id,
+        )
+
+        try:
+            SessionLocal = get_session_local()
+            session = SessionLocal()
+
+            try:
+                while True:
+                    # Find old traces to delete
+                    query = session.query(Trace).filter(
+                        Trace.started_at < cutoff_date
+                    )
+
+                    if organization_id:
+                        query = query.filter(Trace.organization_id == organization_id)
+
+                    # Get batch of trace IDs
+                    traces_to_delete = query.limit(batch_size).all()
+
+                    if not traces_to_delete:
+                        break
+
+                    trace_ids = [t.trace_id for t in traces_to_delete]
+
+                    # Count spans being deleted (for stats)
+                    span_count = session.query(func.count(Span.id)).filter(
+                        Span.trace_id.in_(trace_ids)
+                    ).scalar()
+
+                    # Delete traces (cascade will delete spans due to FK)
+                    for trace in traces_to_delete:
+                        session.delete(trace)
+
+                    session.commit()
+
+                    total_traces_deleted += len(trace_ids)
+                    total_spans_deleted += span_count
+
+                    logger.info(
+                        "trace_cleanup_batch_completed",
+                        traces_deleted=len(trace_ids),
+                        spans_deleted=span_count,
+                        total_traces_deleted=total_traces_deleted,
+                    )
+
+            except Exception as e:
+                session.rollback()
+                self._stats["errors"] += 1
+                logger.error("trace_cleanup_batch_failed", error=str(e), exc_info=True)
+                raise
+
+            finally:
+                session.close()
+
+        except Exception as e:
+            self._stats["errors"] += 1
+            logger.error("trace_cleanup_failed", error=str(e), exc_info=True)
+            raise
+
+        # Update stats
+        self._stats["last_cleanup"] = datetime.now(timezone.utc).isoformat()
+        self._stats["traces_deleted"] += total_traces_deleted
+        self._stats["spans_deleted"] += total_spans_deleted
+
+        logger.info(
+            "trace_cleanup_completed",
+            total_traces_deleted=total_traces_deleted,
+            total_spans_deleted=total_spans_deleted,
+            retention_days=self.retention_days,
+        )
+
+        return {
+            "traces_deleted": total_traces_deleted,
+            "spans_deleted": total_spans_deleted,
+            "cutoff_date": cutoff_date.isoformat(),
+        }
+
+    async def get_storage_stats(
+        self,
+        organization_id: Optional[str] = None,
+    ) -> Dict[str, Any]:
+        """
+        Get storage statistics for traces.
+
+        Args:
+            organization_id: Optional org to get stats for
+
+        Returns:
+            Dict with storage statistics
+        """
+        try:
+            SessionLocal = get_session_local()
+            session = SessionLocal()
+
+            try:
+                # Base queries
+                trace_query = session.query(Trace)
+                span_query = session.query(Span)
+
+                if organization_id:
+                    trace_query = trace_query.filter(Trace.organization_id == organization_id)
+                    span_query = span_query.filter(Span.organization_id == organization_id)
+
+                # Total counts
+                total_traces = trace_query.count()
+                total_spans = span_query.count()
+
+                # Counts by status
+                success_count = trace_query.filter(Trace.status == "success").count()
+                error_count = trace_query.filter(Trace.status == "error").count()
+                running_count = trace_query.filter(Trace.status == "running").count()
+
+                # Date range
+                oldest_trace = trace_query.order_by(Trace.started_at.asc()).first()
+                newest_trace = trace_query.order_by(Trace.started_at.desc()).first()
+
+                # Average metrics
+                avg_duration = session.query(func.avg(Trace.duration_ms)).filter(
+                    Trace.duration_ms.isnot(None)
+                )
+                avg_span_count = session.query(func.avg(Trace.span_count))
+
+                if organization_id:
+                    avg_duration = avg_duration.filter(Trace.organization_id == organization_id)
+                    avg_span_count = avg_span_count.filter(Trace.organization_id == organization_id)
+
+                avg_duration_val = avg_duration.scalar()
+                avg_span_count_val = avg_span_count.scalar()
+
+                # Retention info
+                cutoff_date = self.get_cutoff_date()
+                traces_to_expire = trace_query.filter(Trace.started_at < cutoff_date).count()
+
+                return {
+                    "total_traces": total_traces,
+                    "total_spans": total_spans,
+                    "status_breakdown": {
+                        "success": success_count,
+                        "error": error_count,
+                        "running": running_count,
+                    },
+                    "date_range": {
+                        "oldest": oldest_trace.started_at.isoformat() if oldest_trace else None,
+                        "newest": newest_trace.started_at.isoformat() if newest_trace else None,
+                    },
+                    "averages": {
+                        "duration_ms": round(avg_duration_val, 2) if avg_duration_val else None,
+                        "span_count": round(avg_span_count_val, 2) if avg_span_count_val else None,
+                    },
+                    "retention": {
+                        "retention_days": self.retention_days,
+                        "cutoff_date": cutoff_date.isoformat(),
+                        "traces_to_expire": traces_to_expire,
+                    },
+                    "organization_id": organization_id,
+                    "retrieved_at": datetime.now(timezone.utc).isoformat(),
+                }
+
+            finally:
+                session.close()
+
+        except Exception as e:
+            logger.error("get_storage_stats_failed", error=str(e), exc_info=True)
+            raise
+
+    def get_service_stats(self) -> Dict[str, Any]:
+        """Get retention service statistics."""
+        return {
+            **self._stats,
+            "retention_days": self.retention_days,
+        }
+
+
+# Singleton instance
+_retention_service: Optional[TraceRetentionService] = None
+
+
+def get_retention_service() -> TraceRetentionService:
+    """Get the singleton retention service instance."""
+    global _retention_service
+    if _retention_service is None:
+        _retention_service = TraceRetentionService()
+    return _retention_service
+
+
+async def run_cleanup_job():
+    """
+    Run the trace cleanup job.
+
+    This function can be called from a scheduled job (e.g., Temporal workflow,
+    cron job, or APScheduler).
+    """
+    service = get_retention_service()
+    result = await service.cleanup_old_traces()
+
+    logger.info(
+        "scheduled_trace_cleanup_completed",
+        **result
+    )
+
+    return result
+
+
+# Scheduler configuration
+CLEANUP_INTERVAL_HOURS = 6  # Run cleanup every 6 hours
+
+
+async def _retention_scheduler_loop():
+    """
+    Background loop that runs retention cleanup periodically.
+
+    Runs every CLEANUP_INTERVAL_HOURS to clean up old traces.
+    Uses asyncio.sleep for non-blocking scheduling.
+    """
+    import asyncio
+
+    logger.info(
+        "retention_scheduler_started",
+        interval_hours=CLEANUP_INTERVAL_HOURS,
+        retention_days=RETENTION_DAYS,
+    )
+
+    while True:
+        try:
+            # Wait for the interval
+            await asyncio.sleep(CLEANUP_INTERVAL_HOURS * 3600)
+
+            # Run cleanup
+            logger.info("retention_cleanup_starting")
+            result = await run_cleanup_job()
+            logger.info(
+                "retention_cleanup_completed",
+                traces_deleted=result.get("traces_deleted", 0),
+                spans_deleted=result.get("spans_deleted", 0),
+            )
+
+        except asyncio.CancelledError:
+            logger.info("retention_scheduler_cancelled")
+            break
+        except Exception as e:
+            logger.error(
+                "retention_scheduler_error",
+                error=str(e),
+                exc_info=True,
+            )
+            # Continue running despite errors
+            await asyncio.sleep(60)  # Brief pause before retrying
+
+
+async def start_retention_scheduler():
+    """
+    Start the retention cleanup scheduler as a background task.
+
+    Returns:
+        asyncio.Task: The scheduler task (for cancellation on shutdown)
+    """
+    import asyncio
+
+    task = asyncio.create_task(_retention_scheduler_loop())
+    return task
+
+
+async def stop_retention_scheduler(task):
+    """
+    Stop the retention cleanup scheduler.
+
+    Args:
+        task: The scheduler task returned by start_retention_scheduler()
+    """
+    if task and not task.done():
+        task.cancel()
+        try:
+            await task
+        except Exception:
+            pass  # Task cancelled, ignore
+
+    logger.info("retention_scheduler_stopped")

control_plane_api/app/services/worker_queue_metrics_service.py

@@ -0,0 +1,190 @@
+"""
+Worker Queue Metrics Service.
+
+This service provides business logic for calculating worker queue metrics
+including worker health, task statistics, and performance metrics.
+"""
+
+import structlog
+from datetime import datetime, timedelta
+from typing import Optional, Dict
+from sqlalchemy.orm import Session
+from sqlalchemy import func
+
+from control_plane_api.app.models.worker import WorkerQueue, WorkerHeartbeat
+from control_plane_api.app.models.execution import Execution
+from control_plane_api.app.schemas.worker_queue_observability_schemas import WorkerQueueMetricsResponse
+
+logger = structlog.get_logger()
+
+
+class WorkerQueueMetricsService:
+    """Service for calculating worker queue metrics"""
+
+    def __init__(self, db: Session):
+        self.db = db
+
+    async def get_queue_metrics(
+        self,
+        queue_id: str,
+        organization_id: str
+    ) -> WorkerQueueMetricsResponse:
+        """
+        Calculate comprehensive metrics for a worker queue.
+
+        Args:
+            queue_id: Worker queue UUID
+            organization_id: Organization ID
+
+        Returns:
+            WorkerQueueMetricsResponse with calculated metrics
+
+        Raises:
+            ValueError: If queue not found or doesn't belong to organization
+        """
+        # Verify queue exists and belongs to organization
+        queue = self.db.query(WorkerQueue).filter(
+            WorkerQueue.id == queue_id,
+            WorkerQueue.organization_id == organization_id
+        ).first()
+
+        if not queue:
+            raise ValueError("Worker queue not found")
+
+        now = datetime.utcnow()
+
+        # Calculate worker status counts
+        worker_stats = self._get_worker_status_counts(queue_id, now)
+
+        # Calculate 24h task metrics
+        task_metrics = self._get_task_metrics_24h(queue_id, now)
+
+        # Get last activity timestamp
+        last_activity = self._get_last_activity(queue_id)
+
+        # Build response
+        return WorkerQueueMetricsResponse(
+            queue_id=queue_id,
+            active_workers=worker_stats["active"],
+            idle_workers=worker_stats["idle"],
+            busy_workers=worker_stats["busy"],
+            total_workers=worker_stats["total"],
+            tasks_processed_24h=task_metrics["processed"],
+            tasks_failed_24h=task_metrics["failed"],
+            tasks_pending=task_metrics["pending"],
+            avg_task_duration_ms=task_metrics["avg_duration_ms"],
+            error_rate_percent=task_metrics["error_rate"],
+            last_error_at=task_metrics["last_error_at"],
+            task_queue_backlog=0,  # TODO: Implement Temporal queue metrics
+            task_queue_pollers=0,  # TODO: Implement Temporal queue metrics
+            last_activity_at=last_activity,
+            updated_at=now
+        )
+
+    def _get_worker_status_counts(self, queue_id: str, now: datetime) -> Dict[str, int]:
+        """
+        Get worker status counts from WorkerHeartbeat table.
+
+        Workers are considered stale if last_heartbeat > 90 seconds ago.
+        """
+        stale_threshold = now - timedelta(seconds=90)
+
+        # Query recent heartbeats
+        heartbeats = self.db.query(WorkerHeartbeat).filter(
+            WorkerHeartbeat.worker_queue_id == queue_id,
+            WorkerHeartbeat.last_heartbeat > stale_threshold
+        ).all()
+
+        active = sum(1 for hb in heartbeats if hb.status == "active")
+        idle = sum(1 for hb in heartbeats if hb.status == "idle")
+        busy = sum(1 for hb in heartbeats if hb.status == "busy")
+        total = len(heartbeats)
+
+        logger.info(
+            "worker_status_calculated",
+            queue_id=queue_id,
+            active=active,
+            idle=idle,
+            busy=busy,
+            total=total
+        )
+
+        return {
+            "active": active,
+            "idle": idle,
+            "busy": busy,
+            "total": total
+        }
+
+    def _get_task_metrics_24h(self, queue_id: str, now: datetime) -> Dict:
+        """Calculate task metrics for the last 24 hours"""
+        twenty_four_hours_ago = now - timedelta(hours=24)
+
+        # Get executions in last 24h
+        executions_24h = self.db.query(Execution).filter(
+            Execution.worker_queue_id == queue_id,
+            Execution.created_at >= twenty_four_hours_ago
+        ).all()
+
+        # Count processed and failed tasks
+        processed = sum(1 for e in executions_24h if e.status in ["completed", "failed"])
+        failed = sum(1 for e in executions_24h if e.status == "failed")
+
+        # Get pending tasks count
+        pending = self.db.query(Execution).filter(
+            Execution.worker_queue_id == queue_id,
+            Execution.status == "pending"
+        ).count()
+
+        # Calculate average duration for completed tasks
+        completed_executions = [
+            e for e in executions_24h
+            if e.status == "completed" and e.started_at and e.completed_at
+        ]
+
+        if completed_executions:
+            total_duration_ms = sum(
+                (e.completed_at - e.started_at).total_seconds() * 1000
+                for e in completed_executions
+            )
+            avg_duration_ms = total_duration_ms / len(completed_executions)
+        else:
+            avg_duration_ms = 0
+
+        # Calculate error rate
+        error_rate = (failed / processed * 100) if processed > 0 else 0
+
+        # Get last error timestamp
+        last_error = self.db.query(Execution).filter(
+            Execution.worker_queue_id == queue_id,
+            Execution.status == "failed"
+        ).order_by(Execution.completed_at.desc()).first()
+
+        last_error_at = last_error.completed_at if last_error else None
+
+        logger.info(
+            "task_metrics_calculated",
+            queue_id=queue_id,
+            processed=processed,
+            failed=failed,
+            pending=pending,
+            avg_duration_ms=avg_duration_ms,
+            error_rate=error_rate
+        )
+
+        return {
+            "processed": processed,
+            "failed": failed,
+            "pending": pending,
+            "avg_duration_ms": avg_duration_ms,
+            "error_rate": error_rate,
+            "last_error_at": last_error_at
+        }
+
+    def _get_last_activity(self, queue_id: str) -> Optional[datetime]:
+        """Get timestamp of last worker activity"""
+        last_activity = self.db.query(WorkerHeartbeat).filter(
+            WorkerHeartbeat.worker_queue_id == queue_id
+        ).order_by(WorkerHeartbeat.last_heartbeat.desc()).first()
+
+        return last_activity.last_heartbeat if last_activity else None

control_plane_api/app/services/workflow_cancellation_manager.py

@@ -0,0 +1,135 @@
+"""
+Workflow Cancellation Manager - handles workflow-specific cancellation without affecting the agent execution.
+
+This allows cancelling individual workflow tool calls while the agent continues to run.
+"""
+
+from typing import Dict, Set, Any, Optional
+from datetime import datetime, timezone
+import structlog
+import threading
+
+logger = structlog.get_logger()
+
+
+class WorkflowCancellationManager:
+    """
+    Manages cancellation flags for active workflow executions.
+
+    Workflow tool calls check these flags periodically to see if they should stop,
+    without affecting the parent agent/team execution.
+    """
+
+    def __init__(self):
+        # Key: workflow_execution_key (execution_id + workflow_message_id), Value: cancellation time
+        self._cancelled_workflows: Dict[str, str] = {}
+        # Key: workflow_execution_key, Value: threading.Event for immediate cancellation
+        self._cancellation_events: Dict[str, threading.Event] = {}
+        self._lock = threading.Lock()
+
+    def _make_key(self, execution_id: str, workflow_message_id: str) -> str:
+        """Create a unique key for a workflow execution."""
+        return f"{execution_id}:{workflow_message_id}"
+
+    def register_workflow(self, execution_id: str, workflow_message_id: str) -> threading.Event:
+        """
+        Register a workflow execution and get a cancellation event.
+
+        Args:
+            execution_id: The agent execution ID
+            workflow_message_id: The unique workflow message ID
+
+        Returns:
+            threading.Event that will be set when cancellation is requested
+        """
+        with self._lock:
+            key = self._make_key(execution_id, workflow_message_id)
+            event = threading.Event()
+            self._cancellation_events[key] = event
+
+            logger.info(
+                "workflow_registered",
+                execution_id=execution_id[:8],
+                workflow_message_id=workflow_message_id[-12:],
+                key=key
+            )
+
+            return event
+
+    def request_cancellation(self, execution_id: str, workflow_message_id: str) -> bool:
+        """
+        Request cancellation of a specific workflow.
+
+        Args:
+            execution_id: The agent execution ID
+            workflow_message_id: The unique workflow message ID
+
+        Returns:
+            True if cancellation was requested, False if workflow not found
+        """
+        with self._lock:
+            key = self._make_key(execution_id, workflow_message_id)
+            self._cancelled_workflows[key] = datetime.now(timezone.utc).isoformat()
+
+            # Signal the cancellation event immediately
+            if key in self._cancellation_events:
+                self._cancellation_events[key].set()
+                logger.info(
+                    "workflow_cancellation_event_signaled",
+                    execution_id=execution_id[:8],
+                    workflow_message_id=workflow_message_id[-12:],
+                    key=key
+                )
+
+            logger.info(
+                "workflow_cancellation_requested",
+                execution_id=execution_id[:8],
+                workflow_message_id=workflow_message_id[-12:],
+                key=key
+            )
+
+            return True
+
+    def is_cancelled(self, execution_id: str, workflow_message_id: str) -> bool:
+        """
+        Check if a workflow has been cancelled.
+
+        Args:
+            execution_id: The agent execution ID
+            workflow_message_id: The unique workflow message ID
+
+        Returns:
+            True if the workflow has been cancelled
+        """
+        with self._lock:
+            key = self._make_key(execution_id, workflow_message_id)
+            return key in self._cancelled_workflows
+
+    def clear_cancellation(self, execution_id: str, workflow_message_id: str) -> None:
+        """
+        Clear the cancellation flag for a workflow (called when workflow completes/fails).
+
+        Args:
+            execution_id: The agent execution ID
+            workflow_message_id: The unique workflow message ID
+        """
+        with self._lock:
+            key = self._make_key(execution_id, workflow_message_id)
+            if key in self._cancelled_workflows:
+                del self._cancelled_workflows[key]
+            if key in self._cancellation_events:
+                del self._cancellation_events[key]
+            logger.info(
+                "workflow_cancellation_cleared",
+                execution_id=execution_id[:8],
+                workflow_message_id=workflow_message_id[-12:]
+            )
+
+    def get_active_count(self) -> int:
+        """Get number of workflows with pending cancellation."""
+        with self._lock:
+            return len(self._cancelled_workflows)
+
+
+# Global singleton instance
+workflow_cancellation_manager = WorkflowCancellationManager()