fix: resolve PTY session termination hanging issue after exit command

- Add proper SSH channel cleanup when exit command is sent
- Ensure terminal state is fully restored with disable_raw_mode()
- Fix input task blocking by using shorter polling timeout (100ms)
- Add explicit shutdown signal handling for EOF/Close messages
- Force program exit after session ends to prevent hanging

Author: inureyes

src/commands/interactive.rs

@@ -205,6 +205,10 @@ impl InteractiveCommand {
             pty_manager.run_multiplex_sessions(session_ids).await?;
         }
 
+        // Ensure terminal is fully restored after PTY session ends
+        let _ = crossterm::terminal::disable_raw_mode();
+        let _ = std::io::Write::flush(&mut std::io::stdout());
+
         Ok(InteractiveResult {
             duration: start_time.elapsed(),
             commands_executed: 0, // PTY mode doesn't count discrete commands
@@ -618,6 +622,12 @@ impl InteractiveCommand {
             match rl.readline(&prompt) {
                 Ok(line) => {
                     if line.trim() == "exit" {
+                        // Send exit command to remote server before breaking
+                        let mut session_guard = session_arc.lock().await;
+                        session_guard.send_command("exit").await?;
+                        drop(session_guard);
+                        // Give the SSH session a moment to process the exit
+                        tokio::time::sleep(Duration::from_millis(100)).await;
                         break;
                     }
 
@@ -649,7 +659,18 @@ impl InteractiveCommand {
         }
 
         // Clean up
+        shutdown.store(true, Ordering::Relaxed);
         output_reader.abort();
+
+        // Properly close the SSH session
+        let mut session_guard = session_arc.lock().await;
+        if session_guard.is_connected {
+            // Close the SSH channel properly
+            let _ = session_guard.channel.close().await;
+            session_guard.is_connected = false;
+        }
+        drop(session_guard);
+
         let _ = rl.save_history(&history_path);
 
         Ok(commands_executed)

src/main.rs

@@ -379,12 +379,20 @@ async fn main() -> Result<()> {
                     use_pty,
                 };
                 let result = interactive_cmd.execute().await?;
+
+                // Ensure terminal is fully restored before printing
+                let _ = crossterm::terminal::disable_raw_mode();
+                let _ = crossterm::cursor::Show;
+                let _ = std::io::Write::flush(&mut std::io::stdout());
+
                 println!("\nSession ended.");
                 if cli.verbose > 0 {
                     println!("Duration: {}", format_duration(result.duration));
                     println!("Commands executed: {}", result.commands_executed);
                 }
-                Ok(())
+
+                // Force exit to ensure proper termination
+                std::process::exit(0);
             } else {
                 // Determine timeout: CLI argument takes precedence over config
                 let timeout = if cli.timeout > 0 {

src/pty/mod.rs

@@ -135,11 +135,16 @@ impl PtyManager {
 
     /// Run a single PTY session
     pub async fn run_single_session(&mut self, session_id: usize) -> Result<()> {
-        if let Some(session) = self.active_sessions.get_mut(session_id) {
+        let result = if let Some(session) = self.active_sessions.get_mut(session_id) {
             session.run().await
         } else {
             anyhow::bail!("PTY session {session_id} not found")
-        }
+        };
+
+        // Ensure terminal is properly restored after session ends
+        let _ = crossterm::terminal::disable_raw_mode();
+
+        result
     }
 
     /// Run multiple PTY sessions with session switching

src/pty/session.rs

@@ -161,16 +161,19 @@ impl PtySession {
             .ok_or_else(|| anyhow::anyhow!("Message sender not available"))?
             .clone();
         let shutdown_for_input = Arc::clone(&self.shutdown);
-        let timeout_duration = self.config.timeout;
 
         let input_task = tokio::spawn(async move {
             loop {
                 if shutdown_for_input.load(Ordering::Relaxed) {
                     break;
                 }
 
+                // Use a shorter polling timeout for better responsiveness
+                // This allows the task to check the shutdown flag more frequently
+                let poll_timeout = Duration::from_millis(100);
+
                 // Check for input events with timeout
-                if crossterm::event::poll(timeout_duration).unwrap_or(false) {
+                if crossterm::event::poll(poll_timeout).unwrap_or(false) {
                     match crossterm::event::read() {
                         Ok(event) => {
                             if let Some(data) = Self::handle_input_event(event) {
@@ -219,6 +222,8 @@ impl PtySession {
                     }
                     ChannelMsg::Eof | ChannelMsg::Close => {
                         tracing::debug!("SSH channel closed");
+                        // Set shutdown signal before terminating to ensure input task stops
+                        self.shutdown.store(true, Ordering::Relaxed);
                         should_terminate = true;
                     }
                     _ => {}
@@ -271,15 +276,36 @@ impl PtySession {
             }
         }
 
-        // Cleanup tasks
+        // Signal shutdown first
+        self.shutdown.store(true, Ordering::Relaxed);
+
+        // Abort tasks immediately
         resize_task.abort();
         input_task.abort();
 
+        // Wait for tasks to complete their abort
+        let _ = tokio::time::timeout(Duration::from_millis(100), async {
+            while !resize_task.is_finished() || !input_task.is_finished() {
+                tokio::time::sleep(Duration::from_millis(10)).await;
+            }
+        })
+        .await;
+
         // Disable mouse support if we enabled it
         if self.config.enable_mouse {
             let _ = TerminalOps::disable_mouse();
         }
 
+        // IMPORTANT: Explicitly restore terminal state by dropping the guard
+        // This ensures raw mode is disabled before we return
+        self.terminal_guard = None;
+
+        // Ensure terminal is fully restored
+        let _ = crossterm::terminal::disable_raw_mode();
+
+        // Flush stdout to ensure all output is written
+        let _ = io::stdout().flush();
+
         self.state = PtyState::Closed;
         Ok(())
     }