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 SelectionPoint = {
 11    line: number
 12    character: number
 13  }
 14  
 15  export type SelectionData = {
 16    selection: {
 17      start: SelectionPoint
 18      end: SelectionPoint
 19    } | null
 20    text?: string
 21    filePath?: string
 22  }
 23  
 24  export type IDESelection = {
 25    lineCount: number
 26    lineStart?: number
 27    text?: string
 28    filePath?: string
 29  }
 30  
 31  // Define the selection changed notification schema
 32  const SelectionChangedSchema = lazySchema(() =>
 33    z.object({
 34      method: z.literal('selection_changed'),
 35      params: z.object({
 36        selection: z
 37          .object({
 38            start: z.object({
 39              line: z.number(),
 40              character: z.number(),
 41            }),
 42            end: z.object({
 43              line: z.number(),
 44              character: z.number(),
 45            }),
 46          })
 47          .nullable()
 48          .optional(),
 49        text: z.string().optional(),
 50        filePath: z.string().optional(),
 51      }),
 52    }),
 53  )
 54  
 55  /**
 56   * A hook that tracks IDE text selection information by directly registering
 57   * with MCP client notification handlers
 58   */
 59  export function useIdeSelection(
 60    mcpClients: MCPServerConnection[],
 61    onSelect: (selection: IDESelection) => void,
 62  ): void {
 63    const handlersRegistered = useRef(false)
 64    const currentIDERef = useRef<ConnectedMCPServer | null>(null)
 65  
 66    useEffect(() => {
 67      // Find the IDE client from the MCP clients list
 68      const ideClient = getConnectedIdeClient(mcpClients)
 69  
 70      // If the IDE client changed, we need to re-register handlers.
 71      // Normalize undefined to null so the initial ref value (null) matches
 72      // "no IDE found" (undefined), avoiding spurious resets on every MCP update.
 73      if (currentIDERef.current !== (ideClient ?? null)) {
 74        handlersRegistered.current = false
 75        currentIDERef.current = ideClient || null
 76        // Reset the selection when the IDE client changes.
 77        onSelect({
 78          lineCount: 0,
 79          lineStart: undefined,
 80          text: undefined,
 81          filePath: undefined,
 82        })
 83      }
 84  
 85      // Skip if we've already registered handlers for the current IDE or if there's no IDE client
 86      if (handlersRegistered.current || !ideClient) {
 87        return
 88      }
 89  
 90      // Handler function for selection changes
 91      const selectionChangeHandler = (data: SelectionData) => {
 92        if (data.selection?.start && data.selection?.end) {
 93          const { start, end } = data.selection
 94          let lineCount = end.line - start.line + 1
 95          // If on the first character of the line, do not count the line
 96          // as being selected.
 97          if (end.character === 0) {
 98            lineCount--
 99          }
100          const selection = {
101            lineCount,
102            lineStart: start.line,
103            text: data.text,
104            filePath: data.filePath,
105          }
106  
107          onSelect(selection)
108        }
109      }
110  
111      // Register notification handler for selection_changed events
112      ideClient.client.setNotificationHandler(
113        SelectionChangedSchema(),
114        notification => {
115          if (currentIDERef.current !== ideClient) {
116            return
117          }
118  
119          try {
120            // Get the selection data from the notification params
121            const selectionData = notification.params
122  
123            // Process selection data - validate it has required properties
124            if (
125              selectionData.selection &&
126              selectionData.selection.start &&
127              selectionData.selection.end
128            ) {
129              // Handle selection changes
130              selectionChangeHandler(selectionData as SelectionData)
131            } else if (selectionData.text !== undefined) {
132              // Handle empty selection (when text is empty string)
133              selectionChangeHandler({
134                selection: null,
135                text: selectionData.text,
136                filePath: selectionData.filePath,
137              })
138            }
139          } catch (error) {
140            logError(error as Error)
141          }
142        },
143      )
144  
145      // Mark that we've registered handlers
146      handlersRegistered.current = true
147  
148      // No cleanup needed as MCP clients manage their own lifecycle
149    }, [mcpClients, onSelect])
150  }