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

make terminal session lifetime agent-configurable

Browse Source
This commit is contained in:
lda
2026-07-16 20:23:29 +07:00 Verified
parent f06c861bed
commit f2303e3a1f
10 changed files with 341 additions and 54 deletions
Download Patch File Download Diff File Expand all files Collapse all files
wakey-agent/src/config.rs
+49 -3
View File
@@ -4,7 +4,7 @@ use std::fmt;
use std::path::{Path, PathBuf};
use std::sync::LazyLock;
use crate::protocol::DEFAULT_TERMINAL_MAX_SESSIONS;
use crate::protocol::{DEFAULT_TERMINAL_MAX_SESSIONS, DEFAULT_TERMINAL_SESSION_TTL_SECONDS};
pub const DEFAULT_CONFIG_PATH: &str = "/etc/wakey-agent/config.toml";
pub const DEFAULT_PID_FILE: &str = "/var/run/wakey-agent.pid";
@@ -54,6 +54,9 @@ pub struct TerminalConfig {
pub current_dir: Option<PathBuf>,
#[serde(default = "default_terminal_max_sessions")]
pub max_sessions: usize,
/// Maximum lifetime of a PTY process. Zero disables automatic expiry.
#[serde(default = "default_terminal_session_ttl_seconds")]
pub session_ttl_seconds: u64,
}
impl Default for TerminalConfig {
@@ -64,6 +67,7 @@ impl Default for TerminalConfig {
args: Vec::new(),
current_dir: None,
max_sessions: default_terminal_max_sessions(),
session_ttl_seconds: default_terminal_session_ttl_seconds(),
}
}
}
@@ -148,6 +152,23 @@ const fn default_terminal_max_sessions() -> usize {
DEFAULT_TERMINAL_MAX_SESSIONS
}
const fn default_terminal_session_ttl_seconds() -> u64 {
DEFAULT_TERMINAL_SESSION_TTL_SECONDS
}
impl TerminalConfig {
pub(crate) fn session_ttl(&self) -> Result<Option<std::time::Duration>> {
if self.session_ttl_seconds == 0 {
return Ok(None);
}
let ttl = std::time::Duration::from_secs(self.session_ttl_seconds);
std::time::Instant::now()
.checked_add(ttl)
.context("terminal.session_ttl_seconds is too large for the platform timer")?;
Ok(Some(ttl))
}
}
impl AgentConfig {
pub fn local_path_envs(&self) -> Vec<(&'static str, &Path)> {
vec![
@@ -170,8 +191,10 @@ pub fn apply_local_path_env_to_command(cmd: &mut std::process::Command, config:
pub fn load_config(path: &Path) -> Result<AgentConfig> {
let content = std::fs::read_to_string(path)
.with_context(|| format!("failed to read agent config {}", path.display()))?;
toml::from_str(&content)
.with_context(|| format!("failed to parse agent config {}", path.display()))
let config: AgentConfig = toml::from_str(&content)
.with_context(|| format!("failed to parse agent config {}", path.display()))?;
config.terminal.session_ttl()?;
Ok(config)
}
pub fn save_config(path: &Path, config: &AgentConfig) -> Result<()> {
@@ -261,6 +284,7 @@ mod tests {
args: vec!["-l".into()],
current_dir: Some("/tmp".into()),
max_sessions: 2,
session_ttl_seconds: 600,
},
};
@@ -314,6 +338,10 @@ max_sessions = 2
assert!(config.terminal.args.is_empty());
assert!(config.terminal.current_dir.is_none());
assert_eq!(
config.terminal.session_ttl_seconds,
DEFAULT_TERMINAL_SESSION_TTL_SECONDS
);
}
#[test]
@@ -333,4 +361,22 @@ max_session = 67
assert!(error.to_string().contains("max_session"));
}
#[test]
fn terminal_ttl_zero_is_unlimited_and_positive_is_exact() {
let mut terminal = TerminalConfig {
session_ttl_seconds: 0,
..TerminalConfig::default()
};
assert_eq!(terminal.session_ttl().expect("unlimited TTL"), None);
terminal.session_ttl_seconds = 67;
assert_eq!(
terminal.session_ttl().expect("finite TTL"),
Some(std::time::Duration::from_secs(67))
);
terminal.session_ttl_seconds = u64::MAX;
assert!(terminal.session_ttl().is_err());
}
}
wakey-agent/src/enroll.rs
+1
View File
@@ -157,6 +157,7 @@ mod tests {
args: vec!["-l".into()],
current_dir: Some("/root".into()),
max_sessions: 2,
session_ttl_seconds: 600,
},
};
wakey-agent/src/protocol.rs
+47 -1
View File
@@ -66,6 +66,11 @@ pub enum AgentCapability {
}
pub const DEFAULT_TERMINAL_MAX_SESSIONS: usize = 2;
pub const DEFAULT_TERMINAL_SESSION_TTL_SECONDS: u64 = 12 * 60 * 60;
const fn default_terminal_session_ttl_seconds() -> u64 {
DEFAULT_TERMINAL_SESSION_TTL_SECONDS
}
/// Optional parameters attached to advertised agent capabilities.
///
@@ -87,12 +92,17 @@ impl AgentCapabilityOptions {
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct TerminalCapabilityOptions {
pub max_sessions: usize,
#[serde(default = "default_terminal_session_ttl_seconds")]
pub session_ttl_seconds: u64,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct AgentTerminalSession {
pub terminal_id: TerminalId,
pub created_at_unix: u64,
/// The policy captured when this PTY was created. Zero means unlimited.
#[serde(default = "default_terminal_session_ttl_seconds")]
pub session_ttl_seconds: u64,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
@@ -410,6 +420,7 @@ mod tests {
sessions: vec![AgentTerminalSession {
terminal_id: TerminalId::new("term-1").expect("terminal id"),
created_at_unix: 42,
session_ttl_seconds: 600,
}],
};
let json = serde_json::to_string(&inventory).expect("serialize terminal inventory");
@@ -430,11 +441,46 @@ mod tests {
agent_id: "router".into(),
capabilities: vec![AgentCapability::Terminal],
capability_options: AgentCapabilityOptions {
terminal: Some(TerminalCapabilityOptions { max_sessions: 3 }),
terminal: Some(TerminalCapabilityOptions {
max_sessions: 3,
session_ttl_seconds: 600,
}),
},
};
let value = serde_json::to_value(message).expect("serialize hello");
assert_eq!(value["capability_options"]["terminal"]["max_sessions"], 3);
assert_eq!(
value["capability_options"]["terminal"]["session_ttl_seconds"],
600
);
}
#[test]
fn terminal_ttl_defaults_for_old_messages_and_preserves_explicit_zero() {
let old_options: TerminalCapabilityOptions =
serde_json::from_value(serde_json::json!({ "max_sessions": 2 }))
.expect("deserialize old capability options");
assert_eq!(
old_options.session_ttl_seconds,
DEFAULT_TERMINAL_SESSION_TTL_SECONDS
);
let old_session: AgentTerminalSession = serde_json::from_value(serde_json::json!({
"terminal_id": "term-old",
"created_at_unix": 42
}))
.expect("deserialize old terminal inventory");
assert_eq!(
old_session.session_ttl_seconds,
DEFAULT_TERMINAL_SESSION_TTL_SECONDS
);
let unlimited: TerminalCapabilityOptions = serde_json::from_value(serde_json::json!({
"max_sessions": 2,
"session_ttl_seconds": 0
}))
.expect("deserialize unlimited capability options");
assert_eq!(unlimited.session_ttl_seconds, 0);
}
}
wakey-agent/src/session.rs
+1
View File
@@ -91,6 +91,7 @@ async fn run_once(
.enabled
.then_some(TerminalCapabilityOptions {
max_sessions: config.terminal.max_sessions.max(1),
session_ttl_seconds: config.terminal.session_ttl_seconds,
}),
},
},
wakey-agent/src/terminal.rs
+58 -2
View File
@@ -52,6 +52,7 @@ struct ActiveTerminal {
cancel: Option<oneshot::Sender<()>>,
relay_credentials: mpsc::UnboundedSender<String>,
created_at_unix: u64,
session_ttl_seconds: u64,
}
/// Owns PTY workers independently of any individual control-plane connection.
@@ -91,6 +92,7 @@ impl TerminalManager {
anyhow::bail!("terminal capability is disabled");
}
validate_size(rows, cols)?;
let session_ttl = config.terminal.session_ttl()?;
let terminal_key = terminal_id.to_string();
let (cancel_tx, cancel_rx) = oneshot::channel();
@@ -113,6 +115,7 @@ impl TerminalManager {
cancel: Some(cancel_tx),
relay_credentials: relay_tx.clone(),
created_at_unix,
session_ttl_seconds: config.terminal.session_ttl_seconds,
},
);
}
@@ -122,7 +125,16 @@ impl TerminalManager {
let active = Arc::downgrade(&self.active);
let events = self.events.clone();
tokio::spawn(async move {
let result = run_terminal(&config, &terminal_id, rows, cols, cancel_rx, relay_rx).await;
let result = run_terminal(
&config,
&terminal_id,
rows,
cols,
session_ttl,
cancel_rx,
relay_rx,
)
.await;
// Inventory is authoritative. Remove the stopped worker before
// notifying the control session, which may immediately resync it.
remove_completed(&active, terminal_id.as_str());
@@ -176,6 +188,7 @@ impl TerminalManager {
.map(|terminal_id| AgentTerminalSession {
terminal_id,
created_at_unix: active.created_at_unix,
session_ttl_seconds: active.session_ttl_seconds,
})
})
.collect()
@@ -210,6 +223,7 @@ async fn run_terminal(
terminal_id: &TerminalId,
rows: u16,
cols: u16,
session_ttl: Option<Duration>,
mut cancel: oneshot::Receiver<()>,
mut relay_credentials: mpsc::UnboundedReceiver<String>,
) -> Result<()> {
@@ -242,8 +256,10 @@ async fn run_terminal(
let mut terminal_state = TerminalState::new(rows, cols);
let mut output = [0_u8; 16 * 1024];
let mut requested_close = false;
let mut ttl_expired = false;
let mut observed_status = None;
let mut drain_deadline: Option<std::pin::Pin<Box<tokio::time::Sleep>>> = None;
let mut ttl_deadline = session_ttl.map(|ttl| Box::pin(tokio::time::sleep(ttl)));
loop {
tokio::select! {
biased;
@@ -251,6 +267,15 @@ async fn run_terminal(
requested_close = true;
break;
}
_ = async {
match ttl_deadline.as_mut() {
Some(deadline) => deadline.as_mut().await,
None => std::future::pending().await,
}
} => {
ttl_expired = true;
break;
}
status = child.wait(), if observed_status.is_none() => {
observed_status = Some(status.context("failed waiting for terminal child")?);
drain_deadline = Some(Box::pin(tokio::time::sleep(PTY_EXIT_DRAIN_TIMEOUT)));
@@ -363,7 +388,7 @@ async fn run_terminal(
let _ = output.try_send(Message::Text(text.into()));
}
}
info!(terminal_id = %terminal_id, exit_code = ?status.code(), requested_close, "terminal worker exited");
info!(terminal_id = %terminal_id, exit_code = ?status.code(), requested_close, ttl_expired, "terminal worker exited");
Ok(())
}
@@ -511,6 +536,7 @@ async fn run_terminal(
_terminal_id: &TerminalId,
_rows: u16,
_cols: u16,
_session_ttl: Option<Duration>,
_cancel: oneshot::Receiver<()>,
_relay_credentials: mpsc::UnboundedReceiver<String>,
) -> Result<()> {
@@ -607,6 +633,7 @@ mod tests {
cancel: Some(cancel),
relay_credentials,
created_at_unix: 42,
session_ttl_seconds: 600,
},
)]))),
max_sessions: 2,
@@ -660,6 +687,7 @@ mod tests {
cancel: Some(cancel),
relay_credentials,
created_at_unix: 42,
session_ttl_seconds: 600,
},
)])));
let manager = TerminalManager {
@@ -702,6 +730,34 @@ mod tests {
assert!(manager.active.lock().expect("terminal manager").is_empty());
}
#[cfg(unix)]
#[tokio::test]
async fn local_ttl_terminates_pty_without_control_plane_close() {
let mut config = crate::config::DEFAULT_CONFIG.clone();
config.terminal.enabled = true;
config.terminal.shell = "/bin/sh".into();
config.terminal.args = vec!["-c".into(), "sleep 30".into()];
let terminal_id = TerminalId::new("ttl-test").expect("terminal id");
let (_cancel, cancel_rx) = oneshot::channel();
let (_relay, relay_rx) = mpsc::unbounded_channel();
tokio::time::timeout(
Duration::from_secs(3),
run_terminal(
&config,
&terminal_id,
24,
80,
Some(Duration::from_millis(20)),
cancel_rx,
relay_rx,
),
)
.await
.expect("local TTL should stop the PTY promptly")
.expect("TTL cleanup should succeed");
}
#[test]
fn snapshot_reconstructs_screen_content_and_cursor() {
let mut state = TerminalState::new(24, 80);