acf4b45062
packages/ai: - fix: improve GitHub Copilot OAuth polling and Codex stream recovery - fix: allow google-vertex authentication via GOOGLE_CLOUD_API_KEY - fix: send Gemini/Claude provider-specific thinking headers and thought-signature fallbacks correctly - test: added coverage for Gemini CLI alignment, Codex streaming, and stream edge cases packages/coding-agent: - feat: add treeFilterMode setting for the session tree selector default - fix: prefer later-loaded explicit extensions when commands conflict - fix: truncate serialized tool results during compaction summarization to prevent overflow - fix: normalize CRLF in write tool previews - fix: use shell-based external editor launch on Windows - test: added compaction serialization and extension runner precedence coverage packages/tui: - fix: chain slash-command argument autocomplete after tab-completing the command name - fix: normalize pasted tabs in Input using configured indentation - fix: render blockquote lists, tables, and fenced code blocks correctly - fix: ignore unsupported Kitty CSI-u modifiers and enable modifyOtherKeys fallback cleanup - test: added editor, input, keys, and markdown regression coverage
718 lines
22 KiB
TypeScript
718 lines
22 KiB
TypeScript
import { afterEach, beforeEach, describe, expect, test, vi } from "bun:test";
|
|
import * as fs from "node:fs/promises";
|
|
import * as os from "node:os";
|
|
import * as path from "node:path";
|
|
import { AuthCredentialStore, AuthStorage } from "../src/auth-storage";
|
|
import type { UsageLimit, UsageProvider, UsageReport } from "../src/usage";
|
|
import * as oauthUtils from "../src/utils/oauth";
|
|
import type { OAuthCredentials } from "../src/utils/oauth/types";
|
|
|
|
const WEEK_MS = 7 * 24 * 60 * 60 * 1000;
|
|
const HOUR_MS = 60 * 60 * 1000;
|
|
|
|
const FIVE_HOUR_MS = 5 * HOUR_MS;
|
|
|
|
type UsageWindowSpec = {
|
|
usedFraction: number;
|
|
resetInMs: number;
|
|
};
|
|
|
|
type UsageWindowConfig = {
|
|
windowId: string;
|
|
windowLabel: string;
|
|
durationMs: number;
|
|
};
|
|
|
|
function createLimit(args: {
|
|
key: "primary" | "secondary";
|
|
windowId: string;
|
|
windowLabel: string;
|
|
durationMs: number;
|
|
usedFraction: number;
|
|
resetInMs: number;
|
|
}): UsageLimit {
|
|
const clamped = Math.min(Math.max(args.usedFraction, 0), 1);
|
|
const used = clamped * 100;
|
|
return {
|
|
id: `openai-codex:${args.key}`,
|
|
label: args.windowLabel,
|
|
scope: {
|
|
provider: "openai-codex",
|
|
windowId: args.windowId,
|
|
shared: true,
|
|
},
|
|
window: {
|
|
id: args.windowId,
|
|
label: args.windowLabel,
|
|
durationMs: args.durationMs,
|
|
resetsAt: Date.now() + args.resetInMs,
|
|
},
|
|
amount: {
|
|
unit: "percent",
|
|
used,
|
|
limit: 100,
|
|
remaining: 100 - used,
|
|
usedFraction: clamped,
|
|
remainingFraction: Math.max(0, 1 - clamped),
|
|
},
|
|
status: clamped >= 1 ? "exhausted" : clamped >= 0.9 ? "warning" : "ok",
|
|
};
|
|
}
|
|
|
|
function createCodexUsageReport(args: {
|
|
accountId: string;
|
|
primary: UsageWindowSpec;
|
|
secondary: UsageWindowSpec;
|
|
primaryWindow?: UsageWindowConfig;
|
|
secondaryWindow?: UsageWindowConfig;
|
|
}): UsageReport {
|
|
const primaryWindow = args.primaryWindow ?? { windowId: "1h", windowLabel: "1 Hour", durationMs: HOUR_MS };
|
|
const secondaryWindow = args.secondaryWindow ?? { windowId: "7d", windowLabel: "7 Day", durationMs: WEEK_MS };
|
|
return {
|
|
provider: "openai-codex",
|
|
fetchedAt: Date.now(),
|
|
limits: [
|
|
createLimit({
|
|
key: "primary",
|
|
windowId: primaryWindow.windowId,
|
|
windowLabel: primaryWindow.windowLabel,
|
|
durationMs: primaryWindow.durationMs,
|
|
usedFraction: args.primary.usedFraction,
|
|
resetInMs: args.primary.resetInMs,
|
|
}),
|
|
createLimit({
|
|
key: "secondary",
|
|
windowId: secondaryWindow.windowId,
|
|
windowLabel: secondaryWindow.windowLabel,
|
|
durationMs: secondaryWindow.durationMs,
|
|
usedFraction: args.secondary.usedFraction,
|
|
resetInMs: args.secondary.resetInMs,
|
|
}),
|
|
],
|
|
metadata: { accountId: args.accountId },
|
|
};
|
|
}
|
|
|
|
function createCredential(accountId: string, email: string): OAuthCredentials {
|
|
return {
|
|
access: `access-${accountId}`,
|
|
refresh: `refresh-${accountId}`,
|
|
expires: Date.now() + HOUR_MS,
|
|
accountId,
|
|
email,
|
|
};
|
|
}
|
|
|
|
describe("AuthStorage codex oauth ranking", () => {
|
|
let tempDir = "";
|
|
let store: AuthCredentialStore | null = null;
|
|
let authStorage: AuthStorage | null = null;
|
|
const usageByAccount = new Map<string, UsageReport>();
|
|
|
|
const usageProvider: UsageProvider = {
|
|
id: "openai-codex",
|
|
async fetchUsage(params) {
|
|
const accountId = params.credential.accountId;
|
|
if (!accountId) return null;
|
|
return usageByAccount.get(accountId) ?? null;
|
|
},
|
|
};
|
|
|
|
beforeEach(async () => {
|
|
tempDir = await fs.mkdtemp(path.join(os.tmpdir(), "pi-ai-auth-codex-selection-"));
|
|
store = await AuthCredentialStore.open(path.join(tempDir, "agent.db"));
|
|
authStorage = new AuthStorage(store, {
|
|
usageProviderResolver: provider => (provider === "openai-codex" ? usageProvider : undefined),
|
|
});
|
|
usageByAccount.clear();
|
|
vi.spyOn(oauthUtils, "getOAuthApiKey").mockImplementation(async (_provider, credentials) => {
|
|
const credential = credentials["openai-codex"] as OAuthCredentials | undefined;
|
|
if (!credential?.accountId) return null;
|
|
return {
|
|
apiKey: `api-${credential.accountId}`,
|
|
newCredentials: credential,
|
|
};
|
|
});
|
|
});
|
|
|
|
afterEach(async () => {
|
|
vi.restoreAllMocks();
|
|
store?.close();
|
|
store = null;
|
|
authStorage = null;
|
|
if (tempDir) {
|
|
await fs.rm(tempDir, { recursive: true, force: true });
|
|
tempDir = "";
|
|
}
|
|
});
|
|
|
|
test("prefers near-reset weekly account over lower-used far-reset account", async () => {
|
|
if (!authStorage) throw new Error("test setup failed");
|
|
|
|
await authStorage.set("openai-codex", [
|
|
{ type: "oauth", ...createCredential("acct-near", "near@example.com") },
|
|
{ type: "oauth", ...createCredential("acct-far", "far@example.com") },
|
|
]);
|
|
|
|
usageByAccount.set(
|
|
"acct-near",
|
|
createCodexUsageReport({
|
|
accountId: "acct-near",
|
|
primary: { usedFraction: 0.4, resetInMs: 10 * 60 * 1000 },
|
|
secondary: { usedFraction: 0.92, resetInMs: 15 * 60 * 1000 },
|
|
}),
|
|
);
|
|
usageByAccount.set(
|
|
"acct-far",
|
|
createCodexUsageReport({
|
|
accountId: "acct-far",
|
|
primary: { usedFraction: 0.3, resetInMs: 40 * 60 * 1000 },
|
|
secondary: { usedFraction: 0.55, resetInMs: 6 * 24 * 60 * 60 * 1000 },
|
|
}),
|
|
);
|
|
|
|
const apiKey = await authStorage.getApiKey("openai-codex", "session-weekly-reset");
|
|
expect(apiKey).toBe("api-acct-near");
|
|
});
|
|
|
|
test("prioritizes fresh 5h ticker account at 0% usage", async () => {
|
|
if (!authStorage) throw new Error("test setup failed");
|
|
|
|
await authStorage.set("openai-codex", [
|
|
{ type: "oauth", ...createCredential("acct-zero", "zero@example.com") },
|
|
{ type: "oauth", ...createCredential("acct-progress", "progress@example.com") },
|
|
]);
|
|
|
|
const fiveHourWindow: UsageWindowConfig = {
|
|
windowId: "5h",
|
|
windowLabel: "5 Hours",
|
|
durationMs: FIVE_HOUR_MS,
|
|
};
|
|
|
|
usageByAccount.set(
|
|
"acct-zero",
|
|
createCodexUsageReport({
|
|
accountId: "acct-zero",
|
|
primary: { usedFraction: 0, resetInMs: FIVE_HOUR_MS },
|
|
secondary: { usedFraction: 0.8, resetInMs: 2 * HOUR_MS },
|
|
primaryWindow: fiveHourWindow,
|
|
}),
|
|
);
|
|
usageByAccount.set(
|
|
"acct-progress",
|
|
createCodexUsageReport({
|
|
accountId: "acct-progress",
|
|
primary: { usedFraction: 0.05, resetInMs: 4 * HOUR_MS },
|
|
secondary: { usedFraction: 0.1, resetInMs: 6 * 24 * HOUR_MS },
|
|
primaryWindow: fiveHourWindow,
|
|
}),
|
|
);
|
|
|
|
const apiKey = await authStorage.getApiKey("openai-codex", "session-five-hour-start");
|
|
expect(apiKey).toBe("api-acct-zero");
|
|
});
|
|
test("skips exhausted weekly account even when reset is near", async () => {
|
|
if (!authStorage) throw new Error("test setup failed");
|
|
|
|
await authStorage.set("openai-codex", [
|
|
{ type: "oauth", ...createCredential("acct-exhausted", "exhausted@example.com") },
|
|
{ type: "oauth", ...createCredential("acct-healthy", "healthy@example.com") },
|
|
]);
|
|
|
|
usageByAccount.set(
|
|
"acct-exhausted",
|
|
createCodexUsageReport({
|
|
accountId: "acct-exhausted",
|
|
primary: { usedFraction: 1, resetInMs: 5 * 60 * 1000 },
|
|
secondary: { usedFraction: 1, resetInMs: 5 * 60 * 1000 },
|
|
}),
|
|
);
|
|
usageByAccount.set(
|
|
"acct-healthy",
|
|
createCodexUsageReport({
|
|
accountId: "acct-healthy",
|
|
primary: { usedFraction: 0.5, resetInMs: 20 * 60 * 1000 },
|
|
secondary: { usedFraction: 0.4, resetInMs: 3 * 24 * 60 * 60 * 1000 },
|
|
}),
|
|
);
|
|
|
|
const apiKey = await authStorage.getApiKey("openai-codex", "session-exhausted");
|
|
expect(apiKey).toBe("api-acct-healthy");
|
|
});
|
|
|
|
test("falls back to earliest-unblocking account when all exhausted", async () => {
|
|
if (!authStorage) throw new Error("test setup failed");
|
|
|
|
await authStorage.set("openai-codex", [
|
|
{ type: "oauth", ...createCredential("acct-soon", "soon@example.com") },
|
|
{ type: "oauth", ...createCredential("acct-later", "later@example.com") },
|
|
]);
|
|
|
|
usageByAccount.set(
|
|
"acct-soon",
|
|
createCodexUsageReport({
|
|
accountId: "acct-soon",
|
|
primary: { usedFraction: 1, resetInMs: 5 * 60 * 1000 },
|
|
secondary: { usedFraction: 1, resetInMs: 5 * 60 * 1000 },
|
|
}),
|
|
);
|
|
usageByAccount.set(
|
|
"acct-later",
|
|
createCodexUsageReport({
|
|
accountId: "acct-later",
|
|
primary: { usedFraction: 1, resetInMs: 30 * 60 * 1000 },
|
|
secondary: { usedFraction: 1, resetInMs: 30 * 60 * 1000 },
|
|
}),
|
|
);
|
|
|
|
const apiKey = await authStorage.getApiKey("openai-codex", "session-all-exhausted");
|
|
expect(apiKey).toBe("api-acct-soon");
|
|
});
|
|
|
|
test("works with single credential (no ranking)", async () => {
|
|
if (!authStorage) throw new Error("test setup failed");
|
|
|
|
await authStorage.set("openai-codex", [{ type: "oauth", ...createCredential("acct-solo", "solo@example.com") }]);
|
|
|
|
usageByAccount.set(
|
|
"acct-solo",
|
|
createCodexUsageReport({
|
|
accountId: "acct-solo",
|
|
primary: { usedFraction: 0.3, resetInMs: 20 * 60 * 1000 },
|
|
secondary: { usedFraction: 0.2, resetInMs: 5 * 24 * 60 * 60 * 1000 },
|
|
}),
|
|
);
|
|
|
|
const apiKey = await authStorage.getApiKey("openai-codex", "session-single");
|
|
expect(apiKey).toBe("api-acct-solo");
|
|
});
|
|
|
|
test("prefers Pro accounts for codex spark models over Plus accounts", async () => {
|
|
if (!authStorage) throw new Error("test setup failed");
|
|
|
|
await authStorage.set("openai-codex", [
|
|
{ type: "oauth", ...createCredential("acct-plus", "plus@example.com") },
|
|
{ type: "oauth", ...createCredential("acct-pro", "pro@example.com") },
|
|
]);
|
|
|
|
const plusReport = createCodexUsageReport({
|
|
accountId: "acct-plus",
|
|
primary: { usedFraction: 0.05, resetInMs: 30 * 60 * 1000 },
|
|
secondary: { usedFraction: 0.05, resetInMs: 6 * 24 * 60 * 60 * 1000 },
|
|
});
|
|
plusReport.metadata = { ...plusReport.metadata, planType: "plus" };
|
|
usageByAccount.set("acct-plus", plusReport);
|
|
|
|
const proReport = createCodexUsageReport({
|
|
accountId: "acct-pro",
|
|
primary: { usedFraction: 0.2, resetInMs: 30 * 60 * 1000 },
|
|
secondary: { usedFraction: 0.2, resetInMs: 6 * 24 * 60 * 60 * 1000 },
|
|
});
|
|
proReport.metadata = { ...proReport.metadata, planType: "pro" };
|
|
usageByAccount.set("acct-pro", proReport);
|
|
|
|
const apiKey = await authStorage.getApiKey("openai-codex", "session-spark-prefers-pro", {
|
|
modelId: "gpt-5.3-codex-spark",
|
|
});
|
|
expect(apiKey).toBe("api-acct-pro");
|
|
});
|
|
|
|
test("times out slow usage ranking instead of blocking first account selection", async () => {
|
|
if (!store) throw new Error("test setup failed");
|
|
|
|
const slowAuthStorage = new AuthStorage(store, {
|
|
usageProviderResolver: provider =>
|
|
provider === "openai-codex"
|
|
? ({
|
|
id: "openai-codex",
|
|
async fetchUsage(params) {
|
|
const { promise, resolve } = Promise.withResolvers<UsageReport | null>();
|
|
params.signal?.addEventListener("abort", () => resolve(null), { once: true });
|
|
return Promise.race([promise, Bun.sleep(200).then(() => null)]);
|
|
},
|
|
} satisfies UsageProvider)
|
|
: undefined,
|
|
usageRequestTimeoutMs: 10,
|
|
});
|
|
|
|
await slowAuthStorage.set("openai-codex", [
|
|
{ type: "oauth", ...createCredential("acct-first", "first@example.com") },
|
|
{ type: "oauth", ...createCredential("acct-second", "second@example.com") },
|
|
]);
|
|
|
|
const startedAt = Date.now();
|
|
const apiKey = await slowAuthStorage.getApiKey("openai-codex");
|
|
const elapsedMs = Date.now() - startedAt;
|
|
|
|
expect(apiKey).toBe("api-acct-first");
|
|
expect(elapsedMs).toBeLessThan(100);
|
|
});
|
|
|
|
test("sorts 3 accounts by weekly drain rate", async () => {
|
|
if (!authStorage) throw new Error("test setup failed");
|
|
|
|
await authStorage.set("openai-codex", [
|
|
{ type: "oauth", ...createCredential("acct-fast", "fast@example.com") },
|
|
{ type: "oauth", ...createCredential("acct-medium", "medium@example.com") },
|
|
{ type: "oauth", ...createCredential("acct-slow", "slow@example.com") },
|
|
]);
|
|
|
|
usageByAccount.set(
|
|
"acct-slow",
|
|
createCodexUsageReport({
|
|
accountId: "acct-slow",
|
|
primary: { usedFraction: 0.2, resetInMs: 30 * 60 * 1000 },
|
|
secondary: { usedFraction: 0.1, resetInMs: 6 * 24 * 60 * 60 * 1000 },
|
|
}),
|
|
);
|
|
usageByAccount.set(
|
|
"acct-medium",
|
|
createCodexUsageReport({
|
|
accountId: "acct-medium",
|
|
primary: { usedFraction: 0.2, resetInMs: 30 * 60 * 1000 },
|
|
secondary: { usedFraction: 0.3, resetInMs: 5 * 24 * 60 * 60 * 1000 },
|
|
}),
|
|
);
|
|
usageByAccount.set(
|
|
"acct-fast",
|
|
createCodexUsageReport({
|
|
accountId: "acct-fast",
|
|
primary: { usedFraction: 0.2, resetInMs: 30 * 60 * 1000 },
|
|
secondary: { usedFraction: 0.7, resetInMs: 3 * 24 * 60 * 60 * 1000 },
|
|
}),
|
|
);
|
|
|
|
const apiKey = await authStorage.getApiKey("openai-codex", "session-three-accounts");
|
|
expect(apiKey).toBe("api-acct-slow");
|
|
});
|
|
|
|
test("handles usage fetch failure gracefully (null report)", async () => {
|
|
if (!authStorage) throw new Error("test setup failed");
|
|
|
|
await authStorage.set("openai-codex", [
|
|
{ type: "oauth", ...createCredential("acct-null", "null@example.com") },
|
|
{ type: "oauth", ...createCredential("acct-known", "known@example.com") },
|
|
]);
|
|
|
|
// acct-null has no entry in usageByAccount — fetchUsage returns null
|
|
usageByAccount.set(
|
|
"acct-known",
|
|
createCodexUsageReport({
|
|
accountId: "acct-known",
|
|
primary: { usedFraction: 0.2, resetInMs: 30 * 60 * 1000 },
|
|
secondary: { usedFraction: 0.3, resetInMs: 5 * 24 * 60 * 60 * 1000 },
|
|
}),
|
|
);
|
|
|
|
const apiKey = await authStorage.getApiKey("openai-codex", "session-null-usage");
|
|
expect(apiKey).toBe("api-acct-known");
|
|
});
|
|
test("refreshes expired oauth candidates in parallel before selection", async () => {
|
|
if (!authStorage) throw new Error("test setup failed");
|
|
|
|
vi.spyOn(oauthUtils, "getOAuthApiKey").mockImplementation(async (_provider, credentials) => {
|
|
const credential = credentials["openai-codex"] as OAuthCredentials | undefined;
|
|
if (!credential?.accountId) return null;
|
|
|
|
let nextCredential = credential;
|
|
if (Date.now() >= credential.expires) {
|
|
nextCredential = await oauthUtils.refreshOAuthToken("openai-codex", credential);
|
|
}
|
|
|
|
if (nextCredential.accountId === "acct-first" || nextCredential.accountId === "acct-second") {
|
|
return null;
|
|
}
|
|
|
|
return {
|
|
apiKey: nextCredential.access,
|
|
newCredentials: nextCredential,
|
|
};
|
|
});
|
|
|
|
const refreshDelayMs = 75;
|
|
let inFlight = 0;
|
|
let maxConcurrent = 0;
|
|
const refreshStarts: number[] = [];
|
|
vi.spyOn(oauthUtils, "refreshOAuthToken").mockImplementation(async (_provider, credential) => {
|
|
refreshStarts.push(Date.now());
|
|
inFlight += 1;
|
|
maxConcurrent = Math.max(maxConcurrent, inFlight);
|
|
await Bun.sleep(refreshDelayMs);
|
|
inFlight -= 1;
|
|
return {
|
|
...credential,
|
|
access: `refreshed-${credential.accountId}`,
|
|
expires: Date.now() + HOUR_MS,
|
|
};
|
|
});
|
|
|
|
const expiredAt = Date.now() - HOUR_MS;
|
|
await authStorage.set("openai-codex", [
|
|
{ type: "oauth", ...createCredential("acct-first", "first@example.com"), expires: expiredAt },
|
|
{ type: "oauth", ...createCredential("acct-second", "second@example.com"), expires: expiredAt },
|
|
{ type: "oauth", ...createCredential("acct-third", "third@example.com"), expires: expiredAt },
|
|
]);
|
|
|
|
const startedAt = Date.now();
|
|
const apiKey = await authStorage.getApiKey("openai-codex");
|
|
const elapsedMs = Date.now() - startedAt;
|
|
|
|
expect(apiKey).toBe("refreshed-acct-third");
|
|
expect(refreshStarts).toHaveLength(3);
|
|
expect(maxConcurrent).toBe(3);
|
|
expect(elapsedMs).toBeLessThan(refreshDelayMs * 2);
|
|
});
|
|
});
|
|
|
|
// ─────────────────────────────────────────────────────────────────────────────
|
|
// Claude (Anthropic) ranking tests
|
|
// ─────────────────────────────────────────────────────────────────────────────
|
|
|
|
function createClaudeLimit(args: {
|
|
key: "5h" | "7d";
|
|
durationMs: number;
|
|
usedFraction: number;
|
|
resetInMs: number;
|
|
}): UsageLimit {
|
|
const clamped = Math.min(Math.max(args.usedFraction, 0), 1);
|
|
const used = clamped * 100;
|
|
const label = args.key === "5h" ? "Claude 5 Hour" : "Claude 7 Day";
|
|
return {
|
|
id: `anthropic:${args.key}`,
|
|
label,
|
|
scope: {
|
|
provider: "anthropic",
|
|
windowId: args.key,
|
|
shared: true,
|
|
},
|
|
window: {
|
|
id: args.key,
|
|
label,
|
|
durationMs: args.durationMs,
|
|
resetsAt: Date.now() + args.resetInMs,
|
|
},
|
|
amount: {
|
|
unit: "percent",
|
|
used,
|
|
limit: 100,
|
|
remaining: 100 - used,
|
|
usedFraction: clamped,
|
|
remainingFraction: Math.max(0, 1 - clamped),
|
|
},
|
|
status: clamped >= 1 ? "exhausted" : clamped >= 0.9 ? "warning" : "ok",
|
|
};
|
|
}
|
|
|
|
function createClaudeUsageReport(args: {
|
|
accountId: string;
|
|
primary: { usedFraction: number; resetInMs: number };
|
|
secondary: { usedFraction: number; resetInMs: number };
|
|
}): UsageReport {
|
|
return {
|
|
provider: "anthropic",
|
|
fetchedAt: Date.now(),
|
|
limits: [
|
|
createClaudeLimit({
|
|
key: "5h",
|
|
durationMs: FIVE_HOUR_MS,
|
|
usedFraction: args.primary.usedFraction,
|
|
resetInMs: args.primary.resetInMs,
|
|
}),
|
|
createClaudeLimit({
|
|
key: "7d",
|
|
durationMs: WEEK_MS,
|
|
usedFraction: args.secondary.usedFraction,
|
|
resetInMs: args.secondary.resetInMs,
|
|
}),
|
|
],
|
|
metadata: { accountId: args.accountId },
|
|
};
|
|
}
|
|
|
|
describe("AuthStorage claude oauth ranking", () => {
|
|
let tempDir = "";
|
|
let store: AuthCredentialStore | null = null;
|
|
let authStorage: AuthStorage | null = null;
|
|
const usageByAccount = new Map<string, UsageReport>();
|
|
|
|
const usageProvider: UsageProvider = {
|
|
id: "anthropic",
|
|
async fetchUsage(params) {
|
|
const accountId = params.credential.accountId;
|
|
if (!accountId) return null;
|
|
return usageByAccount.get(accountId) ?? null;
|
|
},
|
|
};
|
|
|
|
beforeEach(async () => {
|
|
tempDir = await fs.mkdtemp(path.join(os.tmpdir(), "pi-ai-auth-claude-selection-"));
|
|
store = await AuthCredentialStore.open(path.join(tempDir, "agent.db"));
|
|
authStorage = new AuthStorage(store, {
|
|
usageProviderResolver: provider => (provider === "anthropic" ? usageProvider : undefined),
|
|
});
|
|
usageByAccount.clear();
|
|
vi.spyOn(oauthUtils, "getOAuthApiKey").mockImplementation(async (_provider, credentials) => {
|
|
const credential = credentials.anthropic as OAuthCredentials | undefined;
|
|
if (!credential?.accountId) return null;
|
|
return {
|
|
apiKey: `api-${credential.accountId}`,
|
|
newCredentials: credential,
|
|
};
|
|
});
|
|
});
|
|
|
|
afterEach(async () => {
|
|
vi.restoreAllMocks();
|
|
store?.close();
|
|
store = null;
|
|
authStorage = null;
|
|
if (tempDir) {
|
|
await fs.rm(tempDir, { recursive: true, force: true });
|
|
tempDir = "";
|
|
}
|
|
});
|
|
|
|
test("prefers lower secondary drain rate account", async () => {
|
|
if (!authStorage) throw new Error("test setup failed");
|
|
|
|
await authStorage.set("anthropic", [
|
|
{ type: "oauth", ...createCredential("acct-near", "near@example.com") },
|
|
{ type: "oauth", ...createCredential("acct-far", "far@example.com") },
|
|
]);
|
|
|
|
usageByAccount.set(
|
|
"acct-near",
|
|
createClaudeUsageReport({
|
|
accountId: "acct-near",
|
|
primary: { usedFraction: 0.4, resetInMs: 2 * HOUR_MS },
|
|
secondary: { usedFraction: 0.92, resetInMs: 15 * 60 * 1000 },
|
|
}),
|
|
);
|
|
usageByAccount.set(
|
|
"acct-far",
|
|
createClaudeUsageReport({
|
|
accountId: "acct-far",
|
|
primary: { usedFraction: 0.3, resetInMs: 4 * HOUR_MS },
|
|
secondary: { usedFraction: 0.55, resetInMs: 6 * 24 * HOUR_MS },
|
|
}),
|
|
);
|
|
|
|
const apiKey = await authStorage.getApiKey("anthropic", "session-claude-drain");
|
|
expect(apiKey).toBe("api-acct-near");
|
|
});
|
|
|
|
test("skips exhausted account and picks healthy", async () => {
|
|
if (!authStorage) throw new Error("test setup failed");
|
|
|
|
await authStorage.set("anthropic", [
|
|
{ type: "oauth", ...createCredential("acct-exhausted", "exhausted@example.com") },
|
|
{ type: "oauth", ...createCredential("acct-healthy", "healthy@example.com") },
|
|
]);
|
|
|
|
usageByAccount.set(
|
|
"acct-exhausted",
|
|
createClaudeUsageReport({
|
|
accountId: "acct-exhausted",
|
|
primary: { usedFraction: 1, resetInMs: 5 * 60 * 1000 },
|
|
secondary: { usedFraction: 1, resetInMs: 5 * 60 * 1000 },
|
|
}),
|
|
);
|
|
usageByAccount.set(
|
|
"acct-healthy",
|
|
createClaudeUsageReport({
|
|
accountId: "acct-healthy",
|
|
primary: { usedFraction: 0.5, resetInMs: 3 * HOUR_MS },
|
|
secondary: { usedFraction: 0.4, resetInMs: 3 * 24 * HOUR_MS },
|
|
}),
|
|
);
|
|
|
|
const apiKey = await authStorage.getApiKey("anthropic", "session-claude-exhausted");
|
|
expect(apiKey).toBe("api-acct-healthy");
|
|
});
|
|
|
|
test("falls back to earliest-unblocking when all exhausted", async () => {
|
|
if (!authStorage) throw new Error("test setup failed");
|
|
|
|
await authStorage.set("anthropic", [
|
|
{ type: "oauth", ...createCredential("acct-soon", "soon@example.com") },
|
|
{ type: "oauth", ...createCredential("acct-later", "later@example.com") },
|
|
]);
|
|
|
|
usageByAccount.set(
|
|
"acct-soon",
|
|
createClaudeUsageReport({
|
|
accountId: "acct-soon",
|
|
primary: { usedFraction: 1, resetInMs: 5 * 60 * 1000 },
|
|
secondary: { usedFraction: 1, resetInMs: 5 * 60 * 1000 },
|
|
}),
|
|
);
|
|
usageByAccount.set(
|
|
"acct-later",
|
|
createClaudeUsageReport({
|
|
accountId: "acct-later",
|
|
primary: { usedFraction: 1, resetInMs: 30 * 60 * 1000 },
|
|
secondary: { usedFraction: 1, resetInMs: 30 * 60 * 1000 },
|
|
}),
|
|
);
|
|
|
|
const apiKey = await authStorage.getApiKey("anthropic", "session-claude-all-exhausted");
|
|
expect(apiKey).toBe("api-acct-soon");
|
|
});
|
|
|
|
test("sorts 3 accounts by secondary drain rate", async () => {
|
|
if (!authStorage) throw new Error("test setup failed");
|
|
|
|
await authStorage.set("anthropic", [
|
|
{ type: "oauth", ...createCredential("acct-fast", "fast@example.com") },
|
|
{ type: "oauth", ...createCredential("acct-medium", "medium@example.com") },
|
|
{ type: "oauth", ...createCredential("acct-slow", "slow@example.com") },
|
|
]);
|
|
|
|
usageByAccount.set(
|
|
"acct-slow",
|
|
createClaudeUsageReport({
|
|
accountId: "acct-slow",
|
|
primary: { usedFraction: 0.2, resetInMs: 4 * HOUR_MS },
|
|
secondary: { usedFraction: 0.1, resetInMs: 6 * 24 * HOUR_MS },
|
|
}),
|
|
);
|
|
usageByAccount.set(
|
|
"acct-medium",
|
|
createClaudeUsageReport({
|
|
accountId: "acct-medium",
|
|
primary: { usedFraction: 0.2, resetInMs: 4 * HOUR_MS },
|
|
secondary: { usedFraction: 0.3, resetInMs: 5 * 24 * HOUR_MS },
|
|
}),
|
|
);
|
|
usageByAccount.set(
|
|
"acct-fast",
|
|
createClaudeUsageReport({
|
|
accountId: "acct-fast",
|
|
primary: { usedFraction: 0.2, resetInMs: 4 * HOUR_MS },
|
|
secondary: { usedFraction: 0.7, resetInMs: 3 * 24 * HOUR_MS },
|
|
}),
|
|
);
|
|
|
|
const apiKey = await authStorage.getApiKey("anthropic", "session-claude-three");
|
|
expect(apiKey).toBe("api-acct-slow");
|
|
});
|
|
|
|
test("single credential works without ranking", async () => {
|
|
if (!authStorage) throw new Error("test setup failed");
|
|
|
|
await authStorage.set("anthropic", [{ type: "oauth", ...createCredential("acct-solo", "solo@example.com") }]);
|
|
|
|
usageByAccount.set(
|
|
"acct-solo",
|
|
createClaudeUsageReport({
|
|
accountId: "acct-solo",
|
|
primary: { usedFraction: 0.3, resetInMs: 3 * HOUR_MS },
|
|
secondary: { usedFraction: 0.2, resetInMs: 5 * 24 * HOUR_MS },
|
|
}),
|
|
);
|
|
|
|
const apiKey = await authStorage.getApiKey("anthropic", "session-claude-single");
|
|
expect(apiKey).toBe("api-acct-solo");
|
|
});
|
|
});
|