diff --git a/server/modules/websocket/services/shell-websocket.service.ts b/server/modules/websocket/services/shell-websocket.service.ts index b8f34ae8..a29959be 100644 --- a/server/modules/websocket/services/shell-websocket.service.ts +++ b/server/modules/websocket/services/shell-websocket.service.ts @@ -18,6 +18,7 @@ type ShellIncomingMessage = { provider?: string; initialCommand?: string; isPlainShell?: boolean; + forceRestart?: boolean; }; type PtySessionEntry = { @@ -180,6 +181,7 @@ export function handleShellConnection( const hasSession = readBoolean(data.hasSession); const provider = readString(data.provider, 'claude'); const initialCommand = readString(data.initialCommand); + const forceRestart = readBoolean(data.forceRestart); const isPlainShell = readBoolean(data.isPlainShell) || (!!initialCommand && !hasSession) || @@ -200,7 +202,7 @@ export function handleShellConnection( : ''; ptySessionKey = `${projectPath}_${sessionId ?? 'default'}${commandSuffix}`; - if (isLoginCommand) { + if (isLoginCommand || forceRestart) { const oldSession = ptySessionsMap.get(ptySessionKey); if (oldSession) { if (oldSession.timeoutId) { @@ -211,7 +213,8 @@ export function handleShellConnection( } } - const existingSession = isLoginCommand ? null : ptySessionsMap.get(ptySessionKey); + const existingSession = + isLoginCommand || forceRestart ? null : ptySessionsMap.get(ptySessionKey); if (existingSession) { shellProcess = existingSession.pty; if (existingSession.timeoutId) { @@ -368,6 +371,10 @@ export function handleShellConnection( } const session = ptySessionsMap.get(ptySessionKey); + if (session && session.pty !== shellProcess) { + return; + } + if (session && session.ws && session.ws.readyState === WebSocket.OPEN) { session.ws.send( JSON.stringify({ @@ -451,6 +458,10 @@ export function handleShellConnection( session.ws = null; session.timeoutId = setTimeout(() => { + if (ptySessionsMap.get(ptySessionKey as string) !== session) { + return; + } + session.pty.kill(); ptySessionsMap.delete(ptySessionKey as string); }, PTY_SESSION_TIMEOUT); diff --git a/src/components/shell/hooks/useShellConnection.ts b/src/components/shell/hooks/useShellConnection.ts index 7babed15..918ed76c 100644 --- a/src/components/shell/hooks/useShellConnection.ts +++ b/src/components/shell/hooks/useShellConnection.ts @@ -2,6 +2,7 @@ import { useCallback, useEffect, useRef, useState } from 'react'; import type { MutableRefObject } from 'react'; import type { FitAddon } from '@xterm/addon-fit'; import type { Terminal } from '@xterm/xterm'; + import type { Project, ProjectSession } from '../../../types/app'; import { TERMINAL_INIT_DELAY_MS } from '../constants/constants'; import { getShellWebSocketUrl, parseShellMessage, sendSocketMessage } from '../utils/socket'; @@ -31,8 +32,8 @@ type UseShellConnectionResult = { isConnected: boolean; isConnecting: boolean; closeSocket: () => void; - connectToShell: () => void; - disconnectFromShell: () => void; + connectToShell: (options?: { forceRestart?: boolean }) => void; + disconnectFromShell: (options?: { suppressAutoConnect?: boolean }) => void; }; export function useShellConnection({ @@ -54,6 +55,8 @@ export function useShellConnection({ const [isConnected, setIsConnected] = useState(false); const [isConnecting, setIsConnecting] = useState(false); const connectingRef = useRef(false); + const forceRestartOnInitRef = useRef(false); + const suppressAutoConnectRef = useRef(false); const handleProcessCompletion = useCallback( (output: string) => { @@ -141,6 +144,8 @@ export function useShellConnection({ } currentFitAddon.fit(); + const forceRestart = forceRestartOnInitRef.current; + forceRestartOnInitRef.current = false; sendSocketMessage(socket, { type: 'init', @@ -152,6 +157,7 @@ export function useShellConnection({ rows: currentTerminal.rows, initialCommand: initialCommandRef.current, isPlainShell: isPlainShellRef.current, + forceRestart, }); }, TERMINAL_INIT_DELAY_MS); }; @@ -177,6 +183,7 @@ export function useShellConnection({ setIsConnected(false); setIsConnecting(false); connectingRef.current = false; + forceRestartOnInitRef.current = false; } }, [ @@ -195,27 +202,40 @@ export function useShellConnection({ ], ); - const connectToShell = useCallback(() => { + const connectToShell = useCallback((options?: { forceRestart?: boolean }) => { if (!isInitialized || isConnected || isConnecting || connectingRef.current) { return; } + forceRestartOnInitRef.current = Boolean(options?.forceRestart); + suppressAutoConnectRef.current = false; connectingRef.current = true; setIsConnecting(true); connectWebSocket(true); }, [connectWebSocket, isConnected, isConnecting, isInitialized]); - const disconnectFromShell = useCallback(() => { + const disconnectFromShell = useCallback((options?: { suppressAutoConnect?: boolean }) => { + if (options?.suppressAutoConnect) { + suppressAutoConnectRef.current = true; + } + closeSocket(); clearTerminalScreen(); setIsConnected(false); setIsConnecting(false); connectingRef.current = false; + forceRestartOnInitRef.current = false; setAuthUrl(''); }, [clearTerminalScreen, closeSocket, setAuthUrl]); useEffect(() => { - if (!autoConnect || !isInitialized || isConnecting || isConnected) { + if ( + !autoConnect || + suppressAutoConnectRef.current || + !isInitialized || + isConnecting || + isConnected + ) { return; } diff --git a/src/components/shell/types/types.ts b/src/components/shell/types/types.ts index 14df2ea7..72a19785 100644 --- a/src/components/shell/types/types.ts +++ b/src/components/shell/types/types.ts @@ -1,6 +1,7 @@ import type { MutableRefObject, RefObject } from 'react'; import type { FitAddon } from '@xterm/addon-fit'; import type { Terminal } from '@xterm/xterm'; + import type { Project, ProjectSession } from '../../../types/app'; export type AuthCopyStatus = 'idle' | 'copied' | 'failed'; @@ -15,6 +16,7 @@ export type ShellInitMessage = { rows: number; initialCommand: string | null | undefined; isPlainShell: boolean; + forceRestart?: boolean; }; export type ShellResizeMessage = { @@ -69,8 +71,8 @@ export type UseShellRuntimeResult = { isConnecting: boolean; authUrl: string; authUrlVersion: number; - connectToShell: () => void; - disconnectFromShell: () => void; + connectToShell: (options?: { forceRestart?: boolean }) => void; + disconnectFromShell: (options?: { suppressAutoConnect?: boolean }) => void; openAuthUrlInBrowser: (url?: string) => boolean; copyAuthUrlToClipboard: (url?: string) => Promise; }; diff --git a/src/components/shell/view/Shell.tsx b/src/components/shell/view/Shell.tsx index a31a3319..575b0ac3 100644 --- a/src/components/shell/view/Shell.tsx +++ b/src/components/shell/view/Shell.tsx @@ -1,5 +1,6 @@ import { useCallback, useEffect, useMemo, useRef, useState } from 'react'; import { useTranslation } from 'react-i18next'; + import '@xterm/xterm/css/xterm.css'; import type { Project, ProjectSession } from '../../../types/app'; import { @@ -13,6 +14,7 @@ import { import { useShellRuntime } from '../hooks/useShellRuntime'; import { sendSocketMessage } from '../utils/socket'; import { getSessionDisplayName } from '../utils/auth'; + import ShellConnectionOverlay from './subcomponents/ShellConnectionOverlay'; import ShellEmptyState from './subcomponents/ShellEmptyState'; import ShellHeader from './subcomponents/ShellHeader'; @@ -46,6 +48,8 @@ export default function Shell({ const [isRestarting, setIsRestarting] = useState(false); const [cliPromptOptions, setCliPromptOptions] = useState(null); const promptCheckTimer = useRef | null>(null); + const restartTimerRef = useRef | null>(null); + const restartAfterInitRef = useRef(false); const onOutputRef = useRef<(() => void) | null>(null); const { @@ -140,6 +144,7 @@ export default function Shell({ useEffect(() => { return () => { if (promptCheckTimer.current) clearTimeout(promptCheckTimer.current); + if (restartTimerRef.current) clearTimeout(restartTimerRef.current); }; }, []); @@ -190,12 +195,42 @@ export default function Shell({ ); const handleRestartShell = useCallback(() => { + restartAfterInitRef.current = true; setIsRestarting(true); - window.setTimeout(() => { + if (restartTimerRef.current) { + clearTimeout(restartTimerRef.current); + } + restartTimerRef.current = setTimeout(() => { setIsRestarting(false); + restartTimerRef.current = null; }, SHELL_RESTART_DELAY_MS); }, []); + const handleDisconnectShell = useCallback(() => { + restartAfterInitRef.current = false; + if (restartTimerRef.current) { + clearTimeout(restartTimerRef.current); + restartTimerRef.current = null; + } + setIsRestarting(false); + disconnectFromShell({ suppressAutoConnect: true }); + }, [disconnectFromShell]); + + useEffect(() => { + if ( + !restartAfterInitRef.current || + isRestarting || + !isInitialized || + isConnected || + isConnecting + ) { + return; + } + + restartAfterInitRef.current = false; + connectToShell({ forceRestart: true }); + }, [connectToShell, isConnected, isConnecting, isInitialized, isRestarting]); + if (!selectedProject) { return (
@@ -281,7 +316,7 @@ export default function Shell({ connectLabel={t('shell.actions.connect')} connectTitle={t('shell.actions.connectTitle')} connectingLabel={t('shell.connecting')} - onConnect={connectToShell} + onConnect={handleRestartShell} /> )} diff --git a/src/components/shell/view/subcomponents/ShellConnectionOverlay.tsx b/src/components/shell/view/subcomponents/ShellConnectionOverlay.tsx index 9ed94380..b0bacb9e 100644 --- a/src/components/shell/view/subcomponents/ShellConnectionOverlay.tsx +++ b/src/components/shell/view/subcomponents/ShellConnectionOverlay.tsx @@ -1,3 +1,5 @@ +import { Loader2, RotateCcw } from 'lucide-react'; + type ShellConnectionOverlayProps = { mode: 'loading' | 'connect' | 'connecting'; description: string; @@ -19,40 +21,42 @@ export default function ShellConnectionOverlay({ }: ShellConnectionOverlayProps) { if (mode === 'loading') { return ( -
-
{loadingLabel}
+
+
+
); } if (mode === 'connect') { return ( -
-
+
+
-

{description}

+

{description}

); } return ( -
-
-
-
+
+
+
+
-

{description}

+

{description}

); diff --git a/src/components/shell/view/subcomponents/ShellHeader.tsx b/src/components/shell/view/subcomponents/ShellHeader.tsx index 0495bc9a..187fd2df 100644 --- a/src/components/shell/view/subcomponents/ShellHeader.tsx +++ b/src/components/shell/view/subcomponents/ShellHeader.tsx @@ -1,3 +1,5 @@ +import { RotateCcw, X } from 'lucide-react'; + type ShellHeaderProps = { isConnected: boolean; isInitialized: boolean; @@ -50,34 +52,27 @@ export default function ShellHeader({ {isRestarting && {statusRestartingText}}
-
+
{isConnected && ( )}
diff --git a/src/i18n/locales/en/chat.json b/src/i18n/locales/en/chat.json index 2f5af581..eac18ade 100644 --- a/src/i18n/locales/en/chat.json +++ b/src/i18n/locales/en/chat.json @@ -229,7 +229,7 @@ "disconnect": "Disconnect", "disconnectTitle": "Disconnect from shell", "restart": "Restart", - "restartTitle": "Restart Shell (disconnect first)", + "restartTitle": "Restart Shell", "connect": "Continue in Shell", "connectTitle": "Connect to shell" },