@fitchmultz/pi-mcp-adapter 5.0.0

package/session-recovery.ts

@@ -0,0 +1,316 @@
+// session-recovery.ts
+//
+// Streamable HTTP session recovery.
+//
+// Per the MCP spec (Streamable HTTP transport):
+//   "When a client receives HTTP 404 in response to a request containing an
+//   Mcp-Session-Id, it MUST start a new session by sending a new
+//   InitializeRequest without a session ID attached."
+//
+// A 404 for a request that carried a session ID is therefore the spec's own
+// definition of "this session no longer exists" (e.g. the remote server
+// process restarted and lost its in-memory session table). Because the spec
+// requires the server to reject the request *before* processing it, retrying
+// the same call against a freshly initialized session cannot double-execute
+// the original request.
+//
+// This module intentionally does NOT:
+//   - match broad error messages without a prior session id
+//   - match generic HTTP 400 responses, which are ambiguous and can mean
+//     many things other than "your session is gone"
+//   - treat generic -32000/ConnectionClosed errors as session expiry
+//   - treat AbortError/cancellation as a session failure
+import { AsyncLocalStorage } from "node:async_hooks";
+import { SdkHttpError, SdkErrorCode, ProtocolError, InsufficientScopeError, UnauthorizedError, type FetchLike, type Transport } from "@modelcontextprotocol/client";
+import { logger } from "./logger.ts";
+import { abortable, throwIfAborted } from "./abort.ts";
+import { isServerDisabled, isNonInteractiveOAuth, type McpConfig } from "./types.ts";
+import type { McpServerManager, ServerConnection } from "./server-manager.ts";
+import { supportsOAuth } from "./mcp-auth-flow.ts";
+
+const toolTransportFailures = new WeakSet<object>();
+export function isToolTransportFailure(error: unknown): boolean {
+  return error instanceof Error && toolTransportFailures.has(error);
+}
+type ToolSend = {
+  id: string | number;
+  streaming: boolean;
+  responseReceived: boolean;
+  bodyFailure: Error | undefined;
+  otherFailure: Error | undefined;
+  finishSend: () => void;
+};
+const toolSend = new AsyncLocalStorage<ToolSend | undefined>();
+
+/** Retain fetch/body-error origin without parsing or buffering the SSE response. */
+export const trackToolTransportFailure: FetchLike = async (input, init) => {
+  const pending = toolSend.getStore();
+  if (pending) pending.bodyFailure = undefined;
+  try {
+    const response = await fetch(input, init);
+    if (!pending || !response.ok || !response.body
+      || response.headers.get("content-type")?.split(";", 1)[0]?.trim().toLowerCase() !== "text/event-stream") return response;
+    pending.streaming = true;
+    // The native end hook has no error argument. Only a failed body read proves loss.
+    const reader = response.body.getReader();
+    return new Response(new ReadableStream({
+      async pull(controller) {
+        try {
+          const { done, value } = await reader.read();
+          if (done) { controller.close(); reader.releaseLock(); }
+          else controller.enqueue(value);
+        } catch (error) {
+          // Native intentional closure suppresses onRequestStreamEnd.
+          if (init?.signal?.aborted) pending.finishSend();
+          else if (error instanceof Error) pending.bodyFailure = error;
+          controller.error(error);
+          reader.releaseLock();
+        }
+      },
+      cancel: reason => reader.cancel(reason).finally(() => reader.releaseLock()),
+    }), response);
+  } catch (error) {
+    if (error instanceof Error && !init?.signal?.aborted && init?.method === "POST"
+      && new Headers(init.headers).get("Mcp-Method") === "tools/call") {
+      toolTransportFailures.add(error);
+    }
+    throw error;
+  }
+};
+
+/** Observe only the originating tool send, not OAuth or catalog requests. */
+export function trackToolHttpFailures(transport: Transport): void {
+  const onmessage = transport.onmessage;
+  transport.onmessage = (message, extra) => {
+    const pending = toolSend.getStore();
+    if (pending && "id" in message && message.id === pending.id && ("result" in message || "error" in message)) {
+      pending.responseReceived = true;
+      pending.finishSend();
+    }
+    onmessage?.(message, extra);
+  };
+  const onerror = transport.onerror;
+  transport.onerror = error => {
+    const pending = toolSend.getStore();
+    // Reader failures are observed before the SDK reports them; earlier parser errors are not retryable.
+    if (pending && !pending.bodyFailure) pending.otherFailure ??= error;
+    onerror?.(error);
+  };
+  const send = transport.send.bind(transport);
+  transport.send = async (message, options) => {
+    try {
+      if (!("method" in message && message.method === "tools/call" && "id" in message && options?.requestSignal)) {
+        return await toolSend.run(undefined, () => send(message, options));
+      }
+      const pending: ToolSend = {
+        id: message.id, streaming: false, responseReceived: false,
+        bodyFailure: undefined, otherFailure: undefined,
+        finishSend: () => {},
+      };
+      try {
+        await abortable(new Promise<void>((resolve, reject) => {
+          pending.finishSend = resolve;
+          void toolSend.run(pending, () => send(message, {
+            ...options,
+            onRequestStreamEnd: () => {
+              options.onRequestStreamEnd?.();
+              const error = pending.otherFailure ?? pending.bodyFailure;
+              if (!pending.responseReceived && error) {
+                if (!pending.otherFailure) toolTransportFailures.add(error);
+                reject(error);
+              } else resolve();
+            },
+          })).then(() => { if (!pending.streaming) resolve(); }, reject);
+        }), options.requestSignal);
+      } finally {
+        pending.finishSend = () => {};
+      }
+    } catch (error) {
+      if ("method" in message && message.method === "tools/call"
+        && error instanceof SdkHttpError && error.code === SdkErrorCode.ClientHttpNotImplemented
+        && error.status >= 500 && error.status < 600) {
+        let errorBody = false;
+        try {
+          const body = JSON.parse(String(error.data.text));
+          errorBody = body?.jsonrpc === "2.0" && body.error !== undefined;
+        } catch {}
+        if (!errorBody) toolTransportFailures.add(error);
+      }
+      throw error;
+    }
+  };
+}
+
+/**
+ * True when `err` is a stale Streamable HTTP session signal for a request
+ * sent while carrying an `Mcp-Session-Id`: the spec's 404 transport
+ * response, or the narrowly-known `-32000 Server not initialized` protocol
+ * gate response some servers emit before dispatching to a handler.
+ *
+ * `hadSessionId` must reflect the transport's session id from *before* the
+ * call that produced `err` was made. The SDK does not
+ * clear `transport.sessionId` on a 404 response, so callers must capture it
+ * before the call rather than rely on catch-time transport state.
+ */
+const SERVER_NOT_INITIALIZED_MCP_MESSAGES = new Set([
+  "Server not initialized",
+  "Bad Request: Server not initialized",
+]);
+
+export function isTerminatedSession(err: unknown, hadSessionId: boolean): boolean {
+  if (!hadSessionId) return false;
+  if (err instanceof SdkHttpError) {
+    return err.status === 404
+      || (err.status === 400
+        && /"code"\s*:\s*-32000/.test(err.message)
+        && /"message"\s*:\s*"Bad Request: Server not initialized"/.test(err.message));
+  }
+  return err instanceof ProtocolError
+    && err.code === -32000
+    && SERVER_NOT_INITIALIZED_MCP_MESSAGES.has(err.message);
+}
+
+function hasSessionId(connection: ServerConnection): boolean {
+  // Only StreamableHTTPClientTransport exposes `sessionId`; stdio/SSE
+  // transports (and test doubles that omit `transport` entirely) simply
+  // read as `undefined` here.
+  const transport = connection.transport as { sessionId?: string } | undefined;
+  return transport?.sessionId != null;
+}
+
+export class SessionRecoveryAuthRequiredError extends Error {
+  constructor(readonly serverName: string, readonly authMessage?: string) {
+    super(authMessage ?? `MCP server "${serverName}" requires OAuth authentication after reconnect.`);
+    this.name = "SessionRecoveryAuthRequiredError";
+  }
+}
+
+export interface AuthChallengeContext {
+  connection: ServerConnection;
+  error: Error;
+}
+
+export interface SessionRecoveryDeps {
+  manager: McpServerManager;
+  config: McpConfig;
+  signal?: AbortSignal;
+  /** Only replay-safe tool calls may opt into a same-client retry of a proven modern POST transport failure. */
+  retryOnTransportFailure?: boolean;
+  onNeedsAuth?: (serverName: string, signal?: AbortSignal, challenge?: AuthChallengeContext) => Promise<ServerConnection | undefined>;
+}
+
+/**
+ * Retry once on the same modern client when explicitly enabled, or reconnect
+ * an expired legacy session once. Neither path retries its second failure.
+ */
+export async function withSessionRecovery<T>(
+  deps: SessionRecoveryDeps,
+  serverName: string,
+  fn: (conn: ServerConnection) => Promise<T>,
+): Promise<T> {
+  if (isServerDisabled(deps.config.mcpServers[serverName])) {
+    throw new Error(`MCP server "${serverName}" is disabled`);
+  }
+  throwIfAborted(deps.signal);
+  let connection = deps.manager.getConnection(serverName);
+  if (!connection) {
+    throw new Error(`Server "${serverName}" is not connected`);
+  }
+
+  const reconnect = async (stale: ServerConnection, error: unknown): Promise<ServerConnection> => {
+    const definition = deps.config.mcpServers[serverName];
+    if (!definition) throw error;
+    throwIfAborted(deps.signal);
+    let fresh = stale;
+    if (fresh.status !== "needs-auth") {
+      fresh = deps.signal
+        ? await deps.manager.reconnect(serverName, definition, stale, deps.signal)
+        : await deps.manager.reconnect(serverName, definition, stale);
+    }
+    throwIfAborted(deps.signal);
+    if (fresh.status === "needs-auth" && deps.onNeedsAuth) {
+      fresh = await abortable(deps.onNeedsAuth(serverName, deps.signal), deps.signal) ?? fresh;
+      throwIfAborted(deps.signal);
+    }
+    if (fresh.status === "needs-auth") throw new SessionRecoveryAuthRequiredError(serverName);
+    if (fresh.status !== "connected") throw error;
+    return fresh;
+  };
+
+  const isRecoverableAuth = (error: unknown, failed: ServerConnection, activating = false): error is Error => {
+    const definition = deps.config.mcpServers[serverName];
+    if (!definition || isServerDisabled(definition) || !supportsOAuth(definition)) return false;
+    if (error instanceof SdkHttpError) {
+      // A provider's exhausted native 401/M2M retry is terminal, not another consent cycle.
+      return failed.oauthProvider === false && error.status === 401 && error.code !== SdkErrorCode.ClientHttpAuthentication;
+    }
+    if (failed.oauthProvider && isNonInteractiveOAuth(definition.oauth)) return false;
+    return error instanceof InsufficientScopeError || (activating && error instanceof UnauthorizedError);
+  };
+
+  const terminalDispatch = async (fresh: ServerConnection): Promise<T> => {
+    throwIfAborted(deps.signal);
+    try {
+      return await fn(fresh);
+    } catch (error) {
+      throwIfAborted(deps.signal);
+      if (isRecoverableAuth(error, fresh)) throw new SessionRecoveryAuthRequiredError(serverName);
+      throw error;
+    }
+  };
+
+  const recoverAuth = async (failed: ServerConnection, error: Error): Promise<T> => {
+    const definition = deps.config.mcpServers[serverName]!;
+    const current = deps.manager.getConnection(serverName);
+    if (current && current !== failed && failed.oauthProvider !== false) {
+      return terminalDispatch(current.status === "connected" ? current : await reconnect(current, error));
+    }
+
+    if (failed.oauthProvider === false) {
+      // One anonymous rejection may activate the provider without asking the user to sign in.
+      failed = await deps.manager.reconnect(serverName, definition, failed, deps.signal);
+      throwIfAborted(deps.signal);
+      if (failed.status === "connected") {
+        try {
+          return await fn(failed);
+        } catch (providerError) {
+          throwIfAborted(deps.signal);
+          if (!isRecoverableAuth(providerError, failed, true)) throw providerError;
+          error = providerError;
+        }
+      }
+    }
+    if (isNonInteractiveOAuth(definition.oauth)) throw error;
+    const fresh = await abortable(deps.onNeedsAuth?.(serverName, deps.signal, { connection: failed, error }) ?? Promise.resolve(undefined), deps.signal);
+    throwIfAborted(deps.signal);
+    if (!fresh || fresh === failed || fresh.status !== "connected") throw new SessionRecoveryAuthRequiredError(serverName);
+    return terminalDispatch(fresh);
+  };
+
+  // Nothing has been dispatched yet: wait for replacement/auth begun during UI preparation.
+  if (connection.status !== "connected") {
+    connection = await reconnect(connection, new Error(`Server "${serverName}" is not connected`));
+  }
+  const hadSessionId = hasSessionId(connection);
+
+  try {
+    return await fn(connection);
+  } catch (err) {
+    throwIfAborted(deps.signal);
+    if (isRecoverableAuth(err, connection)) return recoverAuth(connection, err);
+    if (deps.retryOnTransportFailure === true
+      && deps.config.mcpServers[serverName]?.retryOnTransportFailure === true
+      && connection.client.getProtocolEra?.() === "modern"
+      && isToolTransportFailure(err)) {
+      return terminalDispatch(connection);
+    }
+    if (!isTerminatedSession(err, hadSessionId)) {
+      throw err;
+    }
+
+    logger.debug(`MCP session for "${serverName}" expired; reconnecting`, {
+      server: serverName,
+    });
+    return terminalDispatch(await reconnect(connection, err));
+  }
+}

package/skills/mcp-scripting/SKILL.md

@@ -0,0 +1,38 @@
+---
+name: mcp-scripting
+description: Write mcp_script JavaScript for discovering, inspecting, and calling MCP tools.
+---
+
+# MCP scripting
+
+For multi-call MCP work, write ordinary JavaScript with loops, filtering, chaining, fan-out, or other logic between calls. Run that source with `mcp_script`; it is the primary MCP orchestration surface. For a single MCP search, describe, status check, auth action, or tool call, use `mcp` instead.
+
+Write the source naturally, then pass it as `mcp_script`'s `code` argument:
+
+```js
+const { items } = await tools.search({ query: "search issues", server: "github" });
+const candidate = items[0];
+if (!candidate) return { error: "No matching tool" };
+
+const details = await tools.describe({ path: candidate.path });
+if (details.error) return details;
+
+const result = await tools.call(details.path, { query: "is:open label:bug" });
+if (!result.ok) return result;
+emit({ tool: details.path, completed: true });
+return result.data?.structuredContent ?? result.data;
+```
+
+## Workflow
+
+1. Find candidate tools with `await tools.search({ query, server?, limit?, offset? })`. `limit` defaults to 12 and is capped at 100.
+2. Inspect the exact returned path with `await tools.describe({ path })`. Arguments are described by `inputTypeScript` when representable, otherwise by the original JSON Schema in `inputSchema`; only one is included.
+3. Call it with `tools.call(path, args)`.
+
+Calls resolve to `{ ok: true, data }` or `{ ok: false, error }`; handle failed calls instead of expecting them to stop the script. Successful `data` is the raw MCP result: tool calls usually return `{ content, structuredContent? }`, while resource reads return `{ contents }`. Prefer `structuredContent` when present, and only parse a text block when that server documents JSON text. `emit(value)` adds user-visible output before the final `return` value. `console` output is captured too.
+
+`tools` is a non-enumerable proxy: `Object.keys(tools)` throws. Always use `tools.search` for discovery. When a known flat path is a valid identifier, direct calls such as `tools.github_search_issues(args)` are supported; use bracket syntax for hyphenated names: `tools["server_tool-name"](args)`. `search`, `call`, `describe`, and promise/serialization names (`then`, `catch`, `finally`, `toJSON`, `toString`, `valueOf`) are reserved on the proxy; if a flat path collides with one, call it via `tools.call("exact-path", args)`.
+
+`tools.search` and `tools.describe` are asynchronous and must be awaited. The default script timeout is 30 seconds unless the SDK host overrides or disables it (see the registered `timeoutMs` description); explicit per-script timeouts still apply, and the worker is terminated on timeout or cancellation, including for infinite loops. Every invocation still uses normal lazy connection, authentication, output guarding, and approval gates. Raw call results are cloned into the worker so scripts can reduce them before the model-facing guard; request scoped data instead of unnecessarily large payloads. Result details contain a concise `calls` trace with every search, describe, and call operation; each entry includes its query or path, outcome, and duration.
+
+Use plain JavaScript loops and Promise utilities for composition. The sandbox has no Node, filesystem, or network globals such as `process`, `Buffer`, or `fetch`. Fluent helpers such as `tools.find(...).one()`, `tools.parallel(...)`, and `tools.retry(...)` are not provided.

package/state.ts

@@ -0,0 +1,64 @@
+import type { ExtensionContext } from "@earendil-works/pi-coding-agent";
+import type { ConsentManager } from "./consent-manager.ts";
+import type { McpLifecycleManager } from "./lifecycle.ts";
+import type { McpServerManager } from "./server-manager.ts";
+import type { AuthStorageOptions } from "./mcp-auth.ts";
+import type { ToolMetadata, PromptMetadata, McpConfig, McpAdapterOptions, UiSessionMessages, UiStreamSummary } from "./types.ts";
+import type { UiResourceHandler } from "./ui-resource-handler.ts";
+import type { UiServerHandle } from "./ui-server.ts";
+import type { McpRuntimeOwner } from "./runtime-owner.ts";
+import type { McpOAuthRuntime } from "./mcp-auth-flow.ts";
+import type { McpStatusEventBus } from "./mcp-status.ts";
+
+export interface CompletedUiSession {
+  serverName: string;
+  toolName: string;
+  completedAt: Date;
+  reason: string;
+  messages: UiSessionMessages;
+  stream?: UiStreamSummary;
+}
+
+export type SendMessageFn = (
+  message: {
+    customType: string;
+    content: Array<{ type: "text"; text: string }>;
+    display?: string;
+    details?: unknown;
+  },
+  options?: { triggerTurn?: boolean }
+) => void;
+
+export interface McpExtensionState {
+  owner: McpRuntimeOwner;
+  manager: McpServerManager;
+  lifecycle: McpLifecycleManager;
+  toolMetadata: Map<string, ToolMetadata[]>;
+  /** Resource counts retained separately because tool metadata includes resource tools. */
+  resourceCounts: Map<string, number>;
+  promptMetadata: Map<string, PromptMetadata[]>;
+  /** Servers whose prompt inventory came from successful live discovery. */
+  promptMetadataLive: Set<string>;
+  serverInstructions: Map<string, string>;
+  config: McpConfig;
+  programmaticConfig?: boolean;
+  outputDirectory?: string;
+  onToolCall?: McpAdapterOptions["onToolCall"];
+  oauthRuntime: McpOAuthRuntime;
+  authStorageOptions: AuthStorageOptions;
+  metadataCacheEnabled: boolean;
+  failureTracker: Map<string, number>;
+  failureMessages: Map<string, string>;
+  /** Session-only approvals keyed by server and original tool name. */
+  approvedToolCalls: Map<string, true>;
+  uiResourceHandler: UiResourceHandler;
+  consentManager: ConsentManager;
+  uiServer: UiServerHandle | null;
+  completedUiSessions: CompletedUiSession[];
+  activeScripts: number;
+  openBrowser: (url: string) => Promise<void>;
+  ui?: ExtensionContext["ui"];
+  sendMessage?: SendMessageFn;
+  onToolMetadataUpdated?: (serverName: string, reason: string) => void | Promise<void>;
+  statusEvents?: McpStatusEventBus;
+}

package/tool-approval.ts

@@ -0,0 +1,78 @@
+import { abortable } from "./abort.ts";
+import { combineAbortSignals } from "./runtime-owner.ts";
+import type { McpExtensionState } from "./state.ts";
+import {
+  getToolNameCandidates,
+  matchesToolPattern,
+  resolveToolPrefix,
+  type McpConfig,
+  type ToolMetadata,
+} from "./types.ts";
+import { sanitizeTerminalText } from "./utils.ts";
+
+export type ToolCallApprovalResult =
+  | { ok: true }
+  | { ok: false; reason: "denied" | "approval_required_headless" };
+
+export function isToolCallApprovalRequired(
+  config: McpConfig,
+  serverName: string,
+  toolMeta: Pick<ToolMetadata, "originalName">,
+): boolean {
+  const definition = config.mcpServers[serverName];
+  const approval = definition?.approveTools !== undefined
+    ? definition.approveTools
+    : config.settings?.approveTools;
+
+  if (approval === true) return true;
+  if (!Array.isArray(approval) || approval.length === 0) return false;
+
+  const prefix = resolveToolPrefix(definition, config.settings?.toolPrefix);
+  return matchesToolPattern(
+    getToolNameCandidates(toolMeta.originalName, serverName, prefix),
+    approval,
+  );
+}
+
+export async function ensureToolCallApproved(
+  state: McpExtensionState,
+  serverName: string,
+  toolMeta: ToolMetadata,
+  args: Record<string, unknown> | undefined,
+  signal?: AbortSignal,
+): Promise<ToolCallApprovalResult> {
+  if (!isToolCallApprovalRequired(state.config, serverName, toolMeta)) {
+    return { ok: true };
+  }
+
+  const cacheKey = `${serverName}\u0000${toolMeta.originalName}`;
+  if (state.approvedToolCalls.has(cacheKey)) {
+    return { ok: true };
+  }
+
+  if (!state.ui) {
+    return { ok: false, reason: "approval_required_headless" };
+  }
+
+  const json = JSON.stringify(args ?? {}, null, 2);
+  const sanitized = sanitizeTerminalText(json);
+  const preview = sanitized.length > 500 ? `${sanitized.slice(0, 500)}...` : sanitized;
+  const title = `MCP: ${sanitizeTerminalText(serverName)} wants to run ${sanitizeTerminalText(toolMeta.originalName)}`;
+  const ownedSignal = combineAbortSignals(state.owner?.signal, signal);
+  const decision = await abortable(
+    state.ui.select(
+      `${title}\n\nArguments:\n${preview}`,
+      ["Allow once", "Allow for session", "Deny"],
+    ),
+    ownedSignal,
+  );
+
+  if (decision === "Allow once") {
+    return { ok: true };
+  }
+  if (decision === "Allow for session") {
+    state.approvedToolCalls.set(cacheKey, true);
+    return { ok: true };
+  }
+  return { ok: false, reason: "denied" };
+}