@njuptlzf/mcp-adapter 2.29.0-0.0.1

package/sampling-handler.ts

@@ -0,0 +1,284 @@
+import { complete } from "@earendil-works/pi-ai/compat";
+import type { Api, AssistantMessage, Message, Model, ProviderHeaders, TextContent } from "@earendil-works/pi-ai";
+import { truncateAtWord } from "./utils.ts";
+import { throwIfAborted } from "./abort.ts";
+import type { ExtensionUIContext, ModelRegistry } from "@earendil-works/pi-coding-agent";
+import type { Client } from "@modelcontextprotocol/client";
+import {
+  type CreateMessageRequest,
+  type CreateMessageResult,
+  type ModelPreferences,
+  type SamplingMessage,
+  type SamplingMessageContentBlock,
+} from "@modelcontextprotocol/client";
+
+export type SamplingUIContext = Pick<ExtensionUIContext, "confirm">;
+export type SamplingModelRegistry = Pick<ModelRegistry, "getAvailable" | "getApiKeyAndHeaders">;
+
+export interface SamplingHandlerOptions {
+  serverName: string;
+  autoApprove: boolean;
+  ui?: SamplingUIContext;
+  modelRegistry: SamplingModelRegistry;
+  getCurrentModel: () => Model<Api> | undefined;
+  getSignal: () => AbortSignal | undefined;
+}
+
+export type ServerSamplingConfig = Omit<SamplingHandlerOptions, "serverName">;
+
+export function registerSamplingHandler(client: Client, options: SamplingHandlerOptions): void {
+  client.setRequestHandler("sampling/createMessage", request => {
+    return handleSamplingRequest(options, request);
+  });
+}
+
+export async function handleSamplingRequest(
+  options: SamplingHandlerOptions,
+  request: CreateMessageRequest,
+): Promise<CreateMessageResult> {
+  const params = request.params;
+  const signal = options.getSignal();
+  throwIfAborted(signal);
+
+  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);
+  throwIfAborted(signal);
+  await confirmSampling(
+    options,
+    "Approve MCP sampling request",
+    formatRequestApproval(options.serverName, `${model.provider}/${model.id}`, params.systemPrompt, messages),
+  );
+  throwIfAborted(signal);
+
+  const result = await complete(
+    model,
+    {
+      ...(params.systemPrompt !== undefined ? { systemPrompt: params.systemPrompt } : {}),
+      messages,
+    },
+    {
+      ...(apiKey !== undefined ? { apiKey } : {}),
+      ...(headers !== undefined ? { headers } : {}),
+      maxTokens: params.maxTokens,
+      ...(params.temperature !== undefined ? { temperature: params.temperature } : {}),
+      ...(params.metadata !== undefined ? { metadata: params.metadata as Record<string, unknown> } : {}),
+      ...(signal ? { signal } : {}),
+    },
+  );
+
+  const converted = convertAssistantResult(result);
+  throwIfAborted(signal);
+  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: { type: string; text?: string; mimeType?: string; name?: string }) => {
+    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[] = [];
+  const signal = options.getSignal();
+  for (const model of candidates) {
+    throwIfAborted(signal);
+    const auth = await options.modelRegistry.getApiKeyAndHeaders(model);
+    throwIfAborted(signal);
+    if (auth.ok === false) {
+      errors.push(`${model.provider}/${model.id}: ${auth.error}`);
+      continue;
+    }
+    return {
+      model,
+      ...(auth.apiKey !== undefined ? { apiKey: auth.apiKey } : {}),
+      ...(auth.headers !== undefined ? { 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: { type: string; text?: string }) => {
+      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: string | undefined): 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/search-ranking.ts

@@ -0,0 +1,278 @@
+import type { McpExtensionState } from "./state.ts";
+import type { ServerEntry, ToolMetadata, ToolPrefix } from "./types.ts";
+import { getServerPrefix, getToolNameCandidates, isServerDisabled, matchesToolPattern, resolveToolPrefix } from "./types.ts";
+import { isServerInActiveFailureBackoff } from "./failure-backoff.ts";
+
+/**
+ * Shortest field token allowed to stem-match a longer query token.
+ * Real descriptions tokenize possessives into single letters ("project's" -> ["project", "s"]),
+ * which would otherwise make every query starting with that letter a match.
+ */
+const MIN_STEM_LENGTH = 4;
+
+const FIELD_WEIGHTS = {
+  name: 12,
+  originalName: 10,
+  server: 8,
+  description: 5,
+  keywords: 5,
+} as const;
+
+export interface RankedToolMatch {
+  server: string;
+  tool: ToolMetadata;
+  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
+ * set includeTools/excludeTools use — and all matching entries are unioned.
+ */
+export function resolveSearchKeywords(
+  definition: ServerEntry | undefined,
+  toolOriginalName: string,
+  serverName: string,
+  globalPrefix: ToolPrefix,
+): string[] {
+  const map = definition?.searchKeywords;
+  if (!map || typeof map !== "object" || Array.isArray(map)) return [];
+  const candidates = getToolNameCandidates(toolOriginalName, serverName, resolveToolPrefix(definition, globalPrefix));
+  const keywords: string[] = [];
+  const seen = new Set<string>();
+  for (const [pattern, values] of Object.entries(map)) {
+    if (!Array.isArray(values)) continue;
+    if (!matchesToolPattern(candidates, [pattern])) continue;
+    for (const value of values) {
+      if (typeof value !== "string") continue;
+      const trimmed = value.trim();
+      if (!trimmed || seen.has(trimmed)) continue;
+      seen.add(trimmed);
+      keywords.push(trimmed);
+    }
+  }
+  return keywords;
+}
+
+export function normalizeSearchText(value: string): string {
+  return value
+    .replace(/([a-z0-9])([A-Z])/g, "$1 $2")
+    .replace(/[_./:-]+/g, " ")
+    .toLowerCase();
+}
+
+export function tokenize(value: string): string[] {
+  return normalizeSearchText(value).split(/[^a-z0-9]+/).filter(Boolean);
+}
+
+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, fieldTokens] of prepared.fields) {
+    const weight = FIELD_WEIGHTS[field];
+    if (value === normalizedQuery) {
+      score += weight * 14;
+      phraseMatched = true;
+      wholeFieldExact = true;
+    } else if (value.startsWith(normalizedQuery)) {
+      score += weight * 9;
+      phraseMatched = true;
+    } else if (value.includes(normalizedQuery)) {
+      score += weight * 6;
+      phraseMatched = true;
+    }
+
+    for (const token of queryTokens) {
+      if (fieldTokens.includes(token)) {
+        score += weight * 4;
+        matchedTokens.add(token);
+      } else if (fieldTokens.some(fieldToken => fieldToken.startsWith(token) || (fieldToken.length >= MIN_STEM_LENGTH && token.startsWith(fieldToken)))) {
+        score += weight * 2;
+        matchedTokens.add(token);
+      } else if (value.includes(token)) {
+        score += weight;
+        matchedTokens.add(token);
+      }
+    }
+  }
+
+  // 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 (prepared.keywordPhrases.length > 0) {
+    const weight = FIELD_WEIGHTS.keywords;
+    let phraseScore = 0;
+    for (const phrase of prepared.keywordPhrases) {
+      if (phrase === normalizedQuery) {
+        phraseScore = Math.max(phraseScore, weight * 14);
+        phraseMatched = true;
+        wholeFieldExact = true;
+      } else if (phrase.startsWith(normalizedQuery)) {
+        phraseScore = Math.max(phraseScore, weight * 9);
+        phraseMatched = true;
+      } else if (phrase.includes(normalizedQuery)) {
+        phraseScore = Math.max(phraseScore, weight * 6);
+        phraseMatched = true;
+      }
+    }
+    score += phraseScore;
+
+    for (const token of queryTokens) {
+      if (prepared.keywordTokens.includes(token)) {
+        score += weight * 4;
+        matchedTokens.add(token);
+      } 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 (prepared.keywordPhrases.some(phrase => phrase.includes(token))) {
+        score += weight;
+        matchedTokens.add(token);
+      }
+    }
+  }
+
+  const coverage = matchedTokens.size / queryTokens.length;
+  if (!phraseMatched && (queryTokens.length <= 2 ? coverage !== 1 : coverage < 0.6)) return null;
+
+  score += coverage === 1 ? 25 : Math.round(coverage * 10);
+  const firstQueryToken = queryTokens[0];
+  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,
+  server?: string,
+  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;
+    if (isServerInActiveFailureBackoff(state, serverName)) continue;
+    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));
+}
+
+export function paginate<T>(items: T[], offset: number, limit: number): { items: T[]; total: number; hasMore: boolean; nextOffset: number | null } {
+  const safeOffset = Number.isFinite(offset) ? Math.max(0, Math.trunc(offset)) : 0;
+  const safeLimit = Number.isFinite(limit) ? Math.max(1, Math.trunc(limit)) : 1;
+  const total = items.length;
+  const page = items.slice(safeOffset, safeOffset + safeLimit);
+  const nextOffset = safeOffset + page.length;
+  return {
+    items: page,
+    total,
+    hasMore: nextOffset < total,
+    nextOffset: nextOffset < total ? nextOffset : null,
+  };
+}
+
+export function rankSuggestions(state: McpExtensionState, name: string, limit: number): string[] {
+  const stripped = Object.keys(state.config.mcpServers)
+    .flatMap(server => (["server", "short", "mcp"] as const)
+      .map(prefix => getServerPrefix(server, prefix)))
+    .filter((candidate): candidate is string => Boolean(candidate) && name.startsWith(`${candidate}_`))
+    .sort((a, b) => b.length - a.length)
+    .map(candidate => name.slice(candidate.length + 1));
+  const query = stripped[0] ?? name;
+  return rankToolMatches(state, query, undefined, false).slice(0, limit).map(match => match.tool.name);
+}