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.
This commit is contained in:
Aliou Diallo
2026-04-24 20:15:35 +02:00
committed by GitHub
parent 9b103e5e41
commit 39a9784d2a
2 changed files with 118 additions and 4 deletions

View File

@@ -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<string, unknown>)[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.`);

View File

@@ -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", () => {