1 files changed, 7 insertions, 4 deletions
diff --git a/src/index.ts b/src/index.ts
index 0133f65..d69940a 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -4,7 +4,7 @@ import { getBrowserSystemPrompt } from "./prompt";
 import { registerSetupCommand } from "./setup";
 import { registerProfileCommand } from "./profile/command";
 import { readPin } from "./profile/store";
-import { type BrowserState, defaultState, persistState, restoreState } from "./state";
+import { type BrowserState, createStatePersister, defaultState, restoreState } from "./state";
 import { registerAllTools } from "./registry";
 import { cleanupTempDirs } from "./util/truncate";
 import { createDaemonTransport } from "./daemon/transport";
@@ -16,6 +16,7 @@ export default function browserHarnessExtension(pi: ExtensionAPI): void {
   const namespace = flagNs ?? `pi-${Math.random().toString(36).slice(2, 10)}`;
 
   let state: BrowserState = defaultState(namespace);
+  const persister = createStatePersister(pi);
   let client: BrowserClient | null = null;
   let toolsRegistered = false;
 
@@ -83,6 +84,7 @@ export default function browserHarnessExtension(pi: ExtensionAPI): void {
 
   pi.on("session_start", async (_event, ctx) => {
     state = restoreState(ctx, state.namespace);
+    persister.markPersisted(state);
     setDebugClicks(asBoolean(pi.getFlag("browser-debug-clicks")) ?? false);
 
     // Re-read the on-disk profile pin every session start: another pi session may have changed it since this client was created.
@@ -112,7 +114,7 @@ export default function browserHarnessExtension(pi: ExtensionAPI): void {
               ? { harnessWindowId: snap.harnessWindowId }
               : {}),
           };
-          try { persistState(pi, state); } catch {}
+          try { persister.persist(state); } catch {}
         },
       });
     } else {
@@ -137,13 +139,14 @@ export default function browserHarnessExtension(pi: ExtensionAPI): void {
       }
       // Do NOT stop the client or null it out: the surviving daemon connection is what eliminates the per-session "Allow Remote Debugging" prompt.
     }
-    persistState(pi, state);
+    persister.persist(state);
     await cleanupTempDirs();
   });
 
+  // The branch already holds the restored snapshot: re-appending it only grows the session file.
   pi.on("session_tree", async (_event, ctx) => {
     state = restoreState(ctx, client?.namespace);
-    persistState(pi, state);
+    persister.markPersisted(state);
   });
 
   pi.on("before_agent_start", async (event) => {