From f28f9b7d5585e87eb3e6dfeb896e9f92b8d63486 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 2 Oct 2026 22:07:32 +0000 Subject: [PATCH] feat(models): map model tiers per space and let bots and helpers ask for one A space can now map the tiers fast, standard and code to connected models, including several models served by one OpenAI-compatible router endpoint (one credential, many model ids). A bot can ask for a tier, and a single run_subagent call can ask for one. A tier with no model mapped falls back to the usual model choice, so nothing breaks before tiers are set up. - db: SpaceModelTier repo (list, find, set or clear); tier mappings are copied into a new space like model and voice preferences; createBot carries a tier, so duplicates keep it and spawned bots inherit it. - contracts/api: models.tiers.list and models.tiers.set; models.probeOpenAiCompatible can list a connected server's models by credential id with its stored address and key; bots.update accepts modelTier and keeps a pinned model and a tier mutually exclusive, so a model picked on a surface that knows no tiers (mobile) is never silently overridden. Thinking-level validation follows the tier's model. AI consent status discloses tier models; desktop local settings allow the tier procedures. - adapters: selectConfiguredModel runs a mapped tier ahead of the pinned model and the space default; the executor resolves bot tiers for runs and model lookups, and helper tiers through resolveModelTier. run_subagent takes model_tier; an explicit model wins over it. - web: a collapsed Tiers section in model settings (shown once a model is connected) and mapped tiers as choices in the bot Model select. - docs: tiers in self-host model setup. - tests: resolution order and fallbacks (selection, executor, Pi helper), repo, API and contract units, authorization and offline Pi Postgres suites, and a model-tiers web E2E spec. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01MRtGG9ARfqqhma3t85Ro2u --- apps/api/src/ai-consent.test.ts | 42 +++++ apps/api/src/ai-consent.ts | 15 +- apps/api/src/local-settings.test.ts | 1 + apps/api/src/model-tiers.test.ts | 111 +++++++++++++ apps/api/src/model-tiers.ts | 77 +++++++++ apps/api/src/router.ts | 56 +++++-- apps/web/e2e/model-tiers.spec.ts | 113 +++++++++++++ apps/web/src/lib/connected-models.test.ts | 58 +++++++ apps/web/src/lib/connected-models.ts | 46 ++++++ apps/web/src/lib/model-tier-label.ts | 8 + apps/web/src/pages/ModelSettingsOverlay.tsx | 2 + apps/web/src/pages/ModelTierSettings.tsx | 129 +++++++++++++++ apps/web/src/pages/shell/bot-panel.tsx | 137 +++++++++------- docs/self-host.md | 6 + packages/adapter-kit/src/types.ts | 4 +- packages/adapters/src/builtin-tools.ts | 9 ++ packages/adapters/src/executor.test.ts | 109 +++++++++++++ packages/adapters/src/executor.ts | 74 ++++++--- packages/adapters/src/model-selection.test.ts | 44 ++++++ packages/adapters/src/model-selection.ts | 13 +- .../src/pi-runtime-thinking-level.test.ts | 97 ++++++++++++ .../src/pi-runtime-tool-schema.test.ts | 32 ++++ packages/adapters/src/pi-runtime.ts | 20 ++- packages/contracts/src/domain.ts | 9 ++ packages/contracts/src/index.test.ts | 41 +++++ packages/contracts/src/index.ts | 1 + packages/contracts/src/local-settings.js | 2 + packages/contracts/src/model-tiers.ts | 24 +++ packages/contracts/src/rpc.ts | 17 +- packages/db/src/index.ts | 1 + packages/db/src/model-tiers.test.ts | 108 +++++++++++++ packages/db/src/model-tiers.ts | 83 ++++++++++ packages/db/src/repos.ts | 6 +- packages/db/src/spaces.ts | 19 ++- packages/testkit/src/authorization.test.ts | 80 ++++++++++ packages/testkit/src/model-emulator.ts | 7 +- .../testkit/src/pi-offline.postgres.test.ts | 148 ++++++++++++++++++ 37 files changed, 1653 insertions(+), 96 deletions(-) create mode 100644 apps/api/src/model-tiers.test.ts create mode 100644 apps/api/src/model-tiers.ts create mode 100644 apps/web/e2e/model-tiers.spec.ts create mode 100644 apps/web/src/lib/connected-models.test.ts create mode 100644 apps/web/src/lib/connected-models.ts create mode 100644 apps/web/src/lib/model-tier-label.ts create mode 100644 apps/web/src/pages/ModelTierSettings.tsx create mode 100644 packages/contracts/src/model-tiers.ts create mode 100644 packages/db/src/model-tiers.test.ts create mode 100644 packages/db/src/model-tiers.ts diff --git a/apps/api/src/ai-consent.test.ts b/apps/api/src/ai-consent.test.ts index 8b767a0..caf5be5 100644 --- a/apps/api/src/ai-consent.test.ts +++ b/apps/api/src/ai-consent.test.ts @@ -1,3 +1,4 @@ +import { serializeModelSecret } from "@rakazo/adapters"; import type { Actor } from "@rakazo/contracts"; import { AI_DISCLOSURE_VERSION } from "@rakazo/contracts"; import { describe, expect, it, vi } from "vitest"; @@ -15,6 +16,7 @@ function setup() { env: { agentRuntime: "scripted" }, prisma: { spaceModelPreference: { findMany: vi.fn(async () => []), findFirst: vi.fn(async () => null) }, + spaceModelTier: { findMany: vi.fn(async () => []) }, secret: { findMany: vi.fn(async () => []) }, bot: { findMany: vi.fn(async () => []) }, spaceVoicePreference: { @@ -58,6 +60,46 @@ describe("consent grants", () => { expect.objectContaining({ where: expect.objectContaining({ id: { in: ["bot"] } }) }), ); }); + it("discloses the model a bot's tier maps to on a router endpoint", async () => { + const { deps, actor } = setup(); + deps.env.agentRuntime = "pi"; + deps.secrets = { + load: vi.fn(() => + serializeModelSecret({ kind: "openai_compatible", baseUrl: "http://127.0.0.1:8080/v1" }), + ), + } as unknown as RouterDeps["secrets"]; + vi.mocked(deps.prisma.spaceModelTier.findMany).mockResolvedValue([ + { + tier: "code", + credentialId: "router", + modelId: "router-code", + credential: { + id: "router", + userId: actor.userId, + provider: "openai-compatible", + label: "Router", + secretId: "router-secret", + supportsImages: false, + createdAt: new Date(0), + updatedAt: new Date(0), + preferences: [], + }, + }, + ] as never); + vi.mocked(deps.prisma.bot.findMany).mockResolvedValue([ + { modelProvider: null, modelId: null, thinkingLevel: null, modelTier: "code" }, + ] as never); + vi.mocked(deps.prisma.secret.findMany).mockResolvedValue([ + { id: "router-secret", ciphertext: "fake-ciphertext" }, + ] as never); + const status = await aiConsentStatus(deps, actor, { botId: "bot", uses: ["model"] }); + expect(status.recipients).toEqual([ + expect.objectContaining({ + name: expect.stringContaining("http://127.0.0.1:8080"), + detail: "Model: router-code.", + }), + ]); + }); it("discloses only the selected voice provider before a voice action", async () => { const { deps, actor } = setup(); vi.mocked(deps.prisma.spaceVoicePreference.findMany).mockResolvedValue([ diff --git a/apps/api/src/ai-consent.ts b/apps/api/src/ai-consent.ts index 2795f78..8c595b5 100644 --- a/apps/api/src/ai-consent.ts +++ b/apps/api/src/ai-consent.ts @@ -13,6 +13,7 @@ import { findDefaultModelCredential, findDefaultVoiceCredential, findModelCredential, + listSpaceModelTiers, } from "@rakazo/db"; import { isSpaceOwner } from "./memory-provider-config.js"; import type { RouterDeps } from "./router.js"; @@ -44,7 +45,7 @@ export async function aiConsentStatus( : target.memberBotIds : undefined; const includesMemory = uses.includes("memory"); - const [preferences, defaultCredential, bots, voices, memoryOwner, settings, consents] = + const [preferences, tiers, defaultCredential, bots, voices, memoryOwner, settings, consents] = await Promise.all([ modelsEnabled ? deps.prisma.spaceModelPreference.findMany({ @@ -52,6 +53,7 @@ export async function aiConsentStatus( include: { credential: true }, }) : [], + modelsEnabled ? listSpaceModelTiers(deps.prisma, actor) : [], modelsEnabled ? findDefaultModelCredential(deps.prisma, actor) : null, modelsEnabled ? deps.prisma.bot.findMany({ @@ -61,7 +63,7 @@ export async function aiConsentStatus( archivedAt: null, ...(botIds ? { id: { in: botIds } } : {}), }, - select: { modelProvider: true, modelId: true, thinkingLevel: true }, + select: { modelProvider: true, modelId: true, thinkingLevel: true, modelTier: true }, }) : [], uses.includes("voice") @@ -100,6 +102,7 @@ export async function aiConsentStatus( : null; return selectConfiguredModel({ bot, + tierModel: tiers.find((tier) => tier.tier === bot?.modelTier) ?? null, overrideCredential, defaultCredential, settings, @@ -119,6 +122,14 @@ export async function aiConsentStatus( }, thinkingLevel: null, }); + // So are the models mapped to tiers, which helpers can ask for. + for (const tier of tiers) + selected.push({ + provider: tier.credential.provider, + id: tier.modelId, + credential: tier.credential, + thinkingLevel: null, + }); if (deps.env.teamChatJudgeProvider && deps.env.teamChatJudgeModel) { selected.push({ provider: deps.env.teamChatJudgeProvider, diff --git a/apps/api/src/local-settings.test.ts b/apps/api/src/local-settings.test.ts index e5c76e7..bdd9ab0 100644 --- a/apps/api/src/local-settings.test.ts +++ b/apps/api/src/local-settings.test.ts @@ -70,6 +70,7 @@ describe("local settings authority", () => { for (const procedure of [ "models/connect", "models/setDefault", + "models/tiers/set", "integrationSetup/get", "integrationSetup/save", ]) { diff --git a/apps/api/src/model-tiers.test.ts b/apps/api/src/model-tiers.test.ts new file mode 100644 index 0000000..19679d4 --- /dev/null +++ b/apps/api/src/model-tiers.test.ts @@ -0,0 +1,111 @@ +import { serializeModelSecret } from "@rakazo/adapters"; +import type { Actor } from "@rakazo/contracts"; +import type { PrismaClient } from "@rakazo/db"; +import { describe, expect, it, vi } from "vitest"; +import { connectedOpenAiCompatibleServer, setModelTier } from "./model-tiers.js"; + +const actor = { userId: "user", spaceId: "space" } as Actor; + +function setup(provider: string | null) { + const upsert = vi.fn(async () => ({})); + const deleteMany = vi.fn(async () => ({ count: 1 })); + const credentialFindFirst = vi.fn(async () => + provider ? { provider, secretId: "router-secret" } : null, + ); + const prisma = { + userModelCredential: { findFirst: credentialFindFirst }, + spaceModelTier: { upsert, deleteMany, findMany: vi.fn(async () => []) }, + secret: { + findFirst: vi.fn(async () => ({ id: "router-secret", ciphertext: "fake-ciphertext" })), + }, + } as unknown as PrismaClient; + return { prisma, upsert, deleteMany, credentialFindFirst }; +} + +describe("model tier settings", () => { + it("maps a tier to any model an OpenAI-compatible server serves", async () => { + const { prisma, upsert, credentialFindFirst } = setup("openai-compatible"); + + await setModelTier(prisma, actor, { + tier: "code", + model: { credentialId: "router", modelId: "router-code" }, + }); + + expect(credentialFindFirst).toHaveBeenCalledWith( + expect.objectContaining({ where: { id: "router", userId: "user" } }), + ); + expect(upsert).toHaveBeenCalledWith( + expect.objectContaining({ + create: expect.objectContaining({ + tier: "code", + credentialId: "router", + modelId: "router-code", + }), + }), + ); + }); + + it("refuses another user's credential and models a catalog provider does not run", async () => { + const missing = setup(null); + await expect( + setModelTier(missing.prisma, actor, { + tier: "fast", + model: { credentialId: "someone-else", modelId: "router-fast" }, + }), + ).rejects.toMatchObject({ code: "NOT_FOUND" }); + expect(missing.upsert).not.toHaveBeenCalled(); + + const catalog = setup("xai"); + await expect( + setModelTier(catalog.prisma, actor, { + tier: "fast", + model: { credentialId: "xai", modelId: "not-a-model" }, + }), + ).rejects.toMatchObject({ code: "BAD_REQUEST" }); + await setModelTier(catalog.prisma, actor, { + tier: "fast", + model: { credentialId: "xai", modelId: "grok-4.6" }, + }); + expect(catalog.upsert).toHaveBeenCalledTimes(1); + }); + + it("clears a tier without looking up a credential", async () => { + const { prisma, deleteMany, credentialFindFirst } = setup("openai-compatible"); + + await expect(setModelTier(prisma, actor, { tier: "standard", model: null })).resolves.toEqual( + [], + ); + + expect(credentialFindFirst).not.toHaveBeenCalled(); + expect(deleteMany).toHaveBeenCalledWith({ + where: { spaceId: "space", userId: "user", tier: "standard" }, + }); + }); + + it("probes a connected server with its stored address and key, for its owner only", async () => { + const { prisma, credentialFindFirst } = setup("openai-compatible"); + const secrets = { + load: vi.fn(() => + serializeModelSecret({ + kind: "openai_compatible", + baseUrl: "https://router.example.test/v1", + apiKey: "fake-router-key", + }), + ), + }; + + await expect( + connectedOpenAiCompatibleServer({ prisma, secrets }, actor, "router"), + ).resolves.toEqual({ baseUrl: "https://router.example.test/v1", apiKey: "fake-router-key" }); + expect(credentialFindFirst).toHaveBeenCalledWith( + expect.objectContaining({ + where: { id: "router", userId: "user", provider: "openai-compatible" }, + }), + ); + + const missing = setup(null); + await expect( + connectedOpenAiCompatibleServer({ prisma: missing.prisma, secrets }, actor, "router"), + ).rejects.toMatchObject({ code: "NOT_FOUND" }); + }); +}); diff --git a/apps/api/src/model-tiers.ts b/apps/api/src/model-tiers.ts new file mode 100644 index 0000000..faf0b6b --- /dev/null +++ b/apps/api/src/model-tiers.ts @@ -0,0 +1,77 @@ +import { ORPCError } from "@orpc/server"; +import type { EncryptedSecretStore } from "@rakazo/adapters"; +import { isCatalogModelChoice, parseModelSecret } from "@rakazo/adapters"; +import type { Actor, ModelTierMapping, SetModelTierInput } from "@rakazo/contracts"; +import { ModelTierSchema, OPENAI_COMPATIBLE_PROVIDER_ID } from "@rakazo/contracts"; +import type { PrismaClient } from "@rakazo/db"; +import { listSpaceModelTiers, setSpaceModelTier } from "@rakazo/db"; + +const NOT_CONNECTED = "Connect that model provider first"; + +export async function listModelTiers( + prisma: PrismaClient, + actor: Actor, +): Promise { + const rows = await listSpaceModelTiers(prisma, actor); + return rows.flatMap((row) => { + const tier = ModelTierSchema.safeParse(row.tier); + return tier.success + ? [ + { + tier: tier.data, + credentialId: row.credentialId, + provider: row.credential.provider, + modelId: row.modelId, + }, + ] + : []; + }); +} + +export async function setModelTier(prisma: PrismaClient, actor: Actor, input: SetModelTierInput) { + if (input.model) { + const credential = await prisma.userModelCredential.findFirst({ + where: { id: input.model.credentialId, userId: actor.userId }, + select: { provider: true }, + }); + if (!credential) throw new ORPCError("NOT_FOUND", { message: NOT_CONNECTED }); + // An OpenAI-compatible server runs whatever ids it serves; other providers run catalog models. + if ( + credential.provider !== OPENAI_COMPATIBLE_PROVIDER_ID && + !isCatalogModelChoice(credential.provider, input.model.modelId) + ) { + throw new ORPCError("BAD_REQUEST", { message: "Unknown model for that provider" }); + } + } + await setSpaceModelTier(prisma, actor, input.tier, input.model); + return listModelTiers(prisma, actor); +} + +/** The stored address and key of a connected OpenAI-compatible server, for listing its models. */ +export async function connectedOpenAiCompatibleServer( + deps: { prisma: PrismaClient; secrets: Pick }, + actor: Actor, + credentialId: string, +) { + const credential = await deps.prisma.userModelCredential.findFirst({ + where: { id: credentialId, userId: actor.userId, provider: OPENAI_COMPATIBLE_PROVIDER_ID }, + select: { secretId: true }, + }); + const secret = credential + ? await deps.prisma.secret.findFirst({ + where: { id: credential.secretId, userId: actor.userId, spaceId: null }, + select: { id: true, ciphertext: true }, + }) + : null; + if (!secret) throw new ORPCError("NOT_FOUND", { message: NOT_CONNECTED }); + let stored: ReturnType; + try { + stored = parseModelSecret(deps.secrets.load(secret.ciphertext, secret.id)); + } catch { + throw new ORPCError("NOT_FOUND", { message: NOT_CONNECTED }); + } + if (stored.kind !== "openai_compatible") { + throw new ORPCError("NOT_FOUND", { message: NOT_CONNECTED }); + } + return { baseUrl: stored.baseUrl, apiKey: stored.apiKey }; +} diff --git a/apps/api/src/router.ts b/apps/api/src/router.ts index 7067820..fe88a5d 100644 --- a/apps/api/src/router.ts +++ b/apps/api/src/router.ts @@ -118,6 +118,7 @@ import { findDefaultModelCredential, findDefaultVoiceCredential, findModelCredential, + findSpaceModelTier, formatMessagingLinkCode, InvalidProjectWorkspacePathError, InvalidSpaceNameError, @@ -160,6 +161,7 @@ import { persistMemoryProviderConfig, updateMemoryProviderDefaultScope, } from "./memory-provider-config.js"; +import { connectedOpenAiCompatibleServer, listModelTiers, setModelTier } from "./model-tiers.js"; import { chooseFocus, dismissFocus, @@ -892,8 +894,12 @@ export function createRouter(deps: RouterDeps) { }), probeOpenAiCompatible: authed.models.probeOpenAiCompatible.handler( async ({ context, input }) => { + const server = + "credentialId" in input + ? await connectedOpenAiCompatibleServer(deps, context.actor, input.credentialId) + : input; try { - const models = await probeOpenAiCompatibleModels(input, undefined, context.signal); + const models = await probeOpenAiCompatibleModels(server, undefined, context.signal); return { models }; } catch (error) { throw new ORPCError("BAD_REQUEST", { @@ -969,6 +975,14 @@ export function createRouter(deps: RouterDeps) { ); return { ok: true as const }; }), + tiers: { + list: authed.models.tiers.list.handler(({ context }) => + listModelTiers(deps.prisma, context.actor), + ), + set: authed.models.tiers.set.handler(({ context, input }) => + setModelTier(deps.prisma, context.actor, input), + ), + }, }, bots: { list: authed.bots.list.handler(async ({ context }) => repos.listBots(context.actor)), @@ -1001,6 +1015,7 @@ export function createRouter(deps: RouterDeps) { modelProvider: source.modelProvider, modelId: source.modelId, thinkingLevel: source.thinkingLevel, + modelTier: source.modelTier, }) .catch((error: unknown) => { throw mapSpaceLifecycleError(error); @@ -1062,12 +1077,28 @@ export function createRouter(deps: RouterDeps) { } const thinkingLevel = input.thinkingLevel; if (input.thinkingLevel) { + // Mirror what this update writes: a new tier clears the pin and a new pin the tier. const provider = - input.modelProvider !== undefined ? input.modelProvider : existing.modelProvider; - const modelId = input.modelId !== undefined ? input.modelId : existing.modelId; + input.modelProvider !== undefined + ? input.modelProvider + : input.modelTier + ? null + : existing.modelProvider; + const modelId = + input.modelId !== undefined ? input.modelId : input.modelTier ? null : existing.modelId; + const modelTier = + input.modelTier !== undefined + ? input.modelTier + : input.modelProvider + ? null + : existing.modelTier; + const tierModel = modelTier + ? await findSpaceModelTier(deps.prisma, context.actor, modelTier) + : null; const me = await meDto(deps, context.actor); - const effectiveProvider = provider ?? me.defaultProvider; - const effectiveModelId = modelId ?? me.defaultModel; + const effectiveProvider = + tierModel?.credential.provider ?? provider ?? me.defaultProvider; + const effectiveModelId = tierModel?.modelId ?? modelId ?? me.defaultModel; if (effectiveProvider && effectiveModelId) { const entry = listPiCatalog().find( (item) => item.provider === effectiveProvider && item.id === effectiveModelId, @@ -1075,12 +1106,11 @@ export function createRouter(deps: RouterDeps) { let allowed = entry?.thinkingLevels; if (effectiveProvider === OPENAI_COMPATIBLE_PROVIDER_ID) { allowed = ["off"]; - const credential = await findModelCredential( - deps.prisma, - context.actor, - effectiveProvider, - ); - if (credential && credential.defaultModel === effectiveModelId) { + // A tier's server keeps one reasoning setting for every model it serves. + const credential = + tierModel?.credential ?? + (await findModelCredential(deps.prisma, context.actor, effectiveProvider)); + if (credential && (tierModel || credential.defaultModel === effectiveModelId)) { const secret = await deps.prisma.secret.findFirst({ where: { id: credential.secretId, userId: context.actor.userId, spaceId: null }, select: { ciphertext: true }, @@ -1129,6 +1159,10 @@ export function createRouter(deps: RouterDeps) { ? { modelProvider: input.modelProvider, modelId: input.modelId ?? null } : {}), ...(input.thinkingLevel !== undefined ? { thinkingLevel } : {}), + ...(input.modelTier !== undefined ? { modelTier: input.modelTier } : {}), + // A bot runs on a pinned model or a tier, so choosing one clears the other. + ...(input.modelTier ? { modelProvider: null, modelId: null } : {}), + ...(input.modelProvider ? { modelTier: null } : {}), ...(input.teamChatAmbientEnabled !== undefined ? { teamChatAmbientEnabled: input.teamChatAmbientEnabled } : {}), diff --git a/apps/web/e2e/model-tiers.spec.ts b/apps/web/e2e/model-tiers.spec.ts new file mode 100644 index 0000000..ba4abde --- /dev/null +++ b/apps/web/e2e/model-tiers.spec.ts @@ -0,0 +1,113 @@ +import { createServer } from "node:http"; +import type { AddressInfo } from "node:net"; +import { expect, type Page, test } from "@playwright/test"; +import { captureScreenshot, completeOnboarding, openUserSettings, rpc, signup } from "./helpers"; + +// One OpenAI-compatible router endpoint serving several local models. +const ROUTER_MODELS = ["router-standard", "router-fast", "router-code"]; + +async function mappedTiers(page: Page) { + const tiers = await rpc>(page, "models/tiers/list", {}); + return tiers.map((entry) => `${entry.tier}:${entry.modelId}`).sort(); +} + +async function openTiers(page: Page) { + const tiers = page.getByTestId("model-tiers"); + await tiers.getByText("Tiers", { exact: true }).click(); + return tiers; +} + +async function openBotModel(page: Page) { + await page.locator("main").getByRole("button", { name: "Chief", exact: true }).click(); + const settings = page.getByTestId("bot-settings"); + await expect(settings).toBeVisible(); + await settings.getByTestId("bot-settings-advanced").evaluate((element) => { + (element as HTMLDetailsElement).open = true; + }); + return { settings, model: settings.getByRole("combobox", { name: "Model", exact: true }) }; +} + +test("maps tiers to the models of one router endpoint and lets a bot ask for a tier", async ({ + page, +}, testInfo) => { + const server = createServer((request, response) => { + if (request.method === "GET" && request.url === "/v1/models") { + response.writeHead(200, { "content-type": "application/json" }); + response.end(JSON.stringify({ object: "list", data: ROUTER_MODELS.map((id) => ({ id })) })); + return; + } + response.writeHead(404); + response.end(); + }); + await new Promise((resolve, reject) => { + server.once("error", reject); + server.listen(0, "127.0.0.1", resolve); + }); + + try { + const baseUrl = `http://127.0.0.1:${(server.address() as AddressInfo).port}/v1`; + const stamp = Date.now(); + await signup(page, `model-tiers-${stamp}@rakazo.test`, "password12", `Tiers ${stamp}`); + await completeOnboarding(page); + + await openUserSettings(page, "models"); + // Nothing to map before a model is connected. + await expect(page.getByTestId("model-tiers")).toBeHidden(); + await page.getByPlaceholder("Search providers").fill("openai-compatible"); + await page.getByRole("button", { name: /OpenAI-compatible/ }).click(); + await page.getByLabel("OpenAI-compatible server URL").fill(baseUrl); + await page.getByLabel("Model id").fill("router-standard"); + await page.getByRole("button", { name: "Save", exact: true }).click(); + await expect(page.getByText("Saved.", { exact: true })).toBeVisible(); + const [router] = await rpc>(page, "models/credentials", {}); + + const tiers = await openTiers(page); + const fast = tiers.getByRole("combobox", { name: "Fast", exact: true }); + const code = tiers.getByRole("combobox", { name: "Code", exact: true }); + await expect(fast).toHaveValue(""); + // Every model the router lists becomes a choice for each tier. + for (const modelId of ROUTER_MODELS) { + await expect(code.locator(`option[value="${router!.id}::${modelId}"]`)).toHaveCount(1); + } + await fast.selectOption(`${router!.id}::router-fast`); + await code.selectOption(`${router!.id}::router-code`); + await expect.poll(() => mappedTiers(page)).toEqual(["code:router-code", "fast:router-fast"]); + await captureScreenshot(page, testInfo, "model-tiers-settings"); + + await page.reload(); + await openUserSettings(page, "models"); + const reopened = await openTiers(page); + await expect(reopened.getByRole("combobox", { name: "Code", exact: true })).toHaveValue( + `${router!.id}::router-code`, + ); + await reopened.getByRole("combobox", { name: "Fast", exact: true }).selectOption(""); + await expect.poll(() => mappedTiers(page)).toEqual(["code:router-code"]); + await page.getByRole("button", { name: "Close model settings" }).click(); + + // A mapped tier is a model choice for a bot; an unmapped one is not offered. + const { settings, model } = await openBotModel(page); + await expect(model.locator('option[value="tier:code"]')).toHaveText("Code (router-code)"); + await expect(model.locator('option[value="tier:fast"]')).toHaveCount(0); + await model.selectOption("tier:code"); + await model.scrollIntoViewIfNeeded(); + await captureScreenshot(page, testInfo, "model-tiers-bot-settings"); + const saved = page.waitForResponse( + (response) => response.url().includes("/rpc/bots/update") && response.ok(), + ); + await settings.getByRole("button", { name: "Save", exact: true }).click(); + await saved; + const bots = await rpc< + Array<{ name: string; modelTier: string | null; modelProvider: string | null }> + >(page, "bots/list", {}); + expect(bots.find((bot) => bot.name === "Chief")).toMatchObject({ + modelTier: "code", + modelProvider: null, + }); + + await page.reload(); + const after = await openBotModel(page); + await expect(after.model).toHaveValue("tier:code"); + } finally { + await new Promise((resolve) => server.close(() => resolve())); + } +}); diff --git a/apps/web/src/lib/connected-models.test.ts b/apps/web/src/lib/connected-models.test.ts new file mode 100644 index 0000000..ccc0123 --- /dev/null +++ b/apps/web/src/lib/connected-models.test.ts @@ -0,0 +1,58 @@ +import type { ModelCatalogEntry, ModelCredential } from "@rakazo/contracts"; +import { describe, expect, it } from "vitest"; +import { connectedModelOptions } from "./connected-models"; + +const catalog: ModelCatalogEntry[] = [ + { provider: "xai", providerName: "xAI", id: "grok-4.6", label: "Grok 4.6", billing: "" }, + { provider: "xai", providerName: "xAI", id: "grok-mini", label: "Grok mini", billing: "" }, + { + provider: "openai-compatible", + providerName: "OpenAI-compatible", + id: "custom", + label: "Custom", + billing: "", + placeholder: true, + }, +]; + +const credentials: ModelCredential[] = [ + { + id: "router", + provider: "openai-compatible", + label: "OpenAI-compatible", + hasKey: true, + isDefault: true, + modelId: "router-standard", + }, + { id: "xai", provider: "xai", label: "xAI", hasKey: true, isDefault: false, modelId: "grok-4.6" }, +]; + +describe("connected model options", () => { + it("expands catalog providers and keeps a free-form connection's saved model", () => { + expect(connectedModelOptions(credentials, catalog)).toEqual([ + { + credentialId: "router", + provider: "openai-compatible", + modelId: "router-standard", + label: "OpenAI-compatible · router-standard", + }, + { credentialId: "xai", provider: "xai", modelId: "grok-4.6", label: "xAI · Grok 4.6" }, + { credentialId: "xai", provider: "xai", modelId: "grok-mini", label: "xAI · Grok mini" }, + ]); + }); + + it("adds every model a router endpoint lists, once each", () => { + const options = connectedModelOptions(credentials, catalog, { + router: ["router-standard", "router-fast", "router-code"], + }); + + expect(options.filter((option) => option.credentialId === "router")).toEqual( + ["router-standard", "router-fast", "router-code"].map((modelId) => ({ + credentialId: "router", + provider: "openai-compatible", + modelId, + label: `OpenAI-compatible · ${modelId}`, + })), + ); + }); +}); diff --git a/apps/web/src/lib/connected-models.ts b/apps/web/src/lib/connected-models.ts new file mode 100644 index 0000000..0ebb611 --- /dev/null +++ b/apps/web/src/lib/connected-models.ts @@ -0,0 +1,46 @@ +import type { ModelCatalogEntry, ModelCredential } from "@rakazo/contracts"; + +export type ConnectedModelOption = { + credentialId: string; + provider: string; + modelId: string; + label: string; +}; + +/** + * Models the user's connections can run: every catalog model of a catalog provider, and the + * saved id of a free-form connection plus any ids its server lists (keyed by credential id). + */ +export function connectedModelOptions( + credentials: ModelCredential[], + catalog: ModelCatalogEntry[], + serverModels: Record = {}, +): ConnectedModelOption[] { + const options: ConnectedModelOption[] = []; + const seen = new Set(); + for (const credential of credentials) { + const providerModels = catalog.filter( + (entry) => entry.provider === credential.provider && !entry.placeholder, + ); + const listed = serverModels[credential.id] ?? []; + const freeForm = + Boolean(credential.modelId) && + !providerModels.some((entry) => entry.id === credential.modelId); + const candidates = + freeForm || listed.length + ? [credential.modelId, ...listed] + .filter((modelId): modelId is string => Boolean(modelId)) + .map((modelId) => ({ modelId, label: `${credential.label} · ${modelId}` })) + : providerModels.map((entry) => ({ + modelId: entry.id, + label: `${entry.providerName ?? entry.provider} · ${entry.label}`, + })); + for (const candidate of candidates) { + const key = `${credential.provider}::${candidate.modelId}`; + if (seen.has(key)) continue; + seen.add(key); + options.push({ credentialId: credential.id, provider: credential.provider, ...candidate }); + } + } + return options; +} diff --git a/apps/web/src/lib/model-tier-label.ts b/apps/web/src/lib/model-tier-label.ts new file mode 100644 index 0000000..364ddc4 --- /dev/null +++ b/apps/web/src/lib/model-tier-label.ts @@ -0,0 +1,8 @@ +import { t } from "@lingui/core/macro"; +import type { ModelTier } from "@rakazo/contracts"; + +export function modelTierLabel(tier: ModelTier) { + if (tier === "fast") return t`Fast`; + if (tier === "code") return t`Code`; + return t`Standard`; +} diff --git a/apps/web/src/pages/ModelSettingsOverlay.tsx b/apps/web/src/pages/ModelSettingsOverlay.tsx index b2ea54c..6ea570d 100644 --- a/apps/web/src/pages/ModelSettingsOverlay.tsx +++ b/apps/web/src/pages/ModelSettingsOverlay.tsx @@ -40,6 +40,7 @@ import { localizedProviderHint } from "../lib/localized-provider-hint"; import type { ModelCatalogEntry, ModelCredential } from "../lib/model-auth"; import { rpc } from "../lib/rpc"; import { useModelOAuthSignIn } from "../lib/use-model-oauth-signin"; +import { ModelTierSettings } from "./ModelTierSettings"; export function ModelSettingsOverlay({ onClose, @@ -398,6 +399,7 @@ export function ModelSettingsOverlay({ )} +
diff --git a/apps/web/src/pages/ModelTierSettings.tsx b/apps/web/src/pages/ModelTierSettings.tsx new file mode 100644 index 0000000..bd1ddb0 --- /dev/null +++ b/apps/web/src/pages/ModelTierSettings.tsx @@ -0,0 +1,129 @@ +import { Trans, useLingui } from "@lingui/react/macro"; +import type { + ModelCatalogEntry, + ModelCredential, + ModelTier, + ModelTierMapping, +} from "@rakazo/contracts"; +import { MODEL_TIERS, OPENAI_COMPATIBLE_PROVIDER_ID } from "@rakazo/contracts"; +import { NativeSelect, NativeSelectOption } from "@rakazo/ui-web"; +import { useId, useRef, useState } from "react"; +import { connectedModelOptions } from "../lib/connected-models"; +import { modelTierLabel } from "../lib/model-tier-label"; +import { rpc } from "../lib/rpc"; + +function tierKey(model: { credentialId: string; modelId: string }) { + return `${model.credentialId}::${model.modelId}`; +} + +function parseTierKey(key: string) { + const separator = key.indexOf("::"); + if (separator <= 0) return null; + return { credentialId: key.slice(0, separator), modelId: key.slice(separator + 2) }; +} + +/** Maps each tier to one connected model; bots and helpers that ask for a tier run on it. */ +export function ModelTierSettings({ + credentials, + catalog, +}: { + credentials: ModelCredential[]; + catalog: ModelCatalogEntry[]; +}) { + const { t } = useLingui(); + const ids = useId(); + const [tiers, setTiers] = useState(null); + const [serverModels, setServerModels] = useState>({}); + const [error, setError] = useState(null); + const saveQueueRef = useRef(Promise.resolve()); + + function loadTiers() { + return rpc.models.tiers + .list() + .then(setTiers) + .catch((err: unknown) => + setError(err instanceof Error ? err.message : t`Could not load model settings`), + ); + } + + function load() { + void loadTiers(); + // A router endpoint serves several models; list them so each tier can pick one. + for (const credential of credentials) { + if (credential.provider !== OPENAI_COMPATIBLE_PROVIDER_ID) continue; + void rpc.models + .probeOpenAiCompatible({ credentialId: credential.id }) + .then(({ models }) => + setServerModels((current) => ({ ...current, [credential.id]: models })), + ) + // The saved model stays selectable when the server cannot be reached. + .catch(() => undefined); + } + } + + function chooseModel(tier: ModelTier, key: string) { + const model = parseTierKey(key); + const provider = credentials.find((entry) => entry.id === model?.credentialId)?.provider; + setError(null); + setTiers((current) => [ + ...(current ?? []).filter((entry) => entry.tier !== tier), + ...(model && provider ? [{ tier, provider, ...model }] : []), + ]); + // Serialize saves so a slower earlier request cannot land after a newer choice. + saveQueueRef.current = saveQueueRef.current.then(() => + rpc.models.tiers.set({ tier, model }).then( + () => undefined, + (err: unknown) => { + setError(err instanceof Error ? err.message : t`Could not save`); + return loadTiers(); + }, + ), + ); + } + + if (!credentials.length) return null; + const options = connectedModelOptions(credentials, catalog, serverModels); + + return ( +
{ + if (event.currentTarget.open && !tiers) load(); + }} + > + + Tiers + +
+ {MODEL_TIERS.map((tier) => { + const mapped = tiers?.find((entry) => entry.tier === tier); + const value = mapped ? tierKey(mapped) : ""; + return ( + + ); + })} +
+ {error ?

{error}

: null} +
+ ); +} diff --git a/apps/web/src/pages/shell/bot-panel.tsx b/apps/web/src/pages/shell/bot-panel.tsx index 74c336c..ca8eba9 100644 --- a/apps/web/src/pages/shell/bot-panel.tsx +++ b/apps/web/src/pages/shell/bot-panel.tsx @@ -7,6 +7,8 @@ import type { Me, ModelCatalogEntry, ModelCredential, + ModelTier, + ModelTierMapping, ThinkingLevel, VoiceInfo, } from "@rakazo/contracts"; @@ -14,11 +16,14 @@ import { BOT_DESCRIPTION_MAX_LENGTH, BOT_NAME_MAX_LENGTH, BOT_TITLE_MAX_LENGTH, + MODEL_TIERS, + ModelTierSchema, } from "@rakazo/contracts"; import { Button, Input, NativeSelect, + NativeSelectOptGroup, NativeSelectOption, Switch, Textarea, @@ -26,7 +31,9 @@ import { } from "@rakazo/ui-web"; import { X } from "lucide-react"; import { Suspense, useEffect, useId, useRef, useState } from "react"; +import { connectedModelOptions } from "../../lib/connected-models"; import { lazyWithRecovery } from "../../lib/lazy-recovery"; +import { modelTierLabel } from "../../lib/model-tier-label"; import { rpc } from "../../lib/rpc"; import { AvatarStudioPopover } from "./avatar-studio-popover"; @@ -212,6 +219,7 @@ export function BotSettings({ modelProvider?: string | null; modelId?: string | null; thinkingLevel?: ThinkingLevel | null; + modelTier?: ModelTier | null; }) => Promise; onExport: () => Promise; onClear: () => void; @@ -230,11 +238,16 @@ export function BotSettings({ const [voiceId, setVoiceId] = useState(bot.voiceId ?? ""); const [voices, setVoices] = useState([]); const [modelKey, setModelKey] = useState( - bot.modelProvider && bot.modelId ? modelOptionKey(bot.modelProvider, bot.modelId) : "", + bot.modelTier + ? tierOptionKey(bot.modelTier) + : bot.modelProvider && bot.modelId + ? modelOptionKey(bot.modelProvider, bot.modelId) + : "", ); const [thinkingLevel, setThinkingLevel] = useState(bot.thinkingLevel ?? ""); const [credentials, setCredentials] = useState([]); const [catalog, setCatalog] = useState([]); + const [tiers, setTiers] = useState([]); const [me, setMe] = useState(null); const [modelMetaReady, setModelMetaReady] = useState(false); const [saving, setSaving] = useState(false); @@ -254,11 +267,17 @@ export function BotSettings({ .voices({}) .then(setVoices) .catch(() => setVoices([])); - void Promise.all([rpc.models.credentials(), rpc.models.list(), rpc.me()]) - .then(([nextCredentials, nextCatalog, nextMe]) => { + void Promise.all([ + rpc.models.credentials(), + rpc.models.list(), + rpc.me(), + rpc.models.tiers.list().catch(() => []), + ]) + .then(([nextCredentials, nextCatalog, nextMe, nextTiers]) => { setCredentials(nextCredentials); setCatalog(nextCatalog); setMe(nextMe); + setTiers(nextTiers); // Only mark ready on success — a failed catalog load must not clear // an existing thinkingLevel override on save. setModelMetaReady(true); @@ -266,60 +285,49 @@ export function BotSettings({ .catch(() => undefined); }, []); - const connectedOptions: Array<{ - key: string; - provider: string; - modelId: string; - label: string; - }> = []; - const seenOptions = new Set(); - for (const credential of credentials) { - const providerModels = catalog.filter( - (entry) => entry.provider === credential.provider && !entry.placeholder, - ); - const credentialInCatalog = Boolean( - credential.modelId && providerModels.some((entry) => entry.id === credential.modelId), - ); - // Catalog providers expand to every model for that connection. Free-form - // credentials (model id not in the catalog) stay a single connected pair. - const options = - credential.modelId && !credentialInCatalog - ? [ - { - key: modelOptionKey(credential.provider, credential.modelId), - provider: credential.provider, - modelId: credential.modelId, - label: `${credential.label} · ${credential.modelId}`, - }, - ] - : providerModels.map((entry) => ({ - key: modelOptionKey(entry.provider, entry.id), - provider: entry.provider, - modelId: entry.id, - label: `${entry.providerName ?? entry.provider} · ${entry.label}`, - })); - for (const option of options) { - if (seenOptions.has(option.key)) continue; - seenOptions.add(option.key); - connectedOptions.push(option); - } - } - - const effectiveProvider = modelKey - ? parseModelOptionKey(modelKey)?.provider - : (me?.defaultProvider ?? null); - const effectiveModelId = modelKey - ? parseModelOptionKey(modelKey)?.modelId - : (me?.defaultModel ?? null); + const connectedOptions = connectedModelOptions(credentials, catalog).map((option) => ({ + ...option, + key: modelOptionKey(option.provider, option.modelId), + })); + const selectedTier = parseTierOptionKey(modelKey); + const selectedTierModel = tiers.find((entry) => entry.tier === selectedTier); + // A mapped tier runs its model; an unmapped one falls back to the Space default. + const effectiveProvider = selectedTierModel + ? selectedTierModel.provider + : modelKey && !selectedTier + ? parseModelOptionKey(modelKey)?.provider + : (me?.defaultProvider ?? null); + const effectiveModelId = selectedTierModel + ? selectedTierModel.modelId + : modelKey && !selectedTier + ? parseModelOptionKey(modelKey)?.modelId + : (me?.defaultModel ?? null); const effectiveEntry = effectiveProvider && effectiveModelId ? catalog.find( (entry) => entry.provider === effectiveProvider && entry.id === effectiveModelId, ) : undefined; - const effectiveCredential = credentials.find( - (entry) => entry.provider === effectiveProvider && entry.modelId === effectiveModelId, - ); + // One connection serves every model of its tiers with the same reasoning settings. + const effectiveCredential = selectedTierModel + ? credentials.find((entry) => entry.id === selectedTierModel.credentialId) + : credentials.find( + (entry) => entry.provider === effectiveProvider && entry.modelId === effectiveModelId, + ); + const tierOptions = MODEL_TIERS.flatMap((tier) => { + const mapped = tiers.find((entry) => entry.tier === tier); + // Tiers appear once mapped; an unmapped current choice stays visible. + if (!mapped && tier !== selectedTier) return []; + const model = mapped + ? (catalogLabel(catalog, mapped.provider, mapped.modelId) ?? mapped.modelId) + : null; + return [ + { + key: tierOptionKey(tier), + label: model ? `${modelTierLabel(tier)} (${model})` : modelTierLabel(tier), + }, + ]; + }); const thinkingOptions = ( effectiveCredential?.thinkingLevels ?? effectiveEntry?.thinkingLevels ?? @@ -334,7 +342,7 @@ export function BotSettings({ color?: string; notifyOnFinish?: boolean; }) { - const selected = modelKey ? parseModelOptionKey(modelKey) : null; + const selected = modelKey && !selectedTier ? parseModelOptionKey(modelKey) : null; const nextName = (patchOverrides?.name !== undefined ? patchOverrides.name : name).trim(); const nextTitle = (patchOverrides?.title !== undefined ? patchOverrides.title : title).trim(); const nextDescription = ( @@ -364,6 +372,7 @@ export function BotSettings({ voiceId: voiceId || null, modelProvider: selected?.provider ?? null, modelId: selected?.modelId ?? null, + modelTier: selectedTier, ...(modelMetaReady ? { thinkingLevel: thinkingOptions.length @@ -510,7 +519,18 @@ export function BotSettings({ ? ` (${catalogLabel(catalog, me.defaultProvider, me.defaultModel) ?? me.defaultModel})` : ""} - {modelKey && !connectedOptions.some((option) => option.key === modelKey) ? ( + {tierOptions.length ? ( + + {tierOptions.map((option) => ( + + {option.label} + + ))} + + ) : null} + {modelKey && + !selectedTier && + !connectedOptions.some((option) => option.key === modelKey) ? ( {parseModelOptionKey(modelKey)?.modelId ?? modelKey} @@ -635,6 +655,17 @@ function modelOptionKey(provider: string, modelId: string) { return `${provider}::${modelId}`; } +// Model keys always contain "::", so a tier key can never collide with one. +function tierOptionKey(tier: ModelTier) { + return `tier:${tier}`; +} + +function parseTierOptionKey(key: string): ModelTier | null { + if (!key.startsWith("tier:")) return null; + const tier = ModelTierSchema.safeParse(key.slice("tier:".length)); + return tier.success ? tier.data : null; +} + function thinkingLevelLabel(level: ThinkingLevel) { if (level === "xhigh") return t`Extra high`; if (level === "low") return t`Low`; diff --git a/docs/self-host.md b/docs/self-host.md index 8510b82..8349f41 100644 --- a/docs/self-host.md +++ b/docs/self-host.md @@ -223,6 +223,12 @@ screenshot computer tools stay available. Existing connections default to disabl managed endpoints, the deployment-wide fallback remains `RAKAZO_OPENAI_COMPATIBLE_VISION_MODELS=gpt4o-vision,llava`. +**Settings → Models → Tiers** maps the tiers fast, standard, and code to models of your +connections, per space. One router endpoint that serves several models (for example llama.cpp +behind a model router) can cover all three, because each tier offers every model the server +lists. A bot asks for a tier through its **Model** setting, and a helper can ask for one per +call. A tier with no model mapped falls back to the usual model choice. + For servers that accept standard `reasoning_effort`, enable **Supports thinking** under **Advanced** when connecting. The setting is saved on the connection (no env var or restart). Existing connections default to disabled. Reconnect former Qwen-list or deployment-local models diff --git a/packages/adapter-kit/src/types.ts b/packages/adapter-kit/src/types.ts index d9ebb31..d20970e 100644 --- a/packages/adapter-kit/src/types.ts +++ b/packages/adapter-kit/src/types.ts @@ -1,4 +1,4 @@ -import type { ConnectionCatalogItem, SandboxKind } from "@rakazo/contracts"; +import type { ConnectionCatalogItem, ModelTier, SandboxKind } from "@rakazo/contracts"; export interface AdapterContext { operationId: string; @@ -383,6 +383,8 @@ export interface AgentRunRequest { model: AgentRunModel; /** Resolve an explicitly requested helper model within the active user and space scope. */ resolveModel?: (provider: string, modelId: string) => Promise; + /** Resolve the model the space maps a helper's tier to; null when the tier is unmapped. */ + resolveModelTier?: (tier: ModelTier) => Promise; resumeFromCheckpoint?: string; script?: ScriptedTurn[]; /** diff --git a/packages/adapters/src/builtin-tools.ts b/packages/adapters/src/builtin-tools.ts index cc22529..3969b78 100644 --- a/packages/adapters/src/builtin-tools.ts +++ b/packages/adapters/src/builtin-tools.ts @@ -2,11 +2,15 @@ import type { ConnectorTool } from "@rakazo/adapter-kit"; import { BotSecretDestination, BotSecretName, + MODEL_TIERS, SecretAskPurpose, SecretHttpRequest, } from "@rakazo/contracts"; import { z } from "zod"; +export const SUBAGENT_MODEL_TIER_DESCRIPTION = + "Optional model tier: fast for quick simple work, code for programming, standard otherwise. Uses the model this space maps to the tier, or your own model when it maps none. Ignored when model_provider and model_id are set."; + export const DELEGATION_TOOL_NAMES = new Set([ "run_subagent", "spawn_bot", @@ -733,6 +737,11 @@ export const builtinAgentTools: ConnectorTool[] = [ type: "string", description: "Optional connected model ID. Set together with model_provider.", }, + model_tier: { + type: "string", + enum: [...MODEL_TIERS], + description: SUBAGENT_MODEL_TIER_DESCRIPTION, + }, }, required: ["name", "task"], }, diff --git a/packages/adapters/src/executor.test.ts b/packages/adapters/src/executor.test.ts index 5fd1d90..b5c8d10 100644 --- a/packages/adapters/src/executor.test.ts +++ b/packages/adapters/src/executor.test.ts @@ -1676,4 +1676,113 @@ description: Prepare standup notes thinkingLevel: "high", }); }); + + // One router connection serves several local models; tiers pick among them. + function tierPrisma(bot: Record | null, tierRow: unknown) { + const plaintext = serializeModelSecret({ + kind: "openai_compatible", + baseUrl: "http://127.0.0.1:8080/v1", + reasoning: true, + }); + const findUnique = vi.fn(async () => tierRow); + const prisma = { + bot: { findFirst: vi.fn(async () => bot) }, + spaceModelTier: { findUnique }, + spaceModelPreference: { + findFirst: vi.fn(async (args: { where: { credential?: { provider?: string } } }) => + args.where.credential?.provider === "xai" + ? modelPreference({ + provider: "xai", + secretId: "secret-xai", + modelId: "grok-4.6", + isDefault: false, + }) + : null, + ), + }, + userModelCredential: { findFirst: vi.fn(async () => null) }, + deploymentSettings: { findUnique: vi.fn(async () => null) }, + secret: { + findFirst: vi.fn(async (args: { where: { id: string } }) => + args.where.id === "secret-router" ? { id: "secret-router", ciphertext: plaintext } : null, + ), + findUnique: vi.fn(async () => null), + }, + } as unknown as PrismaClient; + const executor = createRunExecutor({ + prisma, + secretStore: { load: vi.fn(() => plaintext), put: vi.fn() }, + } as unknown as Parameters[0]); + return { executor, findUnique }; + } + + const routerTier = { + tier: "code", + credentialId: "credential-router", + modelId: "router-code", + credential: { + id: "credential-router", + userId: "user-1", + provider: "openai-compatible", + label: "Router", + secretId: "secret-router", + supportsImages: false, + createdAt: new Date(0), + updatedAt: new Date(0), + preferences: [{ isDefault: true, modelId: "router-default" }], + }, + }; + + it("runs a bot on the model its space maps to the bot's tier", async () => { + const { executor, findUnique } = tierPrisma( + { modelProvider: "xai", modelId: "grok-4.6", thinkingLevel: "high", modelTier: "code" }, + routerTier, + ); + + const model = await executor.resolveModel({ + userId: "user-1", + spaceId: "ws-1", + botId: "bot-1", + }); + + expect(model).toMatchObject({ + provider: "openai-compatible", + id: "router-code", + baseUrl: "http://127.0.0.1:8080/v1", + reasoning: true, + thinkingLevel: "high", + }); + expect(findUnique).toHaveBeenCalledWith( + expect.objectContaining({ + where: { spaceId_userId_tier: { spaceId: "ws-1", userId: "user-1", tier: "code" } }, + }), + ); + }); + + it("keeps the bot's own model when its tier is not mapped", async () => { + const { executor } = tierPrisma( + { modelProvider: "xai", modelId: "grok-4.6", thinkingLevel: null, modelTier: "code" }, + null, + ); + + await expect( + executor.resolveModel({ userId: "user-1", spaceId: "ws-1", botId: "bot-1" }), + ).resolves.toMatchObject({ provider: "xai", id: "grok-4.6" }); + }); + + it("resolves a helper tier and reports an unmapped one as null", async () => { + const mapped = tierPrisma(null, routerTier); + await expect( + mapped.executor.resolveTierModel({ userId: "user-1", spaceId: "ws-1" }, "code"), + ).resolves.toMatchObject({ + provider: "openai-compatible", + id: "router-code", + baseUrl: "http://127.0.0.1:8080/v1", + }); + + const unmapped = tierPrisma(null, null); + await expect( + unmapped.executor.resolveTierModel({ userId: "user-1", spaceId: "ws-1" }, "fast"), + ).resolves.toBeNull(); + }); }); diff --git a/packages/adapters/src/executor.ts b/packages/adapters/src/executor.ts index 399906e..2df3b95 100644 --- a/packages/adapters/src/executor.ts +++ b/packages/adapters/src/executor.ts @@ -25,7 +25,7 @@ import { routineWakeupJob, runContinueJob, } from "@rakazo/adapter-kit"; -import type { RunStatus } from "@rakazo/contracts"; +import type { ModelTier, RunStatus } from "@rakazo/contracts"; import { ATTACHMENT_MAX_BYTES, BOT_DESCRIPTION_MAX_LENGTH, @@ -91,6 +91,7 @@ import { effectiveMemoryScope, findDefaultModelCredential, findModelCredential, + findSpaceModelTier, InvalidSpaceNameError, isTooManyDatabaseConnections, listSpaceUserBotIds, @@ -634,27 +635,13 @@ export function createRunExecutor(deps: ExecutorDeps) { const web = deps.web ?? createWebProvider(); const browser = deps.browser ?? createBrowserProvider(undefined, { sandbox: deps.sandbox }); const cloudAgent = deps.cloudAgent; - const resolveConnectedModel = async ( + const connectedModel = async ( scope: { userId: string; spaceId: string }, + credential: Parameters[3], provider: string, modelId: string, registerSecrets?: (values: string[]) => void, ): Promise => { - const validationError = await validateConnectedModelChoice( - deps.prisma, - scope, - provider, - modelId, - ); - if (validationError) throw new Error(validationError); - const credential = await findModelCredential(deps.prisma, scope, provider, modelId); - if (!credential) throw new Error("Connect that model provider first"); - // Free-form selections must keep the preference that owns this modelId. A - // intervening delete/change can make findModelCredential fall back to another - // same-provider credential; reject that mismatch instead of mixing baseUrl. - if (!isCatalogModelChoice(provider, modelId) && credential.defaultModel !== modelId) { - throw new Error("Unknown model for that provider"); - } const resolved = await resolveModelKey( deps, scope.userId, @@ -680,8 +667,48 @@ export function createRunExecutor(deps: ExecutorDeps) { : undefined, }; }; + const resolveConnectedModel = async ( + scope: { userId: string; spaceId: string }, + provider: string, + modelId: string, + registerSecrets?: (values: string[]) => void, + ): Promise => { + const validationError = await validateConnectedModelChoice( + deps.prisma, + scope, + provider, + modelId, + ); + if (validationError) throw new Error(validationError); + const credential = await findModelCredential(deps.prisma, scope, provider, modelId); + if (!credential) throw new Error("Connect that model provider first"); + // Free-form selections must keep the preference that owns this modelId. A + // intervening delete/change can make findModelCredential fall back to another + // same-provider credential; reject that mismatch instead of mixing baseUrl. + if (!isCatalogModelChoice(provider, modelId) && credential.defaultModel !== modelId) { + throw new Error("Unknown model for that provider"); + } + return connectedModel(scope, credential, provider, modelId, registerSecrets); + }; + /** The model a space maps this tier to; null when the tier is not mapped. */ + const resolveTierModel = async ( + scope: { userId: string; spaceId: string }, + tier: ModelTier, + registerSecrets?: (values: string[]) => void, + ): Promise => { + const mapped = await findSpaceModelTier(deps.prisma, scope, tier); + if (!mapped) return null; + return connectedModel( + scope, + mapped.credential, + mapped.credential.provider, + mapped.modelId, + registerSecrets, + ); + }; return { resolveConnectedModel, + resolveTierModel, async resolveModel(scope: { userId: string; spaceId: string; @@ -694,19 +721,21 @@ export function createRunExecutor(deps: ExecutorDeps) { userId: scope.userId, spaceId: scope.spaceId, }, - select: { modelProvider: true, modelId: true, thinkingLevel: true }, + select: { modelProvider: true, modelId: true, thinkingLevel: true, modelTier: true }, }) : null; const hasOverride = Boolean(override?.modelProvider && override.modelId); - const [overrideCredential, defaultCredential, settings] = await Promise.all([ + const [overrideCredential, tierModel, defaultCredential, settings] = await Promise.all([ hasOverride ? findModelCredential(deps.prisma, scope, override!.modelProvider!, override!.modelId) : Promise.resolve(null), + override?.modelTier ? findSpaceModelTier(deps.prisma, scope, override.modelTier) : null, findDefaultModelCredential(deps.prisma, scope), deps.prisma.deploymentSettings.findUnique({ where: { id: "default" } }), ]); const selected = selectConfiguredModel({ bot: override, + tierModel, overrideCredential, defaultCredential, settings, @@ -1081,6 +1110,9 @@ export function createRunExecutor(deps: ExecutorDeps) { hasModelOverride && bot.modelProvider ? await findModelCredential(deps.prisma, run, bot.modelProvider, bot.modelId) : null; + const tierModel = bot.modelTier + ? await findSpaceModelTier(deps.prisma, run, bot.modelTier) + : null; runAbortController = new AbortController(); if (!leaseValid) runAbortController.abort(); if (deps.shutdownSignal?.aborted) runAbortController.abort(deps.shutdownSignal.reason); @@ -1245,6 +1277,7 @@ export function createRunExecutor(deps: ExecutorDeps) { const runtimeFallback = runtimeFallbackModel(deps.runtime); const selected = selectConfiguredModel({ bot, + tierModel, overrideCredential, defaultCredential, settings, @@ -3613,6 +3646,9 @@ export function createRunExecutor(deps: ExecutorDeps) { resolveConnectedModel(run, provider, modelId, (values) => runSecrets.push(...values), ), + resolveModelTier: scripted + ? undefined + : (tier) => resolveTierModel(run, tier, (values) => runSecrets.push(...values)), onToolCompleted: (completion) => appendToolCompletionAudit( deps, diff --git a/packages/adapters/src/model-selection.test.ts b/packages/adapters/src/model-selection.test.ts index ad65496..d9e85ca 100644 --- a/packages/adapters/src/model-selection.test.ts +++ b/packages/adapters/src/model-selection.test.ts @@ -12,6 +12,7 @@ function credential(provider: string, defaultModel: string | null) { provider, label: provider, secretId: `secret-${provider}`, + supportsImages: false, createdAt: new Date(0), updatedAt: new Date(0), isDefault: false, @@ -21,6 +22,9 @@ function credential(provider: string, defaultModel: string | null) { const spaceCredential = credential("space-provider", "space-model"); const overrideCredential = credential("bot-provider", "stored-model"); +// One router connection serves every tier model; its saved model is a different one. +const routerCredential = credential("openai-compatible", "router-default"); +const tierModel = { modelId: "router-code", credential: routerCredential }; const bot = { modelProvider: "bot-provider", modelId: "bot-model", thinkingLevel: "high" }; const defaults: SelectionInput = { bot: null, @@ -76,6 +80,46 @@ describe("configured model selection", () => { thinkingLevel: "high", }, }, + { + name: "runs a mapped tier ahead of the bot's own model", + input: { bot, overrideCredential, tierModel }, + expected: { + provider: "openai-compatible", + id: "router-code", + credential: routerCredential, + thinkingLevel: "high", + }, + }, + { + name: "runs a mapped tier ahead of the Space default", + input: { bot: { modelProvider: null, modelId: null, thinkingLevel: null }, tierModel }, + expected: { + provider: "openai-compatible", + id: "router-code", + credential: routerCredential, + thinkingLevel: null, + }, + }, + { + name: "keeps the bot's own model when its tier is unmapped", + input: { bot, overrideCredential, tierModel: null }, + expected: { + provider: "bot-provider", + id: "bot-model", + credential: overrideCredential, + thinkingLevel: "high", + }, + }, + { + name: "keeps the Space default when the tier is unmapped and nothing is pinned", + input: { bot: { modelProvider: null, modelId: null, thinkingLevel: "low" }, tierModel: null }, + expected: { + provider: "space-provider", + id: "space-model", + credential: spaceCredential, + thinkingLevel: "low", + }, + }, { name: "uses settings before deployment defaults without inventing a credential", input: { defaultCredential: null }, diff --git a/packages/adapters/src/model-selection.ts b/packages/adapters/src/model-selection.ts index ba84366..6359a50 100644 --- a/packages/adapters/src/model-selection.ts +++ b/packages/adapters/src/model-selection.ts @@ -48,12 +48,23 @@ export function selectConfiguredModel(input: { modelId: string | null; thinkingLevel: string | null; } | null; + /** The space's model for the bot's tier (findSpaceModelTier); null when the tier is unmapped. */ + tierModel?: { modelId: string; credential: NonNullable } | null; overrideCredential: ModelCredential; defaultCredential: ModelCredential; settings: { defaultModelProvider: string | null; defaultModelId: string | null } | null; deployment: { provider: string; model: string } | null; }) { - const { bot, overrideCredential, defaultCredential, settings, deployment } = input; + const { bot, tierModel, overrideCredential, defaultCredential, settings, deployment } = input; + // A mapped tier wins; an unmapped one leaves the usual choice below untouched. + if (tierModel) { + return { + provider: tierModel.credential.provider, + id: tierModel.modelId, + credential: tierModel.credential, + thinkingLevel: (bot?.thinkingLevel as AgentRunRequest["model"]["thinkingLevel"]) ?? null, + }; + } const hasOverride = Boolean(bot?.modelProvider && bot.modelId); // The override provider, model and credential must win together. const useOverride = Boolean(hasOverride && overrideCredential); diff --git a/packages/adapters/src/pi-runtime-thinking-level.test.ts b/packages/adapters/src/pi-runtime-thinking-level.test.ts index 74839a5..4c3f0a5 100644 --- a/packages/adapters/src/pi-runtime-thinking-level.test.ts +++ b/packages/adapters/src/pi-runtime-thinking-level.test.ts @@ -117,6 +117,10 @@ async function runWithModel( thinkingLevel?: "off" | "minimal" | "low" | "medium" | "high" | "xhigh" | "max" | null; }>, maxImagesPerPrompt?: number, + resolveModelTier?: (tier: "fast" | "standard" | "code") => Promise<{ + provider: string; + id: string; + } | null>, ) { const runtime = new PiAgentRuntime(); for await (const _event of runtime.run( @@ -131,6 +135,7 @@ async function runWithModel( model: { provider, id: modelId, thinkingLevel, maxImagesPerPrompt }, executeTool: vi.fn(async () => ({ ok: true })), resolveModel, + resolveModelTier, }, { operationId: "1", @@ -253,6 +258,98 @@ describe("Pi agent thinking level", () => { }, ); + it("runs a helper on the model its space maps to the requested tier", async () => { + fakeAgentState.subagentArgs = { name: "helper", task: "help", model_tier: "fast" }; + const resolveModelTier = vi.fn(async () => ({ provider: "test", id: "reasoning-model" })); + + await runWithModel( + "plain-model", + "test", + new AbortController().signal, + null, + undefined, + undefined, + resolveModelTier, + ); + + expect(resolveModelTier).toHaveBeenCalledWith("fast"); + expect(fakeAgentState.models.map((model) => model.id)).toEqual([ + "plain-model", + "reasoning-model", + ]); + }); + + it("keeps a helper on the run's model when its tier is not mapped", async () => { + fakeAgentState.subagentArgs = { name: "helper", task: "help", model_tier: "code" }; + const resolveModelTier = vi.fn(async () => null); + + await runWithModel( + "plain-model", + "test", + new AbortController().signal, + null, + undefined, + undefined, + resolveModelTier, + ); + + expect(resolveModelTier).toHaveBeenCalledWith("code"); + expect(fakeAgentState.models.map((model) => model.id)).toEqual(["plain-model", "plain-model"]); + expect(fakeAgentState.lastSubagentResult).not.toMatchObject({ + details: { result: expect.stringContaining("Subagent failed") }, + }); + }); + + it("prefers an explicit helper model over a tier", async () => { + fakeAgentState.subagentArgs = { + name: "helper", + task: "help", + model_provider: "xai", + model_id: "grok-4.6", + model_tier: "fast", + }; + const resolveModel = vi.fn(async (provider: string, modelId: string) => ({ + provider, + id: modelId, + })); + const resolveModelTier = vi.fn(async () => ({ provider: "test", id: "reasoning-model" })); + + await runWithModel( + "plain-model", + "test", + new AbortController().signal, + null, + resolveModel, + undefined, + resolveModelTier, + ); + + expect(resolveModel).toHaveBeenCalledWith("xai", "grok-4.6"); + expect(resolveModelTier).not.toHaveBeenCalled(); + expect(fakeAgentState.models.map((model) => model.id)).toEqual(["plain-model", "grok-4.6"]); + }); + + it("rejects an unknown helper tier without starting the helper", async () => { + fakeAgentState.subagentArgs = { name: "helper", task: "help", model_tier: "huge" }; + const resolveModelTier = vi.fn(); + + await runWithModel( + "plain-model", + "test", + new AbortController().signal, + null, + undefined, + undefined, + resolveModelTier, + ); + + expect(resolveModelTier).not.toHaveBeenCalled(); + expect(fakeAgentState.models).toHaveLength(1); + expect(fakeAgentState.lastSubagentResult).toMatchObject({ + details: { result: "Subagent failed: model_tier must be one of: fast, standard, code" }, + }); + }); + it("rejects an incomplete per-call subagent model pair", async () => { fakeAgentState.subagentArgs = { name: "helper", diff --git a/packages/adapters/src/pi-runtime-tool-schema.test.ts b/packages/adapters/src/pi-runtime-tool-schema.test.ts index a7377c4..03b8e74 100644 --- a/packages/adapters/src/pi-runtime-tool-schema.test.ts +++ b/packages/adapters/src/pi-runtime-tool-schema.test.ts @@ -1,4 +1,6 @@ +import { validateToolArguments } from "@earendil-works/pi-ai"; import { describe, expect, it } from "vitest"; +import { builtinAgentTools } from "./builtin-tools.js"; import { parseConnectorToolArgs } from "./lazy-tool-catalog.js"; import { jsonSchemaParameters, parametersFor } from "./pi-runtime.js"; @@ -106,3 +108,33 @@ describe("jsonSchemaParameters", () => { ).not.toThrow(); }); }); + +describe("run_subagent parameters", () => { + const tool = builtinAgentTools.find((entry) => entry.name === "run_subagent")!; + const parameters = parametersFor(tool); + const call = (args: Record) => ({ + type: "toolCall" as const, + id: "call-1", + name: "run_subagent", + arguments: { name: "scout", task: "look", ...args }, + }); + + it("takes an optional model tier and only the known tiers", () => { + expect(validateToolArguments({ name: "run_subagent", parameters } as never, call({}))).toEqual({ + name: "scout", + task: "look", + }); + expect( + validateToolArguments( + { name: "run_subagent", parameters } as never, + call({ model_tier: "code" }), + ), + ).toMatchObject({ model_tier: "code" }); + expect(() => + validateToolArguments( + { name: "run_subagent", parameters } as never, + call({ model_tier: "huge" }), + ), + ).toThrow(/model_tier/); + }); +}); diff --git a/packages/adapters/src/pi-runtime.ts b/packages/adapters/src/pi-runtime.ts index 13f8f97..b8f0ea3 100644 --- a/packages/adapters/src/pi-runtime.ts +++ b/packages/adapters/src/pi-runtime.ts @@ -14,9 +14,14 @@ import type { AgentToolCompletion, ConnectorTool, } from "@rakazo/adapter-kit"; +import { MODEL_TIERS, ModelTierSchema } from "@rakazo/contracts"; import { getLogger } from "@rakazo/logging"; import { isToolPauseResult } from "./approval-effect.js"; -import { builtinAgentTools, DELEGATION_TOOL_NAMES } from "./builtin-tools.js"; +import { + builtinAgentTools, + DELEGATION_TOOL_NAMES, + SUBAGENT_MODEL_TIER_DESCRIPTION, +} from "./builtin-tools.js"; import { normalizeOpenAiToolParameters, openAiToolParametersNeedNormalization, @@ -648,6 +653,7 @@ function toAgentTool(tool: ConnectorTool, host: ToolHost, exposedName: string): instructions: raw.instructions ? String(raw.instructions) : "", model_provider: raw.model_provider ? String(raw.model_provider) : "", model_id: raw.model_id ? String(raw.model_id) : "", + ...(raw.model_tier ? { model_tier: String(raw.model_tier).trim().toLowerCase() } : {}), }; } if (tool.name === "spawn_bot") { @@ -824,6 +830,7 @@ async function executeSubagent(host: ToolHost, executionId: string, args: Record const requestedProvider = String(args.model_provider ?? "").trim(); const requestedModelId = String(args.model_id ?? "").trim(); + const requestedTier = String(args.model_tier ?? "").trim(); let requestModel = host.request.model; try { if (Boolean(requestedProvider) !== Boolean(requestedModelId)) { @@ -834,6 +841,11 @@ async function executeSubagent(host: ToolHost, executionId: string, args: Record throw new Error("Per-call subagent model selection is unavailable"); } requestModel = await host.request.resolveModel(requestedProvider, requestedModelId); + } else if (requestedTier) { + const tier = ModelTierSchema.safeParse(requestedTier); + if (!tier.success) throw new Error(`model_tier must be one of: ${MODEL_TIERS.join(", ")}`); + // An unmapped tier keeps the helper on this run's model. + requestModel = (await host.request.resolveModelTier?.(tier.data)) ?? requestModel; } } catch (error) { const message = sanitizeError(error instanceof Error ? error.message : String(error)); @@ -1060,6 +1072,12 @@ function builtinParameters(tool: ConnectorTool) { instructions: Type.Optional(Type.String()), model_provider: Type.Optional(Type.String()), model_id: Type.Optional(Type.String()), + model_tier: Type.Optional( + Type.Union( + MODEL_TIERS.map((tier) => Type.Literal(tier)), + { description: SUBAGENT_MODEL_TIER_DESCRIPTION }, + ), + ), }); } if (tool.name === "spawn_bot") { diff --git a/packages/contracts/src/domain.ts b/packages/contracts/src/domain.ts index 18fe1c7..2165f42 100644 --- a/packages/contracts/src/domain.ts +++ b/packages/contracts/src/domain.ts @@ -363,10 +363,19 @@ export const UpdateBotInput = z modelProvider: z.string().trim().min(1).max(80).nullable().optional(), modelId: z.string().trim().min(1).max(200).nullable().optional(), thinkingLevel: ThinkingLevelSchema.nullable().optional(), + modelTier: ModelTierSchema.nullable().optional(), teamChatAmbientEnabled: z.boolean().optional(), teamChatRules: z.string().max(TEAM_CHAT_RULES_MAX_LENGTH).optional(), }) .superRefine((value, ctx) => { + // A bot runs on a pinned model or a tier; setting one clears the other. + if (value.modelTier && (value.modelProvider || value.modelId)) { + ctx.addIssue({ + code: "custom", + message: "Choose a model or a tier, not both", + path: ["modelTier"], + }); + } const providerProvided = value.modelProvider !== undefined; const modelProvided = value.modelId !== undefined; if (!providerProvided && !modelProvided) return; diff --git a/packages/contracts/src/index.test.ts b/packages/contracts/src/index.test.ts index 6f15467..f082d79 100644 --- a/packages/contracts/src/index.test.ts +++ b/packages/contracts/src/index.test.ts @@ -20,6 +20,7 @@ import { ReorderBotsInput, RunActivityRowSchema, RunSchema, + SetModelTierInput, UpdateBotInput, UpdateGroupInput, } from "./index.js"; @@ -186,6 +187,44 @@ describe("contracts", () => { ).toBe(true); }); + it("lets a bot ask for a model tier and a space map or clear one", () => { + expect(UpdateBotInput.parse({ botId: "bot-1", modelTier: "code" }).modelTier).toBe("code"); + expect(UpdateBotInput.parse({ botId: "bot-1", modelTier: null }).modelTier).toBeNull(); + expect(UpdateBotInput.safeParse({ botId: "bot-1", modelTier: "huge" }).success).toBe(false); + expect( + UpdateBotInput.safeParse({ + botId: "bot-1", + modelTier: "code", + modelProvider: "xai", + modelId: "grok-4.6", + }).success, + ).toBe(false); + expect( + UpdateBotInput.safeParse({ + botId: "bot-1", + modelTier: "code", + modelProvider: null, + modelId: null, + }).success, + ).toBe(true); + + expect( + SetModelTierInput.parse({ + tier: "fast", + model: { credentialId: "router", modelId: " qwen3-4b " }, + }), + ).toEqual({ tier: "fast", model: { credentialId: "router", modelId: "qwen3-4b" } }); + expect(SetModelTierInput.parse({ tier: "code", model: null })).toEqual({ + tier: "code", + model: null, + }); + expect( + SetModelTierInput.safeParse({ tier: "code", model: { credentialId: "router", modelId: " " } }) + .success, + ).toBe(false); + expect(SetModelTierInput.safeParse({ tier: "large", model: null }).success).toBe(false); + }); + it("normalizes group names and rejects duplicate members", () => { expect(CreateGroupInput.parse({ name: " Draft team ", botIds: ["bot-1", "bot-2"] })).toEqual({ name: "Draft team", @@ -242,6 +281,8 @@ describe("contracts", () => { expect(appContract.agentSecrets.put).toBeTruthy(); expect(appContract.agentSecrets.remove).toBeTruthy(); expect(appContract.notifications.registerPush).toBeTruthy(); + expect(appContract.models.tiers.list).toBeTruthy(); + expect(appContract.models.tiers.set).toBeTruthy(); expect(ProductEventType.options).toContain("thread.message.created"); expect(ProductEventType.options).toContain("thread.cleared"); expect(ProductEventType.options).toContain("thread.subagent"); diff --git a/packages/contracts/src/index.ts b/packages/contracts/src/index.ts index bd5ffd8..e71bca9 100644 --- a/packages/contracts/src/index.ts +++ b/packages/contracts/src/index.ts @@ -8,6 +8,7 @@ export * from "./ids.js"; export * from "./integration-settings.js"; export * from "./local-settings.js"; export * from "./mcp.js"; +export * from "./model-tiers.js"; export * from "./openai-compatible-ui.js"; export * from "./project-goals.js"; export * from "./project-sessions.js"; diff --git a/packages/contracts/src/local-settings.js b/packages/contracts/src/local-settings.js index 4afaac6..3f85e59 100644 --- a/packages/contracts/src/local-settings.js +++ b/packages/contracts/src/local-settings.js @@ -15,6 +15,8 @@ const procedures = new Set([ "models/finishOAuth", "models/cancelOAuth", "models/setDefault", + "models/tiers/list", + "models/tiers/set", "integrationSetup/get", "integrationSetup/save", ]); diff --git a/packages/contracts/src/model-tiers.ts b/packages/contracts/src/model-tiers.ts new file mode 100644 index 0000000..8856f58 --- /dev/null +++ b/packages/contracts/src/model-tiers.ts @@ -0,0 +1,24 @@ +import * as z from "zod"; +import { ModelTierSchema } from "./domain.js"; +import { Id } from "./ids.js"; + +/** One tier of a space mapped to a model of a connected credential. */ +export const ModelTierMappingSchema = z.object({ + tier: ModelTierSchema, + credentialId: Id, + provider: z.string(), + modelId: z.string(), +}); +export type ModelTierMapping = z.infer; + +export const SetModelTierInput = z.object({ + tier: ModelTierSchema, + /** null removes the mapping, so bots asking for this tier use their usual model again. */ + model: z + .object({ + credentialId: Id, + modelId: z.string().trim().min(1).max(200), + }) + .nullable(), +}); +export type SetModelTierInput = z.infer; diff --git a/packages/contracts/src/rpc.ts b/packages/contracts/src/rpc.ts index b3a7091..136d7e0 100644 --- a/packages/contracts/src/rpc.ts +++ b/packages/contracts/src/rpc.ts @@ -80,6 +80,7 @@ import { IntegrationProviderConfigSchema, IntegrationSetupStateSchema, } from "./integration-settings.js"; +import { ModelTierMappingSchema, SetModelTierInput } from "./model-tiers.js"; import { CreateProjectGoalInput, ProjectGoalSchema, @@ -211,10 +212,14 @@ export const appContract = { connect: oc.input(ModelConnectInputSchema).output(ModelCredentialSchema), probeOpenAiCompatible: oc .input( - z.object({ - baseUrl: z.string(), - apiKey: z.string().optional(), - }), + z.union([ + z.object({ + baseUrl: z.string(), + apiKey: z.string().optional(), + }), + // A connected server, probed with its stored address and key. + z.object({ credentialId: Id }), + ]), ) .output(z.object({ models: z.array(z.string()) })), beginOAuth: oc @@ -245,6 +250,10 @@ export const appContract = { setDefault: oc .input(z.object({ provider: z.string(), modelId: z.string() })) .output(z.object({ ok: z.literal(true) })), + tiers: { + list: oc.output(z.array(ModelTierMappingSchema)), + set: oc.input(SetModelTierInput).output(z.array(ModelTierMappingSchema)), + }, }, bots: { list: oc.output(z.array(BotSchema)), diff --git a/packages/db/src/index.ts b/packages/db/src/index.ts index 7a2cf76..da46350 100644 --- a/packages/db/src/index.ts +++ b/packages/db/src/index.ts @@ -11,6 +11,7 @@ export * from "./memory-config.js"; export * from "./messages.js"; export * from "./messaging.js"; export * from "./model-credentials.js"; +export * from "./model-tiers.js"; export * from "./project-goals.js"; export * from "./project-workspace.js"; export * from "./repos.js"; diff --git a/packages/db/src/model-tiers.test.ts b/packages/db/src/model-tiers.test.ts new file mode 100644 index 0000000..43ac63f --- /dev/null +++ b/packages/db/src/model-tiers.test.ts @@ -0,0 +1,108 @@ +import { describe, expect, it, vi } from "vitest"; +import type { PrismaClient } from "./client.js"; +import { findSpaceModelTier, listSpaceModelTiers, setSpaceModelTier } from "./model-tiers.js"; + +const scope = { userId: "user", spaceId: "space" }; +const credential = { + id: "router", + userId: "user", + provider: "openai-compatible", + label: "Router", + secretId: "router-secret", + supportsImages: true, + createdAt: new Date(0), + updatedAt: new Date(0), +}; + +describe("space model tiers", () => { + it("finds a tier within the user's space with its credential shaped for model resolution", async () => { + const findUnique = vi.fn().mockResolvedValue({ + tier: "code", + credentialId: "router", + modelId: "router-code", + credential: { ...credential, preferences: [{ isDefault: true, modelId: "router-default" }] }, + }); + const prisma = { spaceModelTier: { findUnique } } as unknown as PrismaClient; + + await expect(findSpaceModelTier(prisma, scope, "code")).resolves.toEqual({ + tier: "code", + credentialId: "router", + modelId: "router-code", + // The space's saved model stays the credential default, so its image support does not + // spread to the tier's model. + credential: { ...credential, isDefault: true, defaultModel: "router-default" }, + }); + expect(findUnique).toHaveBeenCalledWith({ + where: { spaceId_userId_tier: { spaceId: "space", userId: "user", tier: "code" } }, + include: { + credential: { + include: { + preferences: { + where: { spaceId: "space", userId: "user" }, + select: { isDefault: true, modelId: true }, + }, + }, + }, + }, + }); + }); + + it("reports an unmapped tier as null", async () => { + const prisma = { + spaceModelTier: { findUnique: vi.fn().mockResolvedValue(null) }, + } as unknown as PrismaClient; + + await expect(findSpaceModelTier(prisma, scope, "fast")).resolves.toBeNull(); + }); + + it("lists only the user's tiers in the space", async () => { + const findMany = vi.fn().mockResolvedValue([ + { + tier: "fast", + credentialId: "router", + modelId: "router-fast", + credential: { ...credential, preferences: [] }, + }, + ]); + const prisma = { spaceModelTier: { findMany } } as unknown as PrismaClient; + + await expect(listSpaceModelTiers(prisma, scope)).resolves.toEqual([ + { + tier: "fast", + credentialId: "router", + modelId: "router-fast", + credential: { ...credential, isDefault: false, defaultModel: null }, + }, + ]); + expect(findMany).toHaveBeenCalledWith( + expect.objectContaining({ where: { spaceId: "space", userId: "user" } }), + ); + }); + + it("maps a tier with an upsert and clears it with a scoped delete", async () => { + const upsert = vi.fn().mockResolvedValue({}); + const deleteMany = vi.fn().mockResolvedValue({ count: 1 }); + const prisma = { spaceModelTier: { upsert, deleteMany } } as unknown as PrismaClient; + + await setSpaceModelTier(prisma, scope, "code", { + credentialId: "router", + modelId: "router-code", + }); + expect(upsert).toHaveBeenCalledWith({ + where: { spaceId_userId_tier: { spaceId: "space", userId: "user", tier: "code" } }, + create: { + spaceId: "space", + userId: "user", + tier: "code", + credentialId: "router", + modelId: "router-code", + }, + update: { credentialId: "router", modelId: "router-code" }, + }); + + await setSpaceModelTier(prisma, scope, "code", null); + expect(deleteMany).toHaveBeenCalledWith({ + where: { spaceId: "space", userId: "user", tier: "code" }, + }); + }); +}); diff --git a/packages/db/src/model-tiers.ts b/packages/db/src/model-tiers.ts new file mode 100644 index 0000000..c6e5ba1 --- /dev/null +++ b/packages/db/src/model-tiers.ts @@ -0,0 +1,83 @@ +import type { Prisma, PrismaClient } from "./client.js"; + +type ModelTierScope = { userId: string; spaceId: string }; + +// The composite (credentialId, userId) foreign key keeps every mapped credential the user's own. +function withCredential(scope: ModelTierScope) { + return { + credential: { + include: { + preferences: { + where: { spaceId: scope.spaceId, userId: scope.userId }, + select: { isDefault: true, modelId: true }, + }, + }, + }, + } as const; +} + +type ModelTierRow = Prisma.SpaceModelTierGetPayload<{ + include: ReturnType; +}>; + +/** + * Shape the credential like findModelCredential does. defaultModel stays the space's saved + * model, not the tier's, so capabilities kept for that saved model never spread to tier models. + */ +function toModelTier(row: ModelTierRow) { + const { preferences, ...credential } = row.credential; + return { + tier: row.tier, + credentialId: row.credentialId, + modelId: row.modelId, + credential: { + ...credential, + isDefault: preferences[0]?.isDefault ?? false, + defaultModel: preferences[0]?.modelId ?? null, + }, + }; +} + +export async function listSpaceModelTiers( + prisma: Pick, + scope: ModelTierScope, +) { + const rows = await prisma.spaceModelTier.findMany({ + where: { spaceId: scope.spaceId, userId: scope.userId }, + include: withCredential(scope), + orderBy: { createdAt: "asc" }, + }); + return rows.map(toModelTier); +} + +/** The connected model a space maps this tier to, or null when the tier is not mapped. */ +export async function findSpaceModelTier( + prisma: Pick, + scope: ModelTierScope, + tier: string, +) { + const row = await prisma.spaceModelTier.findUnique({ + where: { spaceId_userId_tier: { spaceId: scope.spaceId, userId: scope.userId, tier } }, + include: withCredential(scope), + }); + return row ? toModelTier(row) : null; +} + +export async function setSpaceModelTier( + prisma: Pick, + scope: ModelTierScope, + tier: string, + model: { credentialId: string; modelId: string } | null, +) { + if (!model) { + await prisma.spaceModelTier.deleteMany({ + where: { spaceId: scope.spaceId, userId: scope.userId, tier }, + }); + return; + } + await prisma.spaceModelTier.upsert({ + where: { spaceId_userId_tier: { spaceId: scope.spaceId, userId: scope.userId, tier } }, + create: { spaceId: scope.spaceId, userId: scope.userId, tier, ...model }, + update: model, + }); +} diff --git a/packages/db/src/repos.ts b/packages/db/src/repos.ts index c8d1c0c..b60d28b 100644 --- a/packages/db/src/repos.ts +++ b/packages/db/src/repos.ts @@ -404,6 +404,7 @@ export function createRepos(prisma: PrismaClient) { modelProvider?: string | null; modelId?: string | null; thinkingLevel?: string | null; + modelTier?: string | null; initialMessage?: { role: "user" | "bot" | "system"; blocks: MessageBlock[]; @@ -421,6 +422,7 @@ export function createRepos(prisma: PrismaClient) { let modelProvider = input.modelProvider ?? null; let modelId = input.modelId ?? null; let thinkingLevel = input.thinkingLevel ?? null; + let modelTier = input.modelTier ?? null; if (input.parentBotId) { const parent = await prisma.bot.findFirst({ where: { @@ -430,9 +432,10 @@ export function createRepos(prisma: PrismaClient) { }, }); if (!parent) throw new IsolationError(); - if (!modelId) { + if (!modelId && !modelTier) { modelProvider = parent.modelProvider ?? null; modelId = parent.modelId ?? null; + modelTier = parent.modelTier ?? null; } if (thinkingLevel == null) thinkingLevel = parent.thinkingLevel ?? null; } @@ -474,6 +477,7 @@ export function createRepos(prisma: PrismaClient) { modelProvider, modelId, thinkingLevel, + modelTier, }, }); const thread = await tx.thread.create({ diff --git a/packages/db/src/spaces.ts b/packages/db/src/spaces.ts index ab5429e..418e61b 100644 --- a/packages/db/src/spaces.ts +++ b/packages/db/src/spaces.ts @@ -69,6 +69,7 @@ type SpaceClient = Pick< | "space" | "spaceMember" | "spaceModelPreference" + | "spaceModelTier" | "spaceVoicePreference" | "memoryDocument" | "notificationPreference" @@ -132,10 +133,13 @@ async function copyProviderPreferences( prisma: SpaceClient, input: { sourceSpaceId: string; targetSpaceId: string; userId: string; createdAt: Date }, ): Promise { - const [modelPreferences, voicePreferences] = await Promise.all([ + const [modelPreferences, modelTiers, voicePreferences] = await Promise.all([ prisma.spaceModelPreference.findMany({ where: { spaceId: input.sourceSpaceId, userId: input.userId }, }), + prisma.spaceModelTier.findMany({ + where: { spaceId: input.sourceSpaceId, userId: input.userId }, + }), prisma.spaceVoicePreference.findMany({ where: { spaceId: input.sourceSpaceId, userId: input.userId }, }), @@ -154,6 +158,19 @@ async function copyProviderPreferences( })), }) : Promise.resolve(), + modelTiers.length + ? prisma.spaceModelTier.createMany({ + data: modelTiers.map((tier) => ({ + spaceId: input.targetSpaceId, + userId: input.userId, + tier: tier.tier, + credentialId: tier.credentialId, + modelId: tier.modelId, + createdAt: input.createdAt, + updatedAt: input.createdAt, + })), + }) + : Promise.resolve(), voicePreferences.length ? prisma.spaceVoicePreference.createMany({ data: voicePreferences.map((preference) => ({ diff --git a/packages/testkit/src/authorization.test.ts b/packages/testkit/src/authorization.test.ts index 6fe0dd9..5db92e4 100644 --- a/packages/testkit/src/authorization.test.ts +++ b/packages/testkit/src/authorization.test.ts @@ -74,6 +74,9 @@ describeWithDatabase("API authorization and resource isolation", () => { ["models/finishOAuth", { loginId: "missing-login" }], ["models/cancelOAuth", { loginId: "missing-login" }], ["models/setDefault", { provider: "test", modelId: "test/model" }], + ["models/probeOpenAiCompatible", { credentialId: "missing-credential" }], + ["models/tiers/list"], + ["models/tiers/set", { tier: "fast", model: null }], ["spaces/list"], ["spaces/create", { name: "Nope" }], ["spaces/remove", { spaceId: "missing-space" }], @@ -395,7 +398,24 @@ describeWithDatabase("API authorization and resource isolation", () => { answer: "mixed-resource attack", }); + const ownerModel = await rpc(app, owner, "models/connect", { + provider: "scripted", + apiKey: "fake-owner-scripted-key", + }); + const ownerTiers = await rpc(app, owner, "models/tiers/set", { + tier: "code", + model: { credentialId: ownerModel.id, modelId: "scripted" }, + }); + expect(ownerTiers).toEqual([ + { tier: "code", credentialId: ownerModel.id, provider: "scripted", modelId: "scripted" }, + ]); + const resourceIdCalls = [ + [ + "models/tiers/set", + { tier: "fast", model: { credentialId: ownerModel.id, modelId: "scripted" } }, + ], + ["models/probeOpenAiCompatible", { credentialId: ownerModel.id }], ["routines/update", { routineId: ownerRoutine.id, name: "Stolen Routine" }], ["routines/remove", { routineId: ownerRoutine.id }], ["routines/testRun", { routineId: ownerRoutine.id }], @@ -422,6 +442,8 @@ describeWithDatabase("API authorization and resource isolation", () => { resourceIdCalls.map(([procedure, input]) => expectDenied(app, intruder, procedure, input)), ); + expect(await rpc(app, intruder, "models/tiers/list")).toEqual([]); + expect(await rpc(app, owner, "models/tiers/list")).toEqual(ownerTiers); expect(await rpc(app, intruder, "memory/list", { botId: ownerBot.id })).toEqual([]); expect(await rpc(app, intruder, "memory/exportMarkdown", { botId: ownerBot.id })).toBe( "", @@ -1197,6 +1219,64 @@ describeWithDatabase("API authorization and resource isolation", () => { ).resolves.toBe(32); }); + it("maps model tiers per Space and carries them into a new Space", async () => { + const cookie = await signup(app, `model-tiers-${stamp}@rakazo.test`, "Model Tiers"); + const model = await rpc(app, cookie, "models/connect", { + provider: "scripted", + apiKey: "fake-tier-scripted-key", + }); + const mapped = [ + { tier: "fast", credentialId: model.id, provider: "scripted", modelId: "scripted" }, + ]; + await expect( + rpc(app, cookie, "models/tiers/set", { + tier: "fast", + model: { credentialId: model.id, modelId: "scripted" }, + }), + ).resolves.toEqual(mapped); + // Catalog providers only run catalog models; free-form ids belong to OpenAI-compatible servers. + const unknownModel = await raw(app, cookie, "models/tiers/set", { + tier: "code", + model: { credentialId: model.id, modelId: "not-a-catalog-model" }, + }); + expect(unknownModel.status).toBe(400); + + const created = await rpc(app, cookie, "spaces/create", { name: "Tier copy" }); + await expect(rpc(app, cookie, "models/tiers/list", {}, created.id)).resolves.toEqual(mapped); + + await rpc(app, cookie, "models/tiers/set", { tier: "fast", model: null }, created.id); + await expect(rpc(app, cookie, "models/tiers/list", {}, created.id)).resolves.toEqual([]); + await expect(rpc(app, cookie, "models/tiers/list")).resolves.toEqual(mapped); + + const bot = await rpc<{ id: string; modelTier: string | null }>( + app, + cookie, + "bots/create", + botInput("Tiered"), + ); + await expect( + rpc(app, cookie, "bots/update", { botId: bot.id, modelTier: "fast" }), + ).resolves.toMatchObject({ id: bot.id, modelTier: "fast" }); + await expect(rpc(app, cookie, "bots/duplicate", { botId: bot.id })).resolves.toMatchObject({ + modelTier: "fast", + }); + // A pinned model and a tier replace each other, so a client that only knows models + // (mobile) never leaves a tier silently overriding the model it picked. + await expect( + rpc(app, cookie, "bots/update", { + botId: bot.id, + modelProvider: "scripted", + modelId: "scripted", + }), + ).resolves.toMatchObject({ modelProvider: "scripted", modelId: "scripted", modelTier: null }); + await expect( + rpc(app, cookie, "bots/update", { botId: bot.id, modelTier: "code" }), + ).resolves.toMatchObject({ modelProvider: null, modelId: null, modelTier: "code" }); + await expect( + rpc(app, cookie, "bots/update", { botId: bot.id, modelTier: null }), + ).resolves.toMatchObject({ id: bot.id, modelTier: null }); + }); + it("reuses provider credentials and copies their selections into a new Space", async () => { const cookie = await signup(app, `space-provider-copy-${stamp}@rakazo.test`, "Provider Copy"); const actor = await rpc(app, cookie, "me"); diff --git a/packages/testkit/src/model-emulator.ts b/packages/testkit/src/model-emulator.ts index 1db8ebf..b55843b 100644 --- a/packages/testkit/src/model-emulator.ts +++ b/packages/testkit/src/model-emulator.ts @@ -44,9 +44,12 @@ export interface ModelEmulatorStep { export async function startModelEmulator(options: { steps: ModelEmulatorStep[]; modelId?: string; + /** Further ids this endpoint routes, like a router serving several local models. */ + routedModelIds?: string[]; apiKey?: string; }) { const modelId = options.modelId ?? "offline-fixture"; + const routedModelIds = new Set([modelId, ...(options.routedModelIds ?? [])]); const requests: ModelEmulatorRequest[] = []; const failures: Error[] = []; let stepIndex = 0; @@ -64,7 +67,7 @@ export async function startModelEmulator(options: { } const body = JSON.parse(Buffer.concat(chunks).toString("utf8")) as ModelEmulatorRequest; requests.push(body); - assert.equal(body.model, modelId, "Model routing changed"); + assert.ok(routedModelIds.has(body.model), "Model routing changed"); assert.equal(body.stream, true, "Expected streaming model request"); assert.ok(Array.isArray(body.messages), "Expected model messages"); const step = options.steps[stepIndex++]; @@ -87,7 +90,7 @@ export async function startModelEmulator(options: { id: `fixture-${stepIndex}`, object: "chat.completion.chunk", created: 0, - model: modelId, + model: body.model, choices: [{ index: 0, delta, finish_reason: finishReason }], })}\n\n`, ); diff --git a/packages/testkit/src/pi-offline.postgres.test.ts b/packages/testkit/src/pi-offline.postgres.test.ts index 3a38a3b..563a781 100644 --- a/packages/testkit/src/pi-offline.postgres.test.ts +++ b/packages/testkit/src/pi-offline.postgres.test.ts @@ -140,6 +140,154 @@ describe.skipIf(!databaseAvailable)("offline Pi product journey", () => { } } }, 30_000); + + it("runs a bot's tier and a helper's tier on models of one router connection", async () => { + const fixtureKey = "offline-router-fixture-key"; + const lastToolResult = (request: { messages: Array<{ role: string; content?: unknown }> }) => + String(request.messages.findLast((message) => message.role === "tool")?.content); + const model = await startModelEmulator({ + apiKey: fixtureKey, + modelId: "router-standard", + routedModelIds: ["router-fast", "router-code"], + steps: [ + { + expect(request) { + expect(request.model).toBe("router-code"); + const helper = request.tools?.find((tool) => tool.function.name === "run_subagent"); + expect(JSON.stringify(helper?.function.parameters)).toContain("model_tier"); + }, + response: { + type: "tool", + id: "tier-helper", + name: "run_subagent", + arguments: { name: "scout", task: "Count to three.", model_tier: "fast" }, + }, + }, + { + expect(request) { + expect(request.model).toBe("router-fast"); + expect(JSON.stringify(request.messages)).toContain("Count to three."); + }, + response: { type: "text", text: "one two three" }, + }, + { + expect(request) { + expect(request.model).toBe("router-code"); + expect(lastToolResult(request)).toContain("one two three"); + }, + response: { type: "text", text: "The helper counted." }, + }, + { + // With its tier cleared the bot is back on the connection's saved model. + expect(request) { + expect(request.model).toBe("router-standard"); + }, + response: { type: "text", text: "Back on the saved model." }, + }, + ], + }); + const dataDir = await mkdtemp(path.join(tmpdir(), "rakazo-offline-tiers-")); + let stop: (() => Promise) | undefined; + try { + const { createApp } = await import("../../../apps/api/src/app.ts"); + const handles = await createApp({ + databaseUrl: process.env.DATABASE_URL!, + realtimeDatabaseUrl: process.env.DATABASE_URL!, + authUrl: fixtureOrigin, + webOrigin: fixtureOrigin, + dataDir, + sandboxProvider: "fake", + agentRuntime: "pi", + wakeupDriver: "memory", + signupsEnabled: "true", + composio: new ComposioEmulator(), + encryptionKey: "offline-model-fixture-encryption-key", + }); + stop = handles.stop; + const signup = await handles.app.request("/api/auth/sign-up/email", { + method: "POST", + headers: { "content-type": "application/json", origin: fixtureOrigin }, + body: JSON.stringify({ + email: `offline-tiers-${randomUUID()}@rakazo.test`, + password: "password12", + name: "Offline tiers", + }), + }); + expect(signup.status).toBeLessThan(400); + const cookie = sessionCookieHeader(signup); + const router = await rpc<{ id: string }>(handles.app, cookie, "models/connect", { + provider: model.model.provider, + modelId: model.model.id, + baseUrl: model.baseUrl, + apiKey: fixtureKey, + }); + for (const [tier, modelId] of [ + ["code", "router-code"], + ["fast", "router-fast"], + ]) { + await rpc(handles.app, cookie, "models/tiers/set", { + tier, + model: { credentialId: router.id, modelId }, + }); + } + const bot = await rpc<{ id: string }>(handles.app, cookie, "bots/create", { + name: "Tier fixture", + title: "", + description: "", + instructions: "Complete the task.", + notifyOnFinish: false, + }); + await rpc(handles.app, cookie, "bots/update", { botId: bot.id, modelTier: "code" }); + const completedRun = async (text: string) => { + const sent = await rpc<{ runId: string }>(handles.app, cookie, "threads/send", { + botId: bot.id, + text, + }); + await expect + .poll( + async () => { + const run = await handles.prisma.run.findUnique({ + where: { id: sent.runId }, + select: { status: true }, + }); + if (run?.status === "failed") model.assertComplete(); + return run?.status; + }, + { timeout: 15_000, interval: 100 }, + ) + .toBe("completed"); + return handles.prisma.run.findUniqueOrThrow({ + where: { id: sent.runId }, + select: { modelProvider: true, modelId: true }, + }); + }; + + await expect(completedRun("Ask a fast helper to count.")).resolves.toEqual({ + modelProvider: "openai-compatible", + modelId: "router-code", + }); + await rpc(handles.app, cookie, "models/tiers/set", { tier: "code", model: null }); + await expect(completedRun("Answer directly.")).resolves.toEqual({ + modelProvider: "openai-compatible", + modelId: "router-standard", + }); + model.assertComplete(); + const thread = await rpc<{ messages: unknown[] }>(handles.app, cookie, "threads/get", { + botId: bot.id, + }); + const messages = JSON.stringify(thread.messages); + expect(messages).toContain("The helper counted."); + expect(messages).toContain("Back on the saved model."); + expect(messages).not.toContain(fixtureKey); + } finally { + try { + await stop?.(); + } finally { + await model.close(); + await rm(dataDir, { recursive: true, force: true }); + } + } + }, 45_000); }); async function rpc(