@nocobase/ai 2.1.9 → 2.1.10

package/src/mcp-manager/user-context-client-manager.ts

@@ -0,0 +1,135 @@
+/**
+ * This file is part of the NocoBase (R) project.
+ * Copyright (c) 2020-2024 NocoBase Co., Ltd.
+ * Authors: NocoBase Team.
+ *
+ * This project is dual-licensed under AGPL-3.0 and NocoBase Commercial License.
+ * For more information, please refer to: https://www.nocobase.com/agreement.
+ */
+
+import type { Context } from '@nocobase/actions';
+import { MultiServerMCPClient, StdioConnection, StreamableHTTPConnection } from '@langchain/mcp-adapters';
+import { StructuredToolInterface } from '@langchain/core/tools';
+import type { MCPEntry, MCPOptions } from './types';
+import { renderMCPOptions } from './options-renderer';
+
+type MCPConnection = StdioConnection | StreamableHTTPConnection;
+
+type CacheEntry = {
+  client: MultiServerMCPClient;
+  toolsMap: Record<string, StructuredToolInterface[]>;
+  expiresAt: number;
+  lastAccessedAt: number;
+};
+
+export type UserContextMCPClientManagerOptions = {
+  app: any;
+  listEntries: () => Promise<MCPEntry[]>;
+  buildConnection: (options: MCPOptions) => MCPConnection;
+  ttlMs?: number;
+  maxSize?: number;
+};
+
+export class UserContextMCPClientManager {
+  private readonly cache = new Map<string, CacheEntry>();
+  private readonly ttlMs: number;
+  private readonly maxSize: number;
+
+  constructor(private readonly options: UserContextMCPClientManagerOptions) {
+    this.ttlMs = options.ttlMs ?? 5 * 60 * 1000;
+    this.maxSize = options.maxSize ?? 100;
+  }
+
+  async getToolsMap(ctx: Context): Promise<Record<string, StructuredToolInterface[]>> {
+    const currentUser = ctx?.state?.currentUser ?? ctx?.auth?.user;
+    if (!currentUser?.id) {
+      return {};
+    }
+
+    this.evictExpired();
+
+    const entries = (await this.options.listEntries())
+      .filter((entry) => entry.enabled !== false && entry.useUserContext === true && entry.transport !== 'stdio')
+      .sort((left, right) => left.name.localeCompare(right.name));
+    if (!entries.length) {
+      return {};
+    }
+
+    const cacheKey = String(currentUser.id);
+    const cached = this.cache.get(cacheKey);
+    const now = Date.now();
+    if (cached && cached.expiresAt > now) {
+      cached.lastAccessedAt = now;
+      return cached.toolsMap;
+    }
+
+    if (cached) {
+      await this.closeEntry(cached);
+      this.cache.delete(cacheKey);
+    }
+
+    const connections: Record<string, MCPConnection> = {};
+    for (const entry of entries) {
+      const rendered = await renderMCPOptions(entry, this.options.app, ctx);
+      connections[entry.name] = this.options.buildConnection(rendered);
+    }
+
+    const client = new MultiServerMCPClient(connections);
+    const initializedToolsMap = await client.initializeConnections();
+    const toolsMap = Object.fromEntries(
+      Object.entries(initializedToolsMap).map(([serverName, tools]) => [
+        serverName,
+        tools as StructuredToolInterface[],
+      ]),
+    );
+
+    this.cache.set(cacheKey, {
+      client,
+      toolsMap,
+      expiresAt: now + this.ttlMs,
+      lastAccessedAt: now,
+    });
+    await this.evictOversized();
+
+    return toolsMap;
+  }
+
+  async clear() {
+    const entries = [...this.cache.values()];
+    this.cache.clear();
+    await Promise.all(entries.map((entry) => this.closeEntry(entry)));
+  }
+
+  private evictExpired() {
+    const now = Date.now();
+    for (const [key, entry] of this.cache.entries()) {
+      if (entry.expiresAt <= now) {
+        this.cache.delete(key);
+        this.closeEntry(entry).catch((error) => {
+          this.options.app?.log?.warn?.('fail to close expired user-bound mcp client', error);
+        });
+      }
+    }
+  }
+
+  private async evictOversized() {
+    while (this.cache.size > this.maxSize) {
+      const oldest = [...this.cache.entries()].sort(
+        ([, left], [, right]) => left.lastAccessedAt - right.lastAccessedAt,
+      )[0];
+      if (!oldest) {
+        return;
+      }
+      this.cache.delete(oldest[0]);
+      await this.closeEntry(oldest[1]);
+    }
+  }
+
+  private async closeEntry(entry: CacheEntry) {
+    try {
+      await entry.client.close();
+    } catch (error) {
+      this.options.app?.log?.warn?.('fail to close user-bound mcp client', error);
+    }
+  }
+}

package/src/tools-manager/types.ts

@@ -56,4 +56,5 @@ export type ToolsFilter = {
   defaultPermission?: Permission;
   silence?: boolean;
   sessionId?: string;
+  ctx?: Context;
 };