Fix session cleanup and release 0.15.47

2 files changed, 201 insertions, 1 deletions

diff --git a/test-suite/unit/automation-target.test.mjs b/test-suite/unit/automation-target.test.mjs
index 8a6ea2b..f838fa6 100644
--- a/test-suite/unit/automation-target.test.mjs
+++ b/test-suite/unit/automation-target.test.mjs
@@ -81,7 +81,12 @@ function makeChrome(state, { withWindows = true, withStorage = true, withTabGrou
         return { ...tab };
       },
       update: async (id, props = {}) => { const t = tabs.get(id); if (!t) throw new Error(`No tab with id ${id}`); Object.assign(t, props); return { ...t }; },
-      remove: async (id) => { tabs.delete(id); },
+      remove: async (id) => {
+        const tab = tabs.get(id);
+        tabs.delete(id);
+        // Chrome closes a window automatically when its final tab is removed.
+        if (tab && ![...tabs.values()].some((other) => other.windowId === tab.windowId)) windows.delete(tab.windowId);
+      },
       group: async ({ groupId, tabIds = [] } = {}) => {
         let gid = groupId;
         if (typeof gid !== "number") {
@@ -388,6 +393,108 @@ async function run() {
     ok(stale.closedWindowId === null && stale.closedTabId === null, "cleanup: robust when owned window was already closed");
   }
 
+  // Cleanup must never remove a window wholesale, even if tabs move during removal.
+  for (const moveOwnedTab of [false, true]) {
+    const state = makeChromeState();
+    const chrome = makeChrome(state);
+    chrome.windows.remove = async () => { throw new Error("whole-window removal is forbidden"); };
+    const w = loadWorker(chrome);
+    const owned = await w.getOrCreateAutomationTarget(SK);
+    const remove = chrome.tabs.remove;
+    chrome.tabs.remove = async (id) => {
+      // User adds a tab just as cleanup starts; a pre-removal window check is not enough.
+      state.userArticle.windowId = owned.windowId;
+      if (moveOwnedTab) state.tabs.get(owned.id).windowId = state.userWindowId;
+      await remove(id);
+    };
+    const result = await w.dispatch("automation.cleanup", { sessionKey: SK });
+    ok(result.closedTabId === owned.id && result.closedWindowId === null, "mixed window: only owned tab reported closed");
+    ok(state.tabs.has(state.userArticle.id) && state.windows.has(owned.windowId), "mixed window: user tab and its window survive");
+    ok(!state.tabs.has(owned.id), "mixed window: owned tab removed even after moving to another window");
+  }
+
+  // A populated automation window is not disposable merely because Pi created it.
+  {
+    const state = makeChromeState();
+    const w = loadWorker(makeChrome(state));
+    const owned = await w.getOrCreateAutomationTarget(SK);
+    state.userArticle.windowId = owned.windowId;
+    await w.dispatch("automation.cleanup", { sessionKey: SK });
+    ok(state.tabs.has(state.userArticle.id) && state.windows.has(owned.windowId), "populated window: user's moved-in tab survives cleanup");
+  }
+
+  // Created vs adopted ownership survives restart; matching titles do not grant ownership.
+  for (const restart of [false, true]) {
+    const state = makeChromeState();
+    const chrome = makeChrome(state, { withTabGroups: true });
+    let w = loadWorker(chrome);
+    const groupTitle = "Pi Session: shared-name";
+    await w.dispatch("page.navigate", {
+      sessionKey: SK, targetId: String(state.userGmail.id), url: "https://mail.google.com/",
+      waitUntilLoad: false, joinSessionGroup: true, sessionGroupTitle: groupTitle,
+    });
+    await w.dispatch("tab.group", { sessionKey: SK, targetId: String(state.userArticle.id), groupTitle });
+    const created = await w.dispatch("tab.new", { sessionKey: SK, groupTitle });
+    const other = await w.dispatch("tab.new", { sessionKey: "session:other", groupTitle });
+    // User changes the article's group after Pi adopted it. Cleanup must respect that change.
+    const replacementGroup = state.alloc.group();
+    state.userArticle.groupId = replacementGroup;
+    if (restart) w = loadWorker(makeChrome(state, { withTabGroups: true }));
+    const result = await w.dispatch("automation.cleanup", { sessionKey: SK });
+    ok(!state.tabs.has(created.tab.id) && state.tabs.has(other.tab.id), "resources: close only this session's created tabs");
+    ok(state.tabs.has(state.userGmail.id) && state.userGmail.groupId === -1, "resources: adopted user tab ungrouped, not closed");
+    ok(state.userArticle.groupId === replacementGroup, "resources: user's replacement group untouched");
+    ok(result.closedCreatedTabs === 1 && result.ungroupedAdoptedTabs === 1, "resources: counts reflect successful operations");
+    ok(!(SK in state.storage.piChromeSessionTabs), "resources: completed ownership removed from persistence");
+    const again = await w.dispatch("automation.cleanup", { sessionKey: SK });
+    ok(again.closedCreatedTabs === 0 && again.ungroupedAdoptedTabs === 0, "resources: repeated cleanup is idempotent");
+  }
+
+  // A failed close keeps ownership for retry, without claiming success.
+  {
+    const state = makeChromeState();
+    const chrome = makeChrome(state, { withTabGroups: true });
+    const w = loadWorker(chrome);
+    const owned = await w.getOrCreateAutomationTarget(SK);
+    const opened = await w.dispatch("tab.new", { sessionKey: SK });
+    const remove = chrome.tabs.remove;
+    chrome.tabs.remove = async () => { throw new Error("temporary close failure"); };
+    const failed = await w.dispatch("automation.cleanup", { sessionKey: SK });
+    ok(failed.closedCreatedTabs === 0 && failed.closedTabId === null, "retry: failed closes not reported as success");
+    chrome.tabs.remove = remove;
+    const restarted = loadWorker(chrome);
+    await restarted.dispatch("automation.cleanup", { sessionKey: SK });
+    ok(!state.tabs.has(owned.id) && !state.tabs.has(opened.tab.id), "retry: persisted ownership allows retry after restart");
+  }
+
+  // Failed ungrouping also remains retryable; a full browser restart abandons ownership.
+  {
+    const state = makeChromeState();
+    const chrome = makeChrome(state, { withTabGroups: true });
+    const w = loadWorker(chrome);
+    await w.dispatch("tab.group", { sessionKey: SK, targetId: String(state.userGmail.id) });
+    const ungroup = chrome.tabs.ungroup;
+    chrome.tabs.ungroup = async () => { throw new Error("temporary group failure"); };
+    const failed = await w.dispatch("automation.cleanup", { sessionKey: SK });
+    ok(failed.ungroupedAdoptedTabs === 0 && state.userGmail.groupId >= 0, "ungroup retry: failure leaves user tab and ownership intact");
+    chrome.tabs.ungroup = ungroup;
+    await loadWorker(chrome).dispatch("automation.cleanup", { sessionKey: SK });
+    ok(state.userGmail.groupId === -1, "ungroup retry: successful after worker restart");
+    const created = await w.dispatch("tab.new", { sessionKey: SK });
+    for (const key of Object.keys(state.storage)) delete state.storage[key];
+    await loadWorker(chrome).dispatch("automation.cleanup", { sessionKey: SK });
+    ok(state.tabs.has(created.tab.id), "browser restart: cleared storage never reclaims restored tabs by group name");
+  }
+
+  // Runtime tracking still works when storage.session is unavailable.
+  {
+    const state = makeChromeState();
+    const w = loadWorker(makeChrome(state, { withTabGroups: true, withStorage: false }));
+    const opened = await w.dispatch("tab.new", { sessionKey: SK });
+    await w.dispatch("automation.cleanup", { sessionKey: SK });
+    ok(!state.tabs.has(opened.tab.id) && state.tabs.has(state.userGmail.id), "no storage: created tab cleaned up safely");
+  }
+
   console.log(`\n${passes} passed, ${failures} failed`);
   if (failures) process.exit(1);
 }
diff --git a/test-suite/unit/session-cleanup.test.mjs b/test-suite/unit/session-cleanup.test.mjs
new file mode 100644
index 0000000..7e3285f
--- /dev/null
+++ b/test-suite/unit/session-cleanup.test.mjs
@@ -0,0 +1,93 @@
+// Exercise the shipped cleanup helper and shutdown listener, without opening a bridge or Chrome.
+import assert from "node:assert/strict";
+import fs from "node:fs";
+import vm from "node:vm";
+import { stripTypeScriptTypes } from "node:module";
+
+const source = fs.readFileSync(new URL("../../extensions/chrome-profile-bridge/index.ts", import.meta.url), "utf8");
+function section(start, end) {
+  const from = source.indexOf(start);
+  const to = source.indexOf(end, from);
+  assert.ok(from >= 0 && to > from, `Missing source section: ${start}`);
+  return source.slice(from, to);
+}
+const helper = stripTypeScriptTypes(section("const cleanupAutomationTargetBestEffort =", "const lockChromeControl ="));
+const listener = section('pi.on("session_shutdown",', 'pi.on("before_agent_start",');
+
+function harness(send, sessionKey = "session:test") {
+  const timers = new Map();
+  const events = [];
+  let shutdown;
+  const token = Symbol();
+  const globalState = { loaded: { token } };
+  const context = {
+    AbortController, Promise,
+    setTimeout(fn, ms) { const key = Symbol(); timers.set(key, { fn, ms }); return key; },
+    clearTimeout(key) { timers.delete(key); },
+    bridge: { send(...args) { events.push("send"); return send(...args); }, stop() { events.push("stop"); } },
+    sessionCtx: {}, sessionKeyFor: () => sessionKey ?? undefined,
+    clearAuthExpiryTimer() {}, clearCountdownInterval() {},
+    globalState, PI_CHROME_GLOBAL_KEY: "loaded", instanceToken: token,
+    pi: { on(name, handler) { assert.equal(name, "session_shutdown"); shutdown = handler; } },
+  };
+  vm.runInNewContext(`${helper}\n${listener}`, context);
+  return { shutdown, timers, events, globalState };
+}
+
+// Delivery must settle before bridge.stop; reload does not issue cleanup at all.
+{
+  let complete;
+  let request;
+  const h = harness((...args) => {
+    request = args;
+    return new Promise((resolve) => { complete = resolve; });
+  });
+  const done = h.shutdown({ reason: "quit" });
+  assert.deepEqual(h.events, ["send"]);
+  assert.equal(request[0], "automation.cleanup");
+  assert.equal(request[1].sessionKey, "session:test");
+  assert.equal(request[2], 2000);
+  assert.equal(request[3].aborted, false);
+  complete();
+  await done;
+  assert.deepEqual(h.events, ["send", "stop"]);
+  assert.equal(h.timers.size, 0);
+  assert.equal(h.globalState.loaded, undefined);
+}
+{
+  const h = harness(() => { throw new Error("reload must not send"); });
+  await h.shutdown({ reason: "reload" });
+  assert.deepEqual(h.events, ["stop"]);
+  assert.equal(h.timers.size, 0);
+}
+
+// Even an owner request that never settles cannot hold up shutdown past its budget.
+{
+  let signal;
+  const h = harness((_action, _params, _timeout, s) => { signal = s; return new Promise(() => {}); });
+  const done = h.shutdown({});
+  const [timer] = h.timers.values();
+  assert.equal(timer.ms, 2000);
+  timer.fn();
+  await done;
+  assert.equal(signal.aborted, true);
+  assert.deepEqual(h.events, ["send", "stop"]);
+  assert.equal(h.timers.size, 0);
+}
+for (const reason of ["quit", "new", "resume", "fork"]) {
+  const h = harness(() => Promise.resolve({}));
+  await h.shutdown({ reason });
+  assert.deepEqual(h.events, ["send", "stop"]);
+}
+for (const send of [() => Promise.reject(new Error("offline")), () => { throw new Error("offline"); }]) {
+  const h = harness(send);
+  await h.shutdown({ reason: "exit" });
+  assert.deepEqual(h.events, ["send", "stop"]);
+  assert.equal(h.timers.size, 0);
+}
+{
+  const h = harness(() => { throw new Error("unscoped cleanup must not send"); }, null);
+  await h.shutdown({ reason: "exit" });
+  assert.deepEqual(h.events, ["stop"]);
+}
+console.log("session-cleanup: shutdown ordering, reload, deadline, failures, and missing identity passed");