🤖 fix: restore thinking level when model changes (#1136)

## Summary

Fix model-switch thinking level restore and stop Review/diff “refresh”
behavior on Cycle Model.

## Changes

- Refactor `usePersistedState` to use React 18 `useSyncExternalStore`
(stable localStorage-backed store)
- Remove the keyed ThinkingProvider remount workaround (prevents subtree
unmount/remount)
- Add regression tests: model switch restores per-model thinking + does
not remount children

## Risk

Medium: `usePersistedState` is a shared primitive. Mitigated by stable
snapshot caching and added tests.

## Testing

- `make static-check`
- `make typecheck`
- `bun test src/browser`

---
_Generated with `mux` • Model: `anthropic:claude-sonnet-4-20250514` •
Thinking: `low`_

Author: ammar-agent

src/browser/contexts/ThinkingContext.test.tsx

@@ -0,0 +1,117 @@
+import { GlobalWindow } from "happy-dom";
+import { afterEach, beforeEach, describe, expect, test } from "bun:test";
+import { act, cleanup, render, waitFor } from "@testing-library/react";
+import React from "react";
+import { ThinkingProvider } from "./ThinkingContext";
+import { useThinkingLevel } from "@/browser/hooks/useThinkingLevel";
+import { getModelKey, getThinkingLevelByModelKey } from "@/common/constants/storage";
+import { updatePersistedState } from "@/browser/hooks/usePersistedState";
+
+// Setup basic DOM environment for testing-library
+const dom = new GlobalWindow();
+/* eslint-disable @typescript-eslint/no-explicit-any, @typescript-eslint/no-unsafe-member-access */
+(global as any).window = dom.window;
+(global as any).document = dom.window.document;
+
+// Ensure globals exist for instanceof checks inside usePersistedState
+(globalThis as any).StorageEvent = dom.window.StorageEvent;
+(globalThis as any).CustomEvent = dom.window.CustomEvent;
+
+(global as any).console = console;
+/* eslint-enable @typescript-eslint/no-explicit-any, @typescript-eslint/no-unsafe-member-access */
+
+interface TestProps {
+  workspaceId: string;
+}
+
+const TestComponent: React.FC<TestProps> = (props) => {
+  const [thinkingLevel] = useThinkingLevel();
+  return (
+    <div data-testid="thinking">
+      {thinkingLevel}:{props.workspaceId}
+    </div>
+  );
+};
+
+describe("ThinkingContext", () => {
+  // Make getDefaultModel deterministic.
+  // (getDefaultModel reads from the global "model-default" localStorage key.)
+  beforeEach(() => {
+    window.localStorage.clear();
+    window.localStorage.setItem("model-default", JSON.stringify("openai:default"));
+  });
+
+  afterEach(() => {
+    cleanup();
+  });
+
+  test("switching models does not remount children", async () => {
+    const workspaceId = "ws-1";
+
+    updatePersistedState(getModelKey(workspaceId), "openai:gpt-5.2");
+    updatePersistedState(getThinkingLevelByModelKey("openai:gpt-5.2"), "high");
+    updatePersistedState(getThinkingLevelByModelKey("anthropic:claude-3.5"), "low");
+
+    let unmounts = 0;
+
+    const Child: React.FC = () => {
+      React.useEffect(() => {
+        return () => {
+          unmounts += 1;
+        };
+      }, []);
+
+      const [thinkingLevel] = useThinkingLevel();
+      return <div data-testid="child">{thinkingLevel}</div>;
+    };
+
+    const view = render(
+      <ThinkingProvider workspaceId={workspaceId}>
+        <Child />
+      </ThinkingProvider>
+    );
+
+    await waitFor(() => {
+      expect(view.getByTestId("child").textContent).toBe("high");
+    });
+
+    act(() => {
+      updatePersistedState(getModelKey(workspaceId), "anthropic:claude-3.5");
+    });
+
+    await waitFor(() => {
+      expect(view.getByTestId("child").textContent).toBe("low");
+    });
+
+    expect(unmounts).toBe(0);
+  });
+  test("switching models restores the per-model thinking level", async () => {
+    const workspaceId = "ws-1";
+
+    // Model A
+    updatePersistedState(getModelKey(workspaceId), "openai:gpt-5.2");
+    updatePersistedState(getThinkingLevelByModelKey("openai:gpt-5.2"), "high");
+
+    // Model B
+    updatePersistedState(getThinkingLevelByModelKey("anthropic:claude-3.5"), "low");
+
+    const view = render(
+      <ThinkingProvider workspaceId={workspaceId}>
+        <TestComponent workspaceId={workspaceId} />
+      </ThinkingProvider>
+    );
+
+    await waitFor(() => {
+      expect(view.getByTestId("thinking").textContent).toBe("high:ws-1");
+    });
+
+    // Change model -> should restore that model's stored thinking level
+    act(() => {
+      updatePersistedState(getModelKey(workspaceId), "anthropic:claude-3.5");
+    });
+
+    await waitFor(() => {
+      expect(view.getByTestId("thinking").textContent).toBe("low:ws-1");
+    });
+  });
+});

src/browser/contexts/ThinkingContext.tsx

@@ -1,7 +1,7 @@
 import type { ReactNode } from "react";
-import React, { createContext, useContext } from "react";
+import React, { createContext, useContext, useMemo } from "react";
 import type { ThinkingLevel } from "@/common/types/thinking";
-import { usePersistedState, readPersistedState } from "@/browser/hooks/usePersistedState";
+import { usePersistedState } from "@/browser/hooks/usePersistedState";
 import { getThinkingLevelByModelKey, getModelKey } from "@/common/constants/storage";
 import { getDefaultModel } from "@/browser/hooks/useModelsFromSettings";
 import { migrateGatewayModel } from "@/browser/hooks/useGatewayModels";
@@ -20,46 +20,47 @@ interface ThinkingProviderProps {
 }
 
 /**
- * Reads the current model from localStorage for the given scope.
- * Returns canonical model format (after gateway migration).
+ * Hook to get the model key for the current scope.
  */
-function getScopedModel(workspaceId?: string, projectPath?: string): string {
-  const defaultModel = getDefaultModel();
-  // Use workspace-scoped model key if available, otherwise project-scoped
-  const modelKey = workspaceId
+function useModelKey(workspaceId?: string, projectPath?: string): string | null {
+  return workspaceId
     ? getModelKey(workspaceId)
     : projectPath
       ? getModelKey(`__project__/${projectPath}`)
       : null;
-
-  if (!modelKey) {
-    return defaultModel;
-  }
-
-  const rawModel = readPersistedState<string>(modelKey, defaultModel);
-  // Normalize to canonical format (e.g., strip legacy gateway prefix)
-  return migrateGatewayModel(rawModel || defaultModel);
 }
 
 export const ThinkingProvider: React.FC<ThinkingProviderProps> = ({
   workspaceId,
   projectPath,
   children,
 }) => {
-  // Read current model from localStorage (non-reactive, re-reads on each render)
-  const modelString = getScopedModel(workspaceId, projectPath);
-  const key = getThinkingLevelByModelKey(modelString);
-  const [thinkingLevel, setThinkingLevel] = usePersistedState<ThinkingLevel>(
-    key,
-    "off",
-    { listener: true } // Listen for changes from command palette and other sources
-  );
+  const defaultModel = getDefaultModel();
+  const modelKey = useModelKey(workspaceId, projectPath);
+
+  // Subscribe to model changes so we update thinking level when model changes.
+  // This uses a fallback key to satisfy hooks rules; it should be unused in practice
+  // because ThinkingProvider is expected to have either workspaceId or projectPath.
+  const [rawModel] = usePersistedState<string>(modelKey ?? "model:__unused__", defaultModel, {
+    listener: true,
+  });
+
+  const thinkingKey = useMemo(() => {
+    const model = migrateGatewayModel(rawModel || defaultModel);
+    return getThinkingLevelByModelKey(model);
+  }, [rawModel, defaultModel]);
 
-  return (
-    <ThinkingContext.Provider value={{ thinkingLevel, setThinkingLevel }}>
-      {children}
-    </ThinkingContext.Provider>
+  const [thinkingLevel, setThinkingLevel] = usePersistedState<ThinkingLevel>(thinkingKey, "off", {
+    listener: true,
+  });
+
+  // Memoize context value to prevent unnecessary re-renders of consumers.
+  const contextValue = useMemo(
+    () => ({ thinkingLevel, setThinkingLevel }),
+    [thinkingLevel, setThinkingLevel]
   );
+
+  return <ThinkingContext.Provider value={contextValue}>{children}</ThinkingContext.Provider>;
 };
 
 export const useThinking = () => {