diff --git a/packages/app/src/hooks/checkout-status-revalidation.ts b/packages/app/src/hooks/checkout-status-revalidation.ts new file mode 100644 index 000000000..02f379c67 --- /dev/null +++ b/packages/app/src/hooks/checkout-status-revalidation.ts @@ -0,0 +1,23 @@ +export type CheckoutStatusRevalidationParams = { + serverId: string; + agentId: string; + isOpen: boolean; + explorerTab: string; +}; + +export function checkoutStatusRevalidationKey(params: CheckoutStatusRevalidationParams): string | null { + if (!params.agentId) return null; + if (!params.isOpen) return null; + if (params.explorerTab !== "changes") return null; + return `${params.serverId}:${params.agentId}`; +} + +export function nextCheckoutStatusRefetchDecision( + prevKey: string | null, + nextKey: string | null +): { nextSeenKey: string | null; shouldRefetch: boolean } { + if (!nextKey) return { nextSeenKey: null, shouldRefetch: false }; + if (prevKey === nextKey) return { nextSeenKey: prevKey, shouldRefetch: false }; + return { nextSeenKey: nextKey, shouldRefetch: true }; +} + diff --git a/packages/app/src/hooks/use-checkout-status-query.test.ts b/packages/app/src/hooks/use-checkout-status-query.test.ts new file mode 100644 index 000000000..50044202d --- /dev/null +++ b/packages/app/src/hooks/use-checkout-status-query.test.ts @@ -0,0 +1,74 @@ +import { describe, expect, it } from "vitest"; +import { checkoutStatusRevalidationKey, nextCheckoutStatusRefetchDecision } from "./checkout-status-revalidation"; + +describe("useCheckoutStatusQuery", () => { + describe("checkoutStatusRevalidationKey", () => { + it("returns null when sidebar is closed", () => { + expect( + checkoutStatusRevalidationKey({ + serverId: "daemon-1", + agentId: "agent-1", + isOpen: false, + explorerTab: "changes", + }) + ).toBeNull(); + }); + + it("returns null when tab is not changes", () => { + expect( + checkoutStatusRevalidationKey({ + serverId: "daemon-1", + agentId: "agent-1", + isOpen: true, + explorerTab: "files", + }) + ).toBeNull(); + }); + + it("returns a stable key when open and on changes tab", () => { + expect( + checkoutStatusRevalidationKey({ + serverId: "daemon-1", + agentId: "agent-1", + isOpen: true, + explorerTab: "changes", + }) + ).toBe("daemon-1:agent-1"); + }); + }); + + describe("nextCheckoutStatusRefetchDecision", () => { + it("refetches only once per key until reset", () => { + const key = "daemon-1:agent-1"; + + expect(nextCheckoutStatusRefetchDecision(null, key)).toEqual({ + nextSeenKey: key, + shouldRefetch: true, + }); + + expect(nextCheckoutStatusRefetchDecision(key, key)).toEqual({ + nextSeenKey: key, + shouldRefetch: false, + }); + + expect(nextCheckoutStatusRefetchDecision(key, null)).toEqual({ + nextSeenKey: null, + shouldRefetch: false, + }); + + expect(nextCheckoutStatusRefetchDecision(null, key)).toEqual({ + nextSeenKey: key, + shouldRefetch: true, + }); + }); + + it("refetches again when agent changes while active", () => { + expect( + nextCheckoutStatusRefetchDecision("daemon-1:agent-1", "daemon-1:agent-2") + ).toEqual({ + nextSeenKey: "daemon-1:agent-2", + shouldRefetch: true, + }); + }); + }); +}); diff --git a/packages/app/src/hooks/use-checkout-status-query.ts b/packages/app/src/hooks/use-checkout-status-query.ts index bd0044e09..953055bf7 100644 --- a/packages/app/src/hooks/use-checkout-status-query.ts +++ b/packages/app/src/hooks/use-checkout-status-query.ts @@ -1,9 +1,13 @@ import { useQuery } from "@tanstack/react-query"; -import { useEffect } from "react"; +import { useEffect, useMemo, useRef } from "react"; import { UnistylesRuntime } from "react-native-unistyles"; import { useSessionStore } from "@/stores/session-store"; import { usePanelStore } from "@/stores/panel-store"; import type { CheckoutStatusResponse } from "@server/shared/messages"; +import { + checkoutStatusRevalidationKey, + nextCheckoutStatusRefetchDecision, +} from "./checkout-status-revalidation"; const CHECKOUT_STATUS_STALE_TIME = 15_000; @@ -48,12 +52,17 @@ export function useCheckoutStatusQuery({ serverId, agentId }: UseCheckoutStatusQ }); // Revalidate when sidebar is open with "changes" tab active. + const revalidationKey = useMemo( + () => checkoutStatusRevalidationKey({ serverId, agentId, isOpen, explorerTab }), + [serverId, agentId, isOpen, explorerTab] + ); + const lastRevalidationKey = useRef(null); useEffect(() => { - if (!isOpen || explorerTab !== "changes" || !agentId) { - return; - } + const decision = nextCheckoutStatusRefetchDecision(lastRevalidationKey.current, revalidationKey); + lastRevalidationKey.current = decision.nextSeenKey; + if (!decision.shouldRefetch) return; void query.refetch(); - }, [isOpen, explorerTab, agentId, query]); + }, [revalidationKey, query.refetch]); return { status: query.data ?? null, diff --git a/packages/server/src/client/daemon-client-v2.test.ts b/packages/server/src/client/daemon-client-v2.test.ts new file mode 100644 index 000000000..f3402d41f --- /dev/null +++ b/packages/server/src/client/daemon-client-v2.test.ts @@ -0,0 +1,138 @@ +import { afterEach, describe, expect, test, vi } from "vitest"; +import { DaemonClientV2, type DaemonTransport } from "./daemon-client-v2"; + +function createMockLogger() { + return { + debug: vi.fn(), + info: vi.fn(), + warn: vi.fn(), + error: vi.fn(), + }; +} + +function createMockTransport() { + const sent: string[] = []; + + let onMessage: (data: unknown) => void = () => {}; + let onOpen: () => void = () => {}; + let onClose: (_event?: unknown) => void = () => {}; + let onError: (_event?: unknown) => void = () => {}; + + const transport: DaemonTransport = { + send: (data) => sent.push(data), + close: () => {}, + onMessage: (handler) => { + onMessage = handler; + return () => {}; + }, + onOpen: (handler) => { + onOpen = handler; + return () => {}; + }, + onClose: (handler) => { + onClose = handler; + return () => {}; + }, + onError: (handler) => { + onError = handler; + return () => {}; + }, + }; + + return { + transport, + sent, + triggerOpen: () => onOpen(), + triggerClose: (event?: unknown) => onClose(event), + triggerError: (event?: unknown) => onError(event), + triggerMessage: (data: unknown) => onMessage(data), + }; +} + +describe("DaemonClientV2", () => { + const clients: DaemonClientV2[] = []; + + afterEach(async () => { + for (const client of clients) { + await client.close(); + } + clients.length = 0; + }); + + test("dedupes in-flight checkout status requests per agentId", async () => { + const logger = createMockLogger(); + const mock = createMockTransport(); + + const client = new DaemonClientV2({ + url: "ws://test", + logger, + reconnect: { enabled: false }, + transportFactory: () => mock.transport, + }); + clients.push(client); + + const connectPromise = client.connect(); + mock.triggerOpen(); + await connectPromise; + + const p1 = client.getCheckoutStatus("agent-1"); + const p2 = client.getCheckoutStatus("agent-1"); + + expect(mock.sent).toHaveLength(1); + + const request = JSON.parse(mock.sent[0]) as { + type: "session"; + message: { type: "checkout_status_request"; agentId: string; requestId: string }; + }; + + const response = { + type: "session", + message: { + type: "checkout_status_response", + payload: { + agentId: "agent-1", + cwd: "/tmp", + error: null, + requestId: request.message.requestId, + isGit: false, + isPaseoOwnedWorktree: false, + repoRoot: null, + currentBranch: null, + isDirty: null, + baseRef: null, + aheadBehind: null, + }, + }, + }; + + mock.triggerMessage(JSON.stringify(response)); + const [r1, r2] = await Promise.all([p1, p2]); + expect(r1).toMatchObject({ agentId: "agent-1", requestId: request.message.requestId, isGit: false }); + expect(r2).toMatchObject({ agentId: "agent-1", requestId: request.message.requestId, isGit: false }); + + // After completion, a new call should issue a new request. + const p3 = client.getCheckoutStatus("agent-1"); + expect(mock.sent).toHaveLength(2); + + const request2 = JSON.parse(mock.sent[1]) as { + type: "session"; + message: { type: "checkout_status_request"; agentId: string; requestId: string }; + }; + + mock.triggerMessage( + JSON.stringify({ + ...response, + message: { + ...response.message, + payload: { ...response.message.payload, requestId: request2.message.requestId }, + }, + }) + ); + + await expect(p3).resolves.toMatchObject({ + agentId: "agent-1", + requestId: request2.message.requestId, + isGit: false, + }); + }); +}); diff --git a/packages/server/src/client/daemon-client-v2.ts b/packages/server/src/client/daemon-client-v2.ts index 81916a34e..77fea305b 100644 --- a/packages/server/src/client/daemon-client-v2.ts +++ b/packages/server/src/client/daemon-client-v2.ts @@ -227,6 +227,7 @@ export class DaemonClientV2 { > = new Map(); private eventListeners: Set = new Set(); private waiters: Set> = new Set(); + private checkoutStatusInFlight: Map> = new Map(); private connectionListeners: Set<(status: ConnectionState) => void> = new Set(); private reconnectTimeout: ReturnType | null = null; @@ -945,27 +946,48 @@ export class DaemonClientV2 { agentId: string, requestId?: string ): Promise { + if (!requestId) { + const existing = this.checkoutStatusInFlight.get(agentId); + if (existing) { + return existing; + } + } + const resolvedRequestId = this.createRequestId(requestId); const message = SessionInboundMessageSchema.parse({ type: "checkout_status_request", agentId, requestId: resolvedRequestId, }); - const response = this.waitFor( - (msg) => { - if (msg.type !== "checkout_status_response") { - return null; + + const responsePromise = (async () => { + const response = this.waitFor( + (msg) => { + if (msg.type !== "checkout_status_response") { + return null; + } + if (msg.payload.requestId !== resolvedRequestId) { + return null; + } + return msg.payload; + }, + 60000, + { skipQueue: true } + ); + this.sendSessionMessage(message); + return response; + })(); + + if (!requestId) { + this.checkoutStatusInFlight.set(agentId, responsePromise); + responsePromise.finally(() => { + if (this.checkoutStatusInFlight.get(agentId) === responsePromise) { + this.checkoutStatusInFlight.delete(agentId); } - if (msg.payload.requestId !== resolvedRequestId) { - return null; - } - return msg.payload; - }, - 60000, - { skipQueue: true } - ); - this.sendSessionMessage(message); - return response; + }); + } + + return responsePromise; } async getCheckoutDiff(