1  import { Server } from '@modelcontextprotocol/sdk/server/index.js'
  2  import { StdioServerTransport } from '@modelcontextprotocol/sdk/server/stdio.js'
  3  import {
  4    CallToolRequestSchema,
  5    CallToolResultSchema,
  6    ListToolsRequestSchema,
  7    ListToolsResultSchema,
  8    ToolSchema,
  9  } from '@modelcontextprotocol/sdk/types.js'
 10  import { z } from 'zod'
 11  import { zodToJsonSchema } from 'zod-to-json-schema'
 12  import { AgentTool } from '../tools/AgentTool/AgentTool.js'
 13  import { hasPermissionsToUseTool } from '../permissions.js'
 14  import { setCwd } from '../utils/state.js'
 15  import { getSlowAndCapableModel } from '../utils/model.js'
 16  import { logError } from '../utils/log.js'
 17  import { LSTool } from '../tools/lsTool/lsTool.js'
 18  import { BashTool } from '../tools/BashTool/BashTool.js'
 19  import { FileEditTool } from '../tools/FileEditTool/FileEditTool.js'
 20  import { FileReadTool } from '../tools/FileReadTool/FileReadTool.js'
 21  import { GlobTool } from '../tools/GlobTool/GlobTool.js'
 22  import { GrepTool } from '../tools/GrepTool/GrepTool.js'
 23  import { FileWriteTool } from '../tools/FileWriteTool/FileWriteTool.js'
 24  import { Tool } from '../Tool.js'
 25  import { Command } from '../commands.js'
 26  import review from '../commands/review.js'
 27  import { lastX } from '../utils/generators.js'
 28  
 29  type ToolInput = z.infer<typeof ToolSchema.shape.inputSchema>
 30  
 31  const state: {
 32    readFileTimestamps: Record<string, number>
 33  } = {
 34    readFileTimestamps: {},
 35  }
 36  
 37  const MCP_COMMANDS: Command[] = [review]
 38  
 39  const MCP_TOOLS: Tool[] = [
 40    AgentTool,
 41    BashTool,
 42    FileEditTool,
 43    FileReadTool,
 44    GlobTool,
 45    GrepTool,
 46    FileWriteTool,
 47    LSTool,
 48  ]
 49  
 50  export async function startMCPServer(cwd: string): Promise<void> {
 51    await setCwd(cwd)
 52    const server = new Server(
 53      {
 54        name: 'claude/tengu',
 55        version: MACRO.VERSION,
 56      },
 57      {
 58        capabilities: {
 59          tools: {},
 60        },
 61      },
 62    )
 63  
 64    server.setRequestHandler(
 65      ListToolsRequestSchema,
 66      async (): Promise<Zod.infer<typeof ListToolsResultSchema>> => {
 67        const tools = await Promise.all(
 68          MCP_TOOLS.map(async tool => ({
 69            ...tool,
 70            description: await tool.description(z.object({})),
 71            inputSchema: zodToJsonSchema(tool.inputSchema) as ToolInput,
 72          })),
 73        )
 74  
 75        return {
 76          tools,
 77        }
 78      },
 79    )
 80  
 81    server.setRequestHandler(
 82      CallToolRequestSchema,
 83      async (request): Promise<Zod.infer<typeof CallToolResultSchema>> => {
 84        const { name, arguments: args } = request.params
 85        const tool = MCP_TOOLS.find(_ => _.name === name)
 86        if (!tool) {
 87          throw new Error(`Tool ${name} not found`)
 88        }
 89  
 90        // TODO: validate input types with zod
 91        try {
 92          if (!(await tool.isEnabled())) {
 93            throw new Error(`Tool ${name} is not enabled`)
 94          }
 95          const model = await getSlowAndCapableModel()
 96          const validationResult = await tool.validateInput?.(
 97            (args as never) ?? {},
 98            {
 99              abortController: new AbortController(),
100              options: {
101                commands: MCP_COMMANDS,
102                tools: MCP_TOOLS,
103                slowAndCapableModel: model,
104                forkNumber: 0,
105                messageLogName: 'unused',
106                maxThinkingTokens: 0,
107              },
108              messageId: undefined,
109              readFileTimestamps: state.readFileTimestamps,
110            },
111          )
112          if (validationResult && !validationResult.result) {
113            throw new Error(
114              `Tool ${name} input is invalid: ${validationResult.message}`,
115            )
116          }
117          const result = tool.call(
118            (args ?? {}) as never,
119            {
120              abortController: new AbortController(),
121              messageId: undefined,
122              options: {
123                commands: MCP_COMMANDS,
124                tools: MCP_TOOLS,
125                slowAndCapableModel: await getSlowAndCapableModel(),
126                forkNumber: 0,
127                messageLogName: 'unused',
128                maxThinkingTokens: 0,
129              },
130              readFileTimestamps: state.readFileTimestamps,
131            },
132            hasPermissionsToUseTool,
133          )
134  
135          const finalResult = await lastX(result)
136  
137          if (finalResult.type !== 'result') {
138            throw new Error(`Tool ${name} did not return a result`)
139          }
140  
141          return {
142            content: Array.isArray(finalResult)
143              ? finalResult.map(item => ({
144                  type: 'text' as const,
145                  text: 'text' in item ? item.text : JSON.stringify(item),
146                }))
147              : [
148                  {
149                    type: 'text' as const,
150                    text:
151                      typeof finalResult === 'string'
152                        ? finalResult
153                        : JSON.stringify(finalResult.data),
154                  },
155                ],
156          }
157        } catch (error) {
158          logError(error)
159          return {
160            isError: true,
161            content: [
162              {
163                type: 'text',
164                text: `Error: ${error instanceof Error ? error.message : String(error)}`,
165              },
166            ],
167          }
168        }
169      },
170    )
171  
172    async function runServer() {
173      const transport = new StdioServerTransport()
174      await server.connect(transport)
175    }
176  
177    return await runServer()
178  }