@blackbelt-technology/pi-agent-dashboard 0.4.0 → 0.4.1

package/packages/extension/src/__tests__/fork-entryid-timing.test.ts

@@ -1,100 +1,88 @@
 /**
- * Tests for fork entryId timing fix.
+ * Tests for bridge entryId stamping on message_end events.
  *
- * Verifies that the bridge's message_end entryId enrichment captures the
- * correct leaf ID (after pi core persists the entry), not the stale one
- * (before appendMessage runs).
+ * HISTORY: Originally this file modelled pi <0.69's synchronous emit pattern,
+ * where the bridge's `queueMicrotask` deferral ran BEFORE
+ * sessionManager.appendMessage. That design no longer reflects pi 0.70.x:
+ * pi awaits extension handlers inside _emitExtensionEvent, so the microtask
+ * resolves *inside* the awaited dispatcher, before persistence. The fix
+ * (see change: fix-per-message-fork) is `setTimeout(0)` (a macrotask)
+ * combined with reading `event.message.id` (which pi mutates in place
+ * during appendMessage) or a WeakMap populated by the wrapped appendMessage.
  *
- * The bug: pi core emits message_end via _emit() BEFORE calling
- * sessionManager.appendMessage(), so getLeafId() returns the previous leaf.
- * The fix: the bridge defers getLeafId() for message_end using queueMicrotask,
- * allowing appendMessage to run first.
+ * The previous test "message_start should still capture entryId immediately
+ * (no deferral)" codified the off-by-one bug as expected behavior — it has
+ * been REMOVED. The current test suite below models pi 0.70.x semantics
+ * directly. Detailed pi-0.70-specific scenarios live in
+ * `bridge-entry-id-pi-070.test.ts`.
  */
-import { describe, it, expect, vi } from "vitest";
+import { describe, it, expect } from "vitest";
 
-/**
- * Simulates the pi core + bridge interaction for entryId enrichment.
- *
- * Pi core's _processAgentEvent does:
- *   1. _emit(event)          — bridge handler called (async, not awaited)
- *   2. appendMessage(msg)    — updates leafId synchronously
- *
- * The bridge handler (async) should yield via queueMicrotask before reading
- * getLeafId(), so that appendMessage has already run.
- */
-describe("message_end entryId timing", () => {
-  it("deferred getLeafId() captures the post-persist entry ID", async () => {
-    // Simulate sessionManager with mutable leafId
-    let leafId = "user-entry-100"; // stale leaf before appendMessage
+describe("message_end entryId timing on pi 0.70.x", () => {
+  it("setTimeout(0) deferral captures the post-persist entry ID", async () => {
+    // Simulate pi 0.70.x: bridge handler runs awaited, THEN appendMessage runs.
+    let leafId = "user-entry-100";
     const sessionManager = {
       getLeafId: () => leafId,
+      appendMessage: (msg: any) => {
+        msg.id = "assistant-entry-101";
+        leafId = msg.id;
+        return msg.id;
+      },
     };
 
-    let capturedEntryId: string | undefined;
-
-    // Simulate the bridge handler (with the fix: defers via queueMicrotask)
-    const bridgeHandler = async () => {
-      // This is what the fixed bridge does for message_end:
-      await new Promise<void>(resolve => queueMicrotask(resolve));
-      capturedEntryId = sessionManager.getLeafId();
+    const event = { type: "message_end", message: { role: "assistant" } as any };
+    let captured: string | undefined;
+    let sendDone!: () => void;
+    const sentP = new Promise<void>((r) => { sendDone = r; });
+
+    // Bridge handler: schedules a setTimeout(0) and returns synchronously.
+    // The awaited dispatcher then unwinds, appendMessage runs, AND finally
+    // the timeout fires.
+    const bridgeHandler = (ev: any) => {
+      setTimeout(() => {
+        captured = ev.message.id ?? sessionManager.getLeafId();
+        sendDone();
+      }, 0);
     };
 
-    // Simulate pi core's _processAgentEvent:
-    // 1. _emit calls handler (async, NOT awaited)
-    const handlerPromise = bridgeHandler();
-    // 2. appendMessage runs synchronously, updating leafId
-    leafId = "assistant-entry-101";
-
-    // Wait for the deferred handler to complete
-    await handlerPromise;
+    // Simulate the dispatcher: await handler, then call appendMessage.
+    await bridgeHandler(event);
+    sessionManager.appendMessage(event.message);
+    await sentP;
 
-    expect(capturedEntryId).toBe("assistant-entry-101");
+    expect(captured).toBe("assistant-entry-101");
   });
 
-  it("immediate getLeafId() would capture the stale entry ID (demonstrates the bug)", async () => {
+  it("queueMicrotask deferral would NOT work on pi 0.70.x (regression demonstration)", async () => {
+    // Reproduces why we abandoned queueMicrotask: it resolves inside the
+    // awaited dispatcher, before appendMessage runs.
     let leafId = "user-entry-100";
-    const sessionManager = {
-      getLeafId: () => leafId,
-    };
+    let captured: string | undefined;
 
-    let capturedEntryId: string | undefined;
-
-    // Simulate the OLD (buggy) bridge handler: reads getLeafId() immediately
-    const buggyBridgeHandler = async () => {
-      // No deferral — reads leafId before appendMessage runs
-      capturedEntryId = sessionManager.getLeafId();
+    const buggyBridge = async () => {
+      await new Promise<void>((resolve) => queueMicrotask(resolve));
+      captured = leafId;
     };
 
-    // Simulate pi core's _processAgentEvent:
-    const handlerPromise = buggyBridgeHandler();
-    leafId = "assistant-entry-101"; // too late — handler already read it
-
-    await handlerPromise;
+    // Pi 0.70.x: await the bridge, THEN persist. Mirroring the real ordering:
+    const handlerP = buggyBridge();
+    await handlerP;
+    leafId = "assistant-entry-101"; // appendMessage runs after await
 
-    // Bug: captures the stale leaf, not the assistant's entry
-    expect(capturedEntryId).toBe("user-entry-100");
+    expect(captured).toBe("user-entry-100");
+    expect(captured).not.toBe("assistant-entry-101");
   });
 
-  it("message_start should still capture entryId immediately (no deferral)", async () => {
-    let leafId = "previous-assistant-entry-99";
-    const sessionManager = {
-      getLeafId: () => leafId,
-    };
-
-    let capturedEntryId: string | undefined;
-
-    // Simulate bridge handler for message_start (immediate, no deferral)
-    const messageStartHandler = async () => {
-      capturedEntryId = sessionManager.getLeafId();
-    };
-
-    const handlerPromise = messageStartHandler();
-    // User entry gets written after message_start
-    leafId = "user-entry-100";
-
-    await handlerPromise;
+  it("entry_persisted is the back-fill mechanism for user message_start (where event.message.id is unavailable)", () => {
+    // Behavioural assertion (pure data shape): when the bridge sends a
+    // user message_start, it stamps a nonce; later when pi persists the
+    // user entry, the bridge sends entry_persisted { nonce, entryId }.
+    // The reducer pairs them by nonce. See change: fix-per-message-fork.
+    const start = { type: "message_start", message: { role: "user" }, nonce: "n-1" };
+    const persisted = { type: "entry_persisted", entryId: "user-200", nonce: "n-1" };
 
-    // message_start should capture the leaf BEFORE the user entry is written
-    expect(capturedEntryId).toBe("previous-assistant-entry-99");
+    expect(start.nonce).toBe(persisted.nonce);
+    expect(persisted.entryId).toBe("user-200");
   });
 });

package/packages/extension/src/__tests__/multiselect-list.test.ts

@@ -0,0 +1,137 @@
+import { describe, it, expect, vi } from "vitest";
+import { MultiSelectList } from "../multiselect-list.js";
+
+function make(options: string[] = ["A", "B", "C"], title = "Pick", message?: string) {
+  const list = new MultiSelectList(title, options, message);
+  const onConfirm = vi.fn();
+  const onCancel = vi.fn();
+  list.onConfirm = onConfirm;
+  list.onCancel = onCancel;
+  return { list, onConfirm, onCancel };
+}
+
+describe("MultiSelectList", () => {
+  describe("keybindings", () => {
+    it("space toggles the checked state of the current item and nothing else", () => {
+      const { list, onConfirm } = make();
+      list.handleInput(" "); // toggle index 0 → A checked
+      list.handleInput("\r");
+      expect(onConfirm).toHaveBeenCalledWith(["A"]);
+    });
+
+    it("arrow down moves cursor and space toggles the new current item", () => {
+      const { list, onConfirm } = make();
+      list.handleInput("\u001b[B"); // cursor → 1 (B)
+      list.handleInput(" "); // B checked
+      list.handleInput("\r");
+      expect(onConfirm).toHaveBeenCalledWith(["B"]);
+    });
+
+    it("j / k navigation works like arrows", () => {
+      const { list, onConfirm } = make();
+      list.handleInput("j"); // cursor → 1
+      list.handleInput("j"); // cursor → 2
+      list.handleInput(" "); // C checked
+      list.handleInput("k"); // cursor → 1
+      list.handleInput(" "); // B checked
+      list.handleInput("\r");
+      // selected values returned in ORIGINAL order
+      expect(onConfirm).toHaveBeenCalledWith(["B", "C"]);
+    });
+
+    it("enter with nothing checked confirms with []", () => {
+      const { list, onConfirm } = make();
+      list.handleInput("\r");
+      expect(onConfirm).toHaveBeenCalledWith([]);
+    });
+
+    it("escape cancels; no confirm is fired", () => {
+      const { list, onConfirm, onCancel } = make();
+      list.handleInput(" "); // check A
+      list.handleInput("\u001b");
+      expect(onCancel).toHaveBeenCalledTimes(1);
+      expect(onConfirm).not.toHaveBeenCalled();
+    });
+
+    it("pressing 'a' does NOT bulk-toggle (no select-all in TUI)", () => {
+      const { list, onConfirm } = make();
+      list.handleInput("a");
+      list.handleInput("\r");
+      expect(onConfirm).toHaveBeenCalledWith([]);
+    });
+
+    it("cursor does not go below 0", () => {
+      const { list } = make();
+      list.handleInput("k");
+      list.handleInput("k");
+      expect(list.getCursor()).toBe(0);
+    });
+
+    it("cursor does not go past last item", () => {
+      const { list } = make(["A", "B"]);
+      list.handleInput("j");
+      list.handleInput("j");
+      list.handleInput("j");
+      expect(list.getCursor()).toBe(1);
+    });
+
+    it("toggling twice returns item to unchecked", () => {
+      const { list, onConfirm } = make();
+      list.handleInput(" ");
+      list.handleInput(" ");
+      list.handleInput("\r");
+      expect(onConfirm).toHaveBeenCalledWith([]);
+    });
+
+    it("selected order follows original option order, not toggle order", () => {
+      const { list, onConfirm } = make(["A", "B", "C", "D"]);
+      // toggle D first, then A, then C
+      list.handleInput("j");
+      list.handleInput("j");
+      list.handleInput("j");
+      list.handleInput(" "); // D
+      list.handleInput("k");
+      list.handleInput("k");
+      list.handleInput("k");
+      list.handleInput(" "); // A
+      list.handleInput("j");
+      list.handleInput("j");
+      list.handleInput(" "); // C
+      list.handleInput("\r");
+      expect(onConfirm).toHaveBeenCalledWith(["A", "C", "D"]);
+    });
+  });
+
+  describe("render", () => {
+    it("includes footer hint", () => {
+      const { list } = make();
+      const lines = list.render(80);
+      expect(lines.some((l) => l.includes("space toggle"))).toBe(true);
+      expect(lines.some((l) => l.includes("enter confirm"))).toBe(true);
+      expect(lines.some((l) => l.includes("esc cancel"))).toBe(true);
+    });
+
+    it("renders [ ] for unchecked and [x] for checked items", () => {
+      const { list } = make();
+      list.handleInput(" "); // check A
+      const lines = list.render(80);
+      expect(lines.some((l) => l.includes("[x] A"))).toBe(true);
+      expect(lines.some((l) => l.includes("[ ] B"))).toBe(true);
+    });
+
+    it("renders cursor marker on current item", () => {
+      const { list } = make();
+      list.handleInput("j"); // cursor → 1
+      const lines = list.render(80);
+      // Cursor line should start with "▸ " somewhere
+      expect(lines.some((l) => l.startsWith("▸ ") && l.includes("B"))).toBe(true);
+    });
+
+    it("includes title and message when provided", () => {
+      const { list } = make(["A", "B"], "Pick one or more", "Some context");
+      const lines = list.render(80);
+      expect(lines.some((l) => l.includes("Pick one or more"))).toBe(true);
+      expect(lines.some((l) => l.includes("Some context"))).toBe(true);
+    });
+  });
+});

package/packages/extension/src/__tests__/no-session-replacement-calls.test.ts

@@ -0,0 +1,99 @@
+/**
+ * Repo-level invariant: the bridge MUST NOT call pi's session-replacement
+ * APIs (`pi.newSession(...)`, `ctx.fork(...)`, `ctx.switchSession(...)`)
+ * from any code under `packages/extension/src/`.
+ *
+ * Rationale: pi 0.69.0+ invalidates captured pre-replacement `pi`/`ctx`/
+ * session-bound objects on next access after these calls. The bridge
+ * holds long-lived caches (`cachedCtx`, `cachedModelRegistry`,
+ * `cachedHasUI` in `bridge.ts`; `modelRegistry` in `provider-register.ts`)
+ * that depend on pi being the ONLY originator of session replacement, so
+ * we can re-capture inside the resulting `session_start` handler keyed on
+ * `event.reason ∈ {"new","fork","resume"}`.
+ *
+ * If this test fails: do NOT add the call. Either drive the user-facing
+ * action through the dashboard server (which prompts the user, who
+ * triggers replacement via pi's UI), or wrap your post-switch work in
+ * the `withSession` callback that pi 0.69+ exposes on each replacement
+ * API and capture the freshly-emitted ReplacedSessionContext there.
+ *
+ * See change: pi-zero-seventy-compat.
+ */
+import { describe, it, expect } from "vitest";
+import fs from "node:fs/promises";
+import path from "node:path";
+import url from "node:url";
+
+/**
+ * Each pattern matches `<receiver>.<method>(` allowing for whitespace and
+ * tolerating common variations like `await pi.newSession(...)`. Prefixed
+ * with a non-word boundary so we don't flag method names embedded in
+ * longer identifiers.
+ */
+const PATTERNS: ReadonlyArray<{ name: string; re: RegExp }> = [
+  { name: "pi.newSession",     re: /(?:^|[^.\w])pi\.newSession\s*\(/ },
+  { name: "ctx.fork",          re: /(?:^|[^.\w])ctx\.fork\s*\(/ },
+  { name: "ctx.switchSession", re: /(?:^|[^.\w])ctx\.switchSession\s*\(/ },
+];
+
+/**
+ * Per-line opt-out marker. Use only for documented exceptions (e.g. a
+ * future migration cell that intentionally drives a replacement and
+ * fully re-binds via `withSession`):
+ *   await pi.newSession({ withSession: ... }); // ban:session-replacement-ok
+ */
+const OPT_OUT_MARKER = "ban:session-replacement-ok";
+
+async function* walk(dir: string): AsyncGenerator<string> {
+  const entries = await fs.readdir(dir, { withFileTypes: true });
+  for (const entry of entries) {
+    const full = path.join(dir, entry.name);
+    if (entry.isDirectory()) {
+      if (entry.name === "node_modules" || entry.name === "dist" || entry.name === "__tests__") continue;
+      yield* walk(full);
+    } else if (entry.isFile() && /\.(ts|tsx|mts|cts)$/.test(entry.name)) {
+      yield full;
+    }
+  }
+}
+
+describe("no session-replacement API calls in packages/extension/src/", () => {
+  it("bridge code never invokes pi.newSession / ctx.fork / ctx.switchSession", async () => {
+    const here = path.dirname(url.fileURLToPath(import.meta.url));
+    const srcDir = path.resolve(here, "..");
+    const repoRoot = path.resolve(here, "..", "..", "..", "..");
+
+    const violations: Array<{ file: string; line: number; pattern: string; text: string }> = [];
+
+    for await (const file of walk(srcDir)) {
+      const content = await fs.readFile(file, "utf-8");
+      const lines = content.split(/\r?\n/);
+      lines.forEach((line, idx) => {
+        if (line.includes(OPT_OUT_MARKER)) return;
+        for (const { name, re } of PATTERNS) {
+          if (re.test(line)) {
+            violations.push({
+              file: path.relative(repoRoot, file),
+              line: idx + 1,
+              pattern: name,
+              text: line.trim(),
+            });
+          }
+        }
+      });
+    }
+
+    if (violations.length > 0) {
+      const msg =
+        `Bridge code MUST NOT call pi session-replacement APIs.\n` +
+        `pi 0.69.0+ invalidates captured pre-replacement pi/ctx after these calls;\n` +
+        `the bridge relies on pi being the sole originator of replacement so it can\n` +
+        `re-capture state inside the resulting session_start handler.\n\n` +
+        `Offenders (${violations.length}):\n` +
+        violations
+          .map((v) => `  ${v.file}:${v.line}  [${v.pattern}]  ${v.text}`)
+          .join("\n");
+      expect(violations, msg).toEqual([]);
+    }
+  });
+});

package/packages/extension/src/ask-user-tool.ts

@@ -6,7 +6,8 @@
  * register ask_user. Runtime registration bypasses detectExtensionConflicts.
  */
 import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
-import { Type } from "@sinclair/typebox";
+import { Type } from "typebox";
+import { polyfillMultiselect } from "./multiselect-polyfill.js";
 
 // ──────────────────────────────────────────────────────────────────────────
 // Single-question schema arms (reused inside the batch arm's questions array)
@@ -121,7 +122,7 @@ export function registerAskUserTool(pi: ExtensionAPI): void {
     name: "ask_user",
     label: "Ask User",
     description:
-      "Ask the user a question interactively. Use this when you need clarification, confirmation, or a choice from the user before proceeding.",
+      "Ask the user a question interactively. Use this when you need clarification, confirmation, or a choice from the user before proceeding. UI provides a Select all toggle; do not add one.",
     promptSnippet:
       "Ask the user interactive questions (confirm, select, multiselect, input, or batch — multiple related questions at once)",
     promptGuidelines: [
@@ -254,7 +255,7 @@ export function registerAskUserTool(pi: ExtensionAPI): void {
                     `ask_user batch: sub-question method "multiselect" requires a non-empty "options" array.`,
                   );
                 }
-                answer = await (ctx.ui as any).multiselect(subTitle, opts, subMsg);
+                answer = await polyfillMultiselect(ctx, subTitle, opts, subMsg);
                 break;
               }
               case "input":
@@ -336,7 +337,7 @@ export function registerAskUserTool(pi: ExtensionAPI): void {
           result = await ctx.ui.select(title, options, msgOpts);
           break;
         case "multiselect":
-          result = await (ctx.ui as any).multiselect(title, options, msgOpts);
+          result = await polyfillMultiselect(ctx, title, options, msgOpts);
           break;
         case "input":
           result = await ctx.ui.input(title, params.placeholder, msgOpts);

package/packages/extension/src/bridge.ts

@@ -186,6 +186,61 @@ function initBridge(pi: ExtensionAPI) {
   let lastThinkingLevel: string | undefined;
   let promptBus: PromptBus | undefined;
 
+  // ── Per-message entry id tracking (for fix-per-message-fork) ──
+  // Pi 0.69+ awaits extension handlers BEFORE sessionManager.appendMessage runs,
+  // which means getLeafId() at emit time returns the previous leaf, not the
+  // entry id of the message currently being emitted. We solve this by:
+  //  1. Wrapping ctx.sessionManager.appendMessage at session_start to stamp
+  //     the just-generated entry id onto the message object reference.
+  //  2. Deferring the message_end enrichment-and-send via setTimeout(0) so
+  //     the awaited dispatcher unwinds and appendMessage runs in between.
+  //  3. Stamping a nonce on message_start/message_end events; emitting an
+  //     entry_persisted event after appendMessage so the client reducer can
+  //     back-fill user-message ChatMessage.entryId.
+  // See change: fix-per-message-fork.
+  const idByMessage = new WeakMap<object, string>();
+  const pendingNonces = new WeakMap<object, string>();
+  let nonceCounter = 0;
+  const nextNonce = (): string => `n-${++nonceCounter}-${Date.now()}`;
+  let appendMessageWrapped = false;
+  let lastWrappedSm: any = null;
+
+  /**
+   * Wrap ctx.sessionManager.appendMessage once per session so that when pi
+   * generates an entry id we capture it in the WeakMap and emit
+   * entry_persisted to the server.
+   */
+  function wrapAppendMessageForCtx(ctx: any): void {
+    const sm = ctx?.sessionManager;
+    if (!sm || typeof sm.appendMessage !== "function") return;
+    // Re-wrap when sessionManager identity changes (session replacement).
+    if (sm === lastWrappedSm && appendMessageWrapped) return;
+    const original = sm.appendMessage.bind(sm);
+    sm.appendMessage = (msg: any, ...rest: any[]) => {
+      const result = original(msg, ...rest);
+      try {
+        if (msg && typeof msg === "object" && typeof msg.id === "string") {
+          idByMessage.set(msg as object, msg.id);
+          const nonce = pendingNonces.get(msg as object);
+          if (nonce && sessionReady && isActive()) {
+            const ev = {
+              type: "entry_persisted",
+              entryId: msg.id,
+              nonce,
+            };
+            connection.send(mapEventToProtocol(sessionId, ev));
+            pendingNonces.delete(msg as object);
+          }
+        }
+      } catch (err) {
+        console.error("[dashboard] entry_persisted emit failed:", err);
+      }
+      return result;
+    };
+    lastWrappedSm = sm;
+    appendMessageWrapped = true;
+  }
+
   /** Wrap a callback so errors log instead of crashing the host pi agent. */
   function safe<T extends (...args: any[]) => any>(fn: T): T {
     return ((...args: any[]) => {
@@ -612,30 +667,53 @@ function initBridge(pi: ExtensionAPI) {
         }
       }
 
-      // For message_start, enrich with entryId immediately (current leaf)
+      // For message_start: stamp a nonce on the event so the client reducer
+      // can correlate a later entry_persisted back-fill with this bubble.
+      // We do NOT attach entryId here — the message has no id yet on pi
+      // 0.69+ (persistence is deferred to message_end). See change:
+      // fix-per-message-fork.
       if (eventType === "message_start") {
-        const entryId = ctx.sessionManager?.getLeafId?.();
-        if (entryId) {
-          const enriched = { ...event, entryId };
+        wrapAppendMessageForCtx(ctx);
+        const messageRef = (event as any).message;
+        if (messageRef && typeof messageRef === "object") {
+          const nonce = nextNonce();
+          pendingNonces.set(messageRef as object, nonce);
+          const enriched = { ...event, nonce };
           const msg = mapEventToProtocol(sessionId, enriched);
           connection.send(msg);
           return;
         }
       }
 
-      // For message_end, defer getLeafId() so it runs after pi core persists the entry.
-      // Pi core calls _emit (which invokes this handler) BEFORE appendMessage (which updates leafId).
-      // Since _emit doesn't await async handlers, yielding via queueMicrotask lets appendMessage
-      // run first, so getLeafId() returns the correct entry ID for the just-persisted message.
+      // For message_end: defer the SEND via setTimeout(0). Pi 0.69+ runs
+      // sessionManager.appendMessage AFTER the awaited extension dispatcher
+      // returns, so a queueMicrotask deferral is no longer enough. By the
+      // time the macrotask fires, appendMessage has run, pi has mutated
+      // event.message.id in place, and the wrapped appendMessage above has
+      // populated idByMessage. We also stamp a nonce so a downstream
+      // entry_persisted can correlate (covers user message_end where the
+      // earlier message_start nonce is what the reducer is waiting on).
+      // See change: fix-per-message-fork.
       if (eventType === "message_end") {
-        await new Promise<void>(resolve => queueMicrotask(resolve));
-        const entryId = ctx.sessionManager?.getLeafId?.();
-        if (entryId) {
-          const enriched = { ...event, entryId };
-          const msg = mapEventToProtocol(sessionId, enriched);
-          connection.send(msg);
-          return;
+        wrapAppendMessageForCtx(ctx);
+        const messageRef = (event as any).message;
+        const nonce = messageRef && typeof messageRef === "object"
+          ? (pendingNonces.get(messageRef as object) ?? nextNonce())
+          : nextNonce();
+        if (messageRef && typeof messageRef === "object" && !pendingNonces.has(messageRef as object)) {
+          pendingNonces.set(messageRef as object, nonce);
         }
+        setTimeout(() => {
+          if (!isActive() || !sessionReady) return;
+          const entryId =
+            (messageRef && typeof messageRef === "object" && typeof messageRef.id === "string" ? messageRef.id : undefined)
+            ?? (messageRef ? idByMessage.get(messageRef as object) : undefined)
+            ?? ctx.sessionManager?.getLeafId?.();
+          const enriched = { ...event, entryId, nonce };
+          const protoMsg = mapEventToProtocol(sessionId, enriched);
+          connection.send(protoMsg);
+        }, 0);
+        return;
       }
 
       const msg = mapEventToProtocol(sessionId, event);
@@ -694,6 +772,15 @@ function initBridge(pi: ExtensionAPI) {
     cachedCtx = ctx;
     sessionId = newSessionId;
 
+    // Wrap sessionManager.appendMessage so that future message_end events can
+    // recover the just-generated entry id, even when their setTimeout(0)
+    // fires before pi has finished mutating event.message in place. The
+    // helper is idempotent and re-wraps on session replacement.
+    // See change: fix-per-message-fork.
+    appendMessageWrapped = false;
+    lastWrappedSm = null;
+    wrapAppendMessageForCtx(ctx);
+
     // Register ask_user at runtime (not at load time) to avoid static
     // tool-name conflicts with other extensions like pi-flows.
     registerAskUserTool(pi);