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

control_plane_api/worker/services/workflow_executor/event_publisher.py

@@ -0,0 +1,210 @@
+"""
+Workflow event publisher.
+
+This module handles publishing workflow events to the control plane
+with proper type safety and error handling.
+"""
+
+import structlog
+from typing import Optional, Any
+from .models import (
+    WorkflowEvent,
+    WorkflowStartedEvent,
+    WorkflowStepRunningEvent,
+    WorkflowStepOutputEvent,
+    WorkflowStepCompleteEvent,
+    WorkflowCompletedEvent,
+    WorkflowFailedEvent,
+    WorkflowCancelledEvent,
+)
+
+logger = structlog.get_logger()
+
+
+class WorkflowEventPublisher:
+    """
+    Publishes workflow events to the control plane.
+
+    Handles event serialization, error handling, and control plane communication.
+    """
+
+    def __init__(self, control_plane: Any, execution_id: str):
+        """
+        Initialize event publisher.
+
+        Args:
+            control_plane: Control plane client instance
+            execution_id: Agent execution ID
+        """
+        self.control_plane = control_plane
+        self.execution_id = execution_id
+        self.enabled = control_plane is not None and execution_id is not None
+
+    def publish(self, event: WorkflowEvent) -> bool:
+        """
+        Publish a workflow event.
+
+        Args:
+            event: Workflow event to publish
+
+        Returns:
+            True if published successfully, False otherwise
+        """
+        if not self.enabled:
+            logger.debug(
+                "event_publisher_disabled",
+                event_type=event.event_type,
+                reason="control_plane_or_execution_id_missing"
+            )
+            return False
+
+        try:
+            # Convert Pydantic model to dict for control plane
+            event_data = event.model_dump(mode='json', exclude_none=True)
+
+            logger.debug(
+                "publishing_workflow_event",
+                event_type=event.event_type,
+                execution_id=self.execution_id[:8],
+                workflow_name=event.workflow_name
+            )
+
+            self.control_plane.publish_event(
+                execution_id=self.execution_id,
+                event_type=event.event_type,
+                data=event_data
+            )
+
+            logger.debug(
+                "workflow_event_published",
+                event_type=event.event_type,
+                execution_id=self.execution_id[:8]
+            )
+
+            return True
+
+        except Exception as e:
+            logger.error(
+                "failed_to_publish_workflow_event",
+                event_type=event.event_type,
+                error=str(e),
+                execution_id=self.execution_id[:8]
+            )
+            return False
+
+    def publish_started(
+        self,
+        workflow_name: str,
+        workflow_type: str,
+        runner: str,
+        parameters: dict,
+        message_id: str
+    ) -> bool:
+        """Publish workflow started event."""
+        event = WorkflowStartedEvent(
+            workflow_name=workflow_name,
+            workflow_type=workflow_type,
+            runner=runner,
+            parameters=parameters,
+            message_id=message_id,
+            message=f"🚀 Starting workflow: {workflow_name}"
+        )
+        return self.publish(event)
+
+    def publish_step_running(
+        self,
+        workflow_name: str,
+        step_name: str,
+        message_id: str
+    ) -> bool:
+        """Publish step running event."""
+        event = WorkflowStepRunningEvent(
+            workflow_name=workflow_name,
+            step_name=step_name,
+            message_id=message_id
+        )
+        return self.publish(event)
+
+    def publish_step_output(
+        self,
+        workflow_name: str,
+        step_name: str,
+        output: str,
+        message_id: str
+    ) -> bool:
+        """Publish step output event."""
+        event = WorkflowStepOutputEvent(
+            workflow_name=workflow_name,
+            step_name=step_name,
+            output=output,
+            message_id=message_id
+        )
+        return self.publish(event)
+
+    def publish_step_complete(
+        self,
+        workflow_name: str,
+        step_name: str,
+        status: str,
+        message_id: str,
+        output: Optional[str] = None,
+        error: Optional[str] = None,
+        duration: Optional[float] = None
+    ) -> bool:
+        """Publish step complete event."""
+        event = WorkflowStepCompleteEvent(
+            workflow_name=workflow_name,
+            step_name=step_name,
+            status=status,  # type: ignore
+            message_id=message_id,
+            output=output,
+            error=error,
+            duration=duration
+        )
+        return self.publish(event)
+
+    def publish_completed(
+        self,
+        workflow_name: str,
+        duration: float,
+        steps_completed: int,
+        message_id: str
+    ) -> bool:
+        """Publish workflow completed event."""
+        event = WorkflowCompletedEvent(
+            workflow_name=workflow_name,
+            duration=duration,
+            steps_completed=steps_completed,
+            message_id=message_id
+        )
+        return self.publish(event)
+
+    def publish_failed(
+        self,
+        workflow_name: str,
+        error: str,
+        message_id: str,
+        duration: Optional[float] = None
+    ) -> bool:
+        """Publish workflow failed event."""
+        event = WorkflowFailedEvent(
+            workflow_name=workflow_name,
+            error=error,
+            message_id=message_id,
+            duration=duration
+        )
+        return self.publish(event)
+
+    def publish_cancelled(
+        self,
+        workflow_name: str,
+        message: str,
+        message_id: str
+    ) -> bool:
+        """Publish workflow cancelled event."""
+        event = WorkflowCancelledEvent(
+            workflow_name=workflow_name,
+            message=message,
+            message_id=message_id
+        )
+        return self.publish(event)

control_plane_api/worker/services/workflow_executor/executors/__init__.py

@@ -0,0 +1,15 @@
+"""
+Workflow executors.
+
+This module contains workflow execution strategies for different workflow types.
+"""
+
+from .base import BaseWorkflowExecutor
+from .json_executor import JsonWorkflowExecutor
+from .python_executor import PythonWorkflowExecutor
+
+__all__ = [
+    "BaseWorkflowExecutor",
+    "JsonWorkflowExecutor",
+    "PythonWorkflowExecutor",
+]

control_plane_api/worker/services/workflow_executor/executors/base.py

@@ -0,0 +1,270 @@
+"""
+Base workflow executor.
+
+This module defines the abstract base class for workflow executors,
+implementing the strategy pattern for different workflow types.
+"""
+
+from abc import ABC, abstractmethod
+from datetime import datetime
+import structlog
+from typing import Any
+import threading
+
+from ..models import WorkflowExecutionContext, WorkflowResult
+from ..event_publisher import WorkflowEventPublisher
+from ..event_processor import WorkflowEventProcessor
+
+logger = structlog.get_logger()
+
+
+class BaseWorkflowExecutor(ABC):
+    """
+    Abstract base class for workflow executors.
+
+    Implements the template method pattern for workflow execution,
+    allowing subclasses to define workflow-type-specific logic.
+    """
+
+    def __init__(
+        self,
+        kubiya_client: Any,
+        event_publisher: WorkflowEventPublisher,
+        cancellation_manager: Any
+    ):
+        """
+        Initialize workflow executor.
+
+        Args:
+            kubiya_client: Kubiya SDK client instance
+            event_publisher: Event publisher for control plane communication
+            cancellation_manager: Workflow cancellation manager
+        """
+        self.kubiya_client = kubiya_client
+        self.event_publisher = event_publisher
+        self.cancellation_manager = cancellation_manager
+
+    def execute(self, context: WorkflowExecutionContext) -> WorkflowResult:
+        """
+        Execute workflow (template method).
+
+        This method defines the overall execution flow:
+        1. Register for cancellation
+        2. Publish started event
+        3. Execute workflow (subclass-specific)
+        4. Process streaming events
+        5. Publish completion/failure events
+        6. Return result
+
+        Args:
+            context: Workflow execution context
+
+        Returns:
+            Workflow execution result
+
+        Raises:
+            Exception: If workflow execution fails
+        """
+        start_time = datetime.utcnow()
+
+        # Register workflow for cancellation tracking
+        cancellation_event = self.cancellation_manager.register_workflow(
+            context.execution_id,
+            context.workflow_message_id
+        )
+
+        logger.info(
+            "workflow_execution_started",
+            workflow_name=context.workflow_config.name,
+            workflow_type=context.workflow_config.type,
+            execution_id=context.execution_id[:8]
+        )
+
+        # Publish started event
+        self.event_publisher.publish_started(
+            workflow_name=context.workflow_config.name,
+            workflow_type=context.workflow_config.type,
+            runner=context.workflow_config.runner or "default",
+            parameters=context.workflow_config.parameters,
+            message_id=context.workflow_message_id
+        )
+
+        try:
+            # Execute workflow (subclass implements this)
+            response_stream = self._execute_workflow(context)
+
+            # Process streaming events
+            processor = WorkflowEventProcessor(
+                event_publisher=self.event_publisher,
+                workflow_name=context.workflow_config.name,
+                message_id=context.workflow_message_id,
+                stream_callback=None  # Can be passed from context if needed
+            )
+
+            # Process events from stream
+            for event in response_stream:
+                # Check for cancellation first
+                if cancellation_event.is_set():
+                    logger.warning(
+                        "workflow_cancelled_by_user",
+                        workflow_name=context.workflow_config.name,
+                        execution_id=context.execution_id[:8]
+                    )
+
+                    self.cancellation_manager.clear_cancellation(
+                        context.execution_id,
+                        context.workflow_message_id
+                    )
+
+                    self.event_publisher.publish_cancelled(
+                        workflow_name=context.workflow_config.name,
+                        message=f"❌ Workflow '{context.workflow_config.name}' cancelled by user",
+                        message_id=context.workflow_message_id
+                    )
+
+                    # Return cancellation result
+                    duration = (datetime.utcnow() - start_time).total_seconds()
+                    return WorkflowResult(
+                        workflow_name=context.workflow_config.name,
+                        status="cancelled",
+                        duration=duration,
+                        output=f"❌ Workflow execution cancelled by user\n\nWorkflow: {context.workflow_config.name}\nCancelled at: {datetime.utcnow().isoformat()}",
+                        event_count=processor.get_event_count()
+                    )
+
+                # Process event
+                processor.process_event(event)
+
+            # Calculate duration
+            end_time = datetime.utcnow()
+            duration = (end_time - start_time).total_seconds()
+
+            # Build output from accumulated events
+            accumulated_output = processor.get_accumulated_output()
+            output = self._format_result(
+                context=context,
+                duration=duration,
+                accumulated_output=accumulated_output,
+                event_count=processor.get_event_count()
+            )
+
+            # Publish completion event
+            self.event_publisher.publish_completed(
+                workflow_name=context.workflow_config.name,
+                duration=duration,
+                steps_completed=len(accumulated_output),
+                message_id=context.workflow_message_id
+            )
+
+            logger.info(
+                "workflow_execution_completed",
+                workflow_name=context.workflow_config.name,
+                duration=duration,
+                event_count=processor.get_event_count(),
+                execution_id=context.execution_id[:8]
+            )
+
+            # Clear cancellation registration
+            self.cancellation_manager.clear_cancellation(
+                context.execution_id,
+                context.workflow_message_id
+            )
+
+            return WorkflowResult(
+                workflow_name=context.workflow_config.name,
+                status="completed",
+                duration=duration,
+                output=output,
+                event_count=processor.get_event_count()
+            )
+
+        except Exception as e:
+            # Calculate duration before failure
+            duration = (datetime.utcnow() - start_time).total_seconds()
+
+            logger.error(
+                "workflow_execution_failed",
+                workflow_name=context.workflow_config.name,
+                error=str(e),
+                duration=duration,
+                execution_id=context.execution_id[:8]
+            )
+
+            # Publish failure event
+            self.event_publisher.publish_failed(
+                workflow_name=context.workflow_config.name,
+                error=str(e),
+                message_id=context.workflow_message_id,
+                duration=duration
+            )
+
+            # Clear cancellation registration
+            self.cancellation_manager.clear_cancellation(
+                context.execution_id,
+                context.workflow_message_id
+            )
+
+            return WorkflowResult(
+                workflow_name=context.workflow_config.name,
+                status="failed",
+                duration=duration,
+                output="",
+                error=str(e),
+                event_count=0
+            )
+
+    @abstractmethod
+    def _execute_workflow(self, context: WorkflowExecutionContext) -> Any:
+        """
+        Execute workflow and return event stream.
+
+        Subclasses must implement this method to execute their specific
+        workflow type and return an iterable event stream.
+
+        Args:
+            context: Workflow execution context
+
+        Returns:
+            Iterable event stream from Kubiya SDK
+
+        Raises:
+            Exception: If workflow execution fails
+        """
+        pass
+
+    def _format_result(
+        self,
+        context: WorkflowExecutionContext,
+        duration: float,
+        accumulated_output: list,
+        event_count: int
+    ) -> str:
+        """
+        Format final workflow result for agent.
+
+        Args:
+            context: Workflow execution context
+            duration: Total execution duration
+            accumulated_output: List of output lines
+            event_count: Number of events processed
+
+        Returns:
+            Formatted result string
+        """
+        result = f"\n{'='*60}\n"
+        result += f"Workflow Execution: {context.workflow_config.name}\n"
+        result += f"{'='*60}\n\n"
+        result += f"Status: ✅ Completed\n"
+        result += f"Duration: {duration:.2f}s\n"
+        result += f"Runner: {context.workflow_config.runner or 'default'}\n"
+        result += f"Total Events: {event_count}\n"
+
+        if accumulated_output:
+            result += f"\n{'='*60}\n"
+            result += f"Workflow Output:\n"
+            result += f"{'='*60}\n\n"
+            result += "\n".join(accumulated_output)
+        else:
+            result += "\n⚠️  No output generated\n"
+
+        return result

control_plane_api/worker/services/workflow_executor/executors/json_executor.py

@@ -0,0 +1,50 @@
+"""
+JSON workflow executor.
+
+This module implements workflow execution for JSON-defined workflows.
+"""
+
+import structlog
+from typing import Any
+
+from .base import BaseWorkflowExecutor
+from ..models import WorkflowExecutionContext
+
+logger = structlog.get_logger()
+
+
+class JsonWorkflowExecutor(BaseWorkflowExecutor):
+    """
+    Executor for JSON-defined workflows.
+
+    Executes workflows defined in JSON format using the Kubiya SDK.
+    """
+
+    def _execute_workflow(self, context: WorkflowExecutionContext) -> Any:
+        """
+        Execute JSON workflow and return event stream.
+
+        Args:
+            context: Workflow execution context
+
+        Returns:
+            Iterable event stream from Kubiya SDK
+
+        Raises:
+            Exception: If workflow execution fails
+        """
+        logger.info(
+            "executing_json_workflow",
+            workflow_name=context.workflow_config.name,
+            runner=context.workflow_config.runner,
+            execution_id=context.execution_id[:8]
+        )
+
+        # Execute workflow using Kubiya SDK
+        response = self.kubiya_client.execute_workflow(
+            workflow_definition=context.workflow_config.definition,
+            parameters=context.workflow_config.parameters,
+            stream=True
+        )
+
+        return response

control_plane_api/worker/services/workflow_executor/executors/python_executor.py

@@ -0,0 +1,50 @@
+"""
+Python DSL workflow executor.
+
+This module implements workflow execution for Python DSL workflows.
+"""
+
+import structlog
+from typing import Any
+
+from .base import BaseWorkflowExecutor
+from ..models import WorkflowExecutionContext
+
+logger = structlog.get_logger()
+
+
+class PythonWorkflowExecutor(BaseWorkflowExecutor):
+    """
+    Executor for Python DSL workflows.
+
+    Executes workflows defined using Python DSL using the Kubiya SDK.
+    """
+
+    def _execute_workflow(self, context: WorkflowExecutionContext) -> Any:
+        """
+        Execute Python DSL workflow and return event stream.
+
+        Args:
+            context: Workflow execution context
+
+        Returns:
+            Iterable event stream from Kubiya SDK
+
+        Raises:
+            Exception: If workflow execution fails
+        """
+        logger.info(
+            "executing_python_dsl_workflow",
+            workflow_name=context.workflow_config.name,
+            runner=context.workflow_config.runner,
+            execution_id=context.execution_id[:8]
+        )
+
+        # Execute Python DSL workflow using Kubiya SDK
+        response = self.kubiya_client.execute_workflow(
+            workflow_definition=context.workflow_config.definition,
+            parameters=context.workflow_config.parameters,
+            stream=True
+        )
+
+        return response