diff --git a/packages/app/src/contexts/session-context.tsx b/packages/app/src/contexts/session-context.tsx index b9d139ec7..96f2993fc 100644 --- a/packages/app/src/contexts/session-context.tsx +++ b/packages/app/src/contexts/session-context.tsx @@ -49,6 +49,7 @@ import { derivePendingPermissionKey, normalizeAgentSnapshot } from "@/utils/agen import { resolveProjectPlacement } from "@/utils/project-placement"; import { buildDraftStoreKey } from "@/stores/draft-keys"; import type { AttachmentMetadata } from "@/attachments/types"; +import { reconcilePreviousAgentStatuses } from "@/contexts/session-status-tracking"; // Re-export types from session-store and draft-store for backward compatibility export type { DraftInput } from "@/stores/draft-store"; @@ -295,15 +296,10 @@ function SessionProviderInternal({ children, serverId, client }: SessionProvider }, []); useEffect(() => { - if (!sessionAgents) { - previousAgentStatusRef.current.clear(); - return; - } - const nextStatuses = new Map(); - for (const nextAgent of sessionAgents.values()) { - nextStatuses.set(nextAgent.id, nextAgent.status); - } - previousAgentStatusRef.current = nextStatuses; + previousAgentStatusRef.current = reconcilePreviousAgentStatuses( + previousAgentStatusRef.current, + sessionAgents, + ); }, [sessionAgents]); const hydrateWorkspaces = useCallback( diff --git a/packages/app/src/contexts/session-status-tracking.test.ts b/packages/app/src/contexts/session-status-tracking.test.ts new file mode 100644 index 000000000..d5e98ed3a --- /dev/null +++ b/packages/app/src/contexts/session-status-tracking.test.ts @@ -0,0 +1,72 @@ +import { describe, expect, it } from "vitest"; +import type { Agent } from "@/stores/session-store"; +import { reconcilePreviousAgentStatuses } from "./session-status-tracking"; + +function createAgent(status: Agent["status"]): Agent { + return { + serverId: "server-1", + id: "agent-1", + provider: "codex", + status, + createdAt: new Date(0), + updatedAt: new Date(0), + lastUserMessageAt: null, + lastActivityAt: new Date(0), + capabilities: { + supportsStreaming: true, + supportsSessionPersistence: true, + supportsDynamicModes: true, + supportsMcpServers: true, + supportsReasoningStream: true, + supportsToolInvocations: true, + }, + currentModeId: null, + availableModes: [], + pendingPermissions: [], + persistence: null, + title: "Agent", + cwd: "/tmp", + model: null, + labels: {}, + projectPlacement: null, + }; +} + +describe("reconcilePreviousAgentStatuses", () => { + it("preserves previously seen status for existing agents", () => { + const previous = new Map([["agent-1", "running" as const]]); + const sessionAgents = new Map([["agent-1", createAgent("idle")]]); + + const result = reconcilePreviousAgentStatuses(previous, sessionAgents); + + expect(result).toEqual(new Map([["agent-1", "running"]])); + }); + + it("seeds newly seen agents from the current snapshot", () => { + const sessionAgents = new Map([["agent-1", createAgent("idle")]]); + + const result = reconcilePreviousAgentStatuses(new Map(), sessionAgents); + + expect(result).toEqual(new Map([["agent-1", "idle"]])); + }); + + it("removes agents that are no longer present", () => { + const previous = new Map([ + ["agent-1", "running" as const], + ["agent-2", "idle" as const], + ]); + const sessionAgents = new Map([["agent-1", createAgent("idle")]]); + + const result = reconcilePreviousAgentStatuses(previous, sessionAgents); + + expect(result).toEqual(new Map([["agent-1", "running"]])); + }); + + it("clears all tracked statuses when the session is unavailable", () => { + const previous = new Map([["agent-1", "running" as const]]); + + const result = reconcilePreviousAgentStatuses(previous, undefined); + + expect(result).toEqual(new Map()); + }); +}); diff --git a/packages/app/src/contexts/session-status-tracking.ts b/packages/app/src/contexts/session-status-tracking.ts new file mode 100644 index 000000000..58657ae09 --- /dev/null +++ b/packages/app/src/contexts/session-status-tracking.ts @@ -0,0 +1,29 @@ +import type { AgentLifecycleStatus } from "@server/shared/agent-lifecycle"; +import type { Agent } from "@/stores/session-store"; + +export function reconcilePreviousAgentStatuses( + previousStatuses: Map, + sessionAgents: Map | undefined, +): Map { + if (!sessionAgents) { + return new Map(); + } + + const nextStatuses = new Map(previousStatuses); + const seenAgentIds = new Set(); + + for (const agent of sessionAgents.values()) { + seenAgentIds.add(agent.id); + if (!nextStatuses.has(agent.id)) { + nextStatuses.set(agent.id, agent.status); + } + } + + for (const agentId of nextStatuses.keys()) { + if (!seenAgentIds.has(agentId)) { + nextStatuses.delete(agentId); + } + } + + return nextStatuses; +}