@ian-pascoe/pi-mcp 0.1.0

package/src/mcp-server-client.ts

@@ -0,0 +1,375 @@
+import { AsyncLocalStorage } from "node:async_hooks";
+import { resolve } from "node:path";
+import {
+  Client,
+  ProtocolError,
+  ProtocolErrorCode,
+  SSEClientTransport,
+  StreamableHTTPClientTransport,
+  createFetchWithInit,
+  type AuthProvider,
+  type ClientContext,
+  type ClientOptions,
+  type CreateMessageRequest,
+  type CreateMessageResult,
+  type CreateMessageResultWithTools,
+  type ElicitRequest,
+  type ElicitResult,
+  type FetchLike,
+  type Implementation,
+  type ListRootsRequest,
+  type ListRootsResult,
+  type OAuthClientProvider,
+  type RequestOptions,
+} from "@modelcontextprotocol/client";
+import { StdioClientTransport, getDefaultEnvironment } from "@modelcontextprotocol/client/stdio";
+import { MCP_SHUTDOWN_TIMEOUT_MS, type McpServerDefinition } from "./pi-mcp-settings.js";
+
+const MAX_MCP_LIST_PAGES = 1_000;
+
+type McpConnectableServerDefinition =
+  | Pick<
+      Extract<McpServerDefinition, { readonly transport: "stdio" }>,
+      "args" | "command" | "cwd" | "environment" | "transport"
+    >
+  | Pick<
+      Extract<McpServerDefinition, { readonly transport: "http" | "sse" }>,
+      "auth" | "headers" | "transport" | "url"
+    >;
+
+type McpClientTransport = StdioClientTransport | StreamableHTTPClientTransport | SSEClientTransport;
+
+/** Request-scoped Host callbacks used for MCP sampling, elicitation, and roots. */
+export interface McpServerRequestCallbacks<PiContext> {
+  readonly onElicitation?: (
+    request: ElicitRequest,
+    piContext: PiContext,
+    clientContext: ClientContext,
+  ) => Promise<ElicitResult> | ElicitResult;
+  readonly onListRoots?: (
+    request: ListRootsRequest,
+    piContext: PiContext,
+    clientContext: ClientContext,
+  ) => Promise<ListRootsResult> | ListRootsResult;
+  readonly onSampling?: (
+    request: CreateMessageRequest,
+    piContext: PiContext,
+    clientContext: ClientContext,
+  ) =>
+    | Promise<CreateMessageResult | CreateMessageResultWithTools>
+    | CreateMessageResult
+    | CreateMessageResultWithTools;
+}
+
+/** Per-call cancellation, progress, and Pi request context. */
+export interface McpServerRunOptions<PiContext> {
+  readonly callbacks?: McpServerRequestCallbacks<PiContext>;
+  readonly onProgress?: NonNullable<RequestOptions["onprogress"]>;
+  readonly piContext?: PiContext;
+  readonly signal?: AbortSignal;
+}
+
+/** Inputs required to connect one owned MCP Client to one Server Definition. */
+export interface McpServerClientConnectOptions {
+  readonly authProvider?: AuthProvider | OAuthClientProvider;
+  readonly clientInfo: Implementation;
+  readonly connectTimeoutMs: number;
+  readonly definition: McpConnectableServerDefinition;
+  readonly onConnectionClose?: () => void;
+  readonly listChanged?: ClientOptions["listChanged"];
+  readonly onError?: (error: Error) => void;
+  readonly onStderr?: (text: string) => void;
+  readonly piCwd: string;
+  readonly requestTimeoutMs: number;
+}
+
+interface ActiveMcpRunContext {
+  callbacks?: McpServerRequestCallbacks<unknown>;
+  piContext?: unknown;
+}
+
+function requestHeaders(
+  definition: Extract<McpConnectableServerDefinition, { transport: "http" | "sse" }>,
+): Headers {
+  const headers = new Headers(definition.headers);
+  if (definition.auth?.type === "bearer") {
+    headers.set("Authorization", `Bearer ${definition.auth.token}`);
+  }
+  return headers;
+}
+
+async function readMcpAuthToken(
+  authProvider: AuthProvider | OAuthClientProvider,
+): Promise<string | undefined> {
+  if ("token" in authProvider) return authProvider.token();
+  return (await authProvider.tokens())?.access_token;
+}
+
+function createAuthenticatedFetch(
+  authProvider: AuthProvider | OAuthClientProvider,
+  requestInit: RequestInit,
+): FetchLike {
+  const fetchWithHeaders = createFetchWithInit(undefined, requestInit);
+  return async (input, init) => {
+    const token = await readMcpAuthToken(authProvider);
+    if (token === undefined) return fetchWithHeaders(input, init);
+    const headers = new Headers(init?.headers);
+    headers.set("Authorization", `Bearer ${token}`);
+    return fetchWithHeaders(input, { ...init, headers });
+  };
+}
+
+function createMcpTransport(options: McpServerClientConnectOptions): McpClientTransport {
+  const definition = options.definition;
+  if (definition.transport === "stdio") {
+    return new StdioClientTransport({
+      args: [...definition.args],
+      command: definition.command,
+      cwd: resolve(options.piCwd, definition.cwd ?? "."),
+      env: { ...getDefaultEnvironment(), ...definition.environment },
+      stderr: "pipe",
+    });
+  }
+
+  const headers = requestHeaders(definition);
+  const requestInit = { headers } satisfies RequestInit;
+  if (definition.transport === "http") {
+    const transportOptions: NonNullable<
+      ConstructorParameters<typeof StreamableHTTPClientTransport>[1]
+    > = { requestInit };
+    if (options.authProvider !== undefined) transportOptions.authProvider = options.authProvider;
+    return new StreamableHTTPClientTransport(new URL(definition.url), transportOptions);
+  }
+
+  const transportOptions: NonNullable<ConstructorParameters<typeof SSEClientTransport>[1]> = {
+    requestInit,
+  };
+  if ([...headers].length > 0) {
+    transportOptions.eventSourceInit = {
+      fetch:
+        options.authProvider === undefined
+          ? createFetchWithInit(undefined, requestInit)
+          : createAuthenticatedFetch(options.authProvider, requestInit),
+    };
+  }
+  if (options.authProvider !== undefined) transportOptions.authProvider = options.authProvider;
+  return new SSEClientTransport(new URL(definition.url), transportOptions);
+}
+
+function closeBeforeDeadline(
+  close: () => Promise<void>,
+  processId: number | undefined,
+): Promise<void> {
+  return new Promise((resolveClose) => {
+    let settled = false;
+    const finish = () => {
+      if (settled) return;
+      settled = true;
+      clearTimeout(timer);
+      resolveClose();
+    };
+    const timer = setTimeout(() => {
+      if (processId !== undefined) {
+        try {
+          process.kill(processId, "SIGKILL");
+        } catch {
+          // The owned child already exited.
+        }
+      }
+      finish();
+    }, MCP_SHUTDOWN_TIMEOUT_MS);
+    timer.unref();
+    close().then(finish, finish);
+  });
+}
+
+/** Owns one official MCP Client, its transport, request contexts, and bounded cleanup. */
+export class McpServerClient {
+  private readonly activeRuns = new Set<ActiveMcpRunContext>();
+  private callbackTail: Promise<void> = Promise.resolve();
+  private closePromise: Promise<void> | undefined;
+  private closing = false;
+  private connected = false;
+  private readonly requestContext = new AsyncLocalStorage<ActiveMcpRunContext>();
+
+  private constructor(
+    private readonly client: Client,
+    private readonly transport: McpClientTransport,
+    private readonly requestTimeoutMs: number,
+  ) {}
+
+  /** Connect one MCP Client with automatic current/legacy protocol negotiation. */
+  static async connect(options: McpServerClientConnectOptions): Promise<McpServerClient> {
+    const transport = createMcpTransport(options);
+    if (transport instanceof StdioClientTransport && options.onStderr !== undefined) {
+      transport.stderr?.on("data", (chunk: Buffer | string) =>
+        options.onStderr?.(chunk.toString()),
+      );
+    }
+
+    const clientOptions: ClientOptions = {
+      capabilities: {
+        elicitation: { form: {}, url: {} },
+        roots: {},
+        sampling: {},
+      },
+      inputRequired: { maxRounds: 10 },
+      listMaxPages: MAX_MCP_LIST_PAGES,
+      versionNegotiation: {
+        mode: "auto",
+        probe: { timeoutMs: options.connectTimeoutMs },
+      },
+    };
+    if (options.listChanged !== undefined) clientOptions.listChanged = options.listChanged;
+    const client = new Client(options.clientInfo, clientOptions);
+    const owner = new McpServerClient(client, transport, options.requestTimeoutMs);
+    if (options.onError !== undefined) client.onerror = options.onError;
+    client.onclose = () => {
+      const wasUnexpected = owner.connected && !owner.closing;
+      owner.connected = false;
+      if (wasUnexpected) options.onConnectionClose?.();
+    };
+    owner.registerRequestHandlers();
+
+    try {
+      await client.connect(transport, {
+        maxTotalTimeout: options.connectTimeoutMs,
+        timeout: options.connectTimeoutMs,
+      });
+      owner.connected = true;
+      return owner;
+    } catch (cause) {
+      await owner.close();
+      throw cause;
+    }
+  }
+
+  /** Negotiated protocol revision for the connected MCP Server. */
+  get negotiatedProtocolVersion(): string | undefined {
+    return this.client.getNegotiatedProtocolVersion();
+  }
+
+  /** Negotiated current or legacy protocol era. */
+  get protocolEra(): "modern" | "legacy" | undefined {
+    return this.client.getProtocolEra();
+  }
+
+  /** Server Instructions reported during connection. */
+  get instructions(): string | undefined {
+    return this.client.getInstructions();
+  }
+
+  /** Direct stdio child process identifier, when this Client owns one. */
+  get processId(): number | undefined {
+    if (!(this.transport instanceof StdioClientTransport)) return undefined;
+    return this.transport.pid ?? undefined;
+  }
+
+  /** Run one MCP operation with isolated Pi context, cancellation, progress, and timeout reset. */
+  async run<Result, PiContext = undefined>(
+    operation: (client: Client, requestOptions: RequestOptions) => Promise<Result>,
+    options: McpServerRunOptions<PiContext> = {},
+  ): Promise<Result> {
+    if (this.closing || !this.connected) {
+      throw new Error("Pi MCP Client is not connected");
+    }
+    if (options.callbacks !== undefined) {
+      const predecessor = this.callbackTail;
+      let release: () => void = () => undefined;
+      this.callbackTail = new Promise<void>((resolveTail) => {
+        release = resolveTail;
+      });
+      await predecessor;
+      try {
+        return await this.runOperation(operation, options);
+      } finally {
+        release();
+      }
+    }
+    return this.runOperation(operation, options);
+  }
+
+  private async runOperation<Result, PiContext = undefined>(
+    operation: (client: Client, requestOptions: RequestOptions) => Promise<Result>,
+    options: McpServerRunOptions<PiContext>,
+  ): Promise<Result> {
+    const context: ActiveMcpRunContext = {};
+    if (options.callbacks !== undefined) {
+      // SAFETY: The callback and piContext originate from the same generic run invocation and stay paired in this private context.
+      context.callbacks = options.callbacks as McpServerRequestCallbacks<unknown>;
+    }
+    if (options.piContext !== undefined) context.piContext = options.piContext;
+    const requestOptions: RequestOptions = {
+      resetTimeoutOnProgress: true,
+      timeout: this.requestTimeoutMs,
+    };
+    if (options.onProgress !== undefined) requestOptions.onprogress = options.onProgress;
+    if (options.signal !== undefined) requestOptions.signal = options.signal;
+    if (options.callbacks !== undefined) this.activeRuns.add(context);
+    try {
+      return await this.requestContext.run(context, () => operation(this.client, requestOptions));
+    } finally {
+      if (options.callbacks !== undefined) this.activeRuns.delete(context);
+    }
+  }
+
+  /** Close the Client once, forcing its owned stdio process after five seconds. */
+  close(): Promise<void> {
+    if (this.closePromise !== undefined) return this.closePromise;
+    this.closing = true;
+    const processId = this.processId;
+    this.closePromise = closeBeforeDeadline(() => this.client.close(), processId).finally(() => {
+      this.connected = false;
+    });
+    return this.closePromise;
+  }
+
+  private registerRequestHandlers(): void {
+    this.client.setRequestHandler("sampling/createMessage", (request, clientContext) => {
+      const context = this.resolveRequestContext("sampling/createMessage");
+      const callback = context.callbacks?.onSampling;
+      if (callback === undefined) {
+        throw new ProtocolError(
+          ProtocolErrorCode.InternalError,
+          "Pi MCP sampling callback is unavailable",
+        );
+      }
+      return callback(request, context.piContext, clientContext);
+    });
+    this.client.setRequestHandler("elicitation/create", (request, clientContext) => {
+      const context = this.resolveRequestContext("elicitation/create");
+      const callback = context.callbacks?.onElicitation;
+      if (callback === undefined) {
+        throw new ProtocolError(
+          ProtocolErrorCode.InternalError,
+          "Pi MCP elicitation callback is unavailable",
+        );
+      }
+      return callback(request, context.piContext, clientContext);
+    });
+    this.client.setRequestHandler("roots/list", (request, clientContext) => {
+      const context = this.resolveRequestContext("roots/list");
+      const callback = context.callbacks?.onListRoots;
+      if (callback === undefined) {
+        throw new ProtocolError(
+          ProtocolErrorCode.InternalError,
+          "Pi MCP roots callback is unavailable",
+        );
+      }
+      return callback(request, context.piContext, clientContext);
+    });
+  }
+
+  private resolveRequestContext(method: string): ActiveMcpRunContext {
+    const current = this.requestContext.getStore();
+    if (current !== undefined) return current;
+    if (this.activeRuns.size === 1) {
+      const onlyContext = this.activeRuns.values().next().value;
+      if (onlyContext !== undefined) return onlyContext;
+    }
+    throw new ProtocolError(
+      ProtocolErrorCode.InternalError,
+      `Pi MCP ${method} request context is ambiguous`,
+    );
+  }
+}

package/src/mcp-session-files.ts

@@ -0,0 +1,127 @@
+import { chmod, mkdir, mkdtemp, rm, writeFile } from "node:fs/promises";
+import { join } from "node:path";
+
+/** Maximum stderr and MCP logging bytes retained for one MCP Server. */
+export const MAX_MCP_SERVER_LOG_BYTES = 256 * 1024;
+
+/** Retains only the newest MCP Server stderr and logging bytes. */
+export class RetainedMcpServerLog {
+  private content = Buffer.alloc(0);
+
+  /** Append one stderr or logging chunk in arrival order. */
+  append(chunk: string | Uint8Array): void {
+    const combined = Buffer.concat([this.content, Buffer.from(chunk)]);
+    const overflow = Math.max(0, combined.length - MAX_MCP_SERVER_LOG_BYTES);
+    this.content = overflow === 0 ? combined : Buffer.from(combined.subarray(overflow));
+  }
+
+  /** Return the current bounded stderr and logging tail as UTF-8 text. */
+  read(): string {
+    return this.content.toString("utf8");
+  }
+}
+
+/** One private on-disk tail for an MCP Server's stderr and logging notifications. */
+interface McpServerLogFile {
+  readonly log: RetainedMcpServerLog;
+  readonly path: string;
+}
+
+/** Private Result Spill, unsupported-content, and per-server log files owned by one Pi session. */
+export interface McpSessionFiles {
+  /** Private directory removed when the Pi session shuts down. */
+  readonly directoryPath: string;
+  /** Write complete truncated model-facing output to a Result Spill file. */
+  writeResultSpill(output: string): Promise<string>;
+  /** Write unsupported MCP binary or audio bytes to a private session file. */
+  writeUnsupportedContent(content: Uint8Array, mimeType: string): Promise<string>;
+  /** Append stderr or logging bytes, retaining only the newest 256 KB for that server. */
+  appendServerLog(serverName: string, chunk: string | Uint8Array): Promise<void>;
+  /** Read the current bounded stderr and logging tail for one MCP Server. */
+  readServerLog(serverName: string): Promise<string>;
+  /** Remove all session files after queued writes finish. */
+  close(): Promise<void>;
+}
+
+class McpSessionFileStore implements McpSessionFiles {
+  private closed = false;
+  private closePromise: Promise<void> | undefined;
+  private nextFileIndex = 0;
+  private writeQueue: Promise<void> = Promise.resolve();
+  private readonly serverLogs = new Map<string, McpServerLogFile>();
+
+  /** Create a private file store rooted at a Pi session directory. */
+  constructor(readonly directoryPath: string) {}
+
+  writeResultSpill(output: string): Promise<string> {
+    const path = join(this.directoryPath, `result-spill-${this.nextFileIndex++}.txt`);
+    return this.enqueueMcpSessionFileWrite(async () => {
+      await writeFile(path, output, { encoding: "utf8", mode: 0o600 });
+      await chmod(path, 0o600);
+      return path;
+    });
+  }
+
+  writeUnsupportedContent(content: Uint8Array, _mimeType: string): Promise<string> {
+    const path = join(this.directoryPath, `unsupported-content-${this.nextFileIndex++}.bin`);
+    return this.enqueueMcpSessionFileWrite(async () => {
+      await writeFile(path, content, { mode: 0o600 });
+      await chmod(path, 0o600);
+      return path;
+    });
+  }
+
+  appendServerLog(serverName: string, chunk: string | Uint8Array): Promise<void> {
+    const serverLog = this.getMcpServerLogFile(serverName);
+    return this.enqueueMcpSessionFileWrite(async () => {
+      serverLog.log.append(chunk);
+      await writeFile(serverLog.path, serverLog.log.read(), { encoding: "utf8", mode: 0o600 });
+      await chmod(serverLog.path, 0o600);
+    });
+  }
+
+  readServerLog(serverName: string): Promise<string> {
+    return this.enqueueMcpSessionFileWrite(
+      async () => this.serverLogs.get(serverName)?.log.read() ?? "",
+    );
+  }
+
+  close(): Promise<void> {
+    if (this.closePromise !== undefined) return this.closePromise;
+    this.closed = true;
+    this.closePromise = this.writeQueue.then(
+      () => rm(this.directoryPath, { force: true, recursive: true }),
+      () => rm(this.directoryPath, { force: true, recursive: true }),
+    );
+    return this.closePromise;
+  }
+
+  private getMcpServerLogFile(serverName: string): McpServerLogFile {
+    const existing = this.serverLogs.get(serverName);
+    if (existing !== undefined) return existing;
+    const serverLog = {
+      log: new RetainedMcpServerLog(),
+      path: join(this.directoryPath, `server-log-${this.nextFileIndex++}.log`),
+    };
+    this.serverLogs.set(serverName, serverLog);
+    return serverLog;
+  }
+
+  private enqueueMcpSessionFileWrite<T>(write: () => Promise<T>): Promise<T> {
+    if (this.closed) return Promise.reject(new Error("Pi MCP: session files are closed"));
+    const result = this.writeQueue.then(write);
+    this.writeQueue = result.then(
+      () => undefined,
+      () => undefined,
+    );
+    return result;
+  }
+}
+
+/** Create a mode-safe private directory for MCP Result Spills, content, and server logs. */
+export async function createMcpSessionFiles(sessionDirectory: string): Promise<McpSessionFiles> {
+  await mkdir(sessionDirectory, { mode: 0o700, recursive: true });
+  const directoryPath = await mkdtemp(join(sessionDirectory, "pi-mcp-"));
+  await chmod(directoryPath, 0o700);
+  return new McpSessionFileStore(directoryPath);
+}