1  import type { Client } from '@modelcontextprotocol/sdk/client/index.js'
  2  import {
  3    ElicitationCompleteNotificationSchema,
  4    type ElicitRequestParams,
  5    ElicitRequestSchema,
  6    type ElicitResult,
  7  } from '@modelcontextprotocol/sdk/types.js'
  8  import type { AppState } from '../../state/AppState.js'
  9  import {
 10    executeElicitationHooks,
 11    executeElicitationResultHooks,
 12    executeNotificationHooks,
 13  } from '../../utils/hooks.js'
 14  import { logMCPDebug, logMCPError } from '../../utils/log.js'
 15  import { jsonStringify } from '../../utils/slowOperations.js'
 16  import {
 17    type AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS,
 18    logEvent,
 19  } from '../analytics/index.js'
 20  
 21  /** Configuration for the waiting state shown after the user opens a URL. */
 22  export type ElicitationWaitingState = {
 23    /** Button label, e.g. "Retry now" or "Skip confirmation" */
 24    actionLabel: string
 25    /** Whether to show a visible Cancel button (e.g. for error-based retry flow) */
 26    showCancel?: boolean
 27  }
 28  
 29  export type ElicitationRequestEvent = {
 30    serverName: string
 31    /** The JSON-RPC request ID, unique per server connection. */
 32    requestId: string | number
 33    params: ElicitRequestParams
 34    signal: AbortSignal
 35    /**
 36     * Resolves the elicitation. For explicit elicitations, all actions are
 37     * meaningful. For error-based retry (-32042), 'accept' is a no-op —
 38     * the retry is driven by onWaitingDismiss instead.
 39     */
 40    respond: (response: ElicitResult) => void
 41    /** For URL elicitations: shown after user opens the browser. */
 42    waitingState?: ElicitationWaitingState
 43    /** Called when phase 2 (waiting) is dismissed by user action or completion. */
 44    onWaitingDismiss?: (action: 'dismiss' | 'retry' | 'cancel') => void
 45    /** Set to true by the completion notification handler when the server confirms completion. */
 46    completed?: boolean
 47  }
 48  
 49  function getElicitationMode(params: ElicitRequestParams): 'form' | 'url' {
 50    return params.mode === 'url' ? 'url' : 'form'
 51  }
 52  
 53  /** Find a queued elicitation event by server name and elicitationId. */
 54  function findElicitationInQueue(
 55    queue: ElicitationRequestEvent[],
 56    serverName: string,
 57    elicitationId: string,
 58  ): number {
 59    return queue.findIndex(
 60      e =>
 61        e.serverName === serverName &&
 62        e.params.mode === 'url' &&
 63        'elicitationId' in e.params &&
 64        e.params.elicitationId === elicitationId,
 65    )
 66  }
 67  
 68  export function registerElicitationHandler(
 69    client: Client,
 70    serverName: string,
 71    setAppState: (f: (prevState: AppState) => AppState) => void,
 72  ): void {
 73    // Register the elicitation request handler.
 74    // Wrapped in try/catch because setRequestHandler throws if the client wasn't
 75    // created with elicitation capability declared.
 76    try {
 77      client.setRequestHandler(ElicitRequestSchema, async (request, extra) => {
 78        logMCPDebug(
 79          serverName,
 80          `Received elicitation request: ${jsonStringify(request)}`,
 81        )
 82  
 83        const mode = getElicitationMode(request.params)
 84  
 85        logEvent('tengu_mcp_elicitation_shown', {
 86          mode: mode as AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS,
 87        })
 88  
 89        try {
 90          // Run elicitation hooks first - they can provide a response programmatically
 91          const hookResponse = await runElicitationHooks(
 92            serverName,
 93            request.params,
 94            extra.signal,
 95          )
 96          if (hookResponse) {
 97            logMCPDebug(
 98              serverName,
 99              `Elicitation resolved by hook: ${jsonStringify(hookResponse)}`,
100            )
101            logEvent('tengu_mcp_elicitation_response', {
102              mode: mode as AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS,
103              action:
104                hookResponse.action as AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS,
105            })
106            return hookResponse
107          }
108  
109          const elicitationId =
110            mode === 'url' && 'elicitationId' in request.params
111              ? (request.params.elicitationId as string | undefined)
112              : undefined
113  
114          const response = new Promise<ElicitResult>(resolve => {
115            const onAbort = () => {
116              resolve({ action: 'cancel' })
117            }
118  
119            if (extra.signal.aborted) {
120              onAbort()
121              return
122            }
123  
124            const waitingState: ElicitationWaitingState | undefined =
125              elicitationId ? { actionLabel: 'Skip confirmation' } : undefined
126  
127            setAppState(prev => ({
128              ...prev,
129              elicitation: {
130                queue: [
131                  ...prev.elicitation.queue,
132                  {
133                    serverName,
134                    requestId: extra.requestId,
135                    params: request.params,
136                    signal: extra.signal,
137                    waitingState,
138                    respond: (result: ElicitResult) => {
139                      extra.signal.removeEventListener('abort', onAbort)
140                      logEvent('tengu_mcp_elicitation_response', {
141                        mode: mode as AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS,
142                        action:
143                          result.action as AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS,
144                      })
145                      resolve(result)
146                    },
147                  },
148                ],
149              },
150            }))
151  
152            extra.signal.addEventListener('abort', onAbort, { once: true })
153          })
154          const rawResult = await response
155          logMCPDebug(
156            serverName,
157            `Elicitation response: ${jsonStringify(rawResult)}`,
158          )
159          const result = await runElicitationResultHooks(
160            serverName,
161            rawResult,
162            extra.signal,
163            mode,
164            elicitationId,
165          )
166          return result
167        } catch (error) {
168          logMCPError(serverName, `Elicitation error: ${error}`)
169          return { action: 'cancel' as const }
170        }
171      })
172  
173      // Register handler for elicitation completion notifications (URL mode).
174      // Sets `completed: true` on the matching queue event; the dialog reacts to this flag.
175      client.setNotificationHandler(
176        ElicitationCompleteNotificationSchema,
177        notification => {
178          const { elicitationId } = notification.params
179          logMCPDebug(
180            serverName,
181            `Received elicitation completion notification: ${elicitationId}`,
182          )
183          void executeNotificationHooks({
184            message: `MCP server "${serverName}" confirmed elicitation ${elicitationId} complete`,
185            notificationType: 'elicitation_complete',
186          })
187          let found = false
188          setAppState(prev => {
189            const idx = findElicitationInQueue(
190              prev.elicitation.queue,
191              serverName,
192              elicitationId,
193            )
194            if (idx === -1) return prev
195            found = true
196            const queue = [...prev.elicitation.queue]
197            queue[idx] = { ...queue[idx]!, completed: true }
198            return { ...prev, elicitation: { queue } }
199          })
200          if (!found) {
201            logMCPDebug(
202              serverName,
203              `Ignoring completion notification for unknown elicitation: ${elicitationId}`,
204            )
205          }
206        },
207      )
208    } catch {
209      // Client wasn't created with elicitation capability - nothing to register
210      return
211    }
212  }
213  
214  export async function runElicitationHooks(
215    serverName: string,
216    params: ElicitRequestParams,
217    signal: AbortSignal,
218  ): Promise<ElicitResult | undefined> {
219    try {
220      const mode = params.mode === 'url' ? 'url' : 'form'
221      const url = 'url' in params ? (params.url as string) : undefined
222      const elicitationId =
223        'elicitationId' in params
224          ? (params.elicitationId as string | undefined)
225          : undefined
226  
227      const { elicitationResponse, blockingError } =
228        await executeElicitationHooks({
229          serverName,
230          message: params.message,
231          requestedSchema:
232            'requestedSchema' in params
233              ? (params.requestedSchema as Record<string, unknown>)
234              : undefined,
235          signal,
236          mode,
237          url,
238          elicitationId,
239        })
240  
241      if (blockingError) {
242        return { action: 'decline' }
243      }
244  
245      if (elicitationResponse) {
246        return {
247          action: elicitationResponse.action,
248          content: elicitationResponse.content,
249        }
250      }
251  
252      return undefined
253    } catch (error) {
254      logMCPError(serverName, `Elicitation hook error: ${error}`)
255      return undefined
256    }
257  }
258  
259  /**
260   * Run ElicitationResult hooks after the user has responded, then fire a
261   * `elicitation_response` notification. Returns a (potentially modified)
262   * ElicitResult — hooks may override the action/content or block the response.
263   */
264  export async function runElicitationResultHooks(
265    serverName: string,
266    result: ElicitResult,
267    signal: AbortSignal,
268    mode?: 'form' | 'url',
269    elicitationId?: string,
270  ): Promise<ElicitResult> {
271    try {
272      const { elicitationResultResponse, blockingError } =
273        await executeElicitationResultHooks({
274          serverName,
275          action: result.action,
276          content: result.content as Record<string, unknown> | undefined,
277          signal,
278          mode,
279          elicitationId,
280        })
281  
282      if (blockingError) {
283        void executeNotificationHooks({
284          message: `Elicitation response for server "${serverName}": decline`,
285          notificationType: 'elicitation_response',
286        })
287        return { action: 'decline' }
288      }
289  
290      const finalResult = elicitationResultResponse
291        ? {
292            action: elicitationResultResponse.action,
293            content: elicitationResultResponse.content ?? result.content,
294          }
295        : result
296  
297      // Fire a notification for observability
298      void executeNotificationHooks({
299        message: `Elicitation response for server "${serverName}": ${finalResult.action}`,
300        notificationType: 'elicitation_response',
301      })
302  
303      return finalResult
304    } catch (error) {
305      logMCPError(serverName, `ElicitationResult hook error: ${error}`)
306      // Fire notification even on error
307      void executeNotificationHooks({
308        message: `Elicitation response for server "${serverName}": ${result.action}`,
309        notificationType: 'elicitation_response',
310      })
311      return result
312    }
313  }