Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
75 changes: 75 additions & 0 deletions packages/types/src/__tests__/custom-model-info.test.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,75 @@
import { applyCustomModelInfo, customModelInfoSchema, type ModelInfo } from "../model.js"
import { providerIdentifiers, providerSettingsSchemaDiscriminated } from "../index.js"

describe("custom model info", () => {
it("overlays only supported metadata and preserves provider-owned fields", () => {
const model: ModelInfo = {
maxTokens: 4096,
contextWindow: 8192,
supportsImages: false,
supportsPromptCache: false,
inputPrice: 0.1,
outputPrice: 0.2,
description: "Provider metadata",
}

expect(
applyCustomModelInfo(model, {
customModelInfo: {
contextWindow: 128_000,
maxTokens: 16_384,
supportsImages: true,
supportsPromptCache: true,
},
}),
).toEqual({
...model,
contextWindow: 128_000,
maxTokens: 16_384,
supportsImages: true,
supportsPromptCache: true,
})
})

it("does not synthesize model info without a valid context window", () => {
expect(
applyCustomModelInfo(undefined, {
customModelInfo: {
contextWindow: 0,
maxTokens: -1,
supportsImages: true,
},
}),
).toBeUndefined()
})

it("synthesizes safe defaults when only a context window is supplied", () => {
expect(
applyCustomModelInfo(undefined, {
customModelInfo: { contextWindow: 64_000, supportsImages: true },
}),
).toEqual({
maxTokens: undefined,
contextWindow: 64_000,
supportsImages: true,
supportsPromptCache: false,
})
})

it("rejects unsupported pricing fields in the persisted override schema", () => {
expect(customModelInfoSchema.safeParse({ contextWindow: 64_000, inputPrice: 1 }).success).toBe(false)
expect(
providerSettingsSchemaDiscriminated.safeParse({
apiProvider: providerIdentifiers.openrouter,
customModelInfo: { contextWindow: 64_000, outputPrice: 1 },
}).success,
).toBe(false)
})

it("rejects unsafe integer overrides", () => {
const unsafeInteger = Number.MAX_SAFE_INTEGER + 1

expect(customModelInfoSchema.safeParse({ contextWindow: unsafeInteger }).success).toBe(false)
expect(customModelInfoSchema.safeParse({ maxTokens: unsafeInteger }).success).toBe(false)
})
})
80 changes: 80 additions & 0 deletions packages/types/src/model.ts
Original file line number Diff line number Diff line change
Expand Up @@ -183,6 +183,86 @@ export const modelInfoSchema = z.object({

export type ModelInfo = z.infer<typeof modelInfoSchema>

/**
* User-supplied metadata for a model whose discovered metadata is incomplete
* or unavailable. This is intentionally narrower than ModelInfo: prices and
* other accounting fields must remain provider-owned.
*/
const positiveSafeIntegerSchema = z
.number()
.int()
.positive()
.refine(Number.isSafeInteger, { message: "Expected a safe integer" })

export const customModelInfoSchema = z
.object({
maxTokens: positiveSafeIntegerSchema.optional(),
contextWindow: positiveSafeIntegerSchema.optional(),
supportsImages: z.boolean().optional(),
supportsPromptCache: z.boolean().optional(),
})
.strict()
Comment thread
coderabbitai[bot] marked this conversation as resolved.

export type CustomModelInfo = z.infer<typeof customModelInfoSchema>

export type CustomModelInfoSettings = {
customModelInfo?: Partial<CustomModelInfo> | null
}

const isPositiveInteger = (value: unknown): value is number =>
typeof value === "number" && Number.isSafeInteger(value) && value > 0

/**
* Applies the user metadata overlay without allowing invalid values to enter
* model arithmetic or cost/capability fields outside the supported override.
* When no discovered info exists, a context-window override is required to
* synthesize a usable ModelInfo.
*/
export const applyCustomModelInfo = (
info: ModelInfo | undefined,
settings: CustomModelInfoSettings | undefined,
): ModelInfo | undefined => {
const override = settings?.customModelInfo

if (!override) {
return info
}

const validOverride: CustomModelInfo = {}

if (isPositiveInteger(override.contextWindow)) {
validOverride.contextWindow = override.contextWindow
}

if (isPositiveInteger(override.maxTokens)) {
validOverride.maxTokens = override.maxTokens
}

if (typeof override.supportsImages === "boolean") {
validOverride.supportsImages = override.supportsImages
}

if (typeof override.supportsPromptCache === "boolean") {
validOverride.supportsPromptCache = override.supportsPromptCache
}

if (info) {
return Object.keys(validOverride).length > 0 ? { ...info, ...validOverride } : info
}

if (!validOverride.contextWindow) {
Comment thread
everyoneexe marked this conversation as resolved.
return undefined
}

return {
maxTokens: undefined,
contextWindow: validOverride.contextWindow,
supportsImages: false,
supportsPromptCache: false,
...validOverride,
}
}

export type ModelRecord = Record<string, ModelInfo>

export type RouterModels = Record<DynamicProvider | LocalProvider, ModelRecord>
24 changes: 24 additions & 0 deletions packages/types/src/provider-settings.ts
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ import { z } from "zod"

import {
modelInfoSchema,
customModelInfoSchema,
openAiCodexServiceTierSchema,
reasoningEffortSettingSchema,
verbosityLevelsSchema,
Expand Down Expand Up @@ -72,6 +73,28 @@ export type DynamicProvider = (typeof dynamicProviders)[number]
export const isDynamicProvider = (key: string): key is DynamicProvider =>
dynamicProviders.includes(key as DynamicProvider)

/**
* Providers that accept user-supplied `customModelInfo` overrides in the
* settings UI. This is a strict subset of `dynamicProviders` — only providers
* whose model metadata can be manually adjusted by the user.
*
* NOTE: Of these, only vercel-ai-gateway and zoo-gateway apply the overlay
* inside `RouterProvider.resolveModelInfo()`. The others (openrouter, requesty,
* unbound) apply it in their own overridden `getModel()` methods.
*/
export const customModelInfoProviders = [
providerIdentifiers.openrouter,
providerIdentifiers.requesty,
providerIdentifiers.unbound,
providerIdentifiers.vercelAiGateway,
providerIdentifiers.zooGateway,
] as const

export type CustomModelInfoProvider = (typeof customModelInfoProviders)[number]

export const isCustomModelInfoProvider = (key: string): key is CustomModelInfoProvider =>
customModelInfoProviders.includes(key as CustomModelInfoProvider)

/**
* LocalProvider
*
Expand Down Expand Up @@ -185,6 +208,7 @@ const baseProviderSettingsSchema = z.object({
reasoningEffort: reasoningEffortSettingSchema.optional(),
modelMaxTokens: z.number().optional(),
modelMaxThinkingTokens: z.number().optional(),
customModelInfo: customModelInfoSchema.nullish(),

// Model verbosity.
verbosity: verbosityLevelsSchema.optional(),
Expand Down
11 changes: 11 additions & 0 deletions src/api/providers/__tests__/kenari.spec.ts
Original file line number Diff line number Diff line change
Expand Up @@ -77,6 +77,17 @@ describe("KenariHandler", () => {
expect(result.info.supportsPromptCache).toBe(false)
})

it("does not apply gateway-only custom metadata overrides", async () => {
const handler = new KenariHandler({
...mockOptions,
customModelInfo: { contextWindow: 100_000, maxTokens: 10_000 },
})
const result = await handler.fetchModel()

expect(result.info.contextWindow).toBe(1_048_576)
expect(result.info.maxTokens).toBe(32_768)
})

it("falls back to the default model id when none is configured", async () => {
const handler = new KenariHandler({ kenariApiKey: "test-key" })
const result = await handler.fetchModel()
Expand Down
75 changes: 73 additions & 2 deletions src/api/providers/__tests__/openrouter.spec.ts
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@ import { Anthropic } from "@anthropic-ai/sdk"
import OpenAI from "openai"

import { OpenRouterHandler } from "../openrouter"
import { getModelEndpoints } from "../fetchers/modelEndpointCache"
import { Package } from "../../../shared/package"
import { makeApiHandlerOptions } from "../../../test-utils/api"
import { asyncStreamFrom, collectStream } from "../../../test-utils/stream"
Expand Down Expand Up @@ -102,13 +103,22 @@ vitest.mock("../fetchers/modelCache", () => ({
}),
}))

vitest.mock("../fetchers/modelEndpointCache", () => ({
getModelEndpoints: vitest.fn().mockResolvedValue({}),
}))

describe("OpenRouterHandler", () => {
const mockOptions = makeApiHandlerOptions({
openRouterApiKey: "test-key",
openRouterModelId: "anthropic/claude-sonnet-4",
})

beforeEach(() => clearAllMocks())
beforeEach(() => {
clearAllMocks()
// Reset getModelEndpoints to its default empty-object return so per-test
// overrides (e.g. specific-provider test) don't leak into subsequent tests.
vitest.mocked(getModelEndpoints).mockResolvedValue({})
})

it("initializes with correct options", () => {
const handler = new OpenRouterHandler(mockOptions)
Expand Down Expand Up @@ -140,14 +150,75 @@ describe("OpenRouterHandler", () => {
})
})

it("applies custom metadata before deriving request parameters", async () => {
const handler = new OpenRouterHandler({
...mockOptions,
customModelInfo: {
contextWindow: 100_000,
maxTokens: 10_000,
supportsImages: false,
supportsPromptCache: false,
},
})

const result = await handler.fetchModel()

expect(result.info.contextWindow).toBe(100_000)
expect(result.info.maxTokens).toBe(10_000)
expect(result.info.supportsImages).toBe(false)
expect(result.info.supportsPromptCache).toBe(false)
expect(result.maxTokens).toBe(10_000)
})

it("applies custom metadata to a discovered specific-provider endpoint", async () => {
vitest.mocked(getModelEndpoints).mockResolvedValue({
"test-provider": {
contextWindow: 128_000,
maxTokens: 16_384,
supportsImages: true,
supportsPromptCache: true,
},
})

const handler = new OpenRouterHandler({
...mockOptions,
openRouterSpecificProvider: "test-provider",
customModelInfo: { contextWindow: 100_000, maxTokens: 10_000 },
})

const result = await handler.fetchModel()

expect(result.info.contextWindow).toBe(100_000)
expect(result.info.maxTokens).toBe(10_000)
})

it("synthesizes metadata for an unlisted configured model", async () => {
const modelId = "provider/unlisted-model"
const handler = new OpenRouterHandler({
...mockOptions,
openRouterModelId: modelId,
customModelInfo: {
contextWindow: 100_000,
maxTokens: 10_000,
},
})

const result = await handler.fetchModel()

expect(result.id).toBe(modelId)
expect(result.info.contextWindow).toBe(100_000)
expect(result.info.maxTokens).toBe(10_000)
expect(result.maxTokens).toBe(10_000)
})

it("returns default model info when options are not provided", async () => {
const handler = new OpenRouterHandler({})
const result = await handler.fetchModel()
expect(result.id).toBe("anthropic/claude-sonnet-4.5")
expect(result.info.supportsPromptCache).toBe(true)
})

it("honors custom maxTokens for thinking models", async () => {
it("clamps maxTokens to 20% of context window for thinking models", async () => {
const handler = new OpenRouterHandler(
makeApiHandlerOptions({
openRouterApiKey: "test-key",
Expand Down
39 changes: 39 additions & 0 deletions src/api/providers/__tests__/requesty.spec.ts
Original file line number Diff line number Diff line change
Expand Up @@ -159,6 +159,45 @@ describe("RequestyHandler", () => {
})
})

it("applies custom metadata before deriving request parameters", async () => {
const handler = new RequestyHandler({
...mockOptions,
customModelInfo: {
contextWindow: 100_000,
maxTokens: 10_000,
supportsImages: false,
supportsPromptCache: false,
},
})

const result = await handler.fetchModel()

expect(result.info.contextWindow).toBe(100_000)
expect(result.info.maxTokens).toBe(10_000)
expect(result.info.supportsImages).toBe(false)
expect(result.info.supportsPromptCache).toBe(false)
expect(result.maxTokens).toBe(10_000)
})

it("synthesizes metadata for an unlisted configured model", async () => {
const modelId = "provider/unlisted-model"
const handler = new RequestyHandler({
...mockOptions,
requestyModelId: modelId,
customModelInfo: {
contextWindow: 100_000,
maxTokens: 10_000,
},
})

const result = await handler.fetchModel()

expect(result.id).toBe(modelId)
expect(result.info.contextWindow).toBe(100_000)
expect(result.info.maxTokens).toBe(10_000)
expect(result.maxTokens).toBe(10_000)
})

it("returns default model info when options are not provided", async () => {
const handler = new RequestyHandler({})
const result = await handler.fetchModel()
Expand Down
Loading
Loading