support long running session
plus UI reattaching to existing session
This commit is contained in:
@@ -23,7 +23,7 @@ Session establishment follows this sequence:
|
||||
3. The control plane sends an open-terminal control message over the existing authenticated agent WebSocket.
|
||||
4. The agent opens a new outbound terminal WebSocket to the control plane and authenticates it for that terminal ID.
|
||||
5. The browser opens the corresponding protected operator terminal WebSocket.
|
||||
6. The control plane pairs the sockets and relays terminal frames until either side exits or disconnects.
|
||||
6. The control plane pairs the sockets and relays terminal frames while an operator is attached.
|
||||
|
||||
The main agent WebSocket carries terminal creation and cancellation control only. It does not carry PTY input or output.
|
||||
|
||||
@@ -65,13 +65,15 @@ WebSockets already provide ordered, reliable delivery, so terminal frames do not
|
||||
|
||||
## Lifecycle and Cleanup
|
||||
|
||||
Initial terminal sessions are ephemeral and allow only one attached operator.
|
||||
Terminal sessions allow only one attached operator. Their lifetime belongs to the agent process, not to a browser route or a particular control-plane connection.
|
||||
|
||||
When the operator socket disconnects, the control plane keeps the agent socket and PTY alive for a short grace period. A protected operator request may obtain a fresh, single-use attachment credential for that existing session. The control plane retains only a bounded in-memory output buffer during the gap, replays it before live output on reattachment, and closes the session if no operator returns before the grace period expires.
|
||||
When the operator socket disconnects, the session becomes detached. Navigation, browser closure, and network loss do not terminate the PTY. A protected operator request may discover the live session and obtain a fresh, single-use attachment credential. The control plane retains only a bounded in-memory output buffer during the gap and replays it before live output on reattachment.
|
||||
|
||||
An agent terminal socket disconnect closes the session immediately because the control plane can no longer control or observe the PTY. Agent reconnect creates a new agent session and cannot adopt an old terminal.
|
||||
The agent keeps its terminal manager outside the control-WebSocket reconnect loop. If a dedicated terminal relay disconnects, the PTY continues draining output into a bounded local replay buffer while waiting for replacement relay credentials. On every authenticated control connection, the agent reports its in-memory live terminal IDs and creation times. The control plane reconciles that inventory, adopts sessions missing from its volatile registry, and issues fresh relay credentials. This allows sessions to survive control-plane restart without persisting terminal state in SQLite.
|
||||
|
||||
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. Agent reconnect, agent replacement, control-plane restart, token expiry before attachment, and absolute session timeout all close the session.
|
||||
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.
|
||||
|
||||
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.
|
||||
|
||||
@@ -79,14 +81,14 @@ The agent reports a normal exit status or terminating signal when available. The
|
||||
|
||||
Terminal transport must not use unbounded queues.
|
||||
|
||||
Every queue between PTY, agent socket, control-plane relay, and browser socket is bounded. The brief-disconnect replay buffer is bounded as well. When the downstream consumer is slow, the producer waits rather than accumulating output in memory. Backpressure eventually stops PTY reads and allows the kernel PTY buffer to block an abusive producer.
|
||||
Every queue between PTY, agent socket, control-plane relay, and browser socket is bounded. Agent and control-plane replay buffers are bounded as well. The agent keeps draining the PTY while detached so a noisy child cannot stall the agent; output older than the replay bound is discarded.
|
||||
|
||||
The implementation also enforces:
|
||||
|
||||
- maximum terminal frame size;
|
||||
- per-agent concurrent-session limits;
|
||||
- idle and absolute session timeouts;
|
||||
- bounded disconnect grace;
|
||||
- bounded detached-session replay;
|
||||
- bounded control-plane relay buffers; and
|
||||
- cleanup when any relay task exits unexpectedly.
|
||||
|
||||
@@ -104,7 +106,7 @@ Terminal capability and session state are separate from device inventory, known-
|
||||
|
||||
## Deferred Work
|
||||
|
||||
Durable or long-lived terminal sessions are not part of the initial version. A later design may preserve sessions across longer operator absences or control-plane restart. That design must define durable ownership, reconnect credentials, persisted replay bounds, secret handling, and process reconciliation before adding persistence.
|
||||
Persistence across agent restart or machine reboot is not supported. A later design would need an external process owner, reconnectable local IPC, and explicit secret handling; persisted metadata alone cannot recover a PTY.
|
||||
|
||||
File transfer, multi-operator attachment, terminal recording, shell-history storage, and arbitrary process launch are also deferred.
|
||||
|
||||
@@ -130,7 +132,7 @@ Viable, but not selected initially. Its cross-platform abstraction is broader th
|
||||
|
||||
- Each active terminal consumes two additional WebSockets and one PTY process on the agent.
|
||||
- Terminal traffic cannot delay main agent heartbeat, snapshot, or command messages.
|
||||
- Control-plane terminal state is simple and ephemeral; restarting the control plane closes active terminals.
|
||||
- Control-plane terminal state remains ephemeral; agents rebuild it from their in-memory live-session inventory after a control-plane restart.
|
||||
- The agent session protocol gains terminal-open control and capability advertisement but does not become a terminal-data multiplexer.
|
||||
- Browser and agent terminal transports can be tested independently before adding xterm.js.
|
||||
- Future resume support can extend terminal-session lifecycle without changing the basic PTY byte protocol.
|
||||
- Browser navigation and control-plane restart do not terminate an agent-owned PTY.
|
||||
|
||||
@@ -247,6 +247,10 @@ export function createTerminal(
|
||||
});
|
||||
}
|
||||
|
||||
export function listTerminals(): Promise<TerminalSession[]> {
|
||||
return request<TerminalSession[]>("/api/v1/control/terminals");
|
||||
}
|
||||
|
||||
export function attachTerminal(terminalId: string): Promise<TerminalSession> {
|
||||
return request<TerminalSession>(
|
||||
`/api/v1/control/terminals/${encodeURIComponent(terminalId)}/attach`,
|
||||
|
||||
@@ -10,6 +10,7 @@ import {
|
||||
attachTerminal,
|
||||
closeTerminal,
|
||||
createTerminal,
|
||||
listTerminals,
|
||||
} from "@/api";
|
||||
import { AgentSelector, displayAgentLabel } from "@/components/AgentSelector";
|
||||
import { Badge } from "@/components/ui/badge";
|
||||
@@ -174,10 +175,92 @@ export function TerminalPage({
|
||||
const fit = new FitAddon();
|
||||
terminal.loadAddon(fit);
|
||||
terminal.open(hostRef.current);
|
||||
terminal.attachCustomKeyEventHandler((event) => {
|
||||
if (event.type !== "keydown") return true;
|
||||
|
||||
const copy =
|
||||
(event.ctrlKey && event.shiftKey && event.code === "KeyC") ||
|
||||
(event.metaKey && event.code === "KeyC") ||
|
||||
(event.ctrlKey && event.code === "Insert");
|
||||
if (copy) {
|
||||
event.preventDefault();
|
||||
event.stopPropagation();
|
||||
if (terminal.hasSelection()) {
|
||||
if (!navigator.clipboard) {
|
||||
toast.error("Clipboard access requires HTTPS or localhost");
|
||||
} else {
|
||||
void navigator.clipboard
|
||||
.writeText(terminal.getSelection())
|
||||
.catch((error) =>
|
||||
toast.error("Clipboard access was denied", {
|
||||
description: String(error),
|
||||
}),
|
||||
);
|
||||
}
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
const paste =
|
||||
(event.ctrlKey && event.shiftKey && event.code === "KeyV") ||
|
||||
(event.metaKey && event.code === "KeyV") ||
|
||||
(event.shiftKey && event.code === "Insert");
|
||||
if (paste) {
|
||||
event.preventDefault();
|
||||
event.stopPropagation();
|
||||
if (!navigator.clipboard) {
|
||||
toast.error("Clipboard access requires HTTPS or localhost");
|
||||
return false;
|
||||
}
|
||||
void navigator.clipboard
|
||||
.readText()
|
||||
.then((text) => terminal.paste(text))
|
||||
.catch((error) =>
|
||||
toast.error("Clipboard access was denied", {
|
||||
description: String(error),
|
||||
}),
|
||||
);
|
||||
return false;
|
||||
}
|
||||
|
||||
return true;
|
||||
});
|
||||
fit.fit();
|
||||
terminalRef.current = terminal;
|
||||
fitRef.current = fit;
|
||||
|
||||
async function restoreDetachedSession() {
|
||||
try {
|
||||
const sessions = await listTerminals();
|
||||
if (cancelled || sessions.length === 0) return;
|
||||
const candidate =
|
||||
sessions.find((item) => item.agent_id === selectedAgentId) ??
|
||||
sessions[0];
|
||||
setConnection("connecting");
|
||||
|
||||
// The previous route's websocket may still be completing its close
|
||||
// handshake. Brief retries keep navigation from surfacing that race.
|
||||
for (let attempt = 0; attempt < 4; attempt += 1) {
|
||||
try {
|
||||
const attached = await attachTerminal(candidate.terminal_id);
|
||||
if (cancelled) return;
|
||||
setSession(attached);
|
||||
connect(attached);
|
||||
return;
|
||||
} catch (error) {
|
||||
if (attempt === 3) throw error;
|
||||
await new Promise((resolve) => window.setTimeout(resolve, 150));
|
||||
}
|
||||
}
|
||||
} catch (error) {
|
||||
if (cancelled) return;
|
||||
setConnection("idle");
|
||||
toast.error("Could not restore terminal session", {
|
||||
description: String(error),
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
let cancelled = false;
|
||||
void document.fonts.ready.then(() => {
|
||||
if (cancelled) return;
|
||||
@@ -208,6 +291,7 @@ export function TerminalPage({
|
||||
});
|
||||
});
|
||||
resizeObserver.observe(hostRef.current);
|
||||
void restoreDetachedSession();
|
||||
|
||||
return () => {
|
||||
cancelled = true;
|
||||
@@ -219,7 +303,7 @@ export function TerminalPage({
|
||||
terminalRef.current = null;
|
||||
fitRef.current = null;
|
||||
};
|
||||
}, [sendResize]);
|
||||
}, [connect, selectedAgentId, sendResize]);
|
||||
|
||||
async function start() {
|
||||
if (!selectedAgentId || !terminalRef.current) return;
|
||||
@@ -243,6 +327,7 @@ export function TerminalPage({
|
||||
async function reconnect() {
|
||||
if (!session) return;
|
||||
try {
|
||||
terminalRef.current?.reset();
|
||||
const attached = await attachTerminal(session.terminal_id);
|
||||
setSession(attached);
|
||||
connect(attached);
|
||||
|
||||
+5
-2
@@ -465,8 +465,11 @@
|
||||
fitted canvas, so keep it on the same surface as the terminal theme. */
|
||||
background-color: #0b1117;
|
||||
overscroll-behavior: contain;
|
||||
scrollbar-color: color-mix(in oklab, white 28%, transparent) transparent;
|
||||
scrollbar-width: thin;
|
||||
scrollbar-width: none;
|
||||
}
|
||||
|
||||
.terminal-surface .xterm-viewport::-webkit-scrollbar {
|
||||
display: none;
|
||||
}
|
||||
|
||||
@media (max-width: 640px) {
|
||||
|
||||
@@ -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\""));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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",
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
+299
-73
@@ -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)
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -19,7 +19,7 @@ pub use control::{
|
||||
};
|
||||
pub use terminals::{
|
||||
agent_terminal_ws, attach_terminal, close_terminal, create_terminal, get_terminal,
|
||||
operator_terminal_ws,
|
||||
list_terminals, operator_terminal_ws,
|
||||
};
|
||||
|
||||
use axum::response::{IntoResponse, Response};
|
||||
|
||||
@@ -13,8 +13,8 @@ use wakey_agent::protocol::{
|
||||
|
||||
use crate::api::ApiError;
|
||||
use crate::runtime::terminals::{
|
||||
TERMINAL_ABSOLUTE_TIMEOUT, TERMINAL_ATTACH_TIMEOUT, TERMINAL_DISCONNECT_GRACE,
|
||||
TERMINAL_MAX_FRAME_BYTES, TerminalRelayFrame,
|
||||
TERMINAL_ABSOLUTE_TIMEOUT, TERMINAL_ATTACH_TIMEOUT, TERMINAL_MAX_FRAME_BYTES,
|
||||
TerminalRelayFrame, TerminalSummary,
|
||||
};
|
||||
use crate::runtime::{AppState, SessionEvent};
|
||||
use crate::state::AuditEventInput;
|
||||
@@ -142,6 +142,18 @@ pub async fn get_terminal(
|
||||
}))
|
||||
}
|
||||
|
||||
pub async fn list_terminals(State(state): State<AppState>) -> Json<Vec<TerminalSessionResponse>> {
|
||||
Json(
|
||||
state
|
||||
.terminals
|
||||
.summaries()
|
||||
.await
|
||||
.into_iter()
|
||||
.map(terminal_response)
|
||||
.collect(),
|
||||
)
|
||||
}
|
||||
|
||||
pub async fn attach_terminal(
|
||||
State(state): State<AppState>,
|
||||
Path(terminal_id): Path<String>,
|
||||
@@ -284,11 +296,14 @@ async fn handle_agent_terminal_socket(state: AppState, terminal_id: String, mut
|
||||
}
|
||||
}
|
||||
}
|
||||
state.terminals.remove(&terminal_id).await;
|
||||
info!(
|
||||
terminal_id,
|
||||
"agent terminal socket detached; session closed"
|
||||
);
|
||||
if let Some(agent_id) = state.terminals.detach_agent(&terminal_id).await
|
||||
&& let Some(session) = state.sessions.read().await.get(&agent_id)
|
||||
{
|
||||
let _ = session
|
||||
.tx
|
||||
.send(SessionEvent::Message(ServerMessage::SyncTerminalSessions));
|
||||
}
|
||||
info!(terminal_id, "agent terminal relay detached");
|
||||
}
|
||||
|
||||
async fn handle_operator_terminal_socket(
|
||||
@@ -317,10 +332,11 @@ async fn handle_operator_terminal_socket(
|
||||
}
|
||||
};
|
||||
info!(terminal_id, "operator terminal socket attached");
|
||||
if let Some((agent_id, _, _, _)) = state.terminals.summary(&terminal_id).await {
|
||||
let summary = state.terminals.summary(&terminal_id).await;
|
||||
if let Some((agent_id, _, _, _)) = &summary {
|
||||
append_terminal_audit(
|
||||
&state,
|
||||
&agent_id,
|
||||
agent_id,
|
||||
&terminal_id,
|
||||
TerminalAudit {
|
||||
actor_type: "admin_api",
|
||||
@@ -334,6 +350,17 @@ async fn handle_operator_terminal_socket(
|
||||
}
|
||||
|
||||
let (mut write, mut read) = socket.split();
|
||||
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))
|
||||
.await
|
||||
.is_err()
|
||||
{
|
||||
state.terminals.detach_operator(&terminal_id).await;
|
||||
return;
|
||||
}
|
||||
}
|
||||
for frame in replay {
|
||||
if send_relay_frame(&mut write, frame).await.is_err() {
|
||||
return;
|
||||
@@ -367,15 +394,8 @@ async fn handle_operator_terminal_socket(
|
||||
|
||||
if explicit_close {
|
||||
close_registered_terminal(&state, &terminal_id).await;
|
||||
} else if let Some(detached_at) = state.terminals.detach_operator(&terminal_id).await {
|
||||
let terminals = state.terminals.clone();
|
||||
let terminal_id_for_grace = terminal_id.clone();
|
||||
tokio::spawn(async move {
|
||||
tokio::time::sleep(TERMINAL_DISCONNECT_GRACE).await;
|
||||
terminals
|
||||
.remove_if_still_detached(&terminal_id_for_grace, detached_at)
|
||||
.await;
|
||||
});
|
||||
} else {
|
||||
state.terminals.detach_operator(&terminal_id).await;
|
||||
}
|
||||
info!(
|
||||
terminal_id,
|
||||
@@ -383,6 +403,18 @@ async fn handle_operator_terminal_socket(
|
||||
);
|
||||
}
|
||||
|
||||
fn terminal_response(summary: TerminalSummary) -> TerminalSessionResponse {
|
||||
TerminalSessionResponse {
|
||||
websocket_url: operator_ws_path(&summary.terminal_id),
|
||||
terminal_id: summary.terminal_id,
|
||||
agent_id: summary.agent_id,
|
||||
created_at_unix: summary.created_at_unix,
|
||||
agent_attached: summary.agent_attached,
|
||||
operator_attached: summary.operator_attached,
|
||||
attachment_token: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn agent_relay_frame(message: Message) -> Result<Option<TerminalRelayFrame>, &'static str> {
|
||||
match message {
|
||||
Message::Binary(bytes) => Ok(Some(TerminalRelayFrame::Binary(bytes.to_vec()))),
|
||||
|
||||
@@ -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),
|
||||
|
||||
@@ -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());
|
||||
}
|
||||
}
|
||||
|
||||
@@ -9,7 +9,8 @@ use tokio::sync::mpsc;
|
||||
use tracing::{debug, info, info_span, warn};
|
||||
use uuid::Uuid;
|
||||
use wakey_agent::protocol::{
|
||||
AgentCapability, ErrorPayload, RequestId, ServerMessage, TerminalControl, TerminalId,
|
||||
AgentCapability, AgentTerminalSession, ErrorPayload, RequestId, ServerMessage, TerminalControl,
|
||||
TerminalId,
|
||||
};
|
||||
use wakey_core::Device;
|
||||
|
||||
@@ -47,6 +48,9 @@ enum IncomingClientMessage {
|
||||
terminal_id: TerminalId,
|
||||
error: ErrorPayload,
|
||||
},
|
||||
TerminalSessions {
|
||||
sessions: Vec<AgentTerminalSession>,
|
||||
},
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
@@ -150,7 +154,6 @@ async fn handle_agent_socket(state: AppState, socket: WebSocket) {
|
||||
sessions.remove(&agent_id);
|
||||
}
|
||||
drop(sessions);
|
||||
state.terminals.remove_agent(&agent_id).await;
|
||||
if let Err(err) = state
|
||||
.store
|
||||
.append_audit_event(AuditEventInput {
|
||||
@@ -338,6 +341,28 @@ async fn process_agent_text(
|
||||
}
|
||||
warn!(terminal_id = %terminal_id, agent_id, "agent rejected terminal request");
|
||||
}
|
||||
IncomingClientMessage::TerminalSessions { sessions } => {
|
||||
let agent_id = connection
|
||||
.authed_agent_id
|
||||
.as_deref()
|
||||
.ok_or_else(|| anyhow::anyhow!("terminal inventory before auth"))?;
|
||||
ensure_current_session(state, agent_id, connection_id).await?;
|
||||
let credentials = state
|
||||
.terminals
|
||||
.reconcile_agent_sessions(agent_id, &sessions)
|
||||
.await;
|
||||
for (terminal_id, relay_token) in credentials {
|
||||
let _ = tx.send(SessionEvent::Message(ServerMessage::ResumeTerminal {
|
||||
terminal_id,
|
||||
relay_token,
|
||||
}));
|
||||
}
|
||||
info!(
|
||||
agent_id,
|
||||
sessions = sessions.len(),
|
||||
"agent terminal sessions reconciled"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
|
||||
Reference in New Issue
Block a user