refactor: route boot session mapping through accessor (#96225)

jalehman · web-flow · commit acc2a0ee7297 · 2026-06-24T06:54:19.000-07:00
diff --git a/scripts/check-session-accessor-boundary.mjs b/scripts/check-session-accessor-boundary.mjs
@@ -107,6 +107,7 @@ export const migratedSessionAccessorFiles = new Set([
   "src/gateway/sessions-history-http.ts",
   "src/gateway/session-utils.ts",
   "src/gateway/managed-image-attachments.ts",
+  "src/gateway/boot.ts",
   "src/gateway/server-methods/artifacts.ts",
   "src/gateway/server-methods/chat.ts",
   "src/gateway/sessions-resolve.ts",
@@ -163,6 +164,7 @@ export const migratedSessionAccessorWriteFiles = new Set([
   "src/auto-reply/reply/session-usage.ts",
   "src/commands/tasks.ts",
   "src/config/sessions/cleanup-service.ts",
+  "src/gateway/boot.ts",
   "src/gateway/server-node-events.ts",
   "src/gateway/session-compaction-checkpoints.ts",
   "src/plugins/host-hook-cleanup.ts",
diff --git a/src/config/sessions/session-accessor.ts b/src/config/sessions/session-accessor.ts
@@ -541,6 +541,46 @@ export type RestoreSessionFromCompactionCheckpointParams = {
   storePath: string;
 };
 
+export type TemporarySessionMappingPreservationResult<T> = {
+  /** Result returned by the operation while the temporary mapping may exist. */
+  result: T;
+  /** Snapshot failure; callers may continue when temporary cleanup is best-effort. */
+  snapshotFailure?: string;
+  /** Restore/delete failure for the original temporary mapping state. */
+  restoreFailure?: string;
+};
+
+type TemporarySessionMappingSnapshot =
+  | {
+      canRestore: false;
+      sessionKey: string;
+      snapshotFailure: string;
+      storePath: string;
+    }
+  | {
+      canRestore: true;
+      hadEntry: false;
+      sessionKey: string;
+      storePath: string;
+    }
+  | {
+      canRestore: true;
+      entry: SessionEntry;
+      hadEntry: true;
+      sessionKey: string;
+      storePath: string;
+    };
+
+type TemporarySessionMappingOperationResult<T> =
+  | {
+      ok: true;
+      result: T;
+    }
+  | {
+      error: unknown;
+      ok: false;
+    };
+
 export type SessionEntryCreateWithTranscriptContext = {
   /** Current entry under the requested key before creation, if any. */
   existingEntry?: SessionEntry;
@@ -1393,6 +1433,37 @@ export async function applyRestartRecoveryLifecycle<T>(params: {
   return writerResult.result;
 }
 
+/**
+ * Runs an operation while preserving one temporary session mapping.
+ * The storage backend snapshots exactly the named key before the operation and
+ * restores that entry, or deletes it when it did not previously exist, after
+ * the operation finishes. SQLite backends can implement the same named
+ * preservation lifecycle without exposing mutable store access to callers.
+ */
+export async function preserveTemporarySessionMapping<T>(
+  scope: SessionAccessScope,
+  operation: () => Promise<T> | T,
+): Promise<TemporarySessionMappingPreservationResult<T>> {
+  const snapshot = snapshotTemporarySessionMapping(scope);
+  let operationResult: TemporarySessionMappingOperationResult<T>;
+  try {
+    operationResult = { ok: true, result: await operation() };
+  } catch (err) {
+    operationResult = { error: err, ok: false };
+  }
+
+  const restoreFailure = await restoreTemporarySessionMapping(snapshot);
+  if (!operationResult.ok) {
+    throw operationResult.error;
+  }
+
+  return {
+    result: operationResult.result,
+    ...(snapshot.canRestore ? {} : { snapshotFailure: snapshot.snapshotFailure }),
+    ...(restoreFailure ? { restoreFailure } : {}),
+  };
+}
+
 /** Removes entries and orphan transcript artifacts owned by a named session lifecycle. */
 export async function cleanupSessionLifecycleArtifacts(
   params: SessionLifecycleArtifactCleanupParams,
@@ -2515,6 +2586,53 @@ function createFallbackSessionEntry(patch: Partial<SessionEntry>): SessionEntry
   };
 }
 
+function snapshotTemporarySessionMapping(
+  scope: SessionAccessScope,
+): TemporarySessionMappingSnapshot {
+  const storePath = resolveAccessStorePath(scope);
+  try {
+    const store = loadSessionStore(storePath, { skipCache: true });
+    const entry = store[scope.sessionKey];
+    return {
+      canRestore: true,
+      ...(entry ? { entry: structuredClone(entry), hadEntry: true } : { hadEntry: false }),
+      sessionKey: scope.sessionKey,
+      storePath,
+    };
+  } catch (err) {
+    return {
+      canRestore: false,
+      sessionKey: scope.sessionKey,
+      snapshotFailure: formatErrorMessage(err),
+      storePath,
+    };
+  }
+}
+
+async function restoreTemporarySessionMapping(
+  snapshot: TemporarySessionMappingSnapshot,
+): Promise<string | undefined> {
+  if (!snapshot.canRestore) {
+    return undefined;
+  }
+  try {
+    await updateSessionStore(
+      snapshot.storePath,
+      (store) => {
+        if (snapshot.hadEntry) {
+          store[snapshot.sessionKey] = structuredClone(snapshot.entry);
+          return;
+        }
+        delete store[snapshot.sessionKey];
+      },
+      { activeSessionKey: snapshot.sessionKey },
+    );
+    return undefined;
+  } catch (err) {
+    return formatErrorMessage(err);
+  }
+}
+
 function cleanupPreviousResetTranscripts(params: {
   agentId: string;
   previousEntry: SessionEntry;
diff --git a/src/gateway/boot.ts b/src/gateway/boot.ts
@@ -18,8 +18,7 @@ import {
   resolveMainSessionKey,
 } from "../config/sessions/main-session.js";
 import { resolveStorePath } from "../config/sessions/paths.js";
-import { loadSessionStore, updateSessionStore } from "../config/sessions/store.js";
-import type { SessionEntry } from "../config/sessions/types.js";
+import { preserveTemporarySessionMapping } from "../config/sessions/session-accessor.js";
 import type { OpenClawConfig } from "../config/types.openclaw.js";
 import { formatErrorMessage } from "../infra/errors.js";
 import { createSubsystemLogger } from "../logging/subsystem.js";
@@ -33,14 +32,6 @@ function generateBootSessionId(): string {
   return `boot-${ts}-${suffix}`;
 }
 
-type SessionMappingSnapshot = {
-  storePath: string;
-  sessionKey: string;
-  canRestore: boolean;
-  hadEntry: boolean;
-  entry?: SessionEntry;
-};
-
 const log = createSubsystemLogger("gateway/boot");
 const BOOT_FILENAME = "BOOT.md";
 
@@ -101,68 +92,6 @@ async function loadBootFile(
   }
 }
 
-function snapshotSessionMapping(params: {
-  cfg: OpenClawConfig;
-  sessionKey: string;
-}): SessionMappingSnapshot {
-  const agentId = resolveAgentIdFromSessionKey(params.sessionKey);
-  const storePath = resolveStorePath(params.cfg.session?.store, { agentId });
-  try {
-    const store = loadSessionStore(storePath, { skipCache: true });
-    const entry = store[params.sessionKey];
-    if (!entry) {
-      return {
-        storePath,
-        sessionKey: params.sessionKey,
-        canRestore: true,
-        hadEntry: false,
-      };
-    }
-    return {
-      storePath,
-      sessionKey: params.sessionKey,
-      canRestore: true,
-      hadEntry: true,
-      entry: structuredClone(entry),
-    };
-  } catch (err) {
-    log.debug("boot: could not snapshot session mapping", {
-      sessionKey: params.sessionKey,
-      error: String(err),
-    });
-    return {
-      storePath,
-      sessionKey: params.sessionKey,
-      canRestore: false,
-      hadEntry: false,
-    };
-  }
-}
-
-async function restoreSessionMapping(
-  snapshot: SessionMappingSnapshot,
-): Promise<string | undefined> {
-  if (!snapshot.canRestore) {
-    return undefined;
-  }
-  try {
-    await updateSessionStore(
-      snapshot.storePath,
-      (store) => {
-        if (snapshot.hadEntry && snapshot.entry) {
-          store[snapshot.sessionKey] = snapshot.entry;
-          return;
-        }
-        delete store[snapshot.sessionKey];
-      },
-      { activeSessionKey: snapshot.sessionKey },
-    );
-    return undefined;
-  } catch (err) {
-    return formatErrorMessage(err);
-  }
-}
-
 export async function runBootOnce(params: {
   cfg: OpenClawConfig;
   deps: CliDeps;
@@ -193,39 +122,49 @@ export async function runBootOnce(params: {
   const sessionKey = resolveBootSessionKey(mainSessionKey);
   const message = buildBootPrompt(result.content ?? "");
   const sessionId = generateBootSessionId();
-  const mappingSnapshot = snapshotSessionMapping({
-    cfg: params.cfg,
-    sessionKey,
-  });
+  const agentId = resolveAgentIdFromSessionKey(sessionKey);
+  const storePath = resolveStorePath(params.cfg.session?.store, { agentId });
 
-  // Register the boot prompt for the message-tool echo guard so the
-  // tool layer can drop fallback-model echoes that copy substantial
-  // BOOT.md content without preserving the wrapper markers above.
-  // Always cleared in finally so a failed run does not leave a stale
-  // entry that mis-fires on an unrelated subsequent run reusing the
-  // same session key. Refs #53732.
-  setBootEchoContextForSession(sessionKey, message);
-  let agentFailure: string | undefined;
-  try {
-    await agentCommand(
-      {
-        message,
-        sessionKey,
-        sessionId,
-        deliver: false,
-        suppressPromptPersistence: true,
-      },
-      bootRuntime,
-      params.deps,
-    );
-  } catch (err) {
-    agentFailure = formatErrorMessage(err);
-    log.error(`boot: agent run failed: ${agentFailure}`);
-  } finally {
-    clearBootEchoContextForSession(sessionKey);
+  const mappingPreservation = await preserveTemporarySessionMapping(
+    { storePath, sessionKey },
+    async () => {
+      // Register the boot prompt for the message-tool echo guard so the
+      // tool layer can drop fallback-model echoes that copy substantial
+      // BOOT.md content without preserving the wrapper markers above.
+      // Always cleared in finally so a failed run does not leave a stale
+      // entry that mis-fires on an unrelated subsequent run reusing the
+      // same session key. Refs #53732.
+      setBootEchoContextForSession(sessionKey, message);
+      try {
+        await agentCommand(
+          {
+            message,
+            sessionKey,
+            sessionId,
+            deliver: false,
+            suppressPromptPersistence: true,
+          },
+          bootRuntime,
+          params.deps,
+        );
+        return undefined;
+      } catch (err) {
+        const failure = formatErrorMessage(err);
+        log.error(`boot: agent run failed: ${failure}`);
+        return failure;
+      } finally {
+        clearBootEchoContextForSession(sessionKey);
+      }
+    },
+  );
+  const agentFailure = mappingPreservation.result;
+  if (mappingPreservation.snapshotFailure) {
+    log.debug("boot: could not snapshot session mapping", {
+      sessionKey,
+      error: mappingPreservation.snapshotFailure,
+    });
   }
-
-  const mappingRestoreFailure = await restoreSessionMapping(mappingSnapshot);
+  const mappingRestoreFailure = mappingPreservation.restoreFailure;
   if (mappingRestoreFailure) {
     log.error(`boot: failed to restore session mapping: ${mappingRestoreFailure}`);
   }
diff --git a/test/scripts/check-session-accessor-boundary.test.ts b/test/scripts/check-session-accessor-boundary.test.ts
@@ -57,6 +57,7 @@ describe("session accessor boundary guard", () => {
         "src/gateway/sessions-history-http.ts",
         "src/gateway/session-utils.ts",
         "src/gateway/managed-image-attachments.ts",
+        "src/gateway/boot.ts",
         "src/gateway/server-methods/artifacts.ts",
         "src/gateway/server-methods/chat.ts",
         "src/gateway/sessions-resolve.ts",
@@ -121,6 +122,7 @@ describe("session accessor boundary guard", () => {
         "src/auto-reply/reply/session-usage.ts",
         "src/commands/tasks.ts",
         "src/config/sessions/cleanup-service.ts",
+        "src/gateway/boot.ts",
         "src/gateway/server-node-events.ts",
         "src/gateway/session-compaction-checkpoints.ts",
         "src/plugins/host-hook-cleanup.ts",
