Skip to content

Commit c781939

Browse files
ndycodendycode
committed
Merge pull request #570 from ndycode/claude/audit-51-rotation-selection-tests
test: cover chooseAccount selection tiers and cursor discipline
2 parents 74d521b + 7d946f9 commit c781939

1 file changed

Lines changed: 366 additions & 0 deletions

File tree

‎test/rotation-account-selection.test.ts‎

Lines changed: 366 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,366 @@
1+
import { beforeEach, describe, expect, it, vi } from "vitest";
2+
import { AccountManager } from "../lib/accounts.js";
3+
import type { RuntimePolicyDecision } from "../lib/policy/runtime-policy.js";
4+
import { chooseAccount } from "../lib/runtime/rotation-account-selection.js";
5+
import { SessionAffinityStore } from "../lib/session-affinity.js";
6+
import type { AccountMetadataV3, AccountStorageV3 } from "../lib/storage.js";
7+
8+
const NOW = Date.now();
9+
const FAMILY = "gpt-5-codex" as const;
10+
11+
function storageWith(
12+
count: number,
13+
overridesByIndex: Record<number, Partial<AccountMetadataV3>> = {},
14+
): AccountStorageV3 {
15+
return {
16+
version: 3,
17+
activeIndex: 0,
18+
activeIndexByFamily: { [FAMILY]: 0 },
19+
accounts: Array.from({ length: count }, (_unused, index) => ({
20+
email: `account-${index + 1}@example.com`,
21+
accountId: `acc_${index + 1}`,
22+
refreshToken: `refresh-${index + 1}`,
23+
accessToken: `access-${index + 1}`,
24+
expiresAt: NOW + 3_600_000,
25+
addedAt: NOW - 60_000,
26+
lastUsed: NOW - (count - index) * 60_000,
27+
enabled: true,
28+
...overridesByIndex[index],
29+
})),
30+
};
31+
}
32+
33+
function manager(
34+
count = 3,
35+
overridesByIndex: Record<number, Partial<AccountMetadataV3>> = {},
36+
): AccountManager {
37+
return new AccountManager(undefined, storageWith(count, overridesByIndex));
38+
}
39+
40+
function policyWith(blocked: number[] = []): RuntimePolicyDecision {
41+
return {
42+
allowed: true,
43+
statusCode: 200,
44+
errorCode: null,
45+
reasons: [],
46+
projectKey: null,
47+
blockedAccountIndexes: new Set(blocked),
48+
scoreBoostByAccount: {},
49+
budgetEvaluations: [],
50+
};
51+
}
52+
53+
function baseParams(accountManager: AccountManager) {
54+
return {
55+
accountManager,
56+
sessionAffinityStore: null,
57+
sessionKey: null,
58+
family: FAMILY,
59+
model: null,
60+
attemptedIndexes: new Set<number>(),
61+
now: NOW,
62+
policy: null,
63+
pinnedIndex: null,
64+
};
65+
}
66+
67+
beforeEach(() => {
68+
// Circuit breakers and rate-limit trackers are module-level state.
69+
AccountManager.resetVolatileRuntimeState();
70+
});
71+
72+
describe("chooseAccount pinned selection (issue #474)", () => {
73+
it("returns the pinned account without advancing any cursor", () => {
74+
const accountManager = manager();
75+
const markSwitched = vi.spyOn(accountManager, "markSwitched");
76+
77+
const selected = chooseAccount({
78+
...baseParams(accountManager),
79+
pinnedIndex: 1,
80+
});
81+
82+
expect(selected?.index).toBe(1);
83+
// The proxy must not clobber the pin set by the CLI.
84+
expect(markSwitched).not.toHaveBeenCalled();
85+
});
86+
87+
it.each([
88+
[
89+
"already-attempted",
90+
{ pinnedIndex: 1, attemptedIndexes: new Set([1]) },
91+
],
92+
["missing", { pinnedIndex: 7 }],
93+
["policy-blocked", { pinnedIndex: 1, policy: policyWith([1]) }],
94+
] as const)(
95+
"refuses the pin with reason %s instead of falling back",
96+
(reason, overrides) => {
97+
const accountManager = manager();
98+
const skipReasons = new Map<number, string>();
99+
100+
const selected = chooseAccount({
101+
...baseParams(accountManager),
102+
...overrides,
103+
skipReasons,
104+
});
105+
106+
// A pin never falls back to another account; the request fails over
107+
// to the caller with the reason recorded.
108+
expect(selected).toBeNull();
109+
expect(skipReasons.get(overrides.pinnedIndex)).toBe(reason);
110+
},
111+
);
112+
113+
it("refuses a disabled pinned account", () => {
114+
const accountManager = manager(3, { 1: { enabled: false } });
115+
const skipReasons = new Map<number, string>();
116+
117+
const selected = chooseAccount({
118+
...baseParams(accountManager),
119+
pinnedIndex: 1,
120+
skipReasons,
121+
});
122+
123+
expect(selected).toBeNull();
124+
expect(skipReasons.get(1)).toBe("disabled");
125+
});
126+
127+
it("refuses a rate-limited pinned account with the runtime reason", () => {
128+
const accountManager = manager(3, {
129+
1: { rateLimitResetTimes: { [FAMILY]: NOW + 600_000 } },
130+
});
131+
const skipReasons = new Map<number, string>();
132+
133+
const selected = chooseAccount({
134+
...baseParams(accountManager),
135+
pinnedIndex: 1,
136+
skipReasons,
137+
});
138+
139+
expect(selected).toBeNull();
140+
expect(skipReasons.get(1)).toBe("rate-limited");
141+
});
142+
143+
it("refuses a cooling-down pinned account with the tagged reason", () => {
144+
const accountManager = manager();
145+
const pinned = accountManager.getAccountByIndex(1);
146+
if (!pinned) throw new Error("fixture account missing");
147+
accountManager.markAccountCoolingDown(pinned, 60_000, "auth-failure");
148+
const skipReasons = new Map<number, string>();
149+
150+
const selected = chooseAccount({
151+
...baseParams(accountManager),
152+
pinnedIndex: 1,
153+
skipReasons,
154+
});
155+
156+
expect(selected).toBeNull();
157+
expect(skipReasons.get(1)).toBe("cooling-down:auth-failure");
158+
});
159+
160+
it("refuses a pinned account whose circuit breaker is open", () => {
161+
const accountManager = manager();
162+
const pinned = accountManager.getAccountByIndex(1);
163+
if (!pinned) throw new Error("fixture account missing");
164+
// Trip the breaker (threshold is 3 failures).
165+
for (let i = 0; i < 3; i += 1) {
166+
accountManager.recordFailure(pinned, FAMILY, null);
167+
}
168+
const skipReasons = new Map<number, string>();
169+
170+
const selected = chooseAccount({
171+
...baseParams(accountManager),
172+
pinnedIndex: 1,
173+
skipReasons,
174+
});
175+
176+
expect(selected).toBeNull();
177+
expect(skipReasons.get(1)).toBe("circuit-open");
178+
});
179+
});
180+
181+
describe("chooseAccount session affinity tier", () => {
182+
it("prefers the remembered account and commits the cursor", () => {
183+
const accountManager = manager();
184+
const markSwitched = vi.spyOn(accountManager, "markSwitched");
185+
const store = new SessionAffinityStore();
186+
store.remember("sess-1", 2, NOW);
187+
188+
const selected = chooseAccount({
189+
...baseParams(accountManager),
190+
sessionAffinityStore: store,
191+
sessionKey: "sess-1",
192+
});
193+
194+
expect(selected?.index).toBe(2);
195+
expect(markSwitched).toHaveBeenCalledExactlyOnceWith(
196+
selected,
197+
"rotation",
198+
FAMILY,
199+
);
200+
});
201+
202+
it("records the skip reason and falls through when the sticky account is unusable", () => {
203+
const accountManager = manager(3, {
204+
2: { rateLimitResetTimes: { [FAMILY]: NOW + 600_000 } },
205+
});
206+
vi.spyOn(
207+
accountManager,
208+
"getCurrentOrNextForFamilyHybrid",
209+
).mockReturnValue(accountManager.getAccountByIndex(0));
210+
const store = new SessionAffinityStore();
211+
store.remember("sess-1", 2, NOW);
212+
const skipReasons = new Map<number, string>();
213+
214+
const selected = chooseAccount({
215+
...baseParams(accountManager),
216+
sessionAffinityStore: store,
217+
sessionKey: "sess-1",
218+
skipReasons,
219+
});
220+
221+
expect(skipReasons.get(2)).toBe("rate-limited");
222+
expect(selected?.index).toBe(0);
223+
});
224+
});
225+
226+
describe("chooseAccount hybrid tier and linear fallback", () => {
227+
it("returns the hybrid selector's pick without an extra cursor commit", () => {
228+
const accountManager = manager();
229+
vi.spyOn(
230+
accountManager,
231+
"getCurrentOrNextForFamilyHybrid",
232+
).mockReturnValue(accountManager.getAccountByIndex(1));
233+
const markSwitched = vi.spyOn(accountManager, "markSwitched");
234+
235+
const selected = chooseAccount(baseParams(accountManager));
236+
237+
expect(selected?.index).toBe(1);
238+
// The hybrid selector advances its own cursor internally.
239+
expect(markSwitched).not.toHaveBeenCalled();
240+
});
241+
242+
it("falls back to a linear scan that commits the cursor when the hybrid pick was attempted", () => {
243+
const accountManager = manager();
244+
vi.spyOn(
245+
accountManager,
246+
"getCurrentOrNextForFamilyHybrid",
247+
).mockReturnValue(accountManager.getAccountByIndex(0));
248+
const markSwitched = vi.spyOn(accountManager, "markSwitched");
249+
const skipReasons = new Map<number, string>();
250+
251+
const selected = chooseAccount({
252+
...baseParams(accountManager),
253+
attemptedIndexes: new Set([0]),
254+
policy: policyWith([1]),
255+
skipReasons,
256+
});
257+
258+
// 0 was attempted, 1 is policy-blocked, 2 wins and becomes the cursor.
259+
expect(selected?.index).toBe(2);
260+
expect(skipReasons.get(0)).toBe("already-attempted");
261+
expect(skipReasons.get(1)).toBe("policy-blocked");
262+
expect(markSwitched).toHaveBeenCalledExactlyOnceWith(
263+
selected,
264+
"rotation",
265+
FAMILY,
266+
);
267+
});
268+
269+
it("returns null with a reason per account when the pool is exhausted", () => {
270+
const accountManager = manager();
271+
vi.spyOn(
272+
accountManager,
273+
"getCurrentOrNextForFamilyHybrid",
274+
).mockReturnValue(accountManager.getAccountByIndex(0));
275+
const skipReasons = new Map<number, string>();
276+
277+
const selected = chooseAccount({
278+
...baseParams(accountManager),
279+
attemptedIndexes: new Set([0, 1, 2]),
280+
skipReasons,
281+
});
282+
283+
expect(selected).toBeNull();
284+
expect([...skipReasons.entries()]).toEqual([
285+
[0, "already-attempted"],
286+
[1, "already-attempted"],
287+
[2, "already-attempted"],
288+
]);
289+
});
290+
});
291+
292+
describe("chooseAccount sequential mode (issue #509)", () => {
293+
it("follows the drain-first selector and skips the affinity tier", () => {
294+
const accountManager = manager();
295+
const sequential = vi
296+
.spyOn(accountManager, "getCurrentOrNextForFamilySequential")
297+
.mockReturnValue(accountManager.getAccountByIndex(0));
298+
const store = new SessionAffinityStore();
299+
store.remember("sess-1", 2, NOW);
300+
const affinity = vi.spyOn(store, "getPreferredAccountIndex");
301+
302+
const selected = chooseAccount({
303+
...baseParams(accountManager),
304+
sessionAffinityStore: store,
305+
sessionKey: "sess-1",
306+
schedulingStrategy: "sequential",
307+
});
308+
309+
// Every request follows the single active account: no per-chat
310+
// stickiness consulted.
311+
expect(selected?.index).toBe(0);
312+
expect(sequential).toHaveBeenCalledTimes(1);
313+
expect(affinity).not.toHaveBeenCalled();
314+
});
315+
316+
it("routes around a policy-blocked primary without moving it", () => {
317+
const accountManager = manager();
318+
const sequential = vi
319+
.spyOn(accountManager, "getCurrentOrNextForFamilySequential")
320+
.mockReturnValue(accountManager.getAccountByIndex(0));
321+
const markSwitched = vi.spyOn(accountManager, "markSwitched");
322+
const policy = policyWith([0]);
323+
const skipReasons = new Map<number, string>();
324+
325+
const selected = chooseAccount({
326+
...baseParams(accountManager),
327+
policy,
328+
schedulingStrategy: "sequential",
329+
skipReasons,
330+
});
331+
332+
// The blocked set is threaded into the drain-first selector, the linear
333+
// fallback records the block, and the primary pointer stays put.
334+
expect(sequential).toHaveBeenCalledExactlyOnceWith(
335+
FAMILY,
336+
null,
337+
policy.blockedAccountIndexes,
338+
);
339+
expect(selected?.index).toBe(1);
340+
expect(skipReasons.get(0)).toBe("policy-blocked");
341+
expect(markSwitched).not.toHaveBeenCalled();
342+
});
343+
344+
it("retries another account without moving the drain-first primary", () => {
345+
const accountManager = manager();
346+
vi.spyOn(
347+
accountManager,
348+
"getCurrentOrNextForFamilySequential",
349+
).mockReturnValue(accountManager.getAccountByIndex(0));
350+
const markSwitched = vi.spyOn(accountManager, "markSwitched");
351+
const skipReasons = new Map<number, string>();
352+
353+
const selected = chooseAccount({
354+
...baseParams(accountManager),
355+
attemptedIndexes: new Set([0]),
356+
schedulingStrategy: "sequential",
357+
skipReasons,
358+
});
359+
360+
// The transiently failed active account stays primary: the fallback
361+
// only finds an account to TRY, it must not commit the cursor.
362+
expect(selected?.index).toBe(1);
363+
expect(markSwitched).not.toHaveBeenCalled();
364+
expect(skipReasons.get(0)).toBe("already-attempted");
365+
});
366+
});

0 commit comments

Comments
 (0)