From 7c4aa058ff530845400a5ba82b4c8fefc5bfe64d Mon Sep 17 00:00:00 2001 From: Mohamed Boudra Date: Tue, 20 Oct 2026 13:23:04 +0200 Subject: [PATCH] feat(hosts): bind replica and session lifetime to origin authority --- docs/connected-hosts.md | 4 + packages/app/src/contexts/session-context.tsx | 63 +- packages/app/src/data/query-client.ts | 21 +- .../hooks/use-agent-initialization.test.ts | 44 +- .../app/src/hooks/use-agent-initialization.ts | 51 +- .../app/src/hooks/use-archive-agent.test.ts | 3 + packages/app/src/hooks/use-archive-agent.ts | 167 +- packages/app/src/panels/agent-panel.tsx | 55 +- .../directory-sync/agent-replica.test.ts | 32 + .../runtime/directory-sync/agent-replica.ts | 5 +- .../app/src/runtime/directory-sync/index.ts | 38 +- .../directory-sync/internal/agent-store.ts | 46 +- .../directory-sync/workspace-replica.ts | 22 +- .../runtime/host-authority-provider.test.tsx | 44 + .../src/runtime/host-authority-provider.tsx | 26 + packages/app/src/runtime/host-runtime.test.ts | 332 ++- packages/app/src/runtime/host-runtime.ts | 484 +++- .../runtime/replica-cache/row-store.native.ts | 18 +- .../replica-cache/row-store.web.test.ts | 26 +- .../runtime/replica-cache/row-store.web.ts | 6 +- packages/app/src/stores/create-flow-store.ts | 185 +- packages/app/src/stores/session-store.ts | 2363 +++++++++-------- packages/app/src/timeline/replica.test.ts | 113 +- .../src/timeline/session-stream-reducers.ts | 9 +- .../app/src/timeline/viewed-timeline-sync.ts | 78 +- .../app/src/utils/agent-initialization.ts | 146 +- packages/client/package.json | 4 + packages/client/src/hosts/connection.ts | 171 ++ packages/client/src/hosts/index.ts | 171 +- packages/client/src/schedules/index.ts | 18 +- packages/protocol/src/messages.ts | 1 + .../server/src/server/agent/agent-manager.ts | 49 +- .../src/server/agent/agent-prompt.test.ts | 2 + .../src/server/agent/mcp-parity.e2e.test.ts | 126 + .../src/server/agent/mcp-server.test.ts | 309 +-- .../server/src/server/agent/runs/index.ts | 15 +- packages/server/src/server/bootstrap.ts | 1 + .../server/connected-hosts/hosts.e2e.test.ts | 358 ++- .../server/src/server/connected-hosts/link.ts | 234 +- .../hub/execution-session.websocket.test.ts | 34 +- .../hub/test-utils/relationship-harness.ts | 4 + packages/server/src/server/subagents/index.ts | 147 +- .../server/src/server/websocket-server.ts | 46 +- 43 files changed, 3796 insertions(+), 2275 deletions(-) create mode 100644 packages/app/src/runtime/host-authority-provider.test.tsx create mode 100644 packages/app/src/runtime/host-authority-provider.tsx create mode 100644 packages/client/src/hosts/connection.ts diff --git a/docs/connected-hosts.md b/docs/connected-hosts.md index 918fc01ee4..3ef72556e2 100644 --- a/docs/connected-hosts.md +++ b/docs/connected-hosts.md @@ -33,6 +33,10 @@ A managed child has a qualified `{serverId, agentId}` address on every host. The Durable relationships retain a stable authority reference and creation ceiling, never live credentials or execution capabilities. Recovery must resolve current authority through its owner. Derived runtime sessions cannot authorize their own recovery. Execution-scoped admissions and non-owner durable references without a supported authority resolver fail explicitly; do not recover them through the daemon owner's MCP runtime. +Before accepting new writes, the daemon adopts eligible pre-feature local parent labels into authoritative membership. The upgrade records its bounded candidates before qualifying agent metadata, then resumes those candidates after interruption. It never rescans later label edits. Existing membership, including forgotten entries, takes precedence; explicit null or qualified parent pointers do not establish legacy provenance. Adopted entries preserve the old daemon-owned local lifecycle, not an authenticated creator grant, and contain no creation recipe to replay. Recovery still resolves current owner authority within the retained self-only ceiling. + +Provider-owned subagents are descriptors and timelines in `ProviderSubagentStore`, not `AgentStorage` records. They retain their provider observation/cancellation path. Internal managed agents and execution-owned records are separate stored/runtime concepts and are excluded from adoption on both sides of the relationship. + ## Connections and recovery The shared Node connector owns the target grammar and each connection attempt, including a fresh diff --git a/packages/app/src/contexts/session-context.tsx b/packages/app/src/contexts/session-context.tsx index 7190b42462..26d3bd461a 100644 --- a/packages/app/src/contexts/session-context.tsx +++ b/packages/app/src/contexts/session-context.tsx @@ -27,23 +27,19 @@ import { import type { DaemonClient } from "@getpaseo/client/internal/daemon-client"; import type { AgentPermissionResponse } from "@getpaseo/protocol/agent-types"; -import { getHostRuntimeStore, useHostRuntimeIsConnected } from "@/runtime/host-runtime"; +import { useHostRuntimeStore, useHostRuntimeIsConnected } from "@/runtime/host-runtime"; import { useVoiceAudioEngineOptional, useVoiceRuntimeOptional } from "@/contexts/voice-context"; import type { AudioPlaybackSource } from "@/audio"; import { selectAgentTimelineState, toDaemonServerInfo, useSessionStore, + useSessionStoreApi, type SessionState, } from "@/stores/session-store"; import { sendOsNotification } from "@/utils/os-notifications"; import { getIsAppActivelyVisible, getIsAppVisible } from "@/utils/app-visibility"; -import { - getInitKey, - getInitDeferred, - createInitDeferred, - rejectInitDeferred, -} from "@/utils/agent-initialization"; +import { getInitKey } from "@/utils/agent-initialization"; import { derivePendingPermissionKey } from "@/utils/agent-snapshots"; import { useToast } from "@/contexts/toast-context"; import { toErrorMessage } from "@/utils/error-messages"; @@ -201,6 +197,8 @@ export function SessionProvider(props: SessionProviderProps) { } function SessionProviderInternal({ children, serverId, client }: SessionProviderClientProps) { + const sessionStore = useSessionStoreApi(); + const runtime = useHostRuntimeStore(); const { t } = useTranslation(); const voiceRuntime = useVoiceRuntimeOptional(); const voiceAudioEngine = useVoiceAudioEngineOptional(); @@ -265,7 +263,7 @@ function SessionProviderInternal({ children, serverId, client }: SessionProvider notification?: AgentAttentionNotificationPayload; }) => { const appState = appStateRef.current; - const session = useSessionStore.getState().sessions[serverId]; + const session = sessionStore.getState().sessions[serverId]; const attentionFocusedAgentId = session?.focusedAgentId ?? null; if (params.reason === "error") { return; @@ -309,7 +307,7 @@ function SessionProviderInternal({ children, serverId, client }: SessionProvider data: notification.data, }); }, - [serverId], + [serverId, sessionStore], ); useEffect(() => { @@ -377,7 +375,7 @@ function SessionProviderInternal({ children, serverId, client }: SessionProvider { event: streamEvent, seq, epoch, timestamp: parsedTimestamp }, ]); if (turnLiveness.length > 0) { - getHostRuntimeStore().applyAgentTurnLiveness(serverId, agentId, turnLiveness); + runtime.applyAgentTurnLiveness(serverId, agentId, turnLiveness); } sync.enqueueStreamEvent(agentId, { event: streamEvent, @@ -395,14 +393,14 @@ function SessionProviderInternal({ children, serverId, client }: SessionProvider void requestTimelineReplacement( { fetchAgentTimeline: (agentId, request) => - getHostRuntimeStore().fetchAgentTimeline(serverId, agentId, request), + runtime.fetchAgentTimeline(serverId, agentId, request), }, message.payload.agentId, ).catch((error: unknown) => { console.warn("[Session] timeline replacement refresh failed", { serverId, error }); }); }; - const sync = getHostRuntimeStore().createViewedTimelineOwner(serverId, { + const sync = runtime.createViewedTimelineOwner(serverId, { observe: (agentIds) => { const observation = client.observeTimeline(agentIds); observation.subscribe({ @@ -416,7 +414,7 @@ function SessionProviderInternal({ children, serverId, client }: SessionProvider }, readCursor: (agentId) => { const timeline = selectAgentTimelineState( - useSessionStore.getState().sessions[serverId], + sessionStore.getState().sessions[serverId], agentId, ); return timeline.status === "synced" && timeline.range @@ -424,15 +422,19 @@ function SessionProviderInternal({ children, serverId, client }: SessionProvider : undefined; }, fetchPage: async (agentId, request) => { - const session = useSessionStore.getState().sessions[serverId]; + const session = sessionStore.getState().sessions[serverId]; const initKey = getInitKey(serverId, agentId); const shouldInitialize = selectAgentTimelineState(session, agentId).status !== "synced"; if (shouldInitialize) { - if (!getInitDeferred(initKey)) { - const deferred = createInitDeferred(initKey, request.direction ?? "tail"); + if (!sessionStore.initializations.getInitDeferred(initKey)) { + const deferred = sessionStore.initializations.createInitDeferred( + initKey, + request.direction ?? "tail", + ); void deferred.promise.catch(() => undefined); } refreshAgentInitializationTimeout({ + store: sessionStore, key: initKey, agentId, setAgentInitializing, @@ -440,21 +442,29 @@ function SessionProviderInternal({ children, serverId, client }: SessionProvider setAgentInitializing(agentId, true); } try { - const page = await getHostRuntimeStore().fetchAgentTimeline(serverId, agentId, request); - if (shouldInitialize && getInitDeferred(initKey)) { - refreshAgentInitializationTimeout({ key: initKey, agentId, setAgentInitializing }); + const page = await runtime.fetchAgentTimeline(serverId, agentId, request); + if (shouldInitialize && sessionStore.initializations.getInitDeferred(initKey)) { + refreshAgentInitializationTimeout({ + store: sessionStore, + key: initKey, + agentId, + setAgentInitializing, + }); } return page; } catch (error) { if (shouldInitialize) { setAgentInitializing(agentId, false); - rejectInitDeferred(initKey, error instanceof Error ? error : new Error(String(error))); + sessionStore.initializations.rejectInitDeferred( + initKey, + error instanceof Error ? error : new Error(String(error)), + ); } throw error; } }, fetchLatestTail: (agentId) => - getHostRuntimeStore().fetchAgentTimeline(serverId, agentId, planTimelineTailFetch()), + runtime.fetchAgentTimeline(serverId, agentId, planTimelineTailFetch()), reportError: (error) => { console.warn("[Session] viewed timeline synchronization failed", { serverId, error }); }, @@ -478,7 +488,16 @@ function SessionProviderInternal({ children, serverId, client }: SessionProvider setViewedTimelineSync(serverId, null); sync.dispose(); }; - }, [client, serverId, setInitializingAgents, setViewedTimelineSync, voiceRuntime]); + }, [ + client, + serverId, + sessionStore, + runtime, + sessionStore.initializations, + setInitializingAgents, + setViewedTimelineSync, + voiceRuntime, + ]); useEffect(() => { viewedTimelineSyncRef.current?.setConnected(isConnected); diff --git a/packages/app/src/data/query-client.ts b/packages/app/src/data/query-client.ts index 748bf172e2..c001005d12 100644 --- a/packages/app/src/data/query-client.ts +++ b/packages/app/src/data/query-client.ts @@ -1,12 +1,15 @@ import { QueryClient } from "@tanstack/react-query"; -export const queryClient = new QueryClient({ - defaultOptions: { - queries: { - staleTime: Infinity, - refetchOnMount: false, - refetchOnReconnect: false, - refetchOnWindowFocus: false, +export const createQueryClient = () => + new QueryClient({ + defaultOptions: { + queries: { + staleTime: Infinity, + refetchOnMount: false, + refetchOnReconnect: false, + refetchOnWindowFocus: false, + }, }, - }, -}); + }); + +export const queryClient = createQueryClient(); diff --git a/packages/app/src/hooks/use-agent-initialization.test.ts b/packages/app/src/hooks/use-agent-initialization.test.ts index f2726fba02..64fbeea226 100644 --- a/packages/app/src/hooks/use-agent-initialization.test.ts +++ b/packages/app/src/hooks/use-agent-initialization.test.ts @@ -1,5 +1,5 @@ import { afterEach, describe, expect, it, vi } from "vitest"; -import { useSessionStore } from "@/stores/session-store"; +import { createSessionStore, useSessionStore } from "@/stores/session-store"; import { TIMELINE_FETCH_PAGE_SIZE } from "@/timeline/timeline-fetch-policy"; import type { HostRuntimeStore } from "@/runtime/host-runtime"; import { getInitDeferred, getInitKey, resolveInitDeferred } from "@/utils/agent-initialization"; @@ -288,3 +288,45 @@ describe("refreshAgent", () => { ]); }); }); + +it.each(["complete", "fail"] as const)( + "keeps routed initialization pending when direct initialization for the same address %s", + async (outcome) => { + const direct = createSessionStore(); + const routed = createSessionStore(); + for (const store of [direct, routed]) store.getState().initializeSession(serverId, null); + const directRuntime = new FakeTimelineRuntime(); + const routedRuntime = new FakeTimelineRuntime(); + const client = new FakeDaemonClient(); + const start = (store: typeof direct, runtime: FakeTimelineRuntime) => + ensureAgentIsInitialized({ + store, + serverId, + agentId, + client: client as never, + runtime, + setAgentInitializing: createSetAgentInitializing( + serverId, + store.getState().setInitializingAgents, + ), + }); + const directReady = start(direct, directRuntime); + const routedReady = start(routed, routedRuntime); + void directReady.catch(() => {}); + void routedReady.catch(() => {}); + try { + expect(directReady).not.toBe(routedReady); + const key = getInitKey(serverId, agentId); + if (outcome === "complete") direct.initializations.resolveInitDeferred(key); + else direct.initializations.rejectInitDeferred(key, new Error("Direct history failed")); + await Promise.allSettled([directReady]); + expect(routed.initializations.getInitDeferred(key)?.promise).toBe(routedReady); + routed.initializations.resolveInitDeferred(key); + await routedReady; + } finally { + direct.initializations.resolveInitDeferred(getInitKey(serverId, agentId)); + routed.initializations.resolveInitDeferred(getInitKey(serverId, agentId)); + resolveInitDeferred(getInitKey(serverId, agentId)); + } + }, +); diff --git a/packages/app/src/hooks/use-agent-initialization.ts b/packages/app/src/hooks/use-agent-initialization.ts index a19f36bf20..d70a00528c 100644 --- a/packages/app/src/hooks/use-agent-initialization.ts +++ b/packages/app/src/hooks/use-agent-initialization.ts @@ -1,16 +1,14 @@ import { useCallback, useMemo } from "react"; import { useTranslation } from "react-i18next"; import type { DaemonClient } from "@getpaseo/client/internal/daemon-client"; -import { selectAgentTimelineState, useSessionStore } from "@/stores/session-store"; import { - createInitDeferred, - getInitDeferred, - getInitKey, - INIT_TIMEOUT_MS, - rejectInitDeferred, - refreshInitTimeout, -} from "@/utils/agent-initialization"; -import { getHostRuntimeStore, type HostRuntimeStore } from "@/runtime/host-runtime"; + selectAgentTimelineState, + useSessionStore, + useSessionStoreApi, + type SessionStoreApi, +} from "@/stores/session-store"; +import { getInitKey, INIT_TIMEOUT_MS } from "@/utils/agent-initialization"; +import { useHostRuntimeStore, type HostRuntimeStore } from "@/runtime/host-runtime"; import { planTimelineResumeFetch, planTimelineTailFetch } from "@/timeline/timeline-sync-plan"; import { i18n } from "@/i18n/i18next"; @@ -22,19 +20,22 @@ export function createHistorySyncTimeoutError(): Error { export function refreshAgentInitializationTimeout(input: { key: string; + store?: SessionStoreApi; agentId: string; setAgentInitializing: SetAgentInitializing; }): void { - refreshInitTimeout({ + const initializations = (input.store ?? useSessionStore).initializations; + initializations.refreshInitTimeout({ key: input.key, onTimeout: () => { input.setAgentInitializing(input.agentId, false); - rejectInitDeferred(input.key, createHistorySyncTimeoutError()); + initializations.rejectInitDeferred(input.key, createHistorySyncTimeoutError()); }, }); } export interface EnsureAgentIsInitializedInput { + store?: SessionStoreApi; serverId: string; agentId: string; client: Pick | null; @@ -45,25 +46,27 @@ export interface EnsureAgentIsInitializedInput { export function ensureAgentIsInitialized(input: EnsureAgentIsInitializedInput): Promise { const { serverId, agentId, client, setAgentInitializing } = input; + const store = input.store ?? useSessionStore; + const initializations = store.initializations; const key = getInitKey(serverId, agentId); - const existing = getInitDeferred(key); + const existing = initializations.getInitDeferred(key); if (existing) { return existing.promise; } - const timeline = selectAgentTimelineState(useSessionStore.getState().sessions[serverId], agentId); + const timeline = selectAgentTimelineState(store.getState().sessions[serverId], agentId); const timelineRequest = planTimelineResumeFetch( timeline.status === "synced" ? timeline.range : null, ); - const deferred = createInitDeferred(key, timelineRequest.direction); - refreshAgentInitializationTimeout({ key, agentId, setAgentInitializing }); + const deferred = initializations.createInitDeferred(key, timelineRequest.direction); + refreshAgentInitializationTimeout({ key, store, agentId, setAgentInitializing }); setAgentInitializing(agentId, true); if (!client) { setAgentInitializing(agentId, false); - rejectInitDeferred( + initializations.rejectInitDeferred( key, new Error(input.hostDisconnectedMessage ?? i18n.t("workspace.terminal.hostDisconnected")), ); @@ -72,7 +75,10 @@ export function ensureAgentIsInitialized(input: EnsureAgentIsInitializedInput): input.runtime.fetchAgentTimeline(serverId, agentId, timelineRequest).catch((error) => { setAgentInitializing(agentId, false); - rejectInitDeferred(key, error instanceof Error ? error : new Error(String(error))); + initializations.rejectInitDeferred( + key, + error instanceof Error ? error : new Error(String(error)), + ); }); return deferred.promise; @@ -127,6 +133,8 @@ export function useAgentInitialization({ client: DaemonClient | null; }) { const { t } = useTranslation(); + const store = useSessionStoreApi(); + const runtime = useHostRuntimeStore(); const setInitializingAgents = useSessionStore((state) => state.setInitializingAgents); const setAgentInitializing = useMemo( () => createSetAgentInitializing(serverId, setInitializingAgents), @@ -136,14 +144,15 @@ export function useAgentInitialization({ const ensureAgentIsInitializedCallback = useCallback( (agentId: string): Promise => ensureAgentIsInitialized({ + store, serverId, agentId, client, - runtime: getHostRuntimeStore(), + runtime, setAgentInitializing, hostDisconnectedMessage: t("workspace.terminal.hostDisconnected"), }), - [client, serverId, setAgentInitializing, t], + [client, runtime, serverId, setAgentInitializing, store, t], ); const refreshAgentCallback = useCallback( @@ -152,11 +161,11 @@ export function useAgentInitialization({ serverId, agentId, client, - runtime: getHostRuntimeStore(), + runtime, setAgentInitializing, hostDisconnectedMessage: t("workspace.terminal.hostDisconnected"), }), - [client, serverId, setAgentInitializing, t], + [client, runtime, serverId, setAgentInitializing, t], ); return { diff --git a/packages/app/src/hooks/use-archive-agent.test.ts b/packages/app/src/hooks/use-archive-agent.test.ts index df14f2df38..8ce37cdb7e 100644 --- a/packages/app/src/hooks/use-archive-agent.test.ts +++ b/packages/app/src/hooks/use-archive-agent.test.ts @@ -1,3 +1,4 @@ +import { getHostRuntimeStore } from "@/runtime/host-runtime"; import { seedSessionHosts } from "@/test/seed-session"; import type { DaemonClient } from "@getpaseo/client/internal/daemon-client"; import { QueryClient } from "@tanstack/react-query"; @@ -161,6 +162,7 @@ describe("useArchiveAgent", () => { }); applyArchivedAgentCloseResults({ + runtime: getHostRuntimeStore(), queryClient, serverId: "server-a", results: [{ agentId: "agent-1", archivedAt: "2026-04-01T04:00:00.000Z" }], @@ -233,6 +235,7 @@ describe("useArchiveAgent", () => { }); applyArchivedAgentCloseResults({ + runtime: getHostRuntimeStore(), queryClient, serverId: "server-a", results: [{ agentId: "agent-1", archivedAt: "2026-04-01T04:00:00.000Z" }], diff --git a/packages/app/src/hooks/use-archive-agent.ts b/packages/app/src/hooks/use-archive-agent.ts index 92d8e4d2d6..48dd253d8f 100644 --- a/packages/app/src/hooks/use-archive-agent.ts +++ b/packages/app/src/hooks/use-archive-agent.ts @@ -1,14 +1,8 @@ -import { getHostRuntimeStore } from "@/runtime/host-runtime"; +import { useHostRuntimeStore, type HostRuntimeStore } from "@/runtime/host-runtime"; import { useCallback, useMemo } from "react"; -import { - useMutation, - useQuery, - useQueryClient, - type QueryClient, - type QueryKey, -} from "@tanstack/react-query"; +import { useQuery, useQueryClient, type QueryClient, type QueryKey } from "@tanstack/react-query"; import { useTranslation } from "react-i18next"; -import { useSessionStore } from "@/stores/session-store"; +import { useSessionStoreApi, type SessionStoreApi } from "@/stores/session-store"; import { agentHistoryQueryKey, allAgentHistoryQueryRootKey } from "./agent-history-query-key"; export const ARCHIVE_AGENT_PENDING_QUERY_KEY = ["archive-agent-pending"] as const; @@ -241,14 +235,15 @@ interface ArchiveAgentMutationContext { lists: ArchivedAgentListCacheSnapshot; } -function getStoredAgentSnapshot(input: ArchiveAgentInput) { - return useSessionStore.getState().sessions[input.serverId]?.agents.get(input.agentId); +function getStoredAgentSnapshot(store: SessionStoreApi, input: ArchiveAgentInput) { + return store.getState().sessions[input.serverId]?.agents.get(input.agentId); } function restoreAgentSnapshot( + runtime: HostRuntimeStore, input: ArchiveAgentInput & { agent: ReturnType }, ): void { - getHostRuntimeStore().restoreAgentSnapshot(input.serverId, input.agentId, input.agent); + runtime.restoreAgentSnapshot(input.serverId, input.agentId, input.agent); } function getArchivedAgentListCacheSnapshot( @@ -299,20 +294,20 @@ function restoreArchivedAgentListCacheSnapshot( } } -function markAgentArchivedInStore(input: ArchiveAgentInput & { archivedAt: string }): void { +function markAgentArchivedInStore( + runtime: HostRuntimeStore, + input: ArchiveAgentInput & { archivedAt: string }, +): void { const archivedAt = new Date(input.archivedAt); if (Number.isNaN(archivedAt.getTime())) { return; } - getHostRuntimeStore().archiveAgentSnapshot( - input.serverId, - input.agentId, - archivedAt.toISOString(), - ); + runtime.archiveAgentSnapshot(input.serverId, input.agentId, archivedAt.toISOString()); } interface ApplyArchivedAgentCloseResultsInput { + runtime: HostRuntimeStore; queryClient: QueryClient; serverId: string; results: ArchivedAgentCloseResult[]; @@ -325,7 +320,7 @@ export function applyArchivedAgentCloseResults(input: ApplyArchivedAgentCloseRes } for (const result of input.results) { - markAgentArchivedInStore({ + markAgentArchivedInStore(input.runtime, { serverId: input.serverId, agentId: result.agentId, archivedAt: result.archivedAt, @@ -375,87 +370,69 @@ export function usePendingArchiveAgentIds(serverId: string): ReadonlySet ); } +/** The action owns optimistic state and rollback within its captured authority. */ +export function createAgentArchiveAction({ + runtime, + sessionStore, + queryClient, + unavailableMessage, +}: { + runtime: HostRuntimeStore; + sessionStore: SessionStoreApi; + queryClient: QueryClient; + unavailableMessage: string; +}): (input: ArchiveAgentInput) => Promise { + return async (input) => { + const context: ArchiveAgentMutationContext = { + agent: getStoredAgentSnapshot(sessionStore, input), + lists: getArchivedAgentListCacheSnapshot(queryClient, input.serverId), + }; + applyArchivedAgentCloseResults({ + runtime, + queryClient, + serverId: input.serverId, + results: [{ agentId: input.agentId, archivedAt: new Date().toISOString() }], + invalidateQueries: false, + }); + setAgentArchiving({ ...input, queryClient, isArchiving: true }); + try { + const client = runtime.getClient(input.serverId); + if (!client) throw new Error(unavailableMessage); + const result = await client.archiveAgent(input.agentId); + markAgentArchivedInStore(runtime, { ...input, archivedAt: result.archivedAt }); + } catch (error) { + restoreAgentSnapshot(runtime, { ...input, agent: context.agent }); + restoreArchivedAgentListCacheSnapshot(queryClient, input.serverId, context.lists); + throw error; + } finally { + clearArchiveAgentPending({ ...input, queryClient }); + for (const queryKey of [ + ["sidebarAgentsList", input.serverId], + ["allAgents", input.serverId], + agentHistoryQueryKey(input.serverId), + allAgentHistoryQueryRootKey(), + ]) + void queryClient.invalidateQueries({ queryKey }); + } + }; +} + export function useArchiveAgent() { const queryClient = useQueryClient(); + const runtime = useHostRuntimeStore(); + const sessionStore = useSessionStoreApi(); const { t } = useTranslation(); - const pendingQuery = useArchiveAgentPendingQuery(); - - const archiveMutation = useMutation({ - mutationFn: async (input: ArchiveAgentInput): Promise<{ archivedAt: string }> => { - const client = useSessionStore.getState().sessions[input.serverId]?.client ?? null; - if (!client) { - throw new Error(t("common.errors.daemonClientUnavailable")); - } - return await client.archiveAgent(input.agentId); - }, - onMutate: (input) => { - const context: ArchiveAgentMutationContext = { - agent: getStoredAgentSnapshot(input), - lists: getArchivedAgentListCacheSnapshot(queryClient, input.serverId), - }; - const archivedAt = new Date().toISOString(); - - applyArchivedAgentCloseResults({ + const unavailableMessage = t("common.errors.daemonClientUnavailable"); + const archiveAgent = useMemo( + () => + createAgentArchiveAction({ + runtime, + sessionStore, queryClient, - serverId: input.serverId, - results: [{ agentId: input.agentId, archivedAt }], - invalidateQueries: false, - }); - setAgentArchiving({ - queryClient, - serverId: input.serverId, - agentId: input.agentId, - isArchiving: true, - }); - return context; - }, - onSuccess: (result, input) => { - markAgentArchivedInStore({ - serverId: input.serverId, - agentId: input.agentId, - archivedAt: result.archivedAt, - }); - }, - onError: (_error, input, context) => { - if (!context) { - return; - } - restoreAgentSnapshot({ - serverId: input.serverId, - agentId: input.agentId, - agent: context.agent, - }); - restoreArchivedAgentListCacheSnapshot(queryClient, input.serverId, context.lists); - }, - onSettled: (_result, _error, input) => { - clearArchiveAgentPending({ - queryClient, - serverId: input.serverId, - agentId: input.agentId, - }); - void queryClient.invalidateQueries({ - queryKey: ["sidebarAgentsList", input.serverId], - }); - void queryClient.invalidateQueries({ - queryKey: ["allAgents", input.serverId], - }); - void queryClient.invalidateQueries({ - queryKey: agentHistoryQueryKey(input.serverId), - }); - void queryClient.invalidateQueries({ - queryKey: allAgentHistoryQueryRootKey(), - }); - }, - }); - - const archiveMutateAsync = archiveMutation.mutateAsync; - - const archiveAgent = useCallback( - async (input: ArchiveAgentInput): Promise => { - await archiveMutateAsync(input); - }, - [archiveMutateAsync], + unavailableMessage, + }), + [runtime, sessionStore, queryClient, unavailableMessage], ); const isArchivingAgent = useCallback( diff --git a/packages/app/src/panels/agent-panel.tsx b/packages/app/src/panels/agent-panel.tsx index 2fa2624fb4..ae4bb16d4d 100644 --- a/packages/app/src/panels/agent-panel.tsx +++ b/packages/app/src/panels/agent-panel.tsx @@ -1,4 +1,4 @@ -import { getHostRuntimeStore } from "@/runtime/host-runtime"; +import { useHostRuntimeStore, type HostRuntimeStore } from "@/runtime/host-runtime"; import { Button } from "@/components/ui/button"; import { LoadingSpinner } from "@/components/ui/loading-spinner"; import type { DaemonClient } from "@getpaseo/client/internal/daemon-client"; @@ -87,6 +87,8 @@ import { selectAgentTurnPresentation, type Agent, useSessionStore, + useSessionStoreApi, + type SessionStoreApi, } from "@/stores/session-store"; import { useWorkspaceLayoutStore } from "@/stores/workspace-layout-store"; import { buildWorkspaceTabPersistenceKey } from "@/workspace-tabs/model"; @@ -97,7 +99,7 @@ import type { PendingPermission } from "@/types/shared"; import type { StreamItem, TodoEntry } from "@/types/stream"; import type { ViewedTimelineStatus, ViewedTimelineUiBridge } from "@/timeline/viewed-timeline-sync"; import { useArchiveFinishedSubagents, useSubagentsForParent } from "@/subagents"; -import { getInitDeferred, getInitKey } from "@/utils/agent-initialization"; +import { getInitKey } from "@/utils/agent-initialization"; import { derivePendingPermissionKey, normalizeAgentSnapshot } from "@/utils/agent-snapshots"; import { applyLegacyDaemonWorkspaceOwnership } from "@/workspace/legacy-daemon-workspaces"; import type { WorkspaceFileOpenRequest } from "@/workspace/file-open"; @@ -288,6 +290,8 @@ function shouldStoreFetchedAgentInActiveDirectory(agent: Agent): boolean { type FetchAgentResult = Awaited>; function storeFetchedAgentDetail(input: { + runtime: HostRuntimeStore; + sessionStore: SessionStoreApi; serverId: string; result: NonNullable; }): Agent { @@ -299,10 +303,10 @@ function storeFetchedAgentDetail(input: { projectPlacement: input.result.project, }, }); - const store = useSessionStore.getState(); + const store = input.sessionStore.getState(); if (shouldStoreFetchedAgentInActiveDirectory(hydrated)) { - getHostRuntimeStore().acceptAgentSnapshot(input.serverId, hydrated); + input.runtime.acceptAgentSnapshot(input.serverId, hydrated); } else { store.setAgentDetails(input.serverId, (previous) => { const next = new Map(previous); @@ -388,6 +392,7 @@ function AgentPanel() { } function DraftPanel() { + const runtime = useHostRuntimeStore(); const { serverId, workspaceId, @@ -404,10 +409,10 @@ function DraftPanel() { (agentSnapshot: Parameters[0]) => { const normalized = normalizeAgentSnapshot(agentSnapshot, serverId); const agent = applyLegacyDaemonWorkspaceOwnership({ serverId, agent: normalized }); - getHostRuntimeStore().acceptAgentSnapshot(serverId, agent); + runtime.acceptAgentSnapshot(serverId, agent); retargetCurrentTab({ kind: "agent", agentId: agentSnapshot.id }); }, - [retargetCurrentTab, serverId], + [retargetCurrentTab, runtime, serverId], ); return ( @@ -593,10 +598,12 @@ function AgentPanelBody({ onOpenWorkspaceFile?: (request: WorkspaceFileOpenRequest) => void; }) { const { t } = useTranslation(); + const sessionStore = useSessionStoreApi(); + const runtime = useHostRuntimeStore(); const { isArchivingAgent: _isArchivingAgent } = useArchiveAgent(); const hasSession = useSessionStore((state) => Boolean(state.sessions[serverId])); const projectPlacement = useStoreWithEqualityFn( - useSessionStore, + useSessionStoreApi(), (state) => { if (!agentId) { return null; @@ -656,7 +663,7 @@ function AgentPanelBody({ return; } - storeFetchedAgentDetail({ serverId, result }); + storeFetchedAgentDetail({ serverId, result, sessionStore, runtime }); setLookupState({ tag: "idle" }); return; }) @@ -671,7 +678,17 @@ function AgentPanelBody({ } setLookupState({ tag: "error", message }); }); - }, [agentId, agentState.id, client, hasSession, isConnected, lookupState.tag, serverId]); + }, [ + agentId, + agentState.id, + client, + hasSession, + isConnected, + lookupState.tag, + runtime, + serverId, + sessionStore, + ]); if (lookupState.tag === "not_found") { return ( @@ -748,6 +765,8 @@ function ChatAgentContent({ onOpenWorkspaceFile?: (request: WorkspaceFileOpenRequest) => void; }) { const { t } = useTranslation(); + const sessionStore = useSessionStoreApi(); + const runtime = useHostRuntimeStore(); const isPaneVisible = useRetainedPanelActive(); const { api: toastApi, toast: toastState, dismiss: dismissToast } = useToastHost(); const { isArchivingAgent } = useArchiveAgent(); @@ -763,7 +782,7 @@ function ChatAgentContent({ useShallow((state) => selectChatAgentState(state, serverId, agentId)), ); const projectPlacement = useStoreWithEqualityFn( - useSessionStore, + useSessionStoreApi(), (state) => { if (!agentId) { return null; @@ -863,8 +882,8 @@ function ChatAgentContent({ return false; } const initKey = getInitKey(serverId, agentId); - return Boolean(getInitDeferred(initKey)); - }, [agentId, isInitializing, serverId]); + return Boolean(sessionStore.initializations.getInitDeferred(initKey)); + }, [agentId, isInitializing, serverId, sessionStore]); const needsAuthoritativeSync = useMemo(() => { if (!agentId) { return false; @@ -1005,7 +1024,7 @@ function ChatAgentContent({ if (attemptToken !== initAttemptTokenRef.current) { return; } - const currentSession = useSessionStore.getState().sessions[serverId]; + const currentSession = sessionStore.getState().sessions[serverId]; const currentAgent = currentSession?.agents.get(agentId) ?? currentSession?.agentDetails.get(agentId); if (!currentAgent) { @@ -1020,7 +1039,7 @@ function ChatAgentContent({ }); return; } - storeFetchedAgentDetail({ serverId, result }); + storeFetchedAgentDetail({ serverId, result, sessionStore, runtime }); } if (attemptToken !== initAttemptTokenRef.current) { return; @@ -1050,6 +1069,8 @@ function ChatAgentContent({ isPaneVisible, missingAgentState.kind, serverId, + sessionStore, + runtime, viewedTimelineSync, visibilityCatchUpStatus, ]); @@ -1440,7 +1461,7 @@ const AgentStreamSection = memo(function AgentStreamSection({ ); const streamItems = streamItemsRaw ?? EMPTY_STREAM_ITEMS; const pendingPermissionList = useStoreWithEqualityFn( - useSessionStore, + useSessionStoreApi(), (state) => { if (!agentId) { return EMPTY_PENDING_PERMISSION_LIST; @@ -1573,6 +1594,7 @@ function ActiveAgentComposer({ { initialIsBelow: isCompactFormFactor }, ); const paneContext = usePaneContext(); + const sessionStore = useSessionStoreApi(); const openInSidePane = useSettings((settings) => settings.openInSidePane); const { workspaceId, tabId, retargetCurrentTab } = paneContext; const { archiveAgent } = useArchiveAgent(); @@ -1604,7 +1626,7 @@ function ActiveAgentComposer({ const handleClientSlashCommand = useCallback( async (command: ClientSlashCommand) => { - const agent = resolveChatAgentFromSession(useSessionStore.getState(), serverId, agentId); + const agent = resolveChatAgentFromSession(sessionStore.getState(), serverId, agentId); if (!agent) { throw new Error("Agent not found"); } @@ -1634,6 +1656,7 @@ function ActiveAgentComposer({ hideWorkspaceAgent, retargetCurrentTab, serverId, + sessionStore, tabId, unpinWorkspaceAgent, workspaceId, diff --git a/packages/app/src/runtime/directory-sync/agent-replica.test.ts b/packages/app/src/runtime/directory-sync/agent-replica.test.ts index 648c66c590..4476c4a4ed 100644 --- a/packages/app/src/runtime/directory-sync/agent-replica.test.ts +++ b/packages/app/src/runtime/directory-sync/agent-replica.test.ts @@ -2,6 +2,7 @@ import { describe, expect, it } from "vitest"; import type { DaemonClient, FetchAgentsEntry } from "@getpaseo/client/internal/daemon-client"; import type { AgentSnapshotPayload } from "@getpaseo/protocol/messages"; import { + createSessionStore, selectAgentTimelineState, selectAgentTurnPresentation, useSessionStore, @@ -391,3 +392,34 @@ describe("AgentDirectoryReplica", () => { store.clearSession(serverId); }); }); + +it("partitions equal destination identities by the owning authority store, including late snapshots", () => { + const direct = createSessionStore(); + const routed = createSessionStore(); + const serverId = "srv_B"; + direct.getState().initializeSession(serverId, null); + routed.getState().initializeSession(serverId, null); + const directReplica = new AgentDirectoryReplica( + serverId, + () => {}, + () => {}, + direct, + ); + const routedReplica = new AgentDirectoryReplica( + serverId, + () => {}, + () => {}, + routed, + ); + directReplica.commitSnapshot([entry(payload("Broader direct data"))], []); + expect(routed.getState().sessions[serverId]?.agents.size).toBe(0); + routedReplica.commitSnapshot([entry(payload("Origin-authorized data"))], []); + directReplica.commitSnapshot([entry(payload("Late direct snapshot"))], []); + expect(routed.getState().sessions[serverId]?.agents.get("agent")?.title).toBe( + "Origin-authorized data", + ); + expect(direct.getState().sessions[serverId]?.agents.get("agent")?.title).toBe( + "Late direct snapshot", + ); + expect(routedReplica.snapshot().get("agent")?.title).toBe("Origin-authorized data"); +}); diff --git a/packages/app/src/runtime/directory-sync/agent-replica.ts b/packages/app/src/runtime/directory-sync/agent-replica.ts index 37d5a1e4cb..2087c5f49f 100644 --- a/packages/app/src/runtime/directory-sync/agent-replica.ts +++ b/packages/app/src/runtime/directory-sync/agent-replica.ts @@ -2,7 +2,7 @@ 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 type { Agent } from "@/stores/session-store"; +import { useSessionStore, type SessionStoreApi, type Agent } from "@/stores/session-store"; import { normalizeAgentSnapshot, projectAgentSnapshot } from "@/utils/agent-snapshots"; import { type AgentDirectoryDelta } from "@/utils/agent-directory-sync"; import { reconcileAgentDirectory } from "@/utils/agent-directory-reconciliation"; @@ -38,8 +38,9 @@ export class AgentDirectoryReplica { private readonly serverId: string, private readonly onStoppedRunning: (agentId: string) => void, private readonly persist: (mutations: readonly DirectoryReplicaMutation[]) => void, + store: SessionStoreApi = useSessionStore, ) { - this.storeProjection = new AgentStoreProjection(serverId); + this.storeProjection = new AgentStoreProjection(serverId, store); } captureTimeline(agentId: string): AgentLifecycleToken { diff --git a/packages/app/src/runtime/directory-sync/index.ts b/packages/app/src/runtime/directory-sync/index.ts index b59440d782..6903d648b9 100644 --- a/packages/app/src/runtime/directory-sync/index.ts +++ b/packages/app/src/runtime/directory-sync/index.ts @@ -9,6 +9,7 @@ import { normalizeProjectDescriptor, normalizeWorkspaceDescriptor, useSessionStore, + type SessionStoreApi, type Agent, type ProjectDescriptor, type WorkspaceDescriptor, @@ -172,11 +173,17 @@ export class DirectorySync { markAgentError: (error: string) => void; }, private readonly checkpoints?: DirectoryCheckpointStorage, + private readonly store: SessionStoreApi = useSessionStore, ) { const persist = (mutations: readonly DirectoryReplicaMutation[]) => this.checkpoints?.commitDirectoryMutations(this.serverId, mutations); - this.agents = new AgentDirectoryReplica(serverId, callbacks.onAgentStoppedRunning, persist); - this.workspaces = new WorkspaceDirectoryReplica(serverId); + this.agents = new AgentDirectoryReplica( + serverId, + callbacks.onAgentStoppedRunning, + persist, + store, + ); + this.workspaces = new WorkspaceDirectoryReplica(serverId, store); } connectionChanged(connection: DirectoryConnection): boolean { @@ -401,18 +408,18 @@ export class DirectorySync { async loadCachedAgent(agentId: string): Promise { if (!this.checkpoints) return; - if (useSessionStore.getState().sessions[this.serverId]?.agents.has(agentId)) return; + if (this.store.getState().sessions[this.serverId]?.agents.has(agentId)) return; const token = this.agents.captureCache(agentId); const agent = await this.checkpoints.readAgent(this.serverId, agentId); if (!agent) return; - const session = useSessionStore.getState().sessions[this.serverId]; + const session = this.store.getState().sessions[this.serverId]; if (!session || session.agents.has(agentId)) return; this.agents.commitCachedAgent(token, agent); } async prepareAgentRoute(agentId: string): Promise { await this.loadCachedAgent(agentId); - const agent = useSessionStore.getState().sessions[this.serverId]?.agents.get(agentId); + const agent = this.store.getState().sessions[this.serverId]?.agents.get(agentId); if (agent?.workspaceId) await this.loadCachedWorkspace(agent.workspaceId); } @@ -422,12 +429,12 @@ export class DirectorySync { private async loadCachedWorkspace(workspaceId: string): Promise { if (!this.checkpoints) return; - if (useSessionStore.getState().sessions[this.serverId]?.workspaces.has(workspaceId)) return; + if (this.store.getState().sessions[this.serverId]?.workspaces.has(workspaceId)) return; const revision = this.workspaceRevision; const cached = await this.checkpoints.readWorkspace(this.serverId, workspaceId); if (!cached) return; if (this.workspaceRevision !== revision) return; - const session = useSessionStore.getState().sessions[this.serverId]; + const session = this.store.getState().sessions[this.serverId]; if (!session) return; this.workspaces.commitCachedWorkspace(cached.workspace, cached.project); } @@ -439,7 +446,7 @@ export class DirectorySync { const revision = this.revision; const cached = await checkpoints.readDirectory(this.serverId); if (this.cacheAccepted) return; - if (!useSessionStore.getState().sessions[this.serverId]) return; + if (!this.store.getState().sessions[this.serverId]) return; this.agents.commitCached(cached.agents); this.workspaces.commitCached(cached); if (this.revision === revision) this.cursors = cached.checkpoint ?? {}; @@ -495,7 +502,7 @@ export class DirectorySync { if (!onlineConnection) { this.callbacks.markAgentReady(); return { - agents: new Map(useSessionStore.getState().sessions[this.serverId]?.agents), + agents: new Map(this.store.getState().sessions[this.serverId]?.agents), subscriptionId: null, }; } @@ -560,7 +567,7 @@ export class DirectorySync { }); try { await this.waitForSessionMetadata(client, source); - const serverInfo = useSessionStore.getState().sessions[this.serverId]?.serverInfo; + const serverInfo = this.store.getState().sessions[this.serverId]?.serverInfo; const supportsProjectList = serverInfo?.features?.projectList === true; const supportsDirectorySync = serverInfo?.features?.directorySync === true; const projectsReady = supportsProjectList @@ -690,7 +697,7 @@ export class DirectorySync { } markWorkspacesHydrated(hydrated: boolean): void { - useSessionStore.getState().setHasHydratedWorkspaces(this.serverId, hydrated); + this.store.getState().setHasHydratedWorkspaces(this.serverId, hydrated); } private async fetchAgents( @@ -783,8 +790,7 @@ export class DirectorySync { private supportsDirectorySync(): boolean { return ( - useSessionStore.getState().sessions[this.serverId]?.serverInfo?.features?.directorySync === - true + this.store.getState().sessions[this.serverId]?.serverInfo?.features?.directorySync === true ); } @@ -964,7 +970,7 @@ export class DirectorySync { source: DirectorySourceToken, ): Promise { await this.waitForSessionState(client, source, () => { - const session = useSessionStore.getState().sessions[this.serverId]; + const session = this.store.getState().sessions[this.serverId]; return this.hasMatchingSession(client, source) && session?.serverInfo !== null; }); } @@ -995,13 +1001,13 @@ export class DirectorySync { } }; this.abortSessionWaits.add(abort); - unsubscribe = useSessionStore.subscribe(check); + unsubscribe = this.store.subscribe(check); check(); }); } private hasMatchingSession(client: DaemonClient, source: DirectorySourceToken): boolean { - const session = useSessionStore.getState().sessions[this.serverId]; + const session = this.store.getState().sessions[this.serverId]; return session?.client === client && session.clientGeneration === source.clientGeneration; } diff --git a/packages/app/src/runtime/directory-sync/internal/agent-store.ts b/packages/app/src/runtime/directory-sync/internal/agent-store.ts index 97c8208da5..659eaa367c 100644 --- a/packages/app/src/runtime/directory-sync/internal/agent-store.ts +++ b/packages/app/src/runtime/directory-sync/internal/agent-store.ts @@ -1,6 +1,6 @@ import equal from "fast-deep-equal"; import type { FetchAgentsEntry } from "@getpaseo/client/internal/daemon-client"; -import { useSessionStore, type Agent } from "@/stores/session-store"; +import { useSessionStore, type SessionStoreApi, type Agent } from "@/stores/session-store"; import { acceptAgentDirectoryUpdate } from "@/utils/agent-directory-update-policy"; import { buildAgentDirectoryState, type AgentDirectoryDelta } from "@/utils/agent-directory-sync"; import { derivePendingPermissionKey, normalizeAgentSnapshot } from "@/utils/agent-snapshots"; @@ -9,7 +9,7 @@ import { clearArchiveAgentPending } from "@/hooks/use-archive-agent"; import { queryClient } from "@/data/query-client"; import { buildDraftStoreKey } from "@/stores/draft-keys"; import { useDraftStore } from "@/stores/draft-store"; -import { getInitDeferred, getInitKey, rejectInitDeferred } from "@/utils/agent-initialization"; +import { getInitKey } from "@/utils/agent-initialization"; import { reduceTurnLiveness, type TurnLivenessTransition } from "@/timeline/turn-liveness"; function withoutAgent(agentId: string) { @@ -32,14 +32,17 @@ function mergeSnapshotTurn(previous: Agent | undefined, incoming: Agent): Agent } export class AgentStoreProjection { - constructor(private readonly serverId: string) {} + constructor( + private readonly serverId: string, + private readonly store: SessionStoreApi = useSessionStore, + ) {} snapshot(): Map { - return useSessionStore.getState().sessions[this.serverId]?.agents ?? new Map(); + return this.store.getState().sessions[this.serverId]?.agents ?? new Map(); } get(agentId: string): Agent | undefined { - const session = useSessionStore.getState().sessions[this.serverId]; + const session = this.store.getState().sessions[this.serverId]; return session?.agents.get(agentId) ?? session?.agentDetails.get(agentId); } @@ -47,17 +50,17 @@ export class AgentStoreProjection { const merged = new Map(agents); for (const [agentId, agent] of this.snapshot()) merged.set(agentId, agent); for (const agent of merged.values()) this.publishActivity(agent); - useSessionStore.getState().setAgents(this.serverId, merged); + this.store.getState().setAgents(this.serverId, merged); return merged; } publishActivity(agent: Agent): void { - useSessionStore.getState().setAgentLastActivity(agent.id, agent.lastActivityAt); + this.store.getState().setAgentLastActivity(agent.id, agent.lastActivityAt); } accept(agent: Agent): Agent { let accepted = agent; - useSessionStore.getState().setAgents(this.serverId, (current) => { + this.store.getState().setAgents(this.serverId, (current) => { const previous = current.get(agent.id); accepted = acceptAgentDirectoryUpdate(previous, mergeSnapshotTurn(previous, agent)); if (accepted === previous) return current; @@ -78,7 +81,7 @@ export class AgentStoreProjection { return { agentId: delta.agentId, stoppedRunning: false }; } const normalized = normalizeAgentSnapshot(delta.agent, this.serverId); - const session = useSessionStore.getState().sessions[this.serverId]; + const session = this.store.getState().sessions[this.serverId]; const previous = session?.agents.get(normalized.id) ?? session?.agentDetails.get(normalized.id); const legacyWorkspaceId = previous?.workspaceId ?? @@ -98,7 +101,7 @@ export class AgentStoreProjection { clearArchiveAgentPending({ queryClient, serverId: this.serverId, agentId: accepted.id }); } this.replacePendingPermissions(accepted); - useSessionStore.getState().setAgentLastActivity(accepted.id, accepted.lastActivityAt); + this.store.getState().setAgentLastActivity(accepted.id, accepted.lastActivityAt); return { agentId: accepted.id, stoppedRunning: previous?.turn.phase === "open" && accepted.turn.phase === "idle", @@ -111,7 +114,7 @@ export class AgentStoreProjection { serverId: this.serverId, entries, }); - const store = useSessionStore.getState(); + const store = this.store.getState(); const current = store.sessions[this.serverId]?.agents ?? new Map(); const agents = new Map(); for (const [agentId, fetchedAgent] of fetched) { @@ -141,17 +144,17 @@ export class AgentStoreProjection { } replacePendingPermissions(agent: Agent): void { - const pending = new Map(useSessionStore.getState().sessions[this.serverId]?.pendingPermissions); + const pending = new Map(this.store.getState().sessions[this.serverId]?.pendingPermissions); for (const [key, entry] of pending) if (entry.agentId === agent.id) pending.delete(key); for (const request of agent.pendingPermissions) { const key = derivePendingPermissionKey(agent.id, request); pending.set(key, { key, agentId: agent.id, request }); } - useSessionStore.getState().setPendingPermissions(this.serverId, pending); + this.store.getState().setPendingPermissions(this.serverId, pending); } removeFromDirectory(agentId: string): void { - const store = useSessionStore.getState(); + const store = this.store.getState(); const removeKey = withoutAgent(agentId); clearArchiveAgentPending({ queryClient, serverId: this.serverId, agentId }); store.setAgents(this.serverId, removeKey); @@ -163,7 +166,7 @@ export class AgentStoreProjection { for (const [key, pending] of next) if (pending.agentId === agentId) next.delete(key); return next.size === current.size ? current : next; }); - useSessionStore.setState((state) => { + this.store.setState((state) => { if (!state.agentLastActivity.has(agentId)) return state; const agentLastActivity = new Map(state.agentLastActivity); agentLastActivity.delete(agentId); @@ -173,8 +176,11 @@ export class AgentStoreProjection { draftKey: buildDraftStoreKey({ serverId: this.serverId, agentId }), }); const initKey = getInitKey(this.serverId, agentId); - if (getInitDeferred(initKey)) { - rejectInitDeferred(initKey, new Error("Agent was removed during initialization")); + if (this.store.initializations.getInitDeferred(initKey)) { + this.store.initializations.rejectInitDeferred( + initKey, + new Error("Agent was removed during initialization"), + ); } } @@ -182,7 +188,7 @@ export class AgentStoreProjection { this.removeFromDirectory(agentId); // Only entity deletion destroys the transcript. A filtered active-list // response can arrive after a fresh history response for an archived agent. - const store = useSessionStore.getState(); + const store = this.store.getState(); const removeKey = withoutAgent(agentId); store.setAgentTimelineCursor(this.serverId, removeKey); store.setAgentAuthoritativeHistoryApplied(this.serverId, agentId, false); @@ -192,7 +198,7 @@ export class AgentStoreProjection { archive(agentId: string, archivedAt: string): Agent | null { let archived: Agent | null = null; - useSessionStore.getState().setAgents(this.serverId, (current) => { + this.store.getState().setAgents(this.serverId, (current) => { const agent = current.get(agentId); if (!agent) return current; archived = { @@ -213,7 +219,7 @@ export class AgentStoreProjection { ): Agent | null { const transitions = Array.isArray(transition) ? transition : [transition]; let accepted: Agent | null = null; - useSessionStore.getState().setAgents(this.serverId, (current) => { + this.store.getState().setAgents(this.serverId, (current) => { const agent = current.get(agentId); if (!agent) return current; const turn = transitions.reduce(reduceTurnLiveness, agent.turn); diff --git a/packages/app/src/runtime/directory-sync/workspace-replica.ts b/packages/app/src/runtime/directory-sync/workspace-replica.ts index 2eb107ac4f..ff3b80ba2c 100644 --- a/packages/app/src/runtime/directory-sync/workspace-replica.ts +++ b/packages/app/src/runtime/directory-sync/workspace-replica.ts @@ -7,6 +7,7 @@ import { normalizeProjectDescriptor, normalizeWorkspaceDescriptor, useSessionStore, + type SessionStoreApi, type ProjectDescriptor, type WorkspaceDescriptor, } from "@/stores/session-store"; @@ -38,7 +39,10 @@ export class WorkspaceDirectoryReplica { private projects = new Map(); private workspaceIdsByProject = new Map>(); - constructor(private readonly serverId: string) {} + constructor( + private readonly serverId: string, + private readonly store: SessionStoreApi = useSessionStore, + ) {} applyDelta(delta: WorkspaceDirectoryDelta): DirectoryReplicaMutation[] { if (delta.kind === "script_status") return this.applyScriptStatus(delta.update); @@ -55,7 +59,7 @@ export class WorkspaceDirectoryReplica { workspaces: new Map([...input.workspaces, ...this.workspaces]), projects: new Map([...input.projects, ...this.projects]), }); - useSessionStore.getState().setHasWorkspaceDirectorySnapshot(this.serverId, true); + this.store.getState().setHasWorkspaceDirectorySnapshot(this.serverId, true); } commitCachedWorkspace( @@ -73,7 +77,7 @@ export class WorkspaceDirectoryReplica { ): DirectoryReplicaMutation[] { this.replace(snapshot); const mutations = deltas.flatMap((delta) => this.applyDelta(delta)); - useSessionStore.getState().setHasHydratedWorkspaces(this.serverId, true); + this.store.getState().setHasHydratedWorkspaces(this.serverId, true); return mutations; } @@ -106,7 +110,7 @@ export class WorkspaceDirectoryReplica { this.workspaces = workspaces; this.projects = new Map(snapshot.projects); this.rebuildProjectIndex(); - const store = useSessionStore.getState(); + const store = this.store.getState(); store.setWorkspaces(this.serverId, this.workspaces); store.setProjects(this.serverId, this.projects.values()); } @@ -133,7 +137,7 @@ export class WorkspaceDirectoryReplica { { kind: "project", type: "delete", id: delta.projectId }, ]; this.projects.delete(delta.projectId); - useSessionStore.getState().removeProject(this.serverId, delta.projectId); + this.store.getState().removeProject(this.serverId, delta.projectId); for (const workspaceId of this.workspaceIdsByProject.get(delta.projectId) ?? []) { this.deleteWorkspace(workspaceId); mutations.push({ kind: "workspace", type: "delete", id: workspaceId }); @@ -179,7 +183,7 @@ export class WorkspaceDirectoryReplica { } if (delta.removedProjectId) { this.projects.delete(delta.removedProjectId); - useSessionStore.getState().removeProject(this.serverId, delta.removedProjectId); + this.store.getState().removeProject(this.serverId, delta.removedProjectId); mutations.push({ kind: "project", type: "delete", id: delta.removedProjectId }); } clearWorkspaceArchivePending({ serverId: this.serverId, workspaceId: delta.id }); @@ -217,19 +221,19 @@ export class WorkspaceDirectoryReplica { const projectWorkspaces = this.workspaceIdsByProject.get(workspace.projectId) ?? new Set(); projectWorkspaces.add(workspace.id); this.workspaceIdsByProject.set(workspace.projectId, projectWorkspaces); - useSessionStore.getState().mergeWorkspaces(this.serverId, [workspace]); + this.store.getState().mergeWorkspaces(this.serverId, [workspace]); } private deleteWorkspace(workspaceId: string): void { const workspace = this.workspaces.get(workspaceId); if (workspace) this.workspaceIdsByProject.get(workspace.projectId)?.delete(workspaceId); this.workspaces.delete(workspaceId); - useSessionStore.getState().removeWorkspace(this.serverId, workspaceId); + this.store.getState().removeWorkspace(this.serverId, workspaceId); } private setProject(project: ProjectDescriptor): void { this.projects.set(project.projectId, project); - useSessionStore.getState().upsertProject(this.serverId, project); + this.store.getState().upsertProject(this.serverId, project); } private rebuildProjectIndex(): void { diff --git a/packages/app/src/runtime/host-authority-provider.test.tsx b/packages/app/src/runtime/host-authority-provider.test.tsx new file mode 100644 index 0000000000..4d22934787 --- /dev/null +++ b/packages/app/src/runtime/host-authority-provider.test.tsx @@ -0,0 +1,44 @@ +import React from "react"; +import { renderToStaticMarkup } from "react-dom/server"; +import { expect, it } from "vitest"; +import { QueryClient, useQueryClient } from "@tanstack/react-query"; +import { createSessionStore, useSessionStoreApi } from "@/stores/session-store"; +import { useCreateFlowStoreApi } from "@/stores/create-flow-store"; +import { HostRuntimeStore, useHostRuntimeStore, type HostAuthorityBinding } from "./host-runtime"; +import { HostAuthorityProvider } from "./host-authority-provider"; + +it("ordinary consumer hooks retain the runtime, store and creation owner of each authority", () => { + const binding = (originServerId: string): HostAuthorityBinding => { + const sessionStore = createSessionStore(); + return { + reference: { originServerId, serverId: "B" }, + namespace: JSON.stringify(["host-authority-v1", originServerId, "B"]), + sessionStore, + queries: new QueryClient(), + runtime: new HostRuntimeStore({ sessionStore }), + ready: Promise.resolve(), + }; + }; + const direct = binding("B"); + const routed = binding("A"); + function OrdinaryConsumer({ expected }: { expected: HostAuthorityBinding }) { + const runtime = useHostRuntimeStore(); + const store = useSessionStoreApi(); + const creates = useCreateFlowStoreApi(); + expect(useQueryClient()).toBe(expected.queries); + expect(runtime).toBe(expected.runtime); + expect(store).toBe(expected.sessionStore); + expect(creates).toBe(expected.sessionStore.createFlows); + return null; + } + renderToStaticMarkup( + <> + + + + + + + , + ); +}); diff --git a/packages/app/src/runtime/host-authority-provider.tsx b/packages/app/src/runtime/host-authority-provider.tsx new file mode 100644 index 0000000000..36f1e9994f --- /dev/null +++ b/packages/app/src/runtime/host-authority-provider.tsx @@ -0,0 +1,26 @@ +import type { ReactNode } from "react"; +import { QueryClientProvider } from "@tanstack/react-query"; +import { SessionStoreContext } from "@/stores/session-store"; +import { CreateFlowStoreContext } from "@/stores/create-flow-store"; +import { HostRuntimeContext, type HostAuthorityBinding } from "./host-runtime"; + +/** The retained screen owns acquisition/release; descendants receive one bound capability. */ +export function HostAuthorityProvider({ + binding, + children, +}: { + binding: HostAuthorityBinding; + children: ReactNode; +}) { + return ( + + + + + {children} + + + + + ); +} diff --git a/packages/app/src/runtime/host-runtime.test.ts b/packages/app/src/runtime/host-runtime.test.ts index 61285eb2c6..baaf3f12f1 100644 --- a/packages/app/src/runtime/host-runtime.test.ts +++ b/packages/app/src/runtime/host-runtime.test.ts @@ -1,4 +1,6 @@ +import { indexedDB as fakeIndexedDb, IDBKeyRange as FakeIDBKeyRange } from "fake-indexeddb"; import { afterEach, describe, expect, it, vi } from "vitest"; +import { QueryObserver } from "@tanstack/react-query"; import type { AppStateStatus } from "react-native"; import { bindHostRuntimeAppState, @@ -13,12 +15,15 @@ import type { import { DaemonAuthenticationError } from "@getpaseo/client/internal/daemon-client"; import type { ConnectionOffer } from "@getpaseo/protocol/connection-offer"; import type { SessionOutboundMessage } from "@getpaseo/protocol/messages"; +import { MutableDaemonConfigSchema } from "@getpaseo/protocol/messages"; +import { daemonConfigQueryKey } from "@/data/daemon-config"; import type { AgentPermissionRequest } from "@getpaseo/protocol/agent-types"; import type { HostConnection, HostProfile } from "@/types/host-connection"; import { defaultHostAppearance } from "@/hosts/appearance"; import { useSessionStore, type Agent } from "@/stores/session-store"; import { normalizeAgentSnapshot } from "@/utils/agent-snapshots"; import { isAgentArchiving, setAgentArchiving } from "@/hooks/use-archive-agent"; +import { createAgentArchiveAction } from "@/hooks/use-archive-agent"; import { queryClient } from "@/data/query-client"; import { HostRuntimeController, @@ -52,6 +57,16 @@ it("requests the managed connection credential through desktop main without a we }); class FakeDaemonClient { + private readonly eventListeners = new Set<(message: SessionOutboundMessage) => void>(); + publishEvent(message: SessionOutboundMessage): void { + for (const receive of this.eventListeners) receive(message); + } + readonly archiveCalls: string[] = []; + archiveResponse: Promise<{ archivedAt: string }> | null = null; + async archiveAgent(agentId: string): Promise<{ archivedAt: string }> { + this.archiveCalls.push(agentId); + return this.archiveResponse ?? { archivedAt: new Date().toISOString() }; + } supportsBackgroundWorkspaces(): boolean { return false; } @@ -96,7 +111,12 @@ class FakeDaemonClient { } observeEvents() { - return subscriptionFixture(Promise.resolve({}), () => () => {}); + return subscriptionFixture(Promise.resolve({}), (receive) => { + this.eventListeners.add(receive); + return () => { + this.eventListeners.delete(receive); + }; + }); } async waitForAgentUpdates(): Promise { @@ -161,6 +181,29 @@ class FakeDaemonClient { return this.state; } + private readonly hostDirectoryListeners = new Set<(message: SessionOutboundMessage) => void>(); + observeHosts(): ReturnType { + return subscriptionFixture( + Promise.resolve({ self: { serverId: this.serverId, label: "Origin" }, hosts: [] }), + (receive) => { + this.hostDirectoryListeners.add(receive); + return () => { + this.hostDirectoryListeners.delete(receive); + }; + }, + ); + } + hostsChanged(): void { + for (const receive of this.hostDirectoryListeners) + receive({ + type: "hosts.changed", + payload: { + subscriptionId: "hosts", + directory: { self: { serverId: this.serverId, label: "Origin" }, hosts: [] }, + }, + }); + } + public serverId = "srv_test"; public ownedSubscriptions = true; private daemonVersion: string | null = "0.8.0"; @@ -168,7 +211,7 @@ class FakeDaemonClient { if (this.daemonVersion === null) return null; return { status: "server_info", - serverId: "srv_test", + serverId: this.serverId, hostname: "test", version: this.daemonVersion, features: { ownedSubscriptions: this.ownedSubscriptions }, @@ -4440,3 +4483,288 @@ describe("HostRuntimeStore initial connection hint bootstrap", () => { expect(store.getHosts()).toHaveLength(0); }); }); + +it("retains one authority binding for direct and routed consumers without changing a retained direct destination", async () => { + globalThis.indexedDB = fakeIndexedDb; + globalThis.IDBKeyRange = FakeIDBKeyRange; + const prefix = crypto.randomUUID(); + const aId = `A-${prefix}`; + const bId = `B-${prefix}`; + const clients = new Map(); + const resolutions: Array<{ origin: DaemonClient; serverId: string }> = []; + const releases: string[] = []; + const routeClient = makeConnectedProbeClient(5); + routeClient.serverId = bId; + const store = new HostRuntimeStore({ + storage: createMemoryHostRuntimeStorage(), + replicaRowStore: createMemoryReplicaRowStore(), + deps: { + createClient: ({ host }) => { + const client = makeConnectedProbeClient(5); + client.serverId = host.serverId; + clients.set(host.serverId, client); + return client as unknown as DaemonClient; + }, + connectToDaemon: async () => { + throw new Error("Probe is not needed"); + }, + getClientId: async () => "binding-client", + }, + resolveHostConnection: async (origin, serverId) => { + resolutions.push({ origin, serverId }); + const client = + origin === (clients.get(aId) as unknown as DaemonClient) ? routeClient : origin; + return { + client: client as DaemonClient, + release: async () => { + releases.push(serverId); + }, + }; + }, + }); + const hosts = [aId, bId].map((serverId) => + makeHost({ serverId, connections: [makeHost().connections[0]!] }), + ); + const initial = new Map(); + for (const host of hosts) { + const client = makeConnectedProbeClient(5); + client.serverId = host.serverId; + clients.set(host.serverId, client); + initial.set(host.serverId, { + connectionId: host.connections[0]!.id, + existingClient: client as unknown as DaemonClient, + }); + } + store.syncHosts(hosts, { initialConnectionByServerId: initial }); + await onceHostSnapshotMatches(store, aId, (snapshot) => snapshot?.connectionStatus === "online"); + await onceHostSnapshotMatches(store, bId, (snapshot) => snapshot?.connectionStatus === "online"); + const direct = store.acquireBinding({ originServerId: bId, serverId: bId }); + const routed = store.acquireBinding({ originServerId: aId, serverId: bId }); + const sibling = store.acquireBinding({ originServerId: aId, serverId: bId }); + try { + await Promise.all([direct.binding.ready, routed.binding.ready]); + expect(sibling.binding).toBe(routed.binding); + expect(direct.binding.runtime.getClient(bId)).toBe(clients.get(bId)); + expect(routed.binding.runtime.getClient(bId)).toBe(routeClient); + expect(direct.binding.sessionStore).not.toBe(routed.binding.sessionStore); + expect(direct.binding.namespace).not.toBe(routed.binding.namespace); + expect(routed.binding.queries.getDefaultOptions()).toEqual(queryClient.getDefaultOptions()); + const key = ["schedules", bId]; + let directFetches = 0; + let routedFetches = 0; + direct.binding.queries.setQueryData(key, "Direct cache"); + routed.binding.queries.setQueryData(key, "Routed cache"); + queryClient.setQueryData(key, "Registry cache"); + const directObserver = new QueryObserver(direct.binding.queries, { + queryKey: key, + queryFn: async () => { + directFetches++; + return "Direct recovered"; + }, + }); + const routedObserver = new QueryObserver(routed.binding.queries, { + queryKey: key, + queryFn: async () => { + routedFetches++; + return "Routed recovered"; + }, + }); + const stopDirect = directObserver.subscribe(() => {}); + const stopRouted = routedObserver.subscribe(() => {}); + try { + expect([directFetches, routedFetches]).toEqual([0, 0]); + routeClient.setConnectionState({ status: "disconnected", reason: "Recover routed cache" }); + routeClient.setConnectionState({ status: "connected" }); + await expect.poll(() => routed.binding.queries.getQueryData(key)).toBe("Routed recovered"); + expect([directFetches, routedFetches]).toEqual([0, 1]); + expect(direct.binding.queries.getQueryData(key)).toBe("Direct cache"); + expect(queryClient.getQueryState(key)?.isInvalidated).toBe(false); + clients + .get(bId)! + .setConnectionState({ status: "disconnected", reason: "Recover direct cache" }); + clients.get(bId)!.setConnectionState({ status: "connected" }); + await expect.poll(() => direct.binding.queries.getQueryData(key)).toBe("Direct recovered"); + expect([directFetches, routedFetches]).toEqual([1, 1]); + } finally { + stopDirect(); + stopRouted(); + queryClient.removeQueries({ queryKey: key }); + } + const configKey = daemonConfigQueryKey(bId); + const config = MutableDaemonConfigSchema.parse({ + relay: { enabled: false }, + mcp: { injectIntoAgents: true }, + }); + const directConfig = MutableDaemonConfigSchema.parse({ + relay: { enabled: true }, + mcp: { injectIntoAgents: true }, + }); + direct.binding.queries.setQueryData(configKey, directConfig); + routeClient.publishEvent({ + type: "status", + payload: { status: "daemon_config_changed", config }, + }); + await expect.poll(() => routed.binding.queries.getQueryData(configKey)).toEqual(config); + expect(direct.binding.queries.getQueryData(configKey)).toEqual(directConfig); + const agent = replicaAgent( + makeFetchAgentsEntry({ + id: "same-agent", + cwd: "/repo", + title: "Direct", + updatedAt: new Date().toISOString(), + }).agent, + bId, + ); + direct.binding.runtime.acceptAgentSnapshot(bId, agent); + expect(routed.binding.sessionStore.getState().sessions[bId]?.agents.size).toBe(0); + routed.binding.runtime.acceptAgentSnapshot(bId, { ...agent, title: "Routed" }); + direct.binding.runtime.acceptAgentSnapshot(bId, { ...agent, title: "Late direct" }); + expect(routed.binding.sessionStore.getState().sessions[bId]?.agents.get(agent.id)?.title).toBe( + "Routed", + ); + const archive = createAgentArchiveAction({ + runtime: routed.binding.runtime, + sessionStore: routed.binding.sessionStore, + queryClient: routed.binding.queries, + unavailableMessage: "Authority unavailable", + }); + const gate = createDeferred<{ archivedAt: string }>(); + routeClient.archiveResponse = gate.promise; + const archiving = archive({ serverId: bId, agentId: agent.id }); + void archiving.catch(() => {}); + expect(routeClient.archiveCalls).toEqual([agent.id]); + expect(clients.get(bId)!.archiveCalls).toEqual([]); + expect( + isAgentArchiving({ queryClient: routed.binding.queries, serverId: bId, agentId: agent.id }), + ).toBe(true); + expect( + isAgentArchiving({ queryClient: direct.binding.queries, serverId: bId, agentId: agent.id }), + ).toBe(false); + expect( + direct.binding.sessionStore.getState().sessions[bId]?.agents.get(agent.id)?.archivedAt, + ).toBeNull(); + gate.reject(new Error("Rejected by selected authority")); + await expect(archiving).rejects.toThrow("Rejected by selected authority"); + expect( + routed.binding.sessionStore.getState().sessions[bId]?.agents.get(agent.id)?.archivedAt, + ).toBeNull(); + routeClient.setConnectionState({ status: "disconnected", reason: "Route revoked" }); + expect(routed.binding.runtime.getSnapshot(bId)?.connectionStatus).toBe("error"); + expect(direct.binding.runtime.getSnapshot(bId)?.connectionStatus).toBe("online"); + expect(routed.binding.runtime.getClient(bId)).toBe(routeClient); + expect(resolutions).toHaveLength(2); + // Withdrawal must retire command sources even if the old client stays alive. + routeClient.setConnectionState({ status: "connected" }); + store.syncHosts(hosts.filter((host) => host.serverId === bId)); + expect(routed.binding.runtime.getClient(bId)).toBeNull(); + expect(routed.binding.sessionStore.getState().sessions[bId]?.client).toBeNull(); + expect(routeClient.getConnectionState().status).toBe("connected"); + routeClient.publishEvent({ + type: "status", + payload: { status: "daemon_config_changed", config: directConfig }, + }); + expect(routed.binding.queries.getQueryData(configKey)).toEqual(config); + await expect(archive({ serverId: bId, agentId: agent.id })).rejects.toThrow( + "Authority unavailable", + ); + expect(routeClient.archiveCalls).toEqual([agent.id]); + expect(clients.get(bId)!.archiveCalls).toEqual([]); + expect( + routed.binding.sessionStore.getState().sessions[bId]?.agents.get(agent.id)?.archivedAt, + ).toBeNull(); + expect(direct.binding.runtime.getSnapshot(bId)?.connectionStatus).toBe("online"); + expect(direct.binding.sessionStore.getState().sessions[bId]?.client).toBe(clients.get(bId)); + await routed.release(); + expect(releases).toEqual([bId]); + expect(sibling.binding.sessionStore.getState().sessions[bId]?.agents.get(agent.id)?.title).toBe( + "Routed", + ); + } finally { + await Promise.all([direct.release(), routed.release(), sibling.release()]); + store.syncHosts([]); + } + expect(releases).toHaveLength(2); +}); + +it.each(["after failure", "during acquisition"])( + "recovers a cold offline authority when the same origin enables the destination %s", + async (timing) => { + globalThis.indexedDB = fakeIndexedDb; + globalThis.IDBKeyRange = FakeIDBKeyRange; + const aId = `cold-A-${crypto.randomUUID()}`; + const bId = `cold-B-${crypto.randomUUID()}`; + const origin = makeConnectedProbeClient(5); + origin.serverId = aId; + const target = makeConnectedProbeClient(5); + target.serverId = bId; + const gate = createDeferred(); + let attempts = 0; + let enabled = false; + const store = new HostRuntimeStore({ + storage: createMemoryHostRuntimeStorage(), + replicaRowStore: createMemoryReplicaRowStore(), + resolveHostConnection: async (source, destination) => { + expect(source).toBe(origin); + expect(destination).toBe(bId); + attempts++; + if (attempts === 1) { + await gate.promise; + throw new Error("Destination disabled"); + } + if (!enabled) throw new Error("Destination disabled"); + return { client: target as unknown as DaemonClient, release: async () => {} }; + }, + deps: { + createClient: () => { + throw new Error("Origin must remain the same client"); + }, + connectToDaemon: async () => { + throw new Error("No new connection"); + }, + getClientId: async () => "cold-authority", + }, + }); + const host = makeHost({ serverId: aId, connections: [makeHost().connections[0]!] }); + store.syncHosts([host], { + initialConnectionByServerId: new Map([ + [ + aId, + { + connectionId: host.connections[0]!.id, + existingClient: origin as unknown as DaemonClient, + }, + ], + ]), + }); + await onceHostSnapshotMatches( + store, + aId, + (snapshot) => snapshot?.connectionStatus === "online", + ); + origin.setConnectionState({ status: "disconnected", reason: "Cold offline origin" }); + const retained = store.acquireBinding({ originServerId: aId, serverId: bId }); + try { + origin.setConnectionState({ status: "connected" }); + await expect.poll(() => attempts).toBe(1); + if (timing === "after failure") { + gate.resolve(); + await expect + .poll(() => retained.binding.runtime.getSnapshot(bId)?.connectionStatus) + .toBe("error"); + } + enabled = true; + origin.hostsChanged(); + gate.resolve(); + await expect + .poll(() => retained.binding.runtime.getSnapshot(bId)?.connectionStatus) + .toBe("online"); + expect(store.getClient(aId)).toBe(origin); + expect(retained.binding.runtime.getClient(bId)).toBe(target); + expect(attempts).toBeGreaterThanOrEqual(2); + } finally { + gate.resolve(); + await retained.release(); + store.syncHosts([]); + } + }, +); diff --git a/packages/app/src/runtime/host-runtime.ts b/packages/app/src/runtime/host-runtime.ts index 55298ce5a1..ffb8d0419d 100644 --- a/packages/app/src/runtime/host-runtime.ts +++ b/packages/app/src/runtime/host-runtime.ts @@ -1,4 +1,9 @@ -import { useSyncExternalStore, useMemo } from "react"; +import { + acquireHostConnection, + type HostConnectionLease, +} from "@getpaseo/client/internal/host-connection"; +import { createContext, useContext, useSyncExternalStore, useMemo } from "react"; +import { QueryClient } from "@tanstack/react-query"; import AsyncStorage from "@react-native-async-storage/async-storage"; import equal from "fast-deep-equal/es6"; import { @@ -52,6 +57,8 @@ import { CLIENT_CAPS } from "@getpaseo/protocol/client-capabilities"; import { BROWSER_AUTOMATION_COMMAND_NAMES } from "@getpaseo/protocol/browser-automation/rpc-schemas"; import { useSessionStore, + type SessionStoreApi, + createSessionStore, toDaemonServerInfo, type Agent, type WorkspaceDescriptor, @@ -60,7 +67,7 @@ import { import type { TurnLivenessTransition } from "@/timeline/turn-liveness"; import { useWorkspaceSetupStore } from "@/stores/workspace-setup-store"; import { invalidateCheckoutGitQueriesForServer } from "@/git/query-keys"; -import { queryClient } from "@/data/query-client"; +import { queryClient, createQueryClient } from "@/data/query-client"; import { invalidateServerDataQueriesAfterReconnect, mountServerDataPushRouter, @@ -608,21 +615,8 @@ function createDefaultDeps(): HostRuntimeControllerDeps { }), getClientId: () => getOrCreateClientId(), mountClientHandlers: ({ client, host }) => { - const unmountServerData = mountServerDataPushRouter({ - client, - queryClient, - serverId: host.serverId, - }); - if (!browserAutomationCapabilities) { - return unmountServerData; - } - const unmountBrowserAutomation = mountBrowserAutomationDaemonClientHandler(client, { - serverId: host.serverId, - }); - return () => { - unmountBrowserAutomation(); - unmountServerData(); - }; + if (!browserAutomationCapabilities) return () => {}; + return mountBrowserAutomationDaemonClientHandler(client, { serverId: host.serverId }); }, }; } @@ -1486,8 +1480,54 @@ interface AgentDirectoryRefreshInput { page?: FetchAgentsOptions["page"]; } +export interface HostAuthorityReference { + originServerId: string; + serverId: string; +} + +export interface HostAuthorityBinding { + readonly reference: HostAuthorityReference; + readonly namespace: string; + readonly sessionStore: SessionStoreApi; + readonly queries: QueryClient; + readonly runtime: HostRuntimeStore; + /** Initial acquisition only. Current availability comes from runtime snapshots. */ + readonly ready: Promise; +} + +export interface OwnedHostAuthorityBinding { + readonly binding: HostAuthorityBinding; + release(): Promise; +} + +function authorityNamespace(reference: HostAuthorityReference): string { + return JSON.stringify(["host-authority-v1", reference.originServerId, reference.serverId]); +} + +function initialAuthoritySnapshot(serverId: string): HostRuntimeSnapshot { + return { + serverId, + activeConnectionId: null, + activeConnection: null, + connectionStatus: "connecting", + client: null, + lastError: null, + lastOnlineAt: null, + agentDirectoryStatus: "idle", + agentDirectoryError: null, + hasEverLoadedAgentDirectory: false, + probeByConnectionId: new Map(), + clientGeneration: 0, + connectionEpoch: 0, + }; +} + export class HostRuntimeStore { private controllers = new Map(); + private readonly authorityBindings = new Map< + string, + { binding: HostAuthorityBinding; refs: number; close(): Promise } + >(); private serverListeners = new Map void>>(); private globalListeners = new Set<() => void>(); private hostListListeners = new Set<() => void>(); @@ -1510,20 +1550,280 @@ export class HostRuntimeStore { private readonly hostConfirmations = new HostConfirmations(); private storage: HostRuntimeStorage; private replicaCache: ReplicaCache; + private readonly sessionStore: SessionStoreApi; + private readonly queries: QueryClient; + private readonly serverDataClients = new Map(); + private readonly resolveHostConnection: typeof acquireHostConnection; + private readonly snapshots = new Map(); private readonly revokePushNotifications: typeof revokePushNotifications; constructor(input?: { deps?: HostRuntimeControllerDeps; storage?: HostRuntimeStorage; replicaRowStore?: ReplicaRowStore; + sessionStore?: SessionStoreApi; + queries?: QueryClient; + resolveHostConnection?: typeof acquireHostConnection; revokePushNotifications?: typeof revokePushNotifications; }) { this.deps = input?.deps ?? createDefaultDeps(); + this.sessionStore = input?.sessionStore ?? useSessionStore; + this.queries = input?.queries ?? queryClient; + this.resolveHostConnection = input?.resolveHostConnection ?? acquireHostConnection; this.storage = input?.storage ?? AsyncStorage; this.replicaCache = new ReplicaCache(input?.replicaRowStore ?? createReplicaRowStore()); this.revokePushNotifications = input?.revokePushNotifications ?? revokePushNotifications; } + /** A retained screen owns a lease; route identity never grants access by itself. */ + acquireBinding(reference: HostAuthorityReference): OwnedHostAuthorityBinding { + const normalized = { originServerId: reference.originServerId, serverId: reference.serverId }; + const namespace = authorityNamespace(normalized); + let entry = this.authorityBindings.get(namespace); + if (!entry) { + const queries = createQueryClient(); + const scope = this.createAuthorityScope(normalized, namespace, queries); + entry = { + binding: { + reference: normalized, + namespace, + sessionStore: scope.runtime.sessionStore, + queries, + runtime: scope.runtime, + ready: scope.ready, + }, + refs: 0, + close: async () => { + try { + await scope.close(); + } finally { + queries.clear(); + } + }, + }; + this.authorityBindings.set(namespace, entry); + } + entry.refs++; + const retained = entry; + let released: Promise | null = null; + return { + binding: retained.binding, + release: () => { + if (!released) { + retained.refs--; + if (retained.refs === 0) { + if (this.authorityBindings.get(namespace) === retained) + this.authorityBindings.delete(namespace); + released = retained.close(); + } else released = Promise.resolve(); + } + return released; + }, + }; + } + + private createAuthorityScope( + reference: HostAuthorityReference, + namespace: string, + queries: QueryClient, + ): { + runtime: HostRuntimeStore; + ready: Promise; + close(): Promise; + } { + const scoped = new HostRuntimeStore({ + sessionStore: createSessionStore(), + queries, + replicaRowStore: createReplicaRowStore(namespace), + storage: this.storage, + deps: this.deps, + }); + const { serverId, originServerId } = reference; + scoped.sessionStore.getState().initializeSession(serverId, null); + scoped.snapshots.set(serverId, initialAuthoritySnapshot(serverId)); + const patch = (change: Partial) => { + const previous = scoped.snapshots.get(serverId)!; + scoped.snapshots.set(serverId, { ...previous, ...change }); + scoped.emit(serverId); + }; + scoped.initializeReplica(serverId, { + onAgentStoppedRunning: (agentId) => scoped.drainQueuedAgentMessage(serverId, agentId), + markAgentLoading: () => + patch({ agentDirectoryStatus: "initial_loading", agentDirectoryError: null }), + markAgentReady: () => + patch({ + agentDirectoryStatus: "ready", + agentDirectoryError: null, + hasEverLoadedAgentDirectory: true, + }), + markAgentError: (failure) => + patch({ + agentDirectoryStatus: scoped.getSnapshot(serverId)?.hasEverLoadedAgentDirectory + ? "error_after_ready" + : "error_before_first_success", + agentDirectoryError: failure, + }), + }); + let stopped = false; + let origin: DaemonClient | null = null; + let controller = new AbortController(); + let lease: HostConnectionLease | null = null; + let stopTarget: (() => void) | null = null; + let directory: ReturnType | null = null; + let acquiring: Promise | null = null; + let directoryRevision = 0; + let resolveReady!: () => void; + let rejectReady!: (failure: unknown) => void; + const ready = new Promise((resolve, reject) => { + resolveReady = resolve; + rejectReady = reject; + }); + void ready.catch(() => {}); + let generation = 0; + const publish = ( + client: DaemonClient | null, + state: ConnectionState, + failure: string | null = null, + ) => { + const previous = scoped.snapshots.get(serverId)!; + const online = state.status === "connected"; + let connectionStatus: HostRuntimeConnectionStatus = failure ? "error" : "offline"; + if (state.status === "connecting") connectionStatus = "connecting"; + if (online) connectionStatus = "online"; + const snapshot: HostRuntimeSnapshot = { + ...previous, + client, + connectionStatus, + lastError: failure, + clientGeneration: generation, + connectionEpoch: + online && previous.connectionStatus !== "online" + ? previous.connectionEpoch + 1 + : previous.connectionEpoch, + lastOnlineAt: online ? new Date().toISOString() : previous.lastOnlineAt, + }; + scoped.syncSessionReplica(serverId, snapshot); + scoped.syncDirectoryConnection(serverId); + scoped.emit(serverId); + }; + const acquire = () => { + if ( + stopped || + acquiring || + lease || + !origin || + origin.getConnectionState().status !== "connected" + ) + return; + const source = origin; + const signal = controller.signal; + const attempt = generation; + const observedRevision = directoryRevision; + publish(null, { status: "connecting", attempt: 1 }); + acquiring = (async () => { + try { + const acquired = await this.resolveHostConnection(source, serverId, signal); + if (stopped || signal.aborted || generation !== attempt) { + await acquired.release(); + return; + } + lease = acquired; + stopTarget = acquired.client.subscribeConnectionStatus((state) => { + if (!stopped && generation === attempt) + publish(acquired.client, state, acquired.client.lastError); + }); + resolveReady(); + } catch (failure) { + if (!stopped && generation === attempt) { + publish(null, { status: "disconnected" }, toErrorMessage(failure)); + rejectReady(failure); + } + } finally { + if (generation === attempt) { + acquiring = null; + if (!lease && directoryRevision !== observedRevision) acquire(); + } + } + })(); + }; + const retired = new Set>(); + const retire = (release: Promise) => { + retired.add(release); + void release.then( + () => retired.delete(release), + () => retired.delete(release), + ); + }; + const syncOrigin = () => { + if (stopped) return; + const snapshot = this.getSnapshot(originServerId); + const nextOrigin = snapshot?.client ?? null; + if (origin !== nextOrigin) { + generation++; + controller.abort(); + stopTarget?.(); + stopTarget = null; + if (lease) retire(lease.release()); + if (directory) retire(directory.release()); + if (acquiring) retire(acquiring); + lease = null; + directory = null; + acquiring = null; + controller = new AbortController(); + origin = nextOrigin; + publish(null, { status: "disconnected" }); + } + if (origin && origin.getConnectionState().status === "connected" && !directory) { + try { + directory = origin.observeHosts(); + const changed = () => { + directoryRevision++; + acquire(); + }; + directory.subscribe({ snapshot: changed, update: changed }); + void directory.ready.catch(() => {}); + } catch (failure) { + publish(null, { status: "disconnected" }, toErrorMessage(failure)); + } + } + if (snapshot?.connectionStatus === "online") acquire(); + else if (!lease) + publish( + null, + { status: "disconnected" }, + snapshot?.lastError ?? "Origin authority is unavailable", + ); + }; + const stopOrigin = this.subscribe(originServerId, syncOrigin); + syncOrigin(); + return { + runtime: scoped, + ready, + close: async () => { + stopped = true; + generation++; + stopOrigin(); + controller.abort(); + stopTarget?.(); + rejectReady(new Error("Host authority binding is disposed")); + scoped.syncServerDataClient(serverId, null); + scoped.directorySyncByServer.get(serverId)?.dispose(); + scoped.sessionStore.getState().sessions[serverId]?.viewedTimelineSync?.dispose(); + for (const agentId of scoped.sessionStore + .getState() + .sessions[serverId]?.initializingAgents.keys() ?? []) { + scoped.sessionStore.initializations.rejectInitDeferred( + scoped.sessionStore.initializations.getInitKey(serverId, agentId), + new Error("Host authority binding is disposed"), + ); + } + await Promise.allSettled([lease?.release(), directory?.release(), acquiring, ...retired]); + await scoped.replicaCache.flush(); + scoped.sessionStore.getState().clearSession(serverId); + }, + }; + } + // --- Host registry --- getHosts(): HostProfile[] { @@ -1772,6 +2072,7 @@ export class HostRuntimeStore { } rekeyMap(this.controllers, oldServerId, newServerId); + this.snapshots.delete(oldServerId); rekeyMap(this.lastConnectionStatusByServer, oldServerId, newServerId); rekeyMap(this.connectionStatusStartedAtByServer, oldServerId, newServerId); this.replicaCache.reconcileServerId(oldServerId, newServerId); @@ -1779,25 +2080,12 @@ export class HostRuntimeStore { this.directorySyncByServer.get(oldServerId)?.dispose(); this.directorySyncByServer.delete(oldServerId); this.timelineReplicaByServer.delete(oldServerId); - const directory = new DirectorySync( - newServerId, - { - onAgentStoppedRunning: (agentId) => this.drainQueuedAgentMessage(newServerId, agentId), - markAgentLoading: () => controller.markAgentDirectorySyncLoading(), - markAgentReady: () => controller.markAgentDirectorySyncReady(), - markAgentError: (error) => controller.markAgentDirectorySyncError(error), - }, - this.replicaCache, - ); - this.directorySyncByServer.set(newServerId, directory); - this.timelineReplicaByServer.set( - newServerId, - createTimelineReplica({ - serverId: newServerId, - storage: this.replicaCache, - prepareAgent: (agentId) => directory.prepareAgentRoute(agentId), - }), - ); + const directory = this.initializeReplica(newServerId, { + onAgentStoppedRunning: (agentId) => this.drainQueuedAgentMessage(newServerId, agentId), + markAgentLoading: () => controller.markAgentDirectorySyncLoading(), + markAgentReady: () => controller.markAgentDirectorySyncReady(), + markAgentError: (error) => controller.markAgentDirectorySyncError(error), + }); controller.adoptReconciledServerId(newServerId); const snapshot = controller.getSnapshot(); this.clearHostReplica(oldServerId); @@ -2259,6 +2547,7 @@ export class HostRuntimeStore { continue; } this.controllers.delete(serverId); + this.snapshots.delete(serverId); this.lastConnectionStatusByServer.delete(serverId); this.connectionStatusStartedAtByServer.delete(serverId); this.directorySyncByServer.get(serverId)?.dispose(); @@ -2287,27 +2576,15 @@ export class HostRuntimeStore { onReconcileServerId: (oldId, newId) => this.reconcileServerId(oldId, newId), }); this.controllers.set(host.serverId, controller); - useSessionStore.getState().initializeSession(host.serverId, null); - const directory = new DirectorySync( - host.serverId, - { - onAgentStoppedRunning: (agentId) => this.drainQueuedAgentMessage(host.serverId, agentId), - markAgentLoading: () => controller.markAgentDirectorySyncLoading(), - markAgentReady: () => controller.markAgentDirectorySyncReady(), - markAgentError: (error) => controller.markAgentDirectorySyncError(error), - }, - this.replicaCache, - ); - this.directorySyncByServer.set(host.serverId, directory); - this.timelineReplicaByServer.set( - host.serverId, - createTimelineReplica({ - serverId: host.serverId, - storage: this.replicaCache, - prepareAgent: (agentId) => directory.prepareAgentRoute(agentId), - }), - ); + this.sessionStore.getState().initializeSession(host.serverId, null); + this.initializeReplica(host.serverId, { + onAgentStoppedRunning: (agentId) => this.drainQueuedAgentMessage(host.serverId, agentId), + markAgentLoading: () => controller.markAgentDirectorySyncLoading(), + markAgentReady: () => controller.markAgentDirectorySyncReady(), + markAgentError: (error) => controller.markAgentDirectorySyncError(error), + }); const initialSnapshot = controller.getSnapshot(); + this.syncSessionReplica(host.serverId, initialSnapshot); this.lastConnectionStatusByServer.set(host.serverId, initialSnapshot.connectionStatus); this.connectionStatusStartedAtByServer.set(host.serverId, Date.now()); controller.subscribe(() => { @@ -2332,13 +2609,31 @@ export class HostRuntimeStore { } } + private initializeReplica( + serverId: string, + callbacks: ConstructorParameters[1], + ): DirectorySync { + const directory = new DirectorySync(serverId, callbacks, this.replicaCache, this.sessionStore); + this.directorySyncByServer.set(serverId, directory); + this.timelineReplicaByServer.set( + serverId, + createTimelineReplica({ + serverId, + store: this.sessionStore, + storage: this.replicaCache, + prepareAgent: (agentId) => directory.prepareAgentRoute(agentId), + }), + ); + return directory; + } + private syncSessionReplica(serverId: string, snapshot: HostRuntimeSnapshot): void { - if (!snapshot.client) { - return; - } - const sessionStore = useSessionStore.getState(); + this.snapshots.set(serverId, snapshot); + const sessionStore = this.sessionStore.getState(); sessionStore.initializeSession(serverId, snapshot.client, snapshot.clientGeneration); sessionStore.updateSessionClient(serverId, snapshot.client, snapshot.clientGeneration); + this.syncServerDataClient(serverId, snapshot.client); + if (!snapshot.client) return; // A reconnect keeps the same client, so the daemon's handshake (a restart or upgrade can // change its version and features) only reaches the store here. The client clears it while // disconnected; keep the last known value until the next handshake. @@ -2348,19 +2643,31 @@ export class HostRuntimeStore { } } + private syncServerDataClient(serverId: string, client: DaemonClient | null): void { + const current = this.serverDataClients.get(serverId); + if (current?.client === client) return; + current?.release(); + this.serverDataClients.delete(serverId); + if (client) + this.serverDataClients.set(serverId, { + client, + release: mountServerDataPushRouter({ client, serverId, queryClient: this.queries }), + }); + } + private clearHostReplica(serverId: string): void { - useSessionStore.getState().clearSession(serverId); + this.syncServerDataClient(serverId, null); + this.sessionStore.getState().clearSession(serverId); useWorkspaceSetupStore.getState().clearServer(serverId); } private syncDirectoryConnection(serverId: string): void { - const controller = this.controllers.get(serverId); - if (!controller) { + const snapshot = this.snapshots.get(serverId); + if (!snapshot) { this.lastConnectionStatusByServer.delete(serverId); this.connectionStatusStartedAtByServer.delete(serverId); return; } - const snapshot = controller.getSnapshot(); const directory = this.directorySyncByServer.get(serverId); directory?.connectionChanged({ client: snapshot.client, @@ -2383,20 +2690,20 @@ export class HostRuntimeStore { const didTransitionOnline = snapshot.connectionStatus === "online" && previousStatus !== "online"; if (didTransitionOnline) { - useSessionStore.getState().bumpHistorySyncGeneration(serverId); + this.sessionStore.getState().bumpHistorySyncGeneration(serverId); // Checkout git data is push-driven; pushes emitted while disconnected are gone for // good (the daemon dedupes by snapshot fingerprint). Mark the caches stale so active // queries refetch now and evicted ones on their next mount. - void invalidateCheckoutGitQueriesForServer(queryClient, serverId); - invalidateServerDataQueriesAfterReconnect({ queryClient, serverId }); - void queryClient.invalidateQueries({ queryKey: schedulesQueryBaseKey }); + void invalidateCheckoutGitQueriesForServer(this.queries, serverId); + invalidateServerDataQueriesAfterReconnect({ queryClient: this.queries, serverId }); + void this.queries.invalidateQueries({ queryKey: schedulesQueryBaseKey }); } } drainQueuedAgentMessage(serverId: string, agentId: string): void { const drainKey = `${serverId}:${agentId}`; if (this.queuedAgentDrainInFlight.has(drainKey)) return; - const store = useSessionStore.getState(); + const store = this.sessionStore.getState(); const session = store.sessions[serverId]; const queue = session?.queuedMessages.get(agentId); const client = session?.client; @@ -2410,12 +2717,13 @@ export class HostRuntimeStore { messageId: next.id, queue: { read: (queuedAgentId) => - useSessionStore.getState().sessions[serverId]?.queuedMessages.get(queuedAgentId) ?? [], - write: (update) => useSessionStore.getState().setQueuedMessages(serverId, update), + this.sessionStore.getState().sessions[serverId]?.queuedMessages.get(queuedAgentId) ?? [], + write: (update) => this.sessionStore.getState().setQueuedMessages(serverId, update), }, submitMessage: async ({ text, attachments }) => { const supportsForgeAttachments = - useSessionStore.getState().sessions[serverId]?.serverInfo?.features?.forgeSearch === true; + this.sessionStore.getState().sessions[serverId]?.serverInfo?.features?.forgeSearch === + true; await dispatchComposerAgentMessage({ client, agentId, @@ -2500,7 +2808,7 @@ export class HostRuntimeStore { } getSnapshot(serverId: string): HostRuntimeSnapshot | null { - return this.controllers.get(serverId)?.getSnapshot() ?? null; + return this.snapshots.get(serverId) ?? null; } getConnectionStatusSince(serverId: string): number | null { @@ -2512,7 +2820,7 @@ export class HostRuntimeStore { } getClient(serverId: string): DaemonClient | null { - return this.controllers.get(serverId)?.getClient() ?? null; + return this.snapshots.get(serverId)?.client ?? null; } subscribe(serverId: string, listener: () => void): () => void { @@ -2639,9 +2947,9 @@ export class HostRuntimeStore { ): Promise>> { const directory = this.directorySyncByServer.get(serverId); if (!directory) throw new Error(`Unknown host runtime for serverId ${serverId}`); - const owner = useSessionStore.getState().sessions[serverId]?.viewedTimelineSync; + const owner = this.sessionStore.getState().sessions[serverId]?.viewedTimelineSync; const page = await directory.fetchTimeline(agentId, request); - if (owner && useSessionStore.getState().sessions[serverId]?.viewedTimelineSync === owner) { + if (owner && this.sessionStore.getState().sessions[serverId]?.viewedTimelineSync === owner) { owner.flushStreamAgent(agentId); owner.applyTimelineResponse(page); } @@ -2657,6 +2965,7 @@ export class HostRuntimeStore { const replica = this.timelineReplicaByServer.get(serverId); if (!replica) throw new Error(`Unknown host runtime for serverId ${serverId}`); return createViewedTimelineOwner({ + store: this.sessionStore, serverId, replica, replaceDemandedAgentIds: (agentIds) => directory.setAgentRouteDemand(agentIds), @@ -2706,8 +3015,15 @@ export function getHostRuntimeConnectionStatusSince(serverId: string): number | return getHostRuntimeStore().getConnectionStatusSince(serverId); } +export const HostRuntimeContext = createContext(null); + +/** Ordinary conversation hooks read the capability bound to their retained screen. */ +export function useHostRuntimeStore(): HostRuntimeStore { + return useContext(HostRuntimeContext) ?? getHostRuntimeStore(); +} + export function useHostRuntimeSnapshot(serverId: string): HostRuntimeSnapshot | null { - const store = getHostRuntimeStore(); + const store = useHostRuntimeStore(); return useSyncExternalStore( (onStoreChange) => store.subscribe(serverId, onStoreChange), () => store.getSnapshot(serverId), @@ -2716,7 +3032,7 @@ export function useHostRuntimeSnapshot(serverId: string): HostRuntimeSnapshot | } export function useHostRuntimeClient(serverId: string): DaemonClient | null { - const store = getHostRuntimeStore(); + const store = useHostRuntimeStore(); return useSyncExternalStore( (onStoreChange) => store.subscribe(serverId, onStoreChange), () => store.getSnapshot(serverId)?.client ?? null, @@ -2725,7 +3041,7 @@ export function useHostRuntimeClient(serverId: string): DaemonClient | null { } export function useHostRuntimeIsConnected(serverId: string): boolean { - const store = getHostRuntimeStore(); + const store = useHostRuntimeStore(); return useSyncExternalStore( (onStoreChange) => store.subscribe(serverId, onStoreChange), () => isHostRuntimeConnected(store.getSnapshot(serverId)), @@ -2734,7 +3050,7 @@ export function useHostRuntimeIsConnected(serverId: string): boolean { } export function useHostRuntimeConnectionStatus(serverId: string): HostRuntimeConnectionStatus { - const store = getHostRuntimeStore(); + const store = useHostRuntimeStore(); return useSyncExternalStore( (onStoreChange) => store.subscribe(serverId, onStoreChange), () => store.getSnapshot(serverId)?.connectionStatus ?? "connecting", @@ -2745,7 +3061,7 @@ export function useHostRuntimeConnectionStatus(serverId: string): HostRuntimeCon export function useHostRuntimeConnectionStatuses( serverIds: readonly string[], ): ReadonlyMap { - const store = getHostRuntimeStore(); + const store = useHostRuntimeStore(); // The snapshot is the statuses themselves, joined into a string so React compares by // value. A version counter read only for reactivity is dropped by the React Compiler. const readStatuses = () => @@ -2765,7 +3081,7 @@ export function useHostRuntimeConnectionStatuses( } export function useHostRuntimeLastError(serverId: string): string | null { - const store = getHostRuntimeStore(); + const store = useHostRuntimeStore(); return useSyncExternalStore( (onStoreChange) => store.subscribe(serverId, onStoreChange), () => store.getSnapshot(serverId)?.lastError ?? null, @@ -2776,7 +3092,7 @@ export function useHostRuntimeLastError(serverId: string): string | null { export function useHostRuntimeAgentDirectoryStatus( serverId: string, ): HostRuntimeAgentDirectoryStatus { - const store = getHostRuntimeStore(); + const store = useHostRuntimeStore(); return useSyncExternalStore( (onStoreChange) => store.subscribe(serverId, onStoreChange), () => store.getSnapshot(serverId)?.agentDirectoryStatus ?? "idle", @@ -2785,7 +3101,7 @@ export function useHostRuntimeAgentDirectoryStatus( } export function useHostRuntimeIsDirectoryLoading(serverId: string): boolean { - const store = getHostRuntimeStore(); + const store = useHostRuntimeStore(); return useSyncExternalStore( (onStoreChange) => store.subscribe(serverId, onStoreChange), () => isHostRuntimeDirectoryLoading(store.getSnapshot(serverId)), diff --git a/packages/app/src/runtime/replica-cache/row-store.native.ts b/packages/app/src/runtime/replica-cache/row-store.native.ts index bb91eb0aa6..af446c22d3 100644 --- a/packages/app/src/runtime/replica-cache/row-store.native.ts +++ b/packages/app/src/runtime/replica-cache/row-store.native.ts @@ -38,12 +38,14 @@ function createConnection(database: SQLiteDatabase): ReplicaSqliteConnection { }; } -const expoSqliteDriver: ReplicaSqliteDriver = { - async open() { - return createConnection(await openDatabaseAsync(DATABASE_NAME)); - }, -}; - -export function createReplicaRowStore(): ReplicaRowStore { - return createSqliteReplicaRowStore(expoSqliteDriver, REPLICA_ROW_STORE_SCHEMA_VERSION); +export function createReplicaRowStore(authorityNamespace?: string): ReplicaRowStore { + const databaseName = authorityNamespace + ? `paseo-replica-${encodeURIComponent(authorityNamespace)}.db` + : DATABASE_NAME; + const driver: ReplicaSqliteDriver = { + async open() { + return createConnection(await openDatabaseAsync(databaseName)); + }, + }; + return createSqliteReplicaRowStore(driver, REPLICA_ROW_STORE_SCHEMA_VERSION); } diff --git a/packages/app/src/runtime/replica-cache/row-store.web.test.ts b/packages/app/src/runtime/replica-cache/row-store.web.test.ts index 472143ee34..73b3528f4b 100644 --- a/packages/app/src/runtime/replica-cache/row-store.web.test.ts +++ b/packages/app/src/runtime/replica-cache/row-store.web.test.ts @@ -5,7 +5,7 @@ import { } from "fake-indexeddb"; import { expect, it, vi } from "vitest"; import { runReplicaRowStoreContract } from "./row-store.contract"; -import { createIndexedDbReplicaRowStore } from "./row-store.web"; +import { createReplicaRowStore, createIndexedDbReplicaRowStore } from "./row-store.web"; globalThis.indexedDB = fakeIndexedDb; globalThis.IDBKeyRange = FakeIDBKeyRange; @@ -49,3 +49,27 @@ it("uses exact IndexedDB keys for targeted rows instead of scanning the host", a get.mockRestore(); getAll.mockRestore(); }); + +it("restores only rows in the selected authority namespace while retaining the destination server ID", async () => { + const serverId = `B-${crypto.randomUUID()}`; + const origin = `A-${crypto.randomUUID()}`; + const direct = createReplicaRowStore(); + const selected = createReplicaRowStore(origin); + await Promise.all([direct.open(), selected.open()]); + const row = { + serverId, + kind: "agent" as const, + id: "colliding-agent", + payload: "broader-direct", + }; + await direct.apply({ upserts: [row], deletes: [] }); + expect(await selected.read(serverId, ["agent"])).toEqual([]); + await selected.apply({ upserts: [{ ...row, payload: "selected-authority" }], deletes: [] }); + const restored = createReplicaRowStore(origin); + await restored.open(); + expect(await restored.read(serverId, ["agent"])).toEqual([ + { ...row, payload: "selected-authority" }, + ]); + expect(await direct.read(serverId, ["agent"])).toEqual([row]); + await Promise.all([selected.clear(), direct.deleteHost(serverId)]); +}); diff --git a/packages/app/src/runtime/replica-cache/row-store.web.ts b/packages/app/src/runtime/replica-cache/row-store.web.ts index 12f484bd91..d140a88d0b 100644 --- a/packages/app/src/runtime/replica-cache/row-store.web.ts +++ b/packages/app/src/runtime/replica-cache/row-store.web.ts @@ -210,9 +210,11 @@ export function createIndexedDbReplicaRowStore( return { open, read, readAll, apply, deleteHost, renameHost, clear }; } -export function createReplicaRowStore(): ReplicaRowStore { +export function createReplicaRowStore(authorityNamespace?: string): ReplicaRowStore { return createIndexedDbReplicaRowStore({ - databaseName: DATABASE_NAME, + databaseName: authorityNamespace + ? `${DATABASE_NAME}:${encodeURIComponent(authorityNamespace)}` + : DATABASE_NAME, schemaVersion: REPLICA_ROW_STORE_SCHEMA_VERSION, }); } diff --git a/packages/app/src/stores/create-flow-store.ts b/packages/app/src/stores/create-flow-store.ts index 95a85137a5..fa30497024 100644 --- a/packages/app/src/stores/create-flow-store.ts +++ b/packages/app/src/stores/create-flow-store.ts @@ -1,4 +1,5 @@ -import { create } from "zustand"; +import { create, useStore } from "zustand"; +import { createContext, useContext } from "react"; import type { UserMessageImageAttachment } from "@/types/stream"; import type { AgentAttachment } from "@getpaseo/protocol/messages"; @@ -48,89 +49,105 @@ interface CreateFlowState { clearAll: () => void; } -export const useCreateFlowStore = create((set, get) => ({ - pendingByDraftId: {}, - setPending: (pending) => - set((state) => ({ - pendingByDraftId: { - ...state.pendingByDraftId, - [pending.draftId]: { - ...pending, - lifecycle: "active", - }, - }, - })), - trySetPending: (pending) => { - const existing = get().pendingByDraftId[pending.draftId]; - if (existing?.serverId === pending.serverId && existing.lifecycle !== "abandoned") return false; - get().setPending(pending); - return true; - }, - updateAgentId: ({ draftId, agentId }) => - set((state) => { - const current = state.pendingByDraftId[draftId]; - if (!current || current.agentId === agentId) { - return state; - } - return { +export const createCreateFlowStore = () => + create((set, get) => ({ + pendingByDraftId: {}, + setPending: (pending) => + set((state) => ({ pendingByDraftId: { ...state.pendingByDraftId, - [draftId]: { ...current, agentId }, + [pending.draftId]: { + ...pending, + lifecycle: "active", + }, }, - }; - }), - markLifecycle: ({ draftId, lifecycle, errorMessage }) => - set((state) => { - const current = state.pendingByDraftId[draftId]; - if (!current || current.lifecycle === lifecycle) { - return state; - } - return { - pendingByDraftId: { - ...state.pendingByDraftId, - [draftId]: { ...current, lifecycle, errorMessage }, - }, - }; - }), - rekeyDraft: ({ fromDraftId, toDraftId }) => - set((state) => { - const current = state.pendingByDraftId[fromDraftId]; - if (!current) { - return state; - } - if (fromDraftId === toDraftId) { - return state; - } - const { [fromDraftId]: _removed, ...rest } = state.pendingByDraftId; - return { - pendingByDraftId: { - ...rest, - [toDraftId]: { ...current, draftId: toDraftId }, - }, - }; - }), - clear: ({ draftId }) => - set((state) => { - if (!state.pendingByDraftId[draftId]) { - return state; - } - const { [draftId]: _removed, ...rest } = state.pendingByDraftId; - return { pendingByDraftId: rest }; - }), - clearByAgent: ({ serverId, agentId }) => - set((state) => { - const next = Object.fromEntries( - Object.entries(state.pendingByDraftId).filter( - ([, pending]) => - pending.lifecycle !== "sent" || - pending.serverId !== serverId || - pending.agentId !== agentId, - ), - ); - if (Object.keys(next).length === Object.keys(state.pendingByDraftId).length) { - return state; - } - return { pendingByDraftId: next }; - }), - clearAll: () => set({ pendingByDraftId: {} }), -})); + })), + trySetPending: (pending) => { + const existing = get().pendingByDraftId[pending.draftId]; + if (existing?.serverId === pending.serverId && existing.lifecycle !== "abandoned") + return false; + get().setPending(pending); + return true; + }, + updateAgentId: ({ draftId, agentId }) => + set((state) => { + const current = state.pendingByDraftId[draftId]; + if (!current || current.agentId === agentId) { + return state; + } + return { + pendingByDraftId: { + ...state.pendingByDraftId, + [draftId]: { ...current, agentId }, + }, + }; + }), + markLifecycle: ({ draftId, lifecycle, errorMessage }) => + set((state) => { + const current = state.pendingByDraftId[draftId]; + if (!current || current.lifecycle === lifecycle) { + return state; + } + return { + pendingByDraftId: { + ...state.pendingByDraftId, + [draftId]: { ...current, lifecycle, errorMessage }, + }, + }; + }), + rekeyDraft: ({ fromDraftId, toDraftId }) => + set((state) => { + const current = state.pendingByDraftId[fromDraftId]; + if (!current) { + return state; + } + if (fromDraftId === toDraftId) { + return state; + } + const { [fromDraftId]: _removed, ...rest } = state.pendingByDraftId; + return { + pendingByDraftId: { + ...rest, + [toDraftId]: { ...current, draftId: toDraftId }, + }, + }; + }), + clear: ({ draftId }) => + set((state) => { + if (!state.pendingByDraftId[draftId]) { + return state; + } + const { [draftId]: _removed, ...rest } = state.pendingByDraftId; + return { pendingByDraftId: rest }; + }), + clearByAgent: ({ serverId, agentId }) => + set((state) => { + const next = Object.fromEntries( + Object.entries(state.pendingByDraftId).filter( + ([, pending]) => + pending.lifecycle !== "sent" || + pending.serverId !== serverId || + pending.agentId !== agentId, + ), + ); + if (Object.keys(next).length === Object.keys(state.pendingByDraftId).length) { + return state; + } + return { pendingByDraftId: next }; + }), + clearAll: () => set({ pendingByDraftId: {} }), + })); + +export type CreateFlowStoreApi = ReturnType; +export const defaultCreateFlowStore = createCreateFlowStore(); +export const CreateFlowStoreContext = createContext(null); +export function useCreateFlowStoreApi(): CreateFlowStoreApi { + return useContext(CreateFlowStoreContext) ?? defaultCreateFlowStore; +} +export const useCreateFlowStore: CreateFlowStoreApi = Object.assign(function useCreateFlowStore( + selector: (state: CreateFlowState) => T = ((state: CreateFlowState) => state) as ( + state: CreateFlowState, + ) => T, +): T { + return useStore(useCreateFlowStoreApi(), selector); +}, defaultCreateFlowStore); diff --git a/packages/app/src/stores/session-store.ts b/packages/app/src/stores/session-store.ts index f7c54e07b2..b7f8747db6 100644 --- a/packages/app/src/stores/session-store.ts +++ b/packages/app/src/stores/session-store.ts @@ -1,5 +1,11 @@ +import { createCreateFlowStore, defaultCreateFlowStore } from "@/stores/create-flow-store"; +import { + createAgentInitializations, + defaultAgentInitializations, +} from "@/utils/agent-initialization"; import equal from "fast-deep-equal"; -import { create } from "zustand"; +import { create, useStore } from "zustand"; +import { createContext, useContext } from "react"; import { subscribeWithSelector } from "zustand/middleware"; import type { DaemonClient } from "@getpaseo/client/internal/daemon-client"; import type { ViewedTimelineOwner } from "@/timeline/viewed-timeline-sync"; @@ -459,7 +465,11 @@ interface SessionStoreActions { ) => void; clearSession: (serverId: string) => void; getSession: (serverId: string) => SessionState | undefined; - updateSessionClient: (serverId: string, client: DaemonClient, clientGeneration?: number) => void; + updateSessionClient: ( + serverId: string, + client: DaemonClient | null, + clientGeneration?: number, + ) => void; setViewedTimelineSync: (serverId: string, sync: ViewedTimelineOwner | null) => void; updateSessionServerInfo: (serverId: string, info: DaemonServerInfo) => void; @@ -717,1186 +727,1209 @@ function isSessionServerInfoUnchanged(input: { ); } -export const useSessionStore = create()( - subscribeWithSelector((set, get) => { - const commitActivityUpdates: AgentLastActivityCommitter = (updates) => { - set((prev) => { - let nextActivity: Map | null = null; - for (const [agentId, timestamp] of updates.entries()) { - const current = prev.agentLastActivity.get(agentId); - if (current && current.getTime() >= timestamp.getTime()) { - continue; +const createState = () => + create()( + subscribeWithSelector((set, get) => { + const commitActivityUpdates: AgentLastActivityCommitter = (updates) => { + set((prev) => { + let nextActivity: Map | null = null; + for (const [agentId, timestamp] of updates.entries()) { + const current = prev.agentLastActivity.get(agentId); + if (current && current.getTime() >= timestamp.getTime()) { + continue; + } + if (!nextActivity) { + nextActivity = new Map(prev.agentLastActivity); + } + nextActivity.set(agentId, timestamp); } if (!nextActivity) { - nextActivity = new Map(prev.agentLastActivity); - } - nextActivity.set(agentId, timestamp); - } - if (!nextActivity) { - return prev; - } - return { - ...prev, - agentLastActivity: nextActivity, - }; - }); - }; - agentLastActivityCoalescer.setCommitter(commitActivityUpdates); - - return { - sessions: {}, - agentLastActivity: new Map(), - - // Session management - initializeSession: (serverId, client, clientGeneration) => { - set((prev) => { - if (prev.sessions[serverId]) { return prev; } return { ...prev, - sessions: { - ...prev.sessions, - [serverId]: createInitialSessionState(serverId, client, clientGeneration), - }, - }; - }); - }, - - clearSession: (serverId) => { - agentLastActivityCoalescer.flushNow(); - set((prev) => { - const session = prev.sessions[serverId]; - if (!session) { - return prev; - } - const nextSessions = { ...prev.sessions }; - delete nextSessions[serverId]; - let nextActivity = prev.agentLastActivity; - if (session.agents.size > 0 || session.agentDetails.size > 0) { - const candidate = new Map(prev.agentLastActivity); - let changed = false; - for (const agentId of new Set([ - ...session.agents.keys(), - ...session.agentDetails.keys(), - ])) { - if (candidate.delete(agentId)) { - changed = true; - } - agentLastActivityCoalescer.deletePending(agentId); - } - if (changed) { - nextActivity = candidate; - } - } - return { - ...prev, - sessions: nextSessions, agentLastActivity: nextActivity, }; }); - }, + }; + agentLastActivityCoalescer.setCommitter(commitActivityUpdates); - updateSessionClient: (serverId, client, clientGeneration = 0) => { - set((prev) => { - const session = prev.sessions[serverId]; + return { + sessions: {}, + agentLastActivity: new Map(), - if (!session) { - return prev; - } - - if (session.client === client && session.clientGeneration === clientGeneration) { - return prev; - } - - return { - ...prev, - sessions: { - ...prev.sessions, - [serverId]: { - ...session, - client, - clientGeneration, - }, - }, - }; - }); - }, - - setViewedTimelineSync: (serverId, viewedTimelineSync) => { - set((prev) => { - const session = prev.sessions[serverId]; - if (!session || session.viewedTimelineSync === viewedTimelineSync) { - return prev; - } - return { - ...prev, - sessions: { - ...prev.sessions, - [serverId]: { ...session, viewedTimelineSync }, - }, - }; - }); - }, - - updateSessionServerInfo: (serverId, info) => { - set((prev) => { - const session = prev.sessions[serverId]; - if (!session) { - return prev; - } - - const nextHostname = info.hostname?.trim() || null; - const nextVersion = info.version?.trim() || null; - const nextDesktopManaged = info.desktopManaged; - const nextCapabilities = info.capabilities; - const nextFeatures = info.features; - - if ( - isSessionServerInfoUnchanged({ - currentServerInfo: session.serverInfo, - nextHostname, - nextVersion, - nextDesktopManaged, - nextCapabilities, - nextFeatures, - nextServerId: info.serverId, - }) - ) { - return prev; - } - - return { - ...prev, - sessions: { - ...prev.sessions, - [serverId]: { - ...session, - serverInfo: { - serverId: info.serverId, - hostname: nextHostname, - version: nextVersion, - ...(nextDesktopManaged !== undefined - ? { desktopManaged: nextDesktopManaged } - : {}), - ...(nextCapabilities ? { capabilities: nextCapabilities } : {}), - ...(nextFeatures ? { features: nextFeatures } : {}), - }, - }, - }, - }; - }); - }, - - getSession: (serverId) => { - return get().sessions[serverId]; - }, - - // Audio state - setIsPlayingAudio: (serverId, playing) => { - set((prev) => { - const session = prev.sessions[serverId]; - if (!session || session.isPlayingAudio === playing) { - return prev; - } - return { - ...prev, - sessions: { - ...prev.sessions, - [serverId]: { ...session, isPlayingAudio: playing }, - }, - }; - }); - }, - - // Focus - setFocusedAgentId: (serverId, agentId) => { - set((prev) => { - const session = prev.sessions[serverId]; - if (!session || session.focusedAgentId === agentId) { - return prev; - } - return { - ...prev, - sessions: { - ...prev.sessions, - [serverId]: { - ...session, - focusedAgentId: agentId, - }, - }, - }; - }); - }, - - setFocusedTerminalId: (serverId, terminalId) => { - set((prev) => { - const session = prev.sessions[serverId]; - if (!session || session.focusedTerminalId === terminalId) { - return prev; - } - return { - ...prev, - sessions: { - ...prev.sessions, - [serverId]: { - ...session, - focusedTerminalId: terminalId, - }, - }, - }; - }); - }, - - // Stream state (head/tail model) - setAgentStreamTail: (serverId, state) => { - set((prev) => { - const session = prev.sessions[serverId]; - if (!session) { - return prev; - } - const nextState = typeof state === "function" ? state(session.agentStreamTail) : state; - if (session.agentStreamTail === nextState) { - return prev; - } - return { - ...prev, - sessions: { - ...prev.sessions, - [serverId]: { ...session, agentStreamTail: nextState }, - }, - }; - }); - }, - - setAgentStreamHead: (serverId, state) => { - set((prev) => { - const session = prev.sessions[serverId]; - if (!session) { - return prev; - } - const nextState = typeof state === "function" ? state(session.agentStreamHead) : state; - if (session.agentStreamHead === nextState) { - return prev; - } - return { - ...prev, - sessions: { - ...prev.sessions, - [serverId]: { ...session, agentStreamHead: nextState }, - }, - }; - }); - }, - - setAgentStreamState: (serverId, agentId, state) => { - set((prev) => { - const session = prev.sessions[serverId]; - if (!session) { - return prev; - } - - let nextTail = session.agentStreamTail; - let nextHead = session.agentStreamHead; - let changedTail = false; - let changedHead = false; - - if (state.tail !== undefined) { - const existingTail = session.agentStreamTail.get(agentId); - if (existingTail !== state.tail) { - nextTail = new Map(session.agentStreamTail); - nextTail.set(agentId, state.tail); - changedTail = true; - } - } - - if (state.head !== undefined) { - const existingHead = session.agentStreamHead.get(agentId); - const shouldDeleteHead = state.head.length === 0; - if (shouldDeleteHead) { - if (session.agentStreamHead.has(agentId)) { - nextHead = new Map(session.agentStreamHead); - nextHead.delete(agentId); - changedHead = true; - } - } else if (existingHead !== state.head) { - nextHead = new Map(session.agentStreamHead); - nextHead.set(agentId, state.head); - changedHead = true; - } - } - - const currentSubmissions = session.messageSubmissions.get(agentId) ?? []; - const observedSubmissions = observeMessageSubmissionCanonical( - currentSubmissions, - state.acknowledgedClientMessageIds ?? [], - ); - const changedSubmissions = observedSubmissions !== currentSubmissions; - const agentTasks = updateAgentTasks(session.agentTasks, agentId, state.taskSnapshot); - const changedTasks = agentTasks !== session.agentTasks; - - if (!changedTail && !changedHead && !changedSubmissions && !changedTasks) { - return prev; - } - - let messageSubmissions = session.messageSubmissions; - if (changedSubmissions) { - messageSubmissions = new Map(session.messageSubmissions); - if (observedSubmissions.length > 0) { - messageSubmissions.set(agentId, observedSubmissions); - } else { - messageSubmissions.delete(agentId); - } - } - - return { - ...prev, - sessions: { - ...prev.sessions, - [serverId]: { - ...session, - agentStreamTail: nextTail, - agentStreamHead: nextHead, - agentTasks, - messageSubmissions, - }, - }, - }; - }); - }, - - beginAgentMessageSubmission: (serverId, agentId, message) => { - set((prev) => { - const session = prev.sessions[serverId]; - if (!session) return prev; - if (!message.clientMessageId) { - throw new Error("Beginning a message submission requires client identity"); - } - const currentTail = session.agentStreamTail.get(agentId) ?? []; - const currentHead = session.agentStreamHead.get(agentId) ?? []; - const stream = appendSubmittedUserMessage({ - tail: currentTail, - head: currentHead, - message, - }); - const submissions = beginMessageSubmission( - session.messageSubmissions.get(agentId) ?? [], - { clientMessageId: message.clientMessageId }, - ); - const messageSubmissions = new Map(session.messageSubmissions); - messageSubmissions.set(agentId, submissions); - return { - ...prev, - sessions: { - ...prev.sessions, - [serverId]: { - ...session, - agentStreamTail: - stream.tail === currentTail - ? session.agentStreamTail - : new Map(session.agentStreamTail).set(agentId, stream.tail), - agentStreamHead: - stream.head === currentHead - ? session.agentStreamHead - : new Map(session.agentStreamHead).set(agentId, stream.head), - messageSubmissions, - }, - }, - }; - }); - }, - - acceptAgentMessageSubmission: (serverId, agentId, clientMessageId) => { - set((prev) => { - const session = prev.sessions[serverId]; - if (!session) return prev; - const currentSubmissions = session.messageSubmissions.get(agentId) ?? []; - const submissions = acceptMessageSubmission(currentSubmissions, clientMessageId); - if (submissions === currentSubmissions) return prev; - const messageSubmissions = new Map(session.messageSubmissions); - if (submissions.length > 0) { - messageSubmissions.set(agentId, submissions); - } else { - messageSubmissions.delete(agentId); - } - return { - ...prev, - sessions: { - ...prev.sessions, - [serverId]: { ...session, messageSubmissions }, - }, - }; - }); - }, - - rejectAgentMessageSubmission: (serverId, agentId, clientMessageId) => { - let outcome: MessageSubmissionRejectionOutcome = "unknown"; - set((prev) => { - const session = prev.sessions[serverId]; - if (!session) return prev; - const currentTail = session.agentStreamTail.get(agentId) ?? []; - const currentHead = session.agentStreamHead.get(agentId) ?? []; - const currentSubmissions = session.messageSubmissions.get(agentId) ?? []; - const result = rejectMessageSubmission(currentSubmissions, clientMessageId); - outcome = result.outcome; - if (outcome === "unknown") return prev; - const stream = - outcome === "rejected" - ? removeSubmittedUserMessage({ - tail: currentTail, - head: currentHead, - clientMessageId, - }) - : { tail: currentTail, head: currentHead }; - const messageSubmissions = new Map(session.messageSubmissions); - if (result.submissions.length > 0) { - messageSubmissions.set(agentId, result.submissions); - } else { - messageSubmissions.delete(agentId); - } - return { - ...prev, - sessions: { - ...prev.sessions, - [serverId]: { - ...session, - agentStreamTail: - stream.tail === currentTail - ? session.agentStreamTail - : new Map(session.agentStreamTail).set(agentId, stream.tail), - agentStreamHead: - stream.head === currentHead - ? session.agentStreamHead - : new Map(session.agentStreamHead).set(agentId, stream.head), - messageSubmissions, - }, - }, - }; - }); - return outcome; - }, - - handoffCreatedAgentUserMessage: (serverId, agentId, message) => { - let didHandoff = false; - set((prev) => { - const session = prev.sessions[serverId]; - if (!session) { - return prev; - } - - const currentTail = session.agentStreamTail.get(agentId) ?? []; - const currentHead = session.agentStreamHead.get(agentId) ?? []; - const result = handoffCreatedAgentUserMessageToStream({ - tail: currentTail, - head: currentHead, - message, - }); - if (!result.changedTail && !result.changedHead) { - return prev; - } - - const nextTail = result.changedTail - ? new Map(session.agentStreamTail).set(agentId, result.tail) - : session.agentStreamTail; - const nextHead = result.changedHead - ? new Map(session.agentStreamHead).set(agentId, result.head) - : session.agentStreamHead; - didHandoff = true; - - return { - ...prev, - sessions: { - ...prev.sessions, - [serverId]: { - ...session, - agentStreamTail: nextTail, - agentStreamHead: nextHead, - }, - }, - }; - }); - return didHandoff; - }, - - clearAgentStreamHead: (serverId, agentId) => { - set((prev) => { - const session = prev.sessions[serverId]; - if (!session) { - return prev; - } - if (!session.agentStreamHead.has(agentId)) { - return prev; - } - const nextHead = new Map(session.agentStreamHead); - nextHead.delete(agentId); - return { - ...prev, - sessions: { - ...prev.sessions, - [serverId]: { ...session, agentStreamHead: nextHead }, - }, - }; - }); - }, - - setAgentTimelineCursor: (serverId, state) => { - set((prev) => { - const session = prev.sessions[serverId]; - if (!session) { - return prev; - } - const nextState = - typeof state === "function" ? state(session.agentTimelineCursor) : state; - if (session.agentTimelineCursor === nextState) { - return prev; - } - return { - ...prev, - sessions: { - ...prev.sessions, - [serverId]: { ...session, agentTimelineCursor: nextState }, - }, - }; - }); - }, - - setAgentTimelineHasOlder: (serverId, state) => { - set((prev) => { - const session = prev.sessions[serverId]; - if (!session) { - return prev; - } - const nextState = - typeof state === "function" ? state(session.agentTimelineHasOlder) : state; - if (session.agentTimelineHasOlder === nextState) { - return prev; - } - return { - ...prev, - sessions: { - ...prev.sessions, - [serverId]: { ...session, agentTimelineHasOlder: nextState }, - }, - }; - }); - }, - - setAgentTimelineHasNewer: (serverId, state) => { - set((prev) => { - const session = prev.sessions[serverId]; - if (!session) return prev; - const nextState = - typeof state === "function" ? state(session.agentTimelineHasNewer) : state; - if (session.agentTimelineHasNewer === nextState) return prev; - return { - ...prev, - sessions: { - ...prev.sessions, - [serverId]: { ...session, agentTimelineHasNewer: nextState }, - }, - }; - }); - }, - - setAgentTimelineOlderFetchInFlight: (serverId, state) => { - set((prev) => { - const session = prev.sessions[serverId]; - if (!session) { - return prev; - } - const nextState = - typeof state === "function" ? state(session.agentTimelineOlderFetchInFlight) : state; - if (session.agentTimelineOlderFetchInFlight === nextState) { - return prev; - } - return { - ...prev, - sessions: { - ...prev.sessions, - [serverId]: { ...session, agentTimelineOlderFetchInFlight: nextState }, - }, - }; - }); - }, - - bumpHistorySyncGeneration: (serverId) => { - set((prev) => { - const session = prev.sessions[serverId]; - if (!session) { - return prev; - } - const nextGeneration = session.historySyncGeneration + 1; - return { - ...prev, - sessions: { - ...prev.sessions, - [serverId]: { - ...session, - historySyncGeneration: nextGeneration, - }, - }, - }; - }); - }, - - markAgentHistorySynchronized: (serverId, agentId) => { - set((prev) => { - const session = prev.sessions[serverId]; - if (!session) { - return prev; - } - const currentGeneration = session.historySyncGeneration; - const previousGeneration = session.agentHistorySyncGeneration.get(agentId); - if (previousGeneration === currentGeneration) { - return prev; - } - const nextMap = new Map(session.agentHistorySyncGeneration); - nextMap.set(agentId, currentGeneration); - return { - ...prev, - sessions: { - ...prev.sessions, - [serverId]: { - ...session, - agentHistorySyncGeneration: nextMap, - }, - }, - }; - }); - }, - - setAgentAuthoritativeHistoryApplied: (serverId, agentId, applied) => { - set((prev) => { - const session = prev.sessions[serverId]; - if (!session) { - return prev; - } - - const previousApplied = session.agentAuthoritativeHistoryApplied.get(agentId) ?? false; - if (previousApplied === applied) { - return prev; - } - - const nextApplied = new Map(session.agentAuthoritativeHistoryApplied); - if (applied) { - nextApplied.set(agentId, true); - } else { - nextApplied.delete(agentId); - } - - return { - ...prev, - sessions: { - ...prev.sessions, - [serverId]: { - ...session, - agentAuthoritativeHistoryApplied: nextApplied, - }, - }, - }; - }); - }, - - applyAgentTimelineResponseState: (serverId, agentId, state) => { - set((prev) => { - const session = prev.sessions[serverId]; - if (!session) return prev; - - const nextTail = new Map(session.agentStreamTail); - nextTail.set(agentId, state.items); - const nextHead = new Map(session.agentStreamHead); - if (state.head.length > 0) nextHead.set(agentId, state.head); - else nextHead.delete(agentId); - const nextCursor = new Map(session.agentTimelineCursor); - if (state.range) nextCursor.set(agentId, state.range); - else nextCursor.delete(agentId); - const nextHasOlder = new Map(session.agentTimelineHasOlder); - if (state.older !== "unchanged") { - nextHasOlder.set(agentId, state.older === "available"); - } - const nextHasNewer = new Map(session.agentTimelineHasNewer); - nextHasNewer.set(agentId, state.newer); - const nextAuthoritative = new Map(session.agentAuthoritativeHistoryApplied); - const nextSyncGeneration = new Map(session.agentHistorySyncGeneration); - const currentSubmissions = session.messageSubmissions.get(agentId) ?? []; - const observedSubmissions = observeMessageSubmissionCanonical( - currentSubmissions, - state.acknowledgedClientMessageIds, - ); - let messageSubmissions = session.messageSubmissions; - if (observedSubmissions !== currentSubmissions) { - messageSubmissions = new Map(session.messageSubmissions); - if (observedSubmissions.length > 0) { - messageSubmissions.set(agentId, observedSubmissions); - } else { - messageSubmissions.delete(agentId); - } - } - if (state.synchronized) { - nextAuthoritative.set(agentId, true); - nextSyncGeneration.set(agentId, session.historySyncGeneration); - } - const tasks = latestTasksFromStream([...state.items, ...state.head]); - const agentTasks = new Map(session.agentTasks); - if (tasks.length > 0) agentTasks.set(agentId, tasks); - else agentTasks.delete(agentId); - - return { - ...prev, - sessions: { - ...prev.sessions, - [serverId]: { - ...session, - agentStreamTail: nextTail, - agentStreamHead: nextHead, - agentTasks, - agentTimelineCursor: nextCursor, - agentTimelineHasOlder: nextHasOlder, - agentTimelineHasNewer: nextHasNewer, - agentAuthoritativeHistoryApplied: nextAuthoritative, - agentHistorySyncGeneration: nextSyncGeneration, - messageSubmissions, - }, - }, - }; - }); - }, - - // Initializing agents - setInitializingAgents: (serverId, state) => { - set((prev) => { - const session = prev.sessions[serverId]; - if (!session) { - return prev; - } - const nextState = typeof state === "function" ? state(session.initializingAgents) : state; - if (session.initializingAgents === nextState) { - return prev; - } - return { - ...prev, - sessions: { - ...prev.sessions, - [serverId]: { ...session, initializingAgents: nextState }, - }, - }; - }); - }, - - // Agents - setAgents: (serverId, agents) => { - set((prev) => { - const session = prev.sessions[serverId]; - if (!session) { - return prev; - } - const nextAgents = typeof agents === "function" ? agents(session.agents) : agents; - if (session.agents === nextAgents) { - return prev; - } - return { - ...prev, - sessions: { - ...prev.sessions, - [serverId]: { - ...session, - agents: nextAgents, - workspaceAgentActivity: buildWorkspaceAgentActivityIndex( - nextAgents, - session.workspaceAgentActivity, - ), - }, - }, - }; - }); - }, - - setAgentDetails: (serverId, agents) => { - set((prev) => { - const session = prev.sessions[serverId]; - if (!session) { - return prev; - } - const nextAgents = typeof agents === "function" ? agents(session.agentDetails) : agents; - if (session.agentDetails === nextAgents) { - return prev; - } - return { - ...prev, - sessions: { - ...prev.sessions, - [serverId]: { ...session, agentDetails: nextAgents }, - }, - }; - }); - }, - - setWorkspaces: (serverId, workspaces) => { - set((prev) => { - const session = prev.sessions[serverId]; - if (!session) { - return prev; - } - const nextWorkspaces = - typeof workspaces === "function" ? workspaces(session.workspaces) : workspaces; - const preservedWorkspaces = preserveMapIdentity(session.workspaces, nextWorkspaces); - if (session.workspaces === preservedWorkspaces) { - return prev; - } - return { - ...prev, - sessions: { - ...prev.sessions, - [serverId]: { ...session, workspaces: preservedWorkspaces }, - }, - }; - }); - }, - - setProjects: (serverId, projects) => { - const next = new Map(); - for (const project of projects) next.set(project.projectId, project); - set((prev) => { - const session = prev.sessions[serverId]; - if (!session) return prev; - const preservedProjects = preserveMapIdentity(session.projects, next); - if (session.projects === preservedProjects) return prev; - return { - ...prev, - sessions: { - ...prev.sessions, - [serverId]: { ...session, projects: preservedProjects }, - }, - }; - }); - }, - - upsertProject: (serverId, project) => { - set((prev) => { - const session = prev.sessions[serverId]; - if (!session || equal(session.projects.get(project.projectId), project)) return prev; - const projects = new Map(session.projects); - projects.set(project.projectId, project); - return { - ...prev, - sessions: { ...prev.sessions, [serverId]: { ...session, projects } }, - }; - }); - }, - - removeProject: (serverId, projectId) => { - set((prev) => { - const session = prev.sessions[serverId]; - if (!session?.projects.has(projectId)) return prev; - const projects = new Map(session.projects); - projects.delete(projectId); - return { - ...prev, - sessions: { ...prev.sessions, [serverId]: { ...session, projects } }, - }; - }); - }, - - setWorkspaceRestoreStatus: (serverId, workspaceId, status) => { - set((prev) => { - const session = prev.sessions[serverId]; - if (!session) { - return prev; - } - if (session.restoringWorkspaces.get(workspaceId) === status) { - return prev; - } - // A late dir-gone timeout must not override a successful restore: - // only mark failed while still restoring and the descriptor is absent. - if ( - status === "failed" && - (session.restoringWorkspaces.get(workspaceId) !== "restoring" || - session.workspaces.has(workspaceId)) - ) { - return prev; - } - const next = new Map(session.restoringWorkspaces); - next.set(workspaceId, status); - return { - ...prev, - sessions: { - ...prev.sessions, - [serverId]: { ...session, restoringWorkspaces: next }, - }, - }; - }); - }, - - clearWorkspaceRestoreStatus: (serverId, workspaceId) => { - set((prev) => { - const session = prev.sessions[serverId]; - if (!session || !session.restoringWorkspaces.has(workspaceId)) { - return prev; - } - const next = new Map(session.restoringWorkspaces); - next.delete(workspaceId); - return { - ...prev, - sessions: { - ...prev.sessions, - [serverId]: { ...session, restoringWorkspaces: next }, - }, - }; - }); - }, - - mergeWorkspaces: (serverId, workspaces) => { - const nextEntries = Array.from(workspaces); - set((prev) => { - const session = prev.sessions[serverId]; - if (!session || nextEntries.length === 0) { - return prev; - } - const next = new Map(session.workspaces); - let changed = false; - // A descriptor arriving is the success signal for a pending restore: - // clear it at the source so every entry point converges to "ready". - let nextRestoring: Map | null = null; - for (const workspace of nextEntries) { - if (session.restoringWorkspaces.has(workspace.id)) { - nextRestoring ??= new Map(session.restoringWorkspaces); - nextRestoring.delete(workspace.id); - changed = true; - } - const existing = next.get(workspace.id); - const nextWorkspace = preserveWorkspaceDescriptorIdentity(workspace, existing); - if (existing === nextWorkspace) { - continue; - } - next.set(workspace.id, nextWorkspace); - changed = true; - } - if (!changed) { - return prev; - } - return { - ...prev, - sessions: { - ...prev.sessions, - [serverId]: { - ...session, - workspaces: next, - restoringWorkspaces: nextRestoring ?? session.restoringWorkspaces, - }, - }, - }; - }); - }, - - removeWorkspace: (serverId, workspaceId) => { - set((prev) => { - const session = prev.sessions[serverId]; - const workspaceKey = resolveWorkspaceMapKeyByIdentity({ - workspaces: session?.workspaces, - workspaceId, - }); - if (!session || !workspaceKey) { - return prev; - } - const next = new Map(session.workspaces); - next.delete(workspaceKey); - return { - ...prev, - sessions: { - ...prev.sessions, - [serverId]: { ...session, workspaces: next }, - }, - }; - }); - }, - - // Agent activity timestamps (top-level, does NOT mutate session object) - setAgentLastActivity: (agentId, timestamp) => { - agentLastActivityCoalescer.enqueue(agentId, timestamp); - }, - - setAgentLastActivityBatch: (updates) => { - set((prev) => { - const nextActivity = - typeof updates === "function" ? updates(prev.agentLastActivity) : updates; - if (nextActivity === prev.agentLastActivity) { - return prev; - } - if (nextActivity.size === 0) { - if (prev.agentLastActivity.size === 0) { + // Session management + initializeSession: (serverId, client, clientGeneration) => { + set((prev) => { + if (prev.sessions[serverId]) { return prev; } return { ...prev, - agentLastActivity: new Map(), + sessions: { + ...prev.sessions, + [serverId]: createInitialSessionState(serverId, client, clientGeneration), + }, }; - } - let changed = false; - for (const [agentId, timestamp] of nextActivity.entries()) { - const currentTimestamp = prev.agentLastActivity.get(agentId); - if (!currentTimestamp || currentTimestamp.getTime() !== timestamp.getTime()) { - changed = true; - break; - } - } - if (!changed && nextActivity.size === prev.agentLastActivity.size) { - return prev; - } - return { - ...prev, - agentLastActivity: new Map(nextActivity), - }; - }); - }, - - flushAgentLastActivity: () => { - agentLastActivityCoalescer.flushNow(); - }, - - // Permissions - setPendingPermissions: (serverId, perms) => { - set((prev) => { - const session = prev.sessions[serverId]; - if (!session) { - return prev; - } - const nextPerms = typeof perms === "function" ? perms(session.pendingPermissions) : perms; - if (session.pendingPermissions === nextPerms) { - return prev; - } - return { - ...prev, - sessions: { - ...prev.sessions, - [serverId]: { ...session, pendingPermissions: nextPerms }, - }, - }; - }); - }, - - // File explorer - setFileExplorer: (serverId, state) => { - set((prev) => { - const session = prev.sessions[serverId]; - if (!session) { - return prev; - } - const nextState = typeof state === "function" ? state(session.fileExplorer) : state; - if (session.fileExplorer === nextState) { - return prev; - } - return { - ...prev, - sessions: { - ...prev.sessions, - [serverId]: { ...session, fileExplorer: nextState }, - }, - }; - }); - }, - - // Queued messages - setQueuedMessages: (serverId, value) => { - set((prev) => { - const session = prev.sessions[serverId]; - if (!session) { - return prev; - } - const nextValue = typeof value === "function" ? value(session.queuedMessages) : value; - if (session.queuedMessages === nextValue) { - return prev; - } - return { - ...prev, - sessions: { - ...prev.sessions, - [serverId]: { ...session, queuedMessages: nextValue }, - }, - }; - }); - }, - - // Hydration - setHasHydratedAgents: (serverId, hydrated) => { - set((prev) => { - const session = prev.sessions[serverId]; - if (!session || session.hasHydratedAgents === hydrated) { - return prev; - } - return { - ...prev, - sessions: { - ...prev.sessions, - [serverId]: { ...session, hasHydratedAgents: hydrated }, - }, - }; - }); - }, - - setHasHydratedWorkspaces: (serverId, hydrated) => { - set((prev) => { - const session = prev.sessions[serverId]; - if (!session || session.hasHydratedWorkspaces === hydrated) { - return prev; - } - return { - ...prev, - sessions: { - ...prev.sessions, - [serverId]: { ...session, hasHydratedWorkspaces: hydrated }, - }, - }; - }); - }, - - setHasWorkspaceDirectorySnapshot: (serverId, available) => { - set((prev) => { - const session = prev.sessions[serverId]; - if (!session || session.hasWorkspaceDirectorySnapshot === available) return prev; - return { - ...prev, - sessions: { - ...prev.sessions, - [serverId]: { ...session, hasWorkspaceDirectorySnapshot: available }, - }, - }; - }); - }, - - // Agent directory - derived from agents (computed on-demand) - getAgentDirectory: (serverId) => { - const state = get(); - const session = state.sessions[serverId]; - if (!session) { - return undefined; - } - - const entries: AgentDirectoryEntry[] = []; - for (const agent of session.agents.values()) { - // Get lastActivityAt from top-level slice, fallback to agent.lastActivityAt - const lastActivityAt = state.agentLastActivity.get(agent.id) ?? agent.lastActivityAt; - entries.push({ - id: agent.id, - serverId, - title: agent.title ?? null, - status: agent.status, - turn: agent.turn, - lastActivityAt, - cwd: agent.cwd, - provider: agent.provider, - pendingPermissionCount: agent.pendingPermissions.length, - requiresAttention: agent.requiresAttention ?? false, - attentionReason: agent.attentionReason ?? null, - attentionTimestamp: agent.attentionTimestamp ?? null, - createdAt: agent.createdAt, - labels: agent.labels, }); - } - return entries; - }, - }; - }), -); + }, + + clearSession: (serverId) => { + agentLastActivityCoalescer.flushNow(); + set((prev) => { + const session = prev.sessions[serverId]; + if (!session) { + return prev; + } + const nextSessions = { ...prev.sessions }; + delete nextSessions[serverId]; + let nextActivity = prev.agentLastActivity; + if (session.agents.size > 0 || session.agentDetails.size > 0) { + const candidate = new Map(prev.agentLastActivity); + let changed = false; + for (const agentId of new Set([ + ...session.agents.keys(), + ...session.agentDetails.keys(), + ])) { + if (candidate.delete(agentId)) { + changed = true; + } + agentLastActivityCoalescer.deletePending(agentId); + } + if (changed) { + nextActivity = candidate; + } + } + return { + ...prev, + sessions: nextSessions, + agentLastActivity: nextActivity, + }; + }); + }, + + updateSessionClient: (serverId, client, clientGeneration = 0) => { + set((prev) => { + const session = prev.sessions[serverId]; + + if (!session) { + return prev; + } + + if (session.client === client && session.clientGeneration === clientGeneration) { + return prev; + } + + return { + ...prev, + sessions: { + ...prev.sessions, + [serverId]: { + ...session, + client, + clientGeneration, + }, + }, + }; + }); + }, + + setViewedTimelineSync: (serverId, viewedTimelineSync) => { + set((prev) => { + const session = prev.sessions[serverId]; + if (!session || session.viewedTimelineSync === viewedTimelineSync) { + return prev; + } + return { + ...prev, + sessions: { + ...prev.sessions, + [serverId]: { ...session, viewedTimelineSync }, + }, + }; + }); + }, + + updateSessionServerInfo: (serverId, info) => { + set((prev) => { + const session = prev.sessions[serverId]; + if (!session) { + return prev; + } + + const nextHostname = info.hostname?.trim() || null; + const nextVersion = info.version?.trim() || null; + const nextDesktopManaged = info.desktopManaged; + const nextCapabilities = info.capabilities; + const nextFeatures = info.features; + + if ( + isSessionServerInfoUnchanged({ + currentServerInfo: session.serverInfo, + nextHostname, + nextVersion, + nextDesktopManaged, + nextCapabilities, + nextFeatures, + nextServerId: info.serverId, + }) + ) { + return prev; + } + + return { + ...prev, + sessions: { + ...prev.sessions, + [serverId]: { + ...session, + serverInfo: { + serverId: info.serverId, + hostname: nextHostname, + version: nextVersion, + ...(nextDesktopManaged !== undefined + ? { desktopManaged: nextDesktopManaged } + : {}), + ...(nextCapabilities ? { capabilities: nextCapabilities } : {}), + ...(nextFeatures ? { features: nextFeatures } : {}), + }, + }, + }, + }; + }); + }, + + getSession: (serverId) => { + return get().sessions[serverId]; + }, + + // Audio state + setIsPlayingAudio: (serverId, playing) => { + set((prev) => { + const session = prev.sessions[serverId]; + if (!session || session.isPlayingAudio === playing) { + return prev; + } + return { + ...prev, + sessions: { + ...prev.sessions, + [serverId]: { ...session, isPlayingAudio: playing }, + }, + }; + }); + }, + + // Focus + setFocusedAgentId: (serverId, agentId) => { + set((prev) => { + const session = prev.sessions[serverId]; + if (!session || session.focusedAgentId === agentId) { + return prev; + } + return { + ...prev, + sessions: { + ...prev.sessions, + [serverId]: { + ...session, + focusedAgentId: agentId, + }, + }, + }; + }); + }, + + setFocusedTerminalId: (serverId, terminalId) => { + set((prev) => { + const session = prev.sessions[serverId]; + if (!session || session.focusedTerminalId === terminalId) { + return prev; + } + return { + ...prev, + sessions: { + ...prev.sessions, + [serverId]: { + ...session, + focusedTerminalId: terminalId, + }, + }, + }; + }); + }, + + // Stream state (head/tail model) + setAgentStreamTail: (serverId, state) => { + set((prev) => { + const session = prev.sessions[serverId]; + if (!session) { + return prev; + } + const nextState = typeof state === "function" ? state(session.agentStreamTail) : state; + if (session.agentStreamTail === nextState) { + return prev; + } + return { + ...prev, + sessions: { + ...prev.sessions, + [serverId]: { ...session, agentStreamTail: nextState }, + }, + }; + }); + }, + + setAgentStreamHead: (serverId, state) => { + set((prev) => { + const session = prev.sessions[serverId]; + if (!session) { + return prev; + } + const nextState = typeof state === "function" ? state(session.agentStreamHead) : state; + if (session.agentStreamHead === nextState) { + return prev; + } + return { + ...prev, + sessions: { + ...prev.sessions, + [serverId]: { ...session, agentStreamHead: nextState }, + }, + }; + }); + }, + + setAgentStreamState: (serverId, agentId, state) => { + set((prev) => { + const session = prev.sessions[serverId]; + if (!session) { + return prev; + } + + let nextTail = session.agentStreamTail; + let nextHead = session.agentStreamHead; + let changedTail = false; + let changedHead = false; + + if (state.tail !== undefined) { + const existingTail = session.agentStreamTail.get(agentId); + if (existingTail !== state.tail) { + nextTail = new Map(session.agentStreamTail); + nextTail.set(agentId, state.tail); + changedTail = true; + } + } + + if (state.head !== undefined) { + const existingHead = session.agentStreamHead.get(agentId); + const shouldDeleteHead = state.head.length === 0; + if (shouldDeleteHead) { + if (session.agentStreamHead.has(agentId)) { + nextHead = new Map(session.agentStreamHead); + nextHead.delete(agentId); + changedHead = true; + } + } else if (existingHead !== state.head) { + nextHead = new Map(session.agentStreamHead); + nextHead.set(agentId, state.head); + changedHead = true; + } + } + + const currentSubmissions = session.messageSubmissions.get(agentId) ?? []; + const observedSubmissions = observeMessageSubmissionCanonical( + currentSubmissions, + state.acknowledgedClientMessageIds ?? [], + ); + const changedSubmissions = observedSubmissions !== currentSubmissions; + const agentTasks = updateAgentTasks(session.agentTasks, agentId, state.taskSnapshot); + const changedTasks = agentTasks !== session.agentTasks; + + if (!changedTail && !changedHead && !changedSubmissions && !changedTasks) { + return prev; + } + + let messageSubmissions = session.messageSubmissions; + if (changedSubmissions) { + messageSubmissions = new Map(session.messageSubmissions); + if (observedSubmissions.length > 0) { + messageSubmissions.set(agentId, observedSubmissions); + } else { + messageSubmissions.delete(agentId); + } + } + + return { + ...prev, + sessions: { + ...prev.sessions, + [serverId]: { + ...session, + agentStreamTail: nextTail, + agentStreamHead: nextHead, + agentTasks, + messageSubmissions, + }, + }, + }; + }); + }, + + beginAgentMessageSubmission: (serverId, agentId, message) => { + set((prev) => { + const session = prev.sessions[serverId]; + if (!session) return prev; + if (!message.clientMessageId) { + throw new Error("Beginning a message submission requires client identity"); + } + const currentTail = session.agentStreamTail.get(agentId) ?? []; + const currentHead = session.agentStreamHead.get(agentId) ?? []; + const stream = appendSubmittedUserMessage({ + tail: currentTail, + head: currentHead, + message, + }); + const submissions = beginMessageSubmission( + session.messageSubmissions.get(agentId) ?? [], + { clientMessageId: message.clientMessageId }, + ); + const messageSubmissions = new Map(session.messageSubmissions); + messageSubmissions.set(agentId, submissions); + return { + ...prev, + sessions: { + ...prev.sessions, + [serverId]: { + ...session, + agentStreamTail: + stream.tail === currentTail + ? session.agentStreamTail + : new Map(session.agentStreamTail).set(agentId, stream.tail), + agentStreamHead: + stream.head === currentHead + ? session.agentStreamHead + : new Map(session.agentStreamHead).set(agentId, stream.head), + messageSubmissions, + }, + }, + }; + }); + }, + + acceptAgentMessageSubmission: (serverId, agentId, clientMessageId) => { + set((prev) => { + const session = prev.sessions[serverId]; + if (!session) return prev; + const currentSubmissions = session.messageSubmissions.get(agentId) ?? []; + const submissions = acceptMessageSubmission(currentSubmissions, clientMessageId); + if (submissions === currentSubmissions) return prev; + const messageSubmissions = new Map(session.messageSubmissions); + if (submissions.length > 0) { + messageSubmissions.set(agentId, submissions); + } else { + messageSubmissions.delete(agentId); + } + return { + ...prev, + sessions: { + ...prev.sessions, + [serverId]: { ...session, messageSubmissions }, + }, + }; + }); + }, + + rejectAgentMessageSubmission: (serverId, agentId, clientMessageId) => { + let outcome: MessageSubmissionRejectionOutcome = "unknown"; + set((prev) => { + const session = prev.sessions[serverId]; + if (!session) return prev; + const currentTail = session.agentStreamTail.get(agentId) ?? []; + const currentHead = session.agentStreamHead.get(agentId) ?? []; + const currentSubmissions = session.messageSubmissions.get(agentId) ?? []; + const result = rejectMessageSubmission(currentSubmissions, clientMessageId); + outcome = result.outcome; + if (outcome === "unknown") return prev; + const stream = + outcome === "rejected" + ? removeSubmittedUserMessage({ + tail: currentTail, + head: currentHead, + clientMessageId, + }) + : { tail: currentTail, head: currentHead }; + const messageSubmissions = new Map(session.messageSubmissions); + if (result.submissions.length > 0) { + messageSubmissions.set(agentId, result.submissions); + } else { + messageSubmissions.delete(agentId); + } + return { + ...prev, + sessions: { + ...prev.sessions, + [serverId]: { + ...session, + agentStreamTail: + stream.tail === currentTail + ? session.agentStreamTail + : new Map(session.agentStreamTail).set(agentId, stream.tail), + agentStreamHead: + stream.head === currentHead + ? session.agentStreamHead + : new Map(session.agentStreamHead).set(agentId, stream.head), + messageSubmissions, + }, + }, + }; + }); + return outcome; + }, + + handoffCreatedAgentUserMessage: (serverId, agentId, message) => { + let didHandoff = false; + set((prev) => { + const session = prev.sessions[serverId]; + if (!session) { + return prev; + } + + const currentTail = session.agentStreamTail.get(agentId) ?? []; + const currentHead = session.agentStreamHead.get(agentId) ?? []; + const result = handoffCreatedAgentUserMessageToStream({ + tail: currentTail, + head: currentHead, + message, + }); + if (!result.changedTail && !result.changedHead) { + return prev; + } + + const nextTail = result.changedTail + ? new Map(session.agentStreamTail).set(agentId, result.tail) + : session.agentStreamTail; + const nextHead = result.changedHead + ? new Map(session.agentStreamHead).set(agentId, result.head) + : session.agentStreamHead; + didHandoff = true; + + return { + ...prev, + sessions: { + ...prev.sessions, + [serverId]: { + ...session, + agentStreamTail: nextTail, + agentStreamHead: nextHead, + }, + }, + }; + }); + return didHandoff; + }, + + clearAgentStreamHead: (serverId, agentId) => { + set((prev) => { + const session = prev.sessions[serverId]; + if (!session) { + return prev; + } + if (!session.agentStreamHead.has(agentId)) { + return prev; + } + const nextHead = new Map(session.agentStreamHead); + nextHead.delete(agentId); + return { + ...prev, + sessions: { + ...prev.sessions, + [serverId]: { ...session, agentStreamHead: nextHead }, + }, + }; + }); + }, + + setAgentTimelineCursor: (serverId, state) => { + set((prev) => { + const session = prev.sessions[serverId]; + if (!session) { + return prev; + } + const nextState = + typeof state === "function" ? state(session.agentTimelineCursor) : state; + if (session.agentTimelineCursor === nextState) { + return prev; + } + return { + ...prev, + sessions: { + ...prev.sessions, + [serverId]: { ...session, agentTimelineCursor: nextState }, + }, + }; + }); + }, + + setAgentTimelineHasOlder: (serverId, state) => { + set((prev) => { + const session = prev.sessions[serverId]; + if (!session) { + return prev; + } + const nextState = + typeof state === "function" ? state(session.agentTimelineHasOlder) : state; + if (session.agentTimelineHasOlder === nextState) { + return prev; + } + return { + ...prev, + sessions: { + ...prev.sessions, + [serverId]: { ...session, agentTimelineHasOlder: nextState }, + }, + }; + }); + }, + + setAgentTimelineHasNewer: (serverId, state) => { + set((prev) => { + const session = prev.sessions[serverId]; + if (!session) return prev; + const nextState = + typeof state === "function" ? state(session.agentTimelineHasNewer) : state; + if (session.agentTimelineHasNewer === nextState) return prev; + return { + ...prev, + sessions: { + ...prev.sessions, + [serverId]: { ...session, agentTimelineHasNewer: nextState }, + }, + }; + }); + }, + + setAgentTimelineOlderFetchInFlight: (serverId, state) => { + set((prev) => { + const session = prev.sessions[serverId]; + if (!session) { + return prev; + } + const nextState = + typeof state === "function" ? state(session.agentTimelineOlderFetchInFlight) : state; + if (session.agentTimelineOlderFetchInFlight === nextState) { + return prev; + } + return { + ...prev, + sessions: { + ...prev.sessions, + [serverId]: { ...session, agentTimelineOlderFetchInFlight: nextState }, + }, + }; + }); + }, + + bumpHistorySyncGeneration: (serverId) => { + set((prev) => { + const session = prev.sessions[serverId]; + if (!session) { + return prev; + } + const nextGeneration = session.historySyncGeneration + 1; + return { + ...prev, + sessions: { + ...prev.sessions, + [serverId]: { + ...session, + historySyncGeneration: nextGeneration, + }, + }, + }; + }); + }, + + markAgentHistorySynchronized: (serverId, agentId) => { + set((prev) => { + const session = prev.sessions[serverId]; + if (!session) { + return prev; + } + const currentGeneration = session.historySyncGeneration; + const previousGeneration = session.agentHistorySyncGeneration.get(agentId); + if (previousGeneration === currentGeneration) { + return prev; + } + const nextMap = new Map(session.agentHistorySyncGeneration); + nextMap.set(agentId, currentGeneration); + return { + ...prev, + sessions: { + ...prev.sessions, + [serverId]: { + ...session, + agentHistorySyncGeneration: nextMap, + }, + }, + }; + }); + }, + + setAgentAuthoritativeHistoryApplied: (serverId, agentId, applied) => { + set((prev) => { + const session = prev.sessions[serverId]; + if (!session) { + return prev; + } + + const previousApplied = session.agentAuthoritativeHistoryApplied.get(agentId) ?? false; + if (previousApplied === applied) { + return prev; + } + + const nextApplied = new Map(session.agentAuthoritativeHistoryApplied); + if (applied) { + nextApplied.set(agentId, true); + } else { + nextApplied.delete(agentId); + } + + return { + ...prev, + sessions: { + ...prev.sessions, + [serverId]: { + ...session, + agentAuthoritativeHistoryApplied: nextApplied, + }, + }, + }; + }); + }, + + applyAgentTimelineResponseState: (serverId, agentId, state) => { + set((prev) => { + const session = prev.sessions[serverId]; + if (!session) return prev; + + const nextTail = new Map(session.agentStreamTail); + nextTail.set(agentId, state.items); + const nextHead = new Map(session.agentStreamHead); + if (state.head.length > 0) nextHead.set(agentId, state.head); + else nextHead.delete(agentId); + const nextCursor = new Map(session.agentTimelineCursor); + if (state.range) nextCursor.set(agentId, state.range); + else nextCursor.delete(agentId); + const nextHasOlder = new Map(session.agentTimelineHasOlder); + if (state.older !== "unchanged") { + nextHasOlder.set(agentId, state.older === "available"); + } + const nextHasNewer = new Map(session.agentTimelineHasNewer); + nextHasNewer.set(agentId, state.newer); + const nextAuthoritative = new Map(session.agentAuthoritativeHistoryApplied); + const nextSyncGeneration = new Map(session.agentHistorySyncGeneration); + const currentSubmissions = session.messageSubmissions.get(agentId) ?? []; + const observedSubmissions = observeMessageSubmissionCanonical( + currentSubmissions, + state.acknowledgedClientMessageIds, + ); + let messageSubmissions = session.messageSubmissions; + if (observedSubmissions !== currentSubmissions) { + messageSubmissions = new Map(session.messageSubmissions); + if (observedSubmissions.length > 0) { + messageSubmissions.set(agentId, observedSubmissions); + } else { + messageSubmissions.delete(agentId); + } + } + if (state.synchronized) { + nextAuthoritative.set(agentId, true); + nextSyncGeneration.set(agentId, session.historySyncGeneration); + } + const tasks = latestTasksFromStream([...state.items, ...state.head]); + const agentTasks = new Map(session.agentTasks); + if (tasks.length > 0) agentTasks.set(agentId, tasks); + else agentTasks.delete(agentId); + + return { + ...prev, + sessions: { + ...prev.sessions, + [serverId]: { + ...session, + agentStreamTail: nextTail, + agentStreamHead: nextHead, + agentTasks, + agentTimelineCursor: nextCursor, + agentTimelineHasOlder: nextHasOlder, + agentTimelineHasNewer: nextHasNewer, + agentAuthoritativeHistoryApplied: nextAuthoritative, + agentHistorySyncGeneration: nextSyncGeneration, + messageSubmissions, + }, + }, + }; + }); + }, + + // Initializing agents + setInitializingAgents: (serverId, state) => { + set((prev) => { + const session = prev.sessions[serverId]; + if (!session) { + return prev; + } + const nextState = + typeof state === "function" ? state(session.initializingAgents) : state; + if (session.initializingAgents === nextState) { + return prev; + } + return { + ...prev, + sessions: { + ...prev.sessions, + [serverId]: { ...session, initializingAgents: nextState }, + }, + }; + }); + }, + + // Agents + setAgents: (serverId, agents) => { + set((prev) => { + const session = prev.sessions[serverId]; + if (!session) { + return prev; + } + const nextAgents = typeof agents === "function" ? agents(session.agents) : agents; + if (session.agents === nextAgents) { + return prev; + } + return { + ...prev, + sessions: { + ...prev.sessions, + [serverId]: { + ...session, + agents: nextAgents, + workspaceAgentActivity: buildWorkspaceAgentActivityIndex( + nextAgents, + session.workspaceAgentActivity, + ), + }, + }, + }; + }); + }, + + setAgentDetails: (serverId, agents) => { + set((prev) => { + const session = prev.sessions[serverId]; + if (!session) { + return prev; + } + const nextAgents = typeof agents === "function" ? agents(session.agentDetails) : agents; + if (session.agentDetails === nextAgents) { + return prev; + } + return { + ...prev, + sessions: { + ...prev.sessions, + [serverId]: { ...session, agentDetails: nextAgents }, + }, + }; + }); + }, + + setWorkspaces: (serverId, workspaces) => { + set((prev) => { + const session = prev.sessions[serverId]; + if (!session) { + return prev; + } + const nextWorkspaces = + typeof workspaces === "function" ? workspaces(session.workspaces) : workspaces; + const preservedWorkspaces = preserveMapIdentity(session.workspaces, nextWorkspaces); + if (session.workspaces === preservedWorkspaces) { + return prev; + } + return { + ...prev, + sessions: { + ...prev.sessions, + [serverId]: { ...session, workspaces: preservedWorkspaces }, + }, + }; + }); + }, + + setProjects: (serverId, projects) => { + const next = new Map(); + for (const project of projects) next.set(project.projectId, project); + set((prev) => { + const session = prev.sessions[serverId]; + if (!session) return prev; + const preservedProjects = preserveMapIdentity(session.projects, next); + if (session.projects === preservedProjects) return prev; + return { + ...prev, + sessions: { + ...prev.sessions, + [serverId]: { ...session, projects: preservedProjects }, + }, + }; + }); + }, + + upsertProject: (serverId, project) => { + set((prev) => { + const session = prev.sessions[serverId]; + if (!session || equal(session.projects.get(project.projectId), project)) return prev; + const projects = new Map(session.projects); + projects.set(project.projectId, project); + return { + ...prev, + sessions: { ...prev.sessions, [serverId]: { ...session, projects } }, + }; + }); + }, + + removeProject: (serverId, projectId) => { + set((prev) => { + const session = prev.sessions[serverId]; + if (!session?.projects.has(projectId)) return prev; + const projects = new Map(session.projects); + projects.delete(projectId); + return { + ...prev, + sessions: { ...prev.sessions, [serverId]: { ...session, projects } }, + }; + }); + }, + + setWorkspaceRestoreStatus: (serverId, workspaceId, status) => { + set((prev) => { + const session = prev.sessions[serverId]; + if (!session) { + return prev; + } + if (session.restoringWorkspaces.get(workspaceId) === status) { + return prev; + } + // A late dir-gone timeout must not override a successful restore: + // only mark failed while still restoring and the descriptor is absent. + if ( + status === "failed" && + (session.restoringWorkspaces.get(workspaceId) !== "restoring" || + session.workspaces.has(workspaceId)) + ) { + return prev; + } + const next = new Map(session.restoringWorkspaces); + next.set(workspaceId, status); + return { + ...prev, + sessions: { + ...prev.sessions, + [serverId]: { ...session, restoringWorkspaces: next }, + }, + }; + }); + }, + + clearWorkspaceRestoreStatus: (serverId, workspaceId) => { + set((prev) => { + const session = prev.sessions[serverId]; + if (!session || !session.restoringWorkspaces.has(workspaceId)) { + return prev; + } + const next = new Map(session.restoringWorkspaces); + next.delete(workspaceId); + return { + ...prev, + sessions: { + ...prev.sessions, + [serverId]: { ...session, restoringWorkspaces: next }, + }, + }; + }); + }, + + mergeWorkspaces: (serverId, workspaces) => { + const nextEntries = Array.from(workspaces); + set((prev) => { + const session = prev.sessions[serverId]; + if (!session || nextEntries.length === 0) { + return prev; + } + const next = new Map(session.workspaces); + let changed = false; + // A descriptor arriving is the success signal for a pending restore: + // clear it at the source so every entry point converges to "ready". + let nextRestoring: Map | null = null; + for (const workspace of nextEntries) { + if (session.restoringWorkspaces.has(workspace.id)) { + nextRestoring ??= new Map(session.restoringWorkspaces); + nextRestoring.delete(workspace.id); + changed = true; + } + const existing = next.get(workspace.id); + const nextWorkspace = preserveWorkspaceDescriptorIdentity(workspace, existing); + if (existing === nextWorkspace) { + continue; + } + next.set(workspace.id, nextWorkspace); + changed = true; + } + if (!changed) { + return prev; + } + return { + ...prev, + sessions: { + ...prev.sessions, + [serverId]: { + ...session, + workspaces: next, + restoringWorkspaces: nextRestoring ?? session.restoringWorkspaces, + }, + }, + }; + }); + }, + + removeWorkspace: (serverId, workspaceId) => { + set((prev) => { + const session = prev.sessions[serverId]; + const workspaceKey = resolveWorkspaceMapKeyByIdentity({ + workspaces: session?.workspaces, + workspaceId, + }); + if (!session || !workspaceKey) { + return prev; + } + const next = new Map(session.workspaces); + next.delete(workspaceKey); + return { + ...prev, + sessions: { + ...prev.sessions, + [serverId]: { ...session, workspaces: next }, + }, + }; + }); + }, + + // Agent activity timestamps (top-level, does NOT mutate session object) + setAgentLastActivity: (agentId, timestamp) => { + agentLastActivityCoalescer.enqueue(agentId, timestamp); + }, + + setAgentLastActivityBatch: (updates) => { + set((prev) => { + const nextActivity = + typeof updates === "function" ? updates(prev.agentLastActivity) : updates; + if (nextActivity === prev.agentLastActivity) { + return prev; + } + if (nextActivity.size === 0) { + if (prev.agentLastActivity.size === 0) { + return prev; + } + return { + ...prev, + agentLastActivity: new Map(), + }; + } + let changed = false; + for (const [agentId, timestamp] of nextActivity.entries()) { + const currentTimestamp = prev.agentLastActivity.get(agentId); + if (!currentTimestamp || currentTimestamp.getTime() !== timestamp.getTime()) { + changed = true; + break; + } + } + if (!changed && nextActivity.size === prev.agentLastActivity.size) { + return prev; + } + return { + ...prev, + agentLastActivity: new Map(nextActivity), + }; + }); + }, + + flushAgentLastActivity: () => { + agentLastActivityCoalescer.flushNow(); + }, + + // Permissions + setPendingPermissions: (serverId, perms) => { + set((prev) => { + const session = prev.sessions[serverId]; + if (!session) { + return prev; + } + const nextPerms = + typeof perms === "function" ? perms(session.pendingPermissions) : perms; + if (session.pendingPermissions === nextPerms) { + return prev; + } + return { + ...prev, + sessions: { + ...prev.sessions, + [serverId]: { ...session, pendingPermissions: nextPerms }, + }, + }; + }); + }, + + // File explorer + setFileExplorer: (serverId, state) => { + set((prev) => { + const session = prev.sessions[serverId]; + if (!session) { + return prev; + } + const nextState = typeof state === "function" ? state(session.fileExplorer) : state; + if (session.fileExplorer === nextState) { + return prev; + } + return { + ...prev, + sessions: { + ...prev.sessions, + [serverId]: { ...session, fileExplorer: nextState }, + }, + }; + }); + }, + + // Queued messages + setQueuedMessages: (serverId, value) => { + set((prev) => { + const session = prev.sessions[serverId]; + if (!session) { + return prev; + } + const nextValue = typeof value === "function" ? value(session.queuedMessages) : value; + if (session.queuedMessages === nextValue) { + return prev; + } + return { + ...prev, + sessions: { + ...prev.sessions, + [serverId]: { ...session, queuedMessages: nextValue }, + }, + }; + }); + }, + + // Hydration + setHasHydratedAgents: (serverId, hydrated) => { + set((prev) => { + const session = prev.sessions[serverId]; + if (!session || session.hasHydratedAgents === hydrated) { + return prev; + } + return { + ...prev, + sessions: { + ...prev.sessions, + [serverId]: { ...session, hasHydratedAgents: hydrated }, + }, + }; + }); + }, + + setHasHydratedWorkspaces: (serverId, hydrated) => { + set((prev) => { + const session = prev.sessions[serverId]; + if (!session || session.hasHydratedWorkspaces === hydrated) { + return prev; + } + return { + ...prev, + sessions: { + ...prev.sessions, + [serverId]: { ...session, hasHydratedWorkspaces: hydrated }, + }, + }; + }); + }, + + setHasWorkspaceDirectorySnapshot: (serverId, available) => { + set((prev) => { + const session = prev.sessions[serverId]; + if (!session || session.hasWorkspaceDirectorySnapshot === available) return prev; + return { + ...prev, + sessions: { + ...prev.sessions, + [serverId]: { ...session, hasWorkspaceDirectorySnapshot: available }, + }, + }; + }); + }, + + // Agent directory - derived from agents (computed on-demand) + getAgentDirectory: (serverId) => { + const state = get(); + const session = state.sessions[serverId]; + if (!session) { + return undefined; + } + + const entries: AgentDirectoryEntry[] = []; + for (const agent of session.agents.values()) { + // Get lastActivityAt from top-level slice, fallback to agent.lastActivityAt + const lastActivityAt = state.agentLastActivity.get(agent.id) ?? agent.lastActivityAt; + entries.push({ + id: agent.id, + serverId, + title: agent.title ?? null, + status: agent.status, + turn: agent.turn, + lastActivityAt, + cwd: agent.cwd, + provider: agent.provider, + pendingPermissionCount: agent.pendingPermissions.length, + requiresAttention: agent.requiresAttention ?? false, + attentionReason: agent.attentionReason ?? null, + attentionTimestamp: agent.attentionTimestamp ?? null, + createdAt: agent.createdAt, + labels: agent.labels, + }); + } + return entries; + }, + }; + }), + ); + +export const createSessionStore = ( + initializations = createAgentInitializations(), + createFlows = createCreateFlowStore(), +) => Object.assign(createState(), { initializations, createFlows }); + +export type SessionStoreApi = ReturnType; + +const defaultSessionStore = createSessionStore(defaultAgentInitializations, defaultCreateFlowStore); +export const SessionStoreContext = createContext(null); +export function useSessionStoreApi(): SessionStoreApi { + return useContext(SessionStoreContext) ?? defaultSessionStore; +} +export const useSessionStore: SessionStoreApi = Object.assign(function useSessionStore( + selector: (state: SessionStore) => T = ((state: SessionStore) => state) as ( + state: SessionStore, + ) => T, +): T { + return useStore(useSessionStoreApi(), selector); +}, defaultSessionStore); export function useWorkspaceRestoreStatus( serverId: string | null, diff --git a/packages/app/src/timeline/replica.test.ts b/packages/app/src/timeline/replica.test.ts index 47bd36596a..2fbbc7af96 100644 --- a/packages/app/src/timeline/replica.test.ts +++ b/packages/app/src/timeline/replica.test.ts @@ -8,7 +8,12 @@ import { import { afterEach, describe, expect, it } from "vitest"; import type { AgentStreamEventPayload } from "@getpaseo/protocol/messages"; import type { CachedTimeline } from "@/runtime/replica-cache"; -import { selectAgentTimelineState, useSessionStore } from "@/stores/session-store"; +import { + createSessionStore, + type SessionStoreApi, + selectAgentTimelineState, + useSessionStore, +} from "@/stores/session-store"; import { useCreateFlowStore, type PendingCreateAttempt } from "@/stores/create-flow-store"; import type { StreamItem } from "@/types/stream"; import { @@ -40,13 +45,18 @@ function cachedTimeline(): CachedTimeline { }; } -function createOwner(storage: TimelineReplicaStorage): ViewedTimelineOwner { +function createOwner( + storage: TimelineReplicaStorage, + store: SessionStoreApi = useSessionStore, +): ViewedTimelineOwner { const replica = createTimelineReplica({ + store, serverId: SERVER_ID, storage, prepareAgent: async () => undefined, }); return createViewedTimelineOwner({ + store, serverId: SERVER_ID, replica, replaceDemandedAgentIds: () => undefined, @@ -533,8 +543,8 @@ describe("create handoff lifetime", () => { afterEach(() => useCreateFlowStore.getState().clearAll()); - function sentCreateHandoff(): void { - useCreateFlowStore.getState().setPending({ + function sentCreateHandoff(flow = useCreateFlowStore): void { + flow.getState().setPending({ draftId: DRAFT_ID, serverId: SERVER_ID, agentId: AGENT_ID, @@ -542,7 +552,7 @@ describe("create handoff lifetime", () => { text: "Create this chat.", timestamp: 0, }); - useCreateFlowStore.getState().markLifecycle({ draftId: DRAFT_ID, lifecycle: "sent" }); + flow.getState().markLifecycle({ draftId: DRAFT_ID, lifecycle: "sent" }); } function createHandoff(): PendingCreateAttempt | undefined { @@ -575,18 +585,22 @@ describe("create handoff lifetime", () => { release: () => void; } - async function ownerWithHeldCatchUp(): Promise { + async function ownerWithHeldCatchUp( + store: SessionStoreApi = useSessionStore, + ): Promise { let releaseTailFetch!: () => void; let tailFetchStarted = false; const tailFetched = new Promise((resolve) => { releaseTailFetch = resolve; }); const replica = createTimelineReplica({ + store, serverId: SERVER_ID, storage: { readTimeline: async () => undefined, commitTimeline: () => undefined }, prepareAgent: async () => undefined, }); const owner = createViewedTimelineOwner({ + store, serverId: SERVER_ID, replica, replaceDemandedAgentIds: () => undefined, @@ -611,6 +625,30 @@ describe("create handoff lifetime", () => { return { owner, release: releaseTailFetch }; } + it("clears only the completed authority's creation handoff for equal destination identities", async () => { + const direct = createSessionStore(); + const routed = createSessionStore(); + for (const store of [direct, routed]) { + store.getState().initializeSession(SERVER_ID, null); + sentCreateHandoff(store.createFlows); + } + const { owner, release } = await ownerWithHeldCatchUp(direct); + try { + owner.applyTimelineResponse(syncedTailPage()); + release(); + await expect + .poll(() => direct.createFlows.getState().pendingByDraftId[DRAFT_ID]) + .toBeUndefined(); + expect(routed.createFlows.getState().pendingByDraftId[DRAFT_ID]).toMatchObject({ + agentId: AGENT_ID, + lifecycle: "sent", + }); + } finally { + release(); + owner.dispose(); + } + }); + it("keeps the handoff through the first synchronized page and releases it when the chat is current", async () => { useSessionStore.getState().initializeSession(SERVER_ID, null); sentCreateHandoff(); @@ -749,3 +787,66 @@ describe("SQLite baseline restoration", () => { } }); }); + +it.each([null, "Direct history failed"])( + "keeps initialization readiness scoped when another authority returns %s for the same address", + async (failureMessage) => { + const direct = createSessionStore(); + const routed = createSessionStore(); + const key = direct.initializations.getInitKey(SERVER_ID, AGENT_ID); + const directReady = direct.initializations.createInitDeferred(key, "tail").promise; + const routedReady = routed.initializations.createInitDeferred(key, "tail").promise; + void directReady.catch(() => {}); + let routedSettled = false; + void routedReady.then( + () => { + routedSettled = true; + return undefined; + }, + () => { + routedSettled = true; + return undefined; + }, + ); + for (const store of [direct, routed]) { + store.getState().initializeSession(SERVER_ID, null); + store.getState().setInitializingAgents(SERVER_ID, new Map([[AGENT_ID, true]])); + } + const storage: TimelineReplicaStorage = { + readTimeline: async () => undefined, + commitTimeline: () => {}, + }; + const owners = [createOwner(storage, direct), createOwner(storage, routed)]; + const response: Parameters[0] = { + requestId: "direct-tail", + agentId: AGENT_ID, + agent: null, + direction: "tail", + projection: "projected", + reset: false, + epoch: "epoch-1", + window: { minSeq: 1, maxSeq: 0, nextSeq: 1 }, + startCursor: null, + endCursor: null, + entries: [], + error: failureMessage, + hasNewer: false, + hasOlder: false, + staleCursor: false, + gap: false, + }; + try { + owners[0]!.applyTimelineResponse(response); + const directResult = await Promise.allSettled([directReady]); + expect(directResult[0]?.status).toBe(failureMessage ? "rejected" : "fulfilled"); + expect(routedSettled).toBe(false); + expect(routed.initializations.getInitDeferred(key)?.promise).toBe(routedReady); + owners[1]!.applyTimelineResponse({ ...response, requestId: "routed-tail", error: null }); + await routedReady; + expect(routedSettled).toBe(true); + } finally { + for (const store of [direct, routed]) store.initializations.resolveInitDeferred(key); + for (const owner of owners) owner.dispose(); + } + }, +); diff --git a/packages/app/src/timeline/session-stream-reducers.ts b/packages/app/src/timeline/session-stream-reducers.ts index 34a461d893..d7eb82b928 100644 --- a/packages/app/src/timeline/session-stream-reducers.ts +++ b/packages/app/src/timeline/session-stream-reducers.ts @@ -1,6 +1,10 @@ import type { AgentTimelineItem } from "@getpaseo/protocol/agent-types"; import type { AgentStreamEventPayload } from "@getpaseo/protocol/messages"; -import { selectAgentTimelineState, useSessionStore } from "@/stores/session-store"; +import { + selectAgentTimelineState, + useSessionStore, + type SessionStoreApi, +} from "@/stores/session-store"; import type { AssistantMessageItem, StreamItem, TodoEntry } from "@/types/stream"; import type { TurnLivenessTransition } from "@/timeline/turn-liveness"; import { @@ -1829,6 +1833,7 @@ export function deriveAgentStreamTurnLiveness( } export interface CreateSessionAgentStreamReducerQueueInput { + store?: SessionStoreApi; serverId: string; setAgentStreamState: (serverId: string, agentId: string, state: StreamStatePatch) => void; setAgentTimelineCursor: ( @@ -1895,7 +1900,7 @@ export function createSessionAgentStreamReducerQueue( return createAgentStreamReducerQueue({ getSnapshot: (agentId) => { - const session = useSessionStore.getState().sessions[serverId]; + const session = (input.store ?? useSessionStore).getState().sessions[serverId]; const timeline = selectAgentTimelineState(session, agentId); return { currentTail: timeline.status === "cold" ? [] : timeline.items, diff --git a/packages/app/src/timeline/viewed-timeline-sync.ts b/packages/app/src/timeline/viewed-timeline-sync.ts index 5939feabce..37475c8af2 100644 --- a/packages/app/src/timeline/viewed-timeline-sync.ts +++ b/packages/app/src/timeline/viewed-timeline-sync.ts @@ -7,18 +7,13 @@ import type { CachedTimeline } from "@/runtime/replica-cache"; import { selectAgentTimelineState, useSessionStore, + type SessionStoreApi, type AgentTimelineCursorState, type AgentTimelineState, } from "@/stores/session-store"; -import { useCreateFlowStore } from "@/stores/create-flow-store"; import type { SessionOutboundMessage } from "@getpaseo/protocol/messages"; import { getSendingClientMessageIds } from "@/composer/submission/model"; -import { - getInitDeferred, - getInitKey, - rejectInitDeferred, - resolveInitDeferred, -} from "@/utils/agent-initialization"; +import { getInitKey } from "@/utils/agent-initialization"; import { createSessionAgentStreamReducerQueue, type AgentStreamReducerEvent, @@ -34,12 +29,13 @@ export interface TimelineReplicaStorage { } async function prepareCachedTimeline(input: { + store: SessionStoreApi; serverId: string; agentId: string; storage: TimelineReplicaStorage; prepareAgent: (agentId: string) => Promise; }): Promise { - const before = useSessionStore.getState().sessions[input.serverId]; + const before = input.store.getState().sessions[input.serverId]; const beforeTimeline = selectAgentTimelineState(before, input.agentId); if (beforeTimeline.status === "synced") return undefined; const [stored] = await Promise.all([ @@ -47,7 +43,7 @@ async function prepareCachedTimeline(input: { input.prepareAgent(input.agentId), ]); if (!stored) return undefined; - const session = useSessionStore.getState().sessions[input.serverId]; + const session = input.store.getState().sessions[input.serverId]; const currentTimeline = selectAgentTimelineState(session, input.agentId); const currentHead = session?.agentStreamHead.get(input.agentId); if (currentTimeline.status === "synced") return undefined; @@ -77,7 +73,7 @@ async function prepareCachedTimeline(input: { endSeq: null, }, }); - useSessionStore.getState().applyAgentTimelineResponseState(input.serverId, input.agentId, { + input.store.getState().applyAgentTimelineResponseState(input.serverId, input.agentId, { items: replacement.tail, head: replacement.head, range: stored.range, @@ -104,6 +100,7 @@ class TimelineReplicaOwner implements TimelineReplica { private readonly serverId: string, private readonly storage: TimelineReplicaStorage, private readonly prepareAgent: (agentId: string) => Promise, + private readonly store: SessionStoreApi, ) {} async prepare(agentId: string): Promise { @@ -123,6 +120,7 @@ class TimelineReplicaOwner implements TimelineReplica { serverId: this.serverId, agentId, storage: this.storage, + store: this.store, prepareAgent: this.prepareAgent, }); if (!stored) return; @@ -141,7 +139,7 @@ class TimelineReplicaOwner implements TimelineReplica { } timelineUpdated(agentId: string): void { - const session = useSessionStore.getState().sessions[this.serverId]; + const session = this.store.getState().sessions[this.serverId]; const timeline = selectAgentTimelineState(session, agentId); if (timeline.status === "cold") return; if (timeline.status === "synced") this.cachedRanges.delete(agentId); @@ -155,11 +153,17 @@ class TimelineReplicaOwner implements TimelineReplica { } export function createTimelineReplica(input: { + store?: SessionStoreApi; serverId: string; storage: TimelineReplicaStorage; prepareAgent: (agentId: string) => Promise; }): TimelineReplica { - return new TimelineReplicaOwner(input.serverId, input.storage, input.prepareAgent); + return new TimelineReplicaOwner( + input.serverId, + input.storage, + input.prepareAgent, + input.store ?? useSessionStore, + ); } export interface TimelinePageResult { @@ -181,8 +185,12 @@ export function consumeForcedTimelineTailReplacement( return { ...payload, reset: true }; } -function clearAgentInitializingFlag(serverId: string, agentId: string): void { - useSessionStore.getState().setInitializingAgents(serverId, (previous) => { +function clearAgentInitializingFlag( + store: SessionStoreApi, + serverId: string, + agentId: string, +): void { + store.getState().setInitializingAgents(serverId, (previous) => { if (previous.get(agentId) !== true) return previous; const next = new Map(previous); next.set(agentId, false); @@ -191,6 +199,7 @@ function clearAgentInitializingFlag(serverId: string, agentId: string): void { } function commitProcessedTimeline(input: { + store: SessionStoreApi; serverId: string; payload: TimelineResponsePayload; result: ProcessTimelineResponseOutput; @@ -200,7 +209,7 @@ function commitProcessedTimeline(input: { }): void { const { serverId, payload, result, timeline, currentCursor, synchronized } = input; const agentId = payload.agentId; - const store = useSessionStore.getState(); + const store = input.store.getState(); if (result.commit !== "discard") { store.applyAgentTimelineResponseState(serverId, agentId, { items: result.tail, @@ -232,6 +241,7 @@ function commitProcessedTimeline(input: { } function finalizeProcessedTimeline(input: { + store: SessionStoreApi; serverId: string; agentId: string; initKey: string; @@ -244,17 +254,19 @@ function finalizeProcessedTimeline(input: { if (effect.type === "catch_up") input.recoverGap(input.agentId, effect.cursor); } if (input.result.clearInitializing) { - clearAgentInitializingFlag(input.serverId, input.agentId); + clearAgentInitializingFlag(input.store, input.serverId, input.agentId); } if (input.synchronized) { - const session = useSessionStore.getState().sessions[input.serverId]; + const session = input.store.getState().sessions[input.serverId]; const agent = session?.agents.get(input.agentId) ?? session?.agentDetails.get(input.agentId); if (agent && agent.turn.phase === "idle") input.drainQueuedAgentMessage(input.agentId); } - if (input.result.initResolution === "resolve") resolveInitDeferred(input.initKey); + if (input.result.initResolution === "resolve") + input.store.initializations.resolveInitDeferred(input.initKey); } function applyAuthoritativeTimelineResponse(input: { + store: SessionStoreApi; serverId: string; payload: TimelineResponsePayload; cachedCursor?: AgentTimelineCursorState; @@ -264,9 +276,9 @@ function applyAuthoritativeTimelineResponse(input: { const { serverId, payload } = input; const agentId = payload.agentId; const initKey = getInitKey(serverId, agentId); - const session = useSessionStore.getState().sessions[serverId]; + const session = input.store.getState().sessions[serverId]; const timeline = selectAgentTimelineState(session, agentId); - const activeInitDeferred = getInitDeferred(initKey); + const activeInitDeferred = input.store.initializations.getInitDeferred(initKey); const currentCursor = timeline.status === "synced" ? (timeline.range ?? undefined) : input.cachedCursor; const result = processTimelineResponse({ @@ -281,8 +293,9 @@ function applyAuthoritativeTimelineResponse(input: { }); if (result.error) { - if (result.clearInitializing) clearAgentInitializingFlag(serverId, agentId); - if (result.initResolution === "reject") rejectInitDeferred(initKey, new Error(result.error)); + if (result.clearInitializing) clearAgentInitializingFlag(input.store, serverId, agentId); + if (result.initResolution === "reject") + input.store.initializations.rejectInitDeferred(initKey, new Error(result.error)); return false; } @@ -291,8 +304,17 @@ function applyAuthoritativeTimelineResponse(input: { hasNewer: payload.hasNewer, error: payload.error, }); - commitProcessedTimeline({ serverId, payload, result, timeline, currentCursor, synchronized }); + commitProcessedTimeline({ + store: input.store, + serverId, + payload, + result, + timeline, + currentCursor, + synchronized, + }); finalizeProcessedTimeline({ + store: input.store, serverId, agentId, initKey, @@ -358,12 +380,14 @@ export interface ViewedTimelineOwner extends ViewedTimelineSync { } export function createViewedTimelineOwner(input: { + store?: SessionStoreApi; serverId: string; replica: TimelineReplica; replaceDemandedAgentIds: (agentIds: string[]) => void; drainQueuedAgentMessage: (agentId: string) => void; ports: ViewedTimelineOwnerPorts; }): ViewedTimelineOwner { + const store = input.store ?? useSessionStore; const forcedTailReplacements = new Set(); const sync = createViewedTimelineSync({ ...input.ports, @@ -379,12 +403,13 @@ export function createViewedTimelineOwner(input: { readCursor: (agentId) => input.replica.readCursor(agentId) ?? input.ports.readCursor(agentId), replaceDemandedAgentIds: input.replaceDemandedAgentIds, onCatchUpEnded: (agentId) => - useCreateFlowStore.getState().clearByAgent({ serverId: input.serverId, agentId }), + store.createFlows.getState().clearByAgent({ serverId: input.serverId, agentId }), }); const streamQueue = createSessionAgentStreamReducerQueue({ serverId: input.serverId, - setAgentStreamState: (...args) => useSessionStore.getState().setAgentStreamState(...args), - setAgentTimelineCursor: (...args) => useSessionStore.getState().setAgentTimelineCursor(...args), + store, + setAgentStreamState: (...args) => store.getState().setAgentStreamState(...args), + setAgentTimelineCursor: (...args) => store.getState().setAgentTimelineCursor(...args), recoverTimelineGap: (agentId, cursor) => sync.recoverGap(agentId, cursor), onCommitted: (agentId) => input.replica.timelineUpdated(agentId), }); @@ -393,6 +418,7 @@ export function createViewedTimelineOwner(input: { applyTimelineResponse(receivedPayload) { const payload = consumeForcedTimelineTailReplacement(receivedPayload, forcedTailReplacements); const accepted = applyAuthoritativeTimelineResponse({ + store, serverId: input.serverId, payload, cachedCursor: input.replica.readRange(payload.agentId), diff --git a/packages/app/src/utils/agent-initialization.ts b/packages/app/src/utils/agent-initialization.ts index 0d0da8e4df..ddf5a9a3e2 100644 --- a/packages/app/src/utils/agent-initialization.ts +++ b/packages/app/src/utils/agent-initialization.ts @@ -6,74 +6,96 @@ export interface DeferredInit { requestDirection: "tail" | "after"; } -const initPromises = new Map(); export const INIT_TIMEOUT_MS = 65_000; -export function getInitKey(serverId: string, agentId: string): string { - return `${serverId}:${agentId}`; -} +export function createAgentInitializations() { + const initPromises = new Map(); -export function getInitDeferred(key: string): DeferredInit | undefined { - return initPromises.get(key); -} + function getInitKey(serverId: string, agentId: string): string { + return `${serverId}:${agentId}`; + } -export function createInitDeferred(key: string, requestDirection: "tail" | "after"): DeferredInit { - let resolve!: () => void; - let reject!: (error: Error) => void; + function getInitDeferred(key: string): DeferredInit | undefined { + return initPromises.get(key); + } - const promise = new Promise((res, rej) => { - resolve = res; - reject = rej; - }); + function createInitDeferred(key: string, requestDirection: "tail" | "after"): DeferredInit { + let resolve!: () => void; + let reject!: (error: Error) => void; - const deferred: DeferredInit = { - promise, - resolve, - reject, - timeoutId: null, - requestDirection, + const promise = new Promise((res, rej) => { + resolve = res; + reject = rej; + }); + + const deferred: DeferredInit = { + promise, + resolve, + reject, + timeoutId: null, + requestDirection, + }; + initPromises.set(key, deferred); + return deferred; + } + + function refreshInitTimeout(input: { + key: string; + onTimeout: () => void; + timeoutMs?: number; + }): void { + const timeoutId = setTimeout(input.onTimeout, input.timeoutMs ?? INIT_TIMEOUT_MS); + const deferred = initPromises.get(input.key); + if (!deferred) { + clearTimeout(timeoutId); + return; + } + if (deferred.timeoutId) { + clearTimeout(deferred.timeoutId); + } + deferred.timeoutId = timeoutId; + } + + function resolveInitDeferred(key: string): void { + const deferred = initPromises.get(key); + if (!deferred) { + return; + } + if (deferred.timeoutId) { + clearTimeout(deferred.timeoutId); + } + initPromises.delete(key); + deferred.resolve(); + } + + function rejectInitDeferred(key: string, error: Error): void { + const deferred = initPromises.get(key); + if (!deferred) { + return; + } + if (deferred.timeoutId) { + clearTimeout(deferred.timeoutId); + } + initPromises.delete(key); + deferred.reject(error); + } + + return { + getInitKey, + getInitDeferred, + createInitDeferred, + refreshInitTimeout, + resolveInitDeferred, + rejectInitDeferred, }; - initPromises.set(key, deferred); - return deferred; } -export function refreshInitTimeout(input: { - key: string; - onTimeout: () => void; - timeoutMs?: number; -}): void { - const timeoutId = setTimeout(input.onTimeout, input.timeoutMs ?? INIT_TIMEOUT_MS); - const deferred = initPromises.get(input.key); - if (!deferred) { - clearTimeout(timeoutId); - return; - } - if (deferred.timeoutId) { - clearTimeout(deferred.timeoutId); - } - deferred.timeoutId = timeoutId; -} - -export function resolveInitDeferred(key: string): void { - const deferred = initPromises.get(key); - if (!deferred) { - return; - } - if (deferred.timeoutId) { - clearTimeout(deferred.timeoutId); - } - initPromises.delete(key); - deferred.resolve(); -} - -export function rejectInitDeferred(key: string, error: Error): void { - const deferred = initPromises.get(key); - if (!deferred) { - return; - } - if (deferred.timeoutId) { - clearTimeout(deferred.timeoutId); - } - initPromises.delete(key); - deferred.reject(error); -} +export const defaultAgentInitializations = createAgentInitializations(); +export const { + getInitKey, + getInitDeferred, + createInitDeferred, + refreshInitTimeout, + resolveInitDeferred, + rejectInitDeferred, +} = defaultAgentInitializations; diff --git a/packages/client/package.json b/packages/client/package.json index 57c937656d..f0300d740d 100644 --- a/packages/client/package.json +++ b/packages/client/package.json @@ -28,6 +28,10 @@ "./node": { "types": "./dist/node/index.d.ts", "default": "./dist/node/index.js" + }, + "./internal/host-connection": { + "types": "./dist/hosts/connection.d.ts", + "default": "./dist/hosts/connection.js" } }, "publishConfig": { diff --git a/packages/client/src/hosts/connection.ts b/packages/client/src/hosts/connection.ts new file mode 100644 index 0000000000..34e7d90ed5 --- /dev/null +++ b/packages/client/src/hosts/connection.ts @@ -0,0 +1,171 @@ +import { safeRandomId } from "../daemon-client-transport-utils.js"; +import { DaemonClient } from "../daemon-client.js"; +import type { DaemonTransportFactory } from "../daemon-client-transport-types.js"; + +export function createHostChannelTransport( + origin: DaemonClient, + serverId: string, + scopeId: string, + initialRouteId: string | null, + link = false, +): DaemonTransportFactory { + let first: string | null = initialRouteId; + return () => { + const opens = new Set<() => void>(); + const closes = new Set<(event?: unknown) => void>(); + const messages = new Set<(data: unknown, binary: boolean) => void>(); + const errors = new Set<(error: unknown) => void>(); + let routeId: string | null = null; + let closed = false; + const cleanups: (() => void)[] = []; + const close = (error?: unknown) => { + if (closed) return; + closed = true; + if (error instanceof Error) Object.assign(error, { serverId }); + for (const cleanup of cleanups) cleanup(); + try { + if (routeId) origin.closeHostRoute(routeId); + else origin.cancelHostRoute(scopeId); + } catch { + /* Transport already closed. */ + } + if (error) for (const handler of errors) handler(error); + for (const handler of closes) handler(error); + }; + cleanups.push( + origin.subscribeRawMessages((message) => { + if (message.type === "hosts.route.frame" && message.payload.subscriptionId === routeId) { + for (const handler of messages) handler(message.payload.data, false); + } + if (message.type === "hosts.route.closed" && message.payload.subscriptionId === routeId) + close( + Object.assign(new Error(message.payload.message), { + code: message.payload.code, + serverId, + }), + ); + }), + ); + cleanups.push( + origin.subscribeHostFrames((frame) => { + if (frame.routeId === routeId) for (const handler of messages) handler(frame.payload, true); + }), + ); + cleanups.push( + origin.subscribeConnectionStatus((connectionState) => { + if (connectionState.status !== "connected") close(); + }), + ); + let opening: Promise; + if (first) opening = Promise.resolve(first); + else opening = link ? origin.openHostLink(scopeId) : origin.openHostRoute(serverId, scopeId); + first = null; + void opening + .then((id) => { + if (closed) { + origin.closeHostRoute(id); + return; + } + routeId = id; + for (const handler of opens) handler(); + return undefined; + }) + .catch(close); + return { + send(data) { + if (closed || !routeId) + throw Object.assign(new Error(`Host ${serverId} is offline`), { + code: "host_offline", + serverId, + }); + origin.sendHostFrame( + routeId, + typeof data === "string" || data instanceof Uint8Array ? data : new Uint8Array(data), + ); + }, + close: () => close(), + onOpen(handler) { + opens.add(handler); + return () => opens.delete(handler); + }, + onClose(handler) { + closes.add(handler); + return () => closes.delete(handler); + }, + onError(handler) { + errors.add(handler); + return () => errors.delete(handler); + }, + onMessage(handler) { + messages.add(handler); + return () => messages.delete(handler); + }, + }; + }; +} + +/** Internal adapter for clients that own a low-level replica pipeline. */ +export interface HostConnectionLease { + readonly client: DaemonClient; + release(): Promise; +} + +/** The origin alone resolves the destination; leases never probe another authority. */ +export async function acquireHostConnection( + origin: DaemonClient, + serverId: string, + signal: AbortSignal, +): Promise { + const assertActive = () => { + if (signal.aborted) throw new Error("Host connection acquisition is disposed"); + }; + assertActive(); + origin.requireConnectedHosts(); + if ( + origin.getServerId() === serverId && + origin.getLastServerInfoMessage()?.features?.hostSessionScopes !== true + ) + throw Object.assign(new Error("Update the managing daemon to acquire owned session scopes"), { + code: "unsupported_capability", + serverId, + }); + + const scopeId = safeRandomId(); + let client: DaemonClient | null = null; + let release: Promise | null = null; + const cancel = () => { + origin.cancelHostRoute(scopeId); + if (client) void client.close().catch(() => {}); + }; + signal.addEventListener("abort", cancel, { once: true }); + const close = () => { + if (!release) { + signal.removeEventListener("abort", cancel); + release = client ? client.close() : Promise.resolve(); + } + return release; + }; + try { + const routeId = await origin.openHostRoute(serverId, scopeId); + if (signal.aborted) { + origin.closeHostRoute(routeId); + assertActive(); + } + client = new DaemonClient({ + url: "ws://routed-host/ws", + clientId: scopeId, + transportFactory: createHostChannelTransport(origin, serverId, scopeId, routeId), + reconnect: { enabled: false }, + }); + await client.connect(); + assertActive(); + client.setReconnectEnabled(true); + return { client, release: close }; + } catch (error) { + await close().catch(() => {}); + throw Object.assign(error instanceof Error ? error : new Error("Host handshake failed"), { + code: error instanceof Error && "code" in error ? error.code : "host_offline", + serverId, + }); + } +} diff --git a/packages/client/src/hosts/index.ts b/packages/client/src/hosts/index.ts index fc39403e1f..26b657e417 100644 --- a/packages/client/src/hosts/index.ts +++ b/packages/client/src/hosts/index.ts @@ -1,10 +1,10 @@ -import { safeRandomId } from "../daemon-client-transport-utils.js"; import { DaemonClient } from "../daemon-client.js"; -import type { DaemonTransportFactory } from "../daemon-client-transport-types.js"; import type { OwnedSubscription } from "../connection/index.js"; import type { PaseoApi, createPaseoApi } from "../index.js"; import type { HostAdd, HostUpdate, HostSummary, HostDirectory } from "@getpaseo/protocol/hosts"; import { HostRuntime } from "./runtime.js"; +import { acquireHostConnection } from "./connection.js"; +export { createHostChannelTransport } from "./connection.js"; export interface PaseoHostActions { list(): Promise; @@ -16,108 +16,6 @@ export interface PaseoHostActions { runtime(): HostRuntime; } -export function createHostChannelTransport( - origin: DaemonClient, - serverId: string, - scopeId: string, - initialRouteId: string | null, - link = false, -): DaemonTransportFactory { - let first: string | null = initialRouteId; - return () => { - const opens = new Set<() => void>(); - const closes = new Set<(event?: unknown) => void>(); - const messages = new Set<(data: unknown, binary: boolean) => void>(); - const errors = new Set<(error: unknown) => void>(); - let routeId: string | null = null; - let closed = false; - const cleanups: (() => void)[] = []; - const close = (error?: unknown) => { - if (closed) return; - closed = true; - if (error instanceof Error) Object.assign(error, { serverId }); - for (const cleanup of cleanups) cleanup(); - try { - if (routeId) origin.closeHostRoute(routeId); - else origin.cancelHostRoute(scopeId); - } catch { - /* Transport already closed. */ - } - if (error) for (const handler of errors) handler(error); - for (const handler of closes) handler(error); - }; - cleanups.push( - origin.subscribeRawMessages((message) => { - if (message.type === "hosts.route.frame" && message.payload.subscriptionId === routeId) { - for (const handler of messages) handler(message.payload.data, false); - } - if (message.type === "hosts.route.closed" && message.payload.subscriptionId === routeId) - close( - Object.assign(new Error(message.payload.message), { - code: message.payload.code, - serverId, - }), - ); - }), - ); - cleanups.push( - origin.subscribeHostFrames((frame) => { - if (frame.routeId === routeId) for (const handler of messages) handler(frame.payload, true); - }), - ); - cleanups.push( - origin.subscribeConnectionStatus((connectionState) => { - if (connectionState.status !== "connected") close(); - }), - ); - let opening: Promise; - if (first) opening = Promise.resolve(first); - else opening = link ? origin.openHostLink(scopeId) : origin.openHostRoute(serverId, scopeId); - first = null; - void opening - .then((id) => { - if (closed) { - origin.closeHostRoute(id); - return; - } - routeId = id; - for (const handler of opens) handler(); - return undefined; - }) - .catch(close); - return { - send(data) { - if (closed || !routeId) - throw Object.assign(new Error(`Host ${serverId} is offline`), { - code: "host_offline", - serverId, - }); - origin.sendHostFrame( - routeId, - typeof data === "string" || data instanceof Uint8Array ? data : new Uint8Array(data), - ); - }, - close: () => close(), - onOpen(handler) { - opens.add(handler); - return () => opens.delete(handler); - }, - onClose(handler) { - closes.add(handler); - return () => closes.delete(handler); - }, - onError(handler) { - errors.add(handler); - return () => errors.delete(handler); - }, - onMessage(handler) { - messages.add(handler); - return () => messages.delete(handler); - }, - }; - }; -} - export function createHostActions( origin: DaemonClient, owner: { @@ -150,60 +48,23 @@ export function createHostActions( async client(serverId) { assertActive(); origin.requireConnectedHosts(); - if (origin.getLastServerInfoMessage()?.serverId === serverId) { - const api = owner.api(origin, { signal: owner.signal }); - const scope = owner.own(() => ({ release: () => api.dispose() })); - return { ...api, dispose: () => scope.release() }; - } - const scopeId = safeRandomId(); - const cancel = () => origin.cancelHostRoute(scopeId); - owner.signal.addEventListener("abort", cancel, { once: true }); - let routeId: string; + const lease = await acquireHostConnection(origin, serverId, owner.signal); try { - routeId = await origin.openHostRoute(serverId, scopeId); - } catch (error) { - throw Object.assign(error instanceof Error ? error : new Error(String(error)), { - serverId, - }); - } finally { - owner.signal.removeEventListener("abort", cancel); - } - if (owner.signal.aborted) { - origin.closeHostRoute(routeId); assertActive(); - } - const client = new DaemonClient({ - url: "ws://routed-host/ws", - clientId: scopeId, - transportFactory: createHostChannelTransport(origin, serverId, scopeId, routeId), - // Initial acquisition has a deadline; established scopes restore demand on loss. - reconnect: { enabled: false }, - }); - const api = owner.api(client, { signal: owner.signal }); - const scope = owner.own(() => ({ - release: async () => { - try { - await api.dispose(); - } finally { - await client.close(); - } - }, - })); - try { - await client.connect(); - client.setReconnectEnabled(true); - assertActive(); - return { ...api, dispose: () => scope.release() }; - } catch (error) { - const failure = Object.assign( - error instanceof Error ? error : new Error("Host handshake failed"), - { - code: error instanceof Error && "code" in error ? error.code : "host_offline", - serverId, + const api = owner.api(lease.client, { signal: owner.signal }); + const scope = owner.own(() => ({ + release: async () => { + try { + await api.dispose(); + } finally { + await lease.release(); + } }, - ); - await scope.release(); - throw failure; + })); + return { ...api, dispose: () => scope.release() }; + } catch (error) { + await lease.release().catch(() => {}); + throw error; } }, runtime() { diff --git a/packages/client/src/schedules/index.ts b/packages/client/src/schedules/index.ts index 49df7d42b4..c6910a0bcf 100644 --- a/packages/client/src/schedules/index.ts +++ b/packages/client/src/schedules/index.ts @@ -23,14 +23,14 @@ export function createScheduleActions( return method(...args); }; return { - create: action(client.scheduleCreate.bind(client)), - list: action(client.scheduleList.bind(client)), - inspect: action(client.scheduleInspect.bind(client)), - update: action(client.scheduleUpdate.bind(client)), - pause: action(client.schedulePause.bind(client)), - resume: action(client.scheduleResume.bind(client)), - delete: action(client.scheduleDelete.bind(client)), - logs: action(client.scheduleLogs.bind(client)), - runOnce: action(client.scheduleRunOnce.bind(client)), + create: action((...args) => client.scheduleCreate(...args)), + list: action((...args) => client.scheduleList(...args)), + inspect: action((...args) => client.scheduleInspect(...args)), + update: action((...args) => client.scheduleUpdate(...args)), + pause: action((...args) => client.schedulePause(...args)), + resume: action((...args) => client.scheduleResume(...args)), + delete: action((...args) => client.scheduleDelete(...args)), + logs: action((...args) => client.scheduleLogs(...args)), + runOnce: action((...args) => client.scheduleRunOnce(...args)), }; } diff --git a/packages/protocol/src/messages.ts b/packages/protocol/src/messages.ts index 05160d97cd..439192d2df 100644 --- a/packages/protocol/src/messages.ts +++ b/packages/protocol/src/messages.ts @@ -3767,6 +3767,7 @@ export const ServerInfoStatusPayloadSchema = z connectedHosts: z.boolean().optional(), fleetSubagents: z.boolean().optional(), hostMultiplexing: z.boolean().optional(), + hostSessionScopes: z.boolean().optional(), usageSources: z.boolean().optional(), // COMPAT(providersSnapshotCwd): added in v0.3.2, remove gate after 2027-02-10. providersSnapshotCwd: z.boolean().optional(), diff --git a/packages/server/src/server/agent/agent-manager.ts b/packages/server/src/server/agent/agent-manager.ts index 9b98e028f0..c7ca3d22db 100644 --- a/packages/server/src/server/agent/agent-manager.ts +++ b/packages/server/src/server/agent/agent-manager.ts @@ -768,6 +768,7 @@ export class AgentManager { private appendSystemPrompt: string; private onAgentAttention?: AgentAttentionCallback; private readonly agentArchivedListeners = new Set(); + private managedChildLifecycleOwner: AgentArchivedCallback | null = null; private onWorkspaceStateMayHaveChanged?: (params: { cwd: string }) => void; private logger: Logger; private readonly rescueTimeouts: Required; @@ -874,10 +875,19 @@ export class AgentManager { return this.runJournal.get(agentId, submissionId); } - onAgentArchived(callback: AgentArchivedCallback): () => void { + onAgentArchived( + callback: AgentArchivedCallback, + options?: { ownsManagedChildren: boolean }, + ): () => void { + if (options?.ownsManagedChildren) { + if (this.managedChildLifecycleOwner) + throw new Error("Managed child lifecycle already has an owner"); + this.managedChildLifecycleOwner = callback; + } this.agentArchivedListeners.add(callback); return () => { this.agentArchivedListeners.delete(callback); + if (this.managedChildLifecycleOwner === callback) this.managedChildLifecycleOwner = null; }; } @@ -1904,6 +1914,7 @@ export class AgentManager { const parentAgentId = parent.id; const records = await registry.list(); for (const record of records) { + if (this.managedChildLifecycleOwner && !record.internal && !record.owner) continue; if (record.archivedAt) { continue; } @@ -3683,6 +3694,7 @@ export class AgentManager { } } + await this.runJournal.restore(resolvedAgentId); this.assertAcceptingAgentRegistrations(); this.agents.set(resolvedAgentId, managed); registered = true; @@ -4395,22 +4407,14 @@ export class AgentManager { const flags: StreamEventFlags = { shouldDispatchEvent: true, shouldNotifyWaiters: true }; - const runResult = this.getRunResult( - agent.id, - eventTurnId, + const checkpoint = this.captureLiveRunCheckpoint( + agent, event, - options?.fromHistory === true, + eventTurnId, + options, terminalDisposition, ); - if (!options?.fromHistory) { - if (terminalDisposition !== "stale") - await this.runJournal.capture(agent.id, eventTurnId, event, runResult).catch((error) => { - this.logger.error( - { err: error, agentId: agent.id, turnId: eventTurnId }, - "Failed to persist provider checkpoint", - ); - }); - } + if (checkpoint) await checkpoint; const dispatchPromise = this.dispatchStreamEventByType({ agent, event, @@ -4442,6 +4446,23 @@ export class AgentManager { return flags.shouldNotifyWaiters; } + private captureLiveRunCheckpoint( + agent: ActiveManagedAgent, + event: AgentStreamEvent, + turnId: string | undefined, + options: HandleStreamEventOptions | undefined, + disposition: ActiveTurnTerminalDisposition, + ): Promise | undefined { + if (options?.fromHistory || disposition === "stale") return undefined; + const result = this.getRunResult(agent.id, turnId, event, false, disposition); + return this.runJournal.capture(agent.id, turnId, event, result)?.catch((error) => { + this.logger.error( + { err: error, agentId: agent.id, turnId }, + "Failed to persist provider checkpoint", + ); + }); + } + private getRunResult( agentId: string, turnId: string | undefined, diff --git a/packages/server/src/server/agent/agent-prompt.test.ts b/packages/server/src/server/agent/agent-prompt.test.ts index 95838e7157..b85dcdbb2d 100644 --- a/packages/server/src/server/agent/agent-prompt.test.ts +++ b/packages/server/src/server/agent/agent-prompt.test.ts @@ -289,6 +289,7 @@ test("finish notifications tell the parent the child's last assistant message", source: { kind: "agent-notification", agentId: "child-agent", + serverId: "srv_test", title: "Child Agent", event: "finished", }, @@ -325,6 +326,7 @@ test("closing a watched child notifies the caller", async () => { source: { kind: "agent-notification", agentId: "child-agent", + serverId: "srv_test", title: "Child Agent", event: "closed", }, diff --git a/packages/server/src/server/agent/mcp-parity.e2e.test.ts b/packages/server/src/server/agent/mcp-parity.e2e.test.ts index ca2fa1c041..722a01c0e4 100644 --- a/packages/server/src/server/agent/mcp-parity.e2e.test.ts +++ b/packages/server/src/server/agent/mcp-parity.e2e.test.ts @@ -1386,3 +1386,129 @@ test("MCP provider discovery uses the selected host profiles and models", async await b.close(); } }); + +test("real MCP discovery preserves custom catalog modes, disabled-provider isolation, and draft settings", async () => { + const clients = createTestAgentClients(); + const base = clients.codex!; + let disabledAvailabilityChecks = 0; + let disabledCatalogFetches = 0; + const drafts: AgentSessionConfig[] = []; + const dynamicModes = [{ id: "dynamic", label: "Dynamic", description: "Runtime mode" }]; + const custom: AgentClient = { + ...base, + provider: "zai", + capabilities: base.capabilities, + createSession: (...args) => base.createSession(...args), + resumeSession: (...args) => base.resumeSession(...args), + isAvailable: async () => true, + fetchCatalog: async () => ({ + models: [{ provider: "zai", id: "custom-model", label: "Custom model", isDefault: true }], + modes: dynamicModes, + }), + listFeatures: async (config) => { + drafts.push(config); + return [ + { + type: "toggle", + id: "fast_mode", + label: "Fast mode", + value: config.featureValues?.fast_mode === true, + }, + ]; + }, + }; + const disabled: AgentClient = { + ...base, + provider: "codex", + capabilities: base.capabilities, + createSession: (...args) => base.createSession(...args), + resumeSession: (...args) => base.resumeSession(...args), + isAvailable: async () => { + disabledAvailabilityChecks++; + return true; + }, + fetchCatalog: async (options) => { + disabledCatalogFetches++; + return base.fetchCatalog(options); + }, + }; + const destination = await createTestPaseoDaemon({ + agentClients: { ...clients, codex: disabled, zai: custom }, + providerOverrides: { + codex: { enabled: false }, + zai: { extends: "codex", label: "ZAI", enabled: true }, + }, + agentProfiles: [], + }); + const mcp = await createMcpClient(`http://127.0.0.1:${destination.port}/mcp/agents`); + try { + const response = await mcp.callTool({ name: "list_providers", args: {} }); + expect(response.isError, JSON.stringify(response)).not.toBe(true); + const providers = recordArr(getStructuredContent(response)?.providers); + expect(providers).toEqual( + expect.arrayContaining([ + expect.objectContaining({ + id: "zai", + label: "ZAI", + enabled: true, + status: "available", + modes: dynamicModes, + }), + expect.objectContaining({ id: "codex", enabled: false, status: "unavailable", modes: [] }), + ]), + ); + const modelText = + response.content?.map((item) => Reflect.get(item, "text") ?? "").join("\n") ?? ""; + expect(modelText).toContain(`providers_count=${providers.length}`); + expect(modelText).toContain( + `providers_ids=${providers.map((provider) => provider.id).join(",")}`, + ); + expect(modelText).toContain('"providers"'); + expect(await callToolStructured(mcp, "list_profiles")).toMatchObject({ profiles: [] }); + await expectToolError(mcp, "list_models", { provider: "codex" }, /disabled/i); + await expectToolError( + mcp, + "inspect_provider", + { provider: "codex", cwd: destination.staticDir }, + /disabled/i, + ); + expect(disabledAvailabilityChecks).toBe(0); + expect(disabledCatalogFetches).toBe(0); + const inspected = await callToolStructured(mcp, "inspect_provider", { + provider: "zai/custom-model", + cwd: destination.staticDir, + settings: { modeId: "dynamic", thinkingOptionId: "high", features: { fast_mode: true } }, + }); + expect(inspected).toMatchObject({ + provider: "zai", + label: "ZAI", + selectedModel: "custom-model", + modes: dynamicModes, + features: [{ type: "toggle", id: "fast_mode", label: "Fast mode", value: true }], + }); + expect(drafts).toContainEqual( + expect.objectContaining({ + provider: "zai", + cwd: destination.staticDir, + model: "custom-model", + modeId: "dynamic", + thinkingOptionId: "high", + featureValues: { fast_mode: true }, + }), + ); + } finally { + await mcp.close(); + await destination.close(); + } +}); + +test("real MCP without configured profiles returns an empty list", async () => { + const destination = await createTestPaseoDaemon(); + const mcp = await createMcpClient(`http://127.0.0.1:${destination.port}/mcp/agents`); + try { + expect(await callToolStructured(mcp, "list_profiles")).toMatchObject({ profiles: [] }); + } finally { + await mcp.close(); + await destination.close(); + } +}); diff --git a/packages/server/src/server/agent/mcp-server.test.ts b/packages/server/src/server/agent/mcp-server.test.ts index 21d0ef995c..3297e06d0c 100644 --- a/packages/server/src/server/agent/mcp-server.test.ts +++ b/packages/server/src/server/agent/mcp-server.test.ts @@ -67,8 +67,6 @@ import { createTerminalManager, type TerminalManager } from "../../terminal/term import { createWorkerTerminalManager } from "../../terminal/worker-terminal-manager.js"; import { PARENT_AGENT_ID_LABEL } from "@getpaseo/protocol/agent-labels"; import { PASEO_TOOL_NAMES } from "@getpaseo/protocol/paseo-tool-presentation/index"; -import { MutableDaemonConfigSchema, type AgentProfile } from "@getpaseo/protocol/messages"; -import type { DaemonConfigStore } from "../daemon-config-store.js"; import type { BrowserToolsBroker, BrowserToolsExecuteInput } from "../browser-tools/broker.js"; import type { BrowserToolsResponsePayload } from "../browser-tools/errors.js"; import { readPaseoWorktreeMetadata } from "../../utils/worktree-metadata.js"; @@ -4844,7 +4842,7 @@ describe("send_agent_prompt MCP tool", () => { }); expect(childClient.sessions[0]!.prompts).toHaveLength(1); expect(parseAgentMessage(childClient.sessions[0]!.prompts[0])).toMatchObject({ - source: { kind: "agent-message", agentId: parent.id }, + source: { kind: "agent-message", serverId: "srv_test", agentId: parent.id }, text: "Ongoing task", }); expect(parentClient.sessions[0]!.prompts).toHaveLength(0); @@ -5039,7 +5037,7 @@ describe("send_agent_prompt MCP tool", () => { await vi.waitFor(() => expect(childSession.prompts).toHaveLength(1)); expect(parseAgentMessage(childSession.prompts[0])).toEqual({ id: expect.any(String), - source: { kind: "agent-message", agentId: parent.id }, + source: { kind: "agent-message", serverId: "srv_test", agentId: parent.id }, text: "Follow up", }); acknowledgeTurnStart(); @@ -6273,202 +6271,6 @@ describe("schedule_logs MCP tool", () => { }); }); -describe("provider listing MCP tool", () => { - const logger = createTestLogger(); - - it("returns providers from the registry, including custom providers", async () => { - const { agentManager, agentStorage } = createTestDeps(); - const provStub = createProviderSnapshotManagerStub(); - provStub.listRegisteredProviderIds.mockReturnValue(["claude", "zai"]); - provStub.listProviders.mockResolvedValue([ - buildSnapshotEntry({ - provider: "claude", - label: "Claude", - description: "Test provider", - modes: [{ id: "default", label: "Default", description: "Built-in mode" }], - }), - buildSnapshotEntry({ - provider: "zai" as AgentProvider, - label: "ZAI", - description: "Custom Claude profile", - defaultModeId: "default", - modes: [{ id: "default", label: "Default", description: "Custom mode" }], - }), - ]); - - const server = await createAgentMcpServer({ - agentManager, - agentStorage, - providerSnapshotManager: provStub.manager, - logger, - }); - const tool = registeredTool(server, "list_providers"); - const response = await tool.handler({}); - const modelVisibleText = String(response.content[0]?.text); - - expect(response.structuredContent).toEqual({ - providers: [ - { - id: "claude", - label: "Claude", - description: "Test provider", - enabled: true, - status: "available", - modes: [{ id: "default", label: "Default", description: "Built-in mode" }], - }, - { - id: "zai", - label: "ZAI", - status: "available", - description: "Custom Claude profile", - enabled: true, - modes: [{ id: "default", label: "Default", description: "Custom mode" }], - }, - ], - }); - expect(modelVisibleText).toContain("providers_count=2"); - expect(modelVisibleText).toContain("providers_ids=claude,zai"); - expect(modelVisibleText).toContain('"providers"'); - }); - - it("returns provider modes from the shared snapshot catalog", async () => { - const { agentManager, agentStorage } = createTestDeps(); - const provStub = createProviderSnapshotManagerStub(); - provStub.listRegisteredProviderIds.mockReturnValue(["codex"]); - provStub.listProviders.mockResolvedValue([ - buildSnapshotEntry({ - provider: "codex", - label: "Codex", - modes: [{ id: "dynamic", label: "Dynamic", description: "Runtime mode" }], - }), - ]); - const server = await createAgentMcpServer({ - agentManager, - agentStorage, - providerSnapshotManager: provStub.manager, - logger, - }); - const tool = registeredTool(server, "list_providers"); - - const response = await tool.handler({}); - - expect(response.structuredContent.providers).toEqual([ - expect.objectContaining({ - id: "codex", - modes: [{ id: "dynamic", label: "Dynamic", description: "Runtime mode" }], - }), - ]); - }); - - it("returns disabled providers with metadata without checking availability", async () => { - const { agentManager, agentStorage } = createTestDeps(); - const provStub = createProviderSnapshotManagerStub(); - provStub.listRegisteredProviderIds.mockReturnValue(["codex"]); - provStub.listProviders.mockResolvedValue([ - buildSnapshotEntry({ - provider: "codex", - label: "Codex", - description: "OpenAI coding agent", - enabled: false, - modes: [{ id: "read-only", label: "Read Only", description: "No edits" }], - }), - ]); - const server = await createAgentMcpServer({ - agentManager, - agentStorage, - providerSnapshotManager: provStub.manager, - logger, - }); - const tool = registeredTool(server, "list_providers"); - const response = await tool.handler({}); - - expect(response.structuredContent).toEqual({ - providers: [ - { - id: "codex", - label: "Codex", - description: "OpenAI coding agent", - enabled: false, - status: "unavailable", - modes: [], - }, - ], - }); - }); -}); - -function daemonConfigStoreStub(agentProfiles?: AgentProfile[]): Pick { - const config = MutableDaemonConfigSchema.parse({ - relay: { enabled: true }, - mcp: { injectIntoAgents: true }, - ...(agentProfiles !== undefined ? { agentProfiles } : {}), - }); - return { get: () => config }; -} - -describe("agent profile listing MCP tool", () => { - const logger = createTestLogger(); - - it("returns configured profiles, including notes", async () => { - const { agentManager, agentStorage } = createTestDeps(); - const profiles: AgentProfile[] = [ - { - id: "ui-profile", - name: "UI work", - provider: "claude", - model: "claude-test-model", - modeId: "bypassPermissions", - thinkingOptionId: "high", - featureValues: { fast_mode: true }, - notes: "Use for UI work: components, layout, design tokens. Not for backend.", - }, - ]; - const server = await createAgentMcpServer({ - agentManager, - agentStorage, - providerSnapshotManager: createOpenCodeManager().manager, - daemonConfigStore: daemonConfigStoreStub(profiles), - logger, - }); - const tool = registeredTool(server, "list_profiles"); - - const response = await tool.handler({}); - - expect(response.structuredContent).toEqual({ profiles }); - }); - - it("returns an empty array when no profiles are configured", async () => { - const { agentManager, agentStorage } = createTestDeps(); - const server = await createAgentMcpServer({ - agentManager, - agentStorage, - providerSnapshotManager: createOpenCodeManager().manager, - daemonConfigStore: daemonConfigStoreStub(), - logger, - }); - const tool = registeredTool(server, "list_profiles"); - - const response = await tool.handler({}); - - expect(response.structuredContent).toEqual({ profiles: [] }); - }); - - it("returns an empty array when no daemon config store is provided", async () => { - const { agentManager, agentStorage } = createTestDeps(); - const server = await createAgentMcpServer({ - agentManager, - agentStorage, - providerSnapshotManager: createOpenCodeManager().manager, - logger, - }); - const tool = registeredTool(server, "list_profiles"); - - const response = await tool.handler({}); - - expect(response.structuredContent).toEqual({ profiles: [] }); - }); -}); - describe("provider MCP tools", () => { const logger = createTestLogger(); @@ -6483,113 +6285,6 @@ describe("provider MCP tools", () => { expect(lookupTool(server, "list_provider_features")).toBeUndefined(); }); - - it("inspects provider features for a draft agent configuration", async () => { - const { agentManager, agentStorage, spies } = createTestDeps(); - spies.agentManager.listDraftFeatures.mockResolvedValue([ - { - type: "toggle", - id: "fast_mode", - label: "Fast mode", - value: false, - }, - ]); - const provStub = createProviderSnapshotManagerStub(); - provStub.listRegisteredProviderIds.mockReturnValue(["codex"]); - const codexEntry = buildSnapshotEntry({ - provider: "codex", - label: "Codex", - description: "OpenAI coding agent", - modes: [{ id: "full-access", label: "Full Access", description: "Can edit files" }], - }); - provStub.listProviders.mockResolvedValue([codexEntry]); - provStub.getProvider.mockResolvedValue(codexEntry); - const server = await createAgentMcpServer({ - agentManager, - agentStorage, - providerSnapshotManager: provStub.manager, - logger, - }); - const tool = registeredTool(server, "inspect_provider"); - const input = { - provider: "codex/gpt-5.4", - cwd: "~/repo", - settings: { - modeId: "full-access", - thinkingOptionId: "high", - features: { fast_mode: true }, - }, - }; - - const parsed = await tool.inputSchema.safeParseAsync(input); - expect(parsed.success).toBe(true); - - const response = await tool.handler(input); - - expect(spies.agentManager.listDraftFeatures).toHaveBeenCalledWith({ - provider: "codex", - cwd: expect.stringContaining("repo"), - modeId: "full-access", - model: "gpt-5.4", - thinkingOptionId: "high", - featureValues: { fast_mode: true }, - }); - expect(response.structuredContent).toEqual({ - provider: "codex", - label: "Codex", - description: "OpenAI coding agent", - enabled: true, - status: "available", - modes: [{ id: "full-access", label: "Full Access", description: "Can edit files" }], - selectedModel: "gpt-5.4", - features: [ - { - type: "toggle", - id: "fast_mode", - label: "Fast mode", - value: false, - }, - ], - }); - }); - - it("rejects disabled providers without fetching models", async () => { - const { agentManager, agentStorage } = createTestDeps(); - const provStub = createProviderSnapshotManagerStub(); - provStub.listRegisteredProviderIds.mockReturnValue(["codex"]); - provStub.listModels.mockRejectedValue(new Error("Provider 'codex' is disabled")); - const server = await createAgentMcpServer({ - agentManager, - agentStorage, - providerSnapshotManager: provStub.manager, - logger, - }); - const tool = registeredTool(server, "list_models"); - - await expect(tool.handler({ provider: "codex" })).rejects.toThrow( - "Provider 'codex' is disabled", - ); - }); - - it("inspect_provider rejects disabled providers without fetching models", async () => { - const { agentManager, agentStorage } = createTestDeps(); - const provStub = createProviderSnapshotManagerStub(); - provStub.listRegisteredProviderIds.mockReturnValue(["codex"]); - provStub.getProvider.mockResolvedValue( - buildSnapshotEntry({ provider: "codex", label: "Codex", enabled: false }), - ); - const server = await createAgentMcpServer({ - agentManager, - agentStorage, - providerSnapshotManager: provStub.manager, - logger, - }); - const tool = registeredTool(server, "inspect_provider"); - - await expect(tool.handler({ provider: "codex", cwd: "~/repo" })).rejects.toThrow( - "Provider 'codex' is disabled", - ); - }); }); describe("speak MCP tool", () => { diff --git a/packages/server/src/server/agent/runs/index.ts b/packages/server/src/server/agent/runs/index.ts index 3168c42013..189d7045a5 100644 --- a/packages/server/src/server/agent/runs/index.ts +++ b/packages/server/src/server/agent/runs/index.ts @@ -60,6 +60,10 @@ export class AgentRuns { private identity(agentId: string, submissionId: string): string { return JSON.stringify([agentId, submissionId]); } + /** Restore turn ownership before subscribing to live provider evidence. */ + restore(agentId: string): Promise { + return this.mutate(() => this.load(agentId)); + } private async load(agentId: string): Promise { if (this.loadedAgents.has(agentId)) return; let records: StoredRun[] = []; @@ -240,10 +244,15 @@ export class AgentRuns { turnId: string | undefined, event: AgentStreamEvent, result: string, - ): Promise { - if (!turnId) return Promise.resolve(); + ): Promise | undefined { + if (!turnId) return undefined; + if ( + this.loadedAgents.has(agentId) && + !this.submissionsByTurn.has(this.identity(agentId, turnId)) + ) + return undefined; const checkpoint = checkpointFor(event, result); - if (!checkpoint) return Promise.resolve(); + if (!checkpoint) return undefined; return this.mutate(async () => { await this.load(agentId); const submissions = this.submissionsByTurn.get(this.identity(agentId, turnId)); diff --git a/packages/server/src/server/bootstrap.ts b/packages/server/src/server/bootstrap.ts index 80fa80ba0a..afc3536546 100644 --- a/packages/server/src/server/bootstrap.ts +++ b/packages/server/src/server/bootstrap.ts @@ -1779,6 +1779,7 @@ export async function createPaseoDaemon( orchestrationSkills, workspaceLabelService, ); + await wsServer.initializeRelationships(); pluginRuntime.bindPaseoSessionHost(wsServer); await pluginRuntime.start(); providerSnapshotManager.settlePluginProviders(); diff --git a/packages/server/src/server/connected-hosts/hosts.e2e.test.ts b/packages/server/src/server/connected-hosts/hosts.e2e.test.ts index d2dc734438..200498d438 100644 --- a/packages/server/src/server/connected-hosts/hosts.e2e.test.ts +++ b/packages/server/src/server/connected-hosts/hosts.e2e.test.ts @@ -1708,6 +1708,180 @@ test("parent capability owns local and fleet membership, idempotent creation and expect(recovered.address).toEqual(first.address); }); +test("startup upgrade adopts existing ordinary local children once without treating later labels as membership", async () => { + const { AgentStorage } = await import("../agent/agent-storage.js"); + const { default: pino } = await import("pino"); + const root = await mkdtemp(join(tmpdir(), "paseo-legacy-membership-")); + cleanup.push(() => rm(root, { recursive: true, force: true })); + const storage = new AgentStorage(join(root, ".paseo", "agents"), pino({ level: "silent" })); + const parentId = crypto.randomUUID(); + const childId = crypto.randomUUID(); + const sameWorkspaceId = crypto.randomUUID(); + const otherWorkspaceId = crypto.randomUUID(); + const openTabId = crypto.randomUUID(); + const foreignId = crypto.randomUUID(); + const nullParentId = crypto.randomUUID(); + const executedId = crypto.randomUUID(); + const internalParentId = crypto.randomUUID(); + const internalChildId = crypto.randomUUID(); + const seed = async ( + id: string, + overrides: Partial = {}, + ) => { + await storage.upsert({ + id, + provider: "codex", + cwd: root, + workspaceId: "legacy-workspace", + config: { model: "test" }, + createdAt: new Date().toISOString(), + updatedAt: new Date().toISOString(), + lastStatus: "closed", + labels: {}, + ...overrides, + }); + }; + const labels = { "paseo.parent-agent-id": parentId }; + await seed(parentId); + await seed(childId, { labels }); + await seed(sameWorkspaceId, { labels }); + await seed(otherWorkspaceId, { labels, workspaceId: "other-workspace" }); + await seed(openTabId, { labels: { ...labels, "paseo.open-agent-tab.retained": "true" } }); + await seed(foreignId, { labels, parent: { serverId: "another-server", agentId: parentId } }); + await seed(nullParentId, { labels, parent: null }); + await seed(executedId, { + labels, + owner: { kind: "daemon", daemonId: "hub", executionId: "execution" }, + }); + await seed(internalParentId, { internal: true }); + await seed(internalChildId, { labels: { "paseo.parent-agent-id": internalParentId } }); + const first = await createTestPaseoDaemon({ paseoHomeRoot: root, cleanup: false }); + cleanup.push(() => first.close()); + const origin = await api(first); + const parent = origin.paseo.agents.ref(parentId); + const rows = (await parent.subagents.list()).rows; + expect(rows.map((row) => row.child.agentId).sort()).toEqual( + [childId, sameWorkspaceId, otherWorkspaceId, openTabId].sort(), + ); + expect( + rows.every((row) => row.child.serverId === parent.serverId && row.phase === "attached"), + ).toBe(true); + expect((await first.daemon.agentStorage.get(childId))?.parent).toEqual(parent.address); + await parent.subagents.forget({ serverId: parent.serverId, agentId: childId }); + await parent.archive(); + await expect + .poll(async () => (await first.daemon.agentStorage.get(sameWorkspaceId))?.archivedAt) + .toEqual(expect.any(String)); + await expect + .poll(async () => (await first.daemon.agentStorage.get(otherWorkspaceId))?.parent) + .toBeNull(); + await expect + .poll(async () => (await first.daemon.agentStorage.get(openTabId))?.parent) + .toBeNull(); + // Forget leaves destination lifecycle alone; the old label cascade must not run too. + expect((await first.daemon.agentStorage.get(childId))?.archivedAt).toBeUndefined(); + expect((await first.daemon.agentStorage.get(otherWorkspaceId))?.archivedAt).toBeUndefined(); + expect((await first.daemon.agentStorage.get(openTabId))?.archivedAt).toBeUndefined(); + const injectedId = crypto.randomUUID(); + await first.daemon.agentStorage.upsert({ + ...(await first.daemon.agentStorage.get(childId))!, + id: injectedId, + parent: undefined, + }); + await first.close(); + const restarted = await createTestPaseoDaemon({ paseoHomeRoot: root, cleanup: false }); + cleanup.push(() => restarted.close()); + const restored = await api(restarted); + const restoredRows = (await restored.paseo.agents.ref(parentId).subagents.list()).rows; + expect(restoredRows.map((row) => row.child.agentId).sort()).toEqual( + [sameWorkspaceId, otherWorkspaceId, openTabId].sort(), + ); + expect(restoredRows.every((row) => ["archived", "detached"].includes(row.phase))).toBe(true); + expect((await restarted.daemon.agentStorage.get(injectedId))?.parent).toBeUndefined(); +}); + +test("an interrupted legacy upgrade resumes its bounded candidates before admitting new writes", async () => { + const { AgentStorage } = await import("../agent/agent-storage.js"); + const { Subagents } = await import("../subagents/index.js"); + const { getOrCreateServerId } = await import("../server-id.js"); + const { default: pino } = await import("pino"); + const root = await mkdtemp(join(tmpdir(), "paseo-interrupted-legacy-")); + cleanup.push(() => rm(root, { recursive: true, force: true })); + const home = join(root, ".paseo"); + await mkdir(home); + const serverId = getOrCreateServerId(home); + const storage = new AgentStorage(join(home, "agents"), pino({ level: "silent" })); + const parentId = crypto.randomUUID(); + const childIds = [crypto.randomUUID(), crypto.randomUUID()]; + for (const id of [parentId, ...childIds]) + await storage.upsert({ + id, + provider: "codex", + cwd: root, + workspaceId: "legacy-workspace", + config: { model: "test" }, + createdAt: new Date().toISOString(), + updatedAt: new Date().toISOString(), + lastStatus: "closed", + labels: id === parentId ? {} : { "paseo.parent-agent-id": parentId }, + }); + let qualifications = 0; + const interrupted = new Subagents({ + home, + serverId, + agents: { + list: () => storage.list(), + get: (id) => storage.get(id), + upsert: async (record) => { + if (++qualifications === 2) throw new Error("Injected qualification write failure"); + await storage.upsert(record); + }, + }, + resolveAuthority: async () => ({ + principalId: "owner", + permissions: ["daemon.read", "workspace.read", "workspace.write"], + }), + runtime: async () => { + throw new Error("Upgrade must never acquire a client or submit creation"); + }, + }); + await expect(interrupted.initialize()).rejects.toThrow("Injected qualification write failure"); + const pending = JSON.parse(await readFile(join(home, "subagents.json"), "utf8")); + expect(pending.legacyUpgrade).toMatchObject({ complete: false, candidates: expect.any(Array) }); + expect(pending.legacyUpgrade.candidates).toHaveLength(2); + expect((await storage.list()).filter((record) => record.parent)).toHaveLength(1); + const injectedId = crypto.randomUUID(); + await storage.upsert({ + ...(await storage.get(childIds[0]!))!, + id: injectedId, + parent: undefined, + }); + const restarted = await createTestPaseoDaemon({ paseoHomeRoot: root, cleanup: false }); + cleanup.push(() => restarted.close()); + const origin = await api(restarted); + const parent = origin.paseo.agents.ref(parentId); + expect((await parent.subagents.list()).rows.map((row) => row.child.agentId).sort()).toEqual( + childIds.toSorted(), + ); + for (const id of childIds) + expect((await restarted.daemon.agentStorage.get(id))?.parent).toEqual(parent.address); + expect((await restarted.daemon.agentStorage.get(injectedId))?.parent).toBeUndefined(); + const complete = JSON.parse(await readFile(join(home, "subagents.json"), "utf8")); + expect(complete.legacyUpgrade.complete).toBe(true); + expect(complete.legacyUpgrade.candidates).toEqual(pending.legacyUpgrade.candidates); + const archived: string[] = []; + const release = restarted.daemon.agentManager.onAgentArchived((id) => { + archived.push(id); + }); + try { + await parent.archive(); + await expect.poll(() => archived.toSorted()).toEqual([parentId, ...childIds].toSorted()); + expect((await restarted.daemon.agentStorage.get(injectedId))?.archivedAt).toBeUndefined(); + } finally { + release(); + } +}); + test("retained derived sessions cannot supply durable principal authority", async () => { const a = await daemon(); const owner = await api(a); @@ -3188,15 +3362,30 @@ test.each([false, true])( }; for (const attempt of [1, 2]) { const requestId = "legacy-original"; - socket.send(JSON.stringify({ type: "session", message: { ...request, requestId } })); - await expect - .poll( - () => + const receipt = new Promise((resolve, reject) => { + const timeout = setTimeout(() => { + socket.off("message", received); + reject( + new Error( + `Timed out waiting for legacy creation receipt ${attempt}: ${JSON.stringify(messages)}`, + ), + ); + }, 20_000); + function received() { + if ( messages.filter( (message) => message.type === "status" && message.payload.requestId === requestId, - ).length, - ) - .toBe(attempt); + ).length !== attempt + ) + return; + clearTimeout(timeout); + socket.off("message", received); + resolve(); + } + socket.on("message", received); + }); + socket.send(JSON.stringify({ type: "session", message: { ...request, requestId } })); + await receipt; const response = messages.find( (message) => message.type === "status" && message.payload.requestId === requestId, ); @@ -3225,3 +3414,158 @@ test.each([false, true])( expect(prompts).toBe(1); }, ); + +test("internal connection leases share ordinary host resolution and preserve borrowed origin and sibling lifetime", async () => { + const { acquireHostConnection } = await import("@getpaseo/client/internal/host-connection"); + const a = await daemon(); + const b = await daemon(); + const origin = await api(a); + const direct = await api(b); + await origin.paseo.hosts.add({ connection: `127.0.0.1:${b.port}` }); + const directory = await online(origin.paseo); + const controller = new AbortController(); + const [self, first, sibling] = await Promise.all([ + acquireHostConnection(origin.client, directory.self.serverId, controller.signal), + acquireHostConnection(origin.client, b.daemon.getServerId(), controller.signal), + acquireHostConnection(origin.client, b.daemon.getServerId(), controller.signal), + ]); + cleanup.push(() => sibling.release()); + cleanup.push(() => first.release()); + expect(self.client.getServerId()).toBe(origin.client.getServerId()); + await self.release(); + expect(await origin.paseo.hosts.list()).toMatchObject({ + self: { serverId: directory.self.serverId }, + }); + const firstApi = createPaseoApi(first.client); + const siblingApi = createPaseoApi(sibling.client); + cleanup.push(() => siblingApi.dispose()); + cleanup.push(() => firstApi.dispose()); + await first.release(); + await expect(firstApi.workspaces.list()).rejects.toThrow(); + expect((await siblingApi.workspaces.list()).entries).toEqual([]); + await origin.paseo.hosts.update(directory.hosts[0]!.registrationId, { enabled: false }); + await expect(siblingApi.workspaces.list()).rejects.toThrow(); + expect((await direct.paseo.workspaces.list()).entries).toEqual([]); + controller.abort(); + await sibling.release(); + const aborted = new AbortController(); + aborted.abort(); + await expect( + acquireHostConnection(origin.client, directory.self.serverId, aborted.signal), + ).rejects.toThrow("disposed"); + expect((await origin.paseo.hosts.list()).self.serverId).toBe(directory.self.serverId); +}); + +test("released internal self connection capability rejects captured operations while the borrowed origin stays usable", async () => { + const { acquireHostConnection } = await import("@getpaseo/client/internal/host-connection"); + const a = await daemon(); + const origin = await api(a); + const control = new AbortController(); + const self = await acquireHostConnection(origin.client, a.daemon.getServerId(), control.signal); + const scoped = createPaseoApi(self.client); + cleanup.push(() => scoped.dispose()); + await self.release(); + await expect(scoped.workspaces.list()).rejects.toMatchObject({ code: "DAEMON_CONNECTION_LOST" }); + expect((await origin.paseo.workspaces.list()).entries).toEqual([]); +}); + +test("nested self scopes preserve the consumed hop while a direct self scope reaches its registered destination", async () => { + const a = await daemon(); + const b = await daemon(); + const c = await daemon(); + const origin = await api(a); + const directB = await api(b); + const wire = await hostWire(b); + await origin.paseo.hosts.add({ connection: wire.connection }); + await online(origin.paseo); + await directB.paseo.hosts.add({ connection: `127.0.0.1:${c.port}` }); + await online(directB.paseo); + const selfA = await origin.paseo.hosts.client(a.daemon.getServerId()); + cleanup.push(() => selfA.dispose()); + const viaB = await selfA.hosts.client(b.daemon.getServerId()); + cleanup.push(() => viaB.dispose()); + const nestedB = await viaB.hosts.client(b.daemon.getServerId()); + cleanup.push(() => nestedB.dispose()); + const twiceB = await nestedB.hosts.client(b.daemon.getServerId()); + cleanup.push(() => twiceB.dispose()); + expect((await twiceB.workspaces.list()).entries).toEqual([]); + await expect(twiceB.hosts.client(c.daemon.getServerId())).rejects.toMatchObject({ + code: "unsupported_capability", + }); + const directSelfB = await directB.paseo.hosts.client(b.daemon.getServerId()); + cleanup.push(() => directSelfB.dispose()); + const allowedC = await directSelfB.hosts.client(c.daemon.getServerId()); + cleanup.push(() => allowedC.dispose()); + expect((await allowedC.workspaces.list()).entries).toEqual([]); + expect(wire.connections()).toBe(1); + await nestedB.dispose(); + expect((await viaB.workspaces.list()).entries).toEqual([]); + expect((await origin.paseo.workspaces.list()).entries).toEqual([]); +}); + +test("owned self session close settles pending RPC and observation and leaves sibling and physical origin usable", async () => { + const { acquireHostConnection } = await import("@getpaseo/client/internal/host-connection"); + const a = await daemon(); + const wire = await hostWire(a); + const origin = createHostClient(wire.connection, { + clientId: crypto.randomUUID(), + reconnect: { enabled: false }, + }); + cleanup.push(() => origin.close()); + await origin.connect(); + const originApi = createPaseoApi(origin); + cleanup.push(() => originApi.dispose()); + const control = new AbortController(); + const self = await acquireHostConnection(origin, a.daemon.getServerId(), control.signal); + cleanup.push(() => self.release()); + const sibling = await acquireHostConnection( + origin, + a.daemon.getServerId(), + new AbortController().signal, + ); + cleanup.push(() => sibling.release()); + const observed = self.client.observeHosts(); + const siblingObserved = sibling.client.observeHosts(); + let updates = 0; + let siblingUpdates = 0; + observed.subscribe({ + snapshot() {}, + update() { + updates++; + }, + }); + siblingObserved.subscribe({ + snapshot() {}, + update() { + siblingUpdates++; + }, + }); + await Promise.all([observed.ready, siblingObserved.ready]); + const states: string[] = []; + self.client.subscribeConnectionStatus((state) => states.push(state.status)); + const identity = self.client.getServerId(); + const capturedFetch = self.client.fetchWorkspaces.bind(self.client); + const held = wire.holdResponses((message) => message.type === "fetch_workspaces_response"); + const pending = capturedFetch(); + void pending.catch(() => {}); + await held.ready; + try { + await self.release(); + await expect(pending).rejects.toThrow(); + expect(states.at(-1)).toBe("disposed"); + expect(self.client.getServerId()).toBe(identity); + const rejected = capturedFetch(); + expect(rejected).toBeInstanceOf(Promise); + await expect(rejected).rejects.toMatchObject({ code: "DAEMON_CONNECTION_LOST" }); + held.release(); + await originApi.hosts.add({ connection: "127.0.0.1:1", label: "After scope release" }); + await expect.poll(() => siblingUpdates).toBeGreaterThan(0); + expect(updates).toBe(0); + expect((await sibling.client.fetchWorkspaces()).entries).toEqual([]); + expect((await originApi.workspaces.list()).entries).toEqual([]); + control.abort(); + await self.release(); + } finally { + held.release(); + } +}); diff --git a/packages/server/src/server/connected-hosts/link.ts b/packages/server/src/server/connected-hosts/link.ts index 183ae4c220..4754175b3b 100644 --- a/packages/server/src/server/connected-hosts/link.ts +++ b/packages/server/src/server/connected-hosts/link.ts @@ -16,16 +16,23 @@ export class DaemonLinkEndpoint { constructor( private readonly parent: { identity: string; + serverId?: string; + assertScopeAdmission?(): void; permissions(): readonly DaemonPermission[]; routed(): boolean; bufferedAmount(): number | null; delivery: SessionDelivery; - admit(socket: VirtualSessionSocket, permissions: readonly DaemonPermission[]): Promise; + admit( + socket: VirtualSessionSocket, + permissions: readonly DaemonPermission[], + context: { kind: "scope" | "destination-link" }, + ): Promise; }, ) {} handles(message: SessionInboundMessage): boolean { return ( message.type === "hosts.link.open.request" || + (message.type === "hosts.route.open.request" && message.serverId === this.parent.serverId) || (message.type === "hosts.route.frame.request" && this.channels.has(message.routeId)) || (message.type === "hosts.route.close.request" && ((!!message.routeId && this.channels.has(message.routeId)) || @@ -33,105 +40,11 @@ export class DaemonLinkEndpoint { ); } async dispatch(message: SessionInboundMessage): Promise { - if (message.type === "hosts.link.open.request") { - const error = (code: string, text: string) => - this.parent.delivery.reply({ - type: "rpc_error", - payload: { - requestId: message.requestId, - requestType: message.type, - code, - error: text, - }, - }); - if (this.closed || this.parent.routed()) { - error("unsupported_capability", "Connected hosts support one hop"); - return; - } - if (this.channels.size >= 128) { - error("resource_limit", "Too many host channels"); - return; - } - let socket: VirtualSessionSocket; - const owner = this.parent.delivery.begin("host-link", undefined, () => { - this.channels.delete(owner.id); - socket?.close(); - }); - const id = owner.id; - // Capture actual authenticated authority. The caller hello can only reduce it. - const ceiling = [...this.parent.permissions()]; - socket = new VirtualSessionSocket( - async (data) => { - const output: SessionOutboundMessage | null = - typeof data === "string" - ? { - type: "hosts.route.frame", - payload: { subscriptionId: id, data, binary: false }, - } - : null; - const frame = typeof data === "string" ? null : encodeHostRouteFrame(id, data); - const bytes = output - ? Buffer.byteLength(JSON.stringify({ type: "session", message: output })) - : frame!.byteLength; - const buffered = this.parent.bufferedAmount(); - // Account for the exact escaped outer frame, retaining control-frame headroom. - if (buffered === null || buffered + bytes > 64 * 1024 * 1024 - 65536) { - socket.close(1009, "Host channel output limit reached"); - throw new Error("Host channel output limit reached"); - } - try { - if (output) await owner.emitAndWait(output); - else await owner.emitBinaryAndWait(frame!); - } catch (failure) { - socket.close(); - throw failure; - } - }, - (code, reason) => { - this.channels.delete(id); - let failure = "host_offline"; - if (code === 1009) failure = "resource_limit"; - if (code === 1008) failure = "permission_denied"; - owner.emit({ - type: "hosts.route.closed", - payload: { - subscriptionId: id, - code: failure, - message: reason, - }, - }); - void owner.release(); - }, - () => this.parent.bufferedAmount() ?? 0, - (data) => { - const hello = JSON.parse(data); - if (hello.type !== "hello" || typeof hello.clientId !== "string") - throw new Error("Expected hello"); - return JSON.stringify({ - ...hello, - auth: undefined, - routed: true, - clientId: `link-${createHash("sha256") - .update(JSON.stringify([this.parent.identity, hello.clientId])) - .digest("hex")}`, - }); - }, - ); - this.channels.set(id, { scopeId: message.channelId, socket }); - try { - await this.parent.admit(socket, ceiling); - if (this.closed || socket.readyState !== 1) { - socket.close(); - return; - } - this.parent.delivery.reply({ - type: "hosts.link.open.response", - payload: { requestId: message.requestId, routeId: id }, - }); - } catch { - socket.close(); - error("host_offline", "Could not admit host channel"); - } + if ( + message.type === "hosts.link.open.request" || + (message.type === "hosts.route.open.request" && message.serverId === this.parent.serverId) + ) { + await this.openChannel(message); return; } if (message.type === "hosts.route.frame.request" && this.channels.has(message.routeId)) { @@ -146,6 +59,127 @@ export class DaemonLinkEndpoint { } return; } + private async openChannel( + message: Extract< + SessionInboundMessage, + { type: "hosts.link.open.request" | "hosts.route.open.request" } + >, + ): Promise { + const selfScope = message.type === "hosts.route.open.request"; + const error = (code: string, text: string) => + this.parent.delivery.reply({ + type: "rpc_error", + payload: { + requestId: message.requestId, + requestType: message.type, + code, + error: text, + }, + }); + if (this.closed || (!selfScope && this.parent.routed())) { + error("unsupported_capability", "Connected hosts support one hop"); + return; + } + try { + this.parent.assertScopeAdmission?.(); + } catch (failure) { + error( + "unsupported_capability", + failure instanceof Error ? failure.message : "Unsupported scoped admission", + ); + return; + } + if (this.channels.size >= 128) { + error("resource_limit", "Too many host channels"); + return; + } + let socket: VirtualSessionSocket; + const owner = this.parent.delivery.begin("host-link", undefined, () => { + this.channels.delete(owner.id); + socket?.close(); + }); + const id = owner.id; + // Capture actual authenticated authority. The caller hello can only reduce it. + const ceiling = [...this.parent.permissions()]; + socket = new VirtualSessionSocket( + async (data) => { + const output: SessionOutboundMessage | null = + typeof data === "string" + ? { + type: "hosts.route.frame", + payload: { subscriptionId: id, data, binary: false }, + } + : null; + const frame = typeof data === "string" ? null : encodeHostRouteFrame(id, data); + const bytes = output + ? Buffer.byteLength(JSON.stringify({ type: "session", message: output })) + : frame!.byteLength; + const buffered = this.parent.bufferedAmount(); + // Account for the exact escaped outer frame, retaining control-frame headroom. + if (buffered === null || buffered + bytes > 64 * 1024 * 1024 - 65536) { + socket.close(1009, "Host channel output limit reached"); + throw new Error("Host channel output limit reached"); + } + try { + if (output) await owner.emitAndWait(output); + else await owner.emitBinaryAndWait(frame!); + } catch (failure) { + socket.close(); + throw failure; + } + }, + (code, reason) => { + this.channels.delete(id); + let failure = "host_offline"; + if (code === 1009) failure = "resource_limit"; + if (code === 1008) failure = "permission_denied"; + owner.emit({ + type: "hosts.route.closed", + payload: { + subscriptionId: id, + code: failure, + message: reason, + }, + }); + void owner.release(); + }, + () => this.parent.bufferedAmount() ?? 0, + (data) => { + const hello = JSON.parse(data); + if (hello.type !== "hello" || typeof hello.clientId !== "string") + throw new Error("Expected hello"); + return JSON.stringify({ + ...hello, + auth: undefined, + routed: selfScope ? this.parent.routed() : true, + sessionPermissions: ceiling.filter( + (permission) => + !hello.sessionPermissions || hello.sessionPermissions.includes(permission), + ), + clientId: `link-${createHash("sha256") + .update(JSON.stringify([this.parent.identity, hello.clientId])) + .digest("hex")}`, + }); + }, + ); + this.channels.set(id, { scopeId: selfScope ? message.scopeId : message.channelId, socket }); + try { + await this.parent.admit(socket, ceiling, { + kind: selfScope ? "scope" : "destination-link", + }); + if (this.closed || socket.readyState !== 1) { + socket.close(); + return; + } + this.parent.delivery.reply({ + type: selfScope ? "hosts.route.open.response" : "hosts.link.open.response", + payload: { requestId: message.requestId, routeId: id }, + }); + } catch { + socket.close(); + error("host_offline", "Could not admit host channel"); + } + } binary(frame: HostRouteFrame): boolean { const socket = this.channels.get(frame.routeId)?.socket; if (!socket) return false; diff --git a/packages/server/src/server/hub/execution-session.websocket.test.ts b/packages/server/src/server/hub/execution-session.websocket.test.ts index 88f376b3ba..eb815ebda9 100644 --- a/packages/server/src/server/hub/execution-session.websocket.test.ts +++ b/packages/server/src/server/hub/execution-session.websocket.test.ts @@ -62,6 +62,24 @@ test("Hub execute can steer ordinary agents while unrelated administration stays expect(hub.serverInfoPermissions()).toEqual([["hub.execute"]]); }); +test("execution-scoped admission cannot acquire an ordinary self session", async () => { + const hub = await launchRelationship(); + expect( + await hub.requestOrdinary({ + type: "hosts.route.open.request", + requestId: "execution-self-scope", + serverId: hub.serverId(), + scopeId: "execution-self-scope", + }), + ).toMatchObject({ + type: "rpc_error", + payload: { requestId: "execution-self-scope", code: "unsupported_capability" }, + }); + // Rejection leaves the originating execution capability usable. + const created = await hub.createOwnedConcurrently(); + expect(created.duplicate.agentId).toBe(created.first.agentId); +}); + test("Hub execute can title the workspace created with an agent", async () => { const hub = await launchRelationship(); const created = await hub.requestOrdinary({ @@ -165,11 +183,23 @@ test("ordinary Hub create and message retries do not duplicate agents or prompts }; expect(await hub.requestOrdinary({ ...message, requestId: "message-first" })).toEqual({ type: "send_agent_message_response", - payload: { requestId: "message-first", agentId, accepted: true, error: null }, + payload: { + requestId: "message-first", + agentId, + accepted: true, + error: null, + dispatch: { disposition: "turn_started", tracking: "turn" }, + }, }); expect(await hub.requestOrdinary({ ...message, requestId: "message-duplicate" })).toEqual({ type: "send_agent_message_response", - payload: { requestId: "message-duplicate", agentId, accepted: true, error: null }, + payload: { + requestId: "message-duplicate", + agentId, + accepted: true, + error: null, + dispatch: { disposition: "turn_started", tracking: "turn" }, + }, }); expect( hub diff --git a/packages/server/src/server/hub/test-utils/relationship-harness.ts b/packages/server/src/server/hub/test-utils/relationship-harness.ts index 77f75e8bb8..d9e150716c 100644 --- a/packages/server/src/server/hub/test-utils/relationship-harness.ts +++ b/packages/server/src/server/hub/test-utils/relationship-harness.ts @@ -1212,6 +1212,10 @@ export class HubRelationshipHarness { ); } + serverId(): string { + return this.daemon!.getServerId(); + } + probeTrustedHello(): number | null { const socket = this.latestSocket().socket; socket.receiveEnvelope({ diff --git a/packages/server/src/server/subagents/index.ts b/packages/server/src/server/subagents/index.ts index fb608aa699..9b9c4b226d 100644 --- a/packages/server/src/server/subagents/index.ts +++ b/packages/server/src/server/subagents/index.ts @@ -1,6 +1,6 @@ import type { HostDirectory } from "@getpaseo/protocol/hosts"; import { shouldDetachFromArchivedParent } from "@getpaseo/protocol/agent-lifecycle"; -import { hasOpenAgentTab } from "@getpaseo/protocol/agent-labels"; +import { hasOpenAgentTab, getParentAgentIdFromLabels } from "@getpaseo/protocol/agent-labels"; import { WorkspaceAddressSchema, type WorkspaceAddress } from "@getpaseo/protocol/addresses"; import type { AgentSnapshotPayload, CreationSnapshot } from "@getpaseo/protocol/messages"; import { createHash, randomUUID } from "node:crypto"; @@ -19,6 +19,8 @@ import { } from "@getpaseo/protocol/messages"; import { writeJsonFileAtomic } from "../atomic-file.js"; import type { HostRuntimeAuthority } from "../host-runtime/index.js"; +import type { AgentStorage, StoredAgentRecord } from "../agent/agent-storage.js"; +import { buildStoredAgentPayload } from "../agent/agent-projections.js"; /** Stored authority contains identity and a ceiling; live admission belongs to the host. */ export interface SubagentAuthority { @@ -31,11 +33,9 @@ const AuthoritySchema = z.object({ sourceId: z.string(), permissionCeiling: z.array(DaemonPermissionSchema), }); -const RecordSchema = z.object({ +const RelationshipStateSchema = z.object({ row: SubagentRelationshipSchema, authority: AuthoritySchema, - creation: SubagentCreateInputSchema, - fingerprint: z.string(), confirmed: z.enum(["attached", "detached", "archived", "canceled"]).nullable(), intent: z.object({ revision: z.number().int().nonnegative(), @@ -44,8 +44,24 @@ const RecordSchema = z.object({ authority: AuthoritySchema, }), }); +const RecordSchema = z.discriminatedUnion("provenance", [ + RelationshipStateSchema.extend({ + provenance: z.literal("creation"), + creation: SubagentCreateInputSchema, + fingerprint: z.string(), + }), + RelationshipStateSchema.extend({ + provenance: z.literal("legacy-local-lifecycle"), + creation: z.null(), + fingerprint: z.null(), + }), +]); type RelationshipRecord = z.infer; type RelationshipIntent = RelationshipRecord["intent"]; +const LedgerSchema = z.object({ + records: z.array(RecordSchema), + legacyUpgrade: z.object({ complete: z.boolean(), candidates: z.array(z.string()) }), +}); /** Authoritative parent membership, independent of views and destination metadata. */ export class Subagents { @@ -69,11 +85,13 @@ export class Subagents { private mutation: Promise = Promise.resolve(); private readonly ready: Promise; private closed = false; + private legacyUpgrade = { complete: false, candidates: [] as string[] }; constructor( private readonly options: { home: string; serverId: string; + agents: Pick; runtime(authority: HostRuntimeAuthority): Promise; resolveAuthority(principalId: string): ReturnType; }, @@ -83,17 +101,99 @@ export class Subagents { } private async load(): Promise { - let records: RelationshipRecord[]; + let records: RelationshipRecord[] = []; + let existing = false; try { - records = z.array(RecordSchema).parse(JSON.parse(await readFile(this.path(), "utf8"))); + const ledger = LedgerSchema.parse(JSON.parse(await readFile(this.path(), "utf8"))); + records = ledger.records; + this.legacyUpgrade = ledger.legacyUpgrade; + existing = true; } catch (error) { - if ((error as NodeJS.ErrnoException).code === "ENOENT") return; - throw error; + if ((error as NodeJS.ErrnoException).code !== "ENOENT") throw error; } for (const record of records) { record.row.freshness = "stale"; this.records.set(record.row.id, record); } + if (!this.legacyUpgrade.complete) await this.upgradeLegacyMembership(existing); + } + initialize(): Promise { + return this.ready; + } + + private async upgradeLegacyMembership(existing: boolean): Promise { + if (!existing) { + const agents = await this.options.agents.list(); + const byId = new Map(agents.map((agent) => [agent.id, agent])); + const admission = await this.options.resolveAuthority("owner"); + const authority: z.infer = { + principalId: "owner", + sourceId: "owner", + permissionCeiling: admission.permissions.filter((permission) => + ["daemon.read", "workspace.read", "workspace.write"].includes(permission), + ), + }; + for (const child of agents) { + const parentId = getParentAgentIdFromLabels(child.labels); + const parent = parentId ? byId.get(parentId) : undefined; + if ( + !legacyManagedAgent(child) || + child.archivedAt || + !parent || + !legacyManagedAgent(parent) || + child.parent !== undefined + ) + continue; + if ( + [...this.records.values()].some( + (record) => + record.row.child.serverId === this.options.serverId && + record.row.child.agentId === child.id, + ) + ) + continue; + const id = createHash("sha256") + .update(JSON.stringify(["legacy-local-lifecycle", this.options.serverId, child.id])) + .digest("hex"); + const address = { serverId: this.options.serverId, agentId: parent.id }; + this.records.set(id, { + row: { + id, + parent: address, + child: { serverId: this.options.serverId, agentId: child.id }, + ...(child.workspaceId + ? { workspace: { serverId: this.options.serverId, workspaceId: child.workspaceId } } + : {}), + phase: "attached", + freshness: "stale", + observedAt: null, + error: null, + snapshot: buildStoredAgentPayload({ ...child, parent: address }, [child.provider]), + }, + provenance: "legacy-local-lifecycle", + creation: null, + fingerprint: null, + authority, + confirmed: "attached", + intent: { kind: "attached", revision: 0, authority }, + }); + this.legacyUpgrade.candidates.push(id); + } + // Persist the bounded candidate set before qualifying agent metadata. A restart + // resumes these candidates, never a new scan of subsequently edited labels. + await this.persist(); + } + for (const id of this.legacyUpgrade.candidates) { + const record = this.records.get(id)!; + const child = await this.options.agents.get(record.row.child.agentId); + if (!child) continue; + if (child.parent === undefined) + await this.options.agents.upsert({ ...child, parent: record.row.parent }); + else if (!child.parent || !sameAddress(child.parent, record.row.parent)) + throw new Error("Legacy child parent changed during membership upgrade"); + } + this.legacyUpgrade.complete = true; + await this.persist(); } private path(): string { return join(this.options.home, "subagents.json"); @@ -109,7 +209,10 @@ export class Subagents { } private async persist(): Promise { for (const record of this.records.values()) record.row.phase = relationshipPhase(record); - await writeJsonFileAtomic(this.path(), [...this.records.values()]); + await writeJsonFileAtomic(this.path(), { + records: [...this.records.values()], + legacyUpgrade: this.legacyUpgrade, + }); for (const changed of this.listeners) changed(); } async assertHostRemovable(serverId: string): Promise { @@ -262,6 +365,7 @@ export class Subagents { return existing; } const fresh: RelationshipRecord = { + provenance: "creation", row: { id, parent, @@ -309,10 +413,11 @@ export class Subagents { return running; } private creationKey(record: RelationshipRecord): string { - return record.creation.input.idempotencyKey ?? `subagent:${record.row.id}`; + return creationRequest(record).input.idempotencyKey ?? `subagent:${record.row.id}`; } private async createChild(record: RelationshipRecord): Promise { - const { row, creation } = record; + const { row } = record; + const creation = creationRequest(record); const creationIntent = record.intent; try { const runtime = await this.runtime(record.authority); @@ -655,7 +760,7 @@ export class Subagents { return null; } const observation = api.creations.observe({ - kind: record.creation.workspace ? "workspace" : "agent", + kind: creationRequest(record).workspace ? "workspace" : "agent", idempotencyKey: this.creationKey(record), }); let latest: CreationSnapshot | null = null; @@ -705,12 +810,13 @@ export class Subagents { const child = await runtime.agent(record.row.child); let current = await child.refresh(); if (!current) { + const creation = creationRequest(record); const destination = await runtime.client(record.row.child.serverId); try { if (record.intent !== intent) return; const cancellation = await destination.creations.cancel({ - ...(record.creation.workspace - ? { kind: "workspace" as const, workspaceId: record.creation.input.workspaceId! } + ...(creation.workspace + ? { kind: "workspace" as const, workspaceId: creation.input.workspaceId! } : { kind: "agent" as const }), idempotencyKey: this.creationKey(record), agentId: record.row.child.agentId, @@ -842,3 +948,16 @@ function terminalPhase(phase: RelationshipRecord["row"]["phase"]): boolean { phase === "detached" || phase === "archived" || phase === "forgotten" || phase === "canceled" ); } + +function legacyManagedAgent(agent: StoredAgentRecord): boolean { + return !agent.internal && !agent.owner; +} + +function creationRequest(record: RelationshipRecord): SubagentCreateInput { + if (record.provenance !== "creation" || !record.creation) + throw failure( + "agent_not_found", + "The adopted existing child is absent; there is no creation to replay or cancel", + ); + return record.creation; +} diff --git a/packages/server/src/server/websocket-server.ts b/packages/server/src/server/websocket-server.ts index 3b92970d62..da0e4178d3 100644 --- a/packages/server/src/server/websocket-server.ts +++ b/packages/server/src/server/websocket-server.ts @@ -755,6 +755,7 @@ export class VoiceAssistantWebSocketServer { this.subagents = new Subagents({ home: paseoHome, serverId, + agents: this.agentStorage, runtime: (authority) => this.createHostRuntime(authority), resolveAuthority: (principalId) => this.resolveRelationshipAdmission(principalId), }); @@ -766,13 +767,16 @@ export class VoiceAssistantWebSocketServer { this.logger.warn({ err: error }, "Failed to reconcile subagent host availability"), ); }); - this.stopAgentArchived = this.agentManager.onAgentArchived(async (agentId) => { - const record = await this.agentStorage.get(agentId); - await this.subagents.parentArchived( - { serverId, agentId }, - record?.workspaceId ? { serverId, workspaceId: record.workspaceId } : undefined, - ); - }); + this.stopAgentArchived = this.agentManager.onAgentArchived( + async (agentId) => { + const record = await this.agentStorage.get(agentId); + await this.subagents.parentArchived( + { serverId, agentId }, + record?.workspaceId ? { serverId, workspaceId: record.workspaceId } : undefined, + ); + }, + { ownsManagedChildren: true }, + ); this.passwordHash = auth?.password; this.worktreesRoot = daemonRuntimeConfig?.worktreesRoot; this.daemonConfigStore = daemonConfigStore; @@ -1158,6 +1162,10 @@ export class VoiceAssistantWebSocketServer { this.connectionLifecycle = "stopping"; } + public initializeRelationships(): Promise { + return this.subagents.initialize(); + } + public beginAcceptingConnections(): void { if (this.connectionLifecycle === "starting") { this.connectionLifecycle = "accepting"; @@ -2029,6 +2037,7 @@ export class VoiceAssistantWebSocketServer { connectedHosts: true, fleetSubagents: true, hostMultiplexing: true, + hostSessionScopes: true, usageSources: true, ownedSubscriptions: true, agentRequestReceipts: true, @@ -2632,15 +2641,30 @@ export class VoiceAssistantWebSocketServer { if (!link) { link = new DaemonLinkEndpoint({ identity: JSON.stringify([connection.principalId, connection.clientId]), + serverId: this.serverId, + assertScopeAdmission: () => { + if (connection.hubExecutionAgents) + throw Object.assign(new Error("Execution-scoped host acquisition is not supported"), { + code: "unsupported_capability", + }); + }, permissions: () => connection.session.getPermissions(), routed: () => connection.session.isRouted(), bufferedAmount: () => ws.bufferedAmount ?? null, delivery: connection.session.delivery, - admit: (socket, permissions) => - this.attachSocket(socket, undefined, undefined, false, { + admit: (socket, permissions, context) => { + if (connection.hubExecutionAgents) + throw new Error("Execution-scoped host acquisition is not supported"); + this.ownedSockets.add(socket); + if (context.kind === "scope") + this.ownedSocketSources.set(socket, connection.session.getHostSourceId()); + return this.attachSocket(socket, undefined, undefined, false, { principalId: connection.principalId, - permissions, - }), + permissions: permissions.filter((permission) => + connection.session.getPermissions().includes(permission), + ), + }); + }, }); this.hostLinks.set(ws, link); }