omnigent-ai--omnigent/omnigent/inner/pi_native_executor.py

"""Executor that bridges Omnigent messages into a native Pi TUI."""

from __future__ import annotations

import os
from collections.abc import AsyncIterator
from pathlib import Path
from typing import Any

from omnigent.inner.executor import (
    Executor,
    ExecutorConfig,
    ExecutorError,
    ExecutorEvent,
    Message,
    ToolSpec,
    TurnComplete,
)
from omnigent.inner.native_attachments import materialize_attachment
from omnigent.pi_native_bridge import (
    PI_NATIVE_BRIDGE_DIR_ENV_VAR,
    PI_NATIVE_REQUEST_SESSION_ID_ENV_VAR,
    enqueue_user_message,
    refresh_config_auth_headers,
)


class PiNativeExecutor(Executor):
    """
    Harness-side executor for ``omnigent pi`` web UI turns.

    The native Pi process is already running in the session terminal with
    the Omnigent Pi extension loaded. Each turn queues the latest user
    message into the bridge inbox; the extension consumes it and calls
    ``pi.sendUserMessage`` inside the TUI process.

    Policy enforcement for native Pi tool calls is handled entirely by the
    extension: on each ``tool_call`` event the extension POSTs to
    ``POST /v1/sessions/{sessionId}/policies/evaluate`` using the server URL
    and auth headers from its config file. This bypasses the turn-scoped
    ``_policy_evaluator`` round-trip (which would fail because
    ``_current_ctx`` is cleared before Pi ever processes the message) and
    instead routes directly through the same session-level HTTP endpoint that
    the Claude Code native hook uses.
    """

    def __init__(self, bridge_dir: Path | None = None) -> None:
        self._bridge_dir = bridge_dir or _bridge_dir_from_env()
        self._request_session_id = _request_session_id_from_env()

    def supports_streaming(self) -> bool:
        """:returns: ``False`` because output is emitted by the Pi extension."""
        return False

    def supports_live_message_queue(self) -> bool:
        """:returns: ``True`` because messages can be queued for the extension."""
        return True

    async def enqueue_session_message(self, session_key: str, content: Any) -> bool:
        """
        Queue a live steering message for the resident Pi extension.

        :param session_key: Adapter session key. Unused; the bridge is
            per conversation.
        :param content: User-supplied content.
        :returns: ``True`` when the message was queued.
        """
        del session_key
        text = _content_to_text(content, self._bridge_dir)
        if not text:
            return False
        self._refresh_auth_headers()
        enqueue_user_message(self._bridge_dir, text)
        return True

    async def run_turn(
        self,
        messages: list[Message],
        tools: list[ToolSpec],
        system_prompt: str,
        config: ExecutorConfig | None = None,
    ) -> AsyncIterator[ExecutorEvent]:
        """
        Queue the latest user message for Pi.

        :param messages: Conversation history in executor message shape.
        :param tools: Tool schemas from Omnigent. Ignored for now; native
            Pi owns its configured tool surface.
        :param system_prompt: System prompt from the agent spec. Ignored
            because the native Pi terminal controls its own prompt/settings.
        :param config: Per-turn executor config. Unused.
        :yields: :class:`TurnComplete` after the input was queued, or an
            :class:`ExecutorError` when no user text can be sent.
        """
        del tools, system_prompt, config
        text = _latest_user_text(messages, self._bridge_dir)
        if not text:
            yield ExecutorError(message="Pi native turn had no user text to send")
            return
        self._refresh_auth_headers()
        enqueue_user_message(self._bridge_dir, text)
        yield TurnComplete(response=None)

    def _refresh_auth_headers(self) -> None:
        """
        Re-mint the extension's bearer so a long session survives token expiry.

        The bearer baked into ``config.json`` at launch dies with the ~1h
        Databricks OAuth lifetime, and the resident extension re-reads the
        config per request (``freshAuthHeaders``). Refreshing it at each turn
        keeps its policy/MCP POSTs authenticated instead of failing closed —
        the same outcome the refresh-capable runtime auth and the native
        policy hooks already get via re-mint. Minting runs in-runner through
        the same factory those use. Best-effort: any failure (no factory in
        local/unauthenticated runs, transient SDK error) leaves the existing
        headers in place rather than blocking the turn.

        ponytail: refreshes per turn, so a *single* turn running past ~1h
        still risks expiry mid-turn; add a background refresh task keyed to
        the pi terminal if that case ever bites.
        """
        try:
            from omnigent.cli_auth import databricks_request_headers
            from omnigent.runner._entry import _make_auth_token_factory

            factory = _make_auth_token_factory()
            token = factory() if factory is not None else None
            if token:
                # Rebuild the FULL routing header set (not just the bearer) so the
                # per-turn refresh preserves the workspace / deployment routing
                # selectors baked at launch (see runner/app.py). A bearer-only
                # refresh would drop them and re-break routing after the first turn.
                refresh_config_auth_headers(
                    self._bridge_dir,
                    databricks_request_headers(
                        os.environ.get("RUNNER_SERVER_URL", "http://localhost:6767").rstrip("/"),
                        bearer_token=token,
                    ),
                )
        except Exception:  # noqa: BLE001 — best-effort refresh; never block a turn
            pass


def _bridge_dir_from_env() -> Path:
    """
    Resolve the native Pi bridge directory from harness spawn env.

    :returns: Bridge directory path.
    :raises RuntimeError: If the env var is missing.
    """
    raw = os.environ.get(PI_NATIVE_BRIDGE_DIR_ENV_VAR, "").strip()
    if not raw:
        raise RuntimeError(f"{PI_NATIVE_BRIDGE_DIR_ENV_VAR} is required for pi-native harness")
    return Path(raw)


def _request_session_id_from_env() -> str | None:
    """
    Resolve the Omnigent session id that requested this harness process.

    :returns: Omnigent session id, e.g. ``"conv_abc123"``, or ``None``.
    """
    raw = os.environ.get(PI_NATIVE_REQUEST_SESSION_ID_ENV_VAR, "").strip()
    return raw or None


def _latest_user_text(messages: list[Message], bridge_dir: Path) -> str:
    """
    Return the latest user text from executor messages.

    :param messages: Conversation history in executor message shape.
    :param bridge_dir: Bridge directory for materializing attachments.
    :returns: Plain text content, or ``""`` when no user text is present.
    """
    for message in reversed(messages):
        if message.get("role") == "user":
            return _content_to_text(message.get("content"), bridge_dir)
    return ""


def _content_to_text(content: Any, bridge_dir: Path) -> str:
    """
    Normalize executor content into plain text for Pi.

    Text blocks are extracted directly. Image/file blocks are materialized
    to the bridge directory and referenced by path so Pi can inspect them
    with its native filesystem tools.
    """
    if isinstance(content, str):
        return content
    if isinstance(content, list):
        attachment_lines: list[str] = []
        text_parts: list[str] = []
        for block in content:
            if not isinstance(block, dict):
                continue
            block_type = block.get("type", "")
            if block_type == "input_text":
                text = block.get("text")
                if isinstance(text, str):
                    text_parts.append(text)
            elif block_type in ("input_image", "input_file"):
                path = materialize_attachment(block, bridge_dir)
                if path is not None:
                    attachment_lines.append(f"[Attached: {path}]")
        return "\n\n".join([*attachment_lines, *text_parts])
    if content is None:
        return ""
    return str(content)