pi-tmux-session-map 0.4.0 → 0.5.0

package/src/index.ts

@@ -4,19 +4,22 @@ import { open, utimes } from "node:fs/promises";
 import { homedir } from "node:os";
 import { join } from "node:path";
 
-import { cleanupStaleSessionMappings, cleanupStaleTempFiles, cleanupSupersededStatusFiles, ensureDirectory, writeAtomic } from "./domain/atomic-write.ts";
+import { cleanupStaleSessionMappings, cleanupStaleTempFiles, cleanupSupersededStatusFiles, ensureDirectory, writeAtomic, writeAtomicExclusive } from "./domain/atomic-write.ts";
 import { createAsyncQueue } from "./domain/async-queue.ts";
 import { createThrottledErrorReporter } from "./domain/error-reporter.ts";
 import { publishSessionMapping } from "./domain/mapping-publisher.ts";
 import { PANE_KEY_FORMAT, defaultStateDir, statusFileName } from "./domain/pane-key.ts";
 import { parsePaneInfoOutput, parseZellijPaneInfo, mexPaneInfo } from "./domain/pane-info.ts";
+import { ownsInteractivePane, paneOwnershipEvidence } from "./domain/pane-ownership.ts";
 import { buildRestoreAck, restoreAckPath, restoreExpectationFromEnvironment, serializeRestoreAck } from "./domain/restore-ack.ts";
 import { createRuntimeState, nextStatusTransition, refreshStatusTransition, type FinalWorkState, type RuntimeState, type WorkState } from "./domain/state.ts";
-import { buildTwsStatusPayload, serializeStatusPayload, type PaneInfo, type SessionInfo } from "./domain/status.ts";
+import { ContinuationExecutor } from "./continuation-executor.ts";
+import { createSignalEvidenceListener, recordedTerminationSignal } from "./domain/signal-evidence.ts";
+import { buildTwsStatusPayload, serializeStatusPayload, type ContinueStatus, type PaneInfo, type RecoveryEvidence, type RunAnchor, type SessionInfo, type ShutdownCause } from "./domain/status.ts";
 
-function execFileAsync(file: string, args: readonly string[]): Promise<{ stdout: string }> {
+function execFileAsync(file: string, args: readonly string[], timeoutMs = muxTimeoutMs): Promise<{ stdout: string }> {
   return new Promise((resolve, reject) => {
-    execFile(file, args, { timeout: muxTimeoutMs }, (error, stdout) => {
+    execFile(file, args, { timeout: timeoutMs }, (error, stdout) => {
       if (error) {
         reject(error);
         return;
@@ -32,6 +35,16 @@ const twsTriggerFile = process.env.PI_TMUX_SESSION_MAP_TWS_TRIGGER_FILE ?? join(
 const tmuxBin = process.env.PI_TMUX_SESSION_MAP_TMUX_BIN ?? "tmux";
 const zellijBin = process.env.PI_TMUX_SESSION_MAP_ZELLIJ_BIN ?? "zellij";
 const muxTimeoutMs = 1000;
+// A restore recreates many panes at once and keeps the tmux server busy, so a
+// 1 s `display-message` can time out exactly when a restored Pi starts. Without
+// pane info the one-shot restore acknowledgement was never written (observed in
+// a 106-session Tart drill on 2026-11-06). While that acknowledgement is
+// pending, retry with growing timeouts; about 31 s in total, below the restore
+// engine's minimum acknowledgement window of 15 s + 2 s per resumed pane.
+const restorePaneInfoTimeoutsMs = (process.env.PI_TMUX_SESSION_MAP_RESTORE_PANE_TIMEOUTS_MS ?? "1000,2000,4000,8000,16000")
+  .split(",")
+  .map((value) => Number.parseInt(value, 10))
+  .filter((value) => Number.isFinite(value) && value > 0);
 const cleanupMaxAgeMs = 24 * 60 * 60 * 1000;
 const reconcileIntervalMs = Math.max(
   10_000,
@@ -47,6 +60,10 @@ type EventAction = {
   cleanup?: boolean;
   restoreReason?: string;
   refreshPane?: boolean;
+  /** Leaf captured synchronously at agent_start, before the prompt is persisted. */
+  runStartLeafId?: string | null;
+  /** session_shutdown reason; the cause is resolved when the payload is built. */
+  shutdownReason?: string;
 };
 
 type ArtifactCache = {
@@ -54,8 +71,27 @@ type ArtifactCache = {
   paneInfoResolved: boolean;
   mappingSignature: string | null;
   statusSignature: string | null;
+  /** session_start reason of a restore acknowledgement not yet written. */
+  pendingRestoreReason: string | null;
+  /** Anchor of the current agent run; null outside a run. */
+  runAnchor: RunAnchor | null;
+  /** In-process status of the restore acknowledgement, once known. */
+  restoreAckStatus: "resumed" | "mismatch" | undefined;
+  /** Auto-continue executor state for the status sidecar. */
+  continueState: () => ContinueStatus | null;
 };
 
+export function shutdownCauseFor(reason: string, signal: string | null): ShutdownCause {
+  switch (reason) {
+    case "quit":
+      return signal === null ? "quit" : "signal";
+    case "reload":
+      return "reload";
+    default:
+      return "replace";
+  }
+}
+
 type AgentEndMessage = {
   role: string;
   stopReason?: string;
@@ -85,13 +121,14 @@ export function provisionalWorkState(finalState: FinalWorkState): WorkState {
   return finalState === "failed" ? "retrying" : finalState;
 }
 
-export async function resolvePaneInfo(): Promise<PaneInfo | undefined> {
+export async function resolvePaneInfo(timeoutMs = muxTimeoutMs): Promise<PaneInfo | undefined> {
   const tmuxPane = process.env.TMUX_PANE;
   if (process.env.TMUX && tmuxPane) {
     try {
       const { stdout } = await execFileAsync(
         tmuxBin,
         ["display-message", "-p", "-t", tmuxPane, PANE_INFO_FORMAT],
+        timeoutMs,
       );
       return parsePaneInfoOutput(stdout, tmuxPane);
     } catch {
@@ -169,8 +206,13 @@ export function twsStatusSignature(
   state: WorkState,
   startedAtMs: number | null,
   finishedAtMs: number | null,
+  sessionId: string,
+  evidence: RecoveryEvidence = { runAnchor: null, shutdownCause: null },
 ): string {
   return JSON.stringify([
+    evidence.runAnchor,
+    evidence.shutdownCause,
+    evidence.continueState ?? null,
     info.backend,
     info.paneKey,
     info.paneId,
@@ -181,6 +223,7 @@ export function twsStatusSignature(
     session.cwd,
     session.sessionFile,
     session.sessionName,
+    sessionId,
     state,
     startedAtMs,
     finishedAtMs,
@@ -191,19 +234,39 @@ async function writeSessionMapping(info: PaneInfo, session: SessionInfo, session
   return (await publishSessionMapping(stateDir, info, session, sessionId)) !== undefined;
 }
 
+/**
+ * Publish the one-shot restore acknowledgement. The first acknowledgement for a
+ * token wins: a later session_start in the same process (for example after
+ * /new or /reload) must not replace the proof of what the restore resumed.
+ */
 async function writeRestoreAcknowledgement(
   info: PaneInfo,
   session: SessionInfo,
   sessionId: string,
   reason: string,
-): Promise<void> {
+): Promise<"resumed" | "mismatch" | undefined> {
   const expectation = restoreExpectationFromEnvironment();
   if (!expectation) {
-    return;
+    return undefined;
   }
   const ack = buildRestoreAck(expectation, info, session, sessionId, reason, Date.now());
   await ensureDirectory(expectation.ackDir);
-  await writeAtomic(restoreAckPath(expectation), serializeRestoreAck(ack), 0o600);
+  await writeAtomicExclusive(restoreAckPath(expectation), serializeRestoreAck(ack), 0o600);
+  return ack.status;
+}
+
+/** Resolve pane info for a pending restore acknowledgement, with growing timeouts. */
+async function resolvePaneInfoForRestore(): Promise<PaneInfo | undefined> {
+  for (const [attempt, timeoutMs] of restorePaneInfoTimeoutsMs.entries()) {
+    const info = await resolvePaneInfo(timeoutMs);
+    if (info) {
+      return info;
+    }
+    if (attempt + 1 < restorePaneInfoTimeoutsMs.length) {
+      await new Promise((resolve) => setTimeout(resolve, 100));
+    }
+  }
+  return undefined;
 }
 
 function getSessionInfo(ctx: ExtensionContext): SessionInfo {
@@ -244,8 +307,11 @@ async function writeTwsStatus(
   updatedAtMs: number,
   startedAtMs: number | null,
   finishedAtMs: number | null,
+  sessionId: string,
+  evidence: RecoveryEvidence,
 ): Promise<void> {
-  const payload = buildTwsStatusPayload(info, session, state, updatedAtMs, startedAtMs, finishedAtMs);
+  const payload = buildTwsStatusPayload(info, session, state, updatedAtMs, startedAtMs, finishedAtMs,
+    { sessionId, piPid: process.pid }, evidence);
 
   const directory = statusDirectory(info);
   const fileName = statusFileName(info.paneKey);
@@ -271,32 +337,48 @@ async function applyEventAction(
   runtimeState: RuntimeState,
   artifactCache: ArtifactCache,
 ): Promise<RuntimeState> {
+  if (!ownsInteractivePane(paneOwnershipEvidence(ctx.mode, process.stdin.isTTY))) {
+    return runtimeState;
+  }
+
   const session = getSessionInfo(ctx);
   const mappingSessionFile = action.mapping ? session.sessionFile : null;
   if (mappingSessionFile === null && action.status === null) {
     return runtimeState;
   }
 
-  if (action.refreshPane === true || !artifactCache.paneInfoResolved) {
-    artifactCache.paneInfo = await resolvePaneInfo();
+  if (action.restoreReason !== undefined && restoreExpectationFromEnvironment() !== undefined) {
+    artifactCache.pendingRestoreReason = action.restoreReason;
+  }
+  const restorePending = artifactCache.pendingRestoreReason !== null;
+  if (action.refreshPane === true || !artifactCache.paneInfoResolved || (restorePending && !artifactCache.paneInfo)) {
+    const resolved = restorePending ? await resolvePaneInfoForRestore() : await resolvePaneInfo();
+    // A refresh that fails (tmux busy past its timeout) keeps the last known pane of this
+    // process. Dropping it skipped this event's status write, e.g. `working` at agent_start.
+    artifactCache.paneInfo = resolved ?? artifactCache.paneInfo;
     artifactCache.paneInfoResolved = true;
   }
   const info = artifactCache.paneInfo;
   if (!info) {
+    // A pending restore acknowledgement stays pending; the next event or the
+    // periodic reconciliation retries it.
     return runtimeState;
   }
 
   const sessionId = ctx.sessionManager.getSessionId();
+  if (artifactCache.pendingRestoreReason !== null && info.backend === "tmux") {
+    // Acknowledge before the mapping: the restore only waits for this file,
+    // and a mapping failure must not suppress it.
+    artifactCache.restoreAckStatus = await writeRestoreAcknowledgement(info, session, sessionId, artifactCache.pendingRestoreReason);
+    artifactCache.pendingRestoreReason = null;
+  }
+
   const mappingSignature = sessionMappingSignature(info, session, sessionId);
   if (mappingSessionFile !== null && mappingSignature !== null && mappingSignature !== artifactCache.mappingSignature) {
     const published = await writeSessionMapping(info, session, sessionId);
     artifactCache.mappingSignature = published ? mappingSignature : null;
   }
 
-  if (action.restoreReason !== undefined && info.backend === "tmux") {
-    await writeRestoreAcknowledgement(info, session, sessionId, action.restoreReason);
-  }
-
   if (action.status === null) {
     return runtimeState;
   }
@@ -312,18 +394,38 @@ async function applyEventAction(
     return transition;
   }
 
+  // A run starts on the first `working` after a non-running state; retries keep
+  // the anchor. Final states and idle end the run.
+  if (transition.state === "working" && runtimeState.lastState !== "working" && runtimeState.lastState !== "retrying") {
+    artifactCache.runAnchor = {
+      leaf_id: action.runStartLeafId ?? null,
+      started_at_ms: transition.startedAtMs ?? now,
+    };
+  } else if (transition.state !== "working" && transition.state !== "retrying" && transition.state !== "shutdown") {
+    artifactCache.runAnchor = null;
+  }
+  const evidence: RecoveryEvidence = {
+    continueState: artifactCache.continueState(),
+    runAnchor: artifactCache.runAnchor,
+    shutdownCause: action.shutdownReason === undefined
+      ? null
+      : shutdownCauseFor(action.shutdownReason, recordedTerminationSignal()),
+  };
+
   const statusSignature = twsStatusSignature(
     info,
     session,
     transition.state,
     transition.startedAtMs,
     transition.finishedAtMs,
+    sessionId,
+    evidence,
   );
   if (statusSignature === artifactCache.statusSignature) {
     return transition;
   }
 
-  await writeTwsStatus(info, session, transition.state, now, transition.startedAtMs, transition.finishedAtMs);
+  await writeTwsStatus(info, session, transition.state, now, transition.startedAtMs, transition.finishedAtMs, sessionId, evidence);
   artifactCache.statusSignature = statusSignature;
   return transition;
 }
@@ -348,7 +450,13 @@ export default function tmuxSessionMap(pi: ExtensionAPI) {
     paneInfoResolved: false,
     mappingSignature: null,
     statusSignature: null,
+    pendingRestoreReason: null,
+    runAnchor: null,
+    restoreAckStatus: undefined,
+    continueState: () => executor?.status() ?? null,
   };
+  const signalEvidence = createSignalEvidenceListener();
+  let executor: ContinuationExecutor | undefined;
 
   const enqueue = async (ctx: ExtensionContext, action: EventAction) => {
     await queue.enqueue(async () => {
@@ -390,7 +498,11 @@ export default function tmuxSessionMap(pi: ExtensionAPI) {
 
   const startReconciliation = (ctx: ExtensionContext) => {
     stopReconciliation();
-    if (!process.env.TMUX || !process.env.TMUX_PANE) {
+    if (
+      !process.env.TMUX
+      || !process.env.TMUX_PANE
+      || !ownsInteractivePane(paneOwnershipEvidence(ctx.mode, process.stdin.isTTY))
+    ) {
       return;
     }
     reconciliationActive = true;
@@ -404,6 +516,8 @@ export default function tmuxSessionMap(pi: ExtensionAPI) {
     artifactCache.paneInfoResolved = false;
     artifactCache.mappingSignature = null;
     artifactCache.statusSignature = null;
+    artifactCache.pendingRestoreReason = null;
+    artifactCache.runAnchor = null;
     runtimeState = createRuntimeState();
     await enqueue(ctx, {
       mapping: true,
@@ -412,14 +526,69 @@ export default function tmuxSessionMap(pi: ExtensionAPI) {
       restoreReason: event.reason,
       refreshPane: true,
     });
+    // Only the visible interactive Pi owns the pane; print, RPC and headless
+    // children never install a signal listener.
+    const ownsPane = ownsInteractivePane(paneOwnershipEvidence(ctx.mode, process.stdin.isTTY));
+    if (ownsPane) {
+      signalEvidence.install();
+    }
+    // Auto-continue executor: only the restored Pi that the restore engine
+    // planned (PI_TMUX_CONTINUE_PLAN), once, at its startup.
+    if (ownsPane && event.reason === "startup" && executor === undefined && process.env.TMUX) {
+      executor = ContinuationExecutor.fromEnvironment(
+        {
+          ackResumed: () => (artifactCache.restoreAckStatus === undefined ? undefined : artifactCache.restoreAckStatus === "resumed"),
+          paneKey: () => artifactCache.paneInfo?.paneKey ?? null,
+          sessionId: () => ctx.sessionManager.getSessionId(),
+          sessionFile: () => ctx.sessionManager.getSessionFile() ?? null,
+          branch: () => ctx.sessionManager.getBranch() as never,
+          idle: () => ctx.isIdle(),
+          pendingMessages: () => ctx.hasPendingMessages(),
+          editorText: () => ctx.ui.getEditorText() ?? "",
+        },
+        (message, options) => pi.sendMessage(message, options),
+        () => {
+          void enqueue(ctx, { mapping: false, status: "refresh" });
+        },
+      );
+      if (executor) {
+        // Any keystroke before the dispatch vetoes it (terminal replies do not
+        // reach this listener inside tmux; measured in tmux-recovery phase 0).
+        ctx.ui.onTerminalInput(() => {
+          executor?.noteUserActivity();
+          return undefined;
+        });
+        executor.start();
+      }
+    }
     startReconciliation(ctx);
   });
+  pi.on("input", (event) => {
+    if (event.source !== "extension") {
+      executor?.noteUserActivity();
+    }
+    return undefined;
+  });
+  // Pi 1.5 emits ui_prompt_start/end around extension dialogs; older Pi never does.
+  const onUiPrompt = pi.on as unknown as (event: string, handler: () => void) => void;
+  onUiPrompt("ui_prompt_start", () => {
+    executor?.notePrompt(true);
+  });
+  onUiPrompt("ui_prompt_end", () => {
+    executor?.notePrompt(false);
+  });
+  pi.on("before_agent_start", (event) => {
+    executor?.beforeAgentStart(event as never);
+    return undefined;
+  });
   pi.on("session_info_changed", async (_event, ctx) => {
     await enqueue(ctx, { mapping: true, status: "refresh", refreshPane: true });
   });
   pi.on("agent_start", async (_event, ctx) => {
     pendingFinalState = "done";
-    await enqueue(ctx, { mapping: true, status: "working", refreshPane: true });
+    // Read the leaf now: the run's prompt entry is persisted after agent_start.
+    const runStartLeafId = ctx.sessionManager.getLeafId() ?? null;
+    await enqueue(ctx, { mapping: true, status: "working", refreshPane: true, runStartLeafId });
   });
   // agent_end may be followed by automatic retry/compaction. Keep failures
   // provisional until agent_settled confirms that no continuation will run.
@@ -435,15 +604,24 @@ export default function tmuxSessionMap(pi: ExtensionAPI) {
     });
   });
   pi.on("agent_settled", async (_event, ctx) => {
+    executor?.agentSettled(pendingFinalState);
     await enqueue(ctx, { mapping: true, status: pendingFinalState });
   });
-  pi.on("session_shutdown", async (_event, ctx) => {
+  pi.on("session_shutdown", async (event, ctx) => {
     stopReconciliation();
-    await enqueue(ctx, { mapping: false, status: "shutdown" });
+    executor?.shutdown();
+    if (event.reason !== "quit") {
+      // A replacement instance installs its own listener; a quit keeps it so a
+      // signal that arrives during the shutdown is still recorded.
+      signalEvidence.uninstall();
+    }
+    await enqueue(ctx, { mapping: false, status: "shutdown", shutdownReason: event.reason });
     artifactCache.paneInfo = undefined;
     artifactCache.paneInfoResolved = false;
     artifactCache.mappingSignature = null;
     artifactCache.statusSignature = null;
+    artifactCache.pendingRestoreReason = null;
+    artifactCache.runAnchor = null;
     runtimeState = createRuntimeState();
   });
 }