orca/src/shared/terminal-reply-query-scan.ts

import { findCsiFinalByteIndex } from './terminal-reply-query-extraction'
import { parseTerminalOscColorQuery } from './terminal-osc-color-reply'

const ESC = '\x1b'
const MAX_PENDING_QUERY_CHARS = 4096
/* oxlint-disable no-control-regex -- terminal query grammars contain ESC by definition */
const DEVICE_ATTRIBUTES_QUERY_RE = new RegExp('^\\u001b\\[[?>=]?[0-9;]*c$')
const MODE_QUERY_RE = new RegExp('^\\u001b\\[\\??[0-9;]+\\$p$')
/* oxlint-enable no-control-regex */

export type TerminalReplyQuerySequence = {
  data: string
  startSeq: number
  endSeq: number
}

export type TerminalReplyQueryScanState = {
  pending: string
  pendingStartSeq: number | null
}

export const EMPTY_TERMINAL_REPLY_QUERY_SCAN_STATE: TerminalReplyQueryScanState = {
  pending: '',
  pendingStartSeq: null
}

/**
 * Sequences worth replaying to a late-attaching view, because a scrollback snapshot
 * carries only rendered cells and drops the control bytes that produced them.
 *
 * Most entries are reply-eliciting queries the attaching emulator must answer once.
 * The DECSET/DECRST 2031 pair is different: since #9993 nothing answers the subscribe,
 * but a remote renderer derives its color-scheme subscription registry from these very
 * bytes (pty-connection observeLiveMode2031Chunk), so both toggles must be carried —
 * an arm without its withdraw would leave the client subscribed for a TUI that already
 * retired the subscription, and fish rearms/withdraws on every prompt.
 */
function isReplayableCsi(sequence: string): boolean {
  if (DEVICE_ATTRIBUTES_QUERY_RE.test(sequence)) {
    return true
  }
  if (MODE_QUERY_RE.test(sequence)) {
    return true
  }
  return (
    sequence === '\x1b[5n' ||
    sequence === '\x1b[6n' ||
    sequence === '\x1b[?6n' ||
    sequence === '\x1b[?996n' ||
    sequence === '\x1b[>q' ||
    sequence === '\x1b[14t' ||
    sequence === '\x1b[16t' ||
    sequence === '\x1b[18t' ||
    sequence === '\x1b[?u' ||
    sequence === '\x1b[?2031h' ||
    // Why paired with the arm above, exact-form only: replaying a combined DECRST such as
    // `CSI ?2004;2031l` would also toggle unrelated modes in the attaching emulator.
    sequence === '\x1b[?2031l'
  )
}

function boundedPending(input: string, startIndex: number): string {
  return input.slice(startIndex, startIndex + MAX_PENDING_QUERY_CHARS)
}

export function scanTerminalReplyQuerySequences(
  data: string,
  chunkStartSeq: number,
  previous: TerminalReplyQueryScanState
): { queries: TerminalReplyQuerySequence[]; state: TerminalReplyQueryScanState } {
  const continuesPending =
    previous.pendingStartSeq !== null &&
    previous.pendingStartSeq + previous.pending.length === chunkStartSeq
  const pending = continuesPending ? previous.pending : ''
  const input = pending + data
  const inputStartSeq = chunkStartSeq - pending.length
  const queries: TerminalReplyQuerySequence[] = []
  let offset = 0

  while (offset < input.length) {
    const candidateIndex = input.indexOf(ESC, offset)
    if (candidateIndex === -1) {
      break
    }
    if (candidateIndex + 1 >= input.length) {
      const nextPending = boundedPending(input, candidateIndex)
      return {
        queries,
        state: { pending: nextPending, pendingStartSeq: inputStartSeq + candidateIndex }
      }
    }

    let endIndex = -1
    let matches = false
    if (input.startsWith(`${ESC}[`, candidateIndex)) {
      endIndex = findCsiFinalByteIndex(input, candidateIndex + 2)
      if (endIndex !== -1) {
        matches = isReplayableCsi(input.slice(candidateIndex, endIndex + 1))
      }
    } else if (input.startsWith(`${ESC}]`, candidateIndex)) {
      const osc = parseTerminalOscColorQuery(input, candidateIndex)
      if (osc.kind === 'partial') {
        endIndex = -1
      } else if (osc.kind === 'match') {
        endIndex = osc.endIndex - 1
        matches = true
      } else {
        endIndex = candidateIndex + 1
      }
    } else if (input.startsWith(`${ESC}P`, candidateIndex)) {
      const terminatorIndex = input.indexOf(`${ESC}\\`, candidateIndex + 2)
      if (terminatorIndex !== -1) {
        endIndex = terminatorIndex + 1
        const body = input.slice(candidateIndex + 2, terminatorIndex)
        matches = body.startsWith('$q') || body.startsWith('+q')
      }
    } else {
      endIndex = candidateIndex
    }

    if (endIndex === -1) {
      const nextPending = boundedPending(input, candidateIndex)
      return {
        queries,
        state: { pending: nextPending, pendingStartSeq: inputStartSeq + candidateIndex }
      }
    }
    if (matches) {
      queries.push({
        data: input.slice(candidateIndex, endIndex + 1),
        startSeq: inputStartSeq + candidateIndex,
        endSeq: inputStartSeq + endIndex + 1
      })
    }
    offset = endIndex + 1
  }

  return { queries, state: EMPTY_TERMINAL_REPLY_QUERY_SCAN_STATE }
}