Workflow Extension Design — Loops, Branches, and Sub-Workflows for Swarm

exo/_internal/
├── expression.py          # Safe expression evaluator (~200 lines)
├── branch_node.py         # BranchNode for conditional routing (~150 lines)
├── loop_node.py           # LoopNode for iteration (~200 lines)
├── workflow_state.py      # WorkflowState shared dict (~80 lines)
└── ... (existing files unchanged)

┌──────────────────────────────────────────────────────────┐
│                         Swarm                             │
│  flow_order: [a, loop1, branch1, c]                       │
│                                                           │
│  _run_workflow() iterates flow_order:                      │
│    ┌─────┐   ┌──────────┐   ┌────────────┐   ┌─────┐     │
│    │  a  │──▶│ LoopNode │──▶│ BranchNode │──▶│  c  │     │
│    └─────┘   └──────────┘   └────────────┘   └─────┘     │
│                   │                │                       │
│        iterates inner       evaluates condition            │
│        agent N times        routes to true/false agent     │
│                                                           │
│  WorkflowState: shared mutable dict across all nodes       │
└──────────────────────────────────────────────────────────┘

# exo/_internal/workflow_state.py

from __future__ import annotations
from typing import Any

class WorkflowState:
    """Shared mutable state for workflow execution.

    All nodes in a workflow can read and write to this state dict.
    The state persists across the full workflow run.

    Args:
        initial: Optional initial state values.
    """

    def __init__(self, initial: dict[str, Any] | None = None) -> None:
        self._data: dict[str, Any] = dict(initial) if initial else {}

    def get(self, key: str, default: Any = None) -> Any:
        """Get a value from state."""
        return self._data.get(key, default)

    def set(self, key: str, value: Any) -> None:
        """Set a value in state."""
        self._data[key] = value

    def update(self, values: dict[str, Any]) -> None:
        """Merge values into state."""
        self._data.update(values)

    def to_dict(self) -> dict[str, Any]:
        """Return a shallow copy of the state dict."""
        return dict(self._data)

    def __contains__(self, key: str) -> bool:
        return key in self._data

    def __repr__(self) -> str:
        return f"WorkflowState({self._data!r})"

# In Swarm._run_workflow(), changed lines:
state = WorkflowState()
for agent_name in self.flow_order:
    agent = self.agents[agent_name]
    if getattr(agent, "accepts_state", False):
        result = await agent.run(current_input, state=state, ...)
    else:
        # Existing behavior unchanged
        result = await call_runner(agent, current_input, ...)
    current_input = result.output
    # Auto-populate state with agent outputs
    state.set(f"{agent_name}.output", result.output)

# exo/_internal/expression.py

import ast
from typing import Any

MAX_EXPRESSION_LENGTH = 500
MAX_AST_DEPTH = 10
MAX_COLLECTION_SIZE = 1000

SAFE_FUNCTIONS = {"len", "range", "str", "int", "float", "bool", "abs", "min", "max"}

class ExpressionError(Exception):
    """Raised for invalid or unsafe expressions."""

def evaluate(expression: str, variables: dict[str, Any] | None = None) -> Any:
    """Safely evaluate an expression string.

    Supports:
        - Boolean: and, or, not (also &&, ||, !)
        - Comparisons: ==, !=, <, >, <=, >=
        - Arithmetic: +, -, *, /, %, **
        - Subscripting: dict["key"], list[0]
        - Allowlisted functions: len(), str(), int(), etc.
        - Variable resolution from the variables dict

    Security:
        - AST-based parsing, no eval()
        - Blocks dunder access, imports, lambdas
        - Depth and length limits
        - No attribute access on arbitrary objects

    Args:
        expression: String expression to evaluate.
        variables: Variable context for name resolution.

    Returns:
        The evaluated result.

    Raises:
        ExpressionError: If the expression is invalid or unsafe.
    """

# exo/_internal/loop_node.py

from __future__ import annotations
from typing import Any
from exo.types import RunResult

class LoopNode:
    """Iterates an inner agent with three loop modes.

    Modes:
        count: Execute the inner agent N times.
        items: Iterate over an array (from state or literal).
        condition: Repeat while an expression evaluates to truthy.

    The inner agent receives the current iteration context as input.
    Outputs from each iteration are collected and joined.

    Args:
        name: Node name for the flow DSL.
        agent: The agent to execute on each iteration.
        count: Fixed iteration count (count mode).
        items: List or state key resolving to a list (items mode).
        condition: Expression string evaluated each iteration (condition mode).
        max_iterations: Safety limit (default 100).
        separator: String to join iteration outputs.

    Duck-type markers:
        is_loop = True
        accepts_state = True
    """

    def __init__(
        self,
        *,
        name: str,
        agent: Any,
        count: int | None = None,
        items: list[Any] | str | None = None,
        condition: str | None = None,
        max_iterations: int = 100,
        separator: str = "\n",
    ) -> None: ...

    is_loop = True
    accepts_state = True

    async def run(
        self,
        input: str,
        *,
        state: WorkflowState | None = None,
        messages: Sequence[Message] | None = None,
        provider: Any = None,
        max_retries: int = 3,
    ) -> RunResult:
        """Execute the loop.

        Count mode: runs agent `count` times.
        Items mode: resolves items list, runs agent once per item.
        Condition mode: evaluates expression before each iteration.

        Break control: if agent output contains "[BREAK]", loop terminates early.
        The "[BREAK]" marker is stripped from the final output.

        State updates per iteration:
            state["loop.index"] = current iteration index (0-based)
            state["loop.value"] = current item (items mode only)
            state["loop.output"] = previous iteration's output
        """

outputs = []
for i in range(count):
    state["loop.index"] = i
    state["loop.output"] = outputs[-1] if outputs else ""
    result = await call_runner(agent, current_input, ...)
    if "[BREAK]" in result.output:
        outputs.append(result.output.replace("[BREAK]", "").strip())
        break
    outputs.append(result.output)
    current_input = result.output  # chain iterations
return RunResult(output=separator.join(outputs))

i = 0
while i < max_iterations:
    if not evaluate(condition, state.to_dict()):
        break
    state["loop.index"] = i
    result = await call_runner(agent, current_input, ...)
    if "[BREAK]" in result.output:
        break
    i += 1

# exo/_internal/branch_node.py

from __future__ import annotations
from typing import Any
from exo.types import RunResult

class BranchNode:
    """Conditional routing between two agents.

    Evaluates a condition (expression string or callable) and routes
    to either `true_agent` or `false_agent`.

    Args:
        name: Node name for the flow DSL.
        condition: Expression string or callable (receives state dict, returns bool).
        true_agent: Agent to execute when condition is truthy.
        false_agent: Agent to execute when condition is falsy (optional).

    If false_agent is None and condition is falsy, returns RunResult
    with empty output (passthrough).

    Duck-type markers:
        is_branch = True
        accepts_state = True
    """

    def __init__(
        self,
        *,
        name: str,
        condition: str | Callable[[dict[str, Any]], bool],
        true_agent: Any,
        false_agent: Any | None = None,
    ) -> None: ...

    is_branch = True
    accepts_state = True

    async def run(
        self,
        input: str,
        *,
        state: WorkflowState | None = None,
        messages: Sequence[Message] | None = None,
        provider: Any = None,
        max_retries: int = 3,
    ) -> RunResult:
        """Evaluate condition and route to the appropriate agent.

        For string conditions, uses the safe expression evaluator with
        the current WorkflowState as variable context.

        For callable conditions, passes state.to_dict().
        """

if isinstance(condition, str):
    result = evaluate(condition, state.to_dict() if state else {})
else:
    result = condition(state.to_dict() if state else {})

if result:
    return await call_runner(true_agent, input, ...)
elif false_agent is not None:
    return await call_runner(false_agent, input, ...)
else:
    return RunResult(output=input)  # passthrough

# Enhancement to SwarmNode:

class SwarmNode:
    def __init__(
        self,
        *,
        swarm: Any,
        name: str | None = None,
        input_mapping: dict[str, str] | None = None,   # NEW
        output_mapping: dict[str, str] | None = None,   # NEW
        share_state: bool = False,                       # NEW
    ) -> None: ...

# No DSL changes needed — nodes are agents
loop = LoopNode(name="retry", agent=validator, count=3)
branch = BranchNode(
    name="router",
    condition='score > 0.8',
    true_agent=fast_path,
    false_agent=slow_path,
)
sub = SwarmNode(swarm=inner_swarm, name="sub_pipeline")

swarm = Swarm(
    agents=[preprocessor, loop, branch, sub, postprocessor],
    flow="preprocessor >> retry >> router >> sub_pipeline >> postprocessor",
)

async def _run_workflow(self, input, *, messages=None, provider=None, max_retries=3):
    state = WorkflowState()          # NEW: shared state
    current_input = input
    last_result = None

    for agent_name in self.flow_order:
        agent = self.agents[agent_name]

        # NEW: detect state-aware nodes
        if getattr(agent, "accepts_state", False):
            last_result = await agent.run(
                current_input,
                state=state,
                messages=messages,
                provider=provider,
                max_retries=max_retries,
            )
        elif getattr(agent, "is_group", False) or getattr(agent, "is_swarm", False):
            # Existing group/nested swarm behavior — UNCHANGED
            last_result = await agent.run(
                current_input, messages=messages,
                provider=provider, max_retries=max_retries,
            )
        else:
            # Existing agent behavior — UNCHANGED
            last_result = await call_runner(
                agent, current_input, messages=messages,
                provider=provider, max_retries=max_retries,
            )

        current_input = last_result.output
        state.set(f"{agent_name}.output", current_input)   # NEW: auto-populate

    assert last_result is not None
    return last_result