Merge origin/main into pi-native/cost-tracking

Resolve conflicts in the pi-native extension + its tests as a UNION of
#1277 (token-usage tracking: extractPiUsage / external_session_usage /
dedup by message identity from fix 395306a3) and #1239 (streaming text
deltas: external_output_text_delta / finalizeStreamingMessage). Both
feature sets and both test suites coexist.

Co-authored-by: Isaac

omnigent/resources/pi_native/omnigent_pi_native_extension.js

@@ -125,6 +125,51 @@ function messageRole(message) {
   return String(message.role || message.type || "");
 }
 
+function toInt(value) {
+  const n = Number(value);
+  return Number.isFinite(n) && n > 0 ? Math.trunc(n) : 0;
+}
+
+/**
+ * Lift token usage out of one Pi assistant message.
+ *
+ * Mirrors the non-native executor's ``_extract_pi_turn_usage``
+ * (omnigent/inner/pi_executor.py): Pi (``@earendil-works/pi-coding-agent``)
+ * carries a per-message ``usage`` object with ``input`` / ``output`` /
+ * ``cacheRead`` / ``cacheWrite`` / ``totalTokens`` counts, and the message
+ * carries the resolved ``model``. Pi's ``input`` is the NON-cached input
+ * (Anthropic semantics) — ``cacheRead`` / ``cacheWrite`` are separate, so the
+ * full input a turn sent is ``input + cacheRead + cacheWrite``.
+ *
+ * @returns {{input:number,output:number,cacheRead:number,cacheWrite:number,
+ *   total:number,model:(string|null)}|null} the per-message counts, or
+ *   ``null`` when ``message`` is not an assistant message carrying usage.
+ */
+function extractPiUsage(message) {
+  if (!message || typeof message !== "object") return null;
+  if (message.role !== "assistant") return null;
+  const usage = message.usage;
+  if (!usage || typeof usage !== "object") return null;
+  const input = toInt(usage.input);
+  const output = toInt(usage.output);
+  const cacheRead = toInt(usage.cacheRead);
+  const cacheWrite = toInt(usage.cacheWrite);
+  // No countable tokens means Pi emitted an empty usage object — treat as "no
+  // usage" so the server leaves the session unpriced rather than recording a
+  // $0.00 turn (matches _aggregate_pi_turn_usage's empty-usage guard).
+  if (!(input || output || cacheRead || cacheWrite)) return null;
+  const rawModel = message.model;
+  const model = typeof rawModel === "string" && rawModel ? rawModel : null;
+  return {
+    input,
+    output,
+    cacheRead,
+    cacheWrite,
+    total: toInt(usage.totalTokens),
+    model,
+  };
+}
+
 function headers(config) {
   return {
     "content-type": "application/json",
@@ -356,6 +401,88 @@ module.exports = function (pi) {
   // or double-finalize the preview.
   const finalizedTextBlocks = new Set();
 
+  // Cumulative session token usage. Pi reports PER-MESSAGE counts (one
+  // assistant message per LLM call); session billing is their SUM — each call
+  // is billed for the full context it re-sent, so summing per-message inputs is
+  // the correct cumulative input. The server applies vendor pricing to these
+  // cumulative totals and republishes a ``session.usage`` event (the SAME
+  // contract claude-native / codex-native / cursor-native use), so the web
+  // Session-cost badge + per-model token breakdown light up with no
+  // server/frontend changes. Dedup by message fingerprint so a re-emitted
+  // ``message_end`` / ``turn_end`` / ``agent_end`` carrying the same assistant
+  // message never double-counts. ``usageModel`` tracks the latest message's
+  // model (mirrors a mid-session model switch). ``lastPostedUsageKey`` dedups
+  // the POST itself so a flush with no new tokens is skipped.
+  const countedUsageMessages = new Set();
+  let cumulativeInputTokens = 0;
+  let cumulativeOutputTokens = 0;
+  let cumulativeCacheReadTokens = 0;
+  let usageModel = null;
+  let lastPostedUsageKey = "";
+
+  // Build a stable fingerprint for one assistant message so the same message
+  // arriving on multiple lifecycle events is only counted once. Pi's
+  // ``AssistantMessage`` (``@earendil-works/pi-ai``) carries NO ``id`` field
+  // but DOES carry an optional provider ``responseId`` and a required numeric
+  // ``timestamp`` — both stable across the same message's re-emission on
+  // ``message_end`` / ``turn_end`` / ``agent_end``. Prefer those identity
+  // fields (plus a forward-compat ``id``) over the usage-count fingerprint:
+  // hashing counts alone collides two DISTINCT LLM calls that happen to report
+  // identical token counts (e.g. two identical short acks under prompt
+  // caching), which would silently drop the second call's tokens (undercount).
+  // The usage-count fingerprint stays only as a last resort for a message that
+  // carries no identity field at all.
+  function usageMessageKey(message, usage) {
+    if (message && typeof message === "object") {
+      if (typeof message.id === "string" && message.id) return `id:${message.id}`;
+      if (typeof message.responseId === "string" && message.responseId)
+        return `rid:${message.responseId}`;
+      if (typeof message.timestamp === "number")
+        return `ts:${message.timestamp}`;
+    }
+    return `u:${usage.input}-${usage.output}-${usage.cacheRead}-${usage.cacheWrite}-${usage.total}-${usage.model || ""}`;
+  }
+
+  // Fold one assistant message's usage into the cumulative session totals,
+  // deduped by fingerprint. Returns true when it counted (totals advanced).
+  function accumulateUsage(message) {
+    const usage = extractPiUsage(message);
+    if (!usage) return false;
+    const key = usageMessageKey(message, usage);
+    if (countedUsageMessages.has(key)) return false;
+    countedUsageMessages.add(key);
+    // The server's ``cumulative_input_tokens`` is INCLUSIVE of cache reads (it
+    // splits the cache portion back out and prices it at the cache-read rate),
+    // so add cacheRead into the input total. ``cacheWrite`` (cache creation)
+    // has no dedicated cumulative field on the server, so fold it into the
+    // input total too — it is then priced at the input rate rather than the
+    // ~1.25x cache-write rate, a small, documented approximation that never
+    // drops the tokens.
+    cumulativeInputTokens += usage.input + usage.cacheRead + usage.cacheWrite;
+    cumulativeOutputTokens += usage.output;
+    cumulativeCacheReadTokens += usage.cacheRead;
+    if (usage.model) usageModel = usage.model;
+    return true;
+  }
+
+  // POST the cumulative session usage so the server prices it and publishes a
+  // ``session.usage`` event. Cumulative (SET) semantics — the server overwrites
+  // its stored totals each flush. Deduped so a flush with no advance is a
+  // no-op. Fail-open via ``postEvent`` so a failed POST never wedges Pi.
+  async function postSessionUsage() {
+    if (!(cumulativeInputTokens || cumulativeOutputTokens)) return;
+    const postKey = `${cumulativeInputTokens}-${cumulativeOutputTokens}-${cumulativeCacheReadTokens}-${usageModel || ""}`;
+    if (postKey === lastPostedUsageKey) return;
+    lastPostedUsageKey = postKey;
+    const data = {
+      cumulative_input_tokens: cumulativeInputTokens,
+      cumulative_output_tokens: cumulativeOutputTokens,
+      cumulative_cache_read_input_tokens: cumulativeCacheReadTokens,
+    };
+    if (usageModel) data.model = usageModel;
+    await postEvent(config, { type: "external_session_usage", data });
+  }
+
   function rememberContext(ctx) {
     if (ctx) latestContext = ctx;
   }
@@ -609,12 +736,24 @@ module.exports = function (pi) {
     });
   });
 
-  pi.on("agent_end", async (_event, ctx) => {
+  pi.on("agent_end", async (event, ctx) => {
     rememberContext(ctx);
     clearPendingInterrupt();
     agentRunning = false;
     setOmnigentStatus(config, ctx, "idle");
     activeResponseId = null;
+    // Last-chance usage capture from the agent loop's final message set, in
+    // case neither ``message_end`` nor ``turn_end`` carried usage for some
+    // call. ``event.messages`` may hold the whole conversation; the
+    // fingerprint dedup means re-scanning already-counted messages is a no-op,
+    // so a plain forward-scan is safe (no overcount).
+    const messages =
+      event && Array.isArray(event.messages) ? event.messages : [];
+    let changed = false;
+    for (const message of messages) {
+      if (accumulateUsage(message)) changed = true;
+    }
+    if (changed) await postSessionUsage();
     await postEvent(config, {
       type: "external_session_status",
       data: { status: "idle", response_id: `pi-${Date.now()}-${++sequence}` },
@@ -734,6 +873,10 @@ module.exports = function (pi) {
     await finalizeStreamingMessage(responseId);
     streamingMessageOrdinal += 1;
     await mirrorAssistantMessage(message, responseId);
+    // ``message_end`` is the primary usage-capture site (one completed
+    // assistant message per LLM call); fold its token counts into the
+    // cumulative session totals and flush to the server for pricing.
+    if (accumulateUsage(message)) await postSessionUsage();
     const text = textFromMessage(message);
     if (!text) return;
     // The authoritative assistant item. The web UI retires + replaces the
@@ -758,6 +901,11 @@ module.exports = function (pi) {
     replayPendingInterrupt(ctx);
     const responseId = currentResponseId();
     await mirrorAssistantMessage(event && event.message, responseId);
+    // Fallback usage capture: if Pi attached usage to the turn's final
+    // assistant message but no ``message_end`` carried it, fold it in here.
+    // Deduped by fingerprint, so a message already counted on ``message_end``
+    // is a no-op.
+    if (accumulateUsage(event && event.message)) await postSessionUsage();
     const results =
       event && Array.isArray(event.toolResults) ? event.toolResults : [];
     for (const result of results) {

tests/test_pi_native_extension.py

@@ -150,6 +150,343 @@ require(extensionPath)(pi);
     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]