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
docs/adr/0002-remote-terminal-sessions.md
+2
View File
@@ -80,6 +80,8 @@ The agent keeps its terminal manager outside the control-WebSocket reconnect loo
Closing a session must terminate the entire PTY process group, not only the shell process. Cleanup should send a hangup first, then escalate to termination and forced kill after bounded grace periods. Explicit operator close, absolute session timeout, agent process exit, and machine reboot close the session.
The agent configuration sets `terminal.session_ttl_seconds` for newly created PTYs. The default is 43,200 seconds (12 hours); explicit zero disables automatic expiry. The agent owns and enforces this deadline locally because PTYs survive browser and control-plane disconnection. It advertises the policy to the control plane and includes the creation-time TTL in terminal inventory, allowing CC to mirror expiry and preserve it across CC restart. Missing fields from older agents retain the 12-hour default. Positive values are exact and values that cannot fit the platform timer are rejected rather than clamped.
An agent process restart or machine reboot still ends every terminal session because PTY file descriptors cannot be recovered by a new process. Surviving that boundary would require an external session host such as `tmux` or a separate terminal daemon and is not part of this design.
The agent reports a normal exit status or terminating signal when available. The UI must distinguish an exited shell from a transport failure.
ui/src/api.ts
+1
View File
@@ -12,6 +12,7 @@ export type TerminalSession = {
terminal_id: string;
agent_id: string;
created_at_unix: number;
expires_at_unix?: number | null;
agent_attached: boolean;
operator_attached: boolean;
websocket_url: string;
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);
wakey-control-plane/src/api/terminals.rs
+38 -14
View File
@@ -5,16 +5,16 @@ use axum::http::StatusCode;
use axum::response::Response;
use futures_util::{SinkExt, StreamExt};
use serde::{Deserialize, Serialize};
use std::time::Duration;
use tracing::{info, warn};
use wakey_agent::protocol::{
AgentCapability, DEFAULT_TERMINAL_MAX_SESSIONS, ServerMessage, TerminalAgentHandshake,
TerminalControl, TerminalId, TerminalOperatorHandshake,
AgentCapability, DEFAULT_TERMINAL_MAX_SESSIONS, DEFAULT_TERMINAL_SESSION_TTL_SECONDS,
ServerMessage, TerminalAgentHandshake, TerminalControl, TerminalId, TerminalOperatorHandshake,
};
use crate::api::ApiError;
use crate::runtime::terminals::{
TERMINAL_ABSOLUTE_TIMEOUT, TERMINAL_ATTACH_TIMEOUT, TERMINAL_MAX_FRAME_BYTES,
TerminalRelayFrame, TerminalSummary,
TERMINAL_ATTACH_TIMEOUT, TERMINAL_MAX_FRAME_BYTES, TerminalRelayFrame, TerminalSummary,
};
use crate::runtime::{AppState, SessionEvent};
use crate::state::AuditEventInput;
@@ -40,6 +40,7 @@ pub struct TerminalSessionResponse {
pub terminal_id: String,
pub agent_id: String,
pub created_at_unix: u64,
pub expires_at_unix: Option<u64>,
pub agent_attached: bool,
pub operator_attached: bool,
pub websocket_url: String,
@@ -52,7 +53,7 @@ pub async fn create_terminal(
Json(request): Json<CreateTerminalRequest>,
) -> Result<(StatusCode, Json<TerminalSessionResponse>), ApiError> {
validate_size(request.rows, request.cols)?;
let (agent_tx, max_sessions) = {
let (agent_tx, max_sessions, session_ttl_seconds) = {
let sessions = state.sessions.read().await;
let session = sessions.get(&request.agent_id).ok_or_else(|| {
ApiError::new(
@@ -75,12 +76,18 @@ pub async fn create_terminal(
.map(|terminal| terminal.max_sessions)
.unwrap_or(DEFAULT_TERMINAL_MAX_SESSIONS)
.max(1);
(session.tx.clone(), max_sessions)
let session_ttl_seconds = session
.capability_options
.terminal
.as_ref()
.map(|terminal| terminal.session_ttl_seconds)
.unwrap_or(DEFAULT_TERMINAL_SESSION_TTL_SECONDS);
(session.tx.clone(), max_sessions, session_ttl_seconds)
};
let created = state
.terminals
.create_with_limit(request.agent_id.clone(), max_sessions)
.create_with_limits(request.agent_id.clone(), max_sessions, session_ttl_seconds)
.await
.map_err(registry_error)?;
let terminal_id = TerminalId::new(created.terminal_id.clone()).map_err(|message| {
@@ -107,7 +114,12 @@ pub async fn create_terminal(
));
}
spawn_absolute_timeout(state.clone(), terminal_id.clone(), request.agent_id.clone());
spawn_session_timeout(
state.clone(),
terminal_id.clone(),
request.agent_id.clone(),
session_ttl_seconds,
);
append_terminal_audit(
&state,
&request.agent_id,
@@ -129,6 +141,7 @@ pub async fn create_terminal(
terminal_id: terminal_id.to_string(),
agent_id: request.agent_id,
created_at_unix: created.created_at_unix,
expires_at_unix: created.expires_at_unix,
agent_attached: false,
operator_attached: false,
attachment_token: Some(created.attachment_token),
@@ -140,7 +153,7 @@ pub async fn get_terminal(
State(state): State<AppState>,
Path(terminal_id): Path<String>,
) -> Result<Json<TerminalSessionResponse>, ApiError> {
let (agent_id, created_at_unix, agent_attached, operator_attached) = state
let (agent_id, created_at_unix, agent_attached, operator_attached, expires_at_unix) = state
.terminals
.summary(&terminal_id)
.await
@@ -150,6 +163,7 @@ pub async fn get_terminal(
terminal_id,
agent_id,
created_at_unix,
expires_at_unix,
agent_attached,
operator_attached,
attachment_token: None,
@@ -178,7 +192,7 @@ pub async fn attach_terminal(
.issue_attachment_token_for_operator(&terminal_id, &request.operator_id)
.await
.map_err(registry_error)?;
let (agent_id, created_at_unix, agent_attached, operator_attached) = state
let (agent_id, created_at_unix, agent_attached, operator_attached, expires_at_unix) = state
.terminals
.summary(&terminal_id)
.await
@@ -188,6 +202,7 @@ pub async fn attach_terminal(
terminal_id,
agent_id,
created_at_unix,
expires_at_unix,
agent_attached,
operator_attached,
attachment_token: Some(attachment_token),
@@ -350,7 +365,7 @@ async fn handle_operator_terminal_socket(
warn!(terminal_id, code, "failed to request terminal snapshot");
}
let summary = state.terminals.summary(&terminal_id).await;
if let Some((agent_id, _, _, _)) = &summary {
if let Some((agent_id, _, _, _, _)) = &summary {
append_terminal_audit(
&state,
agent_id,
@@ -367,7 +382,7 @@ async fn handle_operator_terminal_socket(
}
let (mut write, mut read) = socket.split();
if summary.is_some_and(|(_, _, agent_attached, _)| agent_attached) {
if summary.is_some_and(|(_, _, agent_attached, _, _)| agent_attached) {
let ready = serde_json::to_string(&TerminalControl::Ready)
.expect("terminal ready control serializes");
if send_relay_frame(&mut write, TerminalRelayFrame::Text(ready))
@@ -421,6 +436,7 @@ fn terminal_response(summary: TerminalSummary) -> TerminalSessionResponse {
terminal_id: summary.terminal_id,
agent_id: summary.agent_id,
created_at_unix: summary.created_at_unix,
expires_at_unix: summary.expires_at_unix,
agent_attached: summary.agent_attached,
operator_attached: summary.operator_attached,
attachment_token: None,
@@ -556,9 +572,17 @@ async fn close_registered_terminal(
Some(agent_id)
}
fn spawn_absolute_timeout(state: AppState, terminal_id: TerminalId, agent_id: String) {
fn spawn_session_timeout(
state: AppState,
terminal_id: TerminalId,
agent_id: String,
session_ttl_seconds: u64,
) {
if session_ttl_seconds == 0 {
return;
}
tokio::spawn(async move {
tokio::time::sleep(TERMINAL_ABSOLUTE_TIMEOUT).await;
tokio::time::sleep(Duration::from_secs(session_ttl_seconds)).await;
if close_registered_terminal(
&state,
terminal_id.as_str(),
wakey-control-plane/src/runtime/terminals.rs
+115 -22
View File
@@ -10,7 +10,6 @@ pub const TERMINAL_RELAY_QUEUE: usize = 32;
pub const TERMINAL_MAX_FRAME_BYTES: usize = 64 * 1024;
pub const TERMINAL_PENDING_AGENT_BYTES: usize = 256 * 1024;
pub const TERMINAL_ATTACH_TIMEOUT: Duration = Duration::from_secs(10);
pub const TERMINAL_ABSOLUTE_TIMEOUT: Duration = Duration::from_secs(12 * 60 * 60);
const TERMINAL_TOMBSTONE_TTL: Duration = Duration::from_secs(5 * 60);
const TERMINAL_MAX_TOMBSTONES: usize = 1024;
@@ -30,7 +29,8 @@ pub struct TerminalRegistry {
struct TerminalSession {
agent_id: String,
created_at_unix: u64,
expires_at: Instant,
expires_at: Option<Instant>,
expires_at_unix: Option<u64>,
agent_confirmed: bool,
relay_token: Option<String>,
attachment_token: Option<String>,
@@ -49,6 +49,7 @@ pub struct CreatedTerminal {
pub relay_token: String,
pub attachment_token: String,
pub created_at_unix: u64,
pub expires_at_unix: Option<u64>,
}
#[derive(Clone, Debug)]
@@ -56,6 +57,7 @@ pub struct TerminalSummary {
pub terminal_id: String,
pub agent_id: String,
pub created_at_unix: u64,
pub expires_at_unix: Option<u64>,
pub agent_attached: bool,
pub operator_attached: bool,
}
@@ -76,9 +78,10 @@ impl TerminalRegistry {
#[cfg(test)]
async fn create(&self, agent_id: String) -> Result<CreatedTerminal, &'static str> {
self.create_with_limit(
self.create_with_limits(
agent_id,
wakey_agent::protocol::DEFAULT_TERMINAL_MAX_SESSIONS,
wakey_agent::protocol::DEFAULT_TERMINAL_SESSION_TTL_SECONDS,
)
.await
}
@@ -86,10 +89,11 @@ impl TerminalRegistry {
/// Creates a session using the limit advertised by the connected agent.
/// The caller snapshots the limit with the agent connection so UI hints and
/// server-side enforcement use the same value.
pub async fn create_with_limit(
pub async fn create_with_limits(
&self,
agent_id: String,
max_sessions: usize,
session_ttl_seconds: u64,
) -> Result<CreatedTerminal, &'static str> {
let mut sessions = self.inner.lock().await;
prune_expired_sessions(&mut sessions);
@@ -109,12 +113,15 @@ impl TerminalRegistry {
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
.as_secs();
let expiry = expiry_for_created_at(created_at_unix, session_ttl_seconds)
.ok_or("terminal_ttl_invalid")?;
sessions.insert(
terminal_id.clone(),
TerminalSession {
agent_id,
created_at_unix,
expires_at: Instant::now() + TERMINAL_ABSOLUTE_TIMEOUT,
expires_at: expiry.deadline,
expires_at_unix: expiry.unix,
agent_confirmed: false,
relay_token: Some(relay_token.clone()),
attachment_token: Some(attachment_token.clone()),
@@ -134,6 +141,7 @@ impl TerminalRegistry {
relay_token,
attachment_token,
created_at_unix,
expires_at_unix: expiry.unix,
})
}
@@ -184,7 +192,10 @@ impl TerminalRegistry {
) -> Vec<(TerminalId, String)> {
let reported_ids = reported
.iter()
.filter(|session| expires_at_for_created_at(session.created_at_unix).is_some())
.filter(|session| {
expiry_for_created_at(session.created_at_unix, session.session_ttl_seconds)
.is_some()
})
.map(|session| session.terminal_id.as_str())
.collect::<std::collections::HashSet<_>>();
let stale_ids = {
@@ -207,8 +218,10 @@ impl TerminalRegistry {
let mut credentials = Vec::new();
let mut sessions = self.inner.lock().await;
for reported_session in reported {
let Some(expires_at) = expires_at_for_created_at(reported_session.created_at_unix)
else {
let Some(expiry) = expiry_for_created_at(
reported_session.created_at_unix,
reported_session.session_ttl_seconds,
) else {
continue;
};
let terminal_id = reported_session.terminal_id.as_str().to_string();
@@ -217,7 +230,8 @@ impl TerminalRegistry {
.or_insert_with(|| TerminalSession {
agent_id: agent_id.to_string(),
created_at_unix: reported_session.created_at_unix,
expires_at,
expires_at: expiry.deadline,
expires_at_unix: expiry.unix,
agent_confirmed: true,
relay_token: None,
attachment_token: None,
@@ -437,7 +451,10 @@ impl TerminalRegistry {
Some(detached_at)
}
pub async fn summary(&self, terminal_id: &str) -> Option<(String, u64, bool, bool)> {
pub async fn summary(
&self,
terminal_id: &str,
) -> Option<(String, u64, bool, bool, Option<u64>)> {
let mut sessions = self.inner.lock().await;
prune_expired_sessions(&mut sessions);
sessions.get(terminal_id).map(|session| {
@@ -446,6 +463,7 @@ impl TerminalRegistry {
session.created_at_unix,
session.agent_tx.is_some(),
session.operator_tx.is_some(),
session.expires_at_unix,
)
})
}
@@ -459,6 +477,7 @@ impl TerminalRegistry {
terminal_id: terminal_id.clone(),
agent_id: session.agent_id.clone(),
created_at_unix: session.created_at_unix,
expires_at_unix: session.expires_at_unix,
agent_attached: session.agent_tx.is_some(),
operator_attached: session.operator_tx.is_some(),
})
@@ -474,7 +493,8 @@ fn active_session<'a>(
) -> Result<&'a mut TerminalSession, &'static str> {
if sessions
.get(terminal_id)
.is_some_and(|session| session.expires_at <= Instant::now())
.and_then(|session| session.expires_at)
.is_some_and(|expires_at| expires_at <= Instant::now())
{
sessions.remove(terminal_id);
return Err("terminal_expired");
@@ -514,22 +534,38 @@ fn prune_tombstones(closed: &mut HashMap<String, Instant>) {
fn prune_expired_sessions(sessions: &mut HashMap<String, TerminalSession>) {
let now = Instant::now();
sessions.retain(|_, session| session.expires_at > now);
sessions.retain(|_, session| session.expires_at.is_none_or(|expires_at| expires_at > now));
}
/// Converts the agent's durable creation timestamp into CC's monotonic
/// deadline. Reconciliation must not grant an old PTY a fresh twelve hours.
fn expires_at_for_created_at(created_at_unix: u64) -> Option<Instant> {
#[derive(Clone, Copy)]
struct TerminalExpiry {
deadline: Option<Instant>,
unix: Option<u64>,
}
/// Converts the agent's creation-time policy into CC's monotonic deadline.
/// The outer `Option` rejects expired or unrepresentable sessions; an inner
/// `None` deadline is the explicit zero-TTL (unlimited) policy.
fn expiry_for_created_at(created_at_unix: u64, session_ttl_seconds: u64) -> Option<TerminalExpiry> {
if session_ttl_seconds == 0 {
return Some(TerminalExpiry {
deadline: None,
unix: None,
});
}
let expires_at_unix = created_at_unix.checked_add(session_ttl_seconds)?;
let now_unix = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
.as_secs();
let age = Duration::from_secs(now_unix.saturating_sub(created_at_unix));
if age >= TERMINAL_ABSOLUTE_TIMEOUT {
if expires_at_unix <= now_unix {
return None;
}
let remaining = TERMINAL_ABSOLUTE_TIMEOUT - age;
Some(Instant::now() + remaining)
let remaining = Duration::from_secs(expires_at_unix - now_unix);
Some(TerminalExpiry {
deadline: Some(Instant::now().checked_add(remaining)?),
unix: Some(expires_at_unix),
})
}
fn relay_frame_size(frame: &TerminalRelayFrame) -> usize {
@@ -543,6 +579,7 @@ fn relay_frame_size(frame: &TerminalRelayFrame) -> usize {
#[cfg(test)]
mod tests {
use super::*;
use wakey_agent::protocol::DEFAULT_TERMINAL_SESSION_TTL_SECONDS;
const OPERATOR_A: &str = "browser-tab-a";
const OPERATOR_B: &str = "browser-tab-b";
@@ -663,14 +700,68 @@ mod tests {
let registry = TerminalRegistry::new();
for _ in 0..3 {
registry
.create_with_limit("router".into(), 3)
.create_with_limits("router".into(), 3, DEFAULT_TERMINAL_SESSION_TTL_SECONDS)
.await
.expect("within advertised limit");
}
let fourth = registry.create_with_limit("router".into(), 3).await;
let fourth = registry
.create_with_limits("router".into(), 3, DEFAULT_TERMINAL_SESSION_TTL_SECONDS)
.await;
assert_eq!(fourth.err(), Some("agent_terminal_limit_reached"));
}
#[tokio::test]
async fn session_ttl_is_captured_and_zero_is_unlimited() {
let registry = TerminalRegistry::new();
let finite = registry
.create_with_limits("router".into(), 2, 67)
.await
.expect("finite session");
assert_eq!(
finite.expires_at_unix,
finite.created_at_unix.checked_add(67)
);
registry.remove(&finite.terminal_id).await;
let unlimited = registry
.create_with_limits("router".into(), 2, 0)
.await
.expect("unlimited session");
assert_eq!(unlimited.expires_at_unix, None);
assert_eq!(
registry
.summary(&unlimited.terminal_id)
.await
.expect("unlimited summary")
.4,
None
);
}
#[tokio::test]
async fn unlimited_inventory_can_adopt_an_old_session() {
let registry = TerminalRegistry::new();
let terminal_id = TerminalId::new("old-unlimited").expect("terminal id");
let reported = AgentTerminalSession {
terminal_id: terminal_id.clone(),
created_at_unix: 0,
session_ttl_seconds: 0,
};
let credentials = registry
.reconcile_agent_sessions("router", &[reported])
.await;
assert_eq!(credentials.len(), 1);
assert_eq!(
registry
.summary(terminal_id.as_str())
.await
.expect("adopted unlimited session")
.4,
None
);
}
#[tokio::test]
async fn removed_session_is_distinct_from_unknown_session() {
let registry = TerminalRegistry::new();
@@ -691,6 +782,7 @@ mod tests {
let reported = AgentTerminalSession {
terminal_id: terminal_id.clone(),
created_at_unix,
session_ttl_seconds: 0,
};
let credentials = registry
@@ -723,7 +815,7 @@ mod tests {
.await
.get_mut(&first.terminal_id)
.expect("first session")
.expires_at = Instant::now();
.expires_at = Some(Instant::now());
registry
.create("router".into())
@@ -739,6 +831,7 @@ mod tests {
let reported = AgentTerminalSession {
terminal_id: terminal_id.clone(),
created_at_unix: 0,
session_ttl_seconds: DEFAULT_TERMINAL_SESSION_TTL_SECONDS,
};
assert!(
wakey-control-plane/src/ws.rs
+29 -12
View File
@@ -4,13 +4,14 @@ use axum::extract::ws::{Message, WebSocket, WebSocketUpgrade};
use axum::response::IntoResponse;
use futures_util::{SinkExt, StreamExt};
use serde::Deserialize;
use std::time::Instant;
use std::time::{Duration, Instant};
use tokio::sync::mpsc;
use tracing::{debug, info, info_span, warn};
use uuid::Uuid;
use wakey_agent::protocol::{
AgentCapability, AgentCapabilityOptions, AgentTerminalSession, DEFAULT_TERMINAL_MAX_SESSIONS,
ErrorPayload, RequestId, ServerMessage, TerminalCapabilityOptions, TerminalControl, TerminalId,
DEFAULT_TERMINAL_SESSION_TTL_SECONDS, ErrorPayload, RequestId, ServerMessage,
TerminalCapabilityOptions, TerminalControl, TerminalId,
};
use wakey_core::Device;
@@ -212,17 +213,33 @@ async fn process_agent_text(
let connect_to_hello_ms = connected_at.elapsed().as_millis() as u64;
info!(agent_id = %agent_id, connect_to_hello_ms, "agent hello received");
connection.hello_agent_id = Some(agent_id);
connection.capability_options = AgentCapabilityOptions {
terminal: capabilities.contains(&AgentCapability::Terminal).then_some(
TerminalCapabilityOptions {
max_sessions: capability_options
.terminal
.map(|terminal| terminal.max_sessions)
.unwrap_or(DEFAULT_TERMINAL_MAX_SESSIONS)
.max(1),
},
),
let terminal = if capabilities.contains(&AgentCapability::Terminal) {
let max_sessions = capability_options
.terminal
.as_ref()
.map(|terminal| terminal.max_sessions)
.unwrap_or(DEFAULT_TERMINAL_MAX_SESSIONS)
.max(1);
let session_ttl_seconds = capability_options
.terminal
.as_ref()
.map(|terminal| terminal.session_ttl_seconds)
.unwrap_or(DEFAULT_TERMINAL_SESSION_TTL_SECONDS);
if session_ttl_seconds != 0
&& Instant::now()
.checked_add(Duration::from_secs(session_ttl_seconds))
.is_none()
{
anyhow::bail!("terminal session TTL is too large for the platform timer");
}
Some(TerminalCapabilityOptions {
max_sessions,
session_ttl_seconds,
})
} else {
None
};
connection.capability_options = AgentCapabilityOptions { terminal };
connection.capabilities = capabilities;
}
IncomingClientMessage::Auth {