1  import { useCallback, useEffect, useMemo, useRef } from 'react'
  2  import { BoundedUUIDSet } from '../bridge/bridgeMessaging.js'
  3  import type { ToolUseConfirm } from '../components/permissions/PermissionRequest.js'
  4  import type { SpinnerMode } from '../components/Spinner/types.js'
  5  import {
  6    type RemotePermissionResponse,
  7    type RemoteSessionConfig,
  8    RemoteSessionManager,
  9  } from '../remote/RemoteSessionManager.js'
 10  import {
 11    createSyntheticAssistantMessage,
 12    createToolStub,
 13  } from '../remote/remotePermissionBridge.js'
 14  import {
 15    convertSDKMessage,
 16    isSessionEndMessage,
 17  } from '../remote/sdkMessageAdapter.js'
 18  import { useSetAppState } from '../state/AppState.js'
 19  import type { AppState } from '../state/AppStateStore.js'
 20  import type { Tool } from '../Tool.js'
 21  import { findToolByName } from '../Tool.js'
 22  import type { Message as MessageType } from '../types/message.js'
 23  import type { PermissionAskDecision } from '../types/permissions.js'
 24  import { logForDebugging } from '../utils/debug.js'
 25  import { truncateToWidth } from '../utils/format.js'
 26  import {
 27    createSystemMessage,
 28    extractTextContent,
 29    handleMessageFromStream,
 30    type StreamingToolUse,
 31  } from '../utils/messages.js'
 32  import { generateSessionTitle } from '../utils/sessionTitle.js'
 33  import type { RemoteMessageContent } from '../utils/teleport/api.js'
 34  import { updateSessionTitle } from '../utils/teleport/api.js'
 35  
 36  // How long to wait for a response before showing a warning
 37  const RESPONSE_TIMEOUT_MS = 60000 // 60 seconds
 38  // Extended timeout during compaction — compact API calls take 5-30s and
 39  // block other SDK messages, so the normal 60s timeout isn't enough when
 40  // compaction itself runs close to the edge.
 41  const COMPACTION_TIMEOUT_MS = 180000 // 3 minutes
 42  
 43  type UseRemoteSessionProps = {
 44    config: RemoteSessionConfig | undefined
 45    setMessages: React.Dispatch<React.SetStateAction<MessageType[]>>
 46    setIsLoading: (loading: boolean) => void
 47    onInit?: (slashCommands: string[]) => void
 48    setToolUseConfirmQueue: React.Dispatch<React.SetStateAction<ToolUseConfirm[]>>
 49    tools: Tool[]
 50    setStreamingToolUses?: React.Dispatch<
 51      React.SetStateAction<StreamingToolUse[]>
 52    >
 53    setStreamMode?: React.Dispatch<React.SetStateAction<SpinnerMode>>
 54    setInProgressToolUseIDs?: (f: (prev: Set<string>) => Set<string>) => void
 55  }
 56  
 57  type UseRemoteSessionResult = {
 58    isRemoteMode: boolean
 59    sendMessage: (
 60      content: RemoteMessageContent,
 61      opts?: { uuid?: string },
 62    ) => Promise<boolean>
 63    cancelRequest: () => void
 64    disconnect: () => void
 65  }
 66  
 67  /**
 68   * Hook for managing a remote CCR session in the REPL.
 69   *
 70   * Handles:
 71   * - WebSocket connection to CCR
 72   * - Converting SDK messages to REPL messages
 73   * - Sending user input to CCR via HTTP POST
 74   * - Permission request/response flow via existing ToolUseConfirm queue
 75   */
 76  export function useRemoteSession({
 77    config,
 78    setMessages,
 79    setIsLoading,
 80    onInit,
 81    setToolUseConfirmQueue,
 82    tools,
 83    setStreamingToolUses,
 84    setStreamMode,
 85    setInProgressToolUseIDs,
 86  }: UseRemoteSessionProps): UseRemoteSessionResult {
 87    const isRemoteMode = !!config
 88  
 89    const setAppState = useSetAppState()
 90    const setConnStatus = useCallback(
 91      (s: AppState['remoteConnectionStatus']) =>
 92        setAppState(prev =>
 93          prev.remoteConnectionStatus === s
 94            ? prev
 95            : { ...prev, remoteConnectionStatus: s },
 96        ),
 97      [setAppState],
 98    )
 99  
100    // Event-sourced count of subagents running inside the remote daemon child.
101    // The viewer's own AppState.tasks is empty — tasks live in a different
102    // process. task_started/task_notification reach us via the bridge WS.
103    const runningTaskIdsRef = useRef(new Set<string>())
104    const writeTaskCount = useCallback(() => {
105      const n = runningTaskIdsRef.current.size
106      setAppState(prev =>
107        prev.remoteBackgroundTaskCount === n
108          ? prev
109          : { ...prev, remoteBackgroundTaskCount: n },
110      )
111    }, [setAppState])
112  
113    // Timer for detecting stuck sessions
114    const responseTimeoutRef = useRef<NodeJS.Timeout | null>(null)
115  
116    // Track whether the remote session is compacting. During compaction the
117    // CLI worker is busy with an API call and won't emit messages for a while;
118    // use a longer timeout and suppress spurious "unresponsive" warnings.
119    const isCompactingRef = useRef(false)
120  
121    const managerRef = useRef<RemoteSessionManager | null>(null)
122  
123    // Track whether we've already updated the session title (for no-initial-prompt sessions)
124    const hasUpdatedTitleRef = useRef(false)
125  
126    // UUIDs of user messages we POSTed locally — the WS echoes them back and
127    // we must filter them out when convertUserTextMessages is on, or the viewer
128    // sees every typed message twice (once from local createUserMessage, once
129    // from the echo). A single POST can echo MULTIPLE times with the same uuid:
130    // the server may broadcast the POST directly to /subscribe, AND the worker
131    // (cowork desktop / CLI daemon) echoes it again on its write path. A
132    // delete-on-first-match Set would let the second echo through — use a
133    // bounded ring instead. Cap is generous: users don't type 50 messages
134    // faster than echoes arrive.
135    // NOTE: this does NOT dedup history-vs-live overlap at attach time (nothing
136    // seeds the set from history UUIDs; only sendMessage populates it).
137    const sentUUIDsRef = useRef(new BoundedUUIDSet(50))
138  
139    // Keep a ref to tools so the WebSocket callback doesn't go stale
140    const toolsRef = useRef(tools)
141    useEffect(() => {
142      toolsRef.current = tools
143    }, [tools])
144  
145    // Initialize and connect to remote session
146    useEffect(() => {
147      // Skip if not in remote mode
148      if (!config) {
149        return
150      }
151  
152      logForDebugging(
153        `[useRemoteSession] Initializing for session ${config.sessionId}`,
154      )
155  
156      const manager = new RemoteSessionManager(config, {
157        onMessage: sdkMessage => {
158          const parts = [`type=${sdkMessage.type}`]
159          if ('subtype' in sdkMessage) parts.push(`subtype=${sdkMessage.subtype}`)
160          if (sdkMessage.type === 'user') {
161            const c = sdkMessage.message?.content
162            parts.push(
163              `content=${Array.isArray(c) ? c.map(b => b.type).join(',') : typeof c}`,
164            )
165          }
166          logForDebugging(`[useRemoteSession] Received ${parts.join(' ')}`)
167  
168          // Clear response timeout on any message received — including the WS
169          // echo of our own POST, which acts as a heartbeat. This must run
170          // BEFORE the echo filter, or slow-to-stream agents (compaction, cold
171          // start) spuriously trip the 60s unresponsive warning + reconnect.
172          if (responseTimeoutRef.current) {
173            clearTimeout(responseTimeoutRef.current)
174            responseTimeoutRef.current = null
175          }
176  
177          // Echo filter: drop user messages we already added locally before POST.
178          // The server and/or worker round-trip our own send back on the WS with
179          // the same uuid we passed to sendEventToRemoteSession. DO NOT delete on
180          // match — the same uuid can echo more than once (server broadcast +
181          // worker echo), and BoundedUUIDSet already caps growth via its ring.
182          if (
183            sdkMessage.type === 'user' &&
184            sdkMessage.uuid &&
185            sentUUIDsRef.current.has(sdkMessage.uuid)
186          ) {
187            logForDebugging(
188              `[useRemoteSession] Dropping echoed user message ${sdkMessage.uuid}`,
189            )
190            return
191          }
192          // Handle init message - extract available slash commands
193          if (
194            sdkMessage.type === 'system' &&
195            sdkMessage.subtype === 'init' &&
196            onInit
197          ) {
198            logForDebugging(
199              `[useRemoteSession] Init received with ${sdkMessage.slash_commands.length} slash commands`,
200            )
201            onInit(sdkMessage.slash_commands)
202          }
203  
204          // Track remote subagent lifecycle for the "N in background" counter.
205          // All task types (Agent/teammate/workflow/bash) flow through
206          // registerTask() → task_started, and complete via task_notification.
207          // Return early — these are status signals, not renderable messages.
208          if (sdkMessage.type === 'system') {
209            if (sdkMessage.subtype === 'task_started') {
210              runningTaskIdsRef.current.add(sdkMessage.task_id)
211              writeTaskCount()
212              return
213            }
214            if (sdkMessage.subtype === 'task_notification') {
215              runningTaskIdsRef.current.delete(sdkMessage.task_id)
216              writeTaskCount()
217              return
218            }
219            if (sdkMessage.subtype === 'task_progress') {
220              return
221            }
222            // Track compaction state. The CLI emits status='compacting' at
223            // the start and status=null when done; compact_boundary also
224            // signals completion. Repeated 'compacting' status messages
225            // (keep-alive ticks) update the ref but don't append to messages.
226            if (sdkMessage.subtype === 'status') {
227              const wasCompacting = isCompactingRef.current
228              isCompactingRef.current = sdkMessage.status === 'compacting'
229              if (wasCompacting && isCompactingRef.current) {
230                return
231              }
232            }
233            if (sdkMessage.subtype === 'compact_boundary') {
234              isCompactingRef.current = false
235            }
236          }
237  
238          // Check if session ended
239          if (isSessionEndMessage(sdkMessage)) {
240            isCompactingRef.current = false
241            setIsLoading(false)
242          }
243  
244          // Clear in-progress tool_use IDs when their tool_result arrives.
245          // Must read the RAW sdkMessage: in non-viewerOnly mode,
246          // convertSDKMessage returns {type:'ignored'} for user messages, so the
247          // delete would never fire post-conversion. Mirrors the add site below
248          // and inProcessRunner.ts; without this the set grows unbounded for the
249          // session lifetime (BQ: CCR cohort shows 5.2x higher RSS slope).
250          if (setInProgressToolUseIDs && sdkMessage.type === 'user') {
251            const content = sdkMessage.message?.content
252            if (Array.isArray(content)) {
253              const resultIds: string[] = []
254              for (const block of content) {
255                if (block.type === 'tool_result') {
256                  resultIds.push(block.tool_use_id)
257                }
258              }
259              if (resultIds.length > 0) {
260                setInProgressToolUseIDs(prev => {
261                  const next = new Set(prev)
262                  for (const id of resultIds) next.delete(id)
263                  return next.size === prev.size ? prev : next
264                })
265              }
266            }
267          }
268  
269          // Convert SDK message to REPL message. In viewerOnly mode, the
270          // remote agent runs BriefTool (SendUserMessage) — its tool_use block
271          // renders empty (userFacingName() === ''), actual content is in the
272          // tool_result. So we must convert tool_results to render them.
273          const converted = convertSDKMessage(
274            sdkMessage,
275            config.viewerOnly
276              ? { convertToolResults: true, convertUserTextMessages: true }
277              : undefined,
278          )
279  
280          if (converted.type === 'message') {
281            // When we receive a complete message, clear streaming tool uses
282            // since the complete message replaces the partial streaming state
283            setStreamingToolUses?.(prev => (prev.length > 0 ? [] : prev))
284  
285            // Mark tool_use blocks as in-progress so the UI shows the correct
286            // spinner state instead of "Waiting…" (queued). In local sessions,
287            // toolOrchestration.ts handles this, but remote sessions receive
288            // pre-built assistant messages without running local tool execution.
289            if (
290              setInProgressToolUseIDs &&
291              converted.message.type === 'assistant'
292            ) {
293              const toolUseIds = converted.message.message.content
294                .filter(block => block.type === 'tool_use')
295                .map(block => block.id)
296              if (toolUseIds.length > 0) {
297                setInProgressToolUseIDs(prev => {
298                  const next = new Set(prev)
299                  for (const id of toolUseIds) {
300                    next.add(id)
301                  }
302                  return next
303                })
304              }
305            }
306  
307            setMessages(prev => [...prev, converted.message])
308            // Note: Don't stop loading on assistant messages - the agent may still be
309            // working (tool use loops). Loading stops only on session end or permission request.
310          } else if (converted.type === 'stream_event') {
311            // Process streaming events to update UI in real-time
312            if (setStreamingToolUses && setStreamMode) {
313              handleMessageFromStream(
314                converted.event,
315                message => setMessages(prev => [...prev, message]),
316                () => {
317                  // No-op for response length - remote sessions don't track this
318                },
319                setStreamMode,
320                setStreamingToolUses,
321              )
322            } else {
323              logForDebugging(
324                `[useRemoteSession] Stream event received but streaming callbacks not provided`,
325              )
326            }
327          }
328          // 'ignored' messages are silently dropped
329        },
330        onPermissionRequest: (request, requestId) => {
331          logForDebugging(
332            `[useRemoteSession] Permission request for tool: ${request.tool_name}`,
333          )
334  
335          // Look up the Tool object by name, or create a stub for unknown tools
336          const tool =
337            findToolByName(toolsRef.current, request.tool_name) ??
338            createToolStub(request.tool_name)
339  
340          const syntheticMessage = createSyntheticAssistantMessage(
341            request,
342            requestId,
343          )
344  
345          const permissionResult: PermissionAskDecision = {
346            behavior: 'ask',
347            message:
348              request.description ?? `${request.tool_name} requires permission`,
349            suggestions: request.permission_suggestions,
350            blockedPath: request.blocked_path,
351          }
352  
353          const toolUseConfirm: ToolUseConfirm = {
354            assistantMessage: syntheticMessage,
355            tool,
356            description:
357              request.description ?? `${request.tool_name} requires permission`,
358            input: request.input,
359            toolUseContext: {} as ToolUseConfirm['toolUseContext'],
360            toolUseID: request.tool_use_id,
361            permissionResult,
362            permissionPromptStartTimeMs: Date.now(),
363            onUserInteraction() {
364              // No-op for remote — classifier runs on the container
365            },
366            onAbort() {
367              const response: RemotePermissionResponse = {
368                behavior: 'deny',
369                message: 'User aborted',
370              }
371              manager.respondToPermissionRequest(requestId, response)
372              setToolUseConfirmQueue(queue =>
373                queue.filter(item => item.toolUseID !== request.tool_use_id),
374              )
375            },
376            onAllow(updatedInput, _permissionUpdates, _feedback) {
377              const response: RemotePermissionResponse = {
378                behavior: 'allow',
379                updatedInput,
380              }
381              manager.respondToPermissionRequest(requestId, response)
382              setToolUseConfirmQueue(queue =>
383                queue.filter(item => item.toolUseID !== request.tool_use_id),
384              )
385              // Resume loading indicator after approving
386              setIsLoading(true)
387            },
388            onReject(feedback?: string) {
389              const response: RemotePermissionResponse = {
390                behavior: 'deny',
391                message: feedback ?? 'User denied permission',
392              }
393              manager.respondToPermissionRequest(requestId, response)
394              setToolUseConfirmQueue(queue =>
395                queue.filter(item => item.toolUseID !== request.tool_use_id),
396              )
397            },
398            async recheckPermission() {
399              // No-op for remote — permission state is on the container
400            },
401          }
402  
403          setToolUseConfirmQueue(queue => [...queue, toolUseConfirm])
404          // Pause loading indicator while waiting for permission
405          setIsLoading(false)
406        },
407        onPermissionCancelled: (requestId, toolUseId) => {
408          logForDebugging(
409            `[useRemoteSession] Permission request cancelled: ${requestId}`,
410          )
411          const idToRemove = toolUseId ?? requestId
412          setToolUseConfirmQueue(queue =>
413            queue.filter(item => item.toolUseID !== idToRemove),
414          )
415          setIsLoading(true)
416        },
417        onConnected: () => {
418          logForDebugging('[useRemoteSession] Connected')
419          setConnStatus('connected')
420        },
421        onReconnecting: () => {
422          logForDebugging('[useRemoteSession] Reconnecting')
423          setConnStatus('reconnecting')
424          // WS gap = we may miss task_notification events. Clear rather than
425          // drift high forever. Undercounts tasks that span the gap; accepted.
426          runningTaskIdsRef.current.clear()
427          writeTaskCount()
428          // Same for tool_use IDs: missed tool_result during the gap would
429          // leave stale spinner state forever.
430          setInProgressToolUseIDs?.(prev => (prev.size > 0 ? new Set() : prev))
431        },
432        onDisconnected: () => {
433          logForDebugging('[useRemoteSession] Disconnected')
434          setConnStatus('disconnected')
435          setIsLoading(false)
436          runningTaskIdsRef.current.clear()
437          writeTaskCount()
438          setInProgressToolUseIDs?.(prev => (prev.size > 0 ? new Set() : prev))
439        },
440        onError: error => {
441          logForDebugging(`[useRemoteSession] Error: ${error.message}`)
442        },
443      })
444  
445      managerRef.current = manager
446      manager.connect()
447  
448      return () => {
449        logForDebugging('[useRemoteSession] Cleanup - disconnecting')
450        // Clear any pending timeout
451        if (responseTimeoutRef.current) {
452          clearTimeout(responseTimeoutRef.current)
453          responseTimeoutRef.current = null
454        }
455        manager.disconnect()
456        managerRef.current = null
457      }
458    }, [
459      config,
460      setMessages,
461      setIsLoading,
462      onInit,
463      setToolUseConfirmQueue,
464      setStreamingToolUses,
465      setStreamMode,
466      setInProgressToolUseIDs,
467      setConnStatus,
468      writeTaskCount,
469    ])
470  
471    // Send a user message to the remote session
472    const sendMessage = useCallback(
473      async (
474        content: RemoteMessageContent,
475        opts?: { uuid?: string },
476      ): Promise<boolean> => {
477        const manager = managerRef.current
478        if (!manager) {
479          logForDebugging('[useRemoteSession] Cannot send - no manager')
480          return false
481        }
482  
483        // Clear any existing timeout
484        if (responseTimeoutRef.current) {
485          clearTimeout(responseTimeoutRef.current)
486        }
487  
488        setIsLoading(true)
489  
490        // Track locally-added message UUIDs so the WS echo can be filtered.
491        // Must record BEFORE the POST to close the race where the echo arrives
492        // before the POST promise resolves.
493        if (opts?.uuid) sentUUIDsRef.current.add(opts.uuid)
494  
495        const success = await manager.sendMessage(content, opts)
496  
497        if (!success) {
498          // No need to undo the pre-POST add — BoundedUUIDSet's ring evicts it.
499          setIsLoading(false)
500          return false
501        }
502  
503        // Update the session title after the first message when no initial prompt was provided.
504        // This gives the session a meaningful title on claude.ai instead of "Background task".
505        // Skip in viewerOnly mode — the remote agent owns the session title.
506        if (
507          !hasUpdatedTitleRef.current &&
508          config &&
509          !config.hasInitialPrompt &&
510          !config.viewerOnly
511        ) {
512          hasUpdatedTitleRef.current = true
513          const sessionId = config.sessionId
514          // Extract plain text from content (may be string or content block array)
515          const description =
516            typeof content === 'string'
517              ? content
518              : extractTextContent(content, ' ')
519          if (description) {
520            // generateSessionTitle never rejects (wraps body in try/catch,
521            // returns null on failure), so no .catch needed on this chain.
522            void generateSessionTitle(
523              description,
524              new AbortController().signal,
525            ).then(title => {
526              void updateSessionTitle(
527                sessionId,
528                title ?? truncateToWidth(description, 75),
529              )
530            })
531          }
532        }
533  
534        // Start timeout to detect stuck sessions. Skip in viewerOnly mode —
535        // the remote agent may be idle-shut and take >60s to respawn.
536        // Use a longer timeout when the remote session is compacting, since
537        // the CLI worker is busy with an API call and won't emit messages.
538        if (!config?.viewerOnly) {
539          const timeoutMs = isCompactingRef.current
540            ? COMPACTION_TIMEOUT_MS
541            : RESPONSE_TIMEOUT_MS
542          responseTimeoutRef.current = setTimeout(
543            (setMessages, manager) => {
544              logForDebugging(
545                '[useRemoteSession] Response timeout - attempting reconnect',
546              )
547              // Add a warning message to the conversation
548              const warningMessage = createSystemMessage(
549                'Remote session may be unresponsive. Attempting to reconnect…',
550                'warning',
551              )
552              setMessages(prev => [...prev, warningMessage])
553  
554              // Attempt to reconnect the WebSocket - the subscription may have become stale
555              manager.reconnect()
556            },
557            timeoutMs,
558            setMessages,
559            manager,
560          )
561        }
562  
563        return success
564      },
565      [config, setIsLoading, setMessages],
566    )
567  
568    // Cancel the current request on the remote session
569    const cancelRequest = useCallback(() => {
570      // Clear any pending timeout
571      if (responseTimeoutRef.current) {
572        clearTimeout(responseTimeoutRef.current)
573        responseTimeoutRef.current = null
574      }
575  
576      // Send interrupt signal to CCR. Skip in viewerOnly mode — Ctrl+C
577      // should never interrupt the remote agent.
578      if (!config?.viewerOnly) {
579        managerRef.current?.cancelSession()
580      }
581  
582      setIsLoading(false)
583    }, [config, setIsLoading])
584  
585    // Disconnect from the session
586    const disconnect = useCallback(() => {
587      // Clear any pending timeout
588      if (responseTimeoutRef.current) {
589        clearTimeout(responseTimeoutRef.current)
590        responseTimeoutRef.current = null
591      }
592      managerRef.current?.disconnect()
593      managerRef.current = null
594    }, [])
595  
596    // All four fields are already stable (boolean derived from a prop that
597    // doesn't change mid-session, three useCallbacks with stable deps). The
598    // result object is consumed by REPL's onSubmit useCallback deps — without
599    // memoization the fresh literal invalidates onSubmit on every REPL render,
600    // which in turn churns PromptInput's props and downstream memoization.
601    return useMemo(
602      () => ({ isRemoteMode, sendMessage, cancelRequest, disconnect }),
603      [isRemoteMode, sendMessage, cancelRequest, disconnect],
604    )
605  }