1/**
2 * Multi-environment usage state.
3 *
4 * Every connected environment answers the same typed query; the client merges
5 * the results. Raw transcripts never leave the machine that produced them.
6 *
7 * @module state/usage
8 */
9import { useAtomValue } from "@effect/atom-react";
10import {
11 AuthDiagnosticsReadScope,
12 USAGE_CONTRACT_VERSION,
13 type EnvironmentId,
14 type UsageBucket,
15 type UsageSummary,
16 type UsageProviderKind,
17 type UsageSummaryInput,
18} from "@t3tools/contracts";
19import { needsCursorKeychainAccess, refreshUsage } from "@t3tools/client-runtime/state/usage";
20import { resolveUsageAccess } from "@t3tools/client-runtime/state/usage-access";
21import * as Option from "effect/Option";
22import { AsyncResult, Atom } from "effect/reactivity";
23import { useCallback, useMemo, useState } from "react";
24
25import { mergeUsage, type EnvironmentUsage, type MergedUsage } from "@t3tools/shared/usageMerge";
26import { appAtomRegistry } from "../rpc/atomRegistry";
27import { environmentPresentations } from "./presentation";
28import { serverEnvironment } from "./server";
29import { environmentSession, readEnvironmentScope } from "./session";
30
31export interface EnvironmentUsageStatus {
32 readonly environmentId: EnvironmentId;
33 readonly label: string;
34 readonly isPending: boolean;
35 readonly canReadDiagnostics: boolean;
36 readonly isConnected: boolean;
37 readonly error: string | null;
38 readonly summary: UsageSummary | null;
39 readonly needsCursorKeychainAccess: boolean;
40}
41
42/**
43 * Reads every environment's summary for one window.
44 *
45 * Keyed by the serialised window so switching ranges does not thrash the atom
46 * cache, and so each environment's query is shared with any other reader of the
47 * same window.
48 */
49const usageByWindowAtom = Atom.family((windowKey: string) =>
50 Atom.make((get): readonly EnvironmentUsageStatus[] => {
51 const input = JSON.parse(windowKey) as UsageSummaryInput;
52 const presentations = get(environmentPresentations.presentationsAtom);
53
54 const statuses: EnvironmentUsageStatus[] = [];
55 for (const [environmentId, presentation] of presentations) {
56 const isConnected = presentation.connection.phase === "connected";
57 const sessionResult = get(environmentSession.sessionStateAtom(environmentId));
58 const access = resolveUsageAccess({
59 connectionPhase: presentation.connection.phase,
60 session: Option.getOrNull(AsyncResult.value(sessionResult)),
61 hasSessionError: sessionResult._tag === "Failure",
62 });
63 if (!access.canReadDiagnostics) {
64 statuses.push({
65 environmentId,
66 label: presentation.entry.target.label,
67 isConnected,
68 ...access,
69 summary: null,
70 needsCursorKeychainAccess: false,
71 });
72 continue;
73 }
74 const result = get(serverEnvironment.usageSummary({ environmentId, input }));
75 const summary = Option.getOrNull(AsyncResult.value(result));
76 statuses.push({
77 environmentId,
78 label: presentation.entry.target.label,
79 isPending: result.waiting,
80 canReadDiagnostics: true,
81 isConnected,
82 error: result._tag === "Failure" ? "This environment could not report usage." : null,
83 summary,
84 needsCursorKeychainAccess: needsCursorKeychainAccess(
85 summary,
86 get(serverEnvironment.providersValueAtom(environmentId)),
87 ),
88 });
89 }
90 return statuses;
91 }).pipe(Atom.withLabel(`web-usage:window:${windowKey}`)),
92);
93
94export interface UsageView {
95 readonly merged: MergedUsage;
96 readonly environments: readonly EnvironmentUsageStatus[];
97 readonly selectedEnvironments: readonly EnvironmentUsageStatus[];
98 /** True until at least one selected environment has answered. */
99 readonly isPending: boolean;
100 /**
101 * The usage to draw: this window's once a selected environment answers,
102 * until then the last window answered for the same selection. Null while
103 * nothing has answered and something still could.
104 */
105 readonly shown: { readonly window: UsageSummaryInput; readonly merged: MergedUsage } | null;
106 /**
107 * True while environments that have not failed are still answering. Failed
108 * environments are reported through their own error rows: totals will not
109 * improve by waiting on them, so they must not read as "still reporting".
110 */
111 readonly isPartial: boolean;
112 readonly refresh: (input?: UsageSummaryInput) => Promise<void>;
113}
114
115/**
116 * Merges every environment that has answered. `keepBucket` narrows the merge,
117 * for example to one model; source ownership still applies, so the result
118 * matches that slice of the full merge. Session counts are per directory and
119 * are not narrowed.
120 */
121export function mergeAnsweredUsage(
122 environments: readonly EnvironmentUsageStatus[],
123 keepBucket?: (bucket: UsageBucket) => boolean,
124): MergedUsage {
125 const answered: EnvironmentUsage[] = environments.flatMap(({ environmentId, label, summary }) =>
126 summary === null
127 ? []
128 : [
129 {
130 environmentId,
131 label,
132 summary:
133 keepBucket === undefined
134 ? summary
135 : { ...summary, buckets: summary.buckets.filter(keepBucket) },
136 },
137 ],
138 );
139 return mergeUsage(answered, USAGE_CONTRACT_VERSION);
140}
141
142const NO_HIDDEN_PROVIDERS: ReadonlySet<UsageProviderKind> = new Set();
143
144/**
145 * Drops hidden providers' buckets and sources before merging, so totals,
146 * shares, and session counts all describe only the visible providers.
147 */
148function withoutProviders(
149 environments: readonly EnvironmentUsageStatus[],
150 hiddenProviders: ReadonlySet<UsageProviderKind>,
151): readonly EnvironmentUsageStatus[] {
152 if (hiddenProviders.size === 0) return environments;
153 return environments.map((environment) =>
154 environment.summary === null
155 ? environment
156 : {
157 ...environment,
158 summary: {
159 ...environment.summary,
160 buckets: environment.summary.buckets.filter(
161 (bucket) => !hiddenProviders.has(bucket.provider),
162 ),
163 sources: environment.summary.sources.filter(
164 (source) => !hiddenProviders.has(source.fingerprint.provider),
165 ),
166 },
167 },
168 );
169}
170
171export function useUsage(
172 input: UsageSummaryInput,
173 selectedEnvironmentIds: ReadonlySet<EnvironmentId> | null = null,
174 hiddenProviders: ReadonlySet<UsageProviderKind> = NO_HIDDEN_PROVIDERS,
175): UsageView {
176 const windowKey = useMemo(
177 () =>
178 JSON.stringify({
179 sinceDay: input.sinceDay,
180 untilDay: input.untilDay,
181 timeZone: input.timeZone,
182 resolution: input.resolution,
183 sinceTime: input.sinceTime,
184 untilTime: input.untilTime,
185 }),
186 [
187 input.sinceDay,
188 input.untilDay,
189 input.timeZone,
190 input.resolution,
191 input.sinceTime,
192 input.untilTime,
193 ],
194 );
195 const atom = usageByWindowAtom(windowKey);
196 const environments = useAtomValue(atom);
197 const selectedEnvironments = useMemo(
198 () =>
199 selectedEnvironmentIds === null
200 ? environments
201 : environments.filter((environment) =>
202 selectedEnvironmentIds.has(environment.environmentId),
203 ),
204 [environments, selectedEnvironmentIds],
205 );
206
207 const refresh = useCallback(
208 (nextInput?: UsageSummaryInput) =>
209 refreshUsage({
210 registry: appAtomRegistry,
211 server: serverEnvironment,
212 presentations: environmentPresentations,
213 // Only environments this connection may read; the others report a
214 // permission error instead of a stale or failed rescan.
215 environmentIds: selectedEnvironments
216 .filter(
217 (environment) =>
218 environment.canReadDiagnostics &&
219 readEnvironmentScope(environment.environmentId, AuthDiagnosticsReadScope),
220 )
221 .map(({ environmentId }) => environmentId),
222 input: nextInput ?? (JSON.parse(windowKey) as UsageSummaryInput),
223 }),
224 [selectedEnvironments, windowKey],
225 );
226
227 const merged = useMemo(
228 () => mergeAnsweredUsage(withoutProviders(selectedEnvironments, hiddenProviders)),
229 [selectedEnvironments, hiddenProviders],
230 );
231
232 const answeredCount = selectedEnvironments.filter(
233 (environment) => environment.summary !== null,
234 ).length;
235 const stillReporting = selectedEnvironments.filter(
236 (environment) => environment.summary === null && environment.error === null,
237 ).length;
238 const isPending = answeredCount === 0 && stillReporting > 0;
239
240 // Stored during render, as React recommends for state that follows props, so
241 // the kept usage is on screen in the same frame the new window starts pending.
242 const [lastAnswered, setLastAnswered] = useState<
243 | (NonNullable<UsageView["shown"]> & {
244 readonly selection: typeof selectedEnvironmentIds;
245 readonly hidden: typeof hiddenProviders;
246 })
247 | null
248 >(null);
249 if (
250 answeredCount > 0 &&
251 (lastAnswered?.merged !== merged ||
252 lastAnswered.window !== input ||
253 lastAnswered.selection !== selectedEnvironmentIds ||
254 lastAnswered.hidden !== hiddenProviders)
255 ) {
256 setLastAnswered({
257 window: input,
258 merged,
259 selection: selectedEnvironmentIds,
260 hidden: hiddenProviders,
261 });
262 }
263 // Kept usage only stands in for the same environments and provider filter.
264 const kept =
265 lastAnswered?.selection === selectedEnvironmentIds && lastAnswered.hidden === hiddenProviders
266 ? lastAnswered
267 : null;
268 // With no answers, even failed ones keep the last answered usage on screen.
269 const shown =
270 answeredCount > 0
271 ? { window: input, merged }
272 : (kept ?? (isPending ? null : { window: input, merged }));
273
274 return {
275 merged,
276 environments,
277 selectedEnvironments,
278 isPending,
279 shown,
280 isPartial: answeredCount > 0 && stillReporting > 0,
281 refresh,
282 };
283}