Add additive custom provider models

This commit is contained in:
Mohamed Boudra
2026-04-22 22:16:54 +07:00
parent 4d5889ff38
commit b133285f6d
6 changed files with 410 additions and 61 deletions
@@ -67,6 +67,7 @@ export const ProviderOverrideSchema = z
command: z.array(z.string().min(1)).min(1).optional(),
env: z.record(z.string()).optional(),
models: z.array(ProviderProfileModelSchema).optional(),
additionalModels: z.array(ProviderProfileModelSchema).optional(),
disallowedTools: z.array(z.string()).optional(),
enabled: z.boolean().optional(),
order: z.number().optional(),
@@ -666,6 +666,193 @@ describe("buildProviderRegistry", () => {
]);
});
test("additional models append to runtime models", async () => {
mockState.runtimeModels.set("claude", [
{
provider: "claude",
id: "runtime-pro",
label: "Runtime Pro",
},
]);
const registry = buildProviderRegistry(logger, {
providerOverrides: {
claude: {
additionalModels: [
{
id: "profile-fast",
label: "Profile Fast",
},
],
},
},
});
const models = await registry.claude.fetchModels({
cwd: "/tmp/registry-models",
force: false,
});
expect(models).toEqual([
{
provider: "claude",
id: "runtime-pro",
label: "Runtime Pro",
},
{
provider: "claude",
id: "profile-fast",
label: "Profile Fast",
},
]);
});
test("additional models merge onto profile replacement models", async () => {
mockState.runtimeModels.set("claude", [
{
provider: "claude",
id: "runtime-pro",
label: "Runtime Pro",
},
]);
const registry = buildProviderRegistry(logger, {
providerOverrides: {
claude: {
models: [
{
id: "profile-curated",
label: "Profile Curated",
},
],
additionalModels: [
{
id: "profile-extra",
label: "Profile Extra",
},
],
},
},
});
const models = await registry.claude.fetchModels({
cwd: "/tmp/registry-models",
force: false,
});
expect(models.map((model) => model.id)).toEqual(["profile-curated", "profile-extra"]);
});
test("additional models override matching runtime models in place", async () => {
mockState.runtimeModels.set("claude", [
{
provider: "claude",
id: "shared-model",
label: "Runtime Label",
description: "Runtime description",
metadata: {
source: "runtime",
},
},
{
provider: "claude",
id: "runtime-only",
label: "Runtime Only",
},
]);
const registry = buildProviderRegistry(logger, {
providerOverrides: {
claude: {
additionalModels: [
{
id: "shared-model",
label: "Profile Label",
},
],
},
},
});
const models = await registry.claude.fetchModels({
cwd: "/tmp/registry-models",
force: false,
});
expect(models).toEqual([
{
provider: "claude",
id: "shared-model",
label: "Profile Label",
description: "Runtime description",
metadata: {
source: "runtime",
},
},
{
provider: "claude",
id: "runtime-only",
label: "Runtime Only",
},
]);
});
test("additional model default overrides runtime default", async () => {
mockState.runtimeModels.set("claude", [
{
provider: "claude",
id: "runtime-default",
label: "Runtime Default",
isDefault: true,
},
{
provider: "claude",
id: "runtime-other",
label: "Runtime Other",
},
]);
const registry = buildProviderRegistry(logger, {
providerOverrides: {
claude: {
additionalModels: [
{
id: "profile-default",
label: "Profile Default",
isDefault: true,
},
],
},
},
});
const models = await registry.claude.fetchModels({
cwd: "/tmp/registry-models",
force: false,
});
expect(models).toEqual([
{
provider: "claude",
id: "runtime-default",
label: "Runtime Default",
isDefault: false,
},
{
provider: "claude",
id: "runtime-other",
label: "Runtime Other",
isDefault: false,
},
{
provider: "claude",
id: "profile-default",
label: "Profile Default",
isDefault: true,
},
]);
});
test("no profile models — runtime models returned as-is", async () => {
mockState.runtimeModels.set("claude", [
{
@@ -63,6 +63,7 @@ type ResolvedProvider = {
definition: AgentProviderDefinition;
runtimeSettings?: ProviderRuntimeSettings;
profileModels: ProviderProfileModel[];
additionalModels: ProviderProfileModel[];
enabled: boolean;
createBaseClient: (logger: Logger) => AgentClient;
};
@@ -224,16 +225,59 @@ function mapModel(provider: AgentProvider, model: AgentModelDefinition): AgentMo
function mergeModels(
provider: AgentProvider,
profileModels: ProviderProfileModel[],
additionalModels: ProviderProfileModel[],
runtimeModels: AgentModelDefinition[],
): AgentModelDefinition[] {
if (profileModels.length === 0) {
return runtimeModels.map((model) => mapModel(provider, model));
const baseModels =
profileModels.length === 0
? runtimeModels.map((model) => mapModel(provider, model))
: profileModels.map((model) => ({
...model,
provider,
}));
if (additionalModels.length === 0) {
return baseModels;
}
return profileModels.map((model) => ({
...model,
provider,
}));
const mergedModels = [...baseModels];
let hasAdditionalDefault = false;
for (const model of additionalModels) {
const additionalModel = {
...model,
provider,
};
hasAdditionalDefault ||= additionalModel.isDefault === true;
const existingIndex = mergedModels.findIndex((candidate) => candidate.id === model.id);
if (existingIndex === -1) {
mergedModels.push(additionalModel);
continue;
}
mergedModels[existingIndex] = {
...mergedModels[existingIndex],
...additionalModel,
};
}
if (!hasAdditionalDefault) {
return mergedModels;
}
const additionalDefaultIds = new Set(
additionalModels.filter((model) => model.isDefault === true).map((model) => model.id),
);
return mergedModels.map((model) =>
additionalDefaultIds.has(model.id)
? model
: {
...model,
isDefault: false,
},
);
}
function wrapSessionProvider(provider: AgentProvider, inner: AgentSession): AgentSession {
@@ -330,7 +374,12 @@ function createRegistryEntry(
return inner.provider === provider ? inner : wrapClientProvider(provider, inner);
},
fetchModels: async (options: ListModelsOptions) =>
mergeModels(provider, resolved.profileModels, await modelClient.listModels(options)),
mergeModels(
provider,
resolved.profileModels,
resolved.additionalModels,
await modelClient.listModels(options),
),
fetchModes: async (options: ListModesOptions) => {
const modes = modelClient.listModes
? await modelClient.listModes(options)
@@ -373,6 +422,7 @@ function buildResolvedBuiltinProviders(
definition: applyOverrideToDefinition(definition, override),
runtimeSettings: mergedRuntimeSettings,
profileModels: override?.models ?? [],
additionalModels: override?.additionalModels ?? [],
enabled: override?.enabled !== false,
createBaseClient: (logger) =>
factory(logger, mergedRuntimeSettings, {
@@ -416,6 +466,7 @@ function addDerivedProviders(
),
runtimeSettings: toRuntimeSettings(override),
profileModels: override.models ?? [],
additionalModels: override.additionalModels ?? [],
enabled: override.enabled !== false,
createBaseClient: (logger) =>
new GenericACPAgentClient({
@@ -445,6 +496,7 @@ function addDerivedProviders(
definition: createDerivedDefinition(providerId, baseDefinition, override),
runtimeSettings: mergedRuntimeSettings,
profileModels: override.models ?? [],
additionalModels: override.additionalModels ?? [],
enabled: override.enabled !== false,
createBaseClient: (logger) => baseFactory(logger, mergedRuntimeSettings),
});
@@ -165,6 +165,34 @@ describe("provider overrides (new format)", () => {
]);
});
test("additionalModels array accepted", () => {
const parsed = PersistedConfigSchema.parse({
agents: {
providers: {
zai: {
extends: "claude",
label: "ZAI",
additionalModels: [
{
id: "zai-fast",
label: "ZAI Fast",
isDefault: true,
},
],
},
},
},
});
expect(parsed.agents?.providers?.zai?.additionalModels).toEqual([
{
id: "zai-fast",
label: "ZAI Fast",
isDefault: true,
},
]);
});
test("order field accepted", () => {
const parsed = PersistedConfigSchema.parse({
agents: {