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

support long running session

Browse Source 
plus UI reattaching to existing session
This commit is contained in:
lda committed 2026-07-15 08:16:10 +07:00 Verified
1 parent 4d793136ba
commit f459bb6914
12 files changed
+751 -148
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

wakey-control-plane/src/runtime/mod.rs
+4 -1
View File
@@ -142,7 +142,10 @@ fn control_api_routes() -> Router<AppState> {
"/api/v1/control/agents/{agent_id}/command",
post(api::run_command),
)
.route("/api/v1/control/terminals", post(api::create_terminal))
.route(
"/api/v1/control/terminals",
get(api::list_terminals).post(api::create_terminal),
)
.route(
"/api/v1/control/terminals/{terminal_id}",
get(api::get_terminal).delete(api::close_terminal),
wakey-control-plane/src/runtime/terminals.rs
+190 -34
View File
@@ -4,14 +4,14 @@ use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
use tokio::sync::{Mutex, mpsc};
use uuid::Uuid;
use wakey_agent::protocol::{AgentTerminalSession, TerminalId};
pub const TERMINAL_RELAY_QUEUE: usize = 32;
pub const TERMINAL_MAX_FRAME_BYTES: usize = 64 * 1024;
pub const TERMINAL_REPLAY_BYTES: usize = 256 * 1024;
pub const TERMINAL_MAX_SESSIONS_PER_AGENT: usize = 2;
pub const TERMINAL_ATTACH_TIMEOUT: Duration = Duration::from_secs(10);
pub const TERMINAL_DISCONNECT_GRACE: Duration = Duration::from_secs(15);
pub const TERMINAL_ABSOLUTE_TIMEOUT: Duration = Duration::from_secs(30 * 60);
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;
@@ -32,6 +32,7 @@ struct TerminalSession {
agent_id: String,
created_at_unix: u64,
expires_at: Instant,
agent_confirmed: bool,
relay_token: Option<String>,
attachment_token: Option<String>,
agent_tx: Option<mpsc::Sender<TerminalRelayFrame>>,
@@ -50,6 +51,15 @@ pub struct CreatedTerminal {
pub created_at_unix: u64,
}
#[derive(Clone, Debug)]
pub struct TerminalSummary {
pub terminal_id: String,
pub agent_id: String,
pub created_at_unix: u64,
pub agent_attached: bool,
pub operator_attached: bool,
}
impl Default for TerminalRegistry {
fn default() -> Self {
Self::new()
@@ -88,6 +98,7 @@ impl TerminalRegistry {
agent_id,
created_at_unix,
expires_at: Instant::now() + TERMINAL_ABSOLUTE_TIMEOUT,
agent_confirmed: false,
relay_token: Some(relay_token.clone()),
attachment_token: Some(attachment_token.clone()),
agent_tx: None,
@@ -146,18 +157,70 @@ impl TerminalRegistry {
}
}
pub async fn remove_agent(&self, agent_id: &str) {
let terminal_ids = {
/// Rebuilds CC's volatile catalog from the PTYs still owned by an agent.
/// Returned credentials tell those workers to establish fresh relay sockets.
pub async fn reconcile_agent_sessions(
&self,
agent_id: &str,
reported: &[AgentTerminalSession],
) -> Vec<(TerminalId, String)> {
let reported_ids = reported
.iter()
.map(|session| session.terminal_id.as_str())
.collect::<std::collections::HashSet<_>>();
let stale_ids = {
let sessions = self.inner.lock().await;
sessions
.iter()
.filter(|(_, session)| session.agent_id == agent_id)
.filter(|(terminal_id, session)| {
session.agent_id == agent_id
&& session.agent_confirmed
&& !reported_ids.contains(terminal_id.as_str())
})
.map(|(terminal_id, _)| terminal_id.clone())
.collect::<Vec<_>>()
};
for terminal_id in terminal_ids {
for terminal_id in stale_ids {
self.remove(&terminal_id).await;
}
let mut credentials = Vec::new();
let mut sessions = self.inner.lock().await;
for reported_session in reported {
let terminal_id = reported_session.terminal_id.as_str().to_string();
let session = sessions
.entry(terminal_id)
.or_insert_with(|| TerminalSession {
agent_id: agent_id.to_string(),
created_at_unix: reported_session.created_at_unix,
expires_at: Instant::now() + TERMINAL_ABSOLUTE_TIMEOUT,
agent_confirmed: true,
relay_token: None,
attachment_token: None,
agent_tx: None,
pending_agent: VecDeque::new(),
pending_agent_bytes: 0,
operator_tx: None,
operator_detached_at: Some(Instant::now()),
replay: VecDeque::new(),
replay_bytes: 0,
});
if session.agent_id != agent_id || session.agent_tx.is_some() {
continue;
}
session.agent_confirmed = true;
let token = new_token();
session.relay_token = Some(token.clone());
credentials.push((reported_session.terminal_id.clone(), token));
}
credentials
}
pub async fn detach_agent(&self, terminal_id: &str) -> Option<String> {
let mut sessions = self.inner.lock().await;
let session = sessions.get_mut(terminal_id)?;
session.agent_tx = None;
Some(session.agent_id.clone())
}
pub async fn issue_attachment_token(&self, terminal_id: &str) -> Result<String, &'static str> {
@@ -189,6 +252,7 @@ impl TerminalRegistry {
return Err("terminal_relay_token_invalid");
}
session.relay_token = None;
session.agent_confirmed = true;
let (tx, rx) = mpsc::channel(TERMINAL_RELAY_QUEUE);
session.agent_tx = Some(tx);
let pending = session.pending_agent.drain(..).collect();
@@ -211,8 +275,9 @@ impl TerminalRegistry {
}
session.attachment_token = None;
session.operator_detached_at = None;
let replay = session.replay.drain(..).collect();
session.replay_bytes = 0;
// Keep the rolling transcript after attachment so a newly mounted
// browser can reconstruct recent terminal state.
let replay = session.replay.iter().cloned().collect();
let (tx, rx) = mpsc::channel(TERMINAL_RELAY_QUEUE);
session.operator_tx = Some(tx);
Ok((rx, replay))
@@ -255,17 +320,17 @@ impl TerminalRegistry {
terminal_id: &str,
frame: TerminalRelayFrame,
) -> Result<(), &'static str> {
let operator_tx = self
.inner
.lock()
.await
.get(terminal_id)
.and_then(|session| session.operator_tx.clone());
let operator_tx = {
let mut sessions = self.inner.lock().await;
let session = active_session(&mut sessions, terminal_id)?;
push_replay(session, frame.clone());
session.operator_tx.clone()
};
if let Some(tx) = operator_tx {
match tx.send(frame).await {
Ok(()) => return Ok(()),
Err(err) => {
Err(_) => {
// The browser task may not have marked itself detached yet.
// Preserve this frame so that race does not kill the PTY.
let mut sessions = self.inner.lock().await;
@@ -274,15 +339,10 @@ impl TerminalRegistry {
session
.operator_detached_at
.get_or_insert_with(Instant::now);
push_replay(session, err.0);
return Ok(());
}
}
}
let mut sessions = self.inner.lock().await;
let session = active_session(&mut sessions, terminal_id)?;
push_replay(session, frame);
Ok(())
}
@@ -303,8 +363,10 @@ impl TerminalRegistry {
}
self.relay_from_agent(terminal_id, TerminalRelayFrame::Text(error_json))
.await?;
self.relay_from_agent(terminal_id, TerminalRelayFrame::Close)
self.remove(terminal_id)
.await
.map(|_| ())
.ok_or("terminal_not_found")
}
pub async fn detach_operator(&self, terminal_id: &str) -> Option<Instant> {
@@ -316,19 +378,6 @@ impl TerminalRegistry {
Some(detached_at)
}
pub async fn remove_if_still_detached(&self, terminal_id: &str, detached_at: Instant) -> bool {
let should_remove = self
.inner
.lock()
.await
.get(terminal_id)
.is_some_and(|session| session.operator_detached_at == Some(detached_at));
if should_remove {
self.remove(terminal_id).await;
}
should_remove
}
pub async fn summary(&self, terminal_id: &str) -> Option<(String, u64, bool, bool)> {
self.inner.lock().await.get(terminal_id).map(|session| {
(
@@ -339,6 +388,22 @@ impl TerminalRegistry {
)
})
}
pub async fn summaries(&self) -> Vec<TerminalSummary> {
let sessions = self.inner.lock().await;
let mut summaries = sessions
.iter()
.map(|(terminal_id, session)| TerminalSummary {
terminal_id: terminal_id.clone(),
agent_id: session.agent_id.clone(),
created_at_unix: session.created_at_unix,
agent_attached: session.agent_tx.is_some(),
operator_attached: session.operator_tx.is_some(),
})
.collect::<Vec<_>>();
summaries.sort_by_key(|session| std::cmp::Reverse(session.created_at_unix));
summaries
}
}
fn active_session<'a>(
@@ -440,6 +505,37 @@ mod tests {
assert!(replay.iter().map(relay_frame_size).sum::<usize>() <= TERMINAL_REPLAY_BYTES);
}
#[tokio::test]
async fn attached_output_remains_available_for_remount() {
let registry = TerminalRegistry::new();
let created = registry.create("router".into()).await.expect("create");
let (mut outbound, _) = registry
.attach_operator(&created.terminal_id, &created.attachment_token)
.await
.expect("attach operator");
let frame = TerminalRelayFrame::Binary(b"recent prompt".to_vec());
registry
.relay_from_agent(&created.terminal_id, frame)
.await
.expect("relay output");
outbound.recv().await.expect("live output");
registry.detach_operator(&created.terminal_id).await;
let token = registry
.issue_attachment_token(&created.terminal_id)
.await
.expect("reattach token");
let (_, replay) = registry
.attach_operator(&created.terminal_id, &token)
.await
.expect("reattach operator");
assert!(matches!(
replay.as_slice(),
[TerminalRelayFrame::Binary(bytes)] if bytes == b"recent prompt"
));
}
#[tokio::test]
async fn operator_input_waits_for_agent_attachment() {
let registry = TerminalRegistry::new();
@@ -478,4 +574,64 @@ mod tests {
assert!(registry.was_closed(&created.terminal_id).await);
assert!(!registry.was_closed("never-existed").await);
}
#[tokio::test]
async fn agent_inventory_adopts_and_reconnects_live_session() {
let registry = TerminalRegistry::new();
let terminal_id = TerminalId::new("survived-cc").expect("terminal id");
let reported = AgentTerminalSession {
terminal_id: terminal_id.clone(),
created_at_unix: 42,
};
let credentials = registry
.reconcile_agent_sessions("router", &[reported])
.await;
assert_eq!(credentials.len(), 1);
let (_, relay_token) = &credentials[0];
registry
.attach_agent(terminal_id.as_str(), "router", relay_token)
.await
.expect("attach adopted agent session");
let summary = registry
.summary(terminal_id.as_str())
.await
.expect("adopted summary");
assert_eq!(summary.0, "router");
assert_eq!(summary.1, 42);
assert!(summary.2);
}
#[tokio::test]
async fn operator_detach_does_not_remove_session() {
let registry = TerminalRegistry::new();
let created = registry.create("router".into()).await.expect("create");
registry
.attach_operator(&created.terminal_id, &created.attachment_token)
.await
.expect("attach operator");
registry.detach_operator(&created.terminal_id).await;
let summary = registry
.summary(&created.terminal_id)
.await
.expect("session remains after detach");
assert!(!summary.3);
registry
.issue_attachment_token(&created.terminal_id)
.await
.expect("detached session can be reattached");
}
#[tokio::test]
async fn reconciliation_does_not_remove_open_request_still_in_flight() {
let registry = TerminalRegistry::new();
let created = registry.create("router".into()).await.expect("create");
registry.reconcile_agent_sessions("router", &[]).await;
assert!(registry.summary(&created.terminal_id).await.is_some());
}
}