omnigent-ai--omnigent/tests/test_pi_native_extension.py

"""End-to-end tests for the generated pi-native bridge extension."""

from __future__ import annotations

import shutil
import subprocess
from pathlib import Path

import pytest


def test_delivery_cap_drops_followup_without_failed_session_status(
    tmp_path: Path,
) -> None:
    """The extension must not terminal-fail a session when follow-up delivery caps.

    This runs the real JavaScript extension under Node with a real inbox payload
    and mocked Pi/fetch boundaries. Five consecutive ``sendUserMessage`` throws
    should consume the inbox file and emit an informational conversation item,
    never ``external_session_status`` with ``status: "failed"``.
    """
    node = shutil.which("node")
    if node is None:
        pytest.skip("node is required for the pi-native extension e2e test")

    extension_path = (
        Path(__file__).resolve().parents[1]
        / "omnigent"
        / "resources"
        / "pi_native"
        / "omnigent_pi_native_extension.js"
    )

    script = r"""
const assert = require("assert").strict;
const fs = require("fs");
const path = require("path");

const extensionPath = process.argv[1];
const tmpDir = process.argv[2];
const inboxDir = path.join(tmpDir, "inbox");
const payloadPath = path.join(inboxDir, "000-msg.json");
const configPath = path.join(tmpDir, "config.json");

fs.mkdirSync(inboxDir, { recursive: true });
fs.writeFileSync(
  payloadPath,
  JSON.stringify({ id: "msg-1", type: "user_message", content: "follow up" }),
);
fs.writeFileSync(
  configPath,
  JSON.stringify({
    serverUrl: "http://omnigent.test",
    sessionId: "session-1",
    inboxDir,
    authHeaders: { authorization: "Bearer test" },
  }),
);

process.env.OMNIGENT_PI_NATIVE_CONFIG = configPath;

const postedEvents = [];
global.fetch = async (_url, request) => {
  postedEvents.push(JSON.parse(request.body));
  return { ok: true };
};

let pollInbox = null;
global.setInterval = (fn, _ms) => {
  pollInbox = fn;
  return { fakeInterval: true };
};

const handlers = {};
const sendAttempts = [];
const pi = {
  registerCommand() {},
  on(eventName, handler) {
    handlers[eventName] = handler;
  },
  sendUserMessage(content, options) {
    sendAttempts.push({ content, options });
    throw new Error("Pi is not ready");
  },
};

require(extensionPath)(pi);

(async () => {
  assert.equal(typeof handlers.session_start, "function");
  await handlers.session_start({}, {
    sessionManager: { getSessionId: () => "native-session-1" },
    ui: { setTitle() {}, setStatus() {}, notify() {} },
  });
  assert.equal(typeof pollInbox, "function");

  for (let attempt = 0; attempt < 5; attempt += 1) {
    pollInbox();
  }
  await new Promise((resolve) => setImmediate(resolve));

  assert.deepEqual(
    sendAttempts,
    Array.from({ length: 5 }, () => ({
      content: "follow up",
      options: { deliverAs: "followUp" },
    })),
  );
  assert.equal(fs.existsSync(payloadPath), false);
  assert.equal(
    postedEvents.some(
      (event) =>
        event.type === "external_session_status" &&
        event.data &&
        event.data.status === "failed",
    ),
    false,
    JSON.stringify(postedEvents),
  );

  const dropNote = postedEvents.find(
    (event) =>
      event.type === "external_conversation_item" &&
      event.data &&
      event.data.item_type === "error" &&
      event.data.item_data &&
      event.data.item_data.code === "pi_followup_delivery_dropped",
  );
  assert.ok(dropNote, JSON.stringify(postedEvents));
  assert.equal(dropNote.data.item_data.source, "execution");
  assert.match(dropNote.data.response_id, /^pi-deliver-dropped-/);
  // The note must be actionable: include the dropped message id and a preview
  // of its content so an operator can identify what was lost.
  assert.match(dropNote.data.item_data.message, /msg-1/);
  assert.match(dropNote.data.item_data.message, /follow up/);
})().catch((error) => {
  console.error(error && error.stack ? error.stack : error);
  process.exit(1);
});
"""

    result = subprocess.run(
        [node, "-e", script, str(extension_path), str(tmp_path)],
        capture_output=True,
        check=False,
        text=True,
        timeout=10,
    )

    assert result.returncode == 0, result.stdout + result.stderr


def _run_extension_script(node: str, extension_path: Path, script: str) -> None:
    """Run a Node test ``script`` against the real extension; fail on nonzero exit."""
    result = subprocess.run(
        [node, "-e", script, str(extension_path)],
        capture_output=True,
        check=False,
        text=True,
        timeout=10,
    )
    assert result.returncode == 0, result.stdout + result.stderr


def _usage_test_preamble() -> str:
    """Shared Node harness: load the extension with a mocked fetch + Pi.

    Exposes ``postedEvents`` (parsed request bodies), ``handlers`` (the
    registered Pi event handlers), and a ``ctx`` stub.
    """
    return r"""
const assert = require("assert").strict;
const path = require("path");

const extensionPath = process.argv[1];
const configPath = path.join(require("os").tmpdir(), `pi-usage-${process.pid}.json`);
require("fs").writeFileSync(
  configPath,
  JSON.stringify({
    serverUrl: "http://omnigent.test",
    sessionId: "session-1",
    authHeaders: { authorization: "Bearer test" },
  }),
);
process.env.OMNIGENT_PI_NATIVE_CONFIG = configPath;

const postedEvents = [];
global.fetch = async (_url, request) => {
  postedEvents.push(JSON.parse(request.body));
  return { ok: true };
};
global.setInterval = () => ({ fakeInterval: true });

const handlers = {};
const pi = {
  registerCommand() {},
  on(eventName, handler) {
    handlers[eventName] = handler;
  },
};
require(extensionPath)(pi);

const ctx = { ui: { setTitle() {}, setStatus() {}, notify() {} } };

function usageEvents() {
  return postedEvents.filter((e) => e.type === "external_session_usage");
}
"""


def test_message_end_posts_external_session_usage(tmp_path: Path) -> None:
    """A ``message_end`` with Pi usage POSTs ``external_session_usage``.

    Asserts the cumulative token fields and model match what the server prices
    (input is INCLUSIVE of cache reads; cache split sent separately).
    """
    node = shutil.which("node")
    if node is None:
        pytest.skip("node is required for the pi-native extension e2e test")
    extension_path = (
        Path(__file__).resolve().parents[1]
        / "omnigent"
        / "resources"
        / "pi_native"
        / "omnigent_pi_native_extension.js"
    )

    script = (
        _usage_test_preamble()
        + r"""
(async () => {
  assert.equal(typeof handlers.message_end, "function");
  await handlers.message_end(
    {
      message: {
        role: "assistant",
        model: "databricks-claude-sonnet-4-6",
        content: [{ type: "text", text: "hi" }],
        usage: {
          input: 100,
          output: 40,
          cacheRead: 30,
          cacheWrite: 10,
          totalTokens: 180,
        },
      },
    },
    ctx,
  );

  const usage = usageEvents();
  assert.equal(usage.length, 1, JSON.stringify(postedEvents));
  const data = usage[0].data;
  // input total is INCLUSIVE of cacheRead + cacheWrite (the server splits the
  // cache-read portion back out): 100 + 30 + 10 = 140.
  assert.equal(data.cumulative_input_tokens, 140);
  assert.equal(data.cumulative_output_tokens, 40);
  assert.equal(data.cumulative_cache_read_input_tokens, 30);
  assert.equal(data.model, "databricks-claude-sonnet-4-6");
})().catch((error) => {
  console.error(error && error.stack ? error.stack : error);
  process.exit(1);
});
"""
    )
    _run_extension_script(node, extension_path, script)


def test_usage_accumulates_and_dedupes_across_messages(tmp_path: Path) -> None:
    """Per-message usage SUMS into cumulative totals; a re-emitted message is deduped."""
    node = shutil.which("node")
    if node is None:
        pytest.skip("node is required for the pi-native extension e2e test")
    extension_path = (
        Path(__file__).resolve().parents[1]
        / "omnigent"
        / "resources"
        / "pi_native"
        / "omnigent_pi_native_extension.js"
    )

    script = (
        _usage_test_preamble()
        + r"""
(async () => {
  const msgA = {
    id: "msg-a",
    role: "assistant",
    model: "databricks-claude-sonnet-4-6",
    usage: { input: 100, output: 40, cacheRead: 0, cacheWrite: 0, totalTokens: 140 },
  };
  const msgB = {
    id: "msg-b",
    role: "assistant",
    model: "databricks-claude-sonnet-4-6",
    usage: { input: 200, output: 60, cacheRead: 50, cacheWrite: 0, totalTokens: 310 },
  };

  await handlers.message_end({ message: msgA }, ctx);
  await handlers.message_end({ message: msgB }, ctx);
  // Re-emit msgB on turn_end (same id) — must NOT double-count.
  await handlers.turn_end({ message: msgB }, ctx);

  const usage = usageEvents();
  // Two distinct flushes (after A, after B); turn_end re-emit is deduped so it
  // neither counts nor re-POSTs.
  assert.equal(usage.length, 2, JSON.stringify(postedEvents));
  const last = usage[usage.length - 1].data;
  // input: (100) + (200 + 50) = 350 ; output: 40 + 60 = 100 ; cacheRead: 50.
  assert.equal(last.cumulative_input_tokens, 350);
  assert.equal(last.cumulative_output_tokens, 100);
  assert.equal(last.cumulative_cache_read_input_tokens, 50);
})().catch((error) => {
  console.error(error && error.stack ? error.stack : error);
  process.exit(1);
});
"""
    )
    _run_extension_script(node, extension_path, script)


def test_no_usage_message_posts_nothing(tmp_path: Path) -> None:
    """A message with no usage (or empty usage / non-assistant role) POSTs no usage event."""
    node = shutil.which("node")
    if node is None:
        pytest.skip("node is required for the pi-native extension e2e test")
    extension_path = (
        Path(__file__).resolve().parents[1]
        / "omnigent"
        / "resources"
        / "pi_native"
        / "omnigent_pi_native_extension.js"
    )

    script = (
        _usage_test_preamble()
        + r"""
(async () => {
  // No usage object.
  await handlers.message_end(
    { message: { role: "assistant", content: [{ type: "text", text: "hi" }] } },
    ctx,
  );
  // Empty usage (all zeros) — treated as "no usage".
  await handlers.message_end(
    {
      message: {
        role: "assistant",
        usage: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, totalTokens: 0 },
      },
    },
    ctx,
  );
  // Non-assistant role.
  await handlers.message_end(
    { message: { role: "user", usage: { input: 5, output: 0 } } },
    ctx,
  );

  assert.equal(usageEvents().length, 0, JSON.stringify(postedEvents));
})().catch((error) => {
  console.error(error && error.stack ? error.stack : error);
  process.exit(1);
});
"""
    )
    _run_extension_script(node, extension_path, script)


def test_distinct_messages_with_identical_usage_are_not_collapsed(
    tmp_path: Path,
) -> None:
    """Two DISTINCT Pi messages with identical token counts each count once.

    Pi's ``AssistantMessage`` (``@earendil-works/pi-ai``) carries NO ``id`` —
    only an optional ``responseId`` and a required numeric ``timestamp``. Two
    genuinely distinct LLM calls can report identical ``usage`` (e.g. two
    identical short acks under prompt caching); keying dedup on the usage
    counts alone would collapse the second call and UNDERCOUNT the session.
    The dedup must key on the message identity (``timestamp`` here), so both
    calls accumulate; re-emitting the SAME message (same ``timestamp``) on
    ``turn_end`` must still dedupe.
    """
    node = shutil.which("node")
    if node is None:
        pytest.skip("node is required for the pi-native extension e2e test")
    extension_path = (
        Path(__file__).resolve().parents[1]
        / "omnigent"
        / "resources"
        / "pi_native"
        / "omnigent_pi_native_extension.js"
    )

    script = (
        _usage_test_preamble()
        + r"""
(async () => {
  // Real Pi shape: no `id`, distinct required `timestamp`, IDENTICAL usage.
  const usage = { input: 100, output: 40, cacheRead: 0, cacheWrite: 0, totalTokens: 140 };
  const msg1 = {
    role: "assistant",
    model: "databricks-claude-sonnet-4-6",
    timestamp: 1000,
    usage: { ...usage },
  };
  const msg2 = {
    role: "assistant",
    model: "databricks-claude-sonnet-4-6",
    timestamp: 2000,
    usage: { ...usage },
  };

  await handlers.message_end({ message: msg1 }, ctx);
  await handlers.message_end({ message: msg2 }, ctx);
  // Re-emit msg2 (same timestamp) on turn_end — must NOT double-count.
  await handlers.turn_end({ message: msg2 }, ctx);

  const events = usageEvents();
  // Two distinct flushes (after msg1, after msg2); the re-emit is deduped.
  assert.equal(events.length, 2, JSON.stringify(postedEvents));
  const last = events[events.length - 1].data;
  // BOTH distinct calls counted despite identical usage: input 100+100=200,
  // output 40+40=80. (A counts-only fingerprint would wrongly stay at 100/40.)
  assert.equal(last.cumulative_input_tokens, 200);
  assert.equal(last.cumulative_output_tokens, 80);
})().catch((error) => {
  console.error(error && error.stack ? error.stack : error);
  process.exit(1);
});
"""
    )
    _run_extension_script(node, extension_path, script)


def test_agent_end_dedupes_real_shaped_messages_by_timestamp(
    tmp_path: Path,
) -> None:
    """The ``agent_end`` whole-conversation re-scan dedupes real Pi messages.

    ``agent_end`` carries the full ``messages`` array and re-scans it as a
    last-chance capture. Real Pi messages have no ``id``, so the dedup keys on
    ``timestamp``; a message already counted on ``message_end`` must be a no-op
    when it reappears in the ``agent_end`` array (no overcount).
    """
    node = shutil.which("node")
    if node is None:
        pytest.skip("node is required for the pi-native extension e2e test")
    extension_path = (
        Path(__file__).resolve().parents[1]
        / "omnigent"
        / "resources"
        / "pi_native"
        / "omnigent_pi_native_extension.js"
    )

    script = (
        _usage_test_preamble()
        + r"""
(async () => {
  const msg = {
    role: "assistant",
    model: "databricks-claude-sonnet-4-6",
    timestamp: 4242,
    usage: { input: 300, output: 50, cacheRead: 20, cacheWrite: 10, totalTokens: 380 },
  };

  // Counted on message_end.
  await handlers.message_end({ message: msg }, ctx);
  // agent_end re-scans the whole conversation including the same message —
  // must NOT re-count it (same timestamp).
  await handlers.agent_end({ messages: [msg] }, ctx);

  const events = usageEvents();
  assert.equal(events.length, 1, JSON.stringify(postedEvents));
  const last = events[events.length - 1].data;
  // input INCLUSIVE of cacheRead + cacheWrite: 300 + 20 + 10 = 330, counted once.
  assert.equal(last.cumulative_input_tokens, 330);
  assert.equal(last.cumulative_output_tokens, 50);
  assert.equal(last.cumulative_cache_read_input_tokens, 20);
})().catch((error) => {
  console.error(error && error.stack ? error.stack : error);
  process.exit(1);
});
"""
    )
    _run_extension_script(node, extension_path, script)


def _extension_path() -> Path:
    return (
        Path(__file__).resolve().parents[1]
        / "omnigent"
        / "resources"
        / "pi_native"
        / "omnigent_pi_native_extension.js"
    )


def _run_node(script: str, *args: str) -> subprocess.CompletedProcess[str]:
    node = shutil.which("node")
    if node is None:
        pytest.skip("node is required for the pi-native extension e2e test")
    return subprocess.run(
        [node, "-e", script, *args],
        capture_output=True,
        check=False,
        text=True,
        timeout=10,
    )


# Shared Node preamble: load the real extension with mocked fetch/setInterval/pi,
# drive its event handlers, and expose the posted event bodies.
_STREAMING_HARNESS = r"""
const assert = require("assert").strict;
const fs = require("fs");
const path = require("path");

const extensionPath = process.argv[1];
const tmpDir = process.argv[2];
const configPath = path.join(tmpDir, "config.json");

fs.writeFileSync(
  configPath,
  JSON.stringify({ serverUrl: "http://omnigent.test", sessionId: "session-1" }),
);
process.env.OMNIGENT_PI_NATIVE_CONFIG = configPath;

const posted = [];
global.fetch = async (_url, request) => {
  posted.push(JSON.parse(request.body));
  return { ok: true, async json() { return { result: "" }; } };
};
global.setInterval = () => ({ fakeInterval: true });

const handlers = {};
const pi = {
  registerCommand() {},
  on(name, handler) { handlers[name] = handler; },
  sendUserMessage() {},
};

require(extensionPath)(pi);

const ctx = { isIdle: () => false, ui: { setTitle() {}, setStatus() {}, notify() {} } };

// Helpers to build the Pi AssistantMessageEvent shapes the extension consumes.
function textDelta(contentIndex, delta) {
  return { type: "text_delta", contentIndex, delta, partial: {} };
}
function textEnd(contentIndex, content) {
  return { type: "text_end", contentIndex, content, partial: {} };
}
async function feed(assistantMessageEvent) {
  await handlers.message_update({ assistantMessageEvent }, ctx);
}
async function endMessage(text) {
  await handlers.message_end(
    {
      message: {
        role: "assistant",
        content: [{ type: "text", text }],
      },
    },
    ctx,
  );
}

function deltas() {
  return posted.filter((e) => e.type === "external_output_text_delta");
}
function items() {
  return posted.filter((e) => e.type === "external_conversation_item");
}
function assistantText(item) {
  return item.data.item_data.content.map((b) => b.text).join("");
}
"""


def test_text_deltas_post_incrementally_with_stable_id() -> None:
    """Each token posts as an external_output_text_delta with a stable id.

    Drives the real extension with a sequence of Pi ``text_delta`` events for
    one assistant message, then ``message_end``. Asserts: every chunk shares one
    ``message_id``, chunk ``index`` is monotonic from 0, the joined deltas equal
    the streamed text, a final marker closes the stream, and the authoritative
    assistant item carries the full text exactly once (no duplication).
    """
    script = (
        _STREAMING_HARNESS
        + r"""
(async () => {
  await handlers.agent_start({}, ctx);
  await handlers.turn_start({ turnIndex: 1 }, ctx);

  const chunks = ["Hello", ", ", "world", "!"];
  for (const c of chunks) await feed(textDelta(0, c));
  await feed(textEnd(0, chunks.join("")));
  await endMessage(chunks.join(""));

  const ds = deltas();
  // One text chunk per delta plus a single final marker.
  const textChunks = ds.filter((d) => d.data.delta !== "");
  const finals = ds.filter((d) => d.data.final === true);
  assert.equal(textChunks.length, chunks.length, JSON.stringify(ds));
  assert.equal(finals.length, 1, JSON.stringify(ds));

  // Stable message_id across every chunk and the final marker.
  const ids = new Set(ds.map((d) => d.data.message_id));
  assert.equal(ids.size, 1, "expected one stable message_id: " + JSON.stringify([...ids]));
  const messageId = [...ids][0];
  assert.ok(typeof messageId === "string" && messageId.length > 0);

  // Monotonic, gapless index from 0.
  const indices = ds.map((d) => d.data.index);
  assert.deepEqual(indices, indices.map((_, i) => i), JSON.stringify(indices));

  // Joined streamed deltas equal the streamed text.
  assert.equal(textChunks.map((d) => d.data.delta).join(""), chunks.join(""));

  // The final marker is last and carries no new text.
  assert.equal(ds[ds.length - 1].data.final, true);
  assert.equal(ds[ds.length - 1].data.delta, "");

  // Exactly one authoritative assistant item, carrying the full text once.
  const assistantItems = items().filter(
    (i) => i.data.item_type === "message" && i.data.item_data.role === "assistant",
  );
  assert.equal(assistantItems.length, 1, JSON.stringify(assistantItems));
  assert.equal(assistantText(assistantItems[0]), chunks.join(""));
})().catch((e) => { console.error(e && e.stack ? e.stack : e); process.exit(1); });
"""
    )
    result = _run_node(script, str(_extension_path()), "/tmp")
    assert result.returncode == 0, result.stdout + result.stderr


def test_multiple_text_blocks_share_one_message_preview(tmp_path: Path) -> None:
    """Multiple text blocks in one message stream under one message_id.

    The web UI finalizes the oldest live preview per authoritative item (FIFO,
    one item per message), so a message with two text blocks (e.g. text → tool
    call → text) must stream as ONE growing preview — not two — or the second
    preview is orphaned. Asserts both blocks' chunks share one ``message_id``
    with a single monotonic index.
    """
    script = (
        _STREAMING_HARNESS
        + r"""
(async () => {
  await handlers.agent_start({}, ctx);
  await handlers.turn_start({ turnIndex: 1 }, ctx);

  // Text block 0, a tool call at index 1, then text block 2.
  await feed(textDelta(0, "First "));
  await feed(textDelta(0, "part."));
  await feed(textEnd(0, "First part."));
  await feed(textDelta(2, " Second "));
  await feed(textDelta(2, "part."));
  await feed(textEnd(2, " Second part."));
  await endMessage("First part. Second part.");

  const ds = deltas();
  const ids = new Set(ds.map((d) => d.data.message_id));
  assert.equal(ids.size, 1, "both blocks must share one id: " + JSON.stringify([...ids]));

  const textChunks = ds.filter((d) => d.data.delta !== "");
  assert.equal(textChunks.length, 4, JSON.stringify(ds));
  // Single monotonic index spanning both blocks.
  const indices = ds.map((d) => d.data.index);
  assert.deepEqual(indices, indices.map((_, i) => i), JSON.stringify(indices));
  assert.equal(
    textChunks.map((d) => d.data.delta).join(""),
    "First part. Second part.",
  );
})().catch((e) => { console.error(e && e.stack ? e.stack : e); process.exit(1); });
"""
    )
    result = _run_node(script, str(_extension_path()), str(tmp_path))
    assert result.returncode == 0, result.stdout + result.stderr


def test_successive_messages_in_turn_get_distinct_ids(tmp_path: Path) -> None:
    """Two assistant messages in one turn stream under distinct message_ids.

    After a tool round-trip Pi begins a fresh assistant message. Its preview
    must NOT reuse the first message's (finalized) id, or the web UI would fold
    the new text into the already-committed first bubble. Asserts the second
    message's deltas carry a different ``message_id`` whose index restarts at 0.
    """
    script = (
        _STREAMING_HARNESS
        + r"""
(async () => {
  await handlers.agent_start({}, ctx);
  await handlers.turn_start({ turnIndex: 1 }, ctx);

  await feed(textDelta(0, "Looking"));
  await feed(textEnd(0, "Looking"));
  await endMessage("Looking");

  // Second assistant message (after a tool round-trip) in the same turn.
  await feed(textDelta(0, "Done"));
  await feed(textEnd(0, "Done"));
  await endMessage("Done");

  const ds = deltas();
  const ids = [...new Set(ds.map((d) => d.data.message_id))];
  assert.equal(ids.size === undefined ? ids.length : ids.length, 2, JSON.stringify(ids));

  // Group indices per id; each must restart at 0 and be monotonic.
  const byId = new Map();
  for (const d of ds) {
    if (!byId.has(d.data.message_id)) byId.set(d.data.message_id, []);
    byId.get(d.data.message_id).push(d.data.index);
  }
  for (const [, idxs] of byId) {
    assert.deepEqual(idxs, idxs.map((_, i) => i), JSON.stringify(idxs));
  }

  // Two authoritative items, one per message, no cross-contamination.
  const assistantItems = items().filter(
    (i) => i.data.item_type === "message" && i.data.item_data.role === "assistant",
  );
  assert.equal(assistantItems.length, 2);
  assert.equal(assistantText(assistantItems[0]), "Looking");
  assert.equal(assistantText(assistantItems[1]), "Done");
})().catch((e) => { console.error(e && e.stack ? e.stack : e); process.exit(1); });
"""
    )
    result = _run_node(script, str(_extension_path()), str(tmp_path))
    assert result.returncode == 0, result.stdout + result.stderr


def test_message_without_streamed_text_posts_no_delta(tmp_path: Path) -> None:
    """A message_end with no preceding text_delta emits no delta events.

    A tool-only assistant message (or a non-streaming path) must not post a
    stray final marker — there is no live preview to close, so a spurious delta
    could create an empty preview bubble in the web UI.
    """
    script = (
        _STREAMING_HARNESS
        + r"""
(async () => {
  await handlers.agent_start({}, ctx);
  await handlers.turn_start({ turnIndex: 1 }, ctx);

  // No text_delta at all — just the authoritative item.
  await endMessage("");

  assert.equal(deltas().length, 0, JSON.stringify(deltas()));
})().catch((e) => { console.error(e && e.stack ? e.stack : e); process.exit(1); });
"""
    )
    result = _run_node(script, str(_extension_path()), str(tmp_path))
    assert result.returncode == 0, result.stdout + result.stderr


def test_registers_omnigent_tools_and_execute_round_trips(tmp_path: Path) -> None:
    """The extension registers config.tools and execute() round-trips via /mcp.

    Drives the real JavaScript extension under Node with a config carrying a
    flat tool list (as the runner now writes). Asserts each tool is registered
    via ``pi.registerTool`` with its schema, and that calling a registered
    tool's ``execute`` POSTs a JSON-RPC ``tools/call`` to the Omnigent server's
    ``/v1/sessions/{id}/mcp`` proxy and returns the tool output to Pi.
    """
    node = shutil.which("node")
    if node is None:
        pytest.skip("node is required for the pi-native extension e2e test")

    extension_path = (
        Path(__file__).resolve().parents[1]
        / "omnigent"
        / "resources"
        / "pi_native"
        / "omnigent_pi_native_extension.js"
    )

    script = r"""
const assert = require("assert").strict;
const fs = require("fs");
const path = require("path");

const extensionPath = process.argv[1];
const tmpDir = process.argv[2];
const inboxDir = path.join(tmpDir, "inbox");
const configPath = path.join(tmpDir, "config.json");

fs.mkdirSync(inboxDir, { recursive: true });
fs.writeFileSync(
  configPath,
  JSON.stringify({
    serverUrl: "http://omnigent.test",
    sessionId: "conv_abc",
    inboxDir,
    authHeaders: { authorization: "Bearer test" },
    tools: [
      {
        name: "sys_os_read",
        description: "Read a file from the OS environment",
        parameters: {
          type: "object",
          properties: { path: { type: "string" } },
          required: ["path"],
        },
      },
      {
        name: "sys_os_shell",
        description: "Run a shell command",
        parameters: { type: "object", properties: {} },
      },
    ],
  }),
);

process.env.OMNIGENT_PI_NATIVE_CONFIG = configPath;

// Capture every fetch so we can assert the execute() round-trip hits /mcp with
// a JSON-RPC tools/call and the right auth headers.
const fetchCalls = [];
global.fetch = async (url, request) => {
  fetchCalls.push({ url, request });
  // Mimic the Omnigent /mcp proxy success envelope.
  return {
    ok: true,
    async json() {
      return {
        jsonrpc: "2.0",
        id: 1,
        result: { content: [{ type: "text", text: "file contents here" }] },
      };
    },
  };
};

global.setInterval = () => ({ fakeInterval: true });

const registered = {};
const pi = {
  registerCommand() {},
  on() {},
  registerTool(spec) {
    registered[spec.name] = spec;
  },
  sendUserMessage() {},
};

require(extensionPath)(pi);

(async () => {
  // Both configured tools must be registered with their schema.
  assert.ok(registered.sys_os_read, "sys_os_read not registered");
  assert.ok(registered.sys_os_shell, "sys_os_shell not registered");
  assert.equal(registered.sys_os_read.name, "sys_os_read");
  assert.equal(registered.sys_os_read.label, "sys_os_read");
  assert.equal(
    registered.sys_os_read.description,
    "Read a file from the OS environment",
  );
  assert.deepEqual(registered.sys_os_read.parameters, {
    type: "object",
    properties: { path: { type: "string" } },
    required: ["path"],
  });
  assert.equal(typeof registered.sys_os_read.execute, "function");

  // execute() must round-trip through the /mcp proxy and return the output.
  const result = await registered.sys_os_read.execute("call-1", {
    path: "/etc/hosts",
  });

  assert.equal(fetchCalls.length, 1, JSON.stringify(fetchCalls));
  const call = fetchCalls[0];
  assert.equal(call.url, "http://omnigent.test/v1/sessions/conv_abc/mcp");
  assert.equal(call.request.method, "POST");
  assert.equal(call.request.headers.authorization, "Bearer test");
  const body = JSON.parse(call.request.body);
  assert.equal(body.jsonrpc, "2.0");
  assert.equal(body.method, "tools/call");
  assert.equal(body.params.name, "sys_os_read");
  assert.deepEqual(body.params.arguments, { path: "/etc/hosts" });

  // The MCP text content is surfaced to Pi as a single text block.
  assert.ok(result && Array.isArray(result.content), JSON.stringify(result));
  assert.equal(result.content[0].type, "text");
  assert.equal(result.content[0].text, "file contents here");
  assert.equal(result.isError, false);
})().catch((error) => {
  console.error(error && error.stack ? error.stack : error);
  process.exit(1);
});
"""

    result = subprocess.run(
        [node, "-e", script, str(extension_path), str(tmp_path)],
        capture_output=True,
        check=False,
        text=True,
        timeout=10,
    )

    assert result.returncode == 0, result.stdout + result.stderr


def test_bridged_tool_call_skips_hook_policy_eval(tmp_path: Path) -> None:
    """The tool_call hook must NOT re-evaluate policy for bridged Omnigent tools.

    Bridged tools are policy-evaluated server-side inside the /mcp proxy when
    execute() runs, so the hook-level ``policies/evaluate`` call would
    double-evaluate (and, for ASK, double-prompt). The hook must skip bridged
    tool names but still evaluate Pi's own built-in tools.
    """
    node = shutil.which("node")
    if node is None:
        pytest.skip("node is required for the pi-native extension e2e test")

    extension_path = (
        Path(__file__).resolve().parents[1]
        / "omnigent"
        / "resources"
        / "pi_native"
        / "omnigent_pi_native_extension.js"
    )

    script = r"""
const assert = require("assert").strict;
const fs = require("fs");
const path = require("path");

const extensionPath = process.argv[1];
const tmpDir = process.argv[2];
const inboxDir = path.join(tmpDir, "inbox");
const configPath = path.join(tmpDir, "config.json");

fs.mkdirSync(inboxDir, { recursive: true });
fs.writeFileSync(
  configPath,
  JSON.stringify({
    serverUrl: "http://omnigent.test",
    sessionId: "conv_abc",
    inboxDir,
    authHeaders: {},
    tools: [
      { name: "sys_os_read", description: "", parameters: { type: "object", properties: {} } },
    ],
  }),
);

process.env.OMNIGENT_PI_NATIVE_CONFIG = configPath;

const policyUrls = [];
global.fetch = async (url, _request) => {
  if (typeof url === "string" && url.indexOf("/policies/evaluate") !== -1) {
    policyUrls.push(url);
  }
  return { ok: true, async json() { return {}; } };
};
global.setInterval = () => ({ fakeInterval: true });

const handlers = {};
const pi = {
  registerCommand() {},
  on(name, fn) { handlers[name] = fn; },
  registerTool() {},
  sendUserMessage() {},
};

require(extensionPath)(pi);

(async () => {
  const ctx = { isIdle: () => false, abort() {} };
  await handlers.session_start({}, ctx);
  await handlers.agent_start({}, ctx);
  await handlers.turn_start({ turnIndex: 1 }, ctx);

  // Bridged tool: hook must NOT call policies/evaluate (server gates it).
  await handlers.tool_call({ toolCallId: "t1", toolName: "sys_os_read", input: {} }, ctx);
  assert.equal(
    policyUrls.length,
    0,
    "bridged tool must not trigger hook-level policy eval: " + JSON.stringify(policyUrls),
  );

  // Pi's own built-in tool (not bridged): hook MUST evaluate policy.
  await handlers.tool_call({ toolCallId: "t2", toolName: "read", input: {} }, ctx);
  assert.equal(
    policyUrls.length,
    1,
    "non-bridged tool must trigger hook-level policy eval: " + JSON.stringify(policyUrls),
  );
})().catch((error) => {
  console.error(error && error.stack ? error.stack : error);
  process.exit(1);
});
"""

    result = subprocess.run(
        [node, "-e", script, str(extension_path), str(tmp_path)],
        capture_output=True,
        check=False,
        text=True,
        timeout=10,
    )

    assert result.returncode == 0, result.stdout + result.stderr


def test_input_required_approve_round_trips_then_executes(tmp_path: Path) -> None:
    """An ASK-gated bridged tool resolves the elicitation and retries once.

    On the initial ``tools/call`` the /mcp proxy returns an MCP ``input_required``
    (MRTR) envelope instead of executing. The extension must:
      1. long-poll ``/policies/evaluate`` for the human verdict (ALLOW here),
      2. retry ``tools/call`` ONCE with ``requestState`` + ``inputResponses``
         carrying the proxy's elicitation id and ``{action: "accept"}``,
      3. surface the executed output to Pi with ``isError: false``.
    It must NEVER hand the raw ``input_required`` envelope back as success.
    """
    node = shutil.which("node")
    if node is None:
        pytest.skip("node is required for the pi-native extension e2e test")

    extension_path = (
        Path(__file__).resolve().parents[1]
        / "omnigent"
        / "resources"
        / "pi_native"
        / "omnigent_pi_native_extension.js"
    )

    script = r"""
const assert = require("assert").strict;
const fs = require("fs");
const path = require("path");

const extensionPath = process.argv[1];
const tmpDir = process.argv[2];
const inboxDir = path.join(tmpDir, "inbox");
const configPath = path.join(tmpDir, "config.json");

fs.mkdirSync(inboxDir, { recursive: true });
fs.writeFileSync(
  configPath,
  JSON.stringify({
    serverUrl: "http://omnigent.test",
    sessionId: "conv_abc",
    inboxDir,
    authHeaders: { authorization: "Bearer test" },
    tools: [
      {
        name: "sys_os_shell",
        description: "Run a shell command",
        parameters: { type: "object", properties: {} },
      },
    ],
  }),
);

process.env.OMNIGENT_PI_NATIVE_CONFIG = configPath;

const ELICIT_ID = "elicit_abc123";
const REQUEST_STATE = JSON.stringify({ elicitation_id: ELICIT_ID, session_id: "conv_abc" });

const mcpCalls = [];
let mcpCallCount = 0;
global.fetch = async (url, request) => {
  if (typeof url === "string" && url.indexOf("/policies/evaluate") !== -1) {
    // The server-side ASK park collapses to a hard ALLOW after the human approves.
    return { ok: true, async json() { return { result: "POLICY_ACTION_ALLOW" }; } };
  }
  // /mcp proxy.
  mcpCallCount += 1;
  mcpCalls.push({ url, body: JSON.parse(request.body) });
  if (mcpCallCount === 1) {
    // First call → ASK → input_required (MRTR envelope).
    return {
      ok: true,
      async json() {
        return {
          jsonrpc: "2.0",
          id: 1,
          result: {
            resultType: "input_required",
            inputRequests: { [ELICIT_ID]: { method: "elicitation/create", params: {} } },
            requestState: REQUEST_STATE,
          },
        };
      },
    };
  }
  // Retry (post-approval) → executed result.
  return {
    ok: true,
    async json() {
      return { jsonrpc: "2.0", id: 2, result: { content: [{ type: "text", text: "ran ok" }] } };
    },
  };
};
global.setInterval = () => ({ fakeInterval: true });

const registered = {};
const pi = {
  registerCommand() {},
  on() {},
  registerTool(spec) { registered[spec.name] = spec; },
  sendUserMessage() {},
};

require(extensionPath)(pi);

(async () => {
  const result = await registered.sys_os_shell.execute("call-1", {});

  // Exactly two /mcp calls: the initial ASK and ONE approval retry.
  assert.equal(mcpCallCount, 2, JSON.stringify(mcpCalls));
  const retry = mcpCalls[1].body;
  assert.equal(retry.id, 2, "retry must use a different JSON-RPC id");
  assert.equal(retry.params.requestState, REQUEST_STATE);
  assert.deepEqual(retry.params.inputResponses, {
    [ELICIT_ID]: { action: "accept" },
  });

  // The executed output surfaces as a normal (non-error) tool result.
  assert.ok(result && Array.isArray(result.content), JSON.stringify(result));
  assert.equal(result.content[0].text, "ran ok");
  assert.equal(result.isError, false);
})().catch((error) => {
  console.error(error && error.stack ? error.stack : error);
  process.exit(1);
});
"""

    result = subprocess.run(
        [node, "-e", script, str(extension_path), str(tmp_path)],
        capture_output=True,
        check=False,
        text=True,
        timeout=10,
    )

    assert result.returncode == 0, result.stdout + result.stderr


def test_mcp_unreachable_fails_closed_without_throwing(tmp_path: Path) -> None:
    """An unreachable Omnigent MCP server resolves to an error, never a throw.

    Boundary discipline at the /mcp call site: a transport failure (connection
    refused) and an HTTP non-2xx must each resolve ``execute`` to a readable
    ``isError: true`` tool result so the Pi agent loop keeps running, rather than
    rejecting the promise and wedging the turn.
    """
    node = shutil.which("node")
    if node is None:
        pytest.skip("node is required for the pi-native extension e2e test")

    extension_path = (
        Path(__file__).resolve().parents[1]
        / "omnigent"
        / "resources"
        / "pi_native"
        / "omnigent_pi_native_extension.js"
    )

    script = r"""
const assert = require("assert").strict;
const fs = require("fs");
const path = require("path");

const extensionPath = process.argv[1];
const tmpDir = process.argv[2];
const inboxDir = path.join(tmpDir, "inbox");
const configPath = path.join(tmpDir, "config.json");

fs.mkdirSync(inboxDir, { recursive: true });
fs.writeFileSync(
  configPath,
  JSON.stringify({
    serverUrl: "http://omnigent.test",
    sessionId: "conv_abc",
    inboxDir,
    authHeaders: {},
    tools: [
      { name: "sys_os_shell", description: "", parameters: { type: "object", properties: {} } },
    ],
  }),
);

process.env.OMNIGENT_PI_NATIVE_CONFIG = configPath;

let mode = "throw";
global.fetch = async () => {
  if (mode === "throw") throw new Error("ECONNREFUSED 127.0.0.1:7284");
  return { ok: false, status: 503, async json() { return {}; } };
};
global.setInterval = () => ({ fakeInterval: true });

const registered = {};
const pi = {
  registerCommand() {},
  on() {},
  registerTool(spec) { registered[spec.name] = spec; },
  sendUserMessage() {},
};

require(extensionPath)(pi);

(async () => {
  const thrown = await registered.sys_os_shell.execute("call-1", {});
  assert.equal(thrown.isError, true, JSON.stringify(thrown));
  assert.ok(thrown.content[0].text.indexOf("ECONNREFUSED") !== -1, thrown.content[0].text);

  mode = "http";
  const http = await registered.sys_os_shell.execute("call-2", {});
  assert.equal(http.isError, true, JSON.stringify(http));
  assert.ok(http.content[0].text.indexOf("503") !== -1, http.content[0].text);
})().catch((error) => {
  console.error(error && error.stack ? error.stack : error);
  process.exit(1);
});
"""

    result = subprocess.run(
        [node, "-e", script, str(extension_path), str(tmp_path)],
        capture_output=True,
        check=False,
        text=True,
        timeout=10,
    )

    assert result.returncode == 0, result.stdout + result.stderr


def test_input_required_denied_fails_closed_not_false_success(tmp_path: Path) -> None:
    """A declined ASK gate fails CLOSED (isError) — never reports false success.

    When the elicitation park collapses to DENY, the extension retries once with
    ``{action: "decline"}``; the proxy returns a -32000 error which surfaces as an
    ``isError: true`` tool result. The raw ``input_required`` envelope must never
    be returned to the model as a successful result.
    """
    node = shutil.which("node")
    if node is None:
        pytest.skip("node is required for the pi-native extension e2e test")

    extension_path = (
        Path(__file__).resolve().parents[1]
        / "omnigent"
        / "resources"
        / "pi_native"
        / "omnigent_pi_native_extension.js"
    )

    script = r"""
const assert = require("assert").strict;
const fs = require("fs");
const path = require("path");

const extensionPath = process.argv[1];
const tmpDir = process.argv[2];
const inboxDir = path.join(tmpDir, "inbox");
const configPath = path.join(tmpDir, "config.json");

fs.mkdirSync(inboxDir, { recursive: true });
fs.writeFileSync(
  configPath,
  JSON.stringify({
    serverUrl: "http://omnigent.test",
    sessionId: "conv_abc",
    inboxDir,
    authHeaders: {},
    tools: [
      { name: "sys_os_shell", description: "", parameters: { type: "object", properties: {} } },
    ],
  }),
);

process.env.OMNIGENT_PI_NATIVE_CONFIG = configPath;

const ELICIT_ID = "elicit_deny";
const REQUEST_STATE = JSON.stringify({ elicitation_id: ELICIT_ID, session_id: "conv_abc" });

let mcpCallCount = 0;
let lastRetryBody = null;
global.fetch = async (url, request) => {
  if (typeof url === "string" && url.indexOf("/policies/evaluate") !== -1) {
    // The human declined → the ASK park collapses to DENY.
    return { ok: true, async json() { return { result: "POLICY_ACTION_DENY", reason: "nope" }; } };
  }
  mcpCallCount += 1;
  if (mcpCallCount === 1) {
    return {
      ok: true,
      async json() {
        return {
          jsonrpc: "2.0",
          id: 1,
          result: {
            resultType: "input_required",
            inputRequests: { [ELICIT_ID]: { method: "elicitation/create", params: {} } },
            requestState: REQUEST_STATE,
          },
        };
      },
    };
  }
  lastRetryBody = JSON.parse(request.body);
  // Server denies the declined retry with the MCP -32000 convention.
  return {
    ok: true,
    async json() {
      return {
        jsonrpc: "2.0",
        id: 2,
        error: { code: -32000, message: "Tool call denied by user" },
      };
    },
  };
};
global.setInterval = () => ({ fakeInterval: true });

const registered = {};
const pi = {
  registerCommand() {},
  on() {},
  registerTool(spec) { registered[spec.name] = spec; },
  sendUserMessage() {},
};

require(extensionPath)(pi);

(async () => {
  const result = await registered.sys_os_shell.execute("call-1", {});

  assert.equal(mcpCallCount, 2, "expected one approval retry");
  assert.deepEqual(lastRetryBody.params.inputResponses, {
    [ELICIT_ID]: { action: "decline" },
  });

  // Must surface as an error, NOT a false success, and must not leak the raw
  // input_required envelope.
  assert.equal(result.isError, true, JSON.stringify(result));
  assert.ok(result.content[0].text.indexOf("denied") !== -1, result.content[0].text);
  assert.equal(result.content[0].text.indexOf("input_required"), -1, result.content[0].text);
})().catch((error) => {
  console.error(error && error.stack ? error.stack : error);
  process.exit(1);
});
"""

    result = subprocess.run(
        [node, "-e", script, str(extension_path), str(tmp_path)],
        capture_output=True,
        check=False,
        text=True,
        timeout=10,
    )

    assert result.returncode == 0, result.stdout + result.stderr


def test_compact_payload_triggers_ctx_compact_and_brackets_spinner(
    tmp_path: Path,
) -> None:
    """A ``compact`` inbox payload calls ``ctx.compact()`` and brackets the spinner.

    Runs the real JavaScript extension under Node. A queued ``compact`` payload
    must (1) call the resident ``ExtensionContext.compact()`` with the custom
    instructions and ``onComplete``/``onError`` callbacks, (2) post
    ``external_compaction_status`` ``in_progress`` BEFORE compact() so the web UI
    spinner is bracketed, and (3) post ``completed`` once Pi invokes
    ``onComplete``. The payload file must be consumed (unlinked) so compaction
    is not re-triggered every poll tick.
    """
    node = shutil.which("node")
    if node is None:
        pytest.skip("node is required for the pi-native extension e2e test")

    extension_path = (
        Path(__file__).resolve().parents[1]
        / "omnigent"
        / "resources"
        / "pi_native"
        / "omnigent_pi_native_extension.js"
    )

    script = r"""
const assert = require("assert").strict;
const fs = require("fs");
const path = require("path");

const extensionPath = process.argv[1];
const tmpDir = process.argv[2];
const inboxDir = path.join(tmpDir, "inbox");
const payloadPath = path.join(inboxDir, "000-compact.json");
const configPath = path.join(tmpDir, "config.json");

fs.mkdirSync(inboxDir, { recursive: true });
fs.writeFileSync(
  payloadPath,
  JSON.stringify({
    id: "compact-1",
    type: "compact",
    custom_instructions: "focus on the refactor",
  }),
);
fs.writeFileSync(
  configPath,
  JSON.stringify({
    serverUrl: "http://omnigent.test",
    sessionId: "session-1",
    inboxDir,
    authHeaders: { authorization: "Bearer test" },
  }),
);

process.env.OMNIGENT_PI_NATIVE_CONFIG = configPath;

const postedEvents = [];
global.fetch = async (_url, request) => {
  postedEvents.push(JSON.parse(request.body));
  return { ok: true };
};

let pollInbox = null;
global.setInterval = (fn, _ms) => {
  pollInbox = fn;
  return { fakeInterval: true };
};

const handlers = {};
const compactCalls = [];
const pi = {
  registerCommand() {},
  on(eventName, handler) {
    handlers[eventName] = handler;
  },
  sendUserMessage() {},
};

require(extensionPath)(pi);

// The resident context the poller compacts. compact() is fire-and-forget; we
// simulate Pi finishing successfully by invoking the onComplete callback.
const ctx = {
  sessionManager: { getSessionId: () => "native-session-1" },
  ui: { setTitle() {}, setStatus() {}, notify() {} },
  abort() {},
  isIdle: () => false,
  compact(options) {
    compactCalls.push(options);
    if (options && typeof options.onComplete === "function") {
      options.onComplete({ summary: "done" });
    }
  },
};

(async () => {
  // session_start remembers ctx and starts the inbox poller.
  await handlers.session_start({}, ctx);
  assert.equal(typeof pollInbox, "function");

  pollInbox();
  await new Promise((resolve) => setImmediate(resolve));

  // 1) ctx.compact() was called exactly once with the custom instructions and
  // both lifecycle callbacks.
  assert.equal(compactCalls.length, 1, JSON.stringify(compactCalls));
  assert.equal(compactCalls[0].customInstructions, "focus on the refactor");
  assert.equal(typeof compactCalls[0].onComplete, "function");
  assert.equal(typeof compactCalls[0].onError, "function");

  // 2/3) The spinner is bracketed: in_progress posted before compact(), then
  // completed from onComplete. No failed edge on the happy path.
  const compactionStatuses = postedEvents
    .filter((event) => event.type === "external_compaction_status")
    .map((event) => event.data && event.data.status);
  assert.deepEqual(
    compactionStatuses,
    ["in_progress", "completed"],
    JSON.stringify(postedEvents),
  );

  // The payload file is consumed so compaction is not re-triggered each tick.
  assert.equal(fs.existsSync(payloadPath), false);

  // A second poll tick must NOT re-trigger compaction (file gone + deduped).
  pollInbox();
  await new Promise((resolve) => setImmediate(resolve));
  assert.equal(compactCalls.length, 1, "compaction must not re-trigger");
})().catch((error) => {
  console.error(error && error.stack ? error.stack : error);
  process.exit(1);
});
"""

    result = subprocess.run(
        [node, "-e", script, str(extension_path), str(tmp_path)],
        capture_output=True,
        check=False,
        text=True,
        timeout=10,
    )

    assert result.returncode == 0, result.stdout + result.stderr


def test_compact_in_progress_reaches_server_before_completed_when_slow(
    tmp_path: Path,
) -> None:
    """The server receives ``in_progress`` before ``completed`` even if it stalls.

    ``ctx.compact()`` is fire-and-forget and may invoke ``onComplete``
    synchronously. The web spinner is raised by the ``in_progress`` SSE and
    dismissed by ``completed``/``failed``, so a ``completed`` that overtakes a
    slow ``in_progress`` POST would strand the spinner. This pins the fix:
    ``triggerCompaction`` AWAITS the ``in_progress`` POST before calling
    ``ctx.compact()``. The mock fetch records each edge on RESOLUTION (server
    receipt) and delays only ``in_progress``, with ``onComplete`` firing
    synchronously — so without the await the server would see
    ``[completed, in_progress]``. Asserts the server saw ``[in_progress,
    completed]``.
    """
    node = shutil.which("node")
    if node is None:
        pytest.skip("node is required for the pi-native extension e2e test")

    script = r"""
const assert = require("assert").strict;
const fs = require("fs");
const path = require("path");

const extensionPath = process.argv[1];
const tmpDir = process.argv[2];
const inboxDir = path.join(tmpDir, "inbox");
const payloadPath = path.join(inboxDir, "000-compact.json");
const configPath = path.join(tmpDir, "config.json");

fs.mkdirSync(inboxDir, { recursive: true });
fs.writeFileSync(payloadPath, JSON.stringify({ id: "compact-1", type: "compact" }));
fs.writeFileSync(
  configPath,
  JSON.stringify({ serverUrl: "http://omnigent.test", sessionId: "session-1", inboxDir }),
);

process.env.OMNIGENT_PI_NATIVE_CONFIG = configPath;

// Record the order the SERVER receives edges (on resolution), and stall only
// the in_progress POST so a non-awaited completed could overtake it.
const serverOrder = [];
global.fetch = async (_url, request) => {
  const body = JSON.parse(request.body);
  if (
    body.type === "external_compaction_status" &&
    body.data &&
    body.data.status === "in_progress"
  ) {
    await new Promise((resolve) => setTimeout(resolve, 40));
  }
  serverOrder.push(body);
  return { ok: true };
};

let pollInbox = null;
global.setInterval = (fn, _ms) => {
  pollInbox = fn;
  return { fakeInterval: true };
};

const handlers = {};
const pi = {
  registerCommand() {},
  on(eventName, handler) { handlers[eventName] = handler; },
  sendUserMessage() {},
};

require(extensionPath)(pi);

const ctx = {
  sessionManager: { getSessionId: () => "native-session-1" },
  ui: { setTitle() {}, setStatus() {}, notify() {} },
  abort() {},
  isIdle: () => false,
  compact(options) {
    // Fire onComplete synchronously: its (immediate) completed POST must not
    // beat the still-pending slow in_progress POST to the server.
    if (options && typeof options.onComplete === "function") {
      options.onComplete({ summary: "done" });
    }
  },
};

(async () => {
  await handlers.session_start({}, ctx);
  pollInbox();
  await new Promise((resolve) => setTimeout(resolve, 120));

  const statuses = serverOrder
    .filter((event) => event.type === "external_compaction_status")
    .map((event) => event.data && event.data.status);
  assert.deepEqual(statuses, ["in_progress", "completed"], JSON.stringify(serverOrder));
})().catch((error) => {
  console.error(error && error.stack ? error.stack : error);
  process.exit(1);
});
"""

    result = _run_node(script, str(_extension_path()), str(tmp_path))
    assert result.returncode == 0, result.stdout + result.stderr


def test_compact_payload_without_ctx_compact_surfaces_error_and_consumes_file(
    tmp_path: Path,
) -> None:
    """A ``compact`` payload with no compactable context surfaces a visible error.

    Runs the real JavaScript extension under Node with a resident context that
    exposes no ``compact`` function (e.g. an older Pi without the extension
    compaction API, or a model that cannot compact). The runner already returned
    200, so the server runs no AP-side fallback — if the extension stayed silent
    the user's /compact would vanish with no feedback. ``triggerCompaction`` must
    instead post a visible ``external_conversation_item`` error
    (``pi_compact_unavailable``) and raise NO spinner edge (zero
    ``external_compaction_status`` events, so the web spinner created only by the
    ``response.compaction.in_progress`` SSE never appears and cannot strand). The
    payload file is still consumed (unlinked) so the poller does not re-read it.
    """
    node = shutil.which("node")
    if node is None:
        pytest.skip("node is required for the pi-native extension e2e test")

    extension_path = (
        Path(__file__).resolve().parents[1]
        / "omnigent"
        / "resources"
        / "pi_native"
        / "omnigent_pi_native_extension.js"
    )

    script = r"""
const assert = require("assert").strict;
const fs = require("fs");
const path = require("path");

const extensionPath = process.argv[1];
const tmpDir = process.argv[2];
const inboxDir = path.join(tmpDir, "inbox");
const payloadPath = path.join(inboxDir, "000-compact.json");
const configPath = path.join(tmpDir, "config.json");

fs.mkdirSync(inboxDir, { recursive: true });
fs.writeFileSync(
  payloadPath,
  JSON.stringify({ id: "compact-1", type: "compact" }),
);
fs.writeFileSync(
  configPath,
  JSON.stringify({
    serverUrl: "http://omnigent.test",
    sessionId: "session-1",
    inboxDir,
  }),
);

process.env.OMNIGENT_PI_NATIVE_CONFIG = configPath;

const postedEvents = [];
global.fetch = async (_url, request) => {
  postedEvents.push(JSON.parse(request.body));
  return { ok: true };
};

let pollInbox = null;
global.setInterval = (fn, _ms) => {
  pollInbox = fn;
  return { fakeInterval: true };
};

const handlers = {};
const pi = {
  registerCommand() {},
  on(eventName, handler) {
    handlers[eventName] = handler;
  },
  sendUserMessage() {},
};

require(extensionPath)(pi);

// Resident context WITHOUT a compact() function: triggerCompaction must
// short-circuit and post nothing.
const ctx = {
  sessionManager: { getSessionId: () => "native-session-1" },
  ui: { setTitle() {}, setStatus() {}, notify() {} },
  abort() {},
  isIdle: () => false,
};

(async () => {
  await handlers.session_start({}, ctx);
  assert.equal(typeof pollInbox, "function");

  // Must not throw even though ctx.compact is absent.
  pollInbox();
  await new Promise((resolve) => setImmediate(resolve));

  const compactionStatuses = postedEvents.filter(
    (event) => event.type === "external_compaction_status",
  );
  // No spinner is ever raised: zero compaction-status edges (most importantly
  // no in_progress), so nothing can strand.
  assert.deepEqual(compactionStatuses, [], JSON.stringify(postedEvents));

  // The failure is surfaced as a visible conversation error item rather than
  // silently swallowed, so the user knows /compact did nothing.
  const unavailable = postedEvents.find(
    (event) =>
      event.type === "external_conversation_item" &&
      event.data &&
      event.data.item_type === "error" &&
      event.data.item_data &&
      event.data.item_data.code === "pi_compact_unavailable",
  );
  assert.ok(unavailable, JSON.stringify(postedEvents));
  assert.equal(unavailable.data.item_data.source, "execution");
  assert.match(unavailable.data.response_id, /^pi-compact-unavailable-/);

  // The payload file is still consumed so the poller does not re-read it.
  assert.equal(fs.existsSync(payloadPath), false);
})().catch((error) => {
  console.error(error && error.stack ? error.stack : error);
  process.exit(1);
});
"""

    result = subprocess.run(
        [node, "-e", script, str(extension_path), str(tmp_path)],
        capture_output=True,
        check=False,
        text=True,
        timeout=10,
    )

    assert result.returncode == 0, result.stdout + result.stderr


def test_compact_payload_synchronous_throw_dismisses_spinner(
    tmp_path: Path,
) -> None:
    """A synchronous throw from ``ctx.compact()`` still dismisses the spinner.

    ``triggerCompaction`` posts ``in_progress`` BEFORE calling the fire-and-forget
    ``ctx.compact()``. If that call throws synchronously (before any
    ``onComplete``/``onError`` can fire), the catch must post ``failed`` so the
    raised spinner is dismissed rather than stranded. Edges must be exactly
    ``[in_progress, failed]`` — distinct from the ``onError`` path, which reaches
    ``failed`` via the callback.
    """
    node = shutil.which("node")
    if node is None:
        pytest.skip("node is required for the pi-native extension e2e test")

    extension_path = (
        Path(__file__).resolve().parents[1]
        / "omnigent"
        / "resources"
        / "pi_native"
        / "omnigent_pi_native_extension.js"
    )

    script = r"""
const assert = require("assert").strict;
const fs = require("fs");
const path = require("path");

const extensionPath = process.argv[1];
const tmpDir = process.argv[2];
const inboxDir = path.join(tmpDir, "inbox");
const payloadPath = path.join(inboxDir, "000-compact.json");
const configPath = path.join(tmpDir, "config.json");

fs.mkdirSync(inboxDir, { recursive: true });
fs.writeFileSync(
  payloadPath,
  JSON.stringify({ id: "compact-1", type: "compact" }),
);
fs.writeFileSync(
  configPath,
  JSON.stringify({
    serverUrl: "http://omnigent.test",
    sessionId: "session-1",
    inboxDir,
  }),
);

process.env.OMNIGENT_PI_NATIVE_CONFIG = configPath;

const postedEvents = [];
global.fetch = async (_url, request) => {
  postedEvents.push(JSON.parse(request.body));
  return { ok: true };
};

let pollInbox = null;
global.setInterval = (fn, _ms) => {
  pollInbox = fn;
  return { fakeInterval: true };
};

const handlers = {};
const pi = {
  registerCommand() {},
  on(eventName, handler) {
    handlers[eventName] = handler;
  },
  sendUserMessage() {},
};

require(extensionPath)(pi);

const ctx = {
  sessionManager: { getSessionId: () => "native-session-1" },
  ui: { setTitle() {}, setStatus() {}, notify() {} },
  abort() {},
  isIdle: () => false,
  compact() {
    throw new Error("compact() exploded synchronously");
  },
};

(async () => {
  await handlers.session_start({}, ctx);
  // Must not throw out of the poller even though compact() throws.
  pollInbox();
  await new Promise((resolve) => setImmediate(resolve));

  const compactionStatuses = postedEvents
    .filter((event) => event.type === "external_compaction_status")
    .map((event) => event.data && event.data.status);
  // in_progress was posted before compact(); the synchronous throw is caught
  // and posts failed to dismiss the spinner.
  assert.deepEqual(
    compactionStatuses,
    ["in_progress", "failed"],
    JSON.stringify(postedEvents),
  );

  // The payload file is consumed so compaction is not re-triggered each tick.
  assert.equal(fs.existsSync(payloadPath), false);
})().catch((error) => {
  console.error(error && error.stack ? error.stack : error);
  process.exit(1);
});
"""

    result = subprocess.run(
        [node, "-e", script, str(extension_path), str(tmp_path)],
        capture_output=True,
        check=False,
        text=True,
        timeout=10,
    )

    assert result.returncode == 0, result.stdout + result.stderr


def test_compact_payload_failure_dismisses_spinner(tmp_path: Path) -> None:
    """When ``ctx.compact()`` reports an error, the extension posts ``failed``.

    Pi's ``compact()`` surfaces failures through the ``onError`` callback. The
    extension must publish ``external_compaction_status`` ``failed`` so the web
    UI's "Compacting…" spinner is dismissed rather than stranded.
    """
    node = shutil.which("node")
    if node is None:
        pytest.skip("node is required for the pi-native extension e2e test")

    extension_path = (
        Path(__file__).resolve().parents[1]
        / "omnigent"
        / "resources"
        / "pi_native"
        / "omnigent_pi_native_extension.js"
    )

    script = r"""
const assert = require("assert").strict;
const fs = require("fs");
const path = require("path");

const extensionPath = process.argv[1];
const tmpDir = process.argv[2];
const inboxDir = path.join(tmpDir, "inbox");
const payloadPath = path.join(inboxDir, "000-compact.json");
const configPath = path.join(tmpDir, "config.json");

fs.mkdirSync(inboxDir, { recursive: true });
fs.writeFileSync(
  payloadPath,
  JSON.stringify({ id: "compact-1", type: "compact" }),
);
fs.writeFileSync(
  configPath,
  JSON.stringify({
    serverUrl: "http://omnigent.test",
    sessionId: "session-1",
    inboxDir,
  }),
);

process.env.OMNIGENT_PI_NATIVE_CONFIG = configPath;

const postedEvents = [];
global.fetch = async (_url, request) => {
  postedEvents.push(JSON.parse(request.body));
  return { ok: true };
};

let pollInbox = null;
global.setInterval = (fn, _ms) => {
  pollInbox = fn;
  return { fakeInterval: true };
};

const handlers = {};
const pi = {
  registerCommand() {},
  on(eventName, handler) {
    handlers[eventName] = handler;
  },
  sendUserMessage() {},
};

require(extensionPath)(pi);

const ctx = {
  sessionManager: { getSessionId: () => "native-session-1" },
  ui: { setTitle() {}, setStatus() {}, notify() {} },
  abort() {},
  isIdle: () => false,
  compact(options) {
    if (options && typeof options.onError === "function") {
      options.onError(new Error("compaction blew up"));
    }
  },
};

(async () => {
  await handlers.session_start({}, ctx);
  pollInbox();
  await new Promise((resolve) => setImmediate(resolve));

  const compactionStatuses = postedEvents
    .filter((event) => event.type === "external_compaction_status")
    .map((event) => event.data && event.data.status);
  // in_progress raised the spinner; failed (from onError) dismisses it.
  // completed must never fire on an errored compaction.
  assert.deepEqual(
    compactionStatuses,
    ["in_progress", "failed"],
    JSON.stringify(postedEvents),
  );
})().catch((error) => {
  console.error(error && error.stack ? error.stack : error);
  process.exit(1);
});
"""

    result = subprocess.run(
        [node, "-e", script, str(extension_path), str(tmp_path)],
        capture_output=True,
        check=False,
        text=True,
        timeout=10,
    )

    assert result.returncode == 0, result.stdout + result.stderr


# ── TOOL_CALL policy evaluation (DENY / ALLOW / ASK elicitation) ──────────
#
# These exercise the real evalNativePolicyHttp park/resolve loop in the
# generated extension by driving its tool_call handler under Node with a
# scripted fetch and assert-rich JS body. Each case supplies a queue of
# responses; the JS harness drives one tool_call and reports the verdict.

# Shared JS preamble: loads the extension, wires a scripted fetch + fake
# timers (so the long park / backoff budgets collapse to instant in test),
# and exposes runToolCall() which fires the tool_call handler and returns the
# verdict the extension would hand back to Pi.
_POLICY_HARNESS_PREAMBLE = r"""
const assert = require("assert").strict;
const fs = require("fs");
const path = require("path");

const extensionPath = process.argv[1];
const tmpDir = process.argv[2];
const configPath = path.join(tmpDir, "config.json");

fs.writeFileSync(
  configPath,
  JSON.stringify({
    serverUrl: "http://omnigent.test",
    sessionId: "session-1",
    authHeaders: { authorization: "Bearer test" },
  }),
);
process.env.OMNIGENT_PI_NATIVE_CONFIG = configPath;

// Captured evaluate-request bodies (parsed) in call order.
const evalBodies = [];
// Queue of responders. Each entry is a function (parsedBody) -> response-like
// object, OR the string "THROW" to simulate a transport error, OR
// "THROW_ABORT" to simulate our own AbortController firing (DOMException-ish).
let responders = [];
let evalCallCount = 0;

function makeJsonResponse(obj, status) {
  return {
    ok: status === undefined || (status >= 200 && status < 300),
    status: status === undefined ? 200 : status,
    json: async () => obj,
  };
}

global.fetch = async (url, request) => {
  const body = JSON.parse(request.body);
  // Only the evaluate endpoint is scripted; postEvent calls (events endpoint)
  // just succeed silently so they never interfere with the verdict assertions.
  if (typeof url === "string" && url.includes("/policies/evaluate")) {
    evalBodies.push(body);
    const idx = evalCallCount;
    evalCallCount += 1;
    const responder = responders[idx];
    if (responder === "THROW") {
      throw new Error("ECONNREFUSED simulated transport error");
    }
    if (responder === "THROW_ABORT") {
      // Simulate our own AbortController firing after holding the connection
      // for the FULL per-attempt park timeout — i.e. a legitimate long-poll
      // re-attach. The extension disambiguates a real park from a genuine
      // transport error by the attempt's elapsed wall-time (a real connect
      // error throws fast; a real park only aborts at the per-attempt
      // timeout), so advance the fake clock by the per-attempt budget before
      // flipping the signal and throwing the AbortError fetch would raise.
      fakeNow += PARK_ATTEMPT_TIMEOUT_MS;
      if (request && request.signal && typeof request.signal._abort === "function") {
        request.signal._abort();
      }
      const err = new Error("The operation was aborted");
      err.name = "AbortError";
      throw err;
    }
    if (typeof responder === "function") return responder(body);
    // Default: allow.
    return makeJsonResponse({ result: "POLICY_ACTION_ALLOW" });
  }
  return { ok: true, status: 200, json: async () => ({}) };
};

// Fake clock + timers. A short delay (sleep/backoff) advances a virtual clock
// by its duration and fires on the next microtask, so the extension's
// wall-clock budgets (transient retry, park ceiling) elapse deterministically
// and instantly — no real 30s wait. The long park abort timer (>= 100s) is
// never fired so a scripted fetch always resolves first; but scheduling it
// still advances nothing (it is cleared in the finally).
let fakeNow = 1_000_000;
// Mirrors _PARK_ATTEMPT_TIMEOUT_MS in the extension. A real long-poll abort
// only fires after the connection is held this long; the THROW_ABORT responder
// advances the fake clock by this amount so the extension classifies it as a
// legitimate re-attach (vs. a fast-failing genuine transport error).
const PARK_ATTEMPT_TIMEOUT_MS = 240000;
const realDateNow = Date.now.bind(Date);
Date.now = () => fakeNow;
global.setTimeout = (fn, ms) => {
  if (typeof ms === "number" && ms >= 100000) {
    return { fakeBig: true };
  }
  if (typeof ms === "number" && ms > 0) fakeNow += ms;
  Promise.resolve().then(fn);
  return { fakeSmall: true };
};
global.clearTimeout = () => {};
// Keep the inbox poller dormant.
global.setInterval = () => ({ fakeInterval: true });

const handlers = {};
const pi = {
  registerCommand() {},
  on(eventName, handler) {
    handlers[eventName] = handler;
  },
  sendUserMessage() {},
};

require(extensionPath)(pi);

const ctx = {
  sessionManager: { getSessionId: () => "native-session-1" },
  ui: { setTitle() {}, setStatus() {}, notify() {} },
  abort() {},
  isIdle() { return false; },
};

async function runToolCall() {
  assert.equal(typeof handlers.tool_call, "function");
  return handlers.tool_call(
    { toolCallId: "call-1", toolName: "Bash", input: { command: "rm -rf /tmp/x" } },
    ctx,
  );
}
"""


def _run_policy_node_script(extension_path: Path, tmp_path: Path, body: str) -> None:
    """Run the extension's tool_call policy path under Node with a scripted fetch.

    :param extension_path: Path to the generated extension JS.
    :param tmp_path: Per-test scratch dir (config is written here).
    :param body: JS test body appended after the shared harness preamble; it
        sets ``responders`` and runs assertions inside an async IIFE.
    """
    node = shutil.which("node")
    if node is None:
        pytest.skip("node is required for the pi-native extension policy test")

    script = _POLICY_HARNESS_PREAMBLE + "\n" + body
    result = subprocess.run(
        [node, "-e", script, str(extension_path), str(tmp_path)],
        capture_output=True,
        check=False,
        text=True,
        timeout=15,
    )
    assert result.returncode == 0, result.stdout + result.stderr


def test_policy_allow_proceeds(tmp_path: Path) -> None:
    """An ALLOW verdict lets the Pi tool call proceed (no block returned)."""
    body = r"""
(async () => {
  responders = [(_b) => makeJsonResponse({ result: "POLICY_ACTION_ALLOW" })];
  const verdict = await runToolCall();
  // tool_call returns undefined (or a non-blocking value) on ALLOW.
  assert.ok(!verdict || verdict.block !== true, JSON.stringify(verdict));
  assert.equal(evalBodies.length, 1);
  // The evaluate body must carry a valid re-attach id and a PHASE_TOOL_CALL.
  assert.match(evalBodies[0]._omnigent_elicitation_id, /^elicit_evaluate_[0-9a-f]{32}$/);
  assert.equal(evalBodies[0].event.type, "PHASE_TOOL_CALL");
  assert.equal(evalBodies[0].event.data.name, "Bash");
})().catch((error) => {
  console.error(error && error.stack ? error.stack : error);
  process.exit(1);
});
"""
    _run_policy_node_script(_extension_path(), tmp_path, body)


def test_policy_deny_blocks(tmp_path: Path) -> None:
    """A DENY verdict blocks the Pi tool call and surfaces the policy reason."""
    body = r"""
(async () => {
  responders = [
    (_b) => makeJsonResponse({ result: "POLICY_ACTION_DENY", reason: "no rm -rf" }),
  ];
  const verdict = await runToolCall();
  assert.deepEqual(verdict, { block: true, reason: "no rm -rf" });
  assert.equal(evalBodies.length, 1);
})().catch((error) => {
  console.error(error && error.stack ? error.stack : error);
  process.exit(1);
});
"""
    _run_policy_node_script(_extension_path(), tmp_path, body)


def test_policy_ask_parks_then_resolves_allow(tmp_path: Path) -> None:
    """A raw ASK re-evaluates (re-attaching) until it resolves to ALLOW.

    The first evaluate returns ASK (gate did not park server-side); the loop
    re-POSTs the SAME elicitation id and the second returns ALLOW, so the tool
    call proceeds. Mirrors the server collapsing ASK to a hard verdict.
    """
    body = r"""
(async () => {
  responders = [
    (_b) => makeJsonResponse({ result: "POLICY_ACTION_ASK", reason: "approve?" }),
    (_b) => makeJsonResponse({ result: "POLICY_ACTION_ALLOW" }),
  ];
  const verdict = await runToolCall();
  assert.ok(!verdict || verdict.block !== true, JSON.stringify(verdict));
  assert.equal(evalBodies.length, 2, "expected park-then-resolve = 2 evaluates");
  // Both POSTs must reuse the SAME elicitation id so the server re-attaches
  // rather than opening a second approval card.
  assert.equal(
    evalBodies[0]._omnigent_elicitation_id,
    evalBodies[1]._omnigent_elicitation_id,
  );
})().catch((error) => {
  console.error(error && error.stack ? error.stack : error);
  process.exit(1);
});
"""
    _run_policy_node_script(_extension_path(), tmp_path, body)


def test_policy_ask_parks_then_resolves_deny(tmp_path: Path) -> None:
    """A raw ASK that resolves to DENY blocks the Pi tool call with the reason."""
    body = r"""
(async () => {
  responders = [
    (_b) => makeJsonResponse({ result: "POLICY_ACTION_ASK", reason: "approve?" }),
    (_b) => makeJsonResponse({ result: "POLICY_ACTION_DENY", reason: "declined" }),
  ];
  const verdict = await runToolCall();
  assert.deepEqual(verdict, { block: true, reason: "declined" });
  assert.equal(evalBodies.length, 2);
  assert.equal(
    evalBodies[0]._omnigent_elicitation_id,
    evalBodies[1]._omnigent_elicitation_id,
  );
})().catch((error) => {
  console.error(error && error.stack ? error.stack : error);
  process.exit(1);
});
"""
    _run_policy_node_script(_extension_path(), tmp_path, body)


def test_policy_park_abort_reattaches_same_id(tmp_path: Path) -> None:
    """An aborted park (our own headers-timeout guard) re-attaches, not re-mints.

    Simulates undici severing a long park: the first attempt throws an
    AbortError (signal.aborted), the loop must re-POST the SAME elicitation id
    immediately (no backoff), and the resolved ALLOW lets the tool proceed.
    """
    body = r"""
(async () => {
  // First call: pretend our AbortController fired (the fake controller below
  // reports aborted=true). Second call: resolved ALLOW.
  responders = ["THROW_ABORT", (_b) => makeJsonResponse({ result: "POLICY_ACTION_ALLOW" })];
  // Fake AbortController whose signal exposes a private _abort() so the fetch
  // mock can flip aborted=true (mimicking our 240s headers-timeout guard
  // firing). The extension's catch reads controller.signal.aborted to choose
  // the re-attach (no-backoff) branch over the transient-error branch.
  global.AbortController = class {
    constructor() {
      let aborted = false;
      const signal = {};
      Object.defineProperty(signal, "aborted", { get() { return aborted; } });
      signal._abort = () => { aborted = true; };
      this.signal = signal;
    }
    abort() { this.signal._abort(); }
  };
  const verdict = await runToolCall();
  assert.ok(!verdict || verdict.block !== true, JSON.stringify(verdict));
  assert.equal(evalBodies.length, 2, "expected re-attach after abort = 2 evaluates");
  assert.equal(
    evalBodies[0]._omnigent_elicitation_id,
    evalBodies[1]._omnigent_elicitation_id,
  );
})().catch((error) => {
  console.error(error && error.stack ? error.stack : error);
  process.exit(1);
});
"""
    _run_policy_node_script(_extension_path(), tmp_path, body)


def test_policy_transient_budget_refreshes_after_park(tmp_path: Path) -> None:
    """A transport blip AFTER a long-poll re-attach is retried, not failed closed.

    Regression guard for the stale transient-budget bug: the entry budget is set
    once at function start. A real human-approval park re-attaches every
    per-attempt timeout, advancing the clock well past that entry budget. Unless
    the re-attach branch refreshes the budget (as the ASK branch does), a genuine
    transport blip during the human's deliberation would compare against an
    already-expired deadline and fail CLOSED with zero retries. Here: park
    (THROW_ABORT) then a transport error (THROW) then ALLOW; with the refresh the
    blip is retried and the tool proceeds (3 evaluates), without it the blip would
    fail closed immediately (2 evaluates, blocked).
    """
    body = r"""
(async () => {
  global.AbortController = class {
    constructor() {
      let aborted = false;
      const signal = {};
      Object.defineProperty(signal, "aborted", { get() { return aborted; } });
      signal._abort = () => { aborted = true; };
      this.signal = signal;
    }
    abort() { this.signal._abort(); }
  };
  responders = [
    "THROW_ABORT",
    "THROW",
    (_b) => makeJsonResponse({ result: "POLICY_ACTION_ALLOW" }),
  ];
  const verdict = await runToolCall();
  assert.ok(
    !verdict || verdict.block !== true,
    "a blip after a park must be retried, got " + JSON.stringify(verdict),
  );
  assert.equal(
    evalBodies.length,
    3,
    "expected park, retried blip, allow = 3 evaluates, got " + evalBodies.length,
  );
})().catch((error) => {
  console.error(error && error.stack ? error.stack : error);
  process.exit(1);
});
"""
    _run_policy_node_script(_extension_path(), tmp_path, body)


def test_policy_transport_error_fails_closed(tmp_path: Path) -> None:
    """A persistent transport error fails CLOSED (finding #1).

    PHASE_TOOL_CALL is the sole enforcement point for a native connector tool,
    so an unevaluable policy must BLOCK, not proceed. Every evaluate POST
    throws a non-abort transport error; after the transient retry budget
    elapses the extension returns a deny verdict (fail closed) rather than
    null, matching omnigent.policies.types.FAIL_CLOSED_PHASES and the Python
    native hook's fail_closed_hook_output(PreToolUse) → deny.
    """
    body = r"""
(async () => {
  // Always throw a transport error; with fake timers collapsing the backoff
  // the transient budget elapses quickly and the loop must fail CLOSED.
  responders = new Array(64).fill("THROW");
  const verdict = await runToolCall();
  // Fail closed → a block verdict with the unreachable-server reason.
  assert.equal(verdict && verdict.block, true, JSON.stringify(verdict));
  assert.match(verdict.reason, /unreachable/);
  assert.match(verdict.reason, /failing closed/);
  // It must have actually retried a few times before giving up (not one-shot).
  assert.ok(evalBodies.length >= 2, "expected retries, got " + evalBodies.length);
})().catch((error) => {
  console.error(error && error.stack ? error.stack : error);
  process.exit(1);
});
"""
    _run_policy_node_script(_extension_path(), tmp_path, body)


def test_policy_server_5xx_fails_closed(tmp_path: Path) -> None:
    """A persistent 5xx response fails CLOSED once the transient budget is spent.

    A 5xx is transient (retried, re-attaching), but if it never clears the
    gate cannot produce a verdict, so PHASE_TOOL_CALL fails closed rather than
    proceeding.
    """
    body = r"""
(async () => {
  responders = new Array(64).fill((_b) => makeJsonResponse({}, 503));
  const verdict = await runToolCall();
  assert.equal(verdict && verdict.block, true, JSON.stringify(verdict));
  assert.match(verdict.reason, /failing closed/);
  assert.ok(evalBodies.length >= 2, "expected retries, got " + evalBodies.length);
})().catch((error) => {
  console.error(error && error.stack ? error.stack : error);
  process.exit(1);
});
"""
    _run_policy_node_script(_extension_path(), tmp_path, body)


def test_policy_4xx_fails_closed(tmp_path: Path) -> None:
    """A 4xx response fails CLOSED immediately (no retry).

    A 4xx is a final error (a bad request won't succeed on retry), so the gate
    has no usable verdict and PHASE_TOOL_CALL must block rather than proceed.
    Unlike a 5xx it is not charged against the transient budget, so it fails
    closed on the first response.
    """
    body = r"""
(async () => {
  responders = [(_b) => makeJsonResponse({}, 400)];
  const verdict = await runToolCall();
  assert.equal(verdict && verdict.block, true, JSON.stringify(verdict));
  assert.match(verdict.reason, /failing closed/);
  // A 4xx is final, not transient: one POST, no retries.
  assert.equal(evalBodies.length, 1, "4xx must not retry, got " + evalBodies.length);
})().catch((error) => {
  console.error(error && error.stack ? error.stack : error);
  process.exit(1);
});
"""
    _run_policy_node_script(_extension_path(), tmp_path, body)


def test_policy_malformed_body_fails_closed(tmp_path: Path) -> None:
    """A 200 response with an unparseable body fails CLOSED.

    A malformed JSON body yields no verdict and is not retryable, so an
    unevaluable PHASE_TOOL_CALL blocks rather than proceeds.
    """
    body = r"""
(async () => {
  responders = [
    (_b) => ({ ok: true, status: 200, json: async () => { throw new Error("bad json"); } }),
  ];
  const verdict = await runToolCall();
  assert.equal(verdict && verdict.block, true, JSON.stringify(verdict));
  assert.match(verdict.reason, /failing closed/);
  assert.equal(evalBodies.length, 1, "malformed body must not retry, got " + evalBodies.length);
})().catch((error) => {
  console.error(error && error.stack ? error.stack : error);
  process.exit(1);
});
"""
    _run_policy_node_script(_extension_path(), tmp_path, body)


def test_policy_raw_ask_never_collapses_fails_closed(tmp_path: Path) -> None:
    """A raw ASK that never collapses fails CLOSED after the round cap (finding #2).

    A read-only caller's gate can return ASK without parking server-side. The
    extension re-evaluates (re-attaching the same id), but a raw ASK that NEVER
    collapses to a hard verdict must not ride the 24h park ceiling and then
    proceed — after _MAX_RAW_ASK_ROUNDS it denies, mirroring the Python native
    hook's stray-ASK-closed behavior.
    """
    body = r"""
(async () => {
  // Always ASK — it never collapses to ALLOW/DENY.
  const askForever = (_b) =>
    makeJsonResponse({ result: "POLICY_ACTION_ASK", reason: "approve?" });
  responders = new Array(256).fill(askForever);
  const verdict = await runToolCall();
  // Fail closed → a block verdict, NOT a proceed (the old 24h-hang-then-allow).
  assert.equal(verdict && verdict.block, true, JSON.stringify(verdict));
  assert.match(verdict.reason, /failing closed/);
  // It re-attached several rounds before giving up, but is bounded (not 24h /
  // unbounded). All POSTs reuse the SAME elicitation id.
  assert.ok(evalBodies.length >= 2, "expected re-evaluation, got " + evalBodies.length);
  assert.ok(evalBodies.length <= 50, "raw ASK rounds must be capped, got " + evalBodies.length);
  const ids = new Set(evalBodies.map((b) => b._omnigent_elicitation_id));
  assert.equal(ids.size, 1, "all raw-ASK rounds must reuse one elicitation id");
})().catch((error) => {
  console.error(error && error.stack ? error.stack : error);
  process.exit(1);
});
"""
    _run_policy_node_script(_extension_path(), tmp_path, body)


def test_policy_fast_error_racing_abort_is_bounded_fails_closed(
    tmp_path: Path,
) -> None:
    """A genuine error racing the abort timer is bounded, not infinite (finding #3).

    Once the per-attempt abort timer has fired, controller.signal.aborted reads
    true even for a real connect reset that raced it. The old code re-attached
    on ``aborted`` alone, so such an error retried unboundedly toward the 24h
    ceiling. The fix disambiguates by elapsed wall-time: a fast-failing error
    (well under the per-attempt timeout) is a genuine transport error charged
    against the short transient budget — so a persistent one fails CLOSED
    instead of looping forever.

    THROW_FAST_ABORT flips ``aborted`` true (as if the timer had fired) but
    throws IMMEDIATELY (no clock advance), modelling the race.
    """
    body = r"""
(async () => {
  // Fake AbortController whose signal can be flipped aborted=true.
  global.AbortController = class {
    constructor() {
      let aborted = false;
      const signal = {};
      Object.defineProperty(signal, "aborted", { get() { return aborted; } });
      signal._abort = () => { aborted = true; };
      this.signal = signal;
    }
    abort() { this.signal._abort(); }
  };
  // Every attempt: flip aborted=true (timer "fired") but throw immediately with
  // NO clock advance — a genuine reset racing the abort timer. The extension
  // must charge these against the transient budget (elapsed ~= 0 << per-attempt
  // timeout) and ultimately fail CLOSED, not re-attach forever.
  let calls = 0;
  responders = new Array(512).fill(null).map(() => (_b) => {
    throw new Error("__USE_FAST_ABORT__");
  });
  // Override fetch's THROW path is awkward; instead simulate via a responder
  // that flips the (current) controller's signal then throws. The controller is
  // recreated each attempt, so reach it through the request signal.
  global.fetch = async (url, request) => {
    const parsed = JSON.parse(request.body);
    if (typeof url === "string" && url.includes("/policies/evaluate")) {
      evalBodies.push(parsed);
      calls += 1;
      if (request && request.signal && typeof request.signal._abort === "function") {
        request.signal._abort(); // timer "fired" — aborted=true
      }
      const err = new Error("ECONNRESET racing abort");
      err.name = "AbortError";
      throw err; // immediate: elapsed ~= 0, NOT a real long-poll
    }
    return { ok: true, status: 200, json: async () => ({}) };
  };
  const verdict = await runToolCall();
  // Bounded by the transient budget → fail CLOSED, not an infinite re-attach.
  assert.equal(verdict && verdict.block, true, JSON.stringify(verdict));
  assert.match(verdict.reason, /failing closed/);
  // Retried a few times (transient budget) but nowhere near unbounded.
  assert.ok(evalBodies.length >= 2, "expected transient retries, got " + evalBodies.length);
  assert.ok(evalBodies.length < 200, "must be bounded, got " + evalBodies.length);
})().catch((error) => {
  console.error(error && error.stack ? error.stack : error);
  process.exit(1);
});
"""
    _run_policy_node_script(_extension_path(), tmp_path, body)


def test_outbound_requests_reread_refreshed_bearer(tmp_path: Path) -> None:
    """Each outbound POST re-reads ``authHeaders``, picking up the runner's re-mint.

    The bearer baked into ``config.json`` at launch dies with the ~1h
    Databricks OAuth lifetime. The runner re-mints it into the config each
    turn; the extension must read the bearer fresh per request (not the
    launch snapshot it closed over at startup) so its POSTs stay
    authenticated mid-session instead of failing closed. Drives two
    ``message_end`` flushes with a config rewrite between them and asserts the
    second POST carries the new bearer.
    """
    node = shutil.which("node")
    if node is None:
        pytest.skip("node is required for the pi-native extension e2e test")

    script = r"""
const assert = require("assert").strict;
const fs = require("fs");
const path = require("path");

const extensionPath = process.argv[1];
const configPath = path.join(require("os").tmpdir(), `pi-reauth-${process.pid}.json`);
function writeConfig(bearer) {
  fs.writeFileSync(
    configPath,
    JSON.stringify({
      serverUrl: "http://omnigent.test",
      sessionId: "session-1",
      authHeaders: { authorization: bearer },
    }),
  );
}
writeConfig("Bearer stale");
process.env.OMNIGENT_PI_NATIVE_CONFIG = configPath;

const sentAuth = [];
global.fetch = async (_url, request) => {
  sentAuth.push(request.headers.authorization);
  return { ok: true };
};
global.setInterval = () => ({ fakeInterval: true });

const handlers = {};
const pi = { registerCommand() {}, on(name, handler) { handlers[name] = handler; } };
require(extensionPath)(pi);
const ctx = { ui: { setTitle() {}, setStatus() {}, notify() {} } };

function usage(id) {
  return {
    message: {
      id,
      role: "assistant",
      model: "databricks-claude-sonnet-4-6",
      usage: { input: 10, output: 1, cacheRead: 0, cacheWrite: 0, totalTokens: 11 },
    },
  };
}

(async () => {
  await handlers.message_end(usage("m1"), ctx);
  // The runner re-mints the bearer into config.json mid-session.
  writeConfig("Bearer fresh");
  await handlers.message_end(usage("m2"), ctx);

  // First POST used the launch token; the second re-read the rewritten config
  // and carried the fresh bearer — not the stale one closed over at startup.
  assert.deepEqual(sentAuth, ["Bearer stale", "Bearer fresh"], JSON.stringify(sentAuth));
})().catch((error) => {
  console.error(error && error.stack ? error.stack : error);
  process.exit(1);
});
"""
    _run_extension_script(node, _extension_path(), script)