From 515cb0a777d5c058cfc866662f45f76bb733afd2 Mon Sep 17 00:00:00 2001 From: Mohamed Boudra Date: Tue, 17 Feb 2026 21:16:31 +0700 Subject: [PATCH] refactor: remove legacy agent update RPCs --- packages/app/src/contexts/session-context.tsx | 106 +++++-- packages/cli/src/utils/client.ts | 1 - .../server/scripts/voice-main-daemon-e2e.ts | 4 +- .../server/scripts/voice-roundtrip-debug.ts | 2 +- .../server/src/client/daemon-client.test.ts | 5 + packages/server/src/client/daemon-client.ts | 44 +-- .../src/server/daemon-client.e2e.test.ts | 4 +- .../daemon-e2e/agent-operations.e2e.test.ts | 8 +- .../daemon-e2e/live-preferences.e2e.test.ts | 2 +- ...hile-running-stuck-claude.real.e2e.test.ts | 8 +- .../send-while-running-stuck.real.e2e.test.ts | 8 +- packages/server/src/server/session.ts | 263 ++++++++++++------ .../server/test-utils/daemon-test-context.ts | 2 +- packages/server/src/shared/messages.ts | 44 +-- 14 files changed, 295 insertions(+), 206 deletions(-) diff --git a/packages/app/src/contexts/session-context.tsx b/packages/app/src/contexts/session-context.tsx index 3136c7122..7e82f3983 100644 --- a/packages/app/src/contexts/session-context.tsx +++ b/packages/app/src/contexts/session-context.tsx @@ -27,6 +27,7 @@ import { useDaemonConnections } from "./daemon-connections-context"; import type { ActiveConnection } from "./daemon-connections-context"; import { useSessionStore, + type Agent, type SessionState, type DaemonConnectionSnapshot, } from "@/stores/session-store"; @@ -850,14 +851,7 @@ export function SessionProvider({ useEffect(() => { if (!connectionSnapshot.isConnected) { hasBootstrappedAgentUpdatesRef.current = false; - const subscriptionId = agentUpdatesSubscriptionIdRef.current; - if (subscriptionId && client) { - try { - client.unsubscribeAgentUpdates(subscriptionId); - } catch { - // no-op - } - } + pendingAgentUpdatesRef.current.clear(); agentUpdatesSubscriptionIdRef.current = null; return; } @@ -866,26 +860,84 @@ export function SessionProvider({ } hasBootstrappedAgentUpdatesRef.current = true; - try { - if (!agentUpdatesSubscriptionIdRef.current) { - agentUpdatesSubscriptionIdRef.current = client.subscribeAgentUpdates({ - subscriptionId: `app:${serverId}`, - filter: { labels: { ui: "true" } }, - }); - } - } catch (err) { - console.error("[Session] subscribeAgentUpdates failed", { serverId, err }); - } + let cancelled = false; + const requestedSubscriptionId = `app:${serverId}`; - // Session bootstrap is now fully event-driven for agent lists. - setInitializingAgents(serverId, new Map()); - setHasHydratedAgents(serverId, true); - updateConnectionStatus(serverId, { - status: "online", - lastOnlineAt: new Date().toISOString(), - agentListReady: true, - }); - }, [connectionSnapshot.isConnected, client, serverId, setHasHydratedAgents, updateConnectionStatus]); + const bootstrapAgentDirectory = async () => { + try { + const payload = await client.fetchAgents({ + filter: { labels: { ui: "true" } }, + subscribe: { subscriptionId: requestedSubscriptionId }, + }); + if (cancelled) { + return; + } + + agentUpdatesSubscriptionIdRef.current = + payload.subscriptionId ?? requestedSubscriptionId; + + const nextAgents = new Map(); + const nextPendingPermissions = new Map< + string, + { key: string; agentId: string; request: AgentPermissionRequest } + >(); + const nextStatuses = new Map(); + + for (const entry of payload.entries) { + const agent = { + ...normalizeAgentSnapshot(entry.agent, serverId), + projectPlacement: entry.project, + }; + nextAgents.set(agent.id, agent); + nextStatuses.set(agent.id, agent.status); + + for (const request of agent.pendingPermissions) { + const key = derivePendingPermissionKey(agent.id, request); + nextPendingPermissions.set(key, { key, agentId: agent.id, request }); + } + } + + previousAgentStatusRef.current = nextStatuses; + pendingAgentUpdatesRef.current.clear(); + setAgents(serverId, nextAgents); + for (const agent of nextAgents.values()) { + setAgentLastActivity(agent.id, agent.lastActivityAt); + } + setPendingPermissions(serverId, nextPendingPermissions); + setInitializingAgents(serverId, new Map()); + setHasHydratedAgents(serverId, true); + updateConnectionStatus(serverId, { + status: "online", + lastOnlineAt: new Date().toISOString(), + agentListReady: true, + }); + } catch (err) { + if (cancelled) { + return; + } + hasBootstrappedAgentUpdatesRef.current = false; + pendingAgentUpdatesRef.current.clear(); + agentUpdatesSubscriptionIdRef.current = null; + console.error("[Session] fetchAgents bootstrap failed", { serverId, err }); + } + }; + + void bootstrapAgentDirectory(); + + return () => { + cancelled = true; + }; + }, [ + connectionSnapshot.isConnected, + client, + serverId, + setAgentLastActivity, + setAgents, + setHasHydratedAgents, + setInitializingAgents, + setPendingPermissions, + updateConnectionStatus, + ]); // Daemon message handlers - directly update Zustand store useEffect(() => { diff --git a/packages/cli/src/utils/client.ts b/packages/cli/src/utils/client.ts index 51786a948..e64482ed1 100644 --- a/packages/cli/src/utils/client.ts +++ b/packages/cli/src/utils/client.ts @@ -62,7 +62,6 @@ export async function connectToDaemon(options?: ConnectOptions): Promise { const client = new DaemonClient({ url: `${params.daemonUrl}/ws` }); await client.connect(); - client.subscribeAgentUpdates({ subscriptionId: `voice-e2e-${randomUUID()}` }); + await client.fetchAgents({ + subscribe: { subscriptionId: `voice-e2e-${randomUUID()}` }, + }); const mode = await client.setVoiceMode(true, params.voiceAgentId); if (!mode.accepted) { diff --git a/packages/server/scripts/voice-roundtrip-debug.ts b/packages/server/scripts/voice-roundtrip-debug.ts index 4c1beef8d..60be54d26 100644 --- a/packages/server/scripts/voice-roundtrip-debug.ts +++ b/packages/server/scripts/voice-roundtrip-debug.ts @@ -48,7 +48,7 @@ async function main(): Promise { try { await client.connect(); - client.subscribeAgentUpdates({ subscriptionId: "voice-debug" }); + await client.fetchAgents({ subscribe: { subscriptionId: "voice-debug" } }); const voiceCwd = mkdtempSync(path.join(tmpdir(), "voice-roundtrip-debug-")); const voiceAgent = await client.createAgent({ diff --git a/packages/server/src/client/daemon-client.test.ts b/packages/server/src/client/daemon-client.test.ts index 2e06f4ffc..4df5f8cb7 100644 --- a/packages/server/src/client/daemon-client.test.ts +++ b/packages/server/src/client/daemon-client.test.ts @@ -535,6 +535,7 @@ describe("DaemonClient", () => { { key: "created_at", direction: "desc" }, ], page: { limit: 25, cursor: "cursor-1" }, + subscribe: { subscriptionId: "sub-1" }, }); expect(mock.sent).toHaveLength(1); @@ -549,6 +550,7 @@ describe("DaemonClient", () => { direction: "asc" | "desc"; }>; page?: { limit: number; cursor?: string }; + subscribe?: { subscriptionId?: string }; }; }; expect(request.message.type).toBe("fetch_agents_request"); @@ -557,6 +559,7 @@ describe("DaemonClient", () => { { key: "created_at", direction: "desc" }, ]); expect(request.message.page).toEqual({ limit: 25, cursor: "cursor-1" }); + expect(request.message.subscribe).toEqual({ subscriptionId: "sub-1" }); mock.triggerMessage( JSON.stringify({ @@ -565,6 +568,7 @@ describe("DaemonClient", () => { type: "fetch_agents_response", payload: { requestId: request.message.requestId, + subscriptionId: "sub-1", entries: [], pageInfo: { nextCursor: null, @@ -578,6 +582,7 @@ describe("DaemonClient", () => { await expect(promise).resolves.toEqual({ requestId: request.message.requestId, + subscriptionId: "sub-1", entries: [], pageInfo: { nextCursor: null, diff --git a/packages/server/src/client/daemon-client.ts b/packages/server/src/client/daemon-client.ts index 91dfa30c2..0355ae193 100644 --- a/packages/server/src/client/daemon-client.ts +++ b/packages/server/src/client/daemon-client.ts @@ -350,10 +350,6 @@ export class DaemonClient { private connectReject: ((error: Error) => void) | null = null; private lastErrorValue: string | null = null; private connectionState: ConnectionState = { status: "idle" }; - private agentUpdateSubscriptions = new Map< - string, - { labels?: Record; agentId?: string } | undefined - >(); private checkoutDiffSubscriptions = new Map< string, { cwd: string; compare: { mode: "uncommitted" | "base"; baseRef?: string } } @@ -475,7 +471,6 @@ export class DaemonClient { this.lastErrorValue = null; this.reconnectAttempt = 0; this.updateConnectionState({ status: "connected" }); - this.resubscribeAgentUpdates(); this.resubscribeCheckoutDiffSubscriptions(); this.resubscribeTerminalDirectorySubscriptions(); this.flushPendingSendQueue(); @@ -997,6 +992,7 @@ export class DaemonClient { ...(options?.filter ? { filter: options.filter } : {}), ...(options?.sort ? { sort: options.sort } : {}), ...(options?.page ? { page: options.page } : {}), + ...(options?.subscribe ? { subscribe: options.subscribe } : {}), }); return this.sendRequest({ requestId: resolvedRequestId, @@ -1043,44 +1039,6 @@ export class DaemonClient { return payload.agent; } - subscribeAgentUpdates(options?: { - subscriptionId?: string; - filter?: { labels?: Record; agentId?: string }; - }): string { - const subscriptionId = options?.subscriptionId ?? crypto.randomUUID(); - this.agentUpdateSubscriptions.set(subscriptionId, options?.filter); - const message = SessionInboundMessageSchema.parse({ - type: "subscribe_agent_updates", - subscriptionId, - ...(options?.filter ? { filter: options.filter } : {}), - }); - this.sendSessionMessage(message); - return subscriptionId; - } - - unsubscribeAgentUpdates(subscriptionId: string): void { - this.agentUpdateSubscriptions.delete(subscriptionId); - const message = SessionInboundMessageSchema.parse({ - type: "unsubscribe_agent_updates", - subscriptionId, - }); - this.sendSessionMessage(message); - } - - private resubscribeAgentUpdates(): void { - if (this.agentUpdateSubscriptions.size === 0) { - return; - } - for (const [subscriptionId, filter] of this.agentUpdateSubscriptions) { - const message = SessionInboundMessageSchema.parse({ - type: "subscribe_agent_updates", - subscriptionId, - ...(filter ? { filter } : {}), - }); - this.sendSessionMessage(message); - } - } - private resubscribeCheckoutDiffSubscriptions(): void { if (this.checkoutDiffSubscriptions.size === 0) { return; diff --git a/packages/server/src/server/daemon-client.e2e.test.ts b/packages/server/src/server/daemon-client.e2e.test.ts index bd2cc5cf3..2e1e4537e 100644 --- a/packages/server/src/server/daemon-client.e2e.test.ts +++ b/packages/server/src/server/daemon-client.e2e.test.ts @@ -291,7 +291,9 @@ describe("daemon client E2E", () => { async () => { const cwd = tmpCwd(); - ctx.client.subscribeAgentUpdates(); + await ctx.client.fetchAgents({ + subscribe: { subscriptionId: "daemon-client-lifecycle" }, + }); const agentUpdatePromise = waitForSignal(15000, (resolve) => { const unsubscribe = ctx.client.on("agent_update", (message) => { diff --git a/packages/server/src/server/daemon-e2e/agent-operations.e2e.test.ts b/packages/server/src/server/daemon-e2e/agent-operations.e2e.test.ts index 47f3091a8..a471c4ba3 100644 --- a/packages/server/src/server/daemon-e2e/agent-operations.e2e.test.ts +++ b/packages/server/src/server/daemon-e2e/agent-operations.e2e.test.ts @@ -143,7 +143,9 @@ describe("daemon E2E", () => { async () => { const cwd = tmpCwd(); - ctx.client.subscribeAgentUpdates(); + await ctx.client.fetchAgents({ + subscribe: { subscriptionId: "agent-operations-cancel" }, + }); // Create Codex agent const agent = await ctx.client.createAgent({ @@ -217,7 +219,9 @@ describe("daemon E2E", () => { async () => { const cwd = tmpCwd(); - ctx.client.subscribeAgentUpdates(); + await ctx.client.fetchAgents({ + subscribe: { subscriptionId: "agent-operations-mode" }, + }); // Create a Codex agent with default mode ("auto") const agent = await ctx.client.createAgent({ diff --git a/packages/server/src/server/daemon-e2e/live-preferences.e2e.test.ts b/packages/server/src/server/daemon-e2e/live-preferences.e2e.test.ts index fd7750088..1aa7efd65 100644 --- a/packages/server/src/server/daemon-e2e/live-preferences.e2e.test.ts +++ b/packages/server/src/server/daemon-e2e/live-preferences.e2e.test.ts @@ -73,7 +73,7 @@ describe("daemon E2E", () => { unsubscribe = ctx.client.subscribeRawMessages((message) => { messages.push(message); }); - ctx.client.subscribeAgentUpdates(); + await ctx.client.fetchAgents({ subscribe: { subscriptionId: "live-preferences" } }); }); afterEach(async () => { diff --git a/packages/server/src/server/daemon-e2e/send-while-running-stuck-claude.real.e2e.test.ts b/packages/server/src/server/daemon-e2e/send-while-running-stuck-claude.real.e2e.test.ts index 98ae27b53..d4614b8f2 100644 --- a/packages/server/src/server/daemon-e2e/send-while-running-stuck-claude.real.e2e.test.ts +++ b/packages/server/src/server/daemon-e2e/send-while-running-stuck-claude.real.e2e.test.ts @@ -32,8 +32,8 @@ describe("daemon E2E (real claude) - send while running recovery", () => { try { await primary.connect(); await secondary.connect(); - primary.subscribeAgentUpdates({ subscriptionId: "primary" }); - secondary.subscribeAgentUpdates({ subscriptionId: "secondary" }); + await primary.fetchAgents({ subscribe: { subscriptionId: "primary" } }); + await secondary.fetchAgents({ subscribe: { subscriptionId: "secondary" } }); const agent = await primary.createAgent({ cwd, @@ -67,7 +67,9 @@ describe("daemon E2E (real claude) - send while running recovery", () => { const reconnected = new DaemonClient({ url: `ws://127.0.0.1:${daemon.port}/ws` }); try { await reconnected.connect(); - reconnected.subscribeAgentUpdates({ subscriptionId: "reconnected" }); + await reconnected.fetchAgents({ + subscribe: { subscriptionId: "reconnected" }, + }); reconnected.on("agent_update", (message) => { if (message.type !== "agent_update" || message.payload.kind !== "upsert") { diff --git a/packages/server/src/server/daemon-e2e/send-while-running-stuck.real.e2e.test.ts b/packages/server/src/server/daemon-e2e/send-while-running-stuck.real.e2e.test.ts index d2a21c728..6710343fd 100644 --- a/packages/server/src/server/daemon-e2e/send-while-running-stuck.real.e2e.test.ts +++ b/packages/server/src/server/daemon-e2e/send-while-running-stuck.real.e2e.test.ts @@ -32,8 +32,8 @@ describe("daemon E2E (real codex) - send while running recovery", () => { try { await primary.connect(); await secondary.connect(); - primary.subscribeAgentUpdates({ subscriptionId: "primary" }); - secondary.subscribeAgentUpdates({ subscriptionId: "secondary" }); + await primary.fetchAgents({ subscribe: { subscriptionId: "primary" } }); + await secondary.fetchAgents({ subscribe: { subscriptionId: "secondary" } }); const agent = await primary.createAgent({ cwd, @@ -64,7 +64,9 @@ describe("daemon E2E (real codex) - send while running recovery", () => { const reconnected = new DaemonClient({ url: `ws://127.0.0.1:${daemon.port}/ws` }); try { await reconnected.connect(); - reconnected.subscribeAgentUpdates({ subscriptionId: "reconnected" }); + await reconnected.fetchAgents({ + subscribe: { subscriptionId: "reconnected" }, + }); reconnected.on("agent_update", (message) => { if (message.type !== "agent_update" || message.payload.kind !== "upsert") { diff --git a/packages/server/src/server/session.ts b/packages/server/src/server/session.ts index a97704b10..f82c3379a 100644 --- a/packages/server/src/server/session.ts +++ b/packages/server/src/server/session.ts @@ -293,6 +293,7 @@ type FetchAgentsRequestMessage = Extract< SessionInboundMessage, { type: "fetch_agents_request" } >; +type FetchAgentsRequestFilter = NonNullable; type FetchAgentsRequestSort = NonNullable[number]; type FetchAgentsResponsePayload = Extract< SessionOutboundMessage, @@ -300,6 +301,17 @@ type FetchAgentsResponsePayload = Extract< >["payload"]; type FetchAgentsResponseEntry = FetchAgentsResponsePayload["entries"][number]; type FetchAgentsResponsePageInfo = FetchAgentsResponsePayload["pageInfo"]; +type AgentUpdatePayload = Extract< + SessionOutboundMessage, + { type: "agent_update" } +>["payload"]; +type AgentUpdatesFilter = FetchAgentsRequestFilter; +type AgentUpdatesSubscriptionState = { + subscriptionId: string; + filter?: AgentUpdatesFilter; + isBootstrapping: boolean; + pendingUpdatesByAgentId: Map; +}; type FetchAgentsCursor = { sort: FetchAgentsRequestSort[]; values: Record; @@ -551,12 +563,7 @@ export class Session { private readonly pushTokenStore: PushTokenStore; private readonly providerRegistry: ReturnType; private unsubscribeAgentEvents: (() => void) | null = null; - private agentUpdatesSubscription: - | { - subscriptionId: string; - filter?: { labels?: Record; agentId?: string }; - } - | null = null; + private agentUpdatesSubscription: AgentUpdatesSubscriptionState | null = null; private clientActivity: { deviceType: "web" | "mobile"; focusedAgentId: string | null; @@ -1081,19 +1088,105 @@ export class Session { } } - private matchesAgentFilter( - agent: AgentSnapshotPayload, - filter?: { labels?: Record; agentId?: string } - ): boolean { - if (filter?.agentId && agent.id !== filter.agentId) { + private matchesAgentFilter(options: { + agent: AgentSnapshotPayload; + project: ProjectPlacementPayload; + filter?: AgentUpdatesFilter; + }): boolean { + const { agent, project, filter } = options; + + if (filter?.labels) { + const matchesLabels = Object.entries(filter.labels).every( + ([key, value]) => agent.labels[key] === value + ); + if (!matchesLabels) { + return false; + } + } + + const includeArchived = filter?.includeArchived ?? false; + if (!includeArchived && agent.archivedAt) { return false; } - if (!filter?.labels) { - return true; + + if (filter?.statuses && filter.statuses.length > 0) { + const statuses = new Set(filter.statuses); + if (!statuses.has(agent.status)) { + return false; + } + } + + if (typeof filter?.requiresAttention === "boolean") { + const requiresAttention = agent.requiresAttention ?? false; + if (requiresAttention !== filter.requiresAttention) { + return false; + } + } + + if (filter?.projectKeys && filter.projectKeys.length > 0) { + const projectKeys = new Set( + filter.projectKeys.filter((item) => item.trim().length > 0) + ); + if (projectKeys.size > 0 && !projectKeys.has(project.projectKey)) { + return false; + } + } + + return true; + } + + private getAgentUpdateTargetId(update: AgentUpdatePayload): string { + return update.kind === "remove" ? update.agentId : update.agent.id; + } + + private bufferOrEmitAgentUpdate( + subscription: AgentUpdatesSubscriptionState, + payload: AgentUpdatePayload + ): void { + if (subscription.isBootstrapping) { + subscription.pendingUpdatesByAgentId.set( + this.getAgentUpdateTargetId(payload), + payload + ); + return; + } + + this.emit({ + type: "agent_update", + payload, + }); + } + + private flushBootstrappedAgentUpdates(options?: { + snapshotUpdatedAtByAgentId?: Map; + }): void { + const subscription = this.agentUpdatesSubscription; + if (!subscription || !subscription.isBootstrapping) { + return; + } + + subscription.isBootstrapping = false; + const pending = Array.from(subscription.pendingUpdatesByAgentId.values()); + subscription.pendingUpdatesByAgentId.clear(); + + for (const payload of pending) { + if (payload.kind === "upsert") { + const snapshotUpdatedAt = options?.snapshotUpdatedAtByAgentId?.get( + payload.agent.id + ); + if (typeof snapshotUpdatedAt === "number") { + const updateUpdatedAt = Date.parse(payload.agent.updatedAt); + if (!Number.isNaN(updateUpdatedAt) && updateUpdatedAt <= snapshotUpdatedAt) { + continue; + } + } + } + + this.emit({ + type: "agent_update", + payload, + }); } - return Object.entries(filter.labels).every( - ([key, value]) => agent.labels[key] === value - ); } private buildFallbackProjectCheckout(cwd: string): ProjectCheckoutLitePayload { @@ -1159,51 +1252,31 @@ export class Session { } const payload = await this.buildAgentPayload(agent); - const matches = this.matchesAgentFilter(payload, subscription.filter); + const project = await this.buildProjectPlacement(payload.cwd); + const matches = this.matchesAgentFilter({ + agent: payload, + project, + filter: subscription.filter, + }); if (matches) { - const project = await this.buildProjectPlacement(payload.cwd); - this.emit({ - type: "agent_update", - payload: { kind: "upsert", agent: payload, project }, + this.bufferOrEmitAgentUpdate(subscription, { + kind: "upsert", + agent: payload, + project, }); return; } - this.emit({ - type: "agent_update", - payload: { kind: "remove", agentId: payload.id }, + this.bufferOrEmitAgentUpdate(subscription, { + kind: "remove", + agentId: payload.id, }); } catch (error) { this.sessionLogger.error({ err: error }, "Failed to emit agent update"); } } - private async emitCurrentAgentUpdatesForSubscription(): Promise { - const subscription = this.agentUpdatesSubscription; - if (!subscription) { - return; - } - - try { - const agents = await this.listAgentPayloads({ - labels: subscription.filter?.labels, - }); - for (const agent of agents) { - const project = await this.buildProjectPlacement(agent.cwd); - this.emit({ - type: "agent_update", - payload: { kind: "upsert", agent, project }, - }); - } - } catch (error) { - this.sessionLogger.error( - { err: error }, - "Failed to emit current agent updates for subscription bootstrap" - ); - } - } - /** * Main entry point for processing session messages */ @@ -1230,22 +1303,6 @@ export class Session { await this.handleFetchAgent(msg.agentId, msg.requestId); break; - case "subscribe_agent_updates": - this.agentUpdatesSubscription = { - subscriptionId: msg.subscriptionId, - filter: msg.filter, - }; - await this.emitCurrentAgentUpdatesForSubscription(); - break; - - case "unsubscribe_agent_updates": - if ( - this.agentUpdatesSubscription?.subscriptionId === msg.subscriptionId - ) { - this.agentUpdatesSubscription = null; - } - break; - case "delete_agent_request": await this.handleDeleteAgentRequest(msg.agentId, msg.requestId); break; @@ -1702,9 +1759,9 @@ export class Session { }); if (this.agentUpdatesSubscription) { - this.emit({ - type: "agent_update", - payload: { kind: "remove", agentId }, + this.bufferOrEmitAgentUpdate(this.agentUpdatesSubscription, { + kind: "remove", + agentId, }); } } @@ -5097,28 +5154,11 @@ export class Session { }> { const filter = request.filter; const sort = this.normalizeFetchAgentsSort(request.sort); - const includeArchived = filter?.includeArchived ?? false; - let agents = await this.listAgentPayloads({ + const agents = await this.listAgentPayloads({ labels: filter?.labels, }); - if (!includeArchived) { - agents = agents.filter((agent) => !agent.archivedAt); - } - - if (filter?.statuses && filter.statuses.length > 0) { - const statuses = new Set(filter.statuses); - agents = agents.filter((agent) => statuses.has(agent.status)); - } - - if (typeof filter?.requiresAttention === "boolean") { - agents = agents.filter( - (agent) => - (agent.requiresAttention ?? false) === filter.requiresAttention - ); - } - const placementByCwd = new Map>(); const getPlacement = (cwd: string): Promise => { const existing = placementByCwd.get(cwd); @@ -5136,11 +5176,13 @@ export class Session { project: await getPlacement(agent.cwd), })) ); - - if (filter?.projectKeys && filter.projectKeys.length > 0) { - const projectKeys = new Set(filter.projectKeys.filter((item) => item.trim().length > 0)); - entries = entries.filter((entry) => projectKeys.has(entry.project.projectKey)); - } + entries = entries.filter((entry) => + this.matchesAgentFilter({ + agent: entry.agent, + project: entry.project, + filter, + }) + ); entries.sort((left, right) => this.compareFetchAgentsEntries(left, right, sort) @@ -5178,16 +5220,55 @@ export class Session { private async handleFetchAgents( request: Extract ): Promise { + const requestedSubscriptionId = request.subscribe?.subscriptionId?.trim(); + const subscriptionId = + request.subscribe + ? requestedSubscriptionId && requestedSubscriptionId.length > 0 + ? requestedSubscriptionId + : uuidv4() + : null; + try { + if (subscriptionId) { + this.agentUpdatesSubscription = { + subscriptionId, + filter: request.filter, + isBootstrapping: true, + pendingUpdatesByAgentId: new Map(), + }; + } + const payload = await this.listFetchAgentsEntries(request); + const snapshotUpdatedAtByAgentId = new Map(); + for (const entry of payload.entries) { + const parsedUpdatedAt = Date.parse(entry.agent.updatedAt); + if (!Number.isNaN(parsedUpdatedAt)) { + snapshotUpdatedAtByAgentId.set(entry.agent.id, parsedUpdatedAt); + } + } + this.emit({ type: "fetch_agents_response", payload: { requestId: request.requestId, + ...(subscriptionId ? { subscriptionId } : {}), ...payload, }, }); + + if ( + subscriptionId && + this.agentUpdatesSubscription?.subscriptionId === subscriptionId + ) { + this.flushBootstrappedAgentUpdates({ snapshotUpdatedAtByAgentId }); + } } catch (error) { + if ( + subscriptionId && + this.agentUpdatesSubscription?.subscriptionId === subscriptionId + ) { + this.agentUpdatesSubscription = null; + } const code = error instanceof SessionRequestError ? error.code : "fetch_agents_failed"; const message = diff --git a/packages/server/src/server/test-utils/daemon-test-context.ts b/packages/server/src/server/test-utils/daemon-test-context.ts index f966547a8..59520b82f 100644 --- a/packages/server/src/server/test-utils/daemon-test-context.ts +++ b/packages/server/src/server/test-utils/daemon-test-context.ts @@ -43,7 +43,7 @@ export async function createDaemonTestContext( url: `ws://127.0.0.1:${daemon.port}/ws`, }); await client.connect(); - client.subscribeAgentUpdates({ subscriptionId: "test" }); + await client.fetchAgents({ subscribe: { subscriptionId: "test" } }); return { daemon, diff --git a/packages/server/src/shared/messages.ts b/packages/server/src/shared/messages.ts index a3d5b94b1..b4c0a5b5b 100644 --- a/packages/server/src/shared/messages.ts +++ b/packages/server/src/shared/messages.ts @@ -405,26 +405,12 @@ export const AudioPlayedMessageSchema = z.object({ id: z.string(), }); -export const RequestAgentListMessageSchema = z.object({ - type: z.literal("request_agent_list"), - requestId: z.string(), - filter: z.object({ - labels: z.record(z.string()).optional(), - }).optional(), -}); - -export const SubscribeAgentUpdatesMessageSchema = z.object({ - type: z.literal("subscribe_agent_updates"), - subscriptionId: z.string(), - filter: z.object({ - labels: z.record(z.string()).optional(), - agentId: z.string().optional(), - }).optional(), -}); - -export const UnsubscribeAgentUpdatesMessageSchema = z.object({ - type: z.literal("unsubscribe_agent_updates"), - subscriptionId: z.string(), +const AgentDirectoryFilterSchema = z.object({ + labels: z.record(z.string()).optional(), + projectKeys: z.array(z.string()).optional(), + statuses: z.array(AgentStatusSchema).optional(), + includeArchived: z.boolean().optional(), + requiresAttention: z.boolean().optional(), }); export const DeleteAgentRequestMessageSchema = z.object({ @@ -472,15 +458,7 @@ export const SendAgentMessageSchema = z.object({ export const FetchAgentsRequestMessageSchema = z.object({ type: z.literal("fetch_agents_request"), requestId: z.string(), - filter: z - .object({ - labels: z.record(z.string()).optional(), - projectKeys: z.array(z.string()).optional(), - statuses: z.array(AgentStatusSchema).optional(), - includeArchived: z.boolean().optional(), - requiresAttention: z.boolean().optional(), - }) - .optional(), + filter: AgentDirectoryFilterSchema.optional(), sort: z .array( z.object({ @@ -495,6 +473,11 @@ export const FetchAgentsRequestMessageSchema = z.object({ cursor: z.string().min(1).optional(), }) .optional(), + subscribe: z + .object({ + subscriptionId: z.string().optional(), + }) + .optional(), }); export const FetchAgentRequestMessageSchema = z.object({ @@ -1045,8 +1028,6 @@ export const SessionInboundMessageSchema = z.discriminatedUnion("type", [ AudioPlayedMessageSchema, FetchAgentsRequestMessageSchema, FetchAgentRequestMessageSchema, - SubscribeAgentUpdatesMessageSchema, - UnsubscribeAgentUpdatesMessageSchema, DeleteAgentRequestMessageSchema, ArchiveAgentRequestMessageSchema, UpdateAgentRequestMessageSchema, @@ -1443,6 +1424,7 @@ export const FetchAgentsResponseMessageSchema = z.object({ type: z.literal("fetch_agents_response"), payload: z.object({ requestId: z.string(), + subscriptionId: z.string().nullable().optional(), entries: z.array( z.object({ agent: AgentSnapshotPayloadSchema,