Support paste markers split across writes. (#11977)

packages/cli/src/ui/components/InputPrompt.test.tsx

@@ -1331,7 +1331,7 @@ describe('InputPrompt', () => {
       await wait();
 
       stdin.write('\x1B');
-      await wait();
+      await wait(100);
 
       expect(props.buffer.setText).toHaveBeenCalledWith('');
       expect(mockCommandCompletion.resetCompletionState).toHaveBeenCalled();
@@ -1372,7 +1372,7 @@ describe('InputPrompt', () => {
       await wait();
 
       stdin.write('\x1B');
-      await wait();
+      await wait(100);
 
       expect(props.setShellModeActive).toHaveBeenCalledWith(false);
       unmount();
@@ -1392,7 +1392,7 @@ describe('InputPrompt', () => {
       await wait();
 
       stdin.write('\x1B');
-      await wait();
+      await wait(100);
 
       expect(mockCommandCompletion.resetCompletionState).toHaveBeenCalled();
       unmount();

packages/cli/src/ui/components/SettingsDialog.test.tsx

@@ -1348,7 +1348,7 @@ describe('SettingsDialog', () => {
 
       // Press Escape to exit
       stdin.write('\u001B');
-      await wait();
+      await wait(100);
 
       expect(onSelect).toHaveBeenCalledWith(undefined, 'User');
 

packages/cli/src/ui/contexts/KeypressContext.test.tsx

@@ -46,7 +46,7 @@ class MockStdin extends EventEmitter {
   pause = vi.fn();
 
   write(text: string) {
-    this.emit('data', Buffer.from(text));
+    this.emit('data', text);
   }
 }
 
@@ -381,6 +381,61 @@ describe('KeypressContext - Kitty Protocol', () => {
         }),
       );
     });
+    it('should paste start code split over multiple writes', async () => {
+      const keyHandler = vi.fn();
+      const pastedText = 'pasted content';
+
+      const { result } = renderHook(() => useKeypressContext(), { wrapper });
+
+      act(() => result.current.subscribe(keyHandler));
+
+      act(() => {
+        // Split PASTE_START into two parts
+        stdin.write(PASTE_START.slice(0, 3));
+        stdin.write(PASTE_START.slice(3));
+        stdin.write(pastedText);
+        stdin.write(PASTE_END);
+      });
+
+      await waitFor(() => {
+        expect(keyHandler).toHaveBeenCalledTimes(1);
+      });
+
+      expect(keyHandler).toHaveBeenCalledWith(
+        expect.objectContaining({
+          paste: true,
+          sequence: pastedText,
+        }),
+      );
+    });
+
+    it('should paste end code split over multiple writes', async () => {
+      const keyHandler = vi.fn();
+      const pastedText = 'pasted content';
+
+      const { result } = renderHook(() => useKeypressContext(), { wrapper });
+
+      act(() => result.current.subscribe(keyHandler));
+
+      act(() => {
+        stdin.write(PASTE_START);
+        stdin.write(pastedText);
+        // Split PASTE_END into two parts
+        stdin.write(PASTE_END.slice(0, 3));
+        stdin.write(PASTE_END.slice(3));
+      });
+
+      await waitFor(() => {
+        expect(keyHandler).toHaveBeenCalledTimes(1);
+      });
+
+      expect(keyHandler).toHaveBeenCalledWith(
+        expect.objectContaining({
+          paste: true,
+          sequence: pastedText,
+        }),
+      );
+    });
   });
 
   describe('debug keystroke logging', () => {

packages/cli/src/ui/contexts/KeypressContext.tsx

@@ -40,10 +40,11 @@ import {
 import { FOCUS_IN, FOCUS_OUT } from '../hooks/useFocus.js';
 
 const ESC = '\u001B';
-export const PASTE_MODE_PREFIX = `${ESC}[200~`;
-export const PASTE_MODE_SUFFIX = `${ESC}[201~`;
+export const PASTE_MODE_START = `${ESC}[200~`;
+export const PASTE_MODE_END = `${ESC}[201~`;
 export const DRAG_COMPLETION_TIMEOUT_MS = 100; // Broadcast full path after 100ms if no more input
 export const KITTY_SEQUENCE_TIMEOUT_MS = 50; // Flush incomplete kitty sequences after 50ms
+export const PASTE_CODE_TIMEOUT_MS = 50; // Flush incomplete paste code after 50ms
 export const SINGLE_QUOTE = "'";
 export const DOUBLE_QUOTE = '"';
 
@@ -353,6 +354,102 @@ function parseKittyPrefix(buffer: string): { key: Key; length: number } | null {
   return null;
 }
 
+/**
+ * Returns the first index before which we are certain there is no paste marker.
+ */
+function earliestPossiblePasteMarker(data: string): number {
+  // Check data for full start-paste or end-paste markers.
+  const startIndex = data.indexOf(PASTE_MODE_START);
+  const endIndex = data.indexOf(PASTE_MODE_END);
+  if (startIndex !== -1 && endIndex !== -1) {
+    return Math.min(startIndex, endIndex);
+  } else if (startIndex !== -1) {
+    return startIndex;
+  } else if (endIndex !== -1) {
+    return endIndex;
+  }
+
+  // data contains no full start-paste or end-paste.
+  // Check if data ends with a prefix of start-paste or end-paste.
+  const codeLength = PASTE_MODE_START.length;
+  for (let i = Math.min(data.length, codeLength - 1); i > 0; i--) {
+    const candidate = data.slice(data.length - i);
+    if (
+      PASTE_MODE_START.indexOf(candidate) === 0 ||
+      PASTE_MODE_END.indexOf(candidate) === 0
+    ) {
+      return data.length - i;
+    }
+  }
+  return data.length;
+}
+
+/**
+ * A generator that takes in data chunks and spits out paste-start and
+ * paste-end keypresses. All non-paste marker data is passed to passthrough.
+ */
+function* pasteMarkerParser(
+  passthrough: PassThrough,
+  keypressHandler: (_: unknown, key: Key) => void,
+): Generator<void, void, string> {
+  while (true) {
+    let data = yield;
+    if (data.length === 0) {
+      continue; // we timed out
+    }
+
+    while (true) {
+      const index = earliestPossiblePasteMarker(data);
+      if (index === data.length) {
+        // no possible paste markers were found
+        passthrough.write(data);
+        break;
+      }
+      if (index > 0) {
+        // snip off and send the part that doesn't have a paste marker
+        passthrough.write(data.slice(0, index));
+        data = data.slice(index);
+      }
+      // data starts with a possible paste marker
+      const codeLength = PASTE_MODE_START.length;
+      if (data.length < codeLength) {
+        // we have a prefix. Concat the next data and try again.
+        const newData = yield;
+        if (newData.length === 0) {
+          // we timed out. Just dump what we have and start over.
+          passthrough.write(data);
+          break;
+        }
+        data += newData;
+      } else if (data.startsWith(PASTE_MODE_START)) {
+        keypressHandler(undefined, {
+          name: 'paste-start',
+          ctrl: false,
+          meta: false,
+          shift: false,
+          paste: false,
+          sequence: '',
+        });
+        data = data.slice(PASTE_MODE_START.length);
+      } else if (data.startsWith(PASTE_MODE_END)) {
+        keypressHandler(undefined, {
+          name: 'paste-end',
+          ctrl: false,
+          meta: false,
+          shift: false,
+          paste: false,
+          sequence: '',
+        });
+        data = data.slice(PASTE_MODE_END.length);
+      } else {
+        // This should never happen.
+        passthrough.write(data);
+        break;
+      }
+    }
+  }
+}
+
 export interface Key {
   name: string;
   ctrl: boolean;
@@ -621,8 +718,8 @@ export function KeypressProvider({
         // Check if this could start a kitty sequence
         const startsWithEsc = key.sequence.startsWith(ESC);
         const isExcluded = [
-          PASTE_MODE_PREFIX,
-          PASTE_MODE_SUFFIX,
+          PASTE_MODE_START,
+          PASTE_MODE_END,
           FOCUS_IN,
           FOCUS_OUT,
         ].some((prefix) => key.sequence.startsWith(prefix));
@@ -766,57 +863,7 @@ export function KeypressProvider({
       broadcast({ ...key, paste: pasteBuffer !== null });
     };
 
-    const handleRawKeypress = (data: Buffer) => {
-      const pasteModePrefixBuffer = Buffer.from(PASTE_MODE_PREFIX);
-      const pasteModeSuffixBuffer = Buffer.from(PASTE_MODE_SUFFIX);
-
-      let pos = 0;
-      while (pos < data.length) {
-        const prefixPos = data.indexOf(pasteModePrefixBuffer, pos);
-        const suffixPos = data.indexOf(pasteModeSuffixBuffer, pos);
-        const isPrefixNext =
-          prefixPos !== -1 && (suffixPos === -1 || prefixPos < suffixPos);
-        const isSuffixNext =
-          suffixPos !== -1 && (prefixPos === -1 || suffixPos < prefixPos);
-
-        let nextMarkerPos = -1;
-        let markerLength = 0;
-
-        if (isPrefixNext) {
-          nextMarkerPos = prefixPos;
-        } else if (isSuffixNext) {
-          nextMarkerPos = suffixPos;
-        }
-        markerLength = pasteModeSuffixBuffer.length;
-
-        if (nextMarkerPos === -1) {
-          keypressStream!.write(data.slice(pos));
-          return;
-        }
-
-        const nextData = data.slice(pos, nextMarkerPos);
-        if (nextData.length > 0) {
-          keypressStream!.write(nextData);
-        }
-        const createPasteKeyEvent = (
-          name: 'paste-start' | 'paste-end',
-        ): Key => ({
-          name,
-          ctrl: false,
-          meta: false,
-          shift: false,
-          paste: false,
-          sequence: '',
-        });
-        if (isPrefixNext) {
-          handleKeypress(undefined, createPasteKeyEvent('paste-start'));
-        } else if (isSuffixNext) {
-          handleKeypress(undefined, createPasteKeyEvent('paste-end'));
-        }
-        pos = nextMarkerPos + markerLength;
-      }
-    };
-
+    let cleanup = () => {};
     let rl: readline.Interface;
     if (keypressStream !== null) {
       rl = readline.createInterface({
@@ -824,22 +871,35 @@ export function KeypressProvider({
         escapeCodeTimeout: 0,
       });
       readline.emitKeypressEvents(keypressStream, rl);
+
+      const parser = pasteMarkerParser(keypressStream, handleKeypress);
+      parser.next(); // prime the generator so it starts listening.
+      let timeoutId: NodeJS.Timeout;
+      const handleRawKeypress = (data: string) => {
+        clearTimeout(timeoutId);
+        parser.next(data);
+        timeoutId = setTimeout(() => parser.next(''), PASTE_CODE_TIMEOUT_MS);
+      };
+
       keypressStream.on('keypress', handleKeypress);
+      process.stdin.setEncoding('utf8'); // so handleRawKeypress gets strings
       stdin.on('data', handleRawKeypress);
+
+      cleanup = () => {
+        keypressStream.removeListener('keypress', handleKeypress);
+        stdin.removeListener('data', handleRawKeypress);
+      };
     } else {
       rl = readline.createInterface({ input: stdin, escapeCodeTimeout: 0 });
       readline.emitKeypressEvents(stdin, rl);
+
       stdin.on('keypress', handleKeypress);
+
+      cleanup = () => stdin.removeListener('keypress', handleKeypress);
     }
 
     return () => {
-      if (keypressStream !== null) {
-        keypressStream.removeListener('keypress', handleKeypress);
-        stdin.removeListener('data', handleRawKeypress);
-      } else {
-        stdin.removeListener('keypress', handleKeypress);
-      }
-
+      cleanup();
       rl.close();
 
       // Restore the terminal to its original state.

packages/cli/src/ui/hooks/useKeypress.test.ts

@@ -34,7 +34,7 @@ class MockStdin extends EventEmitter {
   pause = vi.fn();
 
   write(text: string) {
-    this.emit('data', Buffer.from(text));
+    this.emit('data', text);
   }
 }
 
@@ -187,6 +187,104 @@ describe('useKeypress', () => {
       expect(onKeypress).toHaveBeenCalledTimes(3);
     });
 
+    it('should handle lone pastes', () => {
+      renderHook(() => useKeypress(onKeypress, { isActive: true }), {
+        wrapper,
+      });
+
+      const pasteText = 'pasted';
+      act(() => {
+        stdin.write(PASTE_START);
+        stdin.write(pasteText);
+        stdin.write(PASTE_END);
+      });
+      expect(onKeypress).toHaveBeenCalledWith(
+        expect.objectContaining({ paste: true, sequence: pasteText }),
+      );
+
+      expect(onKeypress).toHaveBeenCalledTimes(1);
+    });
+
+    it('should handle paste false alarm', () => {
+      renderHook(() => useKeypress(onKeypress, { isActive: true }), {
+        wrapper,
+      });
+
+      act(() => {
+        stdin.write(PASTE_START.slice(0, 5));
+        stdin.write('do');
+      });
+      expect(onKeypress).toHaveBeenCalledWith(
+        expect.objectContaining({ code: '[200d' }),
+      );
+      expect(onKeypress).toHaveBeenCalledWith(
+        expect.objectContaining({ sequence: 'o' }),
+      );
+
+      expect(onKeypress).toHaveBeenCalledTimes(2);
+    });
+
+    it('should handle back to back pastes', () => {
+      renderHook(() => useKeypress(onKeypress, { isActive: true }), {
+        wrapper,
+      });
+
+      const pasteText1 = 'herp';
+      const pasteText2 = 'derp';
+      act(() => {
+        stdin.write(
+          PASTE_START +
+            pasteText1 +
+            PASTE_END +
+            PASTE_START +
+            pasteText2 +
+            PASTE_END,
+        );
+      });
+      expect(onKeypress).toHaveBeenCalledWith(
+        expect.objectContaining({ paste: true, sequence: pasteText1 }),
+      );
+      expect(onKeypress).toHaveBeenCalledWith(
+        expect.objectContaining({ paste: true, sequence: pasteText2 }),
+      );
+
+      expect(onKeypress).toHaveBeenCalledTimes(2);
+    });
+
+    it('should handle pastes split across writes', async () => {
+      renderHook(() => useKeypress(onKeypress, { isActive: true }), {
+        wrapper,
+      });
+
+      const keyA = { name: 'a', sequence: 'a' };
+      act(() => stdin.write('a'));
+      expect(onKeypress).toHaveBeenCalledWith(
+        expect.objectContaining({ ...keyA, paste: false }),
+      );
+
+      const pasteText = 'pasted';
+      await act(async () => {
+        stdin.write(PASTE_START.slice(0, 3));
+        await new Promise((r) => setTimeout(r, 50));
+        stdin.write(PASTE_START.slice(3) + pasteText.slice(0, 3));
+        await new Promise((r) => setTimeout(r, 50));
+        stdin.write(pasteText.slice(3) + PASTE_END.slice(0, 3));
+        await new Promise((r) => setTimeout(r, 50));
+        stdin.write(PASTE_END.slice(3));
+      });
+      expect(onKeypress).toHaveBeenCalledWith(
+        expect.objectContaining({ paste: true, sequence: pasteText }),
+      );
+
+      const keyB = { name: 'b', sequence: 'b' };
+      act(() => stdin.write('b'));
+      expect(onKeypress).toHaveBeenCalledWith(
+        expect.objectContaining({ ...keyB, paste: false }),
+      );
+
+      expect(onKeypress).toHaveBeenCalledTimes(3);
+    });
+
     it('should emit partial paste content if unmounted mid-paste', () => {
       const { unmount } = renderHook(
         () => useKeypress(onKeypress, { isActive: true }),