| 1 | import { |
| 2 | EnvironmentId, |
| 3 | ProviderDriverKind, |
| 4 | ProviderInstanceId, |
| 5 | UsageDay, |
| 6 | USAGE_CONTRACT_VERSION, |
| 7 | type ServerProvider, |
| 8 | type UsageSummary, |
| 9 | } from "@t3tools/contracts"; |
| 10 | import * as Effect from "effect/Effect"; |
| 11 | import { AsyncResult, Atom, AtomRegistry } from "effect/reactivity"; |
| 12 | import { afterEach, describe, expect, it, vi } from "vite-plus/test"; |
| 13 | |
| 14 | import type { EnvironmentPresentation } from "../connection/presentation.ts"; |
| 15 | import { EnvironmentRpcUnavailableError } from "../rpc/client.ts"; |
| 16 | import { |
| 17 | cursorKeychainAccessEnvironments, |
| 18 | needsCursorKeychainAccess, |
| 19 | refreshUsage, |
| 20 | refreshUsageLimits, |
| 21 | } from "./usage.ts"; |
| 22 | |
| 23 | const input = { |
| 24 | sinceDay: UsageDay.make("2026-09-05"), |
| 25 | untilDay: UsageDay.make("2026-09-05"), |
| 26 | timeZone: "UTC", |
| 27 | }; |
| 28 | const pricing = { status: "fresh" as const, source: "test", fetchedAt: null, knownModels: 1 }; |
| 29 | const summary: UsageSummary = { |
| 30 | ...input, |
| 31 | contractVersion: USAGE_CONTRACT_VERSION, |
| 32 | readAt: "2026-09-05T12:00:00Z", |
| 33 | buckets: [], |
| 34 | sources: [], |
| 35 | pricing, |
| 36 | scanDurationMs: 1, |
| 37 | }; |
| 38 | const registries: AtomRegistry.AtomRegistry[] = []; |
| 39 | afterEach(() => { |
| 40 | for (const registry of registries.splice(0)) registry.dispose(); |
| 41 | }); |
| 42 | |
| 43 | function harness(ids = ["a"]) { |
| 44 | const registry = AtomRegistry.make(); |
| 45 | registries.push(registry); |
| 46 | const environments = ids.map((id) => { |
| 47 | const environmentId = EnvironmentId.make(id); |
| 48 | const rates = Promise.withResolvers< |
| 49 | AsyncResult.Success<typeof pricing> | AsyncResult.Failure<never, unknown> |
| 50 | >(); |
| 51 | const scan = Promise.withResolvers<UsageSummary>(); |
| 52 | const scanStarted = Promise.withResolvers<void>(); |
| 53 | const presentation = Atom.make({ |
| 54 | connection: { phase: "connected" }, |
| 55 | } as EnvironmentPresentation | null); |
| 56 | const query = Atom.make( |
| 57 | Effect.promise(() => { |
| 58 | scanStarted.resolve(); |
| 59 | return scan.promise; |
| 60 | }), |
| 61 | ); |
| 62 | return { environmentId, rates, scan, scanStarted, presentation, query }; |
| 63 | }); |
| 64 | function get(environmentId: EnvironmentId) { |
| 65 | const environment = environments.find((entry) => entry.environmentId === environmentId); |
| 66 | if (!environment) throw new Error(`Unknown environment: ${environmentId}`); |
| 67 | return environment; |
| 68 | } |
| 69 | const options = { |
| 70 | registry, |
| 71 | environmentIds: environments.map((entry) => entry.environmentId), |
| 72 | input, |
| 73 | server: { |
| 74 | usageSummary: ({ environmentId }: { environmentId: EnvironmentId }) => |
| 75 | get(environmentId).query, |
| 76 | refreshUsageRates: { |
| 77 | label: "test:rates", |
| 78 | requiredScopes: () => [], |
| 79 | permissionAtom: () => Atom.make(true), |
| 80 | run: ( |
| 81 | _registry: AtomRegistry.AtomRegistry, |
| 82 | { environmentId }: { environmentId: EnvironmentId }, |
| 83 | ) => get(environmentId).rates.promise, |
| 84 | }, |
| 85 | }, |
| 86 | presentations: { |
| 87 | presentationAtom: (environmentId: EnvironmentId) => get(environmentId).presentation, |
| 88 | }, |
| 89 | } satisfies Parameters<typeof refreshUsage>[0]; |
| 90 | return { registry, environments, refresh: () => refreshUsage(options) }; |
| 91 | } |
| 92 | |
| 93 | describe("manual usage refresh", () => { |
| 94 | it.each(["success", "failure"])("waits for the rescan after a pricing %s", async (result) => { |
| 95 | const { |
| 96 | environments: [environment], |
| 97 | refresh, |
| 98 | } = harness(); |
| 99 | const entry = environment!; |
| 100 | let finished = false; |
| 101 | const refreshing = refresh().then(() => { |
| 102 | finished = true; |
| 103 | }); |
| 104 | expect(finished).toBe(false); |
| 105 | entry.rates.resolve( |
| 106 | result === "success" |
| 107 | ? AsyncResult.success(pricing) |
| 108 | : AsyncResult.fail(new Error("Pricing offline")), |
| 109 | ); |
| 110 | await entry.scanStarted.promise; |
| 111 | expect(finished).toBe(false); |
| 112 | entry.scan.resolve(summary); |
| 113 | await refreshing; |
| 114 | expect(finished).toBe(true); |
| 115 | }); |
| 116 | |
| 117 | it("settles when an environment disconnects during the rescan", async () => { |
| 118 | const { |
| 119 | registry, |
| 120 | environments: [environment], |
| 121 | refresh, |
| 122 | } = harness(); |
| 123 | const entry = environment!; |
| 124 | const refreshing = refresh(); |
| 125 | entry.rates.resolve(AsyncResult.success(pricing)); |
| 126 | await entry.scanStarted.promise; |
| 127 | registry.set(entry.presentation, null); |
| 128 | await refreshing; |
| 129 | }); |
| 130 | |
| 131 | it("waits for healthy environments without waiting for a recovering environment", async () => { |
| 132 | const { registry, environments, refresh } = harness(["healthy", "recovering"]); |
| 133 | const [healthy, recovering] = environments; |
| 134 | registry.set(recovering!.presentation, null); |
| 135 | let finished = false; |
| 136 | const refreshing = refresh().then(() => { |
| 137 | finished = true; |
| 138 | }); |
| 139 | for (const entry of environments) entry.rates.resolve(AsyncResult.success(pricing)); |
| 140 | await healthy!.scanStarted.promise; |
| 141 | expect(finished).toBe(false); |
| 142 | healthy!.scan.resolve(summary); |
| 143 | await refreshing; |
| 144 | expect(finished).toBe(true); |
| 145 | }); |
| 146 | |
| 147 | it("settles when connected state has no usable RPC session", async () => { |
| 148 | const { |
| 149 | environments: [environment], |
| 150 | refresh, |
| 151 | } = harness(); |
| 152 | const entry = environment!; |
| 153 | const refreshing = refresh(); |
| 154 | entry.rates.resolve( |
| 155 | AsyncResult.fail( |
| 156 | new EnvironmentRpcUnavailableError({ |
| 157 | environmentId: entry.environmentId, |
| 158 | message: "No session", |
| 159 | }), |
| 160 | ), |
| 161 | ); |
| 162 | await refreshing; |
| 163 | }); |
| 164 | |
| 165 | it("replaces a scan that started before pricing was refreshed", async () => { |
| 166 | const { |
| 167 | registry, |
| 168 | environments: [environment], |
| 169 | refresh, |
| 170 | } = harness(); |
| 171 | const entry = environment!; |
| 172 | let reads = 0; |
| 173 | const rescanned = Promise.withResolvers<void>(); |
| 174 | const query = Atom.make( |
| 175 | Effect.promise(() => { |
| 176 | reads += 1; |
| 177 | if (reads > 1) { |
| 178 | rescanned.resolve(); |
| 179 | return Promise.resolve(summary); |
| 180 | } |
| 181 | return new Promise<UsageSummary>(() => {}); |
| 182 | }), |
| 183 | ); |
| 184 | entry.query = query; |
| 185 | const unmount = registry.mount(query); |
| 186 | expect(reads).toBe(1); |
| 187 | const refreshing = refresh(); |
| 188 | entry.rates.resolve(AsyncResult.success(pricing)); |
| 189 | await rescanned.promise; |
| 190 | await refreshing; |
| 191 | expect(reads).toBe(2); |
| 192 | unmount(); |
| 193 | }); |
| 194 | }); |
| 195 | |
| 196 | describe("limits refresh cooldown", () => { |
| 197 | it("runs a fresh check after an in-flight check when settings change", async () => { |
| 198 | const id = EnvironmentId.make("limits-after-enable"); |
| 199 | const oldCheck = Promise.withResolvers<string>(); |
| 200 | const first = refreshUsageLimits(id, () => oldCheck.promise, true); |
| 201 | const newCheck = vi.fn(async () => "new limits"); |
| 202 | const afterEnable = refreshUsageLimits(id, newCheck, false, true); |
| 203 | expect(newCheck).not.toHaveBeenCalled(); |
| 204 | oldCheck.resolve("old limits"); |
| 205 | expect(await first).toBe("old limits"); |
| 206 | expect(await afterEnable).toBe("new limits"); |
| 207 | expect(newCheck).toHaveBeenCalledTimes(1); |
| 208 | }); |
| 209 | |
| 210 | it("joins manual calls and gates automatic refreshes after success or failure", async () => { |
| 211 | const clock = vi.spyOn(Date, "now").mockReturnValue(1_000); |
| 212 | try { |
| 213 | for (const fails of [false, true]) { |
| 214 | const id = EnvironmentId.make(`limits-${fails}`); |
| 215 | const pending = Promise.withResolvers<string>(); |
| 216 | const refresh = vi.fn(() => pending.promise); |
| 217 | const first = refreshUsageLimits(id, refresh, true); |
| 218 | await refreshUsageLimits(id, refresh, true); |
| 219 | const manual = refreshUsageLimits(id, refresh); |
| 220 | const settled = vi.fn(); |
| 221 | void manual.then(settled, settled); |
| 222 | expect(settled).not.toHaveBeenCalled(); |
| 223 | expect(refresh).toHaveBeenCalledTimes(1); |
| 224 | if (fails) { |
| 225 | const firstFailure = expect(first).rejects.toThrow("unavailable"); |
| 226 | const manualFailure = expect(manual).rejects.toThrow("unavailable"); |
| 227 | pending.reject(new Error("unavailable")); |
| 228 | await Promise.all([firstFailure, manualFailure]); |
| 229 | } else { |
| 230 | pending.resolve("quota"); |
| 231 | expect(await first).toBe("quota"); |
| 232 | expect(await manual).toBe("quota"); |
| 233 | } |
| 234 | expect(settled).toHaveBeenCalledTimes(1); |
| 235 | const next = vi.fn(async () => undefined); |
| 236 | clock.mockReturnValue(300_999); |
| 237 | await refreshUsageLimits(id, next, true); |
| 238 | expect(next).not.toHaveBeenCalled(); |
| 239 | clock.mockReturnValue(301_000); |
| 240 | await refreshUsageLimits(id, next, true); |
| 241 | expect(next).toHaveBeenCalledTimes(1); |
| 242 | await refreshUsageLimits(id, next); |
| 243 | expect(next).toHaveBeenCalledTimes(2); |
| 244 | clock.mockReturnValue(1_000); |
| 245 | } |
| 246 | } finally { |
| 247 | clock.mockRestore(); |
| 248 | } |
| 249 | }); |
| 250 | }); |
| 251 | |
| 252 | describe("needsCursorKeychainAccess", () => { |
| 253 | const cursorPrompt: UsageSummary = { |
| 254 | ...summary, |
| 255 | sources: [ |
| 256 | { |
| 257 | fingerprint: { |
| 258 | hostId: "host", |
| 259 | provider: "cursor", |
| 260 | resolvedHomePath: "/Users/me/.cursor/auth.json", |
| 261 | volumeId: "volume", |
| 262 | }, |
| 263 | status: "ok", |
| 264 | scannedFiles: 0, |
| 265 | skippedFiles: 0, |
| 266 | malformedRecords: 0, |
| 267 | distinctSessions: 0, |
| 268 | message: "Cursor account usage is off on this environment.", |
| 269 | action: "enableCursorKeychain", |
| 270 | }, |
| 271 | ], |
| 272 | }; |
| 273 | const cursor = (status: ServerProvider["status"]): ServerProvider => ({ |
| 274 | instanceId: ProviderInstanceId.make("cursor"), |
| 275 | driver: ProviderDriverKind.make("cursor"), |
| 276 | enabled: status !== "disabled", |
| 277 | installed: status === "ready", |
| 278 | version: null, |
| 279 | status, |
| 280 | auth: { status: "unknown" }, |
| 281 | checkedAt: "2026-09-05T12:00:00.000Z", |
| 282 | models: [], |
| 283 | slashCommands: [], |
| 284 | skills: [], |
| 285 | }); |
| 286 | |
| 287 | it("offers access only when Cursor is ready on that environment", () => { |
| 288 | expect(needsCursorKeychainAccess(cursorPrompt, [cursor("ready")])).toBe(true); |
| 289 | expect(needsCursorKeychainAccess(cursorPrompt, [cursor("error")])).toBe(false); |
| 290 | expect(needsCursorKeychainAccess(cursorPrompt, [cursor("disabled")])).toBe(false); |
| 291 | expect(needsCursorKeychainAccess(cursorPrompt, [])).toBe(false); |
| 292 | expect(needsCursorKeychainAccess(cursorPrompt, null)).toBe(false); |
| 293 | expect(needsCursorKeychainAccess(summary, [cursor("ready")])).toBe(false); |
| 294 | }); |
| 295 | |
| 296 | it("stops offering access once any environment reads the Cursor account", () => { |
| 297 | const off = { summary: cursorPrompt, needsCursorKeychainAccess: true }; |
| 298 | const account: UsageSummary = { |
| 299 | ...summary, |
| 300 | sources: [ |
| 301 | { |
| 302 | fingerprint: { |
| 303 | hostId: "cursor.com", |
| 304 | provider: "cursor", |
| 305 | resolvedHomePath: "cursor-account:abc", |
| 306 | volumeId: "abc", |
| 307 | }, |
| 308 | status: "ok", |
| 309 | scannedFiles: 1, |
| 310 | skippedFiles: 0, |
| 311 | malformedRecords: 0, |
| 312 | distinctSessions: 1, |
| 313 | message: null, |
| 314 | }, |
| 315 | ], |
| 316 | }; |
| 317 | expect(cursorKeychainAccessEnvironments([off, off])).toEqual([off, off]); |
| 318 | expect( |
| 319 | cursorKeychainAccessEnvironments([ |
| 320 | off, |
| 321 | { summary: account, needsCursorKeychainAccess: false }, |
| 322 | ]), |
| 323 | ).toEqual([]); |
| 324 | }); |
| 325 | }); |