highway-dsl 0.0.2__py3-none-any.whl → 1.2.0__py3-none-any.whl

highway_dsl/workflow_dsl.py

@@ -1,11 +1,12 @@
 # workflow_dsl.py
-from typing import Any, Dict, List, Optional, Union, Callable, Type
-from enum import Enum
+from abc import ABC
+from collections.abc import Callable
 from datetime import datetime, timedelta
+from enum import Enum
+from typing import Any, Optional, Union
+
 import yaml
-import json
-from abc import ABC, abstractmethod
-from pydantic import BaseModel, Field, model_validator, ConfigDict
+from pydantic import BaseModel, ConfigDict, Field, model_validator
 
 
 class OperatorType(Enum):
@@ -16,6 +17,9 @@ class OperatorType(Enum):
     FOREACH = "foreach"
     SWITCH = "switch"
     TRY_CATCH = "try_catch"
+    WHILE = "while"
+    EMIT_EVENT = "emit_event"
+    WAIT_FOR_EVENT = "wait_for_event"
 
 
 class RetryPolicy(BaseModel):
@@ -27,38 +31,45 @@ class RetryPolicy(BaseModel):
 class TimeoutPolicy(BaseModel):
     timeout: timedelta = Field(..., description="Timeout duration")
     kill_on_timeout: bool = Field(
-        True, description="Whether to kill the task on timeout"
+        True, description="Whether to kill the task on timeout",
     )
 
 
 class BaseOperator(BaseModel, ABC):
     task_id: str
     operator_type: OperatorType
-    dependencies: List[str] = Field(default_factory=list)
-    retry_policy: Optional[RetryPolicy] = None
-    timeout_policy: Optional[TimeoutPolicy] = None
-    metadata: Dict[str, Any] = Field(default_factory=dict)
+    dependencies: list[str] = Field(default_factory=list)
+    retry_policy: RetryPolicy | None = None
+    timeout_policy: TimeoutPolicy | None = None
+    metadata: dict[str, Any] = Field(default_factory=dict)
+    description: str = Field(default="", description="Task description")
+    # Phase 3: Callback hooks
+    on_success_task_id: str | None = Field(None, description="Task to run on success")
+    on_failure_task_id: str | None = Field(None, description="Task to run on failure")
+    is_internal_loop_task: bool = Field(
+        default=False, exclude=True,
+    )  # Mark if task is internal to a loop
 
     model_config = ConfigDict(use_enum_values=True, arbitrary_types_allowed=True)
 
 
 class TaskOperator(BaseOperator):
     function: str
-    args: List[Any] = Field(default_factory=list)
-    kwargs: Dict[str, Any] = Field(default_factory=dict)
-    result_key: Optional[str] = None
+    args: list[Any] = Field(default_factory=list)
+    kwargs: dict[str, Any] = Field(default_factory=dict)
+    result_key: str | None = None
     operator_type: OperatorType = Field(OperatorType.TASK, frozen=True)
 
 
 class ConditionOperator(BaseOperator):
     condition: str
-    if_true: str
-    if_false: str
+    if_true: str | None
+    if_false: str | None
     operator_type: OperatorType = Field(OperatorType.CONDITION, frozen=True)
 
 
 class WaitOperator(BaseOperator):
-    wait_for: Union[timedelta, datetime, str]
+    wait_for: timedelta | datetime | str
     operator_type: OperatorType = Field(OperatorType.WAIT, frozen=True)
 
     @model_validator(mode="before")
@@ -73,7 +84,7 @@ class WaitOperator(BaseOperator):
                     data["wait_for"] = datetime.fromisoformat(wait_for.split(":", 1)[1])
         return data
 
-    def model_dump(self, **kwargs) -> Dict[str, Any]:
+    def model_dump(self, **kwargs: Any) -> dict[str, Any]:
         data = super().model_dump(**kwargs)
         wait_for = data["wait_for"]
         if isinstance(wait_for, timedelta):
@@ -84,44 +95,142 @@ class WaitOperator(BaseOperator):
 
 
 class ParallelOperator(BaseOperator):
-    branches: Dict[str, List[str]] = Field(default_factory=dict)
+    branches: dict[str, list[str]] = Field(default_factory=dict)
+    timeout: int | None = Field(None, description="Optional timeout in seconds for branch execution")
     operator_type: OperatorType = Field(OperatorType.PARALLEL, frozen=True)
 
 
 class ForEachOperator(BaseOperator):
     items: str
-    task_chain: List[str] = Field(default_factory=list)
+    loop_body: list[
+        Union[
+            TaskOperator,
+            ConditionOperator,
+            WaitOperator,
+            ParallelOperator,
+            "ForEachOperator",
+            "WhileOperator",
+            "EmitEventOperator",
+            "WaitForEventOperator",
+            "SwitchOperator",
+        ]
+    ] = Field(default_factory=list)
     operator_type: OperatorType = Field(OperatorType.FOREACH, frozen=True)
 
 
-class Workflow(BaseModel):
-    name: str
-    version: str = "1.0.0"
-    description: str = ""
-    tasks: Dict[
-        str,
+class WhileOperator(BaseOperator):
+    condition: str
+    loop_body: list[
         Union[
             TaskOperator,
             ConditionOperator,
             WaitOperator,
             ParallelOperator,
             ForEachOperator,
-        ],
+            "WhileOperator",
+            "EmitEventOperator",
+            "WaitForEventOperator",
+            "SwitchOperator",
+        ]
+    ] = Field(default_factory=list)
+    operator_type: OperatorType = Field(OperatorType.WHILE, frozen=True)
+
+
+class EmitEventOperator(BaseOperator):
+    """Phase 2: Emit an event that other workflows can wait for."""
+    event_name: str = Field(..., description="Name of the event to emit")
+    payload: dict[str, Any] = Field(default_factory=dict, description="Event payload data")
+    operator_type: OperatorType = Field(OperatorType.EMIT_EVENT, frozen=True)
+
+
+class WaitForEventOperator(BaseOperator):
+    """Phase 2: Wait for an external event with optional timeout."""
+    event_name: str = Field(..., description="Name of the event to wait for")
+    timeout_seconds: int | None = Field(None, description="Timeout in seconds (None = wait forever)")
+    operator_type: OperatorType = Field(OperatorType.WAIT_FOR_EVENT, frozen=True)
+
+
+class SwitchOperator(BaseOperator):
+    """Phase 4: Multi-branch switch/case operator."""
+    switch_on: str = Field(..., description="Expression to evaluate for switch")
+    cases: dict[str, str] = Field(default_factory=dict, description="Map of case values to task IDs")
+    default: str | None = Field(None, description="Default task ID if no case matches")
+    operator_type: OperatorType = Field(OperatorType.SWITCH, frozen=True)
+
+
+class Workflow(BaseModel):
+    name: str
+    version: str = "1.1.0"
+    description: str = ""
+    tasks: dict[
+        str,
+        TaskOperator | ConditionOperator | WaitOperator | ParallelOperator | ForEachOperator | WhileOperator | EmitEventOperator | WaitForEventOperator | SwitchOperator,
     ] = Field(default_factory=dict)
-    variables: Dict[str, Any] = Field(default_factory=dict)
-    start_task: Optional[str] = None
+    variables: dict[str, Any] = Field(default_factory=dict)
+    start_task: str | None = None
+
+    # Phase 1: Scheduling metadata
+    schedule: str | None = Field(None, description="Cron expression for scheduled execution")
+    start_date: datetime | None = Field(None, description="When the schedule becomes active")
+    catchup: bool = Field(False, description="Whether to backfill missed runs")
+    is_paused: bool = Field(False, description="Whether the workflow is paused")
+    tags: list[str] = Field(default_factory=list, description="Workflow categorization tags")
+    max_active_runs: int = Field(1, description="Maximum number of concurrent runs")
+    default_retry_policy: RetryPolicy | None = Field(None, description="Default retry policy for all tasks")
+
+    @model_validator(mode="before")
+    @classmethod
+    def validate_workflow_name_and_version(cls, data: Any) -> Any:
+        """Validate workflow name and version don't contain '__' (double underscore).
+
+        The double underscore is reserved as a separator for display purposes:
+        {workflow_name}__{version}__{step_name}
+
+        Workflow names must match: ^[a-z][a-z0-9_]*$ (lowercase, alphanumeric, single underscore)
+        Workflow versions must match: ^[a-zA-Z0-9._-]+$ (semver compatible)
+        """
+        import re
+
+        if isinstance(data, dict):
+            name = data.get("name", "")
+            version = data.get("version", "")
+
+            # Check for double underscore (reserved separator)
+            if "__" in name:
+                msg = f"Workflow name '{name}' cannot contain '__' (double underscore) - it's reserved as a separator"
+                raise ValueError(msg)
+
+            if "__" in version:
+                msg = f"Workflow version '{version}' cannot contain '__' (double underscore) - it's reserved as a separator"
+                raise ValueError(msg)
+
+            # Validate workflow name format
+            if name and not re.match(r"^[a-z][a-z0-9_]*$", name):
+                msg = f"Workflow name '{name}' must start with lowercase letter and contain only lowercase letters, digits, and single underscores"
+                raise ValueError(msg)
+
+            # Validate workflow version format (semver compatible)
+            if version and not re.match(r"^[a-zA-Z0-9._-]+$", version):
+                msg = f"Workflow version '{version}' must contain only alphanumeric characters, dots, hyphens, and underscores (semver compatible)"
+                raise ValueError(msg)
+
+        return data
 
     @model_validator(mode="before")
     @classmethod
     def validate_tasks(cls, data: Any) -> Any:
         if isinstance(data, dict) and "tasks" in data:
             validated_tasks = {}
-            operator_classes: Dict[str, Type[BaseOperator]] = {
+            operator_classes: dict[str, type[BaseOperator]] = {
                 OperatorType.TASK.value: TaskOperator,
                 OperatorType.CONDITION.value: ConditionOperator,
                 OperatorType.WAIT.value: WaitOperator,
                 OperatorType.PARALLEL.value: ParallelOperator,
                 OperatorType.FOREACH.value: ForEachOperator,
+                OperatorType.WHILE.value: WhileOperator,
+                OperatorType.EMIT_EVENT.value: EmitEventOperator,
+                OperatorType.WAIT_FOR_EVENT.value: WaitForEventOperator,
+                OperatorType.SWITCH.value: SwitchOperator,
             }
             for task_id, task_data in data["tasks"].items():
                 operator_type = task_data.get("operator_type")
@@ -129,24 +238,19 @@ class Workflow(BaseModel):
                     operator_class = operator_classes[operator_type]
                     validated_tasks[task_id] = operator_class.model_validate(task_data)
                 else:
-                    raise ValueError(f"Unknown operator type: {operator_type}")
+                    msg = f"Unknown operator type: {operator_type}"
+                    raise ValueError(msg)
             data["tasks"] = validated_tasks
         return data
 
     def add_task(
         self,
-        task: Union[
-            TaskOperator,
-            ConditionOperator,
-            WaitOperator,
-            ParallelOperator,
-            ForEachOperator,
-        ],
+        task: TaskOperator | ConditionOperator | WaitOperator | ParallelOperator | ForEachOperator | WhileOperator | EmitEventOperator | WaitForEventOperator | SwitchOperator,
     ) -> "Workflow":
         self.tasks[task.task_id] = task
         return self
 
-    def set_variables(self, variables: Dict[str, Any]) -> "Workflow":
+    def set_variables(self, variables: dict[str, Any]) -> "Workflow":
         self.variables.update(variables)
         return self
 
@@ -154,6 +258,42 @@ class Workflow(BaseModel):
         self.start_task = task_id
         return self
 
+    # Phase 1: Scheduling methods
+    def set_schedule(self, cron: str) -> "Workflow":
+        """Set the cron schedule for this workflow."""
+        self.schedule = cron
+        return self
+
+    def set_start_date(self, start_date: datetime) -> "Workflow":
+        """Set when the schedule becomes active."""
+        self.start_date = start_date
+        return self
+
+    def set_catchup(self, enabled: bool) -> "Workflow":
+        """Set whether to backfill missed runs."""
+        self.catchup = enabled
+        return self
+
+    def set_paused(self, paused: bool) -> "Workflow":
+        """Set whether the workflow is paused."""
+        self.is_paused = paused
+        return self
+
+    def add_tags(self, *tags: str) -> "Workflow":
+        """Add tags to the workflow."""
+        self.tags.extend(tags)
+        return self
+
+    def set_max_active_runs(self, count: int) -> "Workflow":
+        """Set maximum number of concurrent runs."""
+        self.max_active_runs = count
+        return self
+
+    def set_default_retry_policy(self, policy: RetryPolicy) -> "Workflow":
+        """Set default retry policy for all tasks."""
+        self.default_retry_policy = policy
+        return self
+
     def to_yaml(self) -> str:
         data = self.model_dump(mode="json", by_alias=True, exclude_none=True)
         return yaml.dump(data, default_flow_style=False)
@@ -161,6 +301,68 @@ class Workflow(BaseModel):
     def to_json(self) -> str:
         return self.model_dump_json(indent=2)
 
+    def to_mermaid(self) -> str:
+        """ convert to mermaid state diagram format """
+        lines = ["stateDiagram-v2"]
+        
+        all_dependencies = {dep for task in self.tasks.values() for dep in task.dependencies}
+        
+        for task_id, task in self.tasks.items():
+            # Add state with description for regular tasks
+            if task.description and not isinstance(task, (ForEachOperator, WhileOperator)):
+                lines.append(f'    state "{task.description}" as {task_id}')
+            
+            # Add dependencies
+            if not task.dependencies:
+                if self.start_task == task_id or not self.start_task:
+                    lines.append(f'    [*] --> {task_id}')
+            else:
+                for dep in task.dependencies:
+                    lines.append(f'    {dep} --> {task_id}')
+            
+            # Add transitions for conditional operator
+            if isinstance(task, ConditionOperator):
+                if task.if_true:
+                    lines.append(f'    {task_id} --> {task.if_true} : True')
+                if task.if_false:
+                    lines.append(f'    {task_id} --> {task.if_false} : False')
+
+            # Add composite state for parallel operator
+            if isinstance(task, ParallelOperator):
+                lines.append(f'    state {task_id} {{')
+                for i, branch in enumerate(task.branches):
+                    lines.append(f'        state "Branch {i+1}" as {branch}')
+                    if i < len(task.branches) - 1:
+                        lines.append('        --')
+                lines.append('    }')
+
+            # Add composite state for foreach operator
+            if isinstance(task, ForEachOperator):
+                lines.append(f'    state {task_id} {{')
+                for sub_task in task.loop_body:
+                    if sub_task.description:
+                        lines.append(f'        state "{sub_task.description}" as {sub_task.task_id}')
+                    else:
+                        lines.append(f'        {sub_task.task_id}')
+                lines.append('    }')
+
+            # Add composite state for while operator
+            if isinstance(task, WhileOperator):
+                lines.append(f'    state {task_id} {{')
+                for sub_task in task.loop_body:
+                    if sub_task.description:
+                        lines.append(f'        state "{sub_task.description}" as {sub_task.task_id}')
+                    else:
+                        lines.append(f'        {sub_task.task_id}')
+                lines.append('    }')
+
+            # End states
+            if task_id not in all_dependencies:
+                if not (isinstance(task, ConditionOperator) and (task.if_true or task.if_false)):
+                    lines.append(f'    {task_id} --> [*]')
+                
+        return "\n".join(lines)
+
     @classmethod
     def from_yaml(cls, yaml_str: str) -> "Workflow":
         data = yaml.safe_load(yaml_str)
@@ -172,66 +374,208 @@ class Workflow(BaseModel):
 
 
 class WorkflowBuilder:
-    def __init__(self, name: str, existing_workflow: Optional[Workflow] = None):
+    def __init__(
+        self,
+        name: str,
+        existing_workflow: Workflow | None = None,
+        parent: Optional["WorkflowBuilder"] = None,
+    ) -> None:
         if existing_workflow:
             self.workflow = existing_workflow
         else:
-            self.workflow = Workflow(name=name)
-        self._current_task: Optional[str] = None
+            self.workflow = Workflow(
+                name=name,
+                version="1.1.0",
+                description="",
+                tasks={},
+                variables={},
+                start_task=None,
+                schedule=None,
+                start_date=None,
+                catchup=False,
+                is_paused=False,
+                tags=[],
+                max_active_runs=1,
+                default_retry_policy=None,
+            )
+        self._current_task: str | None = None
+        self.parent = parent
+
+    def _add_task(
+        self,
+        task: TaskOperator | ConditionOperator | WaitOperator | ParallelOperator | ForEachOperator | WhileOperator | EmitEventOperator | WaitForEventOperator | SwitchOperator,
+        **kwargs: Any,
+    ) -> None:
+        dependencies = kwargs.get("dependencies", [])
+        if self._current_task and not dependencies:
+            dependencies.append(self._current_task)
+
+        task.dependencies = sorted(set(dependencies))
 
-    def task(self, task_id: str, function: str, **kwargs) -> "WorkflowBuilder":
-        task = TaskOperator(task_id=task_id, function=function, **kwargs)
-        if self._current_task:
-            task.dependencies.append(self._current_task)
         self.workflow.add_task(task)
-        self._current_task = task_id
+        self._current_task = task.task_id
+
+    def task(self, task_id: str, function: str, **kwargs: Any) -> "WorkflowBuilder":
+        task = TaskOperator(task_id=task_id, function=function, **kwargs)
+        self._add_task(task, **kwargs)
         return self
 
     def condition(
-        self, task_id: str, condition: str, if_true: str, if_false: str, **kwargs
+        self,
+        task_id: str,
+        condition: str,
+        if_true: Callable[["WorkflowBuilder"], "WorkflowBuilder"],
+        if_false: Callable[["WorkflowBuilder"], "WorkflowBuilder"],
+        **kwargs: Any,
     ) -> "WorkflowBuilder":
+        true_builder = if_true(WorkflowBuilder(f"{task_id}_true", parent=self))
+        false_builder = if_false(WorkflowBuilder(f"{task_id}_false", parent=self))
+
+        true_tasks = list(true_builder.workflow.tasks.keys())
+        false_tasks = list(false_builder.workflow.tasks.keys())
+
         task = ConditionOperator(
             task_id=task_id,
             condition=condition,
-            if_true=if_true,
-            if_false=if_false,
+            if_true=true_tasks[0] if true_tasks else None,
+            if_false=false_tasks[0] if false_tasks else None,
             **kwargs,
         )
-        if self._current_task:
-            task.dependencies.append(self._current_task)
-        self.workflow.add_task(task)
+
+        self._add_task(task, **kwargs)
+
+        for task_obj in true_builder.workflow.tasks.values():
+            # Only add the condition task as dependency, preserve original dependencies
+            if task_id not in task_obj.dependencies:
+                task_obj.dependencies.append(task_id)
+            self.workflow.add_task(task_obj)
+        for task_obj in false_builder.workflow.tasks.values():
+            # Only add the condition task as dependency, preserve original dependencies
+            if task_id not in task_obj.dependencies:
+                task_obj.dependencies.append(task_id)
+            self.workflow.add_task(task_obj)
+
         self._current_task = task_id
         return self
 
     def wait(
-        self, task_id: str, wait_for: Union[timedelta, datetime, str], **kwargs
+        self, task_id: str, wait_for: timedelta | datetime | str, **kwargs: Any,
     ) -> "WorkflowBuilder":
         task = WaitOperator(task_id=task_id, wait_for=wait_for, **kwargs)
-        if self._current_task:
-            task.dependencies.append(self._current_task)
-        self.workflow.add_task(task)
-        self._current_task = task_id
+        self._add_task(task, **kwargs)
         return self
 
     def parallel(
-        self, task_id: str, branches: Dict[str, List[str]], **kwargs
+        self,
+        task_id: str,
+        branches: dict[str, Callable[["WorkflowBuilder"], "WorkflowBuilder"]],
+        **kwargs: Any,
     ) -> "WorkflowBuilder":
-        task = ParallelOperator(task_id=task_id, branches=branches, **kwargs)
-        if self._current_task:
-            task.dependencies.append(self._current_task)
-        self.workflow.add_task(task)
+        branch_builders = {}
+        for name, branch_func in branches.items():
+            branch_builder = branch_func(
+                WorkflowBuilder(f"{task_id}_{name}", parent=self),
+            )
+            branch_builders[name] = branch_builder
+
+        branch_tasks = {
+            name: list(builder.workflow.tasks.keys())
+            for name, builder in branch_builders.items()
+        }
+
+        task = ParallelOperator(task_id=task_id, branches=branch_tasks, **kwargs)
+
+        self._add_task(task, **kwargs)
+
+        for builder in branch_builders.values():
+            for task_obj in builder.workflow.tasks.values():
+                # Only add the parallel task as dependency to non-internal tasks,
+                # preserve original dependencies
+                if (
+                    not getattr(task_obj, "is_internal_loop_task", False)
+                    and task_id not in task_obj.dependencies
+                ):
+                    task_obj.dependencies.append(task_id)
+                self.workflow.add_task(task_obj)
+
         self._current_task = task_id
         return self
 
     def foreach(
-        self, task_id: str, items: str, task_chain: List[str], **kwargs
+        self,
+        task_id: str,
+        items: str,
+        loop_body: Callable[["WorkflowBuilder"], "WorkflowBuilder"],
+        **kwargs: Any,
     ) -> "WorkflowBuilder":
+        # Create a temporary builder for the loop body.
+        temp_builder = WorkflowBuilder(f"{task_id}_loop", parent=self)
+        loop_builder = loop_body(temp_builder)
+        loop_tasks = list(loop_builder.workflow.tasks.values())
+
+        # Mark all loop body tasks as internal to prevent parallel dependency injection
+        for task_obj in loop_tasks:
+            task_obj.is_internal_loop_task = True
+
+        # Create the foreach operator
         task = ForEachOperator(
-            task_id=task_id, items=items, task_chain=task_chain, **kwargs
+            task_id=task_id,
+            items=items,
+            loop_body=loop_tasks,
+            **kwargs,
         )
-        if self._current_task:
-            task.dependencies.append(self._current_task)
-        self.workflow.add_task(task)
+
+        # Add the foreach task to workflow to establish initial dependencies
+        self._add_task(task, **kwargs)
+
+        # Add the foreach task as dependency to the FIRST task in the loop body
+        # and preserve the original dependency chain within the loop
+        if loop_tasks:
+            first_task = loop_tasks[0]
+            if task_id not in first_task.dependencies:
+                first_task.dependencies.append(task_id)
+
+            # Add all loop tasks to workflow
+            for task_obj in loop_tasks:
+                self.workflow.add_task(task_obj)
+
+        self._current_task = task_id
+        return self
+
+    def while_loop(
+        self,
+        task_id: str,
+        condition: str,
+        loop_body: Callable[["WorkflowBuilder"], "WorkflowBuilder"],
+        **kwargs: Any,
+    ) -> "WorkflowBuilder":
+        loop_builder = loop_body(WorkflowBuilder(f"{task_id}_loop", parent=self))
+        loop_tasks = list(loop_builder.workflow.tasks.values())
+
+        # Mark all loop body tasks as internal to prevent parallel dependency injection
+        for task_obj in loop_tasks:
+            task_obj.is_internal_loop_task = True
+
+        task = WhileOperator(
+            task_id=task_id,
+            condition=condition,
+            loop_body=loop_tasks,
+            **kwargs,
+        )
+
+        self._add_task(task, **kwargs)
+
+        # Fix: Only add the while task as dependency to the FIRST task in the loop body
+        # and preserve the original dependency chain within the loop
+        if loop_tasks:
+            first_task = loop_tasks[0]
+            if task_id not in first_task.dependencies:
+                first_task.dependencies.append(task_id)
+
+            # Add all loop tasks to workflow without modifying their dependencies further
+            for task_obj in loop_tasks:
+                self.workflow.add_task(task_obj)
+
         self._current_task = task_id
         return self
 
@@ -242,24 +586,106 @@ class WorkflowBuilder:
         backoff_factor: float = 2.0,
     ) -> "WorkflowBuilder":
         if self._current_task and isinstance(
-            self.workflow.tasks[self._current_task], TaskOperator
+            self.workflow.tasks[self._current_task], TaskOperator,
         ):
             self.workflow.tasks[self._current_task].retry_policy = RetryPolicy(
-                max_retries=max_retries, delay=delay, backoff_factor=backoff_factor
+                max_retries=max_retries, delay=delay, backoff_factor=backoff_factor,
             )
         return self
 
     def timeout(
-        self, timeout: timedelta, kill_on_timeout: bool = True
+        self, timeout: timedelta, kill_on_timeout: bool = True,
     ) -> "WorkflowBuilder":
         if self._current_task and isinstance(
-            self.workflow.tasks[self._current_task], TaskOperator
+            self.workflow.tasks[self._current_task], TaskOperator,
         ):
             self.workflow.tasks[self._current_task].timeout_policy = TimeoutPolicy(
-                timeout=timeout, kill_on_timeout=kill_on_timeout
+                timeout=timeout, kill_on_timeout=kill_on_timeout,
             )
         return self
 
+    # Phase 2: Event-based operators
+    def emit_event(self, task_id: str, event_name: str, **kwargs: Any) -> "WorkflowBuilder":
+        """Emit an event that other workflows can wait for."""
+        task = EmitEventOperator(task_id=task_id, event_name=event_name, **kwargs)
+        self._add_task(task, **kwargs)
+        return self
+
+    def wait_for_event(
+        self, task_id: str, event_name: str, timeout_seconds: int | None = None, **kwargs: Any,
+    ) -> "WorkflowBuilder":
+        """Wait for an external event with optional timeout."""
+        task = WaitForEventOperator(
+            task_id=task_id, event_name=event_name, timeout_seconds=timeout_seconds, **kwargs,
+        )
+        self._add_task(task, **kwargs)
+        return self
+
+    # Phase 3: Callback hooks (applies to current task)
+    def on_success(self, success_task_id: str) -> "WorkflowBuilder":
+        """Set the task to run when the current task succeeds."""
+        if self._current_task:
+            self.workflow.tasks[self._current_task].on_success_task_id = success_task_id
+        return self
+
+    def on_failure(self, failure_task_id: str) -> "WorkflowBuilder":
+        """Set the task to run when the current task fails."""
+        if self._current_task:
+            self.workflow.tasks[self._current_task].on_failure_task_id = failure_task_id
+        return self
+
+    # Phase 4: Switch operator
+    def switch(
+        self,
+        task_id: str,
+        switch_on: str,
+        cases: dict[str, str],
+        default: str | None = None,
+        **kwargs: Any,
+    ) -> "WorkflowBuilder":
+        """Multi-branch switch/case operator."""
+        task = SwitchOperator(
+            task_id=task_id, switch_on=switch_on, cases=cases, default=default, **kwargs,
+        )
+        self._add_task(task, **kwargs)
+        return self
+
+    # Phase 1: Scheduling methods (delegate to Workflow)
+    def set_schedule(self, cron: str) -> "WorkflowBuilder":
+        """Set the cron schedule for this workflow."""
+        self.workflow.set_schedule(cron)
+        return self
+
+    def set_start_date(self, start_date: datetime) -> "WorkflowBuilder":
+        """Set when the schedule becomes active."""
+        self.workflow.set_start_date(start_date)
+        return self
+
+    def set_catchup(self, enabled: bool) -> "WorkflowBuilder":
+        """Set whether to backfill missed runs."""
+        self.workflow.set_catchup(enabled)
+        return self
+
+    def set_paused(self, paused: bool) -> "WorkflowBuilder":
+        """Set whether the workflow is paused."""
+        self.workflow.set_paused(paused)
+        return self
+
+    def add_tags(self, *tags: str) -> "WorkflowBuilder":
+        """Add tags to the workflow."""
+        self.workflow.add_tags(*tags)
+        return self
+
+    def set_max_active_runs(self, count: int) -> "WorkflowBuilder":
+        """Set maximum number of concurrent runs."""
+        self.workflow.set_max_active_runs(count)
+        return self
+
+    def set_default_retry_policy(self, policy: RetryPolicy) -> "WorkflowBuilder":
+        """Set default retry policy for all tasks."""
+        self.workflow.set_default_retry_policy(policy)
+        return self
+
     def build(self) -> Workflow:
         if not self.workflow.start_task and self.workflow.tasks:
             self.workflow.start_task = next(iter(self.workflow.tasks.keys()))