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:
Félix Malfait
2025-11-25 12:10:14 +01:00
committed by GitHub
parent 3c0ae49a23
commit e7ebf51e50
177 changed files with 2720 additions and 3222 deletions
@@ -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 builtin, preconfigured 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',
},
};
@@ -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,
@@ -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,
};
};