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

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@blackbelt-technology/pi-agent-dashboard",
-  "version": "0.4.0",
+  "version": "0.4.2",
   "description": "Web dashboard for monitoring and interacting with pi agent sessions",
   "repository": {
     "type": "git",
@@ -49,9 +49,9 @@
     "postinstall": "node packages/server/scripts/fix-pty-permissions.cjs",
     "dev": "npm run dev --workspace=@blackbelt-technology/pi-dashboard-web",
     "build": "npm run build --workspace=@blackbelt-technology/pi-dashboard-web",
-    "test": "HOME=$(mktemp -d -t pi-test-XXXXXX) vitest run",
-    "test:watch": "HOME=$(mktemp -d -t pi-test-XXXXXX) vitest",
-    "test:bootstrap": "HOME=$(mktemp -d -t pi-test-XXXXXX) vitest run packages/shared/src/__tests__/bootstrap",
+    "test": "HOME=$(mktemp -d -t pi-test-XXXXXX) NODE_OPTIONS=\"--localstorage-file=$(mktemp -t pi-test-ls-XXXXXX)\" vitest run",
+    "test:watch": "HOME=$(mktemp -d -t pi-test-XXXXXX) NODE_OPTIONS=\"--localstorage-file=$(mktemp -t pi-test-ls-XXXXXX)\" vitest",
+    "test:bootstrap": "HOME=$(mktemp -d -t pi-test-XXXXXX) NODE_OPTIONS=\"--localstorage-file=$(mktemp -t pi-test-ls-XXXXXX)\" vitest run packages/shared/src/__tests__/bootstrap",
     "test:bootstrap:watch": "HOME=$(mktemp -d -t pi-test-XXXXXX) vitest packages/shared/src/__tests__/bootstrap",
     "lint": "tsc --noEmit",
     "reload": "./scripts/reload-all.sh",
@@ -66,9 +66,9 @@
     "screenshots": "npm --prefix site run screenshots"
   },
   "dependencies": {
-    "@blackbelt-technology/pi-dashboard-extension": "^0.4.0",
-    "@blackbelt-technology/pi-dashboard-server": "^0.4.0",
-    "@blackbelt-technology/pi-dashboard-web": "^0.4.0"
+    "@blackbelt-technology/pi-dashboard-extension": "^0.4.2",
+    "@blackbelt-technology/pi-dashboard-server": "^0.4.2",
+    "@blackbelt-technology/pi-dashboard-web": "^0.4.2"
   },
   "devDependencies": {
     "jsdom": "^29.0.2",
@@ -83,7 +83,7 @@
     "@oh-my-pi/pi-ai": "*",
     "@oh-my-pi/pi-coding-agent": "*",
     "@oh-my-pi/pi-tui": "*",
-    "@sinclair/typebox": "*"
+    "typebox": "*"
   },
   "peerDependenciesMeta": {
     "@mariozechner/pi-coding-agent": {
@@ -103,6 +103,9 @@
     },
     "@oh-my-pi/pi-tui": {
       "optional": true
+    },
+    "typebox": {
+      "optional": true
     }
   }
 }

package/packages/extension/package.json

@@ -1,8 +1,13 @@
 {
   "name": "@blackbelt-technology/pi-dashboard-extension",
-  "version": "0.4.0",
+  "version": "0.4.2",
   "description": "Pi bridge extension for pi-dashboard",
   "type": "module",
+  "repository": {
+    "type": "git",
+    "url": "https://github.com/BlackBeltTechnology/pi-agent-dashboard",
+    "directory": "packages/extension"
+  },
   "publishConfig": {
     "access": "public"
   },
@@ -19,12 +24,13 @@
     ".pi/skills/pi-dashboard/"
   ],
   "dependencies": {
-    "@blackbelt-technology/pi-dashboard-shared": "^0.4.0",
+    "@blackbelt-technology/pi-dashboard-shared": "^0.4.2",
     "ws": "^8.18.0"
   },
   "peerDependencies": {
     "@mariozechner/pi-coding-agent": "*",
-    "@mariozechner/pi-tui": "*"
+    "@mariozechner/pi-tui": "*",
+    "typebox": "*"
   },
   "peerDependenciesMeta": {
     "@mariozechner/pi-coding-agent": {
@@ -36,6 +42,7 @@
   },
   "devDependencies": {
     "@mariozechner/pi-tui": "*",
-    "@types/ws": "^8.18.1"
+    "@types/ws": "^8.18.1",
+    "typebox": "^1.1.33"
   }
 }

package/packages/extension/src/__tests__/ask-user-schema-discriminator.test.ts

@@ -0,0 +1,141 @@
+/**
+ * Regression guard for the `ask_user` tool's parameters JSON Schema shape.
+ *
+ * History:
+ *   1. Pre-a53933f — schema was `Type.Union` of per-method `Type.Object`
+ *      arms (root-level `anyOf`). Anthropic loved it; OpenAI rejected it
+ *      ("schema must be a JSON Schema of 'type: \"object\"', got
+ *      'type: \"None\"'").
+ *
+ *   2. a53933f — collapsed to a single flat `Type.Object` (root
+ *      `type: "object"`, all fields optional). OpenAI happy. Anthropic
+ *      lost its per-method strictness; Claude started emitting
+ *      multiselect calls without `options` and the dashboard auto-
+ *      cancelled them — the user-visible bug behind
+ *      fix-multiselect-auto-cancel-on-dashboard.
+ *
+ *   3. Layer-2 attempt of fix-multiselect-auto-cancel-on-dashboard —
+ *      tried to keep root `type: "object"` AND attach a body-level
+ *      `oneOf` discriminator over `method`. Anthropic worked. Real
+ *      OpenAI gpt-5 rejected it (verified 2026-04-30) with: "schema
+ *      must have type 'object' and not have 'oneOf' / 'anyOf' / 'allOf'
+ *      / 'enum' / 'not' at the top level." The fallback documented in
+ *      tasks.md §9.7 was taken: Layer 2 dropped, Layer 1 (multiselect
+ *      dashboard routing) ships alone — that's what actually fixes the
+ *      user bug.
+ *
+ * This test pins the post-fallback shape so it cannot regress in
+ * either direction:
+ *   • Root MUST be `type: "object"` (OpenAI rule).
+ *   • Root MUST NOT have `oneOf`, `anyOf`, `allOf`, `enum`, `not`.
+ *   • The same five forbidden keys MUST NOT appear on `SubQuestionSchema`
+ *     either — OpenAI applies the rule recursively per the error message.
+ *
+ * Per-method strictness is enforced by `prepareArguments` rescue + the
+ * `execute` runtime switch (covered by `ask-user-tool.test.ts`).
+ *
+ * See change: fix-multiselect-auto-cancel-on-dashboard.
+ */
+import { describe, it, expect, vi } from "vitest";
+import { registerAskUserTool } from "../ask-user-tool.js";
+
+const FORBIDDEN_TOP_LEVEL_KEYS = ["oneOf", "anyOf", "allOf", "enum", "not"] as const;
+
+interface PiSchema {
+  type?: string;
+  properties?: Record<string, any>;
+  required?: string[];
+  description?: string;
+  oneOf?: unknown;
+  anyOf?: unknown;
+  allOf?: unknown;
+  enum?: unknown;
+  not?: unknown;
+}
+
+function captureRegisteredTool() {
+  const calls: any[] = [];
+  const pi = {
+    registerTool: vi.fn((def: any) => calls.push(def)),
+  };
+  registerAskUserTool(pi as any);
+  expect(calls.length).toBe(1);
+  return calls[0];
+}
+
+function getSubQuestionSchema(parameters: PiSchema): PiSchema {
+  const items = parameters.properties?.questions?.items;
+  expect(items).toBeDefined();
+  return items as PiSchema;
+}
+
+// Walk a Union<Literal> schema (typebox emits it as `anyOf` of `{const: ...}`)
+// into a flat list of literal values.
+function flattenUnionLiterals(s: any): string[] {
+  if (!s) return [];
+  if (Array.isArray(s.anyOf)) return s.anyOf.flatMap(flattenUnionLiterals);
+  if (s.const !== undefined) return [String(s.const)];
+  return [];
+}
+
+describe("ask_user parameters schema — OpenAI strict-mode shape", () => {
+  it("root is type:object", () => {
+    const tool = captureRegisteredTool();
+    const params = tool.parameters as PiSchema;
+    expect(params.type).toBe("object");
+  });
+
+  it.each(FORBIDDEN_TOP_LEVEL_KEYS)(
+    "root has NO `%s` (OpenAI strict mode forbids it at the top level)",
+    (key) => {
+      const tool = captureRegisteredTool();
+      const params = tool.parameters as Record<string, unknown>;
+      expect(
+        params[key],
+        `parameters.${key} would break OpenAI gpt-5: "schema must have type 'object' and not have 'oneOf' / 'anyOf' / 'allOf' / 'enum' / 'not' at the top level."`,
+      ).toBeUndefined();
+    },
+  );
+
+  it("declares the five method literals (regression guard for MethodEnum)", () => {
+    const tool = captureRegisteredTool();
+    const methodSchema = (tool.parameters as PiSchema).properties?.method;
+    expect(methodSchema).toBeDefined();
+    // typebox emits a Union<Literal> as `anyOf` UNDER `properties.method`
+    // (NOT at the schema root). That is OpenAI-legal because it isn't at
+    // the top level. We only assert the literal set is preserved.
+    const methods = flattenUnionLiterals(methodSchema);
+    expect(methods).toEqual(
+      expect.arrayContaining(["confirm", "select", "multiselect", "input", "batch"]),
+    );
+  });
+});
+
+describe("ask_user SubQuestionSchema — OpenAI strict-mode shape", () => {
+  it("sub-question schema is type:object", () => {
+    const tool = captureRegisteredTool();
+    const sq = getSubQuestionSchema(tool.parameters);
+    expect(sq.type).toBe("object");
+  });
+
+  it.each(FORBIDDEN_TOP_LEVEL_KEYS)(
+    "sub-question schema has NO `%s` at top level (OpenAI rule applies recursively)",
+    (key) => {
+      const tool = captureRegisteredTool();
+      const sq = getSubQuestionSchema(tool.parameters) as Record<string, unknown>;
+      expect(sq[key]).toBeUndefined();
+    },
+  );
+
+  it("sub-question excludes 'batch' from method literals (no nesting)", () => {
+    const tool = captureRegisteredTool();
+    const sq = getSubQuestionSchema(tool.parameters);
+    const methodSchema = sq.properties?.method;
+    expect(methodSchema).toBeDefined();
+    const methods = flattenUnionLiterals(methodSchema);
+    expect(methods).toEqual(
+      expect.arrayContaining(["confirm", "select", "multiselect", "input"]),
+    );
+    expect(methods).not.toContain("batch");
+  });
+});

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

@@ -1,7 +1,7 @@
 import { describe, it, expect, vi } from "vitest";
 
 // Mock modules before importing
-vi.mock("@sinclair/typebox", () => ({
+vi.mock("typebox", () => ({
   Type: {
     Object: vi.fn(() => ({})),
     String: vi.fn(() => ({})),
@@ -44,62 +44,119 @@ describe("registerAskUserTool", () => {
     expect(tool.promptGuidelines.length).toBeGreaterThan(0);
   });
 
+  it("description instructs agents not to add a Select all option", () => {
+    const pi = createMockPi();
+    registerAskUserTool(pi as any);
+    const tool = pi.registerTool.mock.calls[0][0];
+    expect(tool.description).toMatch(/UI provides a Select all/i);
+  });
+
   describe("message passthrough", () => {
     function getToolAndMockCtx() {
       const pi = createMockPi();
       registerAskUserTool(pi as any);
       const tool = pi.registerTool.mock.calls[0][0];
+      // `custom` stands in for the multiselect polyfill: it invokes the factory
+      // with a `done` callback; the factory-returned component exposes
+      // onConfirm/onCancel. We auto-confirm with ["A"] to preserve the legacy
+      // mock return value that the multiselect assertions expected.
+      const custom = vi.fn().mockImplementation(async (factory: any) => {
+        return await new Promise<unknown>((resolve) => {
+          const component: any = factory({}, {}, {}, (r: unknown) => resolve(r));
+          component?.onConfirm?.(["A"]);
+        });
+      });
       const ctx = {
         ui: {
           confirm: vi.fn().mockResolvedValue(true),
           select: vi.fn().mockResolvedValue("A"),
           input: vi.fn().mockResolvedValue("hello"),
-          multiselect: vi.fn().mockResolvedValue(["A"]),
+          custom,
         },
       };
-      return { tool, ctx };
+      return { tool, ctx, custom };
     }
 
     it("passes message through opts for input", async () => {
       const { tool, ctx } = getToolAndMockCtx();
       await tool.execute("id", { method: "input", title: "Q", message: "Details here" }, undefined, undefined, ctx);
-      expect(ctx.ui.input).toHaveBeenCalledWith("Q", undefined, { message: "Details here" });
+      // toolCallId is also threaded through opts since change
+      // `fix-interactive-ui-reorder`. Asserts both fields without
+      // pinning property order.
+      expect(ctx.ui.input).toHaveBeenCalledWith(
+        "Q",
+        undefined,
+        expect.objectContaining({ message: "Details here", toolCallId: "id" }),
+      );
     });
 
     it("passes message through opts for select", async () => {
       const { tool, ctx } = getToolAndMockCtx();
       await tool.execute("id", { method: "select", title: "Pick", message: "Context", options: ["A", "B"] }, undefined, undefined, ctx);
-      expect(ctx.ui.select).toHaveBeenCalledWith("Pick", ["A", "B"], { message: "Context" });
+      expect(ctx.ui.select).toHaveBeenCalledWith(
+        "Pick",
+        ["A", "B"],
+        expect.objectContaining({ message: "Context", toolCallId: "id" }),
+      );
     });
 
-    it("passes message through opts for multiselect", async () => {
+    it("dispatches multiselect through the polyfill via ctx.ui.custom", async () => {
       const { tool, ctx } = getToolAndMockCtx();
-      await tool.execute("id", { method: "multiselect", title: "Multi", message: "Info", options: ["A"] }, undefined, undefined, ctx);
-      expect(ctx.ui.multiselect).toHaveBeenCalledWith("Multi", ["A"], { message: "Info" });
+      const result = await tool.execute(
+        "id",
+        { method: "multiselect", title: "Multi", message: "Info", options: ["A"] },
+        undefined,
+        undefined,
+        ctx,
+      );
+      // Polyfill routes via custom(factory); multiselect is not called directly.
+      expect(ctx.ui.custom).toHaveBeenCalledTimes(1);
+      expect(result.details.method).toBe("multiselect");
+      expect(result.details.result).toEqual(["A"]);
     });
 
-    it("does not pass opts when message is undefined", async () => {
+    it("passes only toolCallId through opts when message is undefined", async () => {
       const { tool, ctx } = getToolAndMockCtx();
       await tool.execute("id", { method: "input", title: "Q" }, undefined, undefined, ctx);
-      expect(ctx.ui.input).toHaveBeenCalledWith("Q", undefined, undefined);
+      // Even with no `message`, the wrapper still attaches toolCallId so
+      // the resulting prompt_request can be paired by the client reducer.
+      expect(ctx.ui.input).toHaveBeenCalledWith(
+        "Q",
+        undefined,
+        expect.objectContaining({ toolCallId: "id" }),
+      );
     });
 
     it("falls back to message when title is missing", async () => {
       const { tool, ctx } = getToolAndMockCtx();
       await tool.execute("id", { method: "input", message: "Detailed question" }, undefined, undefined, ctx);
-      expect(ctx.ui.input).toHaveBeenCalledWith("Detailed question", undefined, { message: "Detailed question" });
+      expect(ctx.ui.input).toHaveBeenCalledWith(
+        "Detailed question",
+        undefined,
+        expect.objectContaining({ message: "Detailed question", toolCallId: "id" }),
+      );
     });
 
     it("falls back to 'Question' when both title and message are missing", async () => {
       const { tool, ctx } = getToolAndMockCtx();
       await tool.execute("id", { method: "confirm" }, undefined, undefined, ctx);
-      expect(ctx.ui.confirm).toHaveBeenCalledWith("Question", "");
+      // confirm now also threads toolCallId via 3rd arg.
+      expect(ctx.ui.confirm).toHaveBeenCalledWith(
+        "Question",
+        "",
+        expect.objectContaining({ toolCallId: "id" }),
+      );
     });
 
     it("parses options from JSON string", async () => {
       const { tool, ctx } = getToolAndMockCtx();
       await tool.execute("id", { method: "select", title: "Pick", options: '["A", "B"]' }, undefined, undefined, ctx);
-      expect(ctx.ui.select).toHaveBeenCalledWith("Pick", ["A", "B"], undefined);
+      // No message, no other opts — only toolCallId.
+      expect(ctx.ui.select).toHaveBeenCalledWith(
+        "Pick",
+        ["A", "B"],
+        expect.objectContaining({ toolCallId: "id" }),
+      );
     });
 
     it("throws when select reaches execute with unparseable options string", async () => {
@@ -123,7 +180,7 @@ describe("registerAskUserTool", () => {
       await expect(
         tool.execute("id", { method: "multiselect", title: "Pick" }, undefined, undefined, ctx),
       ).rejects.toThrow(/options/i);
-      expect(ctx.ui.multiselect).not.toHaveBeenCalled();
+      expect(ctx.ui.custom).not.toHaveBeenCalled();
     });
   });
 
@@ -134,6 +191,19 @@ describe("registerAskUserTool", () => {
       return pi.registerTool.mock.calls[0][0];
     }
 
+    it("leaves empty {} args untouched (no synthetic method) so schema rejection still fires", () => {
+      // Regression test for the Opus-emits-empty-args bug seen in session 019dd05c.
+      // Our rescue layer must NOT silently fabricate a method/title when there is
+      // nothing to rescue — the framework's schema validator must still reject {}
+      // so the model is forced to retry with valid args.
+      const tool = getTool();
+      const result = tool.prepareArguments({});
+      expect(result.method).toBeUndefined();
+      expect(result.title).toBeUndefined();
+      expect(result.questions).toBeUndefined();
+      expect(Object.keys(result).filter((k) => k !== "__normalizations")).toHaveLength(0);
+    });
+
     it("parses stringified options array", () => {
       const tool = getTool();
       const result = tool.prepareArguments({ method: "select", title: "Pick", options: '["A", "B"]' });
@@ -311,12 +381,18 @@ describe("registerAskUserTool", () => {
       const pi = createMockPi();
       registerAskUserTool(pi as any);
       const tool = pi.registerTool.mock.calls[0][0];
+      const custom = vi.fn().mockImplementation(async (factory: any) => {
+        return await new Promise<unknown>((resolve) => {
+          const component: any = factory({}, {}, {}, (r: unknown) => resolve(r));
+          component?.onConfirm?.(["A"]);
+        });
+      });
       const ctx = {
         ui: {
           confirm: vi.fn().mockResolvedValue(true),
           select: vi.fn().mockResolvedValue("A"),
           input: vi.fn().mockResolvedValue("hello"),
-          multiselect: vi.fn().mockResolvedValue(["A"]),
+          custom,
         },
       };
       return { tool, ctx };

package/packages/extension/src/__tests__/bridge-entry-id-pi-070.test.ts

@@ -0,0 +1,174 @@
+/**
+ * Tests for the bridge entryId stamping under pi 0.70.x's emit-then-await-then-append
+ * ordering. Pi 0.70.x's _processAgentEvent does (paraphrased):
+ *
+ *   await this._emitExtensionEvent(event);       // <-- bridge runs here, awaited
+ *   this._emit(event);                           // sync legacy listeners
+ *   if (event.type === "message_end") {
+ *     sessionManager.appendMessage(event.message); // <-- entry id GENERATED HERE
+ *   }
+ *
+ * The bridge's old `queueMicrotask` deferral resolves INSIDE the awaited dispatcher,
+ * before appendMessage runs — so getLeafId() still returns the previous leaf. The fix
+ * is `setTimeout(0)` (macrotask) so the entire await chain unwinds and appendMessage
+ * runs first; OR reading `event.message.id` after pi mutates it in-place.
+ *
+ * This test simulates that ordering and asserts the correct mechanisms.
+ */
+import { describe, it, expect } from "vitest";
+
+interface SimMessage {
+  role: string;
+  content: string;
+  id?: string;
+}
+
+/**
+ * Simulate pi 0.70.x's _processAgentEvent ordering. Returns a promise that
+ * resolves when the entire event has been processed (including appendMessage).
+ *
+ * The `bridgeHandler` is registered as an "extension handler" — runs awaited
+ * inside _emitExtensionEvent. It receives the event and a pseudo-ctx with
+ * sessionManager.getLeafId().
+ */
+async function simulatePi070Emit(opts: {
+  event: { type: string; message: SimMessage };
+  state: { leafId: string; nextId: string };
+  appendMessage: (msg: SimMessage) => string; // returns the new id
+  bridgeHandler: (event: any, ctx: any) => Promise<void> | void;
+}): Promise<void> {
+  const ctx = {
+    sessionManager: { getLeafId: () => opts.state.leafId },
+  };
+
+  // Step 1: await _emitExtensionEvent — runs the bridge handler awaited.
+  await opts.bridgeHandler(opts.event, ctx);
+
+  // Step 2: _emit (sync legacy listeners) — no-op in this simulation.
+
+  // Step 3: persistence on message_end.
+  if (opts.event.type === "message_end") {
+    const id = opts.appendMessage(opts.event.message);
+    opts.state.leafId = id;
+  }
+}
+
+describe("pi 0.70 emit/append ordering", () => {
+  it("queueMicrotask deferral DOES NOT capture the post-persist id (the bug)", async () => {
+    const state = { leafId: "prev", nextId: "new-id-42" };
+    let captured: string | undefined;
+
+    const buggyBridge = async (event: any, ctx: any) => {
+      // What the OLD bridge does today:
+      await new Promise<void>((resolve) => queueMicrotask(resolve));
+      captured = ctx.sessionManager.getLeafId();
+    };
+
+    await simulatePi070Emit({
+      event: { type: "message_end", message: { role: "assistant", content: "hi" } },
+      state,
+      appendMessage: (m) => {
+        m.id = state.nextId;
+        return state.nextId;
+      },
+      bridgeHandler: buggyBridge,
+    });
+
+    // Bug: captured is the previous leaf, NOT the just-appended id.
+    expect(captured).toBe("prev");
+    expect(captured).not.toBe("new-id-42");
+  });
+
+  it("setTimeout(0) deferral DOES capture the post-persist id (the fix)", async () => {
+    const state = { leafId: "prev", nextId: "new-id-42" };
+    let capturedFromGetLeafId: string | undefined;
+    let capturedFromMessageId: string | undefined;
+    let sendDone!: () => void;
+    const sendCompleted = new Promise<void>((r) => { sendDone = r; });
+
+    // The bridge schedules and returns synchronously — the only way the
+    // awaited dispatcher can unwind so appendMessage runs before the timeout.
+    const fixedBridge = (event: any, ctx: any) => {
+      setTimeout(() => {
+        capturedFromMessageId = (event.message as any).id;
+        capturedFromGetLeafId = ctx.sessionManager.getLeafId();
+        sendDone();
+      }, 0);
+    };
+
+    await simulatePi070Emit({
+      event: { type: "message_end", message: { role: "assistant", content: "hi" } },
+      state,
+      appendMessage: (m) => {
+        m.id = state.nextId;
+        return state.nextId;
+      },
+      bridgeHandler: fixedBridge,
+    });
+    await sendCompleted;
+
+    // Both signals should now point at the just-persisted entry.
+    expect(capturedFromMessageId).toBe("new-id-42");
+    expect(capturedFromGetLeafId).toBe("new-id-42");
+  });
+
+  it("WeakMap-on-appendMessage captures the id even before the macrotask", async () => {
+    const state = { leafId: "prev", nextId: "new-id-77" };
+    const idByMessage = new WeakMap<object, string>();
+    const wrappedAppend = (m: SimMessage): string => {
+      m.id = state.nextId;
+      idByMessage.set(m as object, m.id);
+      return m.id;
+    };
+
+    let viaWeakMap: string | undefined;
+    let viaMutation: string | undefined;
+    let sendDone!: () => void;
+    const sentP = new Promise<void>((r) => { sendDone = r; });
+
+    // CRITICAL: bridge SCHEDULES the send and RETURNS IMMEDIATELY.
+    // It does NOT await its own setTimeout — that would keep the
+    // outer dispatcher awaiting and we'd be back to the queueMicrotask
+    // bug (timeout fires before appendMessage).
+    const fixedBridge = (event: any) => {
+      setTimeout(() => {
+        viaMutation = (event.message as any).id;
+        viaWeakMap = idByMessage.get(event.message as object);
+        sendDone();
+      }, 0);
+      // Return synchronously — let the awaited dispatcher unwind.
+    };
+
+    await simulatePi070Emit({
+      event: { type: "message_end", message: { role: "assistant", content: "hi" } },
+      state,
+      appendMessage: wrappedAppend,
+      bridgeHandler: fixedBridge,
+    });
+    await sentP;
+
+    expect(viaMutation).toBe("new-id-77");
+    expect(viaWeakMap).toBe("new-id-77");
+  });
+
+  it("user message_start has NO id (pi defers user persistence to message_end)", async () => {
+    const state = { leafId: "prev-assistant", nextId: "new-user-id" };
+    let captured: string | undefined;
+
+    const fixedBridge = async (event: any) => {
+      await new Promise<void>((resolve) => setTimeout(resolve, 0));
+      captured = (event.message as any).id; // still undefined for message_start
+    };
+
+    await simulatePi070Emit({
+      event: { type: "message_start", message: { role: "user", content: "hello" } },
+      state,
+      appendMessage: () => state.nextId, // not called for message_start
+      bridgeHandler: fixedBridge,
+    });
+
+    // No id available at message_start time — must rely on entry_persisted
+    // back-fill (delivered when the message_end of the SAME message fires later).
+    expect(captured).toBeUndefined();
+  });
+});

package/packages/extension/src/__tests__/event-forwarder.test.ts

@@ -74,6 +74,36 @@ describe("mapEventToProtocol", () => {
     expect(result.event.data).toEqual(piEvent);
   });
 
+  it("should map an entry_persisted event (per fix-per-message-fork)", () => {
+    const piEvent = {
+      type: "entry_persisted",
+      entryId: "abc-123",
+      nonce: "n-7",
+    };
+    const result = mapEventToProtocol(sessionId, piEvent);
+    expect(result.type).toBe("event_forward");
+    expect(result.sessionId).toBe(sessionId);
+    expect(result.event.eventType).toBe("entry_persisted");
+    expect(result.event.data).toMatchObject({
+      type: "entry_persisted",
+      entryId: "abc-123",
+      nonce: "n-7",
+    });
+  });
+
+  it("should map a message_end event with nonce (per fix-per-message-fork)", () => {
+    const piEvent = {
+      type: "message_end",
+      message: { role: "assistant", content: "hi", id: "asst-9" },
+      entryId: "asst-9",
+      nonce: "n-7",
+    };
+    const result = mapEventToProtocol(sessionId, piEvent);
+    expect(result.event.eventType).toBe("message_end");
+    expect((result.event.data as any).nonce).toBe("n-7");
+    expect((result.event.data as any).entryId).toBe("asst-9");
+  });
+
   it("should strip non-serializable fields", () => {
     const piEvent = {
       type: "test_event",