"""PiExecutor: run agents through the Pi coding agent's RPC mode.

Spawns Pi (``pi --mode rpc``) as a subprocess and communicates via a JSONL
protocol over stdin/stdout.  Pi manages its own agent loop, tool execution,
context window, and compaction internally.  This executor translates the Pi
event stream into Omnigent ExecutorEvents.

Omnigent tools are bridged into Pi via a generated JavaScript extension
that registers each tool with ``pi.registerTool()``.  Tool execution is
proxied over a local TCP socket to the Omnigent Python process, so
policies, history recording, sub-agents, runtime, and all other Omnigent
features work exactly as they do with the Claude SDK and Codex harnesses.

When ``os_env`` is set, the Pi subprocess is wrapped in the same
sandbox used for other harnesses.

Databricks support: a temporary ``models.json`` is generated with three
providers (OpenAI Responses for GPT, Anthropic Messages for Claude, and
OpenAI Completions for others) and ``PI_CODING_AGENT_DIR`` is set so Pi
picks it up.

Requirements:
    The ``pi`` CLI must be installed and on PATH.

Environment (Databricks):
    DATABRICKS_CONFIG_PROFILE — optional Databricks profile selector
    ~/.databrickscfg          — host + token profile for workspace access
    (or ~/.databrickscfg with a profile containing host + token)
"""

from __future__ import annotations

import asyncio
import base64
import contextlib
import hmac
import json
import logging
import os
import pathlib
import secrets
import shutil
import subprocess
import tempfile
from asyncio import Queue, Task
from collections.abc import AsyncIterator, Awaitable, Callable, Sequence
from dataclasses import dataclass, field
from typing import Any, TypeAlias
from urllib.parse import urlparse as _urlparse

from omnigent.inner.native_attachments import parse_data_uri
from omnigent.llms._usage_observer import notify_from_dict as _notify_usage_from_dict
from omnigent.onboarding.databricks_config import DATABRICKS_CLAUDE_DEFAULT_MODEL
from omnigent.runner.identity import OMNIGENT_SESSION_ENV_VAR
from omnigent.spec.types import RetryPolicy

from ._subprocess_lifecycle import close_subprocess_transport
from .databricks_executor import _read_databrickscfg
from .datamodel import OSEnvSandboxSpec, OSEnvSpec
from .executor import (
    Executor,
    ExecutorConfig,
    ExecutorError,
    ExecutorEvent,
    Message,
    ReasoningChunk,
    TextChunk,
    ToolCallComplete,
    ToolCallRequest,
    ToolCallStatus,
    ToolSpec,
    TurnComplete,
)

logger = logging.getLogger(__name__)

# Each line of Pi's JSONL output; the event schema is owned by the Pi CLI
# not us, and varies across subcommands (response ack, message_update,
# tool_execution_start/end, agent_end, message_end, etc.).
CodexEvent: TypeAlias = dict[str, Any]  # type: ignore[explicit-any]


def _fetch_shell_command_token(command: str) -> str | None:
    """Run an auth helper command and return its stdout token.

    :param command: Shell command that prints a bearer token.
    :returns: The stripped token, or ``None`` when the command fails or
        prints no token.
    """
    result = subprocess.run(
        ["sh", "-c", command],
        check=False,
        capture_output=True,
        text=True,
    )
    token = result.stdout.strip()
    if result.returncode != 0 or not token:
        logger.debug("Pi Databricks auth command failed: %s", result.stderr.strip())
        return None
    return token


# Tool-server callback provided by ``Session._wire_sdk_executor``. Invoked
# with a tool name and argument dict; may return the result dict directly
# or a coroutine/future yielding one.
ToolExecutor: TypeAlias = Callable[  # type: ignore[explicit-any]
    [str, dict[str, Any]],
    Awaitable[dict[str, Any]] | dict[str, Any],
]

# Native-tool policy gate wired by :class:`PiExecutor`. Invoked with a native
# (non-bridged) tool name + argument dict; returns ``{"block": bool, "reason":
# str}``. Pi's native tools (e.g. ``read``, enabled for skill loading) execute
# inside the Pi process and never traverse the bridged ``/mcp`` path, so the
# ``tool_call`` extension hook routes them here for a TOOL_CALL policy verdict.
NativePolicyGate: TypeAlias = Callable[  # type: ignore[explicit-any]
    [str, dict[str, Any]],
    Awaitable[dict[str, Any]] | dict[str, Any],
]

# Plain JSON value — recursive union used by ``_check_blocked`` to walk
# parsed Pi tool-result payloads.
JsonValue: TypeAlias = None | bool | int | float | str | list["JsonValue"] | dict[str, "JsonValue"]

# ---------------------------------------------------------------------------
# TCP tool server — serves Omnigent tools to the Pi extension
# ---------------------------------------------------------------------------


def _safe_dumps(response: dict[str, Any], req_id: str | None) -> str:  # type: ignore[explicit-any]
    """Serialize a tool-server response, never raising on bad payloads.

    Tool callbacks may return values ``json.dumps`` can't encode (a
    ``datetime``, ``set``, ``bytes``, custom object, ...). Encoding happens
    on the response path *outside* :meth:`_ToolServer._execute`'s try, so an
    unguarded ``json.dumps`` would propagate and the connection would close
    with zero bytes written — leaving the JS ``callTool`` promise pending and
    hanging the entire Pi turn. Mirrors codex's ``_result_text`` guard: on a
    serialization failure, fall back to an ``{"error": ...}`` envelope so the
    client always receives a valid frame.

    :param response: The response dict to serialize.
    :param req_id: The originating request id, echoed back on the error
        envelope so the client correlates the failure with its call.
    :returns: A compact JSON string (no trailing newline).
    """
    try:
        return json.dumps(response, separators=(",", ":"))
    except (TypeError, ValueError) as exc:
        # Stringify ``req_id`` so the fallback envelope itself can never raise
        # on a non-serializable id (the caller passes a ``str`` today, but the
        # guard must hold for any future caller). ``None`` stays ``None`` so the
        # client still sees a null id rather than the literal "None".
        safe_id: str | None = req_id if req_id is None or isinstance(req_id, str) else str(req_id)
        return json.dumps(
            {"id": safe_id, "error": f"unserializable tool result: {exc}"},
            separators=(",", ":"),
        )


class _ToolServer:
    """Async TCP server that handles tool-call requests from the Pi extension.

    Protocol (JSONL over TCP):
        Request:  {"id":"...","token":"...","tool":"tool_name","args":{...}}
        Response: {"id":"...","result":{...}} or {"id":"...","error":"..."}

    The loopback socket is reachable by any co-located process, so every
    request must carry :attr:`token` (a per-server secret embedded only in
    the generated Pi extension); frames with a missing or wrong token are
    rejected before dispatch.
    """

    def __init__(self) -> None:
        self._server: asyncio.Server | None = None
        self.port: int = 0
        self._tool_executor: ToolExecutor | None = None
        # Policy gate for native (non-bridged) Pi tool calls, set by
        # :meth:`PiExecutor._ensure_tool_server`. ``None`` (single-process
        # / test paths) means native tool calls are not gated.
        self._policy_gate: NativePolicyGate | None = None
        # Per-server bearer token required on every request. Minted at
        # construction (never None) so auth is always enforced.
        self.token: str = secrets.token_urlsafe(32)

    async def start(self) -> int:
        """Start listening on a random port. Returns the port number."""
        self._server = await asyncio.start_server(
            self._handle_client,
            "127.0.0.1",
            0,
        )
        addr = self._server.sockets[0].getsockname()
        self.port = addr[1]
        logger.debug("PiExecutor tool server listening on port %d", self.port)
        return self.port

    async def stop(self) -> None:
        if self._server is not None:
            self._server.close()
            await self._server.wait_closed()
            self._server = None

    async def _handle_client(
        self,
        reader: asyncio.StreamReader,
        writer: asyncio.StreamWriter,
    ) -> None:
        try:
            while True:
                raw = await reader.readline()
                if not raw:
                    break
                line = raw.decode("utf-8", errors="replace").strip()
                if not line:
                    continue
                try:
                    request = json.loads(line)
                except json.JSONDecodeError:
                    continue
                # Authenticate before any dispatch; close the connection on
                # failure (unlike the malformed-frame ``continue`` below).
                if not self._token_ok(request.get("token")):
                    writer.write(
                        json.dumps(
                            {"id": request.get("id"), "error": "unauthorized"},
                            separators=(",", ":"),
                        ).encode("utf-8")
                        + b"\n"
                    )
                    await writer.drain()
                    break
                raw_req_id = request.get("id")
                raw_tool_name = request.get("tool")
                # The Pi extension always supplies ``id`` and ``tool``
                # on a tool request. Drop malformed frames rather than
                # executing under an empty-string tool name.
                if not isinstance(raw_req_id, str) or not isinstance(raw_tool_name, str):
                    continue
                tool_args = request.get("args", {})
                if request.get("kind") == "policy_eval":
                    # The ``tool_call`` extension hook asks for a TOOL_CALL
                    # verdict on a native (non-bridged) tool — evaluate only,
                    # never execute (Pi runs the tool itself on ALLOW).
                    verdict = await self._evaluate_policy(raw_tool_name, tool_args)
                    response = {"id": raw_req_id, "verdict": verdict}
                else:
                    response = await self._execute(raw_tool_name, tool_args)
                    response["id"] = raw_req_id
                # Serialize defensively: a tool result may carry a value
                # ``json.dumps`` can't encode (e.g. ``datetime``/``set``).
                # If it raises here — outside ``_execute``'s try — the frame
                # is never written and the JS ``callTool`` promise hangs the
                # whole turn (no ``data``/``error`` event). Always emit a JSON
                # frame: a serializable error envelope when the response can't
                # be encoded, mirroring codex's ``_result_text`` guard.
                out = _safe_dumps(response, raw_req_id) + "\n"
                writer.write(out.encode("utf-8"))
                await writer.drain()
        except (asyncio.CancelledError, ConnectionError):
            pass
        finally:
            writer.close()

    def _token_ok(self, presented: JsonValue) -> bool:
        """
        Validate a request's ``token`` against this server's secret.

        :param presented: The ``token`` field from the request JSON — a
            ``str`` when authenticated, any JSON type (or ``None``) on a
            malformed/forged frame, hence the ``isinstance`` guard.
        :returns: ``True`` only when *presented* is a string equal to
            :attr:`token`, compared in constant time.
        """
        return isinstance(presented, str) and hmac.compare_digest(presented, self.token)

    async def _execute(  # type: ignore[explicit-any]
        self,
        name: str,
        args: dict[str, Any],
    ) -> dict[str, Any]:
        if self._tool_executor is None:
            return {"error": f"No tool executor for '{name}'"}
        try:
            raw = self._tool_executor(name, args)
            resolved = await raw if asyncio.iscoroutine(raw) or asyncio.isfuture(raw) else raw
            if not isinstance(resolved, dict):
                resolved = {"result": resolved}
            return {"result": resolved}
        except Exception as exc:  # noqa: BLE001 — tool errors are surfaced via the JSON response envelope
            return {"error": str(exc)}

    async def _evaluate_policy(  # type: ignore[explicit-any]
        self,
        name: str,
        args: dict[str, Any],
    ) -> dict[str, Any]:
        """Evaluate a native (non-bridged) tool call against TOOL_CALL policy.

        Routes through the :attr:`_policy_gate` wired by
        :meth:`PiExecutor._ensure_tool_server`. Returns ``{"block": bool,
        "reason": str}``.

        Fail-open (allow) when no gate is wired or the gate raises: this
        mirrors the runner/scaffold policy-evaluation contract, which also
        defaults to ALLOW on a stalled or failed verdict so a transient
        Omnigent outage can't wedge the agent mid-turn.
        """
        if self._policy_gate is None:
            return {"block": False, "reason": ""}
        try:
            raw = self._policy_gate(name, args)
            resolved = await raw if asyncio.iscoroutine(raw) or asyncio.isfuture(raw) else raw
            if isinstance(resolved, dict):
                return {
                    "block": bool(resolved.get("block")),
                    "reason": str(resolved.get("reason") or ""),
                }
            return {"block": False, "reason": ""}
        except Exception as exc:  # noqa: BLE001 — fail-open; the verdict path must never wedge Pi
            logger.warning("Pi native-tool policy eval failed for %r: %s", name, exc)
            return {"block": False, "reason": ""}


# ---------------------------------------------------------------------------
# Pi extension generator — bridges Omnigent tools into Pi
# ---------------------------------------------------------------------------


def _sanitize_schema(schema: ToolSpec) -> ToolSpec:
    """Strip JSON Schema features unsupported by the OpenAI Responses/Completions APIs.

    Removes ``anyOf``, ``oneOf``, ``allOf``, ``examples``, ``default``,
    ``additionalProperties``, and ``$ref``.  Union keywords are collapsed
    to the first ``object`` branch when one exists (it carries the
    structured properties the model needs), else the first typed branch.
    This ensures the schema is accepted by Databricks serving endpoints.
    """
    if not isinstance(schema, dict):
        return schema
    result: ToolSpec = {}
    for key, value in schema.items():
        if key in ("examples", "default", "$ref", "additionalProperties"):
            continue
        if key in ("anyOf", "oneOf", "allOf"):
            # Prefer the object branch: collapsing to a primitive drops the properties.
            if isinstance(value, list) and value:
                typed = [c for c in value if isinstance(c, dict) and "type" in c]
                for candidate in typed:
                    if candidate["type"] == "object":
                        return _sanitize_schema(candidate)
                if typed:
                    return _sanitize_schema(typed[0])
            continue
        if key == "properties" and isinstance(value, dict):
            result[key] = {k: _sanitize_schema(v) for k, v in value.items()}
        elif key == "items" and isinstance(value, dict):
            result[key] = _sanitize_schema(value)
        else:
            result[key] = value
    return result


def _generate_extension_js(port: int, tool_schemas: list[ToolSpec], token: str) -> str:
    """Generate a JavaScript Pi extension that registers Omnigent tools.

    Each tool is forwarded to the TCP tool server at ``127.0.0.1:<port>``.

    :param port: TCP port the Omnigent tool server is listening on, e.g.
        ``54321``.
    :param tool_schemas: Omnigent tool schemas to register with Pi.
    :param token: The tool server's bearer token
        (:attr:`_ToolServer.token`), embedded in the extension and sent
        on every request so the server can authenticate this Pi process.
    """
    # Build tool descriptors for the JS code. ``ToolSpec`` is the
    # same JSON-shaped ``dict[str, Any]`` we consume.
    descriptors: list[ToolSpec] = []
    for s in tool_schemas:
        raw_name = s.get("name")
        # Pi requires a concrete name on each registered tool; skip
        # malformed entries rather than registering an unnamed tool.
        if not isinstance(raw_name, str) or not raw_name:
            continue
        raw_desc = s.get("description")
        # Description is optional; Pi's JS bridge expects ``str`` on
        # both ``description`` and ``promptSnippet``.
        desc: str = raw_desc if isinstance(raw_desc, str) else ""
        descriptors.append(
            {
                "name": raw_name,
                "description": desc,
                "promptSnippet": desc[:120],
                "parameters": _sanitize_schema(
                    s.get("parameters", {"type": "object", "properties": {}})
                ),
            }
        )
    tools_json = json.dumps(descriptors, indent=2)
    # json.dumps so the secret is correctly JS-string-escaped.
    token_json = json.dumps(token)

    return f"""\
// Auto-generated Omnigent tool bridge extension for Pi.
// Connects to the Omnigent TCP tool server on port {port}.
const net = require("net");

const TOOLS = {tools_json};
const BRIDGED = new Set(TOOLS.map((t) => t.name));
const PORT = {port};
const TOKEN = {token_json};

/** Send a tool call request over TCP and return the result. */
function callTool(toolName, args) {{
  return new Promise((resolve) => {{
    // Idempotent settle: a tool call must resolve exactly once. Route every
    // resolve through finish() so a late "close" after a real "data" response
    // can't clobber the result, and a bare close with no data still resolves
    // (rather than hanging Pi's agent loop forever).
    let settled = false;
    const finish = (v) => {{ if (!settled) {{ settled = true; resolve(v); }} }};
    const errorResult = (text) => ({{
      content: [{{ type: "text", text: JSON.stringify({{ error: text }}) }}],
      isError: true
    }});
    const client = net.createConnection({{ port: PORT, host: "127.0.0.1" }}, () => {{
      const id = Math.random().toString(36).slice(2);
      const req = JSON.stringify({{ id, token: TOKEN, tool: toolName, args }}) + "\\n";
      let buf = "";
      client.on("data", (chunk) => {{
        buf += chunk.toString();
        const nl = buf.indexOf("\\n");
        if (nl !== -1) {{
          try {{
            const resp = JSON.parse(buf.slice(0, nl));
            client.end();
            if (resp.error) {{
              finish(errorResult(resp.error));
            }} else {{
              const text = typeof resp.result === "string"
                ? resp.result
                : JSON.stringify(resp.result);
              const isError = resp.result && (resp.result.error || resp.result.blocked);
              finish({{ content: [{{ type: "text", text }}], isError: !!isError }});
            }}
          }} catch (e) {{
            client.end();
            finish(errorResult(e.message));
          }}
        }}
      }});
      client.on("error", (err) => {{
        finish(errorResult(err.message));
      }});
      // A clean FIN with no bytes (e.g. the server failed to serialize the
      // result and closed) emits neither "data" nor "error"; without this
      // the promise would hang forever. finish() is a no-op if already settled.
      client.on("close", () => {{
        finish(errorResult("tool server closed connection without a response"));
      }});
      client.write(req);
    }});
    client.on("error", (err) => {{
      finish(errorResult(err.message));
    }});
  }});
}}

/**
 * Ask the Omnigent tool server for a TOOL_CALL policy verdict on a native
 * (non-bridged) Pi tool. Resolves to the verdict object {{ block, reason }}
 * or null. Fail-open (null) on any transport error: a wedged native tool
 * would break Pi worse than a missed gate, and bridged tools stay gated
 * server-side regardless.
 */
function evalNativePolicy(toolName, args) {{
  return new Promise((resolve) => {{
    const client = net.createConnection({{ port: PORT, host: "127.0.0.1" }}, () => {{
      const id = Math.random().toString(36).slice(2);
      const frame = {{ id, token: TOKEN, kind: "policy_eval", tool: toolName, args }};
      const req = JSON.stringify(frame) + "\\n";
      let buf = "";
      client.on("data", (chunk) => {{
        buf += chunk.toString();
        const nl = buf.indexOf("\\n");
        if (nl !== -1) {{
          client.end();
          try {{
            const resp = JSON.parse(buf.slice(0, nl));
            resolve(resp && resp.verdict ? resp.verdict : null);
          }} catch (e) {{
            resolve(null);
          }}
        }}
      }});
      client.on("error", () => resolve(null));
      client.write(req);
    }});
    client.on("error", () => resolve(null));
  }});
}}

module.exports = function(pi) {{
  // Gate native (non-bridged) tool calls through Omnigent policy. Pi's
  // native tools (e.g. ``read``, enabled for skill loading) run in-process
  // and never traverse the bridged /mcp path, so without this hook they
  // escape all guardrails. Bridged tools ARE evaluated server-side at /mcp,
  // so skip them here to avoid double-evaluation (and double ASK prompts).
  pi.on("tool_call", async (event) => {{
    if (!event || typeof event.toolName !== "string") return;
    if (BRIDGED.has(event.toolName)) return;
    const verdict = await evalNativePolicy(event.toolName, event.input || {{}});
    if (verdict && verdict.block) {{
      return {{ block: true, reason: verdict.reason || "blocked by Omnigent policy" }};
    }}
  }});

  for (const tool of TOOLS) {{
    // Pi passes tool.parameters directly to the LLM as JSON Schema, so we
    // can use the Omnigent schema as-is without TypeBox conversion.
    pi.registerTool({{
      name: tool.name,
      label: tool.name,
      description: tool.description,
      promptSnippet: tool.promptSnippet || tool.description,
      parameters: tool.parameters || {{ type: "object", properties: {{}} }},
      async execute(_toolCallId, _params, _signal, _onUpdate, _ctx) {{
        return callTool(tool.name, _params);
      }},
    }});
  }}
}};
"""


# ---------------------------------------------------------------------------
# Credential helpers (shared pattern with codex_executor / claude_sdk_executor)
# ---------------------------------------------------------------------------


def _find_pi_cli() -> str | None:
    """Find the ``pi`` CLI on PATH."""
    return shutil.which("pi")


# ---------------------------------------------------------------------------
# Databricks models.json generation
# ---------------------------------------------------------------------------

# ---------------------------------------------------------------------------
# Databricks model definitions for Pi's models.json
# ---------------------------------------------------------------------------
# Databricks exposes API styles at different URL paths. ucode state can
# override these provider URLs; the host-derived defaults remain for legacy
# profile-only usage.

# Each static entry declares ``input: ["text", "image"]`` for the same reason
# the dynamic-registration path does (see _build_models_json): Pi's
# transformMessages strips image blocks unless the model entry advertises
# image input. These are all vision-capable models, and the run model is often
# a static id — in which case _build_models_json's append is skipped, so the
# capability has to be declared here too or attached images are silently
# dropped (#515/#516).
_DATABRICKS_RESPONSES_MODELS = [
    {
        "id": "databricks-gpt-5-4-mini",
        "name": "GPT-5.4 Mini",
        "contextWindow": 1047576,
        "maxTokens": 32768,
        "input": ["text", "image"],
    },
    {
        "id": "databricks-gpt-5-4",
        "name": "GPT-5.4",
        "contextWindow": 1047576,
        "maxTokens": 32768,
        "input": ["text", "image"],
    },
    {
        "id": "databricks-gpt-5-5",
        "name": "GPT-5.5",
        # OSS profile endpoint metadata: 400K total context, 128K max output.
        "contextWindow": 400000,
        "maxTokens": 128000,
        "input": ["text", "image"],
    },
    {
        "id": "databricks-gpt-5-5-pro",
        "name": "GPT-5.5 Pro",
        # OSS profile endpoint metadata: 400K total context, 128K max output.
        "contextWindow": 400000,
        "maxTokens": 128000,
        "input": ["text", "image"],
    },
]

_DATABRICKS_ANTHROPIC_MODELS = [
    {
        "id": "databricks-claude-opus-4-8",
        "name": "Claude Opus 4.8",
        # Gateway-verified caps: >1000000 input rejects, 128001+ output rejects.
        "contextWindow": 1000000,
        "maxTokens": 128000,
        "input": ["text", "image"],
    },
    {
        "id": "databricks-claude-sonnet-4-6",
        "name": "Claude Sonnet 4.6",
        "contextWindow": 1000000,
        "maxTokens": 128000,
        "input": ["text", "image"],
    },
    {
        "id": "databricks-claude-sonnet-4-5",
        "name": "Claude Sonnet 4.5",
        # Gateway rejects this model past ~200k input.
        "contextWindow": 200000,
        "maxTokens": 16384,
        "input": ["text", "image"],
    },
]

# Empty: the only listed endpoint (meta-llama-3.3-70b) no longer exists on
# the gateway. The provider stays so non-Claude/GPT ids keep a routing home.
_DATABRICKS_COMPLETIONS_MODELS: list[dict[str, str | int]] = []

# Prefix-matched env var names allowed into the Pi subprocess. Only
# known-safe categories pass: Pi's own config knobs, proxy settings,
# TLS trust overrides, Node.js runtime knobs (Pi is a Node CLI), and
# locale. Credential families (``DATABRICKS_*``, ``AWS_*``, LLM
# provider API keys, ...) deliberately do NOT match.
_PI_ENV_ALLOW_PREFIXES: tuple[str, ...] = (
    "PI_",
    "HTTP_",
    "HTTPS_",
    "ALL_PROXY",
    "NO_PROXY",
    "SSL_",
    "NODE_",
    "XDG_",
    "LANG",
    "LC_",
)

# Exact-matched env var names allowed into the Pi subprocess: the
# minimal set a POSIX CLI reasonably expects.
_PI_ENV_ALLOW_EXACT: frozenset[str] = frozenset(
    {
        "HOME",
        "PATH",
        "TERM",
        "TMPDIR",
        "TMP",
        "TEMP",
        "USER",
        "LOGNAME",
        "SHELL",
        "TZ",
        OMNIGENT_SESSION_ENV_VAR,  # "inside Omnigent" marker (CLAUDE_CODE/CODEX analog)
    }
)
_STREAM_READ_CHUNK_SIZE = 65536

# CLI flags whose values are sensitive (e.g. the full system prompt) and must
# not be written to logs verbatim. The value following these flags is replaced
# with a length-only placeholder.
_REDACTED_ARGV_FLAGS = frozenset({"--append-system-prompt", "--system-prompt"})


def _redact_argv_for_log(args: Sequence[str]) -> list[str]:
    """
    Return a copy of ``args`` with sensitive flag values redacted for logging.

    The system prompt value (e.g. passed via ``--append-system-prompt``) is
    replaced with a ``[system prompt N chars]`` placeholder so it never leaks
    into debug logs. Two argv forms are handled:

    * the two-token form ``--append-system-prompt <value>`` (what the current
      spawn code emits), and
    * the equals-joined form ``--append-system-prompt=<value>`` (not emitted
      today, but redacted defensively in case a future refactor switches to it).

    All other tokens are preserved so the command remains useful for debugging.
    """
    redacted: list[str] = []
    redact_next = False
    for arg in args:
        if redact_next:
            redacted.append(f"[system prompt {len(arg)} chars]")
            redact_next = False
            continue
        if arg in _REDACTED_ARGV_FLAGS:
            # Two-token form: redact the following value token.
            redacted.append(arg)
            redact_next = True
            continue
        flag, sep, value = arg.partition("=")
        if sep and flag in _REDACTED_ARGV_FLAGS:
            # Equals-joined form: redact the inline value, keep the flag name.
            redacted.append(f"{flag}=[system prompt {len(value)} chars]")
            continue
        redacted.append(arg)
    return redacted


def _build_models_json(
    host: str,
    token: str,
    base_urls: dict[str, str] | None = None,
    model: str | None = None,
) -> dict[str, Any]:  # type: ignore[explicit-any]
    # Pi's models.json mixes str/int/bool/list/dict across provider configs;
    # see _DATABRICKS_*_MODELS and the compat/authHeader shapes below. The
    # schema is owned by the Pi CLI and not worth a TypedDict tree here.
    """Build a Pi ``models.json`` with three gateway providers.

    Each provider targets a different API gateway path and wire format so
    the correct protocol is used for each model family. The static model
    lists cover the known Databricks-gateway ids; *model* additionally
    registers the resolved run model so a gateway model outside those
    lists (an OpenRouter/LiteLLM id like ``moonshotai/kimi-k2.6``, or a
    Databricks id newer than the static list) resolves instead of Pi
    failing with "Model not found" — Pi only accepts ``provider/<model>``
    selectors whose id is registered under that provider.

    :param host: Databricks workspace URL used for legacy profile-only
        defaults.
    :param token: Bearer token to write into Pi's provider entries.
    :param base_urls: Optional provider base URLs keyed by model family,
        e.g. ``{"claude": "...", "openai": "..."}`` — from ucode state or
        a generic (OpenRouter/LiteLLM) provider entry.
    :param model: The resolved model id this run will select, e.g.
        ``"moonshotai/kimi-k2.6"``; registered (bare ``{"id": ...}``, the
        same shape ucode writes) under the provider
        :func:`_pi_provider_for_model` routes it to when absent from the
        static list. ``None`` skips registration (Pi picks its default).
    :returns: Pi ``models.json`` contents.
    """
    h = host.rstrip("/")
    serving_endpoints_url = f"{h}/serving-endpoints"
    raw_openai_base_url = (base_urls or {}).get("openai")
    # ucode's ``openai`` gateway is the Codex Responses gateway
    # (``.../ai-gateway/codex/v1``), which 404s pi's openai-completions
    # ``/chat/completions`` POST. Re-route only that case to serving-endpoints;
    # generic providers (no ``/ai-gateway/codex``) pass through. Gemini rides
    # this path via databricks-completions since pi can't speak generateContent.
    if raw_openai_base_url and "/ai-gateway/codex" in raw_openai_base_url:
        openai_base_url = serving_endpoints_url
    else:
        openai_base_url = raw_openai_base_url or serving_endpoints_url
    claude_base_url = (base_urls or {}).get("claude") or f"{h}/serving-endpoints/anthropic"
    config: dict[str, Any] = {  # type: ignore[explicit-any]  # Pi-owned schema, see note above
        "providers": {
            # GPT models → OpenAI Chat Completions API.
            # We use completions (not responses) because the Databricks
            # Responses API rejects tool-result chaining on subsequent turns.
            # The ``compat`` settings ensure Pi uses ``system`` role (not
            # ``developer``) and avoids other OpenAI-specific features that
            # Databricks doesn't support.
            "databricks": {
                "baseUrl": openai_base_url,
                "apiKey": token,
                "api": "openai-completions",
                "compat": {
                    "supportsDeveloperRole": False,
                    "supportsStore": False,
                    "supportsStrictMode": False,
                    "supportsReasoningEffort": False,
                },
                "models": _DATABRICKS_RESPONSES_MODELS,
            },
            # Claude models → Anthropic Messages API.
            # ``authHeader`` sends ``Authorization: Bearer <token>`` instead
            # of the default Anthropic ``x-api-key`` header.
            "databricks-anthropic": {
                "baseUrl": claude_base_url,
                "apiKey": token,
                "api": "anthropic-messages",
                "authHeader": True,
                "models": _DATABRICKS_ANTHROPIC_MODELS,
            },
            # Everything else (Llama, etc.) → same endpoint, same API
            "databricks-completions": {
                "baseUrl": openai_base_url,
                "apiKey": token,
                "api": "openai-completions",
                "compat": {
                    "supportsDeveloperRole": False,
                    "supportsStore": False,
                    "supportsStrictMode": False,
                    "supportsReasoningEffort": False,
                },
                "models": _DATABRICKS_COMPLETIONS_MODELS,
            },
        },
    }
    if model is not None:
        provider = config["providers"][_pi_provider_for_model(model)]
        if not any(entry.get("id") == model for entry in provider["models"]):
            # Rebind (don't append): the static lists are module-level
            # constants shared across builds, so in-place mutation would
            # leak this run's model id into every later models.json.
            # Declare image input: Pi's transformMessages drops every image
            # block ("model does not support images") unless the model entry
            # advertises it, so a dynamically-registered vision model would
            # silently lose its images (#515). We assert capability for ALL
            # dynamically-routed ids (no per-model vision metadata here): for a
            # genuinely text-only model this turns a silent image drop into a
            # provider-side 400 on image turns — a deliberate trade (loud error
            # over silent loss), since most current gateway models are
            # multimodal and text-only turns are unaffected.
            provider["models"] = [
                *provider["models"],
                {"id": model, "input": ["text", "image"]},
            ]
    return config


def _pi_provider_for_model(model: str) -> str:
    """Return the Pi provider name to use for a given Databricks model."""
    lower = model.lower()
    if "claude" in lower:
        return "databricks-anthropic"
    if "gpt" in lower:
        return "databricks"
    return "databricks-completions"


async def _create_subprocess_exec(*args: Any, **kwargs: Any) -> asyncio.subprocess.Process:  # type: ignore[explicit-any]
    """
    Indirection point for ``asyncio.create_subprocess_exec``.

    Exists so tests can stub subprocess creation without patching
    ``asyncio.create_subprocess_exec`` globally (patching
    ``omnigent.inner.pi_executor.asyncio.create_subprocess_exec``
    walks the dotted path into the real ``asyncio`` module singleton
    and leaks the mock into every other test in the process).

    :param args: Positional argv components forwarded to
        ``asyncio.create_subprocess_exec``.
    :param kwargs: Keyword args (``stdin``, ``stdout``, ``stderr``,
        ``env``, ``cwd``, ...) forwarded as-is.
    :returns: The spawned subprocess handle.
    """
    return await asyncio.create_subprocess_exec(*args, **kwargs)


def _clean_pi_env(extra_allowed: Sequence[str] | None = None) -> dict[str, str]:
    """
    Build a filtered copy of ``os.environ`` for the Pi subprocess.

    Deny-by-default allowlist mirroring the codex path's
    :func:`omnigent.inner.codex_executor._clean_codex_env`. The previous
    additive ``{**os.environ, **env}`` merge leaked every host secret
    — cloud tokens, API keys — into the Pi process, sandboxed or not.
    Credential families are deliberately not allowlisted:
    gateway runs authenticate through the generated ``models.json``,
    and a spec that legitimately needs a variable inside the Pi
    process opts in via ``os_env.sandbox.env_passthrough``.

    :param extra_allowed: Extra exact names to pass through, e.g. the
        spec's ``os_env.sandbox.env_passthrough`` entries
        (``["ANTHROPIC_API_KEY"]`` for an env-key-authenticated
        non-gateway run). ``None`` means no extras.
    :returns: Filtered environment dict, ready to extend with the
        executor's own variables (``PI_CODING_AGENT_DIR``, ...).
    """
    allow_exact = set(_PI_ENV_ALLOW_EXACT)
    if extra_allowed is not None:
        allow_exact.update(extra_allowed)
    return {
        key: value
        for key, value in os.environ.items()
        if key in allow_exact or key.startswith(_PI_ENV_ALLOW_PREFIXES)
    }


# ---------------------------------------------------------------------------
# Pi RPC subprocess wrapper
# ---------------------------------------------------------------------------


@dataclass
class _PiRpcSession:
    """Manages a single Pi subprocess in RPC mode."""

    process: asyncio.subprocess.Process | None = None
    _read_task: Task[None] | None = None
    _stderr_task: Task[None] | None = None
    # Non-optional: initialized eagerly so the reader/writer paths don't
    # need to None-narrow on every access. ``None`` sentinel values are
    # pushed onto the queue to signal EOF to ``read_line``.
    _line_queue: Queue[str | None] = field(default_factory=Queue)
    _stderr_lines: list[str] = field(default_factory=list)
    _tmp_dir: str | None = None

    async def start(
        self,
        pi_path: str,
        *,
        env: dict[str, str],
        cwd: str | None = None,
        model: str | None = None,
        system_prompt: str | None = None,
        extra_args: list[str] | None = None,
    ) -> None:
        """
        Spawn the Pi subprocess in RPC mode and start the I/O readers.

        :param pi_path: Path to spawn — the ``pi`` binary or the
            sandbox launcher script wrapping it.
        :param env: The COMPLETE subprocess environment. Deliberately
            not merged with ``os.environ`` — the caller builds it from
            the :func:`_clean_pi_env` allowlist so host/server secrets
            (e.g. ``DATABRICKS_TOKEN``) never leak into the (possibly
            sandboxed) Pi process.
        :param cwd: Working directory for the subprocess, or ``None``
            to inherit the parent's.
        :param model: Pi model selector, e.g.
            ``"databricks-anthropic/databricks-claude-sonnet-4-6"``.
            ``None`` lets Pi pick its default.
        :param system_prompt: Text appended to Pi's default system
            prompt via ``--append-system-prompt``. ``None`` skips it.
        :param extra_args: Extra CLI tokens (``--extension``,
            ``--tools``, ...). ``None`` appends nothing.
        """
        args = [pi_path, "--mode", "rpc", "--no-session"]
        if model:
            pi_coding_agent_dir = env.get("PI_CODING_AGENT_DIR")
            args.extend(
                [
                    "--model",
                    f"databricks/{model}"
                    if pi_coding_agent_dir is not None and "databricks" in pi_coding_agent_dir
                    else model,
                ]
            )
        if system_prompt:
            # Use --append-system-prompt instead of --system-prompt so Pi
            # keeps its default prompt (which includes tool descriptions from
            # promptSnippet and guidelines).  Using --system-prompt would
            # replace the default prompt entirely, stripping tool awareness.
            args.extend(["--append-system-prompt", system_prompt])
        if extra_args:
            args.extend(extra_args)

        logger.debug("PiExecutor: spawning %s", " ".join(_redact_argv_for_log(args)))
        self.process = await _create_subprocess_exec(
            *args,
            stdin=asyncio.subprocess.PIPE,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
            cwd=cwd,
            env=env,
        )
        self._read_task = asyncio.create_task(self._reader())
        self._stderr_task = asyncio.create_task(self._stderr_reader())

    async def _reader(self) -> None:
        """Background task: read lines from Pi stdout and enqueue them."""
        assert self.process is not None and self.process.stdout is not None
        try:
            async for raw_line in self._iter_stream_lines(self.process.stdout):
                line = raw_line.decode("utf-8", errors="replace").rstrip("\n\r")
                if line:
                    self._line_queue.put_nowait(line)
        except asyncio.CancelledError:
            return
        except Exception as exc:  # noqa: BLE001 — reader loop logs and exits on any unexpected error
            logger.debug("PiExecutor reader error: %s", exc)
        finally:
            self._line_queue.put_nowait(None)

    async def _stderr_reader(self) -> None:
        """Drain stderr in the background."""
        if self.process is None or self.process.stderr is None:
            return
        try:
            async for raw_line in self._iter_stream_lines(self.process.stderr):
                text = raw_line.decode("utf-8", errors="replace").rstrip("\n\r")
                if text:
                    logger.debug("pi stderr: %s", text)
                    if len(self._stderr_lines) < 50:
                        self._stderr_lines.append(text)
        except (asyncio.CancelledError, Exception):  # noqa: BLE001 — stderr drainer is best-effort; never crashes
            pass

    @staticmethod
    async def _iter_stream_chunks(stream: asyncio.StreamReader) -> AsyncIterator[bytes]:
        while True:
            chunk = await stream.read(_STREAM_READ_CHUNK_SIZE)
            if not chunk:
                break
            yield chunk

    @staticmethod
    async def _iter_stream_lines(stream: asyncio.StreamReader) -> AsyncIterator[bytes]:
        buffer = bytearray()
        async for chunk in _PiRpcSession._iter_stream_chunks(stream):
            buffer.extend(chunk)
            while True:
                newline_index = buffer.find(b"\n")
                if newline_index < 0:
                    break
                line = bytes(buffer[: newline_index + 1])
                del buffer[: newline_index + 1]
                yield line
        if buffer:
            yield bytes(buffer)

    async def send_command(self, command: CodexEvent) -> None:
        """Send a JSONL command to Pi's stdin."""
        if self.process is None or self.process.stdin is None:
            raise RuntimeError("Pi process not running")
        line = json.dumps(command, separators=(",", ":")) + "\n"
        self.process.stdin.write(line.encode("utf-8"))
        await self.process.stdin.drain()

    async def read_line(self, timeout: float = 120.0) -> str | None:
        """Read the next JSONL line from Pi's stdout. Returns None on EOF."""
        try:
            return await asyncio.wait_for(self._line_queue.get(), timeout=timeout)
        except asyncio.TimeoutError:
            return None

    async def close(self) -> None:
        for task in (self._read_task, self._stderr_task):
            if task is not None:
                task.cancel()
                with contextlib.suppress(asyncio.CancelledError, Exception):
                    await task
        self._read_task = None
        self._stderr_task = None
        if self.process is not None:
            with contextlib.suppress(ProcessLookupError):
                self.process.terminate()
            try:
                await asyncio.wait_for(self.process.wait(), timeout=2.0)
            except (asyncio.TimeoutError, ProcessLookupError, RuntimeError):
                # RuntimeError can happen when the subprocess was created on a
                # different event loop (e.g. test fixtures that call close() in
                # a fresh loop).  Fall back to synchronous kill.
                with contextlib.suppress(ProcessLookupError):
                    self.process.kill()
            close_subprocess_transport(self.process)
            self.process = None
        if self._tmp_dir is not None:
            shutil.rmtree(self._tmp_dir, ignore_errors=True)
            self._tmp_dir = None


# ---------------------------------------------------------------------------
# PiExecutor
# ---------------------------------------------------------------------------


@dataclass
class _PiSessionState:
    rpc: _PiRpcSession | None = None
    system_prompt: str | None = None
    model: str | None = None
    _has_sent_prompt: bool = False


@dataclass(frozen=True)
class BlockedCheck:
    """Result of inspecting a Pi tool result for a policy-blocked payload.

    :param blocked: ``True`` when the result is (or wraps) a
        ``{"blocked": true, "reason": "..."}`` dict.
    :param reason: The human-readable block reason when ``blocked`` is
        ``True``; an empty string otherwise.
    """

    blocked: bool
    reason: str


@dataclass(frozen=True)
class PiSubprocessConfig:
    """Materialized environment + CLI args for a Pi subprocess.

    :param env: The complete environment for the Pi process — the
        :func:`_clean_pi_env` allowlist base, plus
        ``PI_CODING_AGENT_DIR`` when Databricks model routing is in use.
    :param tmp_dir: Temp directory that owns any ``models.json`` /
        ``omnigent_tools.js`` files generated for this subprocess.  The
        caller is responsible for cleaning it up on shutdown.
    :param extra_args: Extra CLI arguments to append to the Pi invocation,
        e.g. ``["--extension", ".../omnigent_tools.js"]``.
    """

    env: dict[str, str]
    tmp_dir: str
    extra_args: list[str]


def _extract_text(msg: Message) -> str:
    """
    Extract text content from a message dict.

    :param msg: A conversation message dict.
    :returns: Plain text content of the message.
    """
    content = msg.get("content")
    if content is None:
        return ""
    if isinstance(content, str):
        return content
    if isinstance(content, list):
        parts: list[str] = []
        for part in content:
            if not isinstance(part, dict):
                continue
            text = part.get("text")
            if isinstance(text, str):
                parts.append(text)
        return " ".join(parts)
    return str(content)


def _extract_latest_user_content(
    messages: list[Message],
) -> str | list[dict[str, Any]]:
    """
    Extract the latest user message content.

    Returns a plain string for text-only messages. When the
    message carries multimodal content blocks, returns the
    block list so the caller can pass structured input to Pi.

    :param messages: Conversation history.
    :returns: A string prompt or a list of content block dicts.
    """
    for msg in reversed(messages):
        if msg.get("role") == "user":
            content = msg.get("content")
            if content is None:
                return ""
            if isinstance(content, str):
                return content
            if isinstance(content, list):
                return content
            return str(content)
    return ""


def _split_pi_prompt(blocks: list[dict[str, Any]]) -> tuple[str, list[dict[str, str]]]:
    """Split content blocks into Pi's prompt ``message`` text and ``images``.

    Pi's RPC ``prompt`` command carries text in ``message`` and images in a
    native ``images`` field (``ImageContent = {type, data, mimeType}``), which
    the ``openai-completions`` provider converts to ``image_url`` blocks.
    JSON-encoding the blocks into ``message`` instead (the previous behavior)
    makes Pi forward the image data URI as literal text, so the model never
    sees the image (#515).

    :param blocks: Responses-API content block dicts. ``input_text`` →
        message text; ``input_image`` → Pi ``images``; ``input_file`` (a
        resolved attachment data URI) → its text payload is decoded into the
        message for text-like MIME types and skipped (with a warning) for
        binary ones, mirroring the Codex harness — Pi has no native file
        channel, but a text file's content is still useful inline, and the
        old ``json.dumps(prompt)`` path surfaced it as text, so neither a
        silent drop nor a hard failure is the right default here. A genuinely
        unknown block type raises ``ValueError`` (→ ``ExecutorError`` in
        ``run_turn``) so a new shape fails loudly rather than vanishing.
    :returns: ``(message, images)`` — joined text and a list of Pi
        ``ImageContent`` dicts.
    :raises ValueError: On a malformed ``input_image`` or an unhandled
        block type.
    """
    text_parts: list[str] = []
    images: list[dict[str, str]] = []
    for block in blocks:
        block_type = block.get("type")
        if block_type in ("input_text", "output_text", "text"):
            text_parts.append(str(block.get("text") or ""))
        elif block_type == "input_image":
            image_url = block.get("image_url")
            if not isinstance(image_url, str) or not image_url.startswith("data:"):
                raise ValueError(
                    "input_image 'image_url' must be an inline data URI "
                    "('data:<mime>;base64,...'); Pi forwards images inline and "
                    "cannot fetch external URLs or resolve file references "
                    f"(got {str(image_url)[:48]!r})."
                )
            parsed = parse_data_uri(image_url)
            images.append(
                {"type": "image", "data": parsed.base64_payload, "mimeType": parsed.mime_type}
            )
        elif block_type == "input_file":
            # Pi has no file channel, so inline text-like files into the
            # message (the model can still reason about them) and skip binary
            # ones. Matches codex_executor's input_file fallback.
            file_data = block.get("file_data") or ""
            if isinstance(file_data, str) and file_data.startswith("data:"):
                parsed = parse_data_uri(file_data)
                if not parsed.mime_type.startswith("text/"):
                    logger.warning(
                        "Pi harness: skipping non-text input_file (%s) — Pi "
                        "cannot consume binary attachments inline.",
                        parsed.mime_type,
                    )
                    continue
                try:
                    decoded = base64.b64decode(parsed.base64_payload).decode("utf-8", "replace")
                except ValueError as exc:  # binascii.Error subclasses ValueError
                    raise ValueError(
                        f"input_file has an undecodable base64 payload: {exc}"
                    ) from exc
                text_parts.append(decoded)
            elif isinstance(file_data, str) and file_data:
                # Already-inline text (no data: wrapper).
                text_parts.append(file_data)
        else:
            raise ValueError(
                f"Unsupported content block type {block_type!r} for the Pi "
                "harness: only 'input_text', 'input_image', and 'input_file' "
                "blocks can be forwarded. Dropping it silently would lose the "
                "attachment, so the turn fails loudly instead."
            )
    return "\n".join(text_parts), images


def _build_pi_prompt(
    messages: list[Message], *, is_first_turn: bool
) -> str | list[dict[str, Any]]:
    """
    Build the prompt to send to Pi.

    On the first turn with prior history (e.g. sub-agent with
    ``pass_history=True``), serializes the full conversation so
    the Pi process has context. Otherwise returns just the
    latest user message content (may be multimodal).

    :param messages: Omnigent conversation history for the
        turn.
    :param is_first_turn: ``True`` when this is the first turn
        against a freshly-started Pi subprocess, so the full
        history needs to be serialized into the prompt.
    :returns: A string prompt or a list of content block dicts.
    """
    user_messages = [m for m in messages if m.get("role") == "user"]

    if is_first_turn and len(messages) > 1 and len(user_messages) > 1:
        lines = ["Conversation so far:"]
        for msg in messages:
            role = str(msg.get("role") or "user").replace("_", " ")
            text = _extract_text(msg)
            lines.append(f"{role}: {text}")
        lines.append("")
        lines.append(
            "Respond to the latest user message, using the conversation above as context."
        )
        return "\n".join(lines)

    return _extract_latest_user_content(messages)


@dataclass(frozen=True)
class SandboxedPiCli:
    """Result of wrapping the Pi CLI in a sandbox.

    :param launch_path: The path the executor should actually spawn.  Either
        the original ``pi_path`` (sandbox skipped) or a generated wrapper
        script that applies the sandbox before exec-ing Pi.
    :param sandboxed: ``True`` when the wrapper script is active.
    """

    launch_path: str
    sandboxed: bool


def _try_sandbox_pi(
    pi_path: str,
    os_env: OSEnvSpec | None,
    cwd: str | None,
    spawn_env_names: Sequence[str] | None = None,
) -> SandboxedPiCli:
    """Wrap the Pi CLI in a sandbox if ``os_env`` requests it.

    :param pi_path: Path to the resolved ``pi`` binary.
    :param os_env: The agent's ``os_env`` spec.  ``None`` or a ``"none"``
        sandbox skips wrapping.
    :param cwd: Working directory the sandbox should consider the root.
    :param spawn_env_names: Env-var names the executor deliberately
        passes in the Pi subprocess env (the :func:`_clean_pi_env`
        keys plus per-spawn extras like ``PI_CODING_AGENT_DIR``).
        Baked into the launcher policy so :func:`run_launcher` prunes
        anything else its environment picks up — defense in depth
        against host-env leakage into the sandbox. ``None`` skips the
        prune.
    """
    if os_env is None:
        return SandboxedPiCli(launch_path=pi_path, sandboxed=False)
    sandbox_spec = os_env.sandbox or OSEnvSandboxSpec()
    if sandbox_spec.type == "none":
        return SandboxedPiCli(launch_path=pi_path, sandboxed=False)
    try:
        import pathlib

        from .sandbox import (
            create_exec_launcher,
            resolve_sandbox,
            with_additional_read_roots,
            with_additional_write_roots,
            with_spawn_env_allowlist,
        )

        resolved_cwd = pathlib.Path(cwd or os.getcwd()).resolve(strict=False)
        sandbox = resolve_sandbox(os_env, resolved_cwd)
        if not sandbox.active:
            return SandboxedPiCli(launch_path=pi_path, sandboxed=False)
        # Pi needs to read its own installation + node_modules.
        pi_dir = pathlib.Path(pi_path).resolve().parent.parent
        sandbox = with_additional_read_roots(sandbox, [pi_dir])
        # Pi writes to ~/.pi and to /tmp.
        home_pi = pathlib.Path(os.path.expanduser("~/.pi"))
        sandbox = with_additional_write_roots(sandbox, [home_pi, pathlib.Path("/tmp")])
        sandbox = with_spawn_env_allowlist(sandbox, spawn_env_names)
        launcher = create_exec_launcher(pi_path, sandbox)
        return SandboxedPiCli(launch_path=launcher, sandboxed=True)
    except (OSError, ImportError, NotImplementedError) as exc:
        logger.warning("Could not apply sandbox for Pi: %s", exc)
        return SandboxedPiCli(launch_path=pi_path, sandboxed=False)


def _resolve_pi_skill_args(
    skills_filter: str | list[str],
    bundle_dir: pathlib.Path | None,
) -> list[str]:
    """
    Translate ``skills_filter`` into Pi CLI args.

    Pi exposes two skill knobs at the CLI: ``--no-skills`` (suppress
    auto-discovery + loading) and ``--skill <path>`` (explicitly load
    a skill directory; repeatable). This helper composes them based
    on the spec's ``skills_filter``:

    - ``"all"`` → emit ``--skill <path>`` for every bundled skill so
      the agent definitely sees them, AND let Pi's auto-discovery
      run (no ``--no-skills``) so any host-installed skills also
      surface.
    - ``"none"`` → ``["--no-skills"]``. Pi's auto-discovery is
      suppressed and no explicit skills are loaded.
    - ``list[str]`` → ``["--no-skills"]`` (suppress auto-discovery)
      plus one ``--skill <path>`` per named bundle skill that
      exists. Names not present in the bundle are silently skipped
      — matches the SDK convention where a missing skill is no-op,
      not an error.

    Pi's ``--skill`` flag accepts a directory path, not a name, so
    the resolver looks up named skills under ``<bundle>/skills/``.
    Host-installed Pi skills aren't accessible by name without
    knowing Pi's internal extension layout, so the named-list mode
    only resolves bundle skills.

    :param skills_filter: ``"all"`` / ``"none"`` / list of skill
        names from the spec.
    :param bundle_dir: The agent bundle's extracted on-disk path
        (e.g. ``loaded.workdir``). ``None`` when no bundle is
        available — the resolver returns no ``--skill`` flags
        regardless of filter, and only the ``--no-skills`` flag for
        ``"none"``/list cases.
    :returns: A list of CLI tokens to extend ``self._extra_args``
        with. Empty for ``"all"`` when there are no bundle skills.
    """
    bundle_skills: dict[str, pathlib.Path] = {}
    if bundle_dir is not None:
        skills_root = bundle_dir / "skills"
        if skills_root.is_dir():
            for child in sorted(skills_root.iterdir()):
                if child.is_dir() and (child / "SKILL.md").is_file():
                    bundle_skills[child.name] = child

    if skills_filter == "all":
        # Pi auto-discovers host skills; we explicitly add bundled
        # ones so the agent definitely sees them regardless of cwd.
        args: list[str] = []
        for path in bundle_skills.values():
            args.extend(["--skill", str(path)])
        return args
    if skills_filter == "none":
        # ``--no-skills`` suppresses Pi's discovery walk AND loading.
        # No ``--skill`` flags either — explicit paths would override
        # ``--no-skills`` (Pi loads what's explicitly named).
        return ["--no-skills"]
    if isinstance(skills_filter, list):
        args = ["--no-skills"]
        for name in skills_filter:
            if name in bundle_skills:
                args.extend(["--skill", str(bundle_skills[name])])
        return args
    # Unknown shape — fall back to Pi's defaults (auto-discovery on,
    # no explicit skills). The harness wrap's resolver should
    # have validated upstream, so this branch is belt-and-suspenders.
    return []


def _extract_pi_turn_usage(
    message: object,
    fallback_model: str | None,
) -> dict[str, Any] | None:  # type: ignore[explicit-any]
    """Map a Pi assistant message's ``usage`` object onto the wire shape
    that :class:`TurnComplete` consumes, so pi sub-agent cost is priced
    the same way as ``claude-sdk`` and ``codex`` turns.

    Pi (``@earendil-works/pi-coding-agent``) forwards assistant messages
    whose ``usage`` dict carries ``input`` / ``output`` / ``cacheRead`` /
    ``cacheWrite`` / ``totalTokens`` token counts, and the message itself
    carries the resolved ``model``. This translates those into omnigent's
    usage schema (see :class:`omnigent.inner.executor.TurnComplete`).

    :param message: A pi message dict (e.g. ``event["message"]`` from a
        ``message_end`` event, or an entry from ``event["messages"]`` on
        ``agent_end``). Defensive against non-dict shapes.
    :param fallback_model: The executor's configured model, used for
        cost pricing when the assistant message omits ``model``.
    :returns: The mapped usage dict, or ``None`` when ``message`` is not
        an assistant message carrying a ``usage`` dict — callers leave
        ``TurnComplete.usage`` as ``None`` in that case.
    """
    if not isinstance(message, dict):
        return None
    if message.get("role") != "assistant":
        return None
    usage = message.get("usage")
    if not isinstance(usage, dict):
        return None
    raw_model = message.get("model")
    model = raw_model if isinstance(raw_model, str) and raw_model else fallback_model
    return {
        "input_tokens": int(usage.get("input") or 0),
        "output_tokens": int(usage.get("output") or 0),
        "total_tokens": int(usage.get("totalTokens") or 0),
        "cache_read_input_tokens": int(usage.get("cacheRead") or 0),
        "cache_creation_input_tokens": int(usage.get("cacheWrite") or 0),
        "model": model,
    }


def _aggregate_pi_turn_usage(
    message_usages: list[dict[str, Any]],  # type: ignore[explicit-any]
    fallback_model: str | None,
) -> dict[str, Any] | None:  # type: ignore[explicit-any]
    """Aggregate per-message Pi usage into one turn-level usage dict.

    A single Omnigent turn drives Pi's full agent loop, which may make
    several LLM calls (one assistant message per iteration of a tool-use
    loop), each forwarded as its own ``message_end``. Mirrors the
    openai-agents executor's billing/context split (see
    ``openai_agents_sdk_executor``): token counts are SUMMED across those
    calls because each call is billed for the full input it sends, while
    ``context_tokens`` reflects only the LAST call's total — the proxy for
    how full the context window is going into the next request (summing
    inputs would double-count the re-sent history).

    :param message_usages: Per-message usage dicts from
        :func:`_extract_pi_turn_usage`, in arrival order. Empty when pi
        reported no usage for the turn.
    :param fallback_model: The executor's configured model id, used only
        when no captured message carried a ``model``,
        e.g. ``"databricks-claude-sonnet-4-6"``.
    :returns: A turn-level usage dict for ``TurnComplete.usage`` carrying
        summed ``input_tokens`` / ``output_tokens`` / ``total_tokens`` /
        ``cache_read_input_tokens`` / ``cache_creation_input_tokens``, a
        last-call ``context_tokens``, and the pricing ``model``; or
        ``None`` when no usable usage was captured.
    """
    if not message_usages:
        return None
    input_tokens = sum(u["input_tokens"] for u in message_usages)
    output_tokens = sum(u["output_tokens"] for u in message_usages)
    total_tokens = sum(u["total_tokens"] for u in message_usages)
    cache_read = sum(u["cache_read_input_tokens"] for u in message_usages)
    cache_write = sum(u["cache_creation_input_tokens"] for u in message_usages)
    # No countable tokens means pi reported empty usage objects — treat as
    # "no usage" so the server leaves the session unpriced rather than
    # recording a $0.00 turn.
    if not (input_tokens or output_tokens):
        return None
    last = message_usages[-1]
    # context_tokens = the LAST call's total context (the proxy for
    # next-request context fill); recompute from components when the
    # provider omitted ``totalTokens``.
    context_tokens = last["total_tokens"] or (
        last["input_tokens"]
        + last["output_tokens"]
        + last["cache_read_input_tokens"]
        + last["cache_creation_input_tokens"]
    )
    # _extract_pi_turn_usage already applied the per-message
    # message-model-else-fallback precedence, so the last entry's model is
    # the authoritative one to price the turn with.
    model = last["model"] if last["model"] else fallback_model
    return {
        "input_tokens": input_tokens,
        "output_tokens": output_tokens,
        "total_tokens": total_tokens or (input_tokens + output_tokens),
        "cache_read_input_tokens": cache_read,
        "cache_creation_input_tokens": cache_write,
        "context_tokens": context_tokens,
        "model": model,
    }


class PiExecutor(Executor):
    """Execute agent turns via the Pi coding agent (``pi --mode rpc``)."""

    def __init__(
        self,
        *,
        cwd: str | None = None,
        os_env: OSEnvSpec | None = None,
        model: str | None = None,
        pi_path: str | None = None,
        gateway: bool = False,
        databricks_profile: str | None = None,
        gateway_host: str | None = None,
        base_url_override: str | None = None,
        base_urls_override: dict[str, str] | None = None,
        gateway_auth_command: str | None = None,
        retry_policy: RetryPolicy | None = None,
        bundle_dir: pathlib.Path | None = None,
        agent_name: str | None = None,
        skills_filter: str | list[str] = "all",
    ) -> None:
        """Create a PiExecutor.

        :param cwd: Working directory for the Pi subprocess.
        :param os_env: Optional OS environment / sandbox spec.  When set, the
            Pi subprocess is wrapped in the same sandbox other
            harnesses use.
        :param model: Override the model name, e.g. ``"databricks-claude-sonnet-4-6"``.
        :param pi_path: Absolute path to a ``pi`` CLI binary.  When ``None``
            the executor searches ``PATH``.
        :param gateway: When ``True``, write a ``models.json`` pointing Pi
            at a vendor-neutral gateway. The Databricks AI gateway is one
            producer of this transport; generic providers are another.
        :param databricks_profile: Databricks-specific config profile from
            ``~/.databrickscfg``, e.g. ``"<your-profile>"``.  Only used on the
            Databricks producer path. ``None`` falls back to
            ``DATABRICKS_CONFIG_PROFILE`` then the first valid profile.
        :param gateway_host: Gateway workspace host origin, e.g.
            ``"https://example.databricks.com"``.  Set from
            ``HARNESS_PI_GATEWAY_HOST`` (written by the Omnigent workflow layer).
            When set, skips profile host lookup.
        :param base_url_override: Override the workspace host used when
            building Pi's ``models.json``.  Expected to be the Anthropic
            gateway URL from ucode state.json (``base_urls.claude``), e.g.
            ``"https://example.databricks.com/ai-gateway/anthropic"``.
            The host portion is extracted and used as the base for all
            Pi provider URLs.  ``None`` falls back to deriving the host from
            the Databricks profile credentials.
        :param base_urls_override: ucode-provided Pi gateway URLs keyed by
            model family, e.g. ``{"claude": "...", "openai": "..."}``.
        :param gateway_auth_command: Shell command that prints a bearer token,
            e.g.
            ``"databricks auth token --host https://example.databricks.com ..."``
            or ``"printf %s sk-..."``. Set from
            ``HARNESS_PI_GATEWAY_AUTH_COMMAND``.
        :param retry_policy: The spec's ``llm.retry`` budget. Threads
            ``policy.pi.settings()`` into Pi's ``.pi/settings.json``
            before subprocess spawn — Pi natively does exponential
            backoff and Retry-After honoring; this just sets the
            budget. ``None`` resolves to ``RetryPolicy()`` defaults.
            See Phase 1f of ``designs/RETRY_ACROSS_HARNESSES.md``.
        :param bundle_dir: The agent bundle's extracted on-disk path.
            When set, ``<bundle_dir>/skills/<name>/SKILL.md`` files
            are exposed to Pi via ``--skill <path>`` based on
            *skills_filter*. ``None`` skips bundle-skill wiring.
        :param agent_name: Optional agent display name. Reserved for
            future use; currently unused by Pi.
        :param skills_filter: Host-skill filter (``"all"`` / ``"none"``
            / ``list[str]``). ``"all"`` (default) lets Pi's
            auto-discovery run AND adds explicit ``--skill`` flags
            for each bundle skill; ``"none"`` adds ``--no-skills``
            so Pi sees zero skills; a list adds ``--no-skills`` plus
            ``--skill`` for each named bundle skill — names not
            present in the bundle are silently skipped.
        """
        resolved_pi = pi_path or _find_pi_cli()
        if not resolved_pi:
            raise ImportError(
                "PiExecutor requires the 'pi' CLI on PATH. "
                "Install it with: npm install -g @earendil-works/pi-coding-agent"
            )
        self._pi_path = resolved_pi
        self._cwd = cwd
        self._os_env_spec = os_env
        self._model_override = model
        self._gateway = gateway
        self._databricks_profile = databricks_profile
        self._gateway_host_override = gateway_host.rstrip("/") if gateway_host else None
        self._base_url_override = base_url_override
        self._base_urls_override = base_urls_override
        self._gateway_auth_command = gateway_auth_command
        # Retry policy → Pi's .pi/settings.json before subprocess spawn.
        # See ``RetryPolicy.pi.settings()`` for the schema. Pi natively
        # does exponential backoff + Retry-After honoring; we just set
        # the budget.
        self._retry_policy = retry_policy if retry_policy is not None else RetryPolicy()
        # Allowlisted base env for the Pi subprocess — never the full
        # host environ. The spec's os_env.sandbox.env_passthrough is the
        # opt-in for anything beyond the base set (e.g. a provider API
        # key on a direct, non-gateway run).
        self._env: dict[str, str] = _clean_pi_env(
            os_env.sandbox.env_passthrough if os_env is not None and os_env.sandbox else None
        )
        # ``--no-tools`` in pi 0.68.x disables BOTH built-ins AND extension
        # tools by default; we re-enable just the bridged tool names per
        # turn via ``--tools <comma-list>`` in :meth:`_build_env_and_dir`.
        # The combined effect is: pi's native read/bash/edit/write stay
        # off (they don't route through Omnigent policies / history and
        # can 400 against the Databricks Responses API), and the bridge
        # extension's tools are explicitly allowlisted.
        self._extra_args: list[str] = ["--no-tools"]
        self._bundle_dir = bundle_dir
        self._agent_name = agent_name
        self._skills_filter = skills_filter
        # Resolve once at construction (the bundle layout doesn't
        # change across turns within a session). Each turn's
        # ``_build_env_and_dir`` copies ``self._extra_args`` so this
        # extension is read-only after init.
        self._extra_args.extend(_resolve_pi_skill_args(skills_filter, bundle_dir))
        # Set by Session._wire_sdk_executor().
        self._tool_executor: ToolExecutor | None = None

        if gateway:
            creds = None
            if self._gateway_host_override is None:
                creds = _read_databrickscfg(databricks_profile)
                if creds is None:
                    raise OSError(
                        "PiExecutor(gateway=True) requires gateway credentials via "
                        "the gateway base URL / auth command or a valid "
                        "~/.databrickscfg profile."
                    )
            if self._gateway_host_override is not None:
                self._databricks_host = self._gateway_host_override
                if gateway_auth_command is None:
                    raise OSError(
                        "PiExecutor(gateway=True) with a gateway workspace host "
                        "requires a gateway auth command."
                    )
                token = _fetch_shell_command_token(gateway_auth_command)
                if token is None:
                    raise OSError(
                        "PiExecutor(gateway=True) could not fetch a gateway token "
                        "for the workspace host."
                    )
                self._databricks_token = token
            elif base_url_override:
                # Derive the workspace host from the Anthropic gateway URL
                # ucode wrote to state.json.
                _parsed = _urlparse(base_url_override)
                self._databricks_host = f"{_parsed.scheme}://{_parsed.netloc}"
                assert creds is not None
                self._databricks_token = creds.token
            else:
                assert creds is not None
                self._databricks_host = creds.host
                self._databricks_token = creds.token

        # True when the gateway transport was derived from a ~/.databrickscfg
        # profile (no gateway host or base URL supplied directly). Gates the
        # Databricks default model in :meth:`_resolve_model`; on the ucode /
        # generic-provider paths the producer resolves a concrete model.
        self._gateway_uses_databricks_profile = bool(
            gateway and self._gateway_host_override is None and base_url_override is None
        )

        # Apply sandbox.
        # ``PI_CODING_AGENT_DIR`` is added per-spawn by
        # ``_build_env_and_dir`` on the gateway path, so it must be in
        # the launcher's prune allowlist even though it's not in
        # ``self._env`` yet.
        sandboxed = _try_sandbox_pi(
            self._pi_path,
            os_env,
            cwd,
            spawn_env_names=[*self._env, "PI_CODING_AGENT_DIR"],
        )
        self._pi_launch_path = sandboxed.launch_path
        self._sandboxed = sandboxed.sandboxed

        self._session_states: dict[str, _PiSessionState] = {}
        self._tool_server: _ToolServer | None = None

    def supports_streaming(self) -> bool:
        return True

    def supports_tool_calling(self) -> bool:
        return True

    def handles_tools_internally(self) -> bool:
        return True

    def supports_live_message_queue(self) -> bool:
        return True

    async def enqueue_session_message(  # type: ignore[explicit-any]
        self,
        session_key: str,
        content: str | dict[str, Any],
    ) -> bool:
        """Send a steering message to Pi mid-turn.

        Pi's RPC ``steer`` command injects a user message between tool calls
        within the current agent turn, so the model sees it before its next
        response.
        """
        state = self._session_states.get(session_key)
        if state is None or state.rpc is None:
            return False
        text = content if isinstance(content, str) else str(content)
        try:
            await state.rpc.send_command(
                {
                    "type": "steer",
                    "message": text,
                }
            )
            return True
        except Exception as exc:  # noqa: BLE001 — steer is best-effort; any failure surfaces as False
            logger.debug("PiExecutor: failed to enqueue steer message: %s", exc)
            return False

    async def close_session(self, session_key: str) -> None:
        state = self._session_states.pop(session_key, None)
        if state is not None and state.rpc is not None:
            await state.rpc.close()

    async def interrupt_session(self, session_key: str) -> bool:
        state = self._session_states.get(session_key)
        if state is None or state.rpc is None:
            return False
        # Best-effort abort to halt the in-flight turn.
        try:
            await state.rpc.send_command({"type": "abort", "id": "abort"})
        except Exception as exc:  # noqa: BLE001 — abort is best-effort
            logger.debug("PiExecutor: interrupt abort failed: %s", exc)
        # Always drop the session so the next turn starts fresh and replays
        # full history. A resumed subprocess sends only the latest user
        # message, which would bypass the runner's "[System: interrupted]"
        # marker and silently continue the abandoned request. See
        # claude_sdk_executor.interrupt_session for the rationale.
        try:
            await self.close_session(session_key)
            return True
        except Exception as exc:  # noqa: BLE001 — close failures surface as False
            logger.debug("PiExecutor: session close after interrupt failed: %s", exc)
            return False

    async def close(self) -> None:
        keys = list(self._session_states.keys())
        for key in keys:
            await self.close_session(key)
        if self._tool_server is not None:
            await self._tool_server.stop()
            self._tool_server = None

    def _session_key(self, messages: list[Message]) -> str:
        if messages:
            last = messages[-1]
            if last.get("session_id"):
                return str(last["session_id"])
            meta = last.get("metadata", {})
            if meta.get("session_id"):
                return str(meta["session_id"])
        return "__default__"

    def _resolve_model(self, config: ExecutorConfig | None) -> str | None:
        """
        Determine the model name to pass to Pi.

        ``cfg.model`` (per-request /model override) wins over the spec
        default (``HARNESS_PI_MODEL`` → ``self._model_override``). On the
        Databricks-profile gateway path a missing model falls back to
        :data:`DATABRICKS_CLAUDE_DEFAULT_MODEL` — Pi's own default is an
        Anthropic-direct id the gateway rejects. Elsewhere ``None`` falls
        through to let Pi pick its own default.

        :param config: Optional :class:`ExecutorConfig` whose ``model``
            takes precedence when set.
        :returns: The resolved model string, or ``None`` when both
            sources are unset off the Databricks-profile gateway path.
        """
        cfg = config or ExecutorConfig()
        model = cfg.model or self._model_override
        if model is None and self._gateway_uses_databricks_profile:
            return DATABRICKS_CLAUDE_DEFAULT_MODEL
        return model

    async def _ensure_tool_server(self, tools: list[ToolSpec]) -> int | None:
        """Start the TCP tool server if there are Omnigent tools to bridge."""
        if not tools:
            return None
        if self._tool_server is None:
            self._tool_server = _ToolServer()
            await self._tool_server.start()
        self._tool_server._tool_executor = self._tool_executor
        self._tool_server._policy_gate = self._gate_native_tool
        return self._tool_server.port

    async def _gate_native_tool(  # type: ignore[explicit-any]
        self,
        name: str,
        args: dict[str, Any],
    ) -> dict[str, Any]:
        """Evaluate a native Pi tool call against Omnigent TOOL_CALL policy.

        Bridges the tool server's :attr:`_ToolServer._policy_gate` to the
        ``_policy_evaluator`` the harness scaffold installs on this executor
        (the same round-trip the Claude SDK executor uses for LLM_REQUEST /
        LLM_RESPONSE policies). Mirrors the claude-native / codex-native
        PreToolUse hooks: the verdict is computed by the Omnigent server
        against the session's full policy set (inherited parent session
        policies + the agent spec's guardrails).

        :param name: Native tool name, e.g. ``"read"`` or ``"bash"``.
        :param args: The tool's argument dict.
        :returns: ``{"block": bool, "reason": str}`` — block on DENY.
            Allow when no evaluator is wired (single-process / test paths).
            TOOL_CALL ASK is collapsed to ALLOW/DENY server-side before the
            verdict returns, so only DENY blocks here.
        """
        # ``_policy_evaluator`` is installed best-effort by the harness
        # scaffold's executor adapter; absent on single-process / pre-turn
        # paths. Same fetch pattern as the Claude SDK executor.
        evaluator = getattr(self, "_policy_evaluator", None)
        if evaluator is None:
            return {"block": False, "reason": ""}
        verdict = await evaluator("PHASE_TOOL_CALL", {"name": name, "arguments": args})
        if verdict.action == "POLICY_ACTION_DENY":
            return {"block": True, "reason": verdict.reason or "blocked by policy"}
        return {"block": False, "reason": ""}

    def _build_env_and_dir(
        self,
        tools: list[ToolSpec],
        tool_server_port: int | None,
        tool_server_token: str | None,
        model: str | None,
    ) -> PiSubprocessConfig:
        """Build env dict, temp dir, and extra CLI args for a Pi subprocess.

        :param tools: Omnigent tool schemas to bridge into Pi.
        :param tool_server_port: TCP port the Omnigent tool server is
            listening on, or ``None`` if no tools need bridging.
        :param tool_server_token: The tool server's bearer token
            (:attr:`_ToolServer.token`), embedded in the generated
            extension so requests authenticate. Non-``None`` whenever
            *tool_server_port* is.
        :param model: The resolved model id this subprocess will run, e.g.
            ``"moonshotai/kimi-k2.6"``; registered in the generated
            ``models.json`` on the gateway path so the
            ``provider/<model>`` selector resolves. ``None`` when no model
            is pinned (Pi picks its own default).
        """
        env = dict(self._env)
        tmp_dir = tempfile.mkdtemp(prefix="omnigent_pi_")
        extra_args: list[str] = list(self._extra_args)

        if self._gateway:
            models_json = _build_models_json(
                self._databricks_host,
                self._databricks_token,
                self._base_urls_override,
                model=model,
            )
            models_path = os.path.join(tmp_dir, "models.json")
            with open(models_path, "w") as f:
                json.dump(models_json, f)
            env["PI_CODING_AGENT_DIR"] = tmp_dir
            # Gateway mode relocates Pi's agent root — copy the user's global
            # settings (extensions, packages, …) and symlink install trees so
            # ``pi install`` packages still resolve. See #1423.
            from omnigent.inner.pi_settings import prepare_managed_pi_agent_dir

            prepare_managed_pi_agent_dir(
                pathlib.Path(tmp_dir),
                overlay=self._retry_policy.pi.settings(),
            )

        # Pi natively supports retry config via ``.pi/settings.json``
        # (see ``RetryPolicy.pi.settings()`` for schema). On the non-gateway
        # path, merge the retry block into ``<cwd>/.pi/settings.json``.
        # Gateway runs apply retry via :func:`prepare_managed_pi_agent_dir`
        # into the managed agent dir instead.
        if not self._gateway:
            retry_settings = self._retry_policy.pi.settings()
            settings_dir_root = self._cwd or tmp_dir
            settings_path = os.path.join(settings_dir_root, ".pi", "settings.json")
            try:
                if os.path.exists(settings_path):
                    with open(settings_path) as f:
                        existing_settings = json.load(f)
                    if not isinstance(existing_settings, dict):
                        existing_settings = {}
                else:
                    existing_settings = {}
                existing_settings.update(retry_settings)
                os.makedirs(os.path.dirname(settings_path), exist_ok=True)
                with open(settings_path, "w") as f:
                    json.dump(existing_settings, f, indent=2)
            except OSError:
                fallback_path = os.path.join(tmp_dir, ".pi", "settings.json")
                os.makedirs(os.path.dirname(fallback_path), exist_ok=True)
                with open(fallback_path, "w") as f:
                    json.dump(retry_settings, f, indent=2)

        # Generate the Omnigent tool bridge extension if tools are available.
        if tools and tool_server_port is not None:
            if tool_server_token is None:
                # A port without a token would spawn the bridge
                # unauthenticated; fail loud instead.
                raise ValueError("tool_server_token is required when tool_server_port is set")
            ext_path = os.path.join(tmp_dir, "omnigent_tools.js")
            with open(ext_path, "w") as f:
                f.write(_generate_extension_js(tool_server_port, tools, tool_server_token))
            extra_args.extend(["--extension", ext_path])
            # Allowlist the bridged tool names. ``--no-tools`` (set in
            # __init__) disables every tool by default in pi 0.68+;
            # ``--tools`` adds specific names back. Without this pass
            # the bridge extension's tools register but pi never
            # exposes them to the LLM — symptom: model replies "I
            # don't have a calculate tool available."
            tool_names = [
                name for name in (s.get("name") for s in tools) if isinstance(name, str) and name
            ]
            # Pi's ``formatSkillsForPrompt`` (system-prompt.js:33,112)
            # gates skill-index injection on ``selectedTools`` including
            # ``"read"``. Pi's native ``read`` is a local filesystem read
            # that runs in-process and never traverses the bridged /mcp
            # path — so enabling it lets the model see (and load) the skills
            # we wired via ``--skill <path>``. As a native tool it would
            # otherwise escape all guardrails, so the generated extension's
            # ``tool_call`` hook routes it (and any other native tool) through
            # an Omnigent TOOL_CALL policy verdict; see
            # :func:`_generate_extension_js` and
            # :meth:`PiExecutor._gate_native_tool`.
            if self._skills_filter != "none":
                tool_names.append("read")
            if tool_names:
                extra_args.extend(["--tools", ",".join(tool_names)])

        return PiSubprocessConfig(env=env, tmp_dir=tmp_dir, extra_args=extra_args)

    async def _ensure_rpc(
        self,
        session_key: str,
        system_prompt: str,
        model: str | None,
        tools: list[ToolSpec],
    ) -> _PiRpcSession:
        """Get or create a Pi RPC subprocess for the given session."""
        state = self._session_states.setdefault(session_key, _PiSessionState())

        if (
            state.rpc is not None
            and state.rpc.process is not None
            and state.rpc.process.returncode is None
            and state.system_prompt == system_prompt
            and state.model == model
        ):
            return state.rpc

        if state.rpc is not None:
            await state.rpc.close()

        tool_server_port = await self._ensure_tool_server(tools)
        tool_server_token = self._tool_server.token if self._tool_server is not None else None
        subprocess_config = self._build_env_and_dir(
            tools, tool_server_port, tool_server_token, model
        )
        env = subprocess_config.env
        tmp_dir = subprocess_config.tmp_dir
        extra_args = subprocess_config.extra_args

        rpc = _PiRpcSession()
        rpc._tmp_dir = tmp_dir

        # For Databricks models, prefix with the provider name so Pi resolves
        # the model from our custom provider in models.json.
        pi_model: str | None
        if self._gateway and model:
            provider = _pi_provider_for_model(model)
            pi_model = f"{provider}/{model}"
        else:
            pi_model = model

        await rpc.start(
            self._pi_launch_path,
            env=env,
            cwd=self._cwd,
            model=pi_model or None,
            system_prompt=system_prompt or None,
            extra_args=extra_args or None,
        )
        state.rpc = rpc
        state.system_prompt = system_prompt
        state.model = model
        state._has_sent_prompt = False
        return rpc

    async def run_turn(
        self,
        messages: list[Message],
        tools: list[ToolSpec],
        system_prompt: str,
        config: ExecutorConfig | None = None,
    ) -> AsyncIterator[ExecutorEvent]:
        if self._gateway:
            if self._gateway_host_override is None:
                creds = _read_databrickscfg(self._databricks_profile)
                if creds is not None:
                    self._databricks_token = creds.token
            else:
                assert self._gateway_auth_command is not None
                token = _fetch_shell_command_token(self._gateway_auth_command)
                if token:
                    self._databricks_token = token
        session_key = self._session_key(messages)
        model = self._resolve_model(config)

        try:
            rpc = await self._ensure_rpc(session_key, system_prompt, model, tools)
        except Exception as exc:  # noqa: BLE001 — executor boundary surfaces startup errors as ExecutorError
            yield ExecutorError(message=f"Failed to start Pi: {exc}")
            return

        # Build the prompt to send to Pi.  On the first turn of a new Pi
        # process, if there are prior messages (e.g. parent history passed
        # to a sub-agent), we serialize the full conversation so Pi has
        # context.  On subsequent turns the Pi process already has context,
        # so we only send the latest user message.
        state = self._session_states.get(session_key)
        is_first_turn = state is not None and not state._has_sent_prompt

        prompt = _build_pi_prompt(messages, is_first_turn=is_first_turn)

        if not prompt:
            # No prompt built (e.g. empty message list on a resumed session) —
            # end the turn with no assistant text rather than fabricating one.
            yield TurnComplete(response=None)
            return

        if state is not None:
            state._has_sent_prompt = True

        # Send prompt command. Pi's JSONL protocol carries text in ``message``
        # and images in a native ``images`` field. Split multimodal blocks into
        # both (rather than JSON-encoding them into the text) so the model
        # actually receives attached images instead of a base64 text blob (#515).
        message: str
        images: list[dict[str, str]] = []
        if isinstance(prompt, list):
            try:
                message, images = _split_pi_prompt(prompt)
            except Exception as exc:  # noqa: BLE001 — prompt-prep boundary: any failure (malformed/unsupported block, bad data URI) surfaces as ExecutorError, never an uncaught crash, matching the send_command path below
                yield ExecutorError(message=f"Failed to prepare prompt for Pi: {exc}")
                return
        else:
            message = prompt
        cmd_id = f"turn_{id(messages)}"
        command: dict[str, Any] = {"type": "prompt", "message": message, "id": cmd_id}
        if images:
            command["images"] = images
        try:
            await rpc.send_command(command)
        except Exception as exc:  # noqa: BLE001 — executor boundary surfaces prompt-send errors as ExecutorError
            yield ExecutorError(message=f"Failed to send prompt to Pi: {exc}")
            return

        # Read events until agent_end.
        response_text = ""
        streamed_any = False
        # Per-LLM-call token usage captured from each assistant message pi
        # forwards (``message_end`` is the capture site; ``agent_end`` is a
        # fallback). Summed into a turn-level usage dict at completion so a
        # multi-step (tool-loop) turn bills for every call, not just the
        # last. Empty when pi reports no usage — cost tracking is skipped.
        message_usages: list[dict[str, Any]] = []  # type: ignore[explicit-any]

        while True:
            line = await rpc.read_line(timeout=120.0)
            if line is None:
                if not streamed_any and not response_text:
                    stderr = "\n".join(rpc._stderr_lines) if rpc._stderr_lines else ""
                    stderr_suffix = f" Stderr: {stderr}" if stderr else ""
                    yield ExecutorError(
                        message=f"Pi process ended without response.{stderr_suffix}"
                    )
                else:
                    turn_usage = _aggregate_pi_turn_usage(message_usages, model)
                    _notify_usage_from_dict(model=model, usage=turn_usage)
                    yield TurnComplete(response=response_text, usage=turn_usage)
                return

            try:
                event = json.loads(line)
            except json.JSONDecodeError:
                logger.debug("PiExecutor: non-JSON line: %s", line[:200])
                continue

            raw_event_type = event.get("type")
            event_type: str | None = raw_event_type if isinstance(raw_event_type, str) else None

            # Skip the command-ack response.
            if event_type == "response":
                if not event.get("success", True):
                    yield ExecutorError(message=event.get("error", "Pi command failed"))
                    return
                continue

            # Streaming text and thinking deltas.
            if event_type == "message_update":
                ame = event.get("assistantMessageEvent", {})
                ame_type = ame.get("type")
                if ame_type == "text_delta":
                    raw_delta = ame.get("delta")
                    if isinstance(raw_delta, str) and raw_delta:
                        yield TextChunk(text=raw_delta)
                        response_text += raw_delta
                        streamed_any = True
                elif ame_type == "thinking_start":
                    # Anchors the "Thinking…" indicator before the first delta.
                    yield ReasoningChunk(delta="", event_type="reasoning_started")
                elif ame_type == "thinking_delta":
                    raw_delta = ame.get("delta")
                    if isinstance(raw_delta, str) and raw_delta:
                        # Reasoning stays out of response_text — it is not assistant text.
                        yield ReasoningChunk(delta=raw_delta, event_type="reasoning_text")
                continue

            # Tool execution events.
            if event_type == "tool_execution_start":
                tool_name = event.get("toolName", "unknown")
                args = event.get("args", {})
                yield ToolCallRequest(
                    name=tool_name,
                    args=args if isinstance(args, dict) else {},
                )
                continue

            if event_type == "tool_execution_end":
                tool_name = event.get("toolName", "unknown")
                is_error = event.get("isError", False)
                result = event.get("result")
                # Pi may report isError at the top level OR only inside
                # the result dict (result.isError).  Check both.
                if isinstance(result, dict) and result.get("isError"):
                    is_error = True
                # Detect policy-blocked results coming back from the tool
                # server.  The _execute_tool callback returns
                # {"blocked": True, "reason": "..."} when a policy blocks
                # the call.  By the time it reaches us, the result may be:
                #   - a dict with "blocked" key (direct)
                #   - a dict with "content" list containing JSON text (from Pi extension)
                #   - a string containing JSON with "blocked" key
                result_str = str(result) if result is not None else ""
                is_blocked = False

                def _check_blocked(obj: JsonValue) -> BlockedCheck:
                    """Check if obj represents a blocked tool result."""
                    if isinstance(obj, dict):
                        if obj.get("blocked"):
                            return BlockedCheck(blocked=True, reason=str(obj.get("reason", obj)))
                        # Pi extension wraps in {content: [{type:"text", text:"..."}]}
                        content = obj.get("content")
                        if isinstance(content, list):
                            for item in content:
                                if isinstance(item, dict) and item.get("type") == "text":
                                    text = item.get("text")
                                    if not isinstance(text, str):
                                        continue
                                    try:
                                        parsed = json.loads(text)
                                        if isinstance(parsed, dict) and parsed.get("blocked"):
                                            return BlockedCheck(
                                                blocked=True,
                                                reason=str(parsed.get("reason", text)),
                                            )
                                    except (json.JSONDecodeError, TypeError):
                                        pass
                    elif isinstance(obj, str):
                        try:
                            parsed = json.loads(obj)
                            if isinstance(parsed, dict) and parsed.get("blocked"):
                                return BlockedCheck(
                                    blocked=True,
                                    reason=str(parsed.get("reason", obj)),
                                )
                        except (json.JSONDecodeError, TypeError):
                            pass
                    return BlockedCheck(blocked=False, reason="")

                if is_error:
                    check = _check_blocked(result)
                    is_blocked = check.blocked
                    if is_blocked:
                        result_str = check.reason

                if is_blocked:
                    status = ToolCallStatus.BLOCKED
                elif is_error:
                    status = ToolCallStatus.ERROR
                else:
                    status = ToolCallStatus.SUCCESS

                yield ToolCallComplete(
                    name=tool_name,
                    status=status,
                    result=result,
                    error=result_str if (is_error or is_blocked) else "",
                )
                continue

            # Agent ended — the turn is complete.
            if event_type == "agent_end":
                end_messages = event.get("messages", [])
                if not response_text:
                    for m in reversed(end_messages):
                        if m.get("role") == "assistant":
                            content = m.get("content", [])
                            if isinstance(content, str):
                                response_text = content
                            elif isinstance(content, list):
                                text_parts: list[str] = []
                                for part in content:
                                    if not (isinstance(part, dict) and part.get("type") == "text"):
                                        continue
                                    part_text = part.get("text")
                                    if isinstance(part_text, str):
                                        text_parts.append(part_text)
                                response_text = "".join(text_parts)
                            break
                # Fallback usage capture: if no ``message_end`` carried
                # usage, pull it from the last assistant message in
                # ``messages`` (only the last — ``messages`` may hold the
                # whole conversation, so summing it would overcount).
                if not message_usages and isinstance(end_messages, list):
                    for m in reversed(end_messages):
                        captured = _extract_pi_turn_usage(m, model)
                        if captured is not None:
                            message_usages.append(captured)
                            break
                turn_usage = _aggregate_pi_turn_usage(message_usages, model)
                _notify_usage_from_dict(model=model, usage=turn_usage)
                yield TurnComplete(response=response_text, usage=turn_usage)
                return

            # message_end carries one completed assistant message, whose
            # ``usage`` object holds that LLM call's token counts — collect
            # each for the turn-level sum before handling error stop reasons.
            if event_type == "message_end":
                msg = event.get("message", {})
                if isinstance(msg, dict):
                    captured = _extract_pi_turn_usage(msg, model)
                    if captured is not None:
                        message_usages.append(captured)
                    raw_stop = msg.get("stopReason")
                    stop: str | None = raw_stop if isinstance(raw_stop, str) else None
                    if stop in ("error", "aborted"):
                        err = msg.get("errorMessage", stop)
                        yield ExecutorError(message=str(err))
                        return
                continue

            logger.debug("PiExecutor: ignoring event type=%s", event_type)