From a626497f731392ae686a5b5390019c228cf5c512 Mon Sep 17 00:00:00 2001
From: tt-a1i <53142663+tt-a1i@users.noreply.github.com>
Date: Sat, 20 Jun 2026 01:13:28 +0800
Subject: [PATCH] fix(cli): keep keypress handlers current

---
 .../messages/AskUserQuestionDialog.test.tsx   | 51 ++++++++++++++++---
 packages/cli/src/ui/hooks/useKeypress.ts      | 15 ++++--
 2 files changed, 55 insertions(+), 11 deletions(-)

diff --git a/packages/cli/src/ui/components/messages/AskUserQuestionDialog.test.tsx b/packages/cli/src/ui/components/messages/AskUserQuestionDialog.test.tsx
index bdb7d652454..cf5e119be74 100644
--- a/packages/cli/src/ui/components/messages/AskUserQuestionDialog.test.tsx
+++ b/packages/cli/src/ui/components/messages/AskUserQuestionDialog.test.tsx
@@ -13,6 +13,29 @@ import stripAnsi from 'strip-ansi';
 
 const wait = (ms = 50) => new Promise((resolve) => setTimeout(resolve, ms));
 const clean = (value: string | undefined) => stripAnsi(value ?? '');
+const waitForFrame = async (
+  predicate: () => void,
+  options: { timeout?: number; interval?: number } = {},
+) => {
+  const { timeout = 1000, interval = 10 } = options;
+  const start = Date.now();
+  let lastError: unknown;
+
+  while (Date.now() - start < timeout) {
+    try {
+      predicate();
+      return;
+    } catch (error) {
+      lastError = error;
+    }
+    await wait(interval);
+  }
+
+  if (lastError) {
+    throw lastError;
+  }
+  throw new Error('waitForFrame timed out');
+};
 
 const createSingleQuestion = (
   overrides: Partial<
@@ -301,22 +324,36 @@ describe('<AskUserQuestionDialog />', () => {
       await wait();
 
       stdin.write('4'); // Select "Other" custom input
-      await wait(150);
-      expect(clean(lastFrame())).toContain('❯ 4.');
+      await waitForFrame(() => {
+        expect(clean(lastFrame())).toContain('❯ 4.');
+      });
+      await wait();
 
       stdin.write('j');
-      await wait(150);
+      await waitForFrame(() => {
+        const frame = clean(lastFrame());
+        expect(frame).toContain('❯ 4.');
+        expect(frame).toContain('j');
+      });
+
       stdin.write('k');
-      await wait(150);
-      expect(clean(lastFrame())).toContain('❯ 4.');
+      await waitForFrame(() => {
+        const frame = clean(lastFrame());
+        expect(frame).toContain('❯ 4.');
+        expect(frame).toContain('jk');
+      });
 
       stdin.write('\u0010'); // Ctrl+P
       await wait();
-      expect(clean(lastFrame())).toContain('❯ 3. Green');
+      await waitForFrame(() => {
+        expect(clean(lastFrame())).toContain('❯ 3. Green');
+      });
 
       stdin.write('\u000E'); // Ctrl+N
       await wait();
-      expect(clean(lastFrame())).toContain('❯ 4.');
+      await waitForFrame(() => {
+        expect(clean(lastFrame())).toContain('❯ 4.');
+      });
 
       unmount();
     });
diff --git a/packages/cli/src/ui/hooks/useKeypress.ts b/packages/cli/src/ui/hooks/useKeypress.ts
index 1ff3ae2778d..2ed1bc550ed 100644
--- a/packages/cli/src/ui/hooks/useKeypress.ts
+++ b/packages/cli/src/ui/hooks/useKeypress.ts
@@ -4,7 +4,7 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { useEffect } from 'react';
+import { useCallback, useEffect, useRef } from 'react';
 import type { KeypressHandler, Key } from '../contexts/KeypressContext.js';
 import { useKeypressContext } from '../contexts/KeypressContext.js';
 
@@ -22,15 +22,22 @@ export function useKeypress(
   { isActive }: { isActive: boolean },
 ) {
   const { subscribe, unsubscribe } = useKeypressContext();
+  const onKeypressRef = useRef(onKeypress);
+
+  onKeypressRef.current = onKeypress;
+
+  const handleKeypress = useCallback<KeypressHandler>((key) => {
+    onKeypressRef.current(key);
+  }, []);
 
   useEffect(() => {
     if (!isActive) {
       return;
     }
 
-    subscribe(onKeypress);
+    subscribe(handleKeypress);
     return () => {
-      unsubscribe(onKeypress);
+      unsubscribe(handleKeypress);
     };
-  }, [isActive, onKeypress, subscribe, unsubscribe]);
+  }, [isActive, handleKeypress, subscribe, unsubscribe]);
 }