diff --git a/packages/dvmcp-discovery/src/discovery-server.ts b/packages/dvmcp-discovery/src/discovery-server.ts index fc48b68..2d41bf4 100644 --- a/packages/dvmcp-discovery/src/discovery-server.ts +++ b/packages/dvmcp-discovery/src/discovery-server.ts @@ -31,12 +31,16 @@ export class DiscoveryServer { }); this.toolRegistry = new ToolRegistry(this.mcpServer); - this.toolExecutor = new ToolExecutor(this.relayHandler, this.keyManager); + this.toolExecutor = new ToolExecutor( + this.relayHandler, + this.keyManager, + this.toolRegistry + ); this.toolRegistry.setExecutionCallback(async (toolId, args) => { const tool = this.toolRegistry.getTool(toolId); if (!tool) throw new Error('Tool not found'); - return this.toolExecutor.executeTool(tool, args); + return this.toolExecutor.executeTool(toolId, tool, args); }); } @@ -62,7 +66,7 @@ export class DiscoveryServer { for (const tool of announcement.tools) { const toolId = `${tool.name}:${event.pubkey.slice(0, 4)}`; - this.toolRegistry.registerTool(toolId, tool); + this.toolRegistry.registerTool(toolId, tool, event.pubkey); } } catch (error) { console.error('Error processing DVM announcement:', error); diff --git a/packages/dvmcp-discovery/src/tool-executor.ts b/packages/dvmcp-discovery/src/tool-executor.ts index 5fa24cc..da9b79b 100644 --- a/packages/dvmcp-discovery/src/tool-executor.ts +++ b/packages/dvmcp-discovery/src/tool-executor.ts @@ -2,6 +2,7 @@ import { type Event } from 'nostr-tools'; import { RelayHandler } from '@dvmcp/commons/nostr/relay-handler'; import { createKeyManager } from '@dvmcp/commons/nostr/key-manager'; import { type Tool } from '@modelcontextprotocol/sdk/types.js'; +import { ToolRegistry } from './tool-registry.js'; import { TOOL_REQUEST_KIND, TOOL_RESPONSE_KIND, @@ -19,12 +20,17 @@ export class ToolExecutor { constructor( private relayHandler: RelayHandler, - private keyManager: ReturnType + private keyManager: ReturnType, + private toolRegistry: ToolRegistry ) {} - public async executeTool(tool: Tool, params: unknown): Promise { + public async executeTool( + toolId: string, + tool: Tool, + params: unknown + ): Promise { return new Promise((resolve, reject) => { - const request = this.createToolRequest(tool, params); + const request = this.createToolRequest(toolId, tool, params); const executionId = request.id; const context = this.createExecutionContext(executionId); @@ -102,7 +108,11 @@ export class ToolExecutor { } } - private createToolRequest(tool: Tool, params: unknown): Event { + private createToolRequest( + toolId: string, + tool: Tool, + params: unknown + ): Event { const request = this.keyManager.createEventTemplate(TOOL_REQUEST_KIND); const parameters = @@ -111,12 +121,17 @@ export class ToolExecutor { ? {} : params; + const toolInfo = this.toolRegistry.getToolInfo(toolId); + if (!toolInfo) throw new Error(`Tool ${toolId} not found`); + request.content = JSON.stringify({ name: tool.name, parameters, + providerPubkey: toolInfo.providerPubkey, }); request.tags.push(['c', 'execute-tool']); + request.tags.push(['p', toolInfo.providerPubkey]); return this.keyManager.signEvent(request); } } diff --git a/packages/dvmcp-discovery/src/tool-registry.ts b/packages/dvmcp-discovery/src/tool-registry.ts index 15a5b8b..69589b8 100644 --- a/packages/dvmcp-discovery/src/tool-registry.ts +++ b/packages/dvmcp-discovery/src/tool-registry.ts @@ -4,14 +4,19 @@ import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'; import { z } from 'zod'; export class ToolRegistry { - private discoveredTools: Map = new Map(); + private discoveredTools: Map = + new Map(); constructor(private mcpServer: McpServer) {} - public registerTool(toolId: string, tool: Tool): void { + public registerTool( + toolId: string, + tool: Tool, + providerPubkey: string + ): void { try { ToolSchema.parse(tool); - this.discoveredTools.set(toolId, tool); + this.discoveredTools.set(toolId, { tool, providerPubkey }); this.registerWithMcp(toolId, tool); } catch (error) { console.error(`Invalid MCP tool format for ${toolId}:`, error); @@ -19,12 +24,16 @@ export class ToolRegistry { } } - public getTool(toolId: string): Tool | undefined { + public getToolInfo(toolId: string) { return this.discoveredTools.get(toolId); } + public getTool(toolId: string): Tool | undefined { + return this.discoveredTools.get(toolId)?.tool; + } + public listTools(): Tool[] { - return Array.from(this.discoveredTools.values()); + return Array.from(this.discoveredTools.values()).map(({ tool }) => tool); } public clear(): void {