diff --git a/desktop/src/features/agents/activeAgentTurnsStore.test.mjs b/desktop/src/features/agents/activeAgentTurnsStore.test.mjs index 75d216e2eb..4897d73a45 100644 --- a/desktop/src/features/agents/activeAgentTurnsStore.test.mjs +++ b/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", + ); + }); +}); diff --git a/desktop/src/features/agents/activeAgentTurnsStore.ts b/desktop/src/features/agents/activeAgentTurnsStore.ts index 9ffd34db63..07a6e7e01d 100644 --- a/desktop/src/features/agents/activeAgentTurnsStore.ts +++ b/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>; + offsets: Map; + watermarks: Map; + terminals: Map>; +}; + +/** Per-workspace snapshots. Keyed by workspace ID. */ +const savedByWorkspace = new Map(); + +/** + * 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>(); + for (const [agentKey, agentTurns] of activeTurnsByAgent) { + const clonedAgent = new Map(); + 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>(); + 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(); + 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); +} diff --git a/desktop/src/features/workspaces/useWorkspaceInit.ts b/desktop/src/features/workspaces/useWorkspaceInit.ts index d9df6113fb..91d01f076b 100644 --- a/desktop/src/features/workspaces/useWorkspaceInit.ts +++ b/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(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, diff --git a/desktop/src/features/workspaces/useWorkspaces.tsx b/desktop/src/features/workspaces/useWorkspaces.tsx index e984ea4920..e289baa179 100644 --- a/desktop/src/features/workspaces/useWorkspaces.tsx +++ b/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); } }