fix(web): update machines together in auto balance

apps/web/src/components/ChatView.tsx

@@ -437,6 +437,7 @@ import {
 } from "./ui/alert-dialog";
 import { Tooltip, TooltipPopup, TooltipTrigger } from "./ui/tooltip";
 import { ServerUpdateAction } from "./ServerUpdateAction";
+import { useAutoBalanceUpdateBanner } from "./chat/useAutoBalanceUpdateBanner";
 import {
   ComposerServerUpdateIcon,
   ComposerServerUpdateStatus,
@@ -2314,6 +2315,35 @@ export default function ChatView(props: ChatViewProps) {
     advertisedFileAttachmentBytes === null
       ? null
       : clampFileAttachmentUploadBytes(advertisedFileAttachmentBytes);
+  const envLocked = Boolean(
+    activeThread &&
+    (activeThread.messages.length > 0 ||
+      (activeThread.session !== null && activeThread.session.status !== "stopped")),
+  );
+
+  const loadBalancingSettings = useClientSettings();
+  const automaticEnvironment = Boolean(
+    clientSettingsHydrated &&
+    draftId &&
+    !envLocked &&
+    hasMultipleEnvironments &&
+    loadBalancingSettings.loadBalancingEnabled &&
+    draftThread?.environmentSelection !== "manual" &&
+    (!composerHasAttachments || Boolean(draftThread?.loadBalancedEnvironmentId)) &&
+    (!draftThread?.branch || draftThread.environmentSelection === "auto") &&
+    !draftThread?.worktreePath,
+  );
+  const autoUpdateEnvironments = useMemo(
+    () =>
+      automaticEnvironment
+        ? logicalProjectEnvironments.flatMap(({ environmentId }) => {
+            const environment = environmentById.get(environmentId);
+            return environment ? [environment] : [];
+          })
+        : [],
+    [automaticEnvironment, logicalProjectEnvironments, environmentById],
+  );
+  const autoBalanceUpdateBanner = useAutoBalanceUpdateBanner(autoUpdateEnvironments);
   const versionMismatch = resolveServerConfigVersionMismatch(serverConfig);
   const versionMismatchDismissKey =
     versionMismatch && activeThread
@@ -2417,6 +2447,7 @@ export default function ChatView(props: ChatViewProps) {
       }
     }
     if (
+      !automaticEnvironment &&
       serverUpdateEnvironmentId &&
       !reconnectingThroughVersionSkew &&
       (serverUpdateState.status === "idle"
@@ -2495,8 +2526,11 @@ export default function ChatView(props: ChatViewProps) {
             }),
       });
     }
+    if (autoBalanceUpdateBanner) items.push(autoBalanceUpdateBanner);
     return items;
   }, [
+    automaticEnvironment,
+    autoBalanceUpdateBanner,
     activeEnvironmentUnavailableState,
     reconnectWarningGraceElapsed,
     handleReconnectActiveEnvironment,
@@ -3230,24 +3264,6 @@ export default function ChatView(props: ChatViewProps) {
     }
   }, [activeThreadRef, diffOpen, isServerThread, onDiffPanelOpen]);
 
-  const envLocked = Boolean(
-    activeThread &&
-    (activeThread.messages.length > 0 ||
-      (activeThread.session !== null && activeThread.session.status !== "stopped")),
-  );
-
-  const loadBalancingSettings = useClientSettings();
-  const automaticEnvironment = Boolean(
-    clientSettingsHydrated &&
-    draftId &&
-    !envLocked &&
-    hasMultipleEnvironments &&
-    loadBalancingSettings.loadBalancingEnabled &&
-    draftThread?.environmentSelection !== "manual" &&
-    (!composerHasAttachments || Boolean(draftThread?.loadBalancedEnvironmentId)) &&
-    (!draftThread?.branch || draftThread.environmentSelection === "auto") &&
-    !draftThread?.worktreePath,
-  );
   const needsLoadBalancing = automaticEnvironment && !draftThread?.loadBalancedEnvironmentId;
   const loadBalancingCandidates = useMemo(
     () =>

apps/web/src/components/ServerUpdateAction.test.tsx

@@ -1,9 +1,10 @@
-import type { ReactElement } from "react";
+import { act, type ReactElement } from "react";
+import { create, type ReactTestRenderer } from "react-test-renderer";
 import { renderToStaticMarkup } from "react-dom/server";
 import type { EnvironmentId } from "@t3tools/contracts";
 import * as Cause from "effect/Cause";
 import { AsyncResult } from "effect/unstable/reactivity";
-import { beforeEach, describe, expect, it, vi } from "vite-plus/test";
+import { afterEach, beforeEach, describe, expect, it, vi } from "vite-plus/test";
 
 const testState = vi.hoisted(() => ({
   updateServer: vi.fn(),
@@ -30,7 +31,18 @@ vi.mock("./ui/toast", () => ({
   toastManager: { add: testState.toast },
 }));
 
-import { ServerUpdateAction, ServerUpdateProgress } from "./ServerUpdateAction";
+import {
+  readConfirmDialogState,
+  registerConfirmDialogHost,
+  resetConfirmDialogForTests,
+  respondToConfirmDialog,
+} from "~/confirmDialog";
+import {
+  ServerUpdateAction,
+  ServerUpdateProgress,
+  ServerUpdatesAction,
+  type ServerUpdateTarget,
+} from "./ServerUpdateAction";
 
 type ActionElement = ReactElement<{
   readonly onClick?: () => void;
@@ -194,6 +206,149 @@ describe("ServerUpdateAction", () => {
   });
 });
 
+describe("ServerUpdatesAction", () => {
+  let renderer: ReactTestRenderer | undefined;
+  const targets: ReadonlyArray<ServerUpdateTarget> = [
+    {
+      environmentId: "batch-a" as EnvironmentId,
+      serverLabel: "Laptop",
+      selfUpdate: "boot-service",
+      targetVersion: "0.0.31",
+      threadContinuation: true,
+      continueThreadsAfterServerUpdate: true,
+    },
+    {
+      environmentId: "batch-b" as EnvironmentId,
+      serverLabel: "Office",
+      selfUpdate: "respawn",
+      targetVersion: "0.0.31",
+      threadContinuation: true,
+      continueThreadsAfterServerUpdate: false,
+    },
+    {
+      environmentId: "batch-c" as EnvironmentId,
+      serverLabel: "Manual",
+      selfUpdate: null,
+      targetVersion: "0.0.31",
+    },
+  ];
+  const success = AsyncResult.success({ targetVersion: "0.0.31", method: "boot-service" as const });
+
+  async function mount(batch = targets) {
+    await act(async () => {
+      renderer = create(<ServerUpdatesAction targets={batch} />);
+    });
+    return renderer!.root.findByType("button");
+  }
+
+  beforeEach(() => {
+    vi.stubGlobal("IS_REACT_ACT_ENVIRONMENT", true);
+    testState.updateServer.mockReset();
+    testState.toast.mockReset();
+    resetConfirmDialogForTests();
+  });
+  afterEach(async () => {
+    await act(async () => {
+      renderer?.unmount();
+    });
+    renderer = undefined;
+    resetConfirmDialogForTests();
+    vi.unstubAllGlobals();
+  });
+
+  it("updates both supported machines with their own continuation preference and skips the manual machine", async () => {
+    testState.updateServer.mockResolvedValue(success);
+    const button = await mount();
+    await act(async () => {
+      button.props.onClick();
+    });
+
+    expect(testState.updateServer.mock.calls.map(([target]) => target)).toEqual([
+      {
+        environmentId: "batch-a",
+        input: { targetVersion: "0.0.31", continueRunningThreads: true },
+      },
+      { environmentId: "batch-b", input: { targetVersion: "0.0.31" } },
+    ]);
+    expect(testState.toast.mock.calls.map(([toast]) => toast.title)).toEqual([
+      "Laptop updated",
+      "Office updated",
+    ]);
+  });
+
+  it("names a failed machine while letting the other machine complete", async () => {
+    testState.updateServer
+      .mockResolvedValueOnce(AsyncResult.failure(Cause.fail(new Error("Download failed"))))
+      .mockResolvedValueOnce(success);
+    const button = await mount();
+    await act(async () => {
+      button.props.onClick();
+    });
+
+    expect(testState.updateServer).toHaveBeenCalledTimes(2);
+    expect(testState.toast).toHaveBeenCalledWith({
+      type: "error",
+      title: "Laptop update failed",
+      description: "Download failed",
+    });
+    expect(testState.toast).toHaveBeenCalledWith(
+      expect.objectContaining({ type: "success", title: "Office updated" }),
+    );
+    expect(button.props.disabled).toBe(false);
+  });
+
+  it("starts each machine once when double-clicked and disables the action until both finish", async () => {
+    const completions: Array<() => void> = [];
+    testState.updateServer.mockImplementation(
+      () =>
+        new Promise((resolve) => {
+          completions.push(() => resolve(success));
+        }),
+    );
+    const button = await mount();
+    await act(async () => {
+      button.props.onClick();
+      button.props.onClick();
+    });
+    expect(testState.updateServer).toHaveBeenCalledTimes(2);
+    expect(button.props.disabled).toBe(true);
+    await act(async () => {
+      completions[0]!();
+    });
+    expect(button.props.disabled).toBe(true);
+    await act(async () => {
+      completions[1]!();
+    });
+    expect(button.props.disabled).toBe(false);
+    expect(testState.toast).toHaveBeenCalledTimes(2);
+  });
+
+  it("asks once for desktop machines and cancels the entire batch", async () => {
+    registerConfirmDialogHost();
+    const button = await mount(
+      targets.map((target, index) =>
+        index < 2 ? { ...target, selfUpdate: "desktop-managed", desktopAppUpdate: true } : target,
+      ),
+    );
+    await act(async () => {
+      button.props.onClick();
+    });
+    const confirmation = readConfirmDialogState();
+    expect(confirmation).toEqual(
+      expect.objectContaining({
+        status: "confirming",
+        message: expect.stringContaining("Laptop, Office"),
+      }),
+    );
+    expect(testState.updateServer).not.toHaveBeenCalled();
+    await act(async () => {
+      respondToConfirmDialog(false);
+    });
+    expect(testState.updateServer).not.toHaveBeenCalled();
+    expect(button.props.disabled).toBe(false);
+  });
+});
+
 describe("ServerUpdateProgress", () => {
   it("shows one calm status row for the restart wait", () => {
     const markup = renderToStaticMarkup(