From 8936c5d136a9c7a211c7c5a5e70d0c92e355f1e7 Mon Sep 17 00:00:00 2001
From: Josh Palmer <joshp123@users.noreply.github.com>
Date: Tue, 13 Jan 2026 16:29:27 +0100
Subject: [PATCH] fix(coding-agent): resolve api keys by provider

What:
- resolve API keys using provider argument
- add regression test for model-switch key resolution
- update coding-agent changelog
- document provider lookup in SDK

Why:
- prevent model-switch 401s from provider mismatch

Tests:
- npm run check
- npm run test -w @mariozechner/pi-coding-agent -- sdk-api-key.test.ts
---
 packages/coding-agent/CHANGELOG.md            |  4 ++
 packages/coding-agent/src/core/sdk.ts         | 12 ++--
 .../coding-agent/test/sdk-api-key.test.ts     | 56 +++++++++++++++++++
 3 files changed, 67 insertions(+), 5 deletions(-)
 create mode 100644 packages/coding-agent/test/sdk-api-key.test.ts

diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md
index 545939af..2f930548 100644
--- a/packages/coding-agent/CHANGELOG.md
+++ b/packages/coding-agent/CHANGELOG.md
@@ -11,6 +11,10 @@
 - Extension example: `summarize.ts` for summarizing conversations using custom UI and an external model
 - Vercel AI Gateway provider support: set `AI_GATEWAY_API_KEY` and use `--provider vercel-ai-gateway` ([#689](https://github.com/badlogic/pi-mono/pull/689) by [@timolins](https://github.com/timolins))
 
+### Fixed
+
+- Fix API key resolution after model switches by using provider argument ([#691](https://github.com/badlogic/pi-mono/pull/691) by [@joshp123](https://github.com/joshp123))
+
 ## [0.45.3] - 2026-01-13
 
 ## [0.45.2] - 2026-01-13
diff --git a/packages/coding-agent/src/core/sdk.ts b/packages/coding-agent/src/core/sdk.ts
index dd581368..87f550f4 100644
--- a/packages/coding-agent/src/core/sdk.ts
+++ b/packages/coding-agent/src/core/sdk.ts
@@ -628,14 +628,16 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {}
 		steeringMode: settingsManager.getSteeringMode(),
 		followUpMode: settingsManager.getFollowUpMode(),
 		thinkingBudgets: settingsManager.getThinkingBudgets(),
-		getApiKey: async () => {
-			const currentModel = agent.state.model;
-			if (!currentModel) {
+		getApiKey: async (provider) => {
+			// Use the provider argument from the in-flight request;
+			// agent.state.model may already be switched mid-turn.
+			const resolvedProvider = provider || agent.state.model?.provider;
+			if (!resolvedProvider) {
 				throw new Error("No model selected");
 			}
-			const key = await modelRegistry.getApiKey(currentModel);
+			const key = await modelRegistry.authStorage.getApiKey(resolvedProvider);
 			if (!key) {
-				throw new Error(`No API key found for provider "${currentModel.provider}"`);
+				throw new Error(`No API key found for provider "${resolvedProvider}"`);
 			}
 			return key;
 		},
diff --git a/packages/coding-agent/test/sdk-api-key.test.ts b/packages/coding-agent/test/sdk-api-key.test.ts
new file mode 100644
index 00000000..61875c0d
--- /dev/null
+++ b/packages/coding-agent/test/sdk-api-key.test.ts
@@ -0,0 +1,56 @@
+import { mkdirSync, rmSync } from "node:fs";
+import { tmpdir } from "node:os";
+import { join } from "node:path";
+import { getModel } from "@mariozechner/pi-ai";
+import { afterEach, beforeEach, describe, expect, it } from "vitest";
+import { AuthStorage } from "../src/core/auth-storage.js";
+import { ModelRegistry } from "../src/core/model-registry.js";
+import { createAgentSession } from "../src/core/sdk.js";
+import { SessionManager } from "../src/core/session-manager.js";
+import { SettingsManager } from "../src/core/settings-manager.js";
+
+describe("createAgentSession getApiKey", () => {
+	let tempDir: string;
+	let agentDir: string;
+	let projectDir: string;
+
+	beforeEach(() => {
+		tempDir = join(tmpdir(), `pi-test-sdk-${Date.now()}-${Math.random().toString(36).slice(2)}`);
+		agentDir = join(tempDir, "agent");
+		projectDir = join(tempDir, "project");
+		mkdirSync(agentDir, { recursive: true });
+		mkdirSync(join(projectDir, ".pi"), { recursive: true });
+	});
+
+	afterEach(() => {
+		rmSync(tempDir, { recursive: true, force: true });
+	});
+
+	it("uses the provider argument after model switches", async () => {
+		const authStorage = new AuthStorage(join(agentDir, "auth.json"));
+		authStorage.set("anthropic", { type: "api_key", key: "anthropic-key" });
+		authStorage.set("openai-codex", { type: "api_key", key: "codex-key" });
+
+		const modelRegistry = new ModelRegistry(authStorage, join(agentDir, "models.json"));
+		const settingsManager = SettingsManager.create(projectDir, agentDir);
+		const sessionManager = SessionManager.inMemory(projectDir);
+
+		const anthropicModel = getModel("anthropic", "claude-opus-4-5");
+		const codexModel = getModel("openai-codex", "gpt-5.2-codex");
+
+		const { session } = await createAgentSession({
+			cwd: projectDir,
+			agentDir,
+			authStorage,
+			modelRegistry,
+			settingsManager,
+			sessionManager,
+			model: anthropicModel,
+		});
+
+		await session.setModel(codexModel);
+
+		const key = await session.agent.getApiKey?.("anthropic");
+		expect(key).toBe("anthropic-key");
+	});
+});