mirror of
https://github.com/storytold/artcraft.git
synced 2026-11-12 21:02:14 +00:00
mcp: list_models tool and the fake's model catalogues
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
This commit is contained in:
1 parent
bf12fcbf66
commit
f96cd8b915
9 files changed
+645
-3
No files matched your search
@@ -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<GenerationKind, readonly ModelSeed[]> = {
|
||||
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<string, string> = {
|
||||
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<GenerationKind, (seed: ModelSeed) => Record<string, unknown>> = {
|
||||
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<string, unknown> {
|
||||
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<string, unknown> {
|
||||
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<string, unknown>[] {
|
||||
const byProvider = new Map<string, { model: string }[]>();
|
||||
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 }));
|
||||
}
|
||||
@@ -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;
|
||||
}
|
||||
|
||||
@@ -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;
|
||||
}
|
||||
@@ -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<ModelsBody>();
|
||||
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<ModelsBody>();
|
||||
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<ModelsBody>();
|
||||
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<ModelsBody>();
|
||||
expect(audio.models.length).toBe(5);
|
||||
const ignored = await (
|
||||
await createFakeUpstream().request("/v1/omni_gen/models/audio?provider=midjourney")
|
||||
).json<ModelsBody>();
|
||||
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);
|
||||
});
|
||||
});
|
||||
@@ -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<z.ZodRawShape, z.ZodRawShape>[] = [
|
||||
getAccount,
|
||||
getCreditBalance,
|
||||
listModels,
|
||||
];
|
||||
|
||||
export interface McpServerDeps {
|
||||
|
||||
@@ -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<typeof inputSchema, typeof outputSchema> = {
|
||||
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<typeof listModels.handler>[0]["upstream"],
|
||||
kind: GenerationKind,
|
||||
): Promise<RawCatalogue> {
|
||||
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<z.ZodObject<typeof outputSchema>>): 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}.`;
|
||||
}
|
||||
@@ -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();
|
||||
});
|
||||
|
||||
|
||||
@@ -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<typeof modelsResponse>[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);
|
||||
});
|
||||
});
|
||||
@@ -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([]);
|
||||
});
|
||||
|
||||
|
||||
Reference in new issue
Block a user