Open History to see the last commit
1import { useLoadBalancedEnvironment } from "../../hooks/useLoadBalancedEnvironment";
2import { useLogicalProjectEnvironments } from "../../hooks/useLogicalProjectEnvironments";
3import { useScratchProject } from "../../hooks/useScratchProject";
4import { useStableCallback } from "../../hooks/useStableCallback";
5import {
6 type EnvironmentId,
7 type ProviderDriverKind,
8 ProviderInstanceId,
9 resolveProjectProviderInstanceEnabled,
10 type ThreadContextRecord,
11} from "@t3tools/contracts";
12import {
13 scopedThreadKey,
14 scopeProjectRef,
15 scopeThreadRef,
16} from "@t3tools/client-runtime/environment";
17import {
18 isAtomCommandInterrupted,
19 squashAtomCommandFailure,
20} from "@t3tools/client-runtime/state/runtime";
21import { isTriagePrompt, replaceTriagePlaybookLink } from "@t3tools/shared/triagePrompt";
22import { useEffect, useLayoutEffect, useRef, useState } from "react";
23import { flushSync } from "react-dom";
24import { stackedThreadToast, toastManager } from "../ui/toast";
25import {
26 deriveLogicalProjectKeyFromSettings,
27 selectProjectGroupingSettings,
28} from "@t3tools/client-runtime/state/project-grouping";
29import { useComposerDraftStore, DraftId } from "../../composerDraftStore";
30import { type EnvironmentPresentation, useEnvironmentProviders } from "../../state/environments";
31import { currentClientMetadata } from "../../connection/platform";
32import { useProjects } from "../../state/entities";
33import { serverEnvironment } from "../../state/server";
34import { useAtomCommand } from "../../state/use-atom-command";
35import type { DraftSessionState } from "../../composerDraftStore";
36import { loadBalancingSendWarning } from "./composerSend.logic";
37import { appAtomRegistry } from "../../rpc/atomRegistry";
38import { draftFanoutSelectionsAtom } from "./chatViewSend";
39
40interface DraftEnvironmentRoutingInput {
41 readonly draftId: DraftId | null;
42 readonly draftThread: DraftSessionState | null;
43 readonly envLocked: boolean;
44 readonly automaticEnvironment: boolean;
45 readonly logicalProjectEnvironments: ReturnType<typeof useLogicalProjectEnvironments>;
46 readonly environmentById: ReadonlyMap<EnvironmentId, EnvironmentPresentation>;
47 readonly providersByEnvironment: ReturnType<typeof useEnvironmentProviders>;
48 readonly loadBalancingSettings: {
49 readonly loadBalancingWeights: Readonly<Record<string, number>>;
50 };
51 readonly referencedThreadContexts: ReadonlyArray<ThreadContextRecord>;
52 readonly activeProviderInstanceId: ProviderInstanceId | null;
53 readonly selectedProvider: ProviderDriverKind;
54 readonly composerHasAttachments: boolean;
55 readonly activeProjectKey: string | null;
56 readonly activeProjectIsScratch: boolean;
57 readonly projectGroupingSettings: ReturnType<typeof selectProjectGroupingSettings>;
58 readonly sendInFlightRef: { readonly current: boolean };
59}
60
61/**
62 * Which machine an unsent draft runs on: automatic load balancing across the
63 * project's machines, or a manual switch, which may first open a scratch
64 * project on the target machine. A triage draft takes its session along: the
65 * target machine prepares one, and the playbook link follows.
66 */
67export function useDraftEnvironmentRouting(input: DraftEnvironmentRoutingInput) {
68 const {
69 draftId,
70 draftThread,
71 envLocked,
72 automaticEnvironment,
73 logicalProjectEnvironments,
74 environmentById,
75 providersByEnvironment,
76 loadBalancingSettings,
77 referencedThreadContexts,
78 activeProviderInstanceId,
79 selectedProvider,
80 composerHasAttachments,
81 activeProjectKey,
82 activeProjectIsScratch,
83 projectGroupingSettings,
84 sendInFlightRef,
85 } = input;
86 const setDraftThreadContext = useComposerDraftStore((store) => store.setDraftThreadContext);
87 const getDraftSession = useComposerDraftStore((store) => store.getDraftSession);
88 const getComposerDraft = useComposerDraftStore((store) => store.getComposerDraft);
89 const setPrompt = useComposerDraftStore((store) => store.setPrompt);
90 const prepareTriage = useAtomCommand(serverEnvironment.prepareTriage, { reportFailure: false });
91 const setLogicalProjectDraftThreadId = useComposerDraftStore(
92 (store) => store.setLogicalProjectDraftThreadId,
93 );
94 const allProjects = useProjects();
95 const { openScratchProject } = useScratchProject();
96 const needsLoadBalancing = automaticEnvironment && !draftThread?.loadBalancedEnvironmentId;
97 const loadBalancingCandidates = automaticEnvironment
98 ? logicalProjectEnvironments
99 .filter((candidate) => {
100 const environment = environmentById.get(candidate.environmentId);
101 return (
102 environment?.connection.phase === "connected" &&
103 (loadBalancingSettings.loadBalancingWeights[candidate.environmentId] ?? 50) > 0 &&
104 referencedThreadContexts.every(
105 (record) => record.environmentId === candidate.environmentId,
106 ) &&
107 providersByEnvironment
108 .get(candidate.environmentId)
109 ?.some(
110 (provider) =>
111 (activeProviderInstanceId === null ||
112 provider.instanceId === activeProviderInstanceId) &&
113 provider.driver === selectedProvider &&
114 (environment.serverConfig?.settings
115 ? resolveProjectProviderInstanceEnabled(
116 environment.serverConfig.settings,
117 candidate.projectId,
118 provider.instanceId,
119 )
120 : provider.enabled) &&
121 provider.installed &&
122 provider.status !== "error" &&
123 provider.auth.status !== "unauthenticated" &&
124 provider.availability !== "unavailable",
125 )
126 );
127 })
128 .map((candidate) => candidate.environmentId)
129 : [];
130 const loadBalancing = useLoadBalancedEnvironment(
131 needsLoadBalancing ? loadBalancingCandidates : [],
132 loadBalancingSettings.loadBalancingWeights,
133 );
134 const readRouting = useStableCallback(() => ({ input, loadBalancingCandidates }));
135 const refreshBeforeSend = () => {
136 if (
137 !automaticEnvironment ||
138 !draftId ||
139 !draftThread?.loadBalancedEnvironmentId ||
140 draftThread.promotedTo ||
141 composerHasAttachments
142 )
143 return null;
144 const originalDraft = getDraftSession(draftId);
145 const composerDraft = useComposerDraftStore.getState().getComposerDraft(draftId);
146 return loadBalancing.refreshSelection(loadBalancingCandidates).then((environmentId) => {
147 const current = readRouting();
148 if (
149 current.input.draftId !== draftId ||
150 !current.input.automaticEnvironment ||
151 current.input.composerHasAttachments ||
152 current.input.activeProjectKey !== activeProjectKey ||
153 current.input.selectedProvider !== selectedProvider ||
154 current.input.activeProviderInstanceId !== activeProviderInstanceId ||
155 current.input.loadBalancingSettings.loadBalancingWeights !==
156 loadBalancingSettings.loadBalancingWeights ||
157 current.loadBalancingCandidates.join("\n") !== loadBalancingCandidates.join("\n") ||
158 getDraftSession(draftId) !== originalDraft ||
159 useComposerDraftStore.getState().getComposerDraft(draftId) !== composerDraft
160 )
161 return false;
162 const target = current.input.logicalProjectEnvironments.find(
163 (candidate) => candidate.environmentId === environmentId,
164 );
165 if (!target?.projectId) {
166 toastManager.add({
167 type: "warning",
168 ...loadBalancingSendWarning({
169 pending: false,
170 threadContexts: referencedThreadContexts,
171 candidateEnvironmentIds: loadBalancingCandidates,
172 }),
173 });
174 return false;
175 }
176 const projectRef = scopeProjectRef(target.environmentId, target.projectId);
177 if (originalDraft && target.environmentId !== originalDraft.environmentId) {
178 appAtomRegistry.set(
179 draftFanoutSelectionsAtom(
180 scopedThreadKey(scopeThreadRef(target.environmentId, originalDraft.threadId)),
181 ),
182 appAtomRegistry.get(
183 draftFanoutSelectionsAtom(
184 scopedThreadKey(scopeThreadRef(originalDraft.environmentId, originalDraft.threadId)),
185 ),
186 ),
187 );
188 }
189 // The send must read the target project's composer and provider context after retargeting.
190 flushSync(() =>
191 setDraftThreadContext(draftId, {
192 projectRef,
193 environmentSelection: "auto",
194 loadBalancedEnvironmentId: target.environmentId,
195 branch: originalDraft?.branch ?? null,
196 }),
197 );
198 return true;
199 });
200 };
201 // oxlint-disable-next-line t3code/no-unreasoned-state-sync -- writes the resolved load-balanced environment into the draft store once the async pick settles
202 useEffect(() => {
203 if (!needsLoadBalancing || loadBalancing.pending || !draftId || sendInFlightRef.current) return;
204 const target = logicalProjectEnvironments.find(
205 (environment) => environment.environmentId === loadBalancing.environmentId,
206 );
207 if (!target?.projectId) return;
208 setDraftThreadContext(draftId, {
209 projectRef: scopeProjectRef(target.environmentId, target.projectId),
210 environmentSelection: "auto",
211 loadBalancedEnvironmentId: target.environmentId,
212 });
213 }, [
214 needsLoadBalancing,
215 loadBalancing.pending,
216 loadBalancing.environmentId,
217 draftId,
218 logicalProjectEnvironments,
219 sendInFlightRef,
220 setDraftThreadContext,
221 ]);
222 const onAutoEnvironment = () => {
223 if (envLocked || !draftId || sendInFlightRef.current) return;
224 if (composerHasAttachments) {
225 toastManager.add({
226 type: "warning",
227 id: "load-balancing-attachments",
228 title: "Keep attachments on this machine",
229 description:
230 "Remove attachments before choosing automatic routing, then attach them on the selected machine.",
231 });
232 return;
233 }
234 loadBalancing.refresh(
235 logicalProjectEnvironments.map((environment) => environment.environmentId),
236 );
237 setDraftThreadContext(draftId, {
238 environmentSelection: "auto",
239 loadBalancedEnvironmentId: null,
240 branch: null,
241 worktreePath: null,
242 });
243 };
244 const onAutoEnvironmentStable = useStableCallback(onAutoEnvironment);
245 const autoEnvironmentLabel = automaticEnvironment
246 ? draftThread?.loadBalancedEnvironmentId
247 ? "Auto balance"
248 : loadBalancing.pending
249 ? "Checking machines…"
250 : loadBalancing.failed
251 ? "Auto balance unavailable"
252 : "Auto balance"
253 : undefined;
254
255 // A change in flight belongs to the draft and project it started on.
256 const environmentChangeScope = `${draftId ?? ""}:${activeProjectKey ?? ""}`;
257 const environmentChangeRef = useRef<{
258 readonly environmentId: EnvironmentId;
259 readonly scope: string;
260 } | null>(null);
261 const environmentChangeScopeRef = useRef("");
262 // oxlint-disable-next-line t3code/no-unreasoned-state-sync -- mirrors the current change scope into a ref that async environment switches compare against
263 useLayoutEffect(() => {
264 environmentChangeScopeRef.current = environmentChangeScope;
265 });
266 // Leaving the scope, or unmounting, abandons a change still in flight.
267 // oxlint-disable-next-line t3code/no-unreasoned-state-sync -- abandons an in-flight environment change when the draft or project scope changes or unmounts
268 useLayoutEffect(
269 () => () => {
270 if (environmentChangeRef.current?.scope === environmentChangeScope) {
271 environmentChangeRef.current = null;
272 }
273 },
274 [environmentChangeScope],
275 );
276 const readPendingEnvironmentChange = () => {
277 const pending = environmentChangeRef.current;
278 return pending !== null && pending.scope === environmentChangeScopeRef.current ? pending : null;
279 };
280 // The busy state belongs to the request that set it, so an older request settling late
281 // (abandoned, or superseded by a newer switch) can never clear a newer one's.
282 const [busyEnvironmentChange, setBusyEnvironmentChange] = useState<{
283 readonly environmentId: EnvironmentId;
284 readonly scope: string;
285 } | null>(null);
286 const isEnvironmentChanging = busyEnvironmentChange?.scope === environmentChangeScope;
287
288 const onEnvironmentChange = (nextEnvironmentId: EnvironmentId) => {
289 if (envLocked || !draftId || sendInFlightRef.current) return;
290 const originalDraft = getDraftSession(draftId);
291 if (!originalDraft || originalDraft.promotedTo) return;
292 const target = logicalProjectEnvironments.find(
293 (env) => env.environmentId === nextEnvironmentId,
294 );
295 if (!target) return;
296 const request = { environmentId: target.environmentId, scope: environmentChangeScope };
297 environmentChangeRef.current = request;
298 setBusyEnvironmentChange(null);
299 // `triagePrompt` is the prompt the target machine prepared for a triage draft.
300 const retarget = (project: (typeof allProjects)[number], triagePrompt?: string) => {
301 const currentDraft = getDraftSession(draftId);
302 if (
303 environmentChangeRef.current !== request ||
304 environmentChangeScopeRef.current !== request.scope ||
305 sendInFlightRef.current ||
306 !currentDraft ||
307 currentDraft.promotedTo ||
308 currentDraft.environmentId !== originalDraft.environmentId ||
309 currentDraft.projectId !== originalDraft.projectId
310 )
311 return;
312 if (triagePrompt !== undefined) {
313 // Read now, so text typed while the machine prepared the session is kept.
314 const prompt = getComposerDraft(draftId)?.prompt ?? "";
315 const nextPrompt = replaceTriagePlaybookLink(prompt, triagePrompt);
316 if (nextPrompt !== prompt) setPrompt(draftId, nextPrompt);
317 }
318 const projectRef = scopeProjectRef(target.environmentId, project.id);
319 if (activeProjectIsScratch) {
320 setLogicalProjectDraftThreadId(
321 deriveLogicalProjectKeyFromSettings(project, projectGroupingSettings),
322 projectRef,
323 draftId,
324 { environmentSelection: "manual", loadBalancedEnvironmentId: null },
325 );
326 } else {
327 setDraftThreadContext(draftId, {
328 projectRef,
329 environmentSelection: "manual",
330 loadBalancedEnvironmentId: null,
331 });
332 }
333 };
334 const finish = () => {
335 if (environmentChangeRef.current === request) environmentChangeRef.current = null;
336 setBusyEnvironmentChange((busy) => (busy === request ? null : busy));
337 };
338 const knownProject = allProjects.find(
339 (project) =>
340 project.environmentId === target.environmentId && project.id === target.projectId,
341 );
342 // The playbook and context files exist only on the machine that prepared them.
343 const followsTriage = isTriagePrompt(getComposerDraft(draftId)?.prompt ?? "");
344 if (target.projectId !== null && !(knownProject && followsTriage)) {
345 if (knownProject) retarget(knownProject);
346 finish();
347 return;
348 }
349 const prepareTriagePrompt = async () => {
350 const prepared = await prepareTriage({
351 environmentId: target.environmentId,
352 input: { client: currentClientMetadata() },
353 });
354 if (prepared._tag === "Success") return prepared.value.prompt;
355 if (!isAtomCommandInterrupted(prepared)) {
356 const error = squashAtomCommandFailure(prepared);
357 toastManager.add(
358 stackedThreadToast({
359 type: "error",
360 title: "Could not switch machine",
361 description: error instanceof Error ? error.message : "An error occurred.",
362 }),
363 );
364 }
365 return null;
366 };
367 const switchMachine = async () => {
368 const triagePrompt = followsTriage ? await prepareTriagePrompt() : undefined;
369 if (triagePrompt === null) return;
370 const project =
371 knownProject ??
372 (await openScratchProject(target.environmentId, "Could not switch machine"));
373 if (project) retarget(project, triagePrompt);
374 };
375 setBusyEnvironmentChange(request);
376 void switchMachine().finally(finish);
377 };
378 const onEnvironmentChangeStable = useStableCallback(onEnvironmentChange);
379
380 return {
381 needsLoadBalancing,
382 loadBalancing,
383 refreshBeforeSend,
384 onAutoEnvironmentStable,
385 autoEnvironmentLabel,
386 isEnvironmentChanging,
387 readPendingEnvironmentChange,
388 onEnvironmentChangeStable,
389 };
390}