@aria-cli/tools 1.0.0

package/src/registry/registry.ts

@@ -0,0 +1,905 @@
+/**
+ * Tool Registry - Manages tool registration and discovery
+ */
+
+import type {
+  Tool,
+  ToolCategory,
+  ToolContext,
+  ToolResult,
+  ToolMiddleware,
+  RegisterOptions,
+  ToolInfo,
+  Logger,
+} from "../types.js";
+import type { ToolItem } from "@aria/types";
+import { MCPClient } from "../mcp/client.js";
+import type { MCPServerConfig, MCPTool } from "../mcp/types.js";
+import { executeBash } from "../executors/shell.js";
+
+/**
+ * Validate tool input against a JSON Schema object.
+ *
+ * Lightweight validation that covers the schemas produced by our tool factory
+ * (Zod -> JSON Schema): checks top-level type, required properties, property
+ * types, enums, patterns, numeric/string/array bounds, additionalProperties,
+ * and recursively validates nested object/array schemas.
+ * Does not pull in a full JSON Schema validator like Ajv to keep the dependency
+ * footprint minimal.
+ *
+ * @returns null if valid, or a human-readable error string
+ */
+function isPlainObject(value: unknown): value is Record<string, unknown> {
+  return typeof value === "object" && value !== null && !Array.isArray(value);
+}
+
+function valueType(value: unknown): string {
+  if (value === null) return "null";
+  if (Array.isArray(value)) return "array";
+  return typeof value;
+}
+
+function matchesJsonSchemaType(value: unknown, schemaType: string): boolean {
+  if (schemaType === "integer") return typeof value === "number" && Number.isInteger(value);
+  if (schemaType === "number") return typeof value === "number";
+  if (schemaType === "string") return typeof value === "string";
+  if (schemaType === "boolean") return typeof value === "boolean";
+  if (schemaType === "array") return Array.isArray(value);
+  if (schemaType === "object") return isPlainObject(value);
+  if (schemaType === "null") return value === null;
+  return true;
+}
+
+function pathLabel(path: string): string {
+  return path ? `Property "${path}"` : "Value";
+}
+
+function validateValueAgainstSchema(
+  value: unknown,
+  schema: Record<string, unknown>,
+  path: string,
+): string | null {
+  const expectedType = schema.type as string | string[] | undefined;
+  if (expectedType) {
+    const types = Array.isArray(expectedType) ? expectedType : [expectedType];
+    const typeMatches = types.some((typeName) => matchesJsonSchemaType(value, typeName));
+    if (!typeMatches) {
+      return `${pathLabel(path)} expected type ${types.join(" | ")}; got ${valueType(value)}`;
+    }
+  }
+
+  if (Array.isArray(schema.enum) && schema.enum.length > 0) {
+    const allowed = schema.enum;
+    const hasValue = allowed.some((candidate) => Object.is(candidate, value));
+    if (!hasValue) {
+      return `${pathLabel(path)} must be one of: ${allowed.map((item) => String(item)).join(", ")}`;
+    }
+  }
+
+  if ("const" in schema && !Object.is(schema.const, value)) {
+    return `${pathLabel(path)} must equal ${String(schema.const)}`;
+  }
+
+  if (typeof value === "string") {
+    if (typeof schema.minLength === "number" && value.length < schema.minLength) {
+      return `${pathLabel(path)} must have length >= ${schema.minLength}`;
+    }
+    if (typeof schema.maxLength === "number" && value.length > schema.maxLength) {
+      return `${pathLabel(path)} must have length <= ${schema.maxLength}`;
+    }
+    if (typeof schema.pattern === "string") {
+      try {
+        const pattern = new RegExp(schema.pattern);
+        if (!pattern.test(value)) {
+          return `${pathLabel(path)} does not match required pattern`;
+        }
+      } catch {
+        // Invalid schema pattern: ignore rather than failing all validation.
+      }
+    }
+  }
+
+  if (typeof value === "number") {
+    if (typeof schema.minimum === "number" && value < schema.minimum) {
+      return `${pathLabel(path)} must be >= ${schema.minimum}`;
+    }
+    if (typeof schema.maximum === "number" && value > schema.maximum) {
+      return `${pathLabel(path)} must be <= ${schema.maximum}`;
+    }
+    if (typeof schema.exclusiveMinimum === "number" && value <= schema.exclusiveMinimum) {
+      return `${pathLabel(path)} must be > ${schema.exclusiveMinimum}`;
+    }
+    if (typeof schema.exclusiveMaximum === "number" && value >= schema.exclusiveMaximum) {
+      return `${pathLabel(path)} must be < ${schema.exclusiveMaximum}`;
+    }
+    if (typeof schema.multipleOf === "number" && schema.multipleOf > 0) {
+      const quotient = value / schema.multipleOf;
+      if (!Number.isInteger(quotient)) {
+        return `${pathLabel(path)} must be a multiple of ${schema.multipleOf}`;
+      }
+    }
+  }
+
+  if (Array.isArray(value)) {
+    if (typeof schema.minItems === "number" && value.length < schema.minItems) {
+      return `${pathLabel(path)} must contain at least ${schema.minItems} item(s)`;
+    }
+    if (typeof schema.maxItems === "number" && value.length > schema.maxItems) {
+      return `${pathLabel(path)} must contain at most ${schema.maxItems} item(s)`;
+    }
+
+    if (isPlainObject(schema.items)) {
+      for (let index = 0; index < value.length; index++) {
+        const itemError = validateValueAgainstSchema(
+          value[index],
+          schema.items,
+          path ? `${path}[${index}]` : `[${index}]`,
+        );
+        if (itemError) return itemError;
+      }
+    }
+  }
+
+  if (isPlainObject(value)) {
+    const required = Array.isArray(schema.required)
+      ? schema.required.filter((item): item is string => typeof item === "string")
+      : [];
+    if (required.length > 0) {
+      const missing = required.filter((key) => !(key in value) || value[key] === undefined);
+      if (missing.length > 0) {
+        return `${path ? `${path}: ` : ""}Missing required properties: ${missing.join(", ")}`;
+      }
+    }
+
+    const properties = isPlainObject(schema.properties)
+      ? (schema.properties as Record<string, Record<string, unknown>>)
+      : undefined;
+
+    if (properties) {
+      for (const [key, propSchema] of Object.entries(properties)) {
+        if (!(key in value) || value[key] === undefined) continue;
+        if (!isPlainObject(propSchema)) continue;
+        const childPath = path ? `${path}.${key}` : key;
+        const nestedError = validateValueAgainstSchema(value[key], propSchema, childPath);
+        if (nestedError) return nestedError;
+      }
+    }
+
+    const additionalProperties = schema.additionalProperties;
+    if (additionalProperties === false) {
+      const allowed = new Set(properties ? Object.keys(properties) : []);
+      const unknown = Object.keys(value).filter((key) => !allowed.has(key));
+      if (unknown.length > 0) {
+        if (path) {
+          return `${path}: Unknown properties: ${unknown.join(", ")}`;
+        }
+        return `Unknown properties: ${unknown.join(", ")}`;
+      }
+    } else if (isPlainObject(additionalProperties)) {
+      const allowed = new Set(properties ? Object.keys(properties) : []);
+      for (const [extraKey, extraValue] of Object.entries(value)) {
+        if (allowed.has(extraKey)) continue;
+        const childPath = path ? `${path}.${extraKey}` : extraKey;
+        const apError = validateValueAgainstSchema(extraValue, additionalProperties, childPath);
+        if (apError) return apError;
+      }
+    }
+  }
+
+  return null;
+}
+
+export function validateToolInput(
+  input: unknown,
+  schema: Record<string, unknown>,
+  path = "",
+): string | null {
+  // Only validate object-type schemas (all tool parameters are objects)
+  if (schema.type !== "object") return null;
+
+  const prefix = path ? `${path}: ` : "";
+
+  // Input must be an object (or nullish, treated as empty object by many callers)
+  if (input === null || input === undefined) {
+    // If schema has required fields, this is invalid
+    const required = schema.required as string[] | undefined;
+    if (required && required.length > 0) {
+      return `${prefix}Expected an object with required properties: ${required.join(", ")}; got ${input === null ? "null" : "undefined"}`;
+    }
+    return null;
+  }
+
+  if (typeof input !== "object" || Array.isArray(input)) {
+    return `${prefix}Expected an object; got ${Array.isArray(input) ? "array" : typeof input}`;
+  }
+
+  return validateValueAgainstSchema(input, schema, path);
+}
+
+function isMcpToolReadOnly(tool: MCPTool): boolean {
+  return tool.annotations?.readOnlyHint === true;
+}
+
+/**
+ * Memoria-like interface for tool discovery
+ */
+interface MemoriaLike {
+  recallTools(options: {
+    query: string;
+    limit?: number;
+    offset?: number;
+    minConfidence?: number;
+    matchAll?: boolean;
+    updateAccessStats?: boolean;
+  }): Promise<ToolItem[]>;
+  rememberTool(
+    tool: { name: string; description: string } & Record<string, unknown>,
+  ): Promise<string>;
+}
+
+function normalizeDiscoveryIssues(rawFailures: unknown): string[] | undefined {
+  if (!Array.isArray(rawFailures)) return undefined;
+
+  const issues: string[] = [];
+  for (const failure of rawFailures) {
+    if (typeof failure === "string") {
+      const trimmed = failure.trim();
+      if (trimmed) issues.push(trimmed);
+      continue;
+    }
+
+    if (!failure || typeof failure !== "object") continue;
+    const record = failure as Record<string, unknown>;
+
+    const primary =
+      typeof record.error === "string" && record.error.trim()
+        ? record.error.trim()
+        : typeof record.message === "string" && record.message.trim()
+          ? record.message.trim()
+          : undefined;
+
+    if (primary) issues.push(primary);
+  }
+
+  if (issues.length === 0) return undefined;
+  return [...new Set(issues)];
+}
+
+function isShellTemplate(template: string): boolean {
+  return template.trimStart().toLowerCase().startsWith("bash:");
+}
+
+function shellEscape(value: string): string {
+  return `'${value.replace(/'/g, `'\"'\"'`)}'`;
+}
+
+function renderTemplate(
+  template: string,
+  inputObj: Record<string, unknown>,
+  options?: { escapeForShell?: boolean },
+): string {
+  return template.replace(/\{\{(\w+)\}\}/g, (_match: string, key: string) => {
+    const value = inputObj[key];
+    if (value === undefined || value === null) return "";
+    const stringValue = String(value);
+    return options?.escapeForShell ? shellEscape(stringValue) : stringValue;
+  });
+}
+
+export class ToolRegistry {
+  private tools = new Map<string, Tool>();
+  private middleware: ToolMiddleware[] = [];
+  private mcpClient?: MCPClient;
+  private logger: Logger;
+
+  constructor(options?: { logger?: Logger }) {
+    this.logger = options?.logger ?? console;
+  }
+
+  /**
+   * Register a middleware to wrap tool execution.
+   * Middlewares are applied in registration order (first registered = outermost).
+   */
+  use(mw: ToolMiddleware): void {
+    this.middleware.push(mw);
+  }
+
+  /**
+   * Register a tool
+   * @throws Error if tool already registered and override not set
+   */
+  register(tool: Tool, options: RegisterOptions = {}): void {
+    if (this.tools.has(tool.name) && !options.override) {
+      throw new Error(
+        `Tool "${tool.name}" is already registered. Use { override: true } to replace.`,
+      );
+    }
+    this.tools.set(tool.name, tool);
+  }
+
+  /**
+   * Get a tool by name
+   */
+  get(name: string): Tool | undefined {
+    return this.tools.get(name);
+  }
+
+  /**
+   * Check if a tool is registered
+   */
+  has(name: string): boolean {
+    return this.tools.has(name);
+  }
+
+  /**
+   * List all tools, optionally filtered by category
+   */
+  list(category?: ToolCategory): Tool[] {
+    const all = Array.from(this.tools.values());
+    if (!category) return all;
+    return all.filter((tool) => tool.category === category);
+  }
+
+  /**
+   * Unregister a tool
+   * @returns true if tool was removed, false if not found
+   */
+  unregister(name: string): boolean {
+    return this.tools.delete(name);
+  }
+
+  /**
+   * Get count of registered tools
+   */
+  get size(): number {
+    return this.tools.size;
+  }
+
+  /**
+   * Clear all registered tools
+   */
+  clear(): void {
+    this.tools.clear();
+  }
+
+  /**
+   * Get all registered tools
+   */
+  getAll(): Tool[] {
+    return Array.from(this.tools.values());
+  }
+
+  /**
+   * Execute a tool by name with input validation.
+   *
+   * Validates `input` against the tool's JSON Schema `parameters` before
+   * delegating to the tool's handler. Returns a structured error result
+   * for unknown tools or validation failures instead of throwing.
+   *
+   * Tools without a `parameters` schema (or with a non-object schema)
+   * skip validation and pass input directly to the handler.
+   */
+  async execute(name: string, input: unknown, context: ToolContext): Promise<ToolResult> {
+    const tool = this.tools.get(name);
+    if (!tool) {
+      return {
+        success: false,
+        message: `Tool "${name}" is not registered.`,
+      };
+    }
+
+    // Validate input against JSON Schema if the tool declares parameters
+    if (tool.parameters && typeof tool.parameters === "object") {
+      const schema = tool.parameters as Record<string, unknown>;
+      const error = validateToolInput(input, schema);
+      if (error) {
+        return {
+          success: false,
+          message: `Invalid input for tool "${name}": ${error}`,
+        };
+      }
+    }
+
+    // If no middleware, execute directly
+    if (this.middleware.length === 0) {
+      return tool.execute(input, context);
+    }
+
+    // Run through middleware chain (first registered = outermost)
+    const chain = this.middleware.reduceRight(
+      (next: () => Promise<ToolResult>, mw: ToolMiddleware) => () => mw(tool, input, context, next),
+      () => tool.execute(input, context),
+    );
+    return chain();
+  }
+
+  /**
+   * Convert a Tool to ToolInfo for system prompt
+   */
+  private toolToInfo(tool: Tool): ToolInfo {
+    const parameters: ToolInfo["parameters"] = [];
+
+    if (tool.parameters && typeof tool.parameters === "object") {
+      const schema = tool.parameters as {
+        properties?: Record<string, { type?: string; description?: string }>;
+        required?: string[];
+      };
+      if (schema.properties) {
+        for (const [name, prop] of Object.entries(schema.properties)) {
+          parameters.push({
+            name,
+            type: String(prop.type || "unknown"),
+            required: schema.required?.includes(name) ?? false,
+            description: prop.description,
+          });
+        }
+      }
+    }
+
+    return {
+      name: tool.name,
+      description: tool.description,
+      requiresConfirmation: tool.requiresConfirmation ?? tool.riskLevel === "dangerous",
+      parameters,
+    };
+  }
+
+  /**
+   * Get tool info for all tools (for system prompt)
+   */
+  getToolInfos(): ToolInfo[] {
+    return this.getAll().map((tool) => this.toolToInfo(tool));
+  }
+
+  /**
+   * Search for tools by name or description
+   */
+  search(query: string): Tool[] {
+    const queryLower = query.toLowerCase();
+    return Array.from(this.tools.values()).filter(
+      (tool) =>
+        tool.name.toLowerCase().includes(queryLower) ||
+        tool.description.toLowerCase().includes(queryLower),
+    );
+  }
+
+  /**
+   * Discover and load tools from Memoria knowledge base
+   *
+   * @param memoria - Memoria instance to query
+   * @returns Number of tools discovered and loaded
+   */
+  async discoverFromMemoria(memoria: MemoriaLike): Promise<number> {
+    let count = 0;
+    let offset = 0;
+    let pagesLoaded = 0;
+    const seenPageSignatures = new Set<string>();
+    const DISCOVERY_PAGE_SIZE = 200;
+    const MAX_DISCOVERY_PAGES = 250;
+
+    // Valid values for validation
+    const validCategories = new Set([
+      "filesystem",
+      "code",
+      "shell",
+      "web",
+      "data",
+      "memory",
+      "meta",
+      "arion",
+    ]);
+    const validRiskLevels = new Set(["safe", "moderate", "dangerous"]);
+
+    try {
+      while (pagesLoaded < MAX_DISCOVERY_PAGES) {
+        const tools = await memoria.recallTools({
+          query: "",
+          limit: DISCOVERY_PAGE_SIZE,
+          offset,
+          matchAll: true,
+          updateAccessStats: false,
+        });
+        pagesLoaded += 1;
+        if (tools.length === 0) break;
+        const pageSignature = tools.map((tool) => String(tool.id ?? tool.name)).join("|");
+        if (pageSignature && seenPageSignatures.has(pageSignature)) {
+          this.logger.warn(
+            "Memoria discovery received a duplicate page; stopping to avoid pagination loop",
+          );
+          break;
+        }
+        if (pageSignature) {
+          seenPageSignatures.add(pageSignature);
+        }
+
+        for (const k of tools) {
+          try {
+            if (typeof k.name !== "string" || k.name.trim().length === 0) {
+              this.logger.warn("Skipping Memoria tool with missing name");
+              continue;
+            }
+            if (typeof k.description !== "string" || k.description.trim().length === 0) {
+              this.logger.warn(`Skipping Memoria tool "${k.name}" with missing description`);
+              continue;
+            }
+            const category =
+              typeof k.category === "string" && validCategories.has(k.category)
+                ? (k.category as Tool["category"])
+                : "meta";
+            if (category === "meta" && k.category !== "meta") {
+              this.logger.warn(
+                `Memoria tool "${k.name}" has invalid/missing category; defaulting to "meta"`,
+              );
+            }
+            const parameters =
+              k.parameters && typeof k.parameters === "object"
+                ? (k.parameters as Record<string, unknown>)
+                : { type: "object", properties: {}, additionalProperties: true };
+            if (!k.parameters || typeof k.parameters !== "object") {
+              this.logger.warn(
+                `Memoria tool "${k.name}" has invalid/missing parameters; defaulting to permissive schema`,
+              );
+            }
+
+            // Register only tools with explicit response templates.
+            // Knowledge-only records are treated as non-executable metadata.
+            const hasResponseTemplate =
+              typeof k.responseTemplate === "string" && k.responseTemplate.trim().length > 0;
+            if (!hasResponseTemplate) {
+              this.logger.warn(
+                `Skipping Memoria tool "${k.name}" without executable responseTemplate`,
+              );
+              continue;
+            }
+
+            const declaredRiskLevel =
+              typeof k.riskLevel === "string" && validRiskLevels.has(k.riskLevel)
+                ? k.riskLevel
+                : "moderate";
+            const isBashTemplate =
+              typeof k.responseTemplate === "string" && isShellTemplate(k.responseTemplate);
+            // Memoria bash templates execute shell commands and must always
+            // flow through dangerous-risk approval paths.
+            const riskLevel = isBashTemplate ? "dangerous" : declaredRiskLevel;
+            const issues = normalizeDiscoveryIssues(
+              (k as unknown as Record<string, unknown>).failures,
+            );
+
+            const toolDef: Record<string, unknown> = {
+              description: k.description,
+              category,
+              parameters,
+              riskLevel,
+            };
+            if (issues && issues.length > 0) toolDef.issues = issues;
+            // Pass through executor-relevant fields from the tool definition.
+            if (typeof k.responseTemplate === "string")
+              toolDef.responseTemplate = k.responseTemplate;
+
+            const tool: Tool = {
+              name: k.name,
+              description: k.description,
+              category: category as Tool["category"],
+              parameters: parameters as Tool["parameters"],
+              riskLevel: riskLevel as Tool["riskLevel"],
+              issues,
+              loadingTier: "deferred",
+              execute: this.createMemoriaToolExecutor(k.name, toolDef),
+            };
+
+            if (this.tools.has(tool.name)) {
+              // Don't override core tools with Memoria-discovered tools
+              continue;
+            }
+            this.register(tool);
+            count++;
+          } catch {
+            // Skip invalid tool definitions
+            this.logger.warn(`Invalid tool definition in Memoria: ${k.name}`);
+          }
+        }
+
+        offset += tools.length;
+        if (tools.length < DISCOVERY_PAGE_SIZE) break;
+      }
+
+      if (pagesLoaded >= MAX_DISCOVERY_PAGES) {
+        this.logger.warn(
+          `Stopped Memoria discovery after ${MAX_DISCOVERY_PAGES} pages to avoid unbounded paging`,
+        );
+      }
+    } catch (error) {
+      this.logger.error("Error discovering tools from Memoria:", error);
+    }
+
+    return count;
+  }
+
+  /**
+   * Save a tool definition to Memoria
+   *
+   * @param name - Name of the tool to save
+   * @param memoria - Memoria instance to save to
+   */
+  async saveToMemoria(name: string, memoria: MemoriaLike): Promise<void> {
+    const tool = this.get(name);
+    if (!tool) {
+      throw new Error(`Tool "${name}" not found`);
+    }
+
+    const issues =
+      Array.isArray(tool.issues) && tool.issues.length > 0
+        ? tool.issues
+            .map((issue) => (typeof issue === "string" ? issue.trim() : ""))
+            .filter(Boolean)
+        : [];
+    const failures =
+      issues.length > 0
+        ? issues.map((issue) => ({
+            timestamp: new Date(),
+            error: issue,
+          }))
+        : undefined;
+    const rawResponseTemplate = (tool as { responseTemplate?: unknown }).responseTemplate;
+    const responseTemplate =
+      typeof rawResponseTemplate === "string" && rawResponseTemplate.trim().length > 0
+        ? rawResponseTemplate
+        : undefined;
+
+    const memoriaRecord: Record<string, unknown> = {
+      name: tool.name,
+      description: tool.description,
+      category: tool.category,
+      parameters: tool.parameters,
+      riskLevel: tool.riskLevel,
+      ...(failures ? { failures } : {}),
+    };
+
+    if (responseTemplate) {
+      memoriaRecord.responseTemplate = responseTemplate;
+    }
+
+    await memoria.rememberTool(memoriaRecord as Parameters<MemoriaLike["rememberTool"]>[0]);
+  }
+
+  /**
+   * Create an executor for a tool loaded from Memoria.
+   *
+   * Supports template execution only (no eval, no arbitrary code):
+   *
+   * **Response template**: If `toolDef.responseTemplate` is a string,
+   *    `{{key}}` placeholders are replaced with values from the input object.
+   *
+   * Input is validated against the full declared JSON Schema before interpolation.
+   */
+  private createMemoriaToolExecutor(
+    name: string,
+    toolDef: Record<string, unknown>,
+  ): Tool["execute"] {
+    return async (input: unknown, _context: ToolContext): Promise<ToolResult> => {
+      try {
+        const inputObj = (input && typeof input === "object" ? input : {}) as Record<
+          string,
+          unknown
+        >;
+
+        // Validate full input schema if one is present.
+        const params = toolDef.parameters as Record<string, unknown> | undefined;
+        if (params && typeof params === "object") {
+          const validationError = validateToolInput(input, params);
+          if (validationError) {
+            return {
+              success: false,
+              message: `Invalid input for tool "${name}": ${validationError}`,
+            };
+          }
+        }
+
+        // Template mode with {{key}} interpolation
+        if (typeof toolDef.responseTemplate === "string") {
+          const template = toolDef.responseTemplate;
+          const isBashTemplate = isShellTemplate(template);
+          const rendered = renderTemplate(template, inputObj, {
+            escapeForShell: isBashTemplate,
+          });
+          if (isBashTemplate) {
+            const bashPrefixIndex = rendered.toLowerCase().indexOf("bash:");
+            const command = rendered.slice(bashPrefixIndex + "bash:".length).trim();
+            if (!command) {
+              return {
+                success: false,
+                message: `Tool "${name}" has empty bash command template`,
+              };
+            }
+            const bashResult = await executeBash(
+              {
+                command,
+                cwd: _context.workingDir,
+                env: _context.env,
+                timeout: 30_000,
+              },
+              _context,
+            );
+            if (!bashResult.success) {
+              return {
+                success: false,
+                message: `Tool "${name}" command failed: ${bashResult.message}`,
+              };
+            }
+            const shellData =
+              bashResult.data && typeof bashResult.data === "object"
+                ? (bashResult.data as { stdout?: string; stderr?: string })
+                : undefined;
+            const output = [shellData?.stdout, shellData?.stderr].filter(Boolean).join("\n").trim();
+            return {
+              success: true,
+              message: `Tool "${name}" executed command template.`,
+              data: output,
+            };
+          }
+          return {
+            success: true,
+            message: `Tool "${name}" executed with template response.`,
+            data: rendered,
+          };
+        }
+
+        // No executable content at all
+        return {
+          success: false,
+          message:
+            `Tool "${name}" was loaded from Memoria but has no executable content. ` +
+            `Add a "responseTemplate" field to the tool definition.`,
+        };
+      } catch (error) {
+        return {
+          success: false,
+          message: `Tool "${name}" execution failed: ${error instanceof Error ? error.message : String(error)}`,
+        };
+      }
+    };
+  }
+
+  /**
+   * Connect to MCP servers and register their tools
+   */
+  async connectMCP(configs: MCPServerConfig[]): Promise<{
+    tools: number;
+    resources: number;
+    prompts: number;
+  }> {
+    if (this.mcpClient) {
+      await this.disconnectMCP();
+    }
+    this.mcpClient = new MCPClient({ logger: this.logger });
+
+    // Connect to all servers
+    const results = await Promise.allSettled(configs.map((c) => this.mcpClient!.connect(c)));
+    const failures = results
+      .map((r, i) => (r.status === "rejected" ? configs[i]!.name : null))
+      .filter(Boolean);
+    if (failures.length > 0) {
+      this.logger.warn(`Failed to connect MCP servers: ${failures.join(", ")}`);
+    }
+
+    // Register MCP tools
+    const mcpTools = await this.mcpClient.listAllTools();
+    for (const tool of mcpTools) {
+      const isReadOnly = isMcpToolReadOnly(tool);
+      this.register({
+        name: `mcp__${tool.server}__${tool.name}`,
+        description: tool.description || `MCP tool from ${tool.server}`,
+        category: "meta",
+        riskLevel: "moderate",
+        isReadOnly,
+        loadingTier: "always",
+        parameters: tool.inputSchema,
+        execute: async (args, ctx) => {
+          if (ctx.abortSignal?.aborted) {
+            return { success: false, message: "Cancelled" };
+          }
+          return this.mcpClient!.callTool(tool.server, tool.name, args, ctx.abortSignal);
+        },
+      });
+    }
+
+    // Listen for tool list changes
+    this.mcpClient.on("toolsChanged", async (serverName: string) => {
+      try {
+        await this.refreshMCPTools(serverName);
+      } catch (error) {
+        this.logger.error(
+          `Failed to refresh MCP tools for ${serverName}:`,
+          (error as Error).message,
+        );
+      }
+    });
+
+    const resources = await this.mcpClient.listAllResources();
+    const prompts = await this.mcpClient.listAllPrompts();
+
+    return {
+      tools: mcpTools.length,
+      resources: resources.length,
+      prompts: prompts.length,
+    };
+  }
+
+  /**
+   * Refresh tools from a specific MCP server
+   */
+  private async refreshMCPTools(serverName: string): Promise<void> {
+    // Remove existing tools from this server
+    const prefix = `mcp__${serverName}__`;
+    for (const name of this.tools.keys()) {
+      if (name.startsWith(prefix)) {
+        this.unregister(name);
+      }
+    }
+
+    // Re-register tools
+    const tools = await this.mcpClient!.listAllTools();
+    const serverTools = tools.filter((t) => t.server === serverName);
+
+    for (const tool of serverTools) {
+      const isReadOnly = isMcpToolReadOnly(tool);
+      this.register({
+        name: `mcp__${tool.server}__${tool.name}`,
+        description: tool.description || `MCP tool from ${tool.server}`,
+        category: "meta",
+        riskLevel: "moderate",
+        isReadOnly,
+        loadingTier: "always",
+        parameters: tool.inputSchema,
+        execute: async (args, ctx) => {
+          if (ctx.abortSignal?.aborted) {
+            return { success: false, message: "Cancelled" };
+          }
+          return this.mcpClient!.callTool(tool.server, tool.name, args, ctx.abortSignal);
+        },
+      });
+    }
+  }
+
+  /**
+   * Disconnect from all MCP servers
+   */
+  async disconnectMCP(): Promise<void> {
+    // Remove all MCP tools
+    for (const name of this.tools.keys()) {
+      if (name.startsWith("mcp__")) {
+        this.unregister(name);
+      }
+    }
+    await this.mcpClient?.disconnectAll();
+    this.mcpClient = undefined;
+  }
+
+  /**
+   * Get the MCP client for direct resource/prompt access
+   */
+  getMCPClient(): MCPClient | undefined {
+    return this.mcpClient;
+  }
+
+  /**
+   * Dispose of the registry, cleaning up all MCP connections.
+   *
+   * Shuts down every MCP server connection and clears all registered tools.
+   * Safe to call multiple times (idempotent).
+   */
+  async dispose(): Promise<void> {
+    await this.disconnectMCP();
+    this.tools.clear();
+  }
+
+  /**
+   * Async disposable support for `await using registry = new ToolRegistry()`.
+   *
+   * Delegates to dispose() so that ToolRegistry works with the Explicit
+   * Resource Management proposal (using/await using).
+   */
+  async [Symbol.asyncDispose](): Promise<void> {
+    await this.dispose();
+  }
+}