mirror of
https://github.com/diegosouzapw/OmniRoute.git
synced 2026-08-06 07:12:12 +03:00
47468636e4
Unify audit access under the logs experience and add richer active request visibility with sanitized client and provider payload previews. Expose provider warnings from upstream responses in compliance audit logs, surface learned rate-limit header data in health monitoring, and add MCP cache stats and cache flush tools with test coverage. Also add local provider catalog support for SD WebUI and ComfyUI, include video generation in endpoint docs and UI, add eval run storage migrations, and refresh docs and translations to match the expanded feature set.
928 lines
31 KiB
TypeScript
928 lines
31 KiB
TypeScript
/**
|
|
* OmniRoute MCP Server — Model Context Protocol server exposing
|
|
* OmniRoute gateway intelligence as tools for AI agents.
|
|
*
|
|
* Supports two transports:
|
|
* 1. stdio — for IDE integration (VS Code, Cursor, Claude Desktop)
|
|
* 2. HTTP — for remote/programmatic access
|
|
*
|
|
* Tools wrap existing OmniRoute API endpoints and add intelligence
|
|
* such as routing simulation, budget guards, and session snapshots.
|
|
*/
|
|
|
|
import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
|
|
import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js";
|
|
import {
|
|
getComboModelProvider,
|
|
getComboModelString,
|
|
getComboStepTarget,
|
|
} from "../../src/lib/combos/steps.ts";
|
|
|
|
import {
|
|
MCP_TOOLS,
|
|
getHealthInput,
|
|
listCombosInput,
|
|
getComboMetricsInput,
|
|
switchComboInput,
|
|
checkQuotaInput,
|
|
routeRequestInput,
|
|
costReportInput,
|
|
listModelsCatalogInput,
|
|
webSearchInput,
|
|
simulateRouteInput,
|
|
setBudgetGuardInput,
|
|
setRoutingStrategyInput,
|
|
setResilienceProfileInput,
|
|
testComboInput,
|
|
getProviderMetricsInput,
|
|
bestComboForTaskInput,
|
|
explainRouteInput,
|
|
getSessionSnapshotInput,
|
|
dbHealthCheckInput,
|
|
syncPricingInput,
|
|
cacheStatsInput,
|
|
cacheFlushInput,
|
|
} from "./schemas/tools.ts";
|
|
import { startMcpHeartbeat } from "./runtimeHeartbeat.ts";
|
|
|
|
import { closeAuditDb, logToolCall } from "./audit.ts";
|
|
import {
|
|
evaluateToolScopes,
|
|
resolveCallerScopeContext,
|
|
type McpToolExtraLike,
|
|
} from "./scopeEnforcement.ts";
|
|
|
|
import {
|
|
handleSimulateRoute,
|
|
handleSetBudgetGuard,
|
|
handleSetRoutingStrategy,
|
|
handleSetResilienceProfile,
|
|
handleTestCombo,
|
|
handleGetProviderMetrics,
|
|
handleBestComboForTask,
|
|
handleExplainRoute,
|
|
handleGetSessionSnapshot,
|
|
handleDbHealthCheck,
|
|
handleSyncPricing,
|
|
handleCacheStats,
|
|
handleCacheFlush,
|
|
} from "./tools/advancedTools.ts";
|
|
import { memoryTools } from "./tools/memoryTools.ts";
|
|
import { skillTools } from "./tools/skillTools.ts";
|
|
import { normalizeQuotaResponse } from "../../src/shared/contracts/quota.ts";
|
|
import { resolveOmniRouteBaseUrl } from "../../src/shared/utils/resolveOmniRouteBaseUrl.ts";
|
|
|
|
// ============ Configuration ============
|
|
|
|
const OMNIROUTE_BASE_URL = resolveOmniRouteBaseUrl();
|
|
const OMNIROUTE_API_KEY = process.env.OMNIROUTE_API_KEY || "";
|
|
const MCP_ENFORCE_SCOPES = process.env.OMNIROUTE_MCP_ENFORCE_SCOPES === "true";
|
|
const MCP_ALLOWED_SCOPES = new Set(
|
|
(process.env.OMNIROUTE_MCP_SCOPES || "")
|
|
.split(",")
|
|
.map((s) => s.trim())
|
|
.filter(Boolean)
|
|
);
|
|
const TOTAL_MCP_TOOL_COUNT =
|
|
MCP_TOOLS.length + Object.keys(memoryTools).length + Object.keys(skillTools).length;
|
|
|
|
type JsonRecord = Record<string, unknown>;
|
|
|
|
type TextToolResult = {
|
|
content: Array<{ type: "text"; text: string }>;
|
|
isError?: boolean;
|
|
};
|
|
|
|
function toRecord(value: unknown): JsonRecord {
|
|
return value && typeof value === "object" && !Array.isArray(value) ? (value as JsonRecord) : {};
|
|
}
|
|
|
|
function toArray(value: unknown): unknown[] {
|
|
return Array.isArray(value) ? value : [];
|
|
}
|
|
|
|
function toString(value: unknown, fallback = ""): string {
|
|
return typeof value === "string" ? value : fallback;
|
|
}
|
|
|
|
function toNumber(value: unknown, fallback = 0): number {
|
|
return typeof value === "number" && Number.isFinite(value) ? value : fallback;
|
|
}
|
|
|
|
function toStringArray(value: unknown, fallback: string[] = []): string[] {
|
|
const values = toArray(value).filter((entry): entry is string => typeof entry === "string");
|
|
return values.length > 0 ? values : fallback;
|
|
}
|
|
|
|
function normalizeComboModels(
|
|
rawModels: unknown
|
|
): Array<{ provider: string; model: string; priority: number }> {
|
|
return toArray(rawModels).map((rawModel, index) => {
|
|
const modelRecord = toRecord(rawModel);
|
|
const modelString = getComboModelString(rawModel);
|
|
const target = getComboStepTarget(rawModel);
|
|
const provider =
|
|
getComboModelProvider(rawModel) ||
|
|
(modelString ? "unknown" : target ? "combo" : toString(modelRecord.provider, "unknown"));
|
|
|
|
return {
|
|
provider,
|
|
model: modelString || target || toString(modelRecord.model, "unknown"),
|
|
priority: toNumber(modelRecord.priority, index + 1),
|
|
};
|
|
});
|
|
}
|
|
|
|
/**
|
|
* Internal fetch helper that calls OmniRoute API endpoints.
|
|
*/
|
|
async function omniRouteFetch(path: string, options: RequestInit = {}): Promise<unknown> {
|
|
const url = `${OMNIROUTE_BASE_URL}${path}`;
|
|
const headers: Record<string, string> = {
|
|
"Content-Type": "application/json",
|
|
...(OMNIROUTE_API_KEY ? { Authorization: `Bearer ${OMNIROUTE_API_KEY}` } : {}),
|
|
...((options.headers as Record<string, string>) || {}),
|
|
};
|
|
|
|
const signal = options.signal || AbortSignal.timeout(10000);
|
|
const response = await fetch(url, { ...options, headers, signal });
|
|
|
|
if (!response.ok) {
|
|
const errorText = await response.text().catch(() => "Unknown error");
|
|
throw new Error(`OmniRoute API error [${response.status}]: ${errorText}`);
|
|
}
|
|
|
|
return response.json();
|
|
}
|
|
|
|
function withScopeEnforcement(
|
|
toolName: string,
|
|
handler: (args: unknown, extra?: McpToolExtraLike) => Promise<TextToolResult>
|
|
) {
|
|
return async (args: unknown, extra?: McpToolExtraLike): Promise<TextToolResult> => {
|
|
const scopeContext = resolveCallerScopeContext(extra, Array.from(MCP_ALLOWED_SCOPES));
|
|
const scopeCheck = evaluateToolScopes(toolName, scopeContext.scopes, MCP_ENFORCE_SCOPES);
|
|
if (!scopeCheck.allowed) {
|
|
const missingScopes =
|
|
scopeCheck.missing.length > 0 ? scopeCheck.missing.join(", ") : "unavailable";
|
|
const reason = scopeCheck.reason || "scope_check_failed";
|
|
const msg =
|
|
`Insufficient MCP scopes for ${toolName}. ` +
|
|
`Missing: ${missingScopes}. ` +
|
|
`Caller=${scopeContext.callerId}, source=${scopeContext.source}.`;
|
|
const safeArgs = args && typeof args === "object" ? toRecord(args) : { rawArgs: args };
|
|
await logToolCall(
|
|
toolName,
|
|
{
|
|
...safeArgs,
|
|
_scopeCheck: {
|
|
callerId: scopeContext.callerId,
|
|
source: scopeContext.source,
|
|
required: scopeCheck.required,
|
|
provided: scopeCheck.provided,
|
|
missing: scopeCheck.missing,
|
|
},
|
|
},
|
|
null,
|
|
0,
|
|
false,
|
|
`scope_denied:${reason}`
|
|
);
|
|
return {
|
|
content: [{ type: "text" as const, text: `Error: ${msg}` }],
|
|
isError: true,
|
|
};
|
|
}
|
|
|
|
return handler(args, extra);
|
|
};
|
|
}
|
|
|
|
// ============ Tool Handlers ============
|
|
|
|
async function handleGetHealth() {
|
|
const start = Date.now();
|
|
try {
|
|
const [healthRaw, resilienceRaw, rateLimitsRaw] = await Promise.allSettled([
|
|
omniRouteFetch("/api/monitoring/health"),
|
|
omniRouteFetch("/api/resilience"),
|
|
omniRouteFetch("/api/rate-limits"),
|
|
]);
|
|
|
|
const health = healthRaw.status === "fulfilled" ? toRecord(healthRaw.value) : {};
|
|
const resilience = resilienceRaw.status === "fulfilled" ? toRecord(resilienceRaw.value) : {};
|
|
const rateLimits = rateLimitsRaw.status === "fulfilled" ? toRecord(rateLimitsRaw.value) : {};
|
|
const memoryUsageRaw = toRecord(health.memoryUsage);
|
|
const cacheStatsRaw = toRecord(health.cacheStats);
|
|
const resilienceCircuitBreakers = toArray(resilience.circuitBreakers);
|
|
const rateLimitEntries = toArray(rateLimits.limits);
|
|
|
|
const result = {
|
|
uptime: toString(health.uptime, "unknown"),
|
|
version: toString(health.version, "unknown"),
|
|
memoryUsage: {
|
|
heapUsed: toNumber(memoryUsageRaw.heapUsed, 0),
|
|
heapTotal: toNumber(memoryUsageRaw.heapTotal, 0),
|
|
},
|
|
circuitBreakers: resilienceCircuitBreakers,
|
|
rateLimits: rateLimitEntries,
|
|
cacheStats:
|
|
Object.keys(cacheStatsRaw).length > 0
|
|
? {
|
|
hits: toNumber(cacheStatsRaw.hits, 0),
|
|
misses: toNumber(cacheStatsRaw.misses, 0),
|
|
hitRate: toNumber(cacheStatsRaw.hitRate, 0),
|
|
}
|
|
: undefined,
|
|
cryptography: health.cryptography
|
|
? {
|
|
status: toString(toRecord(health.cryptography).status, "missing_or_invalid"),
|
|
provider: toString(toRecord(health.cryptography).provider, "unknown"),
|
|
}
|
|
: undefined,
|
|
};
|
|
|
|
await logToolCall("omniroute_get_health", {}, 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_get_health", {}, null, Date.now() - start, false, msg);
|
|
return { content: [{ type: "text" as const, text: `Error: ${msg}` }], isError: true };
|
|
}
|
|
}
|
|
|
|
async function handleListCombos(args: { includeMetrics?: boolean }) {
|
|
const start = Date.now();
|
|
try {
|
|
const combosRaw = await omniRouteFetch("/api/combos");
|
|
const combosRecord = toRecord(combosRaw);
|
|
const combos = Array.isArray(combosRecord.combos)
|
|
? combosRecord.combos
|
|
: Array.isArray(combosRaw)
|
|
? combosRaw
|
|
: [];
|
|
let metrics: JsonRecord = {};
|
|
if (args.includeMetrics) {
|
|
metrics = toRecord(await omniRouteFetch("/api/combos/metrics").catch(() => ({})));
|
|
}
|
|
|
|
const result = {
|
|
combos: toArray(combos).map((rawCombo) => {
|
|
const combo = toRecord(rawCombo);
|
|
const comboData = toRecord(combo.data);
|
|
const comboId = toString(combo.id, "");
|
|
const modelsSource =
|
|
Array.isArray(combo.models) && combo.models.length > 0 ? combo.models : comboData.models;
|
|
return {
|
|
id: comboId,
|
|
name: toString(combo.name, comboId || "unnamed"),
|
|
models: normalizeComboModels(modelsSource),
|
|
strategy: toString(combo.strategy, toString(comboData.strategy, "priority")),
|
|
enabled: combo.enabled !== false,
|
|
...(args.includeMetrics ? { metrics: metrics[comboId] ?? null } : {}),
|
|
};
|
|
}),
|
|
};
|
|
|
|
await logToolCall("omniroute_list_combos", 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_list_combos", args, null, Date.now() - start, false, msg);
|
|
return { content: [{ type: "text" as const, text: `Error: ${msg}` }], isError: true };
|
|
}
|
|
}
|
|
|
|
async function handleGetComboMetrics(args: { comboId: string }) {
|
|
const start = Date.now();
|
|
try {
|
|
const result = await omniRouteFetch(
|
|
`/api/combos/metrics?comboId=${encodeURIComponent(args.comboId)}`
|
|
);
|
|
await logToolCall("omniroute_get_combo_metrics", 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_get_combo_metrics", args, null, Date.now() - start, false, msg);
|
|
return { content: [{ type: "text" as const, text: `Error: ${msg}` }], isError: true };
|
|
}
|
|
}
|
|
|
|
async function handleSwitchCombo(args: { comboId: string; active: boolean }) {
|
|
const start = Date.now();
|
|
try {
|
|
const result = await omniRouteFetch(`/api/combos/${encodeURIComponent(args.comboId)}`, {
|
|
method: "PUT",
|
|
body: JSON.stringify({ isActive: args.active }),
|
|
});
|
|
await logToolCall("omniroute_switch_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_switch_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 {
|
|
let path = "/api/usage/quota";
|
|
if (args.connectionId) path += `?connectionId=${encodeURIComponent(args.connectionId)}`;
|
|
else if (args.provider) path += `?provider=${encodeURIComponent(args.provider)}`;
|
|
|
|
const result = normalizeQuotaResponse(await omniRouteFetch(path), {
|
|
provider: args.provider || null,
|
|
connectionId: args.connectionId || null,
|
|
});
|
|
|
|
await logToolCall("omniroute_check_quota", 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_check_quota", args, null, Date.now() - start, false, msg);
|
|
return { content: [{ type: "text" as const, text: `Error: ${msg}` }], isError: true };
|
|
}
|
|
}
|
|
|
|
async function handleRouteRequest(args: {
|
|
model: string;
|
|
messages: Array<{ role: string; content: string }>;
|
|
combo?: string;
|
|
budget?: number;
|
|
role?: string;
|
|
stream?: boolean;
|
|
}) {
|
|
const start = Date.now();
|
|
try {
|
|
const body: Record<string, unknown> = {
|
|
model: args.model,
|
|
messages: args.messages,
|
|
stream: false, // MCP tool always returns non-streaming
|
|
};
|
|
if (args.combo) {
|
|
body["x-combo"] = args.combo;
|
|
}
|
|
|
|
const raw = (await omniRouteFetch("/v1/chat/completions", {
|
|
method: "POST",
|
|
body: JSON.stringify(body),
|
|
})) as JsonRecord;
|
|
const choices = toArray(raw.choices);
|
|
const firstChoice = toRecord(choices[0]);
|
|
const firstMessage = toRecord(firstChoice.message);
|
|
const usage = toRecord(raw.usage);
|
|
|
|
const result = {
|
|
response: {
|
|
content: toString(firstMessage.content, ""),
|
|
model: toString(raw.model, args.model),
|
|
tokens: {
|
|
prompt: toNumber(usage.prompt_tokens, 0),
|
|
completion: toNumber(usage.completion_tokens, 0),
|
|
},
|
|
},
|
|
routing: {
|
|
provider: toString(raw.provider, "unknown"),
|
|
combo: raw.combo ?? null,
|
|
fallbacksTriggered: toNumber(raw.fallbacksTriggered, 0),
|
|
cost: toNumber(raw.cost, 0),
|
|
latencyMs: Date.now() - start,
|
|
routingExplanation: toString(
|
|
raw.routingExplanation,
|
|
"Request routed through primary provider"
|
|
),
|
|
},
|
|
};
|
|
|
|
await logToolCall(
|
|
"omniroute_route_request",
|
|
{ model: args.model, messageCount: args.messages.length },
|
|
result.routing,
|
|
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_route_request",
|
|
{ model: args.model },
|
|
null,
|
|
Date.now() - start,
|
|
false,
|
|
msg
|
|
);
|
|
return { content: [{ type: "text" as const, text: `Error: ${msg}` }], isError: true };
|
|
}
|
|
}
|
|
|
|
async function handleCostReport(args: { period?: string }) {
|
|
const start = Date.now();
|
|
try {
|
|
const period = args.period || "session";
|
|
const rangeMap: Record<string, string> = {
|
|
session: "1d",
|
|
day: "1d",
|
|
week: "7d",
|
|
month: "30d",
|
|
};
|
|
const range = rangeMap[period] || "30d";
|
|
const raw = toRecord(
|
|
await omniRouteFetch(`/api/usage/analytics?range=${encodeURIComponent(range)}`)
|
|
);
|
|
const tokenCount = toRecord(raw.tokenCount);
|
|
const budget = toRecord(raw.budget);
|
|
|
|
const result = {
|
|
period,
|
|
totalCost: toNumber(raw.totalCost, 0),
|
|
requestCount: toNumber(raw.requestCount, 0),
|
|
tokenCount: {
|
|
prompt: toNumber(tokenCount.prompt, 0),
|
|
completion: toNumber(tokenCount.completion, 0),
|
|
},
|
|
byProvider: toArray(raw.byProvider),
|
|
byModel: toArray(raw.byModel),
|
|
budget: {
|
|
limit: budget.limit ?? null,
|
|
remaining: budget.remaining ?? null,
|
|
},
|
|
};
|
|
|
|
await logToolCall("omniroute_cost_report", 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_cost_report", args, null, Date.now() - start, false, msg);
|
|
return { content: [{ type: "text" as const, text: `Error: ${msg}` }], isError: true };
|
|
}
|
|
}
|
|
|
|
async function handleListModelsCatalog(args: { provider?: string; capability?: string }) {
|
|
const start = Date.now();
|
|
try {
|
|
let path = "/v1/models";
|
|
let isProviderSpecific = false;
|
|
let source = "local_catalog";
|
|
let warning: string | undefined;
|
|
|
|
if (args.provider && !args.capability) {
|
|
// Use direct provider fetch to get real-time API status
|
|
path = `/api/providers/${encodeURIComponent(args.provider)}/models?excludeHidden=true`;
|
|
isProviderSpecific = true;
|
|
} else {
|
|
const params = new URLSearchParams();
|
|
if (args.provider) params.set("provider", args.provider);
|
|
if (args.capability) params.set("capability", args.capability);
|
|
if (params.toString()) path += `?${params.toString()}`;
|
|
}
|
|
|
|
const raw = toRecord(await omniRouteFetch(path));
|
|
|
|
// If we used the direct provider endpoint
|
|
let rawModels: unknown[] = [];
|
|
if (isProviderSpecific) {
|
|
rawModels = Array.isArray(raw.models) ? raw.models : [];
|
|
source = typeof raw.source === "string" ? raw.source : "api";
|
|
if (raw.warning) warning = String(raw.warning);
|
|
} else {
|
|
rawModels = Array.isArray(raw.data) ? raw.data : [];
|
|
source = "local_catalog";
|
|
// OmniRoute's global /v1/models is always a cached/local catalog
|
|
}
|
|
|
|
const result = {
|
|
models: rawModels.map((rawModel) => {
|
|
const model = toRecord(rawModel);
|
|
return {
|
|
id: toString(model.id, ""),
|
|
provider: toString(model.owned_by, toString(model.provider, args.provider || "unknown")),
|
|
capabilities: toStringArray(model.capabilities, ["chat"]),
|
|
status: toString(model.status, "available"),
|
|
pricing: model.pricing,
|
|
};
|
|
}),
|
|
source,
|
|
...(warning ? { warning } : {}),
|
|
};
|
|
|
|
await logToolCall(
|
|
"omniroute_list_models_catalog",
|
|
args,
|
|
{ modelCount: result.models.length },
|
|
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_list_models_catalog", args, null, Date.now() - start, false, msg);
|
|
return { content: [{ type: "text" as const, text: `Error: ${msg}` }], isError: true };
|
|
}
|
|
}
|
|
|
|
async function handleWebSearch(args: {
|
|
query: string;
|
|
max_results?: number;
|
|
search_type?: "web" | "news";
|
|
provider?:
|
|
| "serper-search"
|
|
| "brave-search"
|
|
| "perplexity-search"
|
|
| "exa-search"
|
|
| "tavily-search"
|
|
| "google-pse-search"
|
|
| "linkup-search"
|
|
| "searchapi-search"
|
|
| "searxng-search";
|
|
}) {
|
|
const start = Date.now();
|
|
try {
|
|
const body: Record<string, unknown> = {
|
|
query: args.query,
|
|
max_results: args.max_results ?? 5,
|
|
search_type: args.search_type ?? "web",
|
|
};
|
|
if (args.provider) body.provider = args.provider;
|
|
|
|
const result = await omniRouteFetch("/v1/search", {
|
|
method: "POST",
|
|
body: JSON.stringify(body),
|
|
signal: AbortSignal.timeout(60000),
|
|
});
|
|
await logToolCall("omniroute_web_search", 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_web_search", args, null, Date.now() - start, false, msg);
|
|
return { content: [{ type: "text" as const, text: `Error: ${msg}` }], isError: true };
|
|
}
|
|
}
|
|
|
|
// ============ MCP Server Setup ============
|
|
|
|
/**
|
|
* Create and configure the OmniRoute MCP Server with all essential tools.
|
|
*/
|
|
export function createMcpServer(): McpServer {
|
|
const server = new McpServer({
|
|
name: "omniroute",
|
|
version: process.env.npm_package_version || "1.8.1",
|
|
});
|
|
|
|
// Register essential tools
|
|
server.registerTool(
|
|
"omniroute_get_health",
|
|
{
|
|
description:
|
|
"Returns OmniRoute health status including uptime, memory, circuit breakers, rate limits, and cache stats",
|
|
inputSchema: getHealthInput,
|
|
},
|
|
withScopeEnforcement("omniroute_get_health", async (args) => {
|
|
getHealthInput.parse(args ?? {});
|
|
return handleGetHealth();
|
|
})
|
|
);
|
|
|
|
server.registerTool(
|
|
"omniroute_list_combos",
|
|
{
|
|
description:
|
|
"Lists all configured combos (model chains) with strategies and optional metrics",
|
|
inputSchema: listCombosInput,
|
|
},
|
|
withScopeEnforcement("omniroute_list_combos", (args) =>
|
|
handleListCombos(listCombosInput.parse(args))
|
|
)
|
|
);
|
|
|
|
server.registerTool(
|
|
"omniroute_get_combo_metrics",
|
|
{
|
|
description: "Returns detailed performance metrics for a specific combo",
|
|
inputSchema: getComboMetricsInput,
|
|
},
|
|
withScopeEnforcement("omniroute_get_combo_metrics", (args) =>
|
|
handleGetComboMetrics(getComboMetricsInput.parse(args))
|
|
)
|
|
);
|
|
|
|
server.registerTool(
|
|
"omniroute_switch_combo",
|
|
{
|
|
description: "Activates or deactivates a combo for routing",
|
|
inputSchema: switchComboInput,
|
|
},
|
|
withScopeEnforcement("omniroute_switch_combo", (args) =>
|
|
handleSwitchCombo(switchComboInput.parse(args))
|
|
)
|
|
);
|
|
|
|
server.registerTool(
|
|
"omniroute_check_quota",
|
|
{
|
|
description: "Checks remaining API quota for one or all providers",
|
|
inputSchema: checkQuotaInput,
|
|
},
|
|
withScopeEnforcement("omniroute_check_quota", (args) =>
|
|
handleCheckQuota(checkQuotaInput.parse(args))
|
|
)
|
|
);
|
|
|
|
server.registerTool(
|
|
"omniroute_route_request",
|
|
{
|
|
description: "Sends a chat completion request through OmniRoute intelligent routing",
|
|
inputSchema: routeRequestInput,
|
|
},
|
|
withScopeEnforcement("omniroute_route_request", (args) =>
|
|
handleRouteRequest(routeRequestInput.parse(args))
|
|
)
|
|
);
|
|
|
|
server.registerTool(
|
|
"omniroute_cost_report",
|
|
{
|
|
description: "Generates a cost report for the specified period",
|
|
inputSchema: costReportInput,
|
|
},
|
|
withScopeEnforcement("omniroute_cost_report", (args) =>
|
|
handleCostReport(costReportInput.parse(args))
|
|
)
|
|
);
|
|
|
|
server.registerTool(
|
|
"omniroute_list_models_catalog",
|
|
{
|
|
description: "Lists all available AI models across providers with capabilities and pricing",
|
|
inputSchema: listModelsCatalogInput,
|
|
},
|
|
withScopeEnforcement("omniroute_list_models_catalog", (args) =>
|
|
handleListModelsCatalog(listModelsCatalogInput.parse(args))
|
|
)
|
|
);
|
|
|
|
// ── Advanced Tools (Phase 3) ──────────────────────────────
|
|
|
|
server.registerTool(
|
|
"omniroute_simulate_route",
|
|
{
|
|
description: "Simulates the routing path a request would take without executing it (dry-run)",
|
|
inputSchema: simulateRouteInput,
|
|
},
|
|
withScopeEnforcement("omniroute_simulate_route", (args) =>
|
|
handleSimulateRoute(simulateRouteInput.parse(args))
|
|
)
|
|
);
|
|
|
|
server.registerTool(
|
|
"omniroute_set_budget_guard",
|
|
{
|
|
description:
|
|
"Sets a session budget limit with configurable action when exceeded (degrade/block/alert)",
|
|
inputSchema: setBudgetGuardInput,
|
|
},
|
|
withScopeEnforcement("omniroute_set_budget_guard", (args) =>
|
|
handleSetBudgetGuard(setBudgetGuardInput.parse(args))
|
|
)
|
|
);
|
|
|
|
server.registerTool(
|
|
"omniroute_set_routing_strategy",
|
|
{
|
|
description:
|
|
"Updates combo routing strategy at runtime (priority/weighted/round-robin/auto/etc.)",
|
|
inputSchema: setRoutingStrategyInput,
|
|
},
|
|
withScopeEnforcement("omniroute_set_routing_strategy", (args) =>
|
|
handleSetRoutingStrategy(setRoutingStrategyInput.parse(args))
|
|
)
|
|
);
|
|
|
|
server.registerTool(
|
|
"omniroute_set_resilience_profile",
|
|
{
|
|
description:
|
|
"Applies a resilience profile controlling circuit breakers, retries, timeouts, and fallback depth",
|
|
inputSchema: setResilienceProfileInput,
|
|
},
|
|
withScopeEnforcement("omniroute_set_resilience_profile", (args) =>
|
|
handleSetResilienceProfile(setResilienceProfileInput.parse(args))
|
|
)
|
|
);
|
|
|
|
server.registerTool(
|
|
"omniroute_test_combo",
|
|
{
|
|
description:
|
|
"Tests each provider in a combo with a real prompt, reporting latency, cost, and success per provider",
|
|
inputSchema: testComboInput,
|
|
},
|
|
withScopeEnforcement("omniroute_test_combo", (args) =>
|
|
handleTestCombo(testComboInput.parse(args))
|
|
)
|
|
);
|
|
|
|
server.registerTool(
|
|
"omniroute_get_provider_metrics",
|
|
{
|
|
description:
|
|
"Returns detailed metrics for a specific provider including latency percentiles and circuit breaker state",
|
|
inputSchema: getProviderMetricsInput,
|
|
},
|
|
withScopeEnforcement("omniroute_get_provider_metrics", (args) =>
|
|
handleGetProviderMetrics(getProviderMetricsInput.parse(args))
|
|
)
|
|
);
|
|
|
|
server.registerTool(
|
|
"omniroute_best_combo_for_task",
|
|
{
|
|
description:
|
|
"Recommends the best combo for a task type based on provider fitness and constraints",
|
|
inputSchema: bestComboForTaskInput,
|
|
},
|
|
withScopeEnforcement("omniroute_best_combo_for_task", (args) =>
|
|
handleBestComboForTask(bestComboForTaskInput.parse(args))
|
|
)
|
|
);
|
|
|
|
server.registerTool(
|
|
"omniroute_explain_route",
|
|
{
|
|
description:
|
|
"Explains why a request was routed to a specific provider, showing scoring factors and fallbacks",
|
|
inputSchema: explainRouteInput,
|
|
},
|
|
withScopeEnforcement("omniroute_explain_route", (args) =>
|
|
handleExplainRoute(explainRouteInput.parse(args))
|
|
)
|
|
);
|
|
|
|
server.registerTool(
|
|
"omniroute_get_session_snapshot",
|
|
{
|
|
description:
|
|
"Returns a full snapshot of the current working session: cost, tokens, top models, errors, budget status",
|
|
inputSchema: getSessionSnapshotInput,
|
|
},
|
|
withScopeEnforcement("omniroute_get_session_snapshot", async (args) => {
|
|
getSessionSnapshotInput.parse(args ?? {});
|
|
return handleGetSessionSnapshot();
|
|
})
|
|
);
|
|
|
|
server.registerTool(
|
|
"omniroute_db_health_check",
|
|
{
|
|
description:
|
|
"Diagnoses or repairs OmniRoute database drift, including broken combo references and orphan quota/domain rows",
|
|
inputSchema: dbHealthCheckInput,
|
|
},
|
|
withScopeEnforcement("omniroute_db_health_check", (args) =>
|
|
handleDbHealthCheck(dbHealthCheckInput.parse(args ?? {}))
|
|
)
|
|
);
|
|
|
|
server.registerTool(
|
|
"omniroute_sync_pricing",
|
|
{
|
|
description:
|
|
"Syncs pricing data from external sources (LiteLLM) into OmniRoute without overwriting user-set prices",
|
|
inputSchema: syncPricingInput,
|
|
},
|
|
withScopeEnforcement("omniroute_sync_pricing", (args) =>
|
|
handleSyncPricing(syncPricingInput.parse(args))
|
|
)
|
|
);
|
|
|
|
server.registerTool(
|
|
"omniroute_web_search",
|
|
{
|
|
description:
|
|
"Performs a web search using OmniRoute's search gateway. Supports multiple providers (Serper, Brave, Perplexity, Exa, Tavily) with automatic failover. Returns search results with titles, URLs, snippets, and position data.",
|
|
inputSchema: webSearchInput,
|
|
},
|
|
withScopeEnforcement("omniroute_web_search", (args) =>
|
|
handleWebSearch(webSearchInput.parse(args))
|
|
)
|
|
);
|
|
|
|
server.registerTool(
|
|
"omniroute_cache_stats",
|
|
{
|
|
description:
|
|
"Returns cache statistics including semantic cache hit rate, prompt cache metrics by provider, and idempotency layer stats.",
|
|
inputSchema: cacheStatsInput,
|
|
},
|
|
withScopeEnforcement("omniroute_cache_stats", () => handleCacheStats())
|
|
);
|
|
|
|
server.registerTool(
|
|
"omniroute_cache_flush",
|
|
{
|
|
description:
|
|
"Flush cache entries. Provide signature to invalidate a single entry, model to invalidate all entries for a model, or omit both to clear all.",
|
|
inputSchema: cacheFlushInput,
|
|
},
|
|
withScopeEnforcement("omniroute_cache_flush", (args) =>
|
|
handleCacheFlush(cacheFlushInput.parse(args))
|
|
)
|
|
);
|
|
|
|
// ── Memory Tools ──────────────────────────────
|
|
Object.values(memoryTools).forEach((toolDef) => {
|
|
server.registerTool(
|
|
toolDef.name,
|
|
{
|
|
description: toolDef.description,
|
|
// @ts-ignore: dynamic zod access
|
|
inputSchema: toolDef.inputSchema,
|
|
},
|
|
withScopeEnforcement(toolDef.name, async (args) => {
|
|
try {
|
|
const parsedArgs = toolDef.inputSchema.parse(args ?? {});
|
|
// @ts-ignore: handler expected specific object
|
|
const result = await toolDef.handler(parsedArgs);
|
|
return { content: [{ type: "text" as const, text: JSON.stringify(result, null, 2) }] };
|
|
} catch (err) {
|
|
const msg = err instanceof Error ? err.message : String(err);
|
|
return { content: [{ type: "text" as const, text: `Error: ${msg}` }], isError: true };
|
|
}
|
|
})
|
|
);
|
|
});
|
|
|
|
// ── Skill Tools ──────────────────────────────
|
|
Object.values(skillTools).forEach((toolDef) => {
|
|
server.registerTool(
|
|
toolDef.name,
|
|
{
|
|
description: toolDef.description,
|
|
// @ts-ignore: dynamic zod access
|
|
inputSchema: toolDef.inputSchema,
|
|
},
|
|
withScopeEnforcement(toolDef.name, async (args) => {
|
|
try {
|
|
const parsedArgs = toolDef.inputSchema.parse(args ?? {});
|
|
// @ts-ignore: handler expected specific object
|
|
const result = await toolDef.handler(parsedArgs);
|
|
return { content: [{ type: "text" as const, text: JSON.stringify(result, null, 2) }] };
|
|
} catch (err) {
|
|
const msg = err instanceof Error ? err.message : String(err);
|
|
return { content: [{ type: "text" as const, text: `Error: ${msg}` }], isError: true };
|
|
}
|
|
})
|
|
);
|
|
});
|
|
|
|
return server;
|
|
}
|
|
|
|
// ============ Main Entry Point (stdio) ============
|
|
|
|
/**
|
|
* Start the MCP server with stdio transport.
|
|
* Called when `omniroute --mcp` is used.
|
|
*/
|
|
export async function startMcpStdio(): Promise<void> {
|
|
const server = createMcpServer();
|
|
const transport = new StdioServerTransport();
|
|
const version = process.env.npm_package_version || "1.8.1";
|
|
const stopHeartbeat = startMcpHeartbeat({
|
|
version,
|
|
scopesEnforced: MCP_ENFORCE_SCOPES,
|
|
allowedScopes: Array.from(MCP_ALLOWED_SCOPES),
|
|
toolCount: TOTAL_MCP_TOOL_COUNT,
|
|
});
|
|
const stopHeartbeatOnce = () => {
|
|
stopHeartbeat();
|
|
};
|
|
process.once("exit", stopHeartbeatOnce);
|
|
process.once("SIGINT", stopHeartbeatOnce);
|
|
process.once("SIGTERM", stopHeartbeatOnce);
|
|
|
|
console.error("[MCP] OmniRoute MCP Server starting (stdio transport)...");
|
|
try {
|
|
await server.connect(transport);
|
|
console.error("[MCP] OmniRoute MCP Server connected and ready.");
|
|
} finally {
|
|
if (closeAuditDb()) {
|
|
console.error("[MCP] Audit database checkpointed and closed.");
|
|
}
|
|
stopHeartbeatOnce();
|
|
process.off("exit", stopHeartbeatOnce);
|
|
process.off("SIGINT", stopHeartbeatOnce);
|
|
process.off("SIGTERM", stopHeartbeatOnce);
|
|
}
|
|
}
|
|
|
|
// If this file is run directly, start stdio server
|
|
if (process.argv[1] && import.meta.url.endsWith(process.argv[1].replace(/\\/g, "/"))) {
|
|
startMcpStdio().catch((err) => {
|
|
console.error("[MCP] Fatal error:", err);
|
|
process.exit(1);
|
|
});
|
|
}
|