MCP 多服务器编排与上下文路由

// 代理聚合服务器核心
import { Server } from "@modelcontextprotocol/sdk/server/index.js";

type Downstream = { name: string; client: any; tools: string[] };

const downstreams: Downstream[] = [
  { name: "db", client: dbClient, tools: ["query_sql", "list_tables"] },
  { name: "search", client: searchClient, tools: ["search_docs", "get_chunk"] },
];

// 工具名 -> 下游服务器 的索引
const toolRoute = new Map<string, Downstream>();
for (const ds of downstreams) {
  for (const tool of ds.tools) toolRoute.set(tool, ds);
}

server.setRequestHandler(CallToolRequestSchema, async (request) => {
  const { name, arguments: args } = request.params;
  const route = toolRoute.get(name);
  if (!route) {
    return {
      content: [{ type: "text", text: `无路由: ${name}` }],
      isError: true,
    };
  }
  // 转发给正确的下游
  const result = await route.client.callTool({ name, arguments: args });
  return result;
});

function namespacedName(dsName: string, toolName: string): string {
  return `${dsName}__${toolName}`; // 如 search__search_docs
}

// 发布给模型时使用命名空间名
const visibleTools = downstreams.flatMap((ds) =>
  ds.tools.map((t) => ({
    name: namespacedName(ds.name, t),
    description: `[${ds.name}] ${t}`,
    inputSchema: ds.getToolSchema(t),
  }))
);

// 收到调用时还原为原始工具名
function resolve(visibleName: string): { ds: Downstream; tool: string } {
  const [dsName, ...rest] = visibleName.split("__");
  const ds = downstreams.find((d) => d.name === dsName)!;
  return { ds, tool: rest.join("__") };
}

if (result.isError) {
  return {
    content: [{ type: "text", text: `[${ds.name}] ${result.content[0]?.text ?? "未知错误"}` }],
    isError: true,
  };
}

# 语义路由：把用户任务映射到最匹配的服务器
def semantic_route(task: str, servers: dict) -> list:
    """按任务与服务器描述的相关性打分，返回 Top-K 服务器"""
    task_vec = embed(task)
    scored = []
    for name, desc in servers.items():
        sim = cosine(task_vec, embed(desc["summary"]))
        scored.append((sim, name))
    scored.sort(reverse=True)
    # 只把 Top-K 服务器的工具暴露给模型
    return [name for _, name in scored[:k]]

def select_tools_for_task(task: str, registry, budget_tokens: int) -> list:
    """根据任务语义挑出命中工具，并控制在 token 预算内"""
    candidates = registry.match(task)   # 基于描述相似度召回
    candidates.sort(key=lambda t: t.priority, reverse=True)

    selected = []
    used = 0
    for tool in candidates:
        schema_cost = estimate_tokens(tool.input_schema)
        if used + schema_cost > budget_tokens:
            break
        selected.append(tool)
        used += schema_cost
    return selected

// 并发调用多个服务器，返回结构化结果
async function fanOut(calls: { name: string; args: any }[]) {
  const results = await Promise.allSettled(
    calls.map(async (c) => {
      const route = toolRoute.get(c.name)!;
      return { tool: c.name, result: await route.client.callTool(c) };
    })
  );
  return results;
}

// 消费端合并
const merged = await fanOut([
  { name: "db__query_sql", args: { sql: "SELECT count(*) FROM orders" } },
  { name: "search__search_docs", args: { query: "订单量" } },
]);
for (const item of merged) {
  if (item.status === "fulfilled") {
    console.log(item.value.tool, item.value.result.content[0].text);
  } else {
    console.warn(item.reason);
  }
}

class CircuitBreaker {
  private failures = 0;
  private openedAt = 0;
  constructor(
    private threshold = 5,
    private cooldownMs = 30_000
  ) {}

  async call<T>(fn: () => Promise<T>, fallback: T): Promise<T> {
    if (this.isOpen()) return fallback;   // 熔断期内直接降级
    try {
      const result = await fn();
      this.failures = 0;
      return result;
    } catch (err) {
      this.failures++;
      if (this.failures >= this.threshold) this.openedAt = Date.now();
      return fallback;
    }
  }

  private isOpen() {
    if (!this.openedAt) return false;
    return Date.now() - this.openedAt < this.cooldownMs;
  }
}

// 带健康检查的连接池
class McpConnectionPool {
  private pools = new Map<string, PoolEntry>();

  async healthCheckAll(): Promise<void> {
    for (const [name, entry] of this.pools) {
      try {
        await entry.client.ping();   // 或发一个最小请求
        entry.degraded = false;
      } catch {
        entry.degraded = true;       // 标记降级，路由时跳过
        console.warn(`服务器 ${name} 健康检查失败`);
      }
    }
  }

  async call(name: string, tool: string, args: any) {
    const entry = this.pools.get(name);
    if (!entry || entry.degraded) {
      throw new Error(`服务器 ${name} 不可用`);
    }
    // 简单的轮询：每次调用取一条可用连接
    const conn = entry.next();
    return conn.callTool({ name: tool, arguments: args });
  }
}

// 每 30 秒跑一轮健康检查
setInterval(() => pool.healthCheckAll(), 30_000);