refactor(ui): extract switch controls logic into reusable hook and terminal restoration utility

Extracts duplicated exit/switch confirmation handling from CodexDisplay and RemoteModeDisplay into a custom useSwitchControls hook. Centralizes terminal state restoration (raw mode, keyboard protocol cleanup) into a restoreTerminalState utility function used across codex modules. Improves code reusability and maintainability.
2026-08-05 06:24:37 +00:00 · 2025-12-22 23:12:24 +08:00
parent 360e3a9919
commit 93810689ad
10 changed files with 316 additions and 139 deletions
@@ -13,6 +13,9 @@
    "cli": {
      "name": "hapi",
      "version": "0.12.0-1",
+      "bin": {
+        "hapi": "./src/index.ts",
+      },
      "dependencies": {
        "@modelcontextprotocol/sdk": "^1.22.0",
        "@types/cross-spawn": "^6.0.6",
@@ -40,6 +43,7 @@
        "dotenv": "^16.6.1",
        "eslint": "^9",
        "eslint-config-prettier": "^10",
+        "react-test-renderer": "^19.2.0",
        "shx": "^0.3.3",
        "ts-node": "^10",
        "tsx": "^4.20.6",
@@ -1515,6 +1519,8 @@

    "react-dom": ["react-dom@19.2.3", "", { "dependencies": { "scheduler": "^0.27.0" }, "peerDependencies": { "react": "^19.2.3" } }, "sha512-yELu4WmLPw5Mr/lmeEpox5rw3RETacE++JgHqQzd2dg+YbJuat3jH4ingc+WPZhxaoFzdv9y33G+F7Nl5O0GBg=="],

+    "react-is": ["react-is@19.2.3", "", {}, "sha512-qJNJfu81ByyabuG7hPFEbXqNcWSU3+eVus+KJs+0ncpGfMyYdvSmxiJxbWR65lYi1I+/0HBcliO029gc4F+PnA=="],
+
    "react-markdown": ["react-markdown@10.1.0", "", { "dependencies": { "@types/hast": "^3.0.0", "@types/mdast": "^4.0.0", "devlop": "^1.0.0", "hast-util-to-jsx-runtime": "^2.0.0", "html-url-attributes": "^3.0.0", "mdast-util-to-hast": "^13.0.0", "remark-parse": "^11.0.0", "remark-rehype": "^11.0.0", "unified": "^11.0.0", "unist-util-visit": "^5.0.0", "vfile": "^6.0.0" }, "peerDependencies": { "@types/react": ">=18", "react": ">=18" } }, "sha512-qKxVopLT/TyA6BX3Ue5NwabOsAzm0Q7kAPwq6L+wWDwisYs7R8vZ0nRXqq6rkueboxpkjvLGU9fWifiX/ZZFxQ=="],

    "react-reconciler": ["react-reconciler@0.33.0", "", { "dependencies": { "scheduler": "^0.27.0" }, "peerDependencies": { "react": "^19.2.0" } }, "sha512-KetWRytFv1epdpJc3J4G75I4WrplZE5jOL7Yq0p34+OVOKF4Se7WrdIdVC45XsSSmUTlht2FM/fM1FZb1mfQeA=="],
@@ -1527,6 +1533,8 @@

    "react-style-singleton": ["react-style-singleton@2.2.3", "", { "dependencies": { "get-nonce": "^1.0.0", "tslib": "^2.0.0" }, "peerDependencies": { "@types/react": "*", "react": "^16.8.0 || ^17.0.0 || ^18.0.0 || ^19.0.0 || ^19.0.0-rc" }, "optionalPeers": ["@types/react"] }, "sha512-b6jSvxvVnyptAiLjbkWLE/lOnR4lfTtDAl+eUC7RZy+QQWc6wRzIV2CE6xBuMmDxc2qIihtDCZD5NPOFl7fRBQ=="],

+    "react-test-renderer": ["react-test-renderer@19.2.3", "", { "dependencies": { "react-is": "^19.2.3", "scheduler": "^0.27.0" }, "peerDependencies": { "react": "^19.2.3" } }, "sha512-TMR1LnSFiWZMJkCgNf5ATSvAheTT2NvKIwiVwdBPHxjBI7n/JbWd4gaZ16DVd9foAXdvDz+sB5yxZTwMjPRxpw=="],
+
    "react-textarea-autosize": ["react-textarea-autosize@8.5.9", "", { "dependencies": { "@babel/runtime": "^7.20.13", "use-composed-ref": "^1.3.0", "use-latest": "^1.2.1" }, "peerDependencies": { "react": "^16.8.0 || ^17.0.0 || ^18.0.0 || ^19.0.0" } }, "sha512-U1DGlIQN5AwgjTyOEnI1oCcMuEr1pv1qOtklB2l4nyMGbHzWrI0eFsYK0zos2YWqAolJyG0IWJaqWmWj5ETh0A=="],

    "read-cache": ["read-cache@1.0.0", "", { "dependencies": { "pify": "^2.3.0" } }, "sha512-Owdv/Ft7IjOgm/i0xvNDZ1LrRANRfew4b2prF3OWMQLxLfu3bS8FVhCsrSCMK4lR56Y9ya+AThoTpDCTxCmpRA=="],
@@ -64,7 +64,8 @@
    "ts-node": "^10",
    "tsx": "^4.20.6",
    "typescript": "^5",
-    "vitest": "^3.2.4"
+    "vitest": "^3.2.4",
+    "react-test-renderer": "^19.2.0"
  },
  "resolutions": {
    "whatwg-url": "14.2.0",
@@ -1,5 +1,6 @@
 import { spawn } from 'node:child_process';
 import { logger } from '@/ui/logger';
+import { restoreTerminalState } from '@/ui/terminalState';

 export async function codexLocal(opts: {
    abort: AbortSignal;
@@ -91,5 +92,6 @@ export async function codexLocal(opts: {
        });
    } finally {
        process.stdin.resume();
+        restoreTerminalState();
    }
 }
@@ -19,6 +19,7 @@ import { startHappyServer } from '@/claude/utils/startHappyServer';
 import { emitReadyIfIdle } from './utils/emitReadyIfIdle';
 import type { CodexSession } from './session';
 import type { EnhancedMode } from './loop';
+import { restoreTerminalState } from '@/ui/terminalState';

 export async function codexRemoteLauncher(session: CodexSession): Promise<'switch' | 'exit'> {
    const hasTTY = process.stdout.isTTY && process.stdin.isTTY;
@@ -467,9 +468,7 @@ export async function codexRemoteLauncher(session: CodexSession): Promise<'switc
        reasoningProcessor.abort();
        diffProcessor.reset();

-        if (process.stdin.isTTY) {
-            try { process.stdin.setRawMode(false); } catch {}
-        }
+        restoreTerminalState();
        if (hasTTY) {
            try { process.stdin.pause(); } catch {}
        }
@@ -4,6 +4,7 @@ import { resolve } from 'node:path';

 import { ApiClient } from '@/api/api';
 import { logger } from '@/ui/logger';
+import { restoreTerminalState } from '@/ui/terminalState';
 import { loop, type EnhancedMode, type PermissionMode } from './loop';
 import { MessageQueue2 } from '@/utils/MessageQueue2';
 import { hashObject } from '@/utils/deterministicJson';
@@ -136,6 +137,7 @@ export async function runCodex(opts: {
        }
        cleanupStarted = true;
        logger.debug('[codex] Cleanup start');
+        restoreTerminalState();
        try {
            if (sessionWrapper) {
                sessionWrapper.stopKeepAlive();
@@ -1,6 +1,7 @@
-import React, { useState, useEffect, useRef, useCallback } from 'react'
-import { Box, Text, useStdout, useInput } from 'ink'
+import React, { useState, useEffect } from 'react'
+import { Box, Text, useStdout } from 'ink'
 import { MessageBuffer, type BufferedMessage } from './messageBuffer'
+import { useSwitchControls } from './useSwitchControls'

 interface CodexDisplayProps {
    messageBuffer: MessageBuffer
@@ -11,9 +12,10 @@ interface CodexDisplayProps {

 export const CodexDisplay: React.FC<CodexDisplayProps> = ({ messageBuffer, logPath, onExit, onSwitchToLocal }) => {
    const [messages, setMessages] = useState<BufferedMessage[]>([])
-    const [confirmationMode, setConfirmationMode] = useState<'exit' | 'switch' | null>(null)
-    const [actionInProgress, setActionInProgress] = useState<'exiting' | 'switching' | null>(null)
-    const confirmationTimeoutRef = useRef<NodeJS.Timeout | null>(null)
+    const { confirmationMode, actionInProgress } = useSwitchControls({
+        onExit,
+        onSwitch: onSwitchToLocal
+    })
    const { stdout } = useStdout()
    const terminalWidth = stdout.columns || 80
    const terminalHeight = stdout.rows || 24
@@ -27,69 +29,9 @@ export const CodexDisplay: React.FC<CodexDisplayProps> = ({ messageBuffer, logPa

        return () => {
            unsubscribe()
-            if (confirmationTimeoutRef.current) {
-                clearTimeout(confirmationTimeoutRef.current)
-            }
        }
    }, [messageBuffer])

-    const resetConfirmation = useCallback(() => {
-        setConfirmationMode(null)
-        if (confirmationTimeoutRef.current) {
-            clearTimeout(confirmationTimeoutRef.current)
-            confirmationTimeoutRef.current = null
-        }
-    }, [])
-
-    const setConfirmationWithTimeout = useCallback((mode: 'exit' | 'switch') => {
-        setConfirmationMode(mode)
-        if (confirmationTimeoutRef.current) {
-            clearTimeout(confirmationTimeoutRef.current)
-        }
-        confirmationTimeoutRef.current = setTimeout(() => {
-            resetConfirmation()
-        }, 15000) // 15 seconds timeout
-    }, [resetConfirmation])
-
-    useInput(useCallback(async (input, key) => {
-        // Don't process input if action is in progress
-        if (actionInProgress) return
-        
-        // Handle Ctrl-C - exits the agent directly instead of switching modes
-        if (key.ctrl && input === 'c') {
-            if (confirmationMode === 'exit') {
-                // Second Ctrl-C, exit
-                resetConfirmation()
-                setActionInProgress('exiting')
-                // Small delay to show the status message
-                await new Promise(resolve => setTimeout(resolve, 100))
-                onExit?.()
-            } else {
-                // First Ctrl-C, show confirmation
-                setConfirmationWithTimeout('exit')
-            }
-            return
-        }
-
-        const isSpace = input === ' ' || key.name === 'space'
-
-        if (isSpace && onSwitchToLocal) {
-            if (confirmationMode === 'switch') {
-                resetConfirmation()
-                setActionInProgress('switching')
-                await new Promise(resolve => setTimeout(resolve, 100))
-                onSwitchToLocal()
-            } else {
-                setConfirmationWithTimeout('switch')
-            }
-            return
-        }
-
-        if (confirmationMode) {
-            resetConfirmation()
-        }
-    }, [confirmationMode, actionInProgress, onExit, onSwitchToLocal, setConfirmationWithTimeout, resetConfirmation]))
-
    const getMessageColor = (type: BufferedMessage['type']): string => {
        switch (type) {
            case 'user': return 'magenta'
@@ -1,6 +1,7 @@
-import React, { useState, useEffect, useRef, useCallback } from 'react'
-import { Box, Text, useStdout, useInput } from 'ink'
+import React, { useState, useEffect } from 'react'
+import { Box, Text, useStdout } from 'ink'
 import { MessageBuffer, type BufferedMessage } from './messageBuffer'
+import { useSwitchControls } from './useSwitchControls'

 interface RemoteModeDisplayProps {
    messageBuffer: MessageBuffer
@@ -11,9 +12,10 @@ interface RemoteModeDisplayProps {

 export const RemoteModeDisplay: React.FC<RemoteModeDisplayProps> = ({ messageBuffer, logPath, onExit, onSwitchToLocal }) => {
    const [messages, setMessages] = useState<BufferedMessage[]>([])
-    const [confirmationMode, setConfirmationMode] = useState<'exit' | 'switch' | null>(null)
-    const [actionInProgress, setActionInProgress] = useState<'exiting' | 'switching' | null>(null)
-    const confirmationTimeoutRef = useRef<NodeJS.Timeout | null>(null)
+    const { confirmationMode, actionInProgress } = useSwitchControls({
+        onExit,
+        onSwitch: onSwitchToLocal
+    })
    const { stdout } = useStdout()
    const terminalWidth = stdout.columns || 80
    const terminalHeight = stdout.rows || 24
@@ -27,74 +29,9 @@ export const RemoteModeDisplay: React.FC<RemoteModeDisplayProps> = ({ messageBuf

        return () => {
            unsubscribe()
-            if (confirmationTimeoutRef.current) {
-                clearTimeout(confirmationTimeoutRef.current)
-            }
        }
    }, [messageBuffer])

-    const resetConfirmation = useCallback(() => {
-        setConfirmationMode(null)
-        if (confirmationTimeoutRef.current) {
-            clearTimeout(confirmationTimeoutRef.current)
-            confirmationTimeoutRef.current = null
-        }
-    }, [])
-
-    const setConfirmationWithTimeout = useCallback((mode: 'exit' | 'switch') => {
-        setConfirmationMode(mode)
-        if (confirmationTimeoutRef.current) {
-            clearTimeout(confirmationTimeoutRef.current)
-        }
-        confirmationTimeoutRef.current = setTimeout(() => {
-            resetConfirmation()
-        }, 15000) // 15 seconds timeout
-    }, [resetConfirmation])
-
-    useInput(useCallback(async (input, key) => {
-        // Don't process input if action is in progress
-        if (actionInProgress) return
-        
-        // Handle Ctrl-C
-        if (key.ctrl && input === 'c') {
-            if (confirmationMode === 'exit') {
-                // Second Ctrl-C, exit
-                resetConfirmation()
-                setActionInProgress('exiting')
-                // Small delay to show the status message
-                await new Promise(resolve => setTimeout(resolve, 100))
-                onExit?.()
-            } else {
-                // First Ctrl-C, show confirmation
-                setConfirmationWithTimeout('exit')
-            }
-            return
-        }
-
-        const isSpace = input === ' ' || key.name === 'space'
-
-        // Handle double space
-        if (isSpace) {
-            if (confirmationMode === 'switch') {
-                // Second space, switch to local
-                resetConfirmation()
-                setActionInProgress('switching')
-                // Small delay to show the status message
-                await new Promise(resolve => setTimeout(resolve, 100))
-                onSwitchToLocal?.()
-            } else {
-                // First space, show confirmation
-                setConfirmationWithTimeout('switch')
-            }
-            return
-        }
-
-        // Any other key cancels confirmation
-        if (confirmationMode) {
-            resetConfirmation()
-        }
-    }, [confirmationMode, actionInProgress, onExit, onSwitchToLocal, setConfirmationWithTimeout, resetConfirmation]))
-
    const getMessageColor = (type: BufferedMessage['type']): string => {
        switch (type) {
            case 'user': return 'magenta'
@@ -0,0 +1,164 @@
+import React, { useEffect } from 'react';
+import TestRenderer, { act, type ReactTestRenderer } from 'react-test-renderer';
+import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
+import { useSwitchControls, type ConfirmationMode, type ActionInProgress } from './useSwitchControls';
+
+type Key = {
+    ctrl?: boolean;
+    name?: string;
+    sequence?: string;
+};
+
+type SwitchState = {
+    confirmationMode: ConfirmationMode;
+    actionInProgress: ActionInProgress;
+};
+
+let inputHandler: ((input: string, key: Key) => void | Promise<void>) | null = null;
+
+vi.mock('ink', () => ({
+    useInput: (handler: (input: string, key: Key) => void | Promise<void>) => {
+        inputHandler = handler;
+    }
+}));
+
+function HookProbe(props: {
+    onExit?: () => void;
+    onSwitch?: () => void;
+    onState: (state: SwitchState) => void;
+}): null {
+    const state = useSwitchControls({
+        onExit: props.onExit,
+        onSwitch: props.onSwitch,
+        confirmationTimeoutMs: 5000
+    });
+
+    useEffect(() => {
+        props.onState(state);
+    }, [props.onState, state]);
+
+    return null;
+}
+
+describe('useSwitchControls', () => {
+    let renderer: ReactTestRenderer | null = null;
+    let latestState: SwitchState | null = null;
+
+    const mount = async (opts: { onExit?: () => void; onSwitch?: () => void }) => {
+        await act(async () => {
+            renderer = TestRenderer.create(
+                React.createElement(HookProbe, {
+                    ...opts,
+                    onState: (state) => {
+                        latestState = state;
+                    }
+                })
+            );
+        });
+    };
+
+    const triggerInput = async (input: string, key: Key) => {
+        if (!inputHandler) {
+            throw new Error('useInput handler was not registered');
+        }
+        await act(async () => {
+            await inputHandler?.(input, key);
+        });
+    };
+
+    const triggerInputWithTimers = async (input: string, key: Key, advanceMs: number) => {
+        if (!inputHandler) {
+            throw new Error('useInput handler was not registered');
+        }
+        await act(async () => {
+            const promise = inputHandler?.(input, key);
+            vi.advanceTimersByTime(advanceMs);
+            await promise;
+        });
+    };
+
+    const advanceTimers = async (advanceMs: number) => {
+        await act(async () => {
+            vi.advanceTimersByTime(advanceMs);
+        });
+    };
+
+    beforeEach(() => {
+        vi.useFakeTimers();
+        inputHandler = null;
+        latestState = null;
+    });
+
+    afterEach(() => {
+        if (renderer) {
+            act(() => {
+                renderer?.unmount();
+            });
+            renderer = null;
+        }
+        vi.runOnlyPendingTimers();
+        vi.useRealTimers();
+    });
+
+    it('forwards Ctrl-C to process when onExit is missing', async () => {
+        const onSwitch = vi.fn();
+        const killSpy = vi.spyOn(process, 'kill').mockImplementation(() => true);
+        await mount({ onSwitch });
+
+        await triggerInput(' ', { name: 'space' });
+        expect(latestState?.confirmationMode).toBe('switch');
+        expect(latestState?.actionInProgress).toBe(null);
+        expect(onSwitch).not.toHaveBeenCalled();
+
+        await triggerInput('c', { ctrl: true });
+        expect(killSpy).toHaveBeenCalledWith(process.pid, 'SIGINT');
+        expect(latestState?.confirmationMode).toBe(null);
+        expect(latestState?.actionInProgress).toBe(null);
+        killSpy.mockRestore();
+    });
+
+    it('confirms exit and invokes callback on second Ctrl-C', async () => {
+        const onExit = vi.fn();
+        await mount({ onExit });
+
+        await triggerInput('c', { ctrl: true });
+        expect(latestState?.confirmationMode).toBe('exit');
+        expect(latestState?.actionInProgress).toBe(null);
+
+        await triggerInputWithTimers('c', { ctrl: true }, 100);
+        expect(latestState?.actionInProgress).toBe('exiting');
+        expect(onExit).toHaveBeenCalledTimes(1);
+    });
+
+    it('ignores key-release sequences so confirmation stays visible', async () => {
+        const onSwitch = vi.fn();
+        await mount({ onSwitch });
+
+        await triggerInput(' ', { name: 'space' });
+        expect(latestState?.confirmationMode).toBe('switch');
+
+        await triggerInput('', { sequence: '\u001b[1:3u' });
+        expect(latestState?.confirmationMode).toBe('switch');
+        expect(latestState?.actionInProgress).toBe(null);
+    });
+
+    it('does not switch on key-release space sequences', async () => {
+        const onSwitch = vi.fn();
+        await mount({ onSwitch });
+
+        await triggerInput('', { sequence: '\u001b[3:3u', name: 'space' });
+        expect(onSwitch).not.toHaveBeenCalled();
+        expect(latestState?.confirmationMode).toBe(null);
+    });
+
+    it('clears confirmation after timeout', async () => {
+        const onSwitch = vi.fn();
+        await mount({ onSwitch });
+
+        await triggerInput(' ', { name: 'space' });
+        expect(latestState?.confirmationMode).toBe('switch');
+
+        await advanceTimers(5000);
+        expect(latestState?.confirmationMode).toBe(null);
+    });
+});
@@ -0,0 +1,108 @@
+import { useCallback, useEffect, useRef, useState } from 'react';
+import { useInput } from 'ink';
+
+export type ConfirmationMode = 'exit' | 'switch' | null;
+export type ActionInProgress = 'exiting' | 'switching' | null;
+
+export function useSwitchControls(opts: {
+    onExit?: () => void;
+    onSwitch?: () => void;
+    confirmationTimeoutMs?: number;
+}): {
+    confirmationMode: ConfirmationMode;
+    actionInProgress: ActionInProgress;
+} {
+    const [confirmationMode, setConfirmationMode] = useState<ConfirmationMode>(null);
+    const [actionInProgress, setActionInProgress] = useState<ActionInProgress>(null);
+    const confirmationTimeoutRef = useRef<NodeJS.Timeout | null>(null);
+    const { onExit, onSwitch } = opts;
+    const confirmationTimeoutMs = opts.confirmationTimeoutMs ?? 15000;
+
+    const resetConfirmation = useCallback(() => {
+        setConfirmationMode(null);
+        if (confirmationTimeoutRef.current) {
+            clearTimeout(confirmationTimeoutRef.current);
+            confirmationTimeoutRef.current = null;
+        }
+    }, []);
+
+    const setConfirmationWithTimeout = useCallback((mode: Exclude<ConfirmationMode, null>) => {
+        setConfirmationMode(mode);
+        if (confirmationTimeoutRef.current) {
+            clearTimeout(confirmationTimeoutRef.current);
+        }
+        confirmationTimeoutRef.current = setTimeout(() => {
+            resetConfirmation();
+        }, confirmationTimeoutMs);
+    }, [confirmationTimeoutMs, resetConfirmation]);
+
+    useEffect(() => {
+        return () => {
+            if (confirmationTimeoutRef.current) {
+                clearTimeout(confirmationTimeoutRef.current);
+            }
+        };
+    }, []);
+
+    useInput(useCallback(async (input, key) => {
+        if (actionInProgress) {
+            return;
+        }
+
+        if (key.ctrl && input === 'c') {
+            if (!onExit) {
+                if (confirmationMode) {
+                    resetConfirmation();
+                }
+                try {
+                    process.kill(process.pid, 'SIGINT');
+                } catch {
+                    process.exit(130);
+                }
+                return;
+            }
+            if (confirmationMode === 'exit') {
+                resetConfirmation();
+                setActionInProgress('exiting');
+                await new Promise(resolve => setTimeout(resolve, 100));
+                onExit();
+            } else {
+                setConfirmationWithTimeout('exit');
+            }
+            return;
+        }
+
+        const sequence = typeof key.sequence === 'string' ? key.sequence : input;
+        const isKeyRelease = typeof sequence === 'string' && /^\u001b\[[0-9;]*:3u$/.test(sequence);
+        const isSpace = Boolean(onSwitch) && !isKeyRelease && (input === ' ' || key.name === 'space');
+        const hasPrintableInput = typeof input === 'string' && input.length > 0;
+
+        if (isSpace) {
+            if (confirmationMode === 'switch') {
+                resetConfirmation();
+                setActionInProgress('switching');
+                await new Promise(resolve => setTimeout(resolve, 100));
+                onSwitch?.();
+            } else {
+                setConfirmationWithTimeout('switch');
+            }
+            return;
+        }
+
+        if (confirmationMode && hasPrintableInput && !isKeyRelease) {
+            resetConfirmation();
+        }
+    }, [
+        actionInProgress,
+        confirmationMode,
+        onExit,
+        onSwitch,
+        resetConfirmation,
+        setConfirmationWithTimeout
+    ]));
+
+    return {
+        confirmationMode,
+        actionInProgress
+    };
+}
@@ -0,0 +1,14 @@
+export function restoreTerminalState(): void {
+    if (process.stdout.isTTY) {
+        // Disable kitty keyboard protocol / CSI u key release reporting if enabled.
+        process.stdout.write('\x1b[>4;0m');
+        process.stdout.write('\x1b[?2004l');
+    }
+    if (process.stdin.isTTY) {
+        try {
+            process.stdin.setRawMode(false);
+        } catch {
+            // Ignore if raw mode is not supported.
+        }
+    }
+}