Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
kirberich wants to merge 1 commit into from
Closed

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions packages/budget-model/package.json
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down
13 changes: 9 additions & 4 deletions packages/budget-model/src/index.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}
Comment on lines +124 to +127
Copy link
Copy Markdown

@greptile-apps greptile-apps Bot May 2, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Headers from getApiKeyAndHeaders are silently discarded

getApiKeyAndHeaders was introduced precisely to carry extra auth headers alongside the API key (e.g. for AWS Bedrock or other header-based auth providers). This wrapper extracts only auth.apiKey and throws away auth.headers, so any provider that requires those headers for authentication will still fail — the returned BudgetModel.apiKey alone won't be enough to make a successful API call for those providers.

Consider either returning the full auth object from getApiKey, or updating the BudgetModel interface to carry headers too:

export interface BudgetModel {
    model: Model<Api>;
    apiKey: string;
    headers?: Record<string, string>;
}

and propagating auth.headers wherever BudgetModel is constructed.

Prompt To Fix With AI 
This is a comment left during a code review.
Path: packages/budget-model/src/index.ts
Line: 124-127

Comment:
**Headers from `getApiKeyAndHeaders` are silently discarded**

`getApiKeyAndHeaders` was introduced precisely to carry extra auth headers alongside the API key (e.g. for AWS Bedrock or other header-based auth providers). This wrapper extracts only `auth.apiKey` and throws away `auth.headers`, so any provider that requires those headers for authentication will still fail — the returned `BudgetModel.apiKey` alone won't be enough to make a successful API call for those providers.

Consider either returning the full auth object from `getApiKey`, or updating the `BudgetModel` interface to carry `headers` too:

```ts
export interface BudgetModel {
    model: Model<Api>;
    apiKey: string;
    headers?: Record<string, string>;
}
```

and propagating `auth.headers` wherever `BudgetModel` is constructed.

How can I resolve this? If you propose a fix, please make it concise.


// --- Strategies ---

async function findSameProvider(
Expand Down Expand Up @@ -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 };
}
Expand Down Expand Up @@ -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 };
}
Expand All @@ -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}"`);
}
Expand Down Expand Up @@ -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,
Expand Down
87 changes: 85 additions & 2 deletions yarn.lock
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand All @@ -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"
Expand Down Expand Up @@ -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"
Expand All @@ -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"
Expand Down Expand Up @@ -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"
Expand Down
Loading