/headroom/pipeline.py

"""Canonical Headroom pipeline lifecycle and extension contracts."""

from __future__ import annotations

import importlib.metadata
import logging
from dataclasses import dataclass, field
from enum import Enum
from typing import Any, Protocol

log = logging.getLogger(__name__)

ENTRY_POINT_GROUP = "headroom.pipeline_extension"


class PipelineStage(str, Enum):
    """Stable lifecycle stages for the canonical Headroom pipeline."""

    SETUP = "setup"
    PRE_START = "pre_start"
    POST_START = "post_start"
    INPUT_RECEIVED = "input_received"
    INPUT_CACHED = "input_cached"
    INPUT_ROUTED = "input_routed"
    INPUT_COMPRESSED = "input_compressed"
    INPUT_REMEMBERED = "input_remembered"
    PRE_SEND = "pre_send"
    POST_SEND = "post_send"
    RESPONSE_RECEIVED = "response_received"


CANONICAL_PIPELINE_STAGES: tuple[PipelineStage, ...] = (
    PipelineStage.SETUP,
    PipelineStage.PRE_START,
    PipelineStage.POST_START,
    PipelineStage.INPUT_RECEIVED,
    PipelineStage.INPUT_CACHED,
    PipelineStage.INPUT_ROUTED,
    PipelineStage.INPUT_COMPRESSED,
    PipelineStage.INPUT_REMEMBERED,
    PipelineStage.PRE_SEND,
    PipelineStage.POST_SEND,
    PipelineStage.RESPONSE_RECEIVED,
)


@dataclass
class PipelineEvent:
    """Event emitted at a canonical pipeline stage.

    Extensions may mutate ``messages``, ``tools``, ``headers``, or ``metadata`` in
    place, or return a replacement ``PipelineEvent`` from ``on_pipeline_event``.
    """

    stage: PipelineStage
    operation: str
    request_id: str = ""
    provider: str = ""
    model: str = ""
    messages: list[dict[str, Any]] | None = None
    tools: list[dict[str, Any]] | None = None
    headers: dict[str, str] | None = None
    response: Any = None
    metadata: dict[str, Any] = field(default_factory=dict)


class PipelineExtension(Protocol):
    """Request lifecycle extension contract for the canonical pipeline."""

    def on_pipeline_event(self, event: PipelineEvent) -> PipelineEvent | None:
        """Handle a canonical pipeline event."""


def discover_pipeline_extensions() -> list[PipelineExtension]:
    """Load registered pipeline extensions from Python entry points."""

    discovered: list[PipelineExtension] = []
    try:
        entries = importlib.metadata.entry_points(group=ENTRY_POINT_GROUP)
    except Exception as exc:  # noqa: BLE001 - importlib metadata varies by runtime
        log.debug("pipeline extensions: entry-point enumeration failed: %s", exc)
        return discovered

    for entry in entries:
        try:
            extension = entry.load()
        except Exception as exc:  # noqa: BLE001 - third-party load failures are isolated
            log.warning("pipeline extension %r failed to load: %s", entry.name, exc)
            continue

        if isinstance(extension, type):
            try:
                extension = extension()
            except Exception as exc:  # noqa: BLE001
                log.warning("pipeline extension %r failed to initialize: %s", entry.name, exc)
                continue

        discovered.append(extension)

    return discovered


def summarize_routing_markers(transforms_applied: list[str]) -> list[str]:
    """Return the routed transform markers emitted by ContentRouter."""

    return [item for item in transforms_applied if item.startswith("router:")]


class PipelineExtensionManager:
    """Dispatch canonical pipeline events to configured extensions."""

    def __init__(
        self,
        *,
        hooks: Any = None,
        extensions: list[Any] | None = None,
        discover: bool = True,
    ) -> None:
        resolved: list[Any] = []
        if hooks is not None and callable(getattr(hooks, "on_pipeline_event", None)):
            resolved.append(hooks)
        if extensions:
            resolved.extend(extensions)
        if discover:
            resolved.extend(discover_pipeline_extensions())
        self._extensions = resolved

    @property
    def enabled(self) -> bool:
        return bool(self._extensions)

    def emit(
        self,
        stage: PipelineStage,
        *,
        operation: str,
        request_id: str = "",
        provider: str = "",
        model: str = "",
        messages: list[dict[str, Any]] | None = None,
        tools: list[dict[str, Any]] | None = None,
        headers: dict[str, str] | None = None,
        response: Any = None,
        metadata: dict[str, Any] | None = None,
    ) -> PipelineEvent:
        """Emit a canonical lifecycle event and return the final event state."""

        event = PipelineEvent(
            stage=stage,
            operation=operation,
            request_id=request_id,
            provider=provider,
            model=model,
            messages=messages,
            tools=tools,
            headers=headers,
            response=response,
            metadata=metadata or {},
        )

        for extension in self._extensions:
            handler = getattr(extension, "on_pipeline_event", None)
            if not callable(handler):
                continue
            try:
                updated = handler(event)
            except Exception as exc:  # noqa: BLE001 - preserve hook fail-open behavior
                log.warning(
                    "pipeline extension %r failed during %s: %s",
                    type(extension).__name__,
                    stage.value,
                    exc,
                )
                continue
            if isinstance(updated, PipelineEvent):
                event = updated

        return event