fix(desktop): preserve agent timer state across workspace switches (#1679)

Co-authored-by: npub1mn7jgtj4w2pd0g0zeuhxsa6jy6p0rewxz4kujt98my82ahfmp72sxjexk7 <dcfd242e557282d7a1e2cf2e6877522682f1e5c6156dc92ca7d90eaedd3b0f95@sprout-oss.stage.blox.sqprod.co>

desktop/src/features/agents/activeAgentTurnsStore.test.mjs

@@ -8,6 +8,9 @@ import {
   getActiveTurnsByChannel,
   resetActiveAgentTurnsStore,
   subscribeActiveAgentTurns,
+  saveActiveAgentTurnsForWorkspace,
+  restoreActiveAgentTurnsForWorkspace,
+  clearSavedWorkspaceSnapshot,
 } from "./activeAgentTurnsStore.ts";
 import {
   injectObserverEventsForE2E,
@@ -1328,3 +1331,243 @@ describe("formatElapsed", () => {
     assert.equal(formatElapsed(-5_000), "0s");
   });
 });
+
+describe("workspace-switch save / restore", () => {
+  beforeEach(() => {
+    resetActiveAgentTurnsStore();
+  });
+
+  it("restores original startedAt timestamps across a round-trip", () => {
+    // Simulate a turn active in workspace A with a known start timestamp.
+    syncAgentTurnsFromEvents(AGENT, [
+      makeEvent({
+        seq: 1,
+        turnId: "t1",
+        channelId: "c1",
+        timestamp: "2024-01-01T00:00:00Z",
+      }),
+    ]);
+    const anchorBefore = getActiveTurnsForAgent(AGENT)[0].anchorAt;
+
+    // Switch away (save A, reset, apply B).
+    saveActiveAgentTurnsForWorkspace("ws-a");
+    resetActiveAgentTurnsStore();
+    assert.equal(
+      getActiveTurnsForAgent(AGENT).length,
+      0,
+      "store must be empty after reset",
+    );
+
+    // Switch back (restore A).
+    restoreActiveAgentTurnsForWorkspace("ws-a");
+    const turns = getActiveTurnsForAgent(AGENT);
+    assert.equal(turns.length, 1, "restored turn must reappear");
+    assert.equal(turns[0].channelId, "c1");
+    assert.equal(
+      turns[0].anchorAt,
+      anchorBefore,
+      "anchorAt must equal the pre-switch value — startedAt and offset are preserved",
+    );
+  });
+
+  it("no-op restore when no snapshot exists for the workspace", () => {
+    // Populate the store so we can verify nothing changes.
+    syncAgentTurnsFromEvents(AGENT, [
+      makeEvent({ seq: 1, turnId: "t1", channelId: "c1" }),
+    ]);
+    let notified = 0;
+    const unsub = subscribeActiveAgentTurns(() => {
+      notified++;
+    });
+    restoreActiveAgentTurnsForWorkspace("ws-unknown");
+    unsub();
+    assert.equal(notified, 0, "no-op restore must not notify listeners");
+    // Store should be unchanged — still has the turn we set up above.
+    assert.equal(getActiveTurnsForAgent(AGENT).length, 1);
+  });
+
+  it("empty store discards a prior snapshot rather than saving an empty one", () => {
+    // First round-trip: save something.
+    syncAgentTurnsFromEvents(AGENT, [
+      makeEvent({ seq: 1, turnId: "t1", channelId: "c1" }),
+    ]);
+    saveActiveAgentTurnsForWorkspace("ws-a");
+
+    // Second pass: the turns ended while we were away, so the store is empty
+    // when we try to save again.
+    resetActiveAgentTurnsStore();
+    saveActiveAgentTurnsForWorkspace("ws-a"); // empty store → delete snapshot
+
+    // Restore must be a no-op now.
+    restoreActiveAgentTurnsForWorkspace("ws-a");
+    assert.equal(
+      getActiveTurnsForAgent(AGENT).length,
+      0,
+      "restoring after an empty save must yield no turns",
+    );
+  });
+
+  it("snapshot is consumed on restore — a second restore is a no-op", () => {
+    syncAgentTurnsFromEvents(AGENT, [
+      makeEvent({ seq: 1, turnId: "t1", channelId: "c1" }),
+    ]);
+    saveActiveAgentTurnsForWorkspace("ws-a");
+    resetActiveAgentTurnsStore();
+
+    restoreActiveAgentTurnsForWorkspace("ws-a");
+    assert.equal(getActiveTurnsForAgent(AGENT).length, 1, "first restore ok");
+
+    // Second restore — snapshot was consumed.
+    resetActiveAgentTurnsStore();
+    restoreActiveAgentTurnsForWorkspace("ws-a");
+    assert.equal(
+      getActiveTurnsForAgent(AGENT).length,
+      0,
+      "second restore must be no-op — snapshot is consumed on first use",
+    );
+  });
+
+  it("watermark is preserved — events already processed before save are not reprocessed", () => {
+    syncAgentTurnsFromEvents(AGENT, [
+      makeEvent({
+        seq: 5,
+        turnId: "t1",
+        channelId: "c1",
+        timestamp: "2024-01-01T00:00:00Z",
+      }),
+    ]);
+    saveActiveAgentTurnsForWorkspace("ws-a");
+    resetActiveAgentTurnsStore();
+    restoreActiveAgentTurnsForWorkspace("ws-a");
+
+    // Replaying an event at or below the watermark must be a no-op.
+    let notified = 0;
+    const unsub = subscribeActiveAgentTurns(() => {
+      notified++;
+    });
+    syncAgentTurnsFromEvents(AGENT, [
+      makeEvent({
+        seq: 3,
+        turnId: "t2",
+        channelId: "c2",
+        timestamp: "2024-01-01T00:00:00Z",
+      }),
+    ]);
+    unsub();
+
+    assert.equal(notified, 0, "stale event after restore must not notify");
+    const channels = new Set(
+      getActiveTurnsForAgent(AGENT).map((s) => s.channelId),
+    );
+    assert.ok(!channels.has("c2"), "stale event must not add a turn");
+  });
+
+  it("terminal tombstones are preserved — a stale liveness cannot resurrect a completed turn", () => {
+    // Complete a turn before saving.
+    syncAgentTurnsFromEvents(AGENT, [
+      makeEvent({
+        seq: 1,
+        turnId: "t1",
+        channelId: "c1",
+        timestamp: "2024-01-01T00:00:00Z",
+      }),
+      makeEvent({
+        seq: 2,
+        kind: "turn_completed",
+        turnId: "t1",
+        channelId: "c1",
+        timestamp: "2024-01-01T00:00:10Z",
+      }),
+    ]);
+    saveActiveAgentTurnsForWorkspace("ws-a");
+    resetActiveAgentTurnsStore();
+    restoreActiveAgentTurnsForWorkspace("ws-a");
+
+    // A liveness stamped before the completion must not resurrect t1.
+    syncAgentTurnsFromEvents(AGENT, [
+      makeEvent({
+        seq: 3,
+        kind: "turn_liveness",
+        turnId: "t1",
+        channelId: "c1",
+        timestamp: "2024-01-01T00:00:05Z",
+      }),
+    ]);
+    assert.equal(
+      getActiveTurnsForAgent(AGENT).length,
+      0,
+      "terminal tombstone must survive save/restore — stale liveness must not resurrect",
+    );
+  });
+
+  it("notifies listeners after restore so UI re-renders with recovered turns", () => {
+    syncAgentTurnsFromEvents(AGENT, [
+      makeEvent({ seq: 1, turnId: "t1", channelId: "c1" }),
+    ]);
+    saveActiveAgentTurnsForWorkspace("ws-a");
+    resetActiveAgentTurnsStore();
+
+    let notified = 0;
+    const unsub = subscribeActiveAgentTurns(() => {
+      notified++;
+    });
+    restoreActiveAgentTurnsForWorkspace("ws-a");
+    unsub();
+
+    assert.ok(notified > 0, "restore must notify listeners");
+  });
+
+  it("multiple workspaces maintain independent snapshots", () => {
+    // Workspace A: agent in c1.
+    syncAgentTurnsFromEvents(AGENT, [
+      makeEvent({ seq: 1, turnId: "t1", channelId: "c1" }),
+    ]);
+    saveActiveAgentTurnsForWorkspace("ws-a");
+
+    // Workspace B: different agent/channel.
+    resetActiveAgentTurnsStore();
+    syncAgentTurnsFromEvents(AGENT_2, [
+      makeEvent({ seq: 1, turnId: "t2", channelId: "c2" }),
+    ]);
+    saveActiveAgentTurnsForWorkspace("ws-b");
+
+    // Switch to A — only A's turns must appear.
+    resetActiveAgentTurnsStore();
+    restoreActiveAgentTurnsForWorkspace("ws-a");
+    const aChannels = new Set(
+      getActiveTurnsForAgent(AGENT).map((s) => s.channelId),
+    );
+    assert.ok(aChannels.has("c1"), "ws-a must restore c1");
+    assert.equal(
+      getActiveTurnsForAgent(AGENT_2).length,
+      0,
+      "ws-b turns must not appear in ws-a restore",
+    );
+
+    // Switch to B — only B's turns.
+    resetActiveAgentTurnsStore();
+    restoreActiveAgentTurnsForWorkspace("ws-b");
+    const bChannels = new Set(
+      getActiveTurnsForAgent(AGENT_2).map((s) => s.channelId),
+    );
+    assert.ok(bChannels.has("c2"), "ws-b must restore c2");
+  });
+
+  it("clearSavedWorkspaceSnapshot discards the snapshot so restore is a no-op", () => {
+    syncAgentTurnsFromEvents(AGENT, [
+      makeEvent({ seq: 1, turnId: "t1", channelId: "c1" }),
+    ]);
+    saveActiveAgentTurnsForWorkspace("ws-a");
+
+    // Simulate workspace deletion: GC the snapshot before switching back.
+    clearSavedWorkspaceSnapshot("ws-a");
+
+    resetActiveAgentTurnsStore();
+    restoreActiveAgentTurnsForWorkspace("ws-a");
+    assert.equal(
+      getActiveTurnsForAgent(AGENT).length,
+      0,
+      "restore must be no-op after clearSavedWorkspaceSnapshot",
+    );
+  });
+});

desktop/src/features/agents/activeAgentTurnsStore.ts

@@ -559,6 +559,11 @@ export function useActiveAgentTurnsBridge(
   }, [agents]);
 }
 
+/**
+ * Clears all live turn state (active turns, offsets, watermarks, tombstones).
+ * Intentionally preserves `savedByWorkspace` — workspace-switch snapshots
+ * must survive the reset that runs between save and restore.
+ */
 export function resetActiveAgentTurnsStore() {
   activeTurnsByAgent.clear();
   lastProcessed.clear();
@@ -568,3 +573,121 @@ export function resetActiveAgentTurnsStore() {
   terminalAtByAgent.clear();
   notifyListeners();
 }
+
+// ---------------------------------------------------------------------------
+// Workspace-switch save / restore
+// ---------------------------------------------------------------------------
+
+type TurnsStoreSnapshot = {
+  turns: Map<string, Map<string, ActiveTurn>>;
+  offsets: Map<string, number>;
+  watermarks: Map<string, ObserverEvent>;
+  terminals: Map<string, Map<string, number>>;
+};
+
+/** Per-workspace snapshots. Keyed by workspace ID. */
+const savedByWorkspace = new Map<string, TurnsStoreSnapshot>();
+
+/**
+ * Snapshot the current active-turns state under `workspaceId` so it can be
+ * restored when the user switches back.  If both the turns map and the
+ * tombstone map are empty there is nothing worth restoring — discard any
+ * previously-saved snapshot instead.
+ *
+ * Deep-clones all four maps so subsequent mutations on the live maps do not
+ * corrupt the snapshot.
+ */
+export function saveActiveAgentTurnsForWorkspace(workspaceId: string): void {
+  if (activeTurnsByAgent.size === 0 && terminalAtByAgent.size === 0) {
+    savedByWorkspace.delete(workspaceId);
+    return;
+  }
+
+  // Deep-clone activeTurnsByAgent: outer map + inner per-agent maps + turn
+  // objects (plain structs, no nested references beyond primitives).
+  const turns = new Map<string, Map<string, ActiveTurn>>();
+  for (const [agentKey, agentTurns] of activeTurnsByAgent) {
+    const clonedAgent = new Map<string, ActiveTurn>();
+    for (const [turnId, turn] of agentTurns) {
+      clonedAgent.set(turnId, { ...turn });
+    }
+    turns.set(agentKey, clonedAgent);
+  }
+
+  // Shallow-clone scalar maps (primitives as values).
+  const offsets = new Map(clockOffsetByAgent);
+  const watermarks = new Map(lastProcessed);
+
+  // Deep-clone terminalAtByAgent: outer map + inner per-agent maps.
+  const terminals = new Map<string, Map<string, number>>();
+  for (const [agentKey, tombstones] of terminalAtByAgent) {
+    terminals.set(agentKey, new Map(tombstones));
+  }
+
+  savedByWorkspace.set(workspaceId, { turns, offsets, watermarks, terminals });
+}
+
+/**
+ * Restore a previously saved active-turns snapshot for `workspaceId` into the
+ * module maps.  No-op when no snapshot exists.
+ *
+ * Clears all four module maps before writing so the function is
+ * self-contained — it replaces rather than merging, regardless of whether the
+ * caller pre-cleared.  At the primary call site (`useWorkspaceInit`) the maps
+ * are already empty after `resetWorkspaceState()`, but this guard makes the
+ * contract explicit.
+ *
+ * Refreshes `lastActivityAt` on every restored turn so the prune interval
+ * doesn't immediately kill turns that were saved more than 25 s ago (the prune
+ * threshold).  New observer events arriving after restore will update
+ * `lastActivityAt` normally via `recordActivity`.
+ *
+ * Consumes the snapshot (deletes it from `savedByWorkspace`) — a given
+ * workspace's snapshot is only usable once per round-trip.
+ */
+export function restoreActiveAgentTurnsForWorkspace(workspaceId: string): void {
+  const snap = savedByWorkspace.get(workspaceId);
+  if (!snap) return;
+  savedByWorkspace.delete(workspaceId);
+
+  // Clear before writing so this is a replace, not a merge.
+  activeTurnsByAgent.clear();
+  clockOffsetByAgent.clear();
+  lastProcessed.clear();
+  terminalAtByAgent.clear();
+
+  const now = Date.now();
+
+  for (const [agentKey, agentTurns] of snap.turns) {
+    const restored = new Map<string, ActiveTurn>();
+    for (const [turnId, turn] of agentTurns) {
+      restored.set(turnId, { ...turn, lastActivityAt: now });
+    }
+    activeTurnsByAgent.set(agentKey, restored);
+  }
+
+  for (const [agentKey, offset] of snap.offsets) {
+    clockOffsetByAgent.set(agentKey, offset);
+  }
+
+  for (const [agentKey, event] of snap.watermarks) {
+    lastProcessed.set(agentKey, event);
+  }
+
+  for (const [agentKey, tombstones] of snap.terminals) {
+    terminalAtByAgent.set(agentKey, new Map(tombstones));
+  }
+
+  cachedTurnSummaries.clear();
+  cachedChannelTurnSummaries = null;
+  notifyListeners();
+}
+
+/**
+ * Discard the saved turn-state snapshot for a workspace that has been
+ * permanently deleted so the entry doesn't sit in memory indefinitely.
+ * Call this alongside the other relay-specific GC in `removeWorkspace`.
+ */
+export function clearSavedWorkspaceSnapshot(workspaceId: string): void {
+  savedByWorkspace.delete(workspaceId);
+}

desktop/src/features/workspaces/useWorkspaceInit.ts

@@ -13,7 +13,11 @@ import {
   initDraftStore,
 } from "@/features/messages/lib/useDrafts";
 import { resetRenderScopedReactionHydration } from "@/features/messages/lib/renderScopedReactions";
-import { resetActiveAgentTurnsStore } from "@/features/agents/activeAgentTurnsStore";
+import {
+  resetActiveAgentTurnsStore,
+  saveActiveAgentTurnsForWorkspace,
+  restoreActiveAgentTurnsForWorkspace,
+} from "@/features/agents/activeAgentTurnsStore";
 import { resetAgentWorkingSignal } from "@/features/agents/agentWorkingSignal";
 import { resetAgentObserverStore } from "@/features/agents/observerRelayStore";
 import { resetSidebarRelayConnectionCardState } from "@/features/sidebar/ui/useSidebarRelayConnectionCard";
@@ -76,6 +80,10 @@ export function useWorkspaceInit(
   // On the initial mount we skip resetting singletons (they're fresh).
   const hasInitializedRef = useRef(false);
 
+  // Track the previously-applied workspace ID so we can save its turn state
+  // before resetting when the user switches to a different workspace.
+  const prevWorkspaceIdRef = useRef<string | null>(null);
+
   // biome-ignore lint/correctness/useExhaustiveDependencies: we intentionally depend on specific properties (id/relayUrl/token/reposDir) — depending on the whole object would trigger resets on name-only changes
   useEffect(() => {
     let cancelled = false;
@@ -127,6 +135,15 @@ export function useWorkspaceInit(
       // On workspace switch (not initial mount), reset module singletons
       // so the new tree starts with a clean slate.
       if (hasInitializedRef.current) {
+        // Save the outgoing workspace's turn state before wiping the store so
+        // timers survive a round-trip (A → B → A keeps A's elapsed time).
+        if (prevWorkspaceIdRef.current) {
+          saveActiveAgentTurnsForWorkspace(prevWorkspaceIdRef.current);
+          // Null out immediately so a rapid workspace switch (A→B→C before
+          // B's applyWorkspace resolves) doesn't re-save the now-empty
+          // store under the outgoing workspace ID and delete its snapshot.
+          prevWorkspaceIdRef.current = null;
+        }
         resetWorkspaceState();
       }
       hasInitializedRef.current = true;
@@ -170,6 +187,12 @@ export function useWorkspaceInit(
         if (activeWorkspace.pubkey) {
           initDraftStore(activeWorkspace.pubkey);
         }
+        // Restore any turn state saved for this workspace (a prior A→B round-
+        // trip). This runs after applyWorkspace succeeds and before the app
+        // renders so components see the restored timers on first render.
+        restoreActiveAgentTurnsForWorkspace(activeWorkspace.id);
+        // Prime the ref so the NEXT switch saves this workspace's state.
+        prevWorkspaceIdRef.current = activeWorkspace.id;
         setResult({
           isReady: true,
           needsSetup: false,

desktop/src/features/workspaces/useWorkspaces.tsx

@@ -19,6 +19,7 @@ import {
 import { removeSelfProfileCachesForRelay } from "@/features/profile/lib/selfProfileStorage";
 import { removeChannelSnapshotForRelay } from "@/features/channels/channelSnapshot";
 import { removeMessageSnapshotsForRelay } from "@/features/messages/lib/messageSnapshot";
+import { clearSavedWorkspaceSnapshot } from "@/features/agents/activeAgentTurnsStore";
 
 export type UseWorkspacesReturn = {
   workspaces: Workspace[];
@@ -120,6 +121,7 @@ function useWorkspacesInternal(): UseWorkspacesReturn {
           removeSelfProfileCachesForRelay(removed.relayUrl);
           removeChannelSnapshotForRelay(removed.relayUrl);
           removeMessageSnapshotsForRelay(removed.relayUrl);
+          clearSavedWorkspaceSnapshot(id);
         }
       }