omnigent-ai--omnigent/tests/test_pi_native_bridge.py

"""Tests for omnigent.pi_native_bridge inbox enqueue contract."""

from __future__ import annotations

import json
import stat
from pathlib import Path

from omnigent import pi_native_bridge


def _inbox_files(bridge_dir: Path) -> list[str]:
    """Return the inbox's ``*.json`` filenames in the extension's read order.

    The Pi extension delivers files via ``readdirSync(...).sort()`` — plain
    lexicographic order — so sorting the names here mirrors delivery order.

    :param bridge_dir: A prepared Pi bridge directory.
    :returns: Sorted inbox filenames.
    """
    return sorted(p.name for p in (bridge_dir / "inbox").glob("*.json"))


def test_enqueue_preserves_send_order_across_types(tmp_path: Path) -> None:
    """Inbox filenames sort in enqueue order regardless of payload type.

    The extension delivers inbox files in lexicographic order. The payload id
    is a random uuid (no time ordering), and an ``interrupt_`` id sorts before
    a ``msg_`` id — so without an ordering prefix a message queued before an
    interrupt would be delivered *after* it. This pins that send order is
    preserved even when a message is followed by an interrupt.
    """
    bridge_dir = tmp_path / "bridge"
    (bridge_dir / "inbox").mkdir(parents=True)

    first = pi_native_bridge.enqueue_user_message(bridge_dir, "first")
    interrupt = pi_native_bridge.enqueue_interrupt(bridge_dir)
    second = pi_native_bridge.enqueue_user_message(bridge_dir, "second")

    files = _inbox_files(bridge_dir)
    assert len(files) == 3, files
    # Delivery order must match enqueue order, not the lexicographic order of
    # the raw ids (where "interrupt_" would jump ahead of "msg_").
    assert first in files[0], files
    assert interrupt in files[1], files
    assert second in files[2], files


def test_enqueue_compact_payload_shape(tmp_path: Path) -> None:
    """A queued compact request round-trips as well-formed JSON the extension reads.

    The extension dispatches on ``payload.type == "compact"`` and feeds
    ``payload.custom_instructions`` (when present) to Pi's
    ``ctx.compact({customInstructions})``. The id must equal the returned
    ``compact_`` id (the extension dedups on it).
    """
    bridge_dir = tmp_path / "bridge"
    (bridge_dir / "inbox").mkdir(parents=True)

    compact_id = pi_native_bridge.enqueue_compact(bridge_dir)

    (path,) = list((bridge_dir / "inbox").glob("*.json"))
    payload = json.loads(path.read_text(encoding="utf-8"))
    assert payload["id"] == compact_id
    assert compact_id.startswith("compact_")
    assert payload["type"] == "compact"
    assert isinstance(payload["created_at"], (int, float))
    # No instructions given → key omitted so the extension uses Pi's default
    # summarisation rather than passing an empty customInstructions string.
    assert "custom_instructions" not in payload


def test_enqueue_compact_includes_custom_instructions(tmp_path: Path) -> None:
    """Non-blank custom instructions ride along; blank ones are dropped.

    Pi's ``compact({customInstructions})`` focuses the summary on the given
    text. An empty/whitespace string must be dropped so it doesn't override
    Pi's default summarisation with a no-op instruction.
    """
    bridge_dir = tmp_path / "bridge"
    (bridge_dir / "inbox").mkdir(parents=True)

    pi_native_bridge.enqueue_compact(bridge_dir, "focus on the refactor")
    pi_native_bridge.enqueue_compact(bridge_dir, "   ")

    payloads = sorted(
        (json.loads(p.read_text(encoding="utf-8")) for p in (bridge_dir / "inbox").glob("*.json")),
        key=lambda payload: payload["created_at"],
    )
    assert len(payloads) == 2, payloads
    with_instructions = [p for p in payloads if "custom_instructions" in p]
    assert len(with_instructions) == 1, payloads
    assert with_instructions[0]["custom_instructions"] == "focus on the refactor"


def test_enqueue_user_message_payload_shape(tmp_path: Path) -> None:
    """A queued user message round-trips as well-formed JSON the extension reads.

    The extension dedups on ``payload.id`` and delivers ``payload.content`` via
    ``pi.sendUserMessage``, so the id must equal the returned ``msg_`` id and
    the content must be preserved verbatim.
    """
    bridge_dir = tmp_path / "bridge"
    (bridge_dir / "inbox").mkdir(parents=True)

    message_id = pi_native_bridge.enqueue_user_message(bridge_dir, "hello world")

    (path,) = list((bridge_dir / "inbox").glob("*.json"))
    payload = json.loads(path.read_text(encoding="utf-8"))
    assert payload["id"] == message_id
    assert payload["type"] == "user_message"
    assert payload["content"] == "hello world"
    assert isinstance(payload["created_at"], (int, float))


def test_enqueue_leaves_no_partial_tmp_files(tmp_path: Path) -> None:
    """Only the final ``.json`` lands in the inbox — no ``.tmp`` residue.

    The atomic temp-write-then-rename exists so the 250 ms poller never reads a
    half-written file. A leftover ``.tmp`` (or a non-``.json`` name) would mean
    the rename contract regressed.
    """
    bridge_dir = tmp_path / "bridge"
    (bridge_dir / "inbox").mkdir(parents=True)

    pi_native_bridge.enqueue_user_message(bridge_dir, "x")
    pi_native_bridge.enqueue_interrupt(bridge_dir)

    names = [p.name for p in (bridge_dir / "inbox").iterdir()]
    assert all(n.endswith(".json") for n in names), names
    assert not any(n.endswith(".tmp") for n in names), names


def test_prepare_bridge_dir_is_owner_only(tmp_path: Path, monkeypatch) -> None:
    """The bridge dir and its inbox are created 0o700 (per-session isolation).

    The bearer token written alongside the inbox makes owner-only perms the
    isolation boundary between sessions sharing ``~/.omnigent/pi-native``.
    """
    monkeypatch.setattr(pi_native_bridge, "_BRIDGE_ROOT", tmp_path / "pi-native")

    bridge_dir = pi_native_bridge.prepare_bridge_dir("conv_perms")

    assert stat.S_IMODE((bridge_dir).stat().st_mode) == 0o700
    assert stat.S_IMODE((bridge_dir / "inbox").stat().st_mode) == 0o700


def test_clear_inbox_drops_leftover_payloads(tmp_path: Path) -> None:
    """clear_inbox empties the queue so a relaunched Pi process can't replay.

    A fresh Pi process starts with an empty dedup set; without clearing, a
    payload a prior process left behind would replay into the new session.
    """
    bridge_dir = tmp_path / "bridge"
    (bridge_dir / "inbox").mkdir(parents=True)
    pi_native_bridge.enqueue_user_message(bridge_dir, "stale")
    pi_native_bridge.enqueue_interrupt(bridge_dir)
    assert _inbox_files(bridge_dir), "precondition: inbox has payloads"

    pi_native_bridge.clear_inbox(bridge_dir)

    assert _inbox_files(bridge_dir) == []
    # The inbox dir itself survives (only its contents are dropped).
    assert (bridge_dir / "inbox").is_dir()


def test_clear_inbox_is_a_noop_without_an_inbox(tmp_path: Path) -> None:
    """clear_inbox tolerates a bridge dir that has no inbox yet."""
    pi_native_bridge.clear_inbox(tmp_path / "nonexistent")  # must not raise


def test_write_extension_files_embeds_tools(tmp_path: Path) -> None:
    """write_extension_files embeds the tool list so the extension can register it.

    The runner builds the session's Omnigent tool surface (sys_* tools) and
    passes it to write_extension_files; the extension reads ``config.tools`` and
    registers each via ``pi.registerTool``. The config must round-trip the list
    verbatim so the schemas reach the Pi agent unchanged.
    """
    bridge_dir = tmp_path / "bridge"
    tools = [
        {
            "name": "sys_os_read",
            "description": "Read a file",
            "parameters": {"type": "object", "properties": {"path": {"type": "string"}}},
        },
        {
            "name": "sys_os_shell",
            "description": "Run a command",
            "parameters": {"type": "object", "properties": {}},
        },
    ]

    _ext, cfg = pi_native_bridge.write_extension_files(
        bridge_dir,
        session_id="conv_abc",
        server_url="http://omnigent.test/",
        conversation_url="http://omnigent.test/c/conv_abc",
        auth_headers={"Authorization": "Bearer t"},
        tools=tools,
    )

    payload = json.loads(cfg.read_text(encoding="utf-8"))
    assert payload["sessionId"] == "conv_abc"
    # serverUrl trailing slash is stripped so the extension can append paths.
    assert payload["serverUrl"] == "http://omnigent.test"
    assert payload["tools"] == tools


def test_write_extension_files_defaults_tools_to_empty(tmp_path: Path) -> None:
    """Omitting tools writes an empty list (Pi runs with its built-ins only)."""
    bridge_dir = tmp_path / "bridge"

    _ext, cfg = pi_native_bridge.write_extension_files(
        bridge_dir,
        session_id="conv_abc",
        server_url="http://omnigent.test",
        conversation_url="http://omnigent.test/c/conv_abc",
    )

    payload = json.loads(cfg.read_text(encoding="utf-8"))
    assert payload["tools"] == []


def test_refresh_config_auth_headers_replaces_only_auth(tmp_path: Path) -> None:
    """Refreshing the bearer rewrites only ``authHeaders``, leaving the rest.

    The baked launch token dies with the ~1h Databricks OAuth lifetime; the
    runner re-mints it into ``config.json`` each turn and the extension
    re-reads per request. The rewrite must touch only ``authHeaders`` so
    ``serverUrl`` / ``tools`` (read once at startup) stay intact.
    """
    bridge_dir = tmp_path / "bridge"
    pi_native_bridge.write_extension_files(
        bridge_dir,
        session_id="conv_abc",
        server_url="http://omnigent.test",
        conversation_url="http://omnigent.test/c/conv_abc",
        auth_headers={"Authorization": "Bearer stale"},
        tools=[{"name": "t"}],
    )

    assert (
        pi_native_bridge.refresh_config_auth_headers(bridge_dir, {"Authorization": "Bearer fresh"})
        is True
    )
    payload = json.loads(pi_native_bridge.config_path(bridge_dir).read_text(encoding="utf-8"))
    assert payload["authHeaders"] == {"Authorization": "Bearer fresh"}
    # Untouched: the static fields the extension reads once at startup.
    assert payload["serverUrl"] == "http://omnigent.test"
    assert payload["tools"] == [{"name": "t"}]


def test_refresh_config_auth_headers_noops(tmp_path: Path) -> None:
    """No-op (returns False) on empty headers, a missing config, or no change."""
    bridge_dir = tmp_path / "bridge"
    # Missing config: nothing to rewrite.
    assert pi_native_bridge.refresh_config_auth_headers(bridge_dir, {"a": "b"}) is False

    pi_native_bridge.write_extension_files(
        bridge_dir,
        session_id="conv_abc",
        server_url="http://omnigent.test",
        conversation_url="http://omnigent.test/c/conv_abc",
        auth_headers={"Authorization": "Bearer x"},
    )
    # Empty headers (local/unauthenticated) must never blank out a live token.
    assert pi_native_bridge.refresh_config_auth_headers(bridge_dir, {}) is False
    payload = json.loads(pi_native_bridge.config_path(bridge_dir).read_text(encoding="utf-8"))
    assert payload["authHeaders"] == {"Authorization": "Bearer x"}
    # Identical headers: skip the rewrite.
    assert (
        pi_native_bridge.refresh_config_auth_headers(bridge_dir, {"Authorization": "Bearer x"})
        is False
    )