@nocobase/ai 2.2.0-beta.6 → 2.2.0-beta.8

package/src/mcp-manager/options-renderer.ts

@@ -0,0 +1,132 @@
+/**
+ * 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 { parse } from '@nocobase/utils';
+import type { MCPOptions } from './types';
+
+const unsafePathSegments = new Set(['__proto__', 'prototype', 'constructor']);
+const currentUserVariableRegExp = /{{\s*(?:(ctx)\.)?(currentUser|\$user)(?:\.([^}]+))?\s*}}/g;
+
+const hasUnsafePathSegment = (path: string) => {
+  return path
+    .split(/[.[\]]+/)
+    .filter(Boolean)
+    .some((segment) => unsafePathSegments.has(segment));
+};
+
+const toPlainObject = (value: any) => {
+  if (!value) {
+    return value;
+  }
+  if (typeof value.toJSON === 'function') {
+    return value.toJSON();
+  }
+  return value;
+};
+
+const getCurrentUserReferencePaths = (value: unknown) => {
+  const text = typeof value === 'string' ? value : JSON.stringify(value ?? {});
+  return Array.from(text.matchAll(currentUserVariableRegExp))
+    .map((match) => match[3]?.trim())
+    .filter((path): path is string => !!path && !hasUnsafePathSegment(path));
+};
+
+const getCurrentUserAppends = (paths: string[], currentUser: any) => {
+  return Array.from(
+    new Set(
+      paths.map((path) => path.split('.')[0]).filter((append) => append && !Reflect.has(currentUser || {}, append)),
+    ),
+  );
+};
+
+async function getCurrentUser(ctx?: Context, template?: unknown) {
+  const currentUser = ctx?.state?.currentUser ?? ctx?.auth?.user;
+  if (!ctx || !currentUser) {
+    return toPlainObject(currentUser);
+  }
+
+  const appends = getCurrentUserAppends(getCurrentUserReferencePaths(template), currentUser);
+  if (!appends.length) {
+    return toPlainObject(currentUser);
+  }
+
+  const user = await ctx.db.getRepository('users').findOne({
+    filterByTk: currentUser.id,
+    appends,
+  });
+  return toPlainObject(user) || toPlainObject(currentUser);
+}
+
+const stringifyRecord = (value: unknown): Record<string, string> => {
+  if (!value || typeof value !== 'object' || Array.isArray(value)) {
+    return {};
+  }
+  return Object.entries(value as Record<string, unknown>).reduce<Record<string, string>>((result, [key, item]) => {
+    if (!key || item == null) {
+      return result;
+    }
+    result[key] = String(item);
+    return result;
+  }, {});
+};
+
+const stringifyArray = (value: unknown): string[] => {
+  if (!Array.isArray(value)) {
+    return [];
+  }
+  return value.filter((item) => item != null).map((item) => String(item));
+};
+
+const getRequestVariables = (ctx?: Context) => ({
+  headers: ctx?.request?.headers ?? {},
+  token: ctx?.getBearerToken?.() ?? '',
+});
+
+const emptyRequestVariables = {
+  headers: {},
+  token: '',
+};
+
+export const normalizeMCPOptions = (options: MCPOptions): MCPOptions => {
+  const normalized: MCPOptions = {
+    ...options,
+    args: stringifyArray(options.args),
+    env: stringifyRecord(options.env),
+    headers: stringifyRecord(options.headers),
+    useUserContext: options.transport === 'stdio' ? false : options.useUserContext === true,
+  };
+
+  if (normalized.transport === 'stdio') {
+    normalized.url = undefined;
+    normalized.headers = {};
+  }
+  return normalized;
+};
+
+export async function renderMCPOptions(
+  options: MCPOptions,
+  app: { environment?: { getVariables?: () => Record<string, unknown> } },
+  ctx?: Context,
+): Promise<MCPOptions> {
+  const currentUser = options.useUserContext ? await getCurrentUser(ctx, options) : undefined;
+  const request = options.useUserContext ? getRequestVariables(ctx) : emptyRequestVariables;
+  const variables = {
+    $env: app.environment?.getVariables?.() ?? {},
+    currentUser,
+    $user: currentUser,
+    request,
+    ctx: {
+      currentUser,
+      request,
+    },
+  };
+
+  return normalizeMCPOptions(parse(options)(variables) as MCPOptions);
+}

package/src/mcp-manager/types.ts

@@ -8,18 +8,20 @@
  */
 
 import type { MultiServerMCPClient } from '@langchain/mcp-adapters';
+import type { Context } from '@nocobase/actions';
 import type { DynamicToolsProvider, Permission } from '../tools-manager/types';
 
 export interface MCPManager extends MCPRegistration {
   init(): Promise<void>;
   getMCP(name: string): Promise<MCPEntry>;
   listMCP(filter: MCPFilter): Promise<MCPEntry[]>;
-  testConnection(options: MCPOptions): Promise<MCPTestResult>;
+  testConnection(options: MCPOptions, ctx?: Context): Promise<MCPTestResult>;
   rebuildClient(): Promise<void>;
   getClient(): MultiServerMCPClient | null;
   getMCPToolsProvider(): DynamicToolsProvider;
-  listMCPTools(): Promise<Record<string, MCPToolEntry[]>>;
+  listMCPTools(ctx?: Context): Promise<Record<string, MCPToolEntry[]>>;
   updateMCPToolPermission(toolName: string, permission: Permission): Promise<void>;
+  clearUserContextCache(): Promise<void>;
 }
 
 export interface MCPRegistration {
@@ -34,6 +36,7 @@ export type MCPOptions = {
   url?: string;
   headers?: Record<string, string>;
   restart?: Record<string, any>;
+  useUserContext?: boolean;
 };
 
 export type MCPEntry = MCPOptions & {
@@ -45,6 +48,7 @@ export type MCPFilter = {
   name?: string;
   enabled?: boolean;
   transport?: MCPTransport;
+  useUserContext?: boolean;
 };
 
 export type MCPTransport = 'stdio' | 'sse' | 'http';

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

@@ -0,0 +1,150 @@
+/**
+ * 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 {};
+    }
+
+    let client: MultiServerMCPClient | null = null;
+    try {
+      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);
+      }
+
+      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,
+      });
+      client = null;
+      await this.evictOversized();
+
+      return toolsMap;
+    } catch (error) {
+      if (client) {
+        await this.closeEntry({
+          client,
+          toolsMap: {},
+          expiresAt: 0,
+          lastAccessedAt: 0,
+        });
+      }
+      this.options.app?.log?.warn?.('fail to get user-bound mcp tools', error);
+      return {};
+    }
+  }
+
+  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;
 };