From c6251a866b1d7d7900c443e4cf615ec148dbe415 Mon Sep 17 00:00:00 2001 From: Mario Zechner Date: Thu, 9 Jul 2026 11:22:35 +0200 Subject: [PATCH] fix(coding-agent): apply modelOverrides to extension providers closes #6367 --- packages/coding-agent/CHANGELOG.md | 1 + packages/coding-agent/docs/models.md | 8 +-- .../coding-agent/src/core/model-registry.ts | 21 ++++++- .../coding-agent/test/model-registry.test.ts | 59 ++++++++++++++++++- 4 files changed, 82 insertions(+), 7 deletions(-) diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index f2b440c7..e0e85154 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -20,6 +20,7 @@ - Fixed the edit tool schema to allow model-invented extra replacement fields instead of rejecting otherwise valid edits ([#6278](https://github.com/earendil-works/pi/issues/6278)). - Fixed new session resets to clear cached label timestamps ([#6354](https://github.com/earendil-works/pi/issues/6354)). - Fixed auto-retry for Bun fetch socket-drop errors reported as `socket connection was closed`, so transient provider disconnects do not end headless runs without retrying ([#6431](https://github.com/earendil-works/pi/issues/6431)). +- Fixed `models.json` `modelOverrides` to apply to extension-registered provider models ([#6367](https://github.com/earendil-works/pi/issues/6367)). ### Removed diff --git a/packages/coding-agent/docs/models.md b/packages/coding-agent/docs/models.md index d17e9928..1183b4d3 100644 --- a/packages/coding-agent/docs/models.md +++ b/packages/coding-agent/docs/models.md @@ -139,7 +139,7 @@ Set `api` at provider level (default for all models) or model level (override pe | `headers` | Custom headers (see value resolution below) | | `authHeader` | Set `true` to add `Authorization: Bearer ` automatically | | `models` | Array of model configurations | -| `modelOverrides` | Per-model overrides for built-in models on this provider | +| `modelOverrides` | Per-model overrides for built-in or extension-registered models on this provider | For providers with `models`, non-built-in provider configs need `baseUrl` and an `api` value at either provider or model level. `apiKey` is not required to load the file: models become available when auth is configured through `/login`/`auth.json`, CLI `--api-key`, or provider `apiKey`. If no auth is configured, the models load but stay unavailable in `/model` and `--list-models`. @@ -293,7 +293,7 @@ Merge semantics: ## Per-model Overrides -Use `modelOverrides` to customize specific built-in models without replacing the provider's full model list. +Use `modelOverrides` to customize built-in models and matching extension-registered models without replacing the provider's full model list. ```json { @@ -314,10 +314,10 @@ Use `modelOverrides` to customize specific built-in models without replacing the } ``` -`modelOverrides` supports these fields per model: `name`, `reasoning`, `input`, `cost` (partial), `contextWindow`, `maxTokens`, `headers`, `compat`. +`modelOverrides` supports these fields per model: `name`, `reasoning`, `thinkingLevelMap`, `input`, `cost` (partial), `contextWindow`, `maxTokens`, `headers`, `compat`. Behavior notes: -- `modelOverrides` are applied to built-in provider models. +- `modelOverrides` are applied to built-in provider models and matching extension-registered provider models. - Unknown model IDs are ignored. - You can combine provider-level `baseUrl`/`headers` with `modelOverrides`. - Overriding `name` changes model matching and secondary detail text only; the footer and primary model lists continue to show the model `id`. diff --git a/packages/coding-agent/src/core/model-registry.ts b/packages/coding-agent/src/core/model-registry.ts index 70f39394..2fa599f1 100644 --- a/packages/coding-agent/src/core/model-registry.ts +++ b/packages/coding-agent/src/core/model-registry.ts @@ -353,6 +353,7 @@ export class ModelRegistry { private models: Model[] = []; private providerRequestConfigs: Map = new Map(); private modelRequestHeaders: Map> = new Map(); + private configModelOverrides: Map> = new Map(); private registeredProviders: Map = new Map(); private loadError: string | undefined = undefined; readonly authStorage: AuthStorage; @@ -406,6 +407,7 @@ export class ModelRegistry { modelOverrides, error, } = this.modelsJsonPath ? this.loadCustomModels(this.modelsJsonPath) : emptyCustomModelsResult(); + this.configModelOverrides = modelOverrides; if (error) { this.loadError = error; @@ -459,6 +461,15 @@ export class ModelRegistry { }); } + private getConfiguredModelOverride(providerName: string, modelId: string): ModelOverride | undefined { + return this.configModelOverrides.get(providerName)?.get(modelId); + } + + private applyConfiguredModelOverride(providerName: string, model: Model): Model { + const modelOverride = this.getConfiguredModelOverride(providerName, model.id); + return modelOverride ? applyModelOverride(model, modelOverride) : model; + } + /** Merge custom models into built-in list by provider+id (custom wins on conflicts). */ private mergeCustomModels(builtInModels: Model[], customModels: Model[]): Model[] { const merged = [...builtInModels]; @@ -921,9 +932,14 @@ export class ModelRegistry { // Parse and add new models for (const modelDef of config.models) { const api = modelDef.api || config.api; - this.storeModelHeaders(providerName, modelDef.id, modelDef.headers); + const modelOverride = this.getConfiguredModelOverride(providerName, modelDef.id); + const headers = + modelDef.headers || modelOverride?.headers + ? { ...modelDef.headers, ...modelOverride?.headers } + : undefined; + this.storeModelHeaders(providerName, modelDef.id, headers); - this.models.push({ + const model = this.applyConfiguredModelOverride(providerName, { id: modelDef.id, name: modelDef.name, api: api as Api, @@ -938,6 +954,7 @@ export class ModelRegistry { headers: undefined, compat: modelDef.compat, } as Model); + this.models.push(model); } // Apply OAuth modifyModels if credentials exist (e.g., to update baseUrl) diff --git a/packages/coding-agent/test/model-registry.test.ts b/packages/coding-agent/test/model-registry.test.ts index fb559e1d..33eb2e00 100644 --- a/packages/coding-agent/test/model-registry.test.ts +++ b/packages/coding-agent/test/model-registry.test.ts @@ -8,7 +8,7 @@ import type { Model, OpenAICompletionsCompat, } from "@earendil-works/pi-ai/compat"; -import { getApiProvider } from "@earendil-works/pi-ai/compat"; +import { getApiProvider, getSupportedThinkingLevels } from "@earendil-works/pi-ai/compat"; import { getOAuthProvider } from "@earendil-works/pi-ai/oauth"; import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; import { AuthStorage } from "../src/core/auth-storage.ts"; @@ -938,6 +938,63 @@ describe("ModelRegistry", () => { expect(registry.getProviderDisplayName("oauth-provider")).toBe("OAuth Provider"); }); + test("modelOverrides apply to dynamically registered provider models", async () => { + writeRawModelsJson({ + "extension-provider": { + modelOverrides: { + "extension-model": { + name: "Overridden Extension Model", + thinkingLevelMap: { + off: null, + minimal: null, + low: null, + medium: null, + xhigh: "max", + }, + headers: { "x-model-override": "enabled" }, + }, + }, + }, + }); + + const registry = ModelRegistry.create(authStorage, modelsJsonPath); + registry.registerProvider("extension-provider", { + baseUrl: "https://provider.test/v1", + apiKey: "test-key", + api: "openai-completions", + models: [ + { + id: "extension-model", + name: "Extension Model", + reasoning: true, + input: ["text"], + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, + contextWindow: 128000, + maxTokens: 4096, + }, + ], + }); + + const model = registry.find("extension-provider", "extension-model"); + expect(model).toBeDefined(); + if (!model) { + throw new Error("extension model was not registered"); + } + expect(model.name).toBe("Overridden Extension Model"); + expect(model.thinkingLevelMap).toEqual({ + off: null, + minimal: null, + low: null, + medium: null, + xhigh: "max", + }); + expect(getSupportedThinkingLevels(model)).toEqual(["high", "xhigh"]); + expect(await registry.getApiKeyAndHeaders(model)).toMatchObject({ + ok: true, + headers: { "x-model-override": "enabled" }, + }); + }); + test("stored API key env propagates to request auth and resolves headers", async () => { authStorage.set("cloudflare-ai-gateway", { type: "api_key",