Use new getApiKeyAndHeaders to fetch the API key from the model registry

Author: kirberich

packages/budget-model/package.json

@@ -37,8 +37,8 @@
     "@mariozechner/pi-coding-agent": "*"
   },
   "devDependencies": {
-    "@mariozechner/pi-ai": "^0.57.0",
-    "@mariozechner/pi-coding-agent": "^0.57.0",
+    "@mariozechner/pi-ai": "^0.63.0",
+    "@mariozechner/pi-coding-agent": "^0.63.0",
     "@types/node": "^25.3.5",
     "tsup": "^8.5.1",
     "typescript": "^5.9.3"

packages/budget-model/src/index.ts

@@ -121,6 +121,11 @@ export async function findBudgetModel(ctx: ExtensionContext, options?: BudgetMod
 	return findSameProvider(ctx, activeModel, opts.costRatio, opts.majorVersions);
 }
 
+async function getApiKey(ctx: ExtensionContext, model: Model<Api>): Promise<string | undefined> {
+	const auth = await ctx.modelRegistry.getApiKeyAndHeaders(model);
+	return auth.ok ? auth.apiKey : undefined;
+}
+
 // --- Strategies ---
 
 async function findSameProvider(
@@ -162,7 +167,7 @@ async function findSameProvider(
 
 	for (const candidate of candidates) {
 		if (candidate.cost.input >= activeModel.cost.input * costRatio) break;
-		const apiKey = await ctx.modelRegistry.getApiKey(candidate);
+		const apiKey = await getApiKey(ctx, candidate);
 		if (apiKey) {
 			return { model: candidate, apiKey };
 		}
@@ -205,7 +210,7 @@ async function findAnyProvider(
 
 	for (const model of allCandidates) {
 		if (model.cost.input >= activeModel.cost.input * costRatio) break;
-		const apiKey = await ctx.modelRegistry.getApiKey(model);
+		const apiKey = await getApiKey(ctx, model);
 		if (apiKey) {
 			return { model, apiKey };
 		}
@@ -226,7 +231,7 @@ async function resolveModelOverride(ctx: ExtensionContext, override: string): Pr
 		throw new NoBudgetModelError(`model override "${override}" not found in registry`);
 	}
 
-	const apiKey = await ctx.modelRegistry.getApiKey(model);
+	const apiKey = await getApiKey(ctx, model);
 	if (!apiKey) {
 		throw new NoBudgetModelError(`no API key for model override "${override}"`);
 	}
@@ -276,7 +281,7 @@ export function findCheapestInMajorVersions(models: Model<Api>[], majorVersions:
 
 /** Build a ModelCandidate from a Model (checks API key availability). */
 async function toCandidate(ctx: ExtensionContext, model: Model<Api>, provider: string): Promise<ModelCandidate> {
-	const apiKey = await ctx.modelRegistry.getApiKey(model);
+	const apiKey = await getApiKey(ctx, model);
 	return {
 		provider,
 		modelId: model.id,

yarn.lock

@@ -1381,6 +1381,15 @@ __metadata:
   languageName: node
   linkType: hard
 
+"@mariozechner/pi-agent-core@npm:^0.63.2":
+  version: 0.63.2
+  resolution: "@mariozechner/pi-agent-core@npm:0.63.2"
+  dependencies:
+    "@mariozechner/pi-ai": "npm:^0.63.2"
+  checksum: 10c0/6a8644e6e54b69e83dd109630d09e4af7a1606fa5026fb207a7d2022934e36682550672c74ed5f48d2b28007d6dd8adb3cac9337e6c0c9059a15cfbb49093808
+  languageName: node
+  linkType: hard
+
 "@mariozechner/pi-ai@npm:*, @mariozechner/pi-ai@npm:^0.57.0":
   version: 0.57.0
   resolution: "@mariozechner/pi-ai@npm:0.57.0"
@@ -1404,6 +1413,29 @@ __metadata:
   languageName: node
   linkType: hard
 
+"@mariozechner/pi-ai@npm:^0.63.0, @mariozechner/pi-ai@npm:^0.63.2":
+  version: 0.63.2
+  resolution: "@mariozechner/pi-ai@npm:0.63.2"
+  dependencies:
+    "@anthropic-ai/sdk": "npm:^0.73.0"
+    "@aws-sdk/client-bedrock-runtime": "npm:^3.983.0"
+    "@google/genai": "npm:^1.40.0"
+    "@mistralai/mistralai": "npm:1.14.1"
+    "@sinclair/typebox": "npm:^0.34.41"
+    ajv: "npm:^8.17.1"
+    ajv-formats: "npm:^3.0.1"
+    chalk: "npm:^5.6.2"
+    openai: "npm:6.26.0"
+    partial-json: "npm:^0.1.7"
+    proxy-agent: "npm:^6.5.0"
+    undici: "npm:^7.19.1"
+    zod-to-json-schema: "npm:^3.24.6"
+  bin:
+    pi-ai: dist/cli.js
+  checksum: 10c0/b1c881858ca44fee1a1749056b21fe90d6902d56dd89c57d470773055c435dfa68af1df37ffe565ca0ffa6fb8cacf91e25c2303eb7d52c60d0fe24b4943edcf9
+  languageName: node
+  linkType: hard
+
 "@mariozechner/pi-coding-agent@npm:*, @mariozechner/pi-coding-agent@npm:^0.57.0":
   version: 0.57.0
   resolution: "@mariozechner/pi-coding-agent@npm:0.57.0"
@@ -1437,6 +1469,40 @@ __metadata:
   languageName: node
   linkType: hard
 
+"@mariozechner/pi-coding-agent@npm:^0.63.0":
+  version: 0.63.2
+  resolution: "@mariozechner/pi-coding-agent@npm:0.63.2"
+  dependencies:
+    "@mariozechner/clipboard": "npm:^0.3.2"
+    "@mariozechner/jiti": "npm:^2.6.2"
+    "@mariozechner/pi-agent-core": "npm:^0.63.2"
+    "@mariozechner/pi-ai": "npm:^0.63.2"
+    "@mariozechner/pi-tui": "npm:^0.63.2"
+    "@silvia-odwyer/photon-node": "npm:^0.3.4"
+    ajv: "npm:^8.17.1"
+    chalk: "npm:^5.5.0"
+    cli-highlight: "npm:^2.1.11"
+    diff: "npm:^8.0.2"
+    extract-zip: "npm:^2.0.1"
+    file-type: "npm:^21.1.1"
+    glob: "npm:^13.0.1"
+    hosted-git-info: "npm:^9.0.2"
+    ignore: "npm:^7.0.5"
+    marked: "npm:^15.0.12"
+    minimatch: "npm:^10.2.3"
+    proper-lockfile: "npm:^4.1.2"
+    strip-ansi: "npm:^7.1.0"
+    undici: "npm:^7.19.1"
+    yaml: "npm:^2.8.2"
+  dependenciesMeta:
+    "@mariozechner/clipboard":
+      optional: true
+  bin:
+    pi: dist/cli.js
+  checksum: 10c0/cd8ffa7e2389979420fc45222a3248213dc46430083d608512a9f833aac48a20be62cb9bbdad87d7f26f617b9bffef0aff70d53baf766600d5901b322e27530e
+  languageName: node
+  linkType: hard
+
 "@mariozechner/pi-tui@npm:^0.57.0":
   version: 0.57.0
   resolution: "@mariozechner/pi-tui@npm:0.57.0"
@@ -1454,6 +1520,23 @@ __metadata:
   languageName: node
   linkType: hard
 
+"@mariozechner/pi-tui@npm:^0.63.2":
+  version: 0.63.2
+  resolution: "@mariozechner/pi-tui@npm:0.63.2"
+  dependencies:
+    "@types/mime-types": "npm:^2.1.4"
+    chalk: "npm:^5.5.0"
+    get-east-asian-width: "npm:^1.3.0"
+    koffi: "npm:^2.9.0"
+    marked: "npm:^15.0.12"
+    mime-types: "npm:^3.0.1"
+  dependenciesMeta:
+    koffi:
+      optional: true
+  checksum: 10c0/576a865593e85642793bde9924559a2e3a511952d0e5bac0e5ed417f5f6e9229bc9f03582b786d1ef475bcb721dab64bcab88e1fb92fc64bcf1cd4dee66b4358
+  languageName: node
+  linkType: hard
+
 "@mistralai/mistralai@npm:1.14.1":
   version: 1.14.1
   resolution: "@mistralai/mistralai@npm:1.14.1"
@@ -4530,8 +4613,8 @@ __metadata:
   version: 0.0.0-use.local
   resolution: "pi-budget-model@workspace:packages/budget-model"
   dependencies:
-    "@mariozechner/pi-ai": "npm:^0.57.0"
-    "@mariozechner/pi-coding-agent": "npm:^0.57.0"
+    "@mariozechner/pi-ai": "npm:^0.63.0"
+    "@mariozechner/pi-coding-agent": "npm:^0.63.0"
     "@types/node": "npm:^25.3.5"
     tsup: "npm:^8.5.1"
     typescript: "npm:^5.9.3"