diff --git a/src/sdk/src/control_socket/server/hub_ops.rs b/src/sdk/src/control_socket/server/hub_ops.rs index 8fe37ac69..569f54d4c 100644 --- a/src/sdk/src/control_socket/server/hub_ops.rs +++ b/src/sdk/src/control_socket/server/hub_ops.rs @@ -188,7 +188,22 @@ impl FleetOps for HubFleetOps { ); let tee = tee_status(status); - let outcome = handle.task_runner().run(request, Some(tee)).await; + let runner = handle.task_runner(); + let (capabilities, abort) = runner + .capabilities_for_dispatch(&request.worker_address, &request.abort_id) + .await; + let screen_kill = match capabilities { + Ok(capabilities) => capabilities.screen_kill, + Err(RunError::Aborted) => { + let outcome = Err(RunError::Aborted); + record_outcome(&activity, &wire_task_id, &outcome); + return outcome; + } + Err(_) => false, + }; + let outcome = runner + .run_negotiated(request, Some(tee), screen_kill, Some(abort), Some(task_id)) + .await; record_outcome(&activity, &wire_task_id, &outcome); outcome } diff --git a/src/sdk/src/daemon/capabilities/mod.rs b/src/sdk/src/daemon/capabilities/mod.rs index 6ed00138e..bce57a1e3 100644 --- a/src/sdk/src/daemon/capabilities/mod.rs +++ b/src/sdk/src/daemon/capabilities/mod.rs @@ -62,6 +62,8 @@ pub async fn probe_capabilities(options: ProbeOptions) -> AgentCapabilities { // Filled in by the daemon from its own workflow store: this probe // describes the harness, not what has been authored on top of it. workflows: Vec::new(), + // Enabled by the embedding worker only when it installs a ScreenRouter. + screen_kill: false, // Deterministic digest of CLAUDE.md/AGENTS.md/README.md — the summary // of last resort so a failed probe still carries project context. summary: dir.fallback_summary.clone(), diff --git a/src/sdk/src/daemon/providers/types.rs b/src/sdk/src/daemon/providers/types.rs index a2899c773..bb1f2bb7a 100644 --- a/src/sdk/src/daemon/providers/types.rs +++ b/src/sdk/src/daemon/providers/types.rs @@ -40,6 +40,7 @@ pub type ExistsOnPath = Box bool + Send + Sync>; #[derive(Clone, Default)] pub struct Abort { flag: Arc, + terminate: Arc, notify: Arc, } @@ -55,6 +56,17 @@ impl Abort { self.notify.notify_waiters(); } + /// Signal cancellation that must also terminate the serving harness. + pub fn terminate(&self) { + self.terminate.store(true, Ordering::SeqCst); + self.abort(); + } + + /// Whether cancellation requested termination of the serving harness. + pub fn is_terminated(&self) -> bool { + self.terminate.load(Ordering::SeqCst) + } + /// Whether cancellation has been signalled. pub fn is_aborted(&self) -> bool { self.flag.load(Ordering::SeqCst) diff --git a/src/sdk/src/daemon/runtime.rs b/src/sdk/src/daemon/runtime.rs index 3909d1021..c326d98cc 100644 --- a/src/sdk/src/daemon/runtime.rs +++ b/src/sdk/src/daemon/runtime.rs @@ -4,7 +4,7 @@ //! machine lives in [`super::task_loop`]. use std::collections::HashMap; -use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering}; +use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering}; use std::sync::{Arc, Mutex as StdMutex}; use tokio::sync::{Mutex as TokioMutex, Notify, Semaphore}; @@ -20,6 +20,13 @@ use super::types::{ }; impl DaemonRuntime { + /// Advertise screen termination support for an embedding that installs the + /// authenticated screen-message router. + pub fn enable_screen_kill(&self) { + self.inner + .screen_kill + .store(true, std::sync::atomic::Ordering::Relaxed); + } /// Build a runtime from `config`, an executor (`run_task`), and a /// lock-serialized `send`. pub fn new(config: DaemonConfig, run_task: RunTaskFn, send: SendFn) -> Self { @@ -40,6 +47,7 @@ impl DaemonRuntime { inflight_count: AtomicUsize::new(0), inflight_idle: Notify::new(), capabilities: TokioMutex::new(None), + screen_kill: AtomicBool::new(false), accessible_dirs: StdMutex::new(accessible_dirs), sessions: crate::sessions::SessionRegistry::default(), }), @@ -185,6 +193,23 @@ impl DaemonRuntime { .and_then(|task| task.session_id.clone()) } + /// Terminate the running task identified by sender, id, and dispatch receipt. + /// + /// The signal remains bound to the task record while the map is locked. The + /// correlation check prevents a delayed request from terminating a later + /// dispatch that reused the same task id. + pub fn terminate_task(&self, from: &str, task_id: &str, correlation_id: &str) -> bool { + let running = self.inner.running.lock().unwrap(); + let Some(task) = running.get(&Self::task_key(from, task_id)) else { + return false; + }; + if task.correlation_id.as_deref() != Some(correlation_id) { + return false; + } + task.abort.terminate(); + true + } + /// Record the session an executor opened for a running task. pub(super) fn record_task_session(&self, key: &str, session_id: String) { if let Some(task) = self.inner.running.lock().unwrap().get_mut(key) { diff --git a/src/sdk/src/daemon/task_loop/probe.rs b/src/sdk/src/daemon/task_loop/probe.rs index 882bea699..77e124c96 100644 --- a/src/sdk/src/daemon/task_loop/probe.rs +++ b/src/sdk/src/daemon/task_loop/probe.rs @@ -10,7 +10,29 @@ impl DaemonRuntime { /// Answer a `capabilities` probe with the cached [`AgentCapabilities`]. pub(super) async fn handle_capabilities(&self, from: String, frame: TaskFrame) { #[cfg_attr(not(feature = "workflows"), allow(unused_mut))] - let mut capabilities = self.get_capabilities().await; + let mut capabilities = if self + .inner + .screen_kill + .load(std::sync::atomic::Ordering::Relaxed) + && self.inner.capabilities.lock().await.is_none() + { + let runtime = self.clone(); + tokio::spawn(async move { + runtime.get_capabilities().await; + }); + AgentCapabilities { + cwd: Some(self.inner.config.workspace.clone()), + providers: self.inner.config.providers.clone(), + screen_kill: true, + ..Default::default() + } + } else { + self.get_capabilities().await + }; + capabilities.screen_kill = self + .inner + .screen_kill + .load(std::sync::atomic::Ordering::Relaxed); // Read fresh rather than cached: the harness probe is expensive and // worth caching, but an operator who just installed a workflow expects // the next probe to advertise it. @@ -63,13 +85,9 @@ impl DaemonRuntime { let abort = Abort::new(); let controller_id = self.register_controller(abort.clone()); let accessible_dirs = self.inner.accessible_dirs.lock().unwrap().clone(); - // Compete for the concurrency budget like a task. - let permit = self - .inner - .slots - .acquire() - .await - .expect("semaphore is never closed"); + // Control-plane negotiation must not wait behind the harness-task slot + // it may be needed to terminate. The probe is separately serialized by + // the capability cache lock and bounded by its own timeout. let capabilities = probe_capabilities(ProbeOptions { provider, run_task: self.inner.run_task.clone(), @@ -91,7 +109,6 @@ impl DaemonRuntime { router: self.inner.config.router.clone(), }) .await; - drop(permit); self.unregister_controller(controller_id); *guard = Some(capabilities.clone()); capabilities diff --git a/src/sdk/src/daemon/types.rs b/src/sdk/src/daemon/types.rs index 2d8e77e35..e8ff48df6 100644 --- a/src/sdk/src/daemon/types.rs +++ b/src/sdk/src/daemon/types.rs @@ -9,7 +9,7 @@ use std::collections::HashMap; use std::future::Future; use std::pin::Pin; -use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering}; +use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering}; use std::sync::{Arc, Mutex as StdMutex}; use tokio::sync::{mpsc, Mutex as TokioMutex, Notify, Semaphore}; @@ -262,6 +262,8 @@ pub(super) struct Inner { pub(super) inflight_idle: Notify, /// Cached capability probe result. pub(super) capabilities: TokioMutex>, + /// Whether this embedding routes authenticated screen kill messages. + pub(super) screen_kill: AtomicBool, /// Workspace roots currently approved for capability advertisement. /// /// Kept outside the immutable config so the daemon TUI can change the diff --git a/src/sdk/src/hub/handle/mod.rs b/src/sdk/src/hub/handle/mod.rs index 5212fc0d5..53bda115c 100644 --- a/src/sdk/src/hub/handle/mod.rs +++ b/src/sdk/src/hub/handle/mod.rs @@ -130,6 +130,27 @@ impl HubHandle { sent } + /// Ask `worker` to kill the harness serving `task_id`. + /// + /// The worker resolves the task against the authenticated sender before it + /// touches a PTY, so this cannot be used to kill another controller's work. + pub async fn kill(&self, worker: &str, task_id: &str) -> Result<(), String> { + let (correlation_id, wire_task_id) = self + .runner + .kill_correlation_for(worker, task_id) + .await + .ok_or_else(|| { + format!("task {task_id} is not running with termination support on {worker}") + })?; + let body = + crate::tinyplace::encode_screen_message(&crate::tinyplace::ScreenMessage::Kill { + task_id: wire_task_id, + correlation_id, + }); + (self.log)(&format!("hub: killing task {task_id} on {worker}")); + self.relay.send(worker, &body).await + } + /// Build a handle from its wiring. pub(super) fn new(wiring: HandleWiring) -> Self { HubHandle { diff --git a/src/sdk/src/hub/runner/capabilities.rs b/src/sdk/src/hub/runner/capabilities.rs index 3741cc31e..ad525eb8d 100644 --- a/src/sdk/src/hub/runner/capabilities.rs +++ b/src/sdk/src/hub/runner/capabilities.rs @@ -5,7 +5,7 @@ use tokio::sync::oneshot; use crate::tinyplace::{encode_task_frame, AgentCapabilities, EncodeFrameInput, TaskFrameKind}; use super::types::CapabilityProbeGuard; -use super::{RunError, TaskRunner, MAX_RESETS}; +use super::{RunError, TaskRunner, CONTACT_POLL, CONTACT_WAIT, MAX_RESETS}; impl Drop for CapabilityProbeGuard { fn drop(&mut self) { @@ -26,6 +26,9 @@ impl TaskRunner { /// (the hub's socket-plane `capabilities_result`) treats any error as "no /// budgets to advertise" and falls open to the static facts. pub async fn capabilities(&self, address: &str) -> Result { + // A failed refresh must not leave support advertised by an older worker + // that previously occupied this address. + self.capabilities.lock().await.remove(address); if !self.relay.contact_accepted(address).await { let _ = self.relay.request_contact(address).await; return Err(RunError::Worker( @@ -75,7 +78,13 @@ impl TaskRunner { return Err(RunError::Transport(error)); } match tokio::time::timeout(self.ack_window, receiver).await { - Ok(Ok(Ok(caps))) => return Ok(caps), + Ok(Ok(Ok(caps))) => { + self.capabilities + .lock() + .await + .insert(address.to_string(), caps.clone()); + return Ok(caps); + } Ok(Ok(Err(error))) => return Err(RunError::Worker(error)), Ok(Err(_)) => { return Err(RunError::Transport( @@ -93,4 +102,58 @@ impl TaskRunner { } } } + + /// Negotiate capabilities while making a backend abort effective before + /// the task itself is registered and sent. + pub async fn capabilities_for_dispatch( + &self, + address: &str, + abort_id: &str, + ) -> ( + Result, + std::sync::Arc, + ) { + let abort = std::sync::Arc::new(tokio::sync::Notify::new()); + self.aborts + .lock() + .expect("aborts lock") + .insert(abort_id.to_string(), abort.clone()); + if !self.relay.contact_accepted(address).await { + let _ = self.relay.request_contact(address).await; + let deadline = std::time::Instant::now() + CONTACT_WAIT; + while std::time::Instant::now() < deadline + && !self.relay.contact_accepted(address).await + { + tokio::select! { + biased; + _ = abort.notified() => { + let mut aborts = self.aborts.lock().expect("aborts lock"); + if aborts.get(abort_id).is_some_and(|current| { + std::sync::Arc::ptr_eq(current, &abort) + }) { + aborts.remove(abort_id); + } + return (Err(RunError::Aborted), abort); + }, + _ = tokio::time::sleep(CONTACT_POLL) => {} + } + } + } + tokio::select! { + biased; + _ = abort.notified() => { + let mut aborts = self.aborts.lock().expect("aborts lock"); + if aborts + .get(abort_id) + .is_some_and(|current| std::sync::Arc::ptr_eq(current, &abort)) + { + aborts.remove(abort_id); + } + (Err(RunError::Aborted), abort) + } + result = self.capabilities(address) => { + (result, abort) + }, + } + } } diff --git a/src/sdk/src/hub/runner/mod.rs b/src/sdk/src/hub/runner/mod.rs index 718ccd1d9..85a7e6d4a 100644 --- a/src/sdk/src/hub/runner/mod.rs +++ b/src/sdk/src/hub/runner/mod.rs @@ -190,6 +190,7 @@ impl TaskRunner { waiters, system_info_waiters, capabilities_waiters, + capabilities: Arc::new(Mutex::new(HashMap::new())), aborts: Arc::new(std::sync::Mutex::new(HashMap::new())), counter: AtomicU64::new(0), ack_window, @@ -221,6 +222,22 @@ impl TaskRunner { true } + /// Return the active dispatch receipt for a worker/task pair. + pub async fn kill_correlation_for( + &self, + worker: &str, + task_id: &str, + ) -> Option<(String, String)> { + self.waiters + .lock() + .await + .iter() + .find(|(_, waiter)| { + waiter.from == worker && waiter.task_id == task_id && waiter.screen_kill + }) + .map(|(correlation, waiter)| (correlation.clone(), waiter.wire_task_id.clone())) + } + /// Cancel every dispatch this runner has in flight. /// /// For a caller that owns a runner serving one piece of work and wants to @@ -270,6 +287,31 @@ impl TaskRunner { &self, req: TaskRequest, status: Option>, + ) -> Result { + self.run_inner(req, status, false, None, None).await + } + + /// Run a dispatch with the screen-control support negotiated specifically + /// for this request. + pub async fn run_negotiated( + &self, + req: TaskRequest, + status: Option>, + screen_kill: bool, + abort: Option>, + visible_task_id: Option, + ) -> Result { + self.run_inner(req, status, screen_kill, abort, visible_task_id) + .await + } + + async fn run_inner( + &self, + req: TaskRequest, + status: Option>, + screen_kill: bool, + prepared_abort: Option>, + visible_task_id: Option, ) -> Result { // Register this dispatch's abort signal FIRST — before the contact wait — // so a `task_abort` that arrives during contact negotiation (up to @@ -278,7 +320,7 @@ impl TaskRunner { // the backend aborts by, and held for the whole call (spanning any // reset+resend retries). The guard removes it on every return path, so a // settled dispatch leaves nothing for a later `task_abort` to match. - let abort = Arc::new(Notify::new()); + let abort = prepared_abort.unwrap_or_else(|| Arc::new(Notify::new())); self.aborts .lock() .expect("aborts lock") @@ -322,6 +364,11 @@ impl TaskRunner { self.waiters.lock().await.insert( cid.clone(), Waiter { + task_id: visible_task_id + .clone() + .unwrap_or_else(|| req.task_id.clone()), + wire_task_id: req.task_id.clone(), + screen_kill, from: req.worker_address.clone(), reply: tx, status: status.clone(), diff --git a/src/sdk/src/hub/runner/types.rs b/src/sdk/src/hub/runner/types.rs index 9c3cc97e8..7cd422cd4 100644 --- a/src/sdk/src/hub/runner/types.rs +++ b/src/sdk/src/hub/runner/types.rs @@ -3,6 +3,12 @@ use super::*; /// A registered dispatch awaiting its terminal frame. pub(super) struct Waiter { + /// The task id this dispatch carries on the wire. + pub(super) task_id: String, + /// Worker-facing task id used by the daemon's running-task registry. + pub(super) wire_task_id: String, + /// Screen termination support negotiated for this exact dispatch. + pub(super) screen_kill: bool, /// The worker address this dispatch was sent to — the only sender whose /// frames may settle it. See [`Probe::from`]. pub(super) from: String, @@ -78,6 +84,8 @@ pub struct TaskRunner { pub(super) system_info_waiters: SystemInfoWaiters, /// Capability probes waiting for a worker's `capabilities_result`. pub(super) capabilities_waiters: CapabilitiesWaiters, + /// Last successfully negotiated capabilities for each worker address. + pub(super) capabilities: Arc>>, /// Abort signals for in-flight dispatches, keyed by orchestrator-facing task /// id; [`abort_task`](Self::abort_task) notifies one to cancel its dispatch. pub(super) aborts: Aborts, diff --git a/src/sdk/src/hub/socket/task_run.rs b/src/sdk/src/hub/socket/task_run.rs index dd12d317e..f312ee0a4 100644 --- a/src/sdk/src/hub/socket/task_run.rs +++ b/src/sdk/src/hub/socket/task_run.rs @@ -162,6 +162,24 @@ pub(super) async fn handle_task_run( return; }; + // Negotiate once before dispatch. Besides informing automatic provider + // selection, this records static control-plane support so an emergency kill + // never waits behind a fresh probe of a wedged worker. + let (capabilities, abort) = match runner + .capabilities_for_dispatch(&worker_address, &task_id) + .await + { + (Ok(capabilities), abort) => (Some(capabilities), abort), + (Err(crate::hub::RunError::Aborted), _) => { + let outcome = Err(crate::hub::RunError::Aborted); + let _ = socket + .emit("medulla:task_result", result_frame(&task_id, &outcome)) + .await; + return; + } + (Err(_), abort) => (None, abort), + }; + // An explicit provider is authoritative. Only an untargeted task consults // the subscription strategy, and a failed/unknown budget probe falls open // to the daemon's own configured default. @@ -174,12 +192,9 @@ pub(super) async fn handle_task_run( if strategy == crate::runtime::SubscriptionRoutingStrategy::Manual { None } else { - match runner.capabilities(&worker_address).await { - Ok(capabilities) => { - super::super::roster::subscription_for_strategy(&capabilities, strategy) - } - Err(_) => None, - } + capabilities.as_ref().and_then(|capabilities| { + super::super::roster::subscription_for_strategy(capabilities, strategy) + }) } } }; @@ -265,7 +280,10 @@ pub(super) async fn handle_task_run( fleet_depth: 0, }; - let outcome = runner.run(req, Some(tx)).await; + let screen_kill = capabilities.is_some_and(|capabilities| capabilities.screen_kill); + let outcome = runner + .run_negotiated(req, Some(tx), screen_kill, Some(abort), None) + .await; match &outcome { Ok(o) => log(&format!( "hub: task {} ok ({} chars)", diff --git a/src/sdk/src/runtime/mod.rs b/src/sdk/src/runtime/mod.rs index c37649b07..e82b51ef4 100644 --- a/src/sdk/src/runtime/mod.rs +++ b/src/sdk/src/runtime/mod.rs @@ -205,6 +205,18 @@ pub trait Runtime: Send + Sync { Box::pin(async { Ok(()) }) } + /// Kill the harness serving `task_id` on `worker`. + /// + /// A no-op success without a hub. Interactive callers are responsible for + /// confirming the destructive action before invoking this method. + fn kill_task( + &self, + _worker: String, + _task_id: String, + ) -> BoxFuture<'static, anyhow::Result<()>> { + Box::pin(async { Ok(()) }) + } + /// Tell the orchestrator a harness has been handed back, with the brief. /// /// An **error** by default rather than a silent success, unlike diff --git a/src/sdk/src/runtime/openhuman/mod.rs b/src/sdk/src/runtime/openhuman/mod.rs index 496709219..8b09cce1e 100644 --- a/src/sdk/src/runtime/openhuman/mod.rs +++ b/src/sdk/src/runtime/openhuman/mod.rs @@ -504,6 +504,18 @@ impl Runtime for OpenHumanRuntime { }) } + fn kill_task(&self, worker: String, task_id: String) -> BoxFuture<'static, anyhow::Result<()>> { + let hub = self.hub(); + Box::pin(async move { + let Some(hub) = hub else { + return Ok(()); + }; + hub.kill(&worker, &task_id) + .await + .map_err(|e| anyhow::anyhow!(e)) + }) + } + fn worker_op(&self, op: crate::runtime::WorkerOp) -> BoxFuture<'static, anyhow::Result<()>> { let hub = self.hub(); Box::pin(async move { diff --git a/src/sdk/src/tinyplace/frames/tests/capabilities.rs b/src/sdk/src/tinyplace/frames/tests/capabilities.rs new file mode 100644 index 000000000..a270e90b6 --- /dev/null +++ b/src/sdk/src/tinyplace/frames/tests/capabilities.rs @@ -0,0 +1,21 @@ +//! Compatibility tests for additive worker capability fields. + +use crate::tinyplace::{parse_agent_capabilities, AgentCapabilities, BudgetWindow}; + +#[test] +fn budget_window_defaults_to_unknown() { + assert_eq!(BudgetWindow::default(), BudgetWindow::Unknown); +} + +#[test] +fn screen_kill_support_is_additive_and_defaults_off_for_older_workers() { + let older = parse_agent_capabilities(r#"{"providers":["claude"]}"#).unwrap(); + assert!(!older.screen_kill); + + let current = AgentCapabilities { + screen_kill: true, + ..Default::default() + }; + let value = serde_json::to_value(¤t).unwrap(); + assert_eq!(value["screenKill"], true); +} diff --git a/src/sdk/src/tinyplace/frames/tests/codec.rs b/src/sdk/src/tinyplace/frames/tests/codec.rs index 87b1fb3dd..e1c762543 100644 --- a/src/sdk/src/tinyplace/frames/tests/codec.rs +++ b/src/sdk/src/tinyplace/frames/tests/codec.rs @@ -347,11 +347,6 @@ fn harness_readiness_omits_reason_when_ready() { assert_eq!(value["reason"], "not authenticated"); } -#[test] -fn budget_window_defaults_to_unknown() { - assert_eq!(BudgetWindow::default(), BudgetWindow::Unknown); -} - #[test] fn parse_agent_capabilities_defaults_missing_arrays() { let caps = parse_agent_capabilities(r#"{"cwd":"/x"}"#).unwrap(); @@ -388,6 +383,7 @@ fn empty_budgets_and_readiness_are_omitted_on_the_wire() { let value = serde_json::to_value(&caps).unwrap(); assert!(value.get("budgets").is_none()); assert!(value.get("readiness").is_none()); + assert!(value.get("screenKill").is_none()); } #[test] diff --git a/src/sdk/src/tinyplace/frames/tests/mod.rs b/src/sdk/src/tinyplace/frames/tests/mod.rs index dcce9b97e..39b028aa0 100644 --- a/src/sdk/src/tinyplace/frames/tests/mod.rs +++ b/src/sdk/src/tinyplace/frames/tests/mod.rs @@ -1,4 +1,5 @@ //! Focused tests for task-frame encoding and restricted tool-mode decoding. +mod capabilities; mod codec; mod tool_mode; diff --git a/src/sdk/src/tinyplace/frames/types.rs b/src/sdk/src/tinyplace/frames/types.rs index d83584855..740e43847 100644 --- a/src/sdk/src/tinyplace/frames/types.rs +++ b/src/sdk/src/tinyplace/frames/types.rs @@ -515,6 +515,16 @@ pub struct AgentCapabilities { /// field. Same backward-compatibility contract as the two vectors above. #[serde(default, skip_serializing_if = "Vec::is_empty")] pub workflows: Vec, + /// Whether this worker accepts task-correlated harness termination requests. + /// + /// Older workers omit the field and therefore deserialize as `false`, so an + /// upgraded controller never sends them an unknown screen control message. + #[serde( + rename = "screenKill", + default, + skip_serializing_if = "std::ops::Not::not" + )] + pub screen_kill: bool, } /// Fleet-safe description of one named custom harness. diff --git a/src/sdk/src/tinyplace/screen/tests.rs b/src/sdk/src/tinyplace/screen/tests.rs index 36ef5cf57..2905b22a7 100644 --- a/src/sdk/src/tinyplace/screen/tests.rs +++ b/src/sdk/src/tinyplace/screen/tests.rs @@ -316,6 +316,10 @@ fn messages_round_trip_through_the_envelope() { ScreenMessage::Unsubscribe { task_id: "w_1".into(), }, + ScreenMessage::Kill { + task_id: "w_1".into(), + correlation_id: "cyc/w_1/0".into(), + }, ScreenMessage::Ack { task_id: "w_1".into(), seq: 418, diff --git a/src/sdk/src/tinyplace/screen/types.rs b/src/sdk/src/tinyplace/screen/types.rs index 204a11871..1aa9f1adb 100644 --- a/src/sdk/src/tinyplace/screen/types.rs +++ b/src/sdk/src/tinyplace/screen/types.rs @@ -178,9 +178,9 @@ pub struct ScreenFrame { /// Everything that can cross this protocol, in both directions. /// -/// There is no `input` and no `resize`: the viewer never reaches back into the -/// session, and the sender's geometry is authoritative. Both are additive later -/// if that changes. +/// There is no `input` and no `resize`: the viewer cannot steer the session, +/// and the sender's geometry is authoritative. The one control operation is an +/// explicit, task-scoped kill used by an operator to recover a hung harness. #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] #[serde(tag = "kind", rename_all = "snake_case")] pub enum ScreenMessage { @@ -199,6 +199,14 @@ pub enum ScreenMessage { /// The task to stop watching. task_id: String, }, + /// Viewer → sender: kill the harness serving an owned running task. + Kill { + /// The task whose harness should be killed. + task_id: String, + /// The unique dispatch receipt, preventing a delayed kill from matching + /// a later dispatch that reused the task id. + correlation_id: String, + }, /// Viewer → sender: the highest sequence the viewer holds, which the next /// diff may be taken from. Ack { diff --git a/src/tui/src/event_loop/cmd_dispatch/mod.rs b/src/tui/src/event_loop/cmd_dispatch/mod.rs index 9687195b7..f216c38e5 100644 --- a/src/tui/src/event_loop/cmd_dispatch/mod.rs +++ b/src/tui/src/event_loop/cmd_dispatch/mod.rs @@ -180,6 +180,17 @@ pub(super) fn run_cmd( } }); } + Cmd::KillTask { worker, task_id } => { + let rt = runtime.clone(); + let tx = msg_tx.clone(); + tokio::spawn(async move { + let status = match rt.kill_task(worker, task_id.clone()).await { + Ok(()) => format!("Kill requested for {task_id}"), + Err(e) => format!("Cannot kill {task_id}: {e}"), + }; + let _ = tx.send(AppMsg::Status(status)); + }); + } Cmd::StartLocalHost { host, index } => { let Some(spawner) = local_hosts.cloned() else { let _ = msg_tx.send(AppMsg::Status( diff --git a/src/tui/src/ui/app/input.rs b/src/tui/src/ui/app/input.rs index 00dd76f2f..7c7403bb6 100644 --- a/src/tui/src/ui/app/input.rs +++ b/src/tui/src/ui/app/input.rs @@ -6,7 +6,7 @@ use crossterm::event::{Event, KeyEventKind, MouseButton, MouseEventKind}; use super::rail::RailRow; -use crate::ui::agents::{agent_row_model, AgentRole, AgentRow}; +use crate::ui::agents::{agent_row_model, AgentRole, AgentRow, TaskStatus}; use crate::ui::composer::{insert_at, normalize_paste, Draft}; use super::types::{App, Cmd, ROUTING_SUBPAGES, SETTINGS_SUBPAGES, TOKENMAXXING_SUBPAGES}; @@ -101,6 +101,9 @@ impl App { /// Handle scroll and left-click mouse events for the active tab. pub(super) fn on_mouse(&mut self, m: crossterm::event::MouseEvent) -> Option { + if self.kill_armed.take().is_some() { + self.set_status("Harness kill cancelled"); + } // A modal swallows the mouse, the same way it swallows the keyboard. // The harness picker is one: a click that navigated the rail behind it // left an overlay on screen describing a row nobody was pointing at. @@ -553,7 +556,7 @@ impl App { } /// The `(worker address, task id)` the current selection asks to watch. - fn watch_target(&self) -> Option<(String, String)> { + pub(super) fn watch_target(&self) -> Option<(String, String)> { // Only on the Agents tab: leaving it releases the subscription. if self.tab() != "Agents" { return None; @@ -587,4 +590,16 @@ impl App { .unwrap_or_else(|| agent_id.to_string()); Some((address, task.task_id.clone())) } + + /// The selected running task eligible for destructive termination. + pub(super) fn kill_target(&self) -> Option<(String, String)> { + let rows = self.rail_rows(); + let row = rows.get(self.agent_index.min(rows.len().saturating_sub(1)))?; + let RailRow::Agent(AgentRow::Sub { task, .. }) = row else { + return None; + }; + (task.status == TaskStatus::Running) + .then(|| self.watch_target()) + .flatten() + } } diff --git a/src/tui/src/ui/app/keys/agents.rs b/src/tui/src/ui/app/keys/agents.rs index 92e38dfd2..05078a443 100644 --- a/src/tui/src/ui/app/keys/agents.rs +++ b/src/tui/src/ui/app/keys/agents.rs @@ -82,6 +82,14 @@ impl App { let alt = k.modifiers.contains(KeyModifiers::ALT); match k.code { + KeyCode::Char('K') => { + if let Some(target) = self.kill_target() { + self.arm_kill(target); + } else { + self.set_status("Select a running harness task first"); + } + AgentsKey::Handled(None) + } // The bare arrows are the point of having focus at all. KeyCode::Up | KeyCode::Down => { self.agent_scroll = 0; diff --git a/src/tui/src/ui/app/keys/mod.rs b/src/tui/src/ui/app/keys/mod.rs index ecbfe98bf..0cbdc2bef 100644 --- a/src/tui/src/ui/app/keys/mod.rs +++ b/src/tui/src/ui/app/keys/mod.rs @@ -44,6 +44,17 @@ impl App { let shift = k.modifiers.contains(KeyModifiers::SHIFT); let alt = k.modifiers.contains(KeyModifiers::ALT); + // Killing a harness can lose in-progress work. Once armed, the prompt + // owns exactly one keypress: only a deliberate `y` proceeds. + if let Some((worker, task_id)) = self.kill_armed.take() { + if k.code == KeyCode::Char('y') && k.modifiers.is_empty() { + self.set_status(format!("Killing harness for {task_id}…")); + return Some(Cmd::KillTask { worker, task_id }); + } + self.set_status("Harness kill cancelled"); + return None; + } + // The hand-back question outranks even the attached harness, and has to: // it is asked *while still attached*, because releasing the keyboard // before it is answered would hide the pane the question is about. So diff --git a/src/tui/src/ui/app/render/agents/composer.rs b/src/tui/src/ui/app/render/agents/composer.rs index 80eee4e1d..0b2e9fadc 100644 --- a/src/tui/src/ui/app/render/agents/composer.rs +++ b/src/tui/src/ui/app/render/agents/composer.rs @@ -72,7 +72,7 @@ impl App { // that cannot send it leaves the rail unreachable, so the key that // always works is named right where the cursor is. let caption = if self.agents_rail_focused() { - "↑↓ walk agents · Enter or type to write".to_string() + "↑↓ walk agents · K kill harness · Enter or type to write".to_string() } else { format!("› {target} · Esc to pick an agent") }; diff --git a/src/tui/src/ui/app/state.rs b/src/tui/src/ui/app/state.rs index ac8aa1eb2..f87ef0b9f 100644 --- a/src/tui/src/ui/app/state.rs +++ b/src/tui/src/ui/app/state.rs @@ -66,6 +66,7 @@ impl App { context_index: 0, agent_index: 0, watching: None, + kill_armed: None, agents_focus: super::types::AgentsFocus::default(), agent_scroll: 0, chat_scroll: 0, @@ -372,9 +373,19 @@ impl App { } pub fn set_status(&mut self, s: impl Into) { + // A destructive confirmation is valid only while its question remains + // visible. Any asynchronous status replacement cancels it. + self.kill_armed = None; self.status = s.into(); } + /// Show and arm the harness-kill confirmation as one invariant-preserving + /// state transition. + pub(super) fn arm_kill(&mut self, target: (String, String)) { + self.set_status("Kill this harness? y confirm · any other key cancels"); + self.kill_armed = Some(target); + } + /// Replace the Context-tab chunks. pub fn set_contexts(&mut self, c: Vec) { self.contexts = c; diff --git a/src/tui/src/ui/app/tests.rs b/src/tui/src/ui/app/tests.rs index cd698a67d..27b03c19e 100644 --- a/src/tui/src/ui/app/tests.rs +++ b/src/tui/src/ui/app/tests.rs @@ -22,6 +22,24 @@ fn app() -> App { App::new(rt, loaded) } +fn app_with_running_task() -> App { + let rt = MockRuntime::demo(); + let loaded = { + let mut loaded = LoadedConfig::defaults("medulla.tui.json".into()); + loaded.config.tinyplace = Some(medulla::config::TinyplaceConfig::default()); + loaded + }; + let mut app = App::new(Arc::new(rt), loaded); + app.snapshot.events.retain(|envelope| { + !matches!( + &envelope.event, + crate::ui::events::TuiEvent::TaskComplete { digest } + if digest.task_id == "task-1" + ) + }); + app +} + /// The index of the tab named `name`. Looked up rather than written down: the /// tab bar's order is a product decision that has changed before. fn tab(name: &str) -> usize { @@ -292,6 +310,53 @@ fn selecting_a_task_asks_to_watch_it() { assert!(app.watching.is_some(), "the target is remembered"); } +#[test] +fn killing_a_watched_harness_requires_confirmation() { + let mut app = app_with_running_task(); + select_first_task(&mut app).expect("the fixture has a selectable task"); + app.focus_agents_rail(); + + let armed = app.on_key(KeyEvent::new(KeyCode::Char('K'), KeyModifiers::SHIFT)); + assert!(armed.is_none(), "arming must not kill the harness"); + assert!(app.kill_armed.is_some()); + assert!(app.status().contains("y confirm")); + + let cmd = app.on_key(KeyEvent::new(KeyCode::Char('y'), KeyModifiers::NONE)); + assert!(matches!(cmd, Some(Cmd::KillTask { .. }))); + assert!(app.kill_armed.is_none()); +} + +#[test] +fn killing_resolves_the_current_rail_selection_instead_of_the_cached_watch() { + let mut app = app_with_running_task(); + select_first_task(&mut app).expect("the fixture has a selectable task"); + let selected = app.watch_target().expect("the selected task is watchable"); + app.watching = Some(("stale-worker".into(), "stale-task".into())); + app.focus_agents_rail(); + + app.on_key(KeyEvent::new(KeyCode::Char('K'), KeyModifiers::SHIFT)); + assert_eq!(app.kill_armed.as_ref(), Some(&selected)); + + let cmd = app.on_key(KeyEvent::new(KeyCode::Char('y'), KeyModifiers::NONE)); + let Some(Cmd::KillTask { worker, task_id }) = cmd else { + panic!("confirming should kill the selected task"); + }; + assert_eq!((worker, task_id), selected); +} + +#[test] +fn any_other_key_cancels_a_harness_kill() { + let mut app = app_with_running_task(); + select_first_task(&mut app).expect("the fixture has a selectable task"); + app.focus_agents_rail(); + app.on_key(KeyEvent::new(KeyCode::Char('K'), KeyModifiers::SHIFT)); + + let cmd = app.on_key(KeyEvent::new(KeyCode::Char('n'), KeyModifiers::NONE)); + assert!(cmd.is_none()); + assert!(app.kill_armed.is_none()); + assert!(app.status().contains("cancelled")); +} + #[test] fn reselecting_the_same_task_does_not_resubscribe() { // Every subscribe carries `resync: true`, so re-issuing one on each diff --git a/src/tui/src/ui/app/types.rs b/src/tui/src/ui/app/types.rs index 736691f8a..fde0b5957 100644 --- a/src/tui/src/ui/app/types.rs +++ b/src/tui/src/ui/app/types.rs @@ -352,6 +352,13 @@ pub enum Cmd { /// The `(worker address, task id)` to start streaming, if any. start: Option<(String, String)>, }, + /// Kill the harness serving a watched task after UI confirmation. + KillTask { + /// The worker address that owns the harness. + worker: String, + /// The dispatched task whose harness should be killed. + task_id: String, + }, /// Push a handoff brief for a harness the operator just gave back. /// /// Off the render thread because it does two things that must not block a @@ -770,6 +777,8 @@ pub struct App { /// new one: a subscription nobody is looking at costs the worker a sample, /// a ratchet advance and a send on every tick. pub(super) watching: Option<(String, String)>, + /// The watched `(worker, task)` awaiting destructive-action confirmation. + pub(super) kill_armed: Option<(String, String)>, /// Which half of the Agents tab the keyboard is driving. pub(super) agents_focus: AgentsFocus, pub(super) agent_scroll: usize, diff --git a/src/tui/src/worker/executor/run.rs b/src/tui/src/worker/executor/run.rs index 8f8362079..c69d7109f 100644 --- a/src/tui/src/worker/executor/run.rs +++ b/src/tui/src/worker/executor/run.rs @@ -277,8 +277,7 @@ impl PtySessionExecutor { /// composer, which is the failure that produces confidently wrong answers /// rather than an error. fn stop_turn(&self, id: &str) { - let _ = self.sessions.write(id, &[0x03]); - self.sessions.close(id); + self.sessions.stop_if_orchestrator(id); } /// Decide which session serves this task: reuse an idle one, or launch. @@ -481,9 +480,14 @@ impl PtySessionExecutor { )); } if abort.is_aborted() { - // A real interrupt, not a kill: Ctrl-C reaches the harness the - // same way the operator's would, and the session survives it. - let _ = self.sessions.write(id, &[0x03]); + if abort.is_terminated() { + self.stop_turn(id); + } else { + // A requester abort is an interrupt: Ctrl-C reaches the + // harness the same way the operator's would, and the + // reusable session survives it. + let _ = self.sessions.write(id, &[0x03]); + } return Err(format!("{} task aborted", provider.as_str())); } if !self diff --git a/src/tui/src/worker/pty/handle/control.rs b/src/tui/src/worker/pty/handle/control.rs index 8929c0f5a..ef2b8f561 100644 --- a/src/tui/src/worker/pty/handle/control.rs +++ b/src/tui/src/worker/pty/handle/control.rs @@ -29,10 +29,24 @@ impl SessionHandle { /// running that turn. Clearing `busy` on handback would advertise a harness /// as free while it was still finishing someone else's work. pub(in super::super) fn set_control(&self, control: HarnessControl) { + let _cold = lock(&self.cold); self.operator_held .store(control == HarnessControl::User, Ordering::Release); } + /// Interrupt and close only while orchestration still owns the session. + /// Serialized with control handoff so takeover and termination cannot race. + pub(in super::super) fn stop_if_orchestrator(&self) -> bool { + let _cold = lock(&self.cold); + if self.operator_held.load(Ordering::Acquire) { + return false; + } + let _ = self.write(&[0x03]); + self.kill(); + self.mark_closed(); + true + } + /// Whether this session may serve `label`'s next turn. /// /// Its own label, or the synthetic `you:` one of a handed-back diff --git a/src/tui/src/worker/pty/manager/session.rs b/src/tui/src/worker/pty/manager/session.rs index 2c9eb6706..bed9c65c3 100644 --- a/src/tui/src/worker/pty/manager/session.rs +++ b/src/tui/src/worker/pty/manager/session.rs @@ -186,6 +186,12 @@ impl PtyManager { true } + /// Interrupt and close a session only if the operator has not taken it. + pub fn stop_if_orchestrator(&self, id: &str) -> bool { + self.handle(id) + .is_some_and(|session| session.stop_if_orchestrator()) + } + /// Drop an exited session's record and screen. /// /// Refuses while the child is alive, so a forgotten session can never leave diff --git a/src/tui/src/worker/stream/router.rs b/src/tui/src/worker/stream/router.rs index 55619d6f7..c042eab00 100644 --- a/src/tui/src/worker/stream/router.rs +++ b/src/tui/src/worker/stream/router.rs @@ -2,9 +2,10 @@ //! subscriber may watch, and starting or stopping the stream that answers. //! //! This is the half of the protocol the worker *receives*. The vocabulary is -//! deliberately tiny — subscribe, unsubscribe, acknowledge — because the viewer -//! is a passive observer: there is no message here that can type into a session -//! or resize one, so the worst a subscriber can do is read a screen. +//! deliberately tiny — subscribe, unsubscribe, acknowledge, and an explicit kill. +//! There is no message that can type into or resize a session. Kill is resolved +//! through the same task ownership check as subscribe, so a controller can stop +//! only a harness running work that controller dispatched. //! //! Which leaves one question, and the answer is structural rather than a check. //! A subscription names a **task**, and the daemon's running-task record is @@ -95,6 +96,19 @@ impl ScreenRouter { self.log(&format!("screen: {from} unsubscribed from {task_id}")); } } + ScreenMessage::Kill { + task_id, + correlation_id, + } => { + if !self.runtime.terminate_task(from, &task_id, &correlation_id) { + self.log(&format!( + "screen: refused kill from {from} on {task_id} — no such running task for this sender" + )); + return; + } + self.registry.unsubscribe_for(from, &task_id); + self.log(&format!("screen: {from} killed {task_id}")); + } // The sampler chains from the last frame it actually sent, so an // acknowledgement is not needed to decide what to send next. Kept in // the protocol as the viewer's liveness signal and accepted here so diff --git a/src/tui/src/worker/stream/tests.rs b/src/tui/src/worker/stream/tests.rs index 8c1d55531..ba248e366 100644 --- a/src/tui/src/worker/stream/tests.rs +++ b/src/tui/src/worker/stream/tests.rs @@ -423,6 +423,19 @@ async fn an_unsubscribe_for_a_task_nobody_streams_does_nothing() { assert_eq!(router.active(), 0); } +#[tokio::test] +async fn a_kill_for_a_task_this_sender_never_dispatched_is_refused() { + let mut router = empty_router(); + router.handle( + "peerA", + medulla::tinyplace::ScreenMessage::Kill { + task_id: "t1".into(), + correlation_id: "cyc/t1/0".into(), + }, + ); + assert_eq!(router.active(), 0); +} + #[tokio::test] async fn only_the_peer_a_stream_was_opened_for_can_stop_it() { let sessions = super::super::pty::PtyManager::new(); diff --git a/src/tui/src/worker_loop/commands.rs b/src/tui/src/worker_loop/commands.rs index ce7f3592d..5f4308587 100644 --- a/src/tui/src/worker_loop/commands.rs +++ b/src/tui/src/worker_loop/commands.rs @@ -90,6 +90,7 @@ fn start_worker( ) }); let daemon = worker_runtime(start, mode, provider, &transport); + daemon.enable_screen_kill(); // Screens are only worth streaming when there are screens: the headless // executor runs harnesses without a pty, so a subscriber finds nothing. let screens = diff --git a/src/tui/tests/e2e_screen_stream.rs b/src/tui/tests/e2e_screen_stream.rs index fee4caa26..6176738fe 100644 --- a/src/tui/tests/e2e_screen_stream.rs +++ b/src/tui/tests/e2e_screen_stream.rs @@ -69,7 +69,7 @@ fn sh(script: &str, label: &str) -> LaunchSpec { /// The runtime is the piece under test here as much as the sampler: it owns the /// running-task record keyed by `(sender, task id)`, and that record is the only /// thing that lets a subscription resolve. -fn runtime_serving(session_id: String) -> DaemonRuntime { +fn runtime_serving(sessions: PtyManager, session_id: String) -> DaemonRuntime { let config = DaemonConfig { providers: vec![HarnessProvider::Codex], default_provider: HarnessProvider::Codex, @@ -92,13 +92,17 @@ fn runtime_serving(session_id: String) -> DaemonRuntime { }; let run_task = Arc::new(move |options: medulla::daemon::providers::RunTaskOptions| { let session_id = session_id.clone(); + let sessions = sessions.clone(); Box::pin(async move { if let Some(report) = options.on_session { - report(session_id); + report(session_id.clone()); } // Hold the task open so its record — and therefore the subscription // resolving through it — stays live for the duration of the test. - tokio::time::sleep(Duration::from_secs(60)).await; + options.abort.cancelled().await; + if options.abort.is_terminated() { + sessions.close(&session_id); + } Err("never settles".to_string()) }) as std::pin::Pin + Send>> }) as medulla::daemon::providers::RunTaskFn; @@ -158,7 +162,7 @@ async fn a_watched_task_streams_its_real_terminal_to_the_hubs_store() { )) .expect("a pty session"); - let runtime = runtime_serving(session_id.clone()); + let runtime = runtime_serving(sessions.clone(), session_id.clone()); start_task(&runtime, peer, task_id).await; let outbox: Outbox = Arc::new(Mutex::new(Vec::new())); @@ -273,7 +277,7 @@ async fn an_unchanged_screen_costs_nothing_after_the_first_frame() { let session_id = sessions .open(sh("printf 'STILL\\n'; sleep 30", peer)) .expect("a pty session"); - let runtime = runtime_serving(session_id); + let runtime = runtime_serving(sessions.clone(), session_id); start_task(&runtime, peer, task_id).await; let outbox: Outbox = Arc::new(Mutex::new(Vec::new())); @@ -314,3 +318,54 @@ async fn an_unchanged_screen_costs_nothing_after_the_first_frame() { sessions.shutdown(); } + +#[tokio::test(flavor = "multi_thread")] +async fn an_owned_task_kill_stops_its_real_harness() { + let peer = "peerA"; + let task_id = "t3#0"; + let sessions = PtyManager::new(); + let session_id = sessions.open(sh("sleep 30", peer)).expect("a pty session"); + let runtime = runtime_serving(sessions.clone(), session_id.clone()); + start_task(&runtime, peer, task_id).await; + let mut router = ScreenRouter::new(sessions.clone(), runtime, send_fn(|_, _| async {})); + + router.handle( + peer, + ScreenMessage::Kill { + task_id: task_id.to_string(), + correlation_id: "stale-dispatch".to_string(), + }, + ); + assert!( + sessions + .row(&session_id) + .expect("the live session remains inspectable") + .state + .is_running(), + "a stale dispatch receipt must not kill a reused task id" + ); + + router.handle( + peer, + ScreenMessage::Kill { + task_id: task_id.to_string(), + correlation_id: format!("cyc/{task_id}/0"), + }, + ); + + let deadline = Instant::now() + PATIENCE; + while sessions + .row(&session_id) + .expect("the killed session remains inspectable") + .state + .is_running() + { + assert!( + Instant::now() < deadline, + "the task-scoped kill must stop the harness process" + ); + tokio::time::sleep(Duration::from_millis(20)).await; + } + assert_eq!(router.active(), 0, "its screen stream is also stopped"); + sessions.shutdown(); +}