orchestrata_ultimate.orchestrator.pipeline_kernel

Reference › orchestrata_ultimate › orchestrator › pipeline_kernel

PipelineKernelClient is the in-process bridge that lets the WebSocket gateway drive the pipeline. The gateway server is its only importer (gateway module); it in turn runs a PipelineRunner with a fresh ArtifactStore and turns the run into PipelineEvent messages from kernel.pipeline_events.

send_packet accepts a CXP packet and starts a pipeline only when it is addressed to U:Cx (Concierge) with the subject intent; any other packet returns false, as does a client that has not been connected. The run is spawned as a tracked asyncio task per thread id. _execute_pipeline emits an intent event, runs the pipeline, then emits one stage event per stage (status, duration, artifact count, a shortened error) and a final complete event with the stage and artifact counts. _emit_event broadcasts straight to subscribed UI clients through an injected broadcast function and adds the producer identity; there is no Redis stream in this path.

A workflow control facade (spec 014) sits on the same machinery. start_workflow takes a real work type, pause_workflow cancels the in-flight task and marks the thread paused (completed stages are already checkpointed, so resume_workflow re-runs only the one in flight), cancel_workflow and workflow_status complete the set, and per-thread WorkflowState snapshots are the only extra state. _execute_pipeline_guarded ensures an exception in a started workflow still produces a failed complete event instead of leaving the thread running. close cancels all active tasks.

At a glance

Class diagram

Dependency diagram

API

Classes

PipelineKernelClient

Decorators: traces(spec='003', story='US-004', reqs=['REQ-003-US-004-100']) · Line 83

Gateway-compatible kernel client backed by the local PipelineRunner.

Functions

Source

Permalink: https://gitlab.com/MSamy/orchestrata-ultimate-architecture/-/blob/a4cbc2bdf/src/orchestrata_ultimate/orchestrator/pipeline_kernel.py

"""Pipeline Kernel Client — bridges Gateway → PipelineRunner.

Implements the same ``send_packet(CXPPacket)`` interface as the CXP client
so the gateway can dispatch intents from the VizInt UI to the local
PipelineRunner instead of requiring an out-of-process gRPC kernel.

When the gateway is running without a gRPC endpoint, it falls back to this
in-process client which runs the full agent pipeline directly.

Pipeline events (stage transitions) are constructed as SPEC-001
``PipelineEvent``s (``kernel.pipeline_events``) and broadcast directly to
subscribed VizInt UI clients via the gateway-injected ``broadcast``
coroutine -- no Redis stream. Cross-process delivery (out-of-process
kernel deployments) is the kernel's own ``PipelineEventBus``'s job,
forwarded by the gateway's ``_consume_pipeline_events()`` background task.

SPEC-014 US-014-005 (REQ-014-076..079) extends this client with workflow
CONTROL (start/pause/resume/cancel/status) as a facade over the SAME
machinery above -- no duplicate engine. Pause cancels the active task;
``PipelineRunner`` already checkpoints every completed stage (kernel
``PipelineCheckpoint``, wired via ``PipelineRunner._create_checkpoint``), so
resume re-enters ``PipelineRunner.resume(thread_id)`` from that checkpoint,
re-running only the stage that was in flight (D2). Workflow verbs are only
operable through this in-process client (D3) -- a pure gRPC-kernel
deployment never wires them.

Usage (in gateway server.py)::

    from orchestrata_ultimate.orchestrator.pipeline_kernel import PipelineKernelClient

    kernel_client = PipelineKernelClient(pubsub=pubsub, broadcast=broadcast)
    await kernel_client.connect()
    rpc.set_kernel_client(kernel_client)
"""

from __future__ import annotations

import asyncio
import contextlib
import logging
import uuid
from typing import TYPE_CHECKING, Any

from orchestrata_ultimate.core.base.traces import traces
from orchestrata_ultimate.core.packet import CXPPacket
from orchestrata_ultimate.kernel.pipeline_events import (
    PipelineEvent,
    PipelineEventType,
    build_cancelled_event,
    build_complete_event,
    build_intent_event,
    build_paused_event,
    build_resumed_event,
    build_stage_event,
    pipeline_topic,
)

if TYPE_CHECKING:
    from collections.abc import Awaitable, Callable

    from orchestrata_ultimate.gateway.pubsub import PubSubManager

    # gateway (layer 5) types used only for annotations here -- orchestrator
    # (layer 4) never imports gateway at module scope; every runtime use
    # below is a def-nested import (RULED D7, architecture_checker.py),
    # matching this file's existing lazy-import idiom for PipelineRunner/
    # ArtifactStore.
    from orchestrata_ultimate.gateway.workflow_types import (
        WorkflowState,
        WorkflowStatus,
    )

logger = logging.getLogger(__name__)

# SPEC-001/SPEC-003 convergence compat window (design doc D5): dual-emit
# legacy flat fields (merged into the canonical PipelineEvent payload) for
# exactly one deploy cycle so not-yet-reloaded VizInt tabs keep rendering.
# Removal trigger: the first deploy after this epic lands (Packet 5).
_EMIT_LEGACY_COMPAT = True


@traces(spec="003", story="US-004", reqs=["REQ-003-US-004-100"])
class PipelineKernelClient:
    """Gateway-compatible kernel client backed by the local PipelineRunner.

    Sends CXP packets (intents) to the agent pipeline and broadcasts SPEC-001
    ``PipelineEvent``s (``kernel.pipeline_events``) directly to subscribed UI
    clients in real-time -- no Redis stream; the kernel's own
    ``PipelineEventBus`` owns cross-process delivery (design doc D1). Also
    exposes SPEC-014 US-014-005's workflow control facade
    (start/pause/resume/cancel/status) over the exact same
    ``_active_pipelines``/``_emit_event`` machinery.
    """

    def __init__(
        self,
        pubsub: PubSubManager | None = None,
        broadcast: Callable[[str, dict[str, Any]], Awaitable[int]] | None = None,
        driver: str = "inprocess",
    ):
        self._pubsub = pubsub
        # When the gateway is in-process it injects its ``broadcast`` coroutine
        # so pipeline events are pushed straight to subscribed UI connections;
        # ``PubSubManager.publish`` is only a registry counter and never sends.
        self._broadcast = broadcast
        self._connected = False
        self._active_pipelines: dict[str, asyncio.Task] = {}
        # This client's own producer identity (VolkDriver Phase 0, SPEC-001
        # US-006) -- ``_emit_event`` resolves it onto the broadcast payload
        # for any event the caller left unattributed. "inprocess" by
        # default; a Phase 2 headless-harness deployment injects its own.
        self._driver = driver
        # SPEC-014 US-014-005 (D1): the ONLY new state a workflow-control
        # facade adds -- a per-thread status snapshot. No parallel pipeline
        # state machine; PipelineRunner/PipelineCheckpoint remain the single
        # source of truth for stage progress.
        self._workflow_states: dict[str, WorkflowState] = {}

    @traces(spec="003", story="US-004", reqs=["REQ-003-US-004-100"])
    async def connect(self) -> None:
        """Initialize connection (no-op for in-process client)."""
        self._connected = True
        logger.info("PipelineKernelClient connected (in-process)")

    @traces(spec="003", story="US-004", reqs=["REQ-003-US-004-100"])
    async def close(self) -> None:
        """Clean up active pipeline runs."""
        self._connected = False
        for task in self._active_pipelines.values():
            if not task.done():
                task.cancel()
        self._active_pipelines.clear()
        logger.info("PipelineKernelClient closed")

    def _spawn(self, thread_id: str, coro: Any) -> asyncio.Task:
        """Track *coro* as this thread's active task -- the ONE spawn point
        both ``send_packet`` and ``start_workflow`` use (SPEC-014 note), so
        pause/resume/cancel always find the right task in
        ``_active_pipelines``.
        """
        task = asyncio.create_task(coro)
        self._active_pipelines[thread_id] = task
        return task

    @traces(spec="003", story="US-004", reqs=["REQ-003-US-004-101"])
    async def send_packet(self, packet: CXPPacket) -> bool:
        """Dispatch a CXP packet to the agent pipeline.

        When the packet is addressed to U:Cx with subject "intent",
        runs the full PipelineRunner. Otherwise, returns False (not handled).

        Args:
            packet: CXP packet to dispatch.

        Returns:
            True if the packet was dispatched, False otherwise.
        """
        if not self._connected:
            logger.warning("PipelineKernelClient not connected")
            return False

        if packet.to_agent == "U:Cx" and packet.subject == "intent":
            thread_id = packet.thread_id or uuid.uuid4().hex
            self._spawn(thread_id, self._run_pipeline(packet, thread_id))
            logger.info("Pipeline dispatched: thread=%s", thread_id)
            return True

        logger.debug(
            "Packet not routed to pipeline: to=%s subject=%s",
            packet.to_agent,
            packet.subject,
        )
        return False

    @traces(spec="003", story="US-004", reqs=["REQ-003-US-004-102"])
    async def _run_pipeline(self, packet: CXPPacket, thread_id: str) -> None:
        """Run the pipeline and broadcast stage events.

        Args:
            packet: The incoming CXP intent packet.
            thread_id: Pipeline thread ID.
        """
        try:
            # ``packet.payload`` access is inside the try (rather than before
            # it, as it used to be) so a malformed/non-dict payload raises
            # into the ``except`` below and is reported as a graceful
            # ``pipeline.error``-fold COMPLETE event, instead of escaping as
            # an unhandled exception in this background task — which the
            # gateway/UI would never learn about.
            text = packet.payload.get("text", packet.payload.get("requirements", ""))
            await self._execute_pipeline(text, "feature", thread_id)
        except Exception as e:
            logger.exception("Pipeline failed: %s", e)
            await self._emit_event(
                build_complete_event(
                    thread_id,
                    success=False,
                    reason="error",
                    extra={"error": str(e)[:500]},
                )
            )

    async def _execute_pipeline(
        self, intent: str, work_type: str, thread_id: str
    ) -> None:
        """Run a fresh ``PipelineRunner`` for *intent* and emit its events.

        Factored out of ``_run_pipeline`` so ``start_workflow`` (SPEC-014
        REQ-014-076) can drive the SAME run path directly, with an actual
        ``work_type`` instead of ``send_packet``'s hardcoded ``"feature"``.
        Callers are responsible for their own error handling: this raises
        rather than swallowing, so ``_run_pipeline``'s except-block and
        ``start_workflow``'s spawned task (``_execute_pipeline_guarded``,
        below) each report a failed ``COMPLETE`` event their own way.
        """
        from orchestrata_ultimate.orchestrator.artifact_store import ArtifactStore
        from orchestrata_ultimate.orchestrator.pipeline_runner import PipelineRunner

        store = ArtifactStore()
        runner = PipelineRunner(artifact_store=store)

        # Emit INTENT event (folds the legacy pipeline.start name)
        await self._emit_event(build_intent_event(thread_id, intent, volk="U:Cx"))

        # Run the pipeline
        result = await runner.run(
            intent=intent, work_type=work_type, thread_id=thread_id
        )

        # Emit per-stage STAGE events
        for alias, sr in result.stages.items():
            await self._emit_event(
                build_stage_event(
                    thread_id,
                    alias,
                    sr.status.value,
                    sr.duration_ms,
                    len(sr.artifacts),
                    sr.error[:200] if sr.error else None,
                )
            )

        # Emit COMPLETE event
        await self._emit_event(
            build_complete_event(
                thread_id,
                success=result.successful(),
                reason=result.overall_status.value,
                extra={
                    "stages_completed": len(result.stages),
                    "total_artifacts": len(result.artifacts),
                },
            )
        )

    @traces(spec="003", story="US-004", reqs=["REQ-003-US-004-102"])
    async def _execute_pipeline_guarded(
        self, intent: str, work_type: str, thread_id: str
    ) -> None:
        """Guard ``_execute_pipeline`` for ``start_workflow``'s spawn site,
        mirroring ``_execute_resume``'s try/except shape below: unlike
        ``send_packet``/``_run_pipeline`` (which wrap ``_execute_pipeline``
        in their own except-block), ``start_workflow`` spawns the task
        directly (REQ-014-076), so without this guard an ``_execute_pipeline``
        exception propagates unretrieved into the background task -- no
        terminal ``COMPLETE`` event, thread stuck ``RUNNING`` forever.
        Reports a failed ``COMPLETE`` event instead of raising into the
        spawned task.
        """
        try:
            await self._execute_pipeline(intent, work_type, thread_id)
        except Exception as e:
            logger.exception("Workflow pipeline failed: %s", e)
            await self._emit_event(
                build_complete_event(
                    thread_id,
                    success=False,
                    reason="error",
                    extra={"error": str(e)[:500]},
                )
            )

    @traces(spec="014", story="US-014-005", reqs=["REQ-014-078"])
    async def _execute_resume(self, thread_id: str) -> None:
        """Re-enter ``PipelineRunner.resume(thread_id)`` (the existing
        checkpoint-driven resume, pipeline_runner.py:385) and emit its
        events, mirroring ``_execute_pipeline``'s shape. Reports a failed
        ``COMPLETE`` event on failure instead of raising into the spawned
        task.
        """
        from orchestrata_ultimate.orchestrator.artifact_store import ArtifactStore
        from orchestrata_ultimate.orchestrator.pipeline_runner import PipelineRunner

        try:
            store = ArtifactStore()
            runner = PipelineRunner(artifact_store=store)
            result = await runner.resume(thread_id)
            for alias, sr in result.stages.items():
                await self._emit_event(
                    build_stage_event(
                        thread_id,
                        alias,
                        sr.status.value,
                        sr.duration_ms,
                        len(sr.artifacts),
                        sr.error[:200] if sr.error else None,
                    )
                )
            await self._emit_event(
                build_complete_event(
                    thread_id,
                    success=result.successful(),
                    reason=result.overall_status.value,
                    extra={
                        "stages_completed": len(result.stages),
                        "total_artifacts": len(result.artifacts),
                    },
                )
            )
        except Exception as e:
            logger.exception("Pipeline resume failed: %s", e)
            await self._emit_event(
                build_complete_event(
                    thread_id,
                    success=False,
                    reason="error",
                    extra={"error": str(e)[:500]},
                )
            )

    # -- SPEC-014 US-014-005: workflow control facade ---------------------

    @traces(spec="014", story="US-014-005", reqs=["REQ-014-076"])
    async def start_workflow(
        self, intent: str, work_type: str = "feature", thread_id: str | None = None
    ) -> str:
        """workflow.start facade (REQ-014-076): spawn a fresh pipeline run
        through the SAME ``_spawn`` helper ``send_packet`` uses, registering
        it in ``_active_pipelines`` and this client's per-thread
        ``WorkflowState`` map so pause/resume/cancel/status can address it
        by ``thread_id`` without a CXP packet round-trip.

        Args:
            intent: Human intent text.
            work_type: WorkType string (default ``"feature"``).
            thread_id: Optional caller-supplied thread id; generated when
                omitted.

        Returns:
            The (possibly generated) ``thread_id``.
        """
        from orchestrata_ultimate.gateway.workflow_types import WorkflowStatus

        thread_id = thread_id or uuid.uuid4().hex
        self._spawn(
            thread_id,
            self._execute_pipeline_guarded(intent, work_type, thread_id),
        )
        self._set_state(thread_id, WorkflowStatus.RUNNING)
        logger.info("Workflow started: thread=%s work_type=%s", thread_id, work_type)
        return thread_id

    @traces(spec="014", story="US-014-005", reqs=["REQ-014-077"])
    async def pause_workflow(self, thread_id: str) -> WorkflowState:
        """workflow.pause facade (REQ-014-077, D2): cancel the in-flight
        task -- ``PipelineRunner`` has already checkpointed every COMPLETED
        stage via ``PipelineCheckpoint``, so nothing beyond the single
        stage that was mid-flight is lost; ``resume_workflow`` re-runs
        exactly that stage.

        Args:
            thread_id: The pipeline's thread id.

        Returns:
            The paused ``WorkflowState``.
        """
        from orchestrata_ultimate.gateway.workflow_types import WorkflowStatus

        await self._cancel_active_task(thread_id)
        state = self._set_state(thread_id, WorkflowStatus.PAUSED)
        await self._emit_event(build_paused_event(thread_id))
        return state

    @traces(spec="014", story="US-014-005", reqs=["REQ-014-078"])
    async def resume_workflow(self, thread_id: str) -> WorkflowState:
        """workflow.resume facade (REQ-014-078): spawn a fresh task around
        the existing checkpoint-driven ``PipelineRunner.resume(thread_id)``,
        reusing ``_spawn`` so the new task is tracked in
        ``_active_pipelines`` exactly like a fresh run.

        Args:
            thread_id: The pipeline's thread id.

        Returns:
            The resumed (RUNNING) ``WorkflowState``.

        Raises:
            ValueError: *thread_id* has no known workflow state (never
                started via ``start_workflow``, or already cleared) --
                ``WorkflowHandler`` maps this to RPCError -32602.
        """
        from orchestrata_ultimate.gateway.workflow_types import WorkflowStatus

        if thread_id not in self._workflow_states:
            raise ValueError(f"unknown workflow thread_id: {thread_id}")
        self._spawn_resume(thread_id)
        state = self._set_state(thread_id, WorkflowStatus.RUNNING)
        await self._emit_event(build_resumed_event(thread_id))
        return state

    @traces(spec="014", story="US-014-005", reqs=["REQ-014-079"])
    async def cancel_workflow(self, thread_id: str) -> WorkflowState:
        """workflow.cancel facade (REQ-014-079): cancel the active task.
        Idempotent -- a second cancel on an already-cancelled/unknown
        thread just re-confirms CANCELLED, no error.

        Args:
            thread_id: The pipeline's thread id.

        Returns:
            The cancelled ``WorkflowState``.
        """
        from orchestrata_ultimate.gateway.workflow_types import WorkflowStatus

        await self._cancel_active_task(thread_id)
        state = self._set_state(thread_id, WorkflowStatus.CANCELLED)
        await self._emit_event(build_cancelled_event(thread_id))
        return state

    @traces(spec="014", story="US-014-005", reqs=["REQ-014-076"])
    async def workflow_status(self, thread_id: str) -> WorkflowState:
        """Current ``WorkflowState`` for *thread_id* (REQ-014-076).

        Args:
            thread_id: The pipeline's thread id.

        Returns:
            The known ``WorkflowState``, or an ``UNKNOWN`` one when
            *thread_id* was never started (or was cleared).
        """
        from orchestrata_ultimate.gateway.workflow_types import (
            WorkflowState as _WorkflowState,
        )
        from orchestrata_ultimate.gateway.workflow_types import WorkflowStatus

        state = self._workflow_states.get(thread_id)
        if state is None:
            return _WorkflowState(
                thread_id=thread_id,
                status=WorkflowStatus.UNKNOWN,
                last_stage=None,
                updated_at=_now_iso(),
            )
        return state

    @traces(spec="014", story="US-014-005", reqs=["REQ-014-078"])
    def _spawn_resume(self, thread_id: str) -> None:
        """Spawn ``_execute_resume(thread_id)`` as a tracked background task."""
        self._spawn(thread_id, self._execute_resume(thread_id))

    @traces(
        spec="014",
        story="US-014-005",
        reqs=["REQ-014-076", "REQ-014-077", "REQ-014-078", "REQ-014-079"],
    )
    def _set_state(
        self, thread_id: str, status: WorkflowStatus, stage: str | None = None
    ) -> WorkflowState:
        """Record and return this thread's new ``WorkflowState``."""
        from orchestrata_ultimate.gateway.workflow_types import (
            WorkflowState as _WorkflowState,
        )

        state = _WorkflowState(
            thread_id=thread_id,
            status=status,
            last_stage=stage,
            updated_at=_now_iso(),
        )
        self._workflow_states[thread_id] = state
        return state

    @traces(spec="014", story="US-014-005", reqs=["REQ-014-077", "REQ-014-079"])
    async def _cancel_active_task(self, thread_id: str) -> None:
        """Cancel *thread_id*'s active task, if any, and await its
        cancellation -- shared by ``pause_workflow``/``cancel_workflow``."""
        task = self._active_pipelines.get(thread_id)
        if task is not None and not task.done():
            task.cancel()
            with contextlib.suppress(asyncio.CancelledError):
                await task

    async def _emit_event(self, event: PipelineEvent) -> None:
        """Broadcast a canonical SPEC-001 ``PipelineEvent`` to subscribed
        VizInt UI clients.

        Delivered over the gateway-injected ``broadcast`` coroutine on
        ``pipeline.<thread_id>`` -- the SAME topic the kernel's
        ``PipelineEventBus`` forwards onto via its Redis stream + gateway
        pump. Design doc D1: exactly one producer ever broadcasts directly
        (this client); the bus only writes the stream, never both -- so no
        event is ever delivered twice. No ``broadcast`` wired -> no-op;
        there is no Redis fallback here any more, cross-process delivery is
        the kernel bus's job.

        ``PubSubManager.publish`` is intentionally not used — it is only a
        registry counter and never sends to a websocket.

        During the SPEC-001/SPEC-003 convergence compat window (D5), legacy
        flat fields are merged into the canonical payload so not-yet-
        reloaded VizInt tabs keep rendering; see ``_legacy_compat_fields``.

        Resolves ``driver`` to this client's own producer identity
        (``self._driver``, "inprocess" by default) on the OUTPUT payload
        when the caller left it unset — the second production producer of
        the canonical pipeline payload, attributed the same way the kernel
        bus is (``kernel.pipeline_events.PipelineEventBus.emit``). Never
        mutates *event* itself; a caller-supplied driver is never
        overridden.
        """
        if self._broadcast is None:
            return

        payload = event.to_payload()
        payload["driver"] = payload["driver"] or self._driver
        if _EMIT_LEGACY_COMPAT:
            payload = {**payload, **_legacy_compat_fields(event)}

        try:
            await self._broadcast(pipeline_topic(event.thread_id), payload)
        except Exception:
            logger.debug("Pipeline broadcast failed", exc_info=True)


def _legacy_compat_fields(event: PipelineEvent) -> dict[str, Any]:
    """The SPEC-001/SPEC-003 convergence compat window (design doc D5): the
    legacy flat fields this client used to emit directly, folded from the
    canonical *event*. Additive-only merge target -- the colliding keys
    (``thread_id``, ``timestamp``) carry identical values on both shapes,
    so they are left to the canonical payload and not repeated here.

    Deleted in Packet 5 (one deploy cycle after this lands), together with
    ``_EMIT_LEGACY_COMPAT``.
    """
    if event.event_type == PipelineEventType.INTENT:
        return {
            "event": "pipeline.start",
            "stage": event.volk,
            "intent": event.data.get("intent"),
        }
    if event.event_type == PipelineEventType.STAGE:
        return {
            "event": "pipeline.stage",
            "stage": event.data.get("stage"),
            "status": event.data.get("status"),
            "duration_ms": event.data.get("duration_ms"),
            "artifact_count": event.data.get("artifact_count"),
            "error": event.data.get("error"),
        }
    if event.event_type == PipelineEventType.COMPLETE:
        # Two legacy names fold onto COMPLETE (mapping table, design §3.1):
        # the normal pipeline.complete (carries stages_completed) and the
        # error-path pipeline.error (carries error, no stages_completed).
        if "stages_completed" in event.data:
            return {
                "event": "pipeline.complete",
                "overall_status": event.data.get("reason"),
                "stages_completed": event.data.get("stages_completed"),
                "total_artifacts": event.data.get("total_artifacts"),
            }
        return {"event": "pipeline.error", "error": event.data.get("error")}
    if event.event_type == PipelineEventType.PAUSED:
        return {"event": "pipeline.paused"}
    if event.event_type == PipelineEventType.RESUMED:
        return {"event": "pipeline.resumed"}
    if event.event_type == PipelineEventType.CANCELLED:
        return {"event": "pipeline.cancelled"}
    return {}


def _now_iso() -> str:
    from datetime import UTC, datetime

    return datetime.now(UTC).isoformat()