superviseMcp

tcmts.cli.mcp.superviseMcp · #bbeufpus · term · runs on the server

uses node:child_process

Run the MCP server argv as a child and relay the protocol, starting it again when watch (this tcm's own program) is replaced: a new version installed, or a new known-good tcm verified. The new server starts at once: the client's handshake is replayed to it, what the client sends from then on goes to it, and the client is told the tools changed. The old one answers what it was asked before (nothing is waited on: a request never answered can't hold the new tcm back), then is stopped, or after ten minutes all the same. So an agent picks up a new tcm, tools and all, without restarting. A server that stops while the client is still there (it crashed) is started again the same way, three times a minute at most, and what it hadn't answered is answered with an error, to ask again. env is each server's environment, as it is when it starts (a session's branch, which it can move onto another). Resolves with the server's exit code.

import { REINIT, stampOf } from "#codebase/tcmts/cli/mcp";
import { spawn } from "node:child_process";
import type { ChildProcess } from "node:child_process";

export function superviseMcp(
  argv: string[],
  env: NodeJS.ProcessEnv | (() => NodeJS.ProcessEnv),
  cwd: string,
  watch: string,
  poll = 2000,
): Promise<number> {
  return new Promise((resolve) => {
    let current!: ChildProcess;
    let init: string | undefined;
    let initialized: string | undefined;
    // which server each request went to: one replaced still answers what it was asked
    const owner = new Map<string, ChildProcess>();
    // (until the new server has answered the replayed handshake, what the client sends waits for it)
    let ready = true;
    let held: string[] = [];
    let stamp = stampOf(watch);
    // (the client gone: a server stopping then is the end)
    let ended = false;
    const crashes: number[] = [];
    const write = (line: string) => process.stdout.write(line + "\n");
    const toChild = (line: string) => {
      try {
        const msg = JSON.parse(line) as { id?: unknown; method?: string };
        if (msg.method === "initialize") init = line;
        if (msg.method === "notifications/initialized") initialized = line;
        if (msg.id !== undefined && msg.method) owner.set(JSON.stringify(msg.id), current);
      } catch {
        // not ours to understand: passed on as it is
      }
      current.stdin!.write(line + "\n");
    };
    const retire = (old: ChildProcess) => {
      const idle = () => ![...owner.values()].includes(old);
      const check = setInterval(() => {
        if (!idle()) return;
        clearInterval(check);
        old.kill();
      }, 250);
      setTimeout(() => {
        clearInterval(check);
        old.kill();
      }, 10 * 60_000).unref();
    };
    const start = (replay: boolean) => {
      const child = spawn(argv[0]!, argv.slice(1), {
        cwd,
        env: typeof env === "function" ? env() : env,
        stdio: ["pipe", "pipe", "inherit"],
      });
      current = child;
      // (what's written to one that has just stopped is lost, not the supervisor's end: what it hadn't answered is
      // answered when it's started again)
      child.stdin!.on("error", () => {});
      let buf = "";
      child.stdout!.setEncoding("utf8").on("data", (d: string) => {
        buf += d;
        for (let nl = buf.indexOf("\n"); nl >= 0; nl = buf.indexOf("\n")) {
          const line = buf.slice(0, nl);
          buf = buf.slice(nl + 1);
          if (!line.trim()) continue;
          let msg: { id?: unknown; result?: unknown; error?: unknown } = {};
          try {
            msg = JSON.parse(line);
          } catch {
            // passed on as it is
          }
          // the replayed handshake answered: the new server is ready, and the client hears what changed
          if (msg.id === REINIT) {
            if (initialized) child.stdin!.write(initialized + "\n");
            ready = true;
            write(JSON.stringify({ jsonrpc: "2.0", method: "notifications/tools/list_changed" }));
            const waiting = held;
            held = [];
            for (const l of waiting) toChild(l);
            continue;
          }
          if (msg.id !== undefined && (msg.result !== undefined || msg.error !== undefined))
            owner.delete(JSON.stringify(msg.id));
          write(line);
        }
      });
      child.on("exit", (code, signal) => {
        // (a server replaced on purpose isn't the end)
        if (child !== current) return;
        const now = Date.now();
        if (!ended && crashes.filter((t) => t > now - 60_000).length < 3) {
          crashes.push(now);
          for (const [id, to] of owner)
            if (to === child) {
              owner.delete(id);
              write(
                JSON.stringify({
                  jsonrpc: "2.0",
                  id: JSON.parse(id) as unknown,
                  error: {
                    code: -32603,
                    message: `tcm's MCP server stopped (${signal ?? code}) before answering; it's been started again, on the same branch: ask again`,
                  },
                }),
              );
            }
          start(true);
          return;
        }
        clearInterval(timer);
        resolve(code ?? 0);
      });
      if (replay && init) {
        ready = false;
        child.stdin!.write(JSON.stringify({ ...JSON.parse(init), id: REINIT }) + "\n");
      }
    };
    const timer = setInterval(() => {
      const now = stampOf(watch);
      if (!now || now === stamp) return;
      stamp = now;
      const old = current;
      start(true);
      retire(old);
    }, poll);
    start(false);
    let rest = "";
    process.stdin.setEncoding("utf8");
    process.stdin.on("data", (d: string) => {
      rest += d;
      for (let nl = rest.indexOf("\n"); nl >= 0; nl = rest.indexOf("\n")) {
        const line = rest.slice(0, nl);
        rest = rest.slice(nl + 1);
        if (!line.trim()) continue;
        if (ready) toChild(line);
        else held.push(line);
      }
    });
    process.stdin.on("end", () => {
      ended = true;
      current.stdin!.end();
    });
  });
}

uses

used by