@teddysc/mcp-codemode 1.0.0

package/src/server.ts

@@ -0,0 +1,550 @@
+/**
+ * Generic Codemode MCP Server
+ *
+ * Wraps one or more upstream MCP servers (stdio or HTTP) and exposes two tools:
+ *   - search:  LLM writes JS to inspect tool schemas (progressive discovery)
+ *   - execute: LLM writes JS to orchestrate tools via the codemode.* API
+ *
+ * The "search" tool is omitted when addSearch=false, and added automatically
+ * when addSearch=true or when aggregated upstream tools are >= 10.
+ */
+
+import * as vm from "node:vm";
+import { Server } from "@modelcontextprotocol/sdk/server/index.js";
+import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js";
+import {
+  ListToolsRequestSchema,
+  CallToolRequestSchema,
+} from "@modelcontextprotocol/sdk/types.js";
+import { Client } from "@modelcontextprotocol/sdk/client/index.js";
+import { StdioClientTransport } from "@modelcontextprotocol/sdk/client/stdio.js";
+import { SSEClientTransport } from "@modelcontextprotocol/sdk/client/sse.js";
+import { StreamableHTTPClientTransport } from "@modelcontextprotocol/sdk/client/streamableHttp.js";
+
+import {
+  sanitizeToolName,
+  type ToolDescriptor,
+  type JSONSchema,
+} from "./types.js";
+import { LocalExecutor, type ToolFns } from "./executor.js";
+
+export type StdioUpstreamConfig = {
+  id: string;
+  type: "stdio";
+  command: string;
+  args: string[];
+};
+
+export type HttpUpstreamConfig = {
+  id: string;
+  type: "http";
+  url: string;
+  headers?: Record<string, string>;
+};
+
+export type UpstreamConfig = StdioUpstreamConfig | HttpUpstreamConfig;
+
+export interface ServerConfig {
+  /** One or more upstream MCP servers to aggregate */
+  upstreams: UpstreamConfig[];
+  /** Execution timeout in ms (default: 60000) */
+  timeoutMs?: number;
+  /** Server name shown in MCP metadata */
+  serverName?: string;
+  /**
+   * Whether to expose the "search" tool.
+   *   true      – always expose
+   *   false     – never expose (only "execute" is registered)
+   *   undefined – auto: expose when aggregated upstream tools are >= 10
+   */
+  addSearch?: boolean;
+  /**
+   * Whether to include a full available-tools list in the execute tool
+   * description. Defaults to true.
+   */
+  includeExecuteToolList?: boolean;
+}
+
+const AUTO_SEARCH_THRESHOLD = 10;
+
+function buildUpstreamTransport(upstream: UpstreamConfig) {
+  if (upstream.type === "http") {
+    const url = new URL(upstream.url);
+    const requestInit = upstream.headers
+      ? { headers: upstream.headers }
+      : undefined;
+    try {
+      return new StreamableHTTPClientTransport(url, { requestInit });
+    } catch {
+      return new SSEClientTransport(url, { requestInit });
+    }
+  }
+  return new StdioClientTransport({
+    command: upstream.command,
+    args: upstream.args,
+    env: Object.fromEntries(
+      Object.entries(process.env).filter(
+        (e): e is [string, string] => e[1] !== undefined,
+      ),
+    ),
+  });
+}
+
+type ToolBinding = {
+  callName: string;
+  originalName: string;
+  upstreamId: string;
+  description: string;
+  inputSchema: JSONSchema;
+};
+
+type SearchDescriptor = ToolDescriptor & {
+  upstream: string;
+  originalName: string;
+};
+
+type ContentBlock = { type: "text"; text: string };
+type ToolEntry = {
+  name: string;
+  description: string;
+  inputSchema: object;
+  handler: (
+    args: Record<string, unknown>,
+  ) => Promise<{ content: ContentBlock[]; isError?: boolean }>;
+};
+
+function withTimeout<T>(
+  promise: Promise<T>,
+  ms: number,
+  label: string,
+): Promise<T> {
+  return Promise.race([
+    promise,
+    new Promise<never>((_, reject) =>
+      setTimeout(
+        () => reject(new Error(`${label} timed out after ${ms}ms`)),
+        ms,
+      ),
+    ),
+  ]);
+}
+
+function firstLine(text: string | undefined): string {
+  return (text ?? "").split("\n")[0] ?? "";
+}
+
+function toCallName(
+  binding: Pick<ToolBinding, "upstreamId" | "originalName">,
+  preferPrefixedName: boolean,
+  used: Set<string>,
+): string {
+  const base = sanitizeToolName(binding.originalName);
+  const prefixed = sanitizeToolName(
+    `${binding.upstreamId}_${binding.originalName}`,
+  );
+  const candidates = preferPrefixedName
+    ? [prefixed, base]
+    : [base, prefixed];
+
+  for (const candidate of candidates) {
+    if (candidate.length > 0 && !used.has(candidate)) return candidate;
+  }
+
+  const stem = candidates[0] || "_tool";
+  let i = 2;
+  let suffixed = `${stem}_${i}`;
+  while (used.has(suffixed)) {
+    i += 1;
+    suffixed = `${stem}_${i}`;
+  }
+  return suffixed;
+}
+
+export async function runServer(config: ServerConfig): Promise<void> {
+  const {
+    upstreams,
+    timeoutMs = 60_000,
+    serverName = "mcp-codemode",
+    addSearch,
+    includeExecuteToolList = true,
+  } = config;
+
+  if (upstreams.length === 0) {
+    throw new Error("At least one upstream must be configured.");
+  }
+
+  // Shared state populated once upstream is ready
+  let tools: ToolEntry[] = [];
+
+  // Promise that resolves (or rejects) when upstream setup is complete
+  let resolveUpstream!: () => void;
+  const upstreamReady = new Promise<void>((res) => {
+    resolveUpstream = res;
+  });
+  // Prevent unhandled rejection if nobody awaits it before it rejects
+  upstreamReady.catch(() => {});
+
+  // ─── Low-level MCP Server ─────────────────────────────────────────────────
+  // Declare tools capability upfront; actual tool list is resolved lazily.
+
+  const server = new Server(
+    { name: serverName, version: "1.0.0" },
+    { capabilities: { tools: {} } },
+  );
+
+  server.setRequestHandler(ListToolsRequestSchema, async () => {
+    await upstreamReady;
+    return {
+      tools: tools.map(({ name, description, inputSchema }) => ({
+        name,
+        description,
+        inputSchema,
+      })),
+    };
+  });
+
+  server.setRequestHandler(CallToolRequestSchema, async (req) => {
+    await upstreamReady;
+    const tool = tools.find((t) => t.name === req.params.name);
+    if (!tool) {
+      return {
+        content: [
+          { type: "text" as const, text: `Unknown tool: ${req.params.name}` },
+        ],
+        isError: true,
+      };
+    }
+    return tool.handler(
+      (req.params.arguments ?? {}) as Record<string, unknown>,
+    );
+  });
+
+  // ─── Start transport immediately (client handshakes before upstream is up) ──
+
+  const transport = new StdioServerTransport();
+  await server.connect(transport);
+
+  // ─── Connect to upstreams in the background ───────────────────────────────
+
+  let upstreamClients: Client[] = [];
+
+  (async () => {
+    try {
+      const preferPrefixedName = upstreams.length > 1;
+      const results = await Promise.allSettled(
+        upstreams.map(async (upstream) => {
+          const upstreamTransport = buildUpstreamTransport(upstream);
+          const upstreamClientName =
+            upstreams.length === 1
+              ? `${serverName}-upstream`
+              : `${serverName}-upstream-${upstream.id}`;
+          const client = new Client(
+            { name: upstreamClientName, version: "1.0.0" },
+            { capabilities: {} },
+          );
+
+          try {
+            await withTimeout(
+              client.connect(upstreamTransport),
+              timeoutMs,
+              `Upstream "${upstream.id}" connection`,
+            );
+
+            const { tools: mcpTools } = await withTimeout(
+              client.listTools(),
+              timeoutMs,
+              `Upstream "${upstream.id}" listTools`,
+            );
+
+            return { upstream, client, mcpTools };
+          } catch (err) {
+            await client.close().catch(() => {});
+            throw err;
+          }
+        }),
+      );
+
+      const connected: Array<{
+        upstream: UpstreamConfig;
+        client: Client;
+        mcpTools: Array<{
+          name: string;
+          description?: string;
+          inputSchema?: object;
+        }>;
+      }> = [];
+
+      for (const [index, result] of results.entries()) {
+        const upstream = upstreams[index]!;
+        if (result.status === "fulfilled") {
+          connected.push(result.value);
+          continue;
+        }
+        process.stderr.write(
+          `[mcp-codemode] Upstream "${upstream.id}" failed: ${result.reason instanceof Error ? result.reason.message : String(result.reason)}\n`,
+        );
+      }
+
+      if (connected.length === 0) {
+        throw new Error("Failed to connect to all configured upstream servers.");
+      }
+
+      upstreamClients = connected.map((entry) => entry.client);
+
+      const fns: ToolFns = {};
+      const descriptors: Record<string, SearchDescriptor> = {};
+      const usedCallNames = new Set<string>();
+      const bindings: Array<ToolBinding & { client: Client }> = [];
+
+      for (const { upstream, client, mcpTools } of connected) {
+        for (const tool of mcpTools) {
+          const callName = toCallName(
+            { upstreamId: upstream.id, originalName: tool.name },
+            preferPrefixedName,
+            usedCallNames,
+          );
+          usedCallNames.add(callName);
+          bindings.push({
+            callName,
+            originalName: tool.name,
+            upstreamId: upstream.id,
+            description: tool.description ?? "",
+            inputSchema: (tool.inputSchema ?? {}) as JSONSchema,
+            client,
+          });
+        }
+      }
+
+      for (const binding of bindings) {
+        descriptors[binding.callName] = {
+          upstream: binding.upstreamId,
+          originalName: binding.originalName,
+          description: binding.description,
+          inputSchema: binding.inputSchema,
+        };
+        fns[binding.callName] = async (callArgs: unknown) => {
+          const result = await binding.client.callTool({
+            name: binding.originalName,
+            arguments: (callArgs ?? {}) as Record<string, unknown>,
+          });
+          if (result.isError) {
+            const msg = (
+              result.content as Array<{ type: string; text?: string }>
+            )
+              .filter((c) => c.type === "text")
+              .map((c) => c.text)
+              .join("\n");
+            throw new Error(msg || "MCP tool returned an error");
+          }
+          const content = result.content as Array<{
+            type: string;
+            text?: string;
+          }>;
+          if (content.length === 1 && content[0].type === "text")
+            return content[0].text;
+          return content;
+        };
+      }
+
+      // ─── Decide whether to expose search ──────────────────────────────────
+
+      const exposeSearch =
+        addSearch === true
+          ? true
+          : addSearch === false
+            ? false
+            : bindings.length >= AUTO_SEARCH_THRESHOLD;
+
+      // ─── Build tool list for execute description ──────────────────────────
+
+      const toolList = bindings
+        .map(
+          (binding) =>
+            `  ${binding.callName}: [${binding.upstreamId}] ${firstLine(binding.description)}`,
+        )
+        .join("\n");
+
+      // ─── Tool descriptions ────────────────────────────────────────────────
+
+      const SEARCH_DESCRIPTION = `\
+Search aggregated MCP tools and inspect their argument schemas.
+Write a JavaScript async arrow function that queries the \`tools\` object.
+
+\`tools\` is available in scope:
+  Record<toolName, { upstream: string, originalName: string, description: string, inputSchema: object }>`;
+
+      const executeToolListSection = includeExecuteToolList
+        ? `\n\nAvailable tools:\n${toolList || "  (no upstream tools found)"}`
+        : "\n\nAvailable tools list is hidden by configuration. Use the search tool to discover tool names and schemas.";
+
+      const EXECUTE_DESCRIPTION = `\
+Execute JavaScript code to orchestrate upstream MCP tools via the codemode.* API.${exposeSearch ? "\nUse the search tool first to discover the exact argument schemas for the tools you need." : ""}${executeToolListSection}
+
+Rules:
+- Write an async arrow function body that calls codemode.* methods
+- Do NOT use TypeScript syntax (no type annotations, no interfaces, no generics)
+- Always await async calls
+- Return a value summarising what happened`;
+
+      // ─── Search runner (read-only vm sandbox) ──────────────────────────────
+
+      async function runSearch(code: string): Promise<unknown> {
+        const trimmed = code.trim();
+        const normalizedCode = /^async\s*(\(|[a-zA-Z_$])/.test(trimmed)
+          ? trimmed
+          : `async () => {\n${trimmed}\n}`;
+
+        const g = globalThis as Record<string, unknown>;
+        const sandbox = {
+          tools: descriptors,
+          JSON: g["JSON"],
+          Object: g["Object"],
+          Array: g["Array"],
+          Promise: g["Promise"],
+        };
+        const context = vm.createContext(sandbox);
+        const script = new vm.Script(`(${normalizedCode})()`);
+        return Promise.race([
+          script.runInContext(context),
+          new Promise<never>((_, reject) =>
+            setTimeout(
+              () => reject(new Error("Search timed out after 10s")),
+              10_000,
+            ),
+          ),
+        ]);
+      }
+
+      // ─── Execute runner ────────────────────────────────────────────────────
+
+      const executor = new LocalExecutor({ timeoutMs });
+
+      // ─── Populate tool entries ─────────────────────────────────────────────
+
+      tools = [];
+
+      if (exposeSearch) {
+        tools.push({
+          name: "search",
+          description: SEARCH_DESCRIPTION,
+          inputSchema: {
+            type: "object",
+            properties: {
+              code: {
+                type: "string",
+                description:
+                  "JS async arrow function that queries the `tools` object to find schemas",
+              },
+            },
+            required: ["code"],
+          },
+          handler: async ({ code }) => {
+            try {
+              const result = await runSearch(String(code));
+              return {
+                content: [
+                  {
+                    type: "text" as const,
+                    text:
+                      typeof result === "string"
+                        ? result
+                        : JSON.stringify(result, null, 2),
+                  },
+                ],
+              };
+            } catch (err) {
+              return {
+                content: [
+                  {
+                    type: "text" as const,
+                    text: `Error: ${err instanceof Error ? err.message : String(err)}`,
+                  },
+                ],
+                isError: true,
+              };
+            }
+          },
+        });
+      }
+
+      tools.push({
+        name: "execute",
+        description: EXECUTE_DESCRIPTION,
+        inputSchema: {
+          type: "object",
+          properties: {
+            code: {
+              type: "string",
+              description:
+                "An async arrow function that uses codemode.* to call upstream tools",
+            },
+          },
+          required: ["code"],
+        },
+        handler: async ({ code }) => {
+          const result = await executor.execute(String(code), fns);
+          if (result.error) {
+            const parts = [`Error: ${result.error}`];
+            if (result.logs.length > 0)
+              parts.push(`Logs:\n${result.logs.join("\n")}`);
+            return {
+              content: [{ type: "text" as const, text: parts.join("\n\n") }],
+              isError: true,
+            };
+          }
+          const parts: string[] = [];
+          if (result.logs.length > 0)
+            parts.push(`Logs:\n${result.logs.join("\n")}`);
+          parts.push(
+            typeof result.result === "string"
+              ? result.result
+              : JSON.stringify(result.result, null, 2),
+          );
+          return {
+            content: [{ type: "text" as const, text: parts.join("\n\n") }],
+          };
+        },
+      });
+
+      resolveUpstream();
+      // Notify client to re-fetch the tool list
+      await server.sendToolListChanged();
+    } catch (err) {
+      process.stderr.write(
+        `[mcp-codemode] Upstream initialization error: ${err instanceof Error ? err.message : String(err)}\n`,
+      );
+      // Surface error via a stub execute tool
+      tools = [
+        {
+          name: "execute",
+          description: "Upstream MCP server failed to initialize.",
+          inputSchema: {
+            type: "object",
+            properties: { code: { type: "string" } },
+            required: ["code"],
+          },
+          handler: async () => ({
+            content: [
+              {
+                type: "text" as const,
+                text: `Upstream MCP server failed to initialize: ${err instanceof Error ? err.message : String(err)}`,
+              },
+            ],
+            isError: true,
+          }),
+        },
+      ];
+      resolveUpstream();
+      await server.sendToolListChanged();
+    }
+  })();
+
+  // ─── Cleanup ──────────────────────────────────────────────────────────────
+
+  const cleanup = () =>
+    Promise.allSettled(upstreamClients.map((client) => client.close())).finally(
+      () => process.exit(0),
+    );
+  process.on("SIGINT", cleanup);
+  process.on("SIGTERM", cleanup);
+}