lda/wakey
1
0
Fork 0
Code Issues 5 Pull Requests Actions Packages Projects Releases 18 Wiki Activity

support long running session

Browse Source 
plus UI reattaching to existing session
This commit is contained in:
lda
2026-07-15 08:16:10 +07:00 Verified
parent 4d793136ba
commit f459bb6914
12 changed files with 751 additions and 148 deletions
Download Patch File Download Diff File Expand all files Collapse all files
wakey-agent/src/protocol.rs
+31
View File
@@ -47,6 +47,12 @@ pub enum AgentCapability {
Terminal,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct AgentTerminalSession {
pub terminal_id: TerminalId,
pub created_at_unix: u64,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum TerminalControl {
@@ -234,6 +240,9 @@ pub enum ClientMessage {
terminal_id: TerminalId,
error: ErrorPayload,
},
TerminalSessions {
sessions: Vec<AgentTerminalSession>,
},
}
#[derive(Clone, Debug, Serialize, Deserialize)]
@@ -244,6 +253,7 @@ pub enum ServerMessage {
command: AgentCommand,
},
SyncDeviceSnapshot,
SyncTerminalSessions,
OpenTerminal {
terminal_id: TerminalId,
relay_token: String,
@@ -253,6 +263,10 @@ pub enum ServerMessage {
CloseTerminal {
terminal_id: TerminalId,
},
ResumeTerminal {
terminal_id: TerminalId,
relay_token: String,
},
}
#[cfg(test)]
@@ -321,5 +335,22 @@ mod tests {
};
let json = serde_json::to_string(&resize).expect("serialize resize");
assert_eq!(json, r#"{"type":"resize","rows":40,"cols":160}"#);
let inventory = ClientMessage::TerminalSessions {
sessions: vec![AgentTerminalSession {
terminal_id: TerminalId::new("term-1").expect("terminal id"),
created_at_unix: 42,
}],
};
let json = serde_json::to_string(&inventory).expect("serialize terminal inventory");
assert!(json.contains("\"type\":\"terminal_sessions\""));
assert!(json.contains("\"created_at_unix\":42"));
let resume = ServerMessage::ResumeTerminal {
terminal_id: TerminalId::new("term-1").expect("terminal id"),
relay_token: "replacement".into(),
};
let json = serde_json::to_string(&resume).expect("serialize terminal resume");
assert!(json.contains("\"type\":\"resume_terminal\""));
}
}
wakey-agent/src/session.rs
+42 -6
View File
@@ -12,9 +12,12 @@ use crate::protocol::{AgentCapability, AgentCommand, ClientMessage, ErrorPayload
use crate::terminal::TerminalManager;
pub async fn run(config: AgentConfig) -> Result<()> {
// Terminal workers belong to the agent process, not a single control socket.
// Keeping this manager outside the reconnect loop lets PTYs survive CC loss.
let (terminal_manager, mut terminal_events) = TerminalManager::new(&config);
let mut backoff = config.reconnect_base_ms.max(100);
loop {
match run_once(&config).await {
match run_once(&config, &terminal_manager, &mut terminal_events).await {
Ok(()) => {
backoff = config.reconnect_base_ms.max(100);
}
@@ -27,7 +30,13 @@ pub async fn run(config: AgentConfig) -> Result<()> {
}
}
async fn run_once(config: &AgentConfig) -> Result<()> {
async fn run_once(
config: &AgentConfig,
terminal_manager: &TerminalManager,
terminal_events: &mut tokio::sync::mpsc::UnboundedReceiver<
crate::terminal::TerminalManagerEvent,
>,
) -> Result<()> {
let session_id = format!(
"{}-{}",
std::process::id(),
@@ -84,6 +93,7 @@ async fn run_once(config: &AgentConfig) -> Result<()> {
},
)
.await?;
send_terminal_sessions(&mut sink, terminal_manager).await?;
info!(agent_id = %config.agent_id, "agent websocket session authenticated");
let mut heartbeat = interval(Duration::from_secs(30));
@@ -93,15 +103,13 @@ async fn run_once(config: &AgentConfig) -> Result<()> {
));
snapshot_sync.set_missed_tick_behavior(MissedTickBehavior::Skip);
snapshot_sync.reset();
let (terminal_manager, mut terminal_events) = TerminalManager::new(config);
loop {
tokio::select! {
Some(event) = terminal_events.recv() => {
send_json(
&mut sink,
&ClientMessage::TerminalRejected {
terminal_id: event.terminal_id,
terminal_id: event.terminal_id,
error: ErrorPayload {
code: "terminal_worker_failed".into(),
message: event.error,
@@ -134,7 +142,7 @@ async fn run_once(config: &AgentConfig) -> Result<()> {
Message::Text(text) => {
match serde_json::from_str::<ServerMessage>(&text) {
Ok(message) => {
handle_server_message(config, &terminal_manager, &mut sink, &mut snapshot_sync, message).await?;
handle_server_message(config, terminal_manager, &mut sink, &mut snapshot_sync, message).await?;
}
Err(err) => {
// Allow the server to introduce extra frame types without
@@ -161,6 +169,20 @@ async fn run_once(config: &AgentConfig) -> Result<()> {
}
}
async fn send_terminal_sessions<S>(sink: &mut S, manager: &TerminalManager) -> Result<()>
where
S: SinkExt<Message> + Unpin,
<S as futures_util::Sink<Message>>::Error: std::error::Error + Send + Sync + 'static,
{
send_json(
sink,
&ClientMessage::TerminalSessions {
sessions: manager.sessions(),
},
)
.await
}
async fn send_device_snapshot_ws<S>(sink: &mut S, config: &AgentConfig) -> Result<usize>
where
S: SinkExt<Message> + Unpin,
@@ -240,6 +262,9 @@ where
send_device_snapshot_ws(sink, config).await?;
snapshot_sync.reset();
}
ServerMessage::SyncTerminalSessions => {
send_terminal_sessions(sink, terminal_manager).await?;
}
ServerMessage::OpenTerminal {
terminal_id,
relay_token,
@@ -269,6 +294,16 @@ where
info!(terminal_id = %terminal_id, "received terminal close request");
terminal_manager.close(&terminal_id);
}
ServerMessage::ResumeTerminal {
terminal_id,
relay_token,
} => {
info!(terminal_id = %terminal_id, "received terminal relay resume request");
if let Err(err) = terminal_manager.resume(&terminal_id, relay_token) {
warn!(terminal_id = %terminal_id, error = %err, "terminal relay resume rejected");
send_terminal_sessions(sink, terminal_manager).await?;
}
}
}
Ok(())
}
@@ -348,6 +383,7 @@ fn client_message_kind(message: &ClientMessage) -> &'static str {
ClientMessage::Result { .. } => "result",
ClientMessage::Error { .. } => "error",
ClientMessage::TerminalRejected { .. } => "terminal_rejected",
ClientMessage::TerminalSessions { .. } => "terminal_sessions",
}
}
wakey-agent/src/terminal.rs
+299 -73
View File
@@ -1,10 +1,10 @@
use std::collections::HashMap;
use std::collections::{HashMap, VecDeque};
use std::path::Path;
use std::sync::{Arc, Mutex, Weak};
use std::time::Duration;
use crate::config::AgentConfig;
use crate::protocol::{TerminalAgentHandshake, TerminalControl, TerminalId};
use crate::protocol::{AgentTerminalSession, TerminalAgentHandshake, TerminalControl, TerminalId};
use anyhow::{Context, Result};
use futures_util::{SinkExt, StreamExt};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
@@ -13,11 +13,18 @@ use tokio_tungstenite::tungstenite::Message;
use tracing::{info, warn};
const MAX_TERMINAL_FRAME_BYTES: usize = 64 * 1024;
const TERMINAL_REPLAY_BYTES: usize = 256 * 1024;
const PROCESS_SIGNAL_GRACE: Duration = Duration::from_secs(1);
/// Owns cancellation handles for terminal workers started by the control socket.
struct ActiveTerminal {
cancel: oneshot::Sender<()>,
relay_credentials: mpsc::UnboundedSender<String>,
created_at_unix: u64,
}
/// Owns PTY workers independently of any individual control-plane connection.
pub struct TerminalManager {
active: Arc<Mutex<HashMap<String, oneshot::Sender<()>>>>,
active: Arc<Mutex<HashMap<String, ActiveTerminal>>>,
max_sessions: usize,
events: mpsc::UnboundedSender<TerminalManagerEvent>,
}
@@ -54,6 +61,11 @@ impl TerminalManager {
let terminal_key = terminal_id.to_string();
let (cancel_tx, cancel_rx) = oneshot::channel();
let (relay_tx, relay_rx) = mpsc::unbounded_channel();
let created_at_unix = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_secs();
{
let mut active = self.active.lock().expect("terminal manager poisoned");
if active.contains_key(&terminal_key) {
@@ -62,15 +74,23 @@ impl TerminalManager {
if active.len() >= self.max_sessions {
anyhow::bail!("agent terminal session limit reached");
}
active.insert(terminal_key.clone(), cancel_tx);
active.insert(
terminal_key.clone(),
ActiveTerminal {
cancel: cancel_tx,
relay_credentials: relay_tx.clone(),
created_at_unix,
},
);
}
let _ = relay_tx.send(relay_token);
let config = config.clone();
let active = Arc::downgrade(&self.active);
let events = self.events.clone();
tokio::spawn(async move {
if let Err(err) =
run_terminal(&config, &terminal_id, &relay_token, rows, cols, cancel_rx).await
run_terminal(&config, &terminal_id, rows, cols, cancel_rx, relay_rx).await
{
warn!(terminal_id = %terminal_id, error = %err, "terminal worker failed");
let _ = events.send(TerminalManagerEvent {
@@ -88,7 +108,34 @@ impl TerminalManager {
.lock()
.expect("terminal manager poisoned")
.remove(terminal_id.as_str())
.is_some_and(|cancel| cancel.send(()).is_ok())
.is_some_and(|active| active.cancel.send(()).is_ok())
}
pub fn resume(&self, terminal_id: &TerminalId, relay_token: String) -> Result<()> {
let active = self.active.lock().expect("terminal manager poisoned");
let session = active
.get(terminal_id.as_str())
.with_context(|| format!("terminal session {terminal_id} is not active"))?;
session
.relay_credentials
.send(relay_token)
.map_err(|_| anyhow::anyhow!("terminal session {terminal_id} has stopped"))
}
pub fn sessions(&self) -> Vec<AgentTerminalSession> {
self.active
.lock()
.expect("terminal manager poisoned")
.iter()
.filter_map(|(terminal_id, active)| {
TerminalId::new(terminal_id.clone())
.ok()
.map(|terminal_id| AgentTerminalSession {
terminal_id,
created_at_unix: active.created_at_unix,
})
})
.collect()
}
}
@@ -97,14 +144,14 @@ impl Drop for TerminalManager {
if Arc::strong_count(&self.active) == 1
&& let Ok(mut active) = self.active.lock()
{
for (_, cancel) in active.drain() {
let _ = cancel.send(());
for (_, active) in active.drain() {
let _ = active.cancel.send(());
}
}
}
}
fn remove_completed(active: &Weak<Mutex<HashMap<String, oneshot::Sender<()>>>>, terminal_id: &str) {
fn remove_completed(active: &Weak<Mutex<HashMap<String, ActiveTerminal>>>, terminal_id: &str) {
if let Some(active) = active.upgrade()
&& let Ok(mut active) = active.lock()
{
@@ -116,28 +163,11 @@ fn remove_completed(active: &Weak<Mutex<HashMap<String, oneshot::Sender<()>>>>,
async fn run_terminal(
config: &AgentConfig,
terminal_id: &TerminalId,
relay_token: &str,
rows: u16,
cols: u16,
mut cancel: oneshot::Receiver<()>,
mut relay_credentials: mpsc::UnboundedReceiver<String>,
) -> Result<()> {
let ws_url = terminal_websocket_url(&config.server_url, terminal_id)?;
let (stream, _) = tokio::select! {
_ = &mut cancel => return Ok(()),
result = tokio_tungstenite::connect_async(ws_url.as_str()) => {
result.context("failed to connect terminal relay websocket")?
}
};
let (mut sink, mut source) = stream.split();
send_json(
&mut sink,
&TerminalAgentHandshake::Auth {
agent_id: config.agent_id.clone(),
relay_token: relay_token.to_string(),
},
)
.await?;
let terminal = match wakey::wakey_linux::terminal::TerminalPty::spawn(
Path::new(&config.terminal.shell),
rows,
@@ -145,15 +175,6 @@ async fn run_terminal(
) {
Ok(terminal) => terminal,
Err(err) => {
let _ = send_json(
&mut sink,
&TerminalControl::Error {
code: "terminal_spawn_failed".into(),
message: err.to_string(),
},
)
.await;
let _ = sink.send(Message::Close(None)).await;
return Err(err);
}
};
@@ -163,9 +184,14 @@ async fn run_terminal(
mut child,
} = terminal;
let process_group = child.id();
send_json(&mut sink, &TerminalControl::Ready).await?;
info!(terminal_id = %terminal_id, shell = %config.terminal.shell.display(), "terminal PTY ready");
let (relay_input_tx, mut relay_input_rx) = mpsc::unbounded_channel();
let mut relay_output: Option<mpsc::Sender<Message>> = None;
let mut relay_task: Option<tokio::task::JoinHandle<()>> = None;
let mut relay_generation = 0_u64;
let mut replay = VecDeque::new();
let mut replay_bytes = 0_usize;
let mut output = [0_u8; 16 * 1024];
let mut requested_close = false;
let mut observed_status = None;
@@ -182,48 +208,78 @@ async fn run_terminal(
read = reader.read(&mut output) => {
match read {
Ok(0) => break,
Ok(count) => sink
.send(Message::Binary(output[..count].to_vec().into()))
.await
.context("failed to send PTY output")?,
Ok(count) => send_terminal_output(
Message::Binary(output[..count].to_vec().into()),
&mut relay_output,
&mut replay,
&mut replay_bytes,
).await,
// Linux PTY masters commonly report EIO after the slave closes.
Err(err) if err.raw_os_error() == Some(5) => break,
Err(err) => return Err(err).context("failed to read PTY output"),
}
}
incoming = source.next() => {
credential = relay_credentials.recv() => {
let Some(relay_token) = credential else { break; };
if let Some(task) = relay_task.take() {
task.abort();
}
relay_generation = relay_generation.wrapping_add(1);
let generation = relay_generation;
let (output_tx, output_rx) = mpsc::channel(32);
relay_output = None;
let initial_replay = replay.drain(..).collect();
replay_bytes = 0;
let config = config.clone();
let terminal_id = terminal_id.clone();
let relay_input_tx = relay_input_tx.clone();
relay_task = Some(tokio::spawn(async move {
if let Err(err) = run_terminal_relay(RelayConnection {
config,
terminal_id: terminal_id.clone(),
relay_token,
generation,
initial_replay,
output_tx,
output_rx,
input: relay_input_tx.clone(),
}).await {
warn!(terminal_id = %terminal_id, error = %err, "terminal relay disconnected");
}
let _ = relay_input_tx.send(RelayInput::Disconnected { generation });
}));
}
incoming = relay_input_rx.recv() => {
let Some(message) = incoming else { break; };
match message.context("terminal relay websocket receive failed")? {
Message::Binary(bytes) => {
match message {
RelayInput::Binary { generation, bytes } if generation == relay_generation => {
if bytes.len() > MAX_TERMINAL_FRAME_BYTES {
anyhow::bail!("terminal input frame exceeds size limit");
}
writer.write_all(&bytes).await.context("failed to write PTY input")?;
}
Message::Text(text) => {
match serde_json::from_str::<TerminalControl>(&text)
.context("invalid terminal control frame")?
{
TerminalControl::Resize { rows, cols } => {
validate_size(rows, cols)?;
wakey::wakey_linux::terminal::resize_terminal(
&writer, rows, cols,
)?;
}
TerminalControl::Close => {
requested_close = true;
break;
}
_ => anyhow::bail!("terminal control frame has invalid direction"),
}
RelayInput::Resize { generation, rows, cols } if generation == relay_generation => {
validate_size(rows, cols)?;
wakey::wakey_linux::terminal::resize_terminal(&writer, rows, cols)?;
}
Message::Ping(payload) => sink.send(Message::Pong(payload)).await?,
Message::Pong(_) => {}
Message::Close(_) => {
RelayInput::Close { generation } if generation == relay_generation => {
requested_close = true;
break;
}
Message::Frame(_) => {}
RelayInput::Connected { generation, output } if generation == relay_generation => {
relay_output = Some(output.clone());
while let Some(frame) = replay.pop_front() {
replay_bytes = replay_bytes.saturating_sub(message_size(&frame));
if output.send(frame).await.is_err() {
relay_output = None;
break;
}
}
}
RelayInput::Disconnected { generation } if generation == relay_generation => {
relay_output = None;
}
_ => {}
}
}
}
@@ -233,26 +289,173 @@ async fn run_terminal(
Some(status) => status,
None => terminate_process_group(&mut child, process_group).await?,
};
let _ = send_json(
&mut sink,
&TerminalControl::Exited {
if let Some(output) = relay_output {
let control = TerminalControl::Exited {
exit_code: status.code(),
},
)
.await;
let _ = sink.send(Message::Close(None)).await;
};
if let Ok(text) = serde_json::to_string(&control) {
let _ = output.send(Message::Text(text.into())).await;
}
}
info!(terminal_id = %terminal_id, exit_code = ?status.code(), requested_close, "terminal worker exited");
Ok(())
}
enum RelayInput {
Connected {
generation: u64,
output: mpsc::Sender<Message>,
},
Binary {
generation: u64,
bytes: Vec<u8>,
},
Resize {
generation: u64,
rows: u16,
cols: u16,
},
Close {
generation: u64,
},
Disconnected {
generation: u64,
},
}
#[cfg(unix)]
struct RelayConnection {
config: AgentConfig,
terminal_id: TerminalId,
relay_token: String,
generation: u64,
initial_replay: Vec<Message>,
output_tx: mpsc::Sender<Message>,
output_rx: mpsc::Receiver<Message>,
input: mpsc::UnboundedSender<RelayInput>,
}
#[cfg(unix)]
async fn run_terminal_relay(relay: RelayConnection) -> Result<()> {
let ws_url = terminal_websocket_url(&relay.config.server_url, &relay.terminal_id)?;
let (stream, _) = tokio_tungstenite::connect_async(ws_url.as_str())
.await
.context("failed to connect terminal relay websocket")?;
let (mut sink, mut source) = stream.split();
send_json(
&mut sink,
&TerminalAgentHandshake::Auth {
agent_id: relay.config.agent_id.clone(),
relay_token: relay.relay_token,
},
)
.await?;
send_json(&mut sink, &TerminalControl::Ready).await?;
for frame in relay.initial_replay {
sink.send(frame)
.await
.context("failed to replay detached terminal output")?;
}
relay
.input
.send(RelayInput::Connected {
generation: relay.generation,
output: relay.output_tx,
})
.map_err(|_| anyhow::anyhow!("terminal worker stopped"))?;
let mut output = relay.output_rx;
loop {
tokio::select! {
outgoing = output.recv() => {
let Some(message) = outgoing else { break; };
sink.send(message).await.context("failed to send terminal relay output")?;
}
incoming = source.next() => {
let Some(message) = incoming else { break; };
match message.context("terminal relay websocket receive failed")? {
Message::Binary(bytes) => {
relay.input.send(RelayInput::Binary {
generation: relay.generation,
bytes: bytes.to_vec(),
}).map_err(|_| anyhow::anyhow!("terminal worker stopped"))?;
}
Message::Text(text) => match serde_json::from_str::<TerminalControl>(&text)
.context("invalid terminal control frame")?
{
TerminalControl::Resize { rows, cols } => {
relay.input.send(RelayInput::Resize {
generation: relay.generation,
rows,
cols,
})
.map_err(|_| anyhow::anyhow!("terminal worker stopped"))?;
}
TerminalControl::Close => {
let _ = relay.input.send(RelayInput::Close {
generation: relay.generation,
});
break;
}
_ => anyhow::bail!("terminal control frame has invalid direction"),
},
Message::Ping(payload) => sink.send(Message::Pong(payload)).await?,
Message::Pong(_) => {}
// Transport closure only detaches the relay. The agent-owned
// PTY remains alive and waits for replacement credentials.
Message::Close(_) => break,
Message::Frame(_) => {}
}
}
}
}
Ok(())
}
async fn send_terminal_output(
frame: Message,
relay: &mut Option<mpsc::Sender<Message>>,
replay: &mut VecDeque<Message>,
replay_bytes: &mut usize,
) {
if let Some(tx) = relay.as_ref() {
if let Err(error) = tx.send(frame).await {
*relay = None;
push_local_replay(error.0, replay, replay_bytes);
}
} else {
push_local_replay(frame, replay, replay_bytes);
}
}
fn push_local_replay(frame: Message, replay: &mut VecDeque<Message>, replay_bytes: &mut usize) {
*replay_bytes += message_size(&frame);
replay.push_back(frame);
while *replay_bytes > TERMINAL_REPLAY_BYTES {
if let Some(dropped) = replay.pop_front() {
*replay_bytes = replay_bytes.saturating_sub(message_size(&dropped));
} else {
break;
}
}
}
fn message_size(message: &Message) -> usize {
match message {
Message::Text(text) => text.len(),
Message::Binary(bytes) | Message::Ping(bytes) | Message::Pong(bytes) => bytes.len(),
Message::Close(_) | Message::Frame(_) => 0,
}
}
#[cfg(not(unix))]
async fn run_terminal(
_config: &AgentConfig,
_terminal_id: &TerminalId,
_relay_token: &str,
_rows: u16,
_cols: u16,
_cancel: oneshot::Receiver<()>,
_relay_credentials: mpsc::UnboundedReceiver<String>,
) -> Result<()> {
anyhow::bail!("terminal sessions are unsupported on this platform")
}
@@ -342,4 +545,27 @@ mod tests {
"wss://example.com/api/v1/agent/terminals/term-1/ws"
);
}
#[test]
fn detached_replay_drops_oldest_output_at_bound() {
let mut replay = VecDeque::new();
let mut replay_bytes = 0;
for marker in 0_u8..10 {
push_local_replay(
Message::Binary(vec![marker; TERMINAL_REPLAY_BYTES / 4].into()),
&mut replay,
&mut replay_bytes,
);
}
assert!(replay_bytes <= TERMINAL_REPLAY_BYTES);
assert_eq!(replay.len(), 4);
assert_eq!(
replay.front().and_then(|frame| match frame {
Message::Binary(bytes) => bytes.first().copied(),
_ => None,
}),
Some(6)
);
}
}