CoolFace
Datasetpublic

basant307/AI_Governance_Project

sourceHugging Faceapache-2.0updated 2mo agoView on Hugging Face
0likes45downloads
earlyInputCapture.ts384 linesDownload Raw Back to utils
1/**2 * @license3 * Copyright 2026 Qwen Team4 * SPDX-License-Identifier: Apache-2.05 */6 7/**8 * Early Input Capture - Capture user input during REPL initialization9 *10 * Principle: Start raw mode stdin listening at the earliest CLI entry point,11 * then inject buffered content when REPL is ready. Solves the problem of12 * user input being lost during startup.13 */14 15import { createDebugLogger } from '@qwen-code/qwen-code-core';16 17const debugLogger = createDebugLogger('EARLY_INPUT');18 19/** Maximum buffer size (64KB) */20const MAX_BUFFER_SIZE = 64 * 1024;21 22/**23 * Input buffer - collects chunks and concatenates on retrieval to avoid O(n^2) copies.24 */25interface InputBuffer {26  /** Collected raw byte chunks */27  chunks: Buffer[];28  /** Total bytes across all chunks */29  totalBytes: number;30  /** Whether capture is complete */31  captured: boolean;32}33 34let inputBuffer: InputBuffer = {35  chunks: [],36  totalBytes: 0,37  captured: false,38};39 40let captureHandler: ((data: Buffer) => void) | null = null;41let captureStdin: NodeJS.ReadStream | null = null;42let isCapturing = false;43let pendingTerminalResponse = Buffer.alloc(0);44 45type EscapeSequenceClassification = 'terminal' | 'user' | 'incomplete';46 47/**48 * Classify ESC sequences seen during startup capture.49 * - terminal: known terminal response/query payloads that should be filtered50 * - user: known user key sequences that should be preserved51 * - incomplete: prefix too short to classify yet, buffer for next chunk52 *53 * Note: User input function key sequences should be preserved:54 * - ESC [ A/B/C/D - Arrow keys55 * - ESC O P/Q/R/S - F1-F4 (SS3 sequences)56 * - ESC [ 1;5A - Ctrl+arrow and other modified keys57 */58function classifyEscapeSequence(59  data: Buffer,60  startIdx: number,61): EscapeSequenceClassification {62  if (startIdx >= data.length || data[startIdx] !== 0x1b) {63    return 'user';64  }65 66  const nextIdx = startIdx + 1;67  if (nextIdx >= data.length) {68    return 'incomplete';69  }70 71  const nextByte = data[nextIdx];72 73  // Check for special characters directly after ESC74  // P = 0x50 (DCS), _ = 0x5F (APC), ^ = 0x5E (PM), ] = 0x5D (OSC)75  // Note: O = 0x4F is SS3 sequence for function keys, should be preserved76  if (77    nextByte === 0x50 || // P (DCS)78    nextByte === 0x5f || // _ (APC)79    nextByte === 0x5e || // ^ (PM)80    nextByte === 0x5d // ] (OSC)81  ) {82    return 'terminal';83  }84 85  // Check for terminal responses in CSI sequences86  // ESC [ ? ... (DEC private mode response)87  // ESC [ > ... (DA2 response)88  if (nextByte === 0x5b) {89    // CSI sequence, check third character90    const thirdIdx = startIdx + 2;91    if (thirdIdx >= data.length) {92      return 'incomplete';93    }94    const thirdByte = data[thirdIdx];95    if (thirdByte === 0x3f || thirdByte === 0x3e) {96      // ESC [ ? or ESC [ > - this is a terminal response97      return 'terminal';98    }99    return 'user';100  }101 102  return 'user';103}104 105/**106 * Skip terminal response sequence107 * Returns the index position after skipping108 */109function skipTerminalResponse(110  data: Buffer,111  startIdx: number,112): { nextIndex: number; complete: boolean } {113  if (startIdx >= data.length || data[startIdx] !== 0x1b) {114    return { nextIndex: startIdx + 1, complete: true };115  }116 117  const nextIdx = startIdx + 1;118  if (nextIdx >= data.length) {119    return { nextIndex: nextIdx, complete: false };120  }121 122  const nextByte = data[nextIdx];123 124  // OSC sequence: ESC ] ... BEL or ESC ] ... ST125  if (nextByte === 0x5d) {126    let i = startIdx + 2;127    while (i < data.length) {128      // BEL (0x07) or ST (ESC \)129      if (data[i] === 0x07) {130        return { nextIndex: i + 1, complete: true };131      }132      if (data[i] === 0x1b && i + 1 < data.length && data[i + 1] === 0x5c) {133        return { nextIndex: i + 2, complete: true };134      }135      i++;136    }137    return { nextIndex: data.length, complete: false };138  }139 140  // DCS/APC/PM sequences: ESC P/_/^ ... ST141  if (nextByte === 0x50 || nextByte === 0x5f || nextByte === 0x5e) {142    let i = startIdx + 2;143    while (i < data.length) {144      // ST (ESC \)145      if (data[i] === 0x1b && i + 1 < data.length && data[i + 1] === 0x5c) {146        return { nextIndex: i + 2, complete: true };147      }148      i++;149    }150    return { nextIndex: data.length, complete: false };151  }152 153  // CSI sequence: ESC [ ... (ends with 0x40-0x7E)154  if (nextByte === 0x5b) {155    let i = startIdx + 2;156    while (i < data.length) {157      const byte = data[i];158      // CSI sequences end with 0x40-0x7E159      if (byte >= 0x40 && byte <= 0x7e) {160        return { nextIndex: i + 1, complete: true };161      }162      i++;163    }164    return { nextIndex: data.length, complete: false };165  }166 167  return { nextIndex: startIdx + 1, complete: true };168}169 170/**171 * Filter terminal response sequences (like Kitty protocol responses, device attributes, etc.)172 * Preserve user input (including function keys like arrow keys)173 */174function filterTerminalResponses(data: Buffer): {175  filtered: Buffer;176  trailingPartialTerminalResponse: Buffer;177} {178  const result = Buffer.allocUnsafe(data.length);179  let writeIdx = 0;180  let i = 0;181 182  while (i < data.length) {183    // Detect ESC sequences184    if (data[i] === 0x1b) {185      const sequenceType = classifyEscapeSequence(data, i);186      if (sequenceType === 'incomplete') {187        return {188          filtered: result.subarray(0, writeIdx),189          trailingPartialTerminalResponse: data.subarray(i),190        };191      }192      // Check if this is a terminal response (should be filtered out)193      if (sequenceType === 'terminal') {194        // Skip the terminal response sequence195        const skipResult = skipTerminalResponse(data, i);196        if (!skipResult.complete) {197          return {198            filtered: result.subarray(0, writeIdx),199            trailingPartialTerminalResponse: data.subarray(i),200          };201        }202        i = skipResult.nextIndex;203        continue;204      }205      // User input function keys (like arrow keys ESC [A), preserve206    }207    // Preserve current byte208    result[writeIdx++] = data[i];209    i++;210  }211 212  return {213    filtered: result.subarray(0, writeIdx),214    trailingPartialTerminalResponse: Buffer.alloc(0),215  };216}217 218/**219 * Decide whether pending trailing bytes should be replayed when capture stops.220 * Known terminal-response prefixes are dropped; user/ambiguous prefixes are kept.221 */222function shouldReplayPendingAtStop(pending: Buffer): boolean {223  if (pending.length === 0) {224    return false;225  }226  if (pending.length === 1 && pending[0] === 0x1b) {227    return true;228  }229  return classifyEscapeSequence(pending, 0) === 'user';230}231 232/**233 * Start early input capture234 * Call immediately after setting raw mode in gemini.tsx235 */236export function startEarlyInputCapture(): void {237  if (isCapturing || !process.stdin.isTTY) {238    if (!process.stdin.isTTY) {239      debugLogger.debug('Early input capture skipped: stdin is not a TTY');240    }241    return;242  }243 244  // Check if disabled245  if (process.env['QWEN_CODE_DISABLE_EARLY_CAPTURE'] === '1') {246    debugLogger.debug('Early input capture disabled by environment variable');247    return;248  }249 250  isCapturing = true;251  inputBuffer = {252    chunks: [],253    totalBytes: 0,254    captured: false,255  };256  pendingTerminalResponse = Buffer.alloc(0);257 258  debugLogger.debug('Starting early input capture');259 260  captureHandler = (data: Buffer) => {261    if (inputBuffer.captured) {262      return;263    }264 265    // Check buffer size limit266    if (inputBuffer.totalBytes >= MAX_BUFFER_SIZE) {267      debugLogger.warn(268        `Early input capture buffer full (${MAX_BUFFER_SIZE} bytes). Stopping capture; additional keystrokes during startup will be lost.`,269      );270      stopEarlyInputCapture();271      return;272    }273 274    const dataToFilter =275      pendingTerminalResponse.length > 0276        ? Buffer.concat([pendingTerminalResponse, data])277        : data;278    pendingTerminalResponse = Buffer.alloc(0);279 280    // Filter out terminal response sequences (like Kitty protocol responses)281    const { filtered, trailingPartialTerminalResponse } =282      filterTerminalResponses(dataToFilter);283    if (trailingPartialTerminalResponse.length > 0) {284      pendingTerminalResponse = Buffer.from(trailingPartialTerminalResponse);285    }286 287    if (filtered.length > 0) {288      // Limit buffer size289      const newLength = inputBuffer.totalBytes + filtered.length;290      if (newLength > MAX_BUFFER_SIZE) {291        const truncated = filtered.subarray(292          0,293          MAX_BUFFER_SIZE - inputBuffer.totalBytes,294        );295        inputBuffer.chunks.push(Buffer.from(truncated));296        inputBuffer.totalBytes += truncated.length;297        debugLogger.debug(`Buffer truncated at ${MAX_BUFFER_SIZE} bytes`);298      } else {299        inputBuffer.chunks.push(Buffer.from(filtered));300        inputBuffer.totalBytes += filtered.length;301        debugLogger.debug(302          `Captured ${filtered.length} bytes (total: ${inputBuffer.totalBytes})`,303        );304      }305    }306  };307 308  captureStdin = process.stdin;309  captureStdin.on('data', captureHandler);310}311 312/**313 * Stop early input capture314 * Call before KeypressProvider mounts315 */316export function stopEarlyInputCapture(): void {317  if (!isCapturing || !captureHandler || !captureStdin) {318    return;319  }320 321  captureStdin.removeListener('data', captureHandler);322  captureStdin = null;323  captureHandler = null;324  isCapturing = false;325  inputBuffer.captured = true;326 327  debugLogger.debug(328    `Stopped early input capture: ${inputBuffer.totalBytes} bytes`,329  );330}331 332/**333 * Get and clear captured input334 * For use by KeypressContext335 */336export function getAndClearCapturedInput(): Buffer {337  const parts = [...inputBuffer.chunks];338  if (shouldReplayPendingAtStop(pendingTerminalResponse)) {339    parts.push(Buffer.from(pendingTerminalResponse));340  }341  const buffer = parts.length > 0 ? Buffer.concat(parts) : Buffer.alloc(0);342  inputBuffer.chunks = [];343  inputBuffer.totalBytes = 0;344  pendingTerminalResponse = Buffer.alloc(0);345  // Keep captured=true — capture has completed, don't re-arm346  return buffer;347}348 349/**350 * Stop capture and return captured input in one atomic operation.351 * Preferred over calling stopEarlyInputCapture + getAndClearCapturedInput separately.352 */353export function stopAndGetCapturedInput(): Buffer {354  stopEarlyInputCapture();355  return getAndClearCapturedInput();356}357 358/**359 * Check if there is captured input360 */361export function hasCapturedInput(): boolean {362  return inputBuffer.totalBytes > 0;363}364 365/**366 * Reset capture state (for testing only)367 */368export function resetCaptureState(): void {369  if (captureHandler && captureStdin) {370    captureStdin.removeListener('data', captureHandler);371  } else if (captureHandler) {372    process.stdin.removeListener('data', captureHandler);373  }374  captureStdin = null;375  captureHandler = null;376  isCapturing = false;377  inputBuffer = {378    chunks: [],379    totalBytes: 0,380    captured: false,381  };382  pendingTerminalResponse = Buffer.alloc(0);383}384 
basant307/AI_Governance_Project · CoolFace