diff --git a/CHANGELOG.md b/CHANGELOG.md index b25ca29ad6..dba6bc6e65 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,7 +5,11 @@ All notable changes to Agent Relay will be documented in this file. The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). -## [Unreleased] +## [Unreleased - Minor] + +### Added + +- Opt-in persistent broker task providers preserve final results across reconnects and acknowledge callbacks only after durable Relaycast receipts. ## [12.2.2] - 2026-09-15 diff --git a/crates/broker/src/fleet_wire.rs b/crates/broker/src/fleet_wire.rs index dd08342980..a9cdbcd89d 100644 --- a/crates/broker/src/fleet_wire.rs +++ b/crates/broker/src/fleet_wire.rs @@ -80,6 +80,8 @@ pub struct FleetCapability { skip_serializing_if = "Option::is_none" )] pub queue: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub execution_mode: Option, #[serde( default, deserialize_with = "deserialize_optional_presence", @@ -375,6 +377,7 @@ pub struct ActionResult { pub id: Option, pub invocation_id: String, pub result: ActionResultPayload, + pub task: Option, } impl Serialize for ActionResult { @@ -420,6 +423,14 @@ struct ActionResultWire { skip_serializing_if = "Option::is_none" )] pub error: Option, + #[serde(default, rename = "final", skip_serializing_if = "Option::is_none")] + pub final_result: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub execution_id: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub worker_generation: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub accounting: Option>, } fn deserialize_optional_presence<'de, D, T>(deserializer: D) -> Result, D::Error> @@ -483,6 +494,13 @@ impl From<&ActionResult> for ActionResultWire { v: value.v, id: value.id.clone(), invocation_id: value.invocation_id.clone(), + final_result: value.task.as_ref().map(|task| task.final_result), + execution_id: value.task.as_ref().map(|task| task.execution_id.clone()), + worker_generation: value + .task + .as_ref() + .map(|task| task.worker_generation.clone()), + accounting: value.task.as_ref().and_then(|task| task.accounting.clone()), output: Some(output.output.clone()), error: None, }, @@ -490,6 +508,13 @@ impl From<&ActionResult> for ActionResultWire { v: value.v, id: value.id.clone(), invocation_id: value.invocation_id.clone(), + final_result: value.task.as_ref().map(|task| task.final_result), + execution_id: value.task.as_ref().map(|task| task.execution_id.clone()), + worker_generation: value + .task + .as_ref() + .map(|task| task.worker_generation.clone()), + accounting: value.task.as_ref().and_then(|task| task.accounting.clone()), output: None, error: Some(error.error.clone()), }, @@ -509,11 +534,44 @@ impl TryFrom for ActionResult { } }; + let task = if value.final_result.is_some() + || value.execution_id.is_some() + || value.worker_generation.is_some() + || value.accounting.is_some() + { + if value.id.as_ref().is_none_or(|id| id.is_empty()) { + return Err("task result requires a request id".into()); + } + let execution_id = value + .execution_id + .filter(|id| !id.is_empty()) + .ok_or("task result requires execution_id")?; + let worker_generation = value + .worker_generation + .filter(|id| !id.is_empty() && id.len() <= 512) + .ok_or("task result requires worker_generation")?; + if value.accounting.as_ref().is_some_and(|values| { + values + .values() + .any(|n| n.as_f64().is_none_or(|n| !n.is_finite() || n < 0.0)) + }) { + return Err("task accounting must be finite and nonnegative".into()); + } + Some(TaskResultFields { + final_result: value.final_result.ok_or("task result requires final")?, + execution_id, + worker_generation, + accounting: value.accounting, + }) + } else { + None + }; Ok(Self { v: value.v, id: value.id, invocation_id: value.invocation_id, result, + task, }) } } @@ -605,6 +663,8 @@ pub struct ActionInvoke { skip_serializing_if = "Option::is_none" )] pub agent_name: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub task_execution: Option>, } #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] @@ -720,6 +780,38 @@ where } } +// Nested inside the inbound ActionInvoke frame, so this must retain the same +// forward-compatibility rule: a future engine field cannot make the broker +// drop the entire invocation before acknowledging it. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct TaskExecution { + pub execution_id: String, + pub run_id: String, + pub step_id: String, + pub dispatch_id: String, + pub deadline: String, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct ActionAccept { + pub v: FleetWireVersion, + pub id: String, + pub invocation_id: String, + pub execution_id: String, + pub worker_generation: String, +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct TaskResultFields { + #[serde(rename = "final")] + pub final_result: bool, + pub execution_id: String, + pub worker_generation: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub accounting: Option>, +} + #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] #[serde(tag = "type")] pub enum NodeToServer { @@ -737,6 +829,8 @@ pub enum NodeToServer { DeliveryAck(DeliveryAck), #[serde(rename = "action.result")] ActionResult(ActionResult), + #[serde(rename = "action.accept")] + ActionAccept(ActionAccept), #[serde(rename = "inventory.sync")] InventorySync(InventorySync), } @@ -863,6 +957,7 @@ mod tests { #[test] fn action_result_allows_error_payloads() { let msg = BrokerToRelaycast::ActionResult(ActionResult { + task: None, v: FLEET_WIRE_VERSION, id: None, invocation_id: "inv_2".to_string(), @@ -993,6 +1088,7 @@ mod tests { name: "builder-1".to_string(), node_id: "node_1".to_string(), capabilities: vec![FleetCapability { + execution_mode: None, name: "spawn:codex".to_string(), kind: Some("capacity".to_string()), global: None, @@ -1330,4 +1426,29 @@ mod tests { let decoded: RelaycastToBroker = serde_json::from_value(value).unwrap(); assert_eq!(decoded, msg); } + + #[test] + fn action_invoke_accepts_future_nested_task_execution_fields() { + let invoke: RelaycastToBroker = serde_json::from_value(json!({ + "type": "action.invoke", + "v": 1, + "invocation_id": "inv_task_1", + "action": "task.run", + "input": {}, + "task_execution": { + "execution_id": "inv_task_1/1", + "run_id": "run_1", + "step_id": "step_1", + "dispatch_id": "dispatch_1", + "deadline": "2026-09-15T12:00:00.000Z", + "future_engine_field": { "value": 1 } + } + })) + .expect("inbound task execution metadata must be forward compatible"); + + let RelaycastToBroker::ActionInvoke(invoke) = invoke else { + panic!("action invoke") + }; + assert_eq!(invoke.task_execution.unwrap().execution_id, "inv_task_1/1"); + } } diff --git a/crates/broker/src/listen_api.rs b/crates/broker/src/listen_api.rs index 6663acaf60..b7754fb406 100644 --- a/crates/broker/src/listen_api.rs +++ b/crates/broker/src/listen_api.rs @@ -280,12 +280,16 @@ impl std::error::Error for DeliveryRouteError {} #[derive(Debug, Clone, PartialEq, Eq)] pub enum AgentResultRouteError { InvalidToken, + Retryable, + Conflict, } impl std::fmt::Display for AgentResultRouteError { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { match self { AgentResultRouteError::InvalidToken => write!(f, "invalid_result_token"), + AgentResultRouteError::Retryable => write!(f, "task_receipt_pending"), + AgentResultRouteError::Conflict => write!(f, "task_result_conflict"), } } } @@ -1396,6 +1400,16 @@ async fn listen_api_agent_result( match reply_rx.await { Ok(Ok(value)) => (axum::http::StatusCode::OK, axum::Json(value)), + Ok(Err(AgentResultRouteError::Retryable)) => ( + axum::http::StatusCode::SERVICE_UNAVAILABLE, + axum::Json( + json!({ "success": false, "error": "task_receipt_pending", "retryable": true }), + ), + ), + Ok(Err(AgentResultRouteError::Conflict)) => ( + axum::http::StatusCode::CONFLICT, + axum::Json(json!({ "success": false, "error": "task_result_conflict" })), + ), Ok(Err(AgentResultRouteError::InvalidToken)) => ( axum::http::StatusCode::UNAUTHORIZED, axum::Json(json!({ "success": false, "error": "invalid_result_token" })), diff --git a/crates/broker/src/node_control.rs b/crates/broker/src/node_control.rs index b2d8488e34..114ff78d17 100644 --- a/crates/broker/src/node_control.rs +++ b/crates/broker/src/node_control.rs @@ -519,6 +519,7 @@ impl FleetLoadSnapshot { active_agent_names.sort(); active_agent_names.dedup(); capabilities.push(FleetCapability { + execution_mode: None, name: LIVE_AGENT_CAPABILITY_NAME.to_string(), kind: Some("capacity".to_string()), global: None, @@ -1306,6 +1307,7 @@ impl FleetDeliveryBook { pub(crate) fn handler_unavailable_result(invocation_id: &str) -> ActionResult { ActionResult { + task: None, v: FLEET_WIRE_VERSION, id: None, invocation_id: invocation_id.to_string(), @@ -1327,9 +1329,10 @@ pub(crate) fn build_node_register( .iter() .filter(|capability| capability.name != crate::fleet_wire::DELIVERY_CURSOR_CAPABILITY) .map(|capability| FleetCapability { + execution_mode: (capability.name == "task.run").then(|| "task".to_owned()), name: capability.name.clone(), kind: capability.kind.clone(), - global: None, + global: (capability.name == "task.run").then_some(true), queue: None, metadata: capability.metadata.as_ref().map(|metadata| { metadata @@ -1340,6 +1343,7 @@ pub(crate) fn build_node_register( }) .collect::>(); capabilities.push(FleetCapability { + execution_mode: None, name: crate::fleet_wire::DELIVERY_CURSOR_CAPABILITY.to_string(), kind: Some("capacity".to_string()), global: None, @@ -2635,6 +2639,14 @@ where } match frame { RelaycastToBroker::Reply(reply) => { + if reply.id.starts_with(crate::runtime::task_request_prefix()) { + return event_tx + .send(FleetControlEvent::Message(RelaycastToBroker::Reply( + reply, + ))) + .await + .is_ok(); + } if let Some(pending) = pending_deregistrations.remove(&reply.id) { let result = if reply.ok { Ok(()) @@ -2677,6 +2689,14 @@ where } } RelaycastToBroker::Error(error) => { + if error.id.starts_with(crate::runtime::task_request_prefix()) { + return event_tx + .send(FleetControlEvent::Message(RelaycastToBroker::Error( + error, + ))) + .await + .is_ok(); + } if let Some(pending) = pending_deregistrations.remove(&error.id) { let _ = pending.send(Err(format!("{}: {}", error.code, error.message))); @@ -4218,6 +4238,7 @@ mod tests { assert_eq!( register.capabilities.last(), Some(&FleetCapability { + execution_mode: None, name: crate::fleet_wire::DELIVERY_CURSOR_CAPABILITY.to_string(), kind: Some("capacity".to_string()), global: None, @@ -4319,6 +4340,7 @@ mod tests { ws.send(Message::Text( serde_json::to_string(&RelaycastToBroker::ActionInvoke(ActionInvoke { + task_execution: None, v: FLEET_WIRE_VERSION, invocation_id: "inv-1".to_string(), action: "run:test".to_string(), @@ -4361,6 +4383,7 @@ mod tests { command_tx .send(FleetControlCommand::Send(BrokerToRelaycast::ActionResult( ActionResult { + task: None, v: FLEET_WIRE_VERSION, id: None, invocation_id: "inv-1".to_string(), @@ -4827,6 +4850,7 @@ mod tests { result: ActionResultPayload::Output(ActionResultOutput { output: json!({"ok": true}), }), + task: None, }, ))) .await @@ -6195,6 +6219,42 @@ mod tests { consecutive_unauthorized.saturating_add(1) )); } + #[tokio::test] + async fn task_receipts_are_forwarded_without_consuming_agent_registration_waiters() { + let (tx, mut rx) = mpsc::channel(4); + let mut registrations = HashMap::new(); + let mut deregistrations = HashMap::new(); + let mut liveness = ApplicationLiveness::new(Duration::from_secs(1)); + let mut sink = futures_util::sink::drain(); + for raw in [ + serde_json::json!({"v":1,"type":"reply","id":"task_receipt_a","ok":true,"data":{"status":"running"}}), + serde_json::json!({"v":1,"type":"error","id":"task_receipt_b","ok":false,"code":"stale_task_execution","message":"stale"}), + ] { + assert!( + handle_server_message( + Message::Text(raw.to_string()), + &tx, + &mut registrations, + &mut deregistrations, + &mut liveness, + "node-test", + &mut sink, + None, + ) + .await + ); + let event = rx.recv().await.unwrap(); + match event { + FleetControlEvent::Message(RelaycastToBroker::Reply(reply)) => { + assert_eq!(reply.id, "task_receipt_a") + } + FleetControlEvent::Message(RelaycastToBroker::Error(error)) => { + assert_eq!(error.id, "task_receipt_b") + } + other => panic!("unexpected event {other:?}"), + } + } + } } #[cfg(test)] diff --git a/crates/broker/src/runtime/api.rs b/crates/broker/src/runtime/api.rs index 5df01b8e80..a85e30c99a 100644 --- a/crates/broker/src/runtime/api.rs +++ b/crates/broker/src/runtime/api.rs @@ -278,6 +278,23 @@ impl BrokerRuntime { } else { req }; + if let ListenApiRequest::SubmitAgentResult { token, .. } = &req { + if self.task_provider.store.by_token(token).is_some() { + if let ListenApiRequest::SubmitAgentResult { + token, + name, + data, + final_result, + metadata, + reply, + } = req + { + self.handle_task_callback(token, name, data, final_result, metadata, reply) + .await; + return; + } + } + } let local_only = self.degraded.is_some(); let paths = &self.paths; let state = &mut self.state; diff --git a/crates/broker/src/runtime/event_loop.rs b/crates/broker/src/runtime/event_loop.rs index 7db6849ac4..52bd76dd96 100644 --- a/crates/broker/src/runtime/event_loop.rs +++ b/crates/broker/src/runtime/event_loop.rs @@ -271,6 +271,7 @@ pub(crate) struct BrokerRuntime { pub(super) resize_owners: HashMap, pub(super) delivery_states: HashMap, pub(super) agent_result_tokens: HashMap, + pub(super) task_provider: super::tasks::TaskProvider, pub(super) recent_thread_messages: VecDeque, pub(super) shutdown: bool, pub(super) lease_duration: Option, diff --git a/crates/broker/src/runtime/fleet.rs b/crates/broker/src/runtime/fleet.rs index 340907da7a..0fc51c2d17 100644 --- a/crates/broker/src/runtime/fleet.rs +++ b/crates/broker/src/runtime/fleet.rs @@ -158,6 +158,7 @@ pub(super) fn verified_spawn_ready_result( name: &WorkerName, ) -> ActionResult { ActionResult { + task: None, v: FLEET_WIRE_VERSION, id: None, invocation_id, @@ -169,6 +170,7 @@ pub(super) fn verified_spawn_ready_result( pub(super) fn verified_spawn_failed_result(invocation_id: String, error: &str) -> ActionResult { ActionResult { + task: None, v: FLEET_WIRE_VERSION, id: None, invocation_id, @@ -853,9 +855,13 @@ impl BrokerRuntime { FleetControlEvent::Message(RelaycastToBroker::ActionInvoke(invoke)) => { self.handle_fleet_action_invoke(invoke).await; } - FleetControlEvent::Message(RelaycastToBroker::Ping(_)) - | FleetControlEvent::Message(RelaycastToBroker::Reply(_)) - | FleetControlEvent::Message(RelaycastToBroker::Error(_)) => {} + FleetControlEvent::Message(RelaycastToBroker::Reply(reply)) => { + self.handle_task_reply(reply).await + } + FleetControlEvent::Message(RelaycastToBroker::Error(error)) => { + self.handle_task_error(error).await + } + FleetControlEvent::Message(RelaycastToBroker::Ping(_)) => {} } } @@ -1296,6 +1302,10 @@ impl BrokerRuntime { // connection; the broker runs them directly against its PTY runtime. // Capability action handlers live in their own providers and are // dispatched to those sockets by the engine, never here. + if invoke.action == super::task_store::TASK_ACTION { + self.handle_task_invoke(invoke).await; + return; + } let action = invoke.action.as_str(); if action == "spawn" || action.starts_with("spawn:") { self.handle_fleet_action_spawn(invoke).await; @@ -1424,6 +1434,7 @@ impl BrokerRuntime { session_ref, &self.hosted_agent_event_tx, &mut self.pty_observability, + None, ) .await; @@ -1617,6 +1628,7 @@ impl BrokerRuntime { async fn reply_action_output(&self, invocation_id: &str, output: Value) { self.send_fleet_action_result(ActionResult { + task: None, v: FLEET_WIRE_VERSION, id: None, invocation_id: invocation_id.to_string(), @@ -1625,8 +1637,9 @@ impl BrokerRuntime { .await; } - async fn reply_action_error(&self, invocation_id: &str, error: &str) { + pub(super) async fn reply_action_error(&self, invocation_id: &str, error: &str) { self.send_fleet_action_result(ActionResult { + task: None, v: FLEET_WIRE_VERSION, id: None, invocation_id: invocation_id.to_string(), @@ -1646,7 +1659,7 @@ impl BrokerRuntime { .await; } - async fn publish_fleet_load(&self, heartbeat_now: bool) { + pub(super) async fn publish_fleet_load(&self, heartbeat_now: bool) { let active_agents = u32::try_from(self.workers.workers.len()).unwrap_or(u32::MAX); let active_agent_names = self .workers @@ -1832,6 +1845,7 @@ fn fleet_spawn_action_result( }), }; ActionResult { + task: None, v: FLEET_WIRE_VERSION, id: None, invocation_id: invocation_id.to_string(), @@ -3265,6 +3279,7 @@ mod tests { agent_id: Option<&str>, ) -> ActionInvoke { ActionInvoke { + task_execution: None, v: FLEET_WIRE_VERSION, invocation_id: "inv-1".to_string(), action: "spawn".to_string(), diff --git a/crates/broker/src/runtime/init.rs b/crates/broker/src/runtime/init.rs index b844627f90..890d590ce4 100644 --- a/crates/broker/src/runtime/init.rs +++ b/crates/broker/src/runtime/init.rs @@ -307,7 +307,29 @@ pub(crate) async fn run_init(cmd: InitCommand, telemetry: TelemetryClient) -> Re } else { resolve_cached_node_token(&node_id, &node_workspace_id, node_base_url.as_deref()) }; - let node_manifest = bootstrap_node_manifest(&node_name, &node_id, &broker_version); + let task_enabled = std::env::var("AGENT_RELAY_TASK_PROVIDER").as_deref() == Ok("1"); + anyhow::ensure!( + !task_enabled || (paths.persist && !local_only), + "task provider requires persistent hosted broker mode" + ); + let task_provider = if task_enabled { + super::tasks::TaskProvider { + store: super::task_store::TaskStore::open(paths.state.with_extension("tasks.json"))?, + ..Default::default() + } + } else { + super::tasks::TaskProvider::default() + }; + let mut node_manifest = bootstrap_node_manifest(&node_name, &node_id, &broker_version); + if task_enabled { + node_manifest + .capabilities + .push(crate::protocol::NodeCapabilityManifest { + name: super::task_store::TASK_ACTION.to_owned(), + kind: Some("action".to_owned()), + metadata: None, + }); + } // Retain the node name for the runtime: the HTTP `bind_agent_to_node` // fallback (used when node-control `agent.register` is unavailable) binds // spawned agents to this node so they become `via_node` and node delivery @@ -804,6 +826,7 @@ pub(crate) async fn run_init(cmd: InitCommand, telemetry: TelemetryClient) -> Re resize_owners: HashMap::new(), delivery_states, agent_result_tokens, + task_provider, recent_thread_messages, shutdown, lease_duration, diff --git a/crates/broker/src/runtime/maintenance.rs b/crates/broker/src/runtime/maintenance.rs index 2aa92224d6..2a0983c327 100644 --- a/crates/broker/src/runtime/maintenance.rs +++ b/crates/broker/src/runtime/maintenance.rs @@ -4,6 +4,7 @@ use crate::terminal_control::TerminalToCloud; impl BrokerRuntime { pub(super) async fn handle_maintenance_tick(&mut self) { + self.maintain_tasks().await; self.reconcile_identity_cleanups().await; let paths = &self.paths; let state = &mut self.state; diff --git a/crates/broker/src/runtime/mod.rs b/crates/broker/src/runtime/mod.rs index af0e59e4a4..fb9d1c949d 100644 --- a/crates/broker/src/runtime/mod.rs +++ b/crates/broker/src/runtime/mod.rs @@ -88,6 +88,8 @@ mod relaycast_events; mod session; mod spawn_spec; mod system; +mod task_store; +mod tasks; #[cfg(test)] mod tests; mod util; @@ -115,3 +117,7 @@ pub(crate) use session::*; pub(crate) use spawn_spec::*; pub(crate) use system::*; pub(crate) use util::*; + +pub(crate) fn task_request_prefix() -> &'static str { + task_store::TASK_REQUEST_PREFIX +} diff --git a/crates/broker/src/runtime/relaycast_events.rs b/crates/broker/src/runtime/relaycast_events.rs index ce794e53f6..23f5702172 100644 --- a/crates/broker/src/runtime/relaycast_events.rs +++ b/crates/broker/src/runtime/relaycast_events.rs @@ -514,6 +514,7 @@ pub(super) async fn spawn_worker_from_request( session_ref: Option, hosted_agent_event_tx: &mpsc::Sender, pty_observability: &mut HashMap, + task_binding: Option<(AgentResultMcpConfig, Uuid)>, ) -> Result<()> { if workers.identity_cleanups.contains_key(&name) { anyhow::bail!("worker name has pending owned cleanup; complete it before reuse"); @@ -839,15 +840,16 @@ pub(super) async fn spawn_worker_from_request( }; match workers - .spawn( + .spawn_with_generation( spec, Some("Relaycast".to_string()), None, worker_relay_key.clone(), false, Some(workspace_id.clone()), - None, + task_binding.as_ref().map(|(config, _)| config.clone()), commit_attestation, + task_binding.as_ref().map(|(_, generation)| *generation), ) .await { @@ -1280,6 +1282,7 @@ mod tests { None, &hosted_agent_event_tx, &mut pty_observability, + None, ) .await .expect_err("a sidecar that exits during the stability window must fail the spawn"); diff --git a/crates/broker/src/runtime/task_store.rs b/crates/broker/src/runtime/task_store.rs new file mode 100644 index 0000000000..72db21bfd7 --- /dev/null +++ b/crates/broker/src/runtime/task_store.rs @@ -0,0 +1,771 @@ +//! Durable, single-writer task ownership/outbox. RuntimePaths holds the broker lock. +use super::*; +use crate::fleet_wire::{ + ActionAccept, ActionInvoke, ActionResult, ActionResultError, ActionResultOutput, + ActionResultPayload, TaskResultFields, FLEET_WIRE_VERSION, +}; +use std::collections::BTreeMap; + +pub(super) const TASK_ACTION: &str = "task.run"; +pub(super) const TASK_REQUEST_PREFIX: &str = "task_receipt_"; +const TASK_TERMINAL_RETENTION_GRACE_SECS: i64 = 24 * 60 * 60; + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub(super) struct TaskFinal { + pub output: Option, + pub error: Option, + pub accounting: Option>, +} + +// JavaScript JSON receipts can normalize 17.0 to 17. Compare JSON numbers +// by their wire numeric value, not serde_json's internal integer/float variant. +pub(super) fn json_equal(left: &Value, right: &Value) -> bool { + match (left, right) { + (Value::Number(a), Value::Number(b)) => a.as_f64() == b.as_f64(), + (Value::Array(a), Value::Array(b)) => { + a.len() == b.len() && a.iter().zip(b).all(|(a, b)| json_equal(a, b)) + } + (Value::Object(a), Value::Object(b)) => { + a.len() == b.len() + && a.iter() + .all(|(key, value)| b.get(key).is_some_and(|other| json_equal(value, other))) + } + _ => left == right, + } +} +impl PartialEq for TaskFinal { + fn eq(&self, other: &Self) -> bool { + self.error == other.error + && self.accounting == other.accounting + && match (&self.output, &other.output) { + (Some(a), Some(b)) => json_equal(a, b), + (None, None) => true, + _ => false, + } + } +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub(super) struct TaskRecord { + pub invoke: ActionInvoke, + pub name: WorkerName, + pub generation: Uuid, + pub callback_token: String, + /// Written before any attempt to create a process. Never cleared on restart. + pub launch_claimed: bool, + pub final_result: Option, + pub receipt: Option, + #[serde(default)] + pub rejection: Option, +} + +impl TaskRecord { + pub fn execution(&self) -> &crate::fleet_wire::TaskExecution { + self.invoke + .task_execution + .as_ref() + .expect("validated task record") + } + pub fn accept(&self, id: String) -> ActionAccept { + ActionAccept { + v: FLEET_WIRE_VERSION, + id, + invocation_id: self.invoke.invocation_id.clone(), + execution_id: self.execution().execution_id.clone(), + worker_generation: self.generation.to_string(), + } + } + pub fn result(&self, id: String, value: &TaskFinal, final_result: bool) -> ActionResult { + ActionResult { + v: FLEET_WIRE_VERSION, + id: Some(id), + invocation_id: self.invoke.invocation_id.clone(), + result: match &value.error { + Some(error) => ActionResultPayload::Error(ActionResultError { + error: error.clone(), + }), + None => ActionResultPayload::Output(ActionResultOutput { + output: value.output.clone().unwrap_or(Value::Null), + }), + }, + task: Some(TaskResultFields { + final_result, + execution_id: self.execution().execution_id.clone(), + worker_generation: self.generation.to_string(), + accounting: value.accounting.as_ref().map(|values| { + values + .iter() + .map(|(key, value)| { + ( + key.clone(), + serde_json::Number::from_f64(*value) + .expect("validated finite accounting"), + ) + }) + .collect() + }), + }), + } + } + pub fn expired(&self) -> bool { + chrono::DateTime::parse_from_rfc3339(&self.execution().deadline) + .map_or(true, |deadline| deadline <= chrono::Utc::now()) + } +} + +#[derive(Default)] +pub(super) struct TaskStore { + path: Option, + poisoned: bool, + pub records: BTreeMap, +} + +impl TaskStore { + pub fn open(path: PathBuf) -> Result { + let mut records: BTreeMap = match std::fs::read(&path) { + Ok(bytes) => serde_json::from_slice(&bytes).context("invalid durable task ledger")?, + Err(error) if error.kind() == std::io::ErrorKind::NotFound => BTreeMap::new(), + Err(error) => return Err(error.into()), + }; + for (id, record) in &records { + anyhow::ensure!( + id == &record.invoke.invocation_id, + "task ledger identity mismatch" + ); + validate_invoke(&record.invoke)?; + if let Some(receipt) = &record.receipt { + validate_receipt(record, receipt)?; + anyhow::ensure!( + matches!( + receipt.get("status").and_then(Value::as_str), + Some("completed" | "failed") + ), + "nonterminal persisted task receipt" + ); + } + } + let mut compacted = records.clone(); + if compact_terminal_records(&mut compacted, chrono::Utc::now()) { + match persist_records(&path, &compacted) { + Ok(()) => records = compacted, + Err(error) => tracing::warn!( + error = %error, + "terminal task retention could not be persisted during startup" + ), + } + } + Ok(Self { + path: Some(path), + records, + poisoned: false, + }) + } + pub fn enabled(&self) -> bool { + self.path.is_some() && !self.poisoned + } + fn put(&mut self, record: TaskRecord) -> Result<()> { + anyhow::ensure!( + !self.poisoned, + "task ledger requires reopen after a failed durable write" + ); + let mut next = self.records.clone(); + next.insert(record.invoke.invocation_id.clone(), record); + self.replace(next) + } + fn replace(&mut self, next: BTreeMap) -> Result<()> { + anyhow::ensure!( + !self.poisoned, + "task ledger requires reopen after a failed durable write" + ); + let path = self + .path + .as_ref() + .context("durable task provider is disabled")?; + if let Err(error) = persist_records(path, &next) { + self.poisoned = true; + return Err(error); + } + self.records = next; + Ok(()) + } + pub fn prepare(&mut self, invoke: ActionInvoke) -> Result { + anyhow::ensure!(self.enabled(), "durable task provider is unavailable"); + validate_invoke(&invoke)?; + if let Some(existing) = self.records.get(&invoke.invocation_id) { + if existing.invoke == invoke { + anyhow::ensure!(existing.rejection.is_none(), "task invocation was rejected"); + return Ok(existing.clone()); + } + let mut prior = existing.invoke.clone(); + prior.task_execution = invoke.task_execution.clone(); + anyhow::ensure!( + prior == invoke + && existing.execution().run_id + == invoke.task_execution.as_ref().unwrap().run_id + && existing.execution().step_id + == invoke.task_execution.as_ref().unwrap().step_id + && existing.execution().dispatch_id + == invoke.task_execution.as_ref().unwrap().dispatch_id + && existing.execution().deadline + == invoke.task_execution.as_ref().unwrap().deadline, + "task_invocation_conflict" + ); + anyhow::ensure!( + !existing.launch_claimed + && existing.final_result.is_none() + && existing.receipt.is_none(), + "task_execution_already_claimed" + ); + } + let generation = Uuid::new_v4(); + let name = WorkerName::new(format!("task-{}", &generation.simple().to_string()[..16])); + let record = TaskRecord { + invoke, + name, + generation, + callback_token: format!("arr_{}", Uuid::new_v4().simple()), + launch_claimed: false, + final_result: None, + receipt: None, + rejection: None, + }; + self.put(record.clone())?; + Ok(record) + } + pub fn by_token(&self, token: &str) -> Option<&TaskRecord> { + self.records + .values() + .find(|record| record.callback_token == token) + } + pub fn claim_launch(&mut self, id: &str) -> Result> { + let mut record = self.records.get(id).context("unknown task")?.clone(); + if record.launch_claimed + || record.final_result.is_some() + || record.receipt.is_some() + || record.expired() + { + return Ok(None); + } + record.launch_claimed = true; + self.put(record.clone())?; + Ok(Some(record)) + } + pub fn queue_final(&mut self, id: &str, value: TaskFinal) -> Result { + anyhow::ensure!( + value.error.is_some() != value.output.is_some(), + "task result needs output or error" + ); + anyhow::ensure!( + value + .error + .as_ref() + .is_none_or(|error| !error.trim().is_empty()), + "empty task failure" + ); + anyhow::ensure!( + value + .accounting + .as_ref() + .is_none_or(|a| a.values().all(|n| n.is_finite() && *n >= 0.0)), + "invalid task accounting" + ); + let mut record = self.records.get(id).context("unknown task")?.clone(); + if let Some(existing) = &record.final_result { + anyhow::ensure!(existing == &value, "task_result_conflict"); + return Ok(record); + } + anyhow::ensure!(record.receipt.is_none(), "task_already_terminal"); + record.final_result = Some(value); + self.put(record.clone())?; + Ok(record) + } + pub fn reject(&mut self, id: &str, reason: String) -> Result<()> { + let mut record = self.records.get(id).context("unknown task")?.clone(); + record.rejection = Some(reason); + self.put(record) + } + pub fn finish(&mut self, id: &str, mut receipt: Value) -> Result { + let mut record = self.records.get(id).context("unknown task")?.clone(); + validate_receipt(&record, &receipt)?; + if let Some(object) = receipt.as_object_mut() { + object.remove("newly_accepted"); + } + anyhow::ensure!( + matches!( + receipt.get("status").and_then(Value::as_str), + Some("completed" | "failed") + ), + "nonterminal task receipt" + ); + if let Some(previous) = &record.receipt { + anyhow::ensure!(json_equal(previous, &receipt), "task_receipt_conflict"); + return Ok(record); + } + record.receipt = Some(receipt); + let mut next = self.records.clone(); + next.insert(record.invoke.invocation_id.clone(), record.clone()); + compact_terminal_records(&mut next, chrono::Utc::now()); + self.replace(next)?; + Ok(record) + } + + pub fn compact(&mut self) -> Result<()> { + let now = chrono::Utc::now(); + if self + .records + .values() + .any(|record| terminal_record_expired(record, now)) + { + let mut next = self.records.clone(); + compact_terminal_records(&mut next, now); + self.replace(next)?; + } + Ok(()) + } +} + +fn compact_terminal_records( + records: &mut BTreeMap, + now: chrono::DateTime, +) -> bool { + let before = records.len(); + records.retain(|_, record| !terminal_record_expired(record, now)); + records.len() != before +} + +fn terminal_record_expired(record: &TaskRecord, now: chrono::DateTime) -> bool { + // A rejected invocation (stale_task_execution, task_not_found, + // task_result_conflict) is terminal too, but `reject` only sets + // `rejection`, never `receipt` -- a receipt-only check here would let + // rejected records accumulate in the ledger forever. + (record.receipt.is_some() || record.rejection.is_some()) + && chrono::DateTime::parse_from_rfc3339(&record.execution().deadline).is_ok_and( + |deadline| { + deadline.with_timezone(&chrono::Utc) + + chrono::Duration::seconds(TASK_TERMINAL_RETENTION_GRACE_SECS) + <= now + }, + ) +} + +fn persist_records(path: &Path, records: &BTreeMap) -> Result<()> { + crate::util::fs::write_json_atomic(path, records)?; + // Task receipts require the rename itself to be durable, not best effort. + #[cfg(unix)] + std::fs::File::open(path.parent().context("task ledger has no parent")?)?.sync_all()?; + Ok(()) +} + +pub(super) fn validate_invoke(invoke: &ActionInvoke) -> Result<()> { + anyhow::ensure!( + invoke.action == TASK_ACTION && !invoke.invocation_id.is_empty(), + "invalid task action" + ); + let execution = invoke + .task_execution + .as_ref() + .context("missing task execution fence")?; + anyhow::ensure!( + !execution.execution_id.is_empty(), + "missing execution identity" + ); + chrono::DateTime::parse_from_rfc3339(&execution.deadline).context("invalid task deadline")?; + let context = invoke + .input + .get("task_context") + .context("missing task context")?; + for (key, expected) in [ + ("run_id", &execution.run_id), + ("step_id", &execution.step_id), + ("dispatch_id", &execution.dispatch_id), + ] { + anyhow::ensure!( + !expected.is_empty() + && context.get(key).and_then(Value::as_str) == Some(expected.as_str()), + "task correlation mismatch" + ); + } + Ok(()) +} + +pub(super) fn validate_receipt(record: &TaskRecord, receipt: &Value) -> Result<()> { + anyhow::ensure!( + receipt.get("invocation_id").and_then(Value::as_str) + == Some(record.invoke.invocation_id.as_str()), + "receipt invocation mismatch" + ); + anyhow::ensure!( + receipt.get("action_name").and_then(Value::as_str) == Some(TASK_ACTION), + "receipt action mismatch" + ); + let execution = receipt + .get("task_execution") + .context("receipt missing execution")?; + anyhow::ensure!( + execution.get("execution_id").and_then(Value::as_str) + == Some(record.execution().execution_id.as_str()), + "receipt execution mismatch" + ); + for (key, value) in [ + ("run_id", &record.execution().run_id), + ("step_id", &record.execution().step_id), + ("dispatch_id", &record.execution().dispatch_id), + ("deadline", &record.execution().deadline), + ] { + anyhow::ensure!( + execution.get(key).and_then(Value::as_str) == Some(value.as_str()), + "receipt correlation mismatch" + ); + } + let generation = record.generation.to_string(); + // An invocation can expire before acceptance, when no generation is bound. + let expired_unaccepted = receipt.get("status").and_then(Value::as_str) == Some("failed") + && receipt.get("error").and_then(Value::as_str) == Some("task_deadline_exceeded") + && execution.get("worker_generation").is_none(); + anyhow::ensure!( + expired_unaccepted + || execution.get("worker_generation").and_then(Value::as_str) + == Some(generation.as_str()), + "receipt generation mismatch" + ); + Ok(()) +} + +#[cfg(test)] +pub(super) fn fixture_invoke() -> ActionInvoke { + serde_json::from_value(json!({"v":1,"type":"action.invoke","invocation_id":"inv-task","action":"task.run", + "input":{"task":"answer","task_context":{"run_id":"run","step_id":"step","dispatch_id":"dispatch","timeout_ms":120000}}, + "task_execution":{"execution_id":"inv-task/1","run_id":"run","step_id":"step","dispatch_id":"dispatch", + "deadline":(chrono::Utc::now()+chrono::Duration::minutes(2)).to_rfc3339_opts(chrono::SecondsFormat::Millis,true)}}) + .as_object().unwrap().iter().filter(|(key,_)|key.as_str() != "type").map(|(k,v)|(k.clone(),v.clone())).collect::>().into()).unwrap() +} + +#[cfg(test)] +pub(super) fn fixture_receipt(record: &TaskRecord, status: &str) -> Value { + let mut execution = serde_json::to_value(record.execution()).unwrap(); + execution["worker_generation"] = json!(record.generation.to_string()); + execution["accepted_at"] = json!("2026-01-01T00:00:00.000Z"); + if let Some(accounting) = record + .final_result + .as_ref() + .and_then(|result| result.accounting.as_ref()) + { + execution["accounting"] = json!(accounting); + } + json!({"invocation_id":record.invoke.invocation_id,"action_name":"task.run","status":status, + "task_execution":execution,"output":record.final_result.as_ref().and_then(|result|result.output.clone()), + "error":record.final_result.as_ref().and_then(|result|result.error.clone()), + "completed_at":if status == "running" { Value::Null } else { json!("2026-01-01T00:01:00.000Z") }}) +} + +#[cfg(test)] +mod tests { + use super::*; + fn value() -> TaskFinal { + TaskFinal { + output: Some(json!({"answer":42})), + error: None, + accounting: None, + } + } + fn invoke_with_deadline(id: &str, deadline: chrono::DateTime) -> ActionInvoke { + let mut invoke = fixture_invoke(); + invoke.invocation_id = id.to_owned(); + let execution = invoke.task_execution.as_mut().unwrap(); + execution.execution_id = format!("{id}/1"); + execution.deadline = deadline.to_rfc3339_opts(chrono::SecondsFormat::Millis, true); + invoke + } + #[test] + fn task_store_prepared_and_claimed_launch_survive_reopen_without_a_second_claim() { + let directory = tempfile::tempdir().unwrap(); + let path = directory.path().join("tasks.json"); + let mut store = TaskStore::open(path.clone()).unwrap(); + let original = store.prepare(fixture_invoke()).unwrap(); + let mut reopened = TaskStore::open(path.clone()).unwrap(); + assert_eq!( + reopened + .prepare(original.invoke.clone()) + .unwrap() + .generation, + original.generation + ); + assert!(reopened + .claim_launch(&original.invoke.invocation_id) + .unwrap() + .is_some()); + let mut restarted = TaskStore::open(path).unwrap(); + assert!(restarted + .claim_launch(&original.invoke.invocation_id) + .unwrap() + .is_none()); + assert_eq!( + restarted.records[&original.invoke.invocation_id].callback_token, + original.callback_token + ); + } + #[test] + fn task_store_changed_input_and_claimed_attempt_cannot_replace_execution() { + let directory = tempfile::tempdir().unwrap(); + let mut store = TaskStore::open(directory.path().join("tasks.json")).unwrap(); + let original = store.prepare(fixture_invoke()).unwrap(); + let mut changed = original.invoke.clone(); + changed.input["task"] = json!("different"); + assert!(store.prepare(changed).is_err()); + let mut retried = original.invoke.clone(); + retried.task_execution.as_mut().unwrap().execution_id = "inv-task/2".into(); + let replacement = store.prepare(retried).unwrap(); + assert_ne!(replacement.generation, original.generation); + store + .claim_launch(&replacement.invoke.invocation_id) + .unwrap(); + assert!(store.prepare(original.invoke).is_err()); + assert_eq!(store.records["inv-task"].generation, replacement.generation); + } + #[test] + fn task_store_final_outbox_and_receipt_reconcile_after_restart_and_ack_loss() { + let directory = tempfile::tempdir().unwrap(); + let path = directory.path().join("tasks.json"); + let mut store = TaskStore::open(path.clone()).unwrap(); + let original = store.prepare(fixture_invoke()).unwrap(); + store.claim_launch(&original.invoke.invocation_id).unwrap(); + let queued = store + .queue_final(&original.invoke.invocation_id, value()) + .unwrap(); + let mut restarted = TaskStore::open(path.clone()).unwrap(); + assert_eq!( + restarted + .queue_final("inv-task", value()) + .unwrap() + .final_result, + queued.final_result + ); + assert!(restarted + .queue_final( + "inv-task", + TaskFinal { + output: Some(json!(43)), + ..value() + } + ) + .is_err()); + let receipt = fixture_receipt(&queued, "completed"); + restarted.finish("inv-task", receipt.clone()).unwrap(); + let mut reopened = TaskStore::open(path).unwrap(); + let mut accept_replay = receipt.clone(); + accept_replay["newly_accepted"] = json!(false); + assert_eq!( + reopened.finish("inv-task", accept_replay).unwrap().receipt, + Some(receipt) + ); + } + #[test] + fn task_store_rejects_wrong_receipt_fence_generation_or_correlation() { + let directory = tempfile::tempdir().unwrap(); + let mut store = TaskStore::open(directory.path().join("tasks.json")).unwrap(); + let original = store.prepare(fixture_invoke()).unwrap(); + for field in [ + "execution_id", + "worker_generation", + "run_id", + "step_id", + "dispatch_id", + "deadline", + ] { + let mut receipt = fixture_receipt(&original, "completed"); + receipt["task_execution"][field] = json!("changed"); + assert!(store.finish("inv-task", receipt).is_err()); + } + assert!(store.records["inv-task"].receipt.is_none()); + } + #[test] + fn task_store_open_rejects_malformed_or_nonterminal_persisted_receipts() { + let directory = tempfile::tempdir().unwrap(); + let path = directory.path().join("tasks.json"); + let mut store = TaskStore::open(path.clone()).unwrap(); + let original = store.prepare(fixture_invoke()).unwrap(); + + for receipt in [ + json!({"status":"completed"}), + fixture_receipt(&original, "running"), + ] { + let mut record = original.clone(); + record.receipt = Some(receipt); + persist_records( + &path, + &BTreeMap::from([(record.invoke.invocation_id.clone(), record)]), + ) + .unwrap(); + assert!(TaskStore::open(path.clone()).is_err()); + } + } + #[test] + fn task_store_rejected_identical_invoke_stays_rejected() { + let directory = tempfile::tempdir().unwrap(); + let path = directory.path().join("tasks.json"); + let mut store = TaskStore::open(path).unwrap(); + let record = store.prepare(fixture_invoke()).unwrap(); + store + .reject("inv-task", "task_not_found".to_owned()) + .unwrap(); + + assert!(store.prepare(record.invoke).is_err()); + assert_eq!( + store.records["inv-task"].rejection.as_deref(), + Some("task_not_found") + ); + } + #[test] + fn task_store_failed_disk_write_never_claims_launch_and_requires_reopen() { + let directory = tempfile::tempdir().unwrap(); + let path = directory.path().join("tasks.json"); + let mut store = TaskStore::open(path.clone()).unwrap(); + let original = store.prepare(fixture_invoke()).unwrap(); + // Replacing the ledger with a directory makes atomic rename fail deterministically. + std::fs::remove_file(&path).unwrap(); + std::fs::create_dir(&path).unwrap(); + assert!(store.claim_launch("inv-task").is_err()); + assert!(!store.enabled()); + assert!(!store.records["inv-task"].launch_claimed); + assert!(store.prepare(original.invoke).is_err()); + } + #[test] + fn task_store_corrupt_ledger_fails_closed_and_secrets_are_private() { + let directory = tempfile::tempdir().unwrap(); + let path = directory.path().join("tasks.json"); + let mut store = TaskStore::open(path.clone()).unwrap(); + store.prepare(fixture_invoke()).unwrap(); + #[cfg(unix)] + { + use std::os::unix::fs::PermissionsExt; + assert_eq!( + std::fs::metadata(&path).unwrap().permissions().mode() & 0o777, + 0o600 + ); + } + std::fs::write(&path, "not json").unwrap(); + assert!(TaskStore::open(path).is_err()); + } + + #[test] + fn task_store_open_durably_prunes_terminal_receipts_past_the_grace_period() { + let directory = tempfile::tempdir().unwrap(); + let path = directory.path().join("tasks.json"); + let mut store = TaskStore::open(path.clone()).unwrap(); + let mut record = store + .prepare(invoke_with_deadline( + "inv-old", + chrono::Utc::now() - chrono::Duration::hours(25), + )) + .unwrap(); + record.receipt = Some(fixture_receipt(&record, "completed")); + let records = BTreeMap::from([(record.invoke.invocation_id.clone(), record)]); + persist_records(&path, &records).unwrap(); + + let reopened = TaskStore::open(path.clone()).unwrap(); + assert!(reopened.records.is_empty()); + let persisted: BTreeMap = + serde_json::from_slice(&std::fs::read(path).unwrap()).unwrap(); + assert!(persisted.is_empty()); + } + + #[test] + fn task_store_open_durably_prunes_rejected_records_past_the_grace_period() { + // `reject` only ever sets `rejection`, never `receipt` -- a + // receipt-only expiry check would let rejected invocations + // (stale_task_execution, task_not_found, task_result_conflict) + // accumulate in the ledger forever. + let directory = tempfile::tempdir().unwrap(); + let path = directory.path().join("tasks.json"); + let mut store = TaskStore::open(path.clone()).unwrap(); + let mut record = store + .prepare(invoke_with_deadline( + "inv-old", + chrono::Utc::now() - chrono::Duration::hours(25), + )) + .unwrap(); + record.rejection = Some("stale_task_execution".to_owned()); + assert!(record.receipt.is_none()); + let records = BTreeMap::from([(record.invoke.invocation_id.clone(), record)]); + persist_records(&path, &records).unwrap(); + + let reopened = TaskStore::open(path.clone()).unwrap(); + assert!(reopened.records.is_empty()); + let persisted: BTreeMap = + serde_json::from_slice(&std::fs::read(path).unwrap()).unwrap(); + assert!(persisted.is_empty()); + } + + #[cfg(unix)] + #[test] + fn task_store_open_keeps_valid_ledger_when_startup_compaction_cannot_persist() { + use std::os::unix::fs::PermissionsExt; + + if unsafe { nix::libc::getuid() } == 0 { + eprintln!("skipping unwritable ledger test while running as root"); + return; + } + let directory = tempfile::tempdir().unwrap(); + let ledger_dir = directory.path().join("ledger"); + std::fs::create_dir(&ledger_dir).unwrap(); + let path = ledger_dir.join("tasks.json"); + let mut store = TaskStore::open(path.clone()).unwrap(); + let mut record = store + .prepare(invoke_with_deadline( + "inv-old", + chrono::Utc::now() - chrono::Duration::hours(25), + )) + .unwrap(); + record.receipt = Some(fixture_receipt(&record, "completed")); + let records = BTreeMap::from([(record.invoke.invocation_id.clone(), record)]); + persist_records(&path, &records).unwrap(); + let persisted = std::fs::read(&path).unwrap(); + + std::fs::set_permissions(&ledger_dir, std::fs::Permissions::from_mode(0o500)).unwrap(); + let reopened = TaskStore::open(path.clone()); + std::fs::set_permissions(&ledger_dir, std::fs::Permissions::from_mode(0o700)).unwrap(); + + let mut reopened = reopened.expect("opportunistic compaction must not block startup"); + assert!(reopened.enabled()); + assert!(reopened.records.contains_key("inv-old")); + assert_eq!(std::fs::read(&path).unwrap(), persisted); + reopened.compact().unwrap(); + assert!(reopened.records.is_empty()); + } + + #[test] + fn task_store_finish_keeps_grace_replays_and_compacts_older_terminals() { + let directory = tempfile::tempdir().unwrap(); + let path = directory.path().join("tasks.json"); + let mut store = TaskStore::open(path.clone()).unwrap(); + let recent = store + .prepare(invoke_with_deadline( + "inv-recent", + chrono::Utc::now() - chrono::Duration::hours(23), + )) + .unwrap(); + store + .finish("inv-recent", fixture_receipt(&recent, "completed")) + .unwrap(); + assert!(store.records.contains_key("inv-recent")); + + let old = store + .prepare(invoke_with_deadline( + "inv-old", + chrono::Utc::now() - chrono::Duration::hours(25), + )) + .unwrap(); + let finished = store + .finish("inv-old", fixture_receipt(&old, "completed")) + .unwrap(); + assert!(finished.receipt.is_some()); + assert!(!store.records.contains_key("inv-old")); + assert!(store.records.contains_key("inv-recent")); + + let reopened = TaskStore::open(path).unwrap(); + assert_eq!(reopened.records.len(), 1); + assert!(reopened.records.contains_key("inv-recent")); + } +} diff --git a/crates/broker/src/runtime/tasks.rs b/crates/broker/src/runtime/tasks.rs new file mode 100644 index 0000000000..f02e549ef6 --- /dev/null +++ b/crates/broker/src/runtime/tasks.rs @@ -0,0 +1,626 @@ +use super::task_store::{ + json_equal, validate_receipt, TaskFinal, TaskRecord, TaskStore, TASK_REQUEST_PREFIX, +}; +use super::*; +use crate::fleet_wire::{ActionInvoke, BrokerToRelaycast, Reply}; +use crate::listen_api::AgentResultRouteError; +use tokio::sync::oneshot; + +type CallbackReply = oneshot::Sender>; +#[derive(Clone, Copy, PartialEq, Eq)] +pub(super) enum TaskRequestKind { + Accept, + Final, + Interim, +} +pub(super) struct PendingTaskRequest { + pub invocation: String, + pub execution: String, + pub kind: TaskRequestKind, + pub sent_at: Instant, + pub interim_reply: Option, + pub interim_data: Option, +} +#[derive(Default)] +pub(super) struct TaskProvider { + pub store: TaskStore, + pub pending: HashMap, + pub callbacks: HashMap>, + pub last_retry: Option, + pub retry_cursor: usize, +} + +impl BrokerRuntime { + pub(super) async fn handle_task_invoke(&mut self, invoke: ActionInvoke) { + let invocation_id = invoke.invocation_id.clone(); + match self.task_provider.store.prepare(invoke) { + Ok(record) => { + self.send_task_request(&record, TaskRequestKind::Accept, None, None) + .await + } + Err(error) => { + tracing::warn!(error = %error, "task invocation refused before launch"); + self.reply_action_error(&invocation_id, "handler_unavailable") + .await; + } + } + } + + async fn send_task_request( + &mut self, + record: &TaskRecord, + kind: TaskRequestKind, + interim: Option, + reply: Option, + ) { + if !self.node_delivery_connected + || !self.task_provider.store.enabled() + || self.task_provider.pending.len() >= 64 + { + if let Some(reply) = reply { + let _ = reply.send(Err(AgentResultRouteError::Retryable)); + } + return; + } + if kind != TaskRequestKind::Interim + && self.task_provider.pending.values().any(|request| { + request.invocation == record.invoke.invocation_id + && request.execution == record.execution().execution_id + && request.kind == kind + }) + { + return; + } + let id = format!("{TASK_REQUEST_PREFIX}{}", Uuid::new_v4().simple()); + let message = if kind == TaskRequestKind::Accept { + BrokerToRelaycast::ActionAccept(record.accept(id.clone())) + } else { + let value = interim + .as_ref() + .or(record.final_result.as_ref()) + .expect("result request has payload"); + BrokerToRelaycast::ActionResult(record.result( + id.clone(), + value, + kind == TaskRequestKind::Final, + )) + }; + // Never await an engine reply in the broker loop. The fleet reader returns + // the correlated response as an event; the callback waiter is separate. + let pending = PendingTaskRequest { + invocation: record.invoke.invocation_id.clone(), + execution: record.execution().execution_id.clone(), + kind, + sent_at: Instant::now(), + interim_data: interim.and_then(|value| value.output), + interim_reply: reply, + }; + if self + .fleet_control_tx + .try_send(FleetControlCommand::Send(message)) + .is_ok() + { + self.task_provider.pending.insert(id, pending); + } else if let Some(reply) = pending.interim_reply { + let _ = reply.send(Err(AgentResultRouteError::Retryable)); + } + } + + pub(super) async fn handle_task_reply(&mut self, reply: Reply) { + let Some(request) = self.task_provider.pending.remove(&reply.id) else { + return; + }; + let Some(record) = self + .task_provider + .store + .records + .get(&request.invocation) + .cloned() + else { + return; + }; + if record.execution().execution_id != request.execution { + return; + } + let data = reply.data; + if !reply.ok || validate_receipt(&record, &data).is_err() { + if let Some(callback) = request.interim_reply { + let _ = callback.send(Err(AgentResultRouteError::Retryable)); + } + return; + } + let status = data.get("status").and_then(Value::as_str).unwrap_or(""); + if matches!(status, "completed" | "failed") { + let first_receipt = record.receipt.is_none(); + match self.task_provider.store.finish(&request.invocation, data) { + Ok(record) => { + if first_receipt { + let _ = timeout(Duration::from_millis(200), send_event(&self.sdk_out_tx, json!({ + "kind":"agent_result", "name":record.name, "generation":record.generation.to_string(), + "result_id":record.invoke.invocation_id, "data":record.receipt.as_ref().and_then(|r|r.get("output")), + "final":true, "task_receipt":record.receipt + }))).await; + } + if let Some(callback) = request.interim_reply { + let _ = callback.send(Err(AgentResultRouteError::Conflict)); + } + self.finish_task_callbacks(&record); + if record + .receipt + .as_ref() + .and_then(|r| r.get("status")) + .and_then(Value::as_str) + == Some("failed") + { + let _ = self + .workers + .stop_task_generation(record.name.as_str(), record.generation) + .await; + } + } + Err(error) => { + tracing::warn!(error = %error, "task receipt not acknowledged because local durable write failed") + } + } + return; + } + if request.kind == TaskRequestKind::Interim { + let _ = timeout(Duration::from_millis(200), send_event(&self.sdk_out_tx, json!({ + "kind":"agent_result", "name":record.name, "generation":record.generation.to_string(), + "result_id":reply.id, "data":request.interim_data, "final":false + }))).await; + if let Some(callback) = request.interim_reply { + let _ = callback.send(Ok( + json!({"success": true, "name": record.name, "final": false, "receipt": data}), + )); + } + return; + } + if request.kind != TaskRequestKind::Accept || status != "running" { + return; + } + if record.final_result.is_some() { + self.send_task_request(&record, TaskRequestKind::Final, None, None) + .await; + return; + } + if record.expired() { + self.fail_task(&request.invocation, "task_deadline_exceeded") + .await; + return; + } + if record.launch_claimed { + let live = self + .workers + .workers + .get(&record.name) + .is_some_and(|worker| { + worker.generation == record.generation + && self.workers.is_worker_live(&record.name) + }); + if !live { + self.fail_task(&request.invocation, "worker_execution_lost") + .await; + } + return; + } + // Even newly_accepted=false can reconcile a never-claimed launch: this + // ledger proves no process creation was ever attempted. A claimed launch + // with unknown outcome is failed above and is never retried. + match self.task_provider.store.claim_launch(&request.invocation) { + Ok(Some(record)) => self.launch_task(record).await, + Ok(None) => {} + Err(error) => { + tracing::warn!(error = %error, "task launch claim could not be persisted") + } + } + } + + async fn launch_task(&mut self, record: TaskRecord) { + let Some(cli) = record + .invoke + .input + .get("cli") + .and_then(Value::as_str) + .map(str::to_owned) + else { + self.fail_task(&record.invoke.invocation_id, "task_missing_cli") + .await; + return; + }; + let Some(task) = record + .invoke + .input + .get("task") + .and_then(Value::as_str) + .map(str::to_owned) + else { + self.fail_task(&record.invoke.invocation_id, "task_missing_prompt") + .await; + return; + }; + let Some(workspace_id) = self.default_workspace_id.clone() else { + self.fail_task(&record.invoke.invocation_id, "task_workspace_unavailable") + .await; + return; + }; + let workspace = self + .workspace_lookup + .get(&workspace_id) + .cloned() + .unwrap_or_else(|| self.default_workspace.clone()); + let callback = AgentResultMcpConfig { + callback_url: self + .workers + .env_value("AGENT_RELAY_RESULT_URL") + .unwrap_or("http://127.0.0.1:3889/api/agent-result") + .to_owned(), + token: record.callback_token.clone(), + schema: record.invoke.input.get("result_schema").cloned(), + }; + let outcome = super::relaycast_events::spawn_worker_from_request( + record.name.clone(), + cli, + Some(task), + record + .invoke + .input + .get("channel") + .and_then(Value::as_str) + .map(str::to_owned), + record + .invoke + .input + .get("model") + .and_then(Value::as_str) + .map(str::to_owned), + true, + &record.invoke.input, + &workspace_id, + None, + &workspace, + &mut self.workers, + &mut self.state, + &self.paths, + &self.telemetry, + &self.sdk_out_tx, + &mut self.dedup, + &mut self.agent_spawn_count, + &self.fleet_control_tx, + &mut self.fleet_delivery_book, + &mut self.fleet_inventory, + &self.fleet_node_name, + Some(record.invoke.invocation_id.clone()), + None, + &self.hosted_agent_event_tx, + &mut self.pty_observability, + Some((callback, record.generation)), + ) + .await; + // Spawning may register restart state before returning an error or + // before a just-started child exits. Task generations are never + // eligible for ordinary supervisor restart, on either outcome. + self.workers.supervisor.unregister(&record.name); + if outcome.is_err() || !self.workers.is_worker_live(&record.name) { + let _ = self + .workers + .stop_task_generation(record.name.as_str(), record.generation) + .await; + self.fail_task(&record.invoke.invocation_id, "worker_spawn_failed") + .await; + } + self.publish_fleet_load(true).await; + } + + pub(super) async fn fail_task(&mut self, id: &str, error: &str) { + if self + .task_provider + .store + .records + .get(id) + .is_none_or(|r| r.final_result.is_some() || r.receipt.is_some()) + { + return; + } + match self.task_provider.store.queue_final( + id, + TaskFinal { + output: None, + error: Some(error.to_owned()), + accounting: None, + }, + ) { + Ok(record) => { + // Once a record carries a final_result, `maintain_tasks`'s + // expired-claimed sweep skips it (it only chases records that + // are still outcome-less), so this is the only remaining + // chance to stop a worker that is still live -- e.g. an + // accept-reply deadline expiring while the engine reports + // `running`. `stop_task_generation` is a no-op if the worker + // was never spawned or already stopped, so this is safe for + // every fail_task call site. + self.workers.supervisor.unregister(&record.name); + let _ = self + .workers + .stop_task_generation(record.name.as_str(), record.generation) + .await; + self.send_task_request(&record, TaskRequestKind::Accept, None, None) + .await + } + Err(error) => tracing::warn!(error = %error, "task failure could not be persisted"), + } + } + + pub(super) async fn handle_task_error(&mut self, error: crate::fleet_wire::Error) { + let Some(request) = self.task_provider.pending.remove(&error.id) else { + return; + }; + if self + .task_provider + .store + .records + .get(&request.invocation) + .is_none_or(|r| r.execution().execution_id != request.execution) + { + return; + } + if let Some(callback) = request.interim_reply { + let _ = callback.send(Err(AgentResultRouteError::Retryable)); + } + let terminal_rejection = matches!( + error.code.as_str(), + "stale_task_execution" | "task_not_found" | "task_result_conflict" + ); + let rejected_record = self + .task_provider + .store + .records + .get(&request.invocation) + .cloned(); + if terminal_rejection { + if let Err(error) = self + .task_provider + .store + .reject(&request.invocation, error.code) + { + tracing::warn!(error = %error, "task rejection could not be persisted"); + } + if let Some(record) = rejected_record { + self.workers.supervisor.unregister(&record.name); + let _ = self + .workers + .stop_task_generation(record.name.as_str(), record.generation) + .await; + } + if let Some(callbacks) = self.task_provider.callbacks.remove(&request.invocation) { + for (_, callback) in callbacks { + let _ = callback.send(Err(AgentResultRouteError::Conflict)); + } + } + } + } + + pub(super) async fn handle_task_callback( + &mut self, + token: String, + name: Option, + data: Value, + final_result: bool, + metadata: Option, + reply: CallbackReply, + ) { + if !self.task_provider.store.enabled() { + let _ = reply.send(Err(AgentResultRouteError::Retryable)); + return; + } + let Some(record) = self.task_provider.store.by_token(&token).cloned() else { + let _ = reply.send(Err(AgentResultRouteError::InvalidToken)); + return; + }; + if name.as_ref().is_some_and(|name| name != &record.name) { + let _ = reply.send(Err(AgentResultRouteError::InvalidToken)); + return; + } + // Old tokens remain reconcilable only for their original terminal result. + // A replacement worker with the same display name cannot submit for it. + if self + .workers + .workers + .get(&record.name) + .is_some_and(|w| w.generation != record.generation) + || record.rejection.is_some() + { + let _ = reply.send(Err(AgentResultRouteError::Conflict)); + return; + } + let accounting = match metadata.as_ref().and_then(|v| v.get("accounting")) { + Some(value) => match serde_json::from_value::>( + value.clone(), + ) { + Ok(values) if values.values().all(|n| n.is_finite() && *n >= 0.0) => Some(values), + _ => { + let _ = reply.send(Err(AgentResultRouteError::Conflict)); + return; + } + }, + None => None, + }; + let value = TaskFinal { + output: Some(data), + error: None, + accounting, + }; + if !final_result { + self.send_task_request(&record, TaskRequestKind::Interim, Some(value), Some(reply)) + .await; + return; + } + let record = match self + .task_provider + .store + .queue_final(&record.invoke.invocation_id, value) + { + Ok(record) => record, + Err(error) => { + let conflict = error.to_string().contains("conflict") + || error.to_string().contains("terminal"); + let _ = reply.send(Err(if conflict { + AgentResultRouteError::Conflict + } else { + AgentResultRouteError::Retryable + })); + return; + } + }; + let callbacks = self + .task_provider + .callbacks + .entry(record.invoke.invocation_id.clone()) + .or_default(); + if callbacks.len() >= 16 { + let _ = reply.send(Err(AgentResultRouteError::Retryable)); + return; + } + callbacks.push((Instant::now(), reply)); + if record.receipt.is_some() { + self.finish_task_callbacks(&record); + } else { + self.send_task_request(&record, TaskRequestKind::Accept, None, None) + .await; + } + } + + fn finish_task_callbacks(&mut self, record: &TaskRecord) { + let Some(callbacks) = self + .task_provider + .callbacks + .remove(&record.invoke.invocation_id) + else { + return; + }; + let receipt = record.receipt.as_ref().expect("terminal receipt"); + let matches = record.final_result.as_ref().is_some_and(|result| { + result.error.is_none() + && receipt.get("status").and_then(Value::as_str) == Some("completed") + && receipt + .get("output") + .zip(result.output.as_ref()) + .is_some_and(|(a, b)| json_equal(a, b)) + && receipt + .get("task_execution") + .and_then(|e| e.get("accounting")) + .map(|value| { + serde_json::from_value::>( + value.clone(), + ) + }) + .transpose() + .ok() + == Some(result.accounting.clone()) + }); + for (_, callback) in callbacks { + let response = if matches { + Ok(json!({"success": true, "name": record.name, "final": true, "receipt": receipt})) + } else { + Err(AgentResultRouteError::Conflict) + }; + let _ = callback.send(response); + } + } + + pub(super) async fn maintain_tasks(&mut self) { + if self.task_provider.store.enabled() { + if let Err(error) = self.task_provider.store.compact() { + tracing::warn!(error = %error, "terminal task retention could not be persisted"); + } + } + let expired_claimed: Vec<(String, WorkerName, Uuid)> = self + .task_provider + .store + .records + .values() + .filter(|record| { + record.launch_claimed + && record.final_result.is_none() + && record.receipt.is_none() + && record.rejection.is_none() + && record.expired() + }) + .map(|record| { + ( + record.invoke.invocation_id.clone(), + record.name.clone(), + record.generation, + ) + }) + .collect(); + for (invocation, name, generation) in expired_claimed { + self.workers.supervisor.unregister(&name); + let _ = self + .workers + .stop_task_generation(name.as_str(), generation) + .await; + self.fail_task(&invocation, "task_deadline_exceeded").await; + } + let now = Instant::now(); + let expired: Vec = self + .task_provider + .pending + .iter() + .filter(|(_, request)| now.duration_since(request.sent_at) >= Duration::from_secs(5)) + .map(|(id, _)| id.clone()) + .collect(); + for id in expired { + if let Some(request) = self.task_provider.pending.remove(&id) { + if let Some(callback) = request.interim_reply { + let _ = callback.send(Err(AgentResultRouteError::Retryable)); + } + } + } + self.task_provider.callbacks.retain(|_, callbacks| { + let mut retained = Vec::new(); + for (started, callback) in callbacks.drain(..) { + if callback.is_closed() { + continue; + } + if now.duration_since(started) >= Duration::from_secs(5) { + let _ = callback.send(Err(AgentResultRouteError::Retryable)); + } else { + retained.push((started, callback)); + } + } + *callbacks = retained; + !callbacks.is_empty() + }); + if !self.node_delivery_connected + || self + .task_provider + .last_retry + .is_some_and(|last| now.duration_since(last) < Duration::from_secs(2)) + { + return; + } + self.task_provider.last_retry = Some(now); + let records: Vec = self + .task_provider + .store + .records + .values() + .filter(|r| r.receipt.is_none() && r.rejection.is_none()) + .cloned() + .collect(); + if !records.is_empty() { + let start = self.task_provider.retry_cursor % records.len(); + for offset in 0..records.len().min(16) { + self.send_task_request( + &records[(start + offset) % records.len()], + TaskRequestKind::Accept, + None, + None, + ) + .await; + } + self.task_provider.retry_cursor = (start + records.len().min(16)) % records.len(); + } + } +} diff --git a/crates/broker/src/runtime/tests.rs b/crates/broker/src/runtime/tests.rs index 62d1a77895..a57f44d4c5 100644 --- a/crates/broker/src/runtime/tests.rs +++ b/crates/broker/src/runtime/tests.rs @@ -425,6 +425,7 @@ async fn owned_cleanup_waits_off_actor_and_retains_custody_until_confirmed() { .runtime .handle_fleet_control_event(crate::node_control::FleetControlEvent::Message( crate::fleet_wire::RelaycastToBroker::ActionInvoke(crate::fleet_wire::ActionInvoke { + task_execution: None, v: FLEET_WIRE_VERSION, invocation_id: "replacement-attempt".into(), action: "spawn".into(), @@ -665,6 +666,7 @@ fn worker_event_runtime_fixture( resize_owners: HashMap::new(), delivery_states: HashMap::new(), agent_result_tokens: HashMap::new(), + task_provider: super::tasks::TaskProvider::default(), recent_thread_messages: std::collections::VecDeque::new(), shutdown: false, lease_duration: None, @@ -6931,3 +6933,716 @@ async fn assert_http_spawn_metadata_publication(supplied_token: bool, valid_cwd: .contains_key(&name)); } } + +fn durable_task_fixture() -> WorkerEventRuntimeFixture { + let (tx, _rx) = mpsc::channel(16); + let workers = WorkerRegistry::new(tx, Vec::new(), std::env::temp_dir(), Instant::now()); + let mut fixture = worker_event_runtime_fixture(workers, HashMap::new()); + fixture.runtime.task_provider.store = + super::task_store::TaskStore::open(fixture._temp_dir.path().join("tasks.json")).unwrap(); + fixture +} + +async fn next_task_frame( + fixture: &mut WorkerEventRuntimeFixture, +) -> crate::fleet_wire::BrokerToRelaycast { + loop { + if let FleetControlCommand::Send(message) = + tokio::time::timeout(Duration::from_secs(2), fixture.fleet_control_rx.recv()) + .await + .expect("task frame timeout") + .unwrap() + { + return message; + } + } +} + +async fn deliver_task_receipt( + fixture: &mut WorkerEventRuntimeFixture, + request: &str, + status: &str, +) { + let record = fixture.runtime.task_provider.store.records["inv-task"].clone(); + fixture + .runtime + .handle_fleet_control_event(crate::node_control::FleetControlEvent::Message( + crate::fleet_wire::RelaycastToBroker::Reply(crate::fleet_wire::Reply { + v: crate::fleet_wire::FLEET_WIRE_VERSION, + id: request.to_owned(), + ok: true, + data: super::task_store::fixture_receipt(&record, status), + }), + )) + .await; +} + +#[tokio::test] +async fn durable_task_callback_waits_for_engine_final_receipt_and_retries_after_lost_ack() { + use crate::fleet_wire::BrokerToRelaycast; + let mut fixture = durable_task_fixture(); + let record = fixture + .runtime + .task_provider + .store + .prepare(super::task_store::fixture_invoke()) + .unwrap(); + fixture + .runtime + .task_provider + .store + .claim_launch("inv-task") + .unwrap(); + let (reply, mut receiver) = tokio::sync::oneshot::channel(); + fixture + .runtime + .handle_api_request(crate::listen_api::ListenApiRequest::SubmitAgentResult { + token: record.callback_token.clone(), + name: Some(record.name.clone()), + data: json!({"answer":42}), + final_result: true, + metadata: Some(json!({"accounting":{"tokens":17}})), + reply, + }) + .await; + assert!(matches!( + receiver.try_recv(), + Err(tokio::sync::oneshot::error::TryRecvError::Empty) + )); + let BrokerToRelaycast::ActionAccept(accept) = next_task_frame(&mut fixture).await else { + panic!("expected accept reconciliation") + }; + deliver_task_receipt(&mut fixture, &accept.id, "running").await; + let BrokerToRelaycast::ActionResult(result) = next_task_frame(&mut fixture).await else { + panic!("expected fenced final") + }; + assert!(result.task.as_ref().unwrap().final_result); + assert_eq!( + result.task.as_ref().unwrap().worker_generation, + record.generation.to_string() + ); + assert!(matches!( + receiver.try_recv(), + Err(tokio::sync::oneshot::error::TryRecvError::Empty) + )); + // Simulate lost result ACK: a new accept reconciles the same terminal record. + fixture.runtime.task_provider.pending.clear(); + fixture.runtime.task_provider.last_retry = None; + fixture.runtime.maintain_tasks().await; + let BrokerToRelaycast::ActionAccept(reconcile) = next_task_frame(&mut fixture).await else { + panic!("expected reconcile") + }; + deliver_task_receipt(&mut fixture, &reconcile.id, "completed").await; + let response = receiver.await.unwrap().unwrap(); + assert_eq!(response["receipt"]["output"], json!({"answer":42})); + assert_eq!( + response["receipt"]["task_execution"]["accounting"], + json!({"tokens":17.0}) + ); + // Callback retry after HTTP response loss uses the saved receipt, no new worker or send. + let (reply, receiver) = tokio::sync::oneshot::channel(); + fixture + .runtime + .handle_api_request(crate::listen_api::ListenApiRequest::SubmitAgentResult { + token: record.callback_token, + name: None, + data: json!({"answer":42}), + final_result: true, + metadata: Some(json!({"accounting":{"tokens":17}})), + reply, + }) + .await; + assert!(receiver.await.unwrap().is_ok()); + assert!(fixture.runtime.workers.workers.is_empty()); + assert!(fixture.fleet_control_rx.try_recv().is_err()); +} + +#[tokio::test] +async fn durable_task_interim_and_stale_generation_never_complete_callback_as_final() { + use crate::fleet_wire::BrokerToRelaycast; + let mut fixture = durable_task_fixture(); + let record = fixture + .runtime + .task_provider + .store + .prepare(super::task_store::fixture_invoke()) + .unwrap(); + fixture + .runtime + .task_provider + .store + .claim_launch("inv-task") + .unwrap(); + let (reply, mut receiver) = tokio::sync::oneshot::channel(); + fixture + .runtime + .handle_api_request(crate::listen_api::ListenApiRequest::SubmitAgentResult { + token: record.callback_token, + name: None, + data: json!({"ready":true}), + final_result: false, + metadata: None, + reply, + }) + .await; + let BrokerToRelaycast::ActionResult(result) = next_task_frame(&mut fixture).await else { + panic!("interim") + }; + assert!(!result.task.as_ref().unwrap().final_result); + assert!(receiver.try_recv().is_err()); + deliver_task_receipt(&mut fixture, result.id.as_deref().unwrap(), "running").await; + assert_eq!(receiver.await.unwrap().unwrap()["final"], false); + assert!(fixture.runtime.task_provider.store.records["inv-task"] + .final_result + .is_none()); + assert!(fixture.runtime.task_provider.store.records["inv-task"] + .receipt + .is_none()); +} + +#[tokio::test] +async fn durable_task_restart_of_claimed_launch_reports_loss_without_respawn() { + use crate::fleet_wire::BrokerToRelaycast; + let mut fixture = durable_task_fixture(); + let record = fixture + .runtime + .task_provider + .store + .prepare(super::task_store::fixture_invoke()) + .unwrap(); + fixture + .runtime + .task_provider + .store + .claim_launch("inv-task") + .unwrap(); + fixture.runtime.task_provider.store = + super::task_store::TaskStore::open(fixture._temp_dir.path().join("tasks.json")).unwrap(); + fixture.runtime.handle_task_invoke(record.invoke).await; + let BrokerToRelaycast::ActionAccept(accept) = next_task_frame(&mut fixture).await else { + panic!("accept") + }; + deliver_task_receipt(&mut fixture, &accept.id, "running").await; + assert_eq!( + fixture.runtime.task_provider.store.records["inv-task"] + .final_result + .as_ref() + .unwrap() + .error + .as_deref(), + Some("worker_execution_lost") + ); + assert!(fixture.runtime.workers.workers.is_empty()); + let BrokerToRelaycast::ActionAccept(accept) = next_task_frame(&mut fixture).await else { + panic!("failure reconciliation") + }; + deliver_task_receipt(&mut fixture, &accept.id, "running").await; + let BrokerToRelaycast::ActionResult(result) = next_task_frame(&mut fixture).await else { + panic!("failure") + }; + assert!(matches!( + result.result, + crate::fleet_wire::ActionResultPayload::Error(_) + )); + deliver_task_receipt(&mut fixture, result.id.as_deref().unwrap(), "failed").await; + assert_eq!( + fixture.runtime.task_provider.store.records["inv-task"] + .receipt + .as_ref() + .unwrap()["status"], + "failed" + ); +} + +#[tokio::test] +async fn durable_task_receipt_timeout_is_retryable_and_terminal_outbox_survives_disconnect() { + let mut fixture = durable_task_fixture(); + let record = fixture + .runtime + .task_provider + .store + .prepare(super::task_store::fixture_invoke()) + .unwrap(); + fixture + .runtime + .task_provider + .store + .claim_launch("inv-task") + .unwrap(); + fixture.runtime.node_delivery_connected = false; + let (reply, mut receiver) = tokio::sync::oneshot::channel(); + fixture + .runtime + .handle_api_request(crate::listen_api::ListenApiRequest::SubmitAgentResult { + token: record.callback_token, + name: None, + data: json!(42), + final_result: true, + metadata: None, + reply, + }) + .await; + assert!(receiver.try_recv().is_err()); + fixture + .runtime + .task_provider + .callbacks + .get_mut("inv-task") + .unwrap()[0] + .0 = Instant::now() - Duration::from_secs(6); + fixture.runtime.maintain_tasks().await; + assert_eq!( + receiver.await.unwrap(), + Err(crate::listen_api::AgentResultRouteError::Retryable) + ); + assert!(fixture.runtime.task_provider.store.records["inv-task"] + .final_result + .is_some()); + assert!(fixture.runtime.task_provider.store.records["inv-task"] + .receipt + .is_none()); +} + +#[tokio::test] +async fn durable_task_duplicate_invoke_and_launch_failure_are_generation_fenced() { + use crate::fleet_wire::BrokerToRelaycast; + let mut fixture = durable_task_fixture(); + let invoke = super::task_store::fixture_invoke(); + fixture.runtime.handle_task_invoke(invoke.clone()).await; + fixture.runtime.handle_task_invoke(invoke).await; + let BrokerToRelaycast::ActionAccept(accept) = next_task_frame(&mut fixture).await else { + panic!("accept") + }; + assert!(fixture.fleet_control_rx.try_recv().is_err()); + // Missing CLI fails only after acceptance, through the explicit final-error path. + deliver_task_receipt(&mut fixture, &accept.id, "running").await; + let record = &fixture.runtime.task_provider.store.records["inv-task"]; + assert!(record.launch_claimed); + assert_eq!( + record.final_result.as_ref().unwrap().error.as_deref(), + Some("task_missing_cli") + ); + assert!(fixture.runtime.workers.workers.is_empty()); + // A repeated accepted receipt never attempts launch again. + let BrokerToRelaycast::ActionAccept(accept) = next_task_frame(&mut fixture).await else { + panic!("reconcile") + }; + deliver_task_receipt(&mut fixture, &accept.id, "running").await; + assert!(matches!( + next_task_frame(&mut fixture).await, + BrokerToRelaycast::ActionResult(_) + )); +} + +#[tokio::test] +async fn durable_task_terminal_rejection_stops_worker_and_replay_stays_refused() { + use crate::fleet_wire::{ActionResultPayload, BrokerToRelaycast}; + let mut fixture = durable_task_fixture(); + let invoke = super::task_store::fixture_invoke(); + fixture.runtime.handle_task_invoke(invoke.clone()).await; + let BrokerToRelaycast::ActionAccept(accept) = next_task_frame(&mut fixture).await else { + panic!("accept") + }; + let record = fixture.runtime.task_provider.store.records["inv-task"].clone(); + fixture + .runtime + .task_provider + .store + .claim_launch("inv-task") + .unwrap(); + + let mut workers = make_worker_registry_with_worker(record.name.as_str()).await; + workers.workers.get_mut(&record.name).unwrap().generation = record.generation; + let restart_policy = crate::supervisor::RestartPolicy { + cooldown_ms: 0, + ..crate::supervisor::RestartPolicy::default() + }; + let mut spec = workers.workers[&record.name].spec.clone(); + spec.restart_policy = Some(restart_policy.clone()); + workers.supervisor.register( + record.name.as_str(), + crate::supervisor::SupervisedAgent { + spec, + parent: None, + initial_task: None, + skip_relay_prompt: false, + agent_result: None, + }, + restart_policy, + ); + fixture.runtime.workers = workers; + + fixture + .runtime + .handle_task_error(crate::fleet_wire::Error { + v: FLEET_WIRE_VERSION, + id: accept.id, + ok: false, + code: "task_not_found".to_owned(), + message: "task no longer exists".to_owned(), + }) + .await; + assert!(!fixture.runtime.workers.is_worker_live(&record.name)); + assert!(!fixture + .runtime + .workers + .supervisor + .is_supervised(&record.name)); + assert_eq!( + fixture.runtime.task_provider.store.records["inv-task"] + .rejection + .as_deref(), + Some("task_not_found") + ); + + fixture.runtime.handle_task_invoke(invoke).await; + let BrokerToRelaycast::ActionResult(result) = next_task_frame(&mut fixture).await else { + panic!("rejected replay must receive a terminal action result") + }; + let ActionResultPayload::Error(error) = result.result else { + panic!("rejected replay must fail") + }; + assert_eq!(error.error, "handler_unavailable"); +} + +#[tokio::test] +async fn durable_task_maintenance_stops_claimed_worker_at_deadline() { + let mut fixture = durable_task_fixture(); + let record = fixture + .runtime + .task_provider + .store + .prepare(super::task_store::fixture_invoke()) + .unwrap(); + fixture + .runtime + .task_provider + .store + .claim_launch("inv-task") + .unwrap(); + let mut workers = make_worker_registry_with_worker(record.name.as_str()).await; + workers.workers.get_mut(&record.name).unwrap().generation = record.generation; + fixture.runtime.workers = workers; + fixture + .runtime + .task_provider + .store + .records + .get_mut("inv-task") + .unwrap() + .invoke + .task_execution + .as_mut() + .unwrap() + .deadline = (chrono::Utc::now() - chrono::Duration::seconds(1)) + .to_rfc3339_opts(chrono::SecondsFormat::Millis, true); + fixture.runtime.node_delivery_connected = false; + + fixture.runtime.maintain_tasks().await; + + assert!(!fixture.runtime.workers.is_worker_live(&record.name)); + assert_eq!( + fixture.runtime.task_provider.store.records["inv-task"] + .final_result + .as_ref() + .unwrap() + .error + .as_deref(), + Some("task_deadline_exceeded") + ); +} + +#[tokio::test] +async fn durable_task_refusals_return_terminal_action_errors() { + use crate::fleet_wire::{ActionResultPayload, BrokerToRelaycast}; + let mut fixture = durable_task_fixture(); + let invoke = super::task_store::fixture_invoke(); + fixture.runtime.handle_task_invoke(invoke.clone()).await; + assert!(matches!( + next_task_frame(&mut fixture).await, + BrokerToRelaycast::ActionAccept(_) + )); + + let mut conflicting = invoke.clone(); + conflicting.input["task"] = json!("different"); + fixture.runtime.handle_task_invoke(conflicting).await; + let BrokerToRelaycast::ActionResult(conflict) = next_task_frame(&mut fixture).await else { + panic!("conflicting invoke must receive a result") + }; + let ActionResultPayload::Error(error) = conflict.result else { + panic!("conflicting invoke must fail") + }; + assert_eq!(error.error, "handler_unavailable"); + + fixture.runtime.task_provider.store = Default::default(); + fixture.runtime.handle_task_invoke(invoke).await; + let BrokerToRelaycast::ActionResult(disabled) = next_task_frame(&mut fixture).await else { + panic!("disabled provider invoke must receive a result") + }; + let ActionResultPayload::Error(error) = disabled.result else { + panic!("disabled provider invoke must fail") + }; + assert_eq!(error.error, "handler_unavailable"); +} + +#[tokio::test] +async fn durable_task_disk_failure_after_engine_commit_withholds_callback_ack() { + use crate::fleet_wire::BrokerToRelaycast; + let mut fixture = durable_task_fixture(); + let record = fixture + .runtime + .task_provider + .store + .prepare(super::task_store::fixture_invoke()) + .unwrap(); + fixture + .runtime + .task_provider + .store + .claim_launch("inv-task") + .unwrap(); + let (reply, mut receiver) = tokio::sync::oneshot::channel(); + fixture + .runtime + .handle_api_request(crate::listen_api::ListenApiRequest::SubmitAgentResult { + token: record.callback_token, + name: None, + data: json!(42), + final_result: true, + metadata: None, + reply, + }) + .await; + let BrokerToRelaycast::ActionAccept(accept) = next_task_frame(&mut fixture).await else { + panic!("accept") + }; + deliver_task_receipt(&mut fixture, &accept.id, "running").await; + let BrokerToRelaycast::ActionResult(result) = next_task_frame(&mut fixture).await else { + panic!("result") + }; + let path = fixture._temp_dir.path().join("tasks.json"); + let persisted = std::fs::read(&path).unwrap(); + std::fs::remove_file(&path).unwrap(); + std::fs::create_dir(&path).unwrap(); + deliver_task_receipt(&mut fixture, result.id.as_deref().unwrap(), "completed").await; + assert!(receiver.try_recv().is_err()); + assert!(!fixture.runtime.task_provider.store.enabled()); + assert!(fixture.runtime.task_provider.store.records["inv-task"] + .receipt + .is_none()); + // Reopen a recovered durable outbox and reconcile the already committed result. + std::fs::remove_dir(&path).unwrap(); + std::fs::write(&path, persisted).unwrap(); + fixture.runtime.task_provider.store = super::task_store::TaskStore::open(path).unwrap(); + fixture.runtime.task_provider.last_retry = None; + fixture.runtime.maintain_tasks().await; + let BrokerToRelaycast::ActionAccept(accept) = next_task_frame(&mut fixture).await else { + panic!("reconcile") + }; + deliver_task_receipt(&mut fixture, &accept.id, "completed").await; + assert!(receiver.await.unwrap().is_ok()); + assert!(fixture.runtime.workers.workers.is_empty()); +} + +#[tokio::test] +async fn durable_task_expired_unaccepted_receipt_never_launches_worker() { + use crate::fleet_wire::BrokerToRelaycast; + let mut fixture = durable_task_fixture(); + let mut invoke = super::task_store::fixture_invoke(); + invoke.task_execution.as_mut().unwrap().deadline = (chrono::Utc::now() + - chrono::Duration::minutes(1)) + .to_rfc3339_opts(chrono::SecondsFormat::Millis, true); + fixture.runtime.handle_task_invoke(invoke).await; + let BrokerToRelaycast::ActionAccept(accept) = next_task_frame(&mut fixture).await else { + panic!("accept") + }; + let record = fixture.runtime.task_provider.store.records["inv-task"].clone(); + let mut receipt = super::task_store::fixture_receipt(&record, "failed"); + receipt["error"] = json!("task_deadline_exceeded"); + receipt["task_execution"] + .as_object_mut() + .unwrap() + .remove("worker_generation"); + fixture + .runtime + .handle_task_reply(crate::fleet_wire::Reply { + v: FLEET_WIRE_VERSION, + id: accept.id, + ok: true, + data: receipt, + }) + .await; + assert!(fixture.runtime.workers.workers.is_empty()); + assert!(!fixture.runtime.task_provider.store.records["inv-task"].launch_claimed); + assert_eq!( + fixture.runtime.task_provider.store.records["inv-task"] + .receipt + .as_ref() + .unwrap()["error"], + "task_deadline_exceeded" + ); +} + +#[tokio::test] +async fn durable_task_expired_running_receipt_queues_terminal_deadline_failure() { + use crate::fleet_wire::BrokerToRelaycast; + let mut fixture = durable_task_fixture(); + let mut invoke = super::task_store::fixture_invoke(); + invoke.task_execution.as_mut().unwrap().deadline = (chrono::Utc::now() + - chrono::Duration::minutes(1)) + .to_rfc3339_opts(chrono::SecondsFormat::Millis, true); + fixture.runtime.handle_task_invoke(invoke).await; + let BrokerToRelaycast::ActionAccept(accept) = next_task_frame(&mut fixture).await else { + panic!("accept") + }; + + deliver_task_receipt(&mut fixture, &accept.id, "running").await; + let record = &fixture.runtime.task_provider.store.records["inv-task"]; + assert!(!record.launch_claimed); + assert_eq!( + record.final_result.as_ref().unwrap().error.as_deref(), + Some("task_deadline_exceeded") + ); + let BrokerToRelaycast::ActionAccept(reconcile) = next_task_frame(&mut fixture).await else { + panic!("reconcile") + }; + deliver_task_receipt(&mut fixture, &reconcile.id, "running").await; + let BrokerToRelaycast::ActionResult(result) = next_task_frame(&mut fixture).await else { + panic!("result") + }; + let crate::fleet_wire::ActionResultPayload::Error(error) = result.result else { + panic!("deadline failure") + }; + assert_eq!(error.error, "task_deadline_exceeded"); + assert!(result.task.as_ref().unwrap().final_result); +} + +#[tokio::test] +async fn durable_task_reply_deadline_stops_a_live_worker() { + // Once `fail_task` sets `final_result`, `maintain_tasks`'s expired-claimed + // sweep skips the record (it only chases records still outcome-less) -- + // this reply-path deadline check is the only remaining chance to stop a + // worker that is still live when the accept-reply deadline expires while + // the engine reports "running". + use crate::fleet_wire::BrokerToRelaycast; + let mut fixture = durable_task_fixture(); + let invoke = super::task_store::fixture_invoke(); + fixture.runtime.handle_task_invoke(invoke).await; + let BrokerToRelaycast::ActionAccept(accept) = next_task_frame(&mut fixture).await else { + panic!("accept") + }; + let record = fixture.runtime.task_provider.store.records["inv-task"].clone(); + let mut workers = make_worker_registry_with_worker(record.name.as_str()).await; + workers.workers.get_mut(&record.name).unwrap().generation = record.generation; + fixture.runtime.workers = workers; + assert!(fixture.runtime.workers.is_worker_live(&record.name)); + fixture + .runtime + .task_provider + .store + .records + .get_mut("inv-task") + .unwrap() + .invoke + .task_execution + .as_mut() + .unwrap() + .deadline = (chrono::Utc::now() - chrono::Duration::seconds(1)) + .to_rfc3339_opts(chrono::SecondsFormat::Millis, true); + + deliver_task_receipt(&mut fixture, &accept.id, "running").await; + + assert!(!fixture.runtime.workers.is_worker_live(&record.name)); + assert_eq!( + fixture.runtime.task_provider.store.records["inv-task"] + .final_result + .as_ref() + .unwrap() + .error + .as_deref(), + Some("task_deadline_exceeded") + ); +} + +#[tokio::test] +async fn durable_task_old_attempt_receipt_cannot_start_new_generation() { + use crate::fleet_wire::BrokerToRelaycast; + let mut fixture = durable_task_fixture(); + let invoke = super::task_store::fixture_invoke(); + fixture.runtime.handle_task_invoke(invoke.clone()).await; + let original = fixture.runtime.task_provider.store.records["inv-task"].clone(); + let BrokerToRelaycast::ActionAccept(old) = next_task_frame(&mut fixture).await else { + panic!("accept") + }; + let mut next = invoke; + next.task_execution.as_mut().unwrap().execution_id = "inv-task/2".into(); + fixture.runtime.handle_task_invoke(next).await; + fixture + .runtime + .handle_task_reply(crate::fleet_wire::Reply { + v: FLEET_WIRE_VERSION, + id: old.id, + ok: true, + data: super::task_store::fixture_receipt(&original, "running"), + }) + .await; + assert!(!fixture.runtime.task_provider.store.records["inv-task"].launch_claimed); + assert_ne!( + fixture.runtime.task_provider.store.records["inv-task"].generation, + original.generation + ); + assert!(fixture + .runtime + .task_provider + .store + .by_token(&original.callback_token) + .is_none()); + assert!(fixture.runtime.workers.workers.is_empty()); +} + +#[tokio::test] +async fn durable_task_numeric_output_and_accounting_reconcile_javascript_json() { + use crate::fleet_wire::BrokerToRelaycast; + let mut fixture = durable_task_fixture(); + let record = fixture + .runtime + .task_provider + .store + .prepare(super::task_store::fixture_invoke()) + .unwrap(); + fixture + .runtime + .task_provider + .store + .claim_launch("inv-task") + .unwrap(); + let (reply, receiver) = tokio::sync::oneshot::channel(); + fixture + .runtime + .handle_api_request(crate::listen_api::ListenApiRequest::SubmitAgentResult { + token: record.callback_token, + name: None, + data: json!({"answer":42.0}), + final_result: true, + metadata: Some(json!({"accounting":{"tokens":17.0}})), + reply, + }) + .await; + let BrokerToRelaycast::ActionAccept(accept) = next_task_frame(&mut fixture).await else { + panic!("accept") + }; + let record = fixture.runtime.task_provider.store.records["inv-task"].clone(); + let mut receipt = super::task_store::fixture_receipt(&record, "completed"); + receipt["output"] = json!({"answer":42}); + receipt["task_execution"]["accounting"] = json!({"tokens":17}); + fixture + .runtime + .handle_task_reply(crate::fleet_wire::Reply { + v: FLEET_WIRE_VERSION, + id: accept.id, + ok: true, + data: receipt, + }) + .await; + assert!(receiver.await.unwrap().is_ok()); +} diff --git a/crates/broker/src/worker.rs b/crates/broker/src/worker.rs index b582f1c269..ea54f17e37 100644 --- a/crates/broker/src/worker.rs +++ b/crates/broker/src/worker.rs @@ -583,6 +583,33 @@ impl WorkerRegistry { workspace_id: Option, agent_result: Option, commit_attestation: Option, + ) -> Result { + self.spawn_with_generation( + spec, + parent, + idle_threshold_secs, + worker_relay_api_key, + skip_relay_prompt, + workspace_id, + agent_result, + commit_attestation, + None, + ) + .await + } + + #[allow(clippy::too_many_arguments)] + pub(crate) async fn spawn_with_generation( + &mut self, + spec: AgentSpec, + parent: Option, + idle_threshold_secs: Option, + worker_relay_api_key: Option, + skip_relay_prompt: bool, + workspace_id: Option, + agent_result: Option, + commit_attestation: Option, + task_generation: Option, ) -> Result { let mut spec = spec; if self.identity_cleanups.contains_key(&spec.name) { @@ -1289,7 +1316,7 @@ impl WorkerRegistry { let log_file = self.worker_log_path(&spec.name); let startup_log_file = log_file.clone(); - let generation = Uuid::new_v4(); + let generation = task_generation.unwrap_or_else(Uuid::new_v4); spawn_worker_reader( self.event_tx.clone(), spec.name.clone(), @@ -1507,6 +1534,25 @@ impl WorkerRegistry { .await } + /// Stop a terminal task without touching a replacement worker or bypassing + /// the normal reap/owned-identity cleanup path. + pub(crate) async fn stop_task_generation( + &mut self, + name: &str, + generation: Uuid, + ) -> Result { + let Some(handle) = self.workers.get_mut(name) else { + return Ok(false); + }; + if handle.generation != generation { + return Ok(false); + } + self.supervisor.unregister(name); + handle.exit_reason = Some("task_terminal_failure".to_owned()); + terminate_child(&mut handle.child, ORPHAN_REAP_TIMEOUT).await?; + Ok(true) + } + pub(crate) async fn release(&mut self, name: &str) -> Result<()> { tracing::info!(target = "broker::release", name = %name, "releasing worker"); self.initial_tasks.remove(name); @@ -3934,4 +3980,64 @@ sleep 30 assert_eq!(fallback, None); assert_eq!(args, vec!["--model=gpt-5.5".to_string()]); } + #[cfg(any(target_os = "linux", target_os = "macos"))] + #[tokio::test] + async fn task_spawn_uses_preclaimed_generation_and_injects_result_callback() { + let directory = tempfile::tempdir().unwrap(); + let (tx, mut rx) = mpsc::channel(16); + let mut registry = WorkerRegistry::new( + tx, + Vec::new(), + directory.path().join("logs"), + Instant::now(), + ); + let generation = Uuid::new_v4(); + let mut spec = sleeping_native_worker( + "fenced-task", + Some(directory.path().to_string_lossy().into_owned()), + ); + if let Some(ResolvedHarnessConfig::Native(config)) = &mut spec.harness_config { + config.args = vec!["-c".into(), r#"printf '%s' "$AGENT_RELAY_RESULT_TOKEN" > callback-token; printf '{"type":"worker_ready"}\n'; sleep 30"#.into()]; + } + registry + .spawn_with_generation( + spec, + None, + None, + None, + true, + None, + Some(AgentResultMcpConfig { + callback_url: "http://127.0.0.1:1/api/agent-result".into(), + token: "fixture-task-callback".into(), + schema: None, + }), + None, + Some(generation), + ) + .await + .unwrap(); + assert_eq!(registry.workers["fenced-task"].generation, generation); + assert_eq!( + std::fs::read_to_string(directory.path().join("callback-token")).unwrap(), + "fixture-task-callback" + ); + let event = tokio::time::timeout(Duration::from_secs(2), rx.recv()) + .await + .unwrap() + .unwrap(); + assert!(!registry + .stop_task_generation("fenced-task", Uuid::new_v4()) + .await + .unwrap()); + assert!(registry.is_worker_live("fenced-task")); + assert!(registry + .stop_task_generation("fenced-task", generation) + .await + .unwrap()); + registry.release("fenced-task").await.unwrap(); + assert!( + matches!(event, WorkerEvent::Message { generation: observed, .. } if observed == generation) + ); + } } diff --git a/crates/broker/tests/fixtures/fleet-wire/action.accept.json b/crates/broker/tests/fixtures/fleet-wire/action.accept.json new file mode 100644 index 0000000000..5149b44d8d --- /dev/null +++ b/crates/broker/tests/fixtures/fleet-wire/action.accept.json @@ -0,0 +1,8 @@ +{ + "v": 1, + "id": "task_receipt_accept", + "type": "action.accept", + "invocation_id": "inv-task", + "execution_id": "inv-task/1", + "worker_generation": "worker-1" +} diff --git a/crates/broker/tests/fixtures/fleet-wire/action.invoke.task.json b/crates/broker/tests/fixtures/fleet-wire/action.invoke.task.json new file mode 100644 index 0000000000..be11c15183 --- /dev/null +++ b/crates/broker/tests/fixtures/fleet-wire/action.invoke.task.json @@ -0,0 +1,23 @@ +{ + "v": 1, + "type": "action.invoke", + "invocation_id": "inv-task", + "action": "task.run", + "input": { + "cli": "claude", + "task": "Return 42", + "task_context": { + "run_id": "run", + "step_id": "step", + "dispatch_id": "dispatch", + "timeout_ms": 60000 + } + }, + "task_execution": { + "execution_id": "inv-task/1", + "run_id": "run", + "step_id": "step", + "dispatch_id": "dispatch", + "deadline": "2026-01-01T00:01:00.000Z" + } +} diff --git a/crates/broker/tests/fixtures/fleet-wire/action.result.task.json b/crates/broker/tests/fixtures/fleet-wire/action.result.task.json new file mode 100644 index 0000000000..be83050d69 --- /dev/null +++ b/crates/broker/tests/fixtures/fleet-wire/action.result.task.json @@ -0,0 +1,15 @@ +{ + "v": 1, + "id": "task_receipt_final", + "type": "action.result", + "invocation_id": "inv-task", + "execution_id": "inv-task/1", + "worker_generation": "worker-1", + "final": true, + "output": { + "answer": 42 + }, + "accounting": { + "tokens": 17 + } +} diff --git a/crates/broker/tests/fleet_wire_fixtures.rs b/crates/broker/tests/fleet_wire_fixtures.rs index 813647cb11..0da30608ba 100644 --- a/crates/broker/tests/fleet_wire_fixtures.rs +++ b/crates/broker/tests/fleet_wire_fixtures.rs @@ -13,15 +13,19 @@ const NODE_TO_SERVER_TYPES: &[&str] = &[ "agent.deregister", "delivery.ack", "action.result", + "action.accept", "inventory.sync", ]; const SERVER_TO_NODE_TYPES: &[&str] = &["deliver", "action.invoke", "ping", "reply", "error"]; const EXPECTED_FIXTURE_FILES: &[&str] = &[ + "action.accept.json", "action.invoke.json", + "action.invoke.task.json", "action.result.error.json", "action.result.output.json", + "action.result.task.json", "agent.deregister.json", "agent.register.json", "deliver.json", @@ -122,3 +126,23 @@ fn fleet_wire_fixtures_round_trip_semantically() { "fixture coverage mismatch in {FIXTURE_DIR}" ); } + +#[test] +fn task_wire_requires_fences_and_final_flag_without_changing_short_results() { + let path = Path::new(FIXTURE_DIR).join("action.result.task.json"); + let fixture: Value = serde_json::from_str(&fs::read_to_string(path).unwrap()).unwrap(); + for field in ["id", "execution_id", "worker_generation", "final"] { + let mut invalid = fixture.clone(); + invalid.as_object_mut().unwrap().remove(field); + assert!( + serde_json::from_value::(invalid).is_err(), + "{field} must be required" + ); + } + let mut interim = fixture.clone(); + interim["final"] = Value::Bool(false); + assert!(serde_json::from_value::(interim).is_ok()); + let mut invalid = fixture; + invalid["accounting"]["tokens"] = serde_json::json!(-1); + assert!(serde_json::from_value::(invalid).is_err()); +} diff --git a/specs/durable-task-provider.md b/specs/durable-task-provider.md new file mode 100644 index 0000000000..2ff5b1c323 --- /dev/null +++ b/specs/durable-task-provider.md @@ -0,0 +1,106 @@ +# Durable broker task provider + +This provider depends on the task invocation engine contract in +[Relaycast #436](https://github.com/AgentWorkforce/relaycast/pull/436). +Deploy that engine contract before enabling this broker feature. The existing +published Relaycast package alone does not enable the server contract. + +## Enablement + +Set `AGENT_RELAY_TASK_PROVIDER=1` in the environment of a persistent hosted +broker (for example, when starting `agent-relay node up`). The feature is off by +default and refuses ephemeral or `--local-only` mode. It advertises the global +`task.run` action with `execution_mode: "task"`. Enable only one provider for this +global action in each workspace; ordinary node capabilities remain unchanged. + +## Invocation and completion + +An authenticated agent uses the existing `POST /v1/actions/task.run/invoke` +endpoint and a stable `Idempotency-Key`. The invocation input contains: + +```json +{ + "cli": "claude", + "task": "Return the requested result using the injected result tool.", + "task_context": { + "run_id": "run-example", + "step_id": "step-example", + "dispatch_id": "dispatch-example", + "timeout_ms": 300000 + } +} +``` + +Pass this object as the request's `input` field. Optional `model`, `channel`, +`result_schema`, and existing spawn harness options use the ordinary broker +spawn path. The provider chooses a stable worker name and UUID generation and +injects the generation's callback credential before process creation. + +The engine dispatches `action.invoke` with `task_execution` correlation and a +persisted execution attempt. The provider records the invocation, then sends +`action.accept` with its worker generation. Only a matching durable acceptance +permits launch. Spawn registration, readiness, process exit, and interim output +never imply successful task completion. The invoking agent reads +`GET /v1/actions/task.run/invocations/:id` for the authoritative terminal output +or failure; +websocket events are notifications, not a substitute for that read. + +The injected `/api/agent-result` callback keeps its existing request shape: +`data`, `final`, optional `name`, and optional `metadata.accounting` (finite, +nonnegative counters). For this provider, a final callback is recorded in the +local outbox before sending a fenced `action.result`. HTTP success follows only +a matching engine terminal receipt and its successful local durable write. +Interim callbacks forward an observation after engine acknowledgment and remain +nonterminal. + +- HTTP 503 means the receipt is pending or durable storage is unavailable. Retry + the identical final callback; the persisted outbox also reconciles it without + relying on the worker's connection remaining open. +- HTTP 409 means the generation, terminal outcome, or replay payload conflicts. + It is not a successful result acknowledgment. +- An identical final replay returns the same stored receipt. Changed final + output or accounting cannot replace it. Callback tokens remain bound to their + original generation. + +## Restart, failure, and deadline behavior + +A launch claim is persisted before attempting process creation. A restart can +resume a never-claimed launch after reconciling acceptance. If a launch was +claimed and that exact live generation cannot be proven, the provider reports +`worker_execution_lost`; it does not risk a duplicate launch. Automatic ordinary +worker restart is disabled for task generations. Spawn failures become explicit +failed task results. Old execution responses cannot advance a redispatched +invocation. + +Lost result acknowledgments are reconciled through the existing `action.accept` +receipt. The provider resends its identical outbox result only while the engine +still reports the matching execution as running. Terminal engine failures, +including deadlines, are preserved. After receiving a failed terminal receipt, +the provider terminates only a locally tracked matching worker generation. +Disconnected providers cannot immediately observe an engine deadline; a failed +invocation is not proof that an unknown or disconnected process has stopped. + +The task ledger sits alongside persistent broker state, under the same broker +ownership lock. Atomic replacement, file synchronization, and directory +synchronization precede acknowledgment; newly written files have mode 0600 on +Unix. A failed durable write disables further task writes and launches until the +storage problem is repaired and the broker restarted. A corrupt ledger refuses +startup. Keep the ledger with the broker's state: deleting it loses the evidence +needed to prevent duplicate launches. Terminal records are retained for replay; +the broker prunes them durably 24 hours after their execution deadline on +startup, after terminal completion, and during maintenance. This grace period +keeps late callbacks and equivalent receipt replays idempotent without allowing +the in-memory and on-disk ledger to grow without bound. + +## Compatibility and rollout + +Existing short `spawn` actions keep their readiness receipts and existing wire +shape. Existing local result callbacks retain their behavior. No new public +HTTP endpoint or published dependency version is assumed by this provider. + +The ordered rollout is Relaycast #436, this broker provider, then the Flows +adapter tracked in [Flows #397](https://github.com/AgentWorkforce/flows/issues/397) +and [Relay #1766](https://github.com/AgentWorkforce/relay/issues/1766). Preview +proof requires the engine contract to be deployed, a broker built with this +provider enabled, and the adapter using the task invocation lifecycle. Local +mocked tests and wire fixtures do not constitute deployed preview proof. diff --git a/tests/relayflows/cases/1766-durable-task-receipt/case.json b/tests/relayflows/cases/1766-durable-task-receipt/case.json new file mode 100644 index 0000000000..85f2b04b31 --- /dev/null +++ b/tests/relayflows/cases/1766-durable-task-receipt/case.json @@ -0,0 +1,13 @@ +{ + "version": 1, + "id": "1766-durable-task-receipt", + "kind": "feature", + "title": "Broker waits for task acceptance and reconciles a fenced final failure after restart", + "requirements": ["broker-linux-x64"], + "runner": { "command": ["node", "tests/relayflows/cases/1766-durable-task-receipt/run.mjs"] }, + "timeoutSeconds": 180, + "expected": { + "base": { "outcome": "absent", "signature": "task_action_handler_unavailable" }, + "head": { "outcome": "fixed", "signature": "fenced_task_failure_receipt_survives_restart" } + } +} diff --git a/tests/relayflows/cases/1766-durable-task-receipt/engine-fixture.mjs b/tests/relayflows/cases/1766-durable-task-receipt/engine-fixture.mjs new file mode 100644 index 0000000000..9e77cbea55 --- /dev/null +++ b/tests/relayflows/cases/1766-durable-task-receipt/engine-fixture.mjs @@ -0,0 +1,129 @@ +import assert from 'node:assert/strict'; +import { createHash } from 'node:crypto'; +import http from 'node:http'; + +/** Minimal loopback engine wire peer; no external credentials or package install. */ +export async function engineFixture() { + const frames = []; + const sockets = new Set(); + let node; + let failure; + const server = http.createServer(async (request, response) => { + try { + const chunks = []; + for await (const chunk of request) chunks.push(chunk); + const body = chunks.length ? JSON.parse(Buffer.concat(chunks).toString()) : {}; + if (request.method === 'POST' && request.url === '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/v1/agents') { + response.writeHead(201, { 'content-type': 'application/json' }); + response.end( + JSON.stringify({ + ok: true, + data: { + id: 'fixture-agent', + name: body.name, + workspace_id: 'fixture-workspace', + token: 'at_fixture_task_proof', + status: 'online', + created_at: '2026-01-01T00:00:00Z', + }, + }) + ); + } else { + response.writeHead(404, { 'content-type': 'application/json' }); + response.end( + JSON.stringify({ ok: false, error: { code: 'not_found', message: 'fixture route absent' } }) + ); + } + } catch (error) { + failure = error; + response.destroy(); + } + }); + server.on('connection', (socket) => { + sockets.add(socket); + socket.once('close', () => sockets.delete(socket)); + }); + server.on('upgrade', (request, socket, head) => { + if (new URL(request.url, 'http://fixture.invalid').pathname !== '/v1/node/ws') { + socket.end('HTTP/1.1 404 Not Found\r\nConnection: close\r\n\r\n'); + return; + } + try { + assert.equal(request.headers.authorization, 'Bearer nt_fixture_task_proof'); + const accept = createHash('sha1') + .update(request.headers['sec-websocket-key'] + '258EAFA5-E914-47DA-95CA-C5AB0DC85B11') + .digest('base64'); + socket.write( + `HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Accept: ${accept}\r\n\r\n` + ); + node = socket; + let buffer = head; + socket.on('data', (chunk) => { + buffer = Buffer.concat([buffer, chunk]); + try { + while (buffer.length >= 2) { + const frame = readClientFrame(buffer); + if (!frame) return; + buffer = frame.rest; + const { opcode, data } = frame; + if (opcode === 1) frames.push(JSON.parse(data.toString())); + else if (opcode === 9) writeFrame(socket, data, 10); + else if (opcode === 8) socket.end(); + else assert.equal(opcode, 10, 'Unexpected websocket opcode'); + } + } catch (error) { + failure = error; + socket.destroy(); + } + }); + } catch (error) { + failure = error; + socket.destroy(); + } + }); + await new Promise((resolve) => server.listen(0, '127.0.0.1', resolve)); + return { + baseUrl: `http://127.0.0.1:${server.address().port}`, + frames, + check() { + if (failure) throw failure; + }, + send(value) { + assert(node && !node.destroyed, 'Node wire must be connected'); + writeFrame(node, Buffer.from(JSON.stringify(value))); + }, + async close() { + for (const socket of sockets) socket.destroy(); + await new Promise((resolve) => server.close(resolve)); + }, + }; +} + +export function readClientFrame(buffer) { + if (buffer.length < 2) return undefined; + const opcode = buffer[0] & 15; + assert(buffer[0] & 128, 'Fixture expects unfragmented broker frames'); + assert(buffer[1] & 128, 'Client websocket frame must be masked'); + const indicator = buffer[1] & 127; + assert(indicator !== 127, 'Unexpected oversized broker frame'); + let length = indicator; + let offset = 2; + if (indicator === 126) { + if (buffer.length < 4) return undefined; + length = buffer.readUInt16BE(2); + offset = 4; + } + if (buffer.length < offset + 4 + length) return undefined; + const mask = buffer.subarray(offset, offset + 4); + const data = Buffer.from(buffer.subarray(offset + 4, offset + 4 + length)); + for (let i = 0; i < length; i++) data[i] ^= mask[i % 4]; + return { opcode, data, rest: buffer.subarray(offset + 4 + length) }; +} + +function writeFrame(socket, payload, opcode = 1) { + const header = Buffer.alloc(payload.length < 126 ? 2 : 4); + header[0] = 128 | opcode; + header[1] = payload.length < 126 ? payload.length : 126; + if (header.length === 4) header.writeUInt16BE(payload.length, 2); + socket.write(Buffer.concat([header, payload])); +} diff --git a/tests/relayflows/cases/1766-durable-task-receipt/engine-fixture.test.mjs b/tests/relayflows/cases/1766-durable-task-receipt/engine-fixture.test.mjs new file mode 100644 index 0000000000..e4e2a0215e --- /dev/null +++ b/tests/relayflows/cases/1766-durable-task-receipt/engine-fixture.test.mjs @@ -0,0 +1,24 @@ +import assert from 'node:assert/strict'; +import test from 'node:test'; +import { readClientFrame } from './engine-fixture.mjs'; + +function maskedFrame(payload) { + const mask = Buffer.from([0x12, 0x34, 0x56, 0x78]); + const header = Buffer.from([0x81, 0xfe, 0, payload.length]); + const masked = Buffer.from(payload); + for (let i = 0; i < masked.length; i++) masked[i] ^= mask[i % mask.length]; + return Buffer.concat([header, mask, masked]); +} + +test('accepts indicator 126 with a decoded 127-byte payload', () => { + const payload = Buffer.alloc(127, 0x61); + const frame = readClientFrame(maskedFrame(payload)); + assert(frame); + assert.equal(frame.opcode, 1); + assert.deepEqual(frame.data, payload); + assert.equal(frame.rest.length, 0); +}); + +test('rejects the unsupported 64-bit length indicator', () => { + assert.throws(() => readClientFrame(Buffer.from([0x81, 0xff])), /Unexpected oversized broker frame/); +}); diff --git a/tests/relayflows/cases/1766-durable-task-receipt/run.mjs b/tests/relayflows/cases/1766-durable-task-receipt/run.mjs new file mode 100644 index 0000000000..8968c99581 --- /dev/null +++ b/tests/relayflows/cases/1766-durable-task-receipt/run.mjs @@ -0,0 +1,242 @@ +import assert from 'node:assert/strict'; +import { execFileSync, spawn } from 'node:child_process'; +import { access, mkdtemp, mkdir, readFile, rm, writeFile } from 'node:fs/promises'; +import { tmpdir } from 'node:os'; +import path from 'node:path'; +import { fileURLToPath } from 'node:url'; +import { engineFixture } from './engine-fixture.mjs'; + +const caseId = '1766-durable-task-receipt'; +const required = (key) => { + assert(process.env[key], `Missing ${key}`); + return process.env[key]; +}; +const arm = required('RELAY_PR_PROOF_ARM'); +assert(['base', 'head'].includes(arm)); +const binary = path.resolve(required('RELAY_PR_PROOF_BROKER_BINARY')); +const target = required('RELAY_PR_PROOF_TARGET_DIR'); +const harness = required('RELAY_PR_PROOF_HARNESS_DIR'); +const resultPath = required('RELAY_PR_PROOF_RESULT_PATH'); +assert.equal( + execFileSync('git', ['-C', target, 'rev-parse', 'HEAD'], { encoding: 'utf8' }).trim(), + required(arm === 'base' ? 'RELAY_PR_PROOF_BASE_SHA' : 'RELAY_PR_PROOF_HEAD_SHA') +); +const relative = path.relative(path.resolve(harness), fileURLToPath(import.meta.url)); +assert(relative && !relative.startsWith('..') && !path.isAbsolute(relative)); +await access(binary, 1); +const directory = await mkdtemp(path.join(tmpdir(), 'relayflow-task-receipt-')); +const state = path.join(directory, 'state'); +await mkdir(state); +const engine = await engineFixture(); +let broker; +let logs = ''; +let cursor = 0; +const invocation = { + v: 1, + type: 'action.invoke', + invocation_id: 'inv-proof-task', + action: 'task.run', + // Deliberately missing cli: exercise a real terminal provider failure without + // launching a model or interpreting a process exit as a successful task. + input: { + task: 'fixture task', + task_context: { + run_id: 'run-proof', + step_id: 'step-proof', + dispatch_id: 'dispatch-proof', + timeout_ms: 120000, + }, + }, + task_execution: { + execution_id: 'inv-proof-task/1', + run_id: 'run-proof', + step_id: 'step-proof', + dispatch_id: 'dispatch-proof', + deadline: new Date(Date.now() + 120000).toISOString(), + }, +}; + +async function waitFor(probe, description, milliseconds = 15000) { + const deadline = performance.now() + milliseconds; + while (performance.now() < deadline) { + engine.check(); + if (broker?.exitCode !== null && broker?.exitCode !== undefined) + throw new Error(`Broker exited before ${description}: ${logs}`); + const value = await probe(); + if (value) return value; + await new Promise((resolve) => setTimeout(resolve, 20)); + } + throw new Error(`Timed out waiting for ${description}: ${logs}`); +} +async function next(type) { + return waitFor(() => { + while (cursor < engine.frames.length) { + const value = engine.frames[cursor++]; + if (value.type === type) return value; + assert( + !['action.accept', 'action.result'].includes(value.type), + `Unexpected task frame while waiting for ${type}` + ); + } + }, type); +} +async function start() { + logs = ''; + broker = spawn( + binary, + [ + 'init', + '--persist', + '--instance-name', + 'task-proof-node', + '--workspace-key', + 'rk_fixture_task_proof', + '--state-dir', + state, + '--api-port', + '0', + '--channels', + '', + ], + { + cwd: directory, + // Explicit allowlist: no live tokens or dynamic-loader injection inherited. + env: { + PATH: process.env.PATH, + TMPDIR: directory, + // This proof only consumes captured process output. Avoid the default + // rolling file under a runner-owned home directory. + AGENT_RELAY_BROKER_LOG: 'stderr', + RELAYCAST_BASE_URL: engine.baseUrl, + RELAY_BASE_URL: engine.baseUrl, + RELAY_BROKER_API_KEY: 'br_fixture_task_proof', + RELAY_NODE_ID: 'node-fixture-task-proof', + RELAY_NODE_TOKEN: 'nt_fixture_task_proof', + AGENT_RELAY_TASK_PROVIDER: '1', + AGENT_RELAY_TELEMETRY_DISABLED: '1', + AGENT_RELAY_NO_DEBUG_FILES: '1', + }, + stdio: ['ignore', 'pipe', 'pipe'], + } + ); + broker.stdout.on('data', (chunk) => { + logs = (logs + chunk).slice(-10000); + }); + broker.stderr.on('data', (chunk) => { + logs = (logs + chunk).slice(-10000); + }); + const registration = await next('node.register'); + await next('inventory.sync'); + return registration; +} +async function stop(signal = 'SIGTERM') { + if (!broker || broker.exitCode !== null || broker.signalCode !== null) return; + const child = broker; + const exited = new Promise((resolve) => child.once('exit', resolve)); + child.kill(signal); + const timer = setTimeout(() => child.kill('SIGKILL'), 2000); + await exited; + clearTimeout(timer); + broker = undefined; +} +function receipt(accept, status) { + return { + invocation_id: invocation.invocation_id, + action_name: 'task.run', + status, + task_execution: { + ...invocation.task_execution, + worker_generation: accept.worker_generation, + accepted_at: '2026-01-01T00:00:00.000Z', + }, + output: null, + error: status === 'failed' ? 'task_missing_cli' : null, + completed_at: status === 'failed' ? '2026-01-01T00:00:01.000Z' : null, + }; +} +const reply = (id, data) => engine.send({ v: 1, type: 'reply', id, ok: true, data }); +try { + const registration = await start(); + const capability = registration.capabilities.find((value) => value.name === 'task.run'); + if (arm === 'base') { + assert.equal(capability, undefined); + // Base rejects the known action name using its supported short-action wire; + // no missing-test or deserialization failure is counted as feature absence. + const { task_execution: ignored, ...legacy } = invocation; + engine.send(legacy); + const result = await next('action.result'); + assert.equal(result.invocation_id, invocation.invocation_id); + assert.equal(result.error, 'handler_unavailable'); + assert.equal(result.final, undefined); + } else { + assert.equal(capability.execution_mode, 'task'); + assert.equal(capability.global, true); + engine.send(invocation); + const accepted = await next('action.accept'); + assert.equal(accepted.invocation_id, invocation.invocation_id); + assert.equal(accepted.execution_id, invocation.task_execution.execution_id); + assert.match(accepted.worker_generation, /^[0-9a-f-]{36}$/); + engine.send(invocation); // duplicate delivery before the engine accepts + await new Promise((resolve) => setTimeout(resolve, 300)); + engine.check(); + assert( + !engine.frames.some((frame) => frame.type === 'action.result'), + 'Task failed/completed before durable acceptance' + ); + reply(accepted.id, receipt(accepted, 'running')); + const reconciliation = await next('action.accept'); + assert.equal(reconciliation.worker_generation, accepted.worker_generation); + reply(reconciliation.id, receipt(accepted, 'running')); + const terminal = await next('action.result'); + assert.equal(terminal.invocation_id, invocation.invocation_id); + assert.equal(terminal.execution_id, accepted.execution_id); + assert.equal(terminal.worker_generation, accepted.worker_generation); + assert.equal(terminal.final, true); + assert.equal(terminal.error, 'task_missing_cli'); + assert(terminal.id); + // The fixture engine commits the result but loses the ACK. Kill the actual + // broker, preserving its ledger, then reconcile with the terminal receipt. + await stop('SIGKILL'); + cursor = engine.frames.length; + await start(); + const recovered = await next('action.accept'); + assert.equal(recovered.invocation_id, accepted.invocation_id); + assert.equal(recovered.execution_id, accepted.execution_id); + assert.equal(recovered.worker_generation, accepted.worker_generation); + reply(recovered.id, receipt(accepted, 'failed')); + const ledgerPath = path.join(state, 'state-task-proof-node.tasks.json'); + await waitFor(async () => { + const ledger = JSON.parse(await readFile(ledgerPath, 'utf8')); + return ledger[invocation.invocation_id]?.receipt?.status === 'failed'; + }, 'durably stored reconciled terminal receipt'); + const ledger = JSON.parse(await readFile(ledgerPath, 'utf8')); + assert.equal(ledger[invocation.invocation_id].receipt.error, 'task_missing_cli'); + assert.equal(ledger[invocation.invocation_id].generation, accepted.worker_generation); + assert.equal( + engine.frames.filter((frame) => frame.type === 'action.result').length, + 1, + 'Final result duplicated after terminal reconciliation' + ); + } + engine.check(); + await mkdir(path.dirname(resultPath), { recursive: true }); + await writeFile( + resultPath, + JSON.stringify({ + version: 1, + caseId, + arm, + outcome: arm === 'base' ? 'absent' : 'fixed', + signature: + arm === 'base' ? 'task_action_handler_unavailable' : 'fenced_task_failure_receipt_survives_restart', + details: + arm === 'base' + ? 'Actual base broker returned handler_unavailable for task.run.' + : 'Actual head broker waited for acceptance, sent a fenced explicit failure, and durably reconciled the same generation after SIGKILL and lost final ACK. Loopback engine fixture; no deployed engine or model execution claim.', + }) + '\n' + ); +} finally { + await stop(); + await engine.close(); + await rm(directory, { recursive: true, force: true }); +}