feat(ai): support current user variables in MCP clients (#9828)

* feat(plugin-ai): support user-bound MCP clients

* fix(ai): type normalized MCP entries
This commit is contained in:
Drol
2026-06-23 14:18:20 +08:00
committed by GitHub
parent 656338781b
commit 54f58b737c
18 changed files with 1039 additions and 82 deletions
@@ -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<string, any>;
close = vi.fn();
constructor(connections: Record<string, any>) {
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<string, unknown> = {}) =>
({
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);
});
});
+1 -1
View File
@@ -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',
+108 -49
View File
@@ -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<MCPEntry>();
@@ -23,9 +25,15 @@ export class DefaultMCPManager implements MCPManager {
private client: MultiServerMCPClient | null = null;
private toolsMap: Record<string, StructuredToolInterface[]> = {};
private toolsPermissionMap: Record<string, Permission> = {};
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<string, StdioConnection | StreamableHTTPConnection> = {};
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<void> => {
return async (register: ToolsRegistration, filter): Promise<void> => {
// 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<Record<string, MCPToolEntry[]>> {
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<Record<string, MCPToolEntry[]>> {
const toolsMap = {
...this.toolsMap,
...(ctx ? await this.userContextClientManager.getToolsMap(ctx) : {}),
};
return this.formatMCPTools(toolsMap);
}
private formatMCPTools(toolsMap: Record<string, StructuredToolInterface[]>): Record<string, MCPToolEntry[]> {
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<MCPTestResult> {
const { transport } = options;
async clearUserContextCache(): Promise<void> {
await this.userContextClientManager.clear();
}
async testConnection(options: MCPOptions, ctx?: Context): Promise<MCPTestResult> {
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<void> {
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() {
@@ -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<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));
};
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 variables = {
$env: app.environment?.getVariables?.() ?? {},
currentUser,
$user: currentUser,
ctx: {
currentUser,
},
};
return normalizeMCPOptions(parse(options)(variables) as MCPOptions);
}
+6 -2
View File
@@ -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';
@@ -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);
}
}
}
@@ -56,4 +56,5 @@ export type ToolsFilter = {
defaultPermission?: Permission;
silence?: boolean;
sessionId?: string;
ctx?: Context;
};
@@ -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<string, any>) => {
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<string, string>;
headers?: Record<string, string>;
restart?: Record<string, any>;
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> | 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<typeof Variable.TextArea> & {
variableScope?: 'env' | 'user';
};
const MCPVariableInput: React.FC<MCPVariableInputProps> = observer(
(props) => {
const { variableScope = 'env', ...others } = props;
const form = useForm();
const includeCurrentUser = variableScope === 'user' && form.values.useUserContext === true;
const scope = useMCPVariableScope(includeCurrentUser);
return <Variable.TextArea {...others} scope={scope} fieldNames={mcpVariableFieldNames} trim={false} />;
},
{ displayName: 'MCPVariableInput' },
);
type UserContextCheckboxProps = React.ComponentProps<typeof NocoCheckbox> & {
tooltip?: React.ReactNode;
};
const UserContextCheckbox: React.FC<UserContextCheckboxProps> = 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 = <NocoCheckbox {...others} disabled={disabled || props.disabled} />;
return tooltip ? <Tooltip title={tooltip}>{checkbox}</Tooltip> : checkbox;
},
{ displayName: 'UserContextCheckbox' },
);
const AddNew = () => {
const t = useT();
const [visible, setVisible] = useState(false);
@@ -453,7 +591,7 @@ const AddNew = () => {
</Button>
<TestConnectionContext.Provider value={contextValue}>
<SchemaComponent
components={{ TestConnectionButton, TestConnectionResult, Space }}
components={{ TestConnectionButton, TestConnectionResult, Space, MCPVariableInput, UserContextCheckbox }}
scope={{
t,
transportOptions,
@@ -489,7 +627,7 @@ const MCPEditDrawerContent: React.FC = () => {
<CollectionRecordProvider record={record}>
<TestConnectionContext.Provider value={contextValue}>
<SchemaComponent
components={{ TestConnectionButton, TestConnectionResult, Space }}
components={{ TestConnectionButton, TestConnectionResult, Space, MCPVariableInput, UserContextCheckbox }}
scope={{
t,
transportOptions,
@@ -637,6 +775,8 @@ export const MCPSettings: React.FC = () => {
MCPToolsList,
TestConnectionButton,
TestConnectionResult,
MCPVariableInput,
UserContextCheckbox,
TransportTag,
EnabledSwitch,
}}
@@ -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',
@@ -47,6 +47,11 @@ export default {
'x-component': 'Select',
},
},
{
name: 'useUserContext',
type: 'boolean',
defaultValue: false,
},
{
name: 'command',
type: 'string',
@@ -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"
}
@@ -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": "仍然保存"
}
@@ -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<typeof vi.fn>;
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<void>((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<void>;
await action(ctx, next);
expect(listMCPTools).toHaveBeenCalledWith(ctx);
expect(ctx.body).toBe(tools);
expect(next).toHaveBeenCalledTimes(1);
});
});
@@ -1353,7 +1353,7 @@ If information is missing, clearly state it in the summary.</Important>`;
}
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.</Important>`;
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.</Important>`;
}
private listTools(filter?: ToolsFilter) {
return this.toolsManager.listTools(filter);
return this.toolsManager.listTools({
...filter,
ctx: this.ctx,
});
}
private withRunMetadata(config?: any) {
@@ -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`,
@@ -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) {
@@ -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) => {
@@ -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),