mirror of
https://github.com/diegosouzapw/OmniRoute.git
synced 2026-08-08 00:02:20 +03:00
feat(mcp): add omniroute_create_combo tool (#8925)
Validated in local merge-train (tomni-proxmox-113)
This commit is contained in:
committed by
GitHub
parent
7f36b192f0
commit
ebdbe3a38f
193
open-sse/mcp-server/__tests__/createComboTool.test.ts
Normal file
193
open-sse/mcp-server/__tests__/createComboTool.test.ts
Normal file
@@ -0,0 +1,193 @@
|
||||
import { describe, it, expect, vi, beforeEach, afterEach } from "vitest";
|
||||
import { Client } from "@modelcontextprotocol/sdk/client/index.js";
|
||||
import { InMemoryTransport } from "@modelcontextprotocol/sdk/inMemory.js";
|
||||
import { MCP_TOOLS, MCP_TOOL_MAP, createComboInput, createComboTool } from "../schemas/tools.ts";
|
||||
import { createMcpServer } from "../server.ts";
|
||||
|
||||
const mockFetch = vi.fn();
|
||||
vi.stubGlobal("fetch", mockFetch);
|
||||
|
||||
const mockLogToolCall = vi.hoisted(() => vi.fn().mockResolvedValue(undefined));
|
||||
vi.mock("../audit.ts", () => ({
|
||||
logToolCall: mockLogToolCall,
|
||||
}));
|
||||
|
||||
describe("omniroute_create_combo MCP tool schema", () => {
|
||||
it("should be registered in MCP_TOOLS and MCP_TOOL_MAP", () => {
|
||||
const tool = MCP_TOOLS.find((t) => t.name === "omniroute_create_combo");
|
||||
expect(tool).toBeDefined();
|
||||
expect(MCP_TOOL_MAP["omniroute_create_combo"]).toBeDefined();
|
||||
});
|
||||
|
||||
it("should require write:combos scope", () => {
|
||||
expect(createComboTool.scopes).toContain("write:combos");
|
||||
});
|
||||
|
||||
it("should validate a minimal payload (name + models)", () => {
|
||||
const result = createComboInput.safeParse({
|
||||
name: "My Combo",
|
||||
models: [{ provider: "anthropic", model: "claude-sonnet" }],
|
||||
});
|
||||
expect(result.success).toBe(true);
|
||||
});
|
||||
|
||||
it("should validate a full payload with description and strategy", () => {
|
||||
const result = createComboInput.safeParse({
|
||||
name: "My Combo",
|
||||
description: "A test combo",
|
||||
strategy: "priority",
|
||||
models: [
|
||||
{ provider: "anthropic", model: "claude-sonnet" },
|
||||
{ provider: "google", model: "gemini-pro" },
|
||||
],
|
||||
});
|
||||
expect(result.success).toBe(true);
|
||||
});
|
||||
|
||||
it("should reject a payload missing name", () => {
|
||||
const result = createComboInput.safeParse({
|
||||
models: [{ provider: "anthropic", model: "claude-sonnet" }],
|
||||
});
|
||||
expect(result.success).toBe(false);
|
||||
});
|
||||
|
||||
it("should reject a payload with an empty models array", () => {
|
||||
const result = createComboInput.safeParse({ name: "My Combo", models: [] });
|
||||
expect(result.success).toBe(false);
|
||||
});
|
||||
|
||||
it("should reject an unknown strategy value", () => {
|
||||
const result = createComboInput.safeParse({
|
||||
name: "My Combo",
|
||||
strategy: "not-a-real-strategy",
|
||||
models: [{ provider: "anthropic", model: "claude-sonnet" }],
|
||||
});
|
||||
expect(result.success).toBe(false);
|
||||
});
|
||||
});
|
||||
|
||||
describe("omniroute_create_combo handler (via MCP dispatch)", () => {
|
||||
let client: Client;
|
||||
|
||||
beforeEach(async () => {
|
||||
mockFetch.mockReset();
|
||||
mockLogToolCall.mockClear();
|
||||
|
||||
const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair();
|
||||
const server = createMcpServer();
|
||||
await server.connect(serverTransport);
|
||||
client = new Client({ name: "create-combo-test", version: "1.0.0" });
|
||||
await client.connect(clientTransport);
|
||||
});
|
||||
|
||||
afterEach(async () => {
|
||||
await client.close();
|
||||
});
|
||||
|
||||
it("should appear in tools/list after registration", async () => {
|
||||
const { tools } = await client.listTools();
|
||||
const tool = tools.find((t) => t.name === "omniroute_create_combo");
|
||||
expect(tool).toBeDefined();
|
||||
expect(tool?.description).toContain("Registers new combo");
|
||||
});
|
||||
|
||||
it("should POST to /api/combos and return the created combo on success", async () => {
|
||||
mockFetch.mockResolvedValueOnce({
|
||||
ok: true,
|
||||
json: async () => ({
|
||||
success: true,
|
||||
combo: { id: "combo-123", name: "My Combo", strategy: "priority", enabled: true },
|
||||
}),
|
||||
});
|
||||
|
||||
const args = {
|
||||
name: "My Combo",
|
||||
models: [{ provider: "anthropic", model: "claude-sonnet" }],
|
||||
};
|
||||
|
||||
const result = await client.callTool({ name: "omniroute_create_combo", arguments: args });
|
||||
|
||||
expect(result.isError).toBeFalsy();
|
||||
const content = result.content[0] as { type: string; text: string };
|
||||
const data = JSON.parse(content.text);
|
||||
expect(data.success).toBe(true);
|
||||
expect(data.combo.id).toBe("combo-123");
|
||||
expect(data.combo.name).toBe("My Combo");
|
||||
|
||||
expect(mockFetch).toHaveBeenCalledWith(
|
||||
expect.stringContaining("/api/combos"),
|
||||
expect.objectContaining({ method: "POST" })
|
||||
);
|
||||
const [, options] = mockFetch.mock.calls[0];
|
||||
const body = JSON.parse(options.body as string);
|
||||
expect(body.name).toBe("My Combo");
|
||||
expect(body.models).toHaveLength(1);
|
||||
|
||||
// Audit: the invocation must be logged to mcp_audit (via logToolCall).
|
||||
expect(mockLogToolCall).toHaveBeenCalledWith(
|
||||
"omniroute_create_combo",
|
||||
expect.objectContaining({ name: "My Combo" }),
|
||||
expect.objectContaining({ success: true }),
|
||||
expect.any(Number),
|
||||
true
|
||||
);
|
||||
});
|
||||
|
||||
it("should pass through optional description and strategy fields", async () => {
|
||||
mockFetch.mockResolvedValueOnce({
|
||||
ok: true,
|
||||
json: async () => ({
|
||||
success: true,
|
||||
combo: { id: "combo-456", name: "Cost Saver", strategy: "cost-optimized", enabled: true },
|
||||
}),
|
||||
});
|
||||
|
||||
await client.callTool({
|
||||
name: "omniroute_create_combo",
|
||||
arguments: {
|
||||
name: "Cost Saver",
|
||||
description: "Prefers cheaper models",
|
||||
strategy: "cost-optimized",
|
||||
models: [
|
||||
{ provider: "anthropic", model: "claude-haiku" },
|
||||
{ provider: "google", model: "gemini-flash" },
|
||||
],
|
||||
},
|
||||
});
|
||||
|
||||
const [, options] = mockFetch.mock.calls[0];
|
||||
const body = JSON.parse(options.body as string);
|
||||
expect(body.description).toBe("Prefers cheaper models");
|
||||
expect(body.strategy).toBe("cost-optimized");
|
||||
expect(body.models).toHaveLength(2);
|
||||
});
|
||||
|
||||
it("should return isError and log the failure when the backend rejects the combo (e.g. name collision)", async () => {
|
||||
mockFetch.mockResolvedValueOnce({
|
||||
ok: false,
|
||||
status: 409,
|
||||
text: async () => "Combo name already exists",
|
||||
});
|
||||
|
||||
const result = await client.callTool({
|
||||
name: "omniroute_create_combo",
|
||||
arguments: {
|
||||
name: "Duplicate Combo",
|
||||
models: [{ provider: "anthropic", model: "claude-sonnet" }],
|
||||
},
|
||||
});
|
||||
|
||||
expect(result.isError).toBe(true);
|
||||
const content = result.content[0] as { type: string; text: string };
|
||||
expect(content.text).toContain("Error");
|
||||
|
||||
expect(mockLogToolCall).toHaveBeenCalledWith(
|
||||
"omniroute_create_combo",
|
||||
expect.objectContaining({ name: "Duplicate Combo" }),
|
||||
null,
|
||||
expect.any(Number),
|
||||
false,
|
||||
expect.stringContaining("Combo name already exists")
|
||||
);
|
||||
});
|
||||
});
|
||||
@@ -192,6 +192,53 @@ export const switchComboTool: McpToolDefinition<typeof switchComboInput, typeof
|
||||
sourceEndpoints: ["/api/combos"],
|
||||
};
|
||||
|
||||
// --- Tool 4b: omniroute_create_combo ---
|
||||
export const createComboInput = z.object({
|
||||
name: z
|
||||
.string()
|
||||
.trim()
|
||||
.min(1)
|
||||
.max(100)
|
||||
.describe("Unique combo name (letters, numbers, spaces, -, _, /, ., [ and ])"),
|
||||
description: z.string().max(2000).optional().describe("Optional human-readable description"),
|
||||
strategy: z
|
||||
.enum(ROUTING_STRATEGY_VALUES)
|
||||
.optional()
|
||||
.describe("Routing strategy (default: priority)"),
|
||||
models: z
|
||||
.array(
|
||||
z.object({
|
||||
provider: z.string().describe("Provider name (e.g., 'claude', 'gemini')"),
|
||||
model: z.string().describe("Model ID for that provider"),
|
||||
})
|
||||
)
|
||||
.min(1)
|
||||
.describe("Ordered model chain; order defines priority"),
|
||||
});
|
||||
|
||||
export const createComboOutput = z.object({
|
||||
success: z.boolean(),
|
||||
combo: z.object({
|
||||
id: z.string(),
|
||||
name: z.string(),
|
||||
strategy: z.string(),
|
||||
enabled: z.boolean(),
|
||||
}),
|
||||
});
|
||||
|
||||
export const createComboTool: McpToolDefinition<typeof createComboInput, typeof createComboOutput> =
|
||||
{
|
||||
name: "omniroute_create_combo",
|
||||
description:
|
||||
"Registers a new combo (model chain) with a name, ordered model list, and optional routing strategy. Full validation (name collisions, nested-combo DAG, composite tiers) is enforced by the combos API.",
|
||||
inputSchema: createComboInput,
|
||||
outputSchema: createComboOutput,
|
||||
scopes: ["write:combos"],
|
||||
auditLevel: "full",
|
||||
phase: 1,
|
||||
sourceEndpoints: ["/api/combos"],
|
||||
};
|
||||
|
||||
// --- Tool 5: omniroute_check_quota ---
|
||||
export const checkQuotaInput = z.object({
|
||||
provider: z
|
||||
@@ -1460,6 +1507,7 @@ export const MCP_TOOLS = [
|
||||
listCombosTool,
|
||||
getComboMetricsTool,
|
||||
switchComboTool,
|
||||
createComboTool,
|
||||
checkQuotaTool,
|
||||
routeRequestTool,
|
||||
costReportTool,
|
||||
|
||||
@@ -12,6 +12,7 @@ import {
|
||||
listCombosInput,
|
||||
getComboMetricsInput,
|
||||
switchComboInput,
|
||||
createComboInput,
|
||||
checkQuotaInput,
|
||||
routeRequestInput,
|
||||
costReportInput,
|
||||
@@ -393,6 +394,27 @@ async function handleSwitchCombo(args: { comboId: string; active: boolean }) {
|
||||
}
|
||||
}
|
||||
|
||||
async function handleCreateCombo(args: {
|
||||
name: string;
|
||||
description?: string;
|
||||
strategy?: string;
|
||||
models: { provider: string; model: string }[];
|
||||
}) {
|
||||
const start = Date.now();
|
||||
try {
|
||||
const result = await omniRouteFetch("/api/combos", {
|
||||
method: "POST",
|
||||
body: JSON.stringify(args),
|
||||
});
|
||||
await logToolCall("omniroute_create_combo", args, result, Date.now() - start, true);
|
||||
return { content: [{ type: "text" as const, text: JSON.stringify(result, null, 2) }] };
|
||||
} catch (err) {
|
||||
const msg = err instanceof Error ? err.message : String(err);
|
||||
await logToolCall("omniroute_create_combo", args, null, Date.now() - start, false, msg);
|
||||
return { content: [{ type: "text" as const, text: `Error: ${msg}` }], isError: true };
|
||||
}
|
||||
}
|
||||
|
||||
async function handleCheckQuota(args: { provider?: string; connectionId?: string }) {
|
||||
const start = Date.now();
|
||||
try {
|
||||
@@ -738,6 +760,17 @@ export function createMcpServer(): McpServer {
|
||||
)
|
||||
);
|
||||
|
||||
server.registerTool(
|
||||
"omniroute_create_combo",
|
||||
{
|
||||
description: "Registers a new combo (model chain) with name, models, and strategy",
|
||||
inputSchema: createComboInput,
|
||||
},
|
||||
withScopeEnforcement("omniroute_create_combo", (args) =>
|
||||
handleCreateCombo(createComboInput.parse(args))
|
||||
)
|
||||
);
|
||||
|
||||
server.registerTool(
|
||||
"omniroute_check_quota",
|
||||
{
|
||||
|
||||
Reference in New Issue
Block a user