refactor: harden session store updates

Co-authored-by: Tyler Yust <tyler6204@users.noreply.github.com>
2026-05-08 23:28:27 +00:00 · 2026-01-15 23:06:42 +00:00
parent 35492f8513
commit 688a0ce439
24 changed files with 441 additions and 215 deletions
--- a/src/gateway/server-bridge-events.ts
+++ b/src/gateway/server-bridge-events.ts
@@ -2,7 +2,7 @@ import { randomUUID } from "node:crypto";
 import { normalizeChannelId } from "../channels/plugins/index.js";
 import { agentCommand } from "../commands/agent.js";
 import { loadConfig } from "../config/config.js";
-import { saveSessionStore } from "../config/sessions.js";
+import { updateSessionStore } from "../config/sessions.js";
 import { normalizeMainKey } from "../routing/session-key.js";
 import { defaultRuntime } from "../runtime.js";
 import type { BridgeEvent, BridgeHandlersContext } from "./server-bridge-types.js";
@@ -32,22 +32,23 @@ export const handleBridgeEvent = async (
      const cfg = loadConfig();
      const rawMainKey = normalizeMainKey(cfg.session?.mainKey);
      const sessionKey = sessionKeyRaw.length > 0 ? sessionKeyRaw : rawMainKey;
-      const { storePath, store, entry, canonicalKey } = loadSessionEntry(sessionKey);
+      const { storePath, entry, canonicalKey } = loadSessionEntry(sessionKey);
      const now = Date.now();
      const sessionId = entry?.sessionId ?? randomUUID();
-      store[canonicalKey] = {
-        sessionId,
-        updatedAt: now,
-        thinkingLevel: entry?.thinkingLevel,
-        verboseLevel: entry?.verboseLevel,
-        reasoningLevel: entry?.reasoningLevel,
-        systemSent: entry?.systemSent,
-        sendPolicy: entry?.sendPolicy,
-        lastChannel: entry?.lastChannel,
-        lastTo: entry?.lastTo,
-      };
      if (storePath) {
-        await saveSessionStore(storePath, store);
+        await updateSessionStore(storePath, (store) => {
+          store[canonicalKey] = {
+            sessionId,
+            updatedAt: now,
+            thinkingLevel: entry?.thinkingLevel,
+            verboseLevel: entry?.verboseLevel,
+            reasoningLevel: entry?.reasoningLevel,
+            systemSent: entry?.systemSent,
+            sendPolicy: entry?.sendPolicy,
+            lastChannel: entry?.lastChannel,
+            lastTo: entry?.lastTo,
+          };
+        });
      }

      // Ensure chat UI clients refresh when this run completes (even though it wasn't started via chat.send).
@@ -102,22 +103,23 @@ export const handleBridgeEvent = async (

      const sessionKeyRaw = (link?.sessionKey ?? "").trim();
      const sessionKey = sessionKeyRaw.length > 0 ? sessionKeyRaw : `node-${nodeId}`;
-      const { storePath, store, entry, canonicalKey } = loadSessionEntry(sessionKey);
+      const { storePath, entry, canonicalKey } = loadSessionEntry(sessionKey);
      const now = Date.now();
      const sessionId = entry?.sessionId ?? randomUUID();
-      store[canonicalKey] = {
-        sessionId,
-        updatedAt: now,
-        thinkingLevel: entry?.thinkingLevel,
-        verboseLevel: entry?.verboseLevel,
-        reasoningLevel: entry?.reasoningLevel,
-        systemSent: entry?.systemSent,
-        sendPolicy: entry?.sendPolicy,
-        lastChannel: entry?.lastChannel,
-        lastTo: entry?.lastTo,
-      };
      if (storePath) {
-        await saveSessionStore(storePath, store);
+        await updateSessionStore(storePath, (store) => {
+          store[canonicalKey] = {
+            sessionId,
+            updatedAt: now,
+            thinkingLevel: entry?.thinkingLevel,
+            verboseLevel: entry?.verboseLevel,
+            reasoningLevel: entry?.reasoningLevel,
+            systemSent: entry?.systemSent,
+            sendPolicy: entry?.sendPolicy,
+            lastChannel: entry?.lastChannel,
+            lastTo: entry?.lastTo,
+          };
+        });
      }

      void agentCommand(
--- a/src/gateway/server-bridge-methods-sessions.ts
+++ b/src/gateway/server-bridge-methods-sessions.ts
@@ -8,10 +8,9 @@ import {
 } from "../agents/pi-embedded.js";
 import { loadConfig } from "../config/config.js";
 import {
-  loadSessionStore,
  resolveMainSessionKeyFromConfig,
  type SessionEntry,
-  saveSessionStore,
+  updateSessionStore,
 } from "../config/sessions.js";
 import { clearCommandLane } from "../process/command-queue.js";
 import {
@@ -126,19 +125,20 @@ export const handleSessionsBridgeMethods: BridgeMethodHandler = async (
      const cfg = loadConfig();
      const target = resolveGatewaySessionStoreTarget({ cfg, key });
      const storePath = target.storePath;
-      const store = loadSessionStore(storePath);
-      const primaryKey = target.storeKeys[0] ?? key;
-      const existingKey = target.storeKeys.find((candidate) => store[candidate]);
-      if (existingKey && existingKey !== primaryKey && !store[primaryKey]) {
-        store[primaryKey] = store[existingKey];
-        delete store[existingKey];
-      }
-      const applied = await applySessionsPatchToStore({
-        cfg,
-        store,
-        storeKey: primaryKey,
-        patch: p,
-        loadGatewayModelCatalog: ctx.loadGatewayModelCatalog,
+      const applied = await updateSessionStore(storePath, async (store) => {
+        const primaryKey = target.storeKeys[0] ?? key;
+        const existingKey = target.storeKeys.find((candidate) => store[candidate]);
+        if (existingKey && existingKey !== primaryKey && !store[primaryKey]) {
+          store[primaryKey] = store[existingKey];
+          delete store[existingKey];
+        }
+        return await applySessionsPatchToStore({
+          cfg,
+          store,
+          storeKey: primaryKey,
+          patch: p,
+          loadGatewayModelCatalog: ctx.loadGatewayModelCatalog,
+        });
      });
      if (!applied.ok) {
        return {
@@ -150,7 +150,6 @@ export const handleSessionsBridgeMethods: BridgeMethodHandler = async (
          },
        };
      }
-      await saveSessionStore(storePath, store);
      const payload: SessionsPatchResult = {
        ok: true,
        path: storePath,
@@ -182,32 +181,43 @@ export const handleSessionsBridgeMethods: BridgeMethodHandler = async (
        };
      }

-      const { storePath, store, entry } = loadSessionEntry(key);
-      const now = Date.now();
-      const next: SessionEntry = {
-        sessionId: randomUUID(),
-        updatedAt: now,
-        systemSent: false,
-        abortedLastRun: false,
-        thinkingLevel: entry?.thinkingLevel,
-        verboseLevel: entry?.verboseLevel,
-        reasoningLevel: entry?.reasoningLevel,
-        model: entry?.model,
-        contextTokens: entry?.contextTokens,
-        sendPolicy: entry?.sendPolicy,
-        label: entry?.label,
-        displayName: entry?.displayName,
-        chatType: entry?.chatType,
-        channel: entry?.channel,
-        subject: entry?.subject,
-        room: entry?.room,
-        space: entry?.space,
-        lastChannel: entry?.lastChannel,
-        lastTo: entry?.lastTo,
-        skillsSnapshot: entry?.skillsSnapshot,
-      };
-      store[key] = next;
-      await saveSessionStore(storePath, store);
+      const cfg = loadConfig();
+      const target = resolveGatewaySessionStoreTarget({ cfg, key });
+      const storePath = target.storePath;
+      const next = await updateSessionStore(storePath, (store) => {
+        const primaryKey = target.storeKeys[0] ?? key;
+        const existingKey = target.storeKeys.find((candidate) => store[candidate]);
+        if (existingKey && existingKey !== primaryKey && !store[primaryKey]) {
+          store[primaryKey] = store[existingKey];
+          delete store[existingKey];
+        }
+        const entry = store[primaryKey];
+        const now = Date.now();
+        const nextEntry: SessionEntry = {
+          sessionId: randomUUID(),
+          updatedAt: now,
+          systemSent: false,
+          abortedLastRun: false,
+          thinkingLevel: entry?.thinkingLevel,
+          verboseLevel: entry?.verboseLevel,
+          reasoningLevel: entry?.reasoningLevel,
+          model: entry?.model,
+          contextTokens: entry?.contextTokens,
+          sendPolicy: entry?.sendPolicy,
+          label: entry?.label,
+          displayName: entry?.displayName,
+          chatType: entry?.chatType,
+          channel: entry?.channel,
+          subject: entry?.subject,
+          room: entry?.room,
+          space: entry?.space,
+          lastChannel: entry?.lastChannel,
+          lastTo: entry?.lastTo,
+          skillsSnapshot: entry?.skillsSnapshot,
+        };
+        store[primaryKey] = nextEntry;
+        return nextEntry;
+      });
      return {
        ok: true,
        payloadJSON: JSON.stringify({ ok: true, key, entry: next }),
@@ -249,9 +259,11 @@ export const handleSessionsBridgeMethods: BridgeMethodHandler = async (

      const deleteTranscript = typeof p.deleteTranscript === "boolean" ? p.deleteTranscript : true;

-      const { storePath, store, entry } = loadSessionEntry(key);
+      const cfg = loadConfig();
+      const target = resolveGatewaySessionStoreTarget({ cfg, key });
+      const storePath = target.storePath;
+      const { entry } = loadSessionEntry(key);
      const sessionId = entry?.sessionId;
-      const existed = Boolean(store[key]);
      clearCommandLane(resolveEmbeddedSessionLane(key));
      if (sessionId && isEmbeddedPiRunActive(sessionId)) {
        abortEmbeddedPiRun(sessionId);
@@ -266,8 +278,19 @@ export const handleSessionsBridgeMethods: BridgeMethodHandler = async (
          };
        }
      }
-      if (existed) delete store[key];
-      await saveSessionStore(storePath, store);
+      const deletion = await updateSessionStore(storePath, (store) => {
+        const primaryKey = target.storeKeys[0] ?? key;
+        const existingKey = target.storeKeys.find((candidate) => store[candidate]);
+        if (existingKey && existingKey !== primaryKey && !store[primaryKey]) {
+          store[primaryKey] = store[existingKey];
+          delete store[existingKey];
+        }
+        const entryToDelete = store[primaryKey];
+        const existed = Boolean(entryToDelete);
+        if (existed) delete store[primaryKey];
+        return { existed, entry: entryToDelete };
+      });
+      const existed = deletion.existed;

      const archived: string[] = [];
      if (deleteTranscript && sessionId) {
@@ -323,7 +346,20 @@ export const handleSessionsBridgeMethods: BridgeMethodHandler = async (
          ? Math.max(1, Math.floor(p.maxLines))
          : 400;

-      const { storePath, store, entry } = loadSessionEntry(key);
+      const cfg = loadConfig();
+      const target = resolveGatewaySessionStoreTarget({ cfg, key });
+      const storePath = target.storePath;
+      // Resolve entry inside the lock, but compact outside to avoid holding it.
+      const compactTarget = await updateSessionStore(storePath, (store) => {
+        const primaryKey = target.storeKeys[0] ?? key;
+        const existingKey = target.storeKeys.find((candidate) => store[candidate]);
+        if (existingKey && existingKey !== primaryKey && !store[primaryKey]) {
+          store[primaryKey] = store[existingKey];
+          delete store[existingKey];
+        }
+        return { entry: store[primaryKey], primaryKey };
+      });
+      const entry = compactTarget.entry;
      const sessionId = entry?.sessionId;
      if (!sessionId) {
        return {
@@ -373,13 +409,14 @@ export const handleSessionsBridgeMethods: BridgeMethodHandler = async (
      fs.writeFileSync(filePath, `${keptLines.join("\n")}\n`, "utf-8");

      // Token counts no longer match; clear so status + UI reflect reality after the next turn.
-      if (store[key]) {
-        delete store[key].inputTokens;
-        delete store[key].outputTokens;
-        delete store[key].totalTokens;
-        store[key].updatedAt = Date.now();
-        await saveSessionStore(storePath, store);
-      }
+      await updateSessionStore(storePath, (store) => {
+        const entryToUpdate = store[compactTarget.primaryKey];
+        if (!entryToUpdate) return;
+        delete entryToUpdate.inputTokens;
+        delete entryToUpdate.outputTokens;
+        delete entryToUpdate.totalTokens;
+        entryToUpdate.updatedAt = Date.now();
+      });

      return {
        ok: true,
--- a/src/gateway/server-methods/agent.ts
+++ b/src/gateway/server-methods/agent.ts
@@ -6,7 +6,7 @@ import {
  resolveAgentIdFromSessionKey,
  resolveAgentMainSessionKey,
  type SessionEntry,
-  saveSessionStore,
+  updateSessionStore,
 } from "../../config/sessions.js";
 import { registerAgentRunContext } from "../../infra/agent-events.js";
 import { resolveOutboundTarget } from "../../infra/outbound/targets.js";
@@ -136,7 +136,7 @@ export const agentHandlers: GatewayRequestHandlers = {
    let cfgForAgent: ReturnType<typeof loadConfig> | undefined;

    if (requestedSessionKey) {
-      const { cfg, storePath, store, entry, canonicalKey } = loadSessionEntry(requestedSessionKey);
+      const { cfg, storePath, entry, canonicalKey } = loadSessionEntry(requestedSessionKey);
      cfgForAgent = cfg;
      const now = Date.now();
      const sessionId = entry?.sessionId ?? randomUUID();
@@ -178,11 +178,10 @@ export const agentHandlers: GatewayRequestHandlers = {
      const canonicalSessionKey = canonicalKey;
      const agentId = resolveAgentIdFromSessionKey(canonicalSessionKey);
      const mainSessionKey = resolveAgentMainSessionKey({ cfg, agentId });
-      if (store) {
-        store[canonicalSessionKey] = nextEntry;
-        if (storePath) {
-          await saveSessionStore(storePath, store);
-        }
+      if (storePath) {
+        await updateSessionStore(storePath, (store) => {
+          store[canonicalSessionKey] = nextEntry;
+        });
      }
      if (canonicalSessionKey === mainSessionKey || canonicalSessionKey === "global") {
        context.addChatRun(idem, {
--- a/src/gateway/server-methods/sessions.ts
+++ b/src/gateway/server-methods/sessions.ts
@@ -9,10 +9,9 @@ import {
 } from "../../agents/pi-embedded.js";
 import { loadConfig } from "../../config/config.js";
 import {
-  loadSessionStore,
  resolveMainSessionKey,
  type SessionEntry,
-  saveSessionStore,
+  updateSessionStore,
 } from "../../config/sessions.js";
 import { clearCommandLane } from "../../process/command-queue.js";
 import {
@@ -30,6 +29,7 @@ import {
  archiveFileOnDisk,
  listSessionsFromStore,
  loadCombinedSessionStoreForGateway,
+  loadSessionEntry,
  resolveGatewaySessionStoreTarget,
  resolveSessionTranscriptCandidates,
  type SessionsPatchResult,
@@ -106,26 +106,25 @@ export const sessionsHandlers: GatewayRequestHandlers = {
    const cfg = loadConfig();
    const target = resolveGatewaySessionStoreTarget({ cfg, key });
    const storePath = target.storePath;
-    const store = loadSessionStore(storePath);
-
-    const primaryKey = target.storeKeys[0] ?? key;
-    const existingKey = target.storeKeys.find((candidate) => store[candidate]);
-    if (existingKey && existingKey !== primaryKey && !store[primaryKey]) {
-      store[primaryKey] = store[existingKey];
-      delete store[existingKey];
-    }
-    const applied = await applySessionsPatchToStore({
-      cfg,
-      store,
-      storeKey: primaryKey,
-      patch: p,
-      loadGatewayModelCatalog: context.loadGatewayModelCatalog,
+    const applied = await updateSessionStore(storePath, async (store) => {
+      const primaryKey = target.storeKeys[0] ?? key;
+      const existingKey = target.storeKeys.find((candidate) => store[candidate]);
+      if (existingKey && existingKey !== primaryKey && !store[primaryKey]) {
+        store[primaryKey] = store[existingKey];
+        delete store[existingKey];
+      }
+      return await applySessionsPatchToStore({
+        cfg,
+        store,
+        storeKey: primaryKey,
+        patch: p,
+        loadGatewayModelCatalog: context.loadGatewayModelCatalog,
+      });
    });
    if (!applied.ok) {
      respond(false, undefined, applied.error);
      return;
    }
-    await saveSessionStore(storePath, store);
    const result: SessionsPatchResult = {
      ok: true,
      path: storePath,
@@ -156,34 +155,35 @@ export const sessionsHandlers: GatewayRequestHandlers = {
    const cfg = loadConfig();
    const target = resolveGatewaySessionStoreTarget({ cfg, key });
    const storePath = target.storePath;
-    const store = loadSessionStore(storePath);
-    const primaryKey = target.storeKeys[0] ?? key;
-    const existingKey = target.storeKeys.find((candidate) => store[candidate]);
-    if (existingKey && existingKey !== primaryKey && !store[primaryKey]) {
-      store[primaryKey] = store[existingKey];
-      delete store[existingKey];
-    }
-    const entry = store[primaryKey];
-    const now = Date.now();
-    const next: SessionEntry = {
-      sessionId: randomUUID(),
-      updatedAt: now,
-      systemSent: false,
-      abortedLastRun: false,
-      thinkingLevel: entry?.thinkingLevel,
-      verboseLevel: entry?.verboseLevel,
-      reasoningLevel: entry?.reasoningLevel,
-      responseUsage: entry?.responseUsage,
-      model: entry?.model,
-      contextTokens: entry?.contextTokens,
-      sendPolicy: entry?.sendPolicy,
-      label: entry?.label,
-      lastChannel: entry?.lastChannel,
-      lastTo: entry?.lastTo,
-      skillsSnapshot: entry?.skillsSnapshot,
-    };
-    store[primaryKey] = next;
-    await saveSessionStore(storePath, store);
+    const next = await updateSessionStore(storePath, (store) => {
+      const primaryKey = target.storeKeys[0] ?? key;
+      const existingKey = target.storeKeys.find((candidate) => store[candidate]);
+      if (existingKey && existingKey !== primaryKey && !store[primaryKey]) {
+        store[primaryKey] = store[existingKey];
+        delete store[existingKey];
+      }
+      const entry = store[primaryKey];
+      const now = Date.now();
+      const nextEntry: SessionEntry = {
+        sessionId: randomUUID(),
+        updatedAt: now,
+        systemSent: false,
+        abortedLastRun: false,
+        thinkingLevel: entry?.thinkingLevel,
+        verboseLevel: entry?.verboseLevel,
+        reasoningLevel: entry?.reasoningLevel,
+        responseUsage: entry?.responseUsage,
+        model: entry?.model,
+        contextTokens: entry?.contextTokens,
+        sendPolicy: entry?.sendPolicy,
+        label: entry?.label,
+        lastChannel: entry?.lastChannel,
+        lastTo: entry?.lastTo,
+        skillsSnapshot: entry?.skillsSnapshot,
+      };
+      store[primaryKey] = nextEntry;
+      return nextEntry;
+    });
    respond(true, { ok: true, key: target.canonicalKey, entry: next }, undefined);
  },
  "sessions.delete": async ({ params, respond }) => {
@@ -220,14 +220,7 @@ export const sessionsHandlers: GatewayRequestHandlers = {
    const deleteTranscript = typeof p.deleteTranscript === "boolean" ? p.deleteTranscript : true;

    const storePath = target.storePath;
-    const store = loadSessionStore(storePath);
-    const primaryKey = target.storeKeys[0] ?? key;
-    const existingKey = target.storeKeys.find((candidate) => store[candidate]);
-    if (existingKey && existingKey !== primaryKey && !store[primaryKey]) {
-      store[primaryKey] = store[existingKey];
-      delete store[existingKey];
-    }
-    const entry = store[primaryKey];
+    const { entry } = loadSessionEntry(key);
    const sessionId = entry?.sessionId;
    const existed = Boolean(entry);
    clearCommandLane(resolveEmbeddedSessionLane(target.canonicalKey));
@@ -246,8 +239,15 @@ export const sessionsHandlers: GatewayRequestHandlers = {
        return;
      }
    }
-    if (existed) delete store[primaryKey];
-    await saveSessionStore(storePath, store);
+    await updateSessionStore(storePath, (store) => {
+      const primaryKey = target.storeKeys[0] ?? key;
+      const existingKey = target.storeKeys.find((candidate) => store[candidate]);
+      if (existingKey && existingKey !== primaryKey && !store[primaryKey]) {
+        store[primaryKey] = store[existingKey];
+        delete store[existingKey];
+      }
+      if (store[primaryKey]) delete store[primaryKey];
+    });

    const archived: string[] = [];
    if (deleteTranscript && sessionId) {
@@ -295,14 +295,17 @@ export const sessionsHandlers: GatewayRequestHandlers = {
    const cfg = loadConfig();
    const target = resolveGatewaySessionStoreTarget({ cfg, key });
    const storePath = target.storePath;
-    const store = loadSessionStore(storePath);
-    const primaryKey = target.storeKeys[0] ?? key;
-    const existingKey = target.storeKeys.find((candidate) => store[candidate]);
-    if (existingKey && existingKey !== primaryKey && !store[primaryKey]) {
-      store[primaryKey] = store[existingKey];
-      delete store[existingKey];
-    }
-    const entry = store[primaryKey];
+    // Lock + read in a short critical section; transcript work happens outside.
+    const compactTarget = await updateSessionStore(storePath, (store) => {
+      const primaryKey = target.storeKeys[0] ?? key;
+      const existingKey = target.storeKeys.find((candidate) => store[candidate]);
+      if (existingKey && existingKey !== primaryKey && !store[primaryKey]) {
+        store[primaryKey] = store[existingKey];
+        delete store[existingKey];
+      }
+      return { entry: store[primaryKey], primaryKey };
+    });
+    const entry = compactTarget.entry;
    const sessionId = entry?.sessionId;
    if (!sessionId) {
      respond(
@@ -358,13 +361,15 @@ export const sessionsHandlers: GatewayRequestHandlers = {
    const keptLines = lines.slice(-maxLines);
    fs.writeFileSync(filePath, `${keptLines.join("\n")}\n`, "utf-8");

-    if (store[primaryKey]) {
-      delete store[primaryKey].inputTokens;
-      delete store[primaryKey].outputTokens;
-      delete store[primaryKey].totalTokens;
-      store[primaryKey].updatedAt = Date.now();
-      await saveSessionStore(storePath, store);
-    }
+    await updateSessionStore(storePath, (store) => {
+      const entryKey = compactTarget.primaryKey;
+      const entryToUpdate = store[entryKey];
+      if (!entryToUpdate) return;
+      delete entryToUpdate.inputTokens;
+      delete entryToUpdate.outputTokens;
+      delete entryToUpdate.totalTokens;
+      entryToUpdate.updatedAt = Date.now();
+    });

    respond(
      true,
--- a/src/gateway/session-utils.ts
+++ b/src/gateway/session-utils.ts
@@ -285,9 +285,12 @@ export function loadCombinedSessionStoreForGateway(cfg: ClawdbotConfig): {
    const store = loadSessionStore(storePath);
    for (const [key, entry] of Object.entries(store)) {
      const canonicalKey = canonicalizeSessionKeyForAgent(agentId, key);
+      // Merge with existing entry if present (avoid overwriting with less complete data)
+      const existing = combined[canonicalKey];
      combined[canonicalKey] = {
+        ...existing,
        ...entry,
-        spawnedBy: canonicalizeSpawnedByForAgent(agentId, entry.spawnedBy),
+        spawnedBy: canonicalizeSpawnedByForAgent(agentId, entry.spawnedBy ?? existing?.spawnedBy),
      };
    }
  }