@bastani/atomic 0.9.16-alpha.2 → 0.9.16-alpha.4

package/dist/builtin/mcp/sampling-handler.ts

@@ -1,273 +0,0 @@
-import type { ProviderHeaders } from "@bastani/pi-ai";
-import { complete, type Api, type AssistantMessage, type Message, type Model, type TextContent } from "@bastani/pi-ai/compat";
-import { truncateAtWord } from "./utils.js";
-import type { ExtensionUIContext, ModelRegistry } from "@bastani/atomic";
-import type { Client } from "@modelcontextprotocol/sdk/client/index.js";
-import {
-  CreateMessageRequestSchema,
-  type CreateMessageRequest,
-  type CreateMessageResult,
-  type ModelPreferences,
-  type SamplingMessage,
-  type SamplingMessageContentBlock,
-} from "@modelcontextprotocol/sdk/types.js";
-
-export interface SamplingHandlerOptions {
-  serverName: string;
-  autoApprove: boolean;
-  ui?: ExtensionUIContext;
-  modelRegistry: ModelRegistry;
-  getCurrentModel: () => Model<Api> | undefined;
-  getSignal: () => AbortSignal | undefined;
-}
-
-export type ServerSamplingConfig = Omit<SamplingHandlerOptions, "serverName">;
-
-export function registerSamplingHandler(client: Client, options: SamplingHandlerOptions): void {
-  client.setRequestHandler(CreateMessageRequestSchema, (request) => {
-    return handleSamplingRequest(options, request as CreateMessageRequest);
-  });
-}
-
-export async function handleSamplingRequest(
-  options: SamplingHandlerOptions,
-  request: CreateMessageRequest,
-): Promise<CreateMessageResult> {
-  const params = request.params;
-
-  if ("task" in params && params.task) {
-    throw new Error("MCP sampling tasks are not supported");
-  }
-  if (params.includeContext && params.includeContext !== "none") {
-    throw new Error("MCP sampling context inclusion is not supported");
-  }
-  if (params.tools?.length) {
-    throw new Error("MCP sampling tool use is not supported");
-  }
-  if (params.toolChoice) {
-    throw new Error("MCP sampling tool choice is not supported");
-  }
-  if (params.stopSequences?.length) {
-    throw new Error("MCP sampling stop sequences are not supported");
-  }
-
-  const messages = params.messages.map(convertSamplingMessage);
-  const { model, apiKey, headers } = await resolveSamplingModel(options, params.modelPreferences);
-  await confirmSampling(
-    options,
-    "Approve MCP sampling request",
-    formatRequestApproval(options.serverName, `${model.provider}/${model.id}`, params.systemPrompt, messages),
-  );
-
-  const result = await complete(
-    model,
-    {
-      systemPrompt: params.systemPrompt,
-      messages,
-    },
-    {
-      apiKey,
-      headers,
-      maxTokens: params.maxTokens,
-      temperature: params.temperature,
-      metadata: params.metadata as Record<string, unknown> | undefined,
-      signal: options.getSignal(),
-    },
-  );
-
-  const converted = convertAssistantResult(result);
-  await confirmSampling(
-    options,
-    "Return MCP sampling response",
-    formatResponseApproval(options.serverName, converted),
-  );
-  return converted;
-}
-
-function formatRequestApproval(
-  serverName: string,
-  modelName: string,
-  systemPrompt: string | undefined,
-  messages: Message[],
-): string {
-  const lines = [`${serverName} wants to sample ${messages.length} message${messages.length === 1 ? "" : "s"} with ${modelName}.`];
-  if (systemPrompt) {
-    lines.push(`System: ${truncateAtWord(systemPrompt, 400)}`);
-  }
-  for (const [index, message] of messages.entries()) {
-    lines.push(`${index + 1}. ${message.role}: ${truncateAtWord(messageText(message), 400)}`);
-  }
-  return lines.join("\n\n");
-}
-
-function formatResponseApproval(serverName: string, response: CreateMessageResult): string {
-  const text = response.content.type === "text" ? response.content.text : `[${response.content.type} content]`;
-  return `${serverName} will receive this response from ${response.model}:\n\n${truncateAtWord(text, 1000)}`;
-}
-
-function messageText(message: Message): string {
-  if (typeof message.content === "string") return message.content;
-  return message.content.map((block) => {
-    if (block.type === "text") return block.text;
-    if (block.type === "image") return `[image: ${block.mimeType}]`;
-    if (block.type === "thinking") return "[thinking]";
-    if (block.type === "toolCall") return `[tool call: ${block.name}]`;
-    return "[content]";
-  }).join("\n");
-}
-
-async function resolveSamplingModel(
-  options: SamplingHandlerOptions,
-  modelPreferences: ModelPreferences | undefined,
-): Promise<{
-  model: Model<Api>;
-  apiKey?: string;
-  headers?: ProviderHeaders;
-}> {
-  const candidates: Model<Api>[] = [];
-  const availableModels = options.modelRegistry.getAvailable();
-
-  for (const hint of modelPreferences?.hints ?? []) {
-    const normalizedHint = hint.name?.trim().toLowerCase();
-    if (!normalizedHint) continue;
-    for (const model of availableModels) {
-      const searchableNames = [`${model.provider}/${model.id}`, model.id, model.name];
-      if (searchableNames.some((name) => name.toLowerCase().includes(normalizedHint))) {
-        addSamplingCandidate(candidates, model);
-      }
-    }
-  }
-
-  const currentModel = options.getCurrentModel();
-  if (currentModel) addSamplingCandidate(candidates, currentModel);
-
-  for (const model of availableModels) {
-    addSamplingCandidate(candidates, model);
-  }
-
-  const errors: string[] = [];
-  for (const model of candidates) {
-    const auth = await options.modelRegistry.getApiKeyAndHeaders(model);
-    if (auth.ok === false) {
-      errors.push(`${model.provider}/${model.id}: ${auth.error}`);
-      continue;
-    }
-    return {
-      model: auth.baseUrl === undefined || auth.baseUrl === model.baseUrl ? model : { ...model, baseUrl: auth.baseUrl },
-      apiKey: auth.apiKey,
-      headers: auth.headers,
-    };
-  }
-
-  if (errors.length > 0) {
-    throw new Error(`No configured auth for MCP sampling model. ${errors.join("; ")}`);
-  }
-  throw new Error("No Pi model is available for MCP sampling");
-}
-
-function addSamplingCandidate(candidates: Model<Api>[], model: Model<Api>): void {
-  if (!candidates.some((candidate) => candidate.provider === model.provider && candidate.id === model.id)) {
-    candidates.push(model);
-  }
-}
-
-async function confirmSampling(options: SamplingHandlerOptions, title: string, message: string): Promise<void> {
-  if (options.autoApprove) return;
-  if (!options.ui) {
-    throw new Error("MCP sampling requires interactive approval. Set settings.samplingAutoApprove to true to allow it without UI.");
-  }
-  const approved = await options.ui.confirm(title, message);
-  if (!approved) {
-    throw new Error("MCP sampling request was declined");
-  }
-}
-
-function convertSamplingMessage(message: SamplingMessage): Message {
-  const blocks = Array.isArray(message.content) ? message.content : [message.content];
-  if (message.role === "user") {
-    return {
-      role: "user",
-      content: blocks.map(convertUserContent),
-      timestamp: Date.now(),
-    };
-  }
-
-  return {
-    role: "assistant",
-    content: blocks.map(convertAssistantContent),
-    api: "mcp-sampling",
-    provider: "mcp",
-    model: "sampling-request",
-    usage: zeroUsage(),
-    stopReason: "stop",
-    timestamp: Date.now(),
-  };
-}
-
-function convertUserContent(block: SamplingMessageContentBlock): TextContent {
-  if (block.type === "text") {
-    return { type: "text", text: block.text };
-  }
-  throw new Error(`MCP sampling ${block.type} content is not supported`);
-}
-
-function convertAssistantContent(block: SamplingMessageContentBlock): TextContent {
-  if (block.type === "text") {
-    return { type: "text", text: block.text };
-  }
-  throw new Error(`MCP sampling assistant ${block.type} content is not supported`);
-}
-
-function convertAssistantResult(message: AssistantMessage): CreateMessageResult {
-  if (message.stopReason === "error") {
-    throw new Error(message.errorMessage ?? "MCP sampling model call failed");
-  }
-  if (message.stopReason === "aborted") {
-    throw new Error(message.errorMessage ?? "MCP sampling model call was aborted");
-  }
-
-  const text = message.content
-    .map((block) => {
-      if (block.type === "text") return block.text;
-      if (block.type === "thinking") return undefined;
-      throw new Error(`MCP sampling result ${block.type} content is not supported`);
-    })
-    .filter((value): value is string => value !== undefined)
-    .join("\n\n")
-    .trim();
-
-  if (!text) {
-    throw new Error("MCP sampling result did not contain text content");
-  }
-
-  return {
-    role: "assistant",
-    content: { type: "text", text },
-    model: `${message.provider}/${message.model}`,
-    stopReason: mapStopReason(message.stopReason),
-  };
-}
-
-function mapStopReason(reason: AssistantMessage["stopReason"]): CreateMessageResult["stopReason"] {
-  if (reason === "stop") return "endTurn";
-  if (reason === "length") return "maxTokens";
-  if (reason === "toolUse") return "toolUse";
-  return reason;
-}
-
-function zeroUsage(): AssistantMessage["usage"] {
-  return {
-    input: 0,
-    output: 0,
-    cacheRead: 0,
-    cacheWrite: 0,
-    totalTokens: 0,
-    cost: {
-      input: 0,
-      output: 0,
-      cacheRead: 0,
-      cacheWrite: 0,
-      total: 0,
-    },
-  };
-}

package/dist/builtin/mcp/server-manager.ts

@@ -1,372 +0,0 @@
-import { Client } from "@modelcontextprotocol/sdk/client/index.js";
-import { StdioClientTransport } from "@modelcontextprotocol/sdk/client/stdio.js";
-import { StreamableHTTPClientTransport } from "@modelcontextprotocol/sdk/client/streamableHttp.js";
-import { SSEClientTransport } from "@modelcontextprotocol/sdk/client/sse.js";
-import { UnauthorizedError } from "@modelcontextprotocol/sdk/client/auth.js";
-import type { ReadResourceResult } from "@modelcontextprotocol/sdk/types.js";
-import { createChildProcessEnvironment } from "@bastani/atomic";
-import type {
-  McpTool,
-  McpResource,
-  ServerDefinition,
-  ServerStreamResultPatchNotification,
-  Transport,
-} from "./types.ts";
-import { serverStreamResultPatchNotificationSchema } from "./types.ts";
-import { resolveNpxBinary } from "./npx-resolver.ts";
-import { logger } from "./logger.ts";
-import { McpOAuthProvider } from "./mcp-oauth-provider.js";
-import { supportsOAuth } from "./mcp-auth-flow.js";
-import { registerSamplingHandler, type ServerSamplingConfig } from "./sampling-handler.ts";
-import { interpolateEnvRecord, resolveBearerToken, resolveConfigPath } from "./utils.js";
-
-interface ServerConnection {
-  client: Client;
-  transport: Transport;
-  definition: ServerDefinition;
-  tools: McpTool[];
-  resources: McpResource[];
-  lastUsedAt: number;
-  inFlight: number;
-  status: "connected" | "closed" | "needs-auth";
-}
-
-type UiStreamListener = (serverName: string, notification: ServerStreamResultPatchNotification["params"]) => void;
-
-export class McpServerManager {
-  private connections = new Map<string, ServerConnection>();
-  private connectPromises = new Map<string, Promise<ServerConnection>>();
-  private uiStreamListeners = new Map<string, UiStreamListener>();
-  private samplingConfig: ServerSamplingConfig | undefined;
-
-  setSamplingConfig(config: ServerSamplingConfig | undefined): void {
-    this.samplingConfig = config;
-  }
-  
-  async connect(name: string, definition: ServerDefinition): Promise<ServerConnection> {
-    // Dedupe concurrent connection attempts
-    if (this.connectPromises.has(name)) {
-      return this.connectPromises.get(name)!;
-    }
-    
-    // Reuse existing connection if healthy
-    const existing = this.connections.get(name);
-    if (existing?.status === "connected") {
-      existing.lastUsedAt = Date.now();
-      return existing;
-    }
-    
-    const promise = this.createConnection(name, definition);
-    this.connectPromises.set(name, promise);
-    
-    try {
-      const connection = await promise;
-      this.connections.set(name, connection);
-      return connection;
-    } finally {
-      this.connectPromises.delete(name);
-    }
-  }
-  
-  private async createConnection(
-    name: string,
-    definition: ServerDefinition
-  ): Promise<ServerConnection> {
-    const client = this.createClient(name);
-    
-    let transport: Transport;
-    
-    if (definition.command) {
-      let command = definition.command;
-      let args = definition.args ?? [];
-
-      if (command === "npx" || command === "npm") {
-        const resolved = await resolveNpxBinary(command, args);
-        if (resolved) {
-          command = resolved.isJs ? "node" : resolved.binPath;
-          args = resolved.isJs ? [resolved.binPath, ...resolved.extraArgs] : resolved.extraArgs;
-          logger.debug(`${name} resolved to ${resolved.binPath} (skipping npm parent)`);
-        }
-      }
-
-      transport = new StdioClientTransport({
-        command,
-        args,
-        env: resolveEnv(definition.env),
-        cwd: resolveConfigPath(definition.cwd),
-        stderr: definition.debug ? "inherit" : "ignore",
-      });
-    } else if (definition.url) {
-      // HTTP transport with fallback
-      transport = await this.createHttpTransport(definition, name);
-    } else {
-      throw new Error(`Server ${name} has no command or url`);
-    }
-    
-    try {
-      await client.connect(transport);
-      this.attachAdapterNotificationHandlers(name, client);
-
-      // Discover tools and resources
-      const [tools, resources] = await Promise.all([
-        this.fetchAllTools(client),
-        this.fetchAllResources(client),
-      ]);
-      
-      return {
-        client,
-        transport,
-        definition,
-        tools,
-        resources,
-        lastUsedAt: Date.now(),
-        inFlight: 0,
-        status: "connected",
-      };
-    } catch (error) {
-      // Check for UnauthorizedError - server requires OAuth
-      if (error instanceof UnauthorizedError && supportsOAuth(definition)) {
-        // Clean up both client and transport before reporting needs-auth.
-        await client.close().catch(() => {});
-        await transport.close().catch(() => {});
-
-        return {
-          client,
-          transport,
-          definition,
-          tools: [],
-          resources: [],
-          lastUsedAt: Date.now(),
-          inFlight: 0,
-          status: "needs-auth",
-        };
-      }
-      
-      // Clean up both client and transport on any error
-      await client.close().catch(() => {});
-      await transport.close().catch(() => {});
-      throw error;
-    }
-  }
-  
-  private createClient(serverName: string): Client {
-    const client = new Client(
-      { name: `pi-mcp-${serverName}`, version: "1.0.0" },
-      this.samplingConfig ? { capabilities: { sampling: {} } } : undefined,
-    );
-    if (this.samplingConfig) {
-      registerSamplingHandler(client, { ...this.samplingConfig, serverName });
-    }
-    return client;
-  }
-
-  private async createHttpTransport(
-    definition: ServerDefinition, 
-    serverName: string
-  ): Promise<Transport> {
-    const url = new URL(definition.url!);
-    
-    // Build headers first (including any bearer token)
-    const headers = resolveHeaders(definition.headers) ?? {};
-    
-    // For bearer auth, add the token to headers BEFORE creating requestInit
-    if (definition.auth === "bearer") {
-      const token = resolveBearerToken(definition);
-      if (token) {
-        headers["Authorization"] = `Bearer ${token}`;
-      }
-    }
-    
-    // Create request init with headers (Authorization now included for bearer auth)
-    const requestInit = Object.keys(headers).length > 0 ? { headers } : undefined;
-    
-    // For OAuth servers, create an auth provider
-    let authProvider: McpOAuthProvider | undefined;
-    if (supportsOAuth(definition)) {
-      // Extract OAuth config (handles both object and false cases)
-      const oauthConfig = definition.oauth === false ? {} : {
-        grantType: definition.oauth?.grantType,
-        clientId: definition.oauth?.clientId,
-        clientSecret: definition.oauth?.clientSecret,
-        scope: definition.oauth?.scope,
-      };
-      authProvider = new McpOAuthProvider(
-        serverName,
-        definition.url!,
-        oauthConfig,
-        {
-          onRedirect: async (_authUrl) => {
-            // URL is captured by startAuth, no need to log
-          },
-        }
-      );
-    }
-    
-    // Try StreamableHTTP first (modern MCP servers)
-    const streamableTransport = new StreamableHTTPClientTransport(url, { 
-      requestInit,
-      authProvider,
-    });
-    
-    try {
-      // Create a test client to verify the transport works
-      const testClient = new Client({ name: "pi-mcp-probe", version: "2.1.2" });
-      await testClient.connect(streamableTransport);
-      await testClient.close().catch(() => {});
-      // Close probe transport before creating fresh one
-      await streamableTransport.close().catch(() => {});
-      
-      // StreamableHTTP works - create fresh transport for actual use
-      return new StreamableHTTPClientTransport(url, { requestInit, authProvider });
-    } catch (error) {
-      // StreamableHTTP failed, close and try SSE fallback
-      await streamableTransport.close().catch(() => {});
-      
-      // If this was an UnauthorizedError, don't try SSE - the server needs auth
-      if (error instanceof UnauthorizedError) {
-        throw error;
-      }
-      
-      // SSE is the legacy transport
-      return new SSEClientTransport(url, { requestInit, authProvider });
-    }
-  }
-  
-  private async fetchAllTools(client: Client): Promise<McpTool[]> {
-    const allTools: McpTool[] = [];
-    let cursor: string | undefined;
-    
-    do {
-      const result = await client.listTools(cursor ? { cursor } : undefined);
-      allTools.push(...(result.tools ?? []));
-      cursor = result.nextCursor;
-    } while (cursor);
-    
-    return allTools;
-  }
-  
-  private async fetchAllResources(client: Client): Promise<McpResource[]> {
-    try {
-      const allResources: McpResource[] = [];
-      let cursor: string | undefined;
-      
-      do {
-        const result = await client.listResources(cursor ? { cursor } : undefined);
-        allResources.push(...(result.resources ?? []));
-        cursor = result.nextCursor;
-      } while (cursor);
-      
-      return allResources;
-    } catch {
-      // Server may not support resources
-      return [];
-    }
-  }
-
-  private attachAdapterNotificationHandlers(serverName: string, client: Client): void {
-    client.setNotificationHandler(serverStreamResultPatchNotificationSchema, (notification) => {
-      const listener = this.uiStreamListeners.get(notification.params.streamToken);
-      if (!listener) return;
-      listener(serverName, notification.params);
-    });
-  }
-
-  registerUiStreamListener(streamToken: string, listener: UiStreamListener): void {
-    this.uiStreamListeners.set(streamToken, listener);
-  }
-
-  removeUiStreamListener(streamToken: string): void {
-    this.uiStreamListeners.delete(streamToken);
-  }
-
-  async readResource(name: string, uri: string): Promise<ReadResourceResult> {
-    const connection = this.connections.get(name);
-    if (!connection || connection.status !== "connected") {
-      throw new Error(`Server "${name}" is not connected`);
-    }
-
-    try {
-      this.touch(name);
-      this.incrementInFlight(name);
-      return await connection.client.readResource({ uri });
-    } finally {
-      this.decrementInFlight(name);
-      this.touch(name);
-    }
-  }
-  
-  async close(name: string): Promise<void> {
-    const connection = this.connections.get(name);
-    if (!connection) return;
-    
-    // Delete from map BEFORE async cleanup to prevent a race where a
-    // concurrent connect() creates a new connection that our deferred
-    // delete() would then remove, orphaning the new server process.
-    connection.status = "closed";
-    this.connections.delete(name);
-    await connection.client.close().catch(() => {});
-    await connection.transport.close().catch(() => {});
-  }
-  
-  async closeAll(): Promise<void> {
-    const names = [...this.connections.keys()];
-    await Promise.all(names.map(name => this.close(name)));
-  }
-  
-  getConnection(name: string): ServerConnection | undefined {
-    return this.connections.get(name);
-  }
-  
-  getAllConnections(): Map<string, ServerConnection> {
-    return new Map(this.connections);
-  }
-
-  touch(name: string): void {
-    const connection = this.connections.get(name);
-    if (connection) {
-      connection.lastUsedAt = Date.now();
-    }
-  }
-
-  incrementInFlight(name: string): void {
-    const connection = this.connections.get(name);
-    if (connection) {
-      connection.inFlight = (connection.inFlight ?? 0) + 1;
-    }
-  }
-
-  decrementInFlight(name: string): void {
-    const connection = this.connections.get(name);
-    if (connection && connection.inFlight) {
-      connection.inFlight--;
-    }
-  }
-
-  isIdle(name: string, timeoutMs: number): boolean {
-    const connection = this.connections.get(name);
-    if (!connection || connection.status !== "connected") return false;
-    if (connection.inFlight > 0) return false;
-    return (Date.now() - connection.lastUsedAt) > timeoutMs;
-  }
-}
-
-/**
- * Resolve environment variables with interpolation.
- */
-function resolveEnv(env?: Record<string, string>): Record<string, string> {
-  const overrides = env ? interpolateEnvRecord(env) : undefined;
-  const resolved: Record<string, string> = {};
-  for (const [key, value] of Object.entries(createChildProcessEnvironment(overrides))) {
-    if (value !== undefined) {
-      resolved[key] = value;
-    }
-  }
-  return resolved;
-}
-
-/**
- * Resolve headers with environment variable interpolation.
- */
-function resolveHeaders(headers?: Record<string, string>): Record<string, string> | undefined {
-  return interpolateEnvRecord(headers);
-}

package/dist/builtin/mcp/session-cleanup-barrier.ts

@@ -1,43 +0,0 @@
-/**
- * Retain teardown work across MCP lifecycle generations without allowing a
- * non-abortable producer to block every later generation forever.
- */
-type CleanupTask = PromiseLike<object | string | void>;
-
-const DEFAULT_CLEANUP_DEADLINE_MS = 2_000;
-
-function observeWithin(task: CleanupTask, deadlineMs: number, onTimeout: () => void): Promise<void> {
-  const observed = Promise.resolve(task).then(() => undefined, () => undefined);
-  let timer: ReturnType<typeof setTimeout> | undefined;
-  const deadline = new Promise<void>((resolve) => {
-    timer = setTimeout(() => {
-      onTimeout();
-      resolve();
-    }, deadlineMs);
-    timer.unref?.();
-  });
-  return Promise.race([observed, deadline]).finally(() => {
-    if (timer) clearTimeout(timer);
-  });
-}
-
-export class McpSessionCleanupBarrier {
-  private settled: Promise<void> = Promise.resolve();
-
-  constructor(private readonly deadlineMs = DEFAULT_CLEANUP_DEADLINE_MS) {}
-
-  wait(): Promise<void> {
-    return this.settled;
-  }
-
-  retain(tasks: readonly (CleanupTask | null | undefined)[]): Promise<void> {
-    const observed = tasks.filter((task): task is CleanupTask => task !== null && task !== undefined);
-    const prior = this.settled;
-    const next = Promise.all([
-      observeWithin(prior, this.deadlineMs, () => console.error("MCP: prior cleanup exceeded its deadline; continuing with fenced state")),
-      ...observed.map((task) => observeWithin(task, this.deadlineMs, () => console.error("MCP: cleanup task exceeded its deadline; continuing with fenced state"))),
-    ]).then(() => undefined);
-    this.settled = next;
-    return next;
-  }
-}