gsd-pi 2.10.2 → 2.10.5

package/node_modules/@gsd/pi-coding-agent/src/core/agent-session.ts

@@ -869,7 +869,7 @@ export class AgentSession {
 		}
 
 		// Validate API key
-		const apiKey = await this._modelRegistry.getApiKey(this.model);
+		const apiKey = await this._modelRegistry.getApiKey(this.model, this.sessionId);
 		if (!apiKey) {
 			const isOAuth = this._modelRegistry.isUsingOAuth(this.model);
 			if (isOAuth) {
@@ -1309,7 +1309,7 @@ export class AgentSession {
 	 * @throws Error if no API key available for the model
 	 */
 	async setModel(model: Model<any>, options?: { persist?: boolean }): Promise<void> {
-		const apiKey = await this._modelRegistry.getApiKey(model);
+		const apiKey = await this._modelRegistry.getApiKey(model, this.sessionId);
 		if (!apiKey) {
 			throw new Error(`No API key for ${model.provider}/${model.id}`);
 		}
@@ -1351,7 +1351,7 @@ export class AgentSession {
 			if (apiKeysByProvider.has(provider)) {
 				apiKey = apiKeysByProvider.get(provider);
 			} else {
-				apiKey = await this._modelRegistry.getApiKeyForProvider(provider);
+				apiKey = await this._modelRegistry.getApiKeyForProvider(provider, this.sessionId);
 				apiKeysByProvider.set(provider, apiKey);
 			}
 
@@ -1406,7 +1406,7 @@ export class AgentSession {
 		const nextIndex = direction === "forward" ? (currentIndex + 1) % len : (currentIndex - 1 + len) % len;
 		const nextModel = availableModels[nextIndex];
 
-		const apiKey = await this._modelRegistry.getApiKey(nextModel);
+		const apiKey = await this._modelRegistry.getApiKey(nextModel, this.sessionId);
 		if (!apiKey) {
 			throw new Error(`No API key for ${nextModel.provider}/${nextModel.id}`);
 		}
@@ -1560,7 +1560,7 @@ export class AgentSession {
 				throw new Error("No model selected");
 			}
 
-			const apiKey = await this._modelRegistry.getApiKey(this.model);
+			const apiKey = await this._modelRegistry.getApiKey(this.model, this.sessionId);
 			if (!apiKey) {
 				throw new Error(`No API key for ${this.model.provider}`);
 			}
@@ -1780,7 +1780,7 @@ export class AgentSession {
 				return;
 			}
 
-			const apiKey = await this._modelRegistry.getApiKey(this.model);
+			const apiKey = await this._modelRegistry.getApiKey(this.model, this.sessionId);
 			if (!apiKey) {
 				this._emit({ type: "auto_compaction_end", result: undefined, aborted: false, willRetry: false });
 				return;
@@ -2082,7 +2082,7 @@ export class AgentSession {
 				refreshTools: () => this._refreshToolRegistry(),
 				getCommands,
 				setModel: async (model, options) => {
-					const key = await this.modelRegistry.getApiKey(model);
+					const key = await this.modelRegistry.getApiKey(model, this.sessionId);
 					if (!key) return false;
 					await this.setModel(model, options);
 					return true;
@@ -2188,7 +2188,14 @@ export class AgentSession {
 			? this._baseToolsOverride
 			: createAllTools(this._cwd, {
 					read: { autoResizeImages },
-					bash: { commandPrefix: shellCommandPrefix },
+					bash: {
+						commandPrefix: shellCommandPrefix,
+						interceptor: {
+							enabled: this.settingsManager.getBashInterceptorEnabled(),
+							rules: this.settingsManager.getBashInterceptorRules(),
+						},
+						availableToolNames: () => this.getActiveToolNames(),
+					},
 				});
 
 		this._baseToolRegistry = new Map(Object.entries(baseTools).map(([name, tool]) => [name, tool as AgentTool]));
@@ -2275,8 +2282,21 @@ export class AgentSession {
 		);
 	}
 
+	/**
+	 * Classify an error message into a usage-limit error type for credential backoff.
+	 */
+	private _classifyErrorType(errorMessage: string): import("./auth-storage.js").UsageLimitErrorType {
+		const err = errorMessage.toLowerCase();
+		if (/quota|billing|exceeded.*limit|usage.*limit/i.test(err)) return "quota_exhausted";
+		if (/rate.?limit|too many requests|429/i.test(err)) return "rate_limit";
+		if (/500|502|503|504|server.?error|internal.?error|service.?unavailable/i.test(err)) return "server_error";
+		return "unknown";
+	}
+
 	/**
 	 * Handle retryable errors with exponential backoff.
+	 * When multiple credentials are available, marks the failing credential
+	 * as backed off and retries immediately with the next one.
 	 * @returns true if retry was initiated, false if max retries exceeded or disabled
 	 */
 	private async _handleRetryableError(message: AssistantMessage): Promise<boolean> {
@@ -2294,6 +2314,42 @@ export class AgentSession {
 			});
 		}
 
+		// Try credential fallback before counting against retry budget.
+		// If another credential is available, switch to it and retry immediately.
+		if (this.model && message.errorMessage) {
+			const errorType = this._classifyErrorType(message.errorMessage);
+			const hasAlternate = this._modelRegistry.authStorage.markUsageLimitReached(
+				this.model.provider,
+				this.sessionId,
+				{ errorType },
+			);
+
+			if (hasAlternate) {
+				// Remove error message from agent state
+				const messages = this.agent.state.messages;
+				if (messages.length > 0 && messages[messages.length - 1].role === "assistant") {
+					this.agent.replaceMessages(messages.slice(0, -1));
+				}
+
+				this._emit({
+					type: "auto_retry_start",
+					attempt: this._retryAttempt + 1,
+					maxAttempts: settings.maxRetries,
+					delayMs: 0,
+					errorMessage: `${message.errorMessage} (switching credential)`,
+				});
+
+				// Retry immediately with the next credential - don't increment _retryAttempt
+				setTimeout(() => {
+					this.agent.continue().catch(() => {
+						// Retry failed - will be caught by next agent_end
+					});
+				}, 0);
+
+				return true;
+			}
+		}
+
 		this._retryAttempt++;
 
 		if (this._retryAttempt > settings.maxRetries) {
@@ -2750,7 +2806,7 @@ export class AgentSession {
 		let summaryDetails: unknown;
 		if (options.summarize && entriesToSummarize.length > 0 && !extensionSummary) {
 			const model = this.model!;
-			const apiKey = await this._modelRegistry.getApiKey(model);
+			const apiKey = await this._modelRegistry.getApiKey(model, this.sessionId);
 			if (!apiKey) {
 				throw new Error(`No API key for ${model.provider}`);
 			}

package/node_modules/@gsd/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<string, unknown> = {}) {
+	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<string>();
+		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<string>();
+		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");
+	});
+});