diff --git a/crates/broker/src/runtime/api.rs b/crates/broker/src/runtime/api.rs index cde595106..3cd3ff6fd 100644 --- a/crates/broker/src/runtime/api.rs +++ b/crates/broker/src/runtime/api.rs @@ -16,6 +16,16 @@ const DEFAULT_OBSERVER_TOKEN_NAME: &str = "pear-dashboard-observer"; /// round-trip that resolves in well under a second on a healthy worker. const PTY_INPUT_ACK_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(5); +/// `/model` writes use the same worker-owned stdin writer as protocol frames. +/// Never let the runtime actor wait on its completion without a deadline. +const DEFAULT_SET_MODEL_TIMEOUT: Duration = Duration::from_secs(5); + +fn set_model_write_timeout(timeout_ms: Option) -> Duration { + timeout_ms + .map(Duration::from_millis) + .unwrap_or(DEFAULT_SET_MODEL_TIMEOUT) +} + /// Scopes granted to observer tokens minted via `/api/observer-token`: broad /// read access to workspace activity, deliberately excluding anything /// write/spawn-capable (unlike the raw `rk_live_...` workspace key this @@ -736,35 +746,30 @@ impl BrokerRuntime { timeout_ms, reply, } => { - let Some(handle) = workers.workers.get_mut(&name) else { + if !workers.workers.contains_key(&name) { let _ = reply.send(Err(format!("unknown worker '{}'", name))); return; - }; + } let model_command = format!("/model {}\n", model); - let result = async { - handle - .stdin - .write_all(model_command.as_bytes()) - .await - .with_context(|| { - format!("failed writing model command to worker '{}'", name) - })?; - handle - .stdin - .flush() - .await - .with_context(|| format!("failed flushing worker '{}' stdin", name))?; - if let Some(timeout_ms) = timeout_ms { - tracing::info!( - name = %name, - timeout_ms, - "HTTP API set_model timeout_ms is currently advisory only" - ); - } - Ok::<(), anyhow::Error>(()) - } - .await; + let set_model_timeout = set_model_write_timeout(timeout_ms); + // `send_raw_to_worker` completes only after the command enters + // the worker-owned writer queue. Tokio channel sends are + // cancellation-safe, so a timeout means the command was not + // admitted; once admitted, report it as pending rather than + // claiming it failed while the writer can still emit it. + let result = match timeout( + set_model_timeout, + workers.send_raw_to_worker(&name, model_command.into_bytes()), + ) + .await + { + Ok(result) => result, + Err(_) => Err(anyhow::anyhow!( + "set_model timed out after {}ms for '{name}'", + set_model_timeout.as_millis() + )), + }; match result { Ok(()) => { @@ -772,6 +777,8 @@ impl BrokerRuntime { "name": name, "model": model, "success": true, + "accepted": true, + "pending": true, }))); } Err(error) => { @@ -2593,6 +2600,25 @@ mod skill_injection_tests { } } +#[cfg(test)] +mod set_model_timeout_tests { + use super::{set_model_write_timeout, DEFAULT_SET_MODEL_TIMEOUT}; + use std::time::Duration; + + #[test] + fn set_model_timeout_uses_the_requested_deadline_or_a_finite_default() { + assert_eq!(set_model_write_timeout(None), DEFAULT_SET_MODEL_TIMEOUT); + assert_eq!( + set_model_write_timeout(Some(250)), + Duration::from_millis(250) + ); + assert_eq!( + set_model_write_timeout(Some(10_000)), + Duration::from_millis(10_000) + ); + } +} + fn persist_agent_channels( state: &mut broker::BrokerState, name: &str, diff --git a/crates/broker/src/runtime/fleet.rs b/crates/broker/src/runtime/fleet.rs index bf74c6e99..55ac71345 100644 --- a/crates/broker/src/runtime/fleet.rs +++ b/crates/broker/src/runtime/fleet.rs @@ -19,10 +19,6 @@ 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 @@ -42,6 +38,61 @@ pub(super) fn try_send_terminal( .is_ok() } +/// Release a resize lease only when `session_id` is the terminal session that +/// owns it. A terminal session may end because its client closes, its transport +/// disconnects, or a terminal operation fails; all are equivalent to detach +/// for single-resizer ownership. +pub(super) fn release_terminal_resize_ownership( + resize_owners: &mut HashMap, + agent: &WorkerName, + session_id: &str, +) { + if resize_owners + .get(agent) + .is_some_and(|owner| owner.session_id == session_id) + { + resize_owners.remove(agent); + } +} + +pub(super) fn fail_terminal_session( + terminal_control_tx: &mpsc::Sender, + terminal_sessions: &mut HashMap, + terminal_snapshot_requests: &mut HashMap, + terminal_input_requests: &mut HashMap, + session_id: String, + code: &str, + message: String, +) { + 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); + + // Error is useful when the terminal lane has room, but it is non-final and + // deliberately gives way to the reserved close capacity. Queue its close + // directly instead of routing through `send_terminal`: the generic + // backpressure fallback would otherwise produce a second close with a + // different reason. + let _ = try_send_terminal( + terminal_control_tx, + TerminalToCloud::Error { + session_id: session_id.clone(), + code: code.into(), + message: message.clone(), + }, + ); + if !try_send_terminal( + terminal_control_tx, + TerminalToCloud::Closed { + session_id: session_id.clone(), + code: Some(code.into()), + message: Some(message), + }, + ) { + tracing::warn!(target = "relay_broker::terminal", session_id = %session_id, "terminal close could not be queued after session failure"); + } +} + #[derive(Debug, Clone)] pub(super) struct PendingVerifiedSpawn { pub(super) invocation_id: String, @@ -116,6 +167,8 @@ impl BrokerRuntime { // 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.resize_owners + .retain(|_, owner| !self.terminal_sessions.contains_key(&owner.session_id)); self.terminal_sessions.clear(); self.terminal_snapshot_requests.clear(); self.terminal_input_requests.clear(); @@ -153,44 +206,29 @@ impl BrokerRuntime { 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. + // Queue the grid request on the worker-owned stdin + // writer. The event loop never awaits a PTY pipe write; + // the writer serializes complete frames and reports a + // later pipe failure as a terminal session failure. 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(), - ); - } + self.terminal_snapshot_requests.insert( + request_id.clone(), + TerminalSnapshotRequest { + session_id: session_id.clone(), + deadline: Instant::now() + TERMINAL_SNAPSHOT_TIMEOUT, + }, + ); + if let Err(error) = self.workers.try_send_to_worker( + agent_name.as_str(), + "snapshot_pty", + Some(RequestId::new(request_id.clone())), + json!({ "format": "ansi" }), + ) { + self.fail_terminal_session( + session_id, + "snapshot_failed", + error.to_string(), + ); } } } @@ -268,34 +306,20 @@ impl BrokerRuntime { 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(), - ), + self.terminal_input_requests.insert( + request_id.clone(), + TerminalInputRequest { + session_id: session_id.clone(), + deadline: Instant::now() + TERMINAL_INPUT_ACK_TIMEOUT, + }, + ); + if let Err(error) = self.workers.try_send_to_worker( + session.agent.as_str(), + "write_pty", + Some(RequestId::new(request_id.clone())), + json!({ "data": data }), + ) { + self.fail_terminal_session(session_id, "input_failed", error.to_string()); } } TerminalControlEvent::Message(TerminalFromCloud::Resize { @@ -327,30 +351,66 @@ impl BrokerRuntime { }); 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()) + // Remote drive sessions share the same PTY as local HTTP + // attach clients. Use the common lease planner so only one + // live session controls its size at a time. + match plan_resize( + &mut self.resize_owners, + &session.agent, + rows, + cols, + Some(&session_id), + Instant::now(), + ) { + ResizeAction::Reject => { + tracing::debug!( + target = "relay_broker::terminal", + session_id = %session_id, + worker = %session.agent, + "ignoring terminal resize from a non-owner session" + ); + } + ResizeAction::Refresh => { + // `plan_resize` renewed this session's lease. The PTY + // already has these dimensions, so do not repaint it. + } + ResizeAction::Apply => { + if let Err(error) = self.workers.try_send_to_worker( + session.agent.as_str(), + "resize_pty", + None, + json!({ "rows": rows, "cols": cols }), + ) { + self.fail_terminal_session( + session_id, + "resize_failed", + error.to_string(), + ); + } else { + // As with the HTTP path, do not claim the lease + // until the resize was accepted by the worker + // writer. This path deliberately remains + // non-blocking for the runtime actor. + commit_resize_ownership( + &mut self.resize_owners, + &session.agent, + rows, + cols, + Some(session_id), + Instant::now(), + ); + } } - 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); + if let Some(session) = self.terminal_sessions.remove(&session_id) { + release_terminal_resize_ownership( + &mut self.resize_owners, + &session.agent, + &session_id, + ); + } self.terminal_snapshot_requests .retain(|_, pending| pending.session_id != session_id); self.terminal_input_requests @@ -378,7 +438,13 @@ impl BrokerRuntime { 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); + if let Some(session) = self.terminal_sessions.remove(&session_id) { + release_terminal_resize_ownership( + &mut self.resize_owners, + &session.agent, + &session_id, + ); + } self.terminal_snapshot_requests .retain(|_, pending| pending.session_id != session_id); self.terminal_input_requests @@ -399,21 +465,18 @@ impl BrokerRuntime { } 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 { + if let Some(session) = self.terminal_sessions.get(&session_id) { + release_terminal_resize_ownership(&mut self.resize_owners, &session.agent, &session_id); + } + fail_terminal_session( + &self.terminal_control_tx, + &mut self.terminal_sessions, + &mut self.terminal_snapshot_requests, + &mut self.terminal_input_requests, session_id, - code: Some(code.into()), - message: Some(message), - }); + code, + message, + ); } pub(super) async fn handle_fleet_control_event(&mut self, event: FleetControlEvent) { @@ -1701,6 +1764,92 @@ mod tests { )); } + #[test] + fn terminal_failure_queues_one_close_with_the_original_reason_at_reserve() { + 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, + }, + )); + // This leaves exactly the reserved close capacity. The non-final Error + // must be rejected, but the single final close must still carry the + // actual failure rather than an output_backpressure fallback. + fail_terminal_session( + &tx, + &mut HashMap::new(), + &mut HashMap::new(), + &mut HashMap::new(), + "session-a".into(), + "snapshot_failed", + "worker command queue is full".into(), + ); + + assert!(matches!( + rx.try_recv(), + Ok(TerminalControlCommand::Send(TerminalToCloud::Output { .. })) + )); + assert!(matches!( + rx.try_recv(), + Ok(TerminalControlCommand::Send(TerminalToCloud::Closed { + session_id, + code: Some(code), + message: Some(message), + })) if session_id == "session-a" + && code == "snapshot_failed" + && message == "worker command queue is full" + )); + assert!(rx.try_recv().is_err(), "failure must emit only one close"); + } + + #[test] + fn terminal_session_cleanup_releases_its_resize_lease_only() { + let agent = WorkerName::from("agent-a"); + let other_agent = WorkerName::from("agent-b"); + let mut terminal_sessions = HashMap::from([( + "session-a".to_string(), + TerminalSession { + agent: agent.clone(), + mode: TerminalMode::Drive, + ready: true, + pending_output: Vec::new(), + pending_output_bytes: 0, + }, + )]); + let now = Instant::now(); + let mut resize_owners = HashMap::from([ + ( + agent.clone(), + ResizeOwner { + session_id: "session-a".into(), + last_seen: now, + rows: 24, + cols: 80, + }, + ), + ( + other_agent.clone(), + ResizeOwner { + session_id: "other-session".into(), + last_seen: now, + rows: 30, + cols: 100, + }, + ), + ]); + + let session = terminal_sessions + .remove("session-a") + .expect("test session exists"); + release_terminal_resize_ownership(&mut resize_owners, &session.agent, "session-a"); + + assert!(!resize_owners.contains_key(&agent)); + assert!(resize_owners.contains_key(&other_agent)); + } + #[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/maintenance.rs b/crates/broker/src/runtime/maintenance.rs index 208042fa5..356bdf137 100644 --- a/crates/broker/src/runtime/maintenance.rs +++ b/crates/broker/src/runtime/maintenance.rs @@ -1,4 +1,4 @@ -use super::fleet::try_send_terminal; +use super::fleet::{release_terminal_resize_ownership, try_send_terminal}; use super::*; use crate::terminal_control::TerminalToCloud; @@ -47,7 +47,8 @@ impl BrokerRuntime { .collect(); for (request_id, session_id) in expired_terminal_snapshots { terminal_snapshot_requests.remove(&request_id); - if terminal_sessions.remove(&session_id).is_some() { + if let Some(session) = terminal_sessions.remove(&session_id) { + release_terminal_resize_ownership(resize_owners, &session.agent, &session_id); terminal_input_requests.retain(|_, pending| pending.session_id != session_id); if !try_send_terminal( terminal_control_tx, @@ -81,7 +82,8 @@ impl BrokerRuntime { // 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() { + if let Some(session) = terminal_sessions.remove(&session_id) { + release_terminal_resize_ownership(resize_owners, &session.agent, &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( diff --git a/crates/broker/src/runtime/tests.rs b/crates/broker/src/runtime/tests.rs index 6425b4136..c0615466f 100644 --- a/crates/broker/src/runtime/tests.rs +++ b/crates/broker/src/runtime/tests.rs @@ -17,7 +17,9 @@ use crate::protocol::{ HeadlessHarnessDriver, MessageInjectionMode, NativeHarnessConfig, RelayDelivery, ResolvedHarnessConfig, }; -use crate::worker::{AgentWorkState, WorkerEvent, WorkerHandle, WorkerRegistry}; +use crate::worker::{ + spawn_worker_writer, AgentWorkState, WorkerEvent, WorkerHandle, WorkerRegistry, +}; use crate::{ broker::injection_format::format_injection, util::{ @@ -71,7 +73,7 @@ fn env_test_lock() -> &'static Mutex<()> { async fn make_worker_registry_with_worker(name: &str) -> WorkerRegistry { let (tx, _rx) = mpsc::channel::(16); let mut registry = WorkerRegistry::new( - tx, + tx.clone(), Vec::new(), PathBuf::from("/tmp/agent-relay-broker-tests"), Instant::now(), @@ -83,10 +85,13 @@ async fn make_worker_registry_with_worker(name: &str) -> WorkerRegistry { .spawn() .expect("test worker process should spawn"); let stdin = child.stdin.take().expect("test worker stdin should exist"); + let generation = Uuid::new_v4(); + let (command_tx, command_rx) = mpsc::channel(128); + spawn_worker_writer(tx, WorkerName::from(name), generation, stdin, command_rx); registry.workers.insert( WorkerName::from(name), WorkerHandle { - generation: Uuid::new_v4(), + generation, spec: AgentSpec { name: WorkerName::from(name), runtime: AgentRuntime::Pty, @@ -106,7 +111,7 @@ async fn make_worker_registry_with_worker(name: &str) -> WorkerRegistry { parent: None, workspace_id: Some(WorkspaceId::new("ws_demo")), child, - stdin, + command_tx, harness_pid: None, spawned_at: Instant::now(), // Ready, so the orphan sweep's readiness deadline never applies to diff --git a/crates/broker/src/runtime/worker_events.rs b/crates/broker/src/runtime/worker_events.rs index 91dbf7d29..56cfc42f3 100644 --- a/crates/broker/src/runtime/worker_events.rs +++ b/crates/broker/src/runtime/worker_events.rs @@ -1,5 +1,6 @@ use super::fleet::{ - refresh_fleet_inventory_session_ref, try_send_terminal, verified_spawn_ready_result, + fail_terminal_session, refresh_fleet_inventory_session_ref, try_send_terminal, + verified_spawn_ready_result, }; use super::*; use crate::terminal_control::{TerminalControlCommand, TerminalToCloud}; @@ -532,6 +533,54 @@ impl BrokerRuntime { let terminal_input_requests = &mut self.terminal_input_requests; match worker_event { + WorkerEvent::WriterFailed { + name, + generation, + error, + } => { + let current_generation = workers.workers.get(&name).map(|handle| handle.generation); + if !worker_event_is_current(current_generation, generation) { + tracing::debug!( + target = "agent_relay::broker", + worker = %name, + event_generation = %generation, + current_generation = ?current_generation, + "ignoring writer failure from stale worker generation" + ); + return; + } + + tracing::warn!( + target = "relay_broker::terminal", + worker = %name, + error = %error, + "worker command writer failed; closing attached terminals and resetting worker" + ); + let session_ids: Vec = terminal_sessions + .iter() + .filter(|(_, session)| session.agent == name) + .map(|(session_id, _)| session_id.clone()) + .collect(); + for session_id in session_ids { + fail_terminal_session( + terminal_control_tx, + terminal_sessions, + terminal_snapshot_requests, + terminal_input_requests, + session_id, + "worker_write_failed", + format!("worker command writer failed: {error}"), + ); + } + if let Err(release_error) = workers.terminate_after_writer_failure(name.as_str()) { + tracing::warn!( + target = "relay_broker::terminal", + worker = %name, + error = %release_error, + "failed to signal worker after command writer failure" + ); + } + } WorkerEvent::Message { name, generation, @@ -1019,7 +1068,12 @@ impl BrokerRuntime { } }; let snapshot_ready = matches!(&message, TerminalToCloud::Ready { .. }); - let snapshot_failed = matches!(&message, TerminalToCloud::Error { .. }); + let snapshot_failure = match &message { + TerminalToCloud::Error { code, message, .. } => { + Some((code.clone(), message.clone())) + } + _ => None, + }; 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( @@ -1062,8 +1116,16 @@ impl BrokerRuntime { break; } } - } else if snapshot_failed { - terminal_sessions.remove(&session_id); + } else if let Some((code, message)) = snapshot_failure { + end_terminal_session( + terminal_control_tx, + terminal_sessions, + terminal_snapshot_requests, + terminal_input_requests, + &session_id, + &code, + &message, + ); } } else { // Generic worker request/response dispatch. diff --git a/crates/broker/src/worker.rs b/crates/broker/src/worker.rs index 793958fe2..0dce87415 100644 --- a/crates/broker/src/worker.rs +++ b/crates/broker/src/worker.rs @@ -23,7 +23,7 @@ use serde_json::{json, Value}; use tokio::{ io::{AsyncBufReadExt, AsyncWriteExt, BufReader}, process::{Child, ChildStdin, Command}, - sync::mpsc, + sync::{mpsc, oneshot}, time::timeout, }; use uuid::Uuid; @@ -62,6 +62,23 @@ const WORKER_SPAWN_STABILITY_WINDOW: Duration = Duration::from_millis(250); /// up. Bounded so a wrapper stuck in uninterruptible sleep cannot stall the /// maintenance tick, which also drives delivery retries. const ORPHAN_REAP_TIMEOUT: Duration = Duration::from_secs(2); +const WORKER_WRITE_QUEUE_CAPACITY: usize = 128; +/// A full command queue means the worker is already backpressured. Do not +/// retain another normal request indefinitely waiting for capacity. +const WORKER_COMMAND_QUEUE_TIMEOUT: Duration = Duration::from_millis(250); +/// A PTY can transiently stop draining while it handles a large redraw or a +/// slow provider response. The sole stdin writer must still eventually fault +/// rather than wedge the worker lane, but should tolerate that short stall. +const WORKER_WRITE_TIMEOUT: Duration = Duration::from_secs(5); + +/// A complete newline-delimited worker protocol frame. A dedicated task owns +/// each worker's stdin and writes these frames in order, so cancelling a +/// caller can never cancel an in-progress pipe write and leave a partial JSON +/// frame for the next command to corrupt. +pub(crate) struct WorkerWriteCommand { + frame: Vec, + completion: Option>>, +} /// Why a worker was reaped despite its wrapper process still being alive. #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -162,7 +179,7 @@ pub(crate) struct WorkerHandle { pub(crate) parent: Option, pub(crate) workspace_id: Option, pub(crate) child: Child, - pub(crate) stdin: ChildStdin, + pub(crate) command_tx: mpsc::Sender, pub(crate) harness_pid: Option, pub(crate) spawned_at: Instant, /// When the worker reported `worker_ready`. `None` means the harness has @@ -207,6 +224,14 @@ pub(crate) enum WorkerEvent { generation: Uuid, value: Value, }, + /// The worker-owned stdin writer failed after a command was accepted. + /// The runtime must close dependent terminal sessions and terminate this + /// generation; accepting any more frames would only hide the broken pipe. + WriterFailed { + name: WorkerName, + generation: Uuid, + error: String, + }, } pub(crate) struct WorkerRegistry { @@ -219,6 +244,85 @@ pub(crate) struct WorkerRegistry { pub(crate) metrics: MetricsCollector, } +fn encode_worker_frame( + msg_type: &str, + request_id: Option, + payload: Value, +) -> Result> { + let frame = ProtocolEnvelope { + v: PROTOCOL_VERSION, + msg_type: msg_type.to_string(), + request_id, + payload, + }; + let mut encoded = serde_json::to_vec(&frame)?; + encoded.push(b'\n'); + Ok(encoded) +} + +pub(crate) fn spawn_worker_writer( + event_tx: mpsc::Sender, + name: WorkerName, + generation: Uuid, + mut stdin: ChildStdin, + mut command_rx: mpsc::Receiver, +) { + tokio::spawn(async move { + while let Some(mut command) = command_rx.recv().await { + let write_result = timeout(WORKER_WRITE_TIMEOUT, async { + stdin + .write_all(&command.frame) + .await + .context("failed writing frame to worker stdin")?; + stdin + .flush() + .await + .context("failed flushing worker stdin")?; + Ok::<(), anyhow::Error>(()) + }) + .await + .map_err(|_| { + anyhow::anyhow!( + "worker stdin write timed out after {} ms", + WORKER_WRITE_TIMEOUT.as_millis() + ) + }) + .and_then(|result| result) + .map_err(|error| error.to_string()); + + if let Some(completion) = command.completion.take() { + let _ = completion.send(write_result.clone()); + } + + let Err(error) = write_result else { + continue; + }; + + // A failed or timed-out write may have consumed part of the frame. + // Do not let another command reuse this stream. Complete queued + // waiters before reporting the fault: a full worker-event queue + // must not leave those callers blocked behind this dead writer. + while let Ok(mut queued) = command_rx.try_recv() { + if let Some(completion) = queued.completion.take() { + let _ = completion.send(Err(format!( + "worker command writer stopped after write failure: {error}" + ))); + } + } + // The runtime closes dependent terminal sessions and terminates + // this worker generation after the event is delivered. + let _ = event_tx + .send(WorkerEvent::WriterFailed { + name: name.clone(), + generation, + error, + }) + .await; + break; + } + }); +} + impl WorkerRegistry { pub(crate) fn new( event_tx: mpsc::Sender, @@ -993,6 +1097,14 @@ impl WorkerRegistry { false, log_file, ); + let (command_tx, command_rx) = mpsc::channel(WORKER_WRITE_QUEUE_CAPACITY); + spawn_worker_writer( + self.event_tx.clone(), + spec.name.clone(), + generation, + stdin, + command_rx, + ); let handle = WorkerHandle { generation, @@ -1000,7 +1112,7 @@ impl WorkerRegistry { parent, workspace_id, child, - stdin, + command_tx, harness_pid: initial_harness_pid, spawned_at: Instant::now(), ready_at: None, @@ -1074,38 +1186,99 @@ impl WorkerRegistry { request_id: Option, payload: Value, ) -> Result<()> { - let handle = self + let command_tx = self .workers - .get_mut(name) - .with_context(|| format!("unknown worker '{name}'"))?; - - let frame = ProtocolEnvelope { - v: PROTOCOL_VERSION, - msg_type: msg_type.to_string(), - request_id, - payload, - }; - - let encoded = serde_json::to_string(&frame)?; - handle - .stdin - .write_all(encoded.as_bytes()) + .get(name) + .with_context(|| format!("unknown worker '{name}'"))? + .command_tx + .clone(); + let frame = encode_worker_frame(msg_type, request_id, payload)?; + let (completion_tx, completion_rx) = oneshot::channel(); + timeout( + WORKER_COMMAND_QUEUE_TIMEOUT, + command_tx.send(WorkerWriteCommand { + frame, + completion: Some(completion_tx), + }), + ) + .await + .map_err(|_| anyhow::anyhow!("worker command queue timed out for '{name}'"))? + .map_err(|_| anyhow::anyhow!("worker command writer is unavailable for '{name}'")) + .with_context(|| format!("failed writing frame to worker '{name}'"))?; + // Once a command enters the writer queue, do not return a timeout + // before that sole writer resolves it. Reporting an accepted PTY + // write as failed while it can still be emitted would invite callers + // to retry and duplicate terminal input. The writer itself has the + // finite [`WORKER_WRITE_TIMEOUT`] and drains every queued completion + // with an error if it faults. + completion_rx .await + .map_err(|_| { + anyhow::anyhow!("worker command writer stopped before completing '{name}'") + }) + .with_context(|| format!("failed writing frame to worker '{name}'"))? + .map_err(anyhow::Error::msg) .with_context(|| format!("failed writing frame to worker '{name}'"))?; - handle - .stdin - .write_all(b"\n") - .await - .with_context(|| format!("failed writing newline to worker '{name}'"))?; - handle - .stdin - .flush() - .await - .with_context(|| format!("failed flushing worker '{name}' stdin"))?; Ok(()) } + /// Queue an already-framed raw PTY command through the same sole stdin + /// writer used for protocol frames. This completes once the command has + /// been admitted to the writer queue, rather than after the pipe write. + /// That keeps administrative PTY actions such as `/model` serialized with + /// protocol traffic without ever reporting an admitted command as failed + /// while the writer can still emit it. + pub(crate) async fn send_raw_to_worker(&self, name: &str, frame: Vec) -> Result<()> { + let command_tx = self + .workers + .get(name) + .with_context(|| format!("unknown worker '{name}'"))? + .command_tx + .clone(); + command_tx + .send(WorkerWriteCommand { + frame, + completion: None, + }) + .await + .map_err(|_| anyhow::anyhow!("worker command writer is unavailable for '{name}'"))?; + Ok(()) + } + + /// Enqueue a complete worker frame without awaiting its pipe write. This + /// is used by terminal attach traffic, which must remain responsive when a + /// PTY stops draining stdin. The dedicated writer owns the actual write; + /// a later write failure is reported as [`WorkerEvent::WriterFailed`]. + pub(crate) fn try_send_to_worker( + &self, + name: &str, + msg_type: &str, + request_id: Option, + payload: Value, + ) -> Result<()> { + let command_tx = self + .workers + .get(name) + .with_context(|| format!("unknown worker '{name}'"))? + .command_tx + .clone(); + let frame = encode_worker_frame(msg_type, request_id, payload)?; + command_tx + .try_send(WorkerWriteCommand { + frame, + completion: None, + }) + .map_err(|error| match error { + mpsc::error::TrySendError::Full(_) => { + anyhow::anyhow!("worker command queue is full for '{name}'") + } + mpsc::error::TrySendError::Closed(_) => { + anyhow::anyhow!("worker command writer is unavailable for '{name}'") + } + }) + } + pub(crate) async fn deliver(&mut self, name: &str, delivery: RelayDelivery) -> Result<()> { tracing::debug!( target = "broker::deliver", @@ -1139,10 +1312,22 @@ impl WorkerRegistry { request_id: None, payload: json!({"reason":"release","grace_ms": release_grace.as_millis() as u64}), }; - let encoded = serde_json::to_string(&shutdown_frame)?; - let _ = handle.stdin.write_all(encoded.as_bytes()).await; - let _ = handle.stdin.write_all(b"\n").await; - let _ = handle.stdin.flush().await; + let encoded = serde_json::to_vec(&shutdown_frame)?; + let mut frame = encoded; + frame.push(b'\n'); + let (completion_tx, completion_rx) = oneshot::channel(); + if handle + .command_tx + .try_send(WorkerWriteCommand { + frame, + completion: Some(completion_tx), + }) + .is_ok() + { + // Cancelling this wait cannot cancel the worker-owned write; it + // only bounds release before the normal process termination path. + let _ = timeout(WORKER_WRITE_TIMEOUT, completion_rx).await; + } let result = terminate_child(&mut handle.child, release_grace).await; match &result { @@ -1154,6 +1339,25 @@ impl WorkerRegistry { result } + /// Stop a worker whose sole stdin writer has failed, while leaving its + /// registry and supervisor entries intact. The normal reap path observes + /// the exit and applies the configured restart policy; unlike `release`, + /// this is an unexpected fault rather than an intentional teardown. + pub(crate) fn terminate_after_writer_failure(&mut self, name: &str) -> Result<()> { + let handle = self + .workers + .get_mut(name) + .with_context(|| format!("unknown worker '{name}'"))?; + handle.exit_reason = Some("worker_write_failed".into()); + if handle.child.id().is_none() { + return Ok(()); + } + handle + .child + .start_kill() + .with_context(|| format!("failed to terminate worker '{name}' after writer failure")) + } + pub(crate) async fn shutdown_all(&mut self) -> Result<()> { let names: Vec = self.workers.keys().cloned().collect(); for name in names { @@ -2212,6 +2416,15 @@ mod tests { .unwrap(); let pid = child.id().expect("child has a pid"); let stdin = child.stdin.take().expect("piped stdin"); + let generation = Uuid::new_v4(); + let (command_tx, command_rx) = mpsc::channel(WORKER_WRITE_QUEUE_CAPACITY); + spawn_worker_writer( + reg.event_tx.clone(), + WorkerName::from(name), + generation, + stdin, + command_rx, + ); assert!( is_process_alive(pid), "precondition: child must start alive" @@ -2220,12 +2433,12 @@ mod tests { reg.workers.insert( WorkerName::from(name), WorkerHandle { - generation: Uuid::new_v4(), + generation, spec: spec_for_test(name), parent: None, workspace_id: None, child, - stdin, + command_tx, harness_pid: None, spawned_at: Instant::now(), ready_at: None, @@ -2260,16 +2473,25 @@ mod tests { let mut child = Command::new("true").stdin(Stdio::piped()).spawn().unwrap(); let stdin = child.stdin.take().expect("piped stdin"); child.wait().await.expect("child exits immediately"); + let generation = Uuid::new_v4(); + let (command_tx, command_rx) = mpsc::channel(WORKER_WRITE_QUEUE_CAPACITY); + spawn_worker_writer( + reg.event_tx.clone(), + WorkerName::from(name), + generation, + stdin, + command_rx, + ); reg.workers.insert( WorkerName::from(name), WorkerHandle { - generation: Uuid::new_v4(), + generation, spec: spec_for_test(name), parent: None, workspace_id: None, child, - stdin, + command_tx, harness_pid: None, spawned_at: Instant::now(), ready_at: None, @@ -2294,6 +2516,134 @@ mod tests { assert!(!reg.workers.contains_key(&WorkerName::from(name))); } + #[cfg(unix)] + #[tokio::test] + async fn terminal_enqueue_serializes_complete_frames_through_one_writer() { + let (event_tx, _event_rx) = mpsc::channel::(16); + let mut reg = WorkerRegistry::new( + event_tx.clone(), + Vec::new(), + PathBuf::from("/tmp/worker-tests"), + Instant::now(), + ); + let name = "writer-serialization"; + let mut child = Command::new("cat") + .stdin(Stdio::piped()) + .stdout(Stdio::piped()) + .stderr(Stdio::null()) + .spawn() + .unwrap(); + let stdin = child.stdin.take().expect("piped stdin"); + let stdout = child.stdout.take().expect("piped stdout"); + let generation = Uuid::new_v4(); + let (command_tx, command_rx) = mpsc::channel(WORKER_WRITE_QUEUE_CAPACITY); + spawn_worker_writer( + event_tx, + WorkerName::from(name), + generation, + stdin, + command_rx, + ); + reg.workers.insert( + WorkerName::from(name), + WorkerHandle { + generation, + spec: spec_for_test(name), + parent: None, + workspace_id: None, + child, + command_tx, + harness_pid: None, + spawned_at: Instant::now(), + ready_at: None, + last_activity_at: Instant::now(), + context_budget_pct: None, + state: AgentWorkState::Working, + exit_reason: None, + }, + ); + + reg.try_send_to_worker(name, "snapshot_pty", None, json!({ "format": "ansi" })) + .unwrap(); + reg.try_send_to_worker(name, "resize_pty", None, json!({ "rows": 24, "cols": 80 })) + .unwrap(); + + let mut lines = BufReader::new(stdout).lines(); + for expected_type in ["snapshot_pty", "resize_pty"] { + let line = timeout(Duration::from_secs(1), lines.next_line()) + .await + .expect("writer should not block") + .expect("cat stdout should remain open") + .expect("writer should emit a complete newline-delimited frame"); + let frame: Value = serde_json::from_str(&line) + .expect("each serialized worker command must remain valid JSON"); + assert_eq!( + frame.get("type").and_then(Value::as_str), + Some(expected_type) + ); + } + + let handle = reg + .workers + .get_mut(name) + .expect("test worker remains registered"); + terminate_child(&mut handle.child, Duration::from_millis(200)) + .await + .unwrap(); + } + + #[cfg(unix)] + #[tokio::test] + async fn raw_command_returns_after_queue_admission_without_a_writer() { + // `/model` is best-effort. Its API response must mean the command was + // accepted by the worker-owned queue, not that the eventual pipe write + // completed: the latter can stall while the admitted command remains + // eligible to be emitted. + let mut reg = make_registry(vec![]); + let name = "raw-command-admission"; + let child = Command::new("sleep").arg("30").spawn().unwrap(); + let generation = Uuid::new_v4(); + let (command_tx, mut command_rx) = mpsc::channel(1); + reg.workers.insert( + WorkerName::from(name), + WorkerHandle { + generation, + spec: spec_for_test(name), + parent: None, + workspace_id: None, + child, + command_tx, + harness_pid: None, + spawned_at: Instant::now(), + ready_at: None, + last_activity_at: Instant::now(), + context_budget_pct: None, + state: AgentWorkState::Working, + exit_reason: None, + }, + ); + + timeout( + Duration::from_millis(50), + reg.send_raw_to_worker(name, b"/model sonnet\n".to_vec()), + ) + .await + .expect("queue admission must not wait for a pipe writer") + .expect("open worker queue accepts the raw command"); + + let queued = command_rx.recv().await.expect("raw command was queued"); + assert_eq!(queued.frame, b"/model sonnet\n"); + assert!(queued.completion.is_none()); + + let handle = reg + .workers + .get_mut(name) + .expect("test worker remains registered"); + terminate_child(&mut handle.child, Duration::from_millis(200)) + .await + .unwrap(); + } + // The wrapper process can outlive the harness it hosts, so reaping on the // wrapper alone leaves a dead agent listed as `working` forever. mod orphaned_worker {