pipeline/pipeline.py

"""Pipeline — concrete, composable, sequential step runner."""

from __future__ import annotations

import asyncio
import logging
import threading
import time
import warnings
from collections.abc import Callable, Iterable
from concurrent.futures import ThreadPoolExecutor
from typing import Any

from .branch import Branch, MergeStrategy
from .context import StepContext
from .errors import CancellationToken, PipelineCancelled, PipelineConfigError, PipelineOrderError, cancel_token_var
from .protocol import PipelineHook, SampleResult


# ---------------------------------------------------------------------------
# Per-step-class background executor registry
# ---------------------------------------------------------------------------

_executor_lock = threading.Lock()


def _get_class_executor(step_cls: type) -> ThreadPoolExecutor:
    """Return the class-level ThreadPoolExecutor for *step_cls*, creating it lazily.

    The executor is stored on the class itself (``step_cls._executor``) so it
    is shared across all pipeline instances.  ``max_workers`` defaults to 1
    if not declared on the class.
    """
    if not hasattr(step_cls, "_executor") or getattr(step_cls, "_executor") is None:
        with _executor_lock:
            # Double-checked locking
            if (
                not hasattr(step_cls, "_executor")
                or getattr(step_cls, "_executor") is None
            ):
                max_workers = getattr(step_cls, "max_workers", 1)
                setattr(
                    step_cls, "_executor", ThreadPoolExecutor(max_workers=max_workers)
                )
    return getattr(step_cls, "_executor")


# ---------------------------------------------------------------------------
# Pipeline
# ---------------------------------------------------------------------------


class Pipeline:
    """Ordered sequence of steps.  Satisfies StepProtocol — can be nested.

    Build via the fluent API::

        pipe = (
            Pipeline()
            .then(AgentStep())
            .then(EvaluateStep())
            .then(ReflectStep())   # ReflectStep.async_boundary = True
            .then(UpdateStep())
        )

    Fan-out across samples::

        results = pipe.run(samples, workers=4)

    ``requires`` and ``provides`` are inferred from the step chain and kept
    up-to-date as steps are added, so a ``Pipeline`` can itself be used as a
    step inside another pipeline without extra annotation.
    """

    def __init__(
        self,
        steps: list | None = None,
        hooks: list[PipelineHook] | None = None,
    ) -> None:
        self._steps: list = list(steps or [])
        self._hooks: list[PipelineHook] = list(hooks or [])
        self.requires, self.provides = self._infer_contracts(self._steps)
        self._validate_steps(self._steps)

        # Background thread tracking (per Pipeline instance)
        self._bg_threads: list[threading.Thread] = []
        self._bg_lock = threading.Lock()

    # ------------------------------------------------------------------
    # Hook helpers
    # ------------------------------------------------------------------

    def _fire_before(self, step_name: str, ctx: StepContext) -> None:
        for hook in self._hooks:
            try:
                hook.before_step(step_name, ctx)
            except Exception:
                logging.getLogger(__name__).exception(
                    "Hook %s.before_step raised — ignoring", type(hook).__name__
                )

    def _fire_after(self, step_name: str, ctx: StepContext) -> None:
        for hook in self._hooks:
            try:
                hook.after_step(step_name, ctx)
            except Exception:
                logging.getLogger(__name__).exception(
                    "Hook %s.after_step raised — ignoring", type(hook).__name__
                )

    # ------------------------------------------------------------------
    # Contract inference
    # ------------------------------------------------------------------

    @staticmethod
    def _infer_contracts(steps: list) -> tuple[frozenset, frozenset]:
        """Compute (requires, provides) for the full step chain.

        ``requires`` — fields the pipeline needs from the outside
                       (what its first steps need that no earlier inner
                       step provides).
        ``provides`` — union of everything any inner step writes.
        """
        provided_so_far: set[str] = set()
        external_requires: set[str] = set()
        for step in steps:
            step_requires = set(getattr(step, "requires", frozenset()))
            step_provides = set(getattr(step, "provides", frozenset()))
            external_requires |= step_requires - provided_so_far
            provided_so_far |= step_provides
        return frozenset(external_requires), frozenset(provided_so_far)

    # ------------------------------------------------------------------
    # Validation
    # ------------------------------------------------------------------

    @staticmethod
    def _validate_steps(steps: list) -> None:
        """Raise PipelineOrderError or PipelineConfigError for invalid wiring.

        Order check:
            If step B requires field X, and field X is produced by some step
            in the pipeline but that step appears *after* B, raise
            ``PipelineOrderError``.  Fields not produced by any step in the
            pipeline are treated as external inputs — no error.

        Config checks:
            - More than one ``async_boundary = True`` step in the same pipeline.
            - Any ``async_boundary = True`` step inside a Branch child.
            - Warning (not error) when ``async_boundary`` is set on a nested
              Pipeline (the boundary is ignored when the pipeline runs as a step).
        """
        # Pre-compute all fields ever produced internally
        all_provided_internally: set[str] = set()
        for step in steps:
            all_provided_internally |= set(getattr(step, "provides", frozenset()))

        provided_so_far: set[str] = set()
        boundary_count = 0

        for step in steps:
            step_requires = set(getattr(step, "requires", frozenset()))
            step_provides = set(getattr(step, "provides", frozenset()))

            # Ordering: field is produced internally but not yet available
            out_of_order = (step_requires & all_provided_internally) - provided_so_far
            if out_of_order:
                raise PipelineOrderError(
                    f"{type(step).__name__} requires {out_of_order!r} but these "
                    f"are produced by a later step — check step ordering."
                )

            provided_so_far |= step_provides

            # async_boundary: only one per pipeline
            if getattr(step, "async_boundary", False):
                boundary_count += 1
                if boundary_count > 1:
                    raise PipelineConfigError(
                        f"Only one async_boundary step is allowed per pipeline; "
                        f"{type(step).__name__} is a duplicate."
                    )

            # async_boundary inside a Branch child is forbidden
            if isinstance(step, Branch):
                for child in step.pipelines:
                    for child_step in getattr(child, "_steps", []):
                        if getattr(child_step, "async_boundary", False):
                            raise PipelineConfigError(
                                f"async_boundary is not allowed inside a Branch "
                                f"child (found on {type(child_step).__name__})."
                            )

            # Warn when async_boundary is set on a nested Pipeline (ignored)
            if isinstance(step, Pipeline) and getattr(step, "async_boundary", False):
                warnings.warn(
                    f"async_boundary declared on a nested Pipeline "
                    f"({type(step).__name__}) is ignored — the boundary only "
                    f"fires when the pipeline is used as a top-level runner.",
                    stacklevel=4,
                )

    # ------------------------------------------------------------------
    # Fluent builder
    # ------------------------------------------------------------------

    def then(self, step: object) -> "Pipeline":
        """Append *step* and return ``self`` for chaining."""
        new_steps = self._steps + [step]
        # Validate before mutating so errors are raised immediately
        self._validate_steps(new_steps)
        self._steps = new_steps
        self.requires, self.provides = self._infer_contracts(self._steps)
        return self

    def branch(
        self,
        *pipelines: object,
        merge: MergeStrategy | Any = MergeStrategy.RAISE_ON_CONFLICT,
    ) -> "Pipeline":
        """Append a Branch step and return ``self`` for chaining."""
        return self.then(Branch(*pipelines, merge=merge))

    # ------------------------------------------------------------------
    # __call__ — for use as a nested step
    # ------------------------------------------------------------------

    def __call__(self, ctx: StepContext) -> StepContext:
        """Run all steps sequentially (sync).

        When used as a nested step inside another pipeline, ``async_boundary``
        markers are **ignored** (a warning is already emitted at construction
        time).  All steps — sync and async — are executed to completion before
        returning.
        """
        for step in self._steps:
            if asyncio.iscoroutinefunction(step.__call__):
                # Run the coroutine in a new event loop (safe in non-async contexts)
                ctx = asyncio.run(step(ctx))
            elif isinstance(step, Branch):
                # Branch.__call__ is sync (ThreadPoolExecutor)
                ctx = step(ctx)
            else:
                ctx = step(ctx)
        return ctx

    # ------------------------------------------------------------------
    # Async_boundary helpers
    # ------------------------------------------------------------------

    def _find_boundary_index(self) -> int | None:
        """Return the index of the first async_boundary step, or None."""
        for i, step in enumerate(self._steps):
            if getattr(step, "async_boundary", False):
                return i
        return None

    # ------------------------------------------------------------------
    # Background execution
    # ------------------------------------------------------------------

    def _submit_background(
        self,
        ctx: StepContext,
        background_steps: list,
        result: SampleResult,
    ) -> None:
        """Run *background_steps* sequentially in a background thread.

        Each step is submitted to its own class-level executor so concurrency
        across samples is controlled by ``max_workers`` on the step class —
        independent of how many pipeline instances or background tails are
        running.

        ``result`` is mutated in-place when the tail completes (or fails).
        """

        def run_tail() -> None:
            current_ctx = ctx
            for step in background_steps:
                step_cls = type(step)
                executor = _get_class_executor(step_cls)
                try:
                    # Submit to per-step-class pool; block until slot is free
                    future = executor.submit(step, current_ctx)
                    current_ctx = future.result()
                except Exception as exc:
                    result.error = exc
                    result.failed_at = step_cls.__name__
                    result.output = None
                    return
            result.output = current_ctx

        t = threading.Thread(target=run_tail, daemon=True, name="pipeline-bg")
        with self._bg_lock:
            self._bg_threads.append(t)
        t.start()

    def wait_for_background(self, timeout: float | None = None) -> None:
        """Block until all background tasks submitted by this pipeline finish.

        Raises ``TimeoutError`` if *timeout* seconds elapse before all tasks
        complete.  Completed threads are removed from the tracking list.
        """
        with self._bg_lock:
            threads = list(self._bg_threads)

        deadline = None if timeout is None else time.monotonic() + timeout

        for t in threads:
            if deadline is not None:
                remaining = deadline - time.monotonic()
                if remaining <= 0:
                    raise TimeoutError(
                        "Background pipeline steps did not drain within timeout."
                    )
                t.join(timeout=remaining)
                if t.is_alive():
                    raise TimeoutError(
                        "Background pipeline steps did not drain within timeout."
                    )
            else:
                t.join()

        # Remove completed threads
        with self._bg_lock:
            self._bg_threads = [t for t in self._bg_threads if t.is_alive()]

    def background_stats(self) -> dict[str, int]:
        """Return a snapshot of background task progress.

        Returns a dict with ``active`` and ``completed`` counts.  Safe to
        call from any thread while the pipeline is running.
        """
        with self._bg_lock:
            threads = list(self._bg_threads)
        active = sum(1 for t in threads if t.is_alive())
        completed = len(threads) - active
        return {"active": active, "completed": completed}

    # ------------------------------------------------------------------
    # run() — sync entry point
    # ------------------------------------------------------------------

    def run(
        self,
        contexts: Iterable[StepContext],
        workers: int = 1,
        on_sample_done: Callable[[SampleResult], None] | None = None,
        cancel_token: CancellationToken | None = None,
    ) -> list[SampleResult]:
        """Process *contexts* through the pipeline (sync entry point).

        Each item must be a fully-initialized ``StepContext``.  The pipeline
        never wraps or re-creates contexts — it processes what it receives.

        Splits at the first ``async_boundary`` step:
        - Foreground steps run in the calling context (with up to ``workers``
          samples in parallel via a semaphore inside the event loop).
        - Background steps are submitted to per-step-class executors and run
          asynchronously.  Call ``wait_for_background()`` to block until they
          finish and ``SampleResult`` objects are fully populated.

        Every context produces exactly one ``SampleResult`` — nothing is
        dropped silently.

        Args:
            contexts: Input contexts to process.
            workers: Maximum number of contexts processed concurrently.
            on_sample_done: Optional callback invoked after each sample
                completes its foreground steps (or fails).  Receives the
                ``SampleResult``.  Must not block the event loop.
            cancel_token: Optional cancellation signal.  Checked before
                each step and each new sample.  Pass a fresh token per
                invocation; the pipeline object stays reusable.
        """
        return asyncio.run(
            self.run_async(
                contexts,
                workers=workers,
                on_sample_done=on_sample_done,
                cancel_token=cancel_token,
            )
        )

    # ------------------------------------------------------------------
    # run_async() — async entry point
    # ------------------------------------------------------------------

    async def run_async(
        self,
        contexts: Iterable[StepContext],
        workers: int = 1,
        on_sample_done: Callable[[SampleResult], None] | None = None,
        cancel_token: CancellationToken | None = None,
    ) -> list[SampleResult]:
        """Async entry point; use ``await pipe.run_async(contexts)`` from
        coroutine contexts (e.g. inside browser-use tasks).

        Args:
            contexts: Input contexts to process.
            workers: Maximum number of contexts processed concurrently.
            on_sample_done: Optional callback invoked after each sample
                completes its foreground steps (or fails).  Receives the
                ``SampleResult``.  Must not block the event loop.
            cancel_token: Optional cancellation signal.  Checked before
                each step and each new sample.
        """
        boundary_idx = self._find_boundary_index()
        if boundary_idx is None:
            foreground_steps = self._steps
            background_steps: list = []
        else:
            foreground_steps = self._steps[:boundary_idx]
            background_steps = self._steps[boundary_idx:]

        sem = asyncio.Semaphore(workers)

        async def process_one(ctx: StepContext) -> SampleResult:
            async with sem:
                result = SampleResult(
                    sample=ctx.sample, output=None, error=None, failed_at=None
                )
                last_step_name: str | None = None
                try:
                    for step in foreground_steps:
                        step_name = type(step).__name__

                        # Cancel check — before each step
                        if cancel_token is not None and cancel_token.is_cancelled:
                            result.error = PipelineCancelled(
                                f"Cancelled before {step_name}"
                            )
                            result.failed_at = step_name
                            if on_sample_done is not None:
                                on_sample_done(result)
                            return result

                        last_step_name = step_name
                        self._fire_before(step_name, ctx)

                        if asyncio.iscoroutinefunction(step.__call__):
                            ctx = await step(ctx)
                        elif hasattr(step, "__call_async__"):
                            ctx = await step.__call_async__(ctx)
                        else:
                            ctx = await asyncio.to_thread(step, ctx)

                        self._fire_after(step_name, ctx)
                except Exception as exc:
                    result.error = exc
                    result.failed_at = last_step_name
                    if on_sample_done is not None:
                        on_sample_done(result)
                    return result

                if background_steps:
                    # Fire and forget — result updated by background thread
                    self._submit_background(ctx, background_steps, result)
                else:
                    result.output = ctx

                if on_sample_done is not None:
                    on_sample_done(result)
                return result

        # Set the contextvar so code inside steps (e.g. LLM clients) can
        # read the cancel token without explicit parameter passing.
        _reset = cancel_token_var.set(cancel_token)
        try:
            return list(await asyncio.gather(*[process_one(c) for c in contexts]))
        finally:
            cancel_token_var.reset(_reset)