fix(web): a branch restore in one thread no longer blocks another thread

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

@@ -0,0 +1,149 @@
+import { scopedThreadKey, scopeThreadRef } from "@t3tools/client-runtime/environment";
+import { EnvironmentId, ThreadId } from "@t3tools/contracts";
+import { expect, it, vi } from "vite-plus/test";
+import { page } from "vite-plus/test/browser";
+
+const state = vi.hoisted(() => ({
+  switchRef: (async () => undefined) as (input: unknown) => Promise<unknown>,
+}));
+
+vi.mock(import("../../state/session"), async (importOriginal) => ({
+  ...(await importOriginal()),
+  useEnvironmentScope: (() => true) as never,
+  readEnvironmentScope: (() => true) as never,
+}));
+
+vi.mock(import("../../state/use-atom-command"), async (importOriginal) => ({
+  ...(await importOriginal()),
+  useAtomCommand: (() => (input: unknown) => state.switchRef(input)) as never,
+}));
+
+vi.mock(import("../../state/use-orchestration-command"), async (importOriginal) => ({
+  ...(await importOriginal()),
+  useOrchestrationCommand: (() => async () => ({ _tag: "Success", value: undefined })) as never,
+}));
+
+vi.mock(import("../../hooks/useThreadActions"), async (importOriginal) => ({
+  ...(await importOriginal()),
+  useThreadActions: (() => ({})) as never,
+  useAcknowledgeThreadWoke: (() => () => undefined) as never,
+}));
+
+vi.mock(import("../../state/entities"), async (importOriginal) => ({
+  ...(await importOriginal()),
+  useChildThreadInputs: (() => []) as never,
+}));
+
+import { render } from "../../../test/render";
+import { useThreadLifecycle } from "./useThreadLifecycle";
+
+const environmentId = EnvironmentId.make("environment:test");
+
+function setup() {
+  state.switchRef = async () => undefined;
+}
+
+function threadInput(threadId: string, branch: string) {
+  const ref = scopeThreadRef(environmentId, ThreadId.make(threadId));
+  const thread = { id: ref.threadId, environmentId, branch, worktreePath: null };
+  return {
+    routeThreadKey: scopedThreadKey(ref),
+    activeThread: thread as never,
+    activeThreadId: ref.threadId,
+    activeThreadRef: ref,
+    activeThreadKey: scopedThreadKey(ref),
+    localCheckoutBranchMismatch: { threadBranch: branch, currentBranch: "main" },
+  };
+}
+
+type ThreadInput = ReturnType<typeof threadInput>;
+
+function RestoreBranchHarness(props: {
+  readonly thread: ThreadInput;
+  readonly currentRouteThreadKeyRef: { readonly current: string | null };
+  readonly refreshGitStatus: () => void;
+  readonly scheduleComposerFocus: () => void;
+}) {
+  const lifecycle = useThreadLifecycle({
+    ...props.thread,
+    environmentId,
+    currentRouteThreadKeyRef: props.currentRouteThreadKeyRef,
+    canOperateThread: true,
+    clientSettingsHydrated: true,
+    activeThreadShell: null,
+    activeLatestRun: null as never,
+    activeProjectCwd: "/repo",
+    serverConfig: null,
+    runtimeMode: "full-access" as never,
+    isWorking: false,
+    pendingBackgroundTasks: [],
+    activePendingApproval: null as never,
+    activePendingUserInput: null as never,
+    gitStatusQuery: { refresh: props.refreshGitStatus },
+    nowMinute: "",
+    composerRef: { current: null } as never,
+    composerDraftTarget: { kind: "thread", ref: props.thread.activeThreadRef } as never,
+    sendInFlightRef: { current: false },
+    setOptimisticUserMessages: () => undefined,
+    setThreadError: () => undefined,
+    clearUsageLimitsFor: () => undefined,
+    beginLocalDispatch: () => undefined,
+    resetLocalDispatch: () => undefined,
+    scheduleComposerFocus: props.scheduleComposerFocus,
+    scrollToEnd: () => undefined,
+    persistThreadSettingsForNextTurn: async () => ({ _tag: "Success", value: undefined }) as never,
+  });
+  return (
+    <button
+      type="button"
+      disabled={lifecycle.isRestoringThreadBranch}
+      onClick={lifecycle.handleSwitchCheckoutToThreadStable}
+    >
+      {lifecycle.isRestoringThreadBranch ? "Restoring..." : "Restore branch"}
+    </button>
+  );
+}
+
+it("keeps a branch restore in one thread from blocking or refocusing the next thread", async () => {
+  setup();
+  const switchRef = vi.fn((_input: unknown) => new Promise<unknown>(() => undefined));
+  let finishThreadA: (result: unknown) => void = () => undefined;
+  switchRef.mockImplementationOnce(
+    () =>
+      new Promise((resolve) => {
+        finishThreadA = resolve;
+      }),
+  );
+  state.switchRef = switchRef;
+  const threadA = threadInput("thread:a", "feature-a");
+  const threadB = threadInput("thread:b", "feature-b");
+  const currentRouteThreadKeyRef = { current: threadA.routeThreadKey };
+  const refreshGitStatus = vi.fn();
+  const scheduleComposerFocus = vi.fn();
+  const harness = (thread: ThreadInput) => (
+    <RestoreBranchHarness
+      thread={thread}
+      currentRouteThreadKeyRef={currentRouteThreadKeyRef}
+      refreshGitStatus={refreshGitStatus}
+      scheduleComposerFocus={scheduleComposerFocus}
+    />
+  );
+  await using screen = await render(harness(threadA));
+
+  await page.getByRole("button", { name: "Restore branch" }).click();
+  await expect.element(page.getByRole("button", { name: "Restoring..." })).toBeDisabled();
+
+  // ChatView stays mounted across server threads, so the hook sees thread B in place.
+  currentRouteThreadKeyRef.current = threadB.routeThreadKey;
+  await screen.rerender(harness(threadB));
+  await expect.element(page.getByRole("button", { name: "Restore branch" })).toBeEnabled();
+  await page.getByRole("button", { name: "Restore branch" }).click();
+  expect(switchRef).toHaveBeenLastCalledWith(
+    expect.objectContaining({ input: { cwd: "/repo", refName: "feature-b" } }),
+  );
+
+  finishThreadA({ _tag: "Success", value: { refName: "feature-a" } });
+  await vi.waitFor(() => expect(refreshGitStatus).toHaveBeenCalledOnce());
+  expect(scheduleComposerFocus).not.toHaveBeenCalled();
+  await expect.element(page.getByRole("button", { name: "Restoring..." })).toBeDisabled();
+});

apps/web/src/components/chat/useThreadLifecycle.tsx

@@ -294,7 +294,11 @@ export function useThreadLifecycle(input: ThreadLifecycleInput) {
     setUnsnoozingThreadKey((current) => (current === threadKey ? null : current));
   };
   const handleUnsnoozeActiveThreadStable = useStableCallback(handleUnsnoozeActiveThread);
-  const [isRestoringThreadBranch, setIsRestoringThreadBranch] = useState(false);
+  // ChatView stays mounted across server threads, so an in-flight restore is keyed to
+  // the thread that started it.
+  const [restoringBranchThreadKey, setRestoringBranchThreadKey] = useState<string | null>(null);
+  const isRestoringThreadBranch =
+    restoringBranchThreadKey !== null && restoringBranchThreadKey === activeThreadKey;
   const [revealedBranchMismatchKey, setRevealedBranchMismatchKey] = useState<string | null>(null);
   const [, setBranchMismatchDismissTick] = useState(0);
   const composerHasDraftContent = useComposerDraftStore((store) =>
@@ -329,7 +333,10 @@ export function useThreadLifecycle(input: ThreadLifecycleInput) {
     ) {
       return;
     }
-    setIsRestoringThreadBranch(true);
+    const threadKey = scopedThreadKey(scopeThreadRef(environmentId, activeThread.id));
+    const finishRestore = () =>
+      setRestoringBranchThreadKey((current) => (current === threadKey ? null : current));
+    setRestoringBranchThreadKey(threadKey);
     const checkoutResult = await switchGitRef({
       environmentId,
       input: {
@@ -338,7 +345,7 @@ export function useThreadLifecycle(input: ThreadLifecycleInput) {
       },
     });
     if (checkoutResult._tag === "Failure") {
-      setIsRestoringThreadBranch(false);
+      finishRestore();
       if (!isAtomCommandInterrupted(checkoutResult)) {
         toastManager.add(
           stackedThreadToast({
@@ -352,7 +359,7 @@ export function useThreadLifecycle(input: ThreadLifecycleInput) {
     }
 
     if (!readEnvironmentScope(environmentId, AuthSourceControlWriteScope)) {
-      setIsRestoringThreadBranch(false);
+      finishRestore();
       return;
     }
     const nextBranch = checkoutResult.value.refName ?? localCheckoutBranchMismatch.threadBranch;
@@ -362,7 +369,7 @@ export function useThreadLifecycle(input: ThreadLifecycleInput) {
         input: { threadId: activeThread.id, branch: nextBranch, worktreePath: null },
       });
       if (updateResult._tag === "Failure") {
-        setIsRestoringThreadBranch(false);
+        finishRestore();
         if (!isAtomCommandInterrupted(updateResult)) {
           toastManager.add(
             stackedThreadToast({
@@ -377,8 +384,8 @@ export function useThreadLifecycle(input: ThreadLifecycleInput) {
       }
     }
     gitStatusQuery.refresh();
-    setIsRestoringThreadBranch(false);
-    scheduleComposerFocus();
+    finishRestore();
+    if (currentRouteThreadKeyRef.current === routeThreadKey) scheduleComposerFocus();
   };
   const handleSwitchCheckoutToThreadStable = useStableCallback(() => {
     void handleSwitchCheckoutToThread();