fix: lock main session deletion

src/agents/pi-embedded-runner.ts

@@ -98,6 +98,11 @@ type EmbeddedPiQueueHandle = {
 const log = createSubsystemLogger("agent/embedded");
 
 const ACTIVE_EMBEDDED_RUNS = new Map<string, EmbeddedPiQueueHandle>();
+type EmbeddedRunWaiter = {
+  resolve: (ended: boolean) => void;
+  timer: NodeJS.Timeout;
+};
+const EMBEDDED_RUN_WAITERS = new Map<string, Set<EmbeddedRunWaiter>>();
 
 const OAUTH_FILENAME = "oauth.json";
 const DEFAULT_OAUTH_DIR = path.join(CONFIG_DIR, "credentials");
@@ -247,6 +252,43 @@ export function isEmbeddedPiRunStreaming(sessionId: string): boolean {
   return handle.isStreaming();
 }
 
+export function waitForEmbeddedPiRunEnd(
+  sessionId: string,
+  timeoutMs = 15_000,
+): Promise<boolean> {
+  if (!sessionId || !ACTIVE_EMBEDDED_RUNS.has(sessionId))
+    return Promise.resolve(true);
+  return new Promise((resolve) => {
+    const waiters = EMBEDDED_RUN_WAITERS.get(sessionId) ?? new Set();
+    const waiter: EmbeddedRunWaiter = {
+      resolve,
+      timer: setTimeout(() => {
+        waiters.delete(waiter);
+        if (waiters.size === 0) EMBEDDED_RUN_WAITERS.delete(sessionId);
+        resolve(false);
+      }, Math.max(100, timeoutMs)),
+    };
+    waiters.add(waiter);
+    EMBEDDED_RUN_WAITERS.set(sessionId, waiters);
+    if (!ACTIVE_EMBEDDED_RUNS.has(sessionId)) {
+      waiters.delete(waiter);
+      if (waiters.size === 0) EMBEDDED_RUN_WAITERS.delete(sessionId);
+      clearTimeout(waiter.timer);
+      resolve(true);
+    }
+  });
+}
+
+function notifyEmbeddedRunEnded(sessionId: string) {
+  const waiters = EMBEDDED_RUN_WAITERS.get(sessionId);
+  if (!waiters || waiters.size === 0) return;
+  EMBEDDED_RUN_WAITERS.delete(sessionId);
+  for (const waiter of waiters) {
+    clearTimeout(waiter.timer);
+    waiter.resolve(true);
+  }
+}
+
 export function resolveEmbeddedSessionLane(key: string) {
   return resolveSessionLane(key);
 }
@@ -602,6 +644,7 @@ export async function runEmbeddedPiAgent(params: {
           unsubscribe();
           if (ACTIVE_EMBEDDED_RUNS.get(params.sessionId) === queueHandle) {
             ACTIVE_EMBEDDED_RUNS.delete(params.sessionId);
+            notifyEmbeddedRunEnded(params.sessionId);
           }
           session.dispose();
           params.abortSignal?.removeEventListener?.("abort", onAbort);

src/agents/pi-embedded.ts

@@ -8,6 +8,7 @@ export {
   isEmbeddedPiRunActive,
   isEmbeddedPiRunStreaming,
   queueEmbeddedPiMessage,
+  waitForEmbeddedPiRunEnd,
   resolveEmbeddedSessionLane,
   runEmbeddedPiAgent,
 } from "./pi-embedded-runner.js";