kubiya-control-plane-api 0.1.0__py3-none-any.whl → 0.3.4__py3-none-any.whl

control_plane_api/worker/services/approval_tools.py

@@ -0,0 +1,310 @@
+"""
+Approval workflow tools for human-in-the-loop approval gates.
+
+Provides tools for workflows to request approval from authorized users
+and wait for approval/rejection before continuing.
+"""
+import os
+import time
+import asyncio
+import httpx
+import structlog
+from typing import List, Optional, Dict, Any
+from datetime import datetime, timedelta
+
+logger = structlog.get_logger()
+
+
+class ApprovalTools:
+    """
+    Approval workflow tools for human-in-the-loop gates.
+
+    Provides a temporal-native way to wait for approval from authorized users.
+    The workflow pauses execution and polls the control plane for approval status.
+    """
+
+    def __init__(
+        self,
+        control_plane_url: str,
+        api_key: str,
+        execution_id: str,
+        organization_id: str,
+        config: Optional[Dict[str, Any]] = None,
+    ):
+        """
+        Initialize approval tools.
+
+        Args:
+            control_plane_url: Control plane API base URL
+            api_key: API key for authentication
+            execution_id: Current execution ID
+            organization_id: Organization ID
+            config: Optional configuration (timeout, require_reason, etc.)
+        """
+        self.control_plane_url = control_plane_url.rstrip("/")
+        self.api_key = api_key
+        self.execution_id = execution_id
+        self.organization_id = organization_id
+        self.config = config or {}
+
+        # Configuration
+        self.timeout_minutes = self.config.get("timeout_minutes", 1440)  # 24 hours default
+        self.require_approval_reason = self.config.get("require_approval_reason", False)
+        self.poll_interval_seconds = self.config.get("poll_interval_seconds", 5)  # Poll every 5 seconds
+
+        self.client = httpx.AsyncClient(
+            base_url=self.control_plane_url,
+            headers={
+                "Authorization": f"Bearer {self.api_key}",
+                "Content-Type": "application/json",
+            },
+            timeout=30.0,
+        )
+
+    async def wait_for_approval(
+        self,
+        title: str,
+        message: Optional[str] = None,
+        approver_user_ids: Optional[List[str]] = None,
+        approver_user_emails: Optional[List[str]] = None,
+        approver_group_id: Optional[str] = None,
+        context: Optional[Dict[str, Any]] = None,
+    ) -> Dict[str, Any]:
+        """
+        Wait for approval from authorized users.
+
+        This function creates an approval request and polls the control plane
+        until the request is approved, rejected, or times out.
+
+        Args:
+            title: Brief title for the approval request
+            message: Detailed message or reason for approval
+            approver_user_ids: List of user IDs who can approve (optional)
+            approver_user_emails: List of user emails who can approve (optional)
+            approver_group_id: Group ID that can approve (optional)
+            context: Additional context data (optional)
+
+        Returns:
+            Dict with approval result:
+            {
+                "approved": bool,
+                "status": "approved" | "rejected" | "expired",
+                "approval_id": str,
+                "approved_by_email": str (if approved),
+                "rejection_reason": str (if rejected),
+                "resolved_at": str (ISO timestamp)
+            }
+
+        Raises:
+            Exception: If approval request creation fails or times out
+        """
+        logger.info(
+            "wait_for_approval_started",
+            title=title,
+            execution_id=self.execution_id,
+            approver_emails=approver_user_emails,
+        )
+
+        # Validate at least one approver is specified
+        if not approver_user_ids and not approver_user_emails and not approver_group_id:
+            raise ValueError("At least one of approver_user_ids, approver_user_emails, or approver_group_id must be provided")
+
+        try:
+            # Create approval request via control plane API
+            approval_request = {
+                "execution_id": self.execution_id,
+                "title": title,
+                "message": message,
+                "approver_user_ids": approver_user_ids or [],
+                "approver_user_emails": approver_user_emails or [],
+                "approver_group_id": approver_group_id,
+                "timeout_minutes": self.timeout_minutes,
+                "context": context or {},
+            }
+
+            response = await self.client.post(
+                "/api/v1/approvals",
+                json=approval_request,
+            )
+
+            if response.status_code != 201:
+                error_detail = response.text
+                logger.error(
+                    "approval_request_creation_failed",
+                    status_code=response.status_code,
+                    error=error_detail
+                )
+                raise Exception(f"Failed to create approval request: {error_detail}")
+
+            approval_data = response.json()
+            approval_id = approval_data["id"]
+
+            logger.info(
+                "approval_request_created",
+                approval_id=approval_id,
+                title=title,
+                execution_id=self.execution_id,
+            )
+
+            # Calculate timeout
+            start_time = time.time()
+            timeout_seconds = self.timeout_minutes * 60
+            expires_at = time.time() + timeout_seconds
+
+            # Poll for approval status
+            poll_count = 0
+            while time.time() < expires_at:
+                poll_count += 1
+
+                # Get approval status
+                try:
+                    status_response = await self.client.get(
+                        f"/api/v1/approvals/{approval_id}"
+                    )
+
+                    if status_response.status_code == 200:
+                        approval_status = status_response.json()
+
+                        if approval_status["status"] == "approved":
+                            elapsed_minutes = (time.time() - start_time) / 60
+                            logger.info(
+                                "approval_granted",
+                                approval_id=approval_id,
+                                approved_by=approval_status.get("approved_by_email"),
+                                elapsed_minutes=round(elapsed_minutes, 2),
+                            )
+
+                            return {
+                                "approved": True,
+                                "status": "approved",
+                                "approval_id": approval_id,
+                                "approved_by_email": approval_status.get("approved_by_email"),
+                                "approved_by_name": approval_status.get("approved_by_name"),
+                                "resolved_at": approval_status.get("resolved_at"),
+                            }
+
+                        elif approval_status["status"] == "rejected":
+                            elapsed_minutes = (time.time() - start_time) / 60
+                            logger.info(
+                                "approval_rejected",
+                                approval_id=approval_id,
+                                rejected_by=approval_status.get("approved_by_email"),
+                                reason=approval_status.get("rejection_reason"),
+                                elapsed_minutes=round(elapsed_minutes, 2),
+                            )
+
+                            return {
+                                "approved": False,
+                                "status": "rejected",
+                                "approval_id": approval_id,
+                                "rejected_by_email": approval_status.get("approved_by_email"),
+                                "rejected_by_name": approval_status.get("approved_by_name"),
+                                "rejection_reason": approval_status.get("rejection_reason"),
+                                "resolved_at": approval_status.get("resolved_at"),
+                            }
+
+                        elif approval_status["status"] == "expired":
+                            logger.warning(
+                                "approval_expired",
+                                approval_id=approval_id,
+                            )
+
+                            return {
+                                "approved": False,
+                                "status": "expired",
+                                "approval_id": approval_id,
+                                "resolved_at": approval_status.get("resolved_at"),
+                            }
+
+                        # Still pending, continue polling
+                        if poll_count % 12 == 0:  # Log every minute (12 * 5 seconds)
+                            elapsed_minutes = (time.time() - start_time) / 60
+                            remaining_minutes = (expires_at - time.time()) / 60
+                            logger.debug(
+                                "waiting_for_approval",
+                                approval_id=approval_id,
+                                status=approval_status["status"],
+                                elapsed_minutes=round(elapsed_minutes, 2),
+                                remaining_minutes=round(remaining_minutes, 2),
+                            )
+
+                    else:
+                        logger.warning(
+                            "approval_status_check_failed",
+                            approval_id=approval_id,
+                            status_code=status_response.status_code,
+                        )
+
+                except Exception as poll_error:
+                    logger.warning(
+                        "approval_poll_error",
+                        approval_id=approval_id,
+                        error=str(poll_error),
+                    )
+
+                # Wait before next poll
+                await asyncio.sleep(self.poll_interval_seconds)
+
+            # Timeout reached
+            logger.warning(
+                "approval_timeout",
+                approval_id=approval_id,
+                timeout_minutes=self.timeout_minutes,
+            )
+
+            return {
+                "approved": False,
+                "status": "expired",
+                "approval_id": approval_id,
+                "resolved_at": datetime.utcnow().isoformat(),
+            }
+
+        except Exception as e:
+            logger.error(
+                "wait_for_approval_failed",
+                title=title,
+                execution_id=self.execution_id,
+                error=str(e),
+            )
+            raise
+
+        finally:
+            await self.client.aclose()
+
+    def get_tools_schema(self) -> List[Dict[str, Any]]:
+        """
+        Get the tool schema for LLM function calling.
+
+        Returns list of tool definitions that can be provided to LLMs.
+        """
+        return [
+            {
+                "name": "wait_for_approval",
+                "description": "Pause workflow execution and wait for approval from authorized users before continuing. "
+                               "Use this when you need human approval for sensitive operations, decisions, or actions. "
+                               "The workflow will pause until an authorized user approves or rejects the request.",
+                "input_schema": {
+                    "type": "object",
+                    "properties": {
+                        "title": {
+                            "type": "string",
+                            "description": "Brief title for the approval request (e.g., 'Deploy to Production', 'Delete Database')"
+                        },
+                        "message": {
+                            "type": "string",
+                            "description": "Detailed message explaining why approval is needed and what will happen if approved"
+                        },
+                        "approver_user_emails": {
+                            "type": "array",
+                            "items": {"type": "string"},
+                            "description": "List of user email addresses who can approve this request"
+                        },
+                        "context": {
+                            "type": "object",
+                            "description": "Additional context data to help approvers make a decision"
+                        }
+                    },
+                    "required": ["title", "approver_user_emails"]
+                }
+            }
+        ]

control_plane_api/worker/services/approval_tools_agno.py

@@ -0,0 +1,207 @@
+"""
+Agno-compatible approval workflow tools.
+
+Provides wait_for_approval tool as an Agno Tool for seamless integration
+with Agno agent runtimes.
+"""
+import os
+from typing import List, Optional, Dict, Any
+import structlog
+from control_plane_api.worker.services.approval_tools import ApprovalTools
+
+logger = structlog.get_logger()
+
+
+class ApprovalToolkit:
+    """
+    Agno toolkit for approval workflow (human-in-the-loop gates).
+
+    Provides tools for workflows to request approval from authorized users
+    and wait for approval/rejection before continuing.
+    """
+
+    def __init__(
+        self,
+        control_plane_client: Any,
+        config: Optional[Dict[str, Any]] = None,
+    ):
+        """
+        Initialize approval toolkit.
+
+        Args:
+            control_plane_client: Control plane client (provides URL and credentials)
+            config: Optional configuration (timeout, etc.)
+        """
+        self.control_plane_client = control_plane_client
+        self.config = config or {}
+
+        # Store environment variable names - will be accessed at tool execution time
+        self.control_plane_url = None
+        self.api_key = None
+
+        # execution_id and organization_id will be set at runtime
+        self.execution_id = None
+        self.organization_id = None
+
+        logger.info("approval_toolkit_initialized")
+
+    def set_execution_context(self, execution_id: str, organization_id: str):
+        """
+        Set execution context (called by runtime before tool execution).
+
+        Args:
+            execution_id: Current execution ID
+            organization_id: Organization ID
+        """
+        self.execution_id = execution_id
+        self.organization_id = organization_id
+
+    async def wait_for_approval(
+        self,
+        title: str,
+        approver_user_emails: Optional[List[str]] = None,
+        approver_group_id: Optional[str] = None,
+        message: Optional[str] = None,
+        context: Optional[Dict[str, Any]] = None,
+    ) -> str:
+        """
+        Wait for approval from authorized users or groups before continuing.
+
+        This tool creates an approval request and pauses the workflow until
+        an authorized user approves or rejects the request. Use this for:
+        - Deployments to production
+        - Destructive operations (deletions, etc.)
+        - High-value transactions
+        - Policy-gated actions
+
+        APPROVERS: You can specify approvers in three ways:
+        1. Pass approver_user_emails parameter (list of emails)
+        2. Pass approver_group_id parameter (group UUID)
+        3. Use defaults from skill configuration (if configured)
+
+        If you don't specify approvers, the tool will use the default approvers
+        configured in the skill settings.
+
+        Args:
+            title: Brief title for the approval request (e.g., "Deploy to Production")
+            approver_user_emails: List of user email addresses who can approve (optional, uses config default if not provided)
+            approver_group_id: Group UUID whose members can approve (optional, uses config default if not provided)
+            message: Detailed message explaining why approval is needed (optional)
+            context: Additional context data to help approvers decide (optional)
+
+        Returns:
+            str: Approval result message
+
+        Examples:
+            ```python
+            # Approve by specific users
+            result = await wait_for_approval(
+                title="Deploy to Production",
+                message="Deploy version 2.0.0 to production environment",
+                approver_user_emails=["ops-lead@company.com", "cto@company.com"],
+                context={"version": "2.0.0", "environment": "production"}
+            )
+
+            # Approve by group
+            result = await wait_for_approval(
+                title="Delete Customer Data",
+                message="Permanently delete customer data for GDPR request",
+                approver_group_id="550e8400-e29b-41d4-a716-446655440000",  # Admin group UUID
+            )
+
+            # Approve by group OR specific users
+            result = await wait_for_approval(
+                title="Emergency Hotfix",
+                message="Deploy critical security patch",
+                approver_user_emails=["security-lead@company.com"],
+                approver_group_id="550e8400-e29b-41d4-a716-446655440000",  # Ops group UUID
+            )
+            ```
+        """
+        if not self.execution_id or not self.organization_id:
+            return "❌ Error: Approval toolkit not initialized with execution context"
+
+        # Get control plane URL and API key from environment at execution time
+        self.control_plane_url = os.getenv("CONTROL_PLANE_URL")
+        self.api_key = os.getenv("KUBIYA_API_KEY")
+
+        if not self.control_plane_url or not self.api_key:
+            return "❌ Error: CONTROL_PLANE_URL and KUBIYA_API_KEY environment variables must be set"
+
+        # Use config defaults if no approvers specified
+        if not approver_user_emails and not approver_group_id:
+            # Try to get defaults from config
+            approver_user_emails = self.config.get("default_approver_emails")
+            approver_group_id = self.config.get("default_approver_group_id")
+
+            # Still no approvers? Error
+            if not approver_user_emails and not approver_group_id:
+                return (
+                    "❌ Error: No approvers specified. Either:\n"
+                    "1. Pass approver_user_emails or approver_group_id to this tool call, OR\n"
+                    "2. Configure default approvers in the skill configuration"
+                )
+
+        logger.info(
+            "wait_for_approval_tool_called",
+            title=title,
+            approver_emails=approver_user_emails,
+            approver_group_id=approver_group_id,
+            execution_id=self.execution_id,
+        )
+
+        try:
+            # Create approval tools instance for this execution
+            approval_tools = ApprovalTools(
+                control_plane_url=self.control_plane_url,
+                api_key=self.api_key,
+                execution_id=self.execution_id,
+                organization_id=self.organization_id,
+                config=self.config,
+            )
+
+            # Call underlying approval tools
+            result = await approval_tools.wait_for_approval(
+                title=title,
+                message=message,
+                approver_user_emails=approver_user_emails or [],
+                approver_group_id=approver_group_id,
+                context=context,
+            )
+
+            if result["approved"]:
+                approved_by = result.get("approved_by_email", "unknown")
+                return (
+                    f"✅ Approval granted by {approved_by}. "
+                    f"You may proceed with '{title}'."
+                )
+            elif result["status"] == "rejected":
+                rejected_by = result.get("rejected_by_email", "unknown")
+                reason = result.get("rejection_reason", "No reason provided")
+                return (
+                    f"❌ Request rejected by {rejected_by}. "
+                    f"Reason: {reason}. "
+                    f"You must not proceed with '{title}'."
+                )
+            elif result["status"] == "expired":
+                return (
+                    f"⏱️ Approval request expired without response. "
+                    f"You must not proceed with '{title}' without approval."
+                )
+            else:
+                return (
+                    f"⚠️ Approval request ended with status: {result['status']}. "
+                    f"You must not proceed with '{title}' without explicit approval."
+                )
+
+        except Exception as e:
+            logger.error(
+                "wait_for_approval_tool_error",
+                error=str(e),
+                title=title,
+                execution_id=self.execution_id,
+            )
+            return (
+                f"❌ Failed to request approval: {str(e)}. "
+                f"You must not proceed with '{title}' due to approval system error."
+            )

control_plane_api/worker/services/cancellation_manager.py

@@ -0,0 +1,177 @@
+"""Cancellation manager - handles agent/team registry and cancellation"""
+
+from typing import Dict, Any, Optional
+from datetime import datetime, timezone
+import structlog
+
+logger = structlog.get_logger()
+
+
+class CancellationManager:
+    """
+    Manages active agent/team instances for cancellation support.
+
+    Provides a centralized registry and cancellation logic that works
+    with Agno's cancel_run() API.
+    """
+
+    def __init__(self):
+        # Key: execution_id, Value: {agent/team, run_id, started_at}
+        self._registry: Dict[str, Dict[str, Any]] = {}
+
+    def register(
+        self,
+        execution_id: str,
+        instance: Any,  # Agent or Team
+        instance_type: str = "agent"
+    ) -> None:
+        """
+        Register an agent or team for cancellation support.
+
+        Args:
+            execution_id: Unique execution ID
+            instance: Agno Agent or Team instance
+            instance_type: "agent" or "team"
+        """
+        self._registry[execution_id] = {
+            "instance": instance,
+            "instance_type": instance_type,
+            "run_id": None,  # Set when run starts
+            "started_at": datetime.now(timezone.utc).isoformat(),
+        }
+
+        logger.info(
+            f"{instance_type}_registered_for_cancellation",
+            execution_id=execution_id[:8],
+            instance_type=instance_type
+        )
+
+    def set_run_id(self, execution_id: str, run_id: str) -> None:
+        """
+        Set the Agno run_id for an execution (captured from first streaming chunk).
+
+        Args:
+            execution_id: Execution ID
+            run_id: Agno run_id from streaming response
+        """
+        if execution_id in self._registry:
+            self._registry[execution_id]["run_id"] = run_id
+            logger.info(
+                "run_id_captured",
+                execution_id=execution_id[:8],
+                run_id=run_id[:16]
+            )
+
+    def cancel(self, execution_id: str) -> Dict[str, Any]:
+        """
+        Cancel an active execution using Agno's cancel_run API.
+
+        Args:
+            execution_id: Execution to cancel
+
+        Returns:
+            Dict with success status and details
+        """
+        # Check if execution exists in registry
+        if execution_id not in self._registry:
+            logger.warning(
+                "cancel_execution_not_found",
+                execution_id=execution_id[:8]
+            )
+            return {
+                "success": False,
+                "error": "Execution not found or already completed",
+                "execution_id": execution_id,
+            }
+
+        entry = self._registry[execution_id]
+        instance = entry["instance"]
+        run_id = entry.get("run_id")
+        instance_type = entry.get("instance_type", "agent")
+
+        # Check if run has started
+        if not run_id:
+            logger.warning(
+                "cancel_no_run_id",
+                execution_id=execution_id[:8]
+            )
+            return {
+                "success": False,
+                "error": "Execution not started yet",
+                "execution_id": execution_id,
+            }
+
+        logger.info(
+            f"cancelling_{instance_type}_run",
+            execution_id=execution_id[:8],
+            run_id=run_id[:16]
+        )
+
+        try:
+            # Call Agno's cancel_run API
+            success = instance.cancel_run(run_id)
+
+            if success:
+                logger.info(
+                    f"{instance_type}_run_cancelled",
+                    execution_id=execution_id[:8],
+                    run_id=run_id[:16]
+                )
+
+                # Clean up registry
+                del self._registry[execution_id]
+
+                return {
+                    "success": True,
+                    "execution_id": execution_id,
+                    "run_id": run_id,
+                    "instance_type": instance_type,
+                    "cancelled_at": datetime.now(timezone.utc).isoformat(),
+                }
+            else:
+                logger.warning(
+                    f"{instance_type}_cancel_failed",
+                    execution_id=execution_id[:8],
+                    run_id=run_id[:16]
+                )
+                return {
+                    "success": False,
+                    "error": "Cancel failed - run may be completed",
+                    "execution_id": execution_id,
+                    "run_id": run_id,
+                }
+
+        except Exception as e:
+            logger.error(
+                f"{instance_type}_cancel_error",
+                execution_id=execution_id[:8],
+                error=str(e)
+            )
+            return {
+                "success": False,
+                "error": str(e),
+                "execution_id": execution_id,
+            }
+
+    def unregister(self, execution_id: str) -> None:
+        """
+        Remove an execution from the registry (called on completion).
+
+        Args:
+            execution_id: Execution to unregister
+        """
+        if execution_id in self._registry:
+            instance_type = self._registry[execution_id].get("instance_type", "agent")
+            del self._registry[execution_id]
+            logger.info(
+                f"{instance_type}_unregistered",
+                execution_id=execution_id[:8]
+            )
+
+    def get_active_count(self) -> int:
+        """Get number of active executions in registry."""
+        return len(self._registry)
+
+
+# Global singleton instance
+cancellation_manager = CancellationManager()