orca/src/shared/remote-runtime-shared-control-frame-handler.ts

import { parseAuthenticatedFrame, parseReadyFrame } from './remote-runtime-request-frames'
import type { RemoteRuntimeClientError } from './remote-runtime-client-error'
import {
  AGENT_SESSION_BOUNDARY_RUNTIME_CAPABILITY,
  SESSION_TAB_CLOSE_INTENT_RUNTIME_CAPABILITY,
  WORKTREE_VISIBILITY_DEFAULTS_RUNTIME_CAPABILITY,
  WORKTREE_VISIBILITY_SOURCE_DEFAULTS_RUNTIME_CAPABILITY
} from './protocol-version'
import { dispatchSharedControlFrame } from './remote-runtime-shared-control-frame-dispatch'
import { parseSharedControlFrame } from './remote-runtime-shared-control-protocol'
import type { SharedControlRetiredRequestIds } from './remote-runtime-shared-control-retired-request-ids'
import { resolveSharedControlReadyWaiters } from './remote-runtime-shared-control-state'
import type {
  SharedControlConnectionState,
  SharedControlLogicalSubscription,
  SharedControlPendingRequest,
  SharedControlReadyWaiter
} from './remote-runtime-shared-control-types'

export function handleSharedControlTextFrame(args: {
  frame: string
  state: SharedControlConnectionState
  sharedKey: Uint8Array | null
  deviceToken: string
  environmentId?: string
  pendingRequests: Map<string, SharedControlPendingRequest<unknown>>
  subscriptions: Map<string, SharedControlLogicalSubscription<unknown>>
  retiredRequestIds: SharedControlRetiredRequestIds
  readyWaiters: SharedControlReadyWaiter[]
  setState: (state: SharedControlConnectionState) => void
  handleSocketClosed: (error: RemoteRuntimeClientError) => void
  sendEncrypted: (payload: unknown) => boolean
  markReady: () => void
  replaySubscriptions: () => void
}): void {
  if (args.state === 'awaiting_ready') {
    const error = parseReadyFrame(args.frame)
    if (error) {
      args.handleSocketClosed(error)
      return
    }
    args.setState('awaiting_authenticated')
    args.sendEncrypted({
      type: 'e2ee_auth',
      deviceToken: args.deviceToken,
      clientCapabilities: [
        SESSION_TAB_CLOSE_INTENT_RUNTIME_CAPABILITY,
        AGENT_SESSION_BOUNDARY_RUNTIME_CAPABILITY,
        WORKTREE_VISIBILITY_DEFAULTS_RUNTIME_CAPABILITY,
        WORKTREE_VISIBILITY_SOURCE_DEFAULTS_RUNTIME_CAPABILITY
      ]
    })
    return
  }

  const parsed = parseSharedControlFrame(args.frame, args.sharedKey, args.state)
  if (parsed.type === 'auth') {
    const error = parseAuthenticatedFrame(parsed.plaintext)
    if (error) {
      args.handleSocketClosed(error)
      return
    }
    args.setState('ready')
    args.markReady()
    resolveSharedControlReadyWaiters(args.readyWaiters)
    args.replaySubscriptions()
    return
  }

  if (parsed.type === 'error') {
    args.handleSocketClosed(parsed.error)
    return
  }

  dispatchSharedControlFrame({
    environmentId: args.environmentId,
    frame: parsed.frame,
    pendingRequests: args.pendingRequests,
    subscriptions: args.subscriptions,
    retiredRequestIds: args.retiredRequestIds,
    deviceToken: args.deviceToken,
    send: args.sendEncrypted
  })
}