Skip to content

Commit 61e2a2d

Browse files
committed
🤖 fix: rebind thinking state on model changes
- Key ThinkingProviderInner by per-model thinking key to avoid one-render stale thinking state during model switches (e.g. Cycle Model) - Add regression test: switching models restores per-model thinking level
1 parent af267a6 commit 61e2a2d

File tree

2 files changed

+122
-10
lines changed

2 files changed

+122
-10
lines changed
Lines changed: 91 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,91 @@
1+
import { GlobalWindow } from "happy-dom";
2+
3+
// Setup basic DOM environment for testing-library
4+
const dom = new GlobalWindow();
5+
/* eslint-disable @typescript-eslint/no-explicit-any, @typescript-eslint/no-unsafe-member-access */
6+
(global as any).window = dom.window;
7+
(global as any).document = dom.window.document;
8+
9+
// Ensure globals exist for instanceof checks inside usePersistedState
10+
(globalThis as any).StorageEvent = dom.window.StorageEvent;
11+
(globalThis as any).CustomEvent = dom.window.CustomEvent;
12+
13+
// happy-dom's requestAnimationFrame behavior can vary; ensure it's present so
14+
// usePersistedState listener updates (which batch via RAF) are flushed.
15+
if (!globalThis.requestAnimationFrame) {
16+
globalThis.requestAnimationFrame = (cb: FrameRequestCallback) =>
17+
setTimeout(() => cb(Date.now()), 0) as unknown as number;
18+
}
19+
if (!globalThis.cancelAnimationFrame) {
20+
globalThis.cancelAnimationFrame = (id: number) => {
21+
clearTimeout(id as unknown as NodeJS.Timeout);
22+
};
23+
}
24+
(global as any).console = console;
25+
/* eslint-enable @typescript-eslint/no-explicit-any, @typescript-eslint/no-unsafe-member-access */
26+
27+
import { afterEach, beforeEach, describe, expect, test } from "bun:test";
28+
import { act, cleanup, render, waitFor } from "@testing-library/react";
29+
import React from "react";
30+
import { ThinkingProvider } from "./ThinkingContext";
31+
import { useThinkingLevel } from "@/browser/hooks/useThinkingLevel";
32+
import { getModelKey, getThinkingLevelByModelKey } from "@/common/constants/storage";
33+
import { updatePersistedState } from "@/browser/hooks/usePersistedState";
34+
35+
interface TestProps {
36+
workspaceId: string;
37+
}
38+
39+
const TestComponent: React.FC<TestProps> = (props) => {
40+
const [thinkingLevel] = useThinkingLevel();
41+
return (
42+
<div data-testid="thinking">
43+
{thinkingLevel}:{props.workspaceId}
44+
</div>
45+
);
46+
};
47+
48+
describe("ThinkingContext", () => {
49+
// Make getDefaultModel deterministic.
50+
// (getDefaultModel reads from the global "model-default" localStorage key.)
51+
beforeEach(() => {
52+
window.localStorage.setItem("model-default", JSON.stringify("openai:default"));
53+
});
54+
beforeEach(() => {
55+
window.localStorage.clear();
56+
});
57+
58+
afterEach(() => {
59+
cleanup();
60+
});
61+
62+
test("switching models restores the per-model thinking level", async () => {
63+
const workspaceId = "ws-1";
64+
65+
// Model A
66+
updatePersistedState(getModelKey(workspaceId), "openai:gpt-5.2");
67+
updatePersistedState(getThinkingLevelByModelKey("openai:gpt-5.2"), "high");
68+
69+
// Model B
70+
updatePersistedState(getThinkingLevelByModelKey("anthropic:claude-3.5"), "low");
71+
72+
const view = render(
73+
<ThinkingProvider workspaceId={workspaceId}>
74+
<TestComponent workspaceId={workspaceId} />
75+
</ThinkingProvider>
76+
);
77+
78+
await waitFor(() => {
79+
expect(view.getByTestId("thinking").textContent).toBe("high:ws-1");
80+
});
81+
82+
// Change model -> should restore that model's stored thinking level
83+
act(() => {
84+
updatePersistedState(getModelKey(workspaceId), "anthropic:claude-3.5");
85+
});
86+
87+
await waitFor(() => {
88+
expect(view.getByTestId("thinking").textContent).toBe("low:ws-1");
89+
});
90+
});
91+
});

src/browser/contexts/ThinkingContext.tsx

Lines changed: 31 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,30 @@ interface ThinkingContextType {
1313

1414
const ThinkingContext = createContext<ThinkingContextType | undefined>(undefined);
1515

16+
interface ThinkingProviderInnerProps {
17+
thinkingKey: string;
18+
children: ReactNode;
19+
}
20+
21+
const ThinkingProviderInner: React.FC<ThinkingProviderInnerProps> = (props) => {
22+
const [thinkingLevel, setThinkingLevel] = usePersistedState<ThinkingLevel>(
23+
props.thinkingKey,
24+
"off",
25+
{ listener: true }
26+
);
27+
28+
// Memoize context value to prevent unnecessary re-renders of consumers
29+
const contextValue = useMemo(
30+
() => ({ thinkingLevel, setThinkingLevel }),
31+
[thinkingLevel, setThinkingLevel]
32+
);
33+
34+
return (
35+
<ThinkingContext.Provider value={contextValue}>
36+
{props.children}
37+
</ThinkingContext.Provider>
38+
);
39+
};
1640
interface ThinkingProviderProps {
1741
workspaceId?: string; // For existing workspaces
1842
projectPath?: string; // For workspace creation (uses project-scoped model key)
@@ -38,27 +62,24 @@ export const ThinkingProvider: React.FC<ThinkingProviderProps> = ({
3862
const defaultModel = getDefaultModel();
3963
const modelKey = useModelKey(workspaceId, projectPath);
4064

41-
// Subscribe to model changes so we update thinking level when model changes
65+
// Subscribe to model changes so we update thinking level when model changes.
66+
// This uses a fallback key to satisfy hooks rules; it should be unused in practice
67+
// because ThinkingProvider is expected to have either workspaceId or projectPath.
4268
const [rawModel] = usePersistedState<string>(modelKey ?? "model:__unused__", defaultModel, {
4369
listener: true,
4470
});
4571

46-
// Derive the thinking level key from the current model
4772
const thinkingKey = useMemo(() => {
4873
const model = migrateGatewayModel(rawModel || defaultModel);
4974
return getThinkingLevelByModelKey(model);
5075
}, [rawModel, defaultModel]);
5176

52-
const [thinkingLevel, setThinkingLevel] = usePersistedState<ThinkingLevel>(
53-
thinkingKey,
54-
"off",
55-
{ listener: true } // Listen for changes from command palette and other sources
56-
);
57-
77+
// Key the inner provider by thinkingKey so switching models re-mounts and
78+
// synchronously reads the new key's value (avoids one-render stale state).
5879
return (
59-
<ThinkingContext.Provider value={{ thinkingLevel, setThinkingLevel }}>
80+
<ThinkingProviderInner key={thinkingKey} thinkingKey={thinkingKey}>
6081
{children}
61-
</ThinkingContext.Provider>
82+
</ThinkingProviderInner>
6283
);
6384
};
6485

0 commit comments

Comments
 (0)