From 39a9784d2a4fba33e47bf977b4ff9c312daceee2 Mon Sep 17 00:00:00 2001 From: Aliou Diallo Date: Fri, 24 Apr 2026 20:15:35 +0200 Subject: [PATCH] fix(coding-agent): preserve models when re-registering provider with overrides only (#3651) Previously, registerProvider() overwrote the stored config for a provider, so a models registration followed by a baseUrl/headers-only override lost the models. After refresh(), only the override-only config was replayed, causing extension-provided models to disappear. Added upsertRegisteredProvider() that merges defined fields into the existing stored config instead of replacing it. Fields absent from the incoming config are preserved from the stored config. Also adds regression tests for dynamic provider override persistence across refresh. --- .../coding-agent/src/core/model-registry.ts | 21 +++- .../coding-agent/test/model-registry.test.ts | 101 +++++++++++++++++- 2 files changed, 118 insertions(+), 4 deletions(-) diff --git a/packages/coding-agent/src/core/model-registry.ts b/packages/coding-agent/src/core/model-registry.ts index 362d8df5..7153d214 100644 --- a/packages/coding-agent/src/core/model-registry.ts +++ b/packages/coding-agent/src/core/model-registry.ts @@ -758,7 +758,7 @@ export class ModelRegistry { registerProvider(providerName: string, config: ProviderConfigInput): void { this.validateProviderConfig(providerName, config); this.applyProviderConfig(providerName, config); - this.registeredProviders.set(providerName, config); + this.upsertRegisteredProvider(providerName, config); } /** @@ -776,6 +776,25 @@ export class ModelRegistry { this.refresh(); } + /** + * Upsert a provider config into registeredProviders. + * If the provider is already registered, defined values in the incoming config + * override existing ones; undefined values are preserved from the stored config. + * If the provider is not registered, the incoming config is stored as-is. + */ + private upsertRegisteredProvider(providerName: string, config: ProviderConfigInput): void { + const existing = this.registeredProviders.get(providerName); + if (!existing) { + this.registeredProviders.set(providerName, config); + return; + } + for (const k of Object.keys(config) as (keyof ProviderConfigInput)[]) { + if (config[k] !== undefined) { + (existing as Record)[k] = config[k]; + } + } + } + private validateProviderConfig(providerName: string, config: ProviderConfigInput): void { if (config.streamSimple && !config.api) { throw new Error(`Provider ${providerName}: "api" is required when registering streamSimple.`); diff --git a/packages/coding-agent/test/model-registry.test.ts b/packages/coding-agent/test/model-registry.test.ts index 9c2a6ec7..5a1260a6 100644 --- a/packages/coding-agent/test/model-registry.test.ts +++ b/packages/coding-agent/test/model-registry.test.ts @@ -6,7 +6,7 @@ import { getApiProvider } from "@mariozechner/pi-ai"; import { getOAuthProvider } from "@mariozechner/pi-ai/oauth"; import { afterEach, beforeEach, describe, expect, test } from "vitest"; import { AuthStorage } from "../src/core/auth-storage.js"; -import { clearApiKeyCache, ModelRegistry } from "../src/core/model-registry.js"; +import { clearApiKeyCache, ModelRegistry, type ProviderConfigInput } from "../src/core/model-registry.js"; describe("ModelRegistry", () => { let tempDir: string; @@ -32,11 +32,11 @@ describe("ModelRegistry", () => { baseUrl: string, models: Array<{ id: string; name?: string }>, api: string = "anthropic-messages", - ) { + ): ProviderConfigInput { return { baseUrl, apiKey: "TEST_KEY", - api, + api: api as Api, models: models.map((m) => ({ id: m.id, name: m.name ?? m.id, @@ -952,6 +952,101 @@ describe("ModelRegistry", () => { } expect(threwCustomOverrideAfterUnregister).toBe(false); }); + + describe("dynamic provider override persistence", () => { + test("baseUrl-only override keeps built-in provider models after refresh", () => { + const registry = ModelRegistry.create(authStorage, modelsJsonPath); + + registry.registerProvider("anthropic", { baseUrl: "https://proxy.test/anthropic" }); + registry.refresh(); + + const anthropicModels = getModelsForProvider(registry, "anthropic"); + expect(anthropicModels.length).toBeGreaterThan(1); + expect(anthropicModels.every((m) => m.baseUrl === "https://proxy.test/anthropic")).toBe(true); + }); + + test("models-only override replaces built-in provider models after refresh", () => { + const registry = ModelRegistry.create(authStorage, modelsJsonPath); + + registry.registerProvider("anthropic", { + ...providerConfig("https://custom.test/anthropic", [{ id: "custom-claude" }], "anthropic-messages"), + baseUrl: "https://custom.test/anthropic", + }); + registry.refresh(); + + expect(getModelsForProvider(registry, "anthropic").map((m) => m.id)).toEqual(["custom-claude"]); + expect(registry.find("anthropic", "custom-claude")?.baseUrl).toBe("https://custom.test/anthropic"); + }); + + test("models plus baseUrl override replaces built-in provider models after refresh", () => { + const registry = ModelRegistry.create(authStorage, modelsJsonPath); + + registry.registerProvider("anthropic", { + ...providerConfig("https://custom.test/anthropic", [{ id: "custom-claude" }], "anthropic-messages"), + baseUrl: "https://custom.test/anthropic", + }); + registry.registerProvider("anthropic", { baseUrl: "https://proxy.test/anthropic" }); + registry.refresh(); + + expect(getModelsForProvider(registry, "anthropic").map((m) => m.id)).toEqual(["custom-claude"]); + expect(registry.find("anthropic", "custom-claude")?.baseUrl).toBe("https://proxy.test/anthropic"); + }); + + test("models-only custom provider registration survives refresh", () => { + const registry = ModelRegistry.create(authStorage, modelsJsonPath); + + registry.registerProvider( + "custom-provider", + providerConfig("https://custom.test/v1", [{ id: "custom-a" }, { id: "custom-b" }], "openai-completions"), + ); + registry.refresh(); + + expect(getModelsForProvider(registry, "custom-provider").map((m) => m.id)).toEqual([ + "custom-a", + "custom-b", + ]); + }); + + test("baseUrl-only override keeps custom provider models after refresh", () => { + const registry = ModelRegistry.create(authStorage, modelsJsonPath); + + registry.registerProvider( + "custom-provider", + providerConfig("https://custom.test/v1", [{ id: "custom-a" }, { id: "custom-b" }], "openai-completions"), + ); + registry.registerProvider("custom-provider", { baseUrl: "https://proxy.test/custom" }); + registry.refresh(); + + expect(getModelsForProvider(registry, "custom-provider").map((m) => m.id)).toEqual([ + "custom-a", + "custom-b", + ]); + expect( + getModelsForProvider(registry, "custom-provider").every( + (m) => m.baseUrl === "https://proxy.test/custom", + ), + ).toBe(true); + }); + + test("headers-only override keeps custom provider models after refresh", async () => { + const registry = ModelRegistry.create(authStorage, modelsJsonPath); + + registry.registerProvider( + "custom-provider", + providerConfig("https://custom.test/v1", [{ id: "custom-a" }, { id: "custom-b" }], "openai-completions"), + ); + registry.registerProvider("custom-provider", { headers: { "x-proxy": "enabled" } }); + registry.refresh(); + + const models = getModelsForProvider(registry, "custom-provider"); + expect(models.map((m) => m.id)).toEqual(["custom-a", "custom-b"]); + expect(models.every((m) => m.baseUrl === "https://custom.test/v1")).toBe(true); + expect(await registry.getApiKeyAndHeaders(models[0])).toMatchObject({ + ok: true, + headers: { "x-proxy": "enabled" }, + }); + }); + }); }); describe("API key resolution", () => {