1  import { useCallback, useEffect, useMemo, useRef } from 'react'
  2  import type { ToolUseConfirm } from '../components/permissions/PermissionRequest.js'
  3  import type { RemotePermissionResponse } from '../remote/RemoteSessionManager.js'
  4  import {
  5    createSyntheticAssistantMessage,
  6    createToolStub,
  7  } from '../remote/remotePermissionBridge.js'
  8  import {
  9    convertSDKMessage,
 10    isSessionEndMessage,
 11  } from '../remote/sdkMessageAdapter.js'
 12  import {
 13    type DirectConnectConfig,
 14    DirectConnectSessionManager,
 15  } from '../server/directConnectManager.js'
 16  import type { Tool } from '../Tool.js'
 17  import { findToolByName } from '../Tool.js'
 18  import type { Message as MessageType } from '../types/message.js'
 19  import type { PermissionAskDecision } from '../types/permissions.js'
 20  import { logForDebugging } from '../utils/debug.js'
 21  import { gracefulShutdown } from '../utils/gracefulShutdown.js'
 22  import type { RemoteMessageContent } from '../utils/teleport/api.js'
 23  
 24  type UseDirectConnectResult = {
 25    isRemoteMode: boolean
 26    sendMessage: (content: RemoteMessageContent) => Promise<boolean>
 27    cancelRequest: () => void
 28    disconnect: () => void
 29  }
 30  
 31  type UseDirectConnectProps = {
 32    config: DirectConnectConfig | undefined
 33    setMessages: React.Dispatch<React.SetStateAction<MessageType[]>>
 34    setIsLoading: (loading: boolean) => void
 35    setToolUseConfirmQueue: React.Dispatch<React.SetStateAction<ToolUseConfirm[]>>
 36    tools: Tool[]
 37  }
 38  
 39  export function useDirectConnect({
 40    config,
 41    setMessages,
 42    setIsLoading,
 43    setToolUseConfirmQueue,
 44    tools,
 45  }: UseDirectConnectProps): UseDirectConnectResult {
 46    const isRemoteMode = !!config
 47  
 48    const managerRef = useRef<DirectConnectSessionManager | null>(null)
 49    const hasReceivedInitRef = useRef(false)
 50    const isConnectedRef = useRef(false)
 51  
 52    // Keep a ref to tools so the WebSocket callback doesn't go stale
 53    const toolsRef = useRef(tools)
 54    useEffect(() => {
 55      toolsRef.current = tools
 56    }, [tools])
 57  
 58    useEffect(() => {
 59      if (!config) {
 60        return
 61      }
 62  
 63      hasReceivedInitRef.current = false
 64      logForDebugging(`[useDirectConnect] Connecting to ${config.wsUrl}`)
 65  
 66      const manager = new DirectConnectSessionManager(config, {
 67        onMessage: sdkMessage => {
 68          if (isSessionEndMessage(sdkMessage)) {
 69            setIsLoading(false)
 70          }
 71  
 72          // Skip duplicate init messages (server sends one per turn)
 73          if (sdkMessage.type === 'system' && sdkMessage.subtype === 'init') {
 74            if (hasReceivedInitRef.current) {
 75              return
 76            }
 77            hasReceivedInitRef.current = true
 78          }
 79  
 80          const converted = convertSDKMessage(sdkMessage, {
 81            convertToolResults: true,
 82          })
 83          if (converted.type === 'message') {
 84            setMessages(prev => [...prev, converted.message])
 85          }
 86        },
 87        onPermissionRequest: (request, requestId) => {
 88          logForDebugging(
 89            `[useDirectConnect] Permission request for tool: ${request.tool_name}`,
 90          )
 91  
 92          const tool =
 93            findToolByName(toolsRef.current, request.tool_name) ??
 94            createToolStub(request.tool_name)
 95  
 96          const syntheticMessage = createSyntheticAssistantMessage(
 97            request,
 98            requestId,
 99          )
100  
101          const permissionResult: PermissionAskDecision = {
102            behavior: 'ask',
103            message:
104              request.description ?? `${request.tool_name} requires permission`,
105            suggestions: request.permission_suggestions,
106            blockedPath: request.blocked_path,
107          }
108  
109          const toolUseConfirm: ToolUseConfirm = {
110            assistantMessage: syntheticMessage,
111            tool,
112            description:
113              request.description ?? `${request.tool_name} requires permission`,
114            input: request.input,
115            toolUseContext: {} as ToolUseConfirm['toolUseContext'],
116            toolUseID: request.tool_use_id,
117            permissionResult,
118            permissionPromptStartTimeMs: Date.now(),
119            onUserInteraction() {
120              // No-op for remote
121            },
122            onAbort() {
123              const response: RemotePermissionResponse = {
124                behavior: 'deny',
125                message: 'User aborted',
126              }
127              manager.respondToPermissionRequest(requestId, response)
128              setToolUseConfirmQueue(queue =>
129                queue.filter(item => item.toolUseID !== request.tool_use_id),
130              )
131            },
132            onAllow(updatedInput, _permissionUpdates, _feedback) {
133              const response: RemotePermissionResponse = {
134                behavior: 'allow',
135                updatedInput,
136              }
137              manager.respondToPermissionRequest(requestId, response)
138              setToolUseConfirmQueue(queue =>
139                queue.filter(item => item.toolUseID !== request.tool_use_id),
140              )
141              setIsLoading(true)
142            },
143            onReject(feedback?: string) {
144              const response: RemotePermissionResponse = {
145                behavior: 'deny',
146                message: feedback ?? 'User denied permission',
147              }
148              manager.respondToPermissionRequest(requestId, response)
149              setToolUseConfirmQueue(queue =>
150                queue.filter(item => item.toolUseID !== request.tool_use_id),
151              )
152            },
153            async recheckPermission() {
154              // No-op for remote
155            },
156          }
157  
158          setToolUseConfirmQueue(queue => [...queue, toolUseConfirm])
159          setIsLoading(false)
160        },
161        onConnected: () => {
162          logForDebugging('[useDirectConnect] Connected')
163          isConnectedRef.current = true
164        },
165        onDisconnected: () => {
166          logForDebugging('[useDirectConnect] Disconnected')
167          if (!isConnectedRef.current) {
168            // Never connected — connection failure (e.g. auth rejected)
169            process.stderr.write(
170              `\nFailed to connect to server at ${config.wsUrl}\n`,
171            )
172          } else {
173            // Was connected then lost — server process exited or network dropped
174            process.stderr.write('\nServer disconnected.\n')
175          }
176          isConnectedRef.current = false
177          void gracefulShutdown(1)
178          setIsLoading(false)
179        },
180        onError: error => {
181          logForDebugging(`[useDirectConnect] Error: ${error.message}`)
182        },
183      })
184  
185      managerRef.current = manager
186      manager.connect()
187  
188      return () => {
189        logForDebugging('[useDirectConnect] Cleanup - disconnecting')
190        manager.disconnect()
191        managerRef.current = null
192      }
193    }, [config, setMessages, setIsLoading, setToolUseConfirmQueue])
194  
195    const sendMessage = useCallback(
196      async (content: RemoteMessageContent): Promise<boolean> => {
197        const manager = managerRef.current
198        if (!manager) {
199          return false
200        }
201  
202        setIsLoading(true)
203  
204        return manager.sendMessage(content)
205      },
206      [setIsLoading],
207    )
208  
209    // Cancel the current request
210    const cancelRequest = useCallback(() => {
211      // Send interrupt signal to the server
212      managerRef.current?.sendInterrupt()
213  
214      setIsLoading(false)
215    }, [setIsLoading])
216  
217    const disconnect = useCallback(() => {
218      managerRef.current?.disconnect()
219      managerRef.current = null
220      isConnectedRef.current = false
221    }, [])
222  
223    // Same stability concern as useRemoteSession — memoize so consumers
224    // that depend on the result object don't see a fresh reference per render.
225    return useMemo(
226      () => ({ isRemoteMode, sendMessage, cancelRequest, disconnect }),
227      [isRemoteMode, sendMessage, cancelRequest, disconnect],
228    )
229  }