Replace agent handoff system with planning-based router (#16003)
## Overview This PR replaces the dynamic agent handoff system with a more predictable planning-based router that decides upfront how to handle multi-agent coordination. ## Major Changes ### 🔄 Architecture Shift: Handoffs → Planning **Removed:** - `AgentHandoffEntity` and handoff tracking system - `AgentHandoffService` and `AgentHandoffExecutorService` - Dynamic agent-to-agent transfers during execution - Handoff tool generation and description templates **Added:** - `AiRouterService` with two strategies: `simple` (single agent) and `planned` (multi-agent) - `AgentPlanExecutorService` for executing multi-step plans - Plan validation (cycle detection, dependency resolution) - `UnifiedRouterResult` type with discriminated union ### 🤖 New Standard Agents Added two new specialized agents: - **Researcher Agent**: Web search, fact-finding, competitive intelligence - **Code Agent**: TypeScript function generation for serverless workflows ### 🏗️ Router Refactoring (Latest) Split router responsibilities into focused services: - `AiRouterStrategyDeciderService`: Decides simple vs planned strategy - `AiRouterPlanGeneratorService`: Generates and validates execution plans - `AiRouterService`: Coordinates between services (reduced from 426→275 lines) ### ⚙️ Configuration Improvements - Added `outputStrategy` to agent definitions (`direct` vs `synthesize`) - Removed hardcoded special cases for workflow-builder - Added `plannerModel` field to workspace entity - Increased `MAX_STEPS` from 10 to 25 for complex workflows ### 📝 Agent Prompt Refinements Significantly simplified prompts for better clarity: - Workflow Builder: 51→36 lines - Helper: 49→28 lines - Data Manipulator: Enhanced with sorting guidance ### 🔍 Enhanced Debugging - Plan reasoning and step count in data message parts - Router debug info with token usage tracking - Better logging throughout execution pipeline ## Benefits 1. **Simpler Mental Model**: Router decides upfront vs dynamic transfers 2. **Better Predictability**: Users see the plan before execution 3. **Cleaner Architecture**: SRP with focused services 4. **Configuration Over Code**: Agent behavior via config, not hardcoded logic 5. **Plan Validation**: Catches invalid dependencies and cycles ## Migration Notes - Database migration removes `agentHandoff` table - Adds `plannerModel` column to workspace table - No API breaking changes (agent endpoints unchanged) ## Testing - Integration tests updated to remove handoff dependencies - Agent tool test utilities simplified - Plan validation covered by new logic ## Next Steps (Future PRs) - Parallel execution of independent plan steps - Dynamic re-planning based on results - Plan caching for common routing patterns - Error recovery strategies in plan executor
This commit is contained in:
@@ -0,0 +1,10 @@
|
||||
export const MCP_SERVER_METADATA = {
|
||||
metadata: {
|
||||
info: 'Objects structure your business entities in Twenty. **Standard Objects** (e.g. People, Companies, Opportunities) are built‑in, pre‑configured data models. **Custom Objects** let you define entities specific to your needs (like Rockets, Properties, etc.). **Fields** work like spreadsheet columns and can be standard or custom. Always use the `fields` and `objects` parameters to select only the data you need—this **strongly reduces response size and token usage**, improving performance.',
|
||||
},
|
||||
protocolVersion: '2024-11-05',
|
||||
serverInfo: {
|
||||
name: 'Twenty MCP Server',
|
||||
version: '0.0.1',
|
||||
},
|
||||
};
|
||||
+180
@@ -0,0 +1,180 @@
|
||||
import { Test, type TestingModule } from '@nestjs/testing';
|
||||
|
||||
import { McpProtocolService } from 'src/engine/api/mcp/services/mcp-protocol.service';
|
||||
import { type JsonRpc } from 'src/engine/api/mcp/dtos/json-rpc';
|
||||
import { type WorkspaceEntity } from 'src/engine/core-modules/workspace/workspace.entity';
|
||||
import { MCP_SERVER_METADATA } from 'src/engine/api/mcp/constants/mcp.const';
|
||||
import { AccessTokenService } from 'src/engine/core-modules/auth/token/services/access-token.service';
|
||||
import { WorkspaceCacheStorageService } from 'src/engine/workspace-cache-storage/workspace-cache-storage.service';
|
||||
import { HttpExceptionHandlerService } from 'src/engine/core-modules/exception-handler/http-exception-handler.service';
|
||||
import { McpCoreController } from 'src/engine/api/mcp/controllers/mcp-core.controller';
|
||||
|
||||
describe('McpCoreController', () => {
|
||||
let controller: McpCoreController;
|
||||
let mcpProtocolService: jest.Mocked<McpProtocolService>;
|
||||
|
||||
beforeEach(async () => {
|
||||
const mockMcpProtocolService = {
|
||||
handleMCPCoreQuery: jest.fn(),
|
||||
};
|
||||
|
||||
const module: TestingModule = await Test.createTestingModule({
|
||||
controllers: [McpCoreController],
|
||||
providers: [
|
||||
{
|
||||
provide: McpProtocolService,
|
||||
useValue: mockMcpProtocolService,
|
||||
},
|
||||
{
|
||||
provide: AccessTokenService,
|
||||
useValue: jest.fn(),
|
||||
},
|
||||
{
|
||||
provide: WorkspaceCacheStorageService,
|
||||
useValue: jest.fn(),
|
||||
},
|
||||
{
|
||||
provide: HttpExceptionHandlerService,
|
||||
useValue: {
|
||||
handleError: jest.fn(),
|
||||
},
|
||||
},
|
||||
],
|
||||
}).compile();
|
||||
|
||||
controller = module.get<McpCoreController>(McpCoreController);
|
||||
mcpProtocolService = module.get(McpProtocolService);
|
||||
});
|
||||
|
||||
it('should be defined', () => {
|
||||
expect(controller).toBeDefined();
|
||||
});
|
||||
|
||||
describe('handleMcpCore', () => {
|
||||
const mockWorkspace = { id: 'workspace-1' } as WorkspaceEntity;
|
||||
const mockUserWorkspaceId = 'user-workspace-1';
|
||||
const mockApiKey = 'api-key-1';
|
||||
|
||||
it('should call mcpProtocolService.handleMCPCoreQuery with correct parameters', async () => {
|
||||
const mockRequest: JsonRpc = {
|
||||
jsonrpc: '2.0',
|
||||
method: 'tools/call',
|
||||
params: { name: 'testTool', arguments: { arg1: 'value1' } },
|
||||
id: '123',
|
||||
};
|
||||
|
||||
const mockResponse = {
|
||||
id: '123',
|
||||
jsonrpc: '2.0',
|
||||
result: {
|
||||
content: [{ type: 'text', text: '{"result":"success"}' }],
|
||||
isError: false,
|
||||
},
|
||||
};
|
||||
|
||||
mcpProtocolService.handleMCPCoreQuery.mockResolvedValue(mockResponse);
|
||||
|
||||
const result = await controller.handleMcpCore(
|
||||
mockRequest,
|
||||
mockWorkspace,
|
||||
mockApiKey,
|
||||
mockUserWorkspaceId,
|
||||
);
|
||||
|
||||
expect(mcpProtocolService.handleMCPCoreQuery).toHaveBeenCalledWith(
|
||||
mockRequest,
|
||||
{
|
||||
workspace: mockWorkspace,
|
||||
userWorkspaceId: mockUserWorkspaceId,
|
||||
apiKey: mockApiKey,
|
||||
},
|
||||
);
|
||||
expect(result).toEqual(mockResponse);
|
||||
});
|
||||
|
||||
it('should handle initialize method', async () => {
|
||||
const mockRequest: JsonRpc = {
|
||||
jsonrpc: '2.0',
|
||||
method: 'initialize',
|
||||
id: '123',
|
||||
};
|
||||
|
||||
const mockResponse = {
|
||||
id: '123',
|
||||
jsonrpc: '2.0',
|
||||
result: {
|
||||
...MCP_SERVER_METADATA,
|
||||
capabilities: {
|
||||
tools: { listChanged: false },
|
||||
resources: { listChanged: false },
|
||||
prompts: { listChanged: false },
|
||||
},
|
||||
},
|
||||
};
|
||||
|
||||
mcpProtocolService.handleMCPCoreQuery.mockResolvedValue(mockResponse);
|
||||
|
||||
const result = await controller.handleMcpCore(
|
||||
mockRequest,
|
||||
mockWorkspace,
|
||||
mockApiKey,
|
||||
mockUserWorkspaceId,
|
||||
);
|
||||
|
||||
expect(mcpProtocolService.handleMCPCoreQuery).toHaveBeenCalledWith(
|
||||
mockRequest,
|
||||
{
|
||||
workspace: mockWorkspace,
|
||||
userWorkspaceId: mockUserWorkspaceId,
|
||||
apiKey: mockApiKey,
|
||||
},
|
||||
);
|
||||
expect(result).toEqual(mockResponse);
|
||||
});
|
||||
|
||||
it('should handle tools listing', async () => {
|
||||
const mockRequest: JsonRpc = {
|
||||
jsonrpc: '2.0',
|
||||
method: 'tools/list',
|
||||
id: '123',
|
||||
};
|
||||
|
||||
const mockResponse = {
|
||||
id: '123',
|
||||
jsonrpc: '2.0',
|
||||
result: {
|
||||
...MCP_SERVER_METADATA,
|
||||
capabilities: {
|
||||
tools: { listChanged: false },
|
||||
},
|
||||
tools: [
|
||||
{
|
||||
name: 'testTool',
|
||||
description: 'A test tool',
|
||||
inputSchema: { type: 'object', properties: {} },
|
||||
},
|
||||
],
|
||||
},
|
||||
};
|
||||
|
||||
mcpProtocolService.handleMCPCoreQuery.mockResolvedValue(mockResponse);
|
||||
|
||||
const result = await controller.handleMcpCore(
|
||||
mockRequest,
|
||||
mockWorkspace,
|
||||
mockApiKey,
|
||||
mockUserWorkspaceId,
|
||||
);
|
||||
|
||||
expect(mcpProtocolService.handleMCPCoreQuery).toHaveBeenCalledWith(
|
||||
mockRequest,
|
||||
{
|
||||
workspace: mockWorkspace,
|
||||
userWorkspaceId: mockUserWorkspaceId,
|
||||
apiKey: mockApiKey,
|
||||
},
|
||||
);
|
||||
expect(result).toEqual(mockResponse);
|
||||
});
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,48 @@
|
||||
import {
|
||||
Body,
|
||||
Controller,
|
||||
Post,
|
||||
UseFilters,
|
||||
UseGuards,
|
||||
UsePipes,
|
||||
ValidationPipe,
|
||||
} from '@nestjs/common';
|
||||
|
||||
import { JsonRpc } from 'src/engine/api/mcp/dtos/json-rpc';
|
||||
import { McpProtocolService } from 'src/engine/api/mcp/services/mcp-protocol.service';
|
||||
import { RestApiExceptionFilter } from 'src/engine/api/rest/rest-api-exception.filter';
|
||||
import { WorkspaceEntity } from 'src/engine/core-modules/workspace/workspace.entity';
|
||||
import { AuthApiKey } from 'src/engine/decorators/auth/auth-api-key.decorator';
|
||||
import { AuthUserWorkspaceId } from 'src/engine/decorators/auth/auth-user-workspace-id.decorator';
|
||||
import { AuthWorkspace } from 'src/engine/decorators/auth/auth-workspace.decorator';
|
||||
import { JwtAuthGuard } from 'src/engine/guards/jwt-auth.guard';
|
||||
import { NoPermissionGuard } from 'src/engine/guards/no-permission.guard';
|
||||
import { WorkspaceAuthGuard } from 'src/engine/guards/workspace-auth.guard';
|
||||
|
||||
@Controller('mcp')
|
||||
@UseGuards(JwtAuthGuard, WorkspaceAuthGuard, NoPermissionGuard)
|
||||
@UseFilters(RestApiExceptionFilter)
|
||||
export class McpCoreController {
|
||||
constructor(private readonly mcpProtocolService: McpProtocolService) {}
|
||||
|
||||
@Post()
|
||||
@UsePipes(
|
||||
new ValidationPipe({
|
||||
transform: true,
|
||||
whitelist: true,
|
||||
forbidNonWhitelisted: true,
|
||||
}),
|
||||
)
|
||||
async handleMcpCore(
|
||||
@Body() body: JsonRpc,
|
||||
@AuthWorkspace() workspace: WorkspaceEntity,
|
||||
@AuthApiKey() apiKey: string | undefined,
|
||||
@AuthUserWorkspaceId() userWorkspaceId: string | undefined,
|
||||
) {
|
||||
return await this.mcpProtocolService.handleMCPCoreQuery(body, {
|
||||
workspace,
|
||||
userWorkspaceId,
|
||||
apiKey,
|
||||
});
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,15 @@
|
||||
import {
|
||||
ValidatorConstraint,
|
||||
type ValidatorConstraintInterface,
|
||||
} from 'class-validator';
|
||||
|
||||
@ValidatorConstraint({ name: 'string-or-number', async: false })
|
||||
export class IsNumberOrString implements ValidatorConstraintInterface {
|
||||
validate(value: unknown) {
|
||||
return typeof value === 'number' || typeof value === 'string';
|
||||
}
|
||||
|
||||
defaultMessage() {
|
||||
return '($value) must be number or string';
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,33 @@
|
||||
import {
|
||||
IsDefined,
|
||||
IsNotEmpty,
|
||||
IsObject,
|
||||
IsOptional,
|
||||
IsString,
|
||||
Matches,
|
||||
Validate,
|
||||
} from 'class-validator';
|
||||
|
||||
import { IsNumberOrString } from 'src/engine/api/mcp/decorators/string-or-number';
|
||||
|
||||
export class JsonRpc {
|
||||
@IsString()
|
||||
@Matches(/^2\.0$/, { message: 'jsonrpc must be exactly "2.0"' })
|
||||
jsonrpc = '2.0';
|
||||
|
||||
@IsDefined({ message: 'method is required' })
|
||||
@IsString()
|
||||
@IsNotEmpty()
|
||||
method: string;
|
||||
|
||||
@IsOptional()
|
||||
@IsObject()
|
||||
params?: {
|
||||
name: string;
|
||||
arguments: unknown;
|
||||
};
|
||||
|
||||
@IsOptional()
|
||||
@Validate(IsNumberOrString)
|
||||
id: string | number;
|
||||
}
|
||||
@@ -1,36 +1,43 @@
|
||||
import { Module } from '@nestjs/common';
|
||||
import { TypeOrmModule } from '@nestjs/typeorm';
|
||||
|
||||
import { AiModule } from 'src/engine/core-modules/ai/ai.module';
|
||||
import { TokenModule } from 'src/engine/core-modules/auth/token/token.module';
|
||||
import { McpCoreController } from 'src/engine/api/mcp/controllers/mcp-core.controller';
|
||||
import { McpMetadataController } from 'src/engine/api/mcp/controllers/mcp-metadata.controller';
|
||||
import { WorkspaceCacheStorageModule } from 'src/engine/workspace-cache-storage/workspace-cache-storage.module';
|
||||
import { MCPMetadataService } from 'src/engine/api/mcp/services/mcp-metadata.service';
|
||||
import { FeatureFlagModule } from 'src/engine/core-modules/feature-flag/feature-flag.module';
|
||||
import { RoleEntity } from 'src/engine/metadata-modules/role/role.entity';
|
||||
import { RestApiModule } from 'src/engine/api/rest/rest-api.module';
|
||||
import { MetadataQueryBuilderModule } from 'src/engine/api/rest/metadata/query-builder/metadata-query-builder.module';
|
||||
import { MCPMetadataToolsService } from 'src/engine/api/mcp/services/tools/mcp-metadata-tools.service';
|
||||
import { UpdateToolsService } from 'src/engine/api/mcp/services/tools/update.tools.service';
|
||||
import { McpProtocolService } from 'src/engine/api/mcp/services/mcp-protocol.service';
|
||||
import { McpToolExecutorService } from 'src/engine/api/mcp/services/mcp-tool-executor.service';
|
||||
import { CreateToolsService } from 'src/engine/api/mcp/services/tools/create.tools.service';
|
||||
import { DeleteToolsService } from 'src/engine/api/mcp/services/tools/delete.tools.service';
|
||||
import { GetToolsService } from 'src/engine/api/mcp/services/tools/get.tools.service';
|
||||
import { MCPMetadataToolsService } from 'src/engine/api/mcp/services/tools/mcp-metadata-tools.service';
|
||||
import { UpdateToolsService } from 'src/engine/api/mcp/services/tools/update.tools.service';
|
||||
import { MetadataQueryBuilderModule } from 'src/engine/api/rest/metadata/query-builder/metadata-query-builder.module';
|
||||
import { RestApiModule } from 'src/engine/api/rest/rest-api.module';
|
||||
import { AiToolsModule } from 'src/engine/metadata-modules/ai-tools/ai-tools.module';
|
||||
import { TokenModule } from 'src/engine/core-modules/auth/token/token.module';
|
||||
import { FeatureFlagModule } from 'src/engine/core-modules/feature-flag/feature-flag.module';
|
||||
import { MetricsModule } from 'src/engine/core-modules/metrics/metrics.module';
|
||||
import { RoleEntity } from 'src/engine/metadata-modules/role/role.entity';
|
||||
import { UserRoleModule } from 'src/engine/metadata-modules/user-role/user-role.module';
|
||||
import { WorkspaceCacheStorageModule } from 'src/engine/workspace-cache-storage/workspace-cache-storage.module';
|
||||
|
||||
@Module({
|
||||
imports: [
|
||||
TypeOrmModule.forFeature([RoleEntity]),
|
||||
AiModule,
|
||||
AiToolsModule,
|
||||
TokenModule,
|
||||
WorkspaceCacheStorageModule,
|
||||
FeatureFlagModule,
|
||||
RestApiModule,
|
||||
MetadataQueryBuilderModule,
|
||||
MetricsModule,
|
||||
UserRoleModule,
|
||||
],
|
||||
controllers: [McpMetadataController],
|
||||
exports: [],
|
||||
controllers: [McpCoreController, McpMetadataController],
|
||||
exports: [McpProtocolService],
|
||||
providers: [
|
||||
McpProtocolService,
|
||||
McpToolExecutorService,
|
||||
MCPMetadataService,
|
||||
MCPMetadataToolsService,
|
||||
CreateToolsService,
|
||||
|
||||
+425
@@ -0,0 +1,425 @@
|
||||
import { HttpException, HttpStatus } from '@nestjs/common';
|
||||
import { Test, type TestingModule } from '@nestjs/testing';
|
||||
import { getRepositoryToken } from '@nestjs/typeorm';
|
||||
|
||||
import { jsonSchema } from 'ai';
|
||||
|
||||
import { MCP_SERVER_METADATA } from 'src/engine/api/mcp/constants/mcp.const';
|
||||
import { type JsonRpc } from 'src/engine/api/mcp/dtos/json-rpc';
|
||||
import { McpProtocolService } from 'src/engine/api/mcp/services/mcp-protocol.service';
|
||||
import { McpToolExecutorService } from 'src/engine/api/mcp/services/mcp-tool-executor.service';
|
||||
import { ToolService } from 'src/engine/metadata-modules/ai-tools/services/tool.service';
|
||||
import { FeatureFlagKey } from 'src/engine/core-modules/feature-flag/enums/feature-flag-key.enum';
|
||||
import { FeatureFlagService } from 'src/engine/core-modules/feature-flag/services/feature-flag.service';
|
||||
import { type WorkspaceEntity } from 'src/engine/core-modules/workspace/workspace.entity';
|
||||
import { ADMIN_ROLE_LABEL } from 'src/engine/metadata-modules/permissions/constants/admin-role-label.constants';
|
||||
import { RoleEntity } from 'src/engine/metadata-modules/role/role.entity';
|
||||
import { UserRoleService } from 'src/engine/metadata-modules/user-role/user-role.service';
|
||||
|
||||
describe('McpProtocolService', () => {
|
||||
let service: McpProtocolService;
|
||||
let featureFlagService: jest.Mocked<FeatureFlagService>;
|
||||
let toolService: jest.Mocked<ToolService>;
|
||||
let userRoleService: jest.Mocked<UserRoleService>;
|
||||
let mcpToolExecutorService: jest.Mocked<McpToolExecutorService>;
|
||||
|
||||
const mockWorkspace = { id: 'workspace-1' } as WorkspaceEntity;
|
||||
const mockUserWorkspaceId = 'user-workspace-1';
|
||||
const mockRoleId = 'role-1';
|
||||
const mockAdminRoleId = 'admin-role-1';
|
||||
const mockApiKey = 'api-key-1';
|
||||
|
||||
beforeEach(async () => {
|
||||
const mockFeatureFlagService = {
|
||||
isFeatureEnabled: jest.fn(),
|
||||
};
|
||||
|
||||
const mockToolService = {
|
||||
listTools: jest.fn(),
|
||||
};
|
||||
|
||||
const mockUserRoleService = {
|
||||
getRoleIdForUserWorkspace: jest.fn(),
|
||||
};
|
||||
|
||||
const mockMcpToolExecutorService = {
|
||||
handleToolCall: jest.fn(),
|
||||
handleToolsListing: jest.fn(),
|
||||
};
|
||||
|
||||
const mockAdminRole = {
|
||||
id: mockAdminRoleId,
|
||||
label: ADMIN_ROLE_LABEL,
|
||||
} as RoleEntity;
|
||||
|
||||
const module: TestingModule = await Test.createTestingModule({
|
||||
providers: [
|
||||
McpProtocolService,
|
||||
{
|
||||
provide: FeatureFlagService,
|
||||
useValue: mockFeatureFlagService,
|
||||
},
|
||||
{
|
||||
provide: ToolService,
|
||||
useValue: mockToolService,
|
||||
},
|
||||
{
|
||||
provide: UserRoleService,
|
||||
useValue: mockUserRoleService,
|
||||
},
|
||||
{
|
||||
provide: McpToolExecutorService,
|
||||
useValue: mockMcpToolExecutorService,
|
||||
},
|
||||
{
|
||||
provide: getRepositoryToken(RoleEntity),
|
||||
useValue: {
|
||||
find: jest.fn().mockResolvedValue([mockAdminRole]),
|
||||
},
|
||||
},
|
||||
],
|
||||
}).compile();
|
||||
|
||||
service = module.get<McpProtocolService>(McpProtocolService);
|
||||
featureFlagService = module.get(FeatureFlagService);
|
||||
toolService = module.get(ToolService);
|
||||
userRoleService = module.get(UserRoleService);
|
||||
mcpToolExecutorService = module.get(McpToolExecutorService);
|
||||
});
|
||||
|
||||
it('should be defined', () => {
|
||||
expect(service).toBeDefined();
|
||||
});
|
||||
|
||||
describe('checkAiEnabled', () => {
|
||||
it('should not throw when AI is enabled', async () => {
|
||||
featureFlagService.isFeatureEnabled.mockResolvedValue(true);
|
||||
|
||||
await expect(
|
||||
service.checkAiEnabled('workspace-1'),
|
||||
).resolves.not.toThrow();
|
||||
expect(featureFlagService.isFeatureEnabled).toHaveBeenCalledWith(
|
||||
FeatureFlagKey.IS_AI_ENABLED,
|
||||
'workspace-1',
|
||||
);
|
||||
});
|
||||
|
||||
it('should throw when AI is disabled', async () => {
|
||||
featureFlagService.isFeatureEnabled.mockResolvedValue(false);
|
||||
|
||||
await expect(service.checkAiEnabled('workspace-1')).rejects.toThrow(
|
||||
new HttpException(
|
||||
'AI feature is not enabled for this workspace',
|
||||
HttpStatus.FORBIDDEN,
|
||||
),
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
describe('handleInitialize', () => {
|
||||
it('should return correct initialization response', () => {
|
||||
const requestId = '123';
|
||||
const result = service.handleInitialize(requestId);
|
||||
|
||||
expect(result).toMatchObject({
|
||||
id: requestId,
|
||||
jsonrpc: '2.0',
|
||||
result: expect.objectContaining({
|
||||
...MCP_SERVER_METADATA,
|
||||
capabilities: {
|
||||
tools: { listChanged: false },
|
||||
resources: { listChanged: false },
|
||||
prompts: { listChanged: false },
|
||||
},
|
||||
}),
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
describe('getRoleId', () => {
|
||||
it('should return role ID when available', async () => {
|
||||
userRoleService.getRoleIdForUserWorkspace.mockResolvedValue(mockRoleId);
|
||||
|
||||
const result = await service.getRoleId('workspace-1', 'user-workspace-1');
|
||||
|
||||
expect(result).toBe(mockRoleId);
|
||||
expect(userRoleService.getRoleIdForUserWorkspace).toHaveBeenCalledWith({
|
||||
workspaceId: 'workspace-1',
|
||||
userWorkspaceId: 'user-workspace-1',
|
||||
});
|
||||
});
|
||||
|
||||
it('should throw when userWorkspaceId is missing and no apiKey is provided', async () => {
|
||||
await expect(service.getRoleId('workspace-1', undefined)).rejects.toThrow(
|
||||
new HttpException('User workspace ID missing', HttpStatus.FORBIDDEN),
|
||||
);
|
||||
});
|
||||
|
||||
it('should throw when role ID is missing', async () => {
|
||||
userRoleService.getRoleIdForUserWorkspace.mockResolvedValue(undefined);
|
||||
|
||||
await expect(
|
||||
service.getRoleId('workspace-1', 'user-workspace-1'),
|
||||
).rejects.toThrow(
|
||||
new HttpException('Role ID missing', HttpStatus.FORBIDDEN),
|
||||
);
|
||||
});
|
||||
|
||||
it('should return admin role ID when apiKey is provided', async () => {
|
||||
const result = await service.getRoleId(
|
||||
'workspace-1',
|
||||
undefined,
|
||||
mockApiKey,
|
||||
);
|
||||
|
||||
expect(result).toBe(mockAdminRoleId);
|
||||
});
|
||||
});
|
||||
|
||||
describe('handleMCPCoreQuery', () => {
|
||||
it('should handle initialize method', async () => {
|
||||
featureFlagService.isFeatureEnabled.mockResolvedValue(true);
|
||||
|
||||
const mockRequest: JsonRpc = {
|
||||
jsonrpc: '2.0',
|
||||
method: 'initialize',
|
||||
id: '123',
|
||||
};
|
||||
|
||||
const result = await service.handleMCPCoreQuery(mockRequest, {
|
||||
workspace: mockWorkspace,
|
||||
userWorkspaceId: mockUserWorkspaceId,
|
||||
});
|
||||
|
||||
expect(result).toMatchObject({
|
||||
id: '123',
|
||||
jsonrpc: '2.0',
|
||||
result: expect.objectContaining({
|
||||
...MCP_SERVER_METADATA,
|
||||
capabilities: {
|
||||
tools: { listChanged: false },
|
||||
resources: { listChanged: false },
|
||||
prompts: { listChanged: false },
|
||||
},
|
||||
}),
|
||||
});
|
||||
});
|
||||
|
||||
it('should handle tools/call method with userWorkspaceId', async () => {
|
||||
featureFlagService.isFeatureEnabled.mockResolvedValue(true);
|
||||
userRoleService.getRoleIdForUserWorkspace.mockResolvedValue(mockRoleId);
|
||||
|
||||
const mockTool = {
|
||||
description: 'Test tool',
|
||||
inputSchema: jsonSchema({ type: 'object', properties: {} }),
|
||||
execute: jest.fn().mockResolvedValue({ result: 'success' }),
|
||||
};
|
||||
|
||||
const mockToolsMap = {
|
||||
testTool: mockTool,
|
||||
};
|
||||
|
||||
toolService.listTools.mockResolvedValue(mockToolsMap);
|
||||
|
||||
const mockToolCallResponse = {
|
||||
id: '123',
|
||||
jsonrpc: '2.0',
|
||||
result: {
|
||||
...MCP_SERVER_METADATA,
|
||||
content: [
|
||||
{
|
||||
type: 'text',
|
||||
text: JSON.stringify({ result: 'success' }),
|
||||
},
|
||||
],
|
||||
isError: false,
|
||||
},
|
||||
};
|
||||
|
||||
mcpToolExecutorService.handleToolCall.mockResolvedValue(
|
||||
mockToolCallResponse,
|
||||
);
|
||||
|
||||
const mockRequest: JsonRpc = {
|
||||
jsonrpc: '2.0',
|
||||
method: 'tools/call',
|
||||
params: { name: 'testTool', arguments: { arg1: 'value1' } },
|
||||
id: '123',
|
||||
};
|
||||
|
||||
const result = await service.handleMCPCoreQuery(mockRequest, {
|
||||
workspace: mockWorkspace,
|
||||
userWorkspaceId: mockUserWorkspaceId,
|
||||
});
|
||||
|
||||
expect(result).toEqual(mockToolCallResponse);
|
||||
expect(mcpToolExecutorService.handleToolCall).toHaveBeenCalledWith(
|
||||
'123',
|
||||
mockToolsMap,
|
||||
{ name: 'testTool', arguments: { arg1: 'value1' } },
|
||||
);
|
||||
});
|
||||
|
||||
it('should handle tools/call method with apiKey', async () => {
|
||||
featureFlagService.isFeatureEnabled.mockResolvedValue(true);
|
||||
|
||||
const mockTool = {
|
||||
description: 'Test tool',
|
||||
inputSchema: jsonSchema({ type: 'object', properties: {} }),
|
||||
execute: jest.fn().mockResolvedValue({ result: 'success' }),
|
||||
};
|
||||
|
||||
const mockToolsMap = {
|
||||
testTool: mockTool,
|
||||
};
|
||||
|
||||
toolService.listTools.mockResolvedValue(mockToolsMap);
|
||||
|
||||
const mockToolCallResponse = {
|
||||
id: '123',
|
||||
jsonrpc: '2.0',
|
||||
result: {
|
||||
...MCP_SERVER_METADATA,
|
||||
content: [
|
||||
{
|
||||
type: 'text',
|
||||
text: JSON.stringify({ result: 'success' }),
|
||||
},
|
||||
],
|
||||
isError: false,
|
||||
},
|
||||
};
|
||||
|
||||
mcpToolExecutorService.handleToolCall.mockResolvedValue(
|
||||
mockToolCallResponse,
|
||||
);
|
||||
|
||||
const mockRequest: JsonRpc = {
|
||||
jsonrpc: '2.0',
|
||||
method: 'tools/call',
|
||||
params: { name: 'testTool', arguments: { arg1: 'value1' } },
|
||||
id: '123',
|
||||
};
|
||||
|
||||
const result = await service.handleMCPCoreQuery(mockRequest, {
|
||||
workspace: mockWorkspace,
|
||||
apiKey: mockApiKey,
|
||||
});
|
||||
|
||||
expect(result).toEqual(mockToolCallResponse);
|
||||
expect(toolService.listTools).toHaveBeenCalledWith(
|
||||
{ unionOf: [mockAdminRoleId] },
|
||||
mockWorkspace.id,
|
||||
);
|
||||
});
|
||||
|
||||
it('should handle tools listing', async () => {
|
||||
featureFlagService.isFeatureEnabled.mockResolvedValue(true);
|
||||
userRoleService.getRoleIdForUserWorkspace.mockResolvedValue(mockRoleId);
|
||||
|
||||
const mockToolsMap = {
|
||||
testTool: {
|
||||
description: 'Test tool',
|
||||
inputSchema: jsonSchema({ type: 'object', properties: {} }),
|
||||
},
|
||||
};
|
||||
|
||||
toolService.listTools.mockResolvedValue(mockToolsMap);
|
||||
|
||||
const mockToolsListingResponse = {
|
||||
id: '123',
|
||||
jsonrpc: '2.0',
|
||||
result: expect.objectContaining({
|
||||
...MCP_SERVER_METADATA,
|
||||
capabilities: {
|
||||
tools: { listChanged: false },
|
||||
},
|
||||
tools: [
|
||||
{
|
||||
name: 'testTool',
|
||||
description: 'Test tool',
|
||||
inputSchema: { type: 'object', properties: {} },
|
||||
},
|
||||
],
|
||||
}),
|
||||
};
|
||||
|
||||
mcpToolExecutorService.handleToolsListing.mockReturnValue(
|
||||
mockToolsListingResponse,
|
||||
);
|
||||
|
||||
const mockRequest: JsonRpc = {
|
||||
jsonrpc: '2.0',
|
||||
method: 'tools/list',
|
||||
id: '123',
|
||||
};
|
||||
|
||||
const result = await service.handleMCPCoreQuery(mockRequest, {
|
||||
workspace: mockWorkspace,
|
||||
userWorkspaceId: mockUserWorkspaceId,
|
||||
});
|
||||
|
||||
expect(result).toMatchObject(mockToolsListingResponse);
|
||||
});
|
||||
|
||||
it('should handle error when AI is disabled', async () => {
|
||||
featureFlagService.isFeatureEnabled.mockResolvedValue(false);
|
||||
|
||||
const mockRequest: JsonRpc = {
|
||||
jsonrpc: '2.0',
|
||||
method: 'tools/list',
|
||||
id: '123',
|
||||
};
|
||||
|
||||
const result = await service.handleMCPCoreQuery(mockRequest, {
|
||||
workspace: mockWorkspace,
|
||||
userWorkspaceId: mockUserWorkspaceId,
|
||||
});
|
||||
|
||||
expect(result).toEqual({
|
||||
id: '123',
|
||||
jsonrpc: '2.0',
|
||||
error: {
|
||||
...MCP_SERVER_METADATA,
|
||||
code: HttpStatus.FORBIDDEN,
|
||||
message: 'AI feature is not enabled for this workspace',
|
||||
},
|
||||
});
|
||||
});
|
||||
|
||||
it('should handle error when tool is not found', async () => {
|
||||
featureFlagService.isFeatureEnabled.mockResolvedValue(true);
|
||||
userRoleService.getRoleIdForUserWorkspace.mockResolvedValue(mockRoleId);
|
||||
toolService.listTools.mockResolvedValue({});
|
||||
|
||||
mcpToolExecutorService.handleToolCall.mockRejectedValue(
|
||||
new HttpException(
|
||||
"Tool 'nonExistentTool' not found",
|
||||
HttpStatus.NOT_FOUND,
|
||||
),
|
||||
);
|
||||
|
||||
const mockRequest: JsonRpc = {
|
||||
jsonrpc: '2.0',
|
||||
method: 'tools/call',
|
||||
params: { name: 'nonExistentTool', arguments: {} },
|
||||
id: '123',
|
||||
};
|
||||
|
||||
const result = await service.handleMCPCoreQuery(mockRequest, {
|
||||
workspace: mockWorkspace,
|
||||
userWorkspaceId: mockUserWorkspaceId,
|
||||
});
|
||||
|
||||
expect(result).toEqual({
|
||||
id: '123',
|
||||
jsonrpc: '2.0',
|
||||
error: {
|
||||
...MCP_SERVER_METADATA,
|
||||
code: HttpStatus.NOT_FOUND,
|
||||
message: "Tool 'nonExistentTool' not found",
|
||||
},
|
||||
});
|
||||
});
|
||||
});
|
||||
});
|
||||
@@ -8,7 +8,7 @@ import { CreateToolsService } from 'src/engine/api/mcp/services/tools/create.too
|
||||
import { DeleteToolsService } from 'src/engine/api/mcp/services/tools/delete.tools.service';
|
||||
import { GetToolsService } from 'src/engine/api/mcp/services/tools/get.tools.service';
|
||||
import { UpdateToolsService } from 'src/engine/api/mcp/services/tools/update.tools.service';
|
||||
import { wrapJsonRpcResponse } from 'src/engine/core-modules/ai/utils/wrap-jsonrpc-response.util';
|
||||
import { wrapJsonRpcResponse } from 'src/engine/api/mcp/utils/wrap-jsonrpc-response.util';
|
||||
import { FeatureFlagKey } from 'src/engine/core-modules/feature-flag/enums/feature-flag-key.enum';
|
||||
import { FeatureFlagService } from 'src/engine/core-modules/feature-flag/services/feature-flag.service';
|
||||
import { MetricsService } from 'src/engine/core-modules/metrics/metrics.service';
|
||||
|
||||
@@ -0,0 +1,182 @@
|
||||
import { HttpException, HttpStatus, Injectable } from '@nestjs/common';
|
||||
import { InjectRepository } from '@nestjs/typeorm';
|
||||
|
||||
import { Repository } from 'typeorm';
|
||||
|
||||
import { type JsonRpc } from 'src/engine/api/mcp/dtos/json-rpc';
|
||||
import { McpToolExecutorService } from 'src/engine/api/mcp/services/mcp-tool-executor.service';
|
||||
import { wrapJsonRpcResponse } from 'src/engine/api/mcp/utils/wrap-jsonrpc-response.util';
|
||||
import { ToolService } from 'src/engine/metadata-modules/ai-tools/services/tool.service';
|
||||
import { FeatureFlagKey } from 'src/engine/core-modules/feature-flag/enums/feature-flag-key.enum';
|
||||
import { FeatureFlagService } from 'src/engine/core-modules/feature-flag/services/feature-flag.service';
|
||||
import { WorkspaceEntity } from 'src/engine/core-modules/workspace/workspace.entity';
|
||||
import { RoleEntity } from 'src/engine/metadata-modules/role/role.entity';
|
||||
import { UserRoleService } from 'src/engine/metadata-modules/user-role/user-role.service';
|
||||
import { ADMIN_ROLE } from 'src/engine/workspace-manager/workspace-sync-metadata/standard-roles/roles/admin-role';
|
||||
|
||||
@Injectable()
|
||||
export class McpProtocolService {
|
||||
constructor(
|
||||
private readonly featureFlagService: FeatureFlagService,
|
||||
private readonly toolService: ToolService,
|
||||
private readonly userRoleService: UserRoleService,
|
||||
private readonly mcpToolExecutorService: McpToolExecutorService,
|
||||
@InjectRepository(RoleEntity)
|
||||
private readonly roleRepository: Repository<RoleEntity>,
|
||||
) {}
|
||||
|
||||
async checkAiEnabled(workspaceId: string): Promise<void> {
|
||||
const isAiEnabled = await this.featureFlagService.isFeatureEnabled(
|
||||
FeatureFlagKey.IS_AI_ENABLED,
|
||||
workspaceId,
|
||||
);
|
||||
|
||||
if (!isAiEnabled) {
|
||||
throw new HttpException(
|
||||
'AI feature is not enabled for this workspace',
|
||||
HttpStatus.FORBIDDEN,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
handleInitialize(requestId: string | number) {
|
||||
return wrapJsonRpcResponse(requestId, {
|
||||
result: {
|
||||
capabilities: {
|
||||
tools: { listChanged: false },
|
||||
resources: { listChanged: false },
|
||||
prompts: { listChanged: false },
|
||||
},
|
||||
tools: [],
|
||||
resources: [],
|
||||
prompts: [],
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
async getRoleId(
|
||||
workspaceId: string,
|
||||
userWorkspaceId?: string,
|
||||
apiKey?: string,
|
||||
) {
|
||||
if (apiKey) {
|
||||
const roles = await this.roleRepository.find({
|
||||
where: {
|
||||
workspaceId,
|
||||
standardId: ADMIN_ROLE.standardId,
|
||||
},
|
||||
});
|
||||
|
||||
if (roles.length === 0) {
|
||||
throw new HttpException('Admin role not found', HttpStatus.FORBIDDEN);
|
||||
}
|
||||
|
||||
return roles[0].id;
|
||||
}
|
||||
|
||||
if (!userWorkspaceId) {
|
||||
throw new HttpException(
|
||||
'User workspace ID missing',
|
||||
HttpStatus.FORBIDDEN,
|
||||
);
|
||||
}
|
||||
|
||||
const roleId = await this.userRoleService.getRoleIdForUserWorkspace({
|
||||
workspaceId,
|
||||
userWorkspaceId,
|
||||
});
|
||||
|
||||
if (!roleId) {
|
||||
throw new HttpException('Role ID missing', HttpStatus.FORBIDDEN);
|
||||
}
|
||||
|
||||
return roleId;
|
||||
}
|
||||
|
||||
async handleMCPCoreQuery(
|
||||
{ id, method, params }: JsonRpc,
|
||||
{
|
||||
workspace,
|
||||
userWorkspaceId,
|
||||
apiKey,
|
||||
}: {
|
||||
workspace: WorkspaceEntity;
|
||||
userWorkspaceId?: string;
|
||||
apiKey?: string;
|
||||
},
|
||||
): Promise<Record<string, unknown>> {
|
||||
try {
|
||||
await this.checkAiEnabled(workspace.id);
|
||||
|
||||
if (method === 'initialize') {
|
||||
return this.handleInitialize(id);
|
||||
}
|
||||
|
||||
if (method === 'ping') {
|
||||
return wrapJsonRpcResponse(
|
||||
id,
|
||||
{
|
||||
result: {},
|
||||
},
|
||||
true,
|
||||
);
|
||||
}
|
||||
|
||||
const roleId = await this.getRoleId(
|
||||
workspace.id,
|
||||
userWorkspaceId,
|
||||
apiKey,
|
||||
);
|
||||
|
||||
const toolSet = await this.toolService.listTools(
|
||||
{ unionOf: [roleId] },
|
||||
workspace.id,
|
||||
);
|
||||
|
||||
if (method === 'tools/call' && params) {
|
||||
return await this.mcpToolExecutorService.handleToolCall(
|
||||
id,
|
||||
toolSet,
|
||||
params,
|
||||
);
|
||||
}
|
||||
|
||||
if (method === 'tools/list') {
|
||||
return this.mcpToolExecutorService.handleToolsListing(id, toolSet);
|
||||
}
|
||||
|
||||
if (method === 'prompts/list') {
|
||||
return wrapJsonRpcResponse(id, {
|
||||
result: {
|
||||
capabilities: {
|
||||
prompts: { listChanged: false },
|
||||
},
|
||||
prompts: [],
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
if (method === 'resources/list') {
|
||||
return wrapJsonRpcResponse(id, {
|
||||
result: {
|
||||
capabilities: {
|
||||
resources: { listChanged: false },
|
||||
},
|
||||
resources: [],
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
return wrapJsonRpcResponse(id, {
|
||||
result: {},
|
||||
});
|
||||
} catch (error) {
|
||||
return wrapJsonRpcResponse(id, {
|
||||
error: {
|
||||
code: error.status || HttpStatus.INTERNAL_SERVER_ERROR,
|
||||
message: error.message || 'Failed to execute tool',
|
||||
},
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,75 @@
|
||||
import { HttpException, HttpStatus, Injectable } from '@nestjs/common';
|
||||
|
||||
import { type ToolSet } from 'ai';
|
||||
import { isDefined } from 'twenty-shared/utils';
|
||||
|
||||
import { wrapJsonRpcResponse } from 'src/engine/api/mcp/utils/wrap-jsonrpc-response.util';
|
||||
|
||||
@Injectable()
|
||||
export class McpToolExecutorService {
|
||||
async handleToolCall(
|
||||
id: string | number,
|
||||
toolSet: ToolSet,
|
||||
params: Record<string, unknown>,
|
||||
) {
|
||||
const toolName = params.name as keyof typeof toolSet;
|
||||
const tool = toolSet[toolName];
|
||||
|
||||
if (isDefined(tool) && isDefined(tool.execute)) {
|
||||
return wrapJsonRpcResponse(id, {
|
||||
result: {
|
||||
content: [
|
||||
{
|
||||
type: 'text',
|
||||
text: JSON.stringify(
|
||||
await tool.execute(params.arguments, {
|
||||
toolCallId: '1',
|
||||
messages: [],
|
||||
}),
|
||||
),
|
||||
},
|
||||
],
|
||||
isError: false,
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
throw new HttpException(
|
||||
`Tool '${params.name}' not found`,
|
||||
HttpStatus.NOT_FOUND,
|
||||
);
|
||||
}
|
||||
|
||||
handleToolsListing(id: string | number, toolSet: ToolSet) {
|
||||
const toolsArray = Object.entries(toolSet)
|
||||
.filter(([, def]) => !!def.inputSchema)
|
||||
.map(([name, def]) => {
|
||||
// Unwrap the AI SDK's jsonSchema wrapper if present
|
||||
// The AI SDK serializes schemas as { jsonSchema: {...} } but MCP expects {...} directly
|
||||
const inputSchema = def.inputSchema;
|
||||
const unwrappedSchema =
|
||||
inputSchema &&
|
||||
typeof inputSchema === 'object' &&
|
||||
'jsonSchema' in inputSchema
|
||||
? inputSchema.jsonSchema
|
||||
: inputSchema;
|
||||
|
||||
return {
|
||||
name,
|
||||
description: def.description,
|
||||
inputSchema: unwrappedSchema,
|
||||
};
|
||||
});
|
||||
|
||||
return wrapJsonRpcResponse(id, {
|
||||
result: {
|
||||
capabilities: {
|
||||
tools: { listChanged: false },
|
||||
},
|
||||
tools: toolsArray,
|
||||
resources: [],
|
||||
prompts: [],
|
||||
},
|
||||
});
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,28 @@
|
||||
import { MCP_SERVER_METADATA } from 'src/engine/api/mcp/constants/mcp.const';
|
||||
|
||||
export const wrapJsonRpcResponse = (
|
||||
id: string | number,
|
||||
payload:
|
||||
| Record<'result', Record<string, unknown>>
|
||||
| Record<'error', Record<string, unknown>>,
|
||||
omitMetadata = false,
|
||||
) => {
|
||||
const body =
|
||||
'result' in payload
|
||||
? {
|
||||
result: omitMetadata
|
||||
? payload.result
|
||||
: { ...payload.result, ...MCP_SERVER_METADATA },
|
||||
}
|
||||
: {
|
||||
error: omitMetadata
|
||||
? payload.error
|
||||
: { ...payload.error, ...MCP_SERVER_METADATA },
|
||||
};
|
||||
|
||||
return {
|
||||
id,
|
||||
jsonrpc: '2.0',
|
||||
...body,
|
||||
};
|
||||
};
|
||||
Reference in New Issue
Block a user