Push-driven checkout diff subscriptions for changes sidebar (#24)

* Implement push-driven checkout diff subscriptions

* Remove legacy checkout diff polling RPC

* Fix checkout diff watcher coverage and fallback refresh

packages/server/src/client/daemon-client.test.ts

@@ -151,6 +151,185 @@ describe("DaemonClient", () => {
     });
   });
 
+  test("subscribes to checkout diff updates via RPC handshake", async () => {
+    const logger = createMockLogger();
+    const mock = createMockTransport();
+
+    const client = new DaemonClient({
+      url: "ws://test",
+      logger,
+      reconnect: { enabled: false },
+      transportFactory: () => mock.transport,
+    });
+    clients.push(client);
+
+    const connectPromise = client.connect();
+    mock.triggerOpen();
+    await connectPromise;
+
+    const promise = client.subscribeCheckoutDiff(
+      "/tmp/project",
+      { mode: "uncommitted" },
+      { subscriptionId: "checkout-sub-1" }
+    );
+
+    expect(mock.sent).toHaveLength(1);
+    const request = JSON.parse(mock.sent[0]) as {
+      type: "session";
+      message: {
+        type: "subscribe_checkout_diff_request";
+        subscriptionId: string;
+        cwd: string;
+        compare: { mode: "uncommitted" | "base"; baseRef?: string };
+        requestId: string;
+      };
+    };
+    expect(request.message.type).toBe("subscribe_checkout_diff_request");
+    expect(request.message.subscriptionId).toBe("checkout-sub-1");
+    expect(request.message.cwd).toBe("/tmp/project");
+    expect(request.message.compare).toEqual({ mode: "uncommitted" });
+
+    mock.triggerMessage(
+      JSON.stringify({
+        type: "session",
+        message: {
+          type: "subscribe_checkout_diff_response",
+          payload: {
+            subscriptionId: "checkout-sub-1",
+            cwd: "/tmp/project",
+            files: [],
+            error: null,
+            requestId: request.message.requestId,
+          },
+        },
+      })
+    );
+
+    await expect(promise).resolves.toEqual({
+      subscriptionId: "checkout-sub-1",
+      cwd: "/tmp/project",
+      files: [],
+      error: null,
+      requestId: request.message.requestId,
+    });
+  });
+
+  test("getCheckoutDiff uses one-shot subscription protocol", async () => {
+    const logger = createMockLogger();
+    const mock = createMockTransport();
+
+    const client = new DaemonClient({
+      url: "ws://test",
+      logger,
+      reconnect: { enabled: false },
+      transportFactory: () => mock.transport,
+    });
+    clients.push(client);
+
+    const connectPromise = client.connect();
+    mock.triggerOpen();
+    await connectPromise;
+
+    const promise = client.getCheckoutDiff("/tmp/project", { mode: "base", baseRef: "main" });
+
+    expect(mock.sent).toHaveLength(1);
+    const subscribeRequest = JSON.parse(mock.sent[0]) as {
+      type: "session";
+      message: {
+        type: "subscribe_checkout_diff_request";
+        subscriptionId: string;
+        cwd: string;
+        compare: { mode: "uncommitted" | "base"; baseRef?: string };
+        requestId: string;
+      };
+    };
+    expect(subscribeRequest.message.type).toBe("subscribe_checkout_diff_request");
+    expect(subscribeRequest.message.cwd).toBe("/tmp/project");
+    expect(subscribeRequest.message.compare).toEqual({ mode: "base", baseRef: "main" });
+
+    mock.triggerMessage(
+      JSON.stringify({
+        type: "session",
+        message: {
+          type: "subscribe_checkout_diff_response",
+          payload: {
+            subscriptionId: subscribeRequest.message.subscriptionId,
+            cwd: "/tmp/project",
+            files: [],
+            error: null,
+            requestId: subscribeRequest.message.requestId,
+          },
+        },
+      })
+    );
+
+    await expect(promise).resolves.toEqual({
+      cwd: "/tmp/project",
+      files: [],
+      error: null,
+      requestId: subscribeRequest.message.requestId,
+    });
+
+    expect(mock.sent).toHaveLength(2);
+    const unsubscribeRequest = JSON.parse(mock.sent[1]) as {
+      type: "session";
+      message: {
+        type: "unsubscribe_checkout_diff_request";
+        subscriptionId: string;
+      };
+    };
+    expect(unsubscribeRequest.message.type).toBe("unsubscribe_checkout_diff_request");
+    expect(unsubscribeRequest.message.subscriptionId).toBe(
+      subscribeRequest.message.subscriptionId
+    );
+  });
+
+  test("resubscribes checkout diff streams after reconnect", async () => {
+    const logger = createMockLogger();
+    const mock = createMockTransport();
+
+    const client = new DaemonClient({
+      url: "ws://test",
+      logger,
+      reconnect: { enabled: false },
+      transportFactory: () => mock.transport,
+    });
+    clients.push(client);
+
+    const internal = client as unknown as {
+      checkoutDiffSubscriptions: Map<
+        string,
+        { cwd: string; compare: { mode: "uncommitted" | "base"; baseRef?: string } }
+      >;
+    };
+    internal.checkoutDiffSubscriptions.set("checkout-sub-1", {
+      cwd: "/tmp/project",
+      compare: { mode: "base", baseRef: "main" },
+    });
+
+    const connectPromise = client.connect();
+    mock.triggerOpen();
+    await connectPromise;
+
+    expect(mock.sent).toHaveLength(1);
+    const request = JSON.parse(mock.sent[0]) as {
+      type: "session";
+      message: {
+        type: "subscribe_checkout_diff_request";
+        subscriptionId: string;
+        cwd: string;
+        compare: { mode: "uncommitted" | "base"; baseRef?: string };
+        requestId: string;
+      };
+    };
+    expect(request.message.type).toBe("subscribe_checkout_diff_request");
+    expect(request.message.subscriptionId).toBe("checkout-sub-1");
+    expect(request.message.cwd).toBe("/tmp/project");
+    expect(request.message.compare).toEqual({ mode: "base", baseRef: "main" });
+    expect(typeof request.message.requestId).toBe("string");
+    expect(request.message.requestId.length).toBeGreaterThan(0);
+  });
+
   test("fetches project-grouped agents via RPC", async () => {
     const logger = createMockLogger();
     const mock = createMockTransport();

packages/server/src/client/daemon-client.ts

@@ -19,7 +19,6 @@ import type {
   GitSetupOptions,
   HighlightedDiffResponse,
   CheckoutStatusResponse,
-  CheckoutDiffResponse,
   CheckoutCommitResponse,
   CheckoutMergeResponse,
   CheckoutMergeFromBaseResponse,
@@ -182,7 +181,11 @@ export type CreateAgentRequestOptions = {
 type GitDiffPayload = GitDiffResponse["payload"];
 type HighlightedDiffPayload = HighlightedDiffResponse["payload"];
 type CheckoutStatusPayload = CheckoutStatusResponse["payload"];
-type CheckoutDiffPayload = CheckoutDiffResponse["payload"];
+type SubscribeCheckoutDiffPayload = Extract<
+  SessionOutboundMessage,
+  { type: "subscribe_checkout_diff_response" }
+>["payload"];
+type CheckoutDiffPayload = Omit<SubscribeCheckoutDiffPayload, "subscriptionId">;
 type CheckoutCommitPayload = CheckoutCommitResponse["payload"];
 type CheckoutMergePayload = CheckoutMergeResponse["payload"];
 type CheckoutMergeFromBasePayload = CheckoutMergeFromBaseResponse["payload"];
@@ -296,6 +299,10 @@ export class DaemonClient {
     string,
     { labels?: Record<string, string>; agentId?: string } | undefined
   >();
+  private checkoutDiffSubscriptions = new Map<
+    string,
+    { cwd: string; compare: { mode: "uncommitted" | "base"; baseRef?: string } }
+  >();
   private logger: Logger;
   private pendingSendQueue: PendingSend[] = [];
   private relayClientId: string | null = null;
@@ -401,6 +408,7 @@ export class DaemonClient {
           this.reconnectAttempt = 0;
           this.updateConnectionState({ status: "connected" });
           this.resubscribeAgentUpdates();
+          this.resubscribeCheckoutDiffSubscriptions();
           this.flushPendingSendQueue();
           this.resolveConnect();
         }),
@@ -951,6 +959,22 @@ export class DaemonClient {
     }
   }
 
+  private resubscribeCheckoutDiffSubscriptions(): void {
+    if (this.checkoutDiffSubscriptions.size === 0) {
+      return;
+    }
+    for (const [subscriptionId, subscription] of this.checkoutDiffSubscriptions) {
+      const message = SessionInboundMessageSchema.parse({
+        type: "subscribe_checkout_diff_request",
+        subscriptionId,
+        cwd: subscription.cwd,
+        compare: subscription.compare,
+        requestId: this.createRequestId(),
+      });
+      this.sendSessionMessage(message);
+    }
+  }
+
   // ============================================================================
   // Agent Lifecycle
   // ============================================================================
@@ -1521,32 +1545,101 @@ export class DaemonClient {
     return responsePromise;
   }
 
+  private normalizeCheckoutDiffCompare(
+    compare: { mode: "uncommitted" | "base"; baseRef?: string }
+  ): { mode: "uncommitted" | "base"; baseRef?: string } {
+    if (compare.mode === "uncommitted") {
+      return { mode: "uncommitted" };
+    }
+    const trimmedBaseRef = compare.baseRef?.trim();
+    if (!trimmedBaseRef) {
+      return { mode: "base" };
+    }
+    return { mode: "base", baseRef: trimmedBaseRef };
+  }
+
   async getCheckoutDiff(
     cwd: string,
     compare: { mode: "uncommitted" | "base"; baseRef?: string },
     requestId?: string
   ): Promise<CheckoutDiffPayload> {
-    const resolvedRequestId = this.createRequestId(requestId);
-    const message = SessionInboundMessageSchema.parse({
-      type: "checkout_diff_request",
+    const oneShotSubscriptionId = `oneshot-checkout-diff:${crypto.randomUUID()}`;
+    try {
+      const payload = await this.subscribeCheckoutDiff(cwd, compare, {
+        subscriptionId: oneShotSubscriptionId,
+        requestId,
+      });
+      return {
+        cwd: payload.cwd,
+        files: payload.files,
+        error: payload.error,
+        requestId: payload.requestId,
+      };
+    } finally {
+      try {
+        this.unsubscribeCheckoutDiff(oneShotSubscriptionId);
+      } catch {
+        // Ignore disconnect races during one-shot cleanup.
+      }
+    }
+  }
+
+  async subscribeCheckoutDiff(
+    cwd: string,
+    compare: { mode: "uncommitted" | "base"; baseRef?: string },
+    options?: { subscriptionId?: string; requestId?: string }
+  ): Promise<SubscribeCheckoutDiffPayload> {
+    const subscriptionId = options?.subscriptionId ?? crypto.randomUUID();
+    const normalizedCompare = this.normalizeCheckoutDiffCompare(compare);
+    const previousSubscription = this.checkoutDiffSubscriptions.get(subscriptionId) ?? null;
+    this.checkoutDiffSubscriptions.set(subscriptionId, {
       cwd,
-      compare,
+      compare: normalizedCompare,
+    });
+
+    const resolvedRequestId = this.createRequestId(options?.requestId);
+    const message = SessionInboundMessageSchema.parse({
+      type: "subscribe_checkout_diff_request",
+      subscriptionId,
+      cwd,
+      compare: normalizedCompare,
       requestId: resolvedRequestId,
     });
-    return this.sendRequest({
-      requestId: resolvedRequestId,
-      message,
-      timeout: 60000,
-      options: { skipQueue: true },
-      select: (msg) => {
-        if (msg.type !== "checkout_diff_response") {
-          return null;
-        }
-        if (msg.payload.requestId !== resolvedRequestId) {
-          return null;
-        }
-        return msg.payload;
-      },
+
+    try {
+      return await this.sendRequest({
+        requestId: resolvedRequestId,
+        message,
+        timeout: 60000,
+        options: { skipQueue: true },
+        select: (msg) => {
+          if (msg.type !== "subscribe_checkout_diff_response") {
+            return null;
+          }
+          if (msg.payload.requestId !== resolvedRequestId) {
+            return null;
+          }
+          if (msg.payload.subscriptionId !== subscriptionId) {
+            return null;
+          }
+          return msg.payload;
+        },
+      });
+    } catch (error) {
+      if (previousSubscription) {
+        this.checkoutDiffSubscriptions.set(subscriptionId, previousSubscription);
+      } else {
+        this.checkoutDiffSubscriptions.delete(subscriptionId);
+      }
+      throw error;
+    }
+  }
+
+  unsubscribeCheckoutDiff(subscriptionId: string): void {
+    this.checkoutDiffSubscriptions.delete(subscriptionId);
+    this.sendSessionMessage({
+      type: "unsubscribe_checkout_diff_request",
+      subscriptionId,
     });
   }