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

control_plane_api/worker/services/tool_enforcement.py

@@ -0,0 +1,254 @@
+"""
+Tool enforcement service for policy-based tool governance.
+
+This module provides non-blocking enforcement checks that run in parallel
+with tool execution, injecting policy violations into tool results.
+"""
+
+import asyncio
+import structlog
+from typing import Dict, Any, Optional, Tuple
+from datetime import datetime, timezone
+
+logger = structlog.get_logger(__name__)
+
+
+class ToolEnforcementService:
+    """
+    Service for enforcing policies on tool executions.
+
+    Features:
+    - Non-blocking parallel enforcement checks
+    - Violation injection into tool outputs
+    - Context enrichment with user/org/agent metadata
+    - Graceful degradation on enforcer failures
+    """
+
+    def __init__(self, policy_enforcer_client: Optional[Any] = None):
+        """
+        Initialize the tool enforcement service.
+
+        Args:
+            policy_enforcer_client: PolicyEnforcerClient instance (optional)
+        """
+        self.enforcer = policy_enforcer_client
+        self.enabled = policy_enforcer_client is not None
+
+    async def enforce_tool_execution(
+        self,
+        tool_name: str,
+        tool_args: Dict[str, Any],
+        enforcement_context: Dict[str, Any],
+        timeout: float = 2.0,
+    ) -> Tuple[bool, Optional[str], Dict[str, Any]]:
+        """
+        Perform non-blocking enforcement check for tool execution.
+
+        Args:
+            tool_name: Name of tool being executed
+            tool_args: Tool arguments
+            enforcement_context: Full context (user, org, agent, etc.)
+            timeout: Maximum time to wait for enforcement (default 2s)
+
+        Returns:
+            Tuple of (allow, violation_message, metadata)
+            - allow: Whether tool is allowed (True if enforcer unavailable)
+            - violation_message: Error message if denied (None if allowed)
+            - metadata: Enforcement metadata (policies evaluated, etc.)
+        """
+        if not self.enabled:
+            return True, None, {"enforcer": "disabled"}
+
+        try:
+            # Build enforcement request payload
+            enforcement_payload = self._build_enforcement_payload(
+                tool_name=tool_name,
+                tool_args=tool_args,
+                context=enforcement_context,
+            )
+
+            # Call enforcer with timeout (non-blocking)
+            enforcement_result = await asyncio.wait_for(
+                self.enforcer.evaluation.enforce(enforcement_payload),
+                timeout=timeout,
+            )
+
+            allow = enforcement_result.get("allow", True)
+            policies = enforcement_result.get("policies", [])
+
+            if not allow:
+                violation_msg = self._format_violation_message(
+                    tool_name=tool_name,
+                    policies=policies,
+                    enforcement_result=enforcement_result,
+                )
+                return False, violation_msg, {
+                    "enforcer": "blocked",
+                    "policies": policies,
+                    "enforcement_id": enforcement_result.get("id"),
+                }
+
+            return True, None, {
+                "enforcer": "allowed",
+                "policies": policies,
+                "enforcement_id": enforcement_result.get("id"),
+            }
+
+        except asyncio.TimeoutError:
+            logger.warning(
+                "tool_enforcement_timeout",
+                tool_name=tool_name,
+                timeout=timeout,
+            )
+            # Fail open on timeout
+            return True, None, {"enforcer": "timeout"}
+
+        except Exception as e:
+            logger.error(
+                "tool_enforcement_error",
+                tool_name=tool_name,
+                error=str(e),
+                exc_info=True,
+            )
+            # Fail open on error
+            return True, None, {"enforcer": "error", "error": str(e)}
+
+    def _build_enforcement_payload(
+        self,
+        tool_name: str,
+        tool_args: Dict[str, Any],
+        context: Dict[str, Any],
+    ) -> Dict[str, Any]:
+        """
+        Build enforcement request payload with all required context.
+
+        Required fields:
+        - tool_name, tool_arguments
+        - user_email, organization_id, team_id, roles
+        - execution_id, agent_id, environment
+        - tool_source, tool_category, risk_level
+        - timestamp
+        """
+        return {
+            "action": "tool_execution",
+            "tool": {
+                "name": tool_name,
+                "arguments": tool_args,
+                "source": self._determine_tool_source(tool_name),
+                "category": self._determine_tool_category(tool_name),
+                "risk_level": self._determine_risk_level(tool_name, tool_args),
+            },
+            "user": {
+                "email": context.get("user_email"),
+                "id": context.get("user_id"),
+                "roles": context.get("user_roles", []),
+            },
+            "organization": {
+                "id": context.get("organization_id"),
+                "name": context.get("organization_name"),
+            },
+            "team": {
+                "id": context.get("team_id"),
+                "name": context.get("team_name"),
+            },
+            "execution": {
+                "execution_id": context.get("execution_id"),
+                "agent_id": context.get("agent_id"),
+                "environment": context.get("environment", "production"),
+                "timestamp": datetime.now(timezone.utc).isoformat(),
+            },
+            "metadata": context.get("metadata", {}),
+        }
+
+    def _determine_tool_source(self, tool_name: str) -> str:
+        """Determine tool source from name."""
+        if tool_name.startswith("mcp__"):
+            return "mcp"
+        elif tool_name in [
+            "Bash",
+            "Read",
+            "Write",
+            "Edit",
+            "Grep",
+            "Glob",
+            "WebFetch",
+            "WebSearch",
+            "TodoWrite",
+            "AskUserQuestion",
+        ]:
+            return "builtin"
+        else:
+            return "skill"
+
+    def _determine_tool_category(self, tool_name: str) -> str:
+        """Categorize tool by function."""
+        if tool_name in ["Bash", "Shell"]:
+            return "command_execution"
+        elif tool_name in ["Read", "Write", "Edit"]:
+            return "file_operation"
+        elif tool_name in ["Grep", "Glob", "Find"]:
+            return "file_search"
+        elif "api" in tool_name.lower() or "http" in tool_name.lower() or tool_name in [
+            "WebFetch",
+            "WebSearch",
+        ]:
+            return "network"
+        else:
+            return "general"
+
+    def _determine_risk_level(
+        self,
+        tool_name: str,
+        tool_args: Dict[str, Any],
+    ) -> str:
+        """Assess risk level of tool execution."""
+        # High risk: command execution, destructive file ops
+        if tool_name in ["Bash", "Shell"]:
+            command = tool_args.get("command", "")
+            # Critical: Destructive commands
+            if any(
+                cmd in command
+                for cmd in ["rm -rf", "dd if=", "mkfs", "> /dev/", "format"]
+            ):
+                return "critical"
+            # High: Any command execution
+            return "high"
+
+        if tool_name == "Write":
+            return "medium"
+
+        if tool_name == "Edit":
+            return "medium"
+
+        if tool_name == "Read":
+            path = tool_args.get("file_path", "")
+            # High: Sensitive file access
+            if any(
+                sensitive in path
+                for sensitive in ["/etc/passwd", ".ssh", ".env", ".key", "credentials"]
+            ):
+                return "high"
+            return "low"
+
+        # Network operations
+        if tool_name in ["WebFetch", "WebSearch"]:
+            return "medium"
+
+        return "low"
+
+    def _format_violation_message(
+        self,
+        tool_name: str,
+        policies: list,
+        enforcement_result: Dict[str, Any],
+    ) -> str:
+        """Format user-friendly violation message."""
+        policy_names = ", ".join(policies) if policies else "unknown policies"
+        return (
+            f"Tool execution blocked by policy enforcement.\n"
+            f"Tool: {tool_name}\n"
+            f"Blocked by: {policy_names}\n"
+            f"Enforcement ID: {enforcement_result.get('id', 'N/A')}\n"
+            f"\n"
+            f"Contact your administrator for access approval."
+        )

control_plane_api/worker/services/workflow_executor/__init__.py

@@ -0,0 +1,52 @@
+"""
+Workflow Executor Skill
+
+A well-architected skill for executing workflows from agent control plane.
+
+This package provides:
+- Type-safe Pydantic models for configurations and events
+- Proper separation of concerns with dedicated components
+- Strategy pattern for different workflow types
+- Event-driven architecture for real-time streaming
+- Comprehensive error handling and logging
+
+Architecture:
+- models.py: Pydantic models for type safety
+- event_publisher.py: Publishes events to control plane
+- event_processor.py: Processes streaming events from SDK
+- executors/: Strategy pattern for workflow execution
+  - base.py: Abstract base executor (template method pattern)
+  - json_executor.py: JSON workflow execution
+  - python_executor.py: Python DSL workflow execution
+
+Usage:
+    from control_plane_api.worker.services.workflow_executor import WorkflowExecutorTools
+
+    # This will be the main skill class that agents use
+"""
+
+from .models import (
+    WorkflowConfig,
+    WorkflowExecutionContext,
+    WorkflowEvent,
+    WorkflowResult,
+)
+from .event_publisher import WorkflowEventPublisher
+from .event_processor import WorkflowEventProcessor
+from .executors import (
+    BaseWorkflowExecutor,
+    JsonWorkflowExecutor,
+    PythonWorkflowExecutor,
+)
+
+__all__ = [
+    "WorkflowConfig",
+    "WorkflowExecutionContext",
+    "WorkflowEvent",
+    "WorkflowResult",
+    "WorkflowEventPublisher",
+    "WorkflowEventProcessor",
+    "BaseWorkflowExecutor",
+    "JsonWorkflowExecutor",
+    "PythonWorkflowExecutor",
+]

control_plane_api/worker/services/workflow_executor/event_processor.py

@@ -0,0 +1,287 @@
+"""
+Workflow event processor.
+
+This module processes streaming events from the Kubiya SDK,
+handling deduplication, parsing, and state management.
+"""
+
+import json
+import hashlib
+import structlog
+from typing import Dict, Any, Set, Optional, Callable
+from .event_publisher import WorkflowEventPublisher
+
+logger = structlog.get_logger()
+
+
+class WorkflowEventProcessor:
+    """
+    Processes streaming events from workflow execution.
+
+    Handles event parsing, deduplication, and publishes to control plane.
+    """
+
+    def __init__(
+        self,
+        event_publisher: WorkflowEventPublisher,
+        workflow_name: str,
+        message_id: str,
+        stream_callback: Optional[Callable[[str], None]] = None
+    ):
+        """
+        Initialize event processor.
+
+        Args:
+            event_publisher: Event publisher instance
+            workflow_name: Name of the workflow being executed
+            message_id: Unique workflow message ID
+            stream_callback: Optional callback for streaming output
+        """
+        self.event_publisher = event_publisher
+        self.workflow_name = workflow_name
+        self.message_id = message_id
+        self.stream_callback = stream_callback
+
+        # State tracking
+        self.event_count = 0
+        self.accumulated_output: list[str] = []
+        self.seen_events: Set[str] = set()
+        self.current_step: Optional[str] = None
+        self.step_outputs: Dict[str, str] = {}
+
+    def process_event(self, event: Any) -> Optional[Dict[str, Any]]:
+        """
+        Process a single streaming event from the SDK.
+
+        Args:
+            event: Raw event from SDK (string, bytes, or dict)
+
+        Returns:
+            Parsed event data as dict, or None if event should be skipped
+        """
+        self.event_count += 1
+
+        # Skip None/empty events
+        if event is None:
+            logger.debug("skipping_none_event", event_number=self.event_count)
+            return None
+
+        # Parse the event
+        try:
+            event_data = self._parse_event(event)
+            if event_data is None:
+                return None
+
+            # Process based on event type
+            event_type = event_data.get("type", "unknown")
+            logger.debug(
+                "processing_workflow_event",
+                event_type=event_type,
+                event_number=self.event_count
+            )
+
+            # Route to specific handler
+            if event_type == "step_output":
+                self._handle_step_output(event_data)
+            elif event_type == "step_running":
+                self._handle_step_running(event_data)
+            elif event_type == "step_complete":
+                self._handle_step_complete(event_data)
+            elif event_type == "workflow_complete":
+                self._handle_workflow_complete(event_data)
+            elif event_type == "workflow_failed":
+                self._handle_workflow_failed(event_data)
+            else:
+                logger.debug("unknown_event_type", event_type=event_type)
+
+            return event_data
+
+        except Exception as e:
+            logger.error(
+                "failed_to_process_event",
+                error=str(e),
+                event_number=self.event_count
+            )
+            return None
+
+    def _parse_event(self, event: Any) -> Optional[Dict[str, Any]]:
+        """Parse raw event into dict."""
+        try:
+            # Handle bytes
+            if isinstance(event, bytes):
+                if not event:
+                    return None
+                event = event.decode('utf-8')
+
+            # Handle strings
+            if isinstance(event, str):
+                if not event.strip():
+                    return None
+
+                # Handle SSE format: "data: 2:{json}"
+                if event.startswith("data: "):
+                    event = event[6:]  # Remove "data: " prefix
+
+                    # Strip message ID prefix like "2:"
+                    if ":" in event and event.split(":", 1)[0].isdigit():
+                        event = event.split(":", 1)[1]
+
+                # Parse JSON
+                return json.loads(event)
+
+            # Already a dict
+            elif isinstance(event, dict):
+                return event
+
+            # Unknown type
+            else:
+                logger.warning(
+                    "unknown_event_type_treating_as_text",
+                    type_name=type(event).__name__
+                )
+                event_str = str(event)
+                if event_str.strip():
+                    self.accumulated_output.append(event_str)
+                    if self.stream_callback:
+                        self.stream_callback(f"{event_str}\n")
+                return None
+
+        except (json.JSONDecodeError, UnicodeDecodeError) as e:
+            logger.warning("failed_to_parse_event_treating_as_text", error=str(e))
+            event_str = str(event)
+            if event_str.strip():
+                self.accumulated_output.append(event_str)
+                if self.stream_callback:
+                    self.stream_callback(f"{event_str}\n")
+            return None
+
+    def _handle_step_output(self, event_data: Dict[str, Any]) -> None:
+        """Handle step_output event."""
+        step = event_data.get("step", {})
+        step_name = step.get("name", "unknown")
+        step_status = step.get("status", "")
+        output = step.get("output", "")
+
+        if not output.strip():
+            return
+
+        # Deduplicate events - SDK sends same event twice
+        event_hash = hashlib.md5(f"{step_name}:{output}".encode()).hexdigest()
+        if event_hash in self.seen_events:
+            logger.debug("skipping_duplicate_event", step_name=step_name)
+            return
+        self.seen_events.add(event_hash)
+
+        logger.debug("step_output_received", step_name=step_name, output_length=len(output))
+
+        # Store output
+        self.accumulated_output.append(output)
+        self.step_outputs[step_name] = output
+
+        # Stream to callback
+        if self.stream_callback:
+            formatted_output = f"```\n{output}\n```\n"
+            self.stream_callback(formatted_output)
+
+        # Publish to control plane
+        self.event_publisher.publish_step_output(
+            workflow_name=self.workflow_name,
+            step_name=step_name,
+            output=output,
+            message_id=self.message_id
+        )
+
+        # Check if step failed
+        if step_status in ["error", "failed"]:
+            logger.warning("step_failed_detected", step_name=step_name, status=step_status)
+            self.event_publisher.publish_step_complete(
+                workflow_name=self.workflow_name,
+                step_name=step_name,
+                status="failed",
+                message_id=self.message_id,
+                error=output
+            )
+
+    def _handle_step_running(self, event_data: Dict[str, Any]) -> None:
+        """Handle step_running event."""
+        step = event_data.get("step", {})
+        step_name = step.get("name", "unknown")
+
+        logger.debug("step_running", step_name=step_name)
+
+        self.current_step = step_name
+        formatted = f"\n▶️  Step: {step_name}"
+        self.accumulated_output.append(formatted)
+
+        if self.stream_callback:
+            self.stream_callback(f"{formatted}\n")
+
+        self.event_publisher.publish_step_running(
+            workflow_name=self.workflow_name,
+            step_name=step_name,
+            message_id=self.message_id
+        )
+
+    def _handle_step_complete(self, event_data: Dict[str, Any]) -> None:
+        """Handle step_complete event."""
+        step = event_data.get("step", {})
+        step_name = step.get("name", "unknown")
+        status = step.get("status", "unknown")
+
+        logger.debug("step_complete", step_name=step_name, status=status)
+
+        icon = "✅" if status == "finished" else "❌"
+        formatted = f"{icon} Step '{step_name}' {status}"
+        self.accumulated_output.append(formatted)
+        self.current_step = None
+
+        if self.stream_callback:
+            self.stream_callback(f"{formatted}\n")
+
+        self.event_publisher.publish_step_complete(
+            workflow_name=self.workflow_name,
+            step_name=step_name,
+            status="completed" if status == "finished" else "failed",
+            message_id=self.message_id
+        )
+
+    def _handle_workflow_complete(self, event_data: Dict[str, Any]) -> None:
+        """Handle workflow_complete event."""
+        dag_name = event_data.get("dagName", "unknown")
+        status = event_data.get("status", "unknown")
+        success = event_data.get("success", False)
+
+        logger.info("workflow_complete", dag_name=dag_name, status=status, success=success)
+
+        icon = "✅" if success else "❌"
+        formatted = f"\n{icon} Workflow '{dag_name}' {status}"
+        self.accumulated_output.append(formatted)
+
+        if self.stream_callback:
+            self.stream_callback(f"{formatted}\n")
+
+    def _handle_workflow_failed(self, event_data: Dict[str, Any]) -> None:
+        """Handle workflow_failed event."""
+        error = event_data.get("error", "Unknown error")
+
+        logger.error("workflow_failed", error=error)
+
+        formatted = f"\n❌ Workflow failed: {error}"
+        self.accumulated_output.append(formatted)
+
+        if self.stream_callback:
+            self.stream_callback(f"{formatted}\n")
+
+        self.event_publisher.publish_failed(
+            workflow_name=self.workflow_name,
+            error=error,
+            message_id=self.message_id
+        )
+
+    def get_accumulated_output(self) -> list[str]:
+        """Get all accumulated output lines."""
+        return self.accumulated_output
+
+    def get_event_count(self) -> int:
+        """Get total number of events processed."""
+        return self.event_count