diff --git a/packages/core/ai/src/__tests__/mcp-user-context.test.ts b/packages/core/ai/src/__tests__/mcp-user-context.test.ts new file mode 100644 index 00000000000..916196de49c --- /dev/null +++ b/packages/core/ai/src/__tests__/mcp-user-context.test.ts @@ -0,0 +1,265 @@ +/** + * 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. + */ + +const mcpClientMock = vi.hoisted(() => { + const instances: any[] = []; + + class MultiServerMCPClient { + connections: Record; + close = vi.fn(); + + constructor(connections: Record) { + this.connections = connections; + instances.push(this); + } + + async initializeConnections() { + return Object.fromEntries( + Object.keys(this.connections).map((serverName) => [ + serverName, + [ + { + name: 'getProfile', + description: `Get profile from ${serverName}`, + schema: {}, + invoke: vi.fn(async (args) => ({ serverName, args })), + }, + ], + ]), + ); + } + } + + return { + instances, + MultiServerMCPClient, + }; +}); + +vi.mock('@langchain/mcp-adapters', () => ({ + MultiServerMCPClient: mcpClientMock.MultiServerMCPClient, +})); + +import { DefaultMCPManager } from '../mcp-manager'; +import { normalizeMCPOptions, renderMCPOptions } from '../mcp-manager/options-renderer'; +import { UserContextMCPClientManager } from '../mcp-manager/user-context-client-manager'; + +describe('user-bound MCP clients', () => { + const createApp = () => ({ + environment: { + getVariables: () => ({ + MCP_HOST: 'mcp.example.test', + API_TOKEN: 'env-token', + }), + }, + log: { + warn: vi.fn(), + }, + }); + + const createCtx = (id: number, extraUser: Record = {}) => + ({ + state: { + currentUser: { + id, + name: `user-${id}`, + ...extraUser, + }, + }, + auth: { + user: { + id, + name: `user-${id}`, + ...extraUser, + }, + }, + db: { + getRepository: () => ({ + findOne: vi.fn(), + }), + }, + }) as any; + + beforeEach(() => { + mcpClientMock.instances.length = 0; + }); + + afterEach(() => { + vi.useRealTimers(); + }); + + it('normalizes stdio records as not user-bound', () => { + expect( + normalizeMCPOptions({ + transport: 'stdio', + command: 'node', + args: ['server.js'], + env: { TOKEN: '{{ $env.API_TOKEN }}' }, + useUserContext: true, + }), + ).toMatchObject({ + transport: 'stdio', + command: 'node', + args: ['server.js'], + env: { TOKEN: '{{ $env.API_TOKEN }}' }, + useUserContext: false, + headers: {}, + }); + }); + + it('renders environment and current user variables in MCP options', async () => { + const rendered = await renderMCPOptions( + { + transport: 'http', + url: 'https://{{ $env.MCP_HOST }}/users/{{ currentUser.id }}', + headers: { + Authorization: 'Bearer {{ $env.API_TOKEN }}', + 'X-User': '{{ $user.name }}', + }, + useUserContext: true, + }, + createApp(), + createCtx(7), + ); + + expect(rendered).toMatchObject({ + transport: 'http', + url: 'https://mcp.example.test/users/7', + headers: { + Authorization: 'Bearer env-token', + 'X-User': 'user-7', + }, + useUserContext: true, + }); + }); + + it('excludes user-bound records when rebuilding the shared client', async () => { + const manager = new DefaultMCPManager(createApp() as any) as any; + manager.listMCP = vi.fn().mockResolvedValue([]); + + await manager.rebuildClient(); + + expect(manager.listMCP).toHaveBeenCalledWith({ enabled: true, useUserContext: false }); + expect(mcpClientMock.instances).toHaveLength(0); + }); + + it('registers user-bound tools from filter ctx', async () => { + const manager = new DefaultMCPManager(createApp() as any) as any; + manager.listMCP = vi.fn().mockResolvedValue([ + { + name: 'profile', + enabled: true, + transport: 'http', + url: 'https://{{ $env.MCP_HOST }}/{{ currentUser.id }}', + headers: {}, + useUserContext: true, + }, + ]); + const registered: any[] = []; + + await manager.getMCPToolsProvider()( + { + registerTools: (tool) => registered.push(tool), + registerDynamicTools: vi.fn(), + }, + { ctx: createCtx(9) }, + ); + + expect(registered).toHaveLength(1); + expect(registered[0].definition.name).toBe('mcp-profile-getProfile'); + expect(mcpClientMock.instances[0].connections.profile.url).toBe('https://mcp.example.test/9'); + }); + + it('lists user-bound tools from ctx', async () => { + const manager = new DefaultMCPManager(createApp() as any) as any; + manager.listMCP = vi.fn().mockResolvedValue([ + { + name: 'profile', + enabled: true, + transport: 'http', + url: 'https://{{ $env.MCP_HOST }}/{{ currentUser.id }}', + headers: {}, + useUserContext: true, + }, + ]); + + const tools = await manager.listMCPTools(createCtx(11)); + + expect(tools.profile).toEqual([ + { + name: 'mcp-profile-getProfile', + title: 'getProfile', + description: 'Get profile from profile', + serverName: 'profile', + permission: 'ALLOW', + }, + ]); + expect(mcpClientMock.instances[0].connections.profile.url).toBe('https://mcp.example.test/11'); + }); + + it('reuses cached user-bound tools and refreshes them after TTL', async () => { + vi.useFakeTimers(); + vi.setSystemTime(0); + + const manager = new UserContextMCPClientManager({ + app: createApp(), + ttlMs: 10, + maxSize: 10, + listEntries: async () => [ + { + name: 'profile', + enabled: true, + transport: 'http', + url: 'https://{{ $env.MCP_HOST }}/{{ currentUser.id }}', + headers: {}, + useUserContext: true, + }, + ], + buildConnection: (options) => options as any, + }); + + await manager.getToolsMap(createCtx(1)); + await manager.getToolsMap(createCtx(1)); + expect(mcpClientMock.instances).toHaveLength(1); + + vi.setSystemTime(11); + await manager.getToolsMap(createCtx(1)); + + expect(mcpClientMock.instances).toHaveLength(2); + expect(mcpClientMock.instances[0].close).toHaveBeenCalledTimes(1); + }); + + it('evicts the oldest user-bound cache entry when max size is exceeded', async () => { + const manager = new UserContextMCPClientManager({ + app: createApp(), + ttlMs: 1000, + maxSize: 1, + listEntries: async () => [ + { + name: 'profile', + enabled: true, + transport: 'http', + url: 'https://{{ $env.MCP_HOST }}/{{ currentUser.id }}', + headers: {}, + useUserContext: true, + }, + ], + buildConnection: (options) => options as any, + }); + + await manager.getToolsMap(createCtx(1)); + await manager.getToolsMap(createCtx(2)); + + expect(mcpClientMock.instances).toHaveLength(2); + expect(mcpClientMock.instances[0].close).toHaveBeenCalledTimes(1); + + await manager.clear(); + expect(mcpClientMock.instances[1].close).toHaveBeenCalledTimes(1); + }); +}); diff --git a/packages/core/ai/src/__tests__/mcp.test.ts b/packages/core/ai/src/__tests__/mcp.test.ts index d84da030598..32100eb62c2 100644 --- a/packages/core/ai/src/__tests__/mcp.test.ts +++ b/packages/core/ai/src/__tests__/mcp.test.ts @@ -85,7 +85,7 @@ describe('MCP loader test cases', () => { title: 'getForecast', description: 'Get weather forecast', serverName: 'weather', - permission: 'ASK', + permission: 'ALLOW', }, { name: 'mcp-weather-setDefaultCity', diff --git a/packages/core/ai/src/mcp-manager/index.ts b/packages/core/ai/src/mcp-manager/index.ts index eda38e6a9e6..034e8c3f638 100644 --- a/packages/core/ai/src/mcp-manager/index.ts +++ b/packages/core/ai/src/mcp-manager/index.ts @@ -15,6 +15,8 @@ import { StructuredToolInterface } from '@langchain/core/tools'; import { MCPEntry, MCPFilter, MCPManager, MCPOptions, MCPTestResult, MCPToolEntry } from './types'; import type { DynamicToolsProvider, Permission, ToolsRegistration, ToolsOptions } from '../tools-manager/types'; import type { Context } from '@nocobase/actions'; +import { normalizeMCPOptions, renderMCPOptions } from './options-renderer'; +import { UserContextMCPClientManager } from './user-context-client-manager'; export class DefaultMCPManager implements MCPManager { private readonly mcpRegistry = new Registry(); @@ -23,9 +25,15 @@ export class DefaultMCPManager implements MCPManager { private client: MultiServerMCPClient | null = null; private toolsMap: Record = {}; private toolsPermissionMap: Record = {}; + private readonly userContextClientManager: UserContextMCPClientManager; constructor(private readonly app: any) { this.provideCollectionManager = () => app.mainDataSource; + this.userContextClientManager = new UserContextMCPClientManager({ + app, + listEntries: () => this.listMCP({ enabled: true, useUserContext: true }), + buildConnection: (options) => this.buildMCPConnection(options), + }); } async init() { @@ -72,6 +80,14 @@ export class DefaultMCPManager implements MCPManager { if (filter.transport) { where['transport'] = filter.transport; } + if (filter.useUserContext != null) { + where['useUserContext'] = + filter.useUserContext === true + ? true + : { + [Op.or]: [false, null], + }; + } return (await this.aiMcpClientsModel?.findAll({ where }))?.map((item) => item.toJSON() as MCPEntry) ?? []; } @@ -88,7 +104,7 @@ export class DefaultMCPManager implements MCPManager { } // Get all enabled MCP entries - const entries = await this.listMCP({ enabled: true }); + const entries = await this.listMCP({ enabled: true, useUserContext: false }); if (entries.length === 0) { return; @@ -97,7 +113,7 @@ export class DefaultMCPManager implements MCPManager { // Build connections object const connections: Record = {}; for (const entry of entries) { - connections[entry.name] = this.buildMCPConnection(entry); + connections[entry.name] = this.buildMCPConnection(await renderMCPOptions(entry, this.app)); } // Create new client and initialize connections @@ -121,48 +137,82 @@ export class DefaultMCPManager implements MCPManager { } getMCPToolsProvider(): DynamicToolsProvider { - return async (register: ToolsRegistration): Promise => { + return async (register: ToolsRegistration, filter): Promise => { // Use cached tools from rebuildClient for (const [serverName, tools] of Object.entries(this.toolsMap)) { - for (const tool of tools) { - const toolName = `mcp-${serverName}-${tool.name}`; - const toolOptions: ToolsOptions = { - scope: 'GENERAL', - from: 'mcp', - defaultPermission: this.toolsPermissionMap[toolName], - introduction: { - title: tool.name, - about: tool.description, - }, - definition: { - name: toolName, - description: tool.description || `MCP tool: ${tool.name} from ${serverName}`, - schema: tool.schema, - }, - invoke: async (_ctx: Context, args: any) => { - try { - const result = await tool.invoke(args); - return result; - } catch (error: any) { - return { - status: 'error' as const, - content: error?.message || 'Tool invocation failed', - }; - } - }, - }; - register.registerTools(toolOptions); - } + this.registerToolsFromMap(register, serverName, tools); + } + + if (!filter?.ctx) { + return; + } + + const userToolsMap = await this.userContextClientManager.getToolsMap(filter.ctx); + for (const [serverName, tools] of Object.entries(userToolsMap)) { + this.registerToolsFromMap(register, serverName, tools); } }; } - async listMCPTools(): Promise> { + private registerToolsFromMap( + register: ToolsRegistration, + serverName: string, + tools: StructuredToolInterface[], + ): void { + for (const tool of tools) { + const toolName = `mcp-${serverName}-${tool.name}`; + this.ensureToolPermission(toolName, tool.name); + const toolOptions: ToolsOptions = { + scope: 'GENERAL', + from: 'mcp', + defaultPermission: this.toolsPermissionMap[toolName], + introduction: { + title: tool.name, + about: tool.description, + }, + definition: { + name: toolName, + description: tool.description || `MCP tool: ${tool.name} from ${serverName}`, + schema: tool.schema, + }, + invoke: async (_ctx: Context, args: any) => { + try { + const result = await tool.invoke(args); + return result; + } catch (error: any) { + return { + status: 'error' as const, + content: error?.message || 'Tool invocation failed', + }; + } + }, + }; + register.registerTools(toolOptions); + } + } + + private ensureToolPermission(toolName: string, rawToolName: string) { + if (!(toolName in this.toolsPermissionMap)) { + this.toolsPermissionMap[toolName] = rawToolName.startsWith('get') ? 'ALLOW' : 'ASK'; + } + } + + async listMCPTools(ctx?: Context): Promise> { + const toolsMap = { + ...this.toolsMap, + ...(ctx ? await this.userContextClientManager.getToolsMap(ctx) : {}), + }; + + return this.formatMCPTools(toolsMap); + } + + private formatMCPTools(toolsMap: Record): Record { return Object.fromEntries( - Object.entries(this.toolsMap).map(([serverName, tools]) => [ + Object.entries(toolsMap).map(([serverName, tools]) => [ serverName, tools.map((tool) => { const toolName = `mcp-${serverName}-${tool.name}`; + this.ensureToolPermission(toolName, tool.name); return { name: toolName, title: tool.name, @@ -179,8 +229,13 @@ export class DefaultMCPManager implements MCPManager { this.toolsPermissionMap[toolName] = permission; } - async testConnection(options: MCPOptions): Promise { - const { transport } = options; + async clearUserContextCache(): Promise { + await this.userContextClientManager.clear(); + } + + async testConnection(options: MCPOptions, ctx?: Context): Promise { + const renderedOptions = await renderMCPOptions(normalizeMCPOptions(options), this.app, ctx); + const { transport } = renderedOptions; // Validate required fields if (!transport) { @@ -190,14 +245,14 @@ export class DefaultMCPManager implements MCPManager { }; } - if (transport === 'stdio' && !options.command) { + if (transport === 'stdio' && !renderedOptions.command) { return { success: false, error: 'Command is required for stdio transport', }; } - if ((transport === 'http' || transport === 'sse') && !options.url) { + if ((transport === 'http' || transport === 'sse') && !renderedOptions.url) { return { success: false, error: 'URL is required for HTTP/SSE transport', @@ -207,7 +262,7 @@ export class DefaultMCPManager implements MCPManager { let client: MultiServerMCPClient | null = null; try { - const connection = this.buildMCPConnection(options); + const connection = this.buildMCPConnection(renderedOptions); const serverName = 'test-server'; client = new MultiServerMCPClient({ @@ -303,18 +358,20 @@ export class DefaultMCPManager implements MCPManager { } private async persistenceEntry(entry: MCPEntry): Promise { + const normalizedEntry = normalizeMCPOptions(entry) as MCPEntry; await this.sequelize.transaction(async (transaction) => { - const existed = await this.aiMcpClientsModel.findOne({ where: { name: entry.name }, transaction }); + const existed = await this.aiMcpClientsModel.findOne({ where: { name: normalizedEntry.name }, transaction }); if (existed) { await existed.update( { - transport: entry.transport, - command: entry.command, - args: entry.args, - env: entry.env, - url: entry.url, - headers: entry.headers, - restart: entry.restart, + transport: normalizedEntry.transport, + command: normalizedEntry.command, + args: normalizedEntry.args, + env: normalizedEntry.env, + url: normalizedEntry.url, + headers: normalizedEntry.headers, + restart: normalizedEntry.restart, + useUserContext: normalizedEntry.useUserContext, }, { transaction }, ); @@ -323,7 +380,7 @@ export class DefaultMCPManager implements MCPManager { await this.aiMcpClientsModel.create( { - ...entry, + ...normalizedEntry, }, { transaction }, ); @@ -331,13 +388,15 @@ export class DefaultMCPManager implements MCPManager { } private normalizeEntry(name: string, options: MCPOptions): MCPEntry { - return { + const entry: MCPEntry = { name, enabled: true, ...options, args: options.args ?? [], env: options.env ?? {}, + useUserContext: options.useUserContext === true, }; + return normalizeMCPOptions(entry) as MCPEntry; } private get aiMcpClientsCollection() { diff --git a/packages/core/ai/src/mcp-manager/options-renderer.ts b/packages/core/ai/src/mcp-manager/options-renderer.ts new file mode 100644 index 00000000000..1abeeccaea1 --- /dev/null +++ b/packages/core/ai/src/mcp-manager/options-renderer.ts @@ -0,0 +1,119 @@ +/** + * 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 => { + if (!value || typeof value !== 'object' || Array.isArray(value)) { + return {}; + } + return Object.entries(value as Record).reduce>((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)); +}; + +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 } }, + ctx?: Context, +): Promise { + const currentUser = options.useUserContext ? await getCurrentUser(ctx, options) : undefined; + const variables = { + $env: app.environment?.getVariables?.() ?? {}, + currentUser, + $user: currentUser, + ctx: { + currentUser, + }, + }; + + return normalizeMCPOptions(parse(options)(variables) as MCPOptions); +} diff --git a/packages/core/ai/src/mcp-manager/types.ts b/packages/core/ai/src/mcp-manager/types.ts index ea432b0bef8..1c784a579b6 100644 --- a/packages/core/ai/src/mcp-manager/types.ts +++ b/packages/core/ai/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; getMCP(name: string): Promise; listMCP(filter: MCPFilter): Promise; - testConnection(options: MCPOptions): Promise; + testConnection(options: MCPOptions, ctx?: Context): Promise; rebuildClient(): Promise; getClient(): MultiServerMCPClient | null; getMCPToolsProvider(): DynamicToolsProvider; - listMCPTools(): Promise>; + listMCPTools(ctx?: Context): Promise>; updateMCPToolPermission(toolName: string, permission: Permission): Promise; + clearUserContextCache(): Promise; } export interface MCPRegistration { @@ -34,6 +36,7 @@ export type MCPOptions = { url?: string; headers?: Record; restart?: Record; + 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'; diff --git a/packages/core/ai/src/mcp-manager/user-context-client-manager.ts b/packages/core/ai/src/mcp-manager/user-context-client-manager.ts new file mode 100644 index 00000000000..c800601ce6b --- /dev/null +++ b/packages/core/ai/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; + expiresAt: number; + lastAccessedAt: number; +}; + +export type UserContextMCPClientManagerOptions = { + app: any; + listEntries: () => Promise; + buildConnection: (options: MCPOptions) => MCPConnection; + ttlMs?: number; + maxSize?: number; +}; + +export class UserContextMCPClientManager { + private readonly cache = new Map(); + 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> { + 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 = {}; + 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); + } + } +} diff --git a/packages/core/ai/src/tools-manager/types.ts b/packages/core/ai/src/tools-manager/types.ts index b0a5ceec676..b02576bfed9 100644 --- a/packages/core/ai/src/tools-manager/types.ts +++ b/packages/core/ai/src/tools-manager/types.ts @@ -56,4 +56,5 @@ export type ToolsFilter = { defaultPermission?: Permission; silence?: boolean; sessionId?: string; + ctx?: Context; }; diff --git a/packages/plugins/@nocobase/plugin-ai/src/client/ai-employees/admin/mcp/MCPSettings.tsx b/packages/plugins/@nocobase/plugin-ai/src/client/ai-employees/admin/mcp/MCPSettings.tsx index bcd14496849..7dd385d4d69 100644 --- a/packages/plugins/@nocobase/plugin-ai/src/client/ai-employees/admin/mcp/MCPSettings.tsx +++ b/packages/plugins/@nocobase/plugin-ai/src/client/ai-employees/admin/mcp/MCPSettings.tsx @@ -19,10 +19,14 @@ import { useCollectionRecordData, useDataBlockRequest, useDestroyActionProps, + useGlobalVariable, + useCurrentUserVariable, + Variable, + Checkbox as NocoCheckbox, } from '@nocobase/client'; import { css } from '@emotion/css'; import { CheckCircleOutlined, CloseCircleOutlined, PlusOutlined } from '@ant-design/icons'; -import { App, Button, Space, Switch, Tag, Alert, Spin } from 'antd'; +import { App, Button, Space, Switch, Tag, Alert, Spin, Tooltip } from 'antd'; import { createForm } from '@formily/core'; import { observer, useForm } from '@formily/react'; import React, { useContext, useEffect, useMemo, useRef, useState, useCallback, createContext } from 'react'; @@ -77,6 +81,7 @@ const sanitizeMCPValues = (values: Record) => { const transport = values.transport as MCPTransport; const nextValues = { ...values, + useUserContext: transport === 'stdio' ? false : values.useUserContext === true, args: typeof values.args === 'string' ? values.args @@ -114,6 +119,9 @@ const useCreateFormProps = () => { () => ({ enabled: true, transport: 'stdio', + useUserContext: false, + command: '', + url: '', args: '', env: [], headers: [], @@ -152,6 +160,7 @@ interface MCPRecord { env?: Record; headers?: Record; restart?: Record; + useUserContext?: boolean; } const useEditFormProps = () => { @@ -160,6 +169,9 @@ const useEditFormProps = () => { const initialValues = useMemo( () => ({ ...record, + useUserContext: record?.useUserContext === true, + command: record?.command ?? '', + url: record?.url ?? '', args: Array.isArray(record?.args) ? record.args.join(' ') : '', env: mapObjectToEntries(record?.env), headers: mapObjectToEntries(record?.headers), @@ -250,6 +262,21 @@ const useEnsureConnectionBeforeSubmit = () => { }; }; +const useConfirmSaveAfterFailedTest = () => { + const { modal } = App.useApp(); + const t = useT(); + + return async () => + modal.confirm({ + title: t('Connection test failed'), + content: t( + 'The MCP server uses current user variables and the connection test failed. Do you want to save it anyway?', + ), + okText: t('Save anyway'), + cancelText: t('Cancel'), + }); +}; + const useCreateActionProps = () => { const { setVisible } = useActionContext(); const { message } = App.useApp(); @@ -260,6 +287,7 @@ const useCreateActionProps = () => { const { loading: testLoading } = useContext(TestConnectionContext); const { rebuildClient, rebuilding } = useContext(MCPSettingsContext); const ensureConnectionBeforeSubmit = useEnsureConnectionBeforeSubmit(); + const confirmSaveAfterFailedTest = useConfirmSaveAfterFailedTest(); return { type: 'primary', @@ -268,7 +296,13 @@ const useCreateActionProps = () => { await form.submit(); const passed = await ensureConnectionBeforeSubmit(form.values); if (!passed) { - return; + if (sanitizeMCPValues(form.values).useUserContext !== true) { + return; + } + const confirmed = await confirmSaveAfterFailedTest(); + if (!confirmed) { + return; + } } await api.resource('aiMcpClients').create({ values: sanitizeMCPValues(form.values), @@ -293,6 +327,7 @@ const useEditActionProps = () => { const { loading: testLoading } = useContext(TestConnectionContext); const { rebuildClient, rebuilding } = useContext(MCPSettingsContext); const ensureConnectionBeforeSubmit = useEnsureConnectionBeforeSubmit(); + const confirmSaveAfterFailedTest = useConfirmSaveAfterFailedTest(); return { type: 'primary', @@ -301,7 +336,13 @@ const useEditActionProps = () => { await form.submit(); const passed = await ensureConnectionBeforeSubmit(form.values); if (!passed) { - return; + if (sanitizeMCPValues(form.values).useUserContext !== true) { + return; + } + const confirmed = await confirmSaveAfterFailedTest(); + if (!confirmed) { + return; + } } await api.resource('aiMcpClients').update({ values: sanitizeMCPValues(form.values), @@ -422,6 +463,103 @@ const TestConnectionResult: React.FC = observer( { displayName: 'TestConnectionResult' }, ); +const useMCPVariableScope = (includeCurrentUser: boolean) => { + const environmentVariables = useGlobalVariable('$env'); + const { currentUserSettings } = useCurrentUserVariable({ + maxDepth: 3, + noDisabled: true, + }); + + return useMemo( + () => + [ + normalizeMCPVariableOption(environmentVariables), + includeCurrentUser ? normalizeMCPVariableOption(currentUserSettings) : null, + ].filter(Boolean), + [currentUserSettings, environmentVariables, includeCurrentUser], + ); +}; + +type MCPVariableOption = { + name?: React.Key; + title?: React.ReactNode; + value?: React.Key; + label?: React.ReactNode; + key?: React.Key; + children?: MCPVariableOption[]; + loadChildren?: (option: MCPVariableOption, ...args: unknown[]) => Promise | void; + [key: string]: unknown; +}; + +const normalizeMCPVariableOption = (option?: MCPVariableOption | null): MCPVariableOption | null => { + if (!option) { + return null; + } + + const value = option.value ?? option.name; + const label = option.label ?? option.title; + const loadChildren = option.loadChildren + ? async (target: MCPVariableOption, ...args: unknown[]) => { + await option.loadChildren?.(target, ...args); + target.children = target.children?.map((child) => normalizeMCPVariableOption(child)).filter(Boolean); + } + : undefined; + + return { + ...option, + name: option.name ?? value, + title: option.title ?? label, + value, + label, + key: option.key ?? value, + children: option.children?.map((child) => normalizeMCPVariableOption(child)).filter(Boolean), + ...(loadChildren ? { loadChildren } : {}), + }; +}; + +const mcpVariableFieldNames = { + value: 'name', + label: 'title', +}; + +type MCPVariableInputProps = React.ComponentProps & { + variableScope?: 'env' | 'user'; +}; + +const MCPVariableInput: React.FC = observer( + (props) => { + const { variableScope = 'env', ...others } = props; + const form = useForm(); + const includeCurrentUser = variableScope === 'user' && form.values.useUserContext === true; + const scope = useMCPVariableScope(includeCurrentUser); + + return ; + }, + { displayName: 'MCPVariableInput' }, +); + +type UserContextCheckboxProps = React.ComponentProps & { + tooltip?: React.ReactNode; +}; + +const UserContextCheckbox: React.FC = observer( + (props) => { + const { tooltip, ...others } = props; + const form = useForm(); + const disabled = form.values.transport === 'stdio'; + + useEffect(() => { + if (disabled && form.values.useUserContext) { + form.setValuesIn('useUserContext', false); + } + }, [disabled, form]); + + const checkbox = ; + return tooltip ? {checkbox} : checkbox; + }, + { displayName: 'UserContextCheckbox' }, +); + const AddNew = () => { const t = useT(); const [visible, setVisible] = useState(false); @@ -453,7 +591,7 @@ const AddNew = () => { { { MCPToolsList, TestConnectionButton, TestConnectionResult, + MCPVariableInput, + UserContextCheckbox, TransportTag, EnabledSwitch, }} diff --git a/packages/plugins/@nocobase/plugin-ai/src/client/ai-employees/admin/mcp/schemas.ts b/packages/plugins/@nocobase/plugin-ai/src/client/ai-employees/admin/mcp/schemas.ts index 4a9864993be..35f810b08d8 100644 --- a/packages/plugins/@nocobase/plugin-ai/src/client/ai-employees/admin/mcp/schemas.ts +++ b/packages/plugins/@nocobase/plugin-ai/src/client/ai-employees/admin/mcp/schemas.ts @@ -69,11 +69,21 @@ const createMCPFormProperties = (options: { enum: '{{ transportOptions }}', required: true, }, + useUserContext: { + type: 'boolean', + 'x-decorator': 'FormItem', + title: '{{ t("Depends on current user") }}', + 'x-component': 'UserContextCheckbox', + 'x-component-props': { + tooltip: + '{{ t("When enabled, URL and headers can use current user variables, and the MCP server runs per current user. Stdio transport is not supported.") }}', + }, + }, command: { type: 'string', 'x-decorator': 'FormItem', title: '{{ t("Command") }}', - 'x-component': 'Input', + 'x-component': 'MCPVariableInput', 'x-component-props': { placeholder: '{{ t("For example: npx, uvx, node") }}', }, @@ -83,7 +93,7 @@ const createMCPFormProperties = (options: { type: 'string', 'x-decorator': 'FormItem', title: '{{ t("Arguments") }}', - 'x-component': 'Input', + 'x-component': 'MCPVariableInput', 'x-component-props': { placeholder: '{{ t("Space-separated args, e.g.: -u --flag value") }}', }, @@ -120,10 +130,7 @@ const createMCPFormProperties = (options: { value: { type: 'string', 'x-decorator': 'FormItem', - 'x-component': 'TextAreaWithGlobalScope', - 'x-component-props': { - useTypedConstant: true, - }, + 'x-component': 'MCPVariableInput', }, remove: { type: 'void', @@ -147,9 +154,10 @@ const createMCPFormProperties = (options: { type: 'string', 'x-decorator': 'FormItem', title: '{{ t("URL") }}', - 'x-component': 'Input', + 'x-component': 'MCPVariableInput', 'x-component-props': { placeholder: '{{ t("For example: https://example.com/mcp") }}', + variableScope: 'user', }, 'x-reactions': createTransportReaction('remote', true), }, @@ -184,9 +192,9 @@ const createMCPFormProperties = (options: { value: { type: 'string', 'x-decorator': 'FormItem', - 'x-component': 'TextAreaWithGlobalScope', + 'x-component': 'MCPVariableInput', 'x-component-props': { - useTypedConstant: true, + variableScope: 'user', }, }, remove: { @@ -417,6 +425,18 @@ export const mcpSettingsSchema = { }, }, column5: { + type: 'void', + title: '{{ t("Depends on current user") }}', + 'x-component': 'TableV2.Column', + properties: { + useUserContext: { + type: 'boolean', + 'x-component': 'Checkbox', + 'x-read-pretty': true, + }, + }, + }, + column6: { type: 'void', title: '{{ t("Enabled") }}', 'x-component': 'TableV2.Column', @@ -427,7 +447,7 @@ export const mcpSettingsSchema = { }, }, }, - column6: { + column7: { type: 'void', title: '{{ t("Actions") }}', 'x-decorator': 'TableV2.Column.ActionBar', diff --git a/packages/plugins/@nocobase/plugin-ai/src/collections/ai-mcp-clients.ts b/packages/plugins/@nocobase/plugin-ai/src/collections/ai-mcp-clients.ts index f9ca22a835b..7015c16b604 100644 --- a/packages/plugins/@nocobase/plugin-ai/src/collections/ai-mcp-clients.ts +++ b/packages/plugins/@nocobase/plugin-ai/src/collections/ai-mcp-clients.ts @@ -47,6 +47,11 @@ export default { 'x-component': 'Select', }, }, + { + name: 'useUserContext', + type: 'boolean', + defaultValue: false, + }, { name: 'command', type: 'string', diff --git a/packages/plugins/@nocobase/plugin-ai/src/locale/en-US.json b/packages/plugins/@nocobase/plugin-ai/src/locale/en-US.json index 0f75d82f2f4..4ea4111c39b 100644 --- a/packages/plugins/@nocobase/plugin-ai/src/locale/en-US.json +++ b/packages/plugins/@nocobase/plugin-ai/src/locale/en-US.json @@ -449,7 +449,10 @@ "When enabled, this AI employee will always use the selected model and users cannot switch to other models in chat.": "When enabled, this AI employee will always use the selected model and users cannot switch to other models in chat.", "When enabled, this AI employee can only use the selected models, and users can only switch between these models in chat.": "When enabled, this AI employee can only use the selected models, and users can only switch between these models in chat.", "Restrict this AI employee to the selected models.": "Restrict this AI employee to the selected models.", - "Select model": "Select model", - "Select models": "Select models", - "Models": "Models" + "Models": "Models", + "Depends on current user": "Depends on current user", + "When enabled, URL and headers can use current user variables, and the MCP server runs per current user. Stdio transport is not supported.": "When enabled, URL and headers can use current user variables, and the MCP server runs per current user. Stdio transport is not supported.", + "Connection test failed": "Connection test failed", + "The MCP server uses current user variables and the connection test failed. Do you want to save it anyway?": "The MCP server uses current user variables and the connection test failed. Do you want to save it anyway?", + "Save anyway": "Save anyway" } diff --git a/packages/plugins/@nocobase/plugin-ai/src/locale/zh-CN.json b/packages/plugins/@nocobase/plugin-ai/src/locale/zh-CN.json index ded058d4a8c..b62b5f6867c 100644 --- a/packages/plugins/@nocobase/plugin-ai/src/locale/zh-CN.json +++ b/packages/plugins/@nocobase/plugin-ai/src/locale/zh-CN.json @@ -455,7 +455,10 @@ "When enabled, this AI employee will always use the selected model and users cannot switch to other models in chat.": "启用后,该 AI 员工将始终使用所选模型,用户在对话中不能切换到其他模型。", "When enabled, this AI employee can only use the selected models, and users can only switch between these models in chat.": "启用后,该 AI 员工只能使用这里选择的模型,用户在对话中也只能在这些模型之间切换。", "Restrict this AI employee to the selected models.": "限制该 AI 员工只能使用所选模型。", - "Select model": "选择模型", - "Select models": "选择模型", - "Models": "模型" + "Models": "模型", + "Depends on current user": "依赖当前用户信息", + "When enabled, URL and headers can use current user variables, and the MCP server runs per current user. Stdio transport is not supported.": "启用后,可在 URL 和请求头中使用当前用户变量,并按当前用户运行 MCP 服务。Stdio 传输方式不支持。", + "Connection test failed": "连接测试失败", + "The MCP server uses current user variables and the connection test failed. Do you want to save it anyway?": "该 MCP 服务使用当前用户变量,且连接测试失败。是否仍然保存?", + "Save anyway": "仍然保存" } diff --git a/packages/plugins/@nocobase/plugin-ai/src/server/__tests__/mcp-client-events.test.ts b/packages/plugins/@nocobase/plugin-ai/src/server/__tests__/mcp-client-events.test.ts new file mode 100644 index 00000000000..5ee8dd4e57c --- /dev/null +++ b/packages/plugins/@nocobase/plugin-ai/src/server/__tests__/mcp-client-events.test.ts @@ -0,0 +1,154 @@ +/** + * 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 { createMockServer, MockServer } from '@nocobase/test'; +import { afterAll, beforeAll, beforeEach, describe, expect, it, vi } from 'vitest'; +import PluginAIServer from '../plugin'; +import { aiMcpClients } from '../resource/aiMcpClients'; + +describe('MCP client database events', () => { + let app: MockServer; + let plugin: PluginAIServer; + let clearUserContextCache: ReturnType; + + beforeAll(async () => { + app = await createMockServer({ + plugins: ['nocobase'], + }); + await app.pm.enable('ai'); + plugin = app.pm.get('ai') as PluginAIServer; + clearUserContextCache = vi.fn().mockResolvedValue(undefined); + plugin.ai.mcpManager.clearUserContextCache = clearUserContextCache; + }); + + beforeEach(() => { + clearUserContextCache.mockReset(); + clearUserContextCache.mockResolvedValue(undefined); + }); + + afterAll(async () => { + await app.db.clean({ drop: true }); + await app.destroy(); + }); + + it('clears user-bound cache after transaction commit only for user-bound records', async () => { + const repository = app.db.getRepository('aiMcpClients'); + + const sharedTransaction = await app.db.sequelize.transaction(); + await repository.create({ + values: { + name: 'shared', + transport: 'http', + url: 'http://localhost:3000/mcp', + useUserContext: false, + }, + transaction: sharedTransaction, + }); + await sharedTransaction.commit(); + expect(clearUserContextCache).toHaveBeenCalledTimes(0); + + const boundTransaction = await app.db.sequelize.transaction(); + await repository.create({ + values: { + name: 'bound', + transport: 'http', + url: 'http://localhost:3000/mcp', + useUserContext: true, + }, + transaction: boundTransaction, + }); + expect(clearUserContextCache).toHaveBeenCalledTimes(0); + + await boundTransaction.commit(); + expect(clearUserContextCache).toHaveBeenCalledTimes(1); + }); + + it('clears user-bound cache after commit when a user-bound record is disabled', async () => { + const repository = app.db.getRepository('aiMcpClients'); + await repository.create({ + values: { + name: 'bound-to-disable', + transport: 'http', + url: 'http://localhost:3000/mcp', + useUserContext: true, + }, + }); + clearUserContextCache.mockClear(); + + const transaction = await app.db.sequelize.transaction(); + await repository.update({ + filterByTk: 'bound-to-disable', + values: { + useUserContext: false, + }, + transaction, + }); + expect(clearUserContextCache).toHaveBeenCalledTimes(0); + + await transaction.commit(); + expect(clearUserContextCache).toHaveBeenCalledTimes(1); + }); + + it('does not block transaction commit while clearing user-bound cache', async () => { + const repository = app.db.getRepository('aiMcpClients'); + let resolveClearCache: () => void; + const clearCachePromise = new Promise((resolve) => { + resolveClearCache = resolve; + }); + clearUserContextCache.mockReturnValue(clearCachePromise); + + const transaction = await app.db.sequelize.transaction(); + await repository.create({ + values: { + name: 'bound-for-fire-and-go', + transport: 'http', + url: 'http://localhost:3000/mcp', + useUserContext: true, + }, + transaction, + }); + + await expect(transaction.commit()).resolves.toBeUndefined(); + expect(clearUserContextCache).toHaveBeenCalledTimes(1); + + resolveClearCache(); + await clearCachePromise; + }); +}); + +describe('aiMcpClients resource actions', () => { + it('passes request ctx when listing MCP tools', async () => { + const tools = { + profile: [], + }; + const listMCPTools = vi.fn().mockResolvedValue(tools); + const ctx = { + app: { + pm: { + get: vi.fn(() => ({ + ai: { + mcpManager: { + listMCPTools, + }, + }, + })), + }, + }, + body: undefined, + }; + const next = vi.fn().mockResolvedValue(undefined); + const action = aiMcpClients.actions?.listTools as (ctx: typeof ctx, next: typeof next) => Promise; + + await action(ctx, next); + + expect(listMCPTools).toHaveBeenCalledWith(ctx); + expect(ctx.body).toBe(tools); + expect(next).toHaveBeenCalledTimes(1); + }); +}); diff --git a/packages/plugins/@nocobase/plugin-ai/src/server/ai-employees/ai-employee.ts b/packages/plugins/@nocobase/plugin-ai/src/server/ai-employees/ai-employee.ts index 12937356153..4639bd75703 100644 --- a/packages/plugins/@nocobase/plugin-ai/src/server/ai-employees/ai-employee.ts +++ b/packages/plugins/@nocobase/plugin-ai/src/server/ai-employees/ai-employee.ts @@ -1353,7 +1353,7 @@ If information is missing, clearly state it in the summary.`; } const tools: ToolsEntry[] = await this.listTools({ scope: 'GENERAL' }); if (this.webSearch === true) { - const subAgentWebSearch = await this.toolsManager.getTools(SYSTEM_TOOLS.WEB_SEARCH); + const subAgentWebSearch = await this.toolsManager.getTools(SYSTEM_TOOLS.WEB_SEARCH, { ctx: this.ctx }); tools.push(subAgentWebSearch); } const generalToolsNameSet = new Set(tools.map((x) => x.definition.name)); @@ -1361,7 +1361,9 @@ If information is missing, clearly state it in the summary.`; const settingsTools = this.employee.skillSettings?.tools ?? []; const employeeTools = [...settingsTools, ...this.tools]; if (await this.plugin.knowledgeBaseManager.isEnabledKnowledgeBase(this.employee.toJSON() as AIEmployeeType)) { - const knowledgeBaseRetrieveTool = await this.toolsManager.getTools(SYSTEM_TOOLS.KNOWLEDGE_BASE); + const knowledgeBaseRetrieveTool = await this.toolsManager.getTools(SYSTEM_TOOLS.KNOWLEDGE_BASE, { + ctx: this.ctx, + }); if (knowledgeBaseRetrieveTool) { employeeTools.push({ name: SYSTEM_TOOLS.KNOWLEDGE_BASE }); } @@ -1564,7 +1566,10 @@ If information is missing, clearly state it in the summary.`; } private listTools(filter?: ToolsFilter) { - return this.toolsManager.listTools(filter); + return this.toolsManager.listTools({ + ...filter, + ctx: this.ctx, + }); } private withRunMetadata(config?: any) { diff --git a/packages/plugins/@nocobase/plugin-ai/src/server/plugin.ts b/packages/plugins/@nocobase/plugin-ai/src/server/plugin.ts index b24dd467653..2c904047c60 100644 --- a/packages/plugins/@nocobase/plugin-ai/src/server/plugin.ts +++ b/packages/plugins/@nocobase/plugin-ai/src/server/plugin.ts @@ -27,7 +27,7 @@ import * as aiEmployeeActions from './resource/aiEmployees'; import { googleGenAIProviderOptions } from './llm-providers/google-genai'; import { AIEmployeeTrigger } from './workflow/triggers/ai-employee'; import { getWorkflowCallers, createDocsSearchTool, type DocsFsCache } from './tools'; -import { Model } from '@nocobase/database'; +import { Model, Transaction } from '@nocobase/database'; import { anthropicProviderOptions } from './llm-providers/anthropic'; import aiSettings from './resource/aiSettings'; import { dashscopeProviderOptions } from './llm-providers/dashscope'; @@ -54,6 +54,11 @@ import { import { KnowledgeBaseManager } from './ai-employees/ai-knowledge-base'; import { LLMStreamCachedManager } from './manager/llm-stream-manager'; +type MCPClientModel = Model<{ useUserContext?: boolean }>; +type TransactionOptions = { + transaction?: Transaction; +}; + export class PluginAIServer extends Plugin { features = new AIPluginFeatureManagerImpl(); aiManager = new AIManager(this); @@ -152,6 +157,7 @@ export class PluginAIServer extends Plugin { this.registerLLMProviders(); this.registerTools(); this.defineResources(); + this.registerMcpClientEvents(); this.setPermissions(); this.registerWorkflow(); this.registerWorkContextResolveStrategy(); @@ -214,6 +220,41 @@ export class PluginAIServer extends Plugin { }); } + registerMcpClientEvents() { + this.db.on('aiMcpClients.afterCreate', (model: MCPClientModel, { transaction }: TransactionOptions = {}) => { + if (model.get('useUserContext') === true) { + this.clearMcpUserContextCacheAfterCommit(transaction); + } + }); + + this.db.on('aiMcpClients.afterUpdate', (model: MCPClientModel, { transaction }: TransactionOptions = {}) => { + if (model.get('useUserContext') === true || model.previous('useUserContext') === true) { + this.clearMcpUserContextCacheAfterCommit(transaction); + } + }); + + this.db.on('aiMcpClients.afterDestroy', (model: MCPClientModel, { transaction }: TransactionOptions = {}) => { + if (model.get('useUserContext') === true) { + this.clearMcpUserContextCacheAfterCommit(transaction); + } + }); + } + + private clearMcpUserContextCacheAfterCommit(transaction?: Transaction) { + const clearCache = () => { + this.ai.mcpManager.clearUserContextCache().catch((error) => { + this.app.log.warn('fail to clear user-bound mcp client cache', error); + }); + }; + + if (transaction) { + transaction.afterCommit(clearCache); + return; + } + + clearCache(); + } + setPermissions() { this.app.acl.registerSnippet({ name: `pm.${this.name}.llm-services`, diff --git a/packages/plugins/@nocobase/plugin-ai/src/server/resource/aiConversations.ts b/packages/plugins/@nocobase/plugin-ai/src/server/resource/aiConversations.ts index b3efdabd0ef..af2ed03f123 100644 --- a/packages/plugins/@nocobase/plugin-ai/src/server/resource/aiConversations.ts +++ b/packages/plugins/@nocobase/plugin-ai/src/server/resource/aiConversations.ts @@ -707,7 +707,7 @@ export default { toolMessages.map((toolMessage: Model) => [toolMessage.toolCallId, toolMessage]), ); - const toolsList = await plugin.ai.toolsManager.listTools({ sessionId: message.sessionId }); + const toolsList = await plugin.ai.toolsManager.listTools({ sessionId: message.sessionId, ctx }); const toolsMap = new Map(toolsList.map((t) => [t.definition.name, t])); for (const toolCall of toolCalls) { diff --git a/packages/plugins/@nocobase/plugin-ai/src/server/resource/aiMcpClients.ts b/packages/plugins/@nocobase/plugin-ai/src/server/resource/aiMcpClients.ts index 8d98ec801e8..04946e6593e 100644 --- a/packages/plugins/@nocobase/plugin-ai/src/server/resource/aiMcpClients.ts +++ b/packages/plugins/@nocobase/plugin-ai/src/server/resource/aiMcpClients.ts @@ -27,7 +27,7 @@ export const aiMcpClients: ResourceOptions = { } const plugin = ctx.app.pm.get('ai'); - const result = await plugin.ai.mcpManager.testConnection(values); + const result = await plugin.ai.mcpManager.testConnection(values, ctx); ctx.body = result; await next(); @@ -40,7 +40,7 @@ export const aiMcpClients: ResourceOptions = { }, listTools: async (ctx, next) => { const plugin = ctx.app.pm.get('ai'); - ctx.body = await plugin.ai.mcpManager.listMCPTools(); + ctx.body = await plugin.ai.mcpManager.listMCPTools(ctx); await next(); }, updateToolPermission: async (ctx, next) => { diff --git a/packages/plugins/@nocobase/plugin-ai/src/server/resource/aiTools.ts b/packages/plugins/@nocobase/plugin-ai/src/server/resource/aiTools.ts index 52871d722af..e96a476ee67 100644 --- a/packages/plugins/@nocobase/plugin-ai/src/server/resource/aiTools.ts +++ b/packages/plugins/@nocobase/plugin-ai/src/server/resource/aiTools.ts @@ -17,7 +17,10 @@ export const aiTools: ResourceOptions = { list: async (ctx, next) => { const { toolsManager } = ctx.app.aiManager as AIManager; const { filter } = ctx.action.params; - const tools = await toolsManager.listTools(filter); + const tools = await toolsManager.listTools({ + ...filter, + ctx, + }); ctx.body = tools.map((t) => ({ ...t, definition: { @@ -43,7 +46,7 @@ export const aiTools: ResourceOptions = { const bindingToolNames = aiEmployee.skillSettings?.tools?.map((tool) => tool.name) ?? []; const plugin = ctx.app.pm.get('ai') as PluginAIServer; - const tools = await plugin.ai.toolsManager.listTools(); + const tools = await plugin.ai.toolsManager.listTools({ ctx }); const result = tools.filter( (tool) => (tool.scope === 'GENERAL' && tool.from === 'loader') || bindingToolNames.includes(tool.definition.name),