1  import {
  2    CreateMessageRequestSchema as SamplingRequestSchema,
  3    ElicitRequestSchema
  4  } from '@modelcontextprotocol/sdk/types.js'
  5  import { Client } from '@modelcontextprotocol/sdk/client/index.js'
  6  import { StdioClientTransport } from '@modelcontextprotocol/sdk/client/stdio.js'
  7  
  8  import { McpServerConfig, McpClientTransport, McpProgressCallback } from './types'
  9  import { connect } from './connection'
 10  
 11  import { samplingTransferInvoke, elicitationTransferInvoke } from '../index'
 12  import Constants from '../utils/Constants'
 13  
 14  export async function initializeClient(
 15    name: string,
 16    serverConfig: McpServerConfig,
 17    callback?: McpProgressCallback,
 18    idleTimeout: number = 90 // 90 sec by default
 19  ): Promise<McpClientTransport> {
 20    let idleTimer: NodeJS.Timeout
 21    let rejectFn: (_reason: Error) => void
 22  
 23    const resetTimer = () => {
 24      clearTimeout(idleTimer)
 25      idleTimer = setTimeout(() => {
 26        rejectFn(
 27          new Error(
 28            `Initialization of client for ${name} timed out after ${idleTimeout} seconds of inactivity`
 29          )
 30        )
 31      }, idleTimeout * 1000)
 32    }
 33  
 34    const timeoutPromise = new Promise<McpClientTransport>((_, reject) => {
 35      rejectFn = reject
 36      resetTimer() // Init timer
 37    })
 38  
 39    const stdioPromise = initializeStdioClient(name, serverConfig, callback, resetTimer)
 40  
 41    return Promise.race([stdioPromise, timeoutPromise])
 42  }
 43  
 44  async function initializeStdioClient(
 45    name: string,
 46    config: McpServerConfig,
 47    callback?: McpProgressCallback,
 48    onData?: () => void
 49  ): Promise<McpClientTransport> {
 50    const transport = new StdioClientTransport({
 51      ...config,
 52      stderr: 'pipe'
 53    })
 54    const clientName = `${name}-client`
 55    const client = new Client(
 56      {
 57        name: clientName,
 58        version: Constants.APP_VERSION
 59      },
 60      {
 61        capabilities: {
 62          sampling: {},
 63          elicitation: {}
 64        }
 65      }
 66    )
 67  
 68    if (transport.stderr) {
 69      transport.stderr.on('data', (chunk) => {
 70        if (onData) onData()
 71        if (callback) callback(name, chunk.toString(), 'pending')
 72      })
 73    }
 74  
 75    try {
 76      callback(name, 'Staring...', 'pending')
 77      await connect(client, transport)
 78      console.log(`${clientName} connected.`)
 79      callback(name, 'Done', 'success')
 80    } catch (error) {
 81      const errorMsg = error instanceof Error ? error.message : String(error)
 82      callback(name, errorMsg, 'error')
 83      throw error
 84    }
 85  
 86    client.setRequestHandler(SamplingRequestSchema, async (request) => {
 87      console.log('Sampling request received:\n', request)
 88      const response = await samplingTransferInvoke(request)
 89      console.log(response)
 90      return response
 91    })
 92  
 93    client.setRequestHandler(ElicitRequestSchema, async (request) => {
 94      const response = await elicitationTransferInvoke(request)
 95  
 96      console.log('Elicitation request received:\n', JSON.stringify(response, null, 2))
 97  
 98      return response
 99    })
100  
101    return { client, transport }
102  }
103  
104  export async function manageRequests(client: Client, method: string, schema: any, params?: any) {
105    const requestObject = {
106      method,
107      ...(params && { params })
108    }
109  
110    const result = await client.request(requestObject, schema)
111    console.log(result)
112    return result
113  }