From a3bc895755b308291ccb33df45f97cf3ab41208c Mon Sep 17 00:00:00 2001 From: mbecker20 Date: Mon, 27 Oct 2025 21:04:46 -0700 Subject: [PATCH] fix terminal disconnect --- Cargo.lock | 2 ++ bin/cli/src/command/terminal.rs | 3 ++- bin/periphery/Cargo.toml | 5 +++-- bin/periphery/src/api/terminal.rs | 15 +++------------ client/core/rs/src/ws/terminal.rs | 22 +++++++++++++--------- lib/transport/Cargo.toml | 1 + lib/transport/src/channel.rs | 13 +++++++++++++ 7 files changed, 37 insertions(+), 24 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index f87fef3fa..dab154ada 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2986,6 +2986,7 @@ dependencies = [ "serde_yaml_ng", "serror", "shell-escape", + "shlex", "sysinfo", "tokio", "tokio-stream", @@ -5565,6 +5566,7 @@ dependencies = [ "axum", "base64 0.22.1", "bytes", + "colored", "encoding", "futures-util", "noise", diff --git a/bin/cli/src/command/terminal.rs b/bin/cli/src/command/terminal.rs index d9f7511ad..df1e4ee85 100644 --- a/bin/cli/src/command/terminal.rs +++ b/bin/cli/src/command/terminal.rs @@ -251,7 +251,8 @@ async fn handle_terminal_forwarding< future_or_cancel(ws_read.receive_stdout(), &cancel).await { let bytes = match msg { - Ok(bytes) => bytes, + Ok(Some(bytes)) => bytes, + Ok(None) => break, Err(e) => { cancel.cancel(); return Some(e.context("Websocket read error")); diff --git a/bin/periphery/Cargo.toml b/bin/periphery/Cargo.toml index 6d207093e..42e1bd34e 100644 --- a/bin/periphery/Cargo.toml +++ b/bin/periphery/Cargo.toml @@ -54,9 +54,10 @@ sysinfo.workspace = true dotenvy.workspace = true anyhow.workspace = true rustls.workspace = true -tokio.workspace = true -serde.workspace = true bytes.workspace = true +serde.workspace = true +shlex.workspace = true +tokio.workspace = true axum.workspace = true clap.workspace = true envy.workspace = true diff --git a/bin/periphery/src/api/terminal.rs b/bin/periphery/src/api/terminal.rs index b3fdd4f3a..757b919ef 100644 --- a/bin/periphery/src/api/terminal.rs +++ b/bin/periphery/src/api/terminal.rs @@ -1,7 +1,6 @@ use std::sync::Arc; use anyhow::{Context, anyhow}; -use colored::Colorize; use futures_util::{Stream, StreamExt, TryStreamExt}; use komodo_client::entities::{ KOMODO_EXIT_CODE, NoData, @@ -372,10 +371,12 @@ async fn handle_terminal_forwarding( // Forward stdout -> WS let mut stdout = terminal.stdout.resubscribe(); + loop { let res = tokio::select! { res = stdout.recv() => res, _ = terminal.cancel.cancelled() => { + let _ = sender.send_terminal_exited(channel).await; break }, _ = cancel.cancelled() => { @@ -386,17 +387,7 @@ async fn handle_terminal_forwarding( let bytes = match res { Ok(bytes) => bytes, Err(_e) => { - terminal.cancel(); - let _ = sender - .send_terminal( - channel, - Err(anyhow!( - "\n{} {}", - "pty".bold(), - "exited".red().bold() - )), - ) - .await; + let _ = sender.send_terminal_exited(channel).await; break; } }; diff --git a/client/core/rs/src/ws/terminal.rs b/client/core/rs/src/ws/terminal.rs index ff53d1d50..1946b3f93 100644 --- a/client/core/rs/src/ws/terminal.rs +++ b/client/core/rs/src/ws/terminal.rs @@ -1,4 +1,4 @@ -use anyhow::{Context, anyhow}; +use anyhow::Context; use bytes::Bytes; use futures_util::{ SinkExt, StreamExt, TryStreamExt, @@ -146,17 +146,19 @@ impl TerminalWebsocket { .await } - pub async fn receive_stdout(&mut self) -> anyhow::Result { + pub async fn receive_stdout( + &mut self, + ) -> anyhow::Result> { loop { match self.0.try_next().await.context("Websocket read error")? { Some(tungstenite::Message::Binary(bytes)) => { - return Ok(bytes); + return Ok(Some(bytes)); } Some(tungstenite::Message::Text(text)) => { - return Ok(text.into()); + return Ok(Some(text.into())); } Some(tungstenite::Message::Close(_)) | None => { - return Err(anyhow!("Websocket closed")); + return Ok(None); } // Can ignore these message types Some(tungstenite::Message::Ping(_)) @@ -212,17 +214,19 @@ pub type TerminalWebsocketStreamInner = pub struct TerminalWebsocketStream(TerminalWebsocketStreamInner); impl TerminalWebsocketStream { - pub async fn receive_stdout(&mut self) -> anyhow::Result { + pub async fn receive_stdout( + &mut self, + ) -> anyhow::Result> { loop { match self.0.try_next().await.context("Websocket read error")? { Some(tungstenite::Message::Binary(bytes)) => { - return Ok(bytes); + return Ok(Some(bytes)); } Some(tungstenite::Message::Text(text)) => { - return Ok(text.into()); + return Ok(Some(text.into())); } Some(tungstenite::Message::Close(_)) | None => { - return Err(anyhow!("Websocket closed")); + return Ok(None); } // Can ignore these message types Some(tungstenite::Message::Ping(_)) diff --git a/lib/transport/Cargo.toml b/lib/transport/Cargo.toml index d2e0572e3..3a1c7e20a 100644 --- a/lib/transport/Cargo.toml +++ b/lib/transport/Cargo.toml @@ -18,6 +18,7 @@ tokio-tungstenite.workspace = true pin-project-lite.workspace = true futures-util.workspace = true tokio-util.workspace = true +colored.workspace = true tracing.workspace = true anyhow.workspace = true base64.workspace = true diff --git a/lib/transport/src/channel.rs b/lib/transport/src/channel.rs index 107a76ac8..4e07c560a 100644 --- a/lib/transport/src/channel.rs +++ b/lib/transport/src/channel.rs @@ -1,4 +1,5 @@ use anyhow::{Context, anyhow}; +use colored::Colorize as _; use encoding::{ Encode, EncodedJsonMessage, EncodedResponse, JsonMessage, }; @@ -126,6 +127,18 @@ impl Sender { ) -> anyhow::Result<()> { self.send_message(TerminalMessage::new(channel, data)).await } + + pub async fn send_terminal_exited( + &self, + channel: Uuid, + ) -> anyhow::Result<()> { + self + .send_message(TerminalMessage::new( + channel, + Err(anyhow!("\n{} {}", "pty".bold(), "exited".red().bold())), + )) + .await + } } #[derive(Debug)]