fix(shell): disconnect and restart buttons (#831)

Co-authored-by: Simos Mikelatos <simosmik@gmail.com>

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);