@herbertgao/pi-extensions 2026.8.5 → 2026.8.7

package/node_modules/pi-mcp-adapter/request-headers-command.ts

@@ -0,0 +1,330 @@
+import { spawn, spawnSync, type ChildProcess } from "node:child_process";
+import { randomUUID } from "node:crypto";
+import type { FetchLike } from "@modelcontextprotocol/client";
+import type { HttpRequestHeadersCommand } from "./types.ts";
+import { interpolateEnvVars } from "./utils.ts";
+
+const DEFAULT_TIMEOUT_MS = 10_000;
+const MAX_OUTPUT_BYTES = 64 * 1024;
+const USE_PROCESS_GROUP = process.platform !== "win32";
+const CLEANUP_TOKEN_ENV = "PI_MCP_REQUEST_HEADERS_CLEANUP_TOKEN";
+// Busy hosts can exceed spawnSync's 1 MiB default when `ps axeww` dumps each
+// process environment. Keep discovery below a bounded cap instead of letting
+// `ps` get SIGTERM'd and reporting a false cleanup failure.
+const PS_MAX_BUFFER_BYTES = 8 * 1024 * 1024;
+
+function isNoSuchProcessError(error: unknown): boolean {
+  return typeof error === "object" && error !== null && "code" in error && (error as NodeJS.ErrnoException).code === "ESRCH";
+}
+
+function runPosixPs(args: string[]): { status: number | null; signal: NodeJS.Signals | null; stdout: string } {
+  if (process.env.PI_MCP_ADAPTER_TEST_FAIL_PS === "1") return { status: 1, signal: null, stdout: "" };
+  const result = spawnSync("ps", args, { encoding: "utf8", maxBuffer: PS_MAX_BUFFER_BYTES });
+  return { status: result.status, signal: result.signal, stdout: result.stdout };
+}
+
+function psFailureReason(result: { status: number | null; signal: NodeJS.Signals | null }): string {
+  return result.status === null
+    ? `ps was killed by signal ${result.signal ?? "unknown"}`
+    : `ps exited with code ${result.status}`;
+}
+
+function collectPosixProcessPids(rootPid: number, cleanupToken?: string): number[] {
+  const result = runPosixPs(["axeww", "-o", "pid=,ppid=,command="]);
+  if (result.status !== 0) {
+    throw new Error(`HTTP request headers command cleanup failed: ${psFailureReason(result)}`);
+  }
+
+  const childrenByParent = new Map<number, number[]>();
+  const processPids = new Set<number>();
+  const needle = cleanupToken ? `${CLEANUP_TOKEN_ENV}=${cleanupToken}` : undefined;
+  for (const line of result.stdout.split("\n")) {
+    const match = /^\s*(\d+)\s+(\d+)(?:\s+(.*))?$/.exec(line);
+    if (!match) continue;
+    const pid = Number(match[1]);
+    const ppid = Number(match[2]);
+    if (!Number.isInteger(pid) || !Number.isInteger(ppid)) continue;
+    const children = childrenByParent.get(ppid);
+    if (children) children.push(pid);
+    else childrenByParent.set(ppid, [pid]);
+    if (needle && pid !== process.pid && match[3]?.includes(needle)) processPids.add(pid);
+  }
+
+  const stack = [...(childrenByParent.get(rootPid) ?? [])];
+  while (stack.length > 0) {
+    const pid = stack.pop()!;
+    processPids.add(pid);
+    stack.push(...(childrenByParent.get(pid) ?? []));
+  }
+  return [...processPids];
+}
+
+function assertPosixProcessDiscoveryAvailable(): void {
+  const result = runPosixPs(["axeww", "-o", "pid=,ppid=,command="]);
+  if (result.status !== 0) {
+    throw new Error(`HTTP request headers command cleanup failed: ${psFailureReason(result)}`);
+  }
+}
+
+function isTaskkillNoSuchProcess(result: ReturnType<typeof spawnSync>): boolean {
+  return `${result.stdout ?? ""}\n${result.stderr ?? ""}`.toLowerCase().includes("not found");
+}
+
+function signalPid(pid: number, signal: NodeJS.Signals): void {
+  try {
+    process.kill(pid, signal);
+  } catch (error) {
+    if (!isNoSuchProcessError(error)) throw error;
+  }
+}
+
+function signalProcessGroup(pid: number, signal: NodeJS.Signals): void {
+  try {
+    process.kill(-pid, signal);
+  } catch (error) {
+    if (!isNoSuchProcessError(error)) throw error;
+  }
+}
+
+function killRequestHeadersCommand(child: ChildProcess, trackedPosixDescendantPids = new Set<number>(), cleanupToken?: string): void {
+  if (process.platform === "win32" && child.pid !== undefined) {
+    const result = spawnSync("taskkill", ["/pid", String(child.pid), "/T", "/F"], {
+      encoding: "utf8",
+      windowsHide: true,
+    });
+    if (result.status === 0 || isTaskkillNoSuchProcess(result)) return;
+    throw new Error(`HTTP request headers command cleanup failed: taskkill exited with code ${result.status ?? "unknown"}`);
+  }
+
+  if (USE_PROCESS_GROUP && child.pid !== undefined) {
+    const frozenPids = new Set<number>();
+    let cleanupError: Error | undefined;
+    try {
+      signalProcessGroup(child.pid, "SIGSTOP");
+      for (const pid of trackedPosixDescendantPids) {
+        signalPid(pid, "SIGSTOP");
+        frozenPids.add(pid);
+      }
+      let stablePasses = 0;
+      for (let pass = 0; pass < 16; pass++) {
+        const candidates = collectPosixProcessPids(child.pid, cleanupToken);
+        const newPids = candidates.filter(pid => !frozenPids.has(pid));
+        if (newPids.length === 0) {
+          stablePasses++;
+          if (stablePasses >= 2) return;
+          continue;
+        }
+        stablePasses = 0;
+        for (const pid of newPids) {
+          signalPid(pid, "SIGSTOP");
+          frozenPids.add(pid);
+        }
+      }
+      cleanupError = new Error("HTTP request headers command cleanup failed: descendant process tree did not stabilize");
+    } catch (error) {
+      cleanupError = error instanceof Error ? error : new Error(String(error));
+    } finally {
+      signalProcessGroup(child.pid, "SIGKILL");
+      for (const pid of frozenPids) signalPid(pid, "SIGKILL");
+    }
+    if (cleanupError) throw cleanupError;
+    return;
+  }
+
+  child.kill("SIGKILL");
+}
+
+export interface HttpRequestCommandEnvelope {
+  version: 1;
+  method: string;
+  url: string;
+  bodyBase64: string;
+}
+
+type CommandResult =
+  | { status: "error"; error: Error }
+  | { status: "success"; headers: Headers };
+
+function resolvedCommand(config: HttpRequestHeadersCommand): {
+  command: string;
+  args: string[];
+  env: NodeJS.ProcessEnv;
+  timeoutMs: number;
+} {
+  if (!config || typeof config !== "object" || Array.isArray(config)) {
+    throw new Error("HTTP request headers command must be an object");
+  }
+  if (typeof config.command !== "string" || config.command.trim() === "") {
+    throw new Error("HTTP request headers command requires a non-empty command");
+  }
+  if (config.args !== undefined && (!Array.isArray(config.args) || config.args.some(arg => typeof arg !== "string"))) {
+    throw new Error("HTTP request headers command args must be strings");
+  }
+  if (config.env !== undefined && (
+    typeof config.env !== "object"
+    || Array.isArray(config.env)
+    || Object.values(config.env).some(value => typeof value !== "string")
+  )) {
+    throw new Error("HTTP request headers command env values must be strings");
+  }
+  const timeoutMs = config.timeoutMs ?? DEFAULT_TIMEOUT_MS;
+  if (!Number.isInteger(timeoutMs) || timeoutMs <= 0 || timeoutMs > 60_000) {
+    throw new Error("HTTP request headers command timeoutMs must be an integer between 1 and 60000");
+  }
+  return {
+    command: interpolateEnvVars(config.command),
+    args: (config.args ?? []).map(interpolateEnvVars),
+    env: {
+      ...process.env,
+      ...Object.fromEntries(
+        Object.entries(config.env ?? {}).map(([key, value]) => [key, interpolateEnvVars(value)]),
+      ),
+    },
+    timeoutMs,
+  };
+}
+
+async function invokeRequestHeadersCommand(
+  config: HttpRequestHeadersCommand,
+  envelope: HttpRequestCommandEnvelope,
+  signal: AbortSignal,
+): Promise<Headers> {
+  const resolved = resolvedCommand(config);
+  if (USE_PROCESS_GROUP) assertPosixProcessDiscoveryAvailable();
+  return new Promise<Headers>((resolve, reject) => {
+    let stdout = Buffer.alloc(0);
+    let settled = false;
+    const cleanupToken = randomUUID();
+    const child = spawn(resolved.command, resolved.args, {
+      env: { ...resolved.env, [CLEANUP_TOKEN_ENV]: cleanupToken },
+      stdio: ["pipe", "pipe", "ignore"],
+      windowsHide: true,
+      detached: USE_PROCESS_GROUP,
+    });
+
+    const trackedPosixDescendantPids = new Set<number>();
+    let trackingError: Error | undefined;
+    const trackPosixDescendants = () => {
+      if (!USE_PROCESS_GROUP || child.pid === undefined || settled || trackingError) return;
+      try {
+        for (const pid of collectPosixProcessPids(child.pid, cleanupToken)) trackedPosixDescendantPids.add(pid);
+      } catch (error) {
+        trackingError = error instanceof Error ? error : new Error(String(error));
+      }
+    };
+    const descendantTracker = USE_PROCESS_GROUP ? setInterval(trackPosixDescendants, 50) : undefined;
+    descendantTracker?.unref();
+
+    const finish = (result: CommandResult) => {
+      if (settled) return;
+      settled = true;
+      clearTimeout(timer);
+      if (descendantTracker) clearInterval(descendantTracker);
+      signal.removeEventListener("abort", abort);
+      if (result.status === "error") reject(result.error);
+      else resolve(result.headers);
+    };
+    const finishAfterKill = (result: CommandResult) => {
+      try {
+        killRequestHeadersCommand(child, trackedPosixDescendantPids, cleanupToken);
+        finish(result);
+      } catch (cleanupError) {
+        finish({ status: "error", error: cleanupError instanceof Error ? cleanupError : new Error(String(cleanupError)) });
+      }
+    };
+    const failAfterKill = (message: string) => {
+      finishAfterKill({ status: "error", error: trackingError ?? new Error(message) });
+    };
+    const abort = () => {
+      failAfterKill("HTTP request headers command aborted");
+    };
+    const timer = setTimeout(() => {
+      failAfterKill(`HTTP request headers command timed out after ${resolved.timeoutMs}ms`);
+    }, resolved.timeoutMs);
+
+    signal.addEventListener("abort", abort, { once: true });
+    if (signal.aborted) {
+      abort();
+      return;
+    }
+
+    child.on("error", () => finish({ status: "error", error: new Error("HTTP request headers command failed to start") }));
+    child.stdout.on("data", (chunk: Buffer | string) => {
+      if (settled) return;
+      stdout = Buffer.concat([stdout, Buffer.from(chunk)]);
+      if (stdout.byteLength > MAX_OUTPUT_BYTES) {
+        failAfterKill("HTTP request headers command output exceeded 64 KiB");
+      }
+    });
+    child.on("close", code => {
+      if (settled) return;
+      if (trackingError) {
+        failAfterKill(trackingError.message);
+        return;
+      }
+      if (code !== 0) {
+        failAfterKill(`HTTP request headers command exited with code ${code ?? "unknown"}`);
+        return;
+      }
+      let parsed: unknown;
+      try {
+        parsed = JSON.parse(stdout.toString("utf8"));
+      } catch {
+        finishAfterKill({ status: "error", error: new Error("HTTP request headers command returned invalid JSON") });
+        return;
+      }
+      if (!parsed || typeof parsed !== "object" || Array.isArray(parsed)) {
+        finishAfterKill({ status: "error", error: new Error("HTTP request headers command must return a JSON object") });
+        return;
+      }
+      const entries = Object.entries(parsed);
+      if (entries.some(([, value]) => typeof value !== "string")) {
+        finishAfterKill({ status: "error", error: new Error("HTTP request headers command values must be strings") });
+        return;
+      }
+      try {
+        finishAfterKill({ status: "success", headers: new Headers(entries as Array<[string, string]>) });
+      } catch {
+        finishAfterKill({ status: "error", error: new Error("HTTP request headers command returned an invalid header") });
+      }
+    });
+
+    child.stdin.on("error", () => {});
+    child.stdin.end(JSON.stringify(envelope));
+  });
+}
+
+/** Wrap fetch so a trusted command can derive headers from the exact request. */
+export function createRequestHeadersCommandFetch(
+  config: HttpRequestHeadersCommand,
+  delegate: FetchLike = globalThis.fetch,
+): FetchLike {
+  // Validate static configuration before the first request.
+  resolvedCommand(config);
+  return async (input: URL | RequestInfo, init?: RequestInit): Promise<Response> => {
+    const request = new Request(input, init);
+    const body = Buffer.from(await request.clone().arrayBuffer());
+    const derived = await invokeRequestHeadersCommand(config, {
+      version: 1,
+      method: request.method.toUpperCase(),
+      url: request.url,
+      bodyBase64: body.toString("base64"),
+    }, request.signal);
+    const headers = new Headers(request.headers);
+    derived.forEach((value, name) => headers.set(name, value));
+    return delegate(new URL(request.url), {
+      method: request.method,
+      headers,
+      ...(request.method === "GET" || request.method === "HEAD" ? {} : { body }),
+      signal: request.signal,
+      cache: request.cache,
+      credentials: request.credentials,
+      integrity: request.integrity,
+      keepalive: request.keepalive,
+      mode: request.mode,
+      redirect: request.redirect,
+      referrer: request.referrer,
+      referrerPolicy: request.referrerPolicy,
+    });
+  };
+}

package/node_modules/pi-mcp-adapter/search-ranking.ts

@@ -23,6 +23,26 @@ export interface RankedToolMatch {
   score: number;
 }
 
+type SearchField = Exclude<keyof typeof FIELD_WEIGHTS, "keywords">;
+
+interface PreparedToolSearch {
+  tool: ToolMetadata;
+  fields: Array<[SearchField, string, string[]]>;
+  nameTokens: string[];
+  keywordPhrases: string[];
+  keywordTokens: string[];
+}
+
+interface CachedServerSearch {
+  metadata: ToolMetadata[];
+  searchKeywords: ServerEntry["searchKeywords"];
+  toolPrefix: ReturnType<typeof resolveToolPrefix>;
+  withKeywords?: PreparedToolSearch[];
+  withoutKeywords?: PreparedToolSearch[];
+}
+
+const searchCache = new WeakMap<McpExtensionState, Map<string, CachedServerSearch>>();
+
 /**
  * Resolve the configured searchKeywords entries that apply to a tool.
  * Keys match by original name, prefixed name, or glob — the same candidate
@@ -64,25 +84,37 @@ export function tokenize(value: string): string[] {
   return normalizeSearchText(value).split(/[^a-z0-9]+/).filter(Boolean);
 }
 
-export function scoreToolMatch(tool: ToolMetadata, server: string, query: string, keywords?: string[]): number | null {
-  const normalizedQuery = normalizeSearchText(query).trim();
-  const queryTokens = tokenize(query);
-  if (queryTokens.length === 0) return null;
-
+function prepareToolSearch(tool: ToolMetadata, server: string, keywords?: string[]): PreparedToolSearch {
   const fields = {
     name: normalizeSearchText(tool.name),
     originalName: normalizeSearchText(tool.originalName),
     server: normalizeSearchText(server),
     description: normalizeSearchText(tool.description),
   };
+  const preparedFields = (Object.entries(fields) as Array<[SearchField, string]>)
+    .map(([field, value]): [SearchField, string, string[]] => [field, value, tokenize(value)]);
+  const keywordPhrases = keywords?.map(keyword => normalizeSearchText(keyword).trim()).filter(Boolean) ?? [];
+  return {
+    tool,
+    fields: preparedFields,
+    nameTokens: preparedFields[0]![2],
+    keywordPhrases,
+    keywordTokens: keywordPhrases.flatMap(tokenize),
+  };
+}
+
+function scorePreparedToolMatch(
+  prepared: PreparedToolSearch,
+  normalizedQuery: string,
+  queryTokens: string[],
+): number | null {
   let score = 0;
   let phraseMatched = false;
   let wholeFieldExact = false;
   const matchedTokens = new Set<string>();
 
-  for (const [field, value] of Object.entries(fields) as Array<[keyof typeof FIELD_WEIGHTS, string]>) {
+  for (const [field, value, fieldTokens] of prepared.fields) {
     const weight = FIELD_WEIGHTS[field];
-    const fieldTokens = tokenize(value);
     if (value === normalizedQuery) {
       score += weight * 14;
       phraseMatched = true;
@@ -112,11 +144,10 @@ export function scoreToolMatch(tool: ToolMetadata, server: string, query: string
   // Configured keywords are discrete phrases, so the phrase-level bonus is
   // computed per phrase (best match wins) rather than on a joined string,
   // which would phrase-match queries spanning two unrelated keywords.
-  if (keywords !== undefined && keywords.length > 0) {
+  if (prepared.keywordPhrases.length > 0) {
     const weight = FIELD_WEIGHTS.keywords;
-    const phrases = keywords.map(keyword => normalizeSearchText(keyword).trim()).filter(Boolean);
     let phraseScore = 0;
-    for (const phrase of phrases) {
+    for (const phrase of prepared.keywordPhrases) {
       if (phrase === normalizedQuery) {
         phraseScore = Math.max(phraseScore, weight * 14);
         phraseMatched = true;
@@ -131,15 +162,14 @@ export function scoreToolMatch(tool: ToolMetadata, server: string, query: string
     }
     score += phraseScore;
 
-    const keywordTokens = phrases.flatMap(tokenize);
     for (const token of queryTokens) {
-      if (keywordTokens.includes(token)) {
+      if (prepared.keywordTokens.includes(token)) {
         score += weight * 4;
         matchedTokens.add(token);
-      } else if (keywordTokens.some(keywordToken => keywordToken.startsWith(token) || (keywordToken.length >= MIN_STEM_LENGTH && token.startsWith(keywordToken)))) {
+      } else if (prepared.keywordTokens.some(keywordToken => keywordToken.startsWith(token) || (keywordToken.length >= MIN_STEM_LENGTH && token.startsWith(keywordToken)))) {
         score += weight * 2;
         matchedTokens.add(token);
-      } else if (phrases.some(phrase => phrase.includes(token))) {
+      } else if (prepared.keywordPhrases.some(phrase => phrase.includes(token))) {
         score += weight;
         matchedTokens.add(token);
       }
@@ -151,11 +181,52 @@ export function scoreToolMatch(tool: ToolMetadata, server: string, query: string
 
   score += coverage === 1 ? 25 : Math.round(coverage * 10);
   const firstQueryToken = queryTokens[0];
-  if (firstQueryToken !== undefined && tokenize(fields.name).includes(firstQueryToken)) score += 8;
+  if (firstQueryToken !== undefined && prepared.nameTokens.includes(firstQueryToken)) score += 8;
   if (wholeFieldExact) score += 20;
   return score;
 }
 
+export function scoreToolMatch(tool: ToolMetadata, server: string, query: string, keywords?: string[]): number | null {
+  const normalizedQuery = normalizeSearchText(query).trim();
+  const queryTokens = tokenize(query);
+  if (queryTokens.length === 0) return null;
+  return scorePreparedToolMatch(prepareToolSearch(tool, server, keywords), normalizedQuery, queryTokens);
+}
+
+function getPreparedTools(
+  state: McpExtensionState,
+  serverName: string,
+  metadata: ToolMetadata[],
+  definition: ServerEntry | undefined,
+  globalPrefix: ToolPrefix,
+  includeKeywords: boolean,
+): PreparedToolSearch[] {
+  let stateCache = searchCache.get(state);
+  if (!stateCache) {
+    stateCache = new Map();
+    searchCache.set(state, stateCache);
+  }
+  const toolPrefix = resolveToolPrefix(definition, globalPrefix);
+  let cached = stateCache.get(serverName);
+  if (cached?.metadata !== metadata || cached.searchKeywords !== definition?.searchKeywords || cached.toolPrefix !== toolPrefix) {
+    cached = { metadata, searchKeywords: definition?.searchKeywords, toolPrefix };
+    stateCache.set(serverName, cached);
+  }
+  const cacheKey = includeKeywords ? "withKeywords" : "withoutKeywords";
+  let prepared = cached[cacheKey];
+  if (!prepared) {
+    prepared = metadata.map(tool => prepareToolSearch(
+      tool,
+      serverName,
+      includeKeywords && definition?.searchKeywords !== undefined
+        ? resolveSearchKeywords(definition, tool.originalName, serverName, globalPrefix)
+        : undefined,
+    ));
+    cached[cacheKey] = prepared;
+  }
+  return prepared;
+}
+
 export function rankToolMatches(
   state: McpExtensionState,
   query: string,
@@ -163,18 +234,17 @@ export function rankToolMatches(
   includeKeywords = true,
 ): RankedToolMatch[] {
   const matches: RankedToolMatch[] = [];
+  const normalizedQuery = normalizeSearchText(query).trim();
+  const queryTokens = tokenize(query);
+  if (queryTokens.length === 0) return matches;
   const globalPrefix = state.config.settings?.toolPrefix ?? "server";
   for (const [serverName, metadata] of state.toolMetadata.entries()) {
     if (server && serverName !== server) continue;
     const definition = state.config.mcpServers[serverName];
     if (isServerDisabled(definition)) continue;
-    const hasKeywords = includeKeywords && definition?.searchKeywords !== undefined;
-    for (const tool of metadata) {
-      const keywords = hasKeywords
-        ? resolveSearchKeywords(definition, tool.originalName, serverName, globalPrefix)
-        : undefined;
-      const score = scoreToolMatch(tool, serverName, query, keywords);
-      if (score !== null) matches.push({ server: serverName, tool, score });
+    for (const prepared of getPreparedTools(state, serverName, metadata, definition, globalPrefix, includeKeywords)) {
+      const score = scorePreparedToolMatch(prepared, normalizedQuery, queryTokens);
+      if (score !== null) matches.push({ server: serverName, tool: prepared.tool, score });
     }
   }
   return matches.sort((a, b) => b.score - a.score || a.tool.name.localeCompare(b.tool.name));