Skip to content

Commit 0a3da5c

Browse files
steipetesteipete
committed
perf: slim auth profile test imports
1 parent 38caa68 commit 0a3da5c

6 files changed

Lines changed: 58 additions & 55 deletions

‎src/agents/auth-profiles.ensureauthprofilestore.test.ts‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,13 +3,13 @@ import os from "node:os";
33
import path from "node:path";
44
import { afterEach, describe, expect, it, vi } from "vitest";
55
import type { ProviderExternalAuthProfile } from "../plugins/provider-external-auth.types.js";
6+
import { AUTH_STORE_VERSION, log } from "./auth-profiles/constants.js";
67
import {
78
clearRuntimeAuthProfileStoreSnapshots,
89
ensureAuthProfileStore,
910
loadAuthProfileStoreForRuntime,
1011
saveAuthProfileStore,
11-
} from "./auth-profiles.js";
12-
import { AUTH_STORE_VERSION, log } from "./auth-profiles/constants.js";
12+
} from "./auth-profiles/store.js";
1313
import type { AuthProfileCredential } from "./auth-profiles/types.js";
1414

1515
const resolveExternalAuthProfilesWithPluginsMock = vi.hoisted(() =>

‎src/agents/auth-profiles.getsoonestcooldownexpiry.test.ts‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
import { describe, expect, it } from "vitest";
2-
import type { AuthProfileStore } from "./auth-profiles.js";
3-
import { getSoonestCooldownExpiry } from "./auth-profiles.js";
2+
import type { AuthProfileStore } from "./auth-profiles/types.js";
3+
import { getSoonestCooldownExpiry } from "./auth-profiles/usage-state.js";
44

55
function makeStore(usageStats?: AuthProfileStore["usageStats"]): AuthProfileStore {
66
return {

‎src/agents/auth-profiles.store-cache.test.ts‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,8 +2,11 @@ import fs from "node:fs";
22
import os from "node:os";
33
import path from "node:path";
44
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
5-
import { clearRuntimeAuthProfileStoreSnapshots, ensureAuthProfileStore } from "./auth-profiles.js";
65
import { AUTH_STORE_VERSION } from "./auth-profiles/constants.js";
6+
import {
7+
clearRuntimeAuthProfileStoreSnapshots,
8+
ensureAuthProfileStore,
9+
} from "./auth-profiles/store.js";
710
import type { OAuthCredential } from "./auth-profiles/types.js";
811

912
type RuntimeOnlyOverlay = { profileId: string; credential: OAuthCredential };

‎src/agents/auth-profiles/usage-state.ts‎

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -64,6 +64,54 @@ export function isProfileInCooldown(
6464
return unusableUntil ? ts < unusableUntil : false;
6565
}
6666

67+
/**
68+
* Return the soonest `unusableUntil` timestamp (ms epoch) among the given
69+
* profiles, or `null` when no profile has a recorded cooldown. Note: the
70+
* returned timestamp may be in the past if the cooldown has already expired.
71+
*/
72+
export function getSoonestCooldownExpiry(
73+
store: AuthProfileStore,
74+
profileIds: string[],
75+
options?: { now?: number; forModel?: string },
76+
): number | null {
77+
const ts = options?.now ?? Date.now();
78+
let soonest: number | null = null;
79+
let latestMatchingModelCooldown: number | null = null;
80+
for (const id of profileIds) {
81+
const stats = store.usageStats?.[id];
82+
if (!stats) {
83+
continue;
84+
}
85+
if (shouldBypassModelScopedCooldown(stats, ts, options?.forModel)) {
86+
continue;
87+
}
88+
const until = resolveProfileUnusableUntil(stats);
89+
if (typeof until !== "number" || !Number.isFinite(until) || until <= 0) {
90+
continue;
91+
}
92+
const matchingModelScopedCooldown =
93+
options?.forModel &&
94+
stats.cooldownReason === "rate_limit" &&
95+
stats.cooldownModel === options.forModel &&
96+
!isActiveUnusableWindow(stats.disabledUntil, ts);
97+
if (matchingModelScopedCooldown) {
98+
latestMatchingModelCooldown =
99+
latestMatchingModelCooldown === null ? until : Math.max(latestMatchingModelCooldown, until);
100+
continue;
101+
}
102+
if (soonest === null || until < soonest) {
103+
soonest = until;
104+
}
105+
}
106+
if (soonest === null) {
107+
return latestMatchingModelCooldown;
108+
}
109+
if (latestMatchingModelCooldown === null) {
110+
return soonest;
111+
}
112+
return Math.min(soonest, latestMatchingModelCooldown);
113+
}
114+
67115
/**
68116
* Clear expired cooldowns from all profiles in the store.
69117
*

‎src/agents/auth-profiles/usage.ts‎

Lines changed: 1 addition & 49 deletions
Original file line numberDiff line numberDiff line change
@@ -9,10 +9,10 @@ import {
99
isAuthCooldownBypassedForProvider,
1010
isProfileInCooldown,
1111
resolveProfileUnusableUntil,
12-
shouldBypassModelScopedCooldown,
1312
} from "./usage-state.js";
1413
export {
1514
clearExpiredCooldowns,
15+
getSoonestCooldownExpiry,
1616
isProfileInCooldown,
1717
resolveProfileUnusableUntil,
1818
} from "./usage-state.js";
@@ -309,54 +309,6 @@ export function resolveProfilesUnavailableReason(params: {
309309
return best;
310310
}
311311

312-
/**
313-
* Return the soonest `unusableUntil` timestamp (ms epoch) among the given
314-
* profiles, or `null` when no profile has a recorded cooldown. Note: the
315-
* returned timestamp may be in the past if the cooldown has already expired.
316-
*/
317-
export function getSoonestCooldownExpiry(
318-
store: AuthProfileStore,
319-
profileIds: string[],
320-
options?: { now?: number; forModel?: string },
321-
): number | null {
322-
const ts = options?.now ?? Date.now();
323-
let soonest: number | null = null;
324-
let latestMatchingModelCooldown: number | null = null;
325-
for (const id of profileIds) {
326-
const stats = store.usageStats?.[id];
327-
if (!stats) {
328-
continue;
329-
}
330-
if (shouldBypassModelScopedCooldown(stats, ts, options?.forModel)) {
331-
continue;
332-
}
333-
const until = resolveProfileUnusableUntil(stats);
334-
if (typeof until !== "number" || !Number.isFinite(until) || until <= 0) {
335-
continue;
336-
}
337-
const matchingModelScopedCooldown =
338-
options?.forModel &&
339-
stats.cooldownReason === "rate_limit" &&
340-
stats.cooldownModel === options.forModel &&
341-
!isActiveUnusableWindow(stats.disabledUntil, ts);
342-
if (matchingModelScopedCooldown) {
343-
latestMatchingModelCooldown =
344-
latestMatchingModelCooldown === null ? until : Math.max(latestMatchingModelCooldown, until);
345-
continue;
346-
}
347-
if (soonest === null || until < soonest) {
348-
soonest = until;
349-
}
350-
}
351-
if (soonest === null) {
352-
return latestMatchingModelCooldown;
353-
}
354-
if (latestMatchingModelCooldown === null) {
355-
return soonest;
356-
}
357-
return Math.min(soonest, latestMatchingModelCooldown);
358-
}
359-
360312
/**
361313
* Mark a profile as successfully used. Resets error count and updates lastUsed.
362314
* Uses store lock to avoid overwriting concurrent usage updates.

‎src/agents/pi-model-discovery.synthetic-auth.test.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@ import fs from "node:fs/promises";
22
import os from "node:os";
33
import path from "node:path";
44
import { afterEach, beforeAll, beforeEach, describe, expect, it, vi } from "vitest";
5-
import { saveAuthProfileStore } from "./auth-profiles.js";
5+
import { saveAuthProfileStore } from "./auth-profiles/store.js";
66

77
const resolveRuntimeSyntheticAuthProviderRefs = vi.hoisted(() => vi.fn(() => ["claude-cli"]));
88

0 commit comments

Comments
 (0)