Open History to see the last commit
1import type { EnvironmentId, UsageModelFilter, UsageSummaryInput } from "@t3tools/contracts";
2import * as Schema from "effect/Schema";
3import type { AtomRegistry } from "effect/reactivity";
4
5import { EnvironmentRpcUnavailableError } from "../rpc/client.ts";
6import type { createEnvironmentPresentationAtoms } from "./presentation.ts";
7import { executeAtomQuery, runAtomCommand, squashAtomCommandFailure } from "./runtime.ts";
8import type { createServerEnvironmentAtoms } from "./server.ts";
9
10const isEnvironmentRpcUnavailable = Schema.is(EnvironmentRpcUnavailableError);
11
12/** The Usage tab asks each environment with that environment's configured model filter. */
13export function usageSummaryInput(
14 input: UsageSummaryInput,
15 modelFilter: UsageModelFilter | undefined,
16): UsageSummaryInput {
17 return modelFilter === undefined ? input : { ...input, modelFilter };
18}
19
20const limitsRefreshAfter = new Map<EnvironmentId, number>();
21const limitsRefreshes = new Map<EnvironmentId, Promise<unknown>>();
22
23export async function refreshUsageLimits<A>(
24 environmentId: EnvironmentId,
25 refresh: () => Promise<A>,
26 automatic = false,
27): Promise<A | undefined> {
28 const pending = limitsRefreshes.get(environmentId);
29 if (pending !== undefined) {
30 return automatic ? undefined : ((await pending) as A);
31 }
32 const refreshAfter = limitsRefreshAfter.get(environmentId) ?? 0;
33 // @effect-diagnostics-next-line globalDate:off
34 if (automatic && Date.now() < refreshAfter) return;
35 const current = Promise.resolve()
36 .then(refresh)
37 .finally(() => {
38 limitsRefreshes.delete(environmentId);
39 // @effect-diagnostics-next-line globalDate:off
40 limitsRefreshAfter.set(environmentId, Date.now() + 5 * 60_000);
41 });
42 limitsRefreshes.set(environmentId, current);
43 return await current;
44}
45
46export async function refreshUsage({
47 registry,
48 server,
49 presentations,
50 environmentIds,
51 input,
52}: {
53 registry: AtomRegistry.AtomRegistry;
54 server: Pick<
55 ReturnType<typeof createServerEnvironmentAtoms>,
56 "usageSummary" | "refreshUsageRates" | "usageModelFilterAtom"
57 >;
58 presentations: Pick<ReturnType<typeof createEnvironmentPresentationAtoms>, "presentationAtom">;
59 environmentIds: readonly EnvironmentId[];
60 input: UsageSummaryInput;
61}): Promise<void> {
62 await Promise.all(
63 environmentIds.map(async (environmentId) => {
64 const query = server.usageSummary({
65 environmentId,
66 input: usageSummaryInput(input, registry.get(server.usageModelFilterAtom(environmentId))),
67 });
68 const presentation = presentations.presentationAtom(environmentId);
69 const controller = new AbortController();
70 const abortWhenDisconnected = () => {
71 if (registry.get(presentation)?.connection.phase !== "connected") controller.abort();
72 };
73 const unsubscribe = registry.subscribe(presentation, abortWhenDisconnected);
74 abortWhenDisconnected();
75 try {
76 const ratesResult = await runAtomCommand(
77 registry,
78 server.refreshUsageRates,
79 { environmentId, input: {} },
80 { reportFailure: false },
81 );
82 const sessionUnavailable =
83 ratesResult._tag === "Failure" &&
84 isEnvironmentRpcUnavailable(squashAtomCommandFailure(ratesResult));
85 registry.refresh(query);
86 if (sessionUnavailable || controller.signal.aborted) return;
87 await executeAtomQuery(registry, query, {
88 reportFailure: false,
89 signal: controller.signal,
90 });
91 } finally {
92 unsubscribe();
93 }
94 }),
95 );
96}