1import { RegistryContext, useAtomValue } from "@effect/atom-react";
2import { chooseLoadBalancedEnvironment } from "@t3tools/client-runtime/load-balancing";
3import type { EnvironmentId } from "@t3tools/contracts";
4import { Atom } from "effect/reactivity";
5import { useCallback, useContext, useMemo } from "react";
6
7import { serverEnvironment } from "../state/server";
8
9/** Only mounted for unresolved automatic drafts, so idle clients do not poll hosts. */
10export function useLoadBalancedEnvironment(
11 environmentIds: readonly EnvironmentId[],
12 weights: Readonly<Record<string, number>>,
13) {
14 const registry = useContext(RegistryContext);
15 const refresh = useCallback(
16 (ids: readonly EnvironmentId[]) => {
17 for (const environmentId of ids) {
18 registry.refresh(serverEnvironment.hostResources({ environmentId, input: {} }));
19 }
20 },
21 [registry],
22 );
23 const resourcesAtom = useMemo(
24 () =>
25 Atom.make((get) =>
26 environmentIds.map((environmentId) => {
27 const result = get(serverEnvironment.hostResources({ environmentId, input: {} }));
28 return {
29 environmentId,
30 resources: result._tag === "Success" ? result.value : null,
31 receivedAt: result._tag === "Success" ? result.timestamp : 0,
32 pending: result._tag === "Initial" || result.waiting,
33 failed: result._tag === "Failure",
34 };
35 }),
36 ),
37 [environmentIds],
38 );
39 const resources = useAtomValue(resourcesAtom);
40 const pending = resources.some((resource) => resource.pending);
41 const environmentId = chooseLoadBalancedEnvironment(
42 resources.map((resource) => ({
43 ...resource,
44 weight: weights[resource.environmentId] ?? 50,
45 })),
46 Date.now(),
47 ) as EnvironmentId | null;
48 return {
49 refresh,
50 pending,
51 environmentId,
52 failed: !pending && environmentId === null && resources.some((resource) => resource.failed),
53 };
54}