1  import type { SDKMessage } from '../entrypoints/agentSdkTypes.js'
  2  import type {
  3    SDKControlCancelRequest,
  4    SDKControlPermissionRequest,
  5    SDKControlRequest,
  6    SDKControlResponse,
  7  } from '../entrypoints/sdk/controlTypes.js'
  8  import { logForDebugging } from '../utils/debug.js'
  9  import { logError } from '../utils/log.js'
 10  import {
 11    type RemoteMessageContent,
 12    sendEventToRemoteSession,
 13  } from '../utils/teleport/api.js'
 14  import {
 15    SessionsWebSocket,
 16    type SessionsWebSocketCallbacks,
 17  } from './SessionsWebSocket.js'
 18  
 19  /**
 20   * Type guard to check if a message is an SDKMessage (not a control message)
 21   */
 22  function isSDKMessage(
 23    message:
 24      | SDKMessage
 25      | SDKControlRequest
 26      | SDKControlResponse
 27      | SDKControlCancelRequest,
 28  ): message is SDKMessage {
 29    return (
 30      message.type !== 'control_request' &&
 31      message.type !== 'control_response' &&
 32      message.type !== 'control_cancel_request'
 33    )
 34  }
 35  
 36  /**
 37   * Simple permission response for remote sessions.
 38   * This is a simplified version of PermissionResult for CCR communication.
 39   */
 40  export type RemotePermissionResponse =
 41    | {
 42        behavior: 'allow'
 43        updatedInput: Record<string, unknown>
 44      }
 45    | {
 46        behavior: 'deny'
 47        message: string
 48      }
 49  
 50  export type RemoteSessionConfig = {
 51    sessionId: string
 52    getAccessToken: () => string
 53    orgUuid: string
 54    /** True if session was created with an initial prompt that's being processed */
 55    hasInitialPrompt?: boolean
 56    /**
 57     * When true, this client is a pure viewer. Ctrl+C/Escape do NOT send
 58     * interrupt to the remote agent; 60s reconnect timeout is disabled;
 59     * session title is never updated. Used by `claude assistant`.
 60     */
 61    viewerOnly?: boolean
 62  }
 63  
 64  export type RemoteSessionCallbacks = {
 65    /** Called when an SDKMessage is received from the session */
 66    onMessage: (message: SDKMessage) => void
 67    /** Called when a permission request is received from CCR */
 68    onPermissionRequest: (
 69      request: SDKControlPermissionRequest,
 70      requestId: string,
 71    ) => void
 72    /** Called when the server cancels a pending permission request */
 73    onPermissionCancelled?: (
 74      requestId: string,
 75      toolUseId: string | undefined,
 76    ) => void
 77    /** Called when connection is established */
 78    onConnected?: () => void
 79    /** Called when connection is lost and cannot be restored */
 80    onDisconnected?: () => void
 81    /** Called on transient WS drop while reconnect backoff is in progress */
 82    onReconnecting?: () => void
 83    /** Called on error */
 84    onError?: (error: Error) => void
 85  }
 86  
 87  /**
 88   * Manages a remote CCR session.
 89   *
 90   * Coordinates:
 91   * - WebSocket subscription for receiving messages from CCR
 92   * - HTTP POST for sending user messages to CCR
 93   * - Permission request/response flow
 94   */
 95  export class RemoteSessionManager {
 96    private websocket: SessionsWebSocket | null = null
 97    private pendingPermissionRequests: Map<string, SDKControlPermissionRequest> =
 98      new Map()
 99  
100    constructor(
101      private readonly config: RemoteSessionConfig,
102      private readonly callbacks: RemoteSessionCallbacks,
103    ) {}
104  
105    /**
106     * Connect to the remote session via WebSocket
107     */
108    connect(): void {
109      logForDebugging(
110        `[RemoteSessionManager] Connecting to session ${this.config.sessionId}`,
111      )
112  
113      const wsCallbacks: SessionsWebSocketCallbacks = {
114        onMessage: message => this.handleMessage(message),
115        onConnected: () => {
116          logForDebugging('[RemoteSessionManager] Connected')
117          this.callbacks.onConnected?.()
118        },
119        onClose: () => {
120          logForDebugging('[RemoteSessionManager] Disconnected')
121          this.callbacks.onDisconnected?.()
122        },
123        onReconnecting: () => {
124          logForDebugging('[RemoteSessionManager] Reconnecting')
125          this.callbacks.onReconnecting?.()
126        },
127        onError: error => {
128          logError(error)
129          this.callbacks.onError?.(error)
130        },
131      }
132  
133      this.websocket = new SessionsWebSocket(
134        this.config.sessionId,
135        this.config.orgUuid,
136        this.config.getAccessToken,
137        wsCallbacks,
138      )
139  
140      void this.websocket.connect()
141    }
142  
143    /**
144     * Handle messages from WebSocket
145     */
146    private handleMessage(
147      message:
148        | SDKMessage
149        | SDKControlRequest
150        | SDKControlResponse
151        | SDKControlCancelRequest,
152    ): void {
153      // Handle control requests (permission prompts from CCR)
154      if (message.type === 'control_request') {
155        this.handleControlRequest(message)
156        return
157      }
158  
159      // Handle control cancel requests (server cancelling a pending permission prompt)
160      if (message.type === 'control_cancel_request') {
161        const { request_id } = message
162        const pendingRequest = this.pendingPermissionRequests.get(request_id)
163        logForDebugging(
164          `[RemoteSessionManager] Permission request cancelled: ${request_id}`,
165        )
166        this.pendingPermissionRequests.delete(request_id)
167        this.callbacks.onPermissionCancelled?.(
168          request_id,
169          pendingRequest?.tool_use_id,
170        )
171        return
172      }
173  
174      // Handle control responses (acknowledgments)
175      if (message.type === 'control_response') {
176        logForDebugging('[RemoteSessionManager] Received control response')
177        return
178      }
179  
180      // Forward SDK messages to callback (type guard ensures proper narrowing)
181      if (isSDKMessage(message)) {
182        this.callbacks.onMessage(message)
183      }
184    }
185  
186    /**
187     * Handle control requests from CCR (e.g., permission requests)
188     */
189    private handleControlRequest(request: SDKControlRequest): void {
190      const { request_id, request: inner } = request
191  
192      if (inner.subtype === 'can_use_tool') {
193        logForDebugging(
194          `[RemoteSessionManager] Permission request for tool: ${inner.tool_name}`,
195        )
196        this.pendingPermissionRequests.set(request_id, inner)
197        this.callbacks.onPermissionRequest(inner, request_id)
198      } else {
199        // Send an error response for unrecognized subtypes so the server
200        // doesn't hang waiting for a reply that never comes.
201        logForDebugging(
202          `[RemoteSessionManager] Unsupported control request subtype: ${inner.subtype}`,
203        )
204        const response: SDKControlResponse = {
205          type: 'control_response',
206          response: {
207            subtype: 'error',
208            request_id,
209            error: `Unsupported control request subtype: ${inner.subtype}`,
210          },
211        }
212        this.websocket?.sendControlResponse(response)
213      }
214    }
215  
216    /**
217     * Send a user message to the remote session via HTTP POST
218     */
219    async sendMessage(
220      content: RemoteMessageContent,
221      opts?: { uuid?: string },
222    ): Promise<boolean> {
223      logForDebugging(
224        `[RemoteSessionManager] Sending message to session ${this.config.sessionId}`,
225      )
226  
227      const success = await sendEventToRemoteSession(
228        this.config.sessionId,
229        content,
230        opts,
231      )
232  
233      if (!success) {
234        logError(
235          new Error(
236            `[RemoteSessionManager] Failed to send message to session ${this.config.sessionId}`,
237          ),
238        )
239      }
240  
241      return success
242    }
243  
244    /**
245     * Respond to a permission request from CCR
246     */
247    respondToPermissionRequest(
248      requestId: string,
249      result: RemotePermissionResponse,
250    ): void {
251      const pendingRequest = this.pendingPermissionRequests.get(requestId)
252      if (!pendingRequest) {
253        logError(
254          new Error(
255            `[RemoteSessionManager] No pending permission request with ID: ${requestId}`,
256          ),
257        )
258        return
259      }
260  
261      this.pendingPermissionRequests.delete(requestId)
262  
263      const response: SDKControlResponse = {
264        type: 'control_response',
265        response: {
266          subtype: 'success',
267          request_id: requestId,
268          response: {
269            behavior: result.behavior,
270            ...(result.behavior === 'allow'
271              ? { updatedInput: result.updatedInput }
272              : { message: result.message }),
273          },
274        },
275      }
276  
277      logForDebugging(
278        `[RemoteSessionManager] Sending permission response: ${result.behavior}`,
279      )
280  
281      this.websocket?.sendControlResponse(response)
282    }
283  
284    /**
285     * Check if connected to the remote session
286     */
287    isConnected(): boolean {
288      return this.websocket?.isConnected() ?? false
289    }
290  
291    /**
292     * Send an interrupt signal to cancel the current request on the remote session
293     */
294    cancelSession(): void {
295      logForDebugging('[RemoteSessionManager] Sending interrupt signal')
296      this.websocket?.sendControlRequest({ subtype: 'interrupt' })
297    }
298  
299    /**
300     * Get the session ID
301     */
302    getSessionId(): string {
303      return this.config.sessionId
304    }
305  
306    /**
307     * Disconnect from the remote session
308     */
309    disconnect(): void {
310      logForDebugging('[RemoteSessionManager] Disconnecting')
311      this.websocket?.close()
312      this.websocket = null
313      this.pendingPermissionRequests.clear()
314    }
315  
316    /**
317     * Force reconnect the WebSocket.
318     * Useful when the subscription becomes stale after container shutdown.
319     */
320    reconnect(): void {
321      logForDebugging('[RemoteSessionManager] Reconnecting WebSocket')
322      this.websocket?.reconnect()
323    }
324  }
325  
326  /**
327   * Create a remote session config from OAuth tokens
328   */
329  export function createRemoteSessionConfig(
330    sessionId: string,
331    getAccessToken: () => string,
332    orgUuid: string,
333    hasInitialPrompt = false,
334    viewerOnly = false,
335  ): RemoteSessionConfig {
336    return {
337      sessionId,
338      getAccessToken,
339      orgUuid,
340      hasInitialPrompt,
341      viewerOnly,
342    }
343  }