diff --git a/CHANGELOG.md b/CHANGELOG.md index 0c4227804..a33c7ee67 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,7 +9,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Added -- `agent-relay node agent attach --ssh-host ` now provides an explicit SSH fallback for physical fleet nodes without exporting the remote broker or its API key; `--node` remains reserved for canonical fleet-native attach. +- `agent-relay node agent attach --node ` now opens an authenticated terminal session for physical and Daytona fleet nodes, preserving view, drive, and passthrough modes. +- `agent-relay node agent attach --ssh-host ` now provides an explicit SSH fallback for physical fleet nodes without exporting the remote broker or its API key. - Spawned agents now stamp a `Session-Id:` git trailer on commits when the dispatcher supplies a session reference, enabling auditors to trace each commit back to the session that produced it. - `agent-relay node agent attach` now distinguishes between "agent does not exist" and "agent is running on a different fleet node": when a 404 resolves to a workspace-registered agent with a fleet placement, the error names the node (`agent 'X' is registered on node 'finn-mini'; cross-node attach is not yet supported`) instead of the indistinguishable "no agent named 'X'". diff --git a/crates/broker/src/lib.rs b/crates/broker/src/lib.rs index 1037b3a46..f7e9938c2 100644 --- a/crates/broker/src/lib.rs +++ b/crates/broker/src/lib.rs @@ -51,6 +51,7 @@ pub(crate) mod swarm; pub(crate) mod swarm_tui; #[allow(dead_code)] pub(crate) mod telemetry; +pub(crate) mod terminal_control; #[allow(dead_code)] pub(crate) mod types; pub(crate) mod util; diff --git a/crates/broker/src/runtime/api.rs b/crates/broker/src/runtime/api.rs index 59dffdb2c..cde595106 100644 --- a/crates/broker/src/runtime/api.rs +++ b/crates/broker/src/runtime/api.rs @@ -1,4 +1,6 @@ +use super::fleet::try_send_terminal; use super::*; +use crate::terminal_control::TerminalToCloud; use relaycast::{ CreateObserverTokenRequest, ObserverScope, ObserverToken, ObserverTokenFilters, RelayError, }; @@ -260,6 +262,10 @@ impl BrokerRuntime { let resize_owners = &mut self.resize_owners; let delivery_states = &mut self.delivery_states; let agent_result_tokens = &mut self.agent_result_tokens; + let terminal_control_tx = &self.terminal_control_tx; + let terminal_sessions = &mut self.terminal_sessions; + let terminal_snapshot_requests = &mut self.terminal_snapshot_requests; + let terminal_input_requests = &mut self.terminal_input_requests; let dedup = &mut self.dedup; let recent_thread_messages = &mut self.recent_thread_messages; let delivery_retry_interval = self.delivery_retry_interval; @@ -825,6 +831,28 @@ impl BrokerRuntime { fail_pending_requests_for_worker(pending_requests, &name, "agent_released"); resize_owners.remove(&name); pty_observability.remove(&name); + let terminal_session_ids: Vec = terminal_sessions + .iter() + .filter(|(_, session)| session.agent == name) + .map(|(session_id, _)| session_id.clone()) + .collect(); + for session_id in terminal_session_ids { + terminal_sessions.remove(&session_id); + terminal_snapshot_requests + .retain(|_, pending| pending.session_id != session_id); + terminal_input_requests + .retain(|_, pending| pending.session_id != session_id); + if !try_send_terminal( + terminal_control_tx, + TerminalToCloud::Closed { + session_id: session_id.clone(), + code: Some("agent_released".into()), + message: Some("terminal worker was released".into()), + }, + ) { + tracing::warn!(target = "relay_broker::terminal", session_id = %session_id, "terminal queue full or closed while closing released worker session"); + } + } delivery_states.remove(&name); agent_result_tokens.retain(|_, agent| agent != &name); state.agents.remove(&name); diff --git a/crates/broker/src/runtime/event_loop.rs b/crates/broker/src/runtime/event_loop.rs index b92971061..1bb2d1997 100644 --- a/crates/broker/src/runtime/event_loop.rs +++ b/crates/broker/src/runtime/event_loop.rs @@ -204,6 +204,18 @@ pub(crate) struct BrokerRuntime { pub(super) node_delivery_connected: bool, pub(super) fleet_event_rx: mpsc::Receiver, pub(super) fleet_control_open: bool, + /// Independent outbound terminal lane. It never shares the node-control + /// socket, keeping high-volume PTY bytes away from heartbeats/actions. + pub(super) terminal_control_tx: mpsc::Sender, + pub(super) terminal_event_rx: mpsc::Receiver, + pub(super) terminal_control_open: bool, + pub(super) terminal_sessions: HashMap, + /// Worker snapshot RPCs parked for terminal opens. These are kept outside + /// the HTTP request map because their reply belongs on the terminal lane. + pub(super) terminal_snapshot_requests: HashMap, + /// PTY writes are acknowledged only after the worker drainer confirms + /// them; this map correlates that response back to its terminal session. + pub(super) terminal_input_requests: HashMap, pub(super) fleet_delivery_book: FleetDeliveryBook, pub(super) fleet_max_agents: u32, pub(super) fleet_inventory: HashMap, @@ -259,10 +271,35 @@ enum RuntimeEvent { Stdin(std::io::Result>), Relaycast(Option), Fleet(Option), + Terminal(Option), Worker(Option), MaintenanceTick, } +#[derive(Debug, Clone)] +pub(super) struct TerminalSession { + pub(super) agent: WorkerName, + pub(super) mode: TerminalMode, + /// PTY output is withheld until the initial ANSI grid has been delivered, + /// so a client always receives `terminal.ready` before stream chunks. + pub(super) ready: bool, + /// Output observed while the snapshot RPC is in flight. It is forwarded + /// after `terminal.ready`; per-stream offsets let the client discard bytes + /// already represented by the snapshot without losing later bytes. + pub(super) pending_output: Vec<(String, Option)>, + pub(super) pending_output_bytes: usize, +} + +pub(super) struct TerminalSnapshotRequest { + pub(super) session_id: String, + pub(super) deadline: Instant, +} + +pub(super) struct TerminalInputRequest { + pub(super) session_id: String, + pub(super) deadline: Instant, +} + impl BrokerRuntime { pub(super) async fn run(mut self) -> Result<()> { while !self.shutdown { @@ -277,6 +314,7 @@ impl BrokerRuntime { result = self.sdk_lines.next_line(), if self.stdin_open => RuntimeEvent::Stdin(result), message = self.ws_inbound_rx.recv(), if self.relaycast_open => RuntimeEvent::Relaycast(message), event = self.fleet_event_rx.recv(), if self.fleet_control_open => RuntimeEvent::Fleet(event), + event = self.terminal_event_rx.recv(), if self.terminal_control_open => RuntimeEvent::Terminal(event), event = self.worker_event_rx.recv(), if self.worker_events_open => RuntimeEvent::Worker(event), _ = self.reap_tick.tick() => RuntimeEvent::MaintenanceTick, }; @@ -315,6 +353,12 @@ impl BrokerRuntime { RuntimeEvent::Fleet(None) => { self.fleet_control_open = false; } + RuntimeEvent::Terminal(Some(event)) => { + self.handle_terminal_control_event(event).await; + } + RuntimeEvent::Terminal(None) => { + self.terminal_control_open = false; + } RuntimeEvent::Worker(Some(event)) => { self.handle_worker_event(event).await; } @@ -438,6 +482,17 @@ impl BrokerRuntime { { tracing::debug!(error = %error, "failed to send fleet control shutdown signal"); } + // A terminal client can be awaiting a full event queue while its own + // bounded command queue is full. Shutdown must never await that cycle: + // process teardown closes the task/socket if this best-effort enqueue + // cannot fit immediately. + if self + .terminal_control_tx + .try_send(TerminalControlCommand::Shutdown) + .is_err() + { + tracing::debug!("terminal control shutdown signal could not be queued immediately"); + } // Persist any still-pending deliveries so the next start can // redeliver them; only remove the file when nothing is pending. persist_pending_on_shutdown( diff --git a/crates/broker/src/runtime/fleet.rs b/crates/broker/src/runtime/fleet.rs index fdf55674f..bf74c6e99 100644 --- a/crates/broker/src/runtime/fleet.rs +++ b/crates/broker/src/runtime/fleet.rs @@ -6,10 +6,41 @@ use crate::{ RelaycastToBroker, FLEET_WIRE_VERSION, }, node_control::{delivery_ack, handler_unavailable_result, DeliveryDecision}, + terminal_control::{ + TerminalControlCommand, TerminalControlEvent, TerminalFromCloud, TerminalMode, + TerminalToCloud, + }, }; +use base64::{engine::general_purpose::STANDARD as BASE64, Engine as _}; const FLEET_AGENT_REGISTER_TIMEOUT: Duration = Duration::from_secs(30); const VERIFIED_SPAWN_READY_TIMEOUT: Duration = Duration::from_secs(90); +const TERMINAL_INPUT_MAX_BYTES: usize = 64 * 1024; +const TERMINAL_INPUT_MAX_BASE64_BYTES: usize = TERMINAL_INPUT_MAX_BYTES * 4 / 3 + 4; +const TERMINAL_SNAPSHOT_TIMEOUT: Duration = Duration::from_secs(10); +const TERMINAL_INPUT_ACK_TIMEOUT: Duration = Duration::from_secs(5); +// Terminal worker writes share the broker event loop with fleet control and +// worker lifecycle events. A wedged PTY must fail only its own attach instead +// of awaiting an unbounded pipe write in that loop. +const TERMINAL_WORKER_WRITE_TIMEOUT: Duration = Duration::from_millis(250); +const TERMINAL_INPUT_MAX_IN_FLIGHT_PER_SESSION: usize = 16; +// Relaycast currently limits a node to 32 terminal sessions. Keep that many +// slots free from high-volume frames so every affected session can still get a +// terminal.closed notification when the output lane applies backpressure. +const TERMINAL_CLOSE_RESERVE: usize = 32; + +pub(super) fn try_send_terminal( + terminal_control_tx: &mpsc::Sender, + message: TerminalToCloud, +) -> bool { + let is_close = matches!(&message, TerminalToCloud::Closed { .. }); + if !is_close && terminal_control_tx.capacity() <= TERMINAL_CLOSE_RESERVE { + return false; + } + terminal_control_tx + .try_send(TerminalControlCommand::Send(message)) + .is_ok() +} #[derive(Debug, Clone)] pub(super) struct PendingVerifiedSpawn { @@ -67,6 +98,324 @@ fn plan_fleet_delivery(decision: DeliveryDecision) -> FleetDeliveryPlan { } impl BrokerRuntime { + pub(super) async fn handle_terminal_control_event(&mut self, event: TerminalControlEvent) { + match event { + TerminalControlEvent::Connected => { + tracing::info!( + target = "relay_broker::terminal", + "fleet terminal transport connected" + ); + } + TerminalControlEvent::Disconnected => { + tracing::warn!( + target = "relay_broker::terminal", + sessions = self.terminal_sessions.len(), + "fleet terminal transport disconnected; clearing sessions for cloud resync" + ); + // Relaycast re-opens live terminal sessions when this dedicated + // lane reconnects. Drop the old connection's state now so input + // and snapshot replies cannot be routed into a server-side + // session that was invalidated with the old websocket. + self.terminal_sessions.clear(); + self.terminal_snapshot_requests.clear(); + self.terminal_input_requests.clear(); + } + TerminalControlEvent::Message(TerminalFromCloud::Open { + session_id, + agent, + mode, + }) => { + let agent_name = WorkerName::new(agent.clone()); + let runtime = self + .workers + .workers + .get(&agent_name) + .map(|handle| handle.spec.runtime.clone()); + match runtime { + None => self.send_terminal(TerminalToCloud::Error { + session_id, + code: "agent_not_found".to_string(), + message: format!("no worker named '{agent}'"), + }), + Some(AgentRuntime::Headless) => self.send_terminal(TerminalToCloud::Error { + session_id, + code: "unsupported_runtime".to_string(), + message: format!("worker '{agent}' is headless and has no PTY"), + }), + Some(AgentRuntime::Pty) => { + self.terminal_sessions.insert( + session_id.clone(), + TerminalSession { + agent: agent_name.clone(), + mode, + ready: false, + pending_output: Vec::new(), + pending_output_bytes: 0, + }, + ); + // Request the grid asynchronously. The response is routed + // from worker_events to the terminal lane, so this never + // stalls heartbeat/action processing behind a PTY snapshot. + let request_id = format!("terminal_snapshot_{}", Uuid::new_v4().simple()); + match tokio::time::timeout( + TERMINAL_WORKER_WRITE_TIMEOUT, + self.workers.send_to_worker( + agent_name.as_str(), + "snapshot_pty", + Some(RequestId::new(request_id.clone())), + json!({ "format": "ansi" }), + ), + ) + .await + { + Ok(Ok(())) => { + self.terminal_snapshot_requests.insert( + request_id, + TerminalSnapshotRequest { + session_id, + deadline: Instant::now() + TERMINAL_SNAPSHOT_TIMEOUT, + }, + ); + } + Ok(Err(error)) => { + self.fail_terminal_session( + session_id, + "snapshot_failed", + error.to_string(), + ); + } + Err(_) => { + self.fail_terminal_session( + session_id, + "snapshot_timeout", + "terminal snapshot write timed out".into(), + ); + } + } + } + } + } + TerminalControlEvent::Message(TerminalFromCloud::Input { + session_id, + data_base64, + }) => { + let Some(session) = self.terminal_sessions.get(&session_id).cloned() else { + self.send_terminal(TerminalToCloud::Error { + session_id, + code: "session_not_found".into(), + message: "terminal session is not active".into(), + }); + return; + }; + if session.mode == TerminalMode::View { + self.send_terminal(TerminalToCloud::Error { + session_id, + code: "read_only".into(), + message: "view sessions do not accept input".into(), + }); + return; + } + if !session.ready { + self.send_terminal(TerminalToCloud::Error { + session_id, + code: "session_not_ready".into(), + message: "terminal snapshot is not ready".into(), + }); + return; + } + if data_base64.len() > TERMINAL_INPUT_MAX_BASE64_BYTES { + self.send_terminal(TerminalToCloud::Error { + session_id, + code: "invalid_input".into(), + message: "terminal input must be bounded base64 UTF-8".into(), + }); + return; + } + let bytes = match BASE64.decode(data_base64.as_bytes()) { + Ok(bytes) if bytes.len() <= TERMINAL_INPUT_MAX_BYTES => bytes, + _ => { + self.send_terminal(TerminalToCloud::Error { + session_id, + code: "invalid_input".into(), + message: "terminal input must be bounded base64 UTF-8".into(), + }); + return; + } + }; + let data = match String::from_utf8(bytes) { + Ok(data) => data, + Err(_) => { + self.send_terminal(TerminalToCloud::Error { + session_id, + code: "invalid_input".into(), + message: "terminal input must be UTF-8".into(), + }); + return; + } + }; + if self + .terminal_input_requests + .values() + .filter(|pending| pending.session_id == session_id) + .count() + >= TERMINAL_INPUT_MAX_IN_FLIGHT_PER_SESSION + { + self.send_terminal(TerminalToCloud::Error { + session_id, + code: "input_backpressure".into(), + message: "terminal input acknowledgement backlog is full".into(), + }); + return; + } + let request_id = format!("terminal_input_{}", Uuid::new_v4().simple()); + match tokio::time::timeout( + TERMINAL_WORKER_WRITE_TIMEOUT, + self.workers.send_to_worker( + session.agent.as_str(), + "write_pty", + Some(RequestId::new(request_id.clone())), + json!({ "data": data }), + ), + ) + .await + { + Ok(Ok(())) => { + self.terminal_input_requests.insert( + request_id, + TerminalInputRequest { + session_id, + deadline: Instant::now() + TERMINAL_INPUT_ACK_TIMEOUT, + }, + ); + } + Ok(Err(error)) => { + self.fail_terminal_session(session_id, "input_failed", error.to_string()) + } + Err(_) => self.fail_terminal_session( + session_id, + "input_timeout", + "terminal input write timed out".into(), + ), + } + } + TerminalControlEvent::Message(TerminalFromCloud::Resize { + session_id, + rows, + cols, + }) => { + let Some(session) = self.terminal_sessions.get(&session_id).cloned() else { + self.send_terminal(TerminalToCloud::Error { + session_id, + code: "session_not_found".into(), + message: "terminal session is not active".into(), + }); + return; + }; + if session.mode == TerminalMode::View { + self.send_terminal(TerminalToCloud::Error { + session_id, + code: "read_only".into(), + message: "view sessions do not resize the PTY".into(), + }); + return; + } + if !session.ready { + self.send_terminal(TerminalToCloud::Error { + session_id, + code: "session_not_ready".into(), + message: "terminal snapshot is not ready".into(), + }); + return; + } + match tokio::time::timeout( + TERMINAL_WORKER_WRITE_TIMEOUT, + self.workers.send_to_worker( + session.agent.as_str(), + "resize_pty", + None, + json!({ "rows": rows, "cols": cols }), + ), + ) + .await + { + Ok(Ok(())) => {} + Ok(Err(error)) => { + self.fail_terminal_session(session_id, "resize_failed", error.to_string()) + } + Err(_) => self.fail_terminal_session( + session_id, + "resize_timeout", + "terminal resize write timed out".into(), + ), + } + } + TerminalControlEvent::Message(TerminalFromCloud::Close { session_id }) => { + self.terminal_sessions.remove(&session_id); + self.terminal_snapshot_requests + .retain(|_, pending| pending.session_id != session_id); + self.terminal_input_requests + .retain(|_, pending| pending.session_id != session_id); + self.send_terminal(TerminalToCloud::Closed { + session_id, + code: None, + message: None, + }); + } + } + } + + /// Attempts a bounded enqueue onto the dedicated terminal lane. On + /// saturation we tear down the affected session instead of accumulating + /// unbounded PTY output or delaying control traffic. + pub(super) fn send_terminal(&mut self, message: TerminalToCloud) { + let session_id = match &message { + TerminalToCloud::Ready { session_id, .. } + | TerminalToCloud::Output { session_id, .. } + | TerminalToCloud::InputAck { session_id, .. } + | TerminalToCloud::Error { session_id, .. } + | TerminalToCloud::Closed { session_id, .. } => session_id.clone(), + }; + let is_close = matches!(&message, TerminalToCloud::Closed { .. }); + if !try_send_terminal(&self.terminal_control_tx, message) { + tracing::warn!(target = "relay_broker::terminal", session_id = %session_id, "terminal queue full or closed; ending session"); + self.terminal_sessions.remove(&session_id); + self.terminal_snapshot_requests + .retain(|_, pending| pending.session_id != session_id); + self.terminal_input_requests + .retain(|_, pending| pending.session_id != session_id); + if !is_close + && !try_send_terminal( + &self.terminal_control_tx, + TerminalToCloud::Closed { + session_id: session_id.clone(), + code: Some("output_backpressure".into()), + message: Some("terminal output queue is full".into()), + }, + ) + { + tracing::warn!(target = "relay_broker::terminal", session_id = %session_id, "terminal close could not be queued after backpressure"); + } + } + } + + fn fail_terminal_session(&mut self, session_id: String, code: &str, message: String) { + self.terminal_sessions.remove(&session_id); + self.terminal_snapshot_requests + .retain(|_, pending| pending.session_id != session_id); + self.terminal_input_requests + .retain(|_, pending| pending.session_id != session_id); + self.send_terminal(TerminalToCloud::Error { + session_id: session_id.clone(), + code: code.into(), + message: message.clone(), + }); + self.send_terminal(TerminalToCloud::Closed { + session_id, + code: Some(code.into()), + message: Some(message), + }); + } + pub(super) async fn handle_fleet_control_event(&mut self, event: FleetControlEvent) { match event { FleetControlEvent::Connected => { @@ -1312,6 +1661,46 @@ mod tests { ); } + #[test] + fn terminal_close_reserve_survives_output_backpressure() { + let (tx, mut rx) = mpsc::channel(TERMINAL_CLOSE_RESERVE + 1); + assert!(try_send_terminal( + &tx, + TerminalToCloud::Output { + session_id: "session-a".into(), + chunk: "x".into(), + offset: None, + }, + )); + assert!( + !try_send_terminal( + &tx, + TerminalToCloud::Output { + session_id: "session-a".into(), + chunk: "y".into(), + offset: None, + }, + ), + "non-final terminal traffic must preserve close capacity" + ); + assert!(try_send_terminal( + &tx, + TerminalToCloud::Closed { + session_id: "session-a".into(), + code: Some("output_backpressure".into()), + message: Some("queue full".into()), + }, + )); + assert!(matches!( + rx.try_recv(), + Ok(TerminalControlCommand::Send(TerminalToCloud::Output { .. })) + )); + assert!(matches!( + rx.try_recv(), + Ok(TerminalControlCommand::Send(TerminalToCloud::Closed { .. })) + )); + } + #[test] fn classify_fleet_delivery_injects_message_classes_and_acks_receipts() { // Mirrors relaycast parse_inbound_kind message-class alias set: any of diff --git a/crates/broker/src/runtime/init.rs b/crates/broker/src/runtime/init.rs index ea7ffef6f..dd89cfc27 100644 --- a/crates/broker/src/runtime/init.rs +++ b/crates/broker/src/runtime/init.rs @@ -246,6 +246,12 @@ pub(crate) async fn run_init(cmd: InitCommand, telemetry: TelemetryClient) -> Re (!resolved_node_name.trim().is_empty()).then_some(resolved_node_name.as_str()), ); let fleet_ws_url = relaycast::node_control_ws_url(configured_base.as_deref()); + // Keep terminal traffic on a physically separate websocket. Do not append + // terminal frames to the heartbeat/action control endpoint. + let terminal_ws_url = fleet_ws_url + .strip_suffix("/v1/node/ws") + .map(|base| format!("{base}/v1/node/terminal/ws")) + .context("fleet control URL must end with /v1/node/ws to derive the terminal URL")?; let broker_version = format!("relay-broker/{}", crate::util::version::broker_version()); // The broker enrolls as a relaycast node and delivers/injects solely over // /v1/node/ws. A node token is required to open that connection: prefer an @@ -303,6 +309,16 @@ pub(crate) async fn run_init(cmd: InitCommand, telemetry: TelemetryClient) -> Re }); let (fleet_control_tx, fleet_control_rx) = mpsc::channel::(256); let (fleet_event_tx, fleet_event_rx) = mpsc::channel::(256); + // The terminal queue is deliberately bounded. A wedged remote attach must + // fail its session rather than accumulating unbounded PTY output in the + // broker or starving node control. + // Terminal bytes have a burstier profile than control actions. Keep this + // lane bounded, but give a short PTY burst room without impacting the + // independent control queue. + let (terminal_control_tx, terminal_control_rx) = + mpsc::channel::(1024); + let (terminal_event_tx, terminal_event_rx) = + mpsc::channel::(1024); let node_delivery_token_present = node_token.is_some(); tokio::spawn(crate::node_control::run_node_control_client( crate::node_control::FleetControlConfig { @@ -317,6 +333,14 @@ pub(crate) async fn run_init(cmd: InitCommand, telemetry: TelemetryClient) -> Re fleet_control_rx, fleet_event_tx, )); + tokio::spawn(crate::terminal_control::run_terminal_control_client( + crate::terminal_control::TerminalControlConfig { + ws_url: terminal_ws_url, + session_token: session_node_token.clone(), + }, + terminal_control_rx, + terminal_event_tx, + )); // Register this node unconditionally on connect (no sidecar required). This // is the only command that flips the control client out of its idle state // and into the connect loop, so the broker enrolls every startup. @@ -659,6 +683,12 @@ pub(crate) async fn run_init(cmd: InitCommand, telemetry: TelemetryClient) -> Re node_delivery_connected: false, fleet_event_rx, fleet_control_open: true, + terminal_control_tx, + terminal_event_rx, + terminal_control_open: true, + terminal_sessions: HashMap::new(), + terminal_snapshot_requests: HashMap::new(), + terminal_input_requests: HashMap::new(), fleet_delivery_book: FleetDeliveryBook::default(), // Seed the live capacity with the configured max so heartbeats/load // updates keep reporting it (they overwrite load.max_agents from this diff --git a/crates/broker/src/runtime/maintenance.rs b/crates/broker/src/runtime/maintenance.rs index c99689baf..208042fa5 100644 --- a/crates/broker/src/runtime/maintenance.rs +++ b/crates/broker/src/runtime/maintenance.rs @@ -1,4 +1,6 @@ +use super::fleet::try_send_terminal; use super::*; +use crate::terminal_control::TerminalToCloud; impl BrokerRuntime { pub(super) async fn handle_maintenance_tick(&mut self) { @@ -26,12 +28,85 @@ impl BrokerRuntime { let delivery_states = &mut self.delivery_states; let resize_owners = &mut self.resize_owners; let agent_result_tokens = &mut self.agent_result_tokens; + let terminal_control_tx = &self.terminal_control_tx; + let terminal_sessions = &mut self.terminal_sessions; + let terminal_snapshot_requests = &mut self.terminal_snapshot_requests; + let terminal_input_requests = &mut self.terminal_input_requests; let delivery_retry_interval = self.delivery_retry_interval; let shutdown = &self.shutdown; let default_workspace = &self.default_workspace; let now = Instant::now(); + // A worker can disappear before answering `snapshot_pty`. Bound these + // terminal-only RPCs so their sessions cannot remain live forever. + let expired_terminal_snapshots: Vec<(String, String)> = terminal_snapshot_requests + .iter() + .filter(|(_, pending)| pending.deadline <= now) + .map(|(request_id, pending)| (request_id.clone(), pending.session_id.clone())) + .collect(); + for (request_id, session_id) in expired_terminal_snapshots { + terminal_snapshot_requests.remove(&request_id); + if terminal_sessions.remove(&session_id).is_some() { + terminal_input_requests.retain(|_, pending| pending.session_id != session_id); + if !try_send_terminal( + terminal_control_tx, + TerminalToCloud::Error { + session_id: session_id.clone(), + code: "snapshot_timeout".into(), + message: "terminal snapshot timed out".into(), + }, + ) { + tracing::warn!(target = "relay_broker::terminal", session_id = %session_id, "terminal queue full or closed while reporting snapshot timeout"); + } + if !try_send_terminal( + terminal_control_tx, + TerminalToCloud::Closed { + session_id: session_id.clone(), + code: Some("snapshot_timeout".into()), + message: Some("terminal snapshot timed out".into()), + }, + ) { + tracing::warn!(target = "relay_broker::terminal", session_id = %session_id, "terminal close could not be queued after snapshot timeout"); + } + } + } + let expired_terminal_inputs: Vec<(String, String)> = terminal_input_requests + .iter() + .filter(|(_, pending)| pending.deadline <= now) + .map(|(request_id, pending)| (request_id.clone(), pending.session_id.clone())) + .collect(); + for (request_id, session_id) in expired_terminal_inputs { + terminal_input_requests.remove(&request_id); + // A timed-out write may still reach the PTY after cancellation. End + // the whole session before reporting it so a retry cannot duplicate + // user keystrokes against that uncertain write. + if terminal_sessions.remove(&session_id).is_some() { + terminal_snapshot_requests.retain(|_, pending| pending.session_id != session_id); + terminal_input_requests.retain(|_, pending| pending.session_id != session_id); + if !try_send_terminal( + terminal_control_tx, + TerminalToCloud::Error { + session_id: session_id.clone(), + code: "input_timeout".into(), + message: "terminal input acknowledgement timed out".into(), + }, + ) { + tracing::warn!(target = "relay_broker::terminal", session_id = %session_id, "terminal queue full or closed while reporting input timeout"); + } + if !try_send_terminal( + terminal_control_tx, + TerminalToCloud::Closed { + session_id: session_id.clone(), + code: Some("input_timeout".into()), + message: Some("terminal input acknowledgement timed out".into()), + }, + ) { + tracing::warn!(target = "relay_broker::terminal", session_id = %session_id, "terminal close could not be queued after input timeout"); + } + } + } + // Time out worker request/response calls whose worker never // responded. Common cause: worker crashed between us sending // the request frame and it parsing the frame. Without this @@ -210,6 +285,26 @@ impl BrokerRuntime { ); } pty_observability.remove(name); + let terminal_session_ids: Vec = terminal_sessions + .iter() + .filter(|(_, session)| session.agent == *name) + .map(|(session_id, _)| session_id.clone()) + .collect(); + for session_id in terminal_session_ids { + terminal_sessions.remove(&session_id); + terminal_snapshot_requests.retain(|_, pending| pending.session_id != session_id); + terminal_input_requests.retain(|_, pending| pending.session_id != session_id); + if !try_send_terminal( + terminal_control_tx, + TerminalToCloud::Closed { + session_id: session_id.clone(), + code: Some("agent_exited".into()), + message: Some("terminal worker exited".into()), + }, + ) { + tracing::warn!(target = "relay_broker::terminal", session_id = %session_id, "terminal queue full or closed while closing exited worker session"); + } + } // Record crash in insights let (category, description) = crate::crash_insights::CrashInsights::analyze(*code, signal.as_deref()); diff --git a/crates/broker/src/runtime/mod.rs b/crates/broker/src/runtime/mod.rs index 164f15de7..8c661add0 100644 --- a/crates/broker/src/runtime/mod.rs +++ b/crates/broker/src/runtime/mod.rs @@ -43,6 +43,7 @@ use crate::{ }, replay_buffer::{ReplayBuffer, DEFAULT_REPLAY_CAPACITY}, telemetry::{ActionSource, TelemetryClient, TelemetryEvent}, + terminal_control::{TerminalControlCommand, TerminalControlEvent, TerminalMode}, types::{ AgentResultMcpConfig, InboundDeliveryDispatch, InboundDeliveryMode, InboundDeliveryState, PendingRelayMessage, RelaycastDeliveryReceipt, diff --git a/crates/broker/src/runtime/worker_events.rs b/crates/broker/src/runtime/worker_events.rs index b8ed5ed26..91dbf7d29 100644 --- a/crates/broker/src/runtime/worker_events.rs +++ b/crates/broker/src/runtime/worker_events.rs @@ -1,7 +1,105 @@ -use super::fleet::{refresh_fleet_inventory_session_ref, verified_spawn_ready_result}; +use super::fleet::{ + refresh_fleet_inventory_session_ref, try_send_terminal, verified_spawn_ready_result, +}; use super::*; +use crate::terminal_control::{TerminalControlCommand, TerminalToCloud}; use crate::worker::AgentWorkState; +const TERMINAL_PENDING_OUTPUT_MAX_BYTES: usize = 1024 * 1024; + +fn publish_terminal_output( + terminal_control_tx: &mpsc::Sender, + terminal_sessions: &mut HashMap, + terminal_snapshot_requests: &mut HashMap, + terminal_input_requests: &mut HashMap, + name: &WorkerName, + chunk: String, + offset: Option, +) { + if terminal_sessions.is_empty() { + return; + } + let chunk_bytes = chunk.len(); + let mut session_ids = Vec::new(); + let mut saturated_sessions = Vec::new(); + for (session_id, session) in terminal_sessions.iter_mut() { + if session.agent != *name { + continue; + } + if !session.ready { + if session.pending_output_bytes + chunk_bytes > TERMINAL_PENDING_OUTPUT_MAX_BYTES { + saturated_sessions.push(session_id.clone()); + } else { + session.pending_output.push((chunk.clone(), offset)); + session.pending_output_bytes += chunk_bytes; + } + continue; + } + session_ids.push(session_id.clone()); + } + for session_id in saturated_sessions { + tracing::warn!(target = "relay_broker::terminal", session_id = %session_id, "terminal snapshot wait exceeded bounded output buffer; ending session"); + end_terminal_session( + terminal_control_tx, + terminal_sessions, + terminal_snapshot_requests, + terminal_input_requests, + &session_id, + "output_backpressure", + "terminal snapshot wait exceeded bounded output buffer", + ); + } + for session_id in session_ids { + if !try_send_terminal( + terminal_control_tx, + TerminalToCloud::Output { + session_id: session_id.clone(), + chunk: chunk.clone(), + offset, + }, + ) { + tracing::warn!(target = "relay_broker::terminal", session_id = %session_id, "terminal output queue full or closed; ending session"); + end_terminal_session( + terminal_control_tx, + terminal_sessions, + terminal_snapshot_requests, + terminal_input_requests, + &session_id, + "output_backpressure", + "terminal output queue is full", + ); + } + } +} + +fn end_terminal_session( + terminal_control_tx: &mpsc::Sender, + terminal_sessions: &mut HashMap, + terminal_snapshot_requests: &mut HashMap, + terminal_input_requests: &mut HashMap, + session_id: &str, + code: &str, + message: &str, +) { + terminal_sessions.remove(session_id); + terminal_snapshot_requests.retain(|_, pending| pending.session_id != session_id); + terminal_input_requests.retain(|_, pending| pending.session_id != session_id); + if !try_send_terminal( + terminal_control_tx, + TerminalToCloud::Closed { + session_id: session_id.into(), + code: Some(code.into()), + message: Some(message.into()), + }, + ) { + tracing::warn!( + target = "relay_broker::terminal", + session_id, + "terminal close could not be queued after session failure" + ); + } +} + fn enqueue_pty_event( states: &mut HashMap, tx: &mpsc::Sender, @@ -428,6 +526,10 @@ impl BrokerRuntime { let fleet_control_tx = &self.fleet_control_tx; let fleet_inventory = &mut self.fleet_inventory; let delivery_states = &self.delivery_states; + let terminal_control_tx = &self.terminal_control_tx; + let terminal_sessions = &mut self.terminal_sessions; + let terminal_snapshot_requests = &mut self.terminal_snapshot_requests; + let terminal_input_requests = &mut self.terminal_input_requests; match worker_event { WorkerEvent::Message { @@ -797,27 +899,195 @@ impl BrokerRuntime { } let _ = send_event(sdk_out_tx, agent_event).await; } else if msg_type.ends_with("_response") { - // Generic worker request/response dispatch. - // Any frame whose `type` ends in - // `_response` is routed by `request_id` - // into the matching parked `oneshot` in - // `pending_requests`. The pending entry - // owns the format/error decoding logic - // via `worker_request::fulfil_response_frame`. - let routed = - worker_request::fulfil_response_frame(pending_requests, &value); - if !routed { - let req_id = value - .get("request_id") - .and_then(Value::as_str) - .unwrap_or(""); - tracing::debug!( - target = "agent_relay::broker", - worker = %name, - msg_type = %msg_type, - request_id = %req_id, - "worker response with no pending caller — dropping" - ); + let terminal_input_session_id = value + .get("request_id") + .and_then(Value::as_str) + .and_then(|request_id| terminal_input_requests.remove(request_id)) + .map(|request| request.session_id); + let terminal_session_id = value + .get("request_id") + .and_then(Value::as_str) + .and_then(|request_id| terminal_snapshot_requests.remove(request_id)) + .map(|request| request.session_id); + if let Some(session_id) = terminal_input_session_id { + if !terminal_sessions.contains_key(&session_id) { + return; + } + let payload = value.get("payload").cloned().unwrap_or(Value::Null); + let message = if msg_type != "write_pty_response" { + TerminalToCloud::Error { + session_id: session_id.clone(), + code: "input_failed".into(), + message: + "terminal input returned an unexpected worker response" + .into(), + } + } else if let Some(error) = payload.get("error") { + TerminalToCloud::Error { + session_id: session_id.clone(), + code: error + .get("code") + .and_then(Value::as_str) + .unwrap_or("input_failed") + .to_string(), + message: error + .get("message") + .and_then(Value::as_str) + .unwrap_or("terminal input failed") + .to_string(), + } + } else if let Some(bytes_written) = + payload.get("bytes_written").and_then(Value::as_u64) + { + TerminalToCloud::InputAck { + session_id: session_id.clone(), + bytes_written: usize::try_from(bytes_written) + .unwrap_or(usize::MAX), + } + } else { + TerminalToCloud::Error { + session_id: session_id.clone(), + code: "input_failed".into(), + message: "terminal input response was malformed".into(), + } + }; + if !try_send_terminal(terminal_control_tx, message) { + tracing::warn!(target = "relay_broker::terminal", session_id = %session_id, "terminal queue full or closed while reporting PTY input result; ending session"); + end_terminal_session( + terminal_control_tx, + terminal_sessions, + terminal_snapshot_requests, + terminal_input_requests, + &session_id, + "output_backpressure", + "terminal output queue is full", + ); + } + } else if let Some(session_id) = terminal_session_id { + if !terminal_sessions.contains_key(&session_id) { + return; + } + let payload = value.get("payload").cloned().unwrap_or(Value::Null); + let message = if msg_type != "snapshot_response" { + TerminalToCloud::Error { + session_id: session_id.clone(), + code: "snapshot_failed".into(), + message: + "terminal snapshot returned an unexpected worker response" + .into(), + } + } else if let Some(error) = payload.get("error") { + TerminalToCloud::Error { + session_id: session_id.clone(), + code: error + .get("code") + .and_then(Value::as_str) + .unwrap_or("snapshot_failed") + .to_string(), + message: error + .get("message") + .and_then(Value::as_str) + .unwrap_or("terminal snapshot failed") + .to_string(), + } + } else if let (Some(screen), Some(rows), Some(cols)) = ( + payload.get("screen").and_then(Value::as_str), + payload + .get("rows") + .and_then(Value::as_u64) + .and_then(|value| u16::try_from(value).ok()), + payload + .get("cols") + .and_then(Value::as_u64) + .and_then(|value| u16::try_from(value).ok()), + ) { + TerminalToCloud::Ready { + session_id: session_id.clone(), + screen: screen.to_string(), + rows, + cols, + offset: payload + .get("offset") + .and_then(Value::as_u64) + .unwrap_or(0), + } + } else { + TerminalToCloud::Error { + session_id: session_id.clone(), + code: "snapshot_failed".into(), + message: "terminal snapshot response was malformed".into(), + } + }; + let snapshot_ready = matches!(&message, TerminalToCloud::Ready { .. }); + let snapshot_failed = matches!(&message, TerminalToCloud::Error { .. }); + if !try_send_terminal(terminal_control_tx, message) { + tracing::warn!(target = "relay_broker::terminal", session_id = %session_id, "terminal queue full or closed while sending snapshot; ending session"); + end_terminal_session( + terminal_control_tx, + terminal_sessions, + terminal_snapshot_requests, + terminal_input_requests, + &session_id, + "output_backpressure", + "terminal output queue is full", + ); + } else if snapshot_ready { + let pending_output = + if let Some(session) = terminal_sessions.get_mut(&session_id) { + session.ready = true; + session.pending_output_bytes = 0; + std::mem::take(&mut session.pending_output) + } else { + Vec::new() + }; + for (chunk, offset) in pending_output { + if !try_send_terminal( + terminal_control_tx, + TerminalToCloud::Output { + session_id: session_id.clone(), + chunk, + offset, + }, + ) { + tracing::warn!(target = "relay_broker::terminal", session_id = %session_id, "terminal queue full or closed while flushing buffered output; ending session"); + end_terminal_session( + terminal_control_tx, + terminal_sessions, + terminal_snapshot_requests, + terminal_input_requests, + &session_id, + "output_backpressure", + "terminal output queue is full", + ); + break; + } + } + } else if snapshot_failed { + terminal_sessions.remove(&session_id); + } + } else { + // Generic worker request/response dispatch. + // Any frame whose `type` ends in + // `_response` is routed by `request_id` + // into the matching parked `oneshot` in + // `pending_requests`. The pending entry + // owns the format/error decoding logic + // via `worker_request::fulfil_response_frame`. + let routed = + worker_request::fulfil_response_frame(pending_requests, &value); + if !routed { + let req_id = value + .get("request_id") + .and_then(Value::as_str) + .unwrap_or(""); + tracing::debug!( + target = "agent_relay::broker", + worker = %name, + msg_type = %msg_type, + request_id = %req_id, + "worker response with no pending caller — dropping" + ); + } } } else if msg_type == "worker_stream" { let is_pty = workers @@ -831,22 +1101,40 @@ impl BrokerRuntime { if is_pty { publish_pty_busy(pty_observability, hosted_agent_event_tx, &name); } + let payload = value.get("payload"); + let chunk = payload + .and_then(|payload| payload.get("chunk")) + .and_then(Value::as_str) + .unwrap_or_default() + .to_string(); + let offset = payload + .and_then(|payload| payload.get("offset")) + .and_then(Value::as_u64); let mut stream_event = json!({ "kind": "worker_stream", "name": name, - "stream": value.get("payload").and_then(|p| p.get("stream")).cloned().unwrap_or(Value::String("stdout".to_string())), - "chunk": value.get("payload").and_then(|p| p.get("chunk")).cloned().unwrap_or(Value::String(String::new())), + "stream": payload.and_then(|p| p.get("stream")).cloned().unwrap_or(Value::String("stdout".to_string())), + "chunk": chunk.clone(), }); // Forward the per-worker stream offset when present so // attaching clients can correlate the live stream with // a snapshot. Absent for headless workers (no grid). - if let Some(offset) = - value.get("payload").and_then(|p| p.get("offset")).cloned() - { + if let Some(offset) = offset { if let Some(obj) = stream_event.as_object_mut() { - obj.insert("offset".to_string(), offset); + obj.insert("offset".to_string(), Value::from(offset)); } } + if !chunk.is_empty() { + publish_terminal_output( + terminal_control_tx, + terminal_sessions, + terminal_snapshot_requests, + terminal_input_requests, + &name, + chunk, + offset, + ); + } let _ = send_event(sdk_out_tx, stream_event).await; } else if msg_type == "harness_started" { // A running child process proves liveness but not that diff --git a/crates/broker/src/terminal_control.rs b/crates/broker/src/terminal_control.rs new file mode 100644 index 000000000..f302bd608 --- /dev/null +++ b/crates/broker/src/terminal_control.rs @@ -0,0 +1,328 @@ +//! Dedicated Relaycast terminal transport. +//! +//! This deliberately does not share `node_control`: terminal output can be +//! continuous and subject to backpressure, whereas node registration, +//! heartbeats, and action delivery need a small independent control lane. + +use std::{ + sync::{Arc, RwLock}, + time::Duration, +}; + +use futures_util::{SinkExt, StreamExt}; +use relaycast::ORIGIN_ACTOR_HEADER; +use serde::{Deserialize, Serialize}; +use tokio::sync::mpsc; +use tokio_tungstenite::{ + connect_async, + tungstenite::{client::IntoClientRequest, Message}, +}; + +const INITIAL_RECONNECT_DELAY: Duration = Duration::from_secs(1); +const MAX_RECONNECT_DELAY: Duration = Duration::from_secs(30); +const TOKEN_WAIT_DELAY: Duration = Duration::from_secs(1); + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub(crate) enum TerminalMode { + View, + Drive, + Passthrough, +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(tag = "type")] +pub(crate) enum TerminalFromCloud { + #[serde(rename = "terminal.open")] + Open { + session_id: String, + agent: String, + mode: TerminalMode, + }, + #[serde(rename = "terminal.input")] + Input { + session_id: String, + data_base64: String, + }, + #[serde(rename = "terminal.resize")] + Resize { + session_id: String, + rows: u16, + cols: u16, + }, + #[serde(rename = "terminal.close")] + Close { session_id: String }, +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(tag = "type")] +pub(crate) enum TerminalToCloud { + #[serde(rename = "terminal.ready")] + Ready { + session_id: String, + screen: String, + rows: u16, + cols: u16, + offset: u64, + }, + #[serde(rename = "terminal.output")] + Output { + session_id: String, + chunk: String, + #[serde(skip_serializing_if = "Option::is_none")] + offset: Option, + }, + #[serde(rename = "terminal.input_ack")] + InputAck { + session_id: String, + bytes_written: usize, + }, + #[serde(rename = "terminal.error")] + Error { + session_id: String, + code: String, + message: String, + }, + #[serde(rename = "terminal.closed")] + Closed { + session_id: String, + #[serde(skip_serializing_if = "Option::is_none")] + code: Option, + #[serde(skip_serializing_if = "Option::is_none")] + message: Option, + }, +} + +#[derive(Debug)] +pub(crate) enum TerminalControlCommand { + Send(TerminalToCloud), + Shutdown, +} + +#[derive(Debug, Clone, PartialEq)] +pub(crate) enum TerminalControlEvent { + Connected, + Disconnected, + Message(TerminalFromCloud), +} + +#[derive(Clone)] +pub(crate) struct TerminalControlConfig { + pub(crate) ws_url: String, + /// Written through by node-control when it mints or rotates the node + /// credential. This transport never mints itself, avoiding duplicate + /// credential flows while still reconnecting with a fresh token. + pub(crate) session_token: Arc>>, +} + +pub(crate) async fn run_terminal_control_client( + config: TerminalControlConfig, + mut command_rx: mpsc::Receiver, + event_tx: mpsc::Sender, +) { + let mut reconnect_delay = INITIAL_RECONNECT_DELAY; + loop { + let token = config + .session_token + .read() + .ok() + .and_then(|token| token.clone()); + let Some(token) = token.filter(|token| !token.trim().is_empty()) else { + tokio::select! { + command = command_rx.recv() => { + if matches!(command, Some(TerminalControlCommand::Shutdown) | None) { return; } + // Preserve bounded backpressure by dropping commands only + // when the caller itself chose a non-blocking try_send. + } + _ = tokio::time::sleep(TOKEN_WAIT_DELAY) => {} + } + continue; + }; + + let mut request = match config.ws_url.as_str().into_client_request() { + Ok(request) => request, + Err(error) => { + tracing::warn!(target = "relay_broker::terminal", error = %error, "invalid fleet terminal ws url"); + tokio::time::sleep(reconnect_delay).await; + reconnect_delay = (reconnect_delay * 2).min(MAX_RECONNECT_DELAY); + continue; + } + }; + let header = format!("Bearer {}", token.trim()); + let Ok(header) = header.parse() else { + tracing::warn!( + target = "relay_broker::terminal", + "invalid fleet terminal token header" + ); + reconnect_delay = (reconnect_delay * 2).min(MAX_RECONNECT_DELAY); + tokio::time::sleep(reconnect_delay).await; + continue; + }; + request.headers_mut().insert("authorization", header); + if let Ok(value) = crate::telemetry::BROKER_ORIGIN_ACTOR.parse() { + request.headers_mut().insert(ORIGIN_ACTOR_HEADER, value); + } + for (name, value) in crate::telemetry::cloud_identity_headers() { + let Ok(header_name) = name.parse::() else { + continue; + }; + if let Ok(header_value) = value.parse() { + request.headers_mut().insert(header_name, header_value); + } + } + + let (socket, _) = match connect_async(request).await { + Ok(socket) => socket, + Err(error) => { + tracing::warn!(target = "relay_broker::terminal", url = %config.ws_url, error = %error, "fleet terminal ws connect failed"); + tokio::time::sleep(reconnect_delay).await; + reconnect_delay = (reconnect_delay * 2).min(MAX_RECONNECT_DELAY); + continue; + } + }; + reconnect_delay = INITIAL_RECONNECT_DELAY; + let _ = event_tx.send(TerminalControlEvent::Connected).await; + let (mut sink, mut stream) = socket.split(); + let mut connected = true; + while connected { + tokio::select! { + command = command_rx.recv() => match command { + Some(TerminalControlCommand::Send(message)) => { + match serde_json::to_string(&message) { + Ok(encoded) => { + if sink.send(Message::Text(encoded)).await.is_err() { + connected = false; + } + } + Err(_) => connected = false, + } + } + Some(TerminalControlCommand::Shutdown) | None => { + let _ = sink.send(Message::Close(None)).await; + return; + } + }, + inbound = stream.next() => match inbound { + Some(Ok(Message::Text(text))) => match serde_json::from_str::(&text) { + Ok(message) => { if event_tx.send(TerminalControlEvent::Message(message)).await.is_err() { return; } } + Err(error) => tracing::warn!(target = "relay_broker::terminal", error = %error, "invalid fleet terminal frame"), + }, + Some(Ok(Message::Ping(_))) => {} + Some(Ok(Message::Close(_))) | Some(Err(_)) | None => connected = false, + Some(Ok(_)) => {}, + }, + } + } + let _ = event_tx.send(TerminalControlEvent::Disconnected).await; + tokio::time::sleep(reconnect_delay).await; + reconnect_delay = (reconnect_delay * 2).min(MAX_RECONNECT_DELAY); + } +} + +#[cfg(test)] +mod tests { + use super::{TerminalFromCloud, TerminalMode, TerminalToCloud}; + + #[test] + fn terminal_wire_round_trips_without_control_frames() { + let open: TerminalFromCloud = serde_json::from_str( + r#"{"type":"terminal.open","session_id":"s","agent":"Ada","mode":"view"}"#, + ) + .unwrap(); + assert_eq!( + open, + TerminalFromCloud::Open { + session_id: "s".into(), + agent: "Ada".into(), + mode: TerminalMode::View + } + ); + let output = serde_json::to_value(TerminalToCloud::Output { + session_id: "s".into(), + chunk: "x".into(), + offset: Some(2), + }) + .unwrap(); + assert_eq!(output["type"], "terminal.output"); + assert_eq!(output["offset"], 2); + + let output_without_offset = serde_json::to_value(TerminalToCloud::Output { + session_id: "s".into(), + chunk: "x".into(), + offset: None, + }) + .unwrap(); + assert!(output_without_offset.get("offset").is_none()); + + for (wire, expected) in [ + ( + r#"{"type":"terminal.input","session_id":"s","data_base64":"eA=="}"#, + TerminalFromCloud::Input { + session_id: "s".into(), + data_base64: "eA==".into(), + }, + ), + ( + r#"{"type":"terminal.resize","session_id":"s","rows":24,"cols":80}"#, + TerminalFromCloud::Resize { + session_id: "s".into(), + rows: 24, + cols: 80, + }, + ), + ( + r#"{"type":"terminal.close","session_id":"s"}"#, + TerminalFromCloud::Close { + session_id: "s".into(), + }, + ), + ] { + assert_eq!( + serde_json::from_str::(wire).unwrap(), + expected + ); + } + + let ready = serde_json::to_value(TerminalToCloud::Ready { + session_id: "s".into(), + screen: "screen".into(), + rows: 24, + cols: 80, + offset: 3, + }) + .unwrap(); + assert_eq!(ready["type"], "terminal.ready"); + assert_eq!(ready["offset"], 3); + let ack = serde_json::to_value(TerminalToCloud::InputAck { + session_id: "s".into(), + bytes_written: 1, + }) + .unwrap(); + assert_eq!(ack["type"], "terminal.input_ack"); + let error = serde_json::to_value(TerminalToCloud::Error { + session_id: "s".into(), + code: "bad".into(), + message: "nope".into(), + }) + .unwrap(); + assert_eq!(error["type"], "terminal.error"); + assert_eq!(error["code"], "bad"); + let closed = serde_json::to_value(TerminalToCloud::Closed { + session_id: "s".into(), + code: None, + message: None, + }) + .unwrap(); + assert_eq!(closed["type"], "terminal.closed"); + assert!(closed.get("code").is_none()); + assert!(closed.get("message").is_none()); + let closed_with_error = serde_json::to_value(TerminalToCloud::Closed { + session_id: "s".into(), + code: Some("closed".into()), + message: Some("done".into()), + }) + .unwrap(); + assert_eq!(closed_with_error["code"], "closed"); + assert_eq!(closed_with_error["message"], "done"); + } +} diff --git a/packages/cli/src/cli/commands/local-agent.test.ts b/packages/cli/src/cli/commands/local-agent.test.ts index 6d4c2abdc..388ad3d6d 100644 --- a/packages/cli/src/cli/commands/local-agent.test.ts +++ b/packages/cli/src/cli/commands/local-agent.test.ts @@ -27,6 +27,7 @@ function harness(overrides: Partial = {}) { }; const attach = vi.fn(async () => 0); const attachRemote = vi.fn(async () => 0); + const attachNode = vi.fn(async () => 0); const log = vi.fn(); const error = vi.fn(); const exit = vi.fn(); @@ -34,6 +35,7 @@ function harness(overrides: Partial = {}) { connect: vi.fn(async () => client as never), attach, attachRemote, + attachNode, cwd: () => '/tmp/project', log, error, @@ -44,7 +46,7 @@ function harness(overrides: Partial = {}) { program.exitOverride(); const group = program.command('local'); registerLocalAgentCommands(group, deps); - return { program, client, attach, attachRemote, log, error, exit }; + return { program, client, attach, attachRemote, attachNode, log, error, exit }; } describe('local agent subtree', () => { @@ -94,6 +96,58 @@ describe('local agent subtree', () => { expect(attachRemote).toHaveBeenCalledWith('lead', 'view', '', expect.objectContaining({})); }); + it('attach --node opens the canonical authenticated terminal path without invoking SSH', async () => { + const { program, attach, attachRemote, attachNode } = harness(); + await program.parseAsync( + ['local', 'agent', 'attach', 'lead', '--node', 'daytona-live', '--mode', 'passthrough', '--json'], + { from: 'user' } + ); + expect(attach).not.toHaveBeenCalled(); + expect(attachRemote).not.toHaveBeenCalled(); + expect(attachNode).toHaveBeenCalledWith( + 'lead', + 'passthrough', + 'daytona-live', + expect.objectContaining({ json: true }) + ); + }); + + it('attach --node rejects the SSH fallback conflict', async () => { + const { program, attachNode, attachRemote, error, exit } = harness(); + await program.parseAsync(['local', 'agent', 'attach', 'lead', '--node', 'finn', '--ssh-host', 'finn'], { + from: 'user', + }); + expect(attachNode).not.toHaveBeenCalled(); + expect(attachRemote).not.toHaveBeenCalled(); + expect(error).toHaveBeenCalledWith(expect.stringContaining('--node cannot be combined with --ssh-host')); + expect(exit).toHaveBeenCalledWith(1); + }); + + it.each([ + ['--broker-url', 'http://127.0.0.1:7777'], + ['--api-key', 'do-not-forward'], + ['--state-dir', '/tmp/relay-state'], + ])('attach --node rejects local broker option %s', async (flag, value) => { + const { program, attachNode, error, exit } = harness(); + await program.parseAsync(['local', 'agent', 'attach', 'lead', '--node', 'finn', flag, value], { + from: 'user', + }); + expect(attachNode).not.toHaveBeenCalled(); + expect(error).toHaveBeenCalledWith(expect.stringContaining('--node cannot be combined')); + expect(exit).toHaveBeenCalledWith(1); + }); + + it('attach --node prefixes a terminal setup error once', async () => { + const attachNode = vi.fn(async () => { + throw new Error('terminal unavailable'); + }); + const { program, error, exit } = harness({ attachNode }); + await program.parseAsync(['local', 'agent', 'attach', 'lead', '--node', 'finn'], { from: 'user' }); + expect(error).toHaveBeenCalledWith('Error: terminal unavailable'); + expect(error).toHaveBeenCalledTimes(1); + expect(exit).toHaveBeenCalledWith(1); + }); + it.each([ ['--api-key', 'do-not-forward'], ['--api-key', ''], diff --git a/packages/cli/src/cli/commands/local-agent.ts b/packages/cli/src/cli/commands/local-agent.ts index dca3abfd8..2bdd16296 100644 --- a/packages/cli/src/cli/commands/local-agent.ts +++ b/packages/cli/src/cli/commands/local-agent.ts @@ -12,7 +12,9 @@ import type { AttachMode } from '../lib/attach-mode.js'; import { attachNative, isNativeHarness, type NativeAttachOptions } from '../lib/attach-native.js'; import { attachPassthrough } from '../lib/attach-passthrough.js'; import { attachRemoteNode, type RemoteNodeAttachOptions } from '../lib/attach-remote-node.js'; +import { startFleetNodeAttachProxy } from '../lib/attach-fleet-node.js'; import { attachView } from '../lib/attach-view.js'; +import { createBackpressureAwareWriter } from '../lib/attach.js'; import { defaultStateDir, readConnectionFileFromDisk, @@ -76,6 +78,76 @@ export function runAttach(name: string, mode: AttachMode, options: NativeAttachO }); } +/** + * Run an existing attach mode against the short-lived loopback adapter for a + * remote fleet node. Deliberately bypasses native-harness detection: the proxy + * represents a PTY byte stream, while native mode is an event protocol. + */ +export async function attachFleetNode( + name: string, + mode: AttachMode, + node: string, + options: NativeAttachOptions +): Promise { + const proxy = await startFleetNodeAttachProxy({ agent: name, node, mode }); + const jsonWriter = options.json ? createBackpressureAwareWriter(process.stdout) : undefined; + try { + const connectionOptions = { brokerUrl: proxy.brokerUrl, apiKey: proxy.apiKey }; + // Fleet agents expose a PTY stream rather than native-harness envelopes. + // In JSON mode retain the machine-readable contract by serializing every + // rendered stream chunk as NDJSON, including the initial snapshot. + const jsonOutput = (chunk: string) => { + jsonWriter?.write(`${JSON.stringify({ kind: 'worker_stream', name, stream: 'stdout', chunk })}\n`); + }; + if (options.reasoning || options.diagnostics) { + process.stderr.write( + '[node attach] reasoning and diagnostics are unavailable for remote PTY terminal sessions.\n' + ); + } + switch (mode) { + case 'view': + return await attachView( + name, + connectionOptions, + options.json ? { writeChunk: jsonOutput, stdoutIsTty: false } : {} + ); + case 'passthrough': + return await attachPassthrough( + name, + connectionOptions, + options.json + ? { + writeChunk: jsonOutput, + // JSON mode is a non-terminal stream. Suppressing the local + // status line and reset controls keeps only worker bytes in + // the NDJSON records, matching SSH's no-TTY JSON behaviour. + terminal: { getSize: () => null, onResize: () => () => undefined }, + } + : {} + ); + case 'drive': + default: + return await attachDrive( + name, + connectionOptions, + options.json + ? { + writeChunk: jsonOutput, + terminal: { getSize: () => null, onResize: () => () => undefined }, + } + : {} + ); + } + } finally { + // NDJSON is a data contract, unlike an interactive screen: queue every + // locally buffered record before teardown so the shared CLI stdio drain + // can carry it through a backpressured pipe. + jsonWriter?.flush(); + jsonWriter?.dispose(); + await proxy.close(); + } +} + type ExitFn = (code: number) => never; export interface LocalAgentDependencies { @@ -88,6 +160,7 @@ export interface LocalAgentDependencies { node: string, options: RemoteNodeAttachOptions ) => Promise; + attachNode: (name: string, mode: AttachMode, node: string, options: NativeAttachOptions) => Promise; cwd: () => string; readConnectionFile: (stateDir: string) => unknown; getDefaultStateDir: () => string; @@ -109,6 +182,7 @@ function withDefaults(overrides: Partial = {}): LocalAge fetch: globalThis.fetch, attach: runAttach, attachRemote: attachRemoteNode, + attachNode: attachFleetNode, log: (...args: unknown[]) => console.log(...args), error: (...args: unknown[]) => console.error(...args), exit: defaultExit, @@ -491,6 +565,7 @@ export function registerLocalAgentCommands( .description('Attach to a running agent interactively (drive | view | passthrough)') .argument('', 'Agent name') .option('--mode ', 'drive | view | passthrough', 'view') + .option('--node ', 'Canonical authenticated fleet-node terminal attach (physical or Daytona)') .option('--ssh-host ', 'SSH host fallback for a physical fleet node') .option('--broker-url ', 'Broker base URL (overrides RELAY_BROKER_URL and connection.json)') .option('--api-key ', 'Broker API key (overrides RELAY_BROKER_API_KEY and connection.json)') @@ -509,6 +584,40 @@ export function registerLocalAgentCommands( return; } const sshHost = options.sshHost as string | undefined; + const node = options.node as string | undefined; + if (node !== undefined && sshHost !== undefined) { + deps.error( + 'Error: --node cannot be combined with --ssh-host. Use --ssh-host only as the explicit SSH fallback.' + ); + deps.exit(1); + return; + } + if (node !== undefined) { + if ( + options.brokerUrl !== undefined || + options.apiKey !== undefined || + options.stateDir !== undefined + ) { + deps.error( + 'Error: --node cannot be combined with --broker-url, --api-key, or --state-dir. It opens an ephemeral authenticated loopback session.' + ); + deps.exit(1); + return; + } + try { + const code = await deps.attachNode(name, mode, node, { + json: options.json as boolean | undefined, + reasoning: options.reasoning as boolean | undefined, + diagnostics: options.diagnostics as boolean | undefined, + }); + if (code !== 0) deps.exit(code); + } catch (error) { + const message = error instanceof Error ? error.message : String(error); + deps.error(message.startsWith('Error:') ? message : `Error: ${message}`); + deps.exit(1); + } + return; + } if (sshHost !== undefined) { if (options.brokerUrl !== undefined || options.apiKey !== undefined) { deps.error('Error: --ssh-host cannot be combined with --broker-url or --api-key.'); diff --git a/packages/cli/src/cli/lib/attach-fleet-node.ts b/packages/cli/src/cli/lib/attach-fleet-node.ts new file mode 100644 index 000000000..7e02cc27d --- /dev/null +++ b/packages/cli/src/cli/lib/attach-fleet-node.ts @@ -0,0 +1,595 @@ +/** + * Ticketed fleet-node attach adapter. + * + * The established attach clients intentionally continue to speak the local + * broker HTTP/WebSocket contract. This short-lived loopback adapter maps that + * contract onto Relaycast's authenticated terminal session, so view/drive and + * passthrough retain their behaviour without exposing a remote broker listener + * or copying a broker API key off a physical or Daytona node. + */ + +import { createServer, type IncomingMessage, type ServerResponse } from 'node:http'; +import { once } from 'node:events'; +import { Buffer } from 'node:buffer'; +import { randomBytes } from 'node:crypto'; + +import WebSocket, { WebSocketServer } from 'ws'; + +import type { AttachMode } from './attach-mode.js'; +import { resolveBaseUrl, resolveWorkspaceKey } from './sdk-client.js'; + +const MAX_BUFFERED_BYTES = 1024 * 1024; +const SNAPSHOT_WAIT_MS = 10_000; +const INITIAL_RECONNECT_DELAY_MS = 500; +const MAX_RECONNECT_DELAY_MS = 30_000; +const MAX_RECONNECT_ATTEMPTS = 5; + +type FleetSessionResponse = { + ok?: boolean; + data?: { + session_id?: string; + terminal_url?: string; + resume_token?: string; + }; + error?: { code?: string; message?: string }; +}; + +type TerminalFrame = Record & { type?: string; session_id?: string }; + +type TerminalReadiness = { + generation: number; + settled: boolean; + promise: Promise; + resolve: () => void; + reject: (error: Error) => void; +}; + +export interface FleetNodeAttachOptions { + agent: string; + node: string; + mode: AttachMode; + env?: NodeJS.ProcessEnv; + baseUrl?: string; + workspaceKey?: string; + fetch?: typeof globalThis.fetch; +} + +export interface FleetNodeAttachProxy { + brokerUrl: string; + apiKey: string; + close(): Promise; +} + +export class FleetNodeAttachError extends Error { + constructor( + message: string, + readonly code?: string + ) { + super(message); + this.name = 'FleetNodeAttachError'; + } +} + +function json(response: ServerResponse, status: number, payload: unknown): void { + response.writeHead(status, { 'Content-Type': 'application/json', 'Cache-Control': 'no-store' }); + response.end(JSON.stringify(payload)); +} + +function readBody(request: IncomingMessage): Promise> { + return new Promise((resolve) => { + let settled = false; + const finish = (body: Record) => { + if (settled) return; + settled = true; + resolve(body); + }; + let body = ''; + request.setEncoding('utf8'); + request.on('data', (chunk: string) => { + body += chunk; + }); + request.on('end', () => { + try { + const parsed = JSON.parse(body) as unknown; + finish( + parsed && typeof parsed === 'object' && !Array.isArray(parsed) + ? (parsed as Record) + : {} + ); + } catch { + finish({}); + } + }); + request.on('error', () => finish({})); + request.on('aborted', () => finish({})); + }); +} + +function asWsUrl(value: string): string { + return value.replace(/^http/i, 'ws'); +} + +function safeNodePath(node: string): string { + const trimmed = node.trim().replace(/^#/, ''); + if (!trimmed) throw new FleetNodeAttachError('Error: --node requires a node name or id.', 'invalid_node'); + return encodeURIComponent(trimmed); +} + +function parseFrame(data: WebSocket.RawData): TerminalFrame | null { + try { + const parsed = JSON.parse(rawDataToString(data)) as unknown; + return parsed && typeof parsed === 'object' && !Array.isArray(parsed) ? (parsed as TerminalFrame) : null; + } catch { + return null; + } +} + +function rawDataToString(data: WebSocket.RawData): string { + if (Buffer.isBuffer(data)) return data.toString('utf8'); + if (data instanceof ArrayBuffer) return Buffer.from(data).toString('utf8'); + if (Array.isArray(data)) return Buffer.concat(data).toString('utf8'); + return String(data); +} + +/** Start a broker-compatible loopback proxy for one remote terminal session. */ +export async function startFleetNodeAttachProxy( + options: FleetNodeAttachOptions +): Promise { + const env = options.env ?? process.env; + const fetchFn = options.fetch ?? globalThis.fetch; + const workspaceKey = options.workspaceKey ?? resolveWorkspaceKey({ env }); + const baseUrl = (options.baseUrl ?? resolveBaseUrl({ env }) ?? 'https://cast.agentrelay.com').replace( + /\/+$/, + '' + ); + const nodePath = safeNodePath(options.node); + const ticketResponse = await fetchFn(`${baseUrl}/v1/nodes/${nodePath}/terminal/sessions`, { + method: 'POST', + headers: { Authorization: `Bearer ${workspaceKey}`, 'Content-Type': 'application/json' }, + body: JSON.stringify({ agent: options.agent, mode: options.mode }), + }); + const ticketPayload = (await ticketResponse.json().catch(() => ({}))) as FleetSessionResponse; + const terminalUrl = ticketPayload.data?.terminal_url; + const sessionId = ticketPayload.data?.session_id; + const resumeToken = ticketPayload.data?.resume_token; + if (!ticketResponse.ok || !terminalUrl || !sessionId || !resumeToken) { + const code = ticketPayload.error?.code; + const message = + ticketPayload.error?.message ?? `terminal session request failed (HTTP ${ticketResponse.status})`; + throw new FleetNodeAttachError(`Error: ${message}`, code); + } + + let connectionGeneration = 0; + const createReadiness = (): TerminalReadiness => { + let resolve!: () => void; + let reject!: (error: Error) => void; + const promise = new Promise((resolvePromise, rejectPromise) => { + resolve = resolvePromise; + reject = rejectPromise; + }); + // A failure can land before a snapshot request attaches its waiter. Keep + // the rejection observable while avoiding an unhandled-rejection warning. + void promise.catch(() => undefined); + return { generation: ++connectionGeneration, settled: false, promise, resolve, reject }; + }; + let activeReadiness = createReadiness(); + const resolveReadiness = (readiness: TerminalReadiness) => { + if (readiness.settled) return; + readiness.settled = true; + readiness.resolve(); + }; + const rejectReadiness = (readiness: TerminalReadiness, error: Error) => { + if (readiness.settled) return; + readiness.settled = true; + readiness.reject(error); + }; + const waitForCurrentReadiness = async (): Promise => { + for (;;) { + const readiness = activeReadiness; + await readiness.promise; + if (readiness === activeReadiness) return; + } + }; + const snapshot: { screen: string; rows: number; cols: number; offset: number } = { + screen: '', + rows: 24, + cols: 80, + offset: 0, + }; + const eventSockets = new Set(); + const inputSockets = new Set(); + const outputHistory: Array<{ chunk: string; offset?: number }> = []; + let outputHistoryBytes = 0; + let remote: WebSocket | undefined; + let stopped = false; + let terminalEnded = false; + let terminalEverReady = false; + let reconnecting = false; + let reconnectAttempts = 0; + let reconnectTimer: ReturnType | undefined; + const loopbackApiKey = randomBytes(32).toString('base64url'); + const loopbackAuthorized = (headers: IncomingMessage['headers']) => + headers.authorization === `Bearer ${loopbackApiKey}` || headers['x-api-key'] === loopbackApiKey; + + const server = createServer(async (request, response) => { + if (!loopbackAuthorized(request.headers)) { + json(response, 401, { + error: { code: 'unauthorized', message: 'loopback terminal token is required' }, + }); + return; + } + const path = new URL(request.url ?? '/', 'http://127.0.0.1').pathname; + if (request.method === 'GET' && path === `/api/spawned/${encodeURIComponent(options.agent)}/snapshot`) { + try { + await new Promise((resolve, reject) => { + const timer = setTimeout(() => reject(new Error('terminal snapshot timed out')), SNAPSHOT_WAIT_MS); + void waitForCurrentReadiness().then( + () => { + clearTimeout(timer); + resolve(); + }, + (error: Error) => { + clearTimeout(timer); + reject(error); + } + ); + }); + } catch (error) { + const terminalError = error instanceof FleetNodeAttachError ? error : undefined; + const status = + terminalError?.code === 'agent_not_found' + ? 404 + : terminalError?.code === 'unsupported_runtime' + ? 409 + : 503; + json(response, status, { + error: { + code: terminalError?.code ?? 'snapshot_unavailable', + message: + terminalError?.message ?? (error instanceof Error ? error.message : 'snapshot unavailable'), + }, + }); + return; + } + json(response, 200, { format: 'ansi', ...snapshot }); + return; + } + const name = encodeURIComponent(options.agent); + if (path === `/api/spawned/${name}/delivery-mode`) { + if (request.method === 'GET') { + json(response, 200, { mode: options.mode === 'drive' ? 'manual_flush' : 'auto_inject' }); + } else { + await readBody(request); + json(response, 200, { + mode: options.mode === 'drive' ? 'manual_flush' : 'auto_inject', + flushed: 0, + matched: true, + revision: '1', + }); + } + return; + } + if (request.method === 'GET' && path === `/api/spawned/${name}/pending`) { + json(response, 200, { pending: [] }); + return; + } + if (request.method === 'POST' && path === `/api/spawned/${name}/flush`) { + json(response, 200, { flushed: 0 }); + return; + } + if (request.method === 'GET' && path === '/api/spawned') { + json(response, 200, { agents: [{ name: options.agent, workerPid: 1 }] }); + return; + } + if (request.method === 'POST' && path === `/api/resize/${name}`) { + const body = await readBody(request); + if (body.release === true) { + json(response, 200, { name: options.agent, released: true }); + return; + } + const rows = typeof body.rows === 'number' ? body.rows : 0; + const cols = typeof body.cols === 'number' ? body.cols : 0; + if (!Number.isInteger(rows) || !Number.isInteger(cols) || rows < 1 || cols < 1) { + json(response, 400, { + error: { code: 'invalid_dimensions', message: 'rows and cols must be positive integers' }, + }); + return; + } + try { + await new Promise((resolve, reject) => { + const timer = setTimeout(() => reject(new Error('terminal resize timed out')), SNAPSHOT_WAIT_MS); + void waitForCurrentReadiness().then( + () => { + clearTimeout(timer); + resolve(); + }, + (error: Error) => { + clearTimeout(timer); + reject(error); + } + ); + }); + } catch (error) { + json(response, 503, { + error: { + code: 'session_not_ready', + message: error instanceof Error ? error.message : 'terminal session is not ready', + }, + }); + return; + } + if (!remote || remote.readyState !== WebSocket.OPEN || remote.bufferedAmount > MAX_BUFFERED_BYTES) { + json(response, 503, { + error: { code: 'node_unreachable', message: 'terminal transport is unavailable' }, + }); + return; + } + remote.send(JSON.stringify({ type: 'terminal.resize', session_id: sessionId, rows, cols })); + json(response, 200, { name: options.agent, rows, cols, applied: true }); + return; + } + json(response, 404, { error: { code: 'not_found', message: 'loopback terminal endpoint not found' } }); + }); + const websocketServer = new WebSocketServer({ noServer: true }); + + const closeSocket = (socket: WebSocket, code: number, reason: string) => { + try { + socket.close(code, reason); + } catch { + /* connection already gone */ + } + }; + const broadcast = (sockets: Set, payload: unknown): boolean => { + const encoded = JSON.stringify(payload); + let accepted = false; + for (const socket of sockets) { + if (socket.readyState !== WebSocket.OPEN) continue; + if (socket.bufferedAmount > MAX_BUFFERED_BYTES) { + closeSocket(socket, 1013, 'loopback client backpressure exceeded'); + sockets.delete(socket); + continue; + } + try { + socket.send(encoded); + accepted = true; + } catch { + sockets.delete(socket); + } + } + return accepted; + }; + const workerStreamEvent = (chunk: string, offset?: number) => ({ + kind: 'worker_stream', + name: options.agent, + stream: 'stdout', + chunk, + ...(offset === undefined ? {} : { offset }), + }); + const retainOutput = (chunk: string, offset: number | undefined): boolean => { + const bytes = Buffer.byteLength(chunk, 'utf8'); + if (outputHistoryBytes + bytes > MAX_BUFFERED_BYTES) return false; + outputHistory.push({ chunk, ...(offset === undefined ? {} : { offset }) }); + outputHistoryBytes += bytes; + return true; + }; + + websocketServer.on('connection', (socket, request) => { + const path = new URL(request.url ?? '/', 'http://127.0.0.1').pathname; + if (path === '/ws') { + eventSockets.add(socket); + socket.on('close', () => eventSockets.delete(socket)); + let replayed = 0; + for (const event of outputHistory) { + if (socket.readyState !== WebSocket.OPEN || socket.bufferedAmount > MAX_BUFFERED_BYTES) break; + try { + socket.send(JSON.stringify(workerStreamEvent(event.chunk, event.offset))); + } catch { + break; + } + replayed += 1; + } + if (replayed > 0) { + const sentBytes = outputHistory + .slice(0, replayed) + .reduce((total, event) => total + Buffer.byteLength(event.chunk, 'utf8'), 0); + outputHistory.splice(0, replayed); + outputHistoryBytes -= sentBytes; + } + return; + } + if (path === `/api/input/${encodeURIComponent(options.agent)}/stream`) { + inputSockets.add(socket); + socket.on('close', () => inputSockets.delete(socket)); + socket.send(JSON.stringify({ type: 'pty_input_ready', name: options.agent })); + socket.on('message', (data) => { + if (!remote || remote.readyState !== WebSocket.OPEN || remote.bufferedAmount > MAX_BUFFERED_BYTES) { + broadcast(inputSockets, { + type: 'pty_input_error', + code: 'node_unreachable', + message: 'terminal transport is unavailable', + }); + return; + } + const raw = rawDataToString(data); + remote.send( + JSON.stringify({ + type: 'terminal.input', + session_id: sessionId, + data_base64: Buffer.from(raw, 'utf8').toString('base64'), + }) + ); + }); + return; + } + closeSocket(socket, 1008, 'unknown loopback endpoint'); + }); + server.on('upgrade', (request, socket, head) => { + if (!loopbackAuthorized(request.headers)) { + socket.write('HTTP/1.1 401 Unauthorized\r\nConnection: close\r\n\r\n'); + socket.destroy(); + return; + } + websocketServer.handleUpgrade(request, socket, head, (client) => + websocketServer.emit('connection', client, request) + ); + }); + server.listen(0, '127.0.0.1'); + await once(server, 'listening'); + const address = server.address(); + if (!address || typeof address === 'string') + throw new FleetNodeAttachError( + 'Error: could not allocate loopback terminal listener.', + 'loopback_unavailable' + ); + + const resumeUrl = new URL(terminalUrl); + resumeUrl.searchParams.delete('ticket'); + resumeUrl.searchParams.set('session_id', sessionId); + resumeUrl.searchParams.set('resume', resumeToken); + const endTerminal = (error: FleetNodeAttachError) => { + if (terminalEnded) return; + terminalEnded = true; + if (reconnectTimer) { + clearTimeout(reconnectTimer); + reconnectTimer = undefined; + } + const activeRemote = remote; + remote = undefined; + rejectReadiness(activeReadiness, error); + broadcast(inputSockets, { type: 'pty_input_error', code: error.code, message: error.message }); + for (const socket of eventSockets) closeSocket(socket, 1011, error.message); + if (activeRemote && activeRemote.readyState !== WebSocket.CLOSED) { + try { + activeRemote.terminate(); + } catch { + // The socket may have closed between the state check and terminate. + } + } + }; + const failRemote = (message: string) => { + endTerminal(new FleetNodeAttachError(message, 'node_unreachable')); + }; + const connect = (url: string, readiness: TerminalReadiness) => { + if (stopped || terminalEnded) return; + const socket = new WebSocket(asWsUrl(url)); + remote = socket; + socket.on('message', (data) => { + // A late frame from a transport superseded during reconnect must never + // overwrite the fresh snapshot or end the replacement session. + if (remote !== socket || stopped || terminalEnded) return; + const frame = parseFrame(data); + if (!frame || frame.session_id !== sessionId) return; + if (frame.type === 'terminal.ready') { + snapshot.screen = typeof frame.screen === 'string' ? frame.screen : ''; + snapshot.rows = typeof frame.rows === 'number' ? frame.rows : 24; + snapshot.cols = typeof frame.cols === 'number' ? frame.cols : 80; + snapshot.offset = typeof frame.offset === 'number' ? frame.offset : 0; + terminalEverReady = true; + reconnectAttempts = 0; + if (readiness === activeReadiness) { + resolveReadiness(readiness); + // A reconnect gets a fresh ANSI grid but existing local `/ws` + // consumers have already performed their initial HTTP snapshot. + // Re-emit this screen without an offset so they repaint instead of + // retaining a stale pre-reconnect terminal image. + if (readiness.generation > 1 && snapshot.screen) { + broadcast(eventSockets, workerStreamEvent(snapshot.screen)); + } + } + } else if (frame.type === 'terminal.output' && typeof frame.chunk === 'string') { + const offset = typeof frame.offset === 'number' ? frame.offset : undefined; + if (!broadcast(eventSockets, workerStreamEvent(frame.chunk, offset))) { + if (!retainOutput(frame.chunk, offset)) { + endTerminal( + new FleetNodeAttachError( + 'terminal output exceeded the bounded loopback buffer', + 'output_backpressure' + ) + ); + } + } + } else if (frame.type === 'terminal.input_ack') { + broadcast(inputSockets, { + type: 'pty_input_ack', + name: options.agent, + bytes_written: typeof frame.bytes_written === 'number' ? frame.bytes_written : 0, + }); + } else if (frame.type === 'terminal.error') { + const message = typeof frame.message === 'string' ? frame.message : 'remote terminal failed'; + const code = typeof frame.code === 'string' ? frame.code : 'terminal_error'; + if (readiness === activeReadiness && !readiness.settled) { + endTerminal(new FleetNodeAttachError(message, code)); + } else { + broadcast(inputSockets, { type: 'pty_input_error', code, message }); + } + } else if (frame.type === 'terminal.closed') { + endTerminal(new FleetNodeAttachError('remote terminal session closed', 'terminal_closed')); + } + }); + socket.on('error', () => { + // Initial connection failure has no terminal state worth preserving. + // Fail promptly with the canonical unavailable-node error instead of + // letting the HTTP snapshot timeout mask it. Once Ready has been seen, + // the close handler retains the bounded resume/backoff behaviour. + if (remote === socket && readiness === activeReadiness && !readiness.settled && !terminalEverReady) { + failRemote('terminal transport could not connect to the fleet node'); + } + }); + socket.on('close', () => { + if (remote !== socket || stopped || terminalEnded || reconnecting) return; + // Any waiter that observed the prior connection must retry against the + // fresh generation instead of receiving its stale resolved snapshot. + resolveReadiness(readiness); + const nextReadiness = createReadiness(); + activeReadiness = nextReadiness; + if (reconnectAttempts >= MAX_RECONNECT_ATTEMPTS) { + failRemote('terminal transport could not reconnect to the fleet node'); + return; + } + reconnecting = true; + const delay = Math.min(INITIAL_RECONNECT_DELAY_MS * 2 ** reconnectAttempts, MAX_RECONNECT_DELAY_MS); + reconnectAttempts += 1; + reconnectTimer = setTimeout(() => { + reconnectTimer = undefined; + reconnecting = false; + connect(resumeUrl.toString(), nextReadiness); + }, delay); + }); + }; + connect(terminalUrl, activeReadiness); + + return { + brokerUrl: `http://127.0.0.1:${address.port}`, + apiKey: loopbackApiKey, + async close() { + if (stopped) return; + stopped = true; + terminalEnded = true; + if (reconnectTimer) { + clearTimeout(reconnectTimer); + reconnectTimer = undefined; + } + rejectReadiness(activeReadiness, new FleetNodeAttachError('terminal attach closed', 'closed')); + const activeRemote = remote; + remote = undefined; + if (activeRemote && activeRemote.readyState === WebSocket.OPEN) { + try { + activeRemote.send(JSON.stringify({ type: 'terminal.close', session_id: sessionId })); + } catch { + // Best effort; terminate below still prevents a late reconnect. + } + } + if (activeRemote && activeRemote.readyState !== WebSocket.CLOSED) { + try { + activeRemote.terminate(); + } catch { + // Socket is already gone. + } + } + for (const socket of [...eventSockets, ...inputSockets]) + closeSocket(socket, 1000, 'terminal attach closed'); + websocketServer.close(); + await new Promise((resolve) => server.close(() => resolve())); + }, + }; +} diff --git a/packages/cli/src/cli/lib/attach.test.ts b/packages/cli/src/cli/lib/attach.test.ts index cab08628f..a719c32f4 100644 --- a/packages/cli/src/cli/lib/attach.test.ts +++ b/packages/cli/src/cli/lib/attach.test.ts @@ -960,6 +960,33 @@ describe('createBackpressureAwareWriter', () => { expect(written).toEqual(['a', 'bb', 'c']); }); + it('flushes buffered records in order before disposal', () => { + const written: string[] = []; + const accept = false; + let drain: (() => void) | null = null; + const stdout: BackpressureWritable = { + write: (chunk) => { + written.push(chunk); + return accept; + }, + once: (_event, listener) => { + drain = listener; + return undefined; + }, + off: (_event, listener) => { + if (listener === drain) drain = null; + return undefined; + }, + }; + const w = createBackpressureAwareWriter(stdout); + w.write('first'); + w.write('second'); + w.write('third'); + w.flush(); + w.dispose(); + expect(written).toEqual(['first', 'second', 'third']); + }); + it('drops the pending queue and unhooks drain on dispose', () => { const written: string[] = []; let accept = true; diff --git a/packages/cli/src/cli/lib/attach.ts b/packages/cli/src/cli/lib/attach.ts index 21683dbd2..91886d1a3 100644 --- a/packages/cli/src/cli/lib/attach.ts +++ b/packages/cli/src/cli/lib/attach.ts @@ -1223,14 +1223,19 @@ export interface BackpressureWritable { } /** - * A backpressure-aware writer plus a teardown hook. Call the writer with each - * chunk; call {@link BackpressureAwareWriter.dispose} on detach to drop any - * still-queued chunks and unhook the pending `'drain'` listener so nothing - * flushes to stdout after the session tears down. + * A backpressure-aware writer plus teardown hooks. Call the writer with each + * chunk. Interactive attaches call {@link BackpressureAwareWriter.dispose} on + * detach to drop any still-queued terminal bytes; machine-readable callers can + * call {@link BackpressureAwareWriter.flush} first to preserve their records. */ export interface BackpressureAwareWriter { /** Write a chunk, respecting backpressure. No-op after {@link dispose}. */ write: (chunk: string) => void; + /** + * Queue every locally buffered chunk to the underlying stream in FIFO order. + * The stream's own buffered writes are drained by the shared CLI exit path. + */ + flush: () => void; /** Drop the pending queue and unhook the `'drain'` listener. Idempotent. */ dispose: () => void; } @@ -1309,7 +1314,22 @@ export function createBackpressureAwareWriter( off?.call(stdout, 'drain', flushQueue); }; - return { write, dispose }; + const flush = (): void => { + if (disposed || queue.length === 0) return; + const pending = queue; + queue = []; + queuedBytes = 0; + paused = false; + const off = stdout.off ?? stdout.removeListener; + off?.call(stdout, 'drain', flushQueue); + // `stdout.write` queues asynchronously on pipes. Preserving JSON records + // is more important than applying another local backpressure pause here; + // exitAfterFlush writes a final sentinel and waits (bounded) for that + // stream queue after command completion. + for (const chunk of pending) stdout.write(chunk); + }; + + return { write, flush, dispose }; } /** Dependencies for `captureInitialSnapshot`. `captureAndRenderSnapshot`