orca/src/main/daemon/client.ts

/* eslint-disable max-lines -- Why: daemon handshake, RPC, stream events, and reconnect cleanup share one socket lifecycle. */
import { connect, type Socket } from 'node:net'
import { readFileSync } from 'node:fs'
import { randomUUID } from 'node:crypto'
import { StringDecoder } from 'node:string_decoder'
import { encodeNdjson, createNdjsonParser } from './ndjson'
import {
  CLEAN_DISCONNECT_PROTOCOL_VERSION,
  PROTOCOL_VERSION,
  NOTIFY_PREFIX,
  DaemonProtocolError
} from './types'
import type {
  DaemonEndpointIdentity,
  HelloMessage,
  HelloResponse,
  RpcResponse,
  DaemonEvent
} from './types'
import { addNodePtyRecoveryHint } from './node-pty-error-hints'
import { decodeDaemonResponseError } from './daemon-errors'
import {
  CURRENT_DAEMON_HELLO_CAPABILITIES,
  parseDaemonHelloCapabilities,
  sameDaemonHelloCapabilities,
  TERMINAL_SESSION_AUTHORITY_HELLO_CAPABILITY,
  type DaemonHelloCapabilities
} from './daemon-hello-protocol'

const CONNECT_TIMEOUT_MS = 5000
const CONNECTION_ATTEMPT_WAIT_MS = CONNECT_TIMEOUT_MS * 4
const REQUEST_TIMEOUT_MS = 30000

export type DaemonClientOptions = {
  socketPath: string
  tokenPath: string
  protocolVersion?: number
  terminalSessionAuthorityConsumerProofReady?: () => boolean
}

type PendingRequest = {
  resolve: (value: unknown) => void
  reject: (reason: Error) => void
  timer: ReturnType<typeof setTimeout>
}

type AuthenticatedDaemonHello = Readonly<{
  identity: DaemonEndpointIdentity | null
  capabilities: DaemonHelloCapabilities
  tail: Buffer
}>

export class DaemonClient {
  private socketPath: string
  private tokenPath: string
  private protocolVersion: number
  private terminalSessionAuthorityConsumerProofReady: () => boolean
  private clientId = randomUUID()

  private controlSocket: Socket | null = null
  private streamSocket: Socket | null = null
  private connected = false
  private disconnectArmed = false
  // Why: after a disconnect + reconnect (daemon respawn), a stale 'close'
  // event from the old sockets can fire. Without a generation check, that
  // event would tear down the fresh connection. Each doConnect() increments
  // the generation; handleDisconnect ignores events from old generations.
  private connectionGeneration = 0
  // Why: multiple concurrent spawn() calls from simultaneous pane mounts
  // all call ensureConnected(). Without a lock, each starts a separate
  // connection attempt, overwriting sockets and triggering "Connection lost".
  private connectingPromise: Promise<void> | null = null
  private connectionAttemptGeneration = 0
  private daemonIdentity: DaemonEndpointIdentity | null = null
  private negotiatedCapabilities: DaemonHelloCapabilities = {}
  private observedAuthenticatedDisconnect = false

  private pendingRequests = new Map<string, PendingRequest>()
  private eventListeners: ((event: unknown) => void)[] = []
  private disconnectedListeners: (() => void)[] = []
  private requestCounter = 0
  private cleanupSocketListeners: (() => void) | null = null

  constructor(opts: DaemonClientOptions) {
    this.socketPath = opts.socketPath
    this.tokenPath = opts.tokenPath
    this.protocolVersion = opts.protocolVersion ?? PROTOCOL_VERSION
    this.terminalSessionAuthorityConsumerProofReady =
      opts.terminalSessionAuthorityConsumerProofReady ?? (() => false)
  }

  isConnected(): boolean {
    return this.connected
  }

  getDaemonIdentity(): DaemonEndpointIdentity | null {
    return this.daemonIdentity ? { ...this.daemonIdentity } : null
  }

  hasObservedAuthenticatedDisconnect(): boolean {
    return this.observedAuthenticatedDisconnect
  }

  supportsTerminalSessionAuthority(): boolean {
    return (
      this.connected &&
      this.negotiatedCapabilities.terminalSessionAuthority ===
        TERMINAL_SESSION_AUTHORITY_HELLO_CAPABILITY
    )
  }

  terminalSessionAuthorityConsumerProofHostId(): string | null {
    const capability = this.negotiatedCapabilities.terminalAuthorityConsumerProof
    return capability && 'version' in capability ? capability.authorityHostId : null
  }

  terminalSessionAuthorityConsumerRetirementSupported(): boolean {
    const capability = this.negotiatedCapabilities.terminalAuthorityConsumerProof
    return Boolean(capability && 'version' in capability && capability.retirementVersion === 1)
  }

  async ensureConnected(): Promise<void> {
    return this.ensureConnectedWithTimeout(CONNECT_TIMEOUT_MS, false)
  }

  async ensureConnectedWithin(timeoutMs: number): Promise<void> {
    return this.ensureConnectedWithTimeout(timeoutMs, true)
  }

  private async ensureConnectedWithTimeout(
    timeoutMs: number,
    sharedBudget: boolean
  ): Promise<void> {
    if (this.connected) {
      return
    }
    if (this.connectingPromise) {
      // Why: a normal connection may legitimately consume one timeout for each
      // socket and hello; bounded teardown calls instead keep their one shared budget.
      const waiterTimeoutMs = sharedBudget ? timeoutMs : CONNECTION_ATTEMPT_WAIT_MS
      return this.waitForConnectionAttempt(this.connectingPromise, waiterTimeoutMs)
    }

    const attemptGeneration = this.connectionAttemptGeneration
    this.connectingPromise = this.doConnect(timeoutMs, attemptGeneration, sharedBudget)
    try {
      await this.connectingPromise
    } finally {
      this.connectingPromise = null
    }
  }

  private async doConnect(
    timeoutMs: number,
    attemptGeneration: number,
    sharedBudget: boolean
  ): Promise<void> {
    const token = readFileSync(this.tokenPath, 'utf-8').trim()
    const deadlineMs = Date.now() + timeoutMs
    const remainingMs = (): number =>
      sharedBudget ? Math.max(1, deadlineMs - Date.now()) : timeoutMs
    const pendingListenerCleanups: (() => void)[] = []
    const cleanupPendingListeners = (): void => {
      for (const cleanup of pendingListenerCleanups.splice(0)) {
        cleanup()
      }
    }

    try {
      const requestedCapabilities = this.requestedHelloCapabilities()
      // Sequential: control first, then stream
      const pendingControlSocket = await this.connectSocket(remainingMs())
      this.assertConnectionAttemptCurrent(attemptGeneration, pendingControlSocket)
      this.controlSocket = pendingControlSocket
      const controlHello = await this.sendHello(
        this.controlSocket,
        token,
        'control',
        remainingMs(),
        requestedCapabilities
      )
      this.assertConnectionAttemptCurrent(attemptGeneration, this.controlSocket)
      pendingListenerCleanups.push(this.setupControlParser(this.controlSocket, controlHello.tail))

      const pendingStreamSocket = await this.connectSocket(remainingMs())
      this.assertConnectionAttemptCurrent(attemptGeneration, pendingStreamSocket)
      this.streamSocket = pendingStreamSocket
      const streamHello = await this.sendHello(
        this.streamSocket,
        token,
        'stream',
        remainingMs(),
        requestedCapabilities
      )
      this.assertConnectionAttemptCurrent(attemptGeneration, this.streamSocket)
      if (!sameDaemonIdentity(controlHello.identity, streamHello.identity)) {
        throw new DaemonProtocolError('Daemon identity changed during connection')
      }
      if (!sameDaemonHelloCapabilities(controlHello.capabilities, streamHello.capabilities)) {
        throw new DaemonProtocolError('Daemon capabilities changed during connection')
      }
      pendingListenerCleanups.push(this.setupStreamParser(this.streamSocket, streamHello.tail))

      this.assertConnectionAttemptCurrent(attemptGeneration)
      this.connected = true
      this.observedAuthenticatedDisconnect = false
      this.daemonIdentity = controlHello.identity
      this.negotiatedCapabilities = controlHello.capabilities
      this.disconnectArmed = true
      this.connectionGeneration++

      const gen = this.connectionGeneration
      const handleClose = () => this.handleDisconnect(gen)
      const controlSocket = this.controlSocket
      const streamSocket = this.streamSocket
      controlSocket.on('close', handleClose)
      controlSocket.on('error', handleClose)
      streamSocket.on('close', handleClose)
      streamSocket.on('error', handleClose)
      pendingListenerCleanups.push(() => {
        controlSocket.off('close', handleClose)
        controlSocket.off('error', handleClose)
        streamSocket.off('close', handleClose)
        streamSocket.off('error', handleClose)
      })
      this.cleanupSocketListeners = cleanupPendingListeners
    } catch (error) {
      cleanupPendingListeners()
      this.controlSocket?.destroy()
      this.streamSocket?.destroy()
      this.controlSocket = null
      this.streamSocket = null
      this.connected = false
      this.daemonIdentity = null
      this.negotiatedCapabilities = {}
      this.disconnectArmed = false
      throw error
    }
  }

  async request<T = unknown>(
    type: string,
    payload: unknown,
    timeoutMs = REQUEST_TIMEOUT_MS
  ): Promise<T> {
    if (!this.connected || !this.controlSocket) {
      throw new DaemonProtocolError('Not connected')
    }

    const id = `req-${++this.requestCounter}`
    const msg = { id, type, ...(payload !== undefined ? { payload } : {}) }
    const encoded = encodeNdjson(msg)

    return new Promise<T>((resolve, reject) => {
      const timer = setTimeout(() => {
        this.pendingRequests.delete(id)
        reject(new DaemonProtocolError(`Request ${type} timed out after ${timeoutMs}ms`))
      }, timeoutMs)

      this.pendingRequests.set(id, {
        resolve: resolve as (value: unknown) => void,
        reject,
        timer
      })

      this.controlSocket!.write(encoded)
    })
  }

  // Why: fire-and-forget writes need a local delivery signal to trigger dead-endpoint recovery.
  notify(type: string, payload: unknown): boolean {
    if (!this.connected || !this.controlSocket) {
      return false
    }

    const id = `${NOTIFY_PREFIX}${++this.requestCounter}`
    const msg = { id, type, ...(payload !== undefined ? { payload } : {}) }
    try {
      this.controlSocket.write(encodeNdjson(msg))
      return true
    } catch {
      // Notifications are best-effort; an oversized payload must not tear down the caller.
      return false
    }
  }

  onEvent(listener: (event: unknown) => void): () => void {
    this.eventListeners.push(listener)
    return () => {
      const idx = this.eventListeners.indexOf(listener)
      if (idx !== -1) {
        this.eventListeners.splice(idx, 1)
      }
    }
  }

  onDisconnected(listener: () => void): () => void {
    this.disconnectedListeners.push(listener)
    return () => {
      const idx = this.disconnectedListeners.indexOf(listener)
      if (idx !== -1) {
        this.disconnectedListeners.splice(idx, 1)
      }
    }
  }

  disconnect(): void {
    this.connectionAttemptGeneration++
    this.connected = false
    this.daemonIdentity = null
    this.negotiatedCapabilities = {}
    this.disconnectArmed = false
    this.cleanupActiveSocketListeners()

    for (const [id, pending] of this.pendingRequests) {
      clearTimeout(pending.timer)
      pending.reject(new DaemonProtocolError('Disconnected'))
      this.pendingRequests.delete(id)
    }

    this.controlSocket?.destroy()
    this.streamSocket?.destroy()
    this.controlSocket = null
    this.streamSocket = null
  }

  private connectSocket(timeoutMs: number): Promise<Socket> {
    return new Promise((resolve, reject) => {
      const socket = connect(this.socketPath)
      const cleanup = (): void => {
        clearTimeout(timer)
        socket.removeListener('connect', onConnect)
        socket.removeListener('error', onError)
      }
      const onConnect = (): void => {
        cleanup()
        resolve(socket)
      }
      const onError = (err: Error): void => {
        cleanup()
        reject(err)
      }
      const timer = setTimeout(() => {
        cleanup()
        socket.destroy()
        reject(new DaemonProtocolError('Connection timed out'))
      }, timeoutMs)

      socket.on('connect', onConnect)
      socket.on('error', onError)
    })
  }

  private waitForConnectionAttempt(attempt: Promise<void>, timeoutMs: number): Promise<void> {
    return new Promise((resolve, reject) => {
      const timer = setTimeout(() => {
        reject(new DaemonProtocolError('Connection attempt wait timed out'))
      }, timeoutMs)
      attempt.then(
        () => {
          clearTimeout(timer)
          resolve()
        },
        (error) => {
          clearTimeout(timer)
          reject(error)
        }
      )
    })
  }

  private assertConnectionAttemptCurrent(attemptGeneration: number, socket?: Socket): void {
    if (attemptGeneration === this.connectionAttemptGeneration) {
      return
    }
    socket?.destroy()
    throw new DaemonProtocolError('Disconnected')
  }

  private sendHello(
    socket: Socket,
    token: string,
    role: 'control' | 'stream',
    timeoutMs: number,
    requestedCapabilities: DaemonHelloCapabilities = this.requestedHelloCapabilities()
  ): Promise<AuthenticatedDaemonHello> {
    return new Promise((resolve, reject) => {
      const hello: HelloMessage = {
        type: 'hello',
        version: this.protocolVersion,
        token,
        clientId: this.clientId,
        role,
        ...(requestedCapabilities.terminalSessionAuthority !== undefined ||
        requestedCapabilities.terminalAuthorityConsumerProof !== undefined
          ? { capabilities: requestedCapabilities }
          : {})
      }

      let buffer = Buffer.alloc(0)
      let settled = false
      let timer: ReturnType<typeof setTimeout> | null = null
      const cleanup = (): void => {
        if (timer) {
          clearTimeout(timer)
          timer = null
        }
        socket.removeListener('data', onData)
        socket.removeListener('error', onError)
        socket.removeListener('close', onClose)
      }
      const finish = (
        error?: Error,
        hello: AuthenticatedDaemonHello = {
          identity: null,
          capabilities: {},
          tail: Buffer.alloc(0)
        }
      ): void => {
        if (settled) {
          return
        }
        settled = true
        cleanup()
        if (error) {
          reject(error)
          return
        }
        resolve(hello)
      }
      const onData = (chunk: Buffer): void => {
        buffer = Buffer.concat([buffer, chunk])
        const newlineIdx = buffer.indexOf(0x0a)
        if (newlineIdx === -1) {
          return
        }

        const line = buffer.subarray(0, newlineIdx).toString('utf8')
        const tail = buffer.subarray(newlineIdx + 1)
        try {
          const response = JSON.parse(line) as HelloResponse
          if (response.ok) {
            const identity = parseDaemonEndpointIdentity(response.daemonIdentity)
            if (
              (this.protocolVersion >= CLEAN_DISCONNECT_PROTOCOL_VERSION && identity === null) ||
              (response.daemonIdentity !== undefined && identity === null)
            ) {
              finish(new DaemonProtocolError('Invalid daemon identity'))
              return
            }
            finish(undefined, {
              identity,
              capabilities: parseDaemonHelloCapabilities(response.capabilities),
              tail
            })
          } else {
            finish(
              new DaemonProtocolError(addNodePtyRecoveryHint(response.error ?? 'Hello rejected'))
            )
          }
        } catch {
          finish(new DaemonProtocolError('Invalid hello response'))
        }
      }
      const onError = (error: Error): void => finish(error)
      const onClose = (): void =>
        finish(new DaemonProtocolError('Connection closed before hello response'))

      timer = setTimeout(() => {
        // Why: a stale daemon can accept the socket but never answer hello;
        // without a handshake timeout, startup waits forever on ensureConnected().
        finish(new DaemonProtocolError('Hello response timed out'))
        socket.destroy()
      }, timeoutMs)
      socket.on('data', onData)
      socket.on('error', onError)
      socket.on('close', onClose)
      socket.write(encodeNdjson(hello))
    })
  }

  private requestedHelloCapabilities(): DaemonHelloCapabilities {
    try {
      // Why nothing is offered without a proof keypair: authority admission is proof-only, so a
      // client that cannot prove itself must negotiate the legacy non-authority path instead.
      if (!this.terminalSessionAuthorityConsumerProofReady()) {
        return {}
      }
      return {
        terminalSessionAuthority: CURRENT_DAEMON_HELLO_CAPABILITIES.terminalSessionAuthority,
        terminalAuthorityConsumerProof:
          CURRENT_DAEMON_HELLO_CAPABILITIES.terminalAuthorityConsumerProof
      }
    } catch {
      return {}
    }
  }

  private setupControlParser(socket: Socket, initialData: Buffer): () => void {
    // Why: control responses may contain terminal/startup data with multibyte
    // text; keep incomplete UTF-8 bytes until the next socket chunk.
    const decoder = new StringDecoder('utf8')
    const parser = createNdjsonParser(
      (msg) => {
        const response = msg as RpcResponse
        if (response.id) {
          const pending = this.pendingRequests.get(response.id)
          if (pending) {
            this.pendingRequests.delete(response.id)
            clearTimeout(pending.timer)
            if (response.ok) {
              pending.resolve(response.payload)
            } else {
              const decoded = decodeDaemonResponseError(response.error)
              pending.reject(
                decoded instanceof DaemonProtocolError
                  ? new DaemonProtocolError(addNodePtyRecoveryHint(response.error))
                  : decoded
              )
            }
          }
        }
      },
      () => {} // Ignore parse errors on control socket
    )

    const onData = (chunk: Buffer) => parser.feed(decoder.write(chunk))
    socket.on('data', onData)
    if (initialData.length > 0) {
      parser.feed(decoder.write(initialData))
    }
    return () => socket.off('data', onData)
  }

  private setupStreamParser(socket: Socket, initialData: Buffer): () => void {
    // Why: PTY output streams include emoji/box-drawing tables; socket chunks
    // can split those UTF-8 sequences across packets.
    const decoder = new StringDecoder('utf8')
    const parser = createNdjsonParser(
      (msg) => {
        const event = msg as DaemonEvent
        if (event.type === 'event') {
          for (const listener of this.eventListeners) {
            listener(event)
          }
        }
      },
      () => {} // Ignore parse errors on stream socket
    )

    const onData = (chunk: Buffer) => parser.feed(decoder.write(chunk))
    socket.on('data', onData)
    if (initialData.length > 0) {
      parser.feed(decoder.write(initialData))
    }
    return () => socket.off('data', onData)
  }

  private handleDisconnect(generation: number): void {
    if (!this.disconnectArmed || generation !== this.connectionGeneration) {
      return
    }
    this.disconnectArmed = false
    this.connectionAttemptGeneration++
    if (this.daemonIdentity) {
      this.observedAuthenticatedDisconnect = true
    }
    this.connected = false
    this.daemonIdentity = null
    this.negotiatedCapabilities = {}
    this.cleanupActiveSocketListeners()

    for (const [id, pending] of this.pendingRequests) {
      clearTimeout(pending.timer)
      pending.reject(new DaemonProtocolError('Connection lost'))
      this.pendingRequests.delete(id)
    }

    this.controlSocket?.destroy()
    this.streamSocket?.destroy()
    this.controlSocket = null
    this.streamSocket = null

    for (const listener of this.disconnectedListeners) {
      listener()
    }
  }

  private cleanupActiveSocketListeners(): void {
    const cleanup = this.cleanupSocketListeners
    this.cleanupSocketListeners = null
    cleanup?.()
  }
}

function parseDaemonEndpointIdentity(value: unknown): DaemonEndpointIdentity | null {
  if (!value || typeof value !== 'object') {
    return null
  }
  const identity = value as {
    pid?: unknown
    startedAtMs?: unknown
    launchNonce?: unknown
    entryPath?: unknown
    appVersion?: unknown
    spawnerExecPath?: unknown
  }
  if (
    !Number.isSafeInteger(identity.pid) ||
    (identity.pid as number) <= 0 ||
    typeof identity.startedAtMs !== 'number' ||
    !Number.isFinite(identity.startedAtMs) ||
    identity.startedAtMs <= 0 ||
    typeof identity.launchNonce !== 'string' ||
    identity.launchNonce.length === 0
  ) {
    return null
  }
  return {
    pid: identity.pid as number,
    startedAtMs: identity.startedAtMs,
    launchNonce: identity.launchNonce,
    ...(typeof identity.entryPath === 'string' && identity.entryPath.length > 0
      ? { entryPath: identity.entryPath }
      : {}),
    ...(typeof identity.appVersion === 'string' && identity.appVersion.length > 0
      ? { appVersion: identity.appVersion }
      : {}),
    ...(typeof identity.spawnerExecPath === 'string' && identity.spawnerExecPath.length > 0
      ? { spawnerExecPath: identity.spawnerExecPath }
      : {})
  }
}

function sameDaemonIdentity(
  left: DaemonEndpointIdentity | null,
  right: DaemonEndpointIdentity | null
): boolean {
  return (
    (left === null && right === null) ||
    (left !== null &&
      right !== null &&
      left.pid === right.pid &&
      left.startedAtMs === right.startedAtMs &&
      left.launchNonce === right.launchNonce)
  )
}