From f96cd8b9151a8a980f72cb3bc7d83852f762a530 Mon Sep 17 00:00:00 2001 From: Tanish Manku Date: Mon, 24 Aug 2026 21:42:14 -0700 Subject: [PATCH] mcp: list_models tool and the fake's model catalogues MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit fake-upstream/src/catalogue.ts ports infra/fake-storyteller-web's catalogue: the real model ids, creators and display names for image, video, audio, mesh and splat, with one representative capability profile per kind. The image and video endpoints' own `provider` filter (OmniGen*ModelsProvider: "artcraft" — the default — or "all") is mirrored, including a 400 for other values; audio, mesh and splat declare no filter and list everything. All five catalogues validate against the spec. src/mcp/tools/list-models.ts always asks upstream for "all" so the catalogue is complete, applies the caller's GenerationProvider filter client-side via the providers block, and passes each model's capability fields through verbatim — what lets a client build a valid estimate_cost call from one source of truth. Typecheck caught the contract mismatch on the way: the tool first sent GenerationProvider values as the endpoint's query, which the fake accepted and production would reject; the generated types refused it. Tests (5 tool + 8 fake): normalisation and pass-through, "all" query for image/video, client-side provider filtering for every kind, every kind served; spec-valid catalogues, upstream's default and enum, no-filter kinds listing everything. Server and handler suites list the tool under read:catalog. Claude-Session: https://claude.ai/code/session_01H3gfN2NCx3hBXSrZXYHUQ3 --- mcp/fake-upstream/src/catalogue.ts | 294 +++++++++++++++++++++++ mcp/fake-upstream/src/index.ts | 2 + mcp/fake-upstream/src/routes/omni_gen.ts | 33 +++ mcp/fake-upstream/test/omni_gen.test.ts | 69 ++++++ mcp/src/mcp/server.ts | 2 + mcp/src/mcp/tools/list-models.ts | 132 ++++++++++ mcp/test/mcp/handler.test.ts | 13 +- mcp/test/mcp/list-models.test.ts | 97 ++++++++ mcp/test/mcp/server.test.ts | 6 + 9 files changed, 645 insertions(+), 3 deletions(-) create mode 100644 mcp/fake-upstream/src/catalogue.ts create mode 100644 mcp/fake-upstream/src/routes/omni_gen.ts create mode 100644 mcp/fake-upstream/test/omni_gen.test.ts create mode 100644 mcp/src/mcp/tools/list-models.ts create mode 100644 mcp/test/mcp/list-models.test.ts diff --git a/mcp/fake-upstream/src/catalogue.ts b/mcp/fake-upstream/src/catalogue.ts new file mode 100644 index 0000000000..d11000bddd --- /dev/null +++ b/mcp/fake-upstream/src/catalogue.ts @@ -0,0 +1,294 @@ +/** + * The model catalogues served by `/v1/omni_gen/models/*` — ported from + * `infra/fake-storyteller-web` (generation/catalogue.ts). + * + * The real catalogues are hand-tuned per model in Rust. Reproducing every capability flag + * verbatim would rot immediately, so the fake keeps the real model ids, creators and display + * names — what a client shows and branches on — and applies one representative capability + * profile per media type. Optional fields are omitted rather than sent as null, as upstream does. + */ + +export type GenerationKind = "image" | "video" | "audio" | "mesh" | "splat"; + +export const GENERATION_KINDS: readonly GenerationKind[] = [ + "image", + "video", + "audio", + "mesh", + "splat", +]; + +const IMAGE_ASPECT_RATIOS = [ + "square_hd", + "square", + "tall_three_by_four", + "tall_nine_by_sixteen", + "wide_four_by_three", + "wide_sixteen_by_nine", +]; +const VIDEO_ASPECT_RATIOS = ["square", "tall_nine_by_sixteen", "wide_sixteen_by_nine"]; +const RESOLUTIONS = ["one_k", "two_k", "four_k"]; +const QUALITIES = ["low", "medium", "high"]; + +type ModelSeed = readonly [model: string, creator: string, fullName: string]; + +const IMAGE_MODELS: readonly ModelSeed[] = [ + ["flux_1_dev", "black_forest_labs", "FLUX.1 [dev]"], + ["flux_1_schnell", "black_forest_labs", "FLUX.1 [schnell]"], + ["flux_pro_1p1", "black_forest_labs", "FLUX1.1 [pro]"], + ["flux_pro_1p1_ultra", "black_forest_labs", "FLUX1.1 [pro] ultra"], + ["flux_2_lora_angles", "black_forest_labs", "FLUX.2 Angles"], + ["gpt_image_1", "open_ai", "GPT Image 1"], + ["gpt_image_1p5", "open_ai", "GPT Image 1.5"], + ["gpt_image_2", "open_ai", "GPT Image 2"], + ["grok_imagine_image", "grok", "Grok Imagine"], + ["grok_imagine_image_q", "grok", "Grok Imagine Quick"], + ["midjourney_7", "midjourney", "Midjourney v7"], + ["midjourney_7_niji", "midjourney", "Midjourney Niji v7"], + ["midjourney_8", "midjourney", "Midjourney v8"], + ["nano_banana", "google", "Nano Banana"], + ["nano_banana_2", "google", "Nano Banana 2"], + ["nano_banana_pro", "google", "Nano Banana Pro"], + ["qwen_edit_2511_angles", "alibaba", "Qwen Edit Angles"], + ["seedream_4", "bytedance", "Seedream 4"], + ["seedream_4p5", "bytedance", "Seedream 4.5"], + ["seedream_5_lite", "bytedance", "Seedream 5 Lite"], + ["seedream_5p0_pro", "bytedance", "Seedream 5 Pro"], + ["seedream_5p0_pro_u", "bytedance", "Seedream 5 Pro Ultra"], +]; + +const VIDEO_MODELS: readonly ModelSeed[] = [ + ["grok_video", "grok", "Grok Video"], + ["grok_imagine_video", "grok", "Grok Imagine Video"], + ["grok_imagine_video_1p5", "grok", "Grok Imagine Video 1.5"], + ["kling_1p6_pro", "kling", "Kling 1.6 Pro"], + ["kling_2p1_pro", "kling", "Kling 2.1 Pro"], + ["kling_2p1_master", "kling", "Kling 2.1 Master"], + ["kling_2p5_turbo_pro", "kling", "Kling 2.5 Turbo Pro"], + ["kling_2p6_pro", "kling", "Kling 2.6 Pro"], + ["kling_3p0_standard", "kling", "Kling 3.0 Standard"], + ["kling_3p0_pro", "kling", "Kling 3.0 Pro"], + ["happy_horse_1p0", "artcraft", "Happy Horse 1.0"], + ["seedance_1p0_lite", "bytedance", "Seedance 1.0 Lite"], + ["seedance_1p5_pro", "bytedance", "Seedance 1.5 Pro"], + ["seedance_2p0", "bytedance", "Seedance 2.0"], + ["seedance_2p0_fast", "bytedance", "Seedance 2.0 Fast"], + ["seedance_2p0_mini", "bytedance", "Seedance 2.0 Mini"], + ["sora_2", "open_ai", "Sora 2"], + ["sora_2_pro", "open_ai", "Sora 2 Pro"], + ["veo_2", "google", "Veo 2"], + ["veo_3", "google", "Veo 3"], + ["veo_3_fast", "google", "Veo 3 Fast"], + ["veo_3p1", "google", "Veo 3.1"], + ["veo_3p1_fast", "google", "Veo 3.1 Fast"], + ["veo_3p1_lite", "google", "Veo 3.1 Lite"], + ["vidu_q3", "vidu", "Vidu Q3"], + ["vidu_q3_turbo", "vidu", "Vidu Q3 Turbo"], +]; + +const AUDIO_MODELS: readonly ModelSeed[] = [ + ["suno_music", "suno", "Suno Music"], + ["suno_remix", "suno", "Suno Remix"], + ["suno_sounds", "suno", "Suno Sounds"], + ["suno_sample", "suno", "Suno Sample"], + ["seed_audio_1p0", "bytedance", "Seed Audio 1.0"], +]; + +const MESH_MODELS: readonly ModelSeed[] = [ + ["hunyuan_3d_2p0", "tencent", "Hunyuan3D 2.0"], + ["hunyuan_3d_2p1", "tencent", "Hunyuan3D 2.1"], + ["hunyuan_3d_3", "tencent", "Hunyuan3D 3"], + ["hunyuan_3d_3_sketch", "tencent", "Hunyuan3D 3 Sketch"], + ["hunyuan_3d_3p1_pro", "tencent", "Hunyuan3D 3.1 Pro"], + ["hunyuan_3d_3p1_rapid", "tencent", "Hunyuan3D 3.1 Rapid"], + ["tripo3d_h3p1", "tripo", "Tripo3D H3.1"], + ["meshy_v6", "meshy", "Meshy v6"], + ["rodin_2p5_fast", "deemos", "Rodin 2.5 Fast"], +]; + +const SPLAT_MODELS: readonly ModelSeed[] = [ + ["marble_0p1_mini", "world_labs", "Marble 0.1 Mini"], + ["marble_0p1_plus", "world_labs", "Marble 0.1 Plus"], + ["marble_1p0", "world_labs", "Marble 1.0"], + ["marble_1p0_draft", "world_labs", "Marble 1.0 Draft"], + ["marble_1p1", "world_labs", "Marble 1.1"], + ["marble_1p1_plus", "world_labs", "Marble 1.1 Plus"], + ["triposplat", "tripo", "TripoSplat"], +]; + +const SEEDS: Record = { + image: IMAGE_MODELS, + video: VIDEO_MODELS, + audio: AUDIO_MODELS, + mesh: MESH_MODELS, + splat: SPLAT_MODELS, +}; + +/** Which provider each model is routed through, for the `providers` block. */ +const PROVIDER_BY_CREATOR: Record = { + artcraft: "artcraft", + grok: "grok", + midjourney: "midjourney", + open_ai: "sora", + world_labs: "world_labs", +}; + +const EXTRA_INFO = "Served by artcraft-api-fake. Capabilities are representative, not exact."; + +const PROFILES: Record Record> = { + image: (seed) => ({ + ...baseModel(seed), + text_prompt_supported: true, + text_prompt_max_length: 2000, + image_refs_supported: true, + image_refs_max: 4, + has_fixed_editing_aspect_ratio: false, + aspect_ratio_options: IMAGE_ASPECT_RATIOS, + aspect_ratio_default: "square", + resolution_options: RESOLUTIONS, + resolution_default: "one_k", + quality_options: QUALITIES, + default_quality: "high", + batch_size_min: 1, + batch_size_max: 4, + batch_size_options: [1, 2, 3, 4], + batch_size_default: 1, + is_disabled: false, + }), + video: (seed) => ({ + ...baseModel(seed), + extra_info: EXTRA_INFO, + extra_info_short: "Fake backend", + text_to_video_supported: true, + text_prompt_supported: true, + text_prompt_max_length: 2000, + starting_keyframe_supported: true, + starting_keyframe_required: false, + ending_keyframe_supported: true, + image_references_supported: true, + image_references_max: 4, + character_references_supported: true, + character_references_max: 2, + show_generate_with_sound_toggle: true, + aspect_ratio_options: VIDEO_ASPECT_RATIOS, + aspect_ratio_default: "wide_sixteen_by_nine", + resolution_options: ["seven_twenty_p", "ten_eighty_p"], + resolution_default: "seven_twenty_p", + bitrate_options: ["normal", "high"], + bitrate_default: "normal", + quality_options: QUALITIES, + default_quality: "high", + duration_seconds_min: 4, + duration_seconds_max: 10, + duration_seconds_options: [4, 6, 8, 10], + duration_seconds_default: 6, + batch_size_min: 1, + batch_size_max: 2, + batch_size_options: [1, 2], + batch_size_default: 1, + is_disabled: false, + }), + audio: (seed) => ({ + ...baseModel(seed), + extra_info: EXTRA_INFO, + extra_info_short: "Fake backend", + text_prompt_supported: true, + style_prompt_supported: true, + audio_references_supported: true, + audio_references_max: 2, + image_references_supported: false, + keep_lyrics_supported: true, + instrumental_toggle_supported: true, + loopable_toggle_supported: true, + bpm_supported: true, + musical_key_supported: true, + sample_rate_hz_options: [44_100, 48_000], + sample_rate_hz_default: 44_100, + speed_supported: true, + volume_supported: true, + pitch_supported: true, + is_disabled: false, + }), + mesh: (seed) => ({ + ...baseModel(seed), + extra_info: EXTRA_INFO, + extra_info_short: "Fake backend", + text_prompt_supported: true, + image_input_supported: true, + sketch_input_supported: false, + multi_view_supported: true, + mesh_input_supported: false, + mesh_output_types: ["normal", "low_poly", "geometry"], + polygon_types: ["triangle", "quad"], + face_count_supported: true, + pbr_supported: true, + texture_toggle_supported: true, + texture_quality_supported: true, + geometry_quality_supported: true, + is_disabled: false, + }), + splat: (seed) => ({ + ...baseModel(seed), + extra_info: EXTRA_INFO, + extra_info_short: "Fake backend", + text_prompt_supported: true, + image_references_supported: true, + image_references_max: 1, + video_reference_supported: true, + panorama_supported: true, + disable_recaption_supported: true, + is_disabled: false, + }), +}; + +/** The endpoint's own `provider` filter (image and video only): OmniGen*ModelsProvider. */ +export type ModelsFilter = "artcraft" | "all"; + +/** Only the image and video endpoints declare the filter; the others list everything. */ +export function hasProviderFilter(kind: GenerationKind): boolean { + return kind === "image" || kind === "video"; +} + +/** + * The `/v1/omni_gen/models/{kind}` response. For image and video, upstream's filter is + * `artcraft` (only models Artcraft routes itself — the default) or `all`; audio, mesh and + * splat have no filter and always list every model. + */ +export function modelsResponse( + kind: GenerationKind, + filter?: ModelsFilter, +): Record { + const effective = hasProviderFilter(kind) ? (filter ?? "artcraft") : "all"; + const seeds = SEEDS[kind].filter( + (seed) => effective === "all" || providerFor(seed) === "artcraft", + ); + return { + success: true, + models: seeds.map((seed) => PROFILES[kind](seed)), + providers: providerBlock(seeds), + }; +} + +/** Whether a model id exists in any catalogue — the real API rejects unknown models. */ +export function isKnownModel(kind: GenerationKind, model: string): boolean { + return SEEDS[kind].some(([id]) => id === model); +} + +function baseModel([model, creator, fullName]: ModelSeed): Record { + return { model, model_creator: creator, full_name: fullName }; +} + +function providerFor([, creator]: ModelSeed): string { + return PROVIDER_BY_CREATOR[creator] ?? "fal"; +} + +function providerBlock(seeds: readonly ModelSeed[]): Record[] { + const byProvider = new Map(); + for (const seed of seeds) { + const provider = providerFor(seed); + const models = byProvider.get(provider) ?? []; + models.push({ model: seed[0] }); + byProvider.set(provider, models); + } + return [...byProvider].map(([provider, models]) => ({ provider, models })); +} diff --git a/mcp/fake-upstream/src/index.ts b/mcp/fake-upstream/src/index.ts index be63bddf22..1e76d7103e 100644 --- a/mcp/fake-upstream/src/index.ts +++ b/mcp/fake-upstream/src/index.ts @@ -1,6 +1,7 @@ import { Hono } from "hono"; import { billingRoutes } from "./routes/billing"; +import { omniGenRoutes } from "./routes/omni_gen"; import { sessionRoutes } from "./routes/session"; import { createStore, type FakeStore } from "./state"; @@ -18,6 +19,7 @@ export function createFakeUpstream(store: FakeStore = createStore()): Hono { app.get("/_status", (c) => c.json({ success: true, fake: true })); app.route("/", sessionRoutes(store)); app.route("/", billingRoutes(store)); + app.route("/", omniGenRoutes()); return app; } diff --git a/mcp/fake-upstream/src/routes/omni_gen.ts b/mcp/fake-upstream/src/routes/omni_gen.ts new file mode 100644 index 0000000000..46cce0b43b --- /dev/null +++ b/mcp/fake-upstream/src/routes/omni_gen.ts @@ -0,0 +1,33 @@ +import { Hono } from "hono"; + +import { GENERATION_KINDS, hasProviderFilter, modelsResponse } from "../catalogue"; + +/** + * Omni-gen catalogue routes — ported from `infra/fake-storyteller-web` (routes/omni_gen.ts). + * Models need no user upstream (a session only unlocks a few gated models); the fake serves the + * same list to everyone. Cost estimates arrive with the estimate_cost tool. + */ +export function omniGenRoutes(): Hono { + const app = new Hono(); + + for (const kind of GENERATION_KINDS) { + app.get(`/v1/omni_gen/models/${kind}`, (c) => { + // Only image and video declare the query; actix ignores unknown params elsewhere. + const provider = hasProviderFilter(kind) ? c.req.query("provider") : undefined; + if (provider !== undefined && provider !== "artcraft" && provider !== "all") { + return c.json( + { + success: false, + error_code: 400, + error_code_str: "BadInput", + message: `unknown provider filter: ${provider}`, + }, + 400, + ); + } + return c.json(modelsResponse(kind, provider)); + }); + } + + return app; +} diff --git a/mcp/fake-upstream/test/omni_gen.test.ts b/mcp/fake-upstream/test/omni_gen.test.ts new file mode 100644 index 0000000000..8654b20792 --- /dev/null +++ b/mcp/fake-upstream/test/omni_gen.test.ts @@ -0,0 +1,69 @@ +import { describe, expect, it } from "vitest"; + +import { expectValid, type SchemaName } from "../../test/helpers/contract"; +import { GENERATION_KINDS, isKnownModel } from "../src/catalogue"; +import { createFakeUpstream } from "../src/index"; + +const RESPONSE_SCHEMA: Record<(typeof GENERATION_KINDS)[number], SchemaName> = { + image: "OmniGenImageModelsResponse", + video: "OmniGenVideoModelsResponse", + audio: "OmniGenAudioModelsResponse", + mesh: "OmniGenMeshModelsResponse", + splat: "OmniGenSplatModelsResponse", +}; + +interface ModelsBody { + models: { model: string; full_name: string; model_creator: string }[]; + providers: { provider: string; models: { model: string }[] }[]; +} + +describe("model catalogues", () => { + it.each(GENERATION_KINDS)("serves a spec-shaped %s catalogue without a session", async (kind) => { + const response = await createFakeUpstream().request(`/v1/omni_gen/models/${kind}?provider=all`); + expect(response.status).toBe(200); + const body = await response.json(); + expectValid(RESPONSE_SCHEMA[kind], body); + expect(body.models.length).toBeGreaterThan(0); + expect(body.models.every((m) => m.full_name && m.model_creator)).toBe(true); + const listed = new Set(body.providers.flatMap((p) => p.models.map((m) => m.model))); + expect(body.models.every((m) => listed.has(m.model))).toBe(true); + }); + + it("defaults to the models Artcraft routes itself, like upstream", async () => { + const video = await ( + await createFakeUpstream().request("/v1/omni_gen/models/video") + ).json(); + expectValid("OmniGenVideoModelsResponse", video); + expect(video.models.map((m) => m.model)).toEqual(["happy_horse_1p0"]); + expect(video.providers.map((p) => p.provider)).toEqual(["artcraft"]); + const image = await ( + await createFakeUpstream().request("/v1/omni_gen/models/image?provider=artcraft") + ).json(); + expect(image.models).toEqual([]); + }); + + it("lists everything for kinds whose endpoint declares no filter", async () => { + const audio = await ( + await createFakeUpstream().request("/v1/omni_gen/models/audio") + ).json(); + expect(audio.models.length).toBe(5); + const ignored = await ( + await createFakeUpstream().request("/v1/omni_gen/models/audio?provider=midjourney") + ).json(); + expect(ignored.models.length).toBe(5); + }); + + it("rejects a provider value outside the endpoint's enum", async () => { + const response = await createFakeUpstream().request( + "/v1/omni_gen/models/image?provider=midjourney", + ); + expect(response.status).toBe(400); + expectValid("CommonWebError", await response.json()); + }); + + it("knows its own model ids and nothing else", () => { + expect(isKnownModel("video", "seedance_2p0")).toBe(true); + expect(isKnownModel("image", "seedance_2p0")).toBe(false); + expect(isKnownModel("image", "made_up_model")).toBe(false); + }); +}); diff --git a/mcp/src/mcp/server.ts b/mcp/src/mcp/server.ts index 1a57d04e2a..262e5c7b6f 100644 --- a/mcp/src/mcp/server.ts +++ b/mcp/src/mcp/server.ts @@ -12,6 +12,7 @@ import { hasScope, type Principal } from "../tokens/principal"; import type { UpstreamClient } from "../upstream/client"; import { getAccount } from "./tools/get-account"; import { getCreditBalance } from "./tools/get-credit-balance"; +import { listModels } from "./tools/list-models"; import { type ToolDefinition, ToolFailure, UpstreamSessionInvalid } from "./tools/types"; /** @@ -25,6 +26,7 @@ export const SERVER_INFO = { name: "artcraft", version: "0.0.0" } as const; export const TOOLS: readonly ToolDefinition[] = [ getAccount, getCreditBalance, + listModels, ]; export interface McpServerDeps { diff --git a/mcp/src/mcp/tools/list-models.ts b/mcp/src/mcp/tools/list-models.ts new file mode 100644 index 0000000000..16a4d5721c --- /dev/null +++ b/mcp/src/mcp/tools/list-models.ts @@ -0,0 +1,132 @@ +import { z } from "zod"; + +import { READ_ONLY_ANNOTATIONS, type ToolDefinition, unwrapUpstream } from "./types"; + +/** + * The model catalogue for one kind of generation, as upstream publishes it. Each model's + * capability fields (supported inputs, aspect/resolution/duration options and defaults, batch + * sizes…) are passed through under `capabilities` exactly as the API names them — that is + * what lets a client build a valid estimate_cost request without a second source of truth. + */ + +export const GENERATION_KINDS = ["image", "video", "audio", "mesh", "splat"] as const; +export type GenerationKind = (typeof GENERATION_KINDS)[number]; + +/** GenerationProvider in the spec: who a model is routed through. Filtered client-side. */ +export const PROVIDERS = ["artcraft", "fal", "grok", "midjourney", "sora", "world_labs"] as const; + +const inputSchema = { + kind: z + .enum(GENERATION_KINDS) + .describe("Which catalogue: image, video, audio, mesh (3D object), or splat (3D world)."), + provider: z + .enum(PROVIDERS) + .optional() + .describe("Only list models routed through this provider. Omit for every model."), +}; + +const outputSchema = { + kind: z.enum(GENERATION_KINDS), + models: z.array( + z.object({ + id: z.string().describe("The model id to pass to estimate_cost, e.g. seedance_2p0."), + name: z.string().describe("Display name."), + creator: z.string().describe("Who makes the model, e.g. bytedance."), + disabled: z.boolean().describe("True when the model is temporarily unavailable."), + capabilities: z + .record(z.string(), z.unknown()) + .describe( + "Every other field upstream publishes for the model, verbatim: *_supported flags, " + + "*_options lists and *_default values (aspect ratio, resolution, quality, duration, " + + "batch size…), maxima such as text_prompt_max_length, and extra_info notes.", + ), + }), + ), + providers: z + .array(z.object({ provider: z.string(), model_ids: z.array(z.string()) })) + .describe("Which models are routed through which provider."), +}; + +interface RawModel { + model: string; + full_name?: string | null; + model_creator?: string | null; + is_disabled?: boolean | null; + [capability: string]: unknown; +} + +interface RawCatalogue { + models: RawModel[]; + providers: { provider: string; models: { model: string }[] }[]; +} + +export const listModels: ToolDefinition = { + name: "list_models", + title: "List models", + description: + "Lists the generation models Artcraft offers for one kind — image, video, audio, mesh " + + "(3D object) or splat (3D world) — with each model's capabilities and option lists exactly " + + "as the API publishes them. Call it before estimate_cost to pick a valid model id and " + + "arguments, or when the user asks what models exist. Optionally filter by provider.", + requiredScope: "read:catalog", + inputSchema, + outputSchema, + annotations: { ...READ_ONLY_ANNOTATIONS, title: "List models" }, + + async handler({ upstream }, { kind, provider }) { + const raw = await fetchCatalogue(upstream, kind); + const allowed = provider + ? new Set( + raw.providers + .filter((p) => p.provider === provider) + .flatMap((p) => p.models.map((m) => m.model)), + ) + : undefined; + const models = raw.models + .filter((m) => !allowed || allowed.has(m.model)) + .map(({ model, full_name, model_creator, is_disabled, ...capabilities }) => ({ + id: model, + name: full_name ?? model, + creator: model_creator ?? "unknown", + disabled: is_disabled ?? false, + capabilities, + })); + const providers = raw.providers + .filter((p) => !provider || p.provider === provider) + .map((p) => ({ provider: p.provider, model_ids: p.models.map((m) => m.model) })); + + const structured = { kind, models, providers }; + return { structured, text: describe(structured) }; + }, +}; + +/** + * The image and video endpoints take their own `provider` filter — `"artcraft"` (the + * default: only models Artcraft routes itself) or `"all"`. We always ask for `all` so the + * catalogue is complete, and apply the caller's GenerationProvider filter ourselves. + */ +async function fetchCatalogue( + upstream: Parameters[0]["upstream"], + kind: GenerationKind, +): Promise { + const everything = { params: { query: { provider: "all" as const } } }; + switch (kind) { + case "image": + return unwrapUpstream(await upstream.GET("/v1/omni_gen/models/image", everything)); + case "video": + return unwrapUpstream(await upstream.GET("/v1/omni_gen/models/video", everything)); + case "audio": + return unwrapUpstream(await upstream.GET("/v1/omni_gen/models/audio")); + case "mesh": + return unwrapUpstream(await upstream.GET("/v1/omni_gen/models/mesh")); + case "splat": + return unwrapUpstream(await upstream.GET("/v1/omni_gen/models/splat")); + } +} + +function describe(catalogue: z.infer>): string { + const available = catalogue.models.filter((m) => !m.disabled); + const names = available.map((m) => `${m.name} (${m.id})`).join(", "); + const providers = catalogue.providers.map((p) => p.provider).join(", "); + return `${String(available.length)} ${catalogue.kind} models available: ${names}. Providers: ${providers}.`; +} diff --git a/mcp/test/mcp/handler.test.ts b/mcp/test/mcp/handler.test.ts index 626901b54c..f1fc8c82cf 100644 --- a/mcp/test/mcp/handler.test.ts +++ b/mcp/test/mcp/handler.test.ts @@ -135,7 +135,11 @@ describe("the protected MCP route", () => { expect(client.getServerVersion()?.name).toBe("artcraft"); const { tools } = await client.listTools(); - expect(tools.map((t) => t.name).sort()).toEqual(["get_account", "get_credit_balance"]); + expect(tools.map((t) => t.name).sort()).toEqual([ + "get_account", + "get_credit_balance", + "list_models", + ]); const result = await client.callTool({ name: "get_credit_balance", arguments: {} }); expect(result.isError).toBeFalsy(); @@ -153,10 +157,13 @@ describe("the protected MCP route", () => { await client.close(); }); - it("lists nothing for a grant without read:account, but still answers", async () => { + it("lists only catalogue tools for a read:catalog grant", async () => { const { accessToken } = await connectAsSeededUser("read:catalog"); const client = await mcpClient(accessToken); - expect((await client.listTools()).tools).toEqual([]); + expect((await client.listTools()).tools.map((t) => t.name)).toEqual(["list_models"]); + const result = await client.callTool({ name: "list_models", arguments: { kind: "mesh" } }); + expect(result.isError).toBeFalsy(); + expect((result.structuredContent as { models: unknown[] }).models.length).toBeGreaterThan(0); await client.close(); }); diff --git a/mcp/test/mcp/list-models.test.ts b/mcp/test/mcp/list-models.test.ts new file mode 100644 index 0000000000..99e82feba8 --- /dev/null +++ b/mcp/test/mcp/list-models.test.ts @@ -0,0 +1,97 @@ +import { describe, expect, it } from "vitest"; + +import { listModels } from "../../src/mcp/tools/list-models"; +import { principalFromProps } from "../../src/tokens/principal"; +import { createUpstreamClient } from "../../src/upstream/client"; +import { modelsResponse } from "../../fake-upstream/src/catalogue"; +import { SIGNED_SESSION } from "../helpers/fixtures"; + +const principal = principalFromProps({ + credential: { kind: "session", signedSession: SIGNED_SESSION }, + grantIssuedAt: 1_800_000_000_000, + userToken: "user_x", + username: "x", + displayName: "X", + scopes: ["read:catalog"], +}); + +/** Serves the fake's (spec-validated) catalogues, recording the requests. */ +function catalogueUpstream() { + const requests: URL[] = []; + const upstream = createUpstreamClient({ + baseUrl: "https://api.example.test", + use: "read", + credential: principal.credential, + fetch: (input) => { + const url = new URL(new Request(input).url); + requests.push(url); + const kind = url.pathname.split("/").pop() as Parameters[0]; + const filter = url.searchParams.get("provider"); + if (filter !== null && filter !== "all" && filter !== "artcraft") { + return Promise.resolve(new Response("{}", { status: 400 })); + } + const body = modelsResponse(kind, filter ?? undefined); + return Promise.resolve( + new Response(JSON.stringify(body), { headers: { "content-type": "application/json" } }), + ); + }, + }); + return { requests, context: { principal, upstream } }; +} + +describe("list_models", () => { + it("normalises a catalogue and passes the capability fields through verbatim", async () => { + const { requests, context } = catalogueUpstream(); + const result = await listModels.handler(context, { kind: "video" }); + expect(requests.map((u) => u.pathname + u.search)).toEqual([ + "/v1/omni_gen/models/video?provider=all", + ]); + expect(result.structured.kind).toBe("video"); + const seedance = result.structured.models.find((m) => m.id === "seedance_2p0"); + expect(seedance).toMatchObject({ name: "Seedance 2.0", creator: "bytedance", disabled: false }); + expect(seedance?.capabilities).toMatchObject({ + duration_seconds_options: [4, 6, 8, 10], + aspect_ratio_default: "wide_sixteen_by_nine", + text_to_video_supported: true, + }); + expect(seedance?.capabilities).not.toHaveProperty("model"); + expect(seedance?.capabilities).not.toHaveProperty("full_name"); + expect(result.structured.providers.some((p) => p.provider === "fal")).toBe(true); + expect(result.text).toMatch(/^\d+ video models available: .*Seedance 2\.0 \(seedance_2p0\)/); + }); + + it("asks upstream for the complete catalogue and filters by provider itself", async () => { + const { requests, context } = catalogueUpstream(); + const result = await listModels.handler(context, { kind: "image", provider: "midjourney" }); + expect(requests[0]?.search).toBe("?provider=all"); + expect(result.structured.models.map((m) => m.id)).toEqual([ + "midjourney_7", + "midjourney_7_niji", + "midjourney_8", + ]); + expect(result.structured.providers).toEqual([ + { provider: "midjourney", model_ids: ["midjourney_7", "midjourney_7_niji", "midjourney_8"] }, + ]); + }); + + it("filters the same way for kinds whose endpoint takes no query at all", async () => { + const { requests, context } = catalogueUpstream(); + const result = await listModels.handler(context, { kind: "splat", provider: "world_labs" }); + expect(requests[0]?.search).toBe(""); + expect(result.structured.models.every((m) => m.creator === "world_labs")).toBe(true); + expect(result.structured.models.map((m) => m.id)).not.toContain("triposplat"); + }); + + it("serves every kind", async () => { + const { context } = catalogueUpstream(); + for (const kind of ["image", "video", "audio", "mesh", "splat"] as const) { + const result = await listModels.handler(context, { kind }); + expect(result.structured.models.length).toBeGreaterThan(0); + } + }); + + it("is declared read-only and scoped to read:catalog", () => { + expect(listModels.requiredScope).toBe("read:catalog"); + expect(listModels.annotations.readOnlyHint).toBe(true); + }); +}); diff --git a/mcp/test/mcp/server.test.ts b/mcp/test/mcp/server.test.ts index 0a2858edb4..d59b86dc20 100644 --- a/mcp/test/mcp/server.test.ts +++ b/mcp/test/mcp/server.test.ts @@ -85,6 +85,12 @@ describe("MCP server", () => { it("hides tools the principal's scopes do not cover", async () => { const { client } = await connect(["read:catalog"]); const { tools } = await client.listTools(); + expect(tools.map((t) => t.name)).toEqual(["list_models"]); + }); + + it("lists nothing at all for scopes that cover no tool", async () => { + const { client } = await connect([]); + const { tools } = await client.listTools(); expect(tools).toEqual([]); });