diff --git a/backend/src/app.module.ts b/backend/src/app.module.ts index e08627cc..b893b6d0 100644 --- a/backend/src/app.module.ts +++ b/backend/src/app.module.ts @@ -18,6 +18,7 @@ import { TestingSupportModule } from './testing/testing.module'; import { IntegrationsModule } from './integrations/integrations.module'; import { SchedulesModule } from './schedules/schedules.module'; import { AnalyticsModule } from './analytics/analytics.module'; +import { McpModule } from './mcp/mcp.module'; import { ApiKeysModule } from './api-keys/api-keys.module'; import { WebhooksModule } from './webhooks/webhooks.module'; @@ -37,6 +38,7 @@ const coreModules = [ ApiKeysModule, WebhooksModule, HumanInputsModule, + McpModule, ]; const testingModules = process.env.NODE_ENV === 'production' ? [] : [TestingSupportModule]; diff --git a/backend/src/mcp/__tests__/tool-registry.service.spec.ts b/backend/src/mcp/__tests__/tool-registry.service.spec.ts new file mode 100644 index 00000000..6ab315d1 --- /dev/null +++ b/backend/src/mcp/__tests__/tool-registry.service.spec.ts @@ -0,0 +1,228 @@ +import { describe, it, expect, beforeEach } from 'bun:test'; +import { ToolRegistryService } from '../tool-registry.service'; +import type { SecretsEncryptionService } from '../../secrets/secrets.encryption'; + +// Mock Redis +class MockRedis { + private data = new Map>(); + + async hset(key: string, field: string, value: string): Promise { + if (!this.data.has(key)) { + this.data.set(key, new Map()); + } + this.data.get(key)!.set(field, value); + return 1; + } + + async hget(key: string, field: string): Promise { + return this.data.get(key)?.get(field) ?? null; + } + + async hgetall(key: string): Promise> { + const hash = this.data.get(key); + if (!hash) return {}; + return Object.fromEntries(hash.entries()); + } + + async del(key: string): Promise { + this.data.delete(key); + return 1; + } + + async expire(_key: string, _seconds: number): Promise { + return 1; + } + + async quit(): Promise {} +} + +// Mock encryption service +class MockEncryptionService { + async encrypt(value: string): Promise<{ ciphertext: string; keyId: string }> { + return { + ciphertext: Buffer.from(value).toString('base64'), + keyId: 'test-key', + }; + } + + async decrypt(material: { ciphertext: string }): Promise { + return Buffer.from(material.ciphertext, 'base64').toString('utf-8'); + } +} + +describe('ToolRegistryService', () => { + let service: ToolRegistryService; + let redis: MockRedis; + let encryption: MockEncryptionService; + + beforeEach(() => { + redis = new MockRedis(); + encryption = new MockEncryptionService(); + service = new ToolRegistryService(redis as any, encryption as any as SecretsEncryptionService); + }); + + describe('registerComponentTool', () => { + it('registers a component tool with encrypted credentials', async () => { + await service.registerComponentTool({ + runId: 'run-1', + nodeId: 'node-a', + toolName: 'check_ip_reputation', + componentId: 'security.abuseipdb', + description: 'Check IP reputation', + inputSchema: { + type: 'object', + properties: { ipAddress: { type: 'string' } }, + required: ['ipAddress'], + }, + credentials: { apiKey: 'secret-123' }, + }); + + const tool = await service.getTool('run-1', 'node-a'); + expect(tool).not.toBeNull(); + expect(tool?.toolName).toBe('check_ip_reputation'); + expect(tool?.status).toBe('ready'); + expect(tool?.type).toBe('component'); + expect(tool?.encryptedCredentials).toBeDefined(); + }); + }); + + describe('getToolsForRun', () => { + it('returns all tools for a run', async () => { + await service.registerComponentTool({ + runId: 'run-1', + nodeId: 'node-a', + toolName: 'tool_a', + componentId: 'comp.a', + description: 'Tool A', + inputSchema: { type: 'object', properties: {}, required: [] }, + credentials: {}, + }); + + await service.registerComponentTool({ + runId: 'run-1', + nodeId: 'node-b', + toolName: 'tool_b', + componentId: 'comp.b', + description: 'Tool B', + inputSchema: { type: 'object', properties: {}, required: [] }, + credentials: {}, + }); + + const tools = await service.getToolsForRun('run-1'); + expect(tools.length).toBe(2); + expect(tools.map((t) => t.toolName).sort()).toEqual(['tool_a', 'tool_b']); + }); + }); + + describe('getToolByName', () => { + it('finds a tool by name', async () => { + await service.registerComponentTool({ + runId: 'run-1', + nodeId: 'node-a', + toolName: 'my_tool', + componentId: 'comp.a', + description: 'My Tool', + inputSchema: { type: 'object', properties: {}, required: [] }, + credentials: {}, + }); + + const tool = await service.getToolByName('run-1', 'my_tool'); + expect(tool).not.toBeNull(); + expect(tool?.nodeId).toBe('node-a'); + }); + + it('returns null for unknown tool name', async () => { + const tool = await service.getToolByName('run-1', 'unknown'); + expect(tool).toBeNull(); + }); + }); + + describe('getToolCredentials', () => { + it('decrypts and returns credentials', async () => { + await service.registerComponentTool({ + runId: 'run-1', + nodeId: 'node-a', + toolName: 'tool', + componentId: 'comp', + description: 'Tool', + inputSchema: { type: 'object', properties: {}, required: [] }, + credentials: { apiKey: 'secret-value', token: 'another-secret' }, + }); + + const creds = await service.getToolCredentials('run-1', 'node-a'); + expect(creds).toEqual({ apiKey: 'secret-value', token: 'another-secret' }); + }); + }); + + describe('areAllToolsReady', () => { + it('returns true when all required tools are ready', async () => { + await service.registerComponentTool({ + runId: 'run-1', + nodeId: 'node-a', + toolName: 'tool_a', + componentId: 'comp.a', + description: 'Tool A', + inputSchema: { type: 'object', properties: {}, required: [] }, + credentials: {}, + }); + + await service.registerComponentTool({ + runId: 'run-1', + nodeId: 'node-b', + toolName: 'tool_b', + componentId: 'comp.b', + description: 'Tool B', + inputSchema: { type: 'object', properties: {}, required: [] }, + credentials: {}, + }); + + const ready = await service.areAllToolsReady('run-1', ['node-a', 'node-b']); + expect(ready).toBe(true); + }); + + it('returns false when a required tool is missing', async () => { + await service.registerComponentTool({ + runId: 'run-1', + nodeId: 'node-a', + toolName: 'tool_a', + componentId: 'comp.a', + description: 'Tool A', + inputSchema: { type: 'object', properties: {}, required: [] }, + credentials: {}, + }); + + const ready = await service.areAllToolsReady('run-1', ['node-a', 'node-b']); + expect(ready).toBe(false); + }); + }); + + describe('cleanupRun', () => { + it('removes all tools and returns container IDs', async () => { + await service.registerComponentTool({ + runId: 'run-1', + nodeId: 'node-a', + toolName: 'tool_a', + componentId: 'comp.a', + description: 'Tool A', + inputSchema: { type: 'object', properties: {}, required: [] }, + credentials: {}, + }); + + await service.registerLocalMcp({ + runId: 'run-1', + nodeId: 'node-mcp', + toolName: 'steampipe', + description: 'Steampipe MCP', + inputSchema: { type: 'object', properties: {}, required: [] }, + endpoint: 'http://localhost:8080', + containerId: 'container-123', + }); + + const containerIds = await service.cleanupRun('run-1'); + expect(containerIds).toEqual(['container-123']); + + const tools = await service.getToolsForRun('run-1'); + expect(tools.length).toBe(0); + }); + }); +}); diff --git a/backend/src/mcp/index.ts b/backend/src/mcp/index.ts new file mode 100644 index 00000000..3e61a7f9 --- /dev/null +++ b/backend/src/mcp/index.ts @@ -0,0 +1,2 @@ +export * from './mcp.module'; +export * from './tool-registry.service'; diff --git a/backend/src/mcp/mcp.module.ts b/backend/src/mcp/mcp.module.ts new file mode 100644 index 00000000..026b38f3 --- /dev/null +++ b/backend/src/mcp/mcp.module.ts @@ -0,0 +1,25 @@ +import { Global, Module } from '@nestjs/common'; +import Redis from 'ioredis'; +import { ToolRegistryService, TOOL_REGISTRY_REDIS } from './tool-registry.service'; +import { SecretsModule } from '../secrets/secrets.module'; + +@Global() +@Module({ + imports: [SecretsModule], + providers: [ + { + provide: TOOL_REGISTRY_REDIS, + useFactory: () => { + // Use the same Redis URL as terminal or a dedicated one + const url = process.env.TOOL_REGISTRY_REDIS_URL ?? process.env.TERMINAL_REDIS_URL; + if (!url) { + return null; + } + return new Redis(url); + }, + }, + ToolRegistryService, + ], + exports: [ToolRegistryService], +}) +export class McpModule {} diff --git a/backend/src/mcp/tool-registry.service.ts b/backend/src/mcp/tool-registry.service.ts new file mode 100644 index 00000000..47b6a848 --- /dev/null +++ b/backend/src/mcp/tool-registry.service.ts @@ -0,0 +1,366 @@ +/** + * Tool Registry Service + * + * Redis-backed registry for storing tool metadata and credentials during workflow runs. + * This bridges the gap between Temporal workflows (where credentials are resolved) + * and the MCP gateway (where agents call tools). + * + * Redis key pattern: mcp:run:{runId}:tools (Hash) + * TTL: 1 hour (configurable) + */ + +import { Injectable, Logger, Inject, OnModuleDestroy } from '@nestjs/common'; +import type Redis from 'ioredis'; +import { type ToolInputSchema } from '@shipsec/component-sdk'; +import { SecretsEncryptionService } from '../secrets/secrets.encryption'; + +export const TOOL_REGISTRY_REDIS = Symbol('TOOL_REGISTRY_REDIS'); + +/** + * Types of tools that can be registered + */ +export type RegisteredToolType = 'component' | 'remote-mcp' | 'local-mcp'; + +/** + * Status of a registered tool + */ +export type ToolStatus = 'pending' | 'ready' | 'error'; + +/** + * A tool registered in the registry + */ +export interface RegisteredTool { + /** Unique ID of the workflow node */ + nodeId: string; + + /** Tool name exposed to the agent */ + toolName: string; + + /** Type of tool */ + type: RegisteredToolType; + + /** Current status */ + status: ToolStatus; + + /** Component ID (for component tools) */ + componentId?: string; + + /** JSON Schema for action inputs */ + inputSchema: ToolInputSchema; + + /** Tool description for the agent */ + description: string; + + /** Encrypted credentials (for component tools) */ + encryptedCredentials?: string; + + /** MCP endpoint URL (for remote/local MCPs) */ + endpoint?: string; + + /** Docker container ID (for local MCPs) */ + containerId?: string; + + /** Error message if status is 'error' */ + errorMessage?: string; + + /** Timestamp when tool was registered */ + registeredAt: string; +} + +/** + * Input for registering a component tool + */ +export interface RegisterComponentToolInput { + runId: string; + nodeId: string; + toolName: string; + componentId: string; + description: string; + inputSchema: ToolInputSchema; + credentials: Record; +} + +/** + * Input for registering a remote MCP + */ +export interface RegisterRemoteMcpInput { + runId: string; + nodeId: string; + toolName: string; + description: string; + inputSchema: ToolInputSchema; + endpoint: string; + authToken?: string; +} + +/** + * Input for registering a local MCP (stdio container) + */ +export interface RegisterLocalMcpInput { + runId: string; + nodeId: string; + toolName: string; + description: string; + inputSchema: ToolInputSchema; + endpoint: string; + containerId: string; +} + +const REGISTRY_TTL_SECONDS = 60 * 60; // 1 hour + +@Injectable() +export class ToolRegistryService implements OnModuleDestroy { + private readonly logger = new Logger(ToolRegistryService.name); + + constructor( + @Inject(TOOL_REGISTRY_REDIS) private readonly redis: Redis | null, + private readonly encryption: SecretsEncryptionService, + ) {} + + async onModuleDestroy() { + await this.redis?.quit(); + } + + private getRegistryKey(runId: string): string { + return `mcp:run:${runId}:tools`; + } + + /** + * Register a ShipSec component as an agent-callable tool + */ + async registerComponentTool(input: RegisterComponentToolInput): Promise { + if (!this.redis) { + this.logger.warn('Redis not configured, tool registry disabled'); + return; + } + + const { runId, nodeId, toolName, componentId, description, inputSchema, credentials } = input; + + // Encrypt credentials + const credentialsJson = JSON.stringify(credentials); + const encryptionMaterial = await this.encryption.encrypt(credentialsJson); + const encryptedCredentials = JSON.stringify(encryptionMaterial); + + const tool: RegisteredTool = { + nodeId, + toolName, + type: 'component', + status: 'ready', + componentId, + description, + inputSchema, + encryptedCredentials, + registeredAt: new Date().toISOString(), + }; + + const key = this.getRegistryKey(runId); + await this.redis.hset(key, nodeId, JSON.stringify(tool)); + await this.redis.expire(key, REGISTRY_TTL_SECONDS); + + this.logger.log(`Registered component tool: ${toolName} (node: ${nodeId}, run: ${runId})`); + } + + /** + * Register a remote HTTP MCP server + */ + async registerRemoteMcp(input: RegisterRemoteMcpInput): Promise { + if (!this.redis) { + this.logger.warn('Redis not configured, tool registry disabled'); + return; + } + + const { runId, nodeId, toolName, description, inputSchema, endpoint, authToken } = input; + + // Encrypt auth token if provided + let encryptedCredentials: string | undefined; + if (authToken) { + const encryptionMaterial = await this.encryption.encrypt(authToken); + encryptedCredentials = JSON.stringify(encryptionMaterial); + } + + const tool: RegisteredTool = { + nodeId, + toolName, + type: 'remote-mcp', + status: 'ready', + description, + inputSchema, + endpoint, + encryptedCredentials, + registeredAt: new Date().toISOString(), + }; + + const key = this.getRegistryKey(runId); + await this.redis.hset(key, nodeId, JSON.stringify(tool)); + await this.redis.expire(key, REGISTRY_TTL_SECONDS); + + this.logger.log(`Registered remote MCP: ${toolName} (node: ${nodeId}, run: ${runId})`); + } + + /** + * Register a local stdio MCP running in Docker + */ + async registerLocalMcp(input: RegisterLocalMcpInput): Promise { + if (!this.redis) { + this.logger.warn('Redis not configured, tool registry disabled'); + return; + } + + const { runId, nodeId, toolName, description, inputSchema, endpoint, containerId } = input; + + const tool: RegisteredTool = { + nodeId, + toolName, + type: 'local-mcp', + status: 'ready', + description, + inputSchema, + endpoint, + containerId, + registeredAt: new Date().toISOString(), + }; + + const key = this.getRegistryKey(runId); + await this.redis.hset(key, nodeId, JSON.stringify(tool)); + await this.redis.expire(key, REGISTRY_TTL_SECONDS); + + this.logger.log( + `Registered local MCP: ${toolName} (node: ${nodeId}, container: ${containerId}, run: ${runId})`, + ); + } + + /** + * Get all registered tools for a workflow run + */ + async getToolsForRun(runId: string): Promise { + if (!this.redis) { + return []; + } + + const key = this.getRegistryKey(runId); + const toolsHash = await this.redis.hgetall(key); + + return Object.values(toolsHash).map((json) => JSON.parse(json) as RegisteredTool); + } + + /** + * Get a specific tool by node ID + */ + async getTool(runId: string, nodeId: string): Promise { + if (!this.redis) { + return null; + } + + const key = this.getRegistryKey(runId); + const toolJson = await this.redis.hget(key, nodeId); + + if (!toolJson) { + return null; + } + + return JSON.parse(toolJson) as RegisteredTool; + } + + /** + * Get a tool by its tool name + */ + async getToolByName(runId: string, toolName: string): Promise { + const tools = await this.getToolsForRun(runId); + return tools.find((t) => t.toolName === toolName) ?? null; + } + + /** + * Decrypt and return credentials for a tool + */ + async getToolCredentials(runId: string, nodeId: string): Promise | null> { + const tool = await this.getTool(runId, nodeId); + if (!tool?.encryptedCredentials) { + return null; + } + + try { + const encryptionMaterial = JSON.parse(tool.encryptedCredentials); + const decrypted = await this.encryption.decrypt(encryptionMaterial); + return JSON.parse(decrypted); + } catch (error) { + this.logger.error(`Failed to decrypt credentials for tool ${nodeId}:`, error); + return null; + } + } + + /** + * Check if all required tools are ready + */ + async areAllToolsReady(runId: string, requiredNodeIds: string[]): Promise { + if (!this.redis) { + return true; // If Redis is disabled, assume ready + } + + const key = this.getRegistryKey(runId); + + for (const nodeId of requiredNodeIds) { + const toolJson = await this.redis.hget(key, nodeId); + if (!toolJson) { + return false; + } + + const tool = JSON.parse(toolJson) as RegisteredTool; + if (tool.status !== 'ready') { + return false; + } + } + + return true; + } + + /** + * Update tool status (e.g., to 'error') + */ + async updateToolStatus( + runId: string, + nodeId: string, + status: ToolStatus, + errorMessage?: string, + ): Promise { + if (!this.redis) { + return; + } + + const tool = await this.getTool(runId, nodeId); + if (!tool) { + return; + } + + tool.status = status; + if (errorMessage) { + tool.errorMessage = errorMessage; + } + + const key = this.getRegistryKey(runId); + await this.redis.hset(key, nodeId, JSON.stringify(tool)); + } + + /** + * Clean up all tools for a run (called when workflow completes) + * Returns container IDs that need to be stopped + */ + async cleanupRun(runId: string): Promise { + if (!this.redis) { + return []; + } + + const tools = await this.getToolsForRun(runId); + const containerIds = tools + .filter((t) => t.type === 'local-mcp' && t.containerId) + .map((t) => t.containerId!); + + const key = this.getRegistryKey(runId); + await this.redis.del(key); + + this.logger.log( + `Cleaned up tool registry for run ${runId} (${tools.length} tools, ${containerIds.length} containers)`, + ); + + return containerIds; + } +} diff --git a/backend/src/secrets/secrets.module.ts b/backend/src/secrets/secrets.module.ts index 677ece28..f6bac70a 100644 --- a/backend/src/secrets/secrets.module.ts +++ b/backend/src/secrets/secrets.module.ts @@ -10,6 +10,6 @@ import { SecretsService } from './secrets.service'; imports: [DatabaseModule], controllers: [SecretsController], providers: [SecretsService, SecretsRepository, SecretsEncryptionService], - exports: [SecretsService], + exports: [SecretsService, SecretsEncryptionService], }) export class SecretsModule {}