1import {
2 DEFAULT_SERVER_SETTINGS,
3 EnvironmentId,
4 UsageDay,
5 UsageReadError,
6 USAGE_CONTRACT_VERSION,
7 WS_METHODS,
8 type ServerConfig,
9 type ServerConfigStreamEvent,
10 ServerSettings,
11 type UsageSource,
12 type UsageSummary,
13 type UsageSummaryInput,
14} from "@t3tools/contracts";
15import { expect, it } from "@effect/vitest";
16import * as Deferred from "effect/Deferred";
17import * as Effect from "effect/Effect";
18import * as Layer from "effect/Layer";
19import * as Option from "effect/Option";
20import * as Queue from "effect/Queue";
21import * as Result from "effect/Result";
22import * as Schema from "effect/Schema";
23import * as Stream from "effect/Stream";
24import * as SubscriptionRef from "effect/SubscriptionRef";
25import { AsyncResult, Atom, AtomRegistry } from "effect/reactivity";
26
27import {
28 AVAILABLE_CONNECTION_STATE,
29 PrimaryConnectionTarget,
30 type PreparedConnection,
31 type SupervisorConnectionState,
32} from "../connection/model.ts";
33import * as EnvironmentRegistry from "../connection/registry.ts";
34import * as EnvironmentSupervisor from "../connection/supervisor.ts";
35import * as Persistence from "../platform/persistence.ts";
36import type { WsRpcProtocolClient } from "../rpc/protocol.ts";
37import type { RpcSession } from "../rpc/session.ts";
38import { createServerEnvironmentAtoms } from "./server.ts";
39
40const TARGET = new PrimaryConnectionTarget({
41 environmentId: EnvironmentId.make("usage-environment"),
42 label: "Usage environment",
43 httpBaseUrl: "https://environment.example.test",
44 wsBaseUrl: "wss://environment.example.test",
45});
46
47const INPUT: UsageSummaryInput = {
48 sinceDay: UsageDay.make("2026-09-01"),
49 untilDay: UsageDay.make("2026-09-04"),
50 timeZone: "UTC",
51};
52
53const CONFIG = {
54 settings: DEFAULT_SERVER_SETTINGS,
55 environment: { serverVersion: "0.0.1", capabilities: {} },
56} as ServerConfig;
57
58const PRICING = { status: "unavailable" as const, source: "test", fetchedAt: null, knownModels: 0 };
59
60const cursorSource = (refreshing: boolean): UsageSource => ({
61 fingerprint: {
62 hostId: "cursor.com",
63 provider: "cursor",
64 resolvedHomePath: "cursor-account:abc",
65 volumeId: "abc",
66 },
67 status: "ok",
68 scannedFiles: 1,
69 skippedFiles: 0,
70 malformedRecords: 0,
71 distinctSessions: 1,
72 message: null,
73 ...(refreshing ? { refreshing: true as const } : {}),
74});
75
76const sameServerSettings = Schema.toEquivalence(ServerSettings);
77
78const makeHarness = Effect.fn("ServerUsageTest.makeHarness")(function* (
79 beforeRead: (
80 request: number,
81 input: UsageSummaryInput,
82 ) => Effect.Effect<void, UsageReadError> = () => Effect.void,
83 sources: (input: UsageSummaryInput) => readonly UsageSource[] = () => [],
84) {
85 const events = yield* Queue.unbounded<ServerConfigStreamEvent>();
86 let settings = DEFAULT_SERVER_SETTINGS;
87 let requests = 0;
88 const inputs: UsageSummaryInput[] = [];
89 const client = {
90 [WS_METHODS.serverRefreshUsageRates]: () => Effect.succeed(PRICING),
91 [WS_METHODS.subscribeServerConfig]: () =>
92 Stream.concat(
93 Stream.make({ version: 1 as const, type: "snapshot" as const, config: CONFIG }),
94 Stream.fromQueue(events),
95 ),
96 [WS_METHODS.serverGetUsageSummary]: (input: UsageSummaryInput) =>
97 Effect.gen(function* () {
98 requests += 1;
99 inputs.push(input);
100 const price = settings.usagePriceOverrides["custom-model"]?.inputCostPerMillionTokens ?? 0;
101 yield* beforeRead(requests, input);
102 return {
103 contractVersion: USAGE_CONTRACT_VERSION,
104 readAt: "2026-09-04T12:00:00Z",
105 ...input,
106 buckets: [
107 {
108 day: input.sinceDay,
109 provider: "codex",
110 model: "custom-model",
111 totals: {
112 uncachedInputTokens: 1_000_000,
113 cachedInputTokens: 0,
114 cacheCreationTokens: 0,
115 outputTokens: 0,
116 reasoningTokens: 0,
117 },
118 costUsd: price,
119 cacheSavingsUsd: 0,
120 costSource: price === 0 ? "unpriced" : "modelPriced",
121 records: 1,
122 unpricedRecords: price === 0 ? 1 : 0,
123 sessions: 1,
124 },
125 ],
126 sources: sources(input),
127 pricing: PRICING,
128 scanDurationMs: 0,
129 } satisfies UsageSummary;
130 }),
131 } as unknown as WsRpcProtocolClient;
132 const session: RpcSession = {
133 client,
134 initialConfig: Effect.succeed(CONFIG),
135 subscribeServerConfig: (input = {}) => client.subscribeServerConfig(input),
136 ready: Effect.void,
137 probe: Effect.void,
138 inboundFrames: Effect.succeed(0),
139 closed: Effect.never,
140 };
141 const supervisor = EnvironmentSupervisor.EnvironmentSupervisor.of({
142 target: TARGET,
143 state: yield* SubscriptionRef.make<SupervisorConnectionState>({
144 ...AVAILABLE_CONNECTION_STATE,
145 phase: "connected",
146 }),
147 session: yield* SubscriptionRef.make(Option.some(session)),
148 prepared: yield* SubscriptionRef.make(Option.none<PreparedConnection>()),
149 connect: Effect.void,
150 disconnect: Effect.void,
151 retryNow: Effect.void,
152 });
153 const environments = EnvironmentRegistry.EnvironmentRegistry.of({
154 run: (_environmentId, effect) =>
155 Effect.provideService(effect, EnvironmentSupervisor.EnvironmentSupervisor, supervisor),
156 followStream: (_environmentId, stream) =>
157 Stream.provideService(stream, EnvironmentSupervisor.EnvironmentSupervisor, supervisor),
158 } as EnvironmentRegistry.EnvironmentRegistry["Service"]);
159 const cache = Persistence.EnvironmentCacheStore.of({
160 loadShell: () => Effect.succeedNone,
161 saveShell: () => Effect.void,
162 loadThread: () => Effect.succeedNone,
163 saveThread: () => Effect.void,
164 removeThread: () => Effect.void,
165 loadServerConfig: () => Effect.succeedNone,
166 saveServerConfig: () => Effect.void,
167 loadVcsRefs: () => Effect.succeedNone,
168 saveVcsRefs: () => Effect.void,
169 removeVcsRefs: () => Effect.void,
170 clearVcsRefs: () => Effect.void,
171 clear: () => Effect.void,
172 });
173 const runtime = Atom.runtime(
174 Layer.merge(
175 Layer.succeed(EnvironmentRegistry.EnvironmentRegistry, environments),
176 Layer.succeed(Persistence.EnvironmentCacheStore, cache),
177 ),
178 );
179 const initialConfigValueAtom = Atom.make(CONFIG);
180 const atoms = createServerEnvironmentAtoms(runtime, {
181 initialConfigValueAtom: () => initialConfigValueAtom,
182 });
183 const registry = yield* Effect.acquireRelease(Effect.sync(AtomRegistry.make), (registry) =>
184 Effect.sync(() => registry.dispose()),
185 );
186 const updateSettings = Effect.fn("ServerUsageTest.updateSettings")(function* (
187 next: ServerSettings,
188 ) {
189 settings = next;
190 yield* Queue.offer(events, {
191 version: 1,
192 type: "settingsUpdated",
193 payload: { settings },
194 });
195 // An equal settings document keeps the previous object, so match on content.
196 yield* AtomRegistry.toStream(registry, atoms.settingsValueAtom(TARGET.environmentId)).pipe(
197 Stream.filter((current) => current !== null && sameServerSettings(current, next)),
198 Stream.runHead,
199 );
200 });
201 return {
202 registry,
203 atoms,
204 requests: () => requests,
205 inputs: () => inputs,
206 updateSettings,
207 summary: (input = INPUT) => atoms.usageSummary({ environmentId: TARGET.environmentId, input }),
208 };
209});
210
211function settledSummary<E>(
212 registry: AtomRegistry.AtomRegistry,
213 atom: Atom.Atom<AsyncResult.AsyncResult<UsageSummary, E>>,
214) {
215 return AtomRegistry.toStream(registry, atom).pipe(
216 Stream.filterMap((result) =>
217 AsyncResult.isSuccess(result) && !result.waiting
218 ? Result.succeed(result.value)
219 : Result.failVoid,
220 ),
221 Stream.runHead,
222 Effect.map(Option.getOrThrow),
223 );
224}
225
226const makeRefreshingHarness = Effect.fn("ServerUsageTest.makeRefreshingHarness")(function* (
227 followUp: Effect.Effect<void, UsageReadError>,
228) {
229 const followUpStarted = yield* Deferred.make<void>();
230 const harness = yield* makeHarness(
231 (_request, input) =>
232 input.awaitRefresh === true
233 ? Deferred.succeed(followUpStarted, undefined).pipe(Effect.andThen(followUp))
234 : Effect.void,
235 // The saved answer differs from the plain one, so a test can tell which was kept.
236 (input) => [
237 {
238 ...cursorSource(input.awaitRefresh !== true),
239 ...(input.provisional === true ? { scannedFiles: 0 } : {}),
240 },
241 ],
242 );
243 return { ...harness, followUpStarted };
244});
245
246function waitForCost<E>(
247 registry: AtomRegistry.AtomRegistry,
248 atom: Atom.Atom<AsyncResult.AsyncResult<UsageSummary, E>>,
249 cost: number,
250) {
251 return AtomRegistry.toStream(registry, atom).pipe(
252 Stream.filter(
253 (result) =>
254 AsyncResult.isSuccess(result) &&
255 !result.waiting &&
256 result.value.buckets[0]?.costUsd === cost,
257 ),
258 Stream.runHead,
259 );
260}
261
262it.effect("restarts a pending usage read after a price change", () =>
263 Effect.gen(function* () {
264 const started = yield* Deferred.make<void>();
265 const interrupted = yield* Deferred.make<void>();
266 const harness = yield* makeHarness((request) =>
267 request === 1
268 ? Deferred.succeed(started, undefined).pipe(
269 Effect.andThen(Effect.never),
270 Effect.onInterrupt(() => Deferred.succeed(interrupted, undefined)),
271 )
272 : Effect.void,
273 );
274 const summary = harness.summary();
275 const unmount = harness.registry.mount(summary);
276 yield* Deferred.await(started);
277 yield* harness.updateSettings({
278 ...DEFAULT_SERVER_SETTINGS,
279 usagePriceOverrides: {
280 "custom-model": { inputCostPerMillionTokens: 5, outputCostPerMillionTokens: 10 },
281 },
282 });
283 yield* waitForCost(harness.registry, summary, 5);
284 yield* Deferred.await(interrupted);
285 expect(harness.requests()).toBe(2);
286 unmount();
287 }),
288);
289
290it.effect("keeps the cached usage summary when the awaited refresh fails", () =>
291 Effect.gen(function* () {
292 const harness = yield* makeRefreshingHarness(
293 Effect.fail(new UsageReadError({ reason: "scanFailed", detail: "Cursor is down" })),
294 );
295 const summary = harness.summary();
296 const unmount = harness.registry.mount(summary);
297 expect((yield* settledSummary(harness.registry, summary)).sources).toEqual([
298 cursorSource(true),
299 ]);
300 expect(AsyncResult.isSuccess(harness.registry.get(summary))).toBe(true);
301 // The saved answer is followed by a plain read before the one that waits for Cursor.
302 expect(
303 harness.inputs().map(({ provisional, awaitRefresh }) => ({ provisional, awaitRefresh })),
304 ).toEqual([
305 { provisional: true, awaitRefresh: undefined },
306 { provisional: undefined, awaitRefresh: undefined },
307 { provisional: undefined, awaitRefresh: true },
308 ]);
309 // With figures on screen, a refresh reads them exactly instead of asking for saved ones.
310 yield* harness.updateSettings({
311 ...DEFAULT_SERVER_SETTINGS,
312 usagePriceOverrides: {
313 "custom-model": { inputCostPerMillionTokens: 5, outputCostPerMillionTokens: 10 },
314 },
315 });
316 yield* waitForCost(harness.registry, summary, 5);
317 expect(harness.inputs()[3]?.provisional).toBeUndefined();
318 unmount();
319 }),
320);
321
322it.effect("reads usage in full when the saved answer cannot be read", () =>
323 Effect.gen(function* () {
324 const harness = yield* makeHarness((_request, input) =>
325 input.provisional === true
326 ? Effect.fail(new UsageReadError({ reason: "scanFailed", detail: "saved scan unreadable" }))
327 : Effect.void,
328 );
329 const summary = harness.summary();
330 const unmount = harness.registry.mount(summary);
331 yield* settledSummary(harness.registry, summary);
332 expect(harness.inputs().map(({ provisional }) => provisional)).toEqual([true, undefined]);
333 unmount();
334 }),
335);