1  /**
  2   * REPL integration hook for `claude ssh` sessions.
  3   *
  4   * Sibling to useDirectConnect — same shape (isRemoteMode/sendMessage/
  5   * cancelRequest/disconnect), same REPL wiring, but drives an SSH child
  6   * process instead of a WebSocket. Kept separate rather than generalizing
  7   * useDirectConnect because the lifecycle differs: the ssh process and auth
  8   * proxy are created BEFORE this hook runs (during startup, in main.tsx) and
  9   * handed in; useDirectConnect creates its WebSocket inside the effect.
 10   */
 11  
 12  import { randomUUID } from 'crypto'
 13  import { useCallback, useEffect, useMemo, useRef } from 'react'
 14  import type { ToolUseConfirm } from '../components/permissions/PermissionRequest.js'
 15  import {
 16    createSyntheticAssistantMessage,
 17    createToolStub,
 18  } from '../remote/remotePermissionBridge.js'
 19  import {
 20    convertSDKMessage,
 21    isSessionEndMessage,
 22  } from '../remote/sdkMessageAdapter.js'
 23  import type { SSHSession } from '../ssh/createSSHSession.js'
 24  import type { SSHSessionManager } from '../ssh/SSHSessionManager.js'
 25  import type { Tool } from '../Tool.js'
 26  import { findToolByName } from '../Tool.js'
 27  import type { Message as MessageType } from '../types/message.js'
 28  import type { PermissionAskDecision } from '../types/permissions.js'
 29  import { logForDebugging } from '../utils/debug.js'
 30  import { gracefulShutdown } from '../utils/gracefulShutdown.js'
 31  import type { RemoteMessageContent } from '../utils/teleport/api.js'
 32  
 33  type UseSSHSessionResult = {
 34    isRemoteMode: boolean
 35    sendMessage: (content: RemoteMessageContent) => Promise<boolean>
 36    cancelRequest: () => void
 37    disconnect: () => void
 38  }
 39  
 40  type UseSSHSessionProps = {
 41    session: SSHSession | undefined
 42    setMessages: React.Dispatch<React.SetStateAction<MessageType[]>>
 43    setIsLoading: (loading: boolean) => void
 44    setToolUseConfirmQueue: React.Dispatch<React.SetStateAction<ToolUseConfirm[]>>
 45    tools: Tool[]
 46  }
 47  
 48  export function useSSHSession({
 49    session,
 50    setMessages,
 51    setIsLoading,
 52    setToolUseConfirmQueue,
 53    tools,
 54  }: UseSSHSessionProps): UseSSHSessionResult {
 55    const isRemoteMode = !!session
 56  
 57    const managerRef = useRef<SSHSessionManager | null>(null)
 58    const hasReceivedInitRef = useRef(false)
 59    const isConnectedRef = useRef(false)
 60  
 61    const toolsRef = useRef(tools)
 62    useEffect(() => {
 63      toolsRef.current = tools
 64    }, [tools])
 65  
 66    useEffect(() => {
 67      if (!session) return
 68  
 69      hasReceivedInitRef.current = false
 70      logForDebugging('[useSSHSession] wiring SSH session manager')
 71  
 72      const manager = session.createManager({
 73        onMessage: sdkMessage => {
 74          if (isSessionEndMessage(sdkMessage)) {
 75            setIsLoading(false)
 76          }
 77  
 78          // Skip duplicate init messages (one per turn from stream-json mode).
 79          if (sdkMessage.type === 'system' && sdkMessage.subtype === 'init') {
 80            if (hasReceivedInitRef.current) return
 81            hasReceivedInitRef.current = true
 82          }
 83  
 84          const converted = convertSDKMessage(sdkMessage, {
 85            convertToolResults: true,
 86          })
 87          if (converted.type === 'message') {
 88            setMessages(prev => [...prev, converted.message])
 89          }
 90        },
 91        onPermissionRequest: (request, requestId) => {
 92          logForDebugging(
 93            `[useSSHSession] permission request: ${request.tool_name}`,
 94          )
 95  
 96          const tool =
 97            findToolByName(toolsRef.current, request.tool_name) ??
 98            createToolStub(request.tool_name)
 99  
100          const syntheticMessage = createSyntheticAssistantMessage(
101            request,
102            requestId,
103          )
104  
105          const permissionResult: PermissionAskDecision = {
106            behavior: 'ask',
107            message:
108              request.description ?? `${request.tool_name} requires permission`,
109            suggestions: request.permission_suggestions,
110            blockedPath: request.blocked_path,
111          }
112  
113          const toolUseConfirm: ToolUseConfirm = {
114            assistantMessage: syntheticMessage,
115            tool,
116            description:
117              request.description ?? `${request.tool_name} requires permission`,
118            input: request.input,
119            toolUseContext: {} as ToolUseConfirm['toolUseContext'],
120            toolUseID: request.tool_use_id,
121            permissionResult,
122            permissionPromptStartTimeMs: Date.now(),
123            onUserInteraction() {},
124            onAbort() {
125              manager.respondToPermissionRequest(requestId, {
126                behavior: 'deny',
127                message: 'User aborted',
128              })
129              setToolUseConfirmQueue(q =>
130                q.filter(i => i.toolUseID !== request.tool_use_id),
131              )
132            },
133            onAllow(updatedInput) {
134              manager.respondToPermissionRequest(requestId, {
135                behavior: 'allow',
136                updatedInput,
137              })
138              setToolUseConfirmQueue(q =>
139                q.filter(i => i.toolUseID !== request.tool_use_id),
140              )
141              setIsLoading(true)
142            },
143            onReject(feedback) {
144              manager.respondToPermissionRequest(requestId, {
145                behavior: 'deny',
146                message: feedback ?? 'User denied permission',
147              })
148              setToolUseConfirmQueue(q =>
149                q.filter(i => i.toolUseID !== request.tool_use_id),
150              )
151            },
152            async recheckPermission() {},
153          }
154  
155          setToolUseConfirmQueue(q => [...q, toolUseConfirm])
156          setIsLoading(false)
157        },
158        onConnected: () => {
159          logForDebugging('[useSSHSession] connected')
160          isConnectedRef.current = true
161        },
162        onReconnecting: (attempt, max) => {
163          logForDebugging(
164            `[useSSHSession] ssh dropped, reconnecting (${attempt}/${max})`,
165          )
166          isConnectedRef.current = false
167          // Surface a transient system message in the transcript so the user
168          // knows what's happening — the next onConnected clears the state.
169          // Any in-flight request is lost; the remote's --continue reloads
170          // history but there's no turn in progress to resume.
171          setIsLoading(false)
172          const msg: MessageType = {
173            type: 'system',
174            subtype: 'informational',
175            content: `SSH connection dropped — reconnecting (attempt ${attempt}/${max})...`,
176            timestamp: new Date().toISOString(),
177            uuid: randomUUID(),
178            level: 'warning',
179          }
180          setMessages(prev => [...prev, msg])
181        },
182        onDisconnected: () => {
183          logForDebugging('[useSSHSession] ssh process exited (giving up)')
184          const stderr = session.getStderrTail().trim()
185          const connected = isConnectedRef.current
186          const exitCode = session.proc.exitCode
187          isConnectedRef.current = false
188          setIsLoading(false)
189  
190          let msg = connected
191            ? 'Remote session ended.'
192            : 'SSH session failed before connecting.'
193          // Surface remote stderr if it looks like an error (pre-connect always,
194          // post-connect only on nonzero exit — normal --verbose noise otherwise).
195          if (stderr && (!connected || exitCode !== 0)) {
196            msg += `\nRemote stderr (exit ${exitCode ?? 'signal ' + session.proc.signalCode}):\n${stderr}`
197          }
198          void gracefulShutdown(1, 'other', { finalMessage: msg })
199        },
200        onError: error => {
201          logForDebugging(`[useSSHSession] error: ${error.message}`)
202        },
203      })
204  
205      managerRef.current = manager
206      manager.connect()
207  
208      return () => {
209        logForDebugging('[useSSHSession] cleanup')
210        manager.disconnect()
211        session.proxy.stop()
212        managerRef.current = null
213      }
214    }, [session, setMessages, setIsLoading, setToolUseConfirmQueue])
215  
216    const sendMessage = useCallback(
217      async (content: RemoteMessageContent): Promise<boolean> => {
218        const m = managerRef.current
219        if (!m) return false
220        setIsLoading(true)
221        return m.sendMessage(content)
222      },
223      [setIsLoading],
224    )
225  
226    const cancelRequest = useCallback(() => {
227      managerRef.current?.sendInterrupt()
228      setIsLoading(false)
229    }, [setIsLoading])
230  
231    const disconnect = useCallback(() => {
232      managerRef.current?.disconnect()
233      managerRef.current = null
234      isConnectedRef.current = false
235    }, [])
236  
237    return useMemo(
238      () => ({ isRemoteMode, sendMessage, cancelRequest, disconnect }),
239      [isRemoteMode, sendMessage, cancelRequest, disconnect],
240    )
241  }