From 5f29b553b3c08bae5174b37b1d89797865981589 Mon Sep 17 00:00:00 2001 From: Miya Date: Tue, 11 Aug 2026 23:37:13 +0200 Subject: [PATCH 1/6] fix(broker): serialize terminal worker writes --- crates/broker/src/runtime/api.rs | 35 +-- crates/broker/src/runtime/fleet.rs | 233 ++++++++------- crates/broker/src/runtime/tests.rs | 13 +- crates/broker/src/runtime/worker_events.rs | 52 +++- crates/broker/src/worker.rs | 330 ++++++++++++++++++--- 5 files changed, 496 insertions(+), 167 deletions(-) diff --git a/crates/broker/src/runtime/api.rs b/crates/broker/src/runtime/api.rs index cde595106..8c644c3a6 100644 --- a/crates/broker/src/runtime/api.rs +++ b/crates/broker/src/runtime/api.rs @@ -736,35 +736,22 @@ 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>(()) + let result = workers + .send_raw_to_worker(&name, model_command.into_bytes()) + .await; + if let Some(timeout_ms) = timeout_ms { + tracing::info!( + name = %name, + timeout_ms, + "HTTP API set_model timeout_ms is currently advisory only" + ); } - .await; match result { Ok(()) => { diff --git a/crates/broker/src/runtime/fleet.rs b/crates/broker/src/runtime/fleet.rs index bf74c6e99..1d9d182b1 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,44 @@ pub(super) fn try_send_terminal( .is_ok() } +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, @@ -153,44 +187,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 +287,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,26 +332,13 @@ 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()) - } - Err(_) => self.fail_terminal_session( - session_id, - "resize_timeout", - "terminal resize write timed out".into(), - ), + 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()); } } TerminalControlEvent::Message(TerminalFromCloud::Close { session_id }) => { @@ -399,21 +391,15 @@ 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 { + 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 +1687,47 @@ 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 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/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..9197c8ebf 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,55 @@ 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_map(|(session_id, session)| { + (session.agent == name).then(|| 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.release(name.as_str()).await { + tracing::warn!( + target = "relay_broker::terminal", + worker = %name, + error = %release_error, + "failed to terminate worker after command writer failure" + ); + } + } WorkerEvent::Message { name, generation, diff --git a/crates/broker/src/worker.rs b/crates/broker/src/worker.rs index 793958fe2..36c5ef3ad 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,17 @@ 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; +const WORKER_SHUTDOWN_WRITE_TIMEOUT: Duration = Duration::from_millis(250); + +/// 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 +173,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 +218,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 +238,76 @@ 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 = 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(|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 write may have consumed part of the frame. Do not let + // another command reuse this stream: notify the runtime so it can + // close the dependent terminal sessions and terminate the worker. + let _ = event_tx + .send(WorkerEvent::WriterFailed { + name: name.clone(), + generation, + error: error.clone(), + }) + .await; + + 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}" + ))); + } + } + break; + } + }); +} + impl WorkerRegistry { pub(crate) fn new( event_tx: mpsc::Sender, @@ -993,6 +1082,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 +1097,7 @@ impl WorkerRegistry { parent, workspace_id, child, - stdin, + command_tx, harness_pid: initial_harness_pid, spawned_at: Instant::now(), ready_at: None, @@ -1074,38 +1171,95 @@ 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(); + command_tx + .send(WorkerWriteCommand { + frame, + completion: Some(completion_tx), + }) .await + .map_err(|_| anyhow::anyhow!("worker command writer is unavailable for '{name}'")) .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() + completion_rx .await - .with_context(|| format!("failed flushing worker '{name}' stdin"))?; + .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}'"))?; Ok(()) } + /// Queue an already-framed raw PTY command through the same sole stdin + /// writer used for protocol frames. This keeps administrative PTY actions + /// such as `/model` from interleaving with JSON protocol traffic. + 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(); + let (completion_tx, completion_rx) = oneshot::channel(); + command_tx + .send(WorkerWriteCommand { + frame, + completion: Some(completion_tx), + }) + .await + .map_err(|_| anyhow::anyhow!("worker command writer is unavailable for '{name}'"))?; + completion_rx + .await + .map_err(|_| { + anyhow::anyhow!("worker command writer stopped before completing '{name}'") + })? + .map_err(anyhow::Error::msg) + .with_context(|| format!("failed writing raw command to worker '{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 +1293,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_SHUTDOWN_WRITE_TIMEOUT, completion_rx).await; + } let result = terminate_child(&mut handle.child, release_grace).await; match &result { @@ -2212,6 +2378,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 +2395,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 +2435,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 +2478,82 @@ 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(); + } + // 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 { From 7756b9f2d0f4f75eac7f2adf0a0c4351202b62d3 Mon Sep 17 00:00:00 2001 From: Miya Date: Tue, 11 Aug 2026 23:54:32 +0200 Subject: [PATCH 2/6] fix(broker): satisfy current clippy terminal filter lint --- crates/broker/src/runtime/worker_events.rs | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/crates/broker/src/runtime/worker_events.rs b/crates/broker/src/runtime/worker_events.rs index 9197c8ebf..f895959fb 100644 --- a/crates/broker/src/runtime/worker_events.rs +++ b/crates/broker/src/runtime/worker_events.rs @@ -558,9 +558,8 @@ impl BrokerRuntime { ); let session_ids: Vec = terminal_sessions .iter() - .filter_map(|(session_id, session)| { - (session.agent == name).then(|| session_id.clone()) - }) + .filter(|(_, session)| session.agent == name) + .map(|(session_id, _)| session_id.clone()) .collect(); for session_id in session_ids { fail_terminal_session( From 11828325a8e232a210f934ebd532834ec50ba4ad Mon Sep 17 00:00:00 2001 From: Miya Date: Wed, 12 Aug 2026 00:08:25 +0200 Subject: [PATCH 3/6] fix(broker): bound worker writer fault recovery --- crates/broker/src/runtime/worker_events.rs | 23 ++++-- crates/broker/src/worker.rs | 92 +++++++++++++++------- 2 files changed, 82 insertions(+), 33 deletions(-) diff --git a/crates/broker/src/runtime/worker_events.rs b/crates/broker/src/runtime/worker_events.rs index f895959fb..56cfc42f3 100644 --- a/crates/broker/src/runtime/worker_events.rs +++ b/crates/broker/src/runtime/worker_events.rs @@ -572,12 +572,12 @@ impl BrokerRuntime { format!("worker command writer failed: {error}"), ); } - if let Err(release_error) = workers.release(name.as_str()).await { + 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 terminate worker after command writer failure" + "failed to signal worker after command writer failure" ); } } @@ -1068,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( @@ -1111,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 36c5ef3ad..666fd1d88 100644 --- a/crates/broker/src/worker.rs +++ b/crates/broker/src/worker.rs @@ -63,7 +63,7 @@ const WORKER_SPAWN_STABILITY_WINDOW: Duration = Duration::from_millis(250); /// maintenance tick, which also drives delivery retries. const ORPHAN_REAP_TIMEOUT: Duration = Duration::from_secs(2); const WORKER_WRITE_QUEUE_CAPACITY: usize = 128; -const WORKER_SHUTDOWN_WRITE_TIMEOUT: Duration = Duration::from_millis(250); +const WORKER_COMMAND_TIMEOUT: Duration = Duration::from_millis(250); /// A complete newline-delimited worker protocol frame. A dedicated task owns /// each worker's stdin and writes these frames in order, so cancelling a @@ -263,7 +263,7 @@ pub(crate) fn spawn_worker_writer( ) { tokio::spawn(async move { while let Some(mut command) = command_rx.recv().await { - let write_result = async { + let write_result = timeout(WORKER_COMMAND_TIMEOUT, async { stdin .write_all(&command.frame) .await @@ -273,8 +273,15 @@ pub(crate) fn spawn_worker_writer( .await .context("failed flushing worker stdin")?; Ok::<(), anyhow::Error>(()) - } + }) .await + .map_err(|_| { + anyhow::anyhow!( + "worker stdin write timed out after {} ms", + WORKER_COMMAND_TIMEOUT.as_millis() + ) + }) + .and_then(|result| result) .map_err(|error| error.to_string()); if let Some(completion) = command.completion.take() { @@ -285,17 +292,10 @@ pub(crate) fn spawn_worker_writer( continue; }; - // A failed write may have consumed part of the frame. Do not let - // another command reuse this stream: notify the runtime so it can - // close the dependent terminal sessions and terminate the worker. - let _ = event_tx - .send(WorkerEvent::WriterFailed { - name: name.clone(), - generation, - error: error.clone(), - }) - .await; - + // 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!( @@ -303,6 +303,15 @@ pub(crate) fn spawn_worker_writer( ))); } } + // 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; } }); @@ -1179,16 +1188,20 @@ impl WorkerRegistry { .clone(); let frame = encode_worker_frame(msg_type, request_id, payload)?; let (completion_tx, completion_rx) = oneshot::channel(); - command_tx - .send(WorkerWriteCommand { + timeout( + WORKER_COMMAND_TIMEOUT, + command_tx.send(WorkerWriteCommand { frame, completion: Some(completion_tx), - }) - .await - .map_err(|_| anyhow::anyhow!("worker command writer is unavailable for '{name}'")) - .with_context(|| format!("failed writing frame to worker '{name}'"))?; - completion_rx + }), + ) + .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}'"))?; + timeout(WORKER_COMMAND_TIMEOUT, completion_rx) .await + .map_err(|_| anyhow::anyhow!("worker command writer timed out for '{name}'"))? .map_err(|_| { anyhow::anyhow!("worker command writer stopped before completing '{name}'") }) @@ -1210,15 +1223,19 @@ impl WorkerRegistry { .command_tx .clone(); let (completion_tx, completion_rx) = oneshot::channel(); - command_tx - .send(WorkerWriteCommand { + timeout( + WORKER_COMMAND_TIMEOUT, + command_tx.send(WorkerWriteCommand { frame, completion: Some(completion_tx), - }) - .await - .map_err(|_| anyhow::anyhow!("worker command writer is unavailable for '{name}'"))?; - completion_rx + }), + ) + .await + .map_err(|_| anyhow::anyhow!("worker command queue timed out for '{name}'"))? + .map_err(|_| anyhow::anyhow!("worker command writer is unavailable for '{name}'"))?; + timeout(WORKER_COMMAND_TIMEOUT, completion_rx) .await + .map_err(|_| anyhow::anyhow!("worker command writer timed out for '{name}'"))? .map_err(|_| { anyhow::anyhow!("worker command writer stopped before completing '{name}'") })? @@ -1307,7 +1324,7 @@ impl WorkerRegistry { { // Cancelling this wait cannot cancel the worker-owned write; it // only bounds release before the normal process termination path. - let _ = timeout(WORKER_SHUTDOWN_WRITE_TIMEOUT, completion_rx).await; + let _ = timeout(WORKER_COMMAND_TIMEOUT, completion_rx).await; } let result = terminate_child(&mut handle.child, release_grace).await; @@ -1320,6 +1337,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 { From 719c2ba7f91a61a1ec2cdf4591ecaaff3b8a1572 Mon Sep 17 00:00:00 2001 From: Miya Date: Wed, 12 Aug 2026 00:24:53 +0200 Subject: [PATCH 4/6] fix(broker): preserve accepted worker writes --- crates/broker/src/worker.rs | 30 ++++++++++++++++++++---------- 1 file changed, 20 insertions(+), 10 deletions(-) diff --git a/crates/broker/src/worker.rs b/crates/broker/src/worker.rs index 666fd1d88..1d972520e 100644 --- a/crates/broker/src/worker.rs +++ b/crates/broker/src/worker.rs @@ -63,7 +63,13 @@ const WORKER_SPAWN_STABILITY_WINDOW: Duration = Duration::from_millis(250); /// maintenance tick, which also drives delivery retries. const ORPHAN_REAP_TIMEOUT: Duration = Duration::from_secs(2); const WORKER_WRITE_QUEUE_CAPACITY: usize = 128; -const WORKER_COMMAND_TIMEOUT: Duration = Duration::from_millis(250); +/// 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 @@ -263,7 +269,7 @@ pub(crate) fn spawn_worker_writer( ) { tokio::spawn(async move { while let Some(mut command) = command_rx.recv().await { - let write_result = timeout(WORKER_COMMAND_TIMEOUT, async { + let write_result = timeout(WORKER_WRITE_TIMEOUT, async { stdin .write_all(&command.frame) .await @@ -278,7 +284,7 @@ pub(crate) fn spawn_worker_writer( .map_err(|_| { anyhow::anyhow!( "worker stdin write timed out after {} ms", - WORKER_COMMAND_TIMEOUT.as_millis() + WORKER_WRITE_TIMEOUT.as_millis() ) }) .and_then(|result| result) @@ -1189,7 +1195,7 @@ impl WorkerRegistry { let frame = encode_worker_frame(msg_type, request_id, payload)?; let (completion_tx, completion_rx) = oneshot::channel(); timeout( - WORKER_COMMAND_TIMEOUT, + WORKER_COMMAND_QUEUE_TIMEOUT, command_tx.send(WorkerWriteCommand { frame, completion: Some(completion_tx), @@ -1199,9 +1205,14 @@ impl WorkerRegistry { .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}'"))?; - timeout(WORKER_COMMAND_TIMEOUT, completion_rx) + // 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 timed out for '{name}'"))? .map_err(|_| { anyhow::anyhow!("worker command writer stopped before completing '{name}'") }) @@ -1224,7 +1235,7 @@ impl WorkerRegistry { .clone(); let (completion_tx, completion_rx) = oneshot::channel(); timeout( - WORKER_COMMAND_TIMEOUT, + WORKER_COMMAND_QUEUE_TIMEOUT, command_tx.send(WorkerWriteCommand { frame, completion: Some(completion_tx), @@ -1233,9 +1244,8 @@ impl WorkerRegistry { .await .map_err(|_| anyhow::anyhow!("worker command queue timed out for '{name}'"))? .map_err(|_| anyhow::anyhow!("worker command writer is unavailable for '{name}'"))?; - timeout(WORKER_COMMAND_TIMEOUT, completion_rx) + completion_rx .await - .map_err(|_| anyhow::anyhow!("worker command writer timed out for '{name}'"))? .map_err(|_| { anyhow::anyhow!("worker command writer stopped before completing '{name}'") })? @@ -1324,7 +1334,7 @@ impl WorkerRegistry { { // Cancelling this wait cannot cancel the worker-owned write; it // only bounds release before the normal process termination path. - let _ = timeout(WORKER_COMMAND_TIMEOUT, completion_rx).await; + let _ = timeout(WORKER_WRITE_TIMEOUT, completion_rx).await; } let result = terminate_child(&mut handle.child, release_grace).await; From 7fd9f514dbfe58b79156341409b11b8c699bdce9 Mon Sep 17 00:00:00 2001 From: Miya Date: Wed, 12 Aug 2026 10:07:33 +0200 Subject: [PATCH 5/6] fix(broker): bound model writes and share resize leases --- crates/broker/src/runtime/api.rs | 55 +++++++-- crates/broker/src/runtime/fleet.rs | 138 ++++++++++++++++++++++- crates/broker/src/runtime/maintenance.rs | 4 +- 3 files changed, 180 insertions(+), 17 deletions(-) diff --git a/crates/broker/src/runtime/api.rs b/crates/broker/src/runtime/api.rs index 8c644c3a6..e540eb20e 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 @@ -742,16 +752,22 @@ impl BrokerRuntime { } let model_command = format!("/model {}\n", model); - let result = workers - .send_raw_to_worker(&name, model_command.into_bytes()) - .await; - if let Some(timeout_ms) = timeout_ms { - tracing::info!( - name = %name, - timeout_ms, - "HTTP API set_model timeout_ms is currently advisory only" - ); - } + let set_model_timeout = set_model_write_timeout(timeout_ms); + // The worker-owned writer keeps the accepted command serialized + // if this wait expires, but the single runtime actor must not + // remain blocked behind a stalled worker stdin pipe. + 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(()) => { @@ -2580,6 +2596,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 1d9d182b1..b46a07a99 100644 --- a/crates/broker/src/runtime/fleet.rs +++ b/crates/broker/src/runtime/fleet.rs @@ -38,6 +38,29 @@ 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, + terminal_sessions: &HashMap, + session_id: &str, +) { + let Some(agent) = terminal_sessions + .get(session_id) + .map(|session| session.agent.clone()) + else { + return; + }; + 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, @@ -150,6 +173,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(); @@ -332,16 +357,64 @@ impl BrokerRuntime { }); return; } - if let Err(error) = self.workers.try_send_to_worker( - session.agent.as_str(), - "resize_pty", - None, - json!({ "rows": rows, "cols": cols }), + // 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(), ) { - self.fail_terminal_session(session_id, "resize_failed", error.to_string()); + 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(), + ); + } + } } } TerminalControlEvent::Message(TerminalFromCloud::Close { session_id }) => { + release_terminal_resize_ownership( + &mut self.resize_owners, + &self.terminal_sessions, + &session_id, + ); self.terminal_sessions.remove(&session_id); self.terminal_snapshot_requests .retain(|_, pending| pending.session_id != session_id); @@ -370,6 +443,11 @@ 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"); + release_terminal_resize_ownership( + &mut self.resize_owners, + &self.terminal_sessions, + &session_id, + ); self.terminal_sessions.remove(&session_id); self.terminal_snapshot_requests .retain(|_, pending| pending.session_id != session_id); @@ -391,6 +469,11 @@ impl BrokerRuntime { } fn fail_terminal_session(&mut self, session_id: String, code: &str, message: String) { + release_terminal_resize_ownership( + &mut self.resize_owners, + &self.terminal_sessions, + &session_id, + ); fail_terminal_session( &self.terminal_control_tx, &mut self.terminal_sessions, @@ -1728,6 +1811,49 @@ mod tests { 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, + }, + ), + ]); + + release_terminal_resize_ownership(&mut resize_owners, &terminal_sessions, "session-a"); + terminal_sessions.remove("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..ef73f1d44 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,6 +47,7 @@ impl BrokerRuntime { .collect(); for (request_id, session_id) in expired_terminal_snapshots { terminal_snapshot_requests.remove(&request_id); + release_terminal_resize_ownership(resize_owners, terminal_sessions, &session_id); if terminal_sessions.remove(&session_id).is_some() { terminal_input_requests.retain(|_, pending| pending.session_id != session_id); if !try_send_terminal( @@ -81,6 +82,7 @@ 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. + release_terminal_resize_ownership(resize_owners, terminal_sessions, &session_id); 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); From bce0f53aca6ef1bc60501fb08daa9a8404ed2761 Mon Sep 17 00:00:00 2001 From: Miya Date: Wed, 12 Aug 2026 10:41:13 +0200 Subject: [PATCH 6/6] fix(broker): acknowledge admitted model commands --- crates/broker/src/runtime/api.rs | 10 ++- crates/broker/src/runtime/fleet.rs | 52 +++++++-------- crates/broker/src/runtime/maintenance.rs | 8 +-- crates/broker/src/worker.rs | 80 ++++++++++++++++++------ 4 files changed, 97 insertions(+), 53 deletions(-) diff --git a/crates/broker/src/runtime/api.rs b/crates/broker/src/runtime/api.rs index e540eb20e..3cd3ff6fd 100644 --- a/crates/broker/src/runtime/api.rs +++ b/crates/broker/src/runtime/api.rs @@ -753,9 +753,11 @@ impl BrokerRuntime { let model_command = format!("/model {}\n", model); let set_model_timeout = set_model_write_timeout(timeout_ms); - // The worker-owned writer keeps the accepted command serialized - // if this wait expires, but the single runtime actor must not - // remain blocked behind a stalled worker stdin pipe. + // `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()), @@ -775,6 +777,8 @@ impl BrokerRuntime { "name": name, "model": model, "success": true, + "accepted": true, + "pending": true, }))); } Err(error) => { diff --git a/crates/broker/src/runtime/fleet.rs b/crates/broker/src/runtime/fleet.rs index b46a07a99..55ac71345 100644 --- a/crates/broker/src/runtime/fleet.rs +++ b/crates/broker/src/runtime/fleet.rs @@ -44,20 +44,14 @@ pub(super) fn try_send_terminal( /// for single-resizer ownership. pub(super) fn release_terminal_resize_ownership( resize_owners: &mut HashMap, - terminal_sessions: &HashMap, + agent: &WorkerName, session_id: &str, ) { - let Some(agent) = terminal_sessions - .get(session_id) - .map(|session| session.agent.clone()) - else { - return; - }; if resize_owners - .get(&agent) + .get(agent) .is_some_and(|owner| owner.session_id == session_id) { - resize_owners.remove(&agent); + resize_owners.remove(agent); } } @@ -410,12 +404,13 @@ impl BrokerRuntime { } } TerminalControlEvent::Message(TerminalFromCloud::Close { session_id }) => { - release_terminal_resize_ownership( - &mut self.resize_owners, - &self.terminal_sessions, - &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 @@ -443,12 +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"); - release_terminal_resize_ownership( - &mut self.resize_owners, - &self.terminal_sessions, - &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 @@ -469,11 +465,9 @@ impl BrokerRuntime { } fn fail_terminal_session(&mut self, session_id: String, code: &str, message: String) { - release_terminal_resize_ownership( - &mut self.resize_owners, - &self.terminal_sessions, - &session_id, - ); + 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, @@ -1847,8 +1841,10 @@ mod tests { ), ]); - release_terminal_resize_ownership(&mut resize_owners, &terminal_sessions, "session-a"); - terminal_sessions.remove("session-a"); + 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)); diff --git a/crates/broker/src/runtime/maintenance.rs b/crates/broker/src/runtime/maintenance.rs index ef73f1d44..356bdf137 100644 --- a/crates/broker/src/runtime/maintenance.rs +++ b/crates/broker/src/runtime/maintenance.rs @@ -47,8 +47,8 @@ impl BrokerRuntime { .collect(); for (request_id, session_id) in expired_terminal_snapshots { terminal_snapshot_requests.remove(&request_id); - release_terminal_resize_ownership(resize_owners, terminal_sessions, &session_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, @@ -82,8 +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. - release_terminal_resize_ownership(resize_owners, terminal_sessions, &session_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_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/worker.rs b/crates/broker/src/worker.rs index 1d972520e..0dce87415 100644 --- a/crates/broker/src/worker.rs +++ b/crates/broker/src/worker.rs @@ -1224,8 +1224,11 @@ impl WorkerRegistry { } /// Queue an already-framed raw PTY command through the same sole stdin - /// writer used for protocol frames. This keeps administrative PTY actions - /// such as `/model` from interleaving with JSON protocol traffic. + /// 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 @@ -1233,24 +1236,13 @@ impl WorkerRegistry { .with_context(|| format!("unknown worker '{name}'"))? .command_tx .clone(); - let (completion_tx, completion_rx) = oneshot::channel(); - timeout( - WORKER_COMMAND_QUEUE_TIMEOUT, - command_tx.send(WorkerWriteCommand { + 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}'"))?; - completion_rx + completion: None, + }) .await - .map_err(|_| { - anyhow::anyhow!("worker command writer stopped before completing '{name}'") - })? - .map_err(anyhow::Error::msg) - .with_context(|| format!("failed writing raw command to worker '{name}'"))?; + .map_err(|_| anyhow::anyhow!("worker command writer is unavailable for '{name}'"))?; Ok(()) } @@ -2600,6 +2592,58 @@ mod tests { .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 {