Reduce workspace, agent, and chat sync traffic (#2028)

* feat(sync): keep live data scoped and current

Only viewed chats receive live timeline rows, while directory state now uses one subscribed bootstrap followed by ordered deltas. Legacy clients and daemons retain their existing behavior through centralized compatibility gates.

* fix(sync): preserve selective delivery boundaries

Keep selective timeline capability app-owned, union viewed sets across shared sockets, and reconcile directory side effects from accepted state. Visibility and archive suppression now follow their existing authoritative boundaries.

* test(app): align browser fixtures with runtime contracts

* fix(sync): preserve mixed-client delivery guarantees

* fix(sync): finish paged history before advancing

* fix(sync): replay live deltas after refresh failures

* test(server): model socket capability lookup

* test(app): use platform shortcut for split-pane coverage

* fix(app): keep hidden timelines dormant

* test(app): normalize split-pane shortcut setup

* fix(app): preserve sync state across retries

* fix(app): preserve background sync state across races

* fix directory bootstrap reconciliation edge cases

* fix selective timeline compatibility races

* fix superseded directory sync races

* fix(sync): centralize directory replica ordering

Own agent and workspace refresh transactions in HostRuntime so reconnect epochs, buffered updates, and lifecycle invalidation share one ordering boundary. Route timeline metadata and mixed-capability stream delivery through their authoritative runtime sources.

* fix(app): keep timeline requests runtime-scoped

* fix(app): close directory bootstrap races
This commit is contained in:
Mohamed Boudra
2026-07-17 17:03:34 +02:00
committed by GitHub
parent 6c99efae52
commit 1977d330ed
57 changed files with 5783 additions and 1343 deletions
@@ -0,0 +1,86 @@
import { describe, expect, it } from "vitest";
import type { DaemonClient, FetchAgentsEntry } from "@getpaseo/client/internal/daemon-client";
import type { AgentSnapshotPayload } from "@getpaseo/protocol/messages";
import { useSessionStore } from "@/stores/session-store";
import { AgentDirectoryReplica } from "./agent-replica";
function payload(title: string): AgentSnapshotPayload {
return {
id: "agent",
provider: "codex",
cwd: "/repo",
model: null,
createdAt: "2026-07-17T00:00:00.000Z",
updatedAt: "2026-07-17T00:01:00.000Z",
lastUserMessageAt: null,
status: "idle",
capabilities: {
supportsStreaming: true,
supportsSessionPersistence: true,
supportsDynamicModes: true,
supportsMcpServers: true,
supportsReasoningStream: true,
supportsToolInvocations: true,
},
currentModeId: null,
availableModes: [],
pendingPermissions: [],
persistence: null,
title,
labels: {},
};
}
function entry(agent: AgentSnapshotPayload): FetchAgentsEntry {
return {
agent,
project: {
projectKey: "/repo",
projectName: "repo",
checkout: {
cwd: "/repo",
isGit: false,
currentBranch: null,
remoteUrl: null,
worktreeRoot: null,
isPaseoOwnedWorktree: false,
mainRepoRoot: null,
},
},
};
}
describe("AgentDirectoryReplica", () => {
it("keeps membership authoritative across remove, stale timeline, and re-add", () => {
const serverId = "agent-replica";
const store = useSessionStore.getState();
store.initializeSession(serverId, null as unknown as DaemonClient);
const replica = new AgentDirectoryReplica(serverId, () => undefined);
replica.commitSnapshot([entry(payload("directory"))], []);
const directoryPlacement = useSessionStore
.getState()
.sessions[serverId]?.agents.get("agent")?.projectPlacement;
expect(directoryPlacement).toBeDefined();
const staleToken = replica.captureTimeline("agent");
replica.remove("agent");
expect(replica.submitTimelineAgent(staleToken, payload("stale"))).toBe(false);
expect(useSessionStore.getState().sessions[serverId]?.agents.has("agent")).toBe(false);
replica.applyDelta({
kind: "upsert",
agent: payload("re-added"),
project: entry(payload("x")).project,
});
expect(replica.submitTimelineAgent(staleToken, payload("still stale"))).toBe(false);
const currentToken = replica.captureTimeline("agent");
expect(replica.submitTimelineAgent(currentToken, payload("current"))).toBe(true);
expect(useSessionStore.getState().sessions[serverId]?.agents.get("agent")?.title).toBe(
"current",
);
expect(
useSessionStore.getState().sessions[serverId]?.agents.get("agent")?.projectPlacement,
).toEqual(directoryPlacement);
store.clearSession(serverId);
});
});
@@ -0,0 +1,121 @@
import type { FetchAgentsEntry } from "@getpaseo/client/internal/daemon-client";
import type { AgentSnapshotPayload } from "@getpaseo/protocol/messages";
import { clearArchiveAgentPending } from "@/hooks/use-archive-agent";
import { queryClient } from "@/data/query-client";
import { useSessionStore, type Agent } from "@/stores/session-store";
import { normalizeAgentSnapshot } from "@/utils/agent-snapshots";
import {
applyAgentDirectoryDelta,
type AgentDirectoryDelta,
removeAgentDirectoryReplica,
replaceAgentPendingPermissions,
replaceFetchedAgentDirectory,
upsertAgentReplica,
} from "@/utils/agent-directory-sync";
import { reconcileAgentDirectory } from "@/utils/agent-directory-reconciliation";
import { applyLegacyDaemonWorkspaceOwnership } from "@/workspace/legacy-daemon-workspaces";
export interface AgentLifecycleToken {
readonly agentId: string;
readonly version: number;
}
export class AgentDirectoryReplica {
private readonly lifecycleVersions = new Map<string, number>();
private readonly members = new Set<string>();
constructor(
private readonly serverId: string,
private readonly onStoppedRunning: (agentId: string) => void,
) {}
captureTimeline(agentId: string): AgentLifecycleToken {
return { agentId, version: this.lifecycleVersions.get(agentId) ?? 0 };
}
submitTimelineAgent(token: AgentLifecycleToken, payload: AgentSnapshotPayload): boolean {
if (
!this.members.has(token.agentId) ||
token.version !== (this.lifecycleVersions.get(token.agentId) ?? 0)
) {
return false;
}
const existing = useSessionStore.getState().sessions[this.serverId]?.agents.get(token.agentId);
const timelineAgent = applyLegacyDaemonWorkspaceOwnership({
serverId: this.serverId,
agent: normalizeAgentSnapshot(payload, this.serverId),
});
const normalized: Agent = {
...timelineAgent,
projectPlacement: timelineAgent.projectPlacement ?? existing?.projectPlacement,
};
const accepted = upsertAgentReplica(this.serverId, normalized);
replaceAgentPendingPermissions(this.serverId, accepted);
useSessionStore.getState().setAgentLastActivity(accepted.id, accepted.lastActivityAt);
if (accepted.archivedAt) {
clearArchiveAgentPending({ queryClient, serverId: this.serverId, agentId: accepted.id });
}
return true;
}
applyDelta(delta: AgentDirectoryDelta): void {
const before = this.members.has(delta.kind === "remove" ? delta.agentId : delta.agent.id);
const result = applyAgentDirectoryDelta({ serverId: this.serverId, delta });
if (delta.kind === "remove") {
this.members.delete(delta.agentId);
this.advance(delta.agentId);
} else {
this.members.add(delta.agent.id);
if (!before) this.advance(delta.agent.id);
}
if (result.stoppedRunning) this.onStoppedRunning(result.agentId);
}
commitSnapshot(
entries: FetchAgentsEntry[],
deltas: readonly AgentDirectoryDelta[],
): Map<string, Agent> {
const previous = useSessionStore.getState().sessions[this.serverId]?.agents ?? new Map();
const reconciled = reconcileAgentDirectory({ previous, snapshot: entries, deltas });
const nextIds = new Set(reconciled.entries.map((entry) => entry.agent.id));
for (const agentId of this.members) {
if (!nextIds.has(agentId)) this.advance(agentId);
}
for (const agentId of nextIds) {
if (!this.members.has(agentId)) this.advance(agentId);
}
for (const agentId of previous.keys()) {
if (!nextIds.has(agentId)) removeAgentDirectoryReplica(this.serverId, agentId);
}
this.members.clear();
for (const agentId of nextIds) this.members.add(agentId);
const { agents } = replaceFetchedAgentDirectory({
serverId: this.serverId,
entries: reconciled.entries,
});
for (const agentId of reconciled.stoppedRunningAgentIds) this.onStoppedRunning(agentId);
return agents;
}
archive(agentId: string, archivedAt: string): void {
this.advance(agentId);
useSessionStore.getState().setAgents(this.serverId, (current) => {
const agent = current.get(agentId);
if (!agent) return current;
const next = new Map(current);
next.set(agentId, { ...agent, archivedAt: new Date(archivedAt) });
return next;
});
clearArchiveAgentPending({ queryClient, serverId: this.serverId, agentId });
}
remove(agentId: string): void {
this.members.delete(agentId);
this.advance(agentId);
removeAgentDirectoryReplica(this.serverId, agentId);
}
private advance(agentId: string): void {
this.lifecycleVersions.set(agentId, (this.lifecycleVersions.get(agentId) ?? 0) + 1);
}
}
@@ -0,0 +1,113 @@
import { afterEach, describe, expect, it } from "vitest";
import type { DaemonClient } from "@getpaseo/client/internal/daemon-client";
import { useSessionStore } from "@/stores/session-store";
import { DirectoryRefreshSupersededError, DirectorySync } from "./index";
class FakeDirectoryClient {
fetchAgentsCalls = 0;
fetchWorkspacesCalls = 0;
on(): () => void {
return () => undefined;
}
async fetchAgents(): Promise<Awaited<ReturnType<DaemonClient["fetchAgents"]>>> {
this.fetchAgentsCalls += 1;
return {
requestId: "agents",
entries: [],
pageInfo: { hasMore: false, nextCursor: null, prevCursor: null },
};
}
async fetchWorkspaces(): Promise<Awaited<ReturnType<DaemonClient["fetchWorkspaces"]>>> {
this.fetchWorkspacesCalls += 1;
return {
requestId: "workspaces",
entries: [],
emptyProjects: [],
pageInfo: { hasMore: false, nextCursor: null, prevCursor: null },
};
}
}
const serverIds = new Set<string>();
function createDirectory(serverId: string): {
client: FakeDirectoryClient;
directory: DirectorySync;
} {
serverIds.add(serverId);
const client = new FakeDirectoryClient();
const directory = new DirectorySync(serverId, {
drainQueuedAgentMessage: () => undefined,
markAgentLoading: () => undefined,
markAgentReady: () => undefined,
markAgentError: () => undefined,
});
directory.connectionChanged({
client: client as unknown as DaemonClient,
status: "online",
source: { clientGeneration: 1, connectionEpoch: 1 },
});
return { client, directory };
}
afterEach(() => {
for (const serverId of serverIds) useSessionStore.getState().clearSession(serverId);
serverIds.clear();
});
describe("DirectorySync session readiness", () => {
it("waits for workspace capability metadata before choosing the workspace protocol", async () => {
const serverId = "workspace-metadata";
const { client, directory } = createDirectory(serverId);
const refresh = directory.refreshWorkspaces({ subscribe: true });
await Promise.resolve();
expect(client.fetchWorkspacesCalls).toBe(0);
const store = useSessionStore.getState();
store.initializeSession(serverId, client as unknown as DaemonClient, 1);
await Promise.resolve();
expect(client.fetchWorkspacesCalls).toBe(0);
store.updateSessionServerInfo(serverId, {
serverId,
hostname: null,
version: "test",
features: { workspaceMultiplicity: true },
});
await refresh;
expect(client.fetchWorkspacesCalls).toBe(1);
expect(useSessionStore.getState().sessions[serverId]?.hasHydratedWorkspaces).toBe(true);
directory.dispose();
});
it("rejects a session wait on disconnect so the reconnect can refresh", async () => {
const serverId = "session-wait-reconnect";
const { client, directory } = createDirectory(serverId);
const staleRefresh = directory.refreshAgents();
await Promise.resolve();
directory.connectionChanged({
client: null,
status: "offline",
source: { clientGeneration: 1, connectionEpoch: 1 },
});
await expect(staleRefresh).rejects.toBeInstanceOf(DirectoryRefreshSupersededError);
directory.connectionChanged({
client: client as unknown as DaemonClient,
status: "online",
source: { clientGeneration: 1, connectionEpoch: 2 },
});
const currentRefresh = directory.refreshAgents();
useSessionStore.getState().initializeSession(serverId, client as unknown as DaemonClient, 1);
await currentRefresh;
expect(client.fetchAgentsCalls).toBe(1);
directory.dispose();
});
});
@@ -0,0 +1,416 @@
import type {
DaemonClient,
FetchAgentsEntry,
FetchAgentsOptions,
} from "@getpaseo/client/internal/daemon-client";
import { fetchAgentTimelineOnce } from "@/timeline/fetch-agent-timeline-once";
import {
normalizeEmptyProjectDescriptor,
normalizeWorkspaceDescriptor,
useSessionStore,
type Agent,
} from "@/stores/session-store";
import {
readLegacyDaemonWorkspaceDirectory,
buildLegacyWorkspaces,
shouldUseLegacyDaemonWorkspaceDirectory,
stampLegacyWorkspaceIds,
} from "@/workspace/legacy-daemon-workspaces";
import type { AgentDirectoryDelta } from "@/utils/agent-directory-sync";
import { AgentDirectoryReplica } from "./agent-replica";
import {
WorkspaceDirectoryReplica,
type WorkspaceDirectoryDelta,
type WorkspaceDirectorySnapshot,
} from "./workspace-replica";
import {
DirectoryTransactionOwner,
type DirectorySourceToken,
type DirectoryTransaction,
} from "./transaction";
const PAGE_LIMIT = 200;
const AGENT_SORT: NonNullable<FetchAgentsOptions["sort"]> = [
{ key: "updated_at", direction: "desc" },
];
interface AgentSnapshot {
entries: FetchAgentsEntry[];
subscriptionId: string | null;
legacy: boolean;
}
export interface DirectoryConnection {
client: DaemonClient | null;
status: "online" | "offline";
source: DirectorySourceToken;
}
export interface RefreshAgentDirectoryInput {
filter?: FetchAgentsOptions["filter"];
subscribe?: FetchAgentsOptions["subscribe"];
page?: FetchAgentsOptions["page"];
}
export interface RefreshAgentDirectoryResult {
agents: Map<string, Agent>;
subscriptionId: string | null;
}
export class DirectorySync {
private readonly agentTransactions = new DirectoryTransactionOwner<
AgentSnapshot,
AgentDirectoryDelta
>();
private readonly workspaceTransactions = new DirectoryTransactionOwner<
WorkspaceDirectorySnapshot,
WorkspaceDirectoryDelta
>();
private readonly agents: AgentDirectoryReplica;
private readonly workspaces: WorkspaceDirectoryReplica;
private connection: DirectoryConnection = {
client: null,
status: "offline",
source: { clientGeneration: 0, connectionEpoch: 0 },
};
private unsubscribe: (() => void) | null = null;
private readonly abortSessionWaits = new Set<() => void>();
constructor(
private readonly serverId: string,
private readonly callbacks: {
drainQueuedAgentMessage: (agentId: string) => void;
markAgentLoading: () => void;
markAgentReady: () => void;
markAgentError: (error: string) => void;
},
) {
this.agents = new AgentDirectoryReplica(serverId, callbacks.drainQueuedAgentMessage);
this.workspaces = new WorkspaceDirectoryReplica(serverId);
}
connectionChanged(connection: DirectoryConnection): boolean {
const changed =
this.connection.client !== connection.client ||
this.connection.source.clientGeneration !== connection.source.clientGeneration ||
this.connection.source.connectionEpoch !== connection.source.connectionEpoch;
const wentOffline = this.connection.status === "online" && connection.status === "offline";
if (!changed && !wentOffline) {
this.connection = connection;
return false;
}
this.flushAbortedTransactions();
this.unsubscribe?.();
this.unsubscribe = null;
this.connection = connection;
this.abortPendingSessionWaits();
if (!connection.client || connection.status !== "online") return true;
const client = connection.client;
const source = connection.source;
const subscriptions = [
client.on("agent_update", (message) => {
if (message.type !== "agent_update" || !this.isCurrent(client, source)) return;
if (!this.agentTransactions.record(source, message.payload))
this.agents.applyDelta(message.payload);
}),
client.on("workspace_update", (message) => {
if (message.type !== "workspace_update" || !this.isCurrent(client, source)) return;
if (!this.workspaceTransactions.record(source, message.payload)) {
this.workspaces.applyDelta(message.payload);
}
}),
client.on("agent_deleted", (message) => {
if (message.type === "agent_deleted" && this.isCurrent(client, source)) {
this.agents.remove(message.payload.agentId);
}
}),
client.on("agent_archived", (message) => {
if (message.type === "agent_archived" && this.isCurrent(client, source)) {
this.agents.archive(message.payload.agentId, message.payload.archivedAt);
}
}),
];
this.unsubscribe = () => {
for (const unsubscribe of subscriptions) unsubscribe();
};
return true;
}
dispose(): void {
this.flushAbortedTransactions();
this.abortPendingSessionWaits();
this.unsubscribe?.();
this.unsubscribe = null;
}
async fetchTimeline(
agentId: string,
request: Parameters<DaemonClient["fetchAgentTimeline"]>[1],
): Promise<Awaited<ReturnType<DaemonClient["fetchAgentTimeline"]>>> {
const { client } = this.requireOnline();
const token = this.agents.captureTimeline(agentId);
const page = await fetchAgentTimelineOnce(client, agentId, request);
if (page.agent) this.agents.submitTimelineAgent(token, page.agent);
return page;
}
async refreshAgents(
input: RefreshAgentDirectoryInput = {},
): Promise<RefreshAgentDirectoryResult> {
const { client, source } = this.requireOnline();
const transaction = this.agentTransactions.begin(source, () => ({
entries: [],
subscriptionId: null,
legacy: false,
}));
this.callbacks.markAgentLoading();
try {
await this.waitForSession(client, source);
const session = useSessionStore.getState().sessions[this.serverId];
if (!input.filter && shouldUseLegacyDaemonWorkspaceDirectory(session?.serverInfo)) {
const directory = await readLegacyDaemonWorkspaceDirectory({
client,
subscribe: input.subscribe,
page: input.page,
});
if (
!directory ||
!this.agentTransactions.isCurrent(transaction) ||
!this.isCurrent(client, source)
) {
throw new DirectoryRefreshSupersededError("legacy fetch no longer current");
}
transaction.snapshot.entries.push(...stampLegacyWorkspaceIds(directory.entries));
transaction.snapshot.subscriptionId = directory.subscriptionId;
transaction.snapshot.legacy = true;
} else {
await this.fetchAgents(client, source, transaction, input);
}
if (!this.isCurrent(client, source) || !this.hasMatchingSession(client, source)) {
throw new DirectoryRefreshSupersededError("agent completion no longer current");
}
const completion = this.agentTransactions.complete(transaction);
if (completion.kind === "stale") {
throw new DirectoryRefreshSupersededError("agent completion was superseded");
}
if (completion.snapshot.legacy) {
const store = useSessionStore.getState();
store.setWorkspaces(this.serverId, buildLegacyWorkspaces(completion.snapshot.entries));
store.setEmptyProjects(this.serverId, []);
store.setHasHydratedWorkspaces(this.serverId, true);
}
const deltas = completion.snapshot.legacy
? completion.deltas.map((delta) =>
delta.kind === "upsert"
? { ...delta, agent: { ...delta.agent, workspaceId: delta.agent.cwd } }
: delta,
)
: completion.deltas;
const agents = this.agents.commitSnapshot(completion.snapshot.entries, deltas);
this.callbacks.markAgentReady();
return { agents, subscriptionId: completion.snapshot.subscriptionId };
} catch (error) {
const deltas = this.agentTransactions.fail(transaction);
if (deltas) for (const delta of deltas) this.agents.applyDelta(delta);
if (!(error instanceof DirectoryRefreshSupersededError)) {
this.callbacks.markAgentError(error instanceof Error ? error.message : String(error));
}
throw error;
}
}
async refreshWorkspaces(input?: { subscribe?: boolean }): Promise<void> {
const { client, source } = this.requireOnline();
const transaction = this.workspaceTransactions.begin(source, () => ({
workspaces: new Map(),
emptyProjects: new Map(),
}));
try {
await this.waitForSessionMetadata(client, source);
const serverInfo = useSessionStore.getState().sessions[this.serverId]?.serverInfo;
if (serverInfo?.features?.workspaceMultiplicity !== true) {
const deltas = this.workspaceTransactions.fail(transaction);
if (deltas) for (const delta of deltas) this.workspaces.applyDelta(delta);
return;
}
await this.fetchWorkspaceSnapshot(client, source, transaction, input?.subscribe === true);
if (!this.isCurrent(client, source) || !this.hasMatchingSession(client, source)) {
throw new DirectoryRefreshSupersededError("workspace completion no longer current");
}
const completion = this.workspaceTransactions.complete(transaction);
if (completion.kind === "stale") {
throw new DirectoryRefreshSupersededError("workspace completion was superseded");
}
this.workspaces.commitSnapshot(completion.snapshot, completion.deltas);
} catch (error) {
const deltas = this.workspaceTransactions.fail(transaction);
if (deltas) for (const delta of deltas) this.workspaces.applyDelta(delta);
throw error;
}
}
private async fetchWorkspaceSnapshot(
client: DaemonClient,
source: DirectorySourceToken,
transaction: DirectoryTransaction<WorkspaceDirectorySnapshot, WorkspaceDirectoryDelta>,
initialSubscribe: boolean,
): Promise<void> {
let cursor: string | null = null;
let subscribe = initialSubscribe;
while (true) {
const payload = await client.fetchWorkspaces({
sort: [{ key: "activity_at", direction: "desc" }],
...(subscribe ? { subscribe: {} } : {}),
page: cursor ? { limit: PAGE_LIMIT, cursor } : { limit: PAGE_LIMIT },
});
this.assertWorkspaceTransactionCurrent(client, source, transaction);
for (const entry of payload.entries) {
const workspace = normalizeWorkspaceDescriptor(entry);
transaction.snapshot.workspaces.set(workspace.id, workspace);
}
for (const entry of payload.emptyProjects ?? []) {
const project = normalizeEmptyProjectDescriptor(entry);
transaction.snapshot.emptyProjects.set(project.projectId, project);
}
if (!payload.pageInfo.hasMore || !payload.pageInfo.nextCursor) return;
cursor = payload.pageInfo.nextCursor;
subscribe = false;
}
}
async refreshAll(): Promise<void> {
await Promise.all([this.refreshAgents(), this.refreshWorkspaces({ subscribe: true })]);
}
private async fetchAgents(
client: DaemonClient,
source: DirectorySourceToken,
transaction: DirectoryTransaction<AgentSnapshot, AgentDirectoryDelta>,
input: RefreshAgentDirectoryInput,
): Promise<void> {
let cursor = input.page?.cursor ?? null;
let subscribe = input.subscribe;
while (true) {
const limit = input.page?.limit ?? PAGE_LIMIT;
const payload = await client.fetchAgents({
...(input.filter ? { filter: input.filter } : { scope: "active" as const }),
sort: AGENT_SORT,
...(subscribe ? { subscribe } : {}),
page: cursor ? { limit, cursor } : { limit },
});
this.assertAgentTransactionCurrent(client, source, transaction);
transaction.snapshot.entries.push(...payload.entries);
transaction.snapshot.subscriptionId ??= payload.subscriptionId ?? null;
const pageInfo = payload.pageInfo as {
hasMore?: boolean;
hasMoreAfter?: boolean;
nextCursor?: string | null;
afterCursor?: string | null;
};
const hasMore = pageInfo.hasMore ?? pageInfo.hasMoreAfter ?? false;
const nextCursor = pageInfo.nextCursor ?? pageInfo.afterCursor ?? null;
if (!hasMore || !nextCursor) break;
cursor = nextCursor;
subscribe = undefined;
}
}
private assertAgentTransactionCurrent(
client: DaemonClient,
source: DirectorySourceToken,
transaction: DirectoryTransaction<AgentSnapshot, AgentDirectoryDelta>,
): void {
if (!this.agentTransactions.isCurrent(transaction) || !this.isCurrent(client, source)) {
throw new DirectoryRefreshSupersededError("agent page no longer current");
}
}
private assertWorkspaceTransactionCurrent(
client: DaemonClient,
source: DirectorySourceToken,
transaction: DirectoryTransaction<WorkspaceDirectorySnapshot, WorkspaceDirectoryDelta>,
): void {
if (!this.workspaceTransactions.isCurrent(transaction) || !this.isCurrent(client, source)) {
throw new DirectoryRefreshSupersededError("workspace fetch no longer current");
}
}
private requireOnline(): { client: DaemonClient; source: DirectorySourceToken } {
if (!this.connection.client || this.connection.status !== "online") {
throw new Error(`Host ${this.serverId} is not connected`);
}
return { client: this.connection.client, source: this.connection.source };
}
private isCurrent(client: DaemonClient, source: DirectorySourceToken): boolean {
return (
this.connection.client === client &&
this.connection.status === "online" &&
this.connection.source.clientGeneration === source.clientGeneration &&
this.connection.source.connectionEpoch === source.connectionEpoch
);
}
private async waitForSession(client: DaemonClient, source: DirectorySourceToken): Promise<void> {
await this.waitForSessionState(client, source, () => this.hasMatchingSession(client, source));
}
private async waitForSessionMetadata(
client: DaemonClient,
source: DirectorySourceToken,
): Promise<void> {
await this.waitForSessionState(client, source, () => {
const session = useSessionStore.getState().sessions[this.serverId];
return this.hasMatchingSession(client, source) && session?.serverInfo !== null;
});
}
private async waitForSessionState(
client: DaemonClient,
source: DirectorySourceToken,
matches: () => boolean,
): Promise<void> {
if (matches()) return;
await new Promise<void>((resolve, reject) => {
let settled = false;
let unsubscribe: () => void = () => undefined;
const finish = (result: "ready" | "aborted") => {
if (settled) return;
settled = true;
unsubscribe();
this.abortSessionWaits.delete(abort);
if (result === "ready") resolve();
else reject(new DirectoryRefreshSupersededError("session wait no longer current"));
};
const abort = () => finish("aborted");
const check = () => {
if (matches()) {
finish("ready");
} else if (!this.isCurrent(client, source)) {
finish("aborted");
}
};
this.abortSessionWaits.add(abort);
unsubscribe = useSessionStore.subscribe(check);
check();
});
}
private hasMatchingSession(client: DaemonClient, source: DirectorySourceToken): boolean {
const session = useSessionStore.getState().sessions[this.serverId];
return session?.client === client && session.clientGeneration === source.clientGeneration;
}
private flushAbortedTransactions(): void {
for (const delta of this.agentTransactions.abort()) this.agents.applyDelta(delta);
for (const delta of this.workspaceTransactions.abort()) this.workspaces.applyDelta(delta);
}
private abortPendingSessionWaits(): void {
for (const abort of this.abortSessionWaits) abort();
}
}
export class DirectoryRefreshSupersededError extends Error {}
export type { DirectorySourceToken } from "./transaction";
@@ -0,0 +1,100 @@
import { describe, expect, it } from "vitest";
import { DirectoryTransactionOwner, type DirectorySourceToken } from "./transaction";
const epoch = (connectionEpoch: number): DirectorySourceToken => ({
clientGeneration: 1,
connectionEpoch,
});
type Event = "begin-new" | "first" | "second" | "finish-old" | "finish-new";
function permutations(events: Event[]): Event[][] {
if (events.length === 0) return [[]];
return events.flatMap((event, index) =>
permutations(events.filter((_, candidate) => candidate !== index)).map((tail) =>
[event].concat(tail),
),
);
}
describe("DirectoryTransactionOwner", () => {
it("applies every bounded two-refresh interleaving exactly once", () => {
const schedules = permutations([
"begin-new",
"first",
"second",
"finish-old",
"finish-new",
]).filter((events) => events.indexOf("begin-new") < events.indexOf("finish-new"));
expect(schedules).toHaveLength(60);
for (const events of schedules) {
const owner = new DirectoryTransactionOwner<null, string>();
const old = owner.begin(epoch(1), () => null);
let current = old;
const applied: string[] = [];
for (const event of events) {
if (event === "begin-new") current = owner.begin(epoch(1), () => null);
if (event === "first" || event === "second") {
if (!owner.record(epoch(1), event)) applied.push(event);
}
if (event === "finish-old") {
const result = owner.complete(old);
if (result.kind === "current") applied.push(...result.deltas);
}
if (event === "finish-new") {
const result = owner.complete(current);
if (result.kind === "current") applied.push(...result.deltas);
}
}
expect(applied.sort(), JSON.stringify(events)).toEqual(["first", "second"]);
}
});
it.each([
["success", false, false, ["delta"]],
["failure", true, false, ["delta"]],
["superseded success", false, true, ["delta"]],
["superseded failure", true, true, ["delta"]],
] as const)("preserves buffered deltas through %s", (_name, fail, supersede, expected) => {
const owner = new DirectoryTransactionOwner<string[], string>();
const first = owner.begin(epoch(1), () => []);
owner.record(epoch(1), "delta");
const current = supersede ? owner.begin(epoch(1), () => []) : first;
let deltas: readonly string[] | null;
if (fail) {
deltas = owner.fail(current);
} else {
const completion = owner.complete(current);
deltas = completion.kind === "current" ? completion.deltas : null;
}
expect(deltas).toEqual(expected);
if (supersede) expect(owner.complete(first)).toEqual({ kind: "stale" });
});
it("makes a stale completion inert", () => {
const owner = new DirectoryTransactionOwner<null, string>();
const first = owner.begin(epoch(1), () => null);
const second = owner.begin(epoch(1), () => null);
expect(owner.complete(first)).toEqual({ kind: "stale" });
expect(owner.isCurrent(second)).toBe(true);
});
it("isolates reconnect epochs", () => {
const owner = new DirectoryTransactionOwner<null, string>();
const disconnected = owner.begin(epoch(1), () => null);
owner.record(epoch(1), "old");
expect(owner.abort()).toEqual(["old"]);
const reconnected = owner.begin(epoch(2), () => null);
owner.record(epoch(2), "new");
expect(owner.complete(disconnected)).toEqual({ kind: "stale" });
expect(owner.complete(reconnected)).toEqual({
kind: "current",
snapshot: null,
deltas: ["new"],
});
});
});
@@ -0,0 +1,70 @@
export interface DirectorySourceToken {
clientGeneration: number;
connectionEpoch: number;
}
export interface DirectoryTransaction<TSnapshot, TDelta> {
readonly source: DirectorySourceToken;
readonly snapshot: TSnapshot;
readonly deltas: TDelta[];
}
export type DirectoryTransactionCompletion<TSnapshot, TDelta> =
| { kind: "stale" }
| { kind: "current"; snapshot: TSnapshot; deltas: readonly TDelta[] };
function isSameSource(left: DirectorySourceToken, right: DirectorySourceToken): boolean {
return (
left.clientGeneration === right.clientGeneration &&
left.connectionEpoch === right.connectionEpoch
);
}
/** Owns the shared begin/record/supersede/complete lifecycle for a directory replica. */
export class DirectoryTransactionOwner<TSnapshot, TDelta> {
private current: DirectoryTransaction<TSnapshot, TDelta> | null = null;
begin(
source: DirectorySourceToken,
createSnapshot: () => TSnapshot,
): DirectoryTransaction<TSnapshot, TDelta> {
const transaction: DirectoryTransaction<TSnapshot, TDelta> = {
source,
snapshot: createSnapshot(),
deltas:
this.current && isSameSource(this.current.source, source) ? [...this.current.deltas] : [],
};
this.current = transaction;
return transaction;
}
record(source: DirectorySourceToken, delta: TDelta): boolean {
if (!this.current || !isSameSource(this.current.source, source)) return false;
this.current.deltas.push(delta);
return true;
}
isCurrent(transaction: DirectoryTransaction<TSnapshot, TDelta>): boolean {
return this.current === transaction;
}
complete(
transaction: DirectoryTransaction<TSnapshot, TDelta>,
): DirectoryTransactionCompletion<TSnapshot, TDelta> {
if (this.current !== transaction) return { kind: "stale" };
this.current = null;
return { kind: "current", snapshot: transaction.snapshot, deltas: transaction.deltas };
}
fail(transaction: DirectoryTransaction<TSnapshot, TDelta>): readonly TDelta[] | null {
const completion = this.complete(transaction);
return completion.kind === "current" ? completion.deltas : null;
}
abort(): readonly TDelta[] {
if (!this.current) return [];
const deltas = [...this.current.deltas];
this.current = null;
return deltas;
}
}
@@ -0,0 +1,57 @@
import { expect, it } from "vitest";
import type { DaemonClient } from "@getpaseo/client/internal/daemon-client";
import type { WorkspaceDescriptorPayload } from "@getpaseo/protocol/messages";
import {
normalizeEmptyProjectDescriptor,
normalizeWorkspaceDescriptor,
useSessionStore,
} from "@/stores/session-store";
import { WorkspaceDirectoryReplica } from "./workspace-replica";
function workspace(id: string, projectId = "project"): WorkspaceDescriptorPayload {
return {
id,
projectId,
projectDisplayName: projectId,
projectRootPath: `/repo/${projectId}`,
workspaceDirectory: `/repo/${projectId}/${id}`,
projectKind: "git",
workspaceKind: "worktree",
name: id,
title: id,
status: "done",
activityAt: null,
statusEnteredAt: null,
archivingAt: null,
diffStat: null,
scripts: [],
};
}
it("commits workspace and project-parent state with filtered removals", () => {
const serverId = "workspace-replica";
const store = useSessionStore.getState();
store.initializeSession(serverId, null as unknown as DaemonClient);
const replica = new WorkspaceDirectoryReplica(serverId);
const empty = normalizeEmptyProjectDescriptor({
projectId: "empty",
projectDisplayName: "Empty",
projectRootPath: "/repo/empty",
projectKind: "git",
});
replica.commitSnapshot(
{
workspaces: new Map([
["kept", normalizeWorkspaceDescriptor(workspace("kept"))],
["filtered", normalizeWorkspaceDescriptor(workspace("filtered", "filtered-project"))],
]),
emptyProjects: new Map([[empty.projectId, empty]]),
},
[{ kind: "remove", id: "filtered", removedProjectId: "filtered-project" }],
);
const session = useSessionStore.getState().sessions[serverId];
expect(Array.from(session?.workspaces.keys() ?? [])).toEqual(["kept"]);
expect(Array.from(session?.emptyProjects.keys() ?? [])).toEqual(["empty"]);
store.clearSession(serverId);
});
@@ -0,0 +1,93 @@
import type { SessionOutboundMessage } from "@getpaseo/protocol/messages";
import {
normalizeEmptyProjectDescriptor,
normalizeWorkspaceDescriptor,
useSessionStore,
type EmptyProjectDescriptor,
type WorkspaceDescriptor,
} from "@/stores/session-store";
import { useWorkspaceSetupStore } from "@/stores/workspace-setup-store";
import {
clearWorkspaceArchivePending,
shouldSuppressWorkspaceForLocalArchive,
} from "@/contexts/session-workspace-upserts";
export type WorkspaceDirectoryDelta = Extract<
SessionOutboundMessage,
{ type: "workspace_update" }
>["payload"];
export interface WorkspaceDirectorySnapshot {
workspaces: Map<string, WorkspaceDescriptor>;
emptyProjects: Map<string, EmptyProjectDescriptor>;
}
export class WorkspaceDirectoryReplica {
constructor(private readonly serverId: string) {}
applyDelta(delta: WorkspaceDirectoryDelta): void {
const state = this.reconcile(this.read(), [delta]);
this.commit(state, delta.kind === "remove" ? [delta.id] : []);
}
commitSnapshot(
snapshot: WorkspaceDirectorySnapshot,
deltas: readonly WorkspaceDirectoryDelta[],
): void {
const removedWorkspaceIds = deltas.flatMap((delta) =>
delta.kind === "remove" ? [delta.id] : [],
);
this.commit(this.reconcile(snapshot, deltas), removedWorkspaceIds);
}
private read(): WorkspaceDirectorySnapshot {
const session = useSessionStore.getState().sessions[this.serverId];
return {
workspaces: new Map(session?.workspaces),
emptyProjects: new Map(session?.emptyProjects),
};
}
private reconcile(
snapshot: WorkspaceDirectorySnapshot,
deltas: readonly WorkspaceDirectoryDelta[],
): WorkspaceDirectorySnapshot {
const workspaces = new Map(snapshot.workspaces);
const emptyProjects = new Map(snapshot.emptyProjects);
for (const [workspaceId, workspace] of workspaces) {
if (shouldSuppressWorkspaceForLocalArchive({ serverId: this.serverId, workspace })) {
workspaces.delete(workspaceId);
}
}
for (const delta of deltas) {
if (delta.kind === "remove") {
workspaces.delete(delta.id);
if (delta.emptyProject) {
const project = normalizeEmptyProjectDescriptor(delta.emptyProject);
emptyProjects.set(project.projectId, project);
}
if (delta.removedProjectId) emptyProjects.delete(delta.removedProjectId);
continue;
}
const workspace = normalizeWorkspaceDescriptor(delta.workspace);
if (shouldSuppressWorkspaceForLocalArchive({ serverId: this.serverId, workspace })) {
workspaces.delete(workspace.id);
} else {
workspaces.set(workspace.id, workspace);
emptyProjects.delete(workspace.projectId);
}
}
return { workspaces, emptyProjects };
}
private commit(snapshot: WorkspaceDirectorySnapshot, removedWorkspaceIds: string[]): void {
const store = useSessionStore.getState();
store.setWorkspaces(this.serverId, snapshot.workspaces);
store.setEmptyProjects(this.serverId, snapshot.emptyProjects.values());
store.setHasHydratedWorkspaces(this.serverId, true);
for (const workspaceId of removedWorkspaceIds) {
clearWorkspaceArchivePending({ serverId: this.serverId, workspaceId });
useWorkspaceSetupStore.getState().removeWorkspace({ serverId: this.serverId, workspaceId });
}
}
}
File diff suppressed because it is too large Load Diff
+175 -161
View File
@@ -4,7 +4,6 @@ import equal from "fast-deep-equal/es6";
import {
DaemonClient,
type ConnectionState,
type FetchAgentsEntry,
type FetchAgentsOptions,
} from "@getpaseo/client/internal/daemon-client";
import {
@@ -40,12 +39,7 @@ import {
import { getDesktopHost } from "@/desktop/host";
import { CLIENT_CAPS } from "@getpaseo/protocol/client-capabilities";
import { BROWSER_AUTOMATION_COMMAND_NAMES } from "@getpaseo/protocol/browser-automation/rpc-schemas";
import { replaceFetchedAgentDirectory } from "@/utils/agent-directory-sync";
import { useSessionStore } from "@/stores/session-store";
import {
fetchLegacyDaemonWorkspaceDirectory,
shouldUseLegacyDaemonWorkspaceDirectory,
} from "@/workspace/legacy-daemon-workspaces";
import { invalidateCheckoutGitQueriesForServer } from "@/git/query-keys";
import { queryClient } from "@/data/query-client";
import {
@@ -54,6 +48,13 @@ import {
} from "@/data/push-router";
import { mountBrowserAutomationDaemonClientHandler } from "@/browser-automation/handler";
import { schedulesQueryBaseKey } from "@/schedules/aggregated-schedules";
import { sendQueuedComposerMessageNow } from "@/composer/actions";
import {
resolveComposerAttachmentSubmitFormat,
splitComposerAttachmentsForSubmit,
} from "@/composer/attachments/submit";
import { encodeImages } from "@/utils/encode-images";
import { DirectorySync, type RefreshAgentDirectoryResult } from "@/runtime/directory-sync";
export type HostRuntimeConnectionStatus = "idle" | "connecting" | "online" | "offline" | "error";
export type HostRegistryStatus = "loading" | "ready";
@@ -85,6 +86,7 @@ export interface HostRuntimeSnapshot {
hasEverLoadedAgentDirectory: boolean;
probeByConnectionId: Map<string, ConnectionProbeState>;
clientGeneration: number;
connectionEpoch: number;
}
type HostRuntimeSnapshotPatch = Partial<Omit<HostRuntimeSnapshot, "serverId" | "clientGeneration">>;
@@ -176,90 +178,6 @@ const ADAPTIVE_SWITCH_CONSECUTIVE_PROBES = 3;
const DEFAULT_AGENT_DIRECTORY_PAGE_LIMIT = 200;
const CONFIGURED_OVERRIDE_BOOTSTRAP_RETRY_MS = 1_000;
const DEFAULT_AGENT_DIRECTORY_SORT: NonNullable<FetchAgentsOptions["sort"]> = [
{ key: "updated_at", direction: "desc" },
];
function readFetchAgentsHasMore(
pageInfo: Awaited<ReturnType<DaemonClient["fetchAgents"]>>["pageInfo"],
): boolean {
const page = pageInfo as {
hasMore?: boolean;
hasMoreAfter?: boolean;
};
if (typeof page.hasMore === "boolean") {
return page.hasMore;
}
if (typeof page.hasMoreAfter === "boolean") {
return page.hasMoreAfter;
}
return false;
}
function readFetchAgentsNextCursor(
pageInfo: Awaited<ReturnType<DaemonClient["fetchAgents"]>>["pageInfo"],
): string | null {
const page = pageInfo as {
nextCursor?: string | null;
afterCursor?: string | null;
};
if (typeof page.nextCursor === "string" && page.nextCursor.length > 0) {
return page.nextCursor;
}
if (typeof page.afterCursor === "string" && page.afterCursor.length > 0) {
return page.afterCursor;
}
return null;
}
interface AgentDirectoryFetchInput {
client: DaemonClient;
filter?: FetchAgentsOptions["filter"];
subscribe?: FetchAgentsOptions["subscribe"];
page?: FetchAgentsOptions["page"];
}
interface AgentDirectoryFetchResult {
entries: FetchAgentsEntry[];
subscriptionId: string | null;
}
async function fetchCurrentAgentDirectory(
input: AgentDirectoryFetchInput,
): Promise<AgentDirectoryFetchResult> {
const pageLimit = input.page?.limit ?? DEFAULT_AGENT_DIRECTORY_PAGE_LIMIT;
let cursor = input.page?.cursor ?? null;
let includeSubscribe = true;
let subscriptionId: string | null = null;
const entries: FetchAgentsEntry[] = [];
while (true) {
const payload = await input.client.fetchAgents({
scope: input.filter ? undefined : "active",
...(input.filter ? { filter: input.filter } : {}),
sort: DEFAULT_AGENT_DIRECTORY_SORT,
...(includeSubscribe && input.subscribe ? { subscribe: input.subscribe } : {}),
page: cursor ? { limit: pageLimit, cursor } : { limit: pageLimit },
});
entries.push(...payload.entries);
subscriptionId = subscriptionId ?? payload.subscriptionId ?? null;
includeSubscribe = false;
if (!readFetchAgentsHasMore(payload.pageInfo)) {
break;
}
const nextCursor = readFetchAgentsNextCursor(payload.pageInfo);
if (!nextCursor) {
break;
}
cursor = nextCursor;
}
return { entries, subscriptionId };
}
function toActiveConnection(connection: HostConnection): ActiveConnection {
if (connection.type === "directSocket") {
return {
@@ -445,9 +363,15 @@ function nextConnectionMachineState(input: {
function toSnapshotConnectionPatch(
state: HostRuntimeConnectionMachineState,
connectionEpoch: number,
): Pick<
HostRuntimeSnapshot,
"activeConnectionId" | "activeConnection" | "connectionStatus" | "lastError" | "lastOnlineAt"
| "activeConnectionId"
| "activeConnection"
| "connectionStatus"
| "lastError"
| "lastOnlineAt"
| "connectionEpoch"
> {
if (state.tag === "booting") {
return {
@@ -456,6 +380,7 @@ function toSnapshotConnectionPatch(
connectionStatus: "connecting",
lastError: null,
lastOnlineAt: null,
connectionEpoch,
};
}
if (state.tag === "connecting") {
@@ -465,6 +390,7 @@ function toSnapshotConnectionPatch(
connectionStatus: "connecting",
lastError: null,
lastOnlineAt: null,
connectionEpoch,
};
}
if (state.tag === "online") {
@@ -474,6 +400,7 @@ function toSnapshotConnectionPatch(
connectionStatus: "online",
lastError: null,
lastOnlineAt: state.lastOnlineAt,
connectionEpoch,
};
}
if (state.tag === "offline") {
@@ -483,6 +410,7 @@ function toSnapshotConnectionPatch(
connectionStatus: "offline",
lastError: null,
lastOnlineAt: null,
connectionEpoch,
};
}
return {
@@ -491,6 +419,7 @@ function toSnapshotConnectionPatch(
connectionStatus: "error",
lastError: state.message,
lastOnlineAt: null,
connectionEpoch,
};
}
@@ -538,6 +467,10 @@ function createDefaultDeps(): HostRuntimeControllerDeps {
},
}
: undefined;
const appCapabilities = {
[CLIENT_CAPS.selectiveAgentTimeline]: true,
...browserAutomationCapabilities,
};
return {
createClient: ({ host, connection, clientId, runtimeGeneration }) => {
@@ -548,7 +481,7 @@ function createDefaultDeps(): HostRuntimeControllerDeps {
clientType: "mobile" as const,
appVersion: resolveAppVersion() ?? undefined,
runtimeGeneration,
...(browserAutomationCapabilities ? { capabilities: browserAutomationCapabilities } : {}),
capabilities: appCapabilities,
};
if (connection.type === "directSocket" || connection.type === "directPipe") {
return new DaemonClient({
@@ -586,7 +519,7 @@ function createDefaultDeps(): HostRuntimeControllerDeps {
connectToDaemon(connection, {
...(host.serverId ? { serverId: host.serverId } : {}),
...(timeoutMs !== undefined ? { timeoutMs } : {}),
...(browserAutomationCapabilities ? { capabilities: browserAutomationCapabilities } : {}),
capabilities: appCapabilities,
}),
getClientId: () => getOrCreateClientId(),
mountClientHandlers: ({ client, host }) => {
@@ -614,6 +547,7 @@ export class HostRuntimeController {
private deps: HostRuntimeControllerDeps;
private onReconcileServerId: ((oldId: string, newId: string) => void) | null;
private connectionMachineState: HostRuntimeConnectionMachineState;
private connectionEpoch = 0;
private snapshot: HostRuntimeSnapshot;
private listeners = new Set<() => void>();
private activeClient: DaemonClient | null = null;
@@ -644,7 +578,7 @@ export class HostRuntimeController {
};
this.snapshot = {
serverId: this.host.serverId,
...toSnapshotConnectionPatch(this.connectionMachineState),
...toSnapshotConnectionPatch(this.connectionMachineState, this.connectionEpoch),
client: null,
agentDirectoryStatus: "idle",
agentDirectoryError: null,
@@ -712,7 +646,7 @@ export class HostRuntimeController {
}
this.applyConnectionEvent({ type: "stopped" });
this.updateSnapshot({
...toSnapshotConnectionPatch(this.connectionMachineState),
...toSnapshotConnectionPatch(this.connectionMachineState, this.connectionEpoch),
client: null,
});
}
@@ -776,7 +710,7 @@ export class HostRuntimeController {
markStartupError(message: string): void {
this.applyConnectionEvent({ type: "connect_failed", message });
this.updateSnapshot({
...toSnapshotConnectionPatch(this.connectionMachineState),
...toSnapshotConnectionPatch(this.connectionMachineState, this.connectionEpoch),
});
}
@@ -809,7 +743,7 @@ export class HostRuntimeController {
}
this.applyConnectionEvent({ type: "no_connections" });
this.updateSnapshot({
...toSnapshotConnectionPatch(this.connectionMachineState),
...toSnapshotConnectionPatch(this.connectionMachineState, this.connectionEpoch),
probeByConnectionId: new Map(),
});
return;
@@ -1095,6 +1029,9 @@ export class HostRuntimeController {
state: previousState,
event,
});
if (previousState.tag !== "online" && nextState.tag === "online") {
this.connectionEpoch += 1;
}
this.connectionMachineState = nextState;
this.logConnectionTransition({
from: previousState.tag,
@@ -1171,7 +1108,7 @@ export class HostRuntimeController {
message: `Failed to resolve client id: ${message}`,
});
this.updateSnapshot({
...toSnapshotConnectionPatch(this.connectionMachineState),
...toSnapshotConnectionPatch(this.connectionMachineState, this.connectionEpoch),
});
return null;
}
@@ -1269,7 +1206,7 @@ export class HostRuntimeController {
this.snapshot = {
...this.snapshot,
serverId: this.host.serverId,
...toSnapshotConnectionPatch(this.connectionMachineState),
...toSnapshotConnectionPatch(this.connectionMachineState, this.connectionEpoch),
client,
clientGeneration: nextGeneration,
};
@@ -1287,7 +1224,7 @@ export class HostRuntimeController {
lastError: client.lastError,
});
const patch: HostRuntimeSnapshotPatch = {
...toSnapshotConnectionPatch(this.connectionMachineState),
...toSnapshotConnectionPatch(this.connectionMachineState, this.connectionEpoch),
...this.buildAgentDirectoryStatusPatch(),
};
this.updateSnapshot(patch);
@@ -1307,7 +1244,7 @@ export class HostRuntimeController {
message,
});
this.updateSnapshot({
...toSnapshotConnectionPatch(this.connectionMachineState),
...toSnapshotConnectionPatch(this.connectionMachineState, this.connectionEpoch),
});
}
}
@@ -1393,6 +1330,13 @@ function rekeyMap<V>(map: Map<string, V>, oldKey: string, newKey: string): void
map.set(newKey, value);
}
interface AgentDirectoryRefreshInput {
serverId: string;
filter?: FetchAgentsOptions["filter"];
subscribe?: FetchAgentsOptions["subscribe"];
page?: FetchAgentsOptions["page"];
}
export class HostRuntimeStore {
private controllers = new Map<string, HostRuntimeController>();
private serverListeners = new Map<string, Set<() => void>>();
@@ -1406,6 +1350,8 @@ export class HostRuntimeStore {
private deps: HostRuntimeControllerDeps;
private lastConnectionStatusByServer = new Map<string, HostRuntimeConnectionStatus>();
private agentDirectoryBootstrapInFlight = new Map<string, Promise<void>>();
private queuedAgentDrainInFlight = new Set<string>();
private directorySyncByServer = new Map<string, DirectorySync>();
private configuredOverrideBootstrapInFlight: Promise<void> | null = null;
private bootStarted = false;
private storage: HostRuntimeStorage;
@@ -1636,6 +1582,27 @@ export class HostRuntimeStore {
rekeyMap(this.lastConnectionStatusByServer, oldServerId, newServerId);
rekeyMap(this.agentDirectoryBootstrapInFlight, oldServerId, newServerId);
this.directorySyncByServer.get(oldServerId)?.dispose();
this.directorySyncByServer.delete(oldServerId);
const directory = new DirectorySync(newServerId, {
drainQueuedAgentMessage: (agentId) => this.drainQueuedAgentMessage(newServerId, agentId),
markAgentLoading: () => controller.markAgentDirectorySyncLoading(),
markAgentReady: () => {
this.agentDirectoryBootstrapInFlight.delete(newServerId);
controller.markAgentDirectorySyncReady();
},
markAgentError: (error) => controller.markAgentDirectorySyncError(error),
});
this.directorySyncByServer.set(newServerId, directory);
const snapshot = controller.getSnapshot();
directory.connectionChanged({
client: snapshot.client,
status: snapshot.connectionStatus === "online" ? "online" : "offline",
source: {
clientGeneration: snapshot.clientGeneration,
connectionEpoch: snapshot.connectionEpoch,
},
});
const listeners = this.serverListeners.get(oldServerId);
if (listeners) {
@@ -1926,6 +1893,8 @@ export class HostRuntimeStore {
this.controllers.delete(serverId);
this.lastConnectionStatusByServer.delete(serverId);
this.agentDirectoryBootstrapInFlight.delete(serverId);
this.directorySyncByServer.get(serverId)?.dispose();
this.directorySyncByServer.delete(serverId);
void controller.stop();
this.emit(serverId);
}
@@ -1948,6 +1917,19 @@ export class HostRuntimeStore {
onReconcileServerId: (oldId, newId) => this.reconcileServerId(oldId, newId),
});
this.controllers.set(host.serverId, controller);
this.directorySyncByServer.set(
host.serverId,
new DirectorySync(host.serverId, {
drainQueuedAgentMessage: (agentId) =>
this.drainQueuedAgentMessage(host.serverId, agentId),
markAgentLoading: () => controller.markAgentDirectorySyncLoading(),
markAgentReady: () => {
this.agentDirectoryBootstrapInFlight.delete(host.serverId);
controller.markAgentDirectorySyncReady();
},
markAgentError: (error) => controller.markAgentDirectorySyncError(error),
}),
);
this.lastConnectionStatusByServer.set(
host.serverId,
controller.getSnapshot().connectionStatus,
@@ -1980,6 +1962,15 @@ export class HostRuntimeStore {
return;
}
const snapshot = controller.getSnapshot();
const directorySourceChanged =
this.directorySyncByServer.get(serverId)?.connectionChanged({
client: snapshot.client,
status: snapshot.connectionStatus === "online" ? "online" : "offline",
source: {
clientGeneration: snapshot.clientGeneration,
connectionEpoch: snapshot.connectionEpoch,
},
}) ?? false;
const previousStatus = this.lastConnectionStatusByServer.get(serverId);
this.lastConnectionStatusByServer.set(serverId, snapshot.connectionStatus);
const didTransitionOnline =
@@ -2002,17 +1993,20 @@ export class HostRuntimeStore {
if (!didTransitionOnline && snapshot.hasEverLoadedAgentDirectory) {
return;
}
if (this.agentDirectoryBootstrapInFlight.has(serverId)) {
if (this.agentDirectoryBootstrapInFlight.has(serverId) && !directorySourceChanged) {
return;
}
const bootstrap = Promise.resolve()
.then(() =>
this.refreshAgentDirectory({
serverId,
subscribe: { subscriptionId: `app:${serverId}` },
page: { limit: DEFAULT_AGENT_DIRECTORY_PAGE_LIMIT },
}),
Promise.all([
this.refreshAgentDirectory({
serverId,
subscribe: { subscriptionId: `app:${serverId}` },
page: { limit: DEFAULT_AGENT_DIRECTORY_PAGE_LIMIT },
}),
this.refreshWorkspaceDirectory({ serverId, subscribe: true }),
]),
)
.then(() => undefined)
.catch((error) => {
@@ -2031,6 +2025,55 @@ export class HostRuntimeStore {
this.agentDirectoryBootstrapInFlight.set(serverId, bootstrap);
}
drainQueuedAgentMessage(serverId: string, agentId: string): void {
const drainKey = `${serverId}:${agentId}`;
if (this.queuedAgentDrainInFlight.has(drainKey)) return;
const store = useSessionStore.getState();
const session = store.sessions[serverId];
const queue = session?.queuedMessages.get(agentId);
const client = session?.client;
if (!client || !queue?.length || session.initializingAgents.get(agentId) === true) {
return;
}
this.queuedAgentDrainInFlight.add(drainKey);
const next = queue[0];
void sendQueuedComposerMessageNow({
agentId,
messageId: next.id,
queue: {
read: (queuedAgentId) =>
useSessionStore.getState().sessions[serverId]?.queuedMessages.get(queuedAgentId) ?? [],
write: (update) => useSessionStore.getState().setQueuedMessages(serverId, update),
},
submitMessage: async ({ text, attachments }) => {
const supportsForgeAttachments =
useSessionStore.getState().sessions[serverId]?.serverInfo?.features?.forgeSearch === true;
const wirePayload = splitComposerAttachmentsForSubmit(attachments, {
format: resolveComposerAttachmentSubmitFormat({ supportsForgeAttachments }),
});
const images = await encodeImages(wirePayload.images);
await client.sendAgentMessage(agentId, text, {
messageId: next.id,
...(images && images.length > 0 ? { images } : {}),
attachments: wirePayload.attachments,
});
},
})
.then((result) => {
if (result.status === "failed") {
console.error("[HostRuntime] failed to drain queued agent message", {
serverId,
agentId,
error: result.errorMessage,
});
}
return result;
})
.finally(() => {
this.queuedAgentDrainInFlight.delete(drainKey);
});
}
getSnapshot(serverId: string): HostRuntimeSnapshot | null {
return this.controllers.get(serverId)?.getSnapshot() ?? null;
}
@@ -2095,63 +2138,34 @@ export class HostRuntimeStore {
).then(() => undefined);
}
async refreshAgentDirectory(input: {
serverId: string;
filter?: FetchAgentsOptions["filter"];
subscribe?: FetchAgentsOptions["subscribe"];
page?: FetchAgentsOptions["page"];
}): Promise<{
agents: ReturnType<typeof replaceFetchedAgentDirectory>["agents"];
subscriptionId: string | null;
}> {
const controller = this.controllers.get(input.serverId);
if (!controller) {
throw new Error(`Unknown host runtime for serverId ${input.serverId}`);
}
const snapshot = controller.getSnapshot();
const client = controller.getClient();
if (!client || snapshot.connectionStatus !== "online") {
throw new Error(`Host ${input.serverId} is not connected`);
}
async refreshAgentDirectory(
input: AgentDirectoryRefreshInput,
): Promise<RefreshAgentDirectoryResult> {
const directory = this.directorySyncByServer.get(input.serverId);
if (!directory) throw new Error(`Unknown host runtime for serverId ${input.serverId}`);
return directory.refreshAgents(input);
}
controller.markAgentDirectorySyncLoading();
try {
const session = useSessionStore.getState().sessions[input.serverId];
if (!input.filter && shouldUseLegacyDaemonWorkspaceDirectory(session?.serverInfo)) {
const result = await fetchLegacyDaemonWorkspaceDirectory({
client,
serverId: input.serverId,
subscribe: input.subscribe,
page: input.page,
});
controller.markAgentDirectorySyncReady();
return {
agents: result.agents,
subscriptionId: result.subscriptionId,
};
}
async refreshWorkspaceDirectory(input: { serverId: string; subscribe?: boolean }): Promise<void> {
const directory = this.directorySyncByServer.get(input.serverId);
if (!directory) throw new Error(`Unknown host runtime for serverId ${input.serverId}`);
await directory.refreshWorkspaces({ subscribe: input.subscribe });
}
const directory = await fetchCurrentAgentDirectory({
client,
filter: input.filter,
subscribe: input.subscribe,
page: input.page,
});
async refreshDirectories(serverId: string): Promise<void> {
const directory = this.directorySyncByServer.get(serverId);
if (!directory) throw new Error(`Unknown host runtime for serverId ${serverId}`);
await directory.refreshAll();
}
const { agents } = replaceFetchedAgentDirectory({
serverId: input.serverId,
entries: directory.entries,
});
controller.markAgentDirectorySyncReady();
return {
agents,
subscriptionId: directory.subscriptionId,
};
} catch (error) {
controller.markAgentDirectorySyncError(toErrorMessage(error));
throw error;
}
fetchAgentTimeline(
serverId: string,
agentId: string,
request: Parameters<DaemonClient["fetchAgentTimeline"]>[1],
): Promise<Awaited<ReturnType<DaemonClient["fetchAgentTimeline"]>>> {
const directory = this.directorySyncByServer.get(serverId);
if (!directory) throw new Error(`Unknown host runtime for serverId ${serverId}`);
return directory.fetchTimeline(agentId, request);
}
refreshAllAgentDirectories(input?: { serverIds?: string[] }): void {