gobby 0.2.9__py3-none-any.whl → 0.2.11__py3-none-any.whl

gobby/workflows/pipeline_executor.py

@@ -0,0 +1,801 @@
+"""Pipeline executor for running typed pipeline workflows."""
+
+from __future__ import annotations
+
+import json
+import logging
+import secrets
+import shlex
+from typing import TYPE_CHECKING, Any
+
+from gobby.workflows.pipeline_state import (
+    ApprovalRequired,
+    ExecutionStatus,
+    PipelineExecution,
+    StepStatus,
+)
+from gobby.workflows.safe_evaluator import SafeExpressionEvaluator
+
+if TYPE_CHECKING:
+    from gobby.storage.database import DatabaseProtocol
+    from gobby.storage.pipelines import LocalPipelineExecutionManager
+    from gobby.workflows.definitions import PipelineDefinition
+    from gobby.workflows.templates import TemplateEngine
+
+logger = logging.getLogger(__name__)
+
+
+# Type alias for event callback
+PipelineEventCallback = Any  # Callable[[str, str, dict], Awaitable[None]]
+
+
+class PipelineExecutor:
+    """Executor for pipeline workflows with typed data flow between steps.
+
+    Handles:
+    - Creating and tracking execution records
+    - Iterating through steps in order
+    - Building context with inputs and step outputs
+    - Executing exec commands, prompts, and nested pipelines
+    - Webhook notifications
+    - WebSocket event broadcasting for real-time updates
+    """
+
+    def __init__(
+        self,
+        db: DatabaseProtocol,
+        execution_manager: LocalPipelineExecutionManager,
+        llm_service: Any,
+        template_engine: TemplateEngine | None = None,
+        webhook_notifier: Any | None = None,
+        loader: Any | None = None,
+        event_callback: PipelineEventCallback | None = None,
+    ):
+        """Initialize the pipeline executor.
+
+        Args:
+            db: Database connection for transactions
+            execution_manager: Manager for pipeline execution records
+            llm_service: LLM service for prompt steps
+            template_engine: Optional template engine for variable substitution
+            webhook_notifier: Optional notifier for webhook callbacks
+            loader: Optional workflow loader for nested pipelines
+            event_callback: Optional async callback for broadcasting events.
+                           Signature: async def callback(event: str, execution_id: str, **kwargs)
+        """
+        self.db = db
+        self.execution_manager = execution_manager
+        self.llm_service = llm_service
+        self.template_engine = template_engine
+        self.webhook_notifier = webhook_notifier
+        self.loader = loader
+        self.event_callback = event_callback
+
+    async def _emit_event(self, event: str, execution_id: str, **kwargs: Any) -> None:
+        """Emit a pipeline event via the callback if configured.
+
+        Args:
+            event: Event type (pipeline_started, step_completed, etc.)
+            execution_id: Pipeline execution ID
+            **kwargs: Additional event data
+        """
+        if self.event_callback:
+            try:
+                await self.event_callback(event, execution_id, **kwargs)
+            except Exception as e:
+                logger.warning(f"Failed to emit pipeline event {event}: {e}")
+
+    async def execute(
+        self,
+        pipeline: PipelineDefinition,
+        inputs: dict[str, Any],
+        project_id: str,
+        execution_id: str | None = None,
+        session_id: str | None = None,
+    ) -> PipelineExecution:
+        """Execute a pipeline workflow.
+
+        Args:
+            pipeline: The pipeline definition to execute
+            inputs: Input values for the pipeline
+            project_id: Project context for the execution
+            execution_id: Optional existing execution ID (for resuming)
+            session_id: Optional session that triggered the execution
+
+        Returns:
+            The completed PipelineExecution record
+        """
+        # 1. Create or load execution record
+        if execution_id:
+            execution = self.execution_manager.get_execution(execution_id)
+            if not execution:
+                raise ValueError(f"Execution {execution_id} not found")
+        else:
+            execution = self.execution_manager.create_execution(
+                pipeline_name=pipeline.name,
+                inputs_json=json.dumps(inputs),
+                session_id=session_id,
+            )
+
+        # 2. Update status to RUNNING
+        updated = self.execution_manager.update_execution_status(
+            execution_id=execution.id,
+            status=ExecutionStatus.RUNNING,
+        )
+        if updated:
+            execution = updated
+
+        # Emit pipeline_started event
+        await self._emit_event(
+            "pipeline_started",
+            execution.id,
+            pipeline_name=pipeline.name,
+            inputs=inputs,
+            step_count=len(pipeline.steps),
+        )
+
+        # 3. Build execution context
+        context: dict[str, Any] = {
+            "inputs": inputs,
+            "steps": {},  # Will hold step outputs as they complete
+        }
+
+        # Fetch existing steps if resuming
+        existing_steps = {}
+        if execution_id:
+            steps = self.execution_manager.get_steps_for_execution(execution_id)
+            existing_steps = {s.step_id: s for s in steps}
+
+        try:
+            # 4. Iterate through steps in order
+            for step in pipeline.steps:
+                # Check for existing execution
+                step_execution = existing_steps.get(step.id)
+
+                if step_execution:
+                    # If completed, load output into context and skip
+                    if step_execution.status == StepStatus.COMPLETED:
+                        logger.info(f"Skipping completed step {step.id}")
+                        output = None
+                        if step_execution.output_json:
+                            try:
+                                output = json.loads(step_execution.output_json)
+                            except json.JSONDecodeError:
+                                output = step_execution.output_json
+                        context["steps"][step.id] = {"output": output}
+                        continue
+
+                    # If skipped, just skip
+                    if step_execution.status == StepStatus.SKIPPED:
+                        logger.info(f"Skipping previously skipped step {step.id}")
+                        continue
+
+                    # If waiting approval, check if we should check gate again
+                    # If we are resuming, it might have been approved
+                    if step_execution.status == StepStatus.WAITING_APPROVAL:
+                        # If the step is still marked as waiting approval in DB,
+                        # checking the gate will just re-raise ApprovalRequired.
+                        # If it was approved, status should be COMPLETED.
+                        # So we can just proceed to check/execute.
+                        pass
+
+                # Create new step execution if not exists
+                if not step_execution:
+                    step_execution = self.execution_manager.create_step_execution(
+                        execution_id=execution.id,
+                        step_id=step.id,
+                        input_json=json.dumps(context) if context else None,
+                    )
+
+                # Check if step should run based on condition
+                if not self._should_run_step(step, context):
+                    # Skip this step
+                    self.execution_manager.update_step_execution(
+                        step_execution_id=step_execution.id,
+                        status=StepStatus.SKIPPED,
+                    )
+                    logger.info(f"Skipping step {step.id}: condition not met")
+
+                    # Emit step_skipped event
+                    await self._emit_event(
+                        "step_skipped",
+                        execution.id,
+                        step_id=step.id,
+                        step_name=getattr(step, "name", step.id),
+                        reason="condition not met",
+                    )
+                    continue
+
+                # Update step status to RUNNING
+                self.execution_manager.update_step_execution(
+                    step_execution_id=step_execution.id,
+                    status=StepStatus.RUNNING,
+                )
+
+                # Emit step_started event
+                await self._emit_event(
+                    "step_started",
+                    execution.id,
+                    step_id=step.id,
+                    step_name=getattr(step, "name", step.id),
+                )
+
+                # Check for approval gate
+                await self._check_approval_gate(step, execution, step_execution, pipeline)
+
+                # Execute the step
+                step_output = await self._execute_step(step, context, project_id)
+
+                # Store step output in context for subsequent steps
+                context["steps"][step.id] = {"output": step_output}
+
+                # Update step with output and mark completed
+                self.execution_manager.update_step_execution(
+                    step_execution_id=step_execution.id,
+                    status=StepStatus.COMPLETED,
+                    output_json=json.dumps(step_output) if step_output else None,
+                )
+
+                # Emit step_completed event
+                await self._emit_event(
+                    "step_completed",
+                    execution.id,
+                    step_id=step.id,
+                    step_name=getattr(step, "name", step.id),
+                    output=step_output,
+                )
+
+            # 5. Mark execution as completed
+            outputs = self._build_outputs(pipeline, context)
+            completed = self.execution_manager.update_execution_status(
+                execution_id=execution.id,
+                status=ExecutionStatus.COMPLETED,
+                outputs_json=json.dumps(outputs),
+            )
+            if completed:
+                execution = completed
+
+            # Emit pipeline_completed event
+            await self._emit_event(
+                "pipeline_completed",
+                execution.id,
+                pipeline_name=pipeline.name,
+                outputs=outputs,
+            )
+
+        except ApprovalRequired:
+            # Don't treat approval as an error - just re-raise
+            raise
+
+        except Exception as e:
+            logger.error(f"Pipeline execution failed: {e}", exc_info=True)
+            failed = self.execution_manager.update_execution_status(
+                execution_id=execution.id,
+                status=ExecutionStatus.FAILED,
+            )
+            if failed:
+                execution = failed
+
+            # Emit pipeline_failed event
+            await self._emit_event(
+                "pipeline_failed",
+                execution.id,
+                pipeline_name=pipeline.name,
+                error=str(e),
+            )
+            raise
+
+        return execution
+
+    async def _execute_step(
+        self,
+        step: Any,  # PipelineStep
+        context: dict[str, Any],
+        project_id: str,
+    ) -> Any:
+        """Execute a single pipeline step.
+
+        Args:
+            step: The step to execute
+            context: Current execution context with inputs and step outputs
+            project_id: Project context
+
+        Returns:
+            The step's output value
+        """
+        # Render any template variables in the step
+        rendered_step = self._render_step(step, context)
+
+        if step.exec:
+            # Execute shell command
+            return await self._execute_exec_step(rendered_step.exec, context)
+        elif step.prompt:
+            # Execute LLM prompt
+            return await self._execute_prompt_step(rendered_step.prompt, context)
+        elif step.invoke_pipeline:
+            # Execute nested pipeline
+            return await self._execute_nested_pipeline(
+                rendered_step.invoke_pipeline, context, project_id
+            )
+        else:
+            logger.warning(f"Step {step.id} has no action defined")
+            return None
+
+    def _render_step(self, step: Any, context: dict[str, Any]) -> Any:
+        """Render template variables in step fields.
+
+        Args:
+            step: The step to render
+            context: Context with variables for substitution
+
+        Returns:
+            Step with rendered fields
+        """
+        if not self.template_engine:
+            return step
+
+        template_engine = self.template_engine
+
+        import os
+        import re
+
+        # Build render context
+        render_context = {
+            "inputs": context.get("inputs", {}),
+            "steps": context.get("steps", {}),
+            "env": os.environ,
+        }
+
+        def render_string(s: str) -> str:
+            if not s:
+                return s
+            # Replace ${{ ... }} with {{ ... }} for Jinja2
+            # Use dotall to allow multi-line expressions
+            jinja_template = re.sub(r"\$\{\{(.*?)\}\}", r"{{\1}}", s, flags=re.DOTALL)
+
+            # If no changes (no ${{ }}), we might still want to run it through jinja
+            # if we wanted to support direct {{ }} syntax too.
+            # But the requirement highlights ${{ }}.
+            # If the user provides {{ }}, it will also be rendered by Jinja.
+
+            return template_engine.render(jinja_template, render_context)
+
+        # Create a copy of the step to avoid modifying the definition
+        rendered_step = step.model_copy()
+
+        try:
+            if rendered_step.exec:
+                rendered_step.exec = render_string(rendered_step.exec)
+
+            if rendered_step.prompt:
+                rendered_step.prompt = render_string(rendered_step.prompt)
+
+            # We could also render 'input' field if needed, but requirements mention exec and prompt
+
+        except Exception as e:
+            # If rendering fails, we log it but might want to let it bubble up
+            # or fail the step execution later.
+            # For now, let's allow the exception to bubble up so the step fails immediately.
+            raise ValueError(f"Failed to render step {step.id}: {e}") from e
+
+        return rendered_step
+
+    def _should_run_step(self, step: Any, context: dict[str, Any]) -> bool:
+        """Check if a step should run based on its condition.
+
+        Args:
+            step: The step to check
+            context: Current execution context
+
+        Returns:
+            True if step should run, False if it should be skipped
+        """
+        # No condition means always run
+        if not step.condition:
+            return True
+
+        try:
+            # Evaluate the condition using safe AST-based evaluator
+            # This avoids eval() security risks while supporting common expressions
+            eval_context = {
+                "inputs": context.get("inputs", {}),
+                "steps": context.get("steps", {}),
+            }
+            # Allow common helper functions for conditions
+            allowed_funcs: dict[str, Any] = {
+                "len": len,
+                "bool": bool,
+                "str": str,
+                "int": int,
+            }
+            evaluator = SafeExpressionEvaluator(eval_context, allowed_funcs)
+            return evaluator.evaluate(step.condition)
+        except Exception as e:
+            logger.warning(f"Condition evaluation failed for step {step.id}: {e}")
+            # Default to running the step if condition evaluation fails
+            return True
+
+    async def _check_approval_gate(
+        self,
+        step: Any,
+        execution: PipelineExecution,
+        step_execution: Any,
+        pipeline: Any,
+    ) -> None:
+        """Check if a step has an approval gate and handle it.
+
+        If the step requires approval, this method:
+        1. Generates a unique approval token
+        2. Updates the step and execution status to WAITING_APPROVAL
+        3. Calls the webhook notifier if configured
+        4. Raises ApprovalRequired to pause execution
+
+        Args:
+            step: The step to check for approval requirement
+            execution: The current pipeline execution record
+            step_execution: The current step execution record
+            pipeline: The pipeline definition (for webhook config)
+
+        Raises:
+            ApprovalRequired: If the step requires approval
+        """
+        # Check if step has approval gate
+        if not step.approval or not step.approval.required:
+            return
+
+        # Generate unique approval token
+        token = secrets.token_urlsafe(24)
+
+        # Get approval message
+        message = step.approval.message or f"Approval required for step '{step.id}'"
+
+        # Update step status to WAITING_APPROVAL and store token
+        self.execution_manager.update_step_execution(
+            step_execution_id=step_execution.id,
+            status=StepStatus.WAITING_APPROVAL,
+            approval_token=token,
+        )
+
+        # Update execution status to WAITING_APPROVAL
+        self.execution_manager.update_execution_status(
+            execution_id=execution.id,
+            status=ExecutionStatus.WAITING_APPROVAL,
+            resume_token=token,
+        )
+
+        # Call webhook notifier if configured
+        if self.webhook_notifier:
+            try:
+                await self.webhook_notifier.notify_approval_pending(
+                    execution_id=execution.id,
+                    step_id=step.id,
+                    token=token,
+                    message=message,
+                    pipeline=pipeline,
+                )
+            except Exception as e:
+                logger.warning(f"Failed to send approval webhook: {e}")
+
+        # Emit approval_required event
+        await self._emit_event(
+            "approval_required",
+            execution.id,
+            step_id=step.id,
+            step_name=getattr(step, "name", step.id),
+            message=message,
+            token=token,
+        )
+
+        # Raise to pause execution
+        raise ApprovalRequired(
+            execution_id=execution.id,
+            step_id=step.id,
+            token=token,
+            message=message,
+        )
+
+    async def approve(
+        self,
+        token: str,
+        approved_by: str | None = None,
+    ) -> PipelineExecution:
+        """Approve a pipeline execution that is waiting for approval.
+
+        Args:
+            token: The approval token from the ApprovalRequired exception
+            approved_by: Optional identifier of the approver
+
+        Returns:
+            The updated PipelineExecution record
+
+        Raises:
+            ValueError: If the token is invalid or not found
+        """
+        # Find the step by approval token
+        step = self.execution_manager.get_step_by_approval_token(token)
+        if not step:
+            raise ValueError(f"Invalid approval token: {token}")
+
+        # Mark step as approved
+        self.execution_manager.update_step_execution(
+            step_execution_id=step.id,
+            status=StepStatus.COMPLETED,
+            approved_by=approved_by,
+        )
+
+        # Get the execution
+        execution = self.execution_manager.get_execution(step.execution_id)
+        if not execution:
+            raise ValueError(f"Execution {step.execution_id} not found")
+
+        # Resume execution
+        if self.loader:
+            try:
+                pipeline = self.loader.load_pipeline(execution.pipeline_name)
+                if pipeline:
+                    inputs = {}
+                    if execution.inputs_json:
+                        try:
+                            inputs = json.loads(execution.inputs_json)
+                        except json.JSONDecodeError:
+                            pass
+
+                    # Execute (resume)
+                    execution = await self.execute(
+                        pipeline=pipeline,
+                        inputs=inputs,
+                        project_id=execution.project_id,
+                        execution_id=execution.id,
+                    )
+            except ApprovalRequired:
+                # Pipeline paused again for another approval - this is expected
+                # Refresh execution to get latest status
+                execution = self.execution_manager.get_execution(execution.id)
+                if not execution:
+                    raise ValueError(
+                        f"Execution {step.execution_id} not found after resume"
+                    ) from None
+            except Exception as e:
+                logger.error(f"Failed to resume execution after approval: {e}", exc_info=True)
+                # Don't fail the approval if resume fails, but log it
+        else:
+            logger.warning("No loader configured, cannot resume execution automatically")
+
+        return execution
+
+    async def reject(
+        self,
+        token: str,
+        rejected_by: str | None = None,
+    ) -> PipelineExecution:
+        """Reject a pipeline execution that is waiting for approval.
+
+        Args:
+            token: The approval token from the ApprovalRequired exception
+            rejected_by: Optional identifier of the rejector
+
+        Returns:
+            The updated PipelineExecution record
+
+        Raises:
+            ValueError: If the token is invalid or not found
+        """
+        # Find the step by approval token
+        step = self.execution_manager.get_step_by_approval_token(token)
+        if not step:
+            raise ValueError(f"Invalid approval token: {token}")
+
+        # Mark step as failed
+        error_msg = "Rejected"
+        if rejected_by:
+            error_msg += f" by {rejected_by}"
+
+        self.execution_manager.update_step_execution(
+            step_execution_id=step.id,
+            status=StepStatus.FAILED,
+            error=error_msg,
+        )
+
+        # Set execution status to CANCELLED
+        execution = self.execution_manager.update_execution_status(
+            execution_id=step.execution_id,
+            status=ExecutionStatus.CANCELLED,
+        )
+
+        if not execution:
+            raise ValueError(f"Execution {step.execution_id} not found")
+
+        return execution
+
+    async def _execute_exec_step(self, command: str, context: dict[str, Any]) -> dict[str, Any]:
+        """Execute a shell command step.
+
+        Commands are parsed using shlex.split and executed via create_subprocess_exec
+        to avoid shell injection vulnerabilities. The command string is treated as a
+        space-separated list of arguments, not a shell script.
+
+        Note: Pipeline commands are defined by the pipeline author, not end users.
+        This is a defense-in-depth measure.
+
+        Args:
+            command: The command to execute (space-separated arguments)
+            context: Execution context
+
+        Returns:
+            Dict with stdout, stderr, exit_code
+        """
+        import asyncio
+
+        logger.info(f"Executing command: {command}")
+
+        try:
+            # Parse command into arguments to avoid shell injection
+            args = shlex.split(command)
+            if not args:
+                return {
+                    "stdout": "",
+                    "stderr": "Empty command",
+                    "exit_code": 1,
+                }
+
+            # Run command without shell
+            proc = await asyncio.create_subprocess_exec(
+                *args,
+                stdout=asyncio.subprocess.PIPE,
+                stderr=asyncio.subprocess.PIPE,
+            )
+
+            stdout_bytes, stderr_bytes = await proc.communicate()
+
+            return {
+                "stdout": stdout_bytes.decode("utf-8", errors="replace"),
+                "stderr": stderr_bytes.decode("utf-8", errors="replace"),
+                "exit_code": proc.returncode or 0,
+            }
+
+        except Exception as e:
+            logger.error(f"Command execution failed: {e}", exc_info=True)
+            return {
+                "stdout": "",
+                "stderr": str(e),
+                "exit_code": 1,
+            }
+
+    async def _execute_prompt_step(self, prompt: str, context: dict[str, Any]) -> dict[str, Any]:
+        """Execute an LLM prompt step.
+
+        Args:
+            prompt: The prompt to send to the LLM
+            context: Execution context
+
+        Returns:
+            Dict with response text or error
+        """
+        logger.info("Executing prompt step")
+
+        try:
+            response = await self.llm_service.generate(prompt)
+            return {
+                "response": response,
+            }
+        except Exception as e:
+            logger.error(f"LLM prompt execution failed: {e}", exc_info=True)
+            return {
+                "response": "",
+                "error": str(e),
+            }
+
+    async def _execute_nested_pipeline(
+        self,
+        pipeline_name: str,
+        context: dict[str, Any],
+        project_id: str,
+    ) -> dict[str, Any]:
+        """Execute a nested pipeline.
+
+        Args:
+            pipeline_name: Name of the pipeline to invoke
+            context: Execution context (used as inputs)
+            project_id: Project context
+
+        Returns:
+            Dict with nested pipeline outputs
+        """
+        logger.info(f"Invoking nested pipeline: {pipeline_name}")
+
+        # Check if loader is available
+        if not self.loader:
+            logger.warning("No loader configured for nested pipeline execution")
+            return {
+                "pipeline": pipeline_name,
+                "error": "No loader configured for nested pipeline execution",
+            }
+
+        try:
+            # Load the nested pipeline
+            nested_pipeline = self.loader.load_pipeline(pipeline_name)
+
+            if not nested_pipeline:
+                return {
+                    "pipeline": pipeline_name,
+                    "error": f"Pipeline '{pipeline_name}' not found",
+                }
+
+            # Execute the nested pipeline recursively
+            # Use inputs from context as the nested pipeline's inputs
+            nested_inputs = context.get("inputs", {})
+
+            result = await self.execute(
+                pipeline=nested_pipeline,
+                inputs=nested_inputs,
+                project_id=project_id,
+            )
+
+            return {
+                "pipeline": pipeline_name,
+                "execution_id": result.id,
+                "status": result.status.value,
+            }
+
+        except Exception as e:
+            logger.error(f"Nested pipeline execution failed: {e}", exc_info=True)
+            return {
+                "pipeline": pipeline_name,
+                "error": str(e),
+            }
+
+    def _build_outputs(self, pipeline: Any, context: dict[str, Any]) -> dict[str, Any]:
+        """Build pipeline outputs from context.
+
+        Args:
+            pipeline: The pipeline definition
+            context: Final execution context
+
+        Returns:
+            Dict of output name -> value
+        """
+        outputs: dict[str, Any] = {}
+
+        for name, expr in pipeline.outputs.items():
+            if isinstance(expr, str) and expr.startswith("$"):
+                # Resolve $step.output reference
+                value = self._resolve_reference(expr, context)
+                outputs[name] = value
+            else:
+                outputs[name] = expr
+
+        return outputs
+
+    def _resolve_reference(self, ref: str, context: dict[str, Any]) -> Any:
+        """Resolve a $step.output reference from context.
+
+        Args:
+            ref: Reference string like "$step1.output" or "$step1.output.field"
+            context: Execution context
+
+        Returns:
+            The resolved value
+        """
+        import re
+
+        # Parse reference: $step_id.output[.field]
+        match = re.match(r"\$([a-zA-Z_][a-zA-Z0-9_]*)\.output(?:\.(.+))?", ref)
+        if not match:
+            return ref
+
+        step_id = match.group(1)
+        field_path = match.group(2)
+
+        # Get step output from context
+        step_data = context.get("steps", {}).get(step_id, {})
+        output = step_data.get("output")
+
+        if field_path and isinstance(output, dict):
+            # Navigate nested field path
+            for part in field_path.split("."):
+                if isinstance(output, dict):
+                    output = output.get(part)
+                else:
+                    break
+
+        return output