Skip to content

Commit 80ed7cd

Browse files
Silthusclaude
Silthus
and
claude
committed
fix(web): no-project drafts follow load balancing
With load balancing on, a draft without a project always stayed on the primary machine. Automatic routing skipped scratch drafts because the picked machine may not have its no-project folder yet. The balanced pick now goes through the same path as a manual switch, which opens that folder first, and settles on the current machine if it cannot. Co-Authored-By: Claude Opus 5.5 <[email protected]>
1 parent 9e85538 commit 80ed7cd

3 files changed

Lines changed: 202 additions & 47 deletions

File tree

‎apps/web/src/components/ChatView.tsx‎

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1190,7 +1190,6 @@ export default function ChatView(props: ChatViewProps) {
11901190
activeProjectIsScratch && draftId !== null,
11911191
);
11921192
const hasMultipleEnvironments = logicalProjectEnvironments.length > 1;
1193-
const canAutoBalanceEnvironments = hasMultipleEnvironments && !activeProjectIsScratch;
11941193
const activeEnvironmentOption =
11951194
logicalProjectEnvironments.find(
11961195
(environment) => environment.environmentId === activeThread?.environmentId,
@@ -1259,7 +1258,7 @@ export default function ChatView(props: ChatViewProps) {
12591258
clientSettingsHydrated &&
12601259
draftId &&
12611260
!envLocked &&
1262-
canAutoBalanceEnvironments &&
1261+
hasMultipleEnvironments &&
12631262
loadBalancingSettings.loadBalancingEnabled &&
12641263
draftThread?.environmentSelection !== "manual" &&
12651264
(!composerHasAttachments || Boolean(draftThread?.loadBalancedEnvironmentId)) &&
@@ -3779,7 +3778,7 @@ export default function ChatView(props: ChatViewProps) {
37793778
onAutoEnvironment={
37803779
draftId &&
37813780
!envLocked &&
3782-
canAutoBalanceEnvironments &&
3781+
hasMultipleEnvironments &&
37833782
loadBalancingSettings.loadBalancingEnabled
37843783
? onAutoEnvironmentStable
37853784
: undefined
@@ -3873,7 +3872,7 @@ export default function ChatView(props: ChatViewProps) {
38733872
onAutoEnvironment={
38743873
draftId &&
38753874
!envLocked &&
3876-
canAutoBalanceEnvironments &&
3875+
hasMultipleEnvironments &&
38773876
loadBalancingSettings.loadBalancingEnabled
38783877
? onAutoEnvironmentStable
38793878
: undefined

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

Lines changed: 155 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,155 @@
1+
import { EnvironmentId, ProjectId, ProviderDriverKind } from "@t3tools/contracts";
2+
import { describe, expect, it, vi } from "vite-plus/test";
3+
4+
import { render } from "../../../test/render";
5+
import { type DraftId, useComposerDraftStore } from "../../composerDraftStore";
6+
import { useDraftEnvironmentRouting } from "./useDraftEnvironmentRouting";
7+
8+
const mac = EnvironmentId.make("mac");
9+
const box = EnvironmentId.make("box");
10+
const draftId = "draft-1" as DraftId;
11+
const boxScratch = {
12+
environmentId: box,
13+
id: ProjectId.make("scratch-box"),
14+
workspaceRoot: "/home/me/.t3/scratch",
15+
};
16+
17+
const mocks = vi.hoisted(() => ({
18+
macScratch: { environmentId: "mac", id: "scratch-mac", workspaceRoot: "/Users/me/.t3/scratch" },
19+
draft: {} as Record<string, unknown>,
20+
openScratchProject: vi.fn<(environmentId: string) => Promise<unknown>>(),
21+
setDraftThreadContext: vi.fn(),
22+
setLogicalProjectDraftThreadId: vi.fn(),
23+
}));
24+
25+
vi.mock("../../composerDraftStore", () => {
26+
const store = {
27+
getDraftSession: () => ({ ...mocks.draft }),
28+
setDraftThreadContext: mocks.setDraftThreadContext,
29+
setLogicalProjectDraftThreadId: mocks.setLogicalProjectDraftThreadId,
30+
};
31+
return {
32+
DraftId: {},
33+
useComposerDraftStore: (selector: (state: typeof store) => unknown) => selector(store),
34+
};
35+
});
36+
vi.mock("../../hooks/useLoadBalancedEnvironment", () => ({
37+
useLoadBalancedEnvironment: () => ({
38+
refresh: () => {},
39+
pending: false,
40+
environmentId: "box",
41+
failed: false,
42+
}),
43+
}));
44+
vi.mock("../../hooks/useScratchProject", () => ({
45+
useScratchProject: () => ({ openScratchProject: mocks.openScratchProject }),
46+
}));
47+
vi.mock("../../state/entities", () => ({ useProjects: () => [mocks.macScratch] }));
48+
vi.mock("@t3tools/client-runtime/state/project-grouping", () => ({
49+
deriveLogicalProjectKeyFromSettings: (project: typeof boxScratch) =>
50+
`${project.environmentId}:${project.workspaceRoot}`,
51+
}));
52+
53+
function setup(openedProject: typeof boxScratch | null) {
54+
mocks.openScratchProject.mockReset().mockResolvedValue(openedProject);
55+
mocks.draft = {
56+
environmentId: mac,
57+
projectId: mocks.macScratch.id,
58+
promotedTo: null,
59+
loadBalancedEnvironmentId: null,
60+
};
61+
const retarget = (context: Record<string, unknown>) => {
62+
const { projectRef, ...rest } = context as { projectRef?: { environmentId: string } };
63+
mocks.draft = { ...mocks.draft, ...projectRef, ...rest };
64+
};
65+
mocks.setDraftThreadContext.mockReset().mockImplementation((_draftId, context) => {
66+
retarget(context);
67+
});
68+
mocks.setLogicalProjectDraftThreadId
69+
.mockReset()
70+
.mockImplementation((_key, projectRef, _draftId, options) => {
71+
retarget({ projectRef, ...options });
72+
});
73+
}
74+
75+
const codex = ProviderDriverKind.make("codex");
76+
const provider = {
77+
instanceId: "codex",
78+
driver: codex,
79+
enabled: true,
80+
installed: true,
81+
status: "ready",
82+
auth: { status: "authenticated" },
83+
availability: "available",
84+
};
85+
86+
function ScratchDraft() {
87+
const draftThread = useComposerDraftStore((store) => store.getDraftSession(draftId));
88+
useDraftEnvironmentRouting({
89+
draftId,
90+
draftThread,
91+
envLocked: false,
92+
automaticEnvironment: true,
93+
logicalProjectEnvironments: [mac, box].map((environmentId) => ({
94+
environmentId,
95+
projectId: environmentId === mac ? ProjectId.make(mocks.macScratch.id) : null,
96+
label: environmentId,
97+
isPrimary: environmentId === mac,
98+
machine: "desktop" as const,
99+
})),
100+
environmentById: new Map(
101+
[mac, box].map((environmentId) => [
102+
environmentId,
103+
{ connection: { phase: "connected" }, serverConfig: null } as never,
104+
]),
105+
),
106+
providersByEnvironment: new Map([
107+
[mac, [provider]],
108+
[box, [provider]],
109+
]) as never,
110+
loadBalancingSettings: { loadBalancingWeights: {} },
111+
referencedThreadContexts: [],
112+
activeProviderInstanceId: null,
113+
selectedProvider: codex,
114+
composerHasAttachments: false,
115+
activeProjectKey: `${mac}:${mocks.macScratch.workspaceRoot}`,
116+
activeProjectIsScratch: true,
117+
projectGroupingSettings: {} as never,
118+
sendInFlightRef: { current: false },
119+
});
120+
return null;
121+
}
122+
123+
describe("useDraftEnvironmentRouting", () => {
124+
it("moves a no-project draft to the balanced machine, opening its no-project folder first", async () => {
125+
setup(boxScratch);
126+
127+
// The harness only runs the hook and renders nothing routed.
128+
await using _screen = await render(<ScratchDraft />, { router: false });
129+
130+
await vi.waitFor(() =>
131+
expect(mocks.setLogicalProjectDraftThreadId).toHaveBeenCalledWith(
132+
`${box}:${boxScratch.workspaceRoot}`,
133+
{ environmentId: box, projectId: boxScratch.id },
134+
draftId,
135+
{ environmentSelection: "auto", loadBalancedEnvironmentId: box },
136+
),
137+
);
138+
expect(mocks.openScratchProject).toHaveBeenCalledTimes(1);
139+
});
140+
141+
it("keeps a no-project draft on its machine when the balanced machine cannot open one", async () => {
142+
setup(null);
143+
144+
// The harness only runs the hook and renders nothing routed.
145+
await using _screen = await render(<ScratchDraft />, { router: false });
146+
147+
await vi.waitFor(() =>
148+
expect(mocks.setDraftThreadContext).toHaveBeenCalledWith(draftId, {
149+
environmentSelection: "auto",
150+
loadBalancedEnvironmentId: mac,
151+
}),
152+
);
153+
expect(mocks.setLogicalProjectDraftThreadId).not.toHaveBeenCalled();
154+
});
155+
});

‎apps/web/src/components/chat/useDraftEnvironmentRouting.ts‎

Lines changed: 44 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -111,27 +111,6 @@ export function useDraftEnvironmentRouting(input: DraftEnvironmentRoutingInput)
111111
loadBalancingCandidates,
112112
loadBalancingSettings.loadBalancingWeights,
113113
);
114-
// oxlint-disable-next-line t3code/no-unreasoned-state-sync -- writes the resolved load-balanced environment into the draft store once the async pick settles
115-
useEffect(() => {
116-
if (!needsLoadBalancing || loadBalancing.pending || !draftId || sendInFlightRef.current) return;
117-
const target = logicalProjectEnvironments.find(
118-
(environment) => environment.environmentId === loadBalancing.environmentId,
119-
);
120-
if (!target?.projectId) return;
121-
setDraftThreadContext(draftId, {
122-
projectRef: scopeProjectRef(target.environmentId, target.projectId),
123-
environmentSelection: "auto",
124-
loadBalancedEnvironmentId: target.environmentId,
125-
});
126-
}, [
127-
needsLoadBalancing,
128-
loadBalancing.pending,
129-
loadBalancing.environmentId,
130-
draftId,
131-
logicalProjectEnvironments,
132-
sendInFlightRef,
133-
setDraftThreadContext,
134-
]);
135114
const onAutoEnvironment = () => {
136115
if (envLocked || !draftId) return;
137116
if (composerHasAttachments) {
@@ -198,43 +177,40 @@ export function useDraftEnvironmentRouting(input: DraftEnvironmentRoutingInput)
198177
} | null>(null);
199178
const isEnvironmentChanging = busyEnvironmentChange?.scope === environmentChangeScope;
200179

201-
const onEnvironmentChange = (nextEnvironmentId: EnvironmentId) => {
180+
const moveDraft = (target: (typeof logicalProjectEnvironments)[number], automatic: boolean) => {
202181
if (envLocked || !draftId || sendInFlightRef.current) return;
203182
const originalDraft = getDraftSession(draftId);
204183
if (!originalDraft || originalDraft.promotedTo) return;
205-
const target = logicalProjectEnvironments.find(
206-
(env) => env.environmentId === nextEnvironmentId,
207-
);
208-
if (!target) return;
184+
const selection = automatic
185+
? { environmentSelection: "auto" as const, loadBalancedEnvironmentId: target.environmentId }
186+
: { environmentSelection: "manual" as const, loadBalancedEnvironmentId: null };
209187
const request = { environmentId: target.environmentId, scope: environmentChangeScope };
210188
environmentChangeRef.current = request;
211189
setBusyEnvironmentChange(null);
212-
const retarget = (project: (typeof allProjects)[number]) => {
190+
const isCurrent = () => {
213191
const currentDraft = getDraftSession(draftId);
214-
if (
215-
environmentChangeRef.current !== request ||
216-
environmentChangeScopeRef.current !== request.scope ||
217-
sendInFlightRef.current ||
218-
!currentDraft ||
219-
currentDraft.promotedTo ||
220-
currentDraft.environmentId !== originalDraft.environmentId ||
221-
currentDraft.projectId !== originalDraft.projectId
222-
)
223-
return;
192+
return (
193+
environmentChangeRef.current === request &&
194+
environmentChangeScopeRef.current === request.scope &&
195+
!sendInFlightRef.current &&
196+
currentDraft !== null &&
197+
!currentDraft.promotedTo &&
198+
currentDraft.environmentId === originalDraft.environmentId &&
199+
currentDraft.projectId === originalDraft.projectId
200+
);
201+
};
202+
const retarget = (project: (typeof allProjects)[number]) => {
203+
if (!isCurrent()) return;
224204
const projectRef = scopeProjectRef(target.environmentId, project.id);
225205
if (activeProjectIsScratch) {
226206
setLogicalProjectDraftThreadId(
227207
deriveLogicalProjectKeyFromSettings(project, projectGroupingSettings),
228208
projectRef,
229209
draftId,
230-
{ environmentSelection: "manual", loadBalancedEnvironmentId: null },
210+
selection,
231211
);
232212
} else {
233-
setDraftThreadContext(draftId, {
234-
projectRef,
235-
environmentSelection: "manual",
236-
loadBalancedEnvironmentId: null,
237-
});
213+
setDraftThreadContext(draftId, { projectRef, ...selection });
238214
}
239215
};
240216
const finish = () => {
@@ -254,11 +230,36 @@ export function useDraftEnvironmentRouting(input: DraftEnvironmentRoutingInput)
254230
void openScratchProject(target.environmentId, "Could not switch machine")
255231
.then((project) => {
256232
if (project) retarget(project);
233+
// Settle on the current machine so balancing does not retry every render.
234+
else if (automatic && isCurrent()) {
235+
setDraftThreadContext(draftId, {
236+
environmentSelection: "auto",
237+
loadBalancedEnvironmentId: originalDraft.environmentId,
238+
});
239+
}
257240
})
258241
.finally(finish);
259242
};
243+
244+
const onEnvironmentChange = (nextEnvironmentId: EnvironmentId) => {
245+
const target = logicalProjectEnvironments.find(
246+
(env) => env.environmentId === nextEnvironmentId,
247+
);
248+
if (target) moveDraft(target, false);
249+
};
260250
const onEnvironmentChangeStable = useStableCallback(onEnvironmentChange);
261251

252+
// oxlint-disable-next-line t3code/no-unreasoned-state-sync -- moves the draft to the load-balanced environment once the async pick settles
253+
useEffect(() => {
254+
if (!needsLoadBalancing || loadBalancing.pending || !draftId || sendInFlightRef.current) return;
255+
if (readPendingEnvironmentChange() !== null) return;
256+
const target = logicalProjectEnvironments.find(
257+
(environment) => environment.environmentId === loadBalancing.environmentId,
258+
);
259+
// oxlint-disable-next-line react/set-state-in-effect -- marks the picker busy while a missing scratch project opens on the picked machine
260+
if (target) moveDraft(target, true);
261+
});
262+
262263
return {
263264
needsLoadBalancing,
264265
loadBalancing,

0 commit comments

Comments
 (0)