1  import { useEffect, useRef } from 'react'
 2  import { logError } from 'src/utils/log.js'
 3  import { z } from 'zod/v4'
 4  import type {
 5    ConnectedMCPServer,
 6    MCPServerConnection,
 7  } from '../services/mcp/types.js'
 8  import { getConnectedIdeClient } from '../utils/ide.js'
 9  import { lazySchema } from '../utils/lazySchema.js'
10  export type IDEAtMentioned = {
11    filePath: string
12    lineStart?: number
13    lineEnd?: number
14  }
15  
16  const NOTIFICATION_METHOD = 'at_mentioned'
17  
18  const AtMentionedSchema = lazySchema(() =>
19    z.object({
20      method: z.literal(NOTIFICATION_METHOD),
21      params: z.object({
22        filePath: z.string(),
23        lineStart: z.number().optional(),
24        lineEnd: z.number().optional(),
25      }),
26    }),
27  )
28  
29  /**
30   * A hook that tracks IDE at-mention notifications by directly registering
31   * with MCP client notification handlers,
32   */
33  export function useIdeAtMentioned(
34    mcpClients: MCPServerConnection[],
35    onAtMentioned: (atMentioned: IDEAtMentioned) => void,
36  ): void {
37    const ideClientRef = useRef<ConnectedMCPServer | undefined>(undefined)
38  
39    useEffect(() => {
40      // Find the IDE client from the MCP clients list
41      const ideClient = getConnectedIdeClient(mcpClients)
42  
43      if (ideClientRef.current !== ideClient) {
44        ideClientRef.current = ideClient
45      }
46  
47      // If we found a connected IDE client, register our handler
48      if (ideClient) {
49        ideClient.client.setNotificationHandler(
50          AtMentionedSchema(),
51          notification => {
52            if (ideClientRef.current !== ideClient) {
53              return
54            }
55            try {
56              const data = notification.params
57              // Adjust line numbers to be 1-based instead of 0-based
58              const lineStart =
59                data.lineStart !== undefined ? data.lineStart + 1 : undefined
60              const lineEnd =
61                data.lineEnd !== undefined ? data.lineEnd + 1 : undefined
62              onAtMentioned({
63                filePath: data.filePath,
64                lineStart: lineStart,
65                lineEnd: lineEnd,
66              })
67            } catch (error) {
68              logError(error as Error)
69            }
70          },
71        )
72      }
73  
74      // No cleanup needed as MCP clients manage their own lifecycle
75    }, [mcpClients, onAtMentioned])
76  }