diff --git a/packages/app/src/runtime/directory-sync/index.test.ts b/packages/app/src/runtime/directory-sync/index.test.ts index ba8dc97b8e..ad5212ff4e 100644 --- a/packages/app/src/runtime/directory-sync/index.test.ts +++ b/packages/app/src/runtime/directory-sync/index.test.ts @@ -6,7 +6,10 @@ import { type SqliteValue, } from "@/runtime/replica-cache/row-store-sqlite"; import { afterEach, describe, expect, it } from "vitest"; -import type { DaemonClient } from "@getpaseo/client/internal/daemon-client"; +import type { + DaemonClient, + WorkspaceLabelListPayload, +} from "@getpaseo/client/internal/daemon-client"; import type { SessionOutboundMessage } from "@getpaseo/protocol/messages"; import { normalizeProjectDescriptor, @@ -15,6 +18,7 @@ import { } from "@/stores/session-store"; import { normalizeAgentSnapshot } from "@/utils/agent-snapshots"; import { selectWorkspaceDirectoryServerIds } from "@/stores/session-store-hooks/selectors"; +import { useWorkspaceLabels } from "@/workspace-labels"; import type { DirectoryReplicaMutation } from "@/runtime/replica-cache"; import { DirectoryRefreshSupersededError, @@ -29,6 +33,8 @@ type ProjectListResult = Awaited>; type AgentFetchResult = Awaited>; class FakeDirectoryClient { + supportsWorkspaceLabels = false; + listWorkspaceLabelsCalls = 0; fetchAgentsCalls = 0; lastAgentOptions: unknown; fetchWorkspacesCalls = 0; @@ -144,19 +150,36 @@ class FakeDirectoryClient { }); } - getLastServerInfoMessage(): null { - return null; + observeWorkspaceLabels() { + return subscriptionFixture(this.listWorkspaceLabels(), () => () => {}); + } + + async listWorkspaceLabels(): Promise { + this.listWorkspaceLabelsCalls += 1; + return { + requestId: "workspace-labels", + labels: [{ name: "Urgent", color: "red" }], + sync: { mode: "snapshot", removals: [], generation: "generation-1", headSeq: 0 }, + }; + } + + getLastServerInfoMessage() { + return this.supportsWorkspaceLabels ? { features: { workspaceLabels: true } } : null; } } const serverIds = new Set(); -function createDirectory(serverId: string): { +function createDirectory( + serverId: string, + options?: { workspaceLabels?: boolean }, +): { client: FakeDirectoryClient; directory: DirectorySync; } { serverIds.add(serverId); const client = new FakeDirectoryClient(); + client.supportsWorkspaceLabels = options?.workspaceLabels === true; const directory = new DirectorySync(serverId, { onAgentStoppedRunning: () => undefined, markAgentLoading: () => undefined, @@ -207,6 +230,7 @@ function createAgent(serverId: string, id: string) { afterEach(() => { for (const serverId of serverIds) useSessionStore.getState().clearSession(serverId); serverIds.clear(); + useWorkspaceLabels.setState({ hosts: {} }); }); describe("DirectorySync session readiness", () => { @@ -1349,3 +1373,103 @@ it("fills every cached workspace beneath live updates received during the SQLite await cache.flush(); database.close(); }); + +function createWorkspaceLabelDirectory(serverId: string) { + const { client, directory } = createDirectory(serverId, { workspaceLabels: true }); + const store = useSessionStore.getState(); + store.initializeSession(serverId, client as unknown as DaemonClient, 1); + store.updateSessionServerInfo(serverId, { + serverId, + hostname: null, + version: "test", + features: { + directorySync: true, + projectList: true, + workspaceLabels: true, + workspaceMultiplicity: true, + }, + }); + return { client, directory }; +} + +async function flushWorkspaceLabels(): Promise { + for (let index = 0; index < 4; index += 1) await Promise.resolve(); +} + +describe("DirectorySync workspace labels", () => { + it("loads the catalog when an agent route is open before the sidebar", async () => { + const serverId = "workspace-labels-route-then-full"; + const { directory } = createWorkspaceLabelDirectory(serverId); + + directory.setAgentRouteDemand(["agent-1"]); + await directory.refreshDemand(); + directory.setDemand({}, true); + await flushWorkspaceLabels(); + + expect(useWorkspaceLabels.getState().hosts[serverId]).toMatchObject({ + status: "online", + labels: [{ name: "Urgent", color: "red" }], + }); + directory.dispose(); + }); + + it("loads the catalog when the sidebar demands the directory first", async () => { + const serverId = "workspace-labels-full-first"; + const { directory } = createWorkspaceLabelDirectory(serverId); + + directory.setDemand({}, true); + await directory.refreshDemand(); + + expect(useWorkspaceLabels.getState().hosts[serverId]).toMatchObject({ + status: "online", + labels: [{ name: "Urgent", color: "red" }], + }); + directory.dispose(); + }); + + it("keeps the picker online during a full refresh", async () => { + const serverId = "workspace-labels-refresh-no-flicker"; + const { client, directory } = createWorkspaceLabelDirectory(serverId); + directory.setDemand({}, true); + await directory.refreshDemand(); + const callsBeforeRefresh = client.listWorkspaceLabelsCalls; + const statuses: string[] = []; + const unsubscribe = useWorkspaceLabels.subscribe((state) => { + const status = state.hosts[serverId]?.status; + if (status) statuses.push(status); + }); + + await directory.refreshDemand(); + + expect(client.listWorkspaceLabelsCalls).toBeGreaterThan(callsBeforeRefresh); + expect(statuses).not.toContain("offline"); + unsubscribe(); + directory.dispose(); + }); + + it("restores the catalog after reconnecting", async () => { + const serverId = "workspace-labels-reconnect"; + const { client, directory } = createWorkspaceLabelDirectory(serverId); + directory.setDemand({}, true); + await directory.refreshDemand(); + + directory.connectionChanged({ + client: null, + status: "offline", + source: { clientGeneration: 1, connectionEpoch: 1 }, + }); + expect(useWorkspaceLabels.getState().hosts[serverId]?.status).toBe("offline"); + directory.connectionChanged({ + client: client as unknown as DaemonClient, + status: "online", + source: { clientGeneration: 1, connectionEpoch: 2 }, + }); + await flushWorkspaceLabels(); + + expect(useWorkspaceLabels.getState().hosts[serverId]).toMatchObject({ + status: "online", + labels: [{ name: "Urgent", color: "red" }], + }); + directory.dispose(); + }); +}); diff --git a/packages/app/src/runtime/directory-sync/index.ts b/packages/app/src/runtime/directory-sync/index.ts index 51e8b37397..da181ee473 100644 --- a/packages/app/src/runtime/directory-sync/index.ts +++ b/packages/app/src/runtime/directory-sync/index.ts @@ -189,6 +189,8 @@ export class DirectorySync { this.connection = connection; this.abortPendingSessionWaits(); if (!connection.client || connection.status !== "online") return true; + // Reattach labels here because route-only demand can satisfy the epoch before full demand requests them. + void this.connectWorkspaceLabels().catch(() => undefined); if (this.hasDemand()) void this.requestDemandRefresh().catch(() => undefined); return true; } diff --git a/packages/app/src/workspace-labels/index.ts b/packages/app/src/workspace-labels/index.ts index 49a2e6b1a9..95fd5a3a26 100644 --- a/packages/app/src/workspace-labels/index.ts +++ b/packages/app/src/workspace-labels/index.ts @@ -147,6 +147,11 @@ class WorkspaceLabelsController { client: DaemonClient; supportsWorkspaceLabels: boolean; }): Promise { + const existing = this.connections.get(input.serverId); + if (existing?.client === input.client && input.supportsWorkspaceLabels) { + await this.refresh(input.serverId); + return; + } this.disconnect(input.serverId); const replica = this.replicas.get(input.serverId) ?? new HostWorkspaceLabelReplica(); this.replicas.set(input.serverId, replica); diff --git a/packages/app/src/workspace-labels/picker.tsx b/packages/app/src/workspace-labels/picker.tsx index 7afdaeabe6..a0e2348308 100644 --- a/packages/app/src/workspace-labels/picker.tsx +++ b/packages/app/src/workspace-labels/picker.tsx @@ -131,6 +131,7 @@ function WorkspaceLabelPickerPage({ }, [assignedLabels, host?.status, labels, model]); const snapshot = useSyncExternalStore(model.subscribe, model.snapshot, model.snapshot); const offline = !snapshot.online; + const error = snapshot.error ?? (host?.error ? t("workspaceLabels.errors.load") : null); const pending = useMemo(() => new Set(snapshot.pendingNames), [snapshot.pendingNames]); const toggle = useCallback( (label: WorkspaceLabelDefinition, assigned: boolean) => { @@ -159,9 +160,7 @@ function WorkspaceLabelPickerPage({ > {t("workspaceLabels.create")} - {snapshot.error ? ( - {snapshot.error} - ) : null} + {error ? {error} : null} {host?.status === "unsupported" ? ( {t("workspaceLabels.updateHostUse")} ) : null}