1import {
2 EnvironmentId,
3 UsageDay,
4 USAGE_CONTRACT_VERSION,
5 type UsageModelFilter,
6 type UsageSummary,
7 type UsageSummaryInput,
8} from "@t3tools/contracts";
9import * as Effect from "effect/Effect";
10import { AsyncResult, Atom, AtomRegistry } from "effect/reactivity";
11import { afterEach, describe, expect, it, vi } from "vite-plus/test";
12
13import type { EnvironmentPresentation } from "../connection/presentation.ts";
14import { refreshUsage, refreshUsageLimits } from "./usage.ts";
15
16const input = {
17 sinceDay: UsageDay.make("2026-09-05"),
18 untilDay: UsageDay.make("2026-09-05"),
19 timeZone: "UTC",
20};
21const pricing = { status: "fresh" as const, source: "test", fetchedAt: null, knownModels: 1 };
22const summary: UsageSummary = {
23 ...input,
24 contractVersion: USAGE_CONTRACT_VERSION,
25 readAt: "2026-09-05T12:00:00Z",
26 buckets: [],
27 sources: [],
28 pricing,
29 scanDurationMs: 1,
30};
31const registries: AtomRegistry.AtomRegistry[] = [];
32afterEach(() => {
33 for (const registry of registries.splice(0)) registry.dispose();
34});
35
36function harness(ids = ["a"], modelFilters: Record<string, UsageModelFilter> = {}) {
37 const registry = AtomRegistry.make();
38 registries.push(registry);
39 const environments = ids.map((id) => {
40 const environmentId = EnvironmentId.make(id);
41 const rates = Promise.withResolvers<
42 AsyncResult.Success<typeof pricing> | AsyncResult.Failure<never, unknown>
43 >();
44 const scan = Promise.withResolvers<UsageSummary>();
45 const scanStarted = Promise.withResolvers<void>();
46 const presentation = Atom.make({
47 connection: { phase: "connected" },
48 } as EnvironmentPresentation | null);
49 const query = Atom.make(
50 Effect.promise(() => {
51 scanStarted.resolve();
52 return scan.promise;
53 }),
54 );
55 return { environmentId, rates, scan, scanStarted, presentation, query };
56 });
57 function get(environmentId: EnvironmentId) {
58 const environment = environments.find((entry) => entry.environmentId === environmentId);
59 if (!environment) throw new Error(`Unknown environment: ${environmentId}`);
60 return environment;
61 }
62 const requested: { environmentId: EnvironmentId; input: UsageSummaryInput }[] = [];
63 const options = {
64 registry,
65 environmentIds: environments.map((entry) => entry.environmentId),
66 input,
67 server: {
68 usageSummary: (target: { environmentId: EnvironmentId; input: UsageSummaryInput }) => {
69 requested.push(target);
70 return get(target.environmentId).query;
71 },
72 usageModelFilterAtom: (environmentId: EnvironmentId) =>
73 Atom.make<UsageModelFilter | undefined>(modelFilters[environmentId]),
74 refreshUsageRates: {
75 label: "test:rates",
76 requiredScopes: () => [],
77 permissionAtom: () => Atom.make(true),
78 run: (
79 _registry: AtomRegistry.AtomRegistry,
80 { environmentId }: { environmentId: EnvironmentId },
81 ) => get(environmentId).rates.promise,
82 },
83 },
84 presentations: {
85 presentationAtom: (environmentId: EnvironmentId) => get(environmentId).presentation,
86 },
87 } satisfies Parameters<typeof refreshUsage>[0];
88 return { registry, environments, requested, refresh: () => refreshUsage(options) };
89}
90
91describe("manual usage refresh", () => {
92 it.each(["success", "failure"])("waits for the rescan after a pricing %s", async (result) => {
93 const {
94 environments: [environment],
95 refresh,
96 } = harness();
97 const entry = environment!;
98 let finished = false;
99 const refreshing = refresh().then(() => {
100 finished = true;
101 });
102 expect(finished).toBe(false);
103 entry.rates.resolve(
104 result === "success"
105 ? AsyncResult.success(pricing)
106 : AsyncResult.fail(new Error("Pricing offline")),
107 );
108 await entry.scanStarted.promise;
109 expect(finished).toBe(false);
110 entry.scan.resolve(summary);
111 await refreshing;
112 expect(finished).toBe(true);
113 });
114
115 it("rescans each environment with its own configured model filter", async () => {
116 const modelFilter = { include: [{ provider: "claude" as const, model: "orbit-1" }] };
117 const { environments, requested, refresh } = harness(["filtered", "plain"], {
118 filtered: modelFilter,
119 });
120 const refreshing = refresh();
121 for (const entry of environments) {
122 entry.rates.resolve(AsyncResult.success(pricing));
123 entry.scan.resolve(summary);
124 }
125 await refreshing;
126 expect(requested).toEqual([
127 { environmentId: "filtered", input: { ...input, modelFilter } },
128 { environmentId: "plain", input },
129 ]);
130 });
131
132 it("settles when an environment disconnects during the rescan", async () => {
133 const {
134 registry,
135 environments: [environment],
136 refresh,
137 } = harness();
138 const entry = environment!;
139 const refreshing = refresh();
140 entry.rates.resolve(AsyncResult.success(pricing));
141 await entry.scanStarted.promise;
142 registry.set(entry.presentation, null);
143 await refreshing;
144 });
145
146 it("waits for healthy environments without waiting for a recovering environment", async () => {
147 const { registry, environments, refresh } = harness(["healthy", "recovering"]);
148 const [healthy, recovering] = environments;
149 registry.set(recovering!.presentation, null);
150 let finished = false;
151 const refreshing = refresh().then(() => {
152 finished = true;
153 });
154 for (const entry of environments) entry.rates.resolve(AsyncResult.success(pricing));
155 await healthy!.scanStarted.promise;
156 expect(finished).toBe(false);
157 healthy!.scan.resolve(summary);
158 await refreshing;
159 expect(finished).toBe(true);
160 });
161
162 it("replaces a scan that started before pricing was refreshed", async () => {
163 const {
164 registry,
165 environments: [environment],
166 refresh,
167 } = harness();
168 const entry = environment!;
169 let reads = 0;
170 const rescanned = Promise.withResolvers<void>();
171 const query = Atom.make(
172 Effect.promise(() => {
173 reads += 1;
174 if (reads > 1) {
175 rescanned.resolve();
176 return Promise.resolve(summary);
177 }
178 return new Promise<UsageSummary>(() => {});
179 }),
180 );
181 entry.query = query;
182 const unmount = registry.mount(query);
183 expect(reads).toBe(1);
184 const refreshing = refresh();
185 entry.rates.resolve(AsyncResult.success(pricing));
186 await rescanned.promise;
187 await refreshing;
188 expect(reads).toBe(2);
189 unmount();
190 });
191});
192
193describe("limits refresh cooldown", () => {
194 it("shares an in-flight check and skips duplicate automatic refreshes", async () => {
195 const id = EnvironmentId.make("limits-in-flight");
196 const oldCheck = Promise.withResolvers<string>();
197 const first = refreshUsageLimits(id, () => oldCheck.promise, true);
198 const newCheck = vi.fn(async () => "new limits");
199 const manual = refreshUsageLimits(id, newCheck);
200 expect(await refreshUsageLimits(id, newCheck, true)).toBeUndefined();
201 expect(newCheck).not.toHaveBeenCalled();
202 oldCheck.resolve("old limits");
203 expect(await first).toBe("old limits");
204 expect(await manual).toBe("old limits");
205 expect(newCheck).not.toHaveBeenCalled();
206 });
207});