diff --git a/packages/ai-config/src/__tests__/resolve-models.test.ts b/packages/ai-config/src/__tests__/resolve-models.test.ts index 0419046..ee32c8e 100644 --- a/packages/ai-config/src/__tests__/resolve-models.test.ts +++ b/packages/ai-config/src/__tests__/resolve-models.test.ts @@ -4,7 +4,12 @@ import { describe, it, expect } from "vitest"; -import { resolveModels } from "../resolve-models.js"; +import { + isModelDiscoveryEnabled, + isModelIdAllowedByPolicy, + resolveModelIds, + resolveModels, +} from "../index.js"; import type { ModelInfoLike, ModelsBlock, ResolvedConnection } from "../types.js"; function makeModel(id: string, overrides?: Partial): ModelInfoLike { @@ -399,3 +404,57 @@ describe("resolveModels", () => { expect(result[1].id).toBe("custom-1"); }); }); + +describe("model policy helpers", () => { + it("isModelDiscoveryEnabled treats discovery off as the only disabled state", () => { + expect([ + isModelDiscoveryEnabled(undefined), + isModelDiscoveryEnabled({}), + isModelDiscoveryEnabled({ discovery: "auto" }), + isModelDiscoveryEnabled({ discovery: "off" }), + ]).toEqual([true, true, true, false]); + }); + + it("resolveModelIds applies the complete resolver pipeline", () => { + const block: ModelsBlock = { + custom: [ + { + id: "custom-1", + name: "Custom 1", + maxContextLength: 50000, + supportsTools: true, + supportsImages: false, + supportsToolResultImages: false, + supportsWebSearch: false, + }, + ], + allow: ["model-a", "model-b", "custom-1"], + deny: ["model-b"], + }; + + expect(resolveModelIds(block, ["model-a", "model-b", "model-c"])).toEqual([ + "model-a", + "custom-1", + ]); + }); + + it("isModelIdAllowedByPolicy checks the requested id, not unrelated custom declarations", () => { + const block: ModelsBlock = { + custom: [ + { + id: "custom-1", + name: "Custom 1", + maxContextLength: 50000, + supportsTools: true, + supportsImages: false, + supportsToolResultImages: false, + supportsWebSearch: false, + }, + ], + allow: ["custom-1"], + }; + + expect(isModelIdAllowedByPolicy(block, "model-a")).toBe(false); + expect(isModelIdAllowedByPolicy(block, "custom-1")).toBe(true); + }); +}); diff --git a/packages/ai-config/src/index.ts b/packages/ai-config/src/index.ts index be39d6b..74aea17 100644 --- a/packages/ai-config/src/index.ts +++ b/packages/ai-config/src/index.ts @@ -153,7 +153,12 @@ export { } from "./defaults.js"; // --- Resolution helpers (public) ------------------------------------------- -export { resolveModels } from "./resolve-models.js"; +export { + isModelDiscoveryEnabled, + isModelIdAllowedByPolicy, + resolveModelIds, + resolveModels, +} from "./resolve-models.js"; // --- Bare-host base URL correction ------------------------------------------ export { diff --git a/packages/ai-config/src/resolve-models.ts b/packages/ai-config/src/resolve-models.ts index c7cb18b..1ad1aa3 100644 --- a/packages/ai-config/src/resolve-models.ts +++ b/packages/ai-config/src/resolve-models.ts @@ -149,6 +149,36 @@ export function resolveModels( return result.map((e) => attachRouting(e.model, e.userRouting, providerConnection)); } +/** + * Whether a provider should query runtime discovery for models. + */ +export function isModelDiscoveryEnabled(modelsBlock: ModelsBlock | undefined): boolean { + return modelsBlock?.discovery !== "off"; +} + +/** + * Resolve a provider's final model ids from discovered ids plus custom models. + * `discovery: "off"` drops discovered ids only; custom models still participate + * before overrides, allow, and deny. + */ +export function resolveModelIds( + modelsBlock: ModelsBlock | undefined, + discoveredModelIds: readonly string[], +): string[] { + return resolveModels( + modelsBlock, + discoveredModelIds.map((id) => modelIdToModelInfo(id)), + ).map((model) => model.id); +} + +/** Checks whether a raw provider model id survives complete model policy resolution. */ +export function isModelIdAllowedByPolicy( + modelsBlock: ModelsBlock | undefined, + modelId: string, +): boolean { + return resolveModelIds(modelsBlock, [modelId]).includes(modelId); +} + /** * Resolve protocol and endpoint for a model, applying the correct precedence: * @@ -190,6 +220,18 @@ function attachRouting( }; } +function modelIdToModelInfo(id: string): ModelInfoLike { + return { + id, + name: id, + maxContextLength: 0, + supportsTools: false, + supportsImages: false, + supportsToolResultImages: false, + supportsWebSearch: false, + }; +} + /** Convert a custom model definition to a ModelInfoLike. */ function customModelToModelInfo(custom: CustomModel): ModelInfoLike { return {