fix(web): auto-balanced drafts recheck machine load before sending

apps/web/src/components/ChatView.tsx

@@ -1799,6 +1799,7 @@ export default function ChatView(props: ChatViewProps) {
   const {
     needsLoadBalancing,
     loadBalancing,
+    refreshBeforeSend,
     onAutoEnvironmentStable,
     autoEnvironmentLabel,
     isEnvironmentChanging,
@@ -2940,6 +2941,7 @@ export default function ChatView(props: ChatViewProps) {
     clearUsageLimitsFor,
     needsLoadBalancing,
     loadBalancing,
+    refreshBeforeSend,
     readPendingEnvironmentChange,
     pendingApprovals,
     respondingRequestIds,

apps/web/src/components/chat/useComposerSend.ts

@@ -96,6 +96,7 @@ import { environmentServerConfigsAtom } from "../../state/server";
 import { threadEnvironment } from "../../state/threads";
 import { runMobileComposerTransition } from "./draftHeroTransition";
 import { useClientSettingsHydrated } from "../../hooks/useSettings";
+import { useStableCallback } from "../../hooks/useStableCallback";
 import { environmentPresentations } from "../../state/presentation";
 import {
   approvalRedirectNotice,
@@ -216,6 +217,7 @@ interface ComposerSendInput {
   readonly clearUsageLimitsFor: UsageLimitsPanel["clearFor"];
   readonly needsLoadBalancing: boolean;
   readonly loadBalancing: DraftEnvironmentRouting["loadBalancing"];
+  readonly refreshBeforeSend: DraftEnvironmentRouting["refreshBeforeSend"];
   readonly readPendingEnvironmentChange: DraftEnvironmentRouting["readPendingEnvironmentChange"];
   readonly pendingApprovals: PendingThreadRequests["pendingApprovals"];
   readonly respondingRequestIds: PendingThreadRequests["respondingRequestIds"];
@@ -311,6 +313,7 @@ export function useComposerSend(input: ComposerSendInput) {
     clearUsageLimitsFor,
     needsLoadBalancing,
     loadBalancing,
+    refreshBeforeSend,
     readPendingEnvironmentChange,
     pendingApprovals,
     respondingRequestIds,
@@ -1488,5 +1491,19 @@ export function useComposerSend(input: ComposerSendInput) {
     }
   };
 
-  return { onSend };
+  const onSendLatest = useStableCallback(onSend);
+  const onSendWithRouting = async (...args: Parameters<typeof onSend>) => {
+    args[0]?.preventDefault();
+    const refresh =
+      args[3] || sendInFlightRef.current || readPendingEnvironmentChange() !== null
+        ? null
+        : refreshBeforeSend();
+    if (refresh === null) return onSend(...args);
+    sendInFlightRef.current = true;
+    const ready = await refresh.finally(() => {
+      sendInFlightRef.current = false;
+    });
+    if (ready) return onSendLatest(...args);
+  };
+  return { onSend: onSendWithRouting };
 }

apps/web/src/components/chat/useDraftEnvironmentRouting.browser.test.tsx

@@ -0,0 +1,227 @@
+import { PrimaryConnectionTarget } from "@t3tools/client-runtime/connection";
+import * as Option from "effect/Option";
+import type { EnvironmentPresentation } from "../../state/environments";
+import type { ServerProvider } from "@t3tools/contracts";
+import {
+  scopedThreadKey,
+  scopeProjectRef,
+  scopeThreadRef,
+} from "@t3tools/client-runtime/environment";
+import { createModelSelection } from "@t3tools/shared/model";
+import { DEFAULT_UNIFIED_SETTINGS } from "@t3tools/contracts/settings";
+import {
+  EnvironmentId,
+  ProjectId,
+  ProviderDriverKind,
+  ProviderInstanceId,
+} from "@t3tools/contracts";
+import { expect, it, vi } from "vite-plus/test";
+
+import { act, render } from "../../../test/render";
+import { DraftId, useComposerDraftStore } from "../../composerDraftStore";
+import { useDraftEnvironmentRouting } from "./useDraftEnvironmentRouting";
+import { appAtomRegistry } from "../../rpc/atomRegistry";
+import { draftFanoutSelectionsAtom } from "./chatViewSend";
+
+const resources = vi.hoisted(() => ({
+  refreshSelection: vi.fn<(ids: readonly EnvironmentId[]) => Promise<EnvironmentId | null>>(),
+}));
+vi.mock("../../hooks/useLoadBalancedEnvironment", () => ({
+  useLoadBalancedEnvironment: () => ({
+    pending: false,
+    failed: false,
+    environmentId: null,
+    refresh: () => {},
+    refreshSelection: resources.refreshSelection,
+  }),
+}));
+vi.mock("../../state/entities", () => ({ useProjects: () => [] }));
+vi.mock("../../hooks/useScratchProject", () => ({ useScratchProject: () => ({}) }));
+
+function setup() {
+  resources.refreshSelection.mockReset();
+  const a = EnvironmentId.make("machine-a");
+  const b = EnvironmentId.make("machine-b");
+  const draftId = DraftId.make("fresh-routing-draft");
+  const projectId = ProjectId.make("project");
+  const store = useComposerDraftStore.getState();
+  store.clearDraftThread(draftId);
+  store.setProjectDraftThreadId(scopeProjectRef(a, projectId), draftId, {
+    environmentSelection: "auto",
+    loadBalancedEnvironmentId: a,
+  });
+  store.setPrompt(draftId, "Keep this prompt");
+  const provider: ServerProvider = {
+    instanceId: ProviderInstanceId.make("codex"),
+    driver: ProviderDriverKind.make("codex"),
+    enabled: true,
+    installed: true,
+    status: "ready",
+    auth: { status: "authenticated" },
+    availability: "available",
+    version: null,
+    checkedAt: "2026-11-10T01:00:00.000Z",
+    models: [],
+    slashCommands: [],
+    skills: [],
+  };
+  const input: Parameters<typeof useDraftEnvironmentRouting>[0] = {
+    draftId,
+    draftThread: store.getDraftSession(draftId),
+    envLocked: false,
+    automaticEnvironment: true,
+    logicalProjectEnvironments: [a, b].map((environmentId) => ({
+      environmentId,
+      projectId,
+      label: environmentId,
+      isPrimary: environmentId === a,
+      machine: "desktop",
+    })),
+    environmentById: new Map(
+      [a, b].map((id): [EnvironmentId, EnvironmentPresentation] => [
+        id,
+        {
+          environmentId: id,
+          label: id,
+          displayUrl: null,
+          relayManaged: false,
+          entry: {
+            target: new PrimaryConnectionTarget({
+              environmentId: id,
+              label: id,
+              httpBaseUrl: "http://localhost:3000",
+              wsBaseUrl: "ws://localhost:3000",
+            }),
+            profile: Option.none(),
+            enabled: true,
+          },
+          connection: { phase: "connected", error: null, traceId: null },
+          serverConfig: null,
+        },
+      ]),
+    ),
+    providersByEnvironment: new Map([a, b].map((id) => [id, [provider]])),
+    loadBalancingSettings: { loadBalancingWeights: {} },
+    referencedThreadContexts: [],
+    activeProviderInstanceId: provider.instanceId,
+    selectedProvider: provider.driver,
+    composerHasAttachments: false,
+    activeProjectKey: "project",
+    activeProjectIsScratch: false,
+    projectGroupingSettings: DEFAULT_UNIFIED_SETTINGS,
+    sendInFlightRef: { current: false },
+  };
+  resources.refreshSelection.mockImplementation(async (ids) => (ids.includes(b) ? b : null));
+  return { a, b, draftId, input, store };
+}
+
+function Harness({
+  input,
+  onPrepared,
+}: {
+  input: Parameters<typeof useDraftEnvironmentRouting>[0];
+  onPrepared?: (ready: boolean | null) => void;
+}) {
+  const routing = useDraftEnvironmentRouting(input);
+  const draft = useComposerDraftStore((state) => state.getDraftSession(input.draftId!));
+  return (
+    <>
+      <span>{draft?.environmentId}</span>
+      <button
+        onClick={async () => {
+          const ready = await routing.refreshBeforeSend();
+          onPrepared?.(ready);
+        }}
+      >
+        Send
+      </button>
+    </>
+  );
+}
+
+it("refreshes a stored automatic choice before sending an unsent draft", async () => {
+  const { a, b, draftId, input, store } = setup();
+  const selections = ["model-a", "model-b"].map((model) =>
+    createModelSelection(ProviderInstanceId.make("codex"), model),
+  );
+  const threadId = input.draftThread!.threadId;
+  const destinationSelections = draftFanoutSelectionsAtom(
+    scopedThreadKey(scopeThreadRef(b, threadId)),
+  );
+  appAtomRegistry.set(destinationSelections, null);
+  appAtomRegistry.set(
+    draftFanoutSelectionsAtom(scopedThreadKey(scopeThreadRef(a, threadId))),
+    selections,
+  );
+  store.setDraftThreadContext(draftId, {
+    branch: "main",
+    envMode: "worktree",
+    environmentSelection: "auto",
+  });
+  await using screen = await render(
+    <Harness input={{ ...input, draftThread: store.getDraftSession(draftId) }} />,
+  );
+  await expect.element(screen.getByText(a)).toBeInTheDocument();
+  await screen.getByRole("button", { name: "Send" }).click();
+  await expect.element(screen.getByText(b)).toBeInTheDocument();
+  expect(resources.refreshSelection).toHaveBeenCalledWith([a, b]);
+  expect(store.getDraftSession(draftId)).toMatchObject({
+    loadBalancedEnvironmentId: b,
+    branch: "main",
+    envMode: "worktree",
+  });
+  expect(store.getComposerDraft(draftId)?.prompt).toBe("Keep this prompt");
+  expect(appAtomRegistry.get(destinationSelections)).toEqual(selections);
+});
+
+it.each(["manual", "attachments", "promoted"] as const)(
+  "keeps the machine for %s drafts",
+  async (kind) => {
+    const { a, draftId, input, store } = setup();
+    if (kind === "manual") store.setDraftThreadContext(draftId, { environmentSelection: "manual" });
+    const guardedInput = {
+      ...input,
+      automaticEnvironment: kind !== "manual",
+      composerHasAttachments: kind === "attachments",
+      draftThread:
+        kind === "promoted"
+          ? {
+              ...input.draftThread!,
+              promotedTo: { environmentId: a, threadId: input.draftThread!.threadId },
+            }
+          : input.draftThread,
+    };
+    await using screen = await render(<Harness input={guardedInput} />);
+    await screen.getByRole("button", { name: "Send" }).click();
+    expect(resources.refreshSelection).not.toHaveBeenCalled();
+    expect(store.getDraftSession(draftId)?.environmentId).toBe(a);
+  },
+);
+
+it("does not reuse a stored choice when refreshed resources have no winner", async () => {
+  const { a, input } = setup();
+  resources.refreshSelection.mockResolvedValue(null);
+  const onPrepared = vi.fn();
+  await using screen = await render(<Harness input={input} onPrepared={onPrepared} />);
+  await screen.getByRole("button", { name: "Send" }).click();
+  await expect.element(screen.getByText(a)).toBeInTheDocument();
+  await vi.waitFor(() => expect(onPrepared).toHaveBeenCalledWith(false));
+});
+
+it("abandons a refresh when the draft changes while resources are in flight", async () => {
+  const { a, b, draftId, input, store } = setup();
+  let resolve!: (environmentId: EnvironmentId) => void;
+  resources.refreshSelection.mockReturnValue(
+    new Promise<EnvironmentId>((done) => {
+      resolve = done;
+    }),
+  );
+  await using screen = await render(<Harness input={input} />);
+  await screen.getByRole("button", { name: "Send" }).click();
+  await act(async () => {
+    store.setPrompt(draftId, "Edited while checking");
+    resolve(b);
+  });
+  await expect.element(screen.getByText(a)).toBeInTheDocument();
+  expect(store.getComposerDraft(draftId)?.prompt).toBe("Edited while checking");
+});