@huyooo/ai-chat-core 0.2.45 → 0.3.3

package/src/tool-manager/manager.ts

@@ -0,0 +1,562 @@
+/**
+ * ToolRuntimeManager — 工具运行时管理器
+ *
+ * 所有工具来源（InProcessProvider / McpProvider / 未来的 RemoteProvider）
+ * 统一注册到此处，提供跨来源搜索、会话级启用/禁用、变更通知。
+ *
+ * 新增能力：
+ * - init 失败自动清理（不残留 broken provider）
+ * - 工具名称冲突检测（警告日志）
+ * - 精确查找 getDescriptor / hasDescriptor
+ * - 中间件机制（beforeExecute / afterExecute / onError）
+ * - 工具执行统计（调用次数、成功率、耗时）
+ * - 连接状态事件转发
+ * - 重连后自动恢复会话工具
+ */
+
+import type { Tool, ToolContext, ToolResult } from '../types';
+import { normalizeToolResult } from '../types';
+import { validateJsonSchemaArgs } from './define-tool';
+import type {
+  ToolRuntimeManager,
+  ToolProvider,
+  ToolDescriptor,
+  ToolChangeListener,
+  ToolChangeEvent,
+  SessionToolsChangeListener,
+  SessionToolsChangeEvent,
+  ToolSearchQuery,
+  CategorySummary,
+  ProviderInfo,
+  ToolMiddleware,
+  ToolExecutionStats,
+  ConnectionStatusListener,
+  ConnectionStatusEvent,
+  McpConnectionInfo,
+} from './types';
+import { createModuleLogger } from '../logger';
+import { createAliasedTool, getToolIdentityKey, matchesToolDescriptorRef } from './identity';
+
+const logger = createModuleLogger('ToolManager');
+
+/** 关键词匹配（大小写不敏感） */
+function matchesKeyword(d: ToolDescriptor, keyword: string): boolean {
+  const lower = keyword.toLowerCase();
+  if (d.name.toLowerCase().includes(lower)) return true;
+  if (d.assetId.toLowerCase().includes(lower)) return true;
+  if (d.alias.toLowerCase().includes(lower)) return true;
+  if (d.displayName.toLowerCase().includes(lower)) return true;
+  if (d.id.toLowerCase().includes(lower)) return true;
+  if (d.description.toLowerCase().includes(lower)) return true;
+  if (d.tags.some(t => t.toLowerCase().includes(lower))) return true;
+  if (d.category.toLowerCase().includes(lower)) return true;
+  return false;
+}
+
+export function createToolRuntimeManager(): ToolRuntimeManager {
+  const providers = new Map<string, ToolProvider>();
+  const listeners = new Set<ToolChangeListener>();
+  const connectionStatusListeners = new Set<ConnectionStatusListener>();
+  const sessionToolsChangeListeners = new Set<SessionToolsChangeListener>();
+
+  /** descriptor.id → provider id */
+  const idIndex = new Map<string, string>();
+  /** alias → descriptor.id（alias 冲突时保留先注册的） */
+  const aliasIndex = new Map<string, string>();
+
+  /** 会话级原始工具（descriptor.id → 原始 Tool） */
+  const sessionOriginals = new Map<string, Tool>();
+  /** 会话级包装后工具（alias → 对外 Tool） */
+  const sessionTools = new Map<string, Tool>();
+
+  /** 中间件列表 */
+  const middlewares: ToolMiddleware[] = [];
+
+  /** 执行统计 */
+  const stats = new Map<string, ToolExecutionStats>();
+
+  /** 重连期间暂存待恢复的工具 id（provider id → tool ids） */
+  const pendingReEnable = new Map<string, Set<string>>();
+
+  // ══════════════════════ 内部函数 ══════════════════════
+
+  function recordStats(toolName: string, success: boolean, durationMs: number) {
+    const existing = stats.get(toolName) ?? {
+      totalCalls: 0, successCount: 0, errorCount: 0, totalDurationMs: 0,
+    };
+    existing.totalCalls++;
+    if (success) existing.successCount++;
+    else existing.errorCount++;
+    existing.totalDurationMs += durationMs;
+    existing.lastCallAt = Date.now();
+    stats.set(toolName, existing);
+  }
+
+  async function emitSessionToolsChange(reason: SessionToolsChangeEvent['reason']) {
+    const event: SessionToolsChangeEvent = {
+      reason,
+      toolNames: [...sessionTools.keys()],
+    };
+    await Promise.all([...sessionToolsChangeListeners].map(fn => Promise.resolve(fn(event))));
+  }
+
+  function getDescriptorById(id: string): ToolDescriptor | undefined {
+    const providerId = idIndex.get(id);
+    if (!providerId) return undefined;
+    const provider = providers.get(providerId);
+    return provider?.getDescriptors().find(d => d.id === id);
+  }
+
+  function resolveDescriptor(ref: string): ToolDescriptor | undefined {
+    const direct = getDescriptorById(ref);
+    if (direct) return direct;
+
+    const idFromAlias = aliasIndex.get(ref);
+    if (idFromAlias) {
+      const byAlias = getDescriptorById(idFromAlias);
+      if (byAlias) return byAlias;
+    }
+
+    for (const provider of providers.values()) {
+      const descriptor = provider.getDescriptors().find(d => matchesToolDescriptorRef(d, ref));
+      if (descriptor) return descriptor;
+    }
+    return undefined;
+  }
+
+  async function loadDescriptorTool(descriptor: ToolDescriptor): Promise<Tool | null> {
+    const provider = providers.get(descriptor.provider);
+    if (!provider) return null;
+    return provider.loadTool(descriptor.id);
+  }
+
+  /** 用当前中间件包装工具的 execute */
+  function wrapWithMiddleware(tool: Tool): Tool {
+    const toolKey = getToolIdentityKey(tool);
+    return {
+      ...tool,
+      execute: async (args: Record<string, unknown>, context: ToolContext) => {
+        let processedArgs = args;
+
+        for (const m of middlewares) {
+          if (m.beforeExecute) {
+            const modified = await m.beforeExecute(tool.name, processedArgs);
+            if (modified) processedArgs = modified;
+          }
+        }
+
+        const start = Date.now();
+        try {
+          const validatedArgs = validateJsonSchemaArgs(tool.parameters, processedArgs);
+          let result: ToolResult = normalizeToolResult(await tool.execute(validatedArgs, context));
+          const duration = Date.now() - start;
+
+          for (const m of middlewares) {
+            if (m.afterExecute) {
+              const modified = await m.afterExecute(tool.name, validatedArgs, result, duration);
+              if (modified) result = modified;
+            }
+          }
+
+          recordStats(toolKey, true, duration);
+          return result;
+        } catch (err) {
+          const duration = Date.now() - start;
+          recordStats(toolKey, false, duration);
+
+          for (const m of middlewares) {
+            if (m.onError) {
+              await m.onError(tool.name, processedArgs, err).catch(() => { });
+            }
+          }
+          throw err;
+        }
+      },
+    };
+  }
+
+  /** 异步重载会话中的工具（重连 updated 事件触发） */
+  function reloadSessionTool(ref: string) {
+    const descriptor = resolveDescriptor(ref);
+    if (!descriptor) return;
+    loadDescriptorTool(descriptor).then(tool => {
+      if (!tool) return;
+      sessionOriginals.set(descriptor.id, tool);
+      sessionTools.set(descriptor.alias, wrapWithMiddleware(createAliasedTool(tool, descriptor)));
+      void emitSessionToolsChange('reloaded');
+    }).catch(() => { });
+  }
+
+  /** 转发 Provider 变更到 Manager 级别 */
+  function onProviderChange(event: ToolChangeEvent) {
+    let sessionChanged = false;
+
+    if (event.type === 'removed') {
+      // 重连场景：暂存待恢复的工具名
+      if (event.reconnecting) {
+        const names = new Set(
+          event.tools.map(d => d.id).filter(id => sessionOriginals.has(id)),
+        );
+        if (names.size > 0) {
+          pendingReEnable.set(event.provider, names);
+        }
+      }
+
+      for (const d of event.tools) {
+        if (idIndex.get(d.id) === event.provider) idIndex.delete(d.id);
+        if (aliasIndex.get(d.alias) === d.id) aliasIndex.delete(d.alias);
+        if (sessionOriginals.delete(d.id)) {
+          sessionChanged = true;
+        }
+        if (sessionTools.delete(d.alias)) {
+          sessionChanged = true;
+        }
+      }
+    }
+
+    if (event.type === 'added') {
+      for (const d of event.tools) {
+        const existingAlias = aliasIndex.get(d.alias);
+        if (existingAlias && existingAlias !== d.id) {
+          logger.warn(
+            { alias: d.alias, newProvider: event.provider, existingDescriptorId: existingAlias },
+            `工具 alias 冲突: "${d.alias}" (${event.provider}) 与已有 (${existingAlias}) 冲突，保留先注册的`,
+          );
+        } else {
+          aliasIndex.set(d.alias, d.id);
+        }
+        idIndex.set(d.id, event.provider);
+      }
+
+      // 重连后自动恢复会话工具
+      const pending = pendingReEnable.get(event.provider);
+      if (pending) {
+        const toReEnable = event.tools.map(d => d.id).filter(id => pending.has(id));
+        if (toReEnable.length > 0) {
+          manager.enable(toReEnable).catch(() => { });
+        }
+        pendingReEnable.delete(event.provider);
+      }
+    }
+
+    if (event.type === 'updated') {
+      for (const d of event.tools) {
+        idIndex.set(d.id, event.provider);
+        for (const [alias, descriptorId] of aliasIndex) {
+          if (descriptorId === d.id && alias !== d.alias) {
+            aliasIndex.delete(alias);
+          }
+        }
+        aliasIndex.set(d.alias, d.id);
+        if (sessionOriginals.has(d.id)) {
+          for (const [alias, sessionTool] of sessionTools) {
+            if (getToolIdentityKey(sessionTool) === d.id) {
+              sessionTools.delete(alias);
+              sessionChanged = true;
+            }
+          }
+          reloadSessionTool(d.id);
+        }
+      }
+    }
+
+    if (sessionChanged) {
+      void emitSessionToolsChange('provider_changed');
+    }
+
+    for (const fn of listeners) fn(event);
+  }
+
+  /** 转发连接状态事件 */
+  function onConnectionStatusChange(event: ConnectionStatusEvent) {
+    for (const fn of connectionStatusListeners) fn(event);
+  }
+
+  // ══════════════════════ Manager 实现 ══════════════════════
+
+  const manager: ToolRuntimeManager = {
+    async addProvider(provider: ToolProvider) {
+      if (providers.has(provider.id)) {
+        await manager.removeProvider(provider.id);
+      }
+
+      // 先注册 listener，init 中的事件才能被捕获
+      provider.addChangeListener(onProviderChange);
+      if (provider.addConnectionStatusListener) {
+        provider.addConnectionStatusListener(onConnectionStatusChange);
+      }
+
+      try {
+        await provider.init();
+        // init 成功后才加入 map
+        providers.set(provider.id, provider);
+      } catch (e) {
+        // init 失败，清理 listener 和可能残留的索引
+        provider.removeChangeListener(onProviderChange);
+        if (provider.removeConnectionStatusListener) {
+          provider.removeConnectionStatusListener(onConnectionStatusChange);
+        }
+        for (const [toolId, pid] of idIndex) {
+          if (pid === provider.id) idIndex.delete(toolId);
+        }
+        for (const [alias, toolId] of aliasIndex) {
+          if (!idIndex.has(toolId)) aliasIndex.delete(alias);
+        }
+        throw e;
+      }
+    },
+
+    async removeProvider(id: string) {
+      const provider = providers.get(id);
+      if (!provider) return;
+      provider.removeChangeListener(onProviderChange);
+      if (provider.removeConnectionStatusListener) {
+        provider.removeConnectionStatusListener(onConnectionStatusChange);
+      }
+
+      const descriptors = provider.getDescriptors();
+      let sessionChanged = false;
+      for (const d of descriptors) {
+        if (idIndex.get(d.id) === id) idIndex.delete(d.id);
+        if (aliasIndex.get(d.alias) === d.id) aliasIndex.delete(d.alias);
+        if (sessionOriginals.delete(d.id)) {
+          sessionChanged = true;
+        }
+        if (sessionTools.delete(d.alias)) {
+          sessionChanged = true;
+        }
+      }
+
+      await provider.destroy();
+      providers.delete(id);
+      if (sessionChanged) {
+        await emitSessionToolsChange('provider_changed');
+      }
+    },
+
+    getProviders(): ProviderInfo[] {
+      return Array.from(providers.values()).map(p => ({
+        id: p.id,
+        source: p.source,
+        toolCount: p.getDescriptors().length,
+      }));
+    },
+
+    getProvider(id: string): ToolProvider | undefined {
+      return providers.get(id);
+    },
+
+    getDescriptor(ref: string): ToolDescriptor | undefined {
+      return resolveDescriptor(ref);
+    },
+
+    hasDescriptor(ref: string): boolean {
+      return resolveDescriptor(ref) !== undefined;
+    },
+
+    getAllDescriptors(): ToolDescriptor[] {
+      const result: ToolDescriptor[] = [];
+      for (const p of providers.values()) result.push(...p.getDescriptors());
+      return result;
+    },
+
+    search(query?: ToolSearchQuery): ToolDescriptor[] {
+      let results = manager.getAllDescriptors();
+
+      if (query?.keyword) {
+        const kw = query.keyword;
+        results = results.filter(d => matchesKeyword(d, kw));
+      }
+      if (query?.category) {
+        const cat = query.category.toLowerCase();
+        results = results.filter(d => d.category.toLowerCase() === cat);
+      }
+      if (query?.tags?.length) {
+        const requiredTags = query.tags.map(t => t.toLowerCase());
+        results = results.filter(d => {
+          const entryTags = d.tags.map(t => t.toLowerCase());
+          return requiredTags.every(rt => entryTags.includes(rt));
+        });
+      }
+      if (query?.source) {
+        results = results.filter(d => d.source === query.source);
+      }
+
+      return results;
+    },
+
+    async enable(names: string[]): Promise<Tool[]> {
+      const loaded: Tool[] = [];
+      let sessionChanged = false;
+
+      for (const ref of names) {
+        const descriptor = resolveDescriptor(ref);
+        if (!descriptor) continue;
+
+        if (sessionTools.has(descriptor.alias)) {
+          const existing = sessionTools.get(descriptor.alias);
+          if (existing) loaded.push(existing);
+          continue;
+        }
+
+        const tool = await loadDescriptorTool(descriptor);
+        if (!tool) continue;
+
+        sessionOriginals.set(descriptor.id, tool);
+        const wrapped = wrapWithMiddleware(createAliasedTool(tool, descriptor));
+        sessionTools.set(descriptor.alias, wrapped);
+        loaded.push(wrapped);
+        sessionChanged = true;
+      }
+
+      if (sessionChanged) {
+        await emitSessionToolsChange('enabled');
+      }
+      return loaded;
+    },
+
+    disable(names: string[]) {
+      let sessionChanged = false;
+      for (const ref of names) {
+        const descriptor = resolveDescriptor(ref);
+        if (descriptor) {
+          if (sessionOriginals.delete(descriptor.id)) {
+            sessionChanged = true;
+          }
+          if (sessionTools.delete(descriptor.alias)) {
+            sessionChanged = true;
+          }
+          continue;
+        }
+        if (sessionTools.delete(ref)) {
+          sessionChanged = true;
+        }
+      }
+      if (sessionChanged) {
+        void emitSessionToolsChange('disabled');
+      }
+    },
+
+    async enableAll(): Promise<Tool[]> {
+      const loaded: Tool[] = [];
+      let sessionChanged = false;
+
+      // 按 provider 批量加载，避免反复遍历
+      for (const p of providers.values()) {
+        const descriptors = p.getDescriptors();
+        for (const d of descriptors) {
+          if (sessionTools.has(d.alias)) {
+            const existing = sessionTools.get(d.alias);
+            if (existing) loaded.push(existing);
+            continue;
+          }
+          const tool = await p.loadTool(d.id);
+          if (tool) {
+            sessionOriginals.set(d.id, tool);
+            const wrapped = wrapWithMiddleware(createAliasedTool(tool, d));
+            sessionTools.set(d.alias, wrapped);
+            loaded.push(wrapped);
+            sessionChanged = true;
+          }
+        }
+      }
+
+      if (sessionChanged) {
+        await emitSessionToolsChange('enabled');
+      }
+      return loaded;
+    },
+
+    getSessionTools(): Map<string, Tool> {
+      return sessionTools;
+    },
+
+    getSessionToolNames(): string[] {
+      return [...sessionTools.keys()];
+    },
+
+    use(middleware: ToolMiddleware) {
+      middlewares.push(middleware);
+      // 用新的中间件链重新包装所有已启用的工具
+      let sessionChanged = false;
+      for (const [toolId, original] of sessionOriginals) {
+        const descriptor = getDescriptorById(toolId);
+        if (!descriptor) continue;
+        sessionTools.set(descriptor.alias, wrapWithMiddleware(createAliasedTool(original, descriptor)));
+        sessionChanged = true;
+      }
+      if (sessionChanged) {
+        void emitSessionToolsChange('middleware_updated');
+      }
+    },
+
+    addChangeListener(listener: ToolChangeListener) {
+      listeners.add(listener);
+    },
+
+    removeChangeListener(listener: ToolChangeListener) {
+      listeners.delete(listener);
+    },
+
+    addConnectionStatusListener(listener: ConnectionStatusListener) {
+      connectionStatusListeners.add(listener);
+    },
+
+    removeConnectionStatusListener(listener: ConnectionStatusListener) {
+      connectionStatusListeners.delete(listener);
+    },
+
+    addSessionToolsChangeListener(listener: SessionToolsChangeListener) {
+      sessionToolsChangeListeners.add(listener);
+    },
+
+    removeSessionToolsChangeListener(listener: SessionToolsChangeListener) {
+      sessionToolsChangeListeners.delete(listener);
+    },
+
+    categories(): CategorySummary[] {
+      const counts = new Map<string, number>();
+      for (const d of manager.getAllDescriptors()) {
+        counts.set(d.category, (counts.get(d.category) ?? 0) + 1);
+      }
+      return Array.from(counts.entries()).map(([category, count]) => ({ category, count }));
+    },
+
+    getStats(): Map<string, ToolExecutionStats> {
+      return new Map(stats);
+    },
+
+    getToolStats(ref: string): ToolExecutionStats | undefined {
+      const descriptor = resolveDescriptor(ref);
+      if (descriptor) return stats.get(descriptor.id);
+      return stats.get(ref);
+    },
+
+    getMcpConnectionInfos(): McpConnectionInfo[] {
+      const infos: McpConnectionInfo[] = [];
+      for (const p of providers.values()) {
+        if (p.getConnectionInfo) {
+          infos.push(p.getConnectionInfo());
+        }
+      }
+      return infos;
+    },
+
+    async destroy() {
+      const ids = [...providers.keys()];
+      await Promise.allSettled(ids.map(id => manager.removeProvider(id)));
+      sessionOriginals.clear();
+      sessionTools.clear();
+      idIndex.clear();
+      aliasIndex.clear();
+      stats.clear();
+      pendingReEnable.clear();
+      listeners.clear();
+      connectionStatusListeners.clear();
+      await emitSessionToolsChange('destroyed');
+      sessionToolsChangeListeners.clear();
+    },
+  };
+
+  return manager;
+}