1  import { ipcMain, shell, IpcMainEvent, dialog, BrowserWindow } from 'electron'
  2  import Constants from './utils/Constants'
  3  import {
  4    McpServerCapabilitySchemas,
  5    McpClientObject,
  6    McpFeatureObject,
  7    McpMetadataConfig,
  8    McpProgressCallback
  9  } from './mcp/types'
 10  
 11  import { manageRequests } from './mcp/client'
 12  
 13  import { spawn } from 'child_process'
 14  import { existsSync, mkdirSync, writeFileSync, readdirSync } from 'fs'
 15  import { join, resolve, normalize } from 'path'
 16  import { pathToFileURL } from 'url'
 17  
 18  import { initClients } from './mcp/init'
 19  import { disconnect } from './mcp/connection'
 20  import { loadConfig } from './mcp/init'
 21  import { loadLlmFile } from './mcp/config'
 22  import { unpackDxt, getManifest } from './mcp/dxt'
 23  import { McpbManifestAny } from '@anthropic-ai/mcpb'
 24  
 25  import { closeCommandPicker } from './aid/commands'
 26  
 27  import { commandSelectionInvoke, mcpServersProcessCallback } from './index'
 28  import { getCachedText } from './aid/utils'
 29  import { McpClientResponse, CommandResponse, McpInitResponse } from './types'
 30  
 31  import { IpcFileTransferRequest, IpcFileTransferResponse } from '@/types/ipc'
 32  
 33  const handlerRegistry = new Map<string, Function>()
 34  
 35  interface ManifestResponse {
 36    status: 'success' | 'error'
 37    result?: Record<string, McpbManifestAny> // Object with string keys and DXT values
 38    error?: string
 39  }
 40  
 41  /*
 42   * IPC Communications
 43   * */
 44  export default class IPCs {
 45    static clients: McpClientObject[] = []
 46    static currentFeatures: McpFeatureObject[] = []
 47  
 48    static initialize(): void {
 49      // Get application version
 50      ipcMain.handle('msgRequestAppInfo', () => {
 51        return {
 52          version: Constants.APP_VERSION,
 53          homepage: Constants.APP_HOME_PAGE,
 54          platform: process.platform
 55        }
 56      })
 57  
 58      ipcMain.handle('msgRequestGetDxtUrl', () => {
 59        return pathToFileURL(normalize(resolve(Constants.ASSETS_PATH.mcpb))).toString()
 60      })
 61  
 62      ipcMain.handle('msgMcpServersStop', async () => {
 63        IPCs.stopAllServers()
 64  
 65        const configs = await loadConfig()
 66  
 67        const features = configs.map((params) => {
 68          return registerIpcHandlers(params)
 69        })
 70  
 71        this.currentFeatures = features
 72  
 73        return true
 74      })
 75  
 76      ipcMain.handle(
 77        'msgMcpServersInit',
 78        async (_event: IpcMainEvent, metadata: McpMetadataConfig): Promise<McpInitResponse> => {
 79          IPCs.stopAllServers()
 80  
 81          const progressCallback: McpProgressCallback = (name, message, status) => {
 82            mcpServersProcessCallback({ name, message, status })
 83          }
 84  
 85          const configs = await loadConfig()
 86  
 87          try {
 88            const newClients = await initClients(metadata, progressCallback)
 89            const activeClientNames = newClients.map((client) => client.name)
 90            const inactiveConfigs = configs.filter(
 91              (config) => !activeClientNames.includes(config.name)
 92            )
 93  
 94            const features = [
 95              ...newClients.map((params) => registerIpcHandlers(params)),
 96              ...inactiveConfigs.map((params) => registerIpcHandlers(params))
 97            ]
 98            IPCs.updateMCP(features)
 99            this.clients = newClients
100            return {
101              status: 'success'
102            }
103          } catch (error) {
104            const features = configs.map((params) => {
105              return registerIpcHandlers(params)
106            })
107  
108            IPCs.updateMCP(features)
109  
110            return {
111              status: 'error',
112              error: error instanceof Error ? error.message : String(error)
113            }
114          }
115        }
116      )
117  
118      // Open url via web browser
119      ipcMain.on('msgOpenExternalLink', async (event: IpcMainEvent, url: string) => {
120        await shell.openExternal(url)
121      })
122  
123      ipcMain.on('msgOpenDxtFilePath', async (event: IpcMainEvent, name: string) => {
124        shell.openPath(resolve(join(Constants.ASSETS_PATH.mcpb, name)))
125      })
126  
127      ipcMain.on('msgOpenPath', async (event: IpcMainEvent, name: string) => {
128        shell.openPath(resolve(join(Constants.ASSETS_PATH[name])))
129      })
130  
131      ipcMain.on('msgWindowReload', async (event: IpcMainEvent) => {
132        BrowserWindow.fromWebContents(event.sender).reload()
133      })
134  
135      ipcMain.handle('msgGetApiToken', async (event, cli) => {
136        return new Promise((resolve, reject) => {
137          const child = spawn(cli, { shell: true })
138          const cleanup = () => {
139            child.stdout?.destroy()
140            child.stderr?.destroy()
141            if (!child.killed) {
142              child.kill('SIGKILL')
143            }
144          }
145          try {
146            let stdoutData = ''
147  
148            child.stdout?.on('data', (data) => {
149              const output = data.toString()
150              stdoutData += output
151              event.sender.send('renderListenStdioProgress', output.trim()) // send real-time output
152            })
153  
154            child.stderr?.on('data', (data) => {
155              console.error('Error output:', data.toString())
156              reject(data.toString())
157            })
158  
159            child.on('close', (code) => {
160              if (code === 0) {
161                resolve(stdoutData.trim().split('\n').at(-1))
162              } else {
163                reject(new Error(`Process exited with code ${code}`))
164              }
165            })
166  
167            setTimeout(() => {
168              cleanup()
169              reject(new Error('Process timeout'))
170            }, 30000)
171          } catch (error) {
172            console.error('Error fetching token:', error.message)
173            cleanup()
174            reject(error)
175          }
176        })
177      })
178  
179      // Open file
180      ipcMain.handle('msgOpenFile', async (event: IpcMainEvent, filter: string) => {
181        const filters = []
182        if (filter === 'text') {
183          filters.push({ name: 'Text', extensions: ['txt', 'json'] })
184        } else if (filter === 'zip') {
185          filters.push({ name: 'Zip', extensions: ['zip'] })
186        }
187        const dialogResult = await dialog.showOpenDialog({
188          properties: ['openFile'],
189          filters
190        })
191        return dialogResult
192      })
193  
194      ipcMain.on(
195        'msgFileTransferRequest',
196        async (event: IpcMainEvent, { name, data }: IpcFileTransferRequest) => {
197          try {
198            const buffer = Buffer.from(data)
199            const saveOption = Constants.getDxtSource(name)
200            const filePath = saveOption.mcpbPath
201            const dirPath = saveOption.outputDir
202            if (!existsSync(dirPath)) {
203              mkdirSync(dirPath, { recursive: true })
204            }
205            console.log('MCP bundle to be saved in: ', filePath)
206  
207            writeFileSync(filePath, buffer, { encoding: null })
208  
209            console.log(saveOption)
210            await unpackDxt(saveOption)
211            // console.log(getManifest(dirPath))
212  
213            event.reply('msgFileTransferResponse', {
214              name,
215              success: true,
216              path: saveOption.outputDir
217            } as IpcFileTransferResponse)
218          } catch (err) {
219            event.reply('msgFileTransferResponse', {
220              name,
221              success: false,
222              reason: err.message
223            } as IpcFileTransferResponse)
224          }
225        }
226      )
227  
228      ipcMain.on(
229        'msgCommandSelectionResult',
230        async (_event: IpcMainEvent, response: CommandResponse) => {
231          closeCommandPicker()
232          const prompt = response.prompt
233          const request = { prompt, input: getCachedText(response.id) }
234          commandSelectionInvoke(request)
235        }
236      )
237  
238      ipcMain.handle('list-manifests', async (_event: IpcMainEvent): Promise<ManifestResponse> => {
239        const mcpbPath = Constants.ASSETS_PATH.mcpb
240  
241        try {
242          const entries = readdirSync(mcpbPath, { withFileTypes: true })
243          console.log(entries)
244  
245          // Transform the array into an object
246          const manifestsObject = entries
247            .filter((dirent) => dirent.isDirectory())
248            .reduce(
249              (acc, dirent) => {
250                acc[dirent.name] = getManifest(join(mcpbPath, dirent.name))
251                return acc
252              },
253              {} as Record<string, any>
254            ) // You can replace 'any' with your manifest type
255  
256          return {
257            status: 'success',
258            result: manifestsObject
259          }
260        } catch (err) {
261          return {
262            status: 'error',
263            error: err.message
264          }
265        }
266      })
267  
268      ipcMain.handle('list-llms', () => {
269        return loadLlmFile(Constants.ASSETS_PATH.llm)
270      })
271  
272      ipcMain.handle('list-popups', () => {
273        return loadLlmFile(Constants.ASSETS_PATH.popup)
274      })
275  
276      ipcMain.handle('list-startups', () => {
277        return loadLlmFile(Constants.ASSETS_PATH.startup)
278      })
279    }
280  
281    static updateMCP(newFeatures: McpFeatureObject[]): void {
282      this.currentFeatures = newFeatures
283    }
284  
285    static stopAllServers() {
286      this.clients.forEach((client: McpClientObject) => {
287        if (client.connection?.transport) {
288          disconnect(client.connection.transport)
289          delete client.connection.transport
290        }
291      })
292  
293      IPCs.removeAllHandlers()
294    }
295  
296    static removeAllHandlers() {
297      for (const [eventName] of handlerRegistry) {
298        ipcMain.removeHandler(eventName)
299      }
300      console.log(`handlerRegistry clear: ${handlerRegistry.size}`)
301      handlerRegistry.clear()
302    }
303  
304    static initializeMCP(initialFeatures: McpFeatureObject[]): void {
305      this.currentFeatures = initialFeatures
306      ipcMain.handle('list-clients', () => {
307        return this.currentFeatures
308      })
309    }
310  }
311  
312  export function listenOnceForRendererResponse(
313    responseChannel: string,
314    resolve: (_value: McpClientResponse) => void
315  ) {
316    ipcMain.once(responseChannel, (_event, response: McpClientResponse) => {
317      resolve(response)
318    })
319  }
320  
321  export function registerIpcHandlers({
322    name,
323    connection,
324    configJson
325  }: McpClientObject): McpFeatureObject {
326    const feature: McpFeatureObject = {
327      name,
328      config: configJson
329    }
330  
331    if (!connection) {
332      return feature
333    }
334  
335    const registerHandler = (method: string, schema: any) => {
336      const eventName = `${name}-${method}`
337      console.log(`IPC Main ${eventName}`)
338      const handler = async (
339        _event: Electron.IpcMainInvokeEvent,
340        request: { method: string; params: any }
341      ) => {
342        if (request?.method !== method) {
343          console.log(
344            `Request method not registered: ${request?.method}, fallback to use ${method}, please double check the invoker in Renderer.`
345          )
346        }
347        return await manageRequests(connection.client, `${method}`, schema, request?.params)
348      }
349      ipcMain.handle(eventName, handler)
350      handlerRegistry.set(eventName, handler)
351      return eventName
352    }
353  
354    for (const [type, actions] of Object.entries(McpServerCapabilitySchemas)) {
355      const capabilities = connection.client.getServerCapabilities()
356      if (capabilities?.[type]) {
357        feature[type] = {}
358        for (const [action, schema] of Object.entries(actions)) {
359          feature[type][action] = registerHandler(`${type}/${action}`, schema)
360        }
361      }
362    }
363  
364    const instructions = connection.client.getInstructions()
365  
366    const implementation = connection.client.getServerVersion()
367  
368    return {
369      ...feature,
370      description: {
371        instructions,
372        implementation
373      }
374    }
375  }