diff --git a/packages/app/src/components/rewind/use-rewind-agent-mutation.ts b/packages/app/src/components/rewind/use-rewind-agent-mutation.ts index e684c4bf53..037f7c7330 100644 --- a/packages/app/src/components/rewind/use-rewind-agent-mutation.ts +++ b/packages/app/src/components/rewind/use-rewind-agent-mutation.ts @@ -8,6 +8,7 @@ import { useRewindComposerRestore } from "./composer-restore"; import { useSessionStore } from "@/stores/session-store"; import { shouldRestoreComposerForRewindMode } from "./rewind-mode"; import { clearOptimisticUserMessages } from "@/types/stream"; +import { getHostRuntimeStore } from "@/runtime/host-runtime"; interface UseRewindAgentMutationInput { serverId?: string; @@ -47,7 +48,8 @@ export function useRewindAgentMutation(input: UseRewindAgentMutationInput): { .getState() .sessions[input.serverId]?.agentTimelineCursor.get(input.agentId) : undefined; - await input.client.fetchAgentTimeline(input.agentId, { + if (!input.serverId) throw new Error(t("common.errors.daemonClientUnavailable")); + await getHostRuntimeStore().fetchAgentTimeline(input.serverId, input.agentId, { direction: "tail", projection: "projected", ...(cursor ? { cursor: { epoch: cursor.epoch, seq: cursor.endSeq } } : {}), diff --git a/packages/app/src/contexts/session-context.tsx b/packages/app/src/contexts/session-context.tsx index 8edb1f711c..31720b34d7 100644 --- a/packages/app/src/contexts/session-context.tsx +++ b/packages/app/src/contexts/session-context.tsx @@ -6,7 +6,6 @@ import { useTranslation } from "react-i18next"; import { useClientActivity } from "@/hooks/use-client-activity"; import { useAppVisible } from "@/hooks/use-app-visible"; import { usePushTokenRegistration } from "@/hooks/use-push-token-registration"; -import { clearArchiveAgentPending } from "@/hooks/use-archive-agent"; import { createSetAgentInitializing, refreshAgentInitializationTimeout, @@ -21,7 +20,6 @@ import { } from "@/timeline/session-stream-reducers"; import { useCreateFlowStore } from "@/stores/create-flow-store"; import { isTimelineCatchUpComplete } from "@/timeline/timeline-sync-plan"; -import { fetchAgentTimelineOnce } from "@/timeline/fetch-agent-timeline-once"; import { createViewedTimelineSync, type ViewedTimelineSync } from "@/timeline/viewed-timeline-sync"; import type { AgentAttachment, SessionOutboundMessage } from "@getpaseo/protocol/messages"; import { parseServerInfoStatusPayload } from "@getpaseo/protocol/messages"; @@ -37,17 +35,7 @@ import type { AgentPermissionResponse } from "@getpaseo/protocol/agent-types"; import { getHostRuntimeStore, useHostRuntimeIsConnected } from "@/runtime/host-runtime"; import { useVoiceAudioEngineOptional, useVoiceRuntimeOptional } from "@/contexts/voice-context"; import type { AudioPlaybackSource } from "@/voice/audio-engine-types"; -import { - useSessionStore, - type Agent, - type MessageEntry, - type SessionState, - type WorkspaceDescriptor, - type EmptyProjectDescriptor, - normalizeWorkspaceDescriptor, - normalizeEmptyProjectDescriptor, -} from "@/stores/session-store"; -import { useDraftStore } from "@/stores/draft-store"; +import { useSessionStore, type MessageEntry, type SessionState } from "@/stores/session-store"; import { useWorkspaceSetupStore } from "@/stores/workspace-setup-store"; import { sendOsNotification } from "@/utils/os-notifications"; import { getIsAppActivelyVisible } from "@/utils/app-visibility"; @@ -59,27 +47,15 @@ import { rejectInitDeferred, } from "@/utils/agent-initialization"; import { encodeImages } from "@/utils/encode-images"; -import { derivePendingPermissionKey, normalizeAgentSnapshot } from "@/utils/agent-snapshots"; -import { buildDraftStoreKey } from "@/stores/draft-keys"; +import { derivePendingPermissionKey } from "@/utils/agent-snapshots"; import type { AttachmentMetadata } from "@/attachments/types"; import { patchWorkspaceScripts } from "@/contexts/session-workspace-scripts"; -import { - clearWorkspaceArchivePending, - shouldSuppressWorkspaceForLocalArchive, -} from "@/contexts/session-workspace-upserts"; -import { reconcileWorkspaceDirectory } from "@/contexts/workspace-directory-reconciliation"; import { useToast } from "@/contexts/toast-context"; import { toErrorMessage } from "@/utils/error-messages"; import { showProviderNoticeToast } from "@/utils/provider-notice-toast"; import { applyCheckoutStatusUpdateFromEvent } from "@/git/checkout-status-cache"; -import { - applyLegacyDaemonWorkspaceOwnership, - backfillLegacyDaemonWorkspaceDirectoryIfEmpty, -} from "@/workspace/legacy-daemon-workspaces"; import { useProviderSubagentStore } from "@/subagents/provider-store"; import { revalidateSessionAfterResume } from "@/contexts/session-resume-revalidation"; -import { inheritBufferedDirectoryDeltas } from "@/utils/buffered-directory-transaction"; -import { shouldApplyTimelineAgentSnapshot } from "@/utils/agent-directory-sync"; // Re-export types from session-store and draft-store for backward compatibility export type { DraftInput } from "@/stores/draft-store"; @@ -94,22 +70,6 @@ export type { AgentFileExplorerState, } from "@/stores/session-store"; -function hasAgentUsageChanged( - incomingUsage: Agent["lastUsage"] | undefined, - currentUsage: Agent["lastUsage"] | undefined, -): boolean { - const keys: Array> = [ - "inputTokens", - "outputTokens", - "cachedInputTokens", - "totalCostUsd", - "contextWindowMaxTokens", - "contextWindowUsedTokens", - ]; - - return keys.some((key) => incomingUsage?.[key] !== currentUsage?.[key]); -} - type AudioOutputPayload = Extract["payload"]; interface BufferedAudioChunk { @@ -119,76 +79,6 @@ interface BufferedAudioChunk { id: string; } -interface WorkspaceHydrationSnapshot { - workspaces: Map; - emptyProjects: Map; -} - -interface WorkspaceHydrationTransaction { - id: symbol; - client: DaemonClient; - workspaces: Map; - deltas: WorkspaceUpdatePayload[]; -} - -function isWorkspaceHydrationCancelled(isCancelled?: () => boolean): boolean { - return isCancelled?.() ?? false; -} - -type WorkspaceUpdatePayload = Extract< - SessionOutboundMessage, - { type: "workspace_update" } ->["payload"]; - -async function fetchWorkspaceHydrationSnapshot(input: { - client: DaemonClient; - serverId: string; - subscribe: boolean; - isCancelled?: () => boolean; - transaction: WorkspaceHydrationTransaction; - isCurrent: () => boolean; -}): Promise { - const emptyProjects = new Map(); - let cursor: string | null = null; - let includeSubscribe = input.subscribe; - - while (true) { - const payload = await input.client.fetchWorkspaces({ - sort: [{ key: "activity_at", direction: "desc" }], - ...(includeSubscribe ? { subscribe: {} } : {}), - page: cursor ? { limit: 200, cursor } : { limit: 200 }, - }); - if (input.isCancelled?.() || !input.isCurrent()) { - return null; - } - - for (const entry of payload.entries) { - const workspace = normalizeWorkspaceDescriptor(entry); - if (shouldSuppressWorkspaceForLocalArchive({ serverId: input.serverId, workspace })) { - continue; - } - input.transaction.workspaces.set(workspace.id, workspace); - } - - // Project parents with no active workspaces only ride on the first page. - for (const project of payload.emptyProjects ?? []) { - const descriptor = normalizeEmptyProjectDescriptor(project); - emptyProjects.set(descriptor.projectId, descriptor); - } - - if (!payload.pageInfo.hasMore || !payload.pageInfo.nextCursor) { - break; - } - cursor = payload.pageInfo.nextCursor; - includeSubscribe = false; - } - - return { - workspaces: new Map(input.transaction.workspaces), - emptyProjects, - }; -} - function decodeBase64Chunk(base64: string): Uint8Array { return Buffer.from(base64, "base64"); } @@ -507,24 +397,14 @@ function SessionProviderInternal({ children, serverId, client }: SessionProvider const setAgentAuthoritativeHistoryApplied = useSessionStore( (state) => state.setAgentAuthoritativeHistoryApplied, ); - const setHasHydratedAgents = useSessionStore((state) => state.setHasHydratedAgents); - const setHasHydratedWorkspaces = useSessionStore((state) => state.setHasHydratedWorkspaces); const setAgents = useSessionStore((state) => state.setAgents); const setWorkspaces = useSessionStore((state) => state.setWorkspaces); - const setEmptyProjects = useSessionStore((state) => state.setEmptyProjects); - const addEmptyProject = useSessionStore((state) => state.addEmptyProject); - const removeEmptyProject = useSessionStore((state) => state.removeEmptyProject); - const mergeWorkspaces = useSessionStore((state) => state.mergeWorkspaces); - const removeWorkspace = useSessionStore((state) => state.removeWorkspace); - const setAgentLastActivity = useSessionStore((state) => state.setAgentLastActivity); const flushAgentLastActivity = useSessionStore((state) => state.flushAgentLastActivity); const setPendingPermissions = useSessionStore((state) => state.setPendingPermissions); - const clearDraftInput = useDraftStore((state) => state.clearDraftInput); const updateSessionClient = useSessionStore((state) => state.updateSessionClient); const updateSessionServerInfo = useSessionStore((state) => state.updateSessionServerInfo); const setViewedTimelineSync = useSessionStore((state) => state.setViewedTimelineSync); const upsertWorkspaceSetupProgress = useWorkspaceSetupStore((state) => state.upsertProgress); - const removeWorkspaceSetup = useWorkspaceSetupStore((state) => state.removeWorkspace); const clearWorkspaceSetupServer = useWorkspaceSetupStore((state) => state.clearServer); // Track focused agent for heartbeat @@ -540,36 +420,8 @@ function SessionProviderInternal({ children, serverId, client }: SessionProvider const viewedTimelineSyncRef = useRef(null); const audioOutputBuffersRef = useRef>(new Map()); const activeAudioGroupsRef = useRef>(new Set()); - const workspaceHydrationRef = useRef(null); const isAppVisible = useAppVisible(); - const applyWorkspaceUpdatePayload = useCallback( - (payload: WorkspaceUpdatePayload) => { - if (payload.kind === "remove") { - clearWorkspaceArchivePending({ serverId, workspaceId: payload.id }); - removeWorkspaceSetup({ serverId, workspaceId: payload.id }); - removeWorkspace(serverId, payload.id); - if (payload.emptyProject) { - addEmptyProject(serverId, normalizeEmptyProjectDescriptor(payload.emptyProject)); - } - if (payload.removedProjectId) removeEmptyProject(serverId, payload.removedProjectId); - return; - } - const workspace = normalizeWorkspaceDescriptor(payload.workspace); - if (!shouldSuppressWorkspaceForLocalArchive({ serverId, workspace })) { - mergeWorkspaces(serverId, [workspace]); - } - }, - [ - addEmptyProject, - mergeWorkspaces, - removeEmptyProject, - removeWorkspace, - removeWorkspaceSetup, - serverId, - ], - ); - useEffect(() => { const subscription = AppState.addEventListener("change", (nextState) => { appStateRef.current = nextState; @@ -584,173 +436,6 @@ function SessionProviderInternal({ children, serverId, client }: SessionProvider viewedTimelineSyncRef.current?.setActive(isAppVisible); }, [isAppVisible]); - const hydrateWorkspaces = useCallback( - async (options?: { subscribe?: boolean; isCancelled?: () => boolean }) => { - if (!client || !isConnected) { - return; - } - - const transaction: WorkspaceHydrationTransaction = { - id: Symbol("workspace hydration"), - client, - workspaces: new Map(), - deltas: inheritBufferedDirectoryDeltas({ - client, - previous: workspaceHydrationRef.current, - }), - }; - workspaceHydrationRef.current = transaction; - let snapshot: WorkspaceHydrationSnapshot | null; - try { - snapshot = await fetchWorkspaceHydrationSnapshot({ - client, - serverId, - subscribe: options?.subscribe ?? false, - isCancelled: options?.isCancelled, - transaction, - isCurrent: () => workspaceHydrationRef.current?.id === transaction.id, - }); - } catch (error) { - if (workspaceHydrationRef.current === transaction) { - workspaceHydrationRef.current = null; - for (const delta of transaction.deltas) applyWorkspaceUpdatePayload(delta); - } - throw error; - } - if (!snapshot || isWorkspaceHydrationCancelled(options?.isCancelled)) { - if (workspaceHydrationRef.current === transaction) workspaceHydrationRef.current = null; - return; - } - - const didBackfillLegacy = await backfillLegacyDaemonWorkspaceDirectoryIfEmpty({ - client, - serverId, - workspaces: snapshot.workspaces, - emptyProjects: snapshot.emptyProjects, - isCancelled: options?.isCancelled, - }); - if (didBackfillLegacy) { - if (workspaceHydrationRef.current === transaction) workspaceHydrationRef.current = null; - return; - } - - if ( - workspaceHydrationRef.current !== transaction || - isWorkspaceHydrationCancelled(options?.isCancelled) - ) { - return; - } - const deltas = [...transaction.deltas]; - workspaceHydrationRef.current = null; - setWorkspaces( - serverId, - reconcileWorkspaceDirectory({ - serverId, - snapshot: snapshot.workspaces, - deltas, - }), - ); - setEmptyProjects(serverId, snapshot.emptyProjects.values()); - setHasHydratedWorkspaces(serverId, true); - for (const delta of deltas) applyWorkspaceUpdatePayload(delta); - }, - [ - applyWorkspaceUpdatePayload, - client, - isConnected, - serverId, - setEmptyProjects, - setHasHydratedWorkspaces, - setWorkspaces, - ], - ); - - const applyAuthoritativeAgentSnapshot = useCallback( - (agent: Agent) => { - setAgents(serverId, (prev) => { - const current = prev.get(agent.id); - if (current && agent.updatedAt.getTime() < current.updatedAt.getTime()) { - const hasUsageUpdate = hasAgentUsageChanged(agent.lastUsage, current.lastUsage); - if (hasUsageUpdate) { - const next = new Map(prev); - next.set(agent.id, { - ...current, - lastUsage: agent.lastUsage, - }); - return next; - } - return prev; - } - const next = new Map(prev); - next.set(agent.id, agent); - return next; - }); - - if (agent.archivedAt) { - clearArchiveAgentPending({ - queryClient, - serverId, - agentId: agent.id, - }); - } - - setAgentLastActivity(agent.id, agent.lastActivityAt); - - setPendingPermissions(serverId, (prev) => { - const existingKeysForAgent: string[] = []; - for (const [key, pending] of prev.entries()) { - if (pending.agentId === agent.id) { - existingKeysForAgent.push(key); - } - } - - const nextEntries = agent.pendingPermissions.map((request) => ({ - key: derivePendingPermissionKey(agent.id, request), - agentId: agent.id, - request, - })); - - let changed = existingKeysForAgent.length !== nextEntries.length; - if (!changed) { - const existingKeySet = new Set(existingKeysForAgent); - for (const entry of nextEntries) { - const existing = prev.get(entry.key); - if (!existingKeySet.has(entry.key) || !existing) { - changed = true; - break; - } - - const currentRequest = existing.request; - if ( - currentRequest.id !== entry.request.id || - currentRequest.kind !== entry.request.kind || - currentRequest.name !== entry.request.name || - currentRequest.title !== entry.request.title || - currentRequest.description !== entry.request.description - ) { - changed = true; - break; - } - } - } - - if (!changed) { - return prev; - } - - const next = new Map(prev); - for (const key of existingKeysForAgent) { - next.delete(key); - } - for (const entry of nextEntries) { - next.set(entry.key, entry); - } - return next; - }); - }, - [queryClient, serverId, setAgentLastActivity, setAgents, setPendingPermissions], - ); - const recoverTimelineGap = useCallback( (agentId: string, cursor: { epoch: string; endSeq: number }) => { viewedTimelineSyncRef.current?.recoverGap(agentId, cursor); @@ -764,8 +449,7 @@ function SessionProviderInternal({ children, serverId, client }: SessionProvider awayMs, serverId, bumpHistorySyncGeneration, - refreshAgentDirectory: () => getHostRuntimeStore().refreshAgentDirectory({ serverId }), - hydrateWorkspaces, + refreshDirectories: () => getHostRuntimeStore().refreshDirectories(serverId), }).catch((error) => { console.error("[SessionProvider] resume revalidation failed", { serverId, @@ -773,7 +457,7 @@ function SessionProviderInternal({ children, serverId, client }: SessionProvider }); }); }, - [bumpHistorySyncGeneration, hydrateWorkspaces, serverId], + [bumpHistorySyncGeneration, serverId], ); // Client activity tracking (heartbeat, push token registration) @@ -919,28 +603,6 @@ function SessionProviderInternal({ children, serverId, client }: SessionProvider } }, [flushAgentLastActivity, serverId, isConnected, setInitializingAgents]); - useEffect(() => { - if (!client || !isConnected) { - return () => {}; - } - - let cancelled = false; - void (async () => { - try { - await hydrateWorkspaces({ - subscribe: true, - isCancelled: () => cancelled, - }); - } catch (error) { - console.error("[Session] Failed to hydrate workspaces:", error); - } - })(); - - return () => { - cancelled = true; - }; - }, [client, hydrateWorkspaces, isConnected]); - const applyWorkspaceSetupProgress = useCallback( (payload: WorkspaceSetupProgressPayload) => { upsertWorkspaceSetupProgress({ serverId, payload }); @@ -983,19 +645,6 @@ function SessionProviderInternal({ children, serverId, client }: SessionProvider return next; }); - if (payload.agent && shouldApplyTimelineAgentSnapshot(serverId, agentId)) { - const normalized = normalizeAgentSnapshot(payload.agent, serverId); - applyAuthoritativeAgentSnapshot( - applyLegacyDaemonWorkspaceOwnership({ - serverId, - agent: { - ...normalized, - projectPlacement: session?.agents.get(agentId)?.projectPlacement ?? null, - }, - }), - ); - } - // Call pure reducer const result = processTimelineResponse({ payload, @@ -1048,7 +697,6 @@ function SessionProviderInternal({ children, serverId, client }: SessionProvider }); }, [ - applyAuthoritativeAgentSnapshot, clearAgentStreamHead, markAgentHistorySynchronized, recoverTimelineGap, @@ -1088,7 +736,7 @@ function SessionProviderInternal({ children, serverId, client }: SessionProvider setAgentInitializing(agentId, true); } try { - const page = await fetchAgentTimelineOnce(client, agentId, request); + const page = await getHostRuntimeStore().fetchAgentTimeline(serverId, agentId, request); if (getInitDeferred(initKey)) { refreshAgentInitializationTimeout({ key: initKey, agentId, setAgentInitializing }); } @@ -1177,16 +825,6 @@ function SessionProviderInternal({ children, serverId, client }: SessionProvider useProviderSubagentStore.getState().applyUpdate(serverId, message.payload); }); - const unsubWorkspaceUpdate = client.on("workspace_update", (message) => { - if (message.type !== "workspace_update") return; - const hydration = workspaceHydrationRef.current; - if (hydration) { - hydration.deltas.push(message.payload); - return; - } - applyWorkspaceUpdatePayload(message.payload); - }); - const unsubScriptStatusUpdate = client.on("script_status_update", (message) => { if (message.type !== "script_status_update") return; setWorkspaces(serverId, (prev) => patchWorkspaceScripts(prev, message.payload)); @@ -1448,101 +1086,6 @@ function SessionProviderInternal({ children, serverId, client }: SessionProvider voiceRuntime?.onServerSpeechStateChanged(serverId, message.payload.isSpeaking); }); - const unsubAgentDeleted = client.on("agent_deleted", (message) => { - if (message.type !== "agent_deleted") { - return; - } - const { agentId } = message.payload; - clearArchiveAgentPending({ queryClient, serverId, agentId }); - - setAgents(serverId, (prev) => { - if (!prev.has(agentId)) { - return prev; - } - const next = new Map(prev); - next.delete(agentId); - return next; - }); - - // Remove from agentLastActivity slice (top-level) - useSessionStore.setState((state) => { - if (!state.agentLastActivity.has(agentId)) { - return state; - } - const nextActivity = new Map(state.agentLastActivity); - nextActivity.delete(agentId); - return { - ...state, - agentLastActivity: nextActivity, - }; - }); - - setAgentStreamTail(serverId, (prev) => { - if (!prev.has(agentId)) { - return prev; - } - const next = new Map(prev); - next.delete(agentId); - return next; - }); - clearAgentStreamHead(serverId, agentId); - setAgentTimelineCursor(serverId, (prev) => { - if (!prev.has(agentId)) { - return prev; - } - const next = new Map(prev); - next.delete(agentId); - return next; - }); - - // Remove draft input - clearDraftInput({ - draftKey: buildDraftStoreKey({ serverId, agentId }), - }); - - setPendingPermissions(serverId, (prev) => { - let changed = false; - const next = new Map(prev); - for (const [key, pending] of prev.entries()) { - if (pending.agentId === agentId) { - next.delete(key); - changed = true; - } - } - return changed ? next : prev; - }); - - setInitializingAgents(serverId, (prev) => { - if (!prev.has(agentId)) { - return prev; - } - const next = new Map(prev); - next.delete(agentId); - return next; - }); - }); - - const unsubAgentArchived = client.on("agent_archived", (message) => { - if (message.type !== "agent_archived") { - return; - } - const { agentId, archivedAt } = message.payload; - clearArchiveAgentPending({ queryClient, serverId, agentId }); - - setAgents(serverId, (prev) => { - const existing = prev.get(agentId); - if (!existing) { - return prev; - } - const next = new Map(prev); - next.set(agentId, { - ...existing, - archivedAt: new Date(archivedAt), - }); - return next; - }); - }); - const unsubTerminalAttention = client.on("terminal_attention_required", (message) => { if (message.type !== "terminal_attention_required") { return; @@ -1569,7 +1112,6 @@ function SessionProviderInternal({ children, serverId, client }: SessionProvider unsubAgentTimeline(); unsubProviderSubagentUpdate(); unsubAgentAttention(); - unsubWorkspaceUpdate(); unsubScriptStatusUpdate(); unsubCheckoutStatusUpdate(); unsubWorkspaceSetupProgress(); @@ -1582,8 +1124,6 @@ function SessionProviderInternal({ children, serverId, client }: SessionProvider unsubChunk(); unsubTranscription(); unsubVoiceInputState(); - unsubAgentDeleted(); - unsubAgentArchived(); unsubTerminalAttention(); agentStreamReducerQueue.dispose({ flush: true }); }; @@ -1602,18 +1142,9 @@ function SessionProviderInternal({ children, serverId, client }: SessionProvider setInitializingAgents, setAgents, setWorkspaces, - mergeWorkspaces, - removeWorkspace, - removeWorkspaceSetup, - addEmptyProject, - removeEmptyProject, - setAgentLastActivity, setPendingPermissions, - setHasHydratedAgents, - clearDraftInput, notifyAgentAttention, recoverTimelineGap, - applyWorkspaceUpdatePayload, applyWorkspaceSetupProgress, applyTimelineResponse, updateSessionServerInfo, diff --git a/packages/app/src/contexts/session-resume-revalidation.test.ts b/packages/app/src/contexts/session-resume-revalidation.test.ts index 80b3fd25f6..ebc90b68a2 100644 --- a/packages/app/src/contexts/session-resume-revalidation.test.ts +++ b/packages/app/src/contexts/session-resume-revalidation.test.ts @@ -12,14 +12,11 @@ describe("session resume revalidation", () => { awayMs: SESSION_STALE_AFTER_MS, serverId: "server", bumpHistorySyncGeneration: (serverId) => calls.push(`history:${serverId}`), - refreshAgentDirectory: async () => calls.push("agents"), - hydrateWorkspaces: async () => { - calls.push("workspaces"); - }, + refreshDirectories: async () => calls.push("directories"), }); expect(revalidated).toBe(true); - expect(calls).toEqual(["history:server", "agents", "workspaces"]); + expect(calls).toEqual(["history:server", "directories"]); }); it("does nothing after a brief background interval", async () => { @@ -29,10 +26,7 @@ describe("session resume revalidation", () => { awayMs: SESSION_STALE_AFTER_MS - 1, serverId: "server", bumpHistorySyncGeneration: () => calls.push("history"), - refreshAgentDirectory: async () => calls.push("agents"), - hydrateWorkspaces: async () => { - calls.push("workspaces"); - }, + refreshDirectories: async () => calls.push("directories"), }); expect(revalidated).toBe(false); diff --git a/packages/app/src/contexts/session-resume-revalidation.ts b/packages/app/src/contexts/session-resume-revalidation.ts index 23aa803012..8360316f5a 100644 --- a/packages/app/src/contexts/session-resume-revalidation.ts +++ b/packages/app/src/contexts/session-resume-revalidation.ts @@ -4,14 +4,13 @@ export async function revalidateSessionAfterResume(input: { awayMs: number; serverId: string; bumpHistorySyncGeneration: (serverId: string) => void; - refreshAgentDirectory: () => Promise; - hydrateWorkspaces: () => Promise; + refreshDirectories: () => Promise; }): Promise { if (input.awayMs < SESSION_STALE_AFTER_MS) { return false; } input.bumpHistorySyncGeneration(input.serverId); - await Promise.all([input.refreshAgentDirectory(), input.hydrateWorkspaces()]); + await input.refreshDirectories(); return true; } diff --git a/packages/app/src/diagnostics/app-diagnostic-report.test.ts b/packages/app/src/diagnostics/app-diagnostic-report.test.ts index c01244d967..199ca22a69 100644 --- a/packages/app/src/diagnostics/app-diagnostic-report.test.ts +++ b/packages/app/src/diagnostics/app-diagnostic-report.test.ts @@ -79,6 +79,7 @@ describe("app diagnostics report", () => { ["relay:relay.secret.test:443", { status: "available", latencyMs: 8 }], ]), clientGeneration: 1, + connectionEpoch: 1, }; const report = formatHostRuntimeSection({ host, snapshot }); diff --git a/packages/app/src/hooks/use-agent-initialization.ts b/packages/app/src/hooks/use-agent-initialization.ts index af357deed1..c8409b130f 100644 --- a/packages/app/src/hooks/use-agent-initialization.ts +++ b/packages/app/src/hooks/use-agent-initialization.ts @@ -10,7 +10,7 @@ import { rejectInitDeferred, refreshInitTimeout, } from "@/utils/agent-initialization"; -import { fetchAgentTimelineOnce } from "@/timeline/fetch-agent-timeline-once"; +import { getHostRuntimeStore } from "@/runtime/host-runtime"; import { planInitialAgentTimelineSync, planTimelineTailFetch } from "@/timeline/timeline-sync-plan"; import { i18n } from "@/i18n/i18next"; @@ -69,10 +69,12 @@ export function ensureAgentIsInitialized(input: EnsureAgentIsInitializedInput): return deferred.promise; } - fetchAgentTimelineOnce(client, agentId, timelineRequest).catch((error) => { - setAgentInitializing(agentId, false); - rejectInitDeferred(key, error instanceof Error ? error : new Error(String(error))); - }); + getHostRuntimeStore() + .fetchAgentTimeline(serverId, agentId, timelineRequest) + .catch((error) => { + setAgentInitializing(agentId, false); + rejectInitDeferred(key, error instanceof Error ? error : new Error(String(error))); + }); return deferred.promise; } @@ -93,7 +95,11 @@ export async function refreshAgent(input: RefreshAgentInput): Promise { try { await client.refreshAgent(agentId); - await client.fetchAgentTimeline(agentId, planTimelineTailFetch()); + const serverId = Object.entries(useSessionStore.getState().sessions).find( + ([, session]) => session.client === client, + )?.[0]; + if (!serverId) throw new Error("Agent session is no longer connected"); + await getHostRuntimeStore().fetchAgentTimeline(serverId, agentId, planTimelineTailFetch()); } catch (error) { setAgentInitializing(agentId, false); throw error; diff --git a/packages/app/src/hooks/use-load-older-agent-history.ts b/packages/app/src/hooks/use-load-older-agent-history.ts index 775f5f23a3..c2cf7c02c3 100644 --- a/packages/app/src/hooks/use-load-older-agent-history.ts +++ b/packages/app/src/hooks/use-load-older-agent-history.ts @@ -4,6 +4,7 @@ import type { ToastApi } from "@/components/toast-host"; import { i18n } from "@/i18n/i18next"; import { useSessionStore, type AgentTimelineCursorState } from "@/stores/session-store"; import { planTimelineOlderFetch } from "@/timeline/timeline-sync-plan"; +import { getHostRuntimeStore } from "@/runtime/host-runtime"; export interface LoadOlderAgentHistoryClient { fetchAgentTimeline: ( @@ -97,7 +98,12 @@ export function useLoadOlderAgentHistory({ const loadOlder = useCallback(() => { const session = useSessionStore.getState().sessions[serverId]; void loadOlderAgentHistory(agentId, { - client: (session?.client ?? null) as LoadOlderAgentHistoryClient | null, + client: session?.client + ? { + fetchAgentTimeline: (timelineAgentId, request) => + getHostRuntimeStore().fetchAgentTimeline(serverId, timelineAgentId, request), + } + : null, cursor: session?.agentTimelineCursor.get(agentId), hasOlder: session?.agentTimelineHasOlder.get(agentId) === true, isLoadingOlder: session?.agentTimelineOlderFetchInFlight.get(agentId) === true, diff --git a/packages/app/src/runtime/directory-sync/agent-replica.test.ts b/packages/app/src/runtime/directory-sync/agent-replica.test.ts new file mode 100644 index 0000000000..e0b22119da --- /dev/null +++ b/packages/app/src/runtime/directory-sync/agent-replica.test.ts @@ -0,0 +1,79 @@ +import { describe, expect, it } from "vitest"; +import type { DaemonClient, FetchAgentsEntry } from "@getpaseo/client/internal/daemon-client"; +import type { AgentSnapshotPayload } from "@getpaseo/protocol/messages"; +import { useSessionStore } from "@/stores/session-store"; +import { AgentDirectoryReplica } from "./agent-replica"; + +function payload(title: string): AgentSnapshotPayload { + return { + id: "agent", + provider: "codex", + cwd: "/repo", + model: null, + createdAt: "2026-07-17T00:00:00.000Z", + updatedAt: "2026-07-17T00:01:00.000Z", + lastUserMessageAt: null, + status: "idle", + capabilities: { + supportsStreaming: true, + supportsSessionPersistence: true, + supportsDynamicModes: true, + supportsMcpServers: true, + supportsReasoningStream: true, + supportsToolInvocations: true, + }, + currentModeId: null, + availableModes: [], + pendingPermissions: [], + persistence: null, + title, + labels: {}, + }; +} + +function entry(agent: AgentSnapshotPayload): FetchAgentsEntry { + return { + agent, + project: { + projectKey: "/repo", + projectName: "repo", + checkout: { + cwd: "/repo", + isGit: false, + currentBranch: null, + remoteUrl: null, + worktreeRoot: null, + isPaseoOwnedWorktree: false, + mainRepoRoot: null, + }, + }, + }; +} + +describe("AgentDirectoryReplica", () => { + it("keeps membership authoritative across remove, stale timeline, and re-add", () => { + const serverId = "agent-replica"; + const store = useSessionStore.getState(); + store.initializeSession(serverId, null as unknown as DaemonClient); + const replica = new AgentDirectoryReplica(serverId, () => undefined); + replica.commitSnapshot([entry(payload("directory"))], []); + const staleToken = replica.captureTimeline("agent"); + + replica.remove("agent"); + expect(replica.submitTimelineAgent(staleToken, payload("stale"))).toBe(false); + expect(useSessionStore.getState().sessions[serverId]?.agents.has("agent")).toBe(false); + + replica.applyDelta({ + kind: "upsert", + agent: payload("re-added"), + project: entry(payload("x")).project, + }); + expect(replica.submitTimelineAgent(staleToken, payload("still stale"))).toBe(false); + const currentToken = replica.captureTimeline("agent"); + expect(replica.submitTimelineAgent(currentToken, payload("current"))).toBe(true); + expect(useSessionStore.getState().sessions[serverId]?.agents.get("agent")?.title).toBe( + "current", + ); + store.clearSession(serverId); + }); +}); diff --git a/packages/app/src/runtime/directory-sync/agent-replica.ts b/packages/app/src/runtime/directory-sync/agent-replica.ts new file mode 100644 index 0000000000..e1bcab4c19 --- /dev/null +++ b/packages/app/src/runtime/directory-sync/agent-replica.ts @@ -0,0 +1,116 @@ +import type { FetchAgentsEntry } from "@getpaseo/client/internal/daemon-client"; +import type { AgentSnapshotPayload } from "@getpaseo/protocol/messages"; +import { clearArchiveAgentPending } from "@/hooks/use-archive-agent"; +import { queryClient } from "@/data/query-client"; +import { useSessionStore, type Agent } from "@/stores/session-store"; +import { normalizeAgentSnapshot } from "@/utils/agent-snapshots"; +import { + applyAgentDirectoryDelta, + type AgentDirectoryDelta, + removeAgentDirectoryReplica, + replaceAgentPendingPermissions, + replaceFetchedAgentDirectory, + upsertAgentReplica, +} from "@/utils/agent-directory-sync"; +import { reconcileAgentDirectory } from "@/utils/agent-directory-reconciliation"; +import { applyLegacyDaemonWorkspaceOwnership } from "@/workspace/legacy-daemon-workspaces"; + +export interface AgentLifecycleToken { + readonly agentId: string; + readonly version: number; +} + +export class AgentDirectoryReplica { + private readonly lifecycleVersions = new Map(); + private readonly members = new Set(); + + constructor( + private readonly serverId: string, + private readonly onStoppedRunning: (agentId: string) => void, + ) {} + + captureTimeline(agentId: string): AgentLifecycleToken { + return { agentId, version: this.lifecycleVersions.get(agentId) ?? 0 }; + } + + submitTimelineAgent(token: AgentLifecycleToken, payload: AgentSnapshotPayload): boolean { + if ( + !this.members.has(token.agentId) || + token.version !== (this.lifecycleVersions.get(token.agentId) ?? 0) + ) { + return false; + } + const normalized = applyLegacyDaemonWorkspaceOwnership({ + serverId: this.serverId, + agent: normalizeAgentSnapshot(payload, this.serverId), + }); + const accepted = upsertAgentReplica(this.serverId, normalized); + replaceAgentPendingPermissions(this.serverId, accepted); + useSessionStore.getState().setAgentLastActivity(accepted.id, accepted.lastActivityAt); + if (accepted.archivedAt) { + clearArchiveAgentPending({ queryClient, serverId: this.serverId, agentId: accepted.id }); + } + return true; + } + + applyDelta(delta: AgentDirectoryDelta): void { + const before = this.members.has(delta.kind === "remove" ? delta.agentId : delta.agent.id); + const result = applyAgentDirectoryDelta({ serverId: this.serverId, delta }); + if (delta.kind === "remove") { + this.members.delete(delta.agentId); + this.advance(delta.agentId); + } else { + this.members.add(delta.agent.id); + if (!before) this.advance(delta.agent.id); + } + if (result.stoppedRunning) this.onStoppedRunning(result.agentId); + } + + commitSnapshot( + entries: FetchAgentsEntry[], + deltas: readonly AgentDirectoryDelta[], + ): Map { + const previous = useSessionStore.getState().sessions[this.serverId]?.agents ?? new Map(); + const reconciled = reconcileAgentDirectory({ previous, snapshot: entries, deltas }); + const nextIds = new Set(reconciled.entries.map((entry) => entry.agent.id)); + for (const agentId of this.members) { + if (!nextIds.has(agentId)) this.advance(agentId); + } + for (const agentId of nextIds) { + if (!this.members.has(agentId)) this.advance(agentId); + } + for (const agentId of previous.keys()) { + if (!nextIds.has(agentId)) removeAgentDirectoryReplica(this.serverId, agentId); + } + this.members.clear(); + for (const agentId of nextIds) this.members.add(agentId); + const { agents } = replaceFetchedAgentDirectory({ + serverId: this.serverId, + entries: reconciled.entries, + }); + for (const agentId of reconciled.stoppedRunningAgentIds) this.onStoppedRunning(agentId); + return agents; + } + + archive(agentId: string, archivedAt: string): void { + this.advance(agentId); + useSessionStore.getState().setAgents(this.serverId, (current) => { + const agent = current.get(agentId); + if (!agent) return current; + const next = new Map(current); + next.set(agentId, { ...agent, archivedAt: new Date(archivedAt) }); + return next; + }); + clearArchiveAgentPending({ queryClient, serverId: this.serverId, agentId }); + } + + remove(agentId: string): void { + this.members.delete(agentId); + this.advance(agentId); + removeAgentDirectoryReplica(this.serverId, agentId); + } + + private advance(agentId: string): void { + this.lifecycleVersions.set(agentId, (this.lifecycleVersions.get(agentId) ?? 0) + 1); + } +} diff --git a/packages/app/src/runtime/directory-sync/index.ts b/packages/app/src/runtime/directory-sync/index.ts new file mode 100644 index 0000000000..76658af2f8 --- /dev/null +++ b/packages/app/src/runtime/directory-sync/index.ts @@ -0,0 +1,374 @@ +import type { + DaemonClient, + FetchAgentsEntry, + FetchAgentsOptions, +} from "@getpaseo/client/internal/daemon-client"; +import { fetchAgentTimelineOnce } from "@/timeline/fetch-agent-timeline-once"; +import { + normalizeEmptyProjectDescriptor, + normalizeWorkspaceDescriptor, + useSessionStore, + type Agent, +} from "@/stores/session-store"; +import { + readLegacyDaemonWorkspaceDirectory, + buildLegacyWorkspaces, + shouldUseLegacyDaemonWorkspaceDirectory, + stampLegacyWorkspaceIds, +} from "@/workspace/legacy-daemon-workspaces"; +import type { AgentDirectoryDelta } from "@/utils/agent-directory-sync"; +import { AgentDirectoryReplica } from "./agent-replica"; +import { + WorkspaceDirectoryReplica, + type WorkspaceDirectoryDelta, + type WorkspaceDirectorySnapshot, +} from "./workspace-replica"; +import { + DirectoryTransactionOwner, + type DirectorySourceToken, + type DirectoryTransaction, +} from "./transaction"; + +const PAGE_LIMIT = 200; +const AGENT_SORT: NonNullable = [ + { key: "updated_at", direction: "desc" }, +]; + +interface AgentSnapshot { + entries: FetchAgentsEntry[]; + subscriptionId: string | null; + legacy: boolean; +} + +export interface DirectoryConnection { + client: DaemonClient | null; + status: "online" | "offline"; + source: DirectorySourceToken; +} + +export interface RefreshAgentDirectoryInput { + filter?: FetchAgentsOptions["filter"]; + subscribe?: FetchAgentsOptions["subscribe"]; + page?: FetchAgentsOptions["page"]; +} + +export interface RefreshAgentDirectoryResult { + agents: Map; + subscriptionId: string | null; +} + +export class DirectorySync { + private readonly agentTransactions = new DirectoryTransactionOwner< + AgentSnapshot, + AgentDirectoryDelta + >(); + private readonly workspaceTransactions = new DirectoryTransactionOwner< + WorkspaceDirectorySnapshot, + WorkspaceDirectoryDelta + >(); + private readonly agents: AgentDirectoryReplica; + private readonly workspaces: WorkspaceDirectoryReplica; + private connection: DirectoryConnection = { + client: null, + status: "offline", + source: { clientGeneration: 0, connectionEpoch: 0 }, + }; + private unsubscribe: (() => void) | null = null; + + constructor( + private readonly serverId: string, + private readonly callbacks: { + drainQueuedAgentMessage: (agentId: string) => void; + markAgentLoading: () => void; + markAgentReady: () => void; + markAgentError: (error: string) => void; + }, + ) { + this.agents = new AgentDirectoryReplica(serverId, callbacks.drainQueuedAgentMessage); + this.workspaces = new WorkspaceDirectoryReplica(serverId); + } + + connectionChanged(connection: DirectoryConnection): void { + const changed = + this.connection.client !== connection.client || + this.connection.source.clientGeneration !== connection.source.clientGeneration || + this.connection.source.connectionEpoch !== connection.source.connectionEpoch; + const wentOffline = this.connection.status === "online" && connection.status === "offline"; + if (!changed && !wentOffline) { + this.connection = connection; + return; + } + this.flushAbortedTransactions(); + this.unsubscribe?.(); + this.unsubscribe = null; + this.connection = connection; + if (!connection.client || connection.status !== "online") return; + const client = connection.client; + const source = connection.source; + const subscriptions = [ + client.on("agent_update", (message) => { + if (message.type !== "agent_update" || !this.isCurrent(client, source)) return; + if (!this.agentTransactions.record(source, message.payload)) + this.agents.applyDelta(message.payload); + }), + client.on("workspace_update", (message) => { + if (message.type !== "workspace_update" || !this.isCurrent(client, source)) return; + if (!this.workspaceTransactions.record(source, message.payload)) { + this.workspaces.applyDelta(message.payload); + } + }), + client.on("agent_deleted", (message) => { + if (message.type === "agent_deleted" && this.isCurrent(client, source)) { + this.agents.remove(message.payload.agentId); + } + }), + client.on("agent_archived", (message) => { + if (message.type === "agent_archived" && this.isCurrent(client, source)) { + this.agents.archive(message.payload.agentId, message.payload.archivedAt); + } + }), + ]; + this.unsubscribe = () => { + for (const unsubscribe of subscriptions) unsubscribe(); + }; + } + + dispose(): void { + this.flushAbortedTransactions(); + this.unsubscribe?.(); + this.unsubscribe = null; + } + + async fetchTimeline( + agentId: string, + request: Parameters[1], + ): Promise>> { + const { client } = this.requireOnline(); + const token = this.agents.captureTimeline(agentId); + const page = await fetchAgentTimelineOnce(client, agentId, request); + if (page.agent) this.agents.submitTimelineAgent(token, page.agent); + return page; + } + + async refreshAgents( + input: RefreshAgentDirectoryInput = {}, + ): Promise { + const { client, source } = this.requireOnline(); + const transaction = this.agentTransactions.begin(source, () => ({ + entries: [], + subscriptionId: null, + legacy: false, + })); + this.callbacks.markAgentLoading(); + try { + await this.waitForSession(client, source); + const session = useSessionStore.getState().sessions[this.serverId]; + if (!input.filter && shouldUseLegacyDaemonWorkspaceDirectory(session?.serverInfo)) { + const directory = await readLegacyDaemonWorkspaceDirectory({ + client, + subscribe: input.subscribe, + page: input.page, + }); + if ( + !directory || + !this.agentTransactions.isCurrent(transaction) || + !this.isCurrent(client, source) + ) { + throw new DirectoryRefreshSupersededError("legacy fetch no longer current"); + } + transaction.snapshot.entries.push(...stampLegacyWorkspaceIds(directory.entries)); + transaction.snapshot.subscriptionId = directory.subscriptionId; + transaction.snapshot.legacy = true; + } else { + await this.fetchAgents(client, source, transaction, input); + } + if (!this.isCurrent(client, source) || !this.hasMatchingSession(client, source)) { + throw new DirectoryRefreshSupersededError("agent completion no longer current"); + } + const completion = this.agentTransactions.complete(transaction); + if (completion.kind === "stale") { + throw new DirectoryRefreshSupersededError("agent completion was superseded"); + } + if (completion.snapshot.legacy) { + const store = useSessionStore.getState(); + store.setWorkspaces(this.serverId, buildLegacyWorkspaces(completion.snapshot.entries)); + store.setEmptyProjects(this.serverId, []); + store.setHasHydratedWorkspaces(this.serverId, true); + } + const deltas = completion.snapshot.legacy + ? completion.deltas.map((delta) => + delta.kind === "upsert" + ? { ...delta, agent: { ...delta.agent, workspaceId: delta.agent.cwd } } + : delta, + ) + : completion.deltas; + const agents = this.agents.commitSnapshot(completion.snapshot.entries, deltas); + this.callbacks.markAgentReady(); + return { agents, subscriptionId: completion.snapshot.subscriptionId }; + } catch (error) { + const deltas = this.agentTransactions.fail(transaction); + if (deltas) for (const delta of deltas) this.agents.applyDelta(delta); + if (!(error instanceof DirectoryRefreshSupersededError)) { + this.callbacks.markAgentError(error instanceof Error ? error.message : String(error)); + } + throw error; + } + } + + async refreshWorkspaces(input?: { subscribe?: boolean }): Promise { + const serverInfo = useSessionStore.getState().sessions[this.serverId]?.serverInfo; + if (serverInfo?.features?.workspaceMultiplicity !== true) return; + const { client, source } = this.requireOnline(); + const transaction = this.workspaceTransactions.begin(source, () => ({ + workspaces: new Map(), + emptyProjects: new Map(), + })); + try { + await this.fetchWorkspaceSnapshot(client, source, transaction, input?.subscribe === true); + if (!this.isCurrent(client, source) || !this.hasMatchingSession(client, source)) { + throw new DirectoryRefreshSupersededError("workspace completion no longer current"); + } + const completion = this.workspaceTransactions.complete(transaction); + if (completion.kind === "stale") { + throw new DirectoryRefreshSupersededError("workspace completion was superseded"); + } + this.workspaces.commitSnapshot(completion.snapshot, completion.deltas); + } catch (error) { + const deltas = this.workspaceTransactions.fail(transaction); + if (deltas) for (const delta of deltas) this.workspaces.applyDelta(delta); + throw error; + } + } + + private async fetchWorkspaceSnapshot( + client: DaemonClient, + source: DirectorySourceToken, + transaction: DirectoryTransaction, + initialSubscribe: boolean, + ): Promise { + let cursor: string | null = null; + let subscribe = initialSubscribe; + while (true) { + const payload = await client.fetchWorkspaces({ + sort: [{ key: "activity_at", direction: "desc" }], + ...(subscribe ? { subscribe: {} } : {}), + page: cursor ? { limit: PAGE_LIMIT, cursor } : { limit: PAGE_LIMIT }, + }); + this.assertWorkspaceTransactionCurrent(client, source, transaction); + for (const entry of payload.entries) { + const workspace = normalizeWorkspaceDescriptor(entry); + transaction.snapshot.workspaces.set(workspace.id, workspace); + } + for (const entry of payload.emptyProjects ?? []) { + const project = normalizeEmptyProjectDescriptor(entry); + transaction.snapshot.emptyProjects.set(project.projectId, project); + } + if (!payload.pageInfo.hasMore || !payload.pageInfo.nextCursor) return; + cursor = payload.pageInfo.nextCursor; + subscribe = false; + } + } + + async refreshAll(): Promise { + await Promise.all([this.refreshAgents(), this.refreshWorkspaces({ subscribe: true })]); + } + + private async fetchAgents( + client: DaemonClient, + source: DirectorySourceToken, + transaction: DirectoryTransaction, + input: RefreshAgentDirectoryInput, + ): Promise { + let cursor = input.page?.cursor ?? null; + let subscribe = input.subscribe; + while (true) { + const limit = input.page?.limit ?? PAGE_LIMIT; + const payload = await client.fetchAgents({ + ...(input.filter ? { filter: input.filter } : { scope: "active" as const }), + sort: AGENT_SORT, + ...(subscribe ? { subscribe } : {}), + page: cursor ? { limit, cursor } : { limit }, + }); + this.assertAgentTransactionCurrent(client, source, transaction); + transaction.snapshot.entries.push(...payload.entries); + transaction.snapshot.subscriptionId ??= payload.subscriptionId ?? null; + const pageInfo = payload.pageInfo as { + hasMore?: boolean; + hasMoreAfter?: boolean; + nextCursor?: string | null; + afterCursor?: string | null; + }; + const hasMore = pageInfo.hasMore ?? pageInfo.hasMoreAfter ?? false; + const nextCursor = pageInfo.nextCursor ?? pageInfo.afterCursor ?? null; + if (!hasMore || !nextCursor) break; + cursor = nextCursor; + subscribe = undefined; + } + } + + private assertAgentTransactionCurrent( + client: DaemonClient, + source: DirectorySourceToken, + transaction: DirectoryTransaction, + ): void { + if (!this.agentTransactions.isCurrent(transaction) || !this.isCurrent(client, source)) { + throw new DirectoryRefreshSupersededError("agent page no longer current"); + } + } + + private assertWorkspaceTransactionCurrent( + client: DaemonClient, + source: DirectorySourceToken, + transaction: DirectoryTransaction, + ): void { + if (!this.workspaceTransactions.isCurrent(transaction) || !this.isCurrent(client, source)) { + throw new DirectoryRefreshSupersededError("workspace fetch no longer current"); + } + } + + private requireOnline(): { client: DaemonClient; source: DirectorySourceToken } { + if (!this.connection.client || this.connection.status !== "online") { + throw new Error(`Host ${this.serverId} is not connected`); + } + return { client: this.connection.client, source: this.connection.source }; + } + + private isCurrent(client: DaemonClient, source: DirectorySourceToken): boolean { + return ( + this.connection.client === client && + this.connection.status === "online" && + this.connection.source.clientGeneration === source.clientGeneration && + this.connection.source.connectionEpoch === source.connectionEpoch + ); + } + + private async waitForSession(client: DaemonClient, source: DirectorySourceToken): Promise { + const matches = () => this.hasMatchingSession(client, source); + if (matches()) return; + await new Promise((resolve, reject) => { + const unsubscribe = useSessionStore.subscribe(() => { + if (matches()) { + unsubscribe(); + resolve(); + } else if (!this.isCurrent(client, source)) { + unsubscribe(); + reject(new DirectoryRefreshSupersededError("session wait no longer current")); + } + }); + }); + } + + private hasMatchingSession(client: DaemonClient, source: DirectorySourceToken): boolean { + const session = useSessionStore.getState().sessions[this.serverId]; + return session?.client === client && session.clientGeneration === source.clientGeneration; + } + + private flushAbortedTransactions(): void { + for (const delta of this.agentTransactions.abort()) this.agents.applyDelta(delta); + for (const delta of this.workspaceTransactions.abort()) this.workspaces.applyDelta(delta); + } +} + +export class DirectoryRefreshSupersededError extends Error {} + +export type { DirectorySourceToken } from "./transaction"; diff --git a/packages/app/src/runtime/directory-sync/transaction.test.ts b/packages/app/src/runtime/directory-sync/transaction.test.ts new file mode 100644 index 0000000000..bc80027041 --- /dev/null +++ b/packages/app/src/runtime/directory-sync/transaction.test.ts @@ -0,0 +1,100 @@ +import { describe, expect, it } from "vitest"; +import { DirectoryTransactionOwner, type DirectorySourceToken } from "./transaction"; + +const epoch = (connectionEpoch: number): DirectorySourceToken => ({ + clientGeneration: 1, + connectionEpoch, +}); + +type Event = "begin-new" | "first" | "second" | "finish-old" | "finish-new"; + +function permutations(events: Event[]): Event[][] { + if (events.length === 0) return [[]]; + return events.flatMap((event, index) => + permutations(events.filter((_, candidate) => candidate !== index)).map((tail) => + [event].concat(tail), + ), + ); +} + +describe("DirectoryTransactionOwner", () => { + it("applies every bounded two-refresh interleaving exactly once", () => { + const schedules = permutations([ + "begin-new", + "first", + "second", + "finish-old", + "finish-new", + ]).filter((events) => events.indexOf("begin-new") < events.indexOf("finish-new")); + expect(schedules).toHaveLength(60); + for (const events of schedules) { + const owner = new DirectoryTransactionOwner(); + const old = owner.begin(epoch(1), () => null); + let current = old; + const applied: string[] = []; + for (const event of events) { + if (event === "begin-new") current = owner.begin(epoch(1), () => null); + if (event === "first" || event === "second") { + if (!owner.record(epoch(1), event)) applied.push(event); + } + if (event === "finish-old") { + const result = owner.complete(old); + if (result.kind === "current") applied.push(...result.deltas); + } + if (event === "finish-new") { + const result = owner.complete(current); + if (result.kind === "current") applied.push(...result.deltas); + } + } + expect(applied.sort(), JSON.stringify(events)).toEqual(["first", "second"]); + } + }); + + it.each([ + ["success", false, false, ["delta"]], + ["failure", true, false, ["delta"]], + ["superseded success", false, true, ["delta"]], + ["superseded failure", true, true, ["delta"]], + ] as const)("preserves buffered deltas through %s", (_name, fail, supersede, expected) => { + const owner = new DirectoryTransactionOwner(); + const first = owner.begin(epoch(1), () => []); + owner.record(epoch(1), "delta"); + const current = supersede ? owner.begin(epoch(1), () => []) : first; + + let deltas: readonly string[] | null; + if (fail) { + deltas = owner.fail(current); + } else { + const completion = owner.complete(current); + deltas = completion.kind === "current" ? completion.deltas : null; + } + + expect(deltas).toEqual(expected); + if (supersede) expect(owner.complete(first)).toEqual({ kind: "stale" }); + }); + + it("makes a stale completion inert", () => { + const owner = new DirectoryTransactionOwner(); + const first = owner.begin(epoch(1), () => null); + const second = owner.begin(epoch(1), () => null); + expect(owner.complete(first)).toEqual({ kind: "stale" }); + expect(owner.isCurrent(second)).toBe(true); + }); + + it("isolates reconnect epochs", () => { + const owner = new DirectoryTransactionOwner(); + const disconnected = owner.begin(epoch(1), () => null); + owner.record(epoch(1), "old"); + + expect(owner.abort()).toEqual(["old"]); + const reconnected = owner.begin(epoch(2), () => null); + owner.record(epoch(2), "new"); + + expect(owner.complete(disconnected)).toEqual({ kind: "stale" }); + expect(owner.complete(reconnected)).toEqual({ + kind: "current", + snapshot: null, + deltas: ["new"], + }); + }); +}); diff --git a/packages/app/src/runtime/directory-sync/transaction.ts b/packages/app/src/runtime/directory-sync/transaction.ts new file mode 100644 index 0000000000..15851491c6 --- /dev/null +++ b/packages/app/src/runtime/directory-sync/transaction.ts @@ -0,0 +1,70 @@ +export interface DirectorySourceToken { + clientGeneration: number; + connectionEpoch: number; +} + +export interface DirectoryTransaction { + readonly source: DirectorySourceToken; + readonly snapshot: TSnapshot; + readonly deltas: TDelta[]; +} + +export type DirectoryTransactionCompletion = + | { kind: "stale" } + | { kind: "current"; snapshot: TSnapshot; deltas: readonly TDelta[] }; + +function isSameSource(left: DirectorySourceToken, right: DirectorySourceToken): boolean { + return ( + left.clientGeneration === right.clientGeneration && + left.connectionEpoch === right.connectionEpoch + ); +} + +/** Owns the shared begin/record/supersede/complete lifecycle for a directory replica. */ +export class DirectoryTransactionOwner { + private current: DirectoryTransaction | null = null; + + begin( + source: DirectorySourceToken, + createSnapshot: () => TSnapshot, + ): DirectoryTransaction { + const transaction: DirectoryTransaction = { + source, + snapshot: createSnapshot(), + deltas: + this.current && isSameSource(this.current.source, source) ? [...this.current.deltas] : [], + }; + this.current = transaction; + return transaction; + } + + record(source: DirectorySourceToken, delta: TDelta): boolean { + if (!this.current || !isSameSource(this.current.source, source)) return false; + this.current.deltas.push(delta); + return true; + } + + isCurrent(transaction: DirectoryTransaction): boolean { + return this.current === transaction; + } + + complete( + transaction: DirectoryTransaction, + ): DirectoryTransactionCompletion { + if (this.current !== transaction) return { kind: "stale" }; + this.current = null; + return { kind: "current", snapshot: transaction.snapshot, deltas: transaction.deltas }; + } + + fail(transaction: DirectoryTransaction): readonly TDelta[] | null { + const completion = this.complete(transaction); + return completion.kind === "current" ? completion.deltas : null; + } + + abort(): readonly TDelta[] { + if (!this.current) return []; + const deltas = [...this.current.deltas]; + this.current = null; + return deltas; + } +} diff --git a/packages/app/src/runtime/directory-sync/workspace-replica.test.ts b/packages/app/src/runtime/directory-sync/workspace-replica.test.ts new file mode 100644 index 0000000000..48b55c0024 --- /dev/null +++ b/packages/app/src/runtime/directory-sync/workspace-replica.test.ts @@ -0,0 +1,57 @@ +import { expect, it } from "vitest"; +import type { DaemonClient } from "@getpaseo/client/internal/daemon-client"; +import type { WorkspaceDescriptorPayload } from "@getpaseo/protocol/messages"; +import { + normalizeEmptyProjectDescriptor, + normalizeWorkspaceDescriptor, + useSessionStore, +} from "@/stores/session-store"; +import { WorkspaceDirectoryReplica } from "./workspace-replica"; + +function workspace(id: string, projectId = "project"): WorkspaceDescriptorPayload { + return { + id, + projectId, + projectDisplayName: projectId, + projectRootPath: `/repo/${projectId}`, + workspaceDirectory: `/repo/${projectId}/${id}`, + projectKind: "git", + workspaceKind: "worktree", + name: id, + title: id, + status: "done", + activityAt: null, + statusEnteredAt: null, + archivingAt: null, + diffStat: null, + scripts: [], + }; +} + +it("commits workspace and project-parent state with filtered removals", () => { + const serverId = "workspace-replica"; + const store = useSessionStore.getState(); + store.initializeSession(serverId, null as unknown as DaemonClient); + const replica = new WorkspaceDirectoryReplica(serverId); + const empty = normalizeEmptyProjectDescriptor({ + projectId: "empty", + projectDisplayName: "Empty", + projectRootPath: "/repo/empty", + projectKind: "git", + }); + replica.commitSnapshot( + { + workspaces: new Map([ + ["kept", normalizeWorkspaceDescriptor(workspace("kept"))], + ["filtered", normalizeWorkspaceDescriptor(workspace("filtered", "filtered-project"))], + ]), + emptyProjects: new Map([[empty.projectId, empty]]), + }, + [{ kind: "remove", id: "filtered", removedProjectId: "filtered-project" }], + ); + + const session = useSessionStore.getState().sessions[serverId]; + expect(Array.from(session?.workspaces.keys() ?? [])).toEqual(["kept"]); + expect(Array.from(session?.emptyProjects.keys() ?? [])).toEqual(["empty"]); + store.clearSession(serverId); +}); diff --git a/packages/app/src/runtime/directory-sync/workspace-replica.ts b/packages/app/src/runtime/directory-sync/workspace-replica.ts new file mode 100644 index 0000000000..f2b69c8533 --- /dev/null +++ b/packages/app/src/runtime/directory-sync/workspace-replica.ts @@ -0,0 +1,93 @@ +import type { SessionOutboundMessage } from "@getpaseo/protocol/messages"; +import { + normalizeEmptyProjectDescriptor, + normalizeWorkspaceDescriptor, + useSessionStore, + type EmptyProjectDescriptor, + type WorkspaceDescriptor, +} from "@/stores/session-store"; +import { useWorkspaceSetupStore } from "@/stores/workspace-setup-store"; +import { + clearWorkspaceArchivePending, + shouldSuppressWorkspaceForLocalArchive, +} from "@/contexts/session-workspace-upserts"; + +export type WorkspaceDirectoryDelta = Extract< + SessionOutboundMessage, + { type: "workspace_update" } +>["payload"]; + +export interface WorkspaceDirectorySnapshot { + workspaces: Map; + emptyProjects: Map; +} + +export class WorkspaceDirectoryReplica { + constructor(private readonly serverId: string) {} + + applyDelta(delta: WorkspaceDirectoryDelta): void { + const state = this.reconcile(this.read(), [delta]); + this.commit(state, delta.kind === "remove" ? [delta.id] : []); + } + + commitSnapshot( + snapshot: WorkspaceDirectorySnapshot, + deltas: readonly WorkspaceDirectoryDelta[], + ): void { + const removedWorkspaceIds = deltas.flatMap((delta) => + delta.kind === "remove" ? [delta.id] : [], + ); + this.commit(this.reconcile(snapshot, deltas), removedWorkspaceIds); + } + + private read(): WorkspaceDirectorySnapshot { + const session = useSessionStore.getState().sessions[this.serverId]; + return { + workspaces: new Map(session?.workspaces), + emptyProjects: new Map(session?.emptyProjects), + }; + } + + private reconcile( + snapshot: WorkspaceDirectorySnapshot, + deltas: readonly WorkspaceDirectoryDelta[], + ): WorkspaceDirectorySnapshot { + const workspaces = new Map(snapshot.workspaces); + const emptyProjects = new Map(snapshot.emptyProjects); + for (const [workspaceId, workspace] of workspaces) { + if (shouldSuppressWorkspaceForLocalArchive({ serverId: this.serverId, workspace })) { + workspaces.delete(workspaceId); + } + } + for (const delta of deltas) { + if (delta.kind === "remove") { + workspaces.delete(delta.id); + if (delta.emptyProject) { + const project = normalizeEmptyProjectDescriptor(delta.emptyProject); + emptyProjects.set(project.projectId, project); + } + if (delta.removedProjectId) emptyProjects.delete(delta.removedProjectId); + continue; + } + const workspace = normalizeWorkspaceDescriptor(delta.workspace); + if (shouldSuppressWorkspaceForLocalArchive({ serverId: this.serverId, workspace })) { + workspaces.delete(workspace.id); + } else { + workspaces.set(workspace.id, workspace); + emptyProjects.delete(workspace.projectId); + } + } + return { workspaces, emptyProjects }; + } + + private commit(snapshot: WorkspaceDirectorySnapshot, removedWorkspaceIds: string[]): void { + const store = useSessionStore.getState(); + store.setWorkspaces(this.serverId, snapshot.workspaces); + store.setEmptyProjects(this.serverId, snapshot.emptyProjects.values()); + store.setHasHydratedWorkspaces(this.serverId, true); + for (const workspaceId of removedWorkspaceIds) { + clearWorkspaceArchivePending({ serverId: this.serverId, workspaceId }); + useWorkspaceSetupStore.getState().removeWorkspace({ serverId: this.serverId, workspaceId }); + } + } +} diff --git a/packages/app/src/runtime/host-runtime.test.ts b/packages/app/src/runtime/host-runtime.test.ts index b01bd5fda0..e3141fb0a6 100644 --- a/packages/app/src/runtime/host-runtime.test.ts +++ b/packages/app/src/runtime/host-runtime.test.ts @@ -35,6 +35,7 @@ class FakeDaemonClient { > = []; public sentAgentMessages: Array> = []; public sendAgentMessageFailures: Error[] = []; + public sendAgentMessageResponses: Promise[] = []; private agentUpdateListeners = new Set< (message: Extract) => void >(); @@ -81,6 +82,8 @@ class FakeDaemonClient { async sendAgentMessage(...args: Parameters): Promise { this.sentAgentMessages.push(args); for (const waiter of this.sentMessageWaiters) waiter(); + const response = this.sendAgentMessageResponses.shift(); + if (response) await response; const failure = this.sendAgentMessageFailures.shift(); if (failure) throw failure; } @@ -630,7 +633,6 @@ describe("HostRuntimeController", () => { }); await controller.start({ autoProbe: false }); - const snapshot = controller.getSnapshot(); expect(snapshot.activeConnectionId).toBe("direct:lan:6767"); expect(snapshot.connectionStatus).toBe("online"); @@ -1111,6 +1113,7 @@ describe("HostRuntimeController", () => { }); await controller.start({ autoProbe: false }); + const firstEpoch = controller.getSnapshot().connectionEpoch; controller.markAgentDirectorySyncReady(); expect(controller.getSnapshot().agentDirectoryStatus).toBe("ready"); expect(controller.getSnapshot().hasEverLoadedAgentDirectory).toBe(true); @@ -1125,6 +1128,7 @@ describe("HostRuntimeController", () => { clients[0]?.setConnectionState({ status: "connected" }); expect(controller.getSnapshot().connectionStatus).toBe("online"); expect(controller.getSnapshot().agentDirectoryStatus).toBe("ready"); + expect(controller.getSnapshot().connectionEpoch).toBe(firstEpoch + 1); }); it("stores directory sync errors as non-blocking after a successful directory load", async () => { @@ -1577,7 +1581,10 @@ describe("HostRuntimeStore", () => { }); it("drains legacy snapshot and buffered running transitions exactly once", async () => { - const host = makeHost({ serverId: "srv_legacy_transitions" }); + const host = makeHost({ + serverId: "srv_legacy_transitions", + connections: [{ id: "direct:lan:6767", type: "directTcp", endpoint: "lan:6767" }], + }); const fakeClient = new FakeDaemonClient(); fakeClient.setConnectionState({ status: "connected" }); const pageTwo = new Deferred>>(); @@ -1757,7 +1764,10 @@ describe("HostRuntimeStore", () => { }); it("replays agent updates received while a later bootstrap page is loading", async () => { - const host = makeHost({ serverId: "srv_paged_delta" }); + const host = makeHost({ + serverId: "srv_paged_delta", + connections: [{ id: "direct:lan:6767", type: "directTcp", endpoint: "lan:6767" }], + }); const fakeClient = new FakeDaemonClient(); fakeClient.setConnectionState({ status: "connected" }); let finishPageTwo!: (payload: Awaited>) => void; @@ -1909,7 +1919,10 @@ describe("HostRuntimeStore", () => { }); it("rejects a superseded refresh without overwriting the newer replica", async () => { - const host = makeHost({ serverId: "srv_overlap" }); + const host = makeHost({ + serverId: "srv_overlap", + connections: [{ id: "direct:lan:6767", type: "directTcp", endpoint: "lan:6767" }], + }); const fakeClient = new FakeDaemonClient(); fakeClient.setConnectionState({ status: "connected" }); fakeClient.fetchAgentsResponses.push(makeFetchAgentsPayload({ entries: [] })); @@ -1974,7 +1987,10 @@ describe("HostRuntimeStore", () => { }); it("replays inherited deltas when a superseding refresh fails", async () => { - const host = makeHost({ serverId: "srv_overlap_failure" }); + const host = makeHost({ + serverId: "srv_overlap_failure", + connections: [{ id: "direct:lan:6767", type: "directTcp", endpoint: "lan:6767" }], + }); const fakeClient = new FakeDaemonClient(); fakeClient.setConnectionState({ status: "connected" }); fakeClient.fetchAgentsResponses.push(makeFetchAgentsPayload({ entries: [] })); @@ -2026,7 +2042,10 @@ describe("HostRuntimeStore", () => { }); it("rejects a refresh when the session generation changes before commit", async () => { - const host = makeHost({ serverId: "srv_stale_generation" }); + const host = makeHost({ + serverId: "srv_stale_generation", + connections: [{ id: "direct:lan:6767", type: "directTcp", endpoint: "lan:6767" }], + }); const fakeClient = new FakeDaemonClient(); fakeClient.setConnectionState({ status: "connected" }); const existingEntry = makeFetchAgentsEntry({ @@ -2079,7 +2098,10 @@ describe("HostRuntimeStore", () => { }); it("drains queued messages once for snapshot and buffered running transitions", async () => { - const host = makeHost({ serverId: "srv_queued_transitions" }); + const host = makeHost({ + serverId: "srv_queued_transitions", + connections: [{ id: "direct:lan:6767", type: "directTcp", endpoint: "lan:6767" }], + }); const fakeClient = new FakeDaemonClient(); fakeClient.setConnectionState({ status: "connected" }); const pageTwo = new Deferred>>(); @@ -2211,6 +2233,43 @@ describe("HostRuntimeStore", () => { useSessionStore.getState().clearSession(host.serverId); }); + it("serializes queued-message drains for the same agent", async () => { + const host = makeHost({ serverId: "srv_serialized_queue_drain" }); + const fakeClient = new FakeDaemonClient(); + const send = new Deferred(); + fakeClient.sendAgentMessageResponses.push(send.promise); + const store = new HostRuntimeStore({ + deps: { + createClient: () => fakeClient as unknown as DaemonClient, + connectToDaemon: async () => ({ + client: fakeClient as unknown as DaemonClient, + serverId: host.serverId, + hostname: null, + }), + getClientId: async () => "cid_serialized_queue_drain", + }, + }); + const sessionStore = useSessionStore.getState(); + sessionStore.initializeSession(host.serverId, fakeClient as unknown as DaemonClient, 1); + sessionStore.setQueuedMessages( + host.serverId, + new Map([["agent", [{ id: "first", text: "send once", attachments: [] }]]]), + ); + + store.drainQueuedAgentMessage(host.serverId, "agent"); + store.drainQueuedAgentMessage(host.serverId, "agent"); + await fakeClient.waitForSentMessages(1); + expect(fakeClient.sentAgentMessages).toHaveLength(1); + + send.resolve(); + await vi.waitFor(() => { + expect( + useSessionStore.getState().sessions[host.serverId]?.queuedMessages.get("agent"), + ).toEqual([]); + }); + useSessionStore.getState().clearSession(host.serverId); + }); + it("uses legacy GitHub attachments when draining a queue for an old daemon", async () => { const host = makeHost({ serverId: "srv_legacy_queue_attachment" }); const fakeClient = new FakeDaemonClient(); @@ -2283,7 +2342,10 @@ describe("HostRuntimeStore", () => { }); it("applies buffered stale side effects from the accepted page agent", async () => { - const host = makeHost({ serverId: "srv_buffered_stale_side_effects" }); + const host = makeHost({ + serverId: "srv_buffered_stale_side_effects", + connections: [{ id: "direct:lan:6767", type: "directTcp", endpoint: "lan:6767" }], + }); const fakeClient = new FakeDaemonClient(); fakeClient.setConnectionState({ status: "connected" }); const pageTwo = new Deferred>>(); diff --git a/packages/app/src/runtime/host-runtime.ts b/packages/app/src/runtime/host-runtime.ts index 70d39b7824..d300dc9687 100644 --- a/packages/app/src/runtime/host-runtime.ts +++ b/packages/app/src/runtime/host-runtime.ts @@ -4,7 +4,6 @@ import equal from "fast-deep-equal/es6"; import { DaemonClient, type ConnectionState, - type FetchAgentsEntry, type FetchAgentsOptions, } from "@getpaseo/client/internal/daemon-client"; import { @@ -40,19 +39,7 @@ import { import { getDesktopHost } from "@/desktop/host"; import { CLIENT_CAPS } from "@getpaseo/protocol/client-capabilities"; import { BROWSER_AUTOMATION_COMMAND_NAMES } from "@getpaseo/protocol/browser-automation/rpc-schemas"; -import { - applyAgentDirectoryDelta, - replaceFetchedAgentDirectory, - type AgentDirectoryDelta, -} from "@/utils/agent-directory-sync"; -import { reconcileAgentDirectory } from "@/utils/agent-directory-reconciliation"; -import { useSessionStore, type Agent } from "@/stores/session-store"; -import { - readLegacyDaemonWorkspaceDirectory, - replaceLegacyDaemonWorkspaceDirectory, - shouldUseLegacyDaemonWorkspaceDirectory, - stampLegacyWorkspaceIds, -} from "@/workspace/legacy-daemon-workspaces"; +import { useSessionStore } from "@/stores/session-store"; import { invalidateCheckoutGitQueriesForServer } from "@/git/query-keys"; import { queryClient } from "@/data/query-client"; import { @@ -67,7 +54,7 @@ import { splitComposerAttachmentsForSubmit, } from "@/composer/attachments/submit"; import { encodeImages } from "@/utils/encode-images"; -import { inheritBufferedDirectoryDeltas } from "@/utils/buffered-directory-transaction"; +import { DirectorySync, type RefreshAgentDirectoryResult } from "@/runtime/directory-sync"; export type HostRuntimeConnectionStatus = "idle" | "connecting" | "online" | "offline" | "error"; export type HostRegistryStatus = "loading" | "ready"; @@ -99,6 +86,7 @@ export interface HostRuntimeSnapshot { hasEverLoadedAgentDirectory: boolean; probeByConnectionId: Map; clientGeneration: number; + connectionEpoch: number; } type HostRuntimeSnapshotPatch = Partial>; @@ -190,97 +178,6 @@ const ADAPTIVE_SWITCH_CONSECUTIVE_PROBES = 3; const DEFAULT_AGENT_DIRECTORY_PAGE_LIMIT = 200; const CONFIGURED_OVERRIDE_BOOTSTRAP_RETRY_MS = 1_000; -const DEFAULT_AGENT_DIRECTORY_SORT: NonNullable = [ - { key: "updated_at", direction: "desc" }, -]; - -function readFetchAgentsHasMore( - pageInfo: Awaited>["pageInfo"], -): boolean { - const page = pageInfo as { - hasMore?: boolean; - hasMoreAfter?: boolean; - }; - if (typeof page.hasMore === "boolean") { - return page.hasMore; - } - if (typeof page.hasMoreAfter === "boolean") { - return page.hasMoreAfter; - } - return false; -} - -function readFetchAgentsNextCursor( - pageInfo: Awaited>["pageInfo"], -): string | null { - const page = pageInfo as { - nextCursor?: string | null; - afterCursor?: string | null; - }; - if (typeof page.nextCursor === "string" && page.nextCursor.length > 0) { - return page.nextCursor; - } - if (typeof page.afterCursor === "string" && page.afterCursor.length > 0) { - return page.afterCursor; - } - return null; -} - -interface AgentDirectoryFetchInput { - client: DaemonClient; - filter?: FetchAgentsOptions["filter"]; - subscribe?: FetchAgentsOptions["subscribe"]; - page?: FetchAgentsOptions["page"]; - transaction: AgentDirectoryTransaction; - isCurrent: () => boolean; -} - -interface AgentDirectoryFetchResult { - entries: FetchAgentsEntry[]; - deltas: readonly AgentDirectoryDelta[]; - subscriptionId: string | null; -} - -async function fetchCurrentAgentDirectory( - input: AgentDirectoryFetchInput, -): Promise { - const pageLimit = input.page?.limit ?? DEFAULT_AGENT_DIRECTORY_PAGE_LIMIT; - let cursor = input.page?.cursor ?? null; - let includeSubscribe = true; - let subscriptionId: string | null = null; - - while (true) { - const payload = await input.client.fetchAgents({ - scope: input.filter ? undefined : "active", - ...(input.filter ? { filter: input.filter } : {}), - sort: DEFAULT_AGENT_DIRECTORY_SORT, - ...(includeSubscribe && input.subscribe ? { subscribe: input.subscribe } : {}), - page: cursor ? { limit: pageLimit, cursor } : { limit: pageLimit }, - }); - - if (!input.isCurrent()) throw new AgentDirectoryRefreshSupersededError(); - input.transaction.entries.push(...payload.entries); - subscriptionId = subscriptionId ?? payload.subscriptionId ?? null; - includeSubscribe = false; - - if (!readFetchAgentsHasMore(payload.pageInfo)) { - break; - } - - const nextCursor = readFetchAgentsNextCursor(payload.pageInfo); - if (!nextCursor) { - break; - } - cursor = nextCursor; - } - - return { - entries: input.transaction.entries, - deltas: input.transaction.deltas, - subscriptionId, - }; -} - function toActiveConnection(connection: HostConnection): ActiveConnection { if (connection.type === "directSocket") { return { @@ -466,9 +363,15 @@ function nextConnectionMachineState(input: { function toSnapshotConnectionPatch( state: HostRuntimeConnectionMachineState, + connectionEpoch: number, ): Pick< HostRuntimeSnapshot, - "activeConnectionId" | "activeConnection" | "connectionStatus" | "lastError" | "lastOnlineAt" + | "activeConnectionId" + | "activeConnection" + | "connectionStatus" + | "lastError" + | "lastOnlineAt" + | "connectionEpoch" > { if (state.tag === "booting") { return { @@ -477,6 +380,7 @@ function toSnapshotConnectionPatch( connectionStatus: "connecting", lastError: null, lastOnlineAt: null, + connectionEpoch, }; } if (state.tag === "connecting") { @@ -486,6 +390,7 @@ function toSnapshotConnectionPatch( connectionStatus: "connecting", lastError: null, lastOnlineAt: null, + connectionEpoch, }; } if (state.tag === "online") { @@ -495,6 +400,7 @@ function toSnapshotConnectionPatch( connectionStatus: "online", lastError: null, lastOnlineAt: state.lastOnlineAt, + connectionEpoch, }; } if (state.tag === "offline") { @@ -504,6 +410,7 @@ function toSnapshotConnectionPatch( connectionStatus: "offline", lastError: null, lastOnlineAt: null, + connectionEpoch, }; } return { @@ -512,6 +419,7 @@ function toSnapshotConnectionPatch( connectionStatus: "error", lastError: state.message, lastOnlineAt: null, + connectionEpoch, }; } @@ -639,6 +547,7 @@ export class HostRuntimeController { private deps: HostRuntimeControllerDeps; private onReconcileServerId: ((oldId: string, newId: string) => void) | null; private connectionMachineState: HostRuntimeConnectionMachineState; + private connectionEpoch = 0; private snapshot: HostRuntimeSnapshot; private listeners = new Set<() => void>(); private activeClient: DaemonClient | null = null; @@ -669,7 +578,7 @@ export class HostRuntimeController { }; this.snapshot = { serverId: this.host.serverId, - ...toSnapshotConnectionPatch(this.connectionMachineState), + ...toSnapshotConnectionPatch(this.connectionMachineState, this.connectionEpoch), client: null, agentDirectoryStatus: "idle", agentDirectoryError: null, @@ -737,7 +646,7 @@ export class HostRuntimeController { } this.applyConnectionEvent({ type: "stopped" }); this.updateSnapshot({ - ...toSnapshotConnectionPatch(this.connectionMachineState), + ...toSnapshotConnectionPatch(this.connectionMachineState, this.connectionEpoch), client: null, }); } @@ -801,7 +710,7 @@ export class HostRuntimeController { markStartupError(message: string): void { this.applyConnectionEvent({ type: "connect_failed", message }); this.updateSnapshot({ - ...toSnapshotConnectionPatch(this.connectionMachineState), + ...toSnapshotConnectionPatch(this.connectionMachineState, this.connectionEpoch), }); } @@ -834,7 +743,7 @@ export class HostRuntimeController { } this.applyConnectionEvent({ type: "no_connections" }); this.updateSnapshot({ - ...toSnapshotConnectionPatch(this.connectionMachineState), + ...toSnapshotConnectionPatch(this.connectionMachineState, this.connectionEpoch), probeByConnectionId: new Map(), }); return; @@ -1120,6 +1029,9 @@ export class HostRuntimeController { state: previousState, event, }); + if (previousState.tag !== "online" && nextState.tag === "online") { + this.connectionEpoch += 1; + } this.connectionMachineState = nextState; this.logConnectionTransition({ from: previousState.tag, @@ -1196,7 +1108,7 @@ export class HostRuntimeController { message: `Failed to resolve client id: ${message}`, }); this.updateSnapshot({ - ...toSnapshotConnectionPatch(this.connectionMachineState), + ...toSnapshotConnectionPatch(this.connectionMachineState, this.connectionEpoch), }); return null; } @@ -1294,7 +1206,7 @@ export class HostRuntimeController { this.snapshot = { ...this.snapshot, serverId: this.host.serverId, - ...toSnapshotConnectionPatch(this.connectionMachineState), + ...toSnapshotConnectionPatch(this.connectionMachineState, this.connectionEpoch), client, clientGeneration: nextGeneration, }; @@ -1312,7 +1224,7 @@ export class HostRuntimeController { lastError: client.lastError, }); const patch: HostRuntimeSnapshotPatch = { - ...toSnapshotConnectionPatch(this.connectionMachineState), + ...toSnapshotConnectionPatch(this.connectionMachineState, this.connectionEpoch), ...this.buildAgentDirectoryStatusPatch(), }; this.updateSnapshot(patch); @@ -1332,7 +1244,7 @@ export class HostRuntimeController { message, }); this.updateSnapshot({ - ...toSnapshotConnectionPatch(this.connectionMachineState), + ...toSnapshotConnectionPatch(this.connectionMachineState, this.connectionEpoch), }); } } @@ -1418,14 +1330,6 @@ function rekeyMap(map: Map, oldKey: string, newKey: string): void map.set(newKey, value); } -interface AgentDirectoryTransaction { - id: symbol; - client: DaemonClient; - clientGeneration: number; - entries: FetchAgentsEntry[]; - deltas: AgentDirectoryDelta[]; -} - interface AgentDirectoryRefreshInput { serverId: string; filter?: FetchAgentsOptions["filter"]; @@ -1433,13 +1337,6 @@ interface AgentDirectoryRefreshInput { page?: FetchAgentsOptions["page"]; } -interface AgentDirectoryRefreshResult { - agents: ReturnType["agents"]; - subscriptionId: string | null; -} - -class AgentDirectoryRefreshSupersededError extends Error {} - export class HostRuntimeStore { private controllers = new Map(); private serverListeners = new Map void>>(); @@ -1453,11 +1350,8 @@ export class HostRuntimeStore { private deps: HostRuntimeControllerDeps; private lastConnectionStatusByServer = new Map(); private agentDirectoryBootstrapInFlight = new Map>(); - private agentDirectoryTransactions = new Map(); - private agentDirectoryListeners = new Map< - string, - { client: DaemonClient; generation: number; unsubscribe: () => void } - >(); + private queuedAgentDrainInFlight = new Set(); + private directorySyncByServer = new Map(); private configuredOverrideBootstrapInFlight: Promise | null = null; private bootStarted = false; private storage: HostRuntimeStorage; @@ -1688,6 +1582,27 @@ export class HostRuntimeStore { rekeyMap(this.lastConnectionStatusByServer, oldServerId, newServerId); rekeyMap(this.agentDirectoryBootstrapInFlight, oldServerId, newServerId); + this.directorySyncByServer.get(oldServerId)?.dispose(); + this.directorySyncByServer.delete(oldServerId); + const directory = new DirectorySync(newServerId, { + drainQueuedAgentMessage: (agentId) => this.drainQueuedAgentMessage(newServerId, agentId), + markAgentLoading: () => controller.markAgentDirectorySyncLoading(), + markAgentReady: () => { + this.agentDirectoryBootstrapInFlight.delete(newServerId); + controller.markAgentDirectorySyncReady(); + }, + markAgentError: (error) => controller.markAgentDirectorySyncError(error), + }); + this.directorySyncByServer.set(newServerId, directory); + const snapshot = controller.getSnapshot(); + directory.connectionChanged({ + client: snapshot.client, + status: snapshot.connectionStatus === "online" ? "online" : "offline", + source: { + clientGeneration: snapshot.clientGeneration, + connectionEpoch: snapshot.connectionEpoch, + }, + }); const listeners = this.serverListeners.get(oldServerId); if (listeners) { @@ -1978,9 +1893,8 @@ export class HostRuntimeStore { this.controllers.delete(serverId); this.lastConnectionStatusByServer.delete(serverId); this.agentDirectoryBootstrapInFlight.delete(serverId); - this.agentDirectoryTransactions.delete(serverId); - this.agentDirectoryListeners.get(serverId)?.unsubscribe(); - this.agentDirectoryListeners.delete(serverId); + this.directorySyncByServer.get(serverId)?.dispose(); + this.directorySyncByServer.delete(serverId); void controller.stop(); this.emit(serverId); } @@ -2003,6 +1917,19 @@ export class HostRuntimeStore { onReconcileServerId: (oldId, newId) => this.reconcileServerId(oldId, newId), }); this.controllers.set(host.serverId, controller); + this.directorySyncByServer.set( + host.serverId, + new DirectorySync(host.serverId, { + drainQueuedAgentMessage: (agentId) => + this.drainQueuedAgentMessage(host.serverId, agentId), + markAgentLoading: () => controller.markAgentDirectorySyncLoading(), + markAgentReady: () => { + this.agentDirectoryBootstrapInFlight.delete(host.serverId); + controller.markAgentDirectorySyncReady(); + }, + markAgentError: (error) => controller.markAgentDirectorySyncError(error), + }), + ); this.lastConnectionStatusByServer.set( host.serverId, controller.getSnapshot().connectionStatus, @@ -2035,10 +1962,13 @@ export class HostRuntimeStore { return; } const snapshot = controller.getSnapshot(); - this.installAgentDirectoryListener({ - serverId, + this.directorySyncByServer.get(serverId)?.connectionChanged({ client: snapshot.client, - generation: snapshot.clientGeneration, + status: snapshot.connectionStatus === "online" ? "online" : "offline", + source: { + clientGeneration: snapshot.clientGeneration, + connectionEpoch: snapshot.connectionEpoch, + }, }); const previousStatus = this.lastConnectionStatusByServer.get(serverId); this.lastConnectionStatusByServer.set(serverId, snapshot.connectionStatus); @@ -2066,11 +1996,17 @@ export class HostRuntimeStore { return; } - const bootstrap = this.refreshAgentDirectory({ - serverId, - subscribe: { subscriptionId: `app:${serverId}` }, - page: { limit: DEFAULT_AGENT_DIRECTORY_PAGE_LIMIT }, - }) + const bootstrap = Promise.resolve() + .then(() => + Promise.all([ + this.refreshAgentDirectory({ + serverId, + subscribe: { subscriptionId: `app:${serverId}` }, + page: { limit: DEFAULT_AGENT_DIRECTORY_PAGE_LIMIT }, + }), + this.refreshWorkspaceDirectory({ serverId, subscribe: true }), + ]), + ) .then(() => undefined) .catch((error) => { console.error("[HostRuntime] agent directory bootstrap failed", { @@ -2088,46 +2024,9 @@ export class HostRuntimeStore { this.agentDirectoryBootstrapInFlight.set(serverId, bootstrap); } - private installAgentDirectoryListener(input: { - serverId: string; - client: DaemonClient | null; - generation: number; - }): void { - const current = this.agentDirectoryListeners.get(input.serverId); - if (current?.client === input.client && current.generation === input.generation) return; - current?.unsubscribe(); - this.agentDirectoryListeners.delete(input.serverId); - this.agentDirectoryTransactions.delete(input.serverId); - if (!input.client) return; - const client = input.client; - const unsubscribe = client.on("agent_update", (message) => { - if (message.type !== "agent_update") return; - const installed = this.agentDirectoryListeners.get(input.serverId); - if (installed?.client !== client || installed.generation !== input.generation) return; - const transaction = this.agentDirectoryTransactions.get(input.serverId); - if (transaction?.client === client && transaction.clientGeneration === input.generation) { - transaction.deltas.push(message.payload); - return; - } - const session = useSessionStore.getState().sessions[input.serverId]; - if (session?.client !== client || session.clientGeneration !== input.generation) return; - this.applyAgentDirectoryDelta(input.serverId, message.payload); - }); - this.agentDirectoryListeners.set(input.serverId, { - client, - generation: input.generation, - unsubscribe, - }); - } - - private applyAgentDirectoryDelta(serverId: string, delta: AgentDirectoryDelta): void { - const result = applyAgentDirectoryDelta({ serverId, delta }); - if (result.stoppedRunning) { - this.drainQueuedAgentMessage(serverId, result.agentId); - } - } - drainQueuedAgentMessage(serverId: string, agentId: string): void { + const drainKey = `${serverId}:${agentId}`; + if (this.queuedAgentDrainInFlight.has(drainKey)) return; const store = useSessionStore.getState(); const session = store.sessions[serverId]; const queue = session?.queuedMessages.get(agentId); @@ -2135,6 +2034,7 @@ export class HostRuntimeStore { if (!client || !queue?.length || session.initializingAgents.get(agentId) === true) { return; } + this.queuedAgentDrainInFlight.add(drainKey); const next = queue[0]; void sendQueuedComposerMessageNow({ agentId, @@ -2157,16 +2057,20 @@ export class HostRuntimeStore { attachments: wirePayload.attachments, }); }, - }).then((result) => { - if (result.status === "failed") { - console.error("[HostRuntime] failed to drain queued agent message", { - serverId, - agentId, - error: result.errorMessage, - }); - } - return result; - }); + }) + .then((result) => { + if (result.status === "failed") { + console.error("[HostRuntime] failed to drain queued agent message", { + serverId, + agentId, + error: result.errorMessage, + }); + } + return result; + }) + .finally(() => { + this.queuedAgentDrainInFlight.delete(drainKey); + }); } getSnapshot(serverId: string): HostRuntimeSnapshot | null { @@ -2235,213 +2139,32 @@ export class HostRuntimeStore { async refreshAgentDirectory( input: AgentDirectoryRefreshInput, - ): Promise { - const controller = this.controllers.get(input.serverId); - if (!controller) { - throw new Error(`Unknown host runtime for serverId ${input.serverId}`); - } - const snapshot = controller.getSnapshot(); - const client = controller.getClient(); - if (!client || snapshot.connectionStatus !== "online") { - throw new Error(`Host ${input.serverId} is not connected`); - } - - const previousTransaction = this.agentDirectoryTransactions.get(input.serverId); - const transaction: AgentDirectoryTransaction = { - id: Symbol("agent directory refresh"), - client, - clientGeneration: snapshot.clientGeneration, - entries: [], - deltas: - previousTransaction?.clientGeneration === snapshot.clientGeneration - ? inheritBufferedDirectoryDeltas({ client, previous: previousTransaction }) - : [], - }; - this.agentDirectoryTransactions.set(input.serverId, transaction); - - const hasMatchingSession = () => { - const session = useSessionStore.getState().sessions[input.serverId]; - return ( - session?.client === client && session.clientGeneration === transaction.clientGeneration - ); - }; - if (!hasMatchingSession()) { - await new Promise((resolve, reject) => { - const unsubscribeStore = useSessionStore.subscribe((state) => { - const session = state.sessions[input.serverId]; - if ( - session?.client === client && - session.clientGeneration === transaction.clientGeneration - ) { - unsubscribeStore(); - unsubscribeController(); - resolve(); - } - }); - const unsubscribeController = controller.subscribe(() => { - if (controller.getClient() !== client) { - unsubscribeStore(); - unsubscribeController(); - reject(new Error(`Host ${input.serverId} client changed before directory bootstrap`)); - } - }); - }); - } - - controller.markAgentDirectorySyncLoading(); - try { - const session = useSessionStore.getState().sessions[input.serverId]; - if (!input.filter && shouldUseLegacyDaemonWorkspaceDirectory(session?.serverInfo)) { - return await this.refreshLegacyAgentDirectory({ - input, - controller, - client, - transaction, - hasMatchingSession, - }); - } - return await this.refreshCurrentAgentDirectory({ - input, - controller, - client, - transaction, - hasMatchingSession, - }); - } catch (error) { - if (!(error instanceof AgentDirectoryRefreshSupersededError)) { - if ( - this.agentDirectoryTransactions.get(input.serverId)?.id === transaction.id && - hasMatchingSession() - ) { - for (const delta of transaction.deltas) { - this.applyAgentDirectoryDelta(input.serverId, delta); - } - } - controller.markAgentDirectorySyncError(toErrorMessage(error)); - } - throw error; - } finally { - if (this.agentDirectoryTransactions.get(input.serverId)?.id === transaction.id) { - this.agentDirectoryTransactions.delete(input.serverId); - } - } + ): Promise { + const directory = this.directorySyncByServer.get(input.serverId); + if (!directory) throw new Error(`Unknown host runtime for serverId ${input.serverId}`); + return directory.refreshAgents(input); } - private async refreshLegacyAgentDirectory(context: { - input: AgentDirectoryRefreshInput; - controller: HostRuntimeController; - client: DaemonClient; - transaction: AgentDirectoryTransaction; - hasMatchingSession: () => boolean; - }): Promise { - const { input, controller, client, transaction, hasMatchingSession } = context; - const directory = await readLegacyDaemonWorkspaceDirectory({ - client, - subscribe: input.subscribe, - page: input.page, - }); - if ( - !directory || - this.agentDirectoryTransactions.get(input.serverId)?.id !== transaction.id || - !hasMatchingSession() - ) { - throw new AgentDirectoryRefreshSupersededError(); - } - const previous = useSessionStore.getState().sessions[input.serverId]?.agents ?? new Map(); - const stampedSnapshot = stampLegacyWorkspaceIds(directory.entries); - const reconciled = reconcileAgentDirectory({ - previous, - snapshot: stampedSnapshot, - deltas: transaction.deltas, - }); - const result = replaceLegacyDaemonWorkspaceDirectory({ - serverId: input.serverId, - entries: reconciled.entries, - }); - this.applyAgentDirectoryCommitSideEffects({ - serverId: input.serverId, - previous, - committedEntries: reconciled.entries, - directory: { - entries: stampedSnapshot, - deltas: transaction.deltas, - subscriptionId: directory.subscriptionId, - }, - stoppedRunningAgentIds: reconciled.stoppedRunningAgentIds, - }); - this.agentDirectoryBootstrapInFlight.delete(input.serverId); - controller.markAgentDirectorySyncReady(); - return { agents: result.agents, subscriptionId: directory.subscriptionId }; + async refreshWorkspaceDirectory(input: { serverId: string; subscribe?: boolean }): Promise { + const directory = this.directorySyncByServer.get(input.serverId); + if (!directory) throw new Error(`Unknown host runtime for serverId ${input.serverId}`); + await directory.refreshWorkspaces({ subscribe: input.subscribe }); } - private async refreshCurrentAgentDirectory(context: { - input: AgentDirectoryRefreshInput; - controller: HostRuntimeController; - client: DaemonClient; - transaction: AgentDirectoryTransaction; - hasMatchingSession: () => boolean; - }): Promise { - const { input, controller, client, transaction, hasMatchingSession } = context; - const directory = await fetchCurrentAgentDirectory({ - client, - filter: input.filter, - subscribe: input.subscribe, - page: input.page, - transaction, - isCurrent: () => - this.agentDirectoryTransactions.get(input.serverId)?.id === transaction.id && - controller.getSnapshot().client === client && - controller.getSnapshot().clientGeneration === transaction.clientGeneration, - }); - if ( - this.agentDirectoryTransactions.get(input.serverId)?.id !== transaction.id || - controller.getSnapshot().client !== client || - controller.getSnapshot().clientGeneration !== transaction.clientGeneration || - !hasMatchingSession() - ) { - throw new AgentDirectoryRefreshSupersededError(); - } - const previous = useSessionStore.getState().sessions[input.serverId]?.agents ?? new Map(); - const reconciled = reconcileAgentDirectory({ - previous, - snapshot: directory.entries, - deltas: directory.deltas, - }); - const { agents } = replaceFetchedAgentDirectory({ - serverId: input.serverId, - entries: reconciled.entries, - }); - this.applyAgentDirectoryCommitSideEffects({ - serverId: input.serverId, - previous, - committedEntries: reconciled.entries, - directory, - stoppedRunningAgentIds: reconciled.stoppedRunningAgentIds, - }); - this.agentDirectoryBootstrapInFlight.delete(input.serverId); - controller.markAgentDirectorySyncReady(); - return { agents, subscriptionId: directory.subscriptionId }; + async refreshDirectories(serverId: string): Promise { + const directory = this.directorySyncByServer.get(serverId); + if (!directory) throw new Error(`Unknown host runtime for serverId ${serverId}`); + await directory.refreshAll(); } - private applyAgentDirectoryCommitSideEffects(input: { - serverId: string; - previous: ReadonlyMap; - committedEntries: FetchAgentsEntry[]; - directory: AgentDirectoryFetchResult; - stoppedRunningAgentIds: string[]; - }): void { - const snapshotAgentIds = new Set(input.committedEntries.map((entry) => entry.agent.id)); - for (const agentId of input.previous.keys()) { - if (!snapshotAgentIds.has(agentId)) { - applyAgentDirectoryDelta({ serverId: input.serverId, delta: { kind: "remove", agentId } }); - } - } - for (const delta of input.directory.deltas) { - applyAgentDirectoryDelta({ serverId: input.serverId, delta }); - } - for (const agentId of input.stoppedRunningAgentIds) { - this.drainQueuedAgentMessage(input.serverId, agentId); - } + fetchAgentTimeline( + serverId: string, + agentId: string, + request: Parameters[1], + ): Promise>> { + const directory = this.directorySyncByServer.get(serverId); + if (!directory) throw new Error(`Unknown host runtime for serverId ${serverId}`); + return directory.fetchTimeline(agentId, request); } refreshAllAgentDirectories(input?: { serverIds?: string[] }): void { diff --git a/packages/app/src/screens/workspace/workspace-screen.tsx b/packages/app/src/screens/workspace/workspace-screen.tsx index 49eb175594..3de6329904 100644 --- a/packages/app/src/screens/workspace/workspace-screen.tsx +++ b/packages/app/src/screens/workspace/workspace-screen.tsx @@ -2766,7 +2766,7 @@ function WorkspaceScreenContent({ // dropped against the stale cursor. const sessionState = useSessionStore.getState().sessions[normalizedServerId]; const currentCursor = sessionState?.agentTimelineCursor.get(agentId); - await client.fetchAgentTimeline(agentId, { + await getHostRuntimeStore().fetchAgentTimeline(normalizedServerId, agentId, { direction: "tail", projection: "projected", ...(currentCursor diff --git a/packages/app/src/utils/agent-directory-sync.test.ts b/packages/app/src/utils/agent-directory-sync.test.ts index b357652e75..1ec16bde26 100644 --- a/packages/app/src/utils/agent-directory-sync.test.ts +++ b/packages/app/src/utils/agent-directory-sync.test.ts @@ -7,11 +7,7 @@ import { useSessionStore } from "@/stores/session-store"; import { normalizeAgentSnapshot } from "@/utils/agent-snapshots"; import { isAgentArchiving, setAgentArchiving } from "@/hooks/use-archive-agent"; import { queryClient } from "@/data/query-client"; -import { - applyAgentDirectoryDelta, - replaceFetchedAgentDirectory, - shouldApplyTimelineAgentSnapshot, -} from "./agent-directory-sync"; +import { applyAgentDirectoryDelta, replaceFetchedAgentDirectory } from "./agent-directory-sync"; function createAgentPayload( input: Partial> & { @@ -158,7 +154,6 @@ describe("replaceFetchedAgentDirectory", () => { cursor: session?.agentTimelineCursor.has(agentId), permissions: session?.pendingPermissions.size, initializing: session?.initializingAgents.has(agentId), - acceptsStaleTimelineSnapshot: shouldApplyTimelineAgentSnapshot(serverId, agentId), archivePending: isAgentArchiving({ queryClient, serverId, agentId }), }).toEqual({ agents: false, @@ -167,7 +162,6 @@ describe("replaceFetchedAgentDirectory", () => { cursor: false, permissions: 0, initializing: false, - acceptsStaleTimelineSnapshot: false, archivePending: false, }); diff --git a/packages/app/src/utils/agent-directory-sync.ts b/packages/app/src/utils/agent-directory-sync.ts index fd785ba972..24e14c4208 100644 --- a/packages/app/src/utils/agent-directory-sync.ts +++ b/packages/app/src/utils/agent-directory-sync.ts @@ -6,6 +6,9 @@ import type { SessionOutboundMessage } from "@getpaseo/protocol/messages"; import { clearArchiveAgentPending } from "@/hooks/use-archive-agent"; import { queryClient } from "@/data/query-client"; import { acceptAgentDirectoryUpdate } from "@/utils/agent-directory-update-policy"; +import { buildDraftStoreKey } from "@/stores/draft-keys"; +import { useDraftStore } from "@/stores/draft-store"; +import { getInitDeferred, getInitKey, rejectInitDeferred } from "@/utils/agent-initialization"; type AgentDirectoryFetchEntry = FetchAgentsEntry; export type AgentDirectoryDelta = Extract< @@ -24,11 +27,6 @@ export function applyAgentDirectoryDelta(input: { serverId: string; delta: Agent return upsertAgentDirectoryReplica(input.serverId, input.delta); } -export function shouldApplyTimelineAgentSnapshot(serverId: string, agentId: string): boolean { - const session = useSessionStore.getState().sessions[serverId]; - return session?.agents.has(agentId) === true || session?.initializingAgents.get(agentId) === true; -} - type AgentUpsertDelta = Extract; function upsertAgentDirectoryReplica( @@ -65,7 +63,7 @@ function upsertAgentDirectoryReplica( }; } -function upsertAgentReplica(serverId: string, agent: Agent): Agent { +export function upsertAgentReplica(serverId: string, agent: Agent): Agent { let acceptedAgent = agent; useSessionStore.getState().setAgents(serverId, (current) => { const currentAgent = current.get(agent.id); @@ -78,7 +76,7 @@ function upsertAgentReplica(serverId: string, agent: Agent): Agent { return acceptedAgent; } -function replaceAgentPendingPermissions(serverId: string, agent: Agent): void { +export function replaceAgentPendingPermissions(serverId: string, agent: Agent): void { const pendingPermissions = new Map( useSessionStore.getState().sessions[serverId]?.pendingPermissions, ); @@ -92,7 +90,7 @@ function replaceAgentPendingPermissions(serverId: string, agent: Agent): void { useSessionStore.getState().setPendingPermissions(serverId, pendingPermissions); } -function removeAgentDirectoryReplica(serverId: string, agentId: string): void { +export function removeAgentDirectoryReplica(serverId: string, agentId: string): void { const store = useSessionStore.getState(); clearArchiveAgentPending({ queryClient, serverId, agentId }); const removeKey = (current: Map): Map => { @@ -114,6 +112,21 @@ function removeAgentDirectoryReplica(serverId: string, agentId: string): void { return next.size === current.size ? current : next; }); store.setAgentAuthoritativeHistoryApplied(serverId, agentId, false); + store.setAgentStreamTail(serverId, removeKey); + store.clearAgentStreamHead(serverId, agentId); + useSessionStore.setState((state) => { + if (!state.agentLastActivity.has(agentId)) return state; + const agentLastActivity = new Map(state.agentLastActivity); + agentLastActivity.delete(agentId); + return { ...state, agentLastActivity }; + }); + useDraftStore.getState().clearDraftInput({ + draftKey: buildDraftStoreKey({ serverId, agentId }), + }); + const initKey = getInitKey(serverId, agentId); + if (getInitDeferred(initKey)) { + rejectInitDeferred(initKey, new Error("Agent was removed during initialization")); + } } interface PendingPermissionEntry { diff --git a/packages/app/src/utils/buffered-directory-transaction.test.ts b/packages/app/src/utils/buffered-directory-transaction.test.ts deleted file mode 100644 index ec06ca9d53..0000000000 --- a/packages/app/src/utils/buffered-directory-transaction.test.ts +++ /dev/null @@ -1,23 +0,0 @@ -import { describe, expect, it } from "vitest"; -import { inheritBufferedDirectoryDeltas } from "./buffered-directory-transaction"; - -describe("inheritBufferedDirectoryDeltas", () => { - it("hands buffered updates to a superseding transaction for the same client", () => { - const client = {}; - const previous = { client, deltas: ["first"] }; - - const inherited = inheritBufferedDirectoryDeltas({ client, previous }); - previous.deltas.push("later"); - - expect(inherited).toEqual(["first"]); - }); - - it("does not carry updates across clients", () => { - expect( - inheritBufferedDirectoryDeltas({ - client: {}, - previous: { client: {}, deltas: ["stale"] }, - }), - ).toEqual([]); - }); -}); diff --git a/packages/app/src/utils/buffered-directory-transaction.ts b/packages/app/src/utils/buffered-directory-transaction.ts deleted file mode 100644 index 05592fafa9..0000000000 --- a/packages/app/src/utils/buffered-directory-transaction.ts +++ /dev/null @@ -1,12 +0,0 @@ -export interface BufferedDirectoryTransaction { - client: TClient; - deltas: TDelta[]; -} - -export function inheritBufferedDirectoryDeltas(input: { - client: TClient; - previous: BufferedDirectoryTransaction | null | undefined; -}): TDelta[] { - if (input.previous?.client !== input.client) return []; - return [...input.previous.deltas]; -} diff --git a/packages/app/src/workspace/legacy-daemon-workspaces.ts b/packages/app/src/workspace/legacy-daemon-workspaces.ts index 5979edfa6e..820e373688 100644 --- a/packages/app/src/workspace/legacy-daemon-workspaces.ts +++ b/packages/app/src/workspace/legacy-daemon-workspaces.ts @@ -255,7 +255,9 @@ export function stampLegacyWorkspaceIds(entries: FetchAgentsEntry[]): FetchAgent }); } -function buildLegacyWorkspaces(entries: FetchAgentsEntry[]): Map { +export function buildLegacyWorkspaces( + entries: FetchAgentsEntry[], +): Map { const workspaces = new Map(); for (const entry of entries) { const workspaceId = entry.agent.workspaceId ?? resolveLegacyWorkspaceId(entry); diff --git a/packages/server/src/server/session.test.ts b/packages/server/src/server/session.test.ts index 8f0054dc65..9f05447e85 100644 --- a/packages/server/src/server/session.test.ts +++ b/packages/server/src/server/session.test.ts @@ -4721,6 +4721,58 @@ test("unions viewed timelines across socket sources and removes detached sources ).toEqual(["agent-b"]); }); +test("keeps selective delivery scoped per socket when a retained session also has a legacy socket", async () => { + const messages: SessionOutboundMessage[] = []; + const targetedMessages: Array<{ source: object; message: SessionOutboundMessage }> = []; + const agentEventListeners: Array<(event: AgentManagerEvent) => void> = []; + const session = createSessionForTest({ + messages, + targetedMessages, + agentManager: { + subscribe: vi.fn((listener: (event: AgentManagerEvent) => void) => { + agentEventListeners.push(listener); + return () => {}; + }), + }, + }); + const legacySocket = {}; + const selectiveSocket = {}; + session.updateClientCapabilities(null, legacySocket); + session.updateClientCapabilities({ selective_agent_timeline: true }, selectiveSocket); + await session.handleMessage( + { + type: "agent.timeline.set_subscription.request", + agentIds: ["viewed-agent"], + requestId: "timeline-subscription-selective", + }, + selectiveSocket, + ); + targetedMessages.length = 0; + + const listener = agentEventListeners[0]; + if (!listener) throw new Error("Agent event listener was not installed"); + listener({ + type: "agent_stream", + agentId: "not-viewed-agent", + event: { + type: "timeline", + provider: "mock", + item: { type: "assistant_message", messageId: "message-global", text: "global" }, + }, + }); + + expect(messages).toEqual([]); + expect(targetedMessages).toEqual([ + { + source: legacySocket, + message: expect.objectContaining({ + type: "agent_stream", + payload: expect.objectContaining({ agentId: "not-viewed-agent" }), + }), + }, + ]); +}); + describe("agent config setters", () => { test("set_agent_mode_request: success emits accepted response carrying the notice", async () => { const messages: SessionOutboundMessage[] = []; diff --git a/packages/server/src/server/session.ts b/packages/server/src/server/session.ts index 9d0f23800d..ee010eb295 100644 --- a/packages/server/src/server/session.ts +++ b/packages/server/src/server/session.ts @@ -1019,6 +1019,61 @@ export class Session { return true; } + private forwardAgentStream( + event: Extract, + serializedEvent: Extract["payload"]["event"], + ): void { + if (this.selectiveTimelineCapabilityBySource.size === 0 || !this.onMessageToSource) { + if (this.usesSelectiveTimelineDelivery() && serializedEvent.type === "attention_required") { + this.emit({ + type: "agent_attention_required", + payload: { + agentId: event.agentId, + reason: serializedEvent.reason, + timestamp: serializedEvent.timestamp, + shouldNotify: serializedEvent.shouldNotify, + ...(serializedEvent.notification ? { notification: serializedEvent.notification } : {}), + }, + }); + } else if ( + !this.usesSelectiveTimelineDelivery() || + this.viewedTimelineAgentIds.has(event.agentId) + ) { + this.emit({ + type: "agent_stream", + payload: this.buildAgentStreamPayload(event, serializedEvent), + }); + } + return; + } + + for (const [source, supportsSelectiveDelivery] of this.selectiveTimelineCapabilityBySource) { + if (supportsSelectiveDelivery && serializedEvent.type === "attention_required") { + this.onMessageToSource(source, { + type: "agent_attention_required", + payload: { + agentId: event.agentId, + reason: serializedEvent.reason, + timestamp: serializedEvent.timestamp, + shouldNotify: serializedEvent.shouldNotify, + ...(serializedEvent.notification ? { notification: serializedEvent.notification } : {}), + }, + }); + continue; + } + if ( + supportsSelectiveDelivery && + !this.viewedTimelineAgentIdsBySource.get(source)?.has(event.agentId) + ) { + continue; + } + this.onMessageToSource(source, { + type: "agent_stream", + payload: this.buildAgentStreamPayload(event, serializedEvent), + }); + } + } + supports(capability: ClientCapability): boolean { return this.clientCapabilities.has(capability); } @@ -1340,28 +1395,7 @@ export class Session { "agent.session.forward_stream", ); - if (this.usesSelectiveTimelineDelivery() && serializedEvent.type === "attention_required") { - this.emit({ - type: "agent_attention_required", - payload: { - agentId: event.agentId, - reason: serializedEvent.reason, - timestamp: serializedEvent.timestamp, - shouldNotify: serializedEvent.shouldNotify, - ...(serializedEvent.notification - ? { notification: serializedEvent.notification } - : {}), - }, - }); - } else if ( - !this.usesSelectiveTimelineDelivery() || - this.viewedTimelineAgentIds.has(event.agentId) - ) { - this.emit({ - type: "agent_stream", - payload: this.buildAgentStreamPayload(event, serializedEvent), - }); - } + this.forwardAgentStream(event, serializedEvent); if (event.event.type === "permission_requested") { this.emit({