paseo/plugins/muse-provider/server/connection.ts

import type { ProviderLaunch } from "@getpaseo/plugin/server/provider";
import { spawnProcess, terminateProcess } from "@getpaseo/plugin/server";
import { randomBytes } from "node:crypto";
import { setTimeout as delay } from "node:timers/promises";
import { z } from "zod";
import { MuseError, classifyExit } from "./errors.js";
import { frameSchema, initializedSchema, fingerprint } from "./wire.js";

export function commandId(): string {
  const bytes = randomBytes(16);
  bytes.writeUIntBE(Date.now(), 0, 6);
  bytes[6] = 0x70 | (bytes[6]! & 0x0f);
  bytes[8] = 0x80 | (bytes[8]! & 0x3f);
  const hex = bytes.toString("hex");
  return `${hex.slice(0, 8)}-${hex.slice(8, 12)}-${hex.slice(12, 16)}-${hex.slice(16, 20)}-${hex.slice(20)}`;
}

/**
 * Splits MSP output on "\n" only. `node:readline` also ends lines at U+2028 and U+2029, which JSON
 * allows unescaped inside strings, so it cut frames quoting either character in two.
 */
function readLines(
  input: NodeJS.ReadableStream,
  onLine: (line: string) => void,
): { close(): void } {
  // Decoding in the stream keeps a multibyte character whole when it straddles two chunks.
  input.setEncoding("utf8");
  let partial = "";
  const emit = (line: string) => onLine(line.endsWith("\r") ? line.slice(0, -1) : line);
  const onData = (chunk: string) => {
    let start = 0;
    for (let end = chunk.indexOf("\n"); end !== -1; end = chunk.indexOf("\n", start)) {
      emit(partial + chunk.slice(start, end));
      partial = "";
      start = end + 1;
    }
    partial += chunk.slice(start);
  };
  const onEnd = () => {
    if (partial) emit(partial);
  };
  input.on("data", onData);
  input.on("end", onEnd);
  return {
    close() {
      input.off("data", onData);
      input.off("end", onEnd);
    },
  };
}

interface PendingRequest {
  resolve(value: unknown): void;
  reject(error: Error): void;
  timer: ReturnType<typeof setTimeout>;
}
export interface Notification {
  method: string;
  params: unknown;
}

export class MspConnection {
  private readonly child;
  private readonly pending = new Map<string | number, PendingRequest>();
  private readonly listeners = new Set<(notification: Notification) => void>();
  private readonly exitListeners = new Set<(error: MuseError) => void>();
  private sequence = 0;
  private stderr = "";
  private failure: MuseError | null = null;
  private closing = false;
  private closeResult: Promise<void> | null = null;
  private readonly exited: Promise<void>;

  constructor(
    private readonly options: {
      launch: ProviderLaunch;
      cwd?: string;
      timeoutMs?: number;
      startupTimeoutMs?: number;
      serveArgs?: string[];
    },
  ) {
    const { launch, cwd } = options;
    this.child = spawnProcess(
      launch.command,
      [...launch.args, "serve", ...(options.serveArgs ?? [])],
      {
        env: launch.env,
        cwd,
        stdio: "pipe",
      },
    );
    this.child.stderr.setEncoding("utf8");
    this.child.stderr.on("data", (chunk: string) => {
      this.stderr = (this.stderr + chunk).slice(-8192);
    });
    const lines = readLines(this.child.stdout, (line) => {
      try {
        this.receive(frameSchema.parse(JSON.parse(line)));
      } catch (error) {
        const message = error instanceof Error ? error.message : String(error);
        this.fail(new MuseError("invalidFrame", `Invalid MSP frame: ${message}`));
        void this.close().catch(() => {});
      }
    });
    this.child.stdin.on("error", () => {
      /* Process exit owns the diagnostic and pending requests. */
    });
    this.child.on("error", (error) => this.fail(new MuseError("spawn", error.message)));
    this.exited = new Promise((resolve) => {
      this.child.once("close", (code, signal) => {
        lines.close();
        this.fail(classifyExit(code, signal, this.stderr));
        resolve();
      });
    });
  }

  async initialize(): Promise<void> {
    const response = await this.request(
      "initialize",
      {
        clientInfo: { name: "paseo", version: "1" },
        capabilities: {
          experimentalApi: true,
          requestedCapabilities: ["sessionMcp", "sessionListStream"],
        },
      },
      initializedSchema,
      this.options.startupTimeoutMs ?? 30000,
    );
    if (response.schema.fingerprint !== fingerprint) {
      process.stderr.write(
        `Muse MSP fingerprint mismatch: expected ${fingerprint}, received ${response.schema.fingerprint}\n`,
      );
    }
    this.write({ jsonrpc: "2.0", method: "initialized", params: {} });
  }

  onNotification(listener: (notification: Notification) => void): () => void {
    this.listeners.add(listener);
    return () => this.listeners.delete(listener);
  }
  onExit(listener: (error: MuseError) => void): () => void {
    this.exitListeners.add(listener);
    return () => this.exitListeners.delete(listener);
  }
  async command<T, P extends object>(method: string, params: P, schema: z.ZodType<T>): Promise<T> {
    const id =
      "commandId" in params && typeof params.commandId === "string"
        ? params.commandId
        : commandId();
    const input = { ...params, commandId: id };
    for (let attempt = 0; ; attempt++) {
      try {
        return await this.request(method, input, schema);
      } catch (error) {
        const retryable =
          error instanceof MuseError && ["overloaded", "backpressured"].includes(error.kind);
        if (!retryable || attempt === 2) throw error;
        await delay(100 * (attempt + 1));
      }
    }
  }
  async request<T>(
    method: string,
    params: object,
    schema: z.ZodType<T>,
    timeoutMs = this.options.timeoutMs ?? 10000,
  ): Promise<T> {
    if (this.failure) throw this.failure;
    if (this.closing) throw new MuseError("closed", "Muse host is closing");
    const id = ++this.sequence;
    const response = await new Promise<unknown>((resolve, reject) => {
      const timer = setTimeout(() => {
        this.pending.delete(id);
        reject(new MuseError("timeout", `Muse ${method} timed out`));
      }, timeoutMs);
      this.pending.set(id, { resolve, reject, timer });
      this.write({ jsonrpc: "2.0", id, method, params });
    });
    return schema.parse(response);
  }
  close(): Promise<void> {
    if (this.closeResult) return this.closeResult;
    this.closing = true;
    this.closeResult = this.finishClose();
    return this.closeResult;
  }
  private async finishClose(): Promise<void> {
    this.child.stdin.end();
    const drained = await Promise.race([
      this.exited.then(() => true),
      delay(1000, undefined, { ref: false }).then(() => false),
    ]);
    if (drained) return;
    await terminateProcess(this.child);
    const stopped = await Promise.race([
      this.exited.then(() => true),
      delay(5000, undefined, { ref: false }).then(() => false),
    ]);
    if (!stopped) throw new MuseError("timeout", "Muse host did not close after termination");
  }
  private write(frame: object): void {
    this.child.stdin.write(JSON.stringify(frame) + "\n");
  }
  private receive(frame: z.infer<typeof frameSchema>): void {
    if (frame.method) {
      if (frame.id !== undefined) {
        if (["approval/request", "userInput/request"].includes(frame.method)) {
          this.write({ jsonrpc: "2.0", id: frame.id, result: {} });
        } else {
          this.write({
            jsonrpc: "2.0",
            id: frame.id,
            error: { code: -32601, message: "Unsupported server request" },
          });
        }
        return;
      }
      for (const listener of this.listeners)
        listener({ method: frame.method, params: frame.params });
      return;
    }
    if (frame.id === undefined) throw new MuseError("invalidFrame", "MSP response has no id");
    if (frame.id === null) {
      if (!frame.error) throw new MuseError("invalidFrame", "MSP result has a null id");
      this.fail(new MuseError(frame.error.data?.kind ?? "rpc", frame.error.message));
      void this.close().catch(() => {});
      return;
    }
    const pending = this.pending.get(frame.id);
    if (!pending) return;
    clearTimeout(pending.timer);
    this.pending.delete(frame.id);
    if (frame.error) {
      let kind = frame.error.data?.kind ?? "rpc";
      if (frame.error.code === -32020) kind = "sessionNotFound";
      pending.reject(new MuseError(kind, frame.error.message, undefined, frame.error.data?.reason));
    } else pending.resolve(frame.result);
  }
  private fail(error: MuseError): void {
    if (this.failure) return;
    this.failure = error;
    for (const pending of this.pending.values()) {
      clearTimeout(pending.timer);
      pending.reject(error);
    }
    this.pending.clear();
    if (!this.closing) for (const listener of this.exitListeners) listener(error);
  }
}