UNPKG

@juspay/neurolink

Version:

Universal AI Development Platform with working MCP integration, multi-provider support, voice (TTS/STT/realtime), and professional CLI. 58+ external MCP servers discoverable, multimodal file processing, RAG pipelines. Build, test, and deploy AI applicatio

750 lines (749 loc) 33 kB
/** * MCP Tool Registry - Extended Registry with Tool Management * Updated to match industry standard camelCase interfaces */ import { MCPRegistry } from "./registry.js"; import { registryLogger } from "../utils/logger.js"; import { randomUUID } from "crypto"; import { shouldDisableBuiltinTools } from "../utils/toolUtils.js"; import { directAgentTools } from "../agent/directTools.js"; import { detectCategory, createMCPServerInfo } from "../utils/mcpDefaults.js"; import { FlexibleToolValidator } from "./flexibleToolValidator.js"; import { ErrorFactory } from "../utils/errorHandling.js"; import { HITLUserRejectedError, HITLTimeoutError } from "../hitl/hitlErrors.js"; import { withSpan, tracers, ATTR } from "../telemetry/index.js"; import { SpanStatusCode } from "@opentelemetry/api"; import { getAuthContext } from "../auth/authContext.js"; export class MCPToolRegistry extends MCPRegistry { tools = new Map(); toolImplementations = new Map(); // Store actual tool implementations toolExecutionStats = new Map(); builtInServerInfos = []; // DIRECT storage for MCPServerInfo hitlManager; // Optional HITL manager for safety mechanisms constructor() { super(); if (!shouldDisableBuiltinTools()) { this.registerDirectTools(); } } /** * Set HITL manager for human-in-the-loop safety mechanisms * @param hitlManager - HITL manager instance (optional, can be undefined to disable) */ setHITLManager(hitlManager) { this.hitlManager = hitlManager; if (hitlManager && hitlManager.isEnabled()) { registryLogger.info("HITL safety mechanisms enabled for tool execution"); } else { registryLogger.debug("HITL safety mechanisms disabled or not configured"); } } /** * Get current HITL manager */ getHITLManager() { return this.hitlManager; } /** * Register all direct tools from directAgentTools */ registerDirectTools() { registryLogger.debug("Auto-registering direct tools..."); for (const [toolName, toolDef] of Object.entries(directAgentTools)) { // Skip undefined tools if (!toolDef) { registryLogger.warn(`Skipping undefined tool during registration: ${toolName}`); continue; } const toolId = `direct.${toolName}`; const toolInfo = { name: toolName, description: toolDef.description || `Direct tool: ${toolName}`, inputSchema: {}, serverId: "direct", category: detectCategory({ isBuiltIn: true, serverId: "direct" }), }; this.tools.set(toolId, toolInfo); this.toolImplementations.set(toolId, { execute: async (params, context) => { try { // Direct tools from AI SDK expect their specific parameter structure // Each tool validates its own parameters, so we safely pass them through const result = await toolDef.execute(params, { toolCallId: context?.sessionId || "unknown", messages: [], }); // Return the result wrapped in our standard format return { success: true, data: result, metadata: { toolName, serverId: "direct", executionTime: 0, }, }; } catch (error) { return { success: false, error: error instanceof Error ? error.message : String(error), metadata: { toolName, serverId: "direct", executionTime: 0, }, }; } }, description: toolDef.description, inputSchema: {}, }); registryLogger.debug(`Registered direct tool: ${toolName} as ${toolId}`); } registryLogger.debug(`Auto-registered ${Object.keys(directAgentTools).length} direct tools`); } async registerServer(serverInfoOrId, _serverConfigOrContext, _context) { // Handle both signatures for backward compatibility let serverInfo; if (typeof serverInfoOrId === "string") { // Legacy signature: registerServer(serverId, serverConfig, context) const serverId = serverInfoOrId; // Convert legacy call to MCPServerInfo format using smart defaults serverInfo = createMCPServerInfo({ id: serverId, name: serverId, tools: [], isExternal: true, }); } else { // New signature: registerServer(serverInfo, context) serverInfo = serverInfoOrId; } const serverId = serverInfo.id; // Use MCPServerInfo.tools array directly - ZERO conversions! const toolsObject = {}; for (const tool of serverInfo.tools) { toolsObject[tool.name] = { execute: tool.execute || (async () => { throw new Error(`Tool ${tool.name} has no execute function`); }), description: tool.description, inputSchema: tool.inputSchema, category: detectCategory({ existingCategory: serverInfo.metadata?.category, serverId: serverInfo.id, }), }; } const plugin = { metadata: { name: serverInfo.name, description: serverInfo.description, category: detectCategory({ existingCategory: serverInfo.metadata?.category, serverId: serverInfo.id, }), }, tools: toolsObject, configuration: {}, }; // Call the parent register method this.register(plugin); // Use MCPServerInfo.tools array directly - ZERO conversions! const tools = serverInfo.tools; for (const tool of tools) { // For custom tools, use just the tool name to avoid redundant serverId.toolName format // For other tools, use fully-qualified serverId.toolName to avoid collisions const isCustomTool = serverId.startsWith("custom-tool-"); const toolId = isCustomTool ? tool.name : `${serverId}.${tool.name}`; const toolTimeoutMs = serverInfo.metadata?.toolTimeoutMs; const toolMaxRetries = serverInfo.metadata?.toolMaxRetries; const toolInfo = { name: tool.name, description: tool.description, inputSchema: tool.inputSchema, outputSchema: undefined, // MCPServerInfo.tools doesn't have outputSchema serverId, category: detectCategory({ existingCategory: serverInfo.metadata?.category, serverId: serverInfo.id, }), permissions: [], // MCPServerInfo.tools doesn't have permissions ...(toolTimeoutMs !== undefined && { timeoutMs: toolTimeoutMs }), ...(toolMaxRetries !== undefined && { maxRetries: toolMaxRetries }), }; // Register only with fully-qualified toolId to avoid collisions this.tools.set(toolId, toolInfo); // Store the actual tool implementation for execution using toolId as key this.toolImplementations.set(toolId, { execute: tool.execute || (async () => { throw new Error(`Tool ${tool.name} has no execute function`); }), description: tool.description, inputSchema: tool.inputSchema, category: detectCategory({ existingCategory: serverInfo.metadata?.category, serverId: serverInfo.id, }), ...(toolTimeoutMs !== undefined && { timeoutMs: toolTimeoutMs }), ...(toolMaxRetries !== undefined && { maxRetries: toolMaxRetries }), }); // Tool registered successfully } // Store MCPServerInfo directly - NO recreation needed! if (tools.length > 0) { const category = detectCategory({ existingCategory: serverInfo.metadata?.category, serverId: serverInfo.id, }); // Only store in builtInServerInfos if it's a real in-memory MCP server // Do NOT create fake servers for built-in direct tools if (category === "in-memory") { // Use the original MCPServerInfo directly - ZERO conversions! this.builtInServerInfos.push(serverInfo); registryLogger.debug(`Added ${category} server to builtInServerInfos: ${serverId} with ${tools.length} tools`); } } } /** * Execute a tool with enhanced context and automatic result wrapping * * This method handles both raw return values and ToolResult objects: * - Raw values (primitives, objects) are automatically wrapped in ToolResult format * - Existing ToolResult objects are enhanced with execution metadata * - All results include execution timing and context information * * @param toolName - Name of the tool to execute * @param args - Parameters to pass to the tool execution function * @param context - Execution context with session, user, and environment info * @returns Promise resolving to ToolResult object with data, metadata, and usage info * @throws Error if tool is not found or execution fails * * @example * ```typescript * // Tool that returns raw value * const result = await toolRegistry.executeTool("calculator", { a: 5, b: 3, op: "add" }); * // result.data === 8, result.metadata contains execution info * * // Tool that returns ToolResult * const result = await toolRegistry.executeTool("complexTool", { input: "test" }); * // result is enhanced ToolResult with additional metadata * ``` */ async executeTool(toolName, args, context) { const startTime = Date.now(); // Resolve serverId eagerly for span attributes let preResolvedServerId; const toolEntry = this.tools.get(toolName); if (toolEntry) { preResolvedServerId = toolEntry.serverId; } else { for (const toolInfo of this.tools.values()) { if (toolInfo.name === toolName) { preResolvedServerId = toolInfo.serverId; break; } } } return withSpan({ name: "neurolink.tool.registry.execute", tracer: tracers.mcp, attributes: { [ATTR.GEN_AI_TOOL_NAME]: toolName, [ATTR.MCP_SERVER_ID]: preResolvedServerId || "builtin", // Curator P1-3: registry-level wrapper — duplicates ai.toolCall in // Langfuse. Retained for OTel/metrics; skipped for Langfuse export. "langfuse.internal": true, }, }, async (span) => { try { registryLogger.info(`🔧 [TOOL_EXECUTION] Starting execution: ${toolName}`, { hasArgs: args !== undefined, hasContext: context !== undefined, sessionId: context?.sessionId, }); const { tool, toolId } = this.resolveToolExecutionTarget(toolName); if (!tool) { throw new Error(`Tool '${toolName}' not found in registry`); } // Classify tool type for observability const serverId = tool.serverId || "unknown"; const toolType = serverId === "direct" ? "builtin" : serverId.startsWith("custom-tool-") ? "custom" : "mcp"; span.setAttribute("tool.type", toolType); span.setAttribute(ATTR.MCP_SERVER_ID, serverId); const execContext = this.createExecutionContext(context); // Get the tool implementation using the resolved toolId const toolImpl = this.toolImplementations.get(toolId); registryLogger.debug(`Looking for tool '${toolName}' (toolId: '${toolId}'), found: ${!!toolImpl}, type: ${typeof toolImpl?.execute}`); registryLogger.debug(`Available tools:`, Array.from(this.toolImplementations.keys())); if (!toolImpl || typeof toolImpl?.execute !== "function") { throw new Error(`Tool '${toolName}' implementation not found or not executable`); } // Capture argument metadata (avoid logging raw values which may contain secrets) let argsStr; try { argsStr = JSON.stringify(args).slice(0, 4096); } catch { argsStr = "[unserializable]"; } span.setAttribute("tool.arguments_present", args !== undefined); span.setAttribute("tool.arguments_size", argsStr.length); // HITL Safety Check: Request confirmation if required let finalArgs = args; const HITLState = context?.hitlState; if (!HITLState?.triggered && this.hitlManager && this.hitlManager.isEnabled()) { const requiresConfirmation = this.hitlManager.requiresConfirmation(toolName, args); if (requiresConfirmation) { registryLogger.info(`Tool '${toolName}' requires HITL confirmation`); span.addEvent("tool.hitl_requested"); try { if (HITLState) { HITLState.triggered = true; } const confirmationResult = await this.hitlManager.requestConfirmation(toolName, args, { serverId: tool.serverId, sessionId: execContext.sessionId, userId: execContext.userId, }); if (!confirmationResult.approved) { // User rejected the tool execution span.addEvent("tool.hitl_rejected"); throw new HITLUserRejectedError(`Tool execution rejected by user: ${confirmationResult.reason || "No reason provided"}`, toolName, confirmationResult.reason); } span.addEvent("tool.hitl_approved"); // User approved - use modified arguments if provided if (confirmationResult.modifiedArguments !== undefined) { finalArgs = confirmationResult.modifiedArguments; registryLogger.info(`Tool '${toolName}' arguments modified by user`); } registryLogger.info(`Tool '${toolName}' approved for execution (response time: ${confirmationResult.responseTime}ms)`); } catch (error) { if (error instanceof HITLTimeoutError) { // Timeout occurred - user didn't respond in time registryLogger.warn(`Tool '${toolName}' execution timed out waiting for user confirmation`); throw error; } else if (error instanceof HITLUserRejectedError) { // User explicitly rejected registryLogger.info(`Tool '${toolName}' execution rejected by user`); throw error; } else { // Other HITL error (configuration, system error, etc.) registryLogger.error(`HITL confirmation failed for tool '${toolName}':`, error); throw new Error(`HITL confirmation failed: ${error instanceof Error ? error.message : String(error)}`, { cause: error }); } } } else { registryLogger.debug(`Tool '${toolName}' does not require HITL confirmation`); } } // Execute the actual tool (with potentially modified arguments) registryLogger.debug(`Executing tool '${toolName}' with args:`, finalArgs); const toolResult = await toolImpl.execute(finalArgs, execContext); // Properly wrap raw results in ToolResult format let result; // Check if result is already a ToolResult object if (toolResult && typeof toolResult === "object" && "success" in toolResult && typeof toolResult.success === "boolean") { // Result is already a ToolResult, enhance with metadata const toolResultObj = toolResult; result = { ...toolResultObj, usage: { ...(toolResultObj.usage || {}), executionTime: Date.now() - startTime, }, metadata: { ...(toolResultObj.metadata || {}), toolName, serverId: tool.serverId, sessionId: execContext.sessionId, executionTime: Date.now() - startTime, }, }; } else { // Result is a raw value, wrap it in ToolResult format result = { success: true, data: toolResult, usage: { executionTime: Date.now() - startTime, }, metadata: { toolName, serverId: tool.serverId, sessionId: execContext.sessionId, executionTime: Date.now() - startTime, }, }; } // Update statistics const duration = Date.now() - startTime; this.updateStats(toolName, duration); // Record success on span let resultStr; try { resultStr = JSON.stringify(result.data) ?? "undefined"; } catch { resultStr = "[unserializable]"; } span.setAttribute("tool.result_length", resultStr.length); span.setAttribute("tool.success", true); registryLogger.debug(`Tool '${toolName}' executed successfully in ${duration}ms`); return result; } catch (error) { registryLogger.error(`Tool execution failed: ${toolName}`, error); // Record failure on span span.setAttribute("tool.success", false); // Rethrow precondition errors (tool not found, not executable) const errMsg = error instanceof Error ? error.message : String(error); if (errMsg.includes("not found in registry") || errMsg.includes("not executable")) { throw error; } // Explicitly set ERROR status — we're returning (not throwing) so withSpan // won't automatically detect this as an error (gap T3 fix) span.setStatus({ code: SpanStatusCode.ERROR, message: errMsg, }); if (error instanceof Error) { span.recordException(error); } // Return runtime execution errors in ToolResult format const errorResult = { success: false, data: null, error: error instanceof Error ? error.message : String(error), usage: { executionTime: Date.now() - startTime, }, metadata: { toolName, sessionId: context?.sessionId, }, }; return errorResult; } }); } async listTools(filterOrContext) { // FIXED: Return unique tools (avoid duplicates from dual registration) const uniqueTools = new Map(); for (const tool of this.tools.values()) { const key = `${tool.serverId || "unknown"}.${tool.name}`; if (!uniqueTools.has(key)) { uniqueTools.set(key, tool); } } let result = Array.from(uniqueTools.values()); // Determine if parameter is a filter object or just context let filter; if (filterOrContext) { // Check if it's a filter object (has filter-specific properties) or just context if ("sessionId" in filterOrContext || "userId" in filterOrContext) { // It's an ExecutionContext, treat as no filter filter = undefined; } else { // It's a filter object filter = filterOrContext; } } // Apply filters if provided if (filter) { if (filter.category) { result = result.filter((tool) => tool.category === filter.category); } if (filter.serverId) { result = result.filter((tool) => tool.serverId === filter.serverId); } if (filter.serverCategory) { result = result.filter((tool) => { const server = this.get(tool.serverId || ""); return server?.metadata?.category === filter.serverCategory; }); } if (filter.permissions && filter.permissions.length > 0) { result = result.filter((tool) => { const toolPermissions = tool.permissions || []; return (filter.permissions?.some((perm) => toolPermissions.includes(perm)) ?? false); }); } } registryLogger.debug(`Listed ${result.length} unique tools (${filter ? "filtered" : "unfiltered"})`); return result; } resolveToolExecutionTarget(toolName) { let tool = this.tools.get(toolName); registryLogger.info(`🔍 [TOOL_LOOKUP] Direct lookup result for '${toolName}':`, !!tool); let toolId = toolName; if (!tool) { const matches = Array.from(this.tools.entries()).filter(([, toolInfo]) => toolInfo.name === toolName); if (matches.length > 1) { throw ErrorFactory.toolExecutionFailed(toolName, new Error(`Ambiguous tool name '${toolName}'. Use fully-qualified name 'serverId.${toolName}'.`)); } if (matches.length === 1) { [toolId, tool] = matches[0]; } } return { tool, toolId }; } createExecutionContext(context) { let authUserId; try { authUserId = getAuthContext()?.user?.id; } catch { // Auth context not available — that's fine } return { ...context, sessionId: context?.sessionId ?? randomUUID(), userId: context?.userId ?? authUserId, }; } /** * Get tool information with server details */ getToolInfo(toolName) { // Try to find the tool by fully-qualified name first let tool = this.tools.get(toolName); // If not found, search for tool by name across all entries (for backward compatibility) if (!tool) { for (const toolInfo of this.tools.values()) { if (toolInfo.name === toolName) { tool = toolInfo; break; } } } if (!tool) { return undefined; } return { tool, server: { id: tool.serverId || "unknown-server", }, }; } /** * Update execution statistics */ updateStats(toolName, executionTime) { const stats = this.toolExecutionStats.get(toolName) || { count: 0, totalTime: 0, }; stats.count += 1; stats.totalTime += executionTime; this.toolExecutionStats.set(toolName, stats); } /** * Get execution statistics */ getExecutionStats() { const result = {}; for (const [toolName, stats] of this.toolExecutionStats.entries()) { result[toolName] = { count: stats.count, totalTime: stats.totalTime, averageTime: stats.totalTime / stats.count, }; } return result; } /** * Clear execution statistics */ clearStats() { this.toolExecutionStats.clear(); } /** * Get built-in servers * @returns Array of MCPServerInfo for built-in tools */ getBuiltInServerInfos() { return this.builtInServerInfos; } /** * Get tools by category */ getToolsByCategory(category) { // Return unique tools by fully-qualified toolId const uniqueTools = new Map(); for (const [toolId, tool] of this.tools.entries()) { if (tool.category === category && !uniqueTools.has(toolId)) { uniqueTools.set(toolId, tool); } } return Array.from(uniqueTools.values()); } /** * NL-001: Get available tools, filtering out those with OPEN circuit breakers. * Returns both the filtered tools and the list of unavailable tool names. */ getAvailableTools(circuitBreakers) { const allTools = Array.from(this.tools.values()); const unavailableTools = []; const tools = []; for (const tool of allTools) { const breakerKey = `${tool.serverId || "unknown"}.${tool.name}`; const breaker = circuitBreakers.get(breakerKey); if (breaker && breaker.getState() === "open") { unavailableTools.push(tool.name); } else { tools.push(tool); } } return { tools, unavailableTools }; } /** * Check if tool exists */ hasTool(toolName) { // Check by fully-qualified name first, then fallback to first matching tool name if (this.tools.has(toolName)) { return true; } for (const tool of this.tools.values()) { if (tool.name === toolName) { return true; } } return false; } /** * Register a tool with implementation directly * This is used for external MCP server tools */ async registerTool(toolId, toolInfo, toolImpl) { registryLogger.debug(`Registering tool: ${toolId}`); // Universal safety validation using FlexibleToolValidator // Only blocks truly dangerous cases to support maximum MCP tool compatibility const validation = FlexibleToolValidator.validateToolInfo(toolId, { description: toolInfo.description, serverId: toolInfo.serverId, }); if (!validation.isValid) { registryLogger.error(`Tool registration failed for ${toolId}: ${validation.error}`); throw new Error(`Tool validation failed: ${validation.error}`); } // Log any warnings but allow registration to proceed if (validation.warnings && validation.warnings.length > 0) { registryLogger.warn(`Tool registration warnings for ${toolId}:`, validation.warnings); } registryLogger.debug(`✅ Tool '${toolId}' passed flexible validation - registration proceeding`); this.tools.set(toolId, toolInfo); this.toolImplementations.set(toolId, toolImpl); registryLogger.debug(`Successfully registered tool: ${toolId}`); } /** * Remove a tool */ removeTool(toolName) { // Remove by fully-qualified name first, then fallback to first matching tool name let removed = false; if (this.tools.has(toolName)) { this.tools.delete(toolName); this.toolImplementations.delete(toolName); // Fix memory leak this.toolExecutionStats.delete(toolName); registryLogger.info(`Removed tool: ${toolName}`); removed = true; } else { // Remove all tools with matching name for (const [toolId, tool] of Array.from(this.tools.entries())) { if (tool.name === toolName) { this.tools.delete(toolId); this.toolImplementations.delete(toolId); // Fix memory leak this.toolExecutionStats.delete(toolId); registryLogger.info(`Removed tool: ${toolId}`); removed = true; } } } return removed; } /** * Get tool count */ getToolCount() { return this.tools.size; } /** * Get comprehensive statistics */ getStats() { const servers = this.list(); // Get all registered servers const allTools = Array.from(this.tools.values()); // Count servers by category const serversByCategory = {}; for (const server of servers) { const category = server.metadata?.category || "uncategorized"; serversByCategory[category] = (serversByCategory[category] || 0) + 1; } // Count tools by category const toolsByCategory = {}; for (const tool of allTools) { const category = tool.category || "uncategorized"; toolsByCategory[category] = (toolsByCategory[category] || 0) + 1; } return { totalServers: servers.length, totalTools: allTools.length, serversByCategory, toolsByCategory, executionStats: this.getExecutionStats(), }; } /** * Unregister a server */ unregisterServer(serverId) { // Remove all tools for this server const removedTools = []; for (const [toolId, tool] of this.tools.entries()) { if (tool.serverId === serverId) { this.tools.delete(toolId); this.toolImplementations.delete(toolId); // Fix memory leak this.toolExecutionStats.delete(toolId); // Fix memory leak removedTools.push(toolId); } } // Remove from builtInServerInfos storage const originalLength = this.builtInServerInfos.length; this.builtInServerInfos = this.builtInServerInfos.filter((server) => server.id !== serverId); const removedFromBuiltIn = originalLength > this.builtInServerInfos.length; // Remove from parent registry const removed = this.unregister(serverId); registryLogger.info(`Unregistered server ${serverId}, removed ${removedTools.length} tools${removedFromBuiltIn ? " and server from builtInServerInfos" : ""}`); return removed; } } // Create default instance export const toolRegistry = new MCPToolRegistry(); export const defaultToolRegistry = toolRegistry;