fix(pi): restore the session's mode after OMP session switches

OMP keeps Pi's legacy session_switch and session_branch events instead of re-emitting session_start on /new, /resume, /fork and /branch, so a new session inherited the previous session's mode. Restore the mode on those events, like session_tree.

Author: andrebrait

pi-extension/index.js

@@ -200,9 +200,13 @@ export default function ponytailExtension(pi) {
     }
   });
 
-  pi.on("session_tree", async (_event, ctx) => {
-    restoreSessionMode(ctx);
-  });
+  // Both hosts emit session_tree. Pi re-emits session_start on /new, /resume and
+  // /fork; OMP emits session_switch and session_branch instead.
+  for (const eventName of ["session_tree", "session_switch", "session_branch"]) {
+    pi.on(eventName, async (_event, ctx) => {
+      restoreSessionMode(ctx);
+    });
+  }
 
   pi.on("agent_start", async (_event, ctx) => {
     isActive = true;

pi-extension/test/extension.test.js

@@ -181,6 +181,25 @@ test("session_tree restores the selected branch's mode without persisting it", a
   assert.equal(notifications.length, notificationCount, "navigation must not repeat the startup toast");
 }));
 
+test("OMP session_switch and session_branch restore the target session's mode", async () => withTempConfig(async () => {
+  // OMP keeps the legacy Pi events instead of re-emitting session_start on /new, /resume, /fork and /branch.
+  const { events, appendedEntries } = createPiHarness();
+  const lite = { type: "custom", customType: "ponytail-mode", data: { mode: "lite" } };
+  const off = { type: "custom", customType: "ponytail-mode", data: { mode: "off" } };
+  let branch = [off];
+  const ctx = createCommandContext({ sessionManager: { getBranch: () => branch } });
+
+  await events.get("session_start")({ type: "session_start" }, ctx);
+  for (const [type, entries, mode] of [["session_switch", [lite], "lite"], ["session_branch", [off], "off"], ["session_switch", [], "full"]]) {
+    branch = entries;
+    await events.get(type)?.({ type }, ctx);
+    const result = await events.get("before_agent_start")({ systemPrompt: ["BASE"] }, ctx);
+    if (mode === "off") assert.equal(result, undefined, `${type} should restore off`);
+    else assert.ok(result?.systemPrompt.at(-1).startsWith(`PONYTAIL MODE ACTIVE — level: ${mode}\n`), `${type} should restore ${mode}`);
+  }
+  assert.deepEqual(appendedEntries, [], "restoring must not append a new mode entry");
+}));
+
 test("skill alias commands delegate to Pi skill commands", async () => {
   const { commands, sentUserMessages } = createPiHarness();
   const ctx = createCommandContext();