diff --git a/packages/pi-coding-agent/src/core/auth-storage.test.ts b/packages/pi-coding-agent/src/core/auth-storage.test.ts new file mode 100644 index 000000000..50b7ffedc --- /dev/null +++ b/packages/pi-coding-agent/src/core/auth-storage.test.ts @@ -0,0 +1,194 @@ +import { describe, it } from "node:test"; +import assert from "node:assert/strict"; +import { AuthStorage } from "./auth-storage.js"; + +// ─── helpers ────────────────────────────────────────────────────────────────── + +function makeKey(key: string) { + return { type: "api_key" as const, key }; +} + +function inMemory(data: Record = {}) { + return AuthStorage.inMemory(data as any); +} + +// ─── single credential (backward compat) ───────────────────────────────────── + +describe("AuthStorage — single credential (backward compat)", () => { + it("returns the api key for a provider with one key", async () => { + const storage = inMemory({ anthropic: makeKey("sk-abc") }); + const key = await storage.getApiKey("anthropic"); + assert.equal(key, "sk-abc"); + }); + + it("returns undefined for unknown provider", async () => { + const storage = inMemory({}); + const key = await storage.getApiKey("unknown"); + assert.equal(key, undefined); + }); + + it("runtime override takes precedence over stored key", async () => { + const storage = inMemory({ anthropic: makeKey("sk-stored") }); + storage.setRuntimeApiKey("anthropic", "sk-runtime"); + const key = await storage.getApiKey("anthropic"); + assert.equal(key, "sk-runtime"); + }); +}); + +// ─── multiple credentials ───────────────────────────────────────────────────── + +describe("AuthStorage — multiple credentials", () => { + it("round-robins across multiple api keys without sessionId", async () => { + const storage = inMemory({ + anthropic: [makeKey("sk-1"), makeKey("sk-2"), makeKey("sk-3")], + }); + + const keys = new Set(); + for (let i = 0; i < 6; i++) { + const k = await storage.getApiKey("anthropic"); + assert.ok(k, `call ${i} should return a key`); + keys.add(k); + } + // All three keys should have been selected across 6 calls + assert.deepEqual(keys, new Set(["sk-1", "sk-2", "sk-3"])); + }); + + it("session-sticky: same sessionId always picks the same key", async () => { + const storage = inMemory({ + anthropic: [makeKey("sk-1"), makeKey("sk-2"), makeKey("sk-3")], + }); + + const sessionId = "sess-abc"; + const first = await storage.getApiKey("anthropic", sessionId); + for (let i = 0; i < 5; i++) { + const k = await storage.getApiKey("anthropic", sessionId); + assert.equal(k, first, `call ${i} should be sticky to first selection`); + } + }); + + it("different sessionIds may select different keys", async () => { + const storage = inMemory({ + anthropic: [makeKey("sk-1"), makeKey("sk-2"), makeKey("sk-3")], + }); + + const results = new Set(); + for (let i = 0; i < 20; i++) { + const k = await storage.getApiKey("anthropic", `sess-${i}`); + if (k) results.add(k); + } + // With 20 different sessions and 3 keys, we should see more than one key + assert.ok(results.size > 1, "multiple sessions should hash to different keys"); + }); +}); + +// ─── login accumulation ─────────────────────────────────────────────────────── + +describe("AuthStorage — login accumulation", () => { + it("accumulates api keys on repeated set()", () => { + const storage = inMemory({}); + storage.set("anthropic", makeKey("sk-1")); + storage.set("anthropic", makeKey("sk-2")); + const creds = storage.getCredentialsForProvider("anthropic"); + assert.equal(creds.length, 2); + assert.deepEqual( + creds.map((c) => (c.type === "api_key" ? c.key : null)), + ["sk-1", "sk-2"], + ); + }); + + it("deduplicates identical api keys", () => { + const storage = inMemory({}); + storage.set("anthropic", makeKey("sk-1")); + storage.set("anthropic", makeKey("sk-1")); + const creds = storage.getCredentialsForProvider("anthropic"); + assert.equal(creds.length, 1); + }); +}); + +// ─── backoff / markUsageLimitReached ───────────────────────────────────────── + +describe("AuthStorage — rate-limit backoff", () => { + it("returns true when a backed-off credential has an alternate", async () => { + const storage = inMemory({ + anthropic: [makeKey("sk-1"), makeKey("sk-2")], + }); + + // Use sk-1 via round-robin (first call, index 0) + await storage.getApiKey("anthropic"); + + // Mark it as rate-limited; sk-2 should still be available + const hasAlternate = storage.markUsageLimitReached("anthropic"); + assert.equal(hasAlternate, true); + }); + + it("returns false when all credentials are backed off", async () => { + const storage = inMemory({ + anthropic: [makeKey("sk-1"), makeKey("sk-2")], + }); + + // Back off both keys + await storage.getApiKey("anthropic"); // uses index 0 + storage.markUsageLimitReached("anthropic"); // backs off index 0 + await storage.getApiKey("anthropic"); // uses index 1 + const hasAlternate = storage.markUsageLimitReached("anthropic"); // backs off index 1 + assert.equal(hasAlternate, false); + }); + + it("backed-off credential is skipped; next available key is returned", async () => { + const storage = inMemory({ + anthropic: [makeKey("sk-1"), makeKey("sk-2")], + }); + + // First call → sk-1 (round-robin index 0) + const first = await storage.getApiKey("anthropic"); + assert.equal(first, "sk-1"); + + // Back off sk-1 + storage.markUsageLimitReached("anthropic"); + + // Next call should skip backed-off sk-1 and return sk-2 + const second = await storage.getApiKey("anthropic"); + assert.equal(second, "sk-2"); + }); + + it("single credential: markUsageLimitReached returns false", async () => { + const storage = inMemory({ anthropic: makeKey("sk-only") }); + await storage.getApiKey("anthropic"); + const hasAlternate = storage.markUsageLimitReached("anthropic"); + assert.equal(hasAlternate, false); + }); + + it("session-sticky: marks the correct credential as backed off", async () => { + const storage = inMemory({ + anthropic: [makeKey("sk-1"), makeKey("sk-2")], + }); + + const sessionId = "sess-xyz"; + const chosen = await storage.getApiKey("anthropic", sessionId); + assert.ok(chosen); + + // Back off the chosen credential for this session + const hasAlternate = storage.markUsageLimitReached("anthropic", sessionId); + assert.equal(hasAlternate, true); + + // Next call with same session should return the other key + const next = await storage.getApiKey("anthropic", sessionId); + assert.ok(next); + assert.notEqual(next, chosen); + }); +}); + +// ─── getAll truncation ──────────────────────────────────────────────────────── + +describe("AuthStorage — getAll()", () => { + it("returns first credential only for providers with multiple keys", () => { + const storage = inMemory({ + anthropic: [makeKey("sk-1"), makeKey("sk-2")], + openai: makeKey("sk-openai"), + }); + const all = storage.getAll(); + assert.ok(all["anthropic"]?.type === "api_key"); + assert.equal((all["anthropic"] as any).key, "sk-1"); + assert.equal((all["openai"] as any).key, "sk-openai"); + }); +}); diff --git a/packages/pi-coding-agent/src/core/auth-storage.ts b/packages/pi-coding-agent/src/core/auth-storage.ts index d4578a53c..30beef551 100644 --- a/packages/pi-coding-agent/src/core/auth-storage.ts +++ b/packages/pi-coding-agent/src/core/auth-storage.ts @@ -431,6 +431,10 @@ export class AuthStorage { * Get all credentials (for passing to getOAuthApiKey). * Returns normalized format where each provider has a single credential * (the first one) for backward compatibility with OAuth refresh. + * + * NOTE: For providers with multiple API keys, only the first credential is + * returned. This is intentional — callers use this for OAuth refresh only, + * which is always single-credential. Do not use for API key enumeration. */ getAll(): Record { const result: Record = {}; @@ -542,7 +546,11 @@ export class AuthStorage { usedIndex = hashString(sessionId) % credentials.length; } else { // Round-robin was already incremented in getApiKey, so the last-used - // index is (current - 1) + // index is (current - 1). Note: in a concurrent scenario where another + // getApiKey call fires between the original request and this backoff call, + // we may back off the wrong credential index. This is acceptable because: + // (a) pi runs single-threaded event loop, (b) backing off the wrong key + // is safe — it self-heals when the backoff expires. const current = this.providerRoundRobinIndex.get(provider) ?? 0; usedIndex = ((current - 1) % credentials.length + credentials.length) % credentials.length; }