From 78f3ed73b4037dacac96f75851540156829e123b Mon Sep 17 00:00:00 2001 From: Tautui Date: Sat, 31 Oct 2026 10:11:42 -0400 Subject: [PATCH 01/15] feat(cp): add [control_plane] config and runtime client modules Port of the openab-core runtime control-plane client from feat/cp-runtime-client-main 9455b504 (PR 1471) onto origin/main 6214700a. Plan: .agents/tasks/phase0-cp-runtime/plan.md, Commit 1. Taken as-is: - control_plane/mod.rs, client.rs, executor.rs (new files) - config.rs: ControlPlaneConfig, CpAgentType, the [control_plane] validation, and the control_plane_* config tests. The AgentCoreConfig whitespace hunk in the same commit is not taken. - lib.rs: pub mod control_plane - crates/openab-core/Cargo.toml: runtime openab-cp dependency with default-features = false. The dev-dependencies (openab-cp server, axum) are left for commit 6, because only the cp_client test needs them. - cron.rs, directives.rs: the two format! to to_string clippy hunks Applied 3-way onto main, keeping main code: - acp/pool.rs: see deviation below - adapter.rs: PromptExecution and the stream_prompt_blocks return value - dispatch.rs: the .map(|_| ()) adapter Deviations from the plan: - acp/pool.rs: main already has discard_session returning Result<()>, which the dispatch rollback uses. The branch version returns bool and does not send session/cancel, so it is added as discard_single_use_session. Main's discard_session is unchanged. The executor is the only caller and was updated. - crates/openab-cp/src/proto.rs: MIN_RUNTIME_FRAME_BYTES (4 lines) ported from the branch. The client uses it, and main does not define it. The plan said #1588 added it, but main proto.rs has no change since the merge base. - Cargo.lock: cargo added one line, "openab-cp", under openab-core. - executor.rs: rustfmt applied. It is the only new fmt diff this commit adds. Verification (macmini, cargo and rustc 1.99.1, clone ~/repo/openab-phase0): - cargo check -p openab-core -p openab: pass - cargo clippy -p openab-core -p openab -- -D warnings: pass - cargo test -p openab-core --lib control_plane: 50 passed, 0 failed - cargo test -p openab-core --lib: 754 passed, 1 failed: secrets::tests::resolve_exec_nonzero_exit. The same test fails on a clean origin/main export (6214700a), so it predates this commit. Not fixed here. - cargo fmt --check: no new diffs in files this commit touches, apart from executor.rs, which rustfmt was applied to. Main also has fmt diffs. - Linux run: not needed for this commit. --- Cargo.lock | 1 + crates/openab-core/Cargo.toml | 4 + crates/openab-core/src/acp/pool.rs | 29 + crates/openab-core/src/adapter.rs | 72 +- crates/openab-core/src/config.rs | 331 ++++- .../openab-core/src/control_plane/client.rs | 977 +++++++++++++ .../openab-core/src/control_plane/executor.rs | 1237 +++++++++++++++++ crates/openab-core/src/control_plane/mod.rs | 51 + crates/openab-core/src/cron.rs | 2 +- crates/openab-core/src/directives.rs | 2 +- crates/openab-core/src/dispatch.rs | 4 + crates/openab-core/src/lib.rs | 3 + crates/openab-cp/src/proto.rs | 4 + 13 files changed, 2711 insertions(+), 6 deletions(-) create mode 100644 crates/openab-core/src/control_plane/client.rs create mode 100644 crates/openab-core/src/control_plane/executor.rs create mode 100644 crates/openab-core/src/control_plane/mod.rs diff --git a/Cargo.lock b/Cargo.lock index 2e3928104..af36789ac 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2522,6 +2522,7 @@ dependencies = [ "http 1.4.2", "image", "libc", + "openab-cp", "pulldown-cmark", "rand 0.8.6", "regex", diff --git a/crates/openab-core/Cargo.toml b/crates/openab-core/Cargo.toml index 51aca393f..10d5811dc 100644 --- a/crates/openab-core/Cargo.toml +++ b/crates/openab-core/Cargo.toml @@ -54,6 +54,10 @@ http = { version = "1", optional = true } # axum listener) was removed with the per-session proxy. rmcp = { version = "1.7", default-features = false, optional = true } tokio-util = { version = "0.7", optional = true } +# Control-plane wire types ONLY (`openab_cp::proto`). `default-features = false` +# keeps the CP's `server` feature — axum, its own tokio flavour, the registry — +# out of the runtime build; the runtime is a WebSocket *client* of the CP. +openab-cp = { path = "../openab-cp", default-features = false } [target.'cfg(unix)'.dependencies] libc = "0.2" diff --git a/crates/openab-core/src/acp/pool.rs b/crates/openab-core/src/acp/pool.rs index a3a28f87c..219482e9d 100644 --- a/crates/openab-core/src/acp/pool.rs +++ b/crates/openab-core/src/acp/pool.rs @@ -917,6 +917,35 @@ impl SessionPool { Ok(had_state) } + /// Drop a session and all its bookkeeping WITHOUT sending + /// `session/cancel` first. Returns `true` when there was an active + /// connection to drop. + /// + /// [`Self::reset_session`] is the same teardown *plus* a cancel and an + /// error when the session is unknown, which suits the interactive + /// `/reset` it serves. Non-interactive owners of single-use sessions — + /// the control-plane executor, which runs one fresh session per + /// delegation — need neither: after a completed turn there is nothing to + /// cancel, and after a cancelled one the cancel has already been sent. + /// Calling `reset_session` there would emit a spurious `session/cancel` + /// at the agent and log an error for the benign already-gone case. + /// + /// The ACP process exits once the last `Arc` to its connection drops, so + /// removing the map entry is what reclaims the pool slot. + pub async fn discard_single_use_session(&self, thread_id: &str) -> bool { + let mut state = self.state.write().await; + let had_active = state.active.remove(thread_id).is_some(); + purge_session_entries(&mut state, thread_id); + #[cfg(feature = "acp-mcp")] + revoke_facade_token_for_key(&mut state, thread_id, self.session_registrar.as_ref()); + self.save_mapping(&state.persisted); + self.save_meta(&state.session_workdirs); + if had_active { + info!(thread_id = %crate::redact::redact_session_ids(thread_id), "session discarded"); + } + had_active + } + pub async fn cleanup_idle(&self, ttl_secs: u64) { let cutoff = Instant::now() - std::time::Duration::from_secs(ttl_secs); let hung_threshold = std::time::Duration::from_secs(self.hung_threshold_secs); diff --git a/crates/openab-core/src/adapter.rs b/crates/openab-core/src/adapter.rs index e899a586e..156836193 100644 --- a/crates/openab-core/src/adapter.rs +++ b/crates/openab-core/src/adapter.rs @@ -466,6 +466,26 @@ pub trait ChatAdapter: Send + Sync + 'static { // --- AdapterRouter --- +/// Outcome of one ACP turn driven by [`AdapterRouter::stream_prompt_blocks`]. +/// +/// Platform callers discard this — for them a turn either delivered (`Ok`) or +/// did not (`Err`), which is what they always consumed. It exists for callers +/// that are not a chat platform and therefore cannot read the reply back off a +/// channel: the control-plane executor has to turn one turn into a +/// `cp/delegate_result` status, and needs the text plus the two failure signals +/// the rendered message merely *hints* at. +#[derive(Debug, Clone, Default)] +pub struct PromptExecution { + /// The final content as delivered (post table-conversion, pre-chunking). + pub final_text: String, + /// The agent/broker-level error that ended the turn, if any — the raw + /// message, without the `⚠️` presentation prefix. + pub terminal_error: Option, + /// The turn produced no content and reported zero output tokens: a + /// provider/model/auth failure masquerading as an empty success. + pub silent_failure: bool, +} + /// Shared logic for routing messages to ACP agents, managing sessions, /// streaming edits, and controlling reactions. Platform-independent. pub struct AdapterRouter { @@ -701,11 +721,20 @@ impl AdapterRouter { None, ) .await + // This path delivers to a chat platform: the turn summary has no + // consumer here, only its success/failure. + .map(|_| ()) } /// Drive one ACP turn with the given pre-packed ContentBlocks. /// Called by both `handle_message` (per-message mode) and `dispatch::dispatch_batch` /// (batched mode). + /// + /// Returns the [`PromptExecution`] summary of the turn. Platform callers + /// ignore it (`.map(|_| ())`) — they only ever cared about `Ok`/`Err` — but + /// the control-plane executor has no channel to read the reply back out of, + /// so it needs the turn's outcome as a value. `Err` still means exactly what + /// it meant before: the user-visible delivery is incomplete. #[allow(clippy::too_many_arguments)] pub async fn stream_prompt_blocks( &self, @@ -716,7 +745,7 @@ impl AdapterRouter { reactions: Arc, other_bot_present: bool, recipient: Option<(String, String)>, - ) -> Result<()> { + ) -> Result { let adapter = adapter.clone(); let thread_channel = thread_channel.clone(); let message_limit = reply_message_limit(&thread_channel.platform, adapter.message_limit()); @@ -1124,6 +1153,11 @@ impl AdapterRouter { // it was emitted alongside. let (directives, text_buf) = split_delivery(&text_buf, answer_start, keep_full_text); + // Classify from the raw answer body BEFORE tool summaries, + // reset notices, and presentation diagnostics are composed. + // Otherwise a tool-only silent failure looks non-empty and + // is incorrectly reported to a delegation initiator as success. + let silent_failure = is_silent_execution(&text_buf, &turn_result); // The session-reset notice lives at the head of the buffer; a // tool advancing answer_start past it would drop it from the // slice, so re-prepend it to the (directive-stripped) body in @@ -1134,6 +1168,12 @@ impl AdapterRouter { // Build final content let final_content = display_for(platform_is_acp, &tool_lines, &text_buf, false, tool_display); + // Captured for `PromptExecution` before the composition + // below consumes `response_error` and rewrites the body: + // a caller that has no channel to read (the control-plane + // executor) needs the *classification*, not the rendered + // "⚠️ …" prefix, to decide Completed vs Failed. + let terminal_error = response_error.clone(); let final_content = if final_content.is_empty() { if turn_result.is_silent_failure() { warn!( @@ -1370,7 +1410,11 @@ impl AdapterRouter { "streaming finalization had delivery failures; user view is incomplete" )) } else { - Ok(()) + Ok(PromptExecution { + final_text: final_content, + terminal_error, + silent_failure, + }) } }) }) @@ -1506,6 +1550,12 @@ fn render_group(title: &str, state: ToolState, count: usize) -> String { /// Message to show the consumer when a silent failure is detected. pub(crate) const SILENT_FAILURE_MSG: &str = "⚠️ The agent did not produce a response. This usually indicates a backend configuration issue — not an intentional empty reply. Please try again later."; +/// Whether the raw final answer represents a silent provider/model failure. +/// Tool summaries and other presentation text must never influence this signal. +pub(crate) fn is_silent_execution(answer_text: &str, turn_result: &TurnResult) -> bool { + answer_text.trim().is_empty() && turn_result.is_silent_failure() +} + /// Classify what to display when the composed body is empty. /// Returns the final content string for the consumer. pub(crate) fn classify_empty_turn( @@ -2364,7 +2414,7 @@ mod tests { #[cfg(test)] mod directive_tests { use super::parse_output_directives; - use super::{classify_empty_turn, SILENT_FAILURE_MSG}; + use super::{classify_empty_turn, is_silent_execution, SILENT_FAILURE_MSG}; use crate::acp::TurnResult; #[test] @@ -2520,6 +2570,22 @@ mod directive_tests { // --- classify_empty_turn: adapter-level finalization tests --- + #[test] + fn silent_failure_uses_raw_answer_not_rendered_tool_text() { + let tr = TurnResult { + stop_reason: Some("end_turn".into()), + output_tokens: Some(0), + input_tokens: Some(12), + total_tokens: Some(12), + }; + assert!(is_silent_execution("", &tr)); + assert!(is_silent_execution(" ", &tr)); + assert!( + !is_silent_execution("actual answer", &tr), + "real answer text defeats a contradictory zero-token report" + ); + } + #[test] fn empty_turn_silent_failure_produces_diagnostic() { let tr = TurnResult { diff --git a/crates/openab-core/src/config.rs b/crates/openab-core/src/config.rs index fdd89dcb3..e2a3ec21d 100644 --- a/crates/openab-core/src/config.rs +++ b/crates/openab-core/src/config.rs @@ -1,8 +1,9 @@ use crate::markdown::TableMode; use regex::Regex; use serde::Deserialize; -use std::collections::HashMap; +use std::collections::{BTreeMap, HashMap}; use std::path::Path; +use tokio_tungstenite::tungstenite::client::IntoClientRequest; /// Controls how incoming messages are dispatched to ACP turns. /// @@ -132,6 +133,121 @@ fn default_mcp_listen() -> String { "127.0.0.1:8848".to_string() } +/// `[control_plane]` — enrol this runtime with an OpenAB Agent Control Plane +/// (`docs/adr/agent-control-plane.md`). Presence is the opt-in signal, exactly +/// like [`McpFacadeConfig`]: absent section = no outbound connection, no +/// delegation serving, no behaviour change. There is deliberately **no cargo +/// feature** for it — the config section is the switch, so one binary and one +/// image serve both plain chat brokers and control-plane members. +/// +/// **Strict.** An unknown key here is a hard startup failure rather than a +/// silently-defaulted one: a mistyped `max_delegated_sessions` would otherwise +/// look effective while the runtime advertised the default budget of 1. +#[derive(Clone, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct ControlPlaneConfig { + /// CP WebSocket endpoint, e.g. `wss://cp.internal:9800/cp`. The server + /// mounts the socket at `/cp`. + pub url: String, + /// Explicit opt-in for sending the bearer key over cleartext `ws://` to a + /// non-loopback host. Defaults to false; loopback `ws://` and all `wss://` + /// endpoints do not need this escape hatch. + #[serde(default)] + pub allow_insecure_transport: bool, + /// Bearer key presented on the upgrade request. The CP maps it to the + /// immutable identity claims below and rejects any mismatch, so this value + /// is the whole credential: it is never logged, and never reaches the + /// agent subprocess (the child env is `env_clear()`ed and this key is not + /// in `[agent].env`). + pub auth_key: String, + /// Asserted namespace. Verified against the key's bound claims by the CP. + pub namespace: String, + /// Asserted logical agent name. Verified against the key's bound claims. + pub name: String, + /// Asserted role. Only `primary` and `worker` are constructible here: + /// `observer` is a read-only lobby client (a separate, non-runtime + /// consumer of the same protocol), so an OAB runtime must never be able to + /// register as one by editing its own config. + #[serde(rename = "type")] + pub agent_type: CpAgentType, + /// Advertised selector labels (e.g. `backend = "kiro"`), used by + /// `cp/delegate` label targeting. + #[serde(default)] + pub labels: BTreeMap, + /// Concurrency budget advertised at registration. The CP may clamp it + /// (the ack may return a lower value); the runtime enforces whatever the + /// ack returns. + #[serde(default = "default_max_delegated_sessions")] + pub max_delegated_sessions: u32, +} + +/// Runtime-side agent role for `[control_plane].type`. +/// +/// Intentionally NOT `openab_cp::proto::AgentType`: that enum also has +/// `Observer`, which a runtime must not be able to claim. Keeping a separate, +/// smaller enum makes that unrepresentable in config instead of relying on a +/// validation check somebody can delete. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Deserialize)] +#[serde(rename_all = "lowercase")] +pub enum CpAgentType { + Primary, + Worker, +} + +impl std::fmt::Display for CpAgentType { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::Primary => write!(f, "primary"), + Self::Worker => write!(f, "worker"), + } + } +} + +fn default_max_delegated_sessions() -> u32 { + 1 +} + +impl std::fmt::Debug for ControlPlaneConfig { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("ControlPlaneConfig") + .field("url", &self.url) + .field("allow_insecure_transport", &self.allow_insecure_transport) + .field("auth_key", &"[REDACTED]") + .field("namespace", &self.namespace) + .field("name", &self.name) + .field("agent_type", &self.agent_type) + .field("labels", &self.labels) + .field("max_delegated_sessions", &self.max_delegated_sessions) + .finish() + } +} + +fn control_plane_transport_is_insecure(cp: &ControlPlaneConfig) -> anyhow::Result { + let request = cp + .url + .as_str() + .into_client_request() + .map_err(|e| anyhow::anyhow!("control_plane.url is not a valid WebSocket URL: {e}"))?; + let uri = request.uri(); + let scheme = uri + .scheme_str() + .ok_or_else(|| anyhow::anyhow!("control_plane.url must use ws:// or wss://"))?; + if scheme == "wss" { + return Ok(false); + } + anyhow::ensure!(scheme == "ws", "control_plane.url must use ws:// or wss://"); + let host = uri + .host() + .ok_or_else(|| anyhow::anyhow!("control_plane.url must include a host"))?; + let host = host.trim_start_matches('[').trim_end_matches(']'); + let loopback = host.eq_ignore_ascii_case("localhost") + || host + .parse::() + .map(|ip| ip.is_loopback()) + .unwrap_or(false); + Ok(!loopback) +} + #[derive(Debug, Clone, Deserialize)] pub struct AgentCoreConfig { /// AgentCore Runtime ARN (required) pub runtime_arn: String, @@ -248,6 +364,10 @@ pub struct Config { /// OAB MCP Facade (`[mcp]` — OAB MCP Adapter ADR §6.2/§6.3). Presence is /// the opt-in signal: absent = no facade, no listener, no new behavior. pub mcp: Option, + /// Agent Control Plane membership (`[control_plane]` — Agent Control Plane + /// ADR). Presence is the opt-in signal: absent = no CP connection, no + /// delegation serving, no new behavior. + pub control_plane: Option, #[serde(default)] pub agent: AgentConfig, #[serde(default)] @@ -2346,6 +2466,36 @@ fn parse_config_inner(expanded: &str, source: &str) -> anyhow::Result { "pool.liveness_check_secs must be > 0 (zero would spin the recv loop)" ); + // `[control_plane]`: every field is load-bearing at registration time, and + // the CP rejects a mismatch by closing the socket — which the client then + // retries forever. Failing here turns that silent reconnect loop into one + // startup error naming the empty field. `auth_key` is checked for emptiness + // only; its value is never echoed. + if let Some(ref cp) = config.control_plane { + for (field, value) in [ + ("url", &cp.url), + ("auth_key", &cp.auth_key), + ("namespace", &cp.namespace), + ("name", &cp.name), + ] { + anyhow::ensure!( + !value.trim().is_empty(), + "control_plane.{field} must not be empty" + ); + } + let insecure_non_loopback = control_plane_transport_is_insecure(cp)?; + anyhow::ensure!( + !insecure_non_loopback || cp.allow_insecure_transport, + "control_plane.url uses cleartext ws:// on a non-loopback host; use wss://, \ + a loopback endpoint, or explicitly set allow_insecure_transport = true" + ); + anyhow::ensure!( + cp.max_delegated_sessions > 0, + "control_plane.max_delegated_sessions must be > 0 \ + (zero would advertise capacity the runtime can never serve)" + ); + } + Ok(config) } @@ -2560,6 +2710,185 @@ mod tests { .unwrap(); assert_eq!(cfg.mcp.unwrap().listen, "127.0.0.1:9000"); } + + // --- [control_plane] (Agent Control Plane ADR) --- + + /// The whole opt-in contract: no section, no membership. + #[test] + fn control_plane_absent_by_default() { + let cfg = parse_config_str("[discord]\nbot_token = \"x\"\n", "test").unwrap(); + assert!(cfg.control_plane.is_none()); + } + + #[test] + fn control_plane_full_section_parses_with_the_type_rename() { + std::env::set_var("AB_TEST_CP_KEY", "s3cr3t-from-env"); + let cfg = parse_config( + r#" +[discord] +bot_token = "x" + +[control_plane] +url = "wss://cp.internal:9800/cp" +auth_key = "${AB_TEST_CP_KEY}" +namespace = "prod" +name = "worker-1" +type = "worker" +max_delegated_sessions = 4 + +[control_plane.labels] +backend = "kiro" +tier = "batch" +"#, + "test", + ) + .unwrap(); + std::env::remove_var("AB_TEST_CP_KEY"); + let cp = cfg + .control_plane + .expect("[control_plane] presence is the opt-in"); + assert_eq!(cp.url, "wss://cp.internal:9800/cp"); + // `${ENV}` expansion comes free with the shared loader — the config file + // itself never has to hold the credential. + assert_eq!(cp.auth_key, "s3cr3t-from-env"); + assert_eq!(cp.namespace, "prod"); + assert_eq!(cp.name, "worker-1"); + assert_eq!(cp.agent_type, CpAgentType::Worker); + assert_eq!(cp.max_delegated_sessions, 4); + assert_eq!(cp.labels.get("backend").map(String::as_str), Some("kiro")); + assert_eq!(cp.labels.get("tier").map(String::as_str), Some("batch")); + } + + #[test] + fn control_plane_defaults_are_conservative() { + let cfg = parse_config_str( + "[discord]\nbot_token = \"x\"\n[control_plane]\nurl = \"wss://cp:9800/cp\"\n\ + auth_key = \"k\"\nnamespace = \"prod\"\nname = \"koudu\"\ntype = \"primary\"\n", + "test", + ) + .unwrap(); + let cp = cfg.control_plane.unwrap(); + assert_eq!(cp.agent_type, CpAgentType::Primary); + assert!(cp.labels.is_empty(), "no labels unless asked for"); + assert_eq!( + cp.max_delegated_sessions, 1, + "one at a time until an operator says otherwise" + ); + } + + /// `observer` is a read-only lobby client, not a runtime role. A config + /// claiming it must fail loudly rather than registering as something else. + #[test] + fn control_plane_refuses_the_observer_role() { + let err = parse_config_str( + "[control_plane]\nurl = \"wss://cp:9800/cp\"\nauth_key = \"k\"\n\ + namespace = \"prod\"\nname = \"lobby\"\ntype = \"observer\"\n", + "test", + ) + .expect_err("observer is not a runtime role"); + let msg = err.to_string(); + assert!( + msg.contains("observer") || msg.contains("primary"), + "the error must name the offending value, got: {msg}" + ); + } + + #[test] + fn control_plane_unknown_key_is_a_hard_failure() { + let err = parse_config_str( + "[control_plane]\nurl = \"wss://cp:9800/cp\"\nauth_key = \"k\"\n\ + namespace = \"prod\"\nname = \"w\"\ntype = \"worker\"\nmax_delegated_session = 4\n", + "test", + ) + .expect_err("a mistyped key must not default silently"); + assert!(err.to_string().contains("max_delegated_session")); + } + + #[test] + fn control_plane_empty_fields_are_rejected_by_name() { + for (field, body) in [ + ( + "url", + "url = \"\"\nauth_key = \"k\"\nnamespace = \"p\"\nname = \"w\"\ntype = \"worker\"", + ), + ( + "auth_key", + "url = \"ws://c/cp\"\nauth_key = \" \"\nnamespace = \"p\"\nname = \"w\"\ntype = \"worker\"", + ), + ( + "namespace", + "url = \"ws://c/cp\"\nauth_key = \"k\"\nnamespace = \"\"\nname = \"w\"\ntype = \"worker\"", + ), + ( + "name", + "url = \"ws://c/cp\"\nauth_key = \"k\"\nnamespace = \"p\"\nname = \"\"\ntype = \"worker\"", + ), + ] { + let err = parse_config_str(&format!("[control_plane]\n{body}\n"), "test").expect_err( + "an empty identity field would fail at registration and reconnect forever", + ); + let msg = err.to_string(); + assert!( + msg.contains(&format!("control_plane.{field}")), + "the error must name the empty field; got: {msg}" + ); + // An unexpanded `${VAR}` for a secret that is not set expands to the + // empty string, which is exactly this case. + assert!(msg.contains("must not be empty"), "got: {msg}"); + } + } + + #[test] + fn control_plane_zero_capacity_is_rejected() { + let err = parse_config_str( + "[control_plane]\nurl = \"wss://c/cp\"\nauth_key = \"k\"\nnamespace = \"p\"\n\ + name = \"w\"\ntype = \"worker\"\nmax_delegated_sessions = 0\n", + "test", + ) + .expect_err("advertising zero capacity would make every delegation fail on arrival"); + assert!(err + .to_string() + .contains("control_plane.max_delegated_sessions must be > 0")); + } + #[test] + fn control_plane_transport_requires_tls_or_an_explicit_exception() { + let common = "auth_key = \"k\"\nnamespace = \"prod\"\nname = \"w\"\ntype = \"worker\"\n"; + parse_config_str( + &format!("[control_plane]\nurl = \"ws://127.0.0.1:9800/cp\"\n{common}"), + "test", + ) + .expect("loopback cleartext"); + parse_config_str( + &format!("[control_plane]\nurl = \"ws://[::1]:9800/cp\"\n{common}"), + "test", + ) + .expect("IPv6 loopback cleartext"); + parse_config_str( + &format!("[control_plane]\nurl = \"wss://cp.example/cp\"\n{common}"), + "test", + ) + .expect("TLS remote"); + let err = parse_config_str( + &format!("[control_plane]\nurl = \"ws://cp.example/cp\"\n{common}"), + "test", + ) + .expect_err("remote cleartext"); + assert!(err.to_string().contains("allow_insecure_transport")); + let cfg = parse_config_str(&format!("[control_plane]\nurl = \"ws://cp.internal/cp\"\nallow_insecure_transport = true\n{common}"), "test").unwrap(); + assert!(cfg.control_plane.unwrap().allow_insecure_transport); + } + + #[test] + fn control_plane_debug_always_redacts_the_auth_key() { + let cfg = parse_config_str("[control_plane]\nurl = \"wss://cp.example/cp\"\nauth_key = \"top-secret-key\"\nnamespace = \"prod\"\nname = \"w\"\ntype = \"worker\"\n", "test").unwrap().control_plane.unwrap(); + let debug = format!("{cfg:?}"); + assert!( + !debug.contains("top-secret-key"), + "credential leaked: {debug}" + ); + assert!(debug.contains("[REDACTED]")); + } + use std::io::Write; #[test] diff --git a/crates/openab-core/src/control_plane/client.rs b/crates/openab-core/src/control_plane/client.rs new file mode 100644 index 000000000..169f85bcd --- /dev/null +++ b/crates/openab-core/src/control_plane/client.rs @@ -0,0 +1,977 @@ +//! Control-plane connection state machine. +//! +//! ```text +//! connect ──► cp/register ──► ack ──► serve loop ──► close/error ──► backoff ──┐ +//! ▲ │ +//! └────────────────────────── re-register ─────────────────────────────────┘ +//! ``` +//! +//! One instance id for the whole process, reused across reconnects (it +//! distinguishes replicas of the same logical agent, not connections). One task +//! owns the WebSocket sink, so there is no lock around it and no interleaved +//! frame: everything the runtime says to the CP — heartbeats, delegation +//! results, request acks — is produced by a single `select!`. +//! +//! Backoff follows the gateway adapter's shape (1/2/4/8/16/30s, shutdown-aware), +//! because the failure modes are the same: a hub that is briefly down, a rolling +//! restart, or a config the CP rejects. A rejected registration keeps retrying +//! rather than exiting — a runtime that has a chat platform must keep serving it, +//! and the loud log line is what an operator acts on. + +use std::sync::Arc; +use std::time::Duration; + +use futures_util::{SinkExt, StreamExt}; +use openab_cp::proto::{ + codes, methods, DelegateForward, DelegateResultParams, ErrorObject, HeartbeatParams, + JsonRpcErrorResponse, JsonRpcMessage, JsonRpcRequest, JsonRpcResponse, RegisterAck, + RegisterParams, MIN_RUNTIME_FRAME_BYTES, PROTOCOL_VERSION, +}; +use rand::Rng; +use tokio::net::TcpStream; +use tokio::sync::watch; +use tokio_tungstenite::tungstenite::client::IntoClientRequest; +use tokio_tungstenite::tungstenite::http::HeaderValue; +use tokio_tungstenite::tungstenite::protocol::WebSocketConfig; +use tokio_tungstenite::tungstenite::Message; +use tokio_tungstenite::{MaybeTlsStream, WebSocketStream}; +use tracing::{debug, error, info, warn}; + +use crate::config::{ControlPlaneConfig, CpAgentType}; +use crate::control_plane::executor::{cap_text, DelegationExecutor, PromptRunner}; + +/// Backoff ceiling, matching the gateway adapter. +const MAX_BACKOFF_SECS: u64 = 30; +/// A session must live this long before a clean close resets the reconnect +/// backoff to 1s; shorter sessions escalate instead (anti reconnect-storm). +const STABLE_SESSION_SECS: u64 = 60; + +/// How long a lost connection's in-flight delegations get to unwind (cancel the +/// agent, drop the session) before their tasks are aborted outright. +const DRAIN_TIMEOUT: Duration = Duration::from_secs(5); + +/// Bidirectional WS message/frame ceiling, mirroring the CP server's own +/// `max_frame_bytes` default (1 MiB). The executor caps result/error fields; +/// [`send`] additionally caps each complete serialized outbound frame. +/// +/// Note (F51): The server default is 1 MiB but configurable hub-side. If the +/// hub is configured with a larger ceiling, frames exceeding 1 MiB are +/// rejected at the transport layer; future protocol revisions may negotiate +/// effective limits in `RegisterAck`. +const MAX_FRAME_BYTES: usize = MIN_RUNTIME_FRAME_BYTES; + +/// Bound on DNS/TCP/TLS/WebSocket establishment. The outer shutdown select can +/// cancel this sooner, but a blackholed endpoint must still enter backoff. +const CONNECT_TIMEOUT: Duration = Duration::from_secs(10); + +/// Bound on the wait for the `cp/register` ack, mirroring the CP's own +/// `register_timeout_secs` default. A CP that upgrades the socket but never +/// acks must land in backoff, not hang the client until shutdown. +const REGISTER_TIMEOUT: Duration = Duration::from_secs(10); + +type Ws = WebSocketStream>; +/// Write half. Split from the read half because the serve loop must be able to +/// write from a handler while the read future is still alive — one `select!` +/// cannot hold `&mut` to the whole socket in two branches. +type WsSink = futures_util::stream::SplitSink; +type WsStream = futures_util::stream::SplitStream; + +/// Runtime client for the OpenAB Agent Control Plane. +pub struct ControlPlaneClient { + cfg: ControlPlaneConfig, + /// Process-lifetime instance id, reused across reconnects. + instance_id: String, + executor: Arc, + /// Monotonic JSON-RPC request id for frames this client originates. + next_id: std::sync::atomic::AtomicU64, +} + +impl ControlPlaneClient { + pub fn new( + cfg: ControlPlaneConfig, + runner: Arc, + prompt_hard_timeout: Duration, + ) -> Self { + if cfg.allow_insecure_transport && cfg.url.starts_with("ws://") { + warn!(url = %cfg.url, "control-plane cleartext transport explicitly enabled; the bearer key is not encrypted by this connection"); + } + let instance_id = uuid::Uuid::new_v4().to_string(); + // Advertised budget until the first ack tells us the effective one. + let executor = Arc::new(DelegationExecutor::new( + runner, + instance_id.clone(), + cfg.max_delegated_sessions, + prompt_hard_timeout, + )); + Self { + cfg, + instance_id, + executor, + next_id: std::sync::atomic::AtomicU64::new(1), + } + } + + pub fn instance_id(&self) -> &str { + &self.instance_id + } + + fn next_id(&self) -> u64 { + self.next_id + .fetch_add(1, std::sync::atomic::Ordering::Relaxed) + } + + /// Connect, register, serve — forever, until `shutdown` flips. + pub async fn run(self: Arc, mut shutdown: watch::Receiver) { + let mut backoff = 1u64; + loop { + if *shutdown.borrow() { + info!("control-plane client shutting down"); + return; + } + info!( + agent = %format!("{}/{}", self.cfg.namespace, self.cfg.name), + r#type = %self.cfg.agent_type, + instance = %self.instance_id, + "connecting to control plane" + ); + let session_started = tokio::time::Instant::now(); + // connect_and_serve owns shutdown for each phase: connect/register + // race it directly, then the serve loop alone performs + // cancel/drain/abort. No outer arm may drop serve mid-cleanup. + let served = self.connect_and_serve(&mut shutdown).await; + match served { + Ok(Outcome::Shutdown) => { + info!("control-plane client shutting down"); + return; + } + Ok(Outcome::Disconnected) => { + // Reset the backoff only after a session that genuinely + // served for a while. A CP that accepts registration and + // then promptly closes (lease misconfig, crash loop, + // rolling deploys) would otherwise reconnect every second + // forever — a clean Close frame is not evidence of health. + backoff = backoff_after_session(backoff, session_started.elapsed()); + warn!( + backoff_secs = backoff, + "control-plane connection closed — reconnecting" + ); + } + Err(e) => { + error!(error = %format!("{e:#}"), backoff_secs = backoff, "control-plane connection failed"); + } + } + let retry_delay = jittered_backoff(backoff); + debug!( + base_backoff_secs = backoff, + retry_delay_ms = retry_delay.as_millis(), + "control-plane reconnect delay selected" + ); + tokio::select! { + _ = tokio::time::sleep(retry_delay) => {} + _ = shutdown.changed() => { + info!("control-plane client shutting down"); + return; + } + } + backoff = next_backoff(backoff); + } + } + + async fn connect_and_serve( + &self, + shutdown: &mut watch::Receiver, + ) -> anyhow::Result { + let ws = tokio::select! { + result = self.connect_with_timeout() => result?, + _ = shutdown.changed() => return Ok(Outcome::Shutdown), + }; + let (mut sink, mut stream) = ws.split(); + let ack = tokio::select! { + result = tokio::time::timeout(REGISTER_TIMEOUT, self.register(&mut sink, &mut stream)) => { + result.map_err(|_| anyhow::anyhow!( + "control plane did not ack registration within {}s", + REGISTER_TIMEOUT.as_secs() + ))?? + } + _ = shutdown.changed() => return Ok(Outcome::Shutdown), + }; + self.executor + .set_effective_max(ack.effective_max_delegated_sessions); + info!( + instance = %self.instance_id, + heartbeat_secs = ack.heartbeat_interval_secs, + lease_secs = ack.lease_expiry_secs, + max_delegated_sessions = ack.effective_max_delegated_sessions, + "registered with control plane" + ); + self.serve(sink, stream, &ack, shutdown).await + } + + async fn connect_with_timeout(&self) -> anyhow::Result { + tokio::time::timeout(CONNECT_TIMEOUT, self.connect()) + .await + .map_err(|_| { + anyhow::anyhow!( + "control-plane connection timed out after {}s", + CONNECT_TIMEOUT.as_secs() + ) + })? + } + + /// Dial the CP. The key travels in the `Authorization` header, never the + /// URL — the CP's own contract, so it cannot leak into an access log. + async fn connect(&self) -> anyhow::Result { + let mut request = self.cfg.url.as_str().into_client_request()?; + let bearer = format!("Bearer {}", self.cfg.auth_key); + let mut value = HeaderValue::from_str(&bearer) + .map_err(|_| anyhow::anyhow!("control_plane.auth_key is not a valid header value"))?; + // Belt and braces: the key must not surface in a `{:?}` of the request. + value.set_sensitive(true); + request.headers_mut().insert("Authorization", value); + // Mirror the CP's accept-side transport cap (`max_frame_bytes`, + // default 1 MiB). Without this the client would buffer tungstenite's + // 64 MiB default from an anomalous or misconfigured hub before any + // parsing runs. + let ws_config = WebSocketConfig { + max_message_size: Some(MAX_FRAME_BYTES), + max_frame_size: Some(MAX_FRAME_BYTES), + ..Default::default() + }; + let (ws, _resp) = + tokio_tungstenite::connect_async_with_config(request, Some(ws_config), false) + .await + .map_err(|e| anyhow::anyhow!("control-plane handshake failed: {e}"))?; + Ok(ws) + } + + /// Send the mandatory `cp/register` first frame and await its ack. + async fn register( + &self, + sink: &mut WsSink, + stream: &mut WsStream, + ) -> anyhow::Result { + let id = self.next_id(); + let params = RegisterParams { + protocol_version: PROTOCOL_VERSION, + namespace: self.cfg.namespace.clone(), + name: self.cfg.name.clone(), + agent_type: self.cfg.agent_type.into(), + instance_id: self.instance_id.clone(), + labels: self.cfg.labels.clone(), + max_delegated_sessions: self.cfg.max_delegated_sessions, + }; + let frame = + JsonRpcRequest::new(id, methods::REGISTER, Some(serde_json::to_value(¶ms)?)); + send(sink, &frame).await?; + + // Anything other than the ack to this id is a protocol violation at + // this point: registration is the first frame in both directions. + loop { + let Some(msg) = stream.next().await else { + anyhow::bail!("control plane closed the connection before acking cp/register"); + }; + match msg? { + Message::Text(text) => { + let parsed: JsonRpcMessage = serde_json::from_str(&text) + .map_err(|e| anyhow::anyhow!("malformed cp/register reply: {e}"))?; + if parsed.id != Some(id) { + warn!("ignoring an unexpected frame received before the register ack"); + continue; + } + if let Some(err) = parsed.error { + // Identity/version rejections are operator errors: name + // the code so the log line is actionable, and let the + // caller back off rather than exiting the process. + anyhow::bail!( + "control plane rejected cp/register: {} (code {})", + err.message, + err.code + ); + } + let result = parsed + .result + .ok_or_else(|| anyhow::anyhow!("cp/register reply carried no result"))?; + return validate_register_ack(result); + } + Message::Close(_) => { + anyhow::bail!("control plane closed the connection during registration") + } + _ => continue, + } + } + } + + /// The serve loop. Single owner of the sink; every outbound frame is + /// produced here. + /// + /// Finished delegations report themselves through an mpsc channel rather + /// than a `JoinSet` polled in the `select!`: the inbound branch has to + /// *spawn* while the completion branch is still armed, and one `select!` + /// cannot lend the same `JoinSet` to both. + async fn serve( + &self, + mut sink: WsSink, + mut stream: WsStream, + ack: &RegisterAck, + shutdown: &mut watch::Receiver, + ) -> anyhow::Result { + let mut heartbeat = tokio::time::interval(Duration::from_secs( + // A zero interval would spin; the CP's own default is 15s. + ack.heartbeat_interval_secs.max(1), + )); + heartbeat.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); + // The first tick is immediate: skip it, registration just happened. + heartbeat.tick().await; + + let (result_tx, mut result_rx) = tokio::sync::mpsc::channel::( + // One slot per delegation this runtime can ever admit, so a task + // never blocks handing its result over. + (ack.effective_max_delegated_sessions as usize).max(1), + ); + let mut serving: Vec> = Vec::new(); + + let outcome = loop { + tokio::select! { + _ = shutdown.changed() => break Outcome::Shutdown, + _ = heartbeat.tick() => { + let params = HeartbeatParams { + instance_id: self.instance_id.clone(), + active_delegated_sessions: self.executor.active(), + }; + let frame = JsonRpcRequest::new( + self.next_id(), + methods::HEARTBEAT, + Some(serde_json::to_value(¶ms)?), + ); + if send(&mut sink, &frame).await.is_err() { + break Outcome::Disconnected; + } + } + // A finished delegation reports itself. Emitted by the runtime + // when the turn ends, never by the model: this is the only + // frame that closes the initiator's wait. + Some(result) = result_rx.recv() => { + let frame = delegate_result_request(self.next_id(), result)?; + if send(&mut sink, &frame).await.is_err() { + break Outcome::Disconnected; + } + } + inbound = stream.next() => { + let Some(msg) = inbound else { break Outcome::Disconnected }; + match msg { + Ok(Message::Text(text)) => { + match self.handle_frame(&text) { + FrameAction::Serve { ack, forward } => { + let executor = Arc::clone(&self.executor); + let tx = result_tx.clone(); + serving.push(tokio::spawn(async move { + let result = executor.serve(forward).await; + // A closed channel means the connection + // that would carry this result is gone; + // the CP fails it as target_disconnected. + let _ = tx.send(result).await; + })); + // Prune finished tasks so a long-lived + // connection does not accumulate handles. + serving.retain(|h| !h.is_finished()); + if sink.send(Message::Text(ack)).await.is_err() { + break Outcome::Disconnected; + } + } + FrameAction::Reply(reply) => { + if sink.send(Message::Text(reply)).await.is_err() { + break Outcome::Disconnected; + } + } + FrameAction::Ignore => {} + } + } + Ok(Message::Ping(p)) => { + if sink.send(Message::Pong(p)).await.is_err() { + break Outcome::Disconnected; + } + } + Ok(Message::Close(_)) => { + // The CP closes on lease expiry: reconnecting and + // re-registering is the recovery, since + // registration is first-frame-only. + info!("control plane closed the connection"); + break Outcome::Disconnected; + } + Ok(_) => {} + Err(e) => { + warn!(error = %e, "control-plane WebSocket error"); + break Outcome::Disconnected; + } + } + } + } + }; + + // Whatever ended the session, nothing local may keep running: this + // connection is the only route a result could travel, and on the CP + // side these delegations are already (or about to be) failed as + // `target_disconnected`. Cancelling lets each task stop its agent and + // drop its session; the results themselves are deliberately dropped. + let in_flight = self.executor.active(); + if in_flight > 0 { + warn!( + in_flight, + "cancelling in-flight delegations; the control plane reports them as \ + target_disconnected" + ); + } + self.executor.cancel_all(); + let drained = tokio::time::timeout(DRAIN_TIMEOUT, async { + for handle in &mut serving { + let _ = handle.await; + } + }) + .await; + if drained.is_err() { + warn!("delegation tasks did not unwind within the drain window; aborting them"); + for handle in &serving { + handle.abort(); + } + } + let _ = sink.send(Message::Close(None)).await; + let _ = sink.close().await; + Ok(outcome) + } + + /// Classify one inbound frame. Never spawns and never writes: the serve + /// loop owns both, so this stays a pure-enough function to unit-test (it + /// does signal cancellation, which has no other home). + /// + /// CP-issued requests get a JSON-RPC result ack so the hub sees a + /// well-formed conversation; the *outcome* of a delegation never travels in + /// that ack — it comes later as `cp/delegate_result`, correlated by + /// `delegation_id`. + fn handle_frame(&self, text: &str) -> FrameAction { + let msg: JsonRpcMessage = match serde_json::from_str(text) { + Ok(m) => m, + Err(e) => { + warn!(error = %e, "malformed frame from the control plane"); + return FrameAction::Ignore; + } + }; + let Some(method) = msg.method.clone() else { + // A response to one of our own requests (heartbeat, delegate_result + // ack). Errors are worth a line; successes are noise. + if let Some(err) = msg.error { + warn!(code = err.code, message = %err.message, "control plane returned an error"); + } + return FrameAction::Ignore; + }; + let id = match msg.require_request_envelope() { + Ok(id) => id, + Err(err) => { + warn!(code = err.code, message = %err.message, "invalid request envelope from the control plane"); + return error_reply(msg.id.unwrap_or(0), err); + } + }; + + match method.as_str() { + methods::DELEGATE => { + if self.cfg.agent_type != CpAgentType::Worker { + return error_reply( + id, + ErrorObject::new( + codes::POLICY_DENIED, + "this runtime is registered as primary and does not serve delegations", + ), + ); + } + let forward: Option = + msg.params.and_then(|p| serde_json::from_value(p).ok()); + match forward { + Some(forward) => match ok_reply_text(id) { + Some(ack) => FrameAction::Serve { ack, forward }, + None => FrameAction::Ignore, + }, + None => error_reply( + id, + ErrorObject::new(codes::INVALID_PARAMS, "invalid cp/delegate params"), + ), + } + } + methods::CANCEL => { + let params: Option = + msg.params.and_then(|p| serde_json::from_value(p).ok()); + match params { + Some(params) => { + let known = self + .executor + .cancel(¶ms.delegation_id, params.admission); + info!( + delegation_id = %params.delegation_id, + reason_bytes = params.reason.len(), + known, + "cp/cancel received" + ); + // Acked either way: an unknown id means the delegation + // already finished here, which is not an error the CP + // can act on. + ok_reply(id) + } + None => error_reply( + id, + ErrorObject::new(codes::INVALID_PARAMS, "invalid cp/cancel params"), + ), + } + } + other => { + debug!(method = other, "unsupported control-plane method"); + error_reply( + id, + ErrorObject::new( + codes::METHOD_NOT_FOUND, + format!("runtime does not serve {other}"), + ), + ) + } + } + } +} + +/// What the serve loop should do with one inbound frame. +enum FrameAction { + /// Ack the request and start serving the delegation. + Serve { + ack: String, + forward: DelegateForward, + }, + /// Write this frame back. + Reply(String), + /// Nothing to say. + Ignore, +} + +fn delegate_result_request( + id: u64, + mut result: DelegateResultParams, +) -> anyhow::Result { + let build = |result: &DelegateResultParams| -> anyhow::Result { + Ok(JsonRpcRequest::new( + id, + methods::DELEGATE_RESULT, + Some(serde_json::to_value(result)?), + )) + }; + let frame = build(&result)?; + if serde_json::to_string(&frame)?.len() <= MAX_FRAME_BYTES { + return Ok(frame); + } + + let (payload, is_result) = if let Some(text) = result.result.take() { + (text, true) + } else if let Some(text) = result.error.take() { + (text, false) + } else { + anyhow::bail!("delegate_result envelope exceeds the transport limit without a payload"); + }; + if is_result { + result.result = Some(String::new()); + } else { + result.error = Some(String::new()); + } + let empty_frame = build(&result)?; + let envelope_bytes = serde_json::to_string(&empty_frame)?.len(); + anyhow::ensure!( + envelope_bytes < MAX_FRAME_BYTES, + "delegate_result envelope ({envelope_bytes} bytes) exceeds transport limit ({MAX_FRAME_BYTES} bytes)" + ); + let fitted = cap_text(payload, MAX_FRAME_BYTES - envelope_bytes); + if is_result { + result.result = Some(fitted); + } else { + result.error = Some(fitted); + } + let frame = build(&result)?; + anyhow::ensure!( + serde_json::to_string(&frame)?.len() <= MAX_FRAME_BYTES, + "delegate_result payload could not be fitted to the transport limit" + ); + Ok(frame) +} + +async fn send(sink: &mut WsSink, frame: &JsonRpcRequest) -> anyhow::Result<()> { + let text = serde_json::to_string(frame)?; + if text.len() > MAX_FRAME_BYTES { + anyhow::bail!( + "outbound frame ({} bytes) exceeds transport limit ({} bytes)", + text.len(), + MAX_FRAME_BYTES + ); + } + sink.send(Message::Text(text)).await?; + Ok(()) +} + +fn validate_register_ack(result: serde_json::Value) -> anyhow::Result { + let ack: RegisterAck = serde_json::from_value(result)?; + anyhow::ensure!( + ack.protocol_version == PROTOCOL_VERSION, + "control plane acknowledged protocol version {}, but this runtime requires {}", + ack.protocol_version, + PROTOCOL_VERSION + ); + Ok(ack) +} + +fn jittered_backoff(base_secs: u64) -> Duration { + let max_ms = base_secs.saturating_mul(1000).max(1); + Duration::from_millis(rand::thread_rng().gen_range(max_ms / 2..=max_ms)) +} + +fn next_backoff(current: u64) -> u64 { + current.saturating_mul(2).min(MAX_BACKOFF_SECS) +} + +fn backoff_after_session(current: u64, elapsed: Duration) -> u64 { + if elapsed >= Duration::from_secs(STABLE_SESSION_SECS) { + 1 + } else { + current + } +} + +fn ok_reply_text(id: u64) -> Option { + serde_json::to_string(&JsonRpcResponse::new(id, serde_json::json!({"ok": true}))).ok() +} + +fn ok_reply(id: u64) -> FrameAction { + match ok_reply_text(id) { + Some(text) => FrameAction::Reply(text), + None => FrameAction::Ignore, + } +} + +fn error_reply(id: u64, error: ErrorObject) -> FrameAction { + match serde_json::to_string(&JsonRpcErrorResponse::new(id, error)) { + Ok(text) => FrameAction::Reply(text), + Err(_) => FrameAction::Ignore, + } +} + +/// Why a connection's serve loop ended. +enum Outcome { + /// The process is shutting down; do not reconnect. + Shutdown, + /// The socket ended (close, error, or EOF); reconnect and re-register. + Disconnected, +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::config::CpAgentType; + use crate::control_plane::executor::PromptOutcome; + use async_trait::async_trait; + + struct NoopRunner; + + #[async_trait] + impl PromptRunner for NoopRunner { + async fn run( + &self, + _session_key: &str, + _forward: &DelegateForward, + ) -> anyhow::Result { + Ok(PromptOutcome::default()) + } + async fn cancel(&self, _session_key: &str) {} + async fn discard(&self, _session_key: &str) {} + } + + fn cfg() -> ControlPlaneConfig { + toml::from_str( + r#" +url = "ws://127.0.0.1:1/cp" +auth_key = "k" +namespace = "prod" +name = "worker-1" +type = "worker" +max_delegated_sessions = 3 +"#, + ) + .unwrap() + } + + fn client() -> Arc { + Arc::new(ControlPlaneClient::new( + cfg(), + Arc::new(NoopRunner), + Duration::from_secs(60), + )) + } + + #[test] + fn the_instance_id_is_a_uuid_and_is_stable_for_the_process() { + let c = client(); + assert_eq!(c.instance_id().len(), 36, "uuid v4, hyphenated"); + assert_eq!(c.instance_id(), c.instance_id()); + assert_ne!( + client().instance_id(), + c.instance_id(), + "a second process is a different replica" + ); + } + + #[test] + fn register_params_mirror_the_config_and_never_carry_the_key() { + let c = client(); + let params = RegisterParams { + protocol_version: PROTOCOL_VERSION, + namespace: c.cfg.namespace.clone(), + name: c.cfg.name.clone(), + agent_type: c.cfg.agent_type.into(), + instance_id: c.instance_id.clone(), + labels: c.cfg.labels.clone(), + max_delegated_sessions: c.cfg.max_delegated_sessions, + }; + let v = serde_json::to_value(¶ms).unwrap(); + assert_eq!(v["type"], "worker"); + assert_eq!(v["namespace"], "prod"); + assert_eq!(v["max_delegated_sessions"], 3); + assert_eq!(v["protocol_version"], PROTOCOL_VERSION); + let text = serde_json::to_string(&v).unwrap(); + assert!( + !text.contains("\"k\""), + "the auth key belongs in the header, never the frame: {text}" + ); + } + + #[test] + fn a_primary_config_registers_as_primary() { + let mut c = cfg(); + c.agent_type = CpAgentType::Primary; + let ty: openab_cp::proto::AgentType = c.agent_type.into(); + assert_eq!(ty, openab_cp::proto::AgentType::Primary); + } + + #[test] + fn a_primary_refuses_worker_side_delegate_frames() { + let mut primary_cfg = cfg(); + primary_cfg.agent_type = CpAgentType::Primary; + let primary = + ControlPlaneClient::new(primary_cfg, Arc::new(NoopRunner), Duration::from_secs(60)); + let frame = serde_json::json!({"jsonrpc":"2.0","id":9,"method":"cp/delegate","params":{ + "delegation_id":"d-1","admission":7,"prompt":"hi", + "deadline":(chrono::Utc::now()+chrono::Duration::seconds(60)).to_rfc3339(), + "from":"prod/koudu","chain":["prod/koudu"]}}) + .to_string(); + let FrameAction::Reply(reply) = primary.handle_frame(&frame) else { + panic!("primary served work"); + }; + let v: serde_json::Value = serde_json::from_str(&reply).unwrap(); + assert_eq!(v["error"]["code"], codes::POLICY_DENIED); + } + + #[test] + fn register_ack_requires_the_exact_protocol_version() { + let good = serde_json::json!({"protocol_version":PROTOCOL_VERSION,"heartbeat_interval_secs":15,"lease_expiry_secs":45,"effective_max_delegated_sessions":1}); + assert!(validate_register_ack(good).is_ok()); + let bad = serde_json::json!({"protocol_version":PROTOCOL_VERSION+1,"heartbeat_interval_secs":15,"lease_expiry_secs":45,"effective_max_delegated_sessions":1}); + assert!(validate_register_ack(bad) + .unwrap_err() + .to_string() + .contains("requires")); + } + + #[test] + fn reconnect_backoff_is_jittered_bounded_and_resets_only_after_stability() { + for base in [1, 2, 4, 8, 16, 30] { + for _ in 0..100 { + let delay = jittered_backoff(base); + assert!(delay >= Duration::from_millis(base * 500)); + assert!(delay <= Duration::from_secs(base)); + } + } + assert_eq!(next_backoff(16), 30); + assert_eq!(next_backoff(30), 30); + assert_eq!( + backoff_after_session(16, Duration::from_secs(STABLE_SESSION_SECS - 1)), + 16 + ); + assert_eq!( + backoff_after_session(16, Duration::from_secs(STABLE_SESSION_SECS)), + 1 + ); + } + + #[test] + fn delegate_result_cap_accounts_for_the_complete_envelope() { + let result = DelegateResultParams { + delegation_id: "d".repeat(600 * 1024), + admission: 1, + status: openab_cp::proto::DelegationStatus::Completed, + result: Some("x".repeat(512 * 1024)), + error: None, + }; + let frame = delegate_result_request(42, result).unwrap(); + let serialized = serde_json::to_string(&frame).unwrap(); + assert!(serialized.len() <= MAX_FRAME_BYTES); + assert!(serialized.contains("truncated by worker")); + } + + #[test] + fn rpc_ids_are_monotonic() { + let c = client(); + let a = c.next_id(); + let b = c.next_id(); + assert!(b > a); + } + + #[test] + fn a_delegate_frame_is_acked_and_yields_a_servable_forward() { + let c = client(); + let frame = serde_json::json!({ + "jsonrpc": "2.0", "id": 9, "method": "cp/delegate", + "params": { + "delegation_id": "d-1", + "admission": 7, + "prompt": "hi", + "deadline": (chrono::Utc::now() + chrono::Duration::seconds(60)).to_rfc3339(), + "from": "prod/koudu", + "chain": ["prod/koudu"] + } + }) + .to_string(); + match c.handle_frame(&frame) { + FrameAction::Serve { ack, forward } => { + let v: serde_json::Value = serde_json::from_str(&ack).unwrap(); + assert_eq!(v["id"], 9); + assert_eq!(v["result"]["ok"], true); + assert!( + v.get("error").is_none(), + "the ack says nothing about the outcome" + ); + assert_eq!(forward.delegation_id, "d-1"); + assert_eq!( + forward.admission, 7, + "the token must survive into the forward" + ); + assert_eq!(forward.from, "prod/koudu"); + assert_eq!(forward.chain, vec!["prod/koudu".to_string()]); + } + _ => panic!("cp/delegate must be served"), + } + } + + #[test] + fn malformed_delegate_params_are_rejected_without_serving() { + let c = client(); + // No deadline: the CP never sends this, but a malformed frame must not + // become an unbounded turn. + let frame = serde_json::json!({ + "jsonrpc": "2.0", "id": 3, "method": "cp/delegate", + "params": {"delegation_id": "d-1", "prompt": "hi", "from": "prod/koudu", "chain": []} + }) + .to_string(); + let FrameAction::Reply(reply) = c.handle_frame(&frame) else { + panic!("expected an error reply, not a served delegation"); + }; + let v: serde_json::Value = serde_json::from_str(&reply).unwrap(); + assert_eq!(v["error"]["code"], codes::INVALID_PARAMS); + } + + #[test] + fn cancel_is_acked_even_for_an_unknown_delegation() { + let c = client(); + let frame = serde_json::json!({ + "jsonrpc": "2.0", "id": 4, "method": "cp/cancel", + "params": {"delegation_id": "gone", "admission": 3, "reason": "initiator gave up"} + }) + .to_string(); + let FrameAction::Reply(reply) = c.handle_frame(&frame) else { + panic!("expected an ack"); + }; + assert!(reply.contains("\"ok\":true")); + } + + #[test] + fn an_unknown_method_gets_method_not_found() { + let c = client(); + let frame = serde_json::json!({ + "jsonrpc": "2.0", "id": 5, "method": "cp/event", "params": {} + }) + .to_string(); + let FrameAction::Reply(reply) = c.handle_frame(&frame) else { + panic!("expected an error reply"); + }; + let v: serde_json::Value = serde_json::from_str(&reply).unwrap(); + assert_eq!(v["error"]["code"], codes::METHOD_NOT_FOUND); + } + + #[test] + fn responses_to_our_own_requests_are_absorbed() { + let c = client(); + for frame in [ + r#"{"jsonrpc":"2.0","id":1,"result":{"ok":true}}"#, + r#"{"jsonrpc":"2.0","id":2,"error":{"code":-32004,"message":"no target"}}"#, + ] { + assert!( + matches!(c.handle_frame(frame), FrameAction::Ignore), + "a response is not answered" + ); + } + } + + #[test] + fn a_notification_shaped_request_is_refused() { + // `cp/*` methods are requests; an id-less one cannot be acked, and the + // CP's own parser enforces the same rule in the other direction. + let c = client(); + let FrameAction::Reply(reply) = + c.handle_frame(r#"{"jsonrpc":"2.0","method":"cp/delegate","params":{}}"#) + else { + panic!("expected an error reply"); + }; + let v: serde_json::Value = serde_json::from_str(&reply).unwrap(); + assert_eq!(v["error"]["code"], codes::INVALID_REQUEST); + } + + #[test] + fn garbage_is_dropped_not_answered() { + let c = client(); + assert!(matches!(c.handle_frame("{not json"), FrameAction::Ignore)); + } + + #[tokio::test(start_paused = true)] + async fn a_stalled_websocket_handshake_hits_the_connect_timeout() { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let server = tokio::spawn(async move { + let (_socket, _) = listener.accept().await.unwrap(); + std::future::pending::<()>().await; + }); + let mut stalled_cfg = cfg(); + stalled_cfg.url = format!("ws://{addr}/cp"); + let stalled = Arc::new(ControlPlaneClient::new( + stalled_cfg, + Arc::new(NoopRunner), + Duration::from_secs(60), + )); + let connecting = tokio::spawn(async move { stalled.connect_with_timeout().await }); + tokio::task::yield_now().await; + tokio::time::advance(CONNECT_TIMEOUT + Duration::from_secs(1)).await; + let error = match connecting.await.unwrap() { + Ok(_) => panic!("unexpected connection"), + Err(e) => e, + }; + assert!(error.to_string().contains("timed out")); + server.abort(); + } + + #[tokio::test] + async fn run_returns_immediately_when_shutdown_is_already_set() { + // The url points at a closed port: if the loop dialled before checking + // shutdown, this would hang for the whole backoff instead. + let (tx, rx) = watch::channel(true); + tokio::time::timeout(Duration::from_secs(1), client().run(rx)) + .await + .expect("shutdown is checked before dialling"); + drop(tx); + } +} diff --git a/crates/openab-core/src/control_plane/executor.rs b/crates/openab-core/src/control_plane/executor.rs new file mode 100644 index 000000000..96fb04939 --- /dev/null +++ b/crates/openab-core/src/control_plane/executor.rs @@ -0,0 +1,1237 @@ +//! Serving side of the control-plane client: turn one `cp/delegate` into one +//! `cp/delegate_result`. +//! +//! ## Invariants +//! +//! - **Admission never executes.** An over-cap or duplicate delegation is +//! answered with `status = failed` and an explanation, without touching the +//! session pool. The CP already fast-fails on its own accounting; this is +//! the runtime's own last word on its capacity, and it must be cheap. +//! - **One fresh session per delegation.** The session key is derived from +//! `(instance_id, delegation_id, admission)`, so no delegation can observe +//! another's conversation, and a re-admission of the same reusable id cannot +//! resume an earlier admission's session. The session is discarded on every +//! terminal outcome — nothing accumulates in the pool. +//! - **Exactly one result per admitted delegation.** Every path through +//! [`DelegationExecutor::serve`] returns a `DelegateResultParams`; the +//! client is what decides whether it can still be sent (on a dead socket it +//! cannot, and the CP synthesizes `target_disconnected` instead). + +use std::collections::BTreeMap; +use std::sync::atomic::{AtomicU32, Ordering}; +use std::sync::{Arc, Mutex, MutexGuard}; +use std::time::Duration; + +use anyhow::Result; +use async_trait::async_trait; +use openab_cp::proto::{AdmissionToken, DelegateForward, DelegateResultParams, DelegationStatus}; +use sha2::{Digest, Sha256}; +use tokio::sync::Notify; +use tracing::{info, warn}; + +/// Session-pool key for one ADMISSION of one delegation. +/// +/// Hashed rather than concatenated so an operator-visible key can never carry +/// a `delegation_id` chosen to collide with a chat thread key (they share one +/// namespace in the pool) and so its length is bounded regardless of what the +/// initiator sent. `instance_id` distinguishes replicas of the same logical +/// agent. The CP admission token is mixed in so a re-admission of the same +/// reusable id can never resume an earlier admission's session — in +/// particular one orphaned by a drain-timeout abort, whose transcript and +/// tool state would otherwise leak into the new run through the pool's +/// get-or-create semantics. +pub fn delegation_session_key( + instance_id: &str, + delegation_id: &str, + admission: AdmissionToken, +) -> String { + let mut hasher = Sha256::new(); + hasher.update(instance_id.as_bytes()); + hasher.update(delegation_id.as_bytes()); + hasher.update(admission.to_be_bytes()); + format!("control-plane:{:x}", hasher.finalize()) +} + +/// Outcome of one delegated prompt as reported by a [`PromptRunner`]. +#[derive(Debug, Clone, Default)] +pub struct PromptOutcome { + /// Reply text to hand back to the initiator. + pub text: String, + /// Agent/broker-level error that ended the turn. + pub error: Option, + /// The turn produced nothing and reported zero output tokens — a + /// provider/model/auth failure that must not be reported as success. + pub silent_failure: bool, +} + +/// The prompt-execution seam. +/// +/// Production is [`RouterPromptRunner`] (ACP session pool + `AdapterRouter`). +/// It exists as a trait so the client's state machine can be exercised against +/// the real CP server without a coding agent on the box: the integration test +/// injects a runner that answers from a script. +#[async_trait] +pub trait PromptRunner: Send + Sync + 'static { + /// Run the delegated prompt to completion in `session_key`. + async fn run(&self, session_key: &str, forward: &DelegateForward) -> Result; + + /// Best-effort interrupt of the in-flight turn for `session_key`. + async fn cancel(&self, session_key: &str); + + /// Drop `session_key` and its bookkeeping. + async fn discard(&self, session_key: &str); +} + +/// Local admission + execution of delegations for one runtime instance. +pub struct DelegationExecutor { + runner: Arc, + instance_id: String, + /// Ceiling from the registration ack (the CP may clamp what we advertised). + /// Updated on every re-register, hence atomic rather than a constructor arg. + effective_max: AtomicU32, + /// Per-turn hard ceiling, from `[pool].prompt_hard_timeout_secs`. The + /// delegation deadline is the other clock; the shorter one wins. + prompt_hard_timeout: Duration, + /// Admitted delegations → their cancel signal. Also the capacity counter: + /// its length is the number of active delegated sessions reported in + /// `cp/heartbeat`. + inflight: Mutex)>>, +} + +/// Why admission refused, as the message sent back in `status = failed`. +#[derive(Debug)] +enum Refusal { + OverCapacity { active: u32, max: u32 }, + Duplicate, +} + +impl Refusal { + fn message(&self) -> String { + match self { + Refusal::OverCapacity { active, max } => format!( + "runtime is at its local delegation capacity ({active}/{max}); \ + the delegation was not started" + ), + Refusal::Duplicate => "delegation_id is already in flight on this runtime; \ + the delegation was not started" + .to_string(), + } + } +} + +impl DelegationExecutor { + pub fn new( + runner: Arc, + instance_id: impl Into, + effective_max: u32, + prompt_hard_timeout: Duration, + ) -> Self { + Self { + runner, + instance_id: instance_id.into(), + effective_max: AtomicU32::new(effective_max), + prompt_hard_timeout, + inflight: Mutex::new(BTreeMap::new()), + } + } + + /// Adopt the budget the CP acked. Called on every (re-)registration. + pub fn set_effective_max(&self, max: u32) { + self.effective_max.store(max, Ordering::Relaxed); + } + + pub fn effective_max(&self) -> u32 { + self.effective_max.load(Ordering::Relaxed) + } + + fn lock_inflight(&self) -> MutexGuard<'_, BTreeMap)>> { + self.inflight.lock().unwrap_or_else(|poisoned| { + warn!("recovering poisoned control-plane inflight mutex"); + poisoned.into_inner() + }) + } + + /// Number of admitted, not-yet-finished delegations. + pub fn active(&self) -> u32 { + self.lock_inflight().len() as u32 + } + + /// Reserve a slot for `delegation_id`, or explain why not. + fn admit( + &self, + delegation_id: &str, + admission: AdmissionToken, + ) -> std::result::Result, Refusal> { + let max = self.effective_max(); + let mut g = self.lock_inflight(); + if g.contains_key(delegation_id) { + return Err(Refusal::Duplicate); + } + let active = g.len() as u32; + if active >= max { + return Err(Refusal::OverCapacity { active, max }); + } + let signal = Arc::new(Notify::new()); + g.insert(delegation_id.to_string(), (admission, Arc::clone(&signal))); + Ok(signal) + } + + fn release(&self, delegation_id: &str) { + self.lock_inflight().remove(delegation_id); + } + + /// Signal cancellation for one delegation (`cp/cancel` from the CP). + /// Returns `false` when the id is not in flight here — the CP's view can + /// legitimately be ahead of ours (it also cancels on deadline). + /// + /// `notify_one` rather than `notify_waiters`: it leaves a permit behind, so + /// a cancel that arrives between admission and the first poll of the + /// serving task is still observed instead of being lost. + pub fn cancel(&self, delegation_id: &str, admission: AdmissionToken) -> bool { + let signal = { + let g = self.lock_inflight(); + match g.get(delegation_id) { + // The token names ONE admission of this reusable id. A stale + // cancel — the CP swept admission A, this worker was already + // re-serving B under the same id — must not abort B: that is + // the worker-side half of the misdelivery the wire token + // exists to close. + Some((adm, s)) if *adm == admission => Some(Arc::clone(s)), + Some((adm, _)) => { + tracing::info!( + delegation_id, + live = *adm, + stale = admission, + "cp/cancel names a superseded admission — ignoring" + ); + None + } + None => None, + } + }; + match signal { + Some(s) => { + s.notify_one(); + true + } + None => false, + } + } + + /// Signal cancellation for every in-flight delegation: connection loss and + /// shutdown. Each task cleans its session up and returns a `Cancelled` + /// result the caller is free to drop — on a dead socket the CP synthesizes + /// `target_disconnected` for the initiator, so sending ours is neither + /// possible nor needed. + pub fn cancel_all(&self) { + let signals: Vec> = self + .lock_inflight() + .values() + .map(|(_, s)| Arc::clone(s)) + .collect(); + for s in signals { + s.notify_one(); + } + } + + /// Admit, run, and classify one forwarded delegation. + /// + /// Always resolves to a result frame payload — the refusal paths included, + /// so the initiator is never left waiting on its deadline for a runtime + /// that had already decided not to run. + pub async fn serve(self: Arc, forward: DelegateForward) -> DelegateResultParams { + let id = forward.delegation_id.clone(); + let cancel = match self.admit(&id, forward.admission) { + Ok(signal) => signal, + Err(refusal) => { + let error = refusal.message(); + warn!(delegation_id = %id, from = %forward.from, %error, "delegation refused"); + return failed(&id, forward.admission, error); + } + }; + // RAII: the slot must free even if this task is ABORTED mid-await — + // the client aborts serving tasks that outlive the drain window on + // disconnect, and a plain post-await release would be skipped there, + // leaking the inflight entry forever (with the default cap of 1, the + // worker would refuse every delegation after reconnecting). + let _slot = SlotGuard { + executor: self.as_ref(), + id: id.clone(), + }; + self.execute(&forward, cancel).await + } + + async fn execute( + &self, + forward: &DelegateForward, + cancel: Arc, + ) -> DelegateResultParams { + let id = &forward.delegation_id; + let session_key = delegation_session_key(&self.instance_id, id, forward.admission); + + // Two clocks bound the turn: the CP-enforced delegation deadline and + // the runtime's own per-turn ceiling. Take the nearer one — an already + // elapsed deadline means there is nothing worth starting. + let Ok(remaining) = (forward.deadline - chrono::Utc::now()).to_std() else { + warn!(delegation_id = %id, deadline = %forward.deadline, "delegation arrived past its deadline"); + return timed_out(id, forward.admission); + }; + let budget = remaining.min(self.prompt_hard_timeout); + info!( + delegation_id = %id, + from = %forward.from, + chain_depth = forward.chain.len(), + budget_secs = budget.as_secs(), + "serving delegation" + ); + + // `notified()` is created BEFORE the run so a cancel racing the first + // poll is not missed: `cancel` leaves a permit (`notify_one`), and this + // future consumes it whenever it is first polled. + let cancelled = cancel.notified(); + let outcome = tokio::select! { + biased; + _ = cancelled => { + info!(delegation_id = %id, "delegation cancelled"); + self.cancel_and_discard(&session_key).await; + return DelegateResultParams { + delegation_id: id.clone(), + admission: forward.admission, + status: DelegationStatus::Cancelled, + result: None, + error: None, + }; + } + run = tokio::time::timeout(budget, self.runner.run(&session_key, forward)) => run, + }; + + match outcome { + Err(_elapsed) => { + warn!(delegation_id = %id, budget_secs = budget.as_secs(), "delegation exceeded its local deadline"); + self.cancel_and_discard(&session_key).await; + timed_out(id, forward.admission) + } + Ok(Err(e)) => { + // The turn could not be driven at all (no session, dead agent). + let error = format!("{e:#}"); + warn!(delegation_id = %id, %error, "delegation failed before completion"); + self.bounded_discard(&session_key).await; + failed(id, forward.admission, error) + } + Ok(Ok(outcome)) => { + if let Some(error) = outcome.error { + self.bounded_discard(&session_key).await; + warn!(delegation_id = %id, %error, "delegation ended in an agent error"); + return failed(id, forward.admission, error); + } + if outcome.silent_failure { + self.bounded_discard(&session_key).await; + warn!(delegation_id = %id, "delegation produced an empty turn (silent failure)"); + return failed( + id, + forward.admission, + "agent returned an empty turn (0 output tokens) — \ + likely a provider/model/auth failure", + ); + } + // Turn completed cleanly. Spawn discard off the critical path + // so a slow or wedged discard does not delay result delivery to + // the initiator (F53). With admission-scoped session keys, a + // deferred discard is safe: a subsequent admission of the same + // delegation id gets a distinct session key. + self.spawn_discard(session_key); + info!(delegation_id = %id, bytes = outcome.text.len(), "delegation completed"); + DelegateResultParams { + delegation_id: id.clone(), + admission: forward.admission, + status: DelegationStatus::Completed, + result: Some(cap_result(outcome.text)), + error: None, + } + } + } + } + + /// Best-effort, BOUNDED session teardown: `session/cancel` writes to the + /// agent's stdin, which can wedge (dead child, full pipe), and the pool's + /// discard takes its write lock, which can be starved. Both share ONE + /// deadline so a wedged teardown cannot burn the client's disconnect drain window (F52). + async fn cancel_and_discard(&self, session_key: &str) { + let deadline = tokio::time::Instant::now() + TEARDOWN_BOUND; + let _ = tokio::time::timeout_at(deadline, self.runner.cancel(session_key)).await; + let remaining = deadline.saturating_duration_since(tokio::time::Instant::now()); + if remaining.is_zero() { + warn!( + session_key, + "session cancel consumed the teardown bound; leaving discard to pool cleanup" + ); + return; + } + if tokio::time::timeout_at(deadline, self.runner.discard(session_key)) + .await + .is_err() + { + warn!( + session_key, + "session discard exceeded its bound; leaving it to pool cleanup" + ); + } + } + + /// Discard with the same bound as cancel; on overrun the session is left + /// to the pool's own idle/hung cleanup rather than blocking this task. + async fn bounded_discard(&self, session_key: &str) { + if tokio::time::timeout(TEARDOWN_BOUND, self.runner.discard(session_key)) + .await + .is_err() + { + warn!( + session_key, + "session discard exceeded its bound; leaving it to pool cleanup" + ); + } + } + + /// Spawn discard off the critical path on success (F53). + fn spawn_discard(&self, session_key: String) { + let runner = Arc::clone(&self.runner); + tokio::spawn(async move { + if tokio::time::timeout(TEARDOWN_BOUND, runner.discard(&session_key)) + .await + .is_err() + { + warn!( + session_key = %session_key, + "session discard exceeded its bound; leaving it to pool cleanup" + ); + } + }); + } +} + +/// Bound on each session-teardown step (cancel, discard). Matches the pool's +/// own cleanup bound and stays under the client's disconnect drain window. +const TEARDOWN_BOUND: std::time::Duration = std::time::Duration::from_secs(4); + +/// Client-side ceiling on a `cp/delegate_result` body. +/// +/// The CP enforces `max_frame_bytes` (default 1 MiB) at the WS transport, +/// BEFORE parsing — its own `max_result_bytes` truncation therefore can never +/// save an oversized frame: the transport drops the connection, and every +/// in-flight delegation on this worker dies as `target_disconnected`. Capping +/// here keeps the frame safely under the transport limit; the CP still applies +/// its (typically smaller) `max_result_bytes` on what arrives. +const MAX_RESULT_BYTES: usize = 512 * 1024; + +/// Ceiling on a `cp/delegate_result` error string. Errors ride the same +/// transport frame as results but are diagnostics, not payloads, so the +/// budget is far tighter. Without this, an unbounded `anyhow` chain or an +/// agent-authored error would hit the CP's pre-parse `max_frame_bytes` and +/// drop the connection — the exact failure `MAX_RESULT_BYTES` closes for the +/// success path. +const MAX_ERROR_BYTES: usize = 64 * 1024; + +/// Returns the byte length of `s` when serialized inside a JSON string literal. +/// +/// In JSON (RFC 8259), `"` and `\` become 2 bytes (`\"`, `\\`), control +/// characters 0x00..=0x1f become 2 bytes (`\b`, `\t`, `\n`, `\f`, `\r`) or +/// 6 bytes (`\u00xx`), and all other UTF-8 bytes pass through unchanged (1 byte +/// each). +pub(crate) fn json_escaped_len(s: &str) -> usize { + s.bytes() + .map(|b| match b { + b'\"' | b'\\' => 2, + b'\x08' | b'\t' | b'\n' | b'\x0c' | b'\r' => 2, + c if c < 0x20 => 6, + _ => 1, + }) + .sum() +} + +pub(crate) fn cap_text(text: String, budget: usize) -> String { + if json_escaped_len(&text) <= budget { + return text; + } + let marker = format!( + "\n...[truncated by worker: {} bytes total exceeded the transport budget]", + text.len() + ); + let marker_escaped_len = json_escaped_len(&marker); + let keep_escaped = budget.saturating_sub(marker_escaped_len); + let mut cut = 0; + let mut current_escaped = 0; + for (idx, ch) in text.char_indices() { + let ch_escaped = match ch { + '\"' | '\\' => 2, + '\x08' | '\t' | '\n' | '\x0c' | '\r' => 2, + c if (c as u32) < 0x20 => 6, + c => c.len_utf8(), + }; + if current_escaped + ch_escaped > keep_escaped { + break; + } + current_escaped += ch_escaped; + cut = idx + ch.len_utf8(); + } + let mut out = String::with_capacity(cut + marker.len()); + out.push_str(&text[..cut]); + out.push_str(&marker); + out +} + +fn cap_result(text: String) -> String { + cap_text(text, MAX_RESULT_BYTES) +} + +/// Frees a delegation's inflight slot on drop — including the drop that +/// happens when the serving task is aborted at an await point. +struct SlotGuard<'a> { + executor: &'a DelegationExecutor, + id: String, +} + +impl Drop for SlotGuard<'_> { + fn drop(&mut self) { + self.executor.release(&self.id); + } +} + +fn failed( + delegation_id: &str, + admission: AdmissionToken, + error: impl Into, +) -> DelegateResultParams { + DelegateResultParams { + delegation_id: delegation_id.to_string(), + // Echoed verbatim from the forward: the CP correlates terminal frames + // per admission, so a late result for a superseded admission of this + // id is dropped instead of completing the wrong delegation. + admission, + status: DelegationStatus::Failed, + result: None, + // Capped HERE, not at call sites: every error source (anyhow chains, + // agent-authored errors) must share the transport-safe bound, and a + // new call site must not be able to forget it. + error: Some(cap_text(error.into(), MAX_ERROR_BYTES)), + } +} + +fn timed_out(delegation_id: &str, admission: AdmissionToken) -> DelegateResultParams { + DelegateResultParams { + delegation_id: delegation_id.to_string(), + admission, + status: DelegationStatus::Timeout, + result: None, + error: Some("delegation deadline elapsed at the serving runtime".into()), + } +} + +// --------------------------------------------------------------------------- +// Production runner: ACP session pool + AdapterRouter +// --------------------------------------------------------------------------- + +use crate::adapter::{AdapterRouter, ChannelRef, ChatAdapter, MessageRef}; +use crate::reactions::StatusReactionController; + +/// Platform label for delegated turns. Not a chat platform: it exists so +/// session keys, logs, and the router's platform switches can tell a +/// delegation apart from a user conversation. +pub const CP_PLATFORM: &str = "control-plane"; + +/// [`PromptRunner`] over the real ACP session pool. +/// +/// Reuses `AdapterRouter::stream_prompt_blocks` — the same turn driver every +/// chat platform uses, so tool events, liveness checks, the hard timeout, and +/// silent-failure classification behave identically here — with a sink adapter +/// standing in for the platform. +pub struct RouterPromptRunner { + router: Arc, +} + +impl RouterPromptRunner { + pub fn new(router: Arc) -> Self { + Self { router } + } +} + +#[async_trait] +impl PromptRunner for RouterPromptRunner { + async fn run(&self, session_key: &str, forward: &DelegateForward) -> Result { + // A delegation is always a fresh session, so this creates one rather + // than resuming; `working_dir` stays the configured default. + self.router.pool().get_or_create(session_key, None).await?; + + let adapter: Arc = Arc::new(SinkAdapter); + let channel = ChannelRef { + platform: CP_PLATFORM.to_string(), + channel_id: forward.delegation_id.clone(), + thread_id: None, + parent_id: None, + origin_event_id: None, + }; + // Reactions are constructed disabled: there is no message to react to. + let reactions = Arc::new(StatusReactionController::new( + false, + Arc::clone(&adapter), + MessageRef { + channel: channel.clone(), + message_id: String::new(), + }, + crate::config::ReactionEmojis::default(), + crate::config::ReactionTiming::default(), + )); + + let blocks = AdapterRouter::pack_arrival_event( + &delegation_context_json(forward), + &forward.prompt, + Vec::new(), + ); + let execution = self + .router + .stream_prompt_blocks( + &adapter, + session_key, + blocks, + &channel, + reactions, + false, // other_bot_present: no channel, no other bots + None, // no native-streaming recipient + ) + .await?; + + Ok(PromptOutcome { + text: execution.final_text, + error: execution.terminal_error, + silent_failure: execution.silent_failure, + }) + } + + async fn cancel(&self, session_key: &str) { + if let Err(e) = self.router.pool().cancel_session(session_key).await { + // Nothing in flight to cancel is the common benign case. + tracing::debug!(error = %e, "cancel_session on a delegated session"); + } + } + + async fn discard(&self, session_key: &str) { + self.router + .pool() + .discard_single_use_session(session_key) + .await; + } +} + +/// The arrival metadata block for a delegated turn. +/// +/// Carried inside the same `` envelope every platform arrival +/// uses (so agents keep one place to look for provenance) but with its own +/// schema: the fields that matter here are the CP-authenticated ones — who +/// asked, through which ancestry, and by when — and `chain`/`deadline` have no +/// counterpart in `openab.sender.v1`. Every value is stamped by the CP, so the +/// agent may trust it. +fn delegation_context_json(forward: &DelegateForward) -> String { + serde_json::json!({ + "schema": "openab.delegation.v1", + "delegation_id": forward.delegation_id, + "from": forward.from, + "chain": forward.chain, + "deadline": forward.deadline.to_rfc3339(), + }) + .to_string() +} + +/// A `ChatAdapter` that delivers nowhere. +/// +/// A delegation's reply travels back over the control-plane socket as +/// `cp/delegate_result`, not to a channel, so every write is dropped and the +/// text is read from the returned `PromptExecution` instead. Forcing +/// send-once (`use_streaming = false`) is what makes that safe: no +/// placeholder is posted, no edit loop is spawned, and the full turn text is +/// composed exactly once at the end. +struct SinkAdapter; + +#[async_trait] +impl ChatAdapter for SinkAdapter { + fn platform(&self) -> &'static str { + CP_PLATFORM + } + + /// No chunking: the delegation result is one payload, and the CP applies + /// its own `max_result_bytes` cap. + fn message_limit(&self) -> usize { + usize::MAX + } + + fn use_streaming(&self, _other_bot_present: bool) -> bool { + false + } + + async fn send_message(&self, channel: &ChannelRef, _content: &str) -> Result { + Ok(MessageRef { + channel: channel.clone(), + message_id: String::new(), + }) + } + + async fn create_thread( + &self, + channel: &ChannelRef, + _trigger_msg: &MessageRef, + _title: &str, + ) -> Result { + Ok(channel.clone()) + } + + async fn add_reaction(&self, _msg: &MessageRef, _emoji: &str) -> Result<()> { + Ok(()) + } + + async fn remove_reaction(&self, _msg: &MessageRef, _emoji: &str) -> Result<()> { + Ok(()) + } + + async fn edit_message(&self, _msg: &MessageRef, _content: &str) -> Result<()> { + Ok(()) + } + + async fn delete_message(&self, _msg: &MessageRef) -> Result<()> { + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::atomic::AtomicUsize; + + fn forward(id: &str, secs: i64) -> DelegateForward { + DelegateForward { + delegation_id: id.into(), + admission: 1, + prompt: "do the thing".into(), + deadline: chrono::Utc::now() + chrono::Duration::seconds(secs), + from: "prod/koudu".into(), + chain: vec!["prod/koudu".into()], + } + } + + /// Scripted runner: records lifecycle calls and answers as configured. + #[derive(Default)] + struct FakeRunner { + /// Reply text on success. + text: String, + /// If set, `run` fails with this message. + run_error: Option, + /// If set, the outcome carries this agent error. + agent_error: Option, + silent_failure: bool, + /// If set, `run` sleeps this long before answering. + delay: Option, + /// If set, `cancel` sleeps this long before answering. + cancel_delay: Option, + /// If set, `discard` sleeps this long before answering. + discard_delay: Option, + started: AtomicUsize, + cancelled: Mutex>, + discarded: Mutex>, + } + + impl FakeRunner { + fn completing(text: &str) -> Arc { + Arc::new(Self { + text: text.into(), + ..Default::default() + }) + } + fn starts(&self) -> usize { + self.started.load(Ordering::Relaxed) + } + fn discarded(&self) -> Vec { + self.discarded.lock().unwrap().clone() + } + async fn wait_discarded(&self, expected_len: usize) -> Vec { + for _ in 0..100 { + let d = self.discarded(); + if d.len() >= expected_len { + return d; + } + tokio::task::yield_now().await; + } + self.discarded() + } + fn cancelled(&self) -> Vec { + self.cancelled.lock().unwrap().clone() + } + } + + #[async_trait] + impl PromptRunner for FakeRunner { + async fn run( + &self, + _session_key: &str, + _forward: &DelegateForward, + ) -> Result { + self.started.fetch_add(1, Ordering::Relaxed); + if let Some(d) = self.delay { + tokio::time::sleep(d).await; + } + if let Some(ref e) = self.run_error { + return Err(anyhow::anyhow!(e.clone())); + } + Ok(PromptOutcome { + text: self.text.clone(), + error: self.agent_error.clone(), + silent_failure: self.silent_failure, + }) + } + + async fn cancel(&self, session_key: &str) { + if let Some(d) = self.cancel_delay { + tokio::time::sleep(d).await; + } + self.cancelled.lock().unwrap().push(session_key.to_string()); + } + + async fn discard(&self, session_key: &str) { + if let Some(d) = self.discard_delay { + tokio::time::sleep(d).await; + } + self.discarded.lock().unwrap().push(session_key.to_string()); + } + } + + fn executor(runner: Arc, max: u32) -> Arc { + Arc::new(DelegationExecutor::new( + runner, + "i-test", + max, + Duration::from_secs(600), + )) + } + + #[tokio::test] + async fn success_maps_to_completed_and_discards_the_session() { + let runner = FakeRunner::completing("here you go"); + let ex = executor(Arc::clone(&runner), 1); + let res = Arc::clone(&ex).serve(forward("d-1", 60)).await; + assert_eq!(res.status, DelegationStatus::Completed); + assert_eq!(res.result.as_deref(), Some("here you go")); + assert!(res.error.is_none()); + assert_eq!( + runner.wait_discarded(1).await, + vec![delegation_session_key("i-test", "d-1", 1)], + "a fresh-per-delegation session must not survive its delegation" + ); + assert_eq!(ex.active(), 0, "the slot is released"); + } + + #[tokio::test] + async fn over_capacity_is_refused_without_executing() { + let runner = FakeRunner::completing("ok"); + let ex = executor(Arc::clone(&runner), 1); + // Occupy the only slot: `admit` is what reserves capacity, so the + // entry stands until the (never-spawned) serving task releases it. + let _held = ex.admit("d-held", 1).expect("first slot"); + let res = Arc::clone(&ex).serve(forward("d-2", 60)).await; + assert_eq!(res.status, DelegationStatus::Failed); + assert!(res.error.unwrap().contains("local delegation capacity")); + assert_eq!(runner.starts(), 0, "a refused delegation never runs"); + assert!(runner.discarded().is_empty(), "and touches no session"); + } + + #[tokio::test] + async fn duplicate_delegation_id_is_refused_without_executing() { + let runner = FakeRunner::completing("ok"); + let ex = executor(Arc::clone(&runner), 4); + let _held = ex.admit("d-3", 1).expect("slot"); + let res = Arc::clone(&ex).serve(forward("d-3", 60)).await; + assert_eq!(res.status, DelegationStatus::Failed); + assert!(res.error.unwrap().contains("already in flight")); + assert_eq!(runner.starts(), 0); + } + + #[tokio::test] + async fn effective_max_from_the_ack_is_what_bounds_admission() { + let runner = FakeRunner::completing("ok"); + let ex = executor(Arc::clone(&runner), 4); + ex.set_effective_max(1); // CP clamped us + let _held = ex.admit("d-a", 1).expect("slot"); + let res = Arc::clone(&ex).serve(forward("d-b", 60)).await; + assert_eq!(res.status, DelegationStatus::Failed); + assert!( + res.error.unwrap().contains("(1/1)"), + "the clamped ceiling, not the advertised one, is enforced" + ); + assert_eq!(runner.starts(), 0); + } + + #[test] + fn a_poisoned_inflight_mutex_is_recovered_even_on_the_drop_path() { + let ex = executor(FakeRunner::completing("ok"), 1); + let poisoned = std::panic::catch_unwind(std::panic::AssertUnwindSafe({ + let ex = Arc::clone(&ex); + move || { + let _guard = ex.inflight.lock().unwrap(); + panic!("poison map"); + } + })); + assert!(poisoned.is_err()); + assert_eq!(ex.active(), 0); + { + let _slot = SlotGuard { + executor: ex.as_ref(), + id: "missing".into(), + }; + } + let _admitted = ex.admit("after-poison", 1).expect("recovered admission"); + } + + #[tokio::test] + async fn agent_error_maps_to_failed() { + let runner = Arc::new(FakeRunner { + agent_error: Some("provider returned HTTP 500".into()), + ..Default::default() + }); + let ex = executor(Arc::clone(&runner), 1); + let res = Arc::clone(&ex).serve(forward("d-4", 60)).await; + assert_eq!(res.status, DelegationStatus::Failed); + assert_eq!(res.error.as_deref(), Some("provider returned HTTP 500")); + assert_eq!(runner.discarded().len(), 1); + } + + #[tokio::test] + async fn silent_failure_maps_to_failed_not_completed() { + let runner = Arc::new(FakeRunner { + // A rendered tool summary must not turn a raw silent failure into + // success; classification is carried separately from display text. + text: "3 tools completed".into(), + silent_failure: true, + ..Default::default() + }); + let ex = executor(Arc::clone(&runner), 1); + let res = Arc::clone(&ex).serve(forward("d-5", 60)).await; + assert_eq!(res.status, DelegationStatus::Failed); + assert!(res.error.unwrap().contains("empty turn")); + } + + #[tokio::test] + async fn broker_error_maps_to_failed() { + let runner = Arc::new(FakeRunner { + run_error: Some("no connection for session".into()), + ..Default::default() + }); + let ex = executor(Arc::clone(&runner), 1); + let res = Arc::clone(&ex).serve(forward("d-6", 60)).await; + assert_eq!(res.status, DelegationStatus::Failed); + assert!(res.error.unwrap().contains("no connection")); + assert_eq!( + runner.discarded().len(), + 1, + "the session is still cleaned up" + ); + } + + #[tokio::test] + async fn a_past_deadline_times_out_without_executing() { + let runner = FakeRunner::completing("ok"); + let ex = executor(Arc::clone(&runner), 1); + let res = Arc::clone(&ex).serve(forward("d-7", -1)).await; + assert_eq!(res.status, DelegationStatus::Timeout); + assert_eq!(runner.starts(), 0); + } + + #[tokio::test(start_paused = true)] + async fn the_local_deadline_times_out_and_cleans_the_session_up() { + let runner = Arc::new(FakeRunner { + delay: Some(Duration::from_secs(300)), + ..Default::default() + }); + let ex = executor(Arc::clone(&runner), 1); + let res = Arc::clone(&ex).serve(forward("d-8", 5)).await; + assert_eq!(res.status, DelegationStatus::Timeout); + assert_eq!(runner.starts(), 1, "it did start"); + let key = delegation_session_key("i-test", "d-8", 1); + assert_eq!(runner.cancelled(), vec![key.clone()]); + assert_eq!(runner.discarded(), vec![key]); + assert_eq!(ex.active(), 0); + } + + #[tokio::test(start_paused = true)] + async fn cancel_mid_flight_maps_to_cancelled_and_drops_the_session() { + let runner = Arc::new(FakeRunner { + delay: Some(Duration::from_secs(300)), + ..Default::default() + }); + let ex = executor(Arc::clone(&runner), 1); + let serving = tokio::spawn({ + let ex = Arc::clone(&ex); + async move { ex.serve(forward("d-9", 600)).await } + }); + // Let the task admit and start before cancelling. + while ex.active() == 0 { + tokio::task::yield_now().await; + } + assert!(ex.cancel("d-9", 1), "the id is in flight"); + let res = serving.await.unwrap(); + assert_eq!(res.status, DelegationStatus::Cancelled); + assert!(res.result.is_none()); + let key = delegation_session_key("i-test", "d-9", 1); + assert_eq!(runner.cancelled(), vec![key.clone()]); + assert_eq!(runner.discarded(), vec![key]); + assert_eq!(ex.active(), 0); + } + + #[tokio::test(start_paused = true)] + async fn cancel_all_ends_every_in_flight_delegation() { + let runner = Arc::new(FakeRunner { + delay: Some(Duration::from_secs(300)), + ..Default::default() + }); + let ex = executor(Arc::clone(&runner), 4); + let mut tasks = Vec::new(); + for id in ["d-x", "d-y"] { + let ex = Arc::clone(&ex); + tasks.push(tokio::spawn( + async move { ex.serve(forward(id, 600)).await }, + )); + } + while ex.active() < 2 { + tokio::task::yield_now().await; + } + ex.cancel_all(); + for t in tasks { + assert_eq!(t.await.unwrap().status, DelegationStatus::Cancelled); + } + assert_eq!(ex.active(), 0); + assert_eq!(runner.discarded().len(), 2); + } + + #[test] + fn cancel_of_an_unknown_id_is_a_no_op() { + let ex = executor(FakeRunner::completing("ok"), 1); + assert!(!ex.cancel("never-seen", 1)); + } + + #[test] + fn session_keys_are_namespaced_bounded_and_instance_scoped() { + let a = delegation_session_key("i-1", "d-1", 1); + let b = delegation_session_key("i-2", "d-1", 1); + assert!(a.starts_with("control-plane:")); + assert_ne!(a, b, "another replica's session never collides"); + assert_eq!(a.len(), "control-plane:".len() + 64); + // A re-admission of the same reusable id gets a fresh session: an + // orphaned session from an aborted admission can never be resumed. + let c = delegation_session_key("i-1", "d-1", 2); + assert_ne!(a, c, "a re-admission of the same id gets a fresh session"); + // A hostile id cannot forge another platform's key shape. + let hostile = delegation_session_key("i-1", "discord:12345", 1); + assert!(hostile.starts_with("control-plane:")); + assert_eq!(hostile.len(), a.len()); + } + + #[test] + fn the_delegation_context_block_carries_the_cp_stamped_provenance() { + let v: serde_json::Value = + serde_json::from_str(&delegation_context_json(&forward("d-10", 30))).unwrap(); + assert_eq!(v["schema"], "openab.delegation.v1"); + assert_eq!(v["delegation_id"], "d-10"); + assert_eq!(v["from"], "prod/koudu"); + assert_eq!(v["chain"][0], "prod/koudu"); + assert!(v["deadline"].as_str().unwrap().contains('T')); + } + + #[test] + fn the_sink_adapter_never_streams() { + // Streaming would post a placeholder to a channel that does not exist + // and split the reply the executor has to return whole. + assert!(!SinkAdapter.use_streaming(false)); + assert!(!SinkAdapter.use_streaming(true)); + assert!(!SinkAdapter.uses_native_streaming(false)); + assert!(!SinkAdapter.uses_assistant_status()); + } + + #[test] + fn oversized_results_are_capped_below_the_transport_limit() { + // An uncapped result larger than the CP's max_frame_bytes (1 MiB) + // would be dropped at the WS transport before the CP's own + // max_result_bytes truncation could run, killing the connection and + // every in-flight delegation with it. + let big = "x".repeat(2 * 1024 * 1024); + let capped = cap_result(big); + assert!(capped.len() <= MAX_RESULT_BYTES); + assert!(capped.ends_with("bytes total exceeded the transport budget]")); + + // Multibyte char straddling the cut must not split a boundary. + let emoji = "\u{1F980}".repeat(MAX_RESULT_BYTES / 4 + 64); + let capped = cap_result(emoji); + assert!(capped.len() <= MAX_RESULT_BYTES); + assert!(std::str::from_utf8(capped.as_bytes()).is_ok()); + + // Under the cap: untouched. + assert_eq!(cap_result("small".into()), "small"); + } + + #[test] + fn oversized_errors_are_capped_below_the_transport_limit() { + // The error field rides the same frame as the result and hits the + // same pre-parse max_frame_bytes ceiling at the CP. failed() must + // bound every error source (anyhow chains, agent-authored errors), + // no matter the call site. + let big = "e".repeat(2 * 1024 * 1024); + let res = failed("d-err", 7, big); + let err = res.error.expect("failed() always carries an error"); + assert!(err.len() <= MAX_ERROR_BYTES); + assert!(err.ends_with("bytes total exceeded the transport budget]")); + assert_eq!(res.admission, 7, "the admission echo survives the cap"); + assert_eq!(res.status, DelegationStatus::Failed); + + // Multibyte char straddling the cut must not split a boundary. + let emoji = "\u{1F980}".repeat(MAX_ERROR_BYTES / 4 + 64); + let res = failed("d-err", 7, emoji); + let err = res.error.expect("failed() always carries an error"); + assert!(err.len() <= MAX_ERROR_BYTES); + assert!(std::str::from_utf8(err.as_bytes()).is_ok()); + + // Under the cap: untouched. + let res = failed("d-err", 7, "short"); + assert_eq!(res.error.as_deref(), Some("short")); + } + + #[tokio::test(start_paused = true)] + async fn an_aborted_serving_task_still_frees_its_slot() { + // The client aborts serving tasks that outlive the drain window on + // disconnect. A plain post-await release would be skipped by the + // abort, leaking the inflight entry: with the default cap of 1 the + // worker would then refuse every delegation after reconnecting. + let runner = Arc::new(FakeRunner { + delay: Some(Duration::from_secs(300)), + ..Default::default() + }); + let ex = executor(Arc::clone(&runner), 1); + let task = { + let ex = Arc::clone(&ex); + tokio::spawn(async move { ex.serve(forward("d-abort", 600)).await }) + }; + while ex.active() < 1 { + tokio::task::yield_now().await; + } + + task.abort(); + let _ = task.await; // JoinError::Cancelled — the abort landed + + assert_eq!(ex.active(), 0, "abort must free the slot via the guard"); + // And the freed slot is genuinely reusable. + let done = FakeRunner::completing("ok"); + let ex2 = executor(done, 1); + let r = ex2.serve(forward("d-after", 600)).await; + assert_eq!(r.status, DelegationStatus::Completed); + } + + #[tokio::test(start_paused = true)] + async fn a_stale_cancel_does_not_abort_a_reserving_admission() { + // Worker-side half of the wire-token contract: the CP swept admission + // A of "d-1" and its best-effort cancel (stamped with A's token) can + // arrive after this worker started serving re-admission B. The cancel + // must not abort B. + let runner = Arc::new(FakeRunner { + delay: Some(Duration::from_secs(300)), + ..Default::default() + }); + let ex = executor(Arc::clone(&runner), 1); + let mut fwd = forward("d-1", 600); + fwd.admission = 42; // B's admission + let task = { + let ex = Arc::clone(&ex); + tokio::spawn(async move { ex.serve(fwd).await }) + }; + while ex.active() < 1 { + tokio::task::yield_now().await; + } + + // A's stale cancel: same id, older token. + assert!(!ex.cancel("d-1", 7), "a stale token must be ignored"); + assert_eq!(ex.active(), 1, "B keeps running"); + + // B's own cancel works. + assert!(ex.cancel("d-1", 42)); + let result = task.await.unwrap(); + assert_eq!(result.status, DelegationStatus::Cancelled); + assert_eq!(result.admission, 42, "the terminal frame names B"); + } + + #[test] + fn json_escaped_len_matches_serde_json() { + let test_cases = vec![ + "", + "hello world", + "escapes: \x08 \t \n \x0c \r", + "control: \x00 \x1b \x1f", + "quotes and slashes: \" \\ \"", + "multibyte: 🦀 日本語 \u{1F980}", + ]; + for case in test_cases { + let expected = serde_json::to_string(case).unwrap().len() - 2; + assert_eq!( + json_escaped_len(case), + expected, + "mismatch for case: {:?}", + case + ); + } + } + + #[test] + fn escape_heavy_results_cannot_inflate_past_the_transport_budget() { + // 500 KiB of ESC characters (\x1b) has raw len <= 512 KiB, but JSON + // serializes each as \u001b (6 bytes), inflating to ~3 MiB. + // Unchecked, this exceeds the CP's pre-parse max_frame_bytes (1 MiB), + // causing the CP to close the WebSocket and kill every co-inflight + // delegation (F46). + let escapes = "\x1b".repeat(500 * 1024); + let capped = cap_result(escapes); + let serialized = serde_json::to_string(&capped).unwrap(); + // The serialized JSON string (including quotes) must remain within + // the transport budget (MAX_RESULT_BYTES + 2 for quotes). + assert!( + serialized.len() <= MAX_RESULT_BYTES + 2, + "serialized length {} exceeds budget", + serialized.len() + ); + assert!(capped.ends_with("bytes total exceeded the transport budget]")); + + // Strings with quotes and backslashes also inflate (2x) and must be bounded. + let slashes = "\\\"".repeat(300 * 1024); + let capped_slashes = cap_result(slashes); + let serialized_slashes = serde_json::to_string(&capped_slashes).unwrap(); + assert!( + serialized_slashes.len() <= MAX_RESULT_BYTES + 2, + "serialized slashes length {} exceeds budget", + serialized_slashes.len() + ); + } + + #[tokio::test(start_paused = true)] + async fn stalled_cancel_and_stalled_discard_share_one_teardown_bound() { + // Both cancel and discard can wedge (dead agent process, starved pool lock). + // cancel_and_discard must bound the SUM of both to TEARDOWN_BOUND (4s), + // strictly below the client's DRAIN_TIMEOUT (5s), so normal teardown + // finishes before the connection-level abort backstop (F52). + let runner = Arc::new(FakeRunner { + cancel_delay: Some(Duration::from_secs(10)), + discard_delay: Some(Duration::from_secs(10)), + ..Default::default() + }); + let ex = executor(Arc::clone(&runner), 1); + let start = tokio::time::Instant::now(); + ex.cancel_and_discard("s-wedged").await; + let elapsed = start.elapsed(); + assert_eq!( + elapsed, TEARDOWN_BOUND, + "teardown must complete at the shared {TEARDOWN_BOUND:?} bound" + ); + assert!( + runner.discarded().is_empty(), + "discard must be skipped when cancel consumes the entire bound" + ); + } +} diff --git a/crates/openab-core/src/control_plane/mod.rs b/crates/openab-core/src/control_plane/mod.rs new file mode 100644 index 000000000..c6c8d82ad --- /dev/null +++ b/crates/openab-core/src/control_plane/mod.rs @@ -0,0 +1,51 @@ +//! Runtime side of the OpenAB Agent Control Plane (`docs/adr/agent-control-plane.md`). +//! +//! The CP is a hub: runtimes dial *out* to it, register, and hold one +//! WebSocket for the process lifetime. This module is that client — the +//! connection state machine ([`client`]) and the delegation-serving side +//! ([`executor`]) — plus the seam between them and the ACP session pool. +//! +//! Three deliberate boundaries: +//! +//! - **Wire types are not redefined here.** They come from +//! `openab_cp::proto`, which `openab-core` depends on with +//! `default-features = false` — the contract without the server. +//! - **Prompt execution is behind a trait** ([`PromptRunner`]), so the client +//! can be driven end-to-end against a real CP without a real coding agent. +//! - **No cargo feature.** `[control_plane]` in config is the switch; a build +//! without the section behaves exactly as before. + +pub mod client; +pub mod executor; + +pub use client::ControlPlaneClient; +pub use executor::{ + delegation_session_key, DelegationExecutor, PromptOutcome, PromptRunner, RouterPromptRunner, +}; + +use crate::config::CpAgentType; +use openab_cp::proto::AgentType; + +/// Widen the runtime's two-variant role into the protocol enum. +/// +/// One-way on purpose: `AgentType::Observer` has no runtime counterpart, so +/// there is no `From for CpAgentType` to accidentally admit it. +impl From for AgentType { + fn from(t: CpAgentType) -> Self { + match t { + CpAgentType::Primary => AgentType::Primary, + CpAgentType::Worker => AgentType::Worker, + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn runtime_roles_map_onto_the_wire_enum() { + assert_eq!(AgentType::from(CpAgentType::Primary), AgentType::Primary); + assert_eq!(AgentType::from(CpAgentType::Worker), AgentType::Worker); + } +} diff --git a/crates/openab-core/src/cron.rs b/crates/openab-core/src/cron.rs index d67ba2769..085b721f2 100644 --- a/crates/openab-core/src/cron.rs +++ b/crates/openab-core/src/cron.rs @@ -217,7 +217,7 @@ fn compact_ordinal_set(sorted: &[u32]) -> String { fn render_run(start: u32, end: u32) -> String { if start == end { - format!("{}", start) + start.to_string() } else { format!("{}-{}", start, end) } diff --git a/crates/openab-core/src/directives.rs b/crates/openab-core/src/directives.rs index c1d5c27f2..3184cc929 100644 --- a/crates/openab-core/src/directives.rs +++ b/crates/openab-core/src/directives.rs @@ -295,7 +295,7 @@ mod tests { fs::write(&file_path, "").unwrap(); let aliases = HashMap::new(); - let result = resolve_workspace(&format!("{}", file_path.display()), &aliases, tmp.path()); + let result = resolve_workspace(&file_path.display().to_string(), &aliases, tmp.path()); assert!(result.is_err()); assert!(result.unwrap_err().contains("not a directory")); } diff --git a/crates/openab-core/src/dispatch.rs b/crates/openab-core/src/dispatch.rs index c3a20d263..1be454078 100644 --- a/crates/openab-core/src/dispatch.rs +++ b/crates/openab-core/src/dispatch.rs @@ -195,6 +195,10 @@ impl DispatchTarget for AdapterRouter { recipient, ) .await + // `DispatchTarget` is the platform-facing seam: keep it at + // `Result<()>` so every existing implementor (and its mock) is + // untouched by the executor's need for the turn summary. + .map(|_| ()) } } diff --git a/crates/openab-core/src/lib.rs b/crates/openab-core/src/lib.rs index 0e61e7cb2..4a02737f0 100644 --- a/crates/openab-core/src/lib.rs +++ b/crates/openab-core/src/lib.rs @@ -5,6 +5,9 @@ pub mod acp_mcp; pub mod redact; pub mod bot_turns; pub mod config; +/// Runtime membership in the Agent Control Plane (`[control_plane]`). +/// Unconditional: the opt-in is the config section, not a cargo feature. +pub mod control_plane; pub mod cron; pub mod directives; pub mod dispatch; diff --git a/crates/openab-cp/src/proto.rs b/crates/openab-cp/src/proto.rs index 5e70b984b..acf57b97b 100644 --- a/crates/openab-cp/src/proto.rs +++ b/crates/openab-cp/src/proto.rs @@ -37,6 +37,10 @@ use serde_json::Value; /// Wire protocol version. Carried in `cp/register`; the CP rejects /// registrations with a version it does not support. pub const PROTOCOL_VERSION: u32 = 1; +/// Minimum frame ceiling supported by the bundled runtime client. The server +/// may configure a larger value, but a smaller one would let an otherwise +/// valid runtime result disconnect the whole worker connection. +pub const MIN_RUNTIME_FRAME_BYTES: usize = 1024 * 1024; // --- JSON-RPC envelopes --- From e68e545eec71e2af6e721899ecea0f1279141134 Mon Sep 17 00:00:00 2001 From: Tautui Date: Sat, 31 Oct 2026 10:13:33 -0400 Subject: [PATCH 02/15] feat(cp): headless openab run with [control_plane] Port of the src/main.rs hunks from feat/cp-runtime-client-main 9455b504 (PR 1471) onto origin/main 6214700a. Plan: .agents/tasks/phase0-cp-runtime/ plan.md, Commit 2. What it adds: - HeadlessMode and headless_run_mode: an adapter-less config picks the full runtime for [control_plane] type = worker, or for [mcp] plus [control_plane]. [mcp] alone stays facade-only. A primary alone stays an error, because its prompts come from a chat platform. - supervise_control_plane and wait_for_control_plane_exit: a client task that exits or fails stops the process after the ordered cleanup. - The control-plane client is spawned after the router and stopped before pool.shutdown, so the CP sees a clean disconnect. - Discord and no-Discord runtime waits race the client task with select!. Main code kept: the main() structure, the pool/router construction, and the existing shutdown order. The only removed lines are the facade block moved into the match, the old no-adapter bail text, and the bare shutdown_signal awaits that now sit inside select!. The bail text now also names [line], [lineworks], [feishu] and [teams], as the branch wrote it. Verification (macmini, cargo and rustc 1.99.1, clone ~/repo/openab-phase0): - cargo check -p openab-core -p openab: pass - cargo clippy -p openab-core -p openab -- -D warnings: pass - cargo test -p openab: 26 passed, 0 failed. Includes tests::nothing_configured_is_still_a_startup_error and the 8 headless mode tests. - cargo test -p openab-core --lib control_plane: 50 passed, 0 failed - cargo test -p openab-core --test cp_client: no such target yet. It is added in commit 6. - cargo fmt --check: 14 diffs in src/main.rs, the same count as main. No new diffs. - Linux run: not needed. Known failure on main, not caused by this series: secrets::tests::resolve_exec_nonzero_exit Fails on origin/main 6214700a with the same assertion. Not fixed or skipped. --- src/main.rs | 291 +++++++++++++++++++++++++++++++++++++++++++++++----- 1 file changed, 266 insertions(+), 25 deletions(-) diff --git a/src/main.rs b/src/main.rs index f49879778..381e748c5 100644 --- a/src/main.rs +++ b/src/main.rs @@ -195,6 +195,73 @@ fn has_unified_platform(cfg: &config::Config) -> bool { || (cfg!(feature = "lineworks") && lineworks_activated(cfg.lineworks.as_ref())) } +/// What `openab run` does when NO chat adapter is configured. +/// +/// Two headless deployments exist, and the order below is the whole rule: +/// +/// - Any configuration that needs the control-plane client and the normal +/// runtime takes `FullRuntime`: a worker can be adapter-less, while a primary +/// needs either a chat adapter or `[mcp]` (the facade supplies its local +/// initiating surface). The full path builds the pool/router and starts both +/// configured services. +/// - `[mcp]` alone stays exactly what it was: the facade in the foreground. +/// +/// `type = "primary"` alone (without `[mcp]`) deliberately does NOT unlock a +/// headless boot. A primary is the *initiating* side of a delegation; its +/// prompts come from a chat platform, so an adapter-less primary has no way to +/// be given work and is a misconfiguration worth failing on. +#[derive(Debug, PartialEq, Eq)] +enum HeadlessMode { + FullRuntime, + FacadeOnly, + None, +} + +fn headless_run_mode(cfg: &config::Config) -> HeadlessMode { + let cp_worker = cfg + .control_plane + .as_ref() + .is_some_and(|cp| cp.agent_type == openab_core::config::CpAgentType::Worker); + if cp_worker { + HeadlessMode::FullRuntime + } else if cfg.mcp.is_some() { + if cfg.control_plane.is_some() { + HeadlessMode::FullRuntime + } else { + HeadlessMode::FacadeOnly + } + } else { + HeadlessMode::None + } +} + +struct SupervisedControlPlaneTask { + abort: tokio::task::AbortHandle, + completion: tokio::sync::oneshot::Receiver>, +} + +fn supervise_control_plane(handle: tokio::task::JoinHandle<()>) -> SupervisedControlPlaneTask { + let abort = handle.abort_handle(); + let (tx, completion) = tokio::sync::oneshot::channel(); + tokio::spawn(async move { + let _ = tx.send(handle.await); + }); + SupervisedControlPlaneTask { abort, completion } +} + +async fn wait_for_control_plane_exit( + task: &mut Option, +) -> anyhow::Error { + let Some(task) = task.as_mut() else { + return std::future::pending::().await; + }; + match (&mut task.completion).await { + Ok(Ok(())) => anyhow::anyhow!("control-plane client exited unexpectedly"), + Ok(Err(error)) => anyhow::anyhow!("control-plane client task failed: {error}"), + Err(_) => anyhow::anyhow!("control-plane client supervisor stopped unexpectedly"), + } +} + /// Single LINE WORKS activation validator: the resolved (config → env → /// default) credentials must be complete and non-empty — the same rule the /// adapter constructor applies. Used by startup preflight AND cron platform @@ -421,26 +488,53 @@ async fn main() -> anyhow::Result<()> { && cfg.telegram.is_none() && !has_unified_platform(&cfg) { - // Facade-only run mode (#1451): an adapter-less config with `[mcp]` - // present is a valid deployment — the broker serves just the OAB MCP - // Facade listener. One entrypoint, config-driven: hosts that only - // need the capability surface (coding-CLI-only users, dev loops, CI - // runners, agent hosts with no chat platform) run the same - // `openab run` with a two-line config instead of a chat token. - if let Some(mcp_cfg) = cfg.mcp.clone() { - tracing::info!( - listen = %mcp_cfg.listen, - "no chat adapter configured — running in facade-only mode ([mcp] present)" - ); - // Foreground, not spawned: the facade IS the workload. A bind - // failure or server exit terminates the process (fail fast). - return openab_mcp::mcp::facade::serve_http(&mcp_cfg.listen) - .await - .map_err(|e| anyhow::anyhow!("OAB MCP facade exited: {e:#}")); + match headless_run_mode(&cfg) { + // Control-plane worker (Agent Control Plane ADR): an adapter-less + // config whose `[control_plane]` says `type = "worker"` is a valid + // deployment — the runtime's work arrives as `cp/delegate` over the + // CP socket instead of as chat messages. Fall through to the normal + // boot path: it builds the pool and router the delegation executor + // needs, spawns the CP client, and (if `[mcp]` is also present) + // starts the facade alongside, exactly as an adapter run would. + HeadlessMode::FullRuntime => { + let cp = cfg + .control_plane + .as_ref() + .expect("control-plane headless mode implies [control_plane]"); + tracing::info!( + agent = %format!("{}/{}", cp.namespace, cp.name), + r#type = %match cp.agent_type { + openab_core::config::CpAgentType::Worker => "worker", + openab_core::config::CpAgentType::Primary => "primary", + }, + mcp = cfg.mcp.is_some(), + "no chat adapter configured — running headless with the control-plane client" + ); + } + // Facade-only run mode (#1451): an adapter-less config with `[mcp]` + // present is a valid deployment — the broker serves just the OAB MCP + // Facade listener. One entrypoint, config-driven: hosts that only + // need the capability surface (coding-CLI-only users, dev loops, CI + // runners, agent hosts with no chat platform) run the same + // `openab run` with a two-line config instead of a chat token. + HeadlessMode::FacadeOnly => { + let mcp_cfg = cfg.mcp.clone().expect("facade-only mode implies [mcp]"); + tracing::info!( + listen = %mcp_cfg.listen, + "no chat adapter configured — running in facade-only mode ([mcp] present)" + ); + // Foreground, not spawned: the facade IS the workload. A bind + // failure or server exit terminates the process (fail fast). + return openab_mcp::mcp::facade::serve_http(&mcp_cfg.listen) + .await + .map_err(|e| anyhow::anyhow!("OAB MCP facade exited: {e:#}")); + } + HeadlessMode::None => { + anyhow::bail!( + "no adapter configured — add [discord], [slack], [telegram], [line], [lineworks], [feishu], [wecom], [googlechat], [teams], or [gateway] to config (or [mcp] for facade-only mode, or [control_plane] with type = \"worker\" for control-plane worker mode), or set the corresponding platform environment variables" + ); + } } - anyhow::bail!( - "no adapter configured — add [discord], [slack], [telegram], [wecom], [googlechat], or [gateway] to config (or [mcp] for facade-only mode), or set platform env vars (TELEGRAM_BOT_TOKEN, etc.)" - ); } // --- Lifecycle hooks: Unix-only. Fail fast on unsupported platforms. --- @@ -559,6 +653,12 @@ async fn main() -> anyhow::Result<()> { }); } + // Taken before the sections below are moved into their runtime components: + // the control-plane client is constructed much later (it needs the router), + // and `cfg.agent` / `cfg.pool` are gone by then. + let control_plane_cfg = cfg.control_plane.clone(); + let prompt_hard_timeout_secs = cfg.pool.prompt_hard_timeout_secs; + let pool_inner = acp::SessionPool::new( cfg.agent, cfg.pool.max_sessions, @@ -877,6 +977,25 @@ async fn main() -> anyhow::Result<()> { // Shutdown signal for Slack adapter let (shutdown_tx, shutdown_rx) = tokio::sync::watch::channel(false); + // --- Agent Control Plane membership (`[control_plane]`) --- + // Spawned here because the client needs the router (its delegation executor + // drives ACP turns through the same seam every platform uses) and the + // shutdown watch. Absent section = no task, no socket, no behaviour change. + // The auth key is used only for the `Authorization` header on the outbound + // upgrade: it is never logged, and never reaches the agent subprocess — + // `[agent].env` plumbing is untouched by this. + let mut cp_task = control_plane_cfg.map(|cp_cfg| { + let runner: Arc = Arc::new( + openab_core::control_plane::RouterPromptRunner::new(router.clone()), + ); + let client = Arc::new(openab_core::control_plane::ControlPlaneClient::new( + cp_cfg, + runner, + std::time::Duration::from_secs(prompt_hard_timeout_secs), + )); + supervise_control_plane(tokio::spawn(client.run(shutdown_rx.clone()))) + }); + let dispatchers: Arc>>> = Arc::new(Mutex::new(Vec::new())); // Spawn cleanup task @@ -1619,6 +1738,11 @@ async fn main() -> anyhow::Result<()> { None }; + // A control-plane task failure is fatal, but it must not bypass the + // ordered cleanup below (notably pool.shutdown and pre_shutdown hooks). + // Record it while the runtime wait ends, clean everything up, then return. + let mut runtime_error: Option = None; + // Run Discord adapter (foreground, blocking) or wait for ctrl_c #[cfg(feature = "discord")] if let Some(discord_cfg) = cfg.discord { @@ -1735,7 +1859,14 @@ async fn main() -> anyhow::Result<()> { }); info!("discord bot running"); - match client.start().await { + let discord_result = tokio::select! { + result = client.start() => result, + error = wait_for_control_plane_exit(&mut cp_task) => { + runtime_error = Some(error); + Ok(()) + } + }; + match discord_result { Err(serenity::Error::Gateway(GatewayError::DisallowedGatewayIntents)) => { error!( "Discord rejected privileged intents. \ @@ -1756,8 +1887,10 @@ async fn main() -> anyhow::Result<()> { } } else { info!("running without discord, press ctrl+c to stop"); - shutdown_signal().await; - info!("shutdown signal received"); + tokio::select! { + _ = shutdown_signal() => info!("shutdown signal received"), + error = wait_for_control_plane_exit(&mut cp_task) => runtime_error = Some(error), + } } // When discord feature is disabled at compile time, use this fallback block. // (When discord feature IS enabled but no [discord] config exists, the `else` @@ -1765,8 +1898,10 @@ async fn main() -> anyhow::Result<()> { #[cfg(not(feature = "discord"))] { info!("running without discord, press ctrl+c to stop"); - shutdown_signal().await; - info!("shutdown signal received"); + tokio::select! { + _ = shutdown_signal() => info!("shutdown signal received"), + error = wait_for_control_plane_exit(&mut cp_task) => runtime_error = Some(error), + } } // Cleanup @@ -1797,6 +1932,26 @@ async fn main() -> anyhow::Result<()> { for d in dispatchers.lock().unwrap().iter() { d.shutdown(); } + // Stop the control-plane client BEFORE the pool: it cancels its in-flight + // delegations (each stopping its agent and dropping its session) and closes + // the socket, so the CP sees a clean disconnect instead of a lease timeout. + // Tearing the pool down first would leave those turns writing to sessions + // that no longer exist. + if let Some(mut task) = cp_task { + match tokio::time::timeout(std::time::Duration::from_secs(10), &mut task.completion).await { + Ok(Ok(Ok(()))) => {} + Ok(Ok(Err(error))) => { + tracing::warn!(%error, "control-plane client failed during shutdown"); + } + Ok(Err(_)) => { + tracing::warn!("control-plane client supervisor stopped during shutdown"); + } + Err(_) => { + tracing::warn!("control-plane client missed the shutdown deadline — aborting"); + task.abort.abort(); + } + } + } let shutdown_pool = pool; shutdown_pool.shutdown().await; if let Some(ref hook) = shutdown_hook { @@ -1805,7 +1960,11 @@ async fn main() -> anyhow::Result<()> { } } info!("openab shut down"); - Ok(()) + if let Some(error) = runtime_error { + Err(error) + } else { + Ok(()) + } } fn parse_id_set(raw: &[String], label: &str) -> anyhow::Result> { @@ -2086,4 +2245,86 @@ agent_id = "1000002" assert_eq!(has_unified_wecom_config(&cfg), cfg!(feature = "wecom")); } + + // --- headless run modes (no chat adapter configured) --- + + fn cp_section(agent_type: &str) -> String { + format!( + "[control_plane]\nurl = \"wss://cp:9800/cp\"\nauth_key = \"k\"\n\ + namespace = \"prod\"\nname = \"w\"\ntype = \"{agent_type}\"\n" + ) + } + + #[test] + fn nothing_configured_is_still_a_startup_error() { + let cfg = config::parse_config_str("", "test").unwrap(); + assert_eq!(headless_run_mode(&cfg), HeadlessMode::None); + } + + #[test] + fn mcp_only_is_still_facade_only() { + // Regression guard for #1451: adding control-plane modes must not + // change what an `[mcp]`-only config does. + let cfg = config::parse_config_str("[mcp]\n", "test").unwrap(); + assert_eq!(headless_run_mode(&cfg), HeadlessMode::FacadeOnly); + } + + #[tokio::test] + async fn an_early_control_plane_exit_is_supervised() { + let mut task = Some(supervise_control_plane(tokio::spawn(async {}))); + let error = tokio::time::timeout( + std::time::Duration::from_secs(1), + wait_for_control_plane_exit(&mut task), + ) + .await + .expect("the completed task is observed"); + assert!(error.to_string().contains("exited unexpectedly")); + } + + #[test] + fn a_control_plane_worker_boots_without_any_adapter() { + let cfg = config::parse_config_str(&cp_section("worker"), "test").unwrap(); + assert_eq!(headless_run_mode(&cfg), HeadlessMode::FullRuntime); + } + + #[test] + fn a_control_plane_primary_alone_does_not_unlock_a_headless_boot() { + // A primary initiates delegations; its prompts come from a chat + // platform, so an adapter-less primary has no way to be given work. + let cfg = config::parse_config_str(&cp_section("primary"), "test").unwrap(); + assert_eq!(headless_run_mode(&cfg), HeadlessMode::None); + } + + #[test] + fn a_primary_with_mcp_takes_the_full_boot_path_so_both_run() { + // Facade-only would foreclose the CP client; with both sections + // present the full boot serves the facade in the background AND + // registers with the control plane. + let cfg = + config::parse_config_str(&format!("[mcp]\n{}", cp_section("primary")), "test").unwrap(); + assert_eq!(headless_run_mode(&cfg), HeadlessMode::FullRuntime); + } + + #[test] + fn mcp_plus_worker_takes_the_worker_path_so_both_run() { + // The worker path falls through to the normal boot, which spawns the + // facade in the background — the facade-only path would return before + // ever reaching the CP client. + let cfg = + config::parse_config_str(&format!("[mcp]\n{}", cp_section("worker")), "test").unwrap(); + assert_eq!(headless_run_mode(&cfg), HeadlessMode::FullRuntime); + } + + /// A configured adapter never consults the headless matrix at all — the + /// `[control_plane]` section simply rides along with it. + #[test] + fn an_adapter_config_with_a_control_plane_section_parses() { + let cfg = config::parse_config_str( + &format!("[discord]\nbot_token = \"x\"\n{}", cp_section("primary")), + "test", + ) + .unwrap(); + assert!(cfg.discord.is_some()); + assert!(cfg.control_plane.is_some()); + } } From 612ef7f20e0c4a0e981e2bf4cd4b068b79cd5131 Mon Sep 17 00:00:00 2001 From: Tautui Date: Sat, 31 Oct 2026 10:30:55 -0400 Subject: [PATCH 03/15] feat(cp): openab agent spawn/list/status/cancel over the local socket Plan commit 3, from reference 3742e03c (#1529), restricted to: - new control_plane/local.rs and primary.rs, taken whole - control_plane/client.rs and mod.rs hunks, applied as a patch. The parent blobs of both files matched HEAD, so the patch applied cleanly. - src/main.rs: Agent subcommand, AgentAction, run_agent_command, parse_agent_labels, validate_deadline_secs, local API supervision and shutdown, and six unit tests from the same commit. Not taken: openab-mcp facade.rs, src/cp_tools.rs, src/acp_tunnel_source.rs, the headless_run_mode change, and the MCP facade wiring. spawn_within_bound and new_delegation_id are copied from cp_tools.rs into main.rs because run_agent_command calls them. Five rustfmt fixes apply to added lines only. Checks, run on macmini (rustc 1.99.1) in ~/repo/openab-phase0: - cargo check -p openab-core -p openab: pass, no warnings - cargo test -p openab-core --lib control_plane: 106 passed, 0 failed - cargo test -p openab: 32 passed, 0 failed - cargo run -p openab -- agent --help: lists spawn, list, status, cancel - cargo fmt -- --check: no hunks in added code; remaining hunks predate this commit - Cargo.lock unchanged --- .../openab-core/src/control_plane/client.rs | 417 ++++- crates/openab-core/src/control_plane/local.rs | 1500 +++++++++++++++++ crates/openab-core/src/control_plane/mod.rs | 12 +- .../openab-core/src/control_plane/primary.rs | 1488 ++++++++++++++++ src/main.rs | 425 ++++- 5 files changed, 3827 insertions(+), 15 deletions(-) create mode 100644 crates/openab-core/src/control_plane/local.rs create mode 100644 crates/openab-core/src/control_plane/primary.rs diff --git a/crates/openab-core/src/control_plane/client.rs b/crates/openab-core/src/control_plane/client.rs index 169f85bcd..d4bc0404c 100644 --- a/crates/openab-core/src/control_plane/client.rs +++ b/crates/openab-core/src/control_plane/client.rs @@ -39,6 +39,7 @@ use tracing::{debug, error, info, warn}; use crate::config::{ControlPlaneConfig, CpAgentType}; use crate::control_plane::executor::{cap_text, DelegationExecutor, PromptRunner}; +use crate::control_plane::primary::{ClientCommand, PrimaryState}; /// Backoff ceiling, matching the gateway adapter. const MAX_BACKOFF_SECS: u64 = 30; @@ -76,6 +77,35 @@ type Ws = WebSocketStream>; type WsSink = futures_util::stream::SplitSink; type WsStream = futures_util::stream::SplitStream; +/// A cloneable handle for submitting primary-side commands to a running +/// [`ControlPlaneClient`]. Obtained via [`ControlPlaneClient::handle`], it is +/// the initiating surface for local callers (the UDS server, an MCP tool, a +/// CLI): the client's serve loop is the single socket owner, so a command +/// carries a `oneshot` reply and is answered from inside that loop. +#[derive(Clone)] +pub struct ControlPlaneHandle { + tx: tokio::sync::mpsc::Sender, +} + +impl ControlPlaneHandle { + /// Submit a command to the serve loop, returning `NotConnected` if the + /// loop is not currently running (channel closed). + pub async fn submit( + &self, + command: ClientCommand, + ) -> Result<(), crate::control_plane::primary::CommandError> { + self.tx + .send(command) + .await + .map_err(|_| crate::control_plane::primary::CommandError::NotConnected) + } +} + +/// How many primary-side commands may be queued to the serve loop before a +/// submitter awaits. Bounded so a wedged loop applies backpressure rather than +/// growing an unbounded queue. +const COMMAND_QUEUE_DEPTH: usize = 64; + /// Runtime client for the OpenAB Agent Control Plane. pub struct ControlPlaneClient { cfg: ControlPlaneConfig, @@ -84,6 +114,13 @@ pub struct ControlPlaneClient { executor: Arc, /// Monotonic JSON-RPC request id for frames this client originates. next_id: std::sync::atomic::AtomicU64, + /// Sender kept so [`ControlPlaneClient::handle`] can hand out clones for + /// the process lifetime, independent of connection state. + command_tx: tokio::sync::mpsc::Sender, + /// Receiver, taken once by [`ControlPlaneClient::run`]. Behind a mutex so + /// `run` (which takes `Arc`) can move it into the loop; a second + /// `run` finds it gone and serves without a command channel. + command_rx: std::sync::Mutex>>, } impl ControlPlaneClient { @@ -103,11 +140,23 @@ impl ControlPlaneClient { cfg.max_delegated_sessions, prompt_hard_timeout, )); + let (command_tx, command_rx) = tokio::sync::mpsc::channel(COMMAND_QUEUE_DEPTH); Self { cfg, instance_id, executor, next_id: std::sync::atomic::AtomicU64::new(1), + command_tx, + command_rx: std::sync::Mutex::new(Some(command_rx)), + } + } + + /// A cloneable command handle. Valid for the process lifetime regardless of + /// connection state: submitting while disconnected succeeds at the channel + /// but the command is answered `NotConnected`/`Disconnected` by the loop. + pub fn handle(&self) -> ControlPlaneHandle { + ControlPlaneHandle { + tx: self.command_tx.clone(), } } @@ -122,6 +171,17 @@ impl ControlPlaneClient { /// Connect, register, serve — forever, until `shutdown` flips. pub async fn run(self: Arc, mut shutdown: watch::Receiver) { + // Take the command receiver once for the process lifetime. A second + // `run` (there should never be one) serves without a command channel. + let mut command_rx = self.command_rx.lock().unwrap().take(); + // Primary-side correlation and tracking is PROCESS-lifetime, not + // per-connection: one instance is created here and passed through every + // connect/serve. A reconnect fails/removes pending and live delegations + // (`fail_all_live`) but PRESERVES terminal history, so a caller that + // recorded a completed delegation before a blip can still read it after + // reconnect. Recreating it per connection would silently discard that + // history and orphan any awaiter that spanned the gap. + let mut primary = PrimaryState::new(); let mut backoff = 1u64; loop { if *shutdown.borrow() { @@ -138,7 +198,9 @@ impl ControlPlaneClient { // connect_and_serve owns shutdown for each phase: connect/register // race it directly, then the serve loop alone performs // cancel/drain/abort. No outer arm may drop serve mid-cleanup. - let served = self.connect_and_serve(&mut shutdown).await; + let served = self + .connect_and_serve(&mut shutdown, &mut command_rx, &mut primary) + .await; match served { Ok(Outcome::Shutdown) => { info!("control-plane client shutting down"); @@ -166,11 +228,27 @@ impl ControlPlaneClient { retry_delay_ms = retry_delay.as_millis(), "control-plane reconnect delay selected" ); - tokio::select! { - _ = tokio::time::sleep(retry_delay) => {} - _ = shutdown.changed() => { - info!("control-plane client shutting down"); - return; + // While disconnected, reject any submitted command promptly rather + // than letting it sit in the buffer until the next session. The + // channel stays open (the handle is process-lifetime), so this only + // answers what is queued; new submissions during the sleep are + // answered too. A closed channel is replaced with `None` so + // `recv_command` becomes a pending future and the sleep proceeds. + let sleep = tokio::time::sleep(retry_delay); + tokio::pin!(sleep); + loop { + tokio::select! { + _ = &mut sleep => break, + _ = shutdown.changed() => { + info!("control-plane client shutting down"); + return; + } + maybe = recv_command(&mut command_rx) => { + match maybe { + Some(command) => reject_command(command), + None => { command_rx = None; } + } + } } } backoff = next_backoff(backoff); @@ -180,6 +258,8 @@ impl ControlPlaneClient { async fn connect_and_serve( &self, shutdown: &mut watch::Receiver, + command_rx: &mut Option>, + primary: &mut PrimaryState, ) -> anyhow::Result { let ws = tokio::select! { result = self.connect_with_timeout() => result?, @@ -204,7 +284,8 @@ impl ControlPlaneClient { max_delegated_sessions = ack.effective_max_delegated_sessions, "registered with control plane" ); - self.serve(sink, stream, &ack, shutdown).await + self.serve(sink, stream, &ack, shutdown, command_rx, primary) + .await } async fn connect_with_timeout(&self) -> anyhow::Result { @@ -314,6 +395,8 @@ impl ControlPlaneClient { mut stream: WsStream, ack: &RegisterAck, shutdown: &mut watch::Receiver, + command_rx: &mut Option>, + primary: &mut PrimaryState, ) -> anyhow::Result { let mut heartbeat = tokio::time::interval(Duration::from_secs( // A zero interval would spin; the CP's own default is 15s. @@ -329,6 +412,13 @@ impl ControlPlaneClient { (ack.effective_max_delegated_sessions as usize).max(1), ); let mut serving: Vec> = Vec::new(); + // Primary-side correlation and tracking for delegations THIS runtime + // initiates is owned by `run` and threaded through every connection, so + // it lives for the process, not the socket. It still lives inside the + // single socket owner — no lock guards it and every transition is + // serialized against frame writes — but terminal history and awaiters + // survive a reconnect. The `&mut` borrow makes that single-owner + // discipline explicit. let outcome = loop { tokio::select! { @@ -347,6 +437,23 @@ impl ControlPlaneClient { break Outcome::Disconnected; } } + // A primary-side command: the local initiating surface (UDS + // server, MCP tool, CLI) asked to spawn/cancel/await/list. The + // serve loop is the single socket owner, so it — not the + // caller — writes the frame and parks the reply. + maybe_command = recv_command(command_rx) => { + match maybe_command { + Some(command) => { + if self.handle_command(command, primary, &mut sink).await.is_err() { + break Outcome::Disconnected; + } + } + // The command channel closed (no handles left). Stop + // selecting on it by replacing it with `None`, which + // `recv_command` treats as a pending future. + None => { *command_rx = None; } + } + } // A finished delegation reports itself. Emitted by the runtime // when the turn ends, never by the model: this is the only // frame that closes the initiator's wait. @@ -360,7 +467,7 @@ impl ControlPlaneClient { let Some(msg) = inbound else { break Outcome::Disconnected }; match msg { Ok(Message::Text(text)) => { - match self.handle_frame(&text) { + match self.classify_inbound(&text, primary) { FrameAction::Serve { ack, forward } => { let executor = Arc::clone(&self.executor); let tx = result_tx.clone(); @@ -422,6 +529,11 @@ impl ControlPlaneClient { ); } self.executor.cancel_all(); + // Primary-side mirror: every delegation THIS runtime initiated and was + // still awaiting can no longer complete over a dead socket, and every + // parked command reply would otherwise hang forever. Answer them all + // `Disconnected` so a caller's `await` returns instead of leaking. + primary.fail_all_live(); let drained = tokio::time::timeout(DRAIN_TIMEOUT, async { for handle in &mut serving { let _ = handle.await; @@ -439,6 +551,159 @@ impl ControlPlaneClient { Ok(outcome) } + /// Turn one primary-side command into an outbound frame, parking its reply + /// in `primary` for later correlation. `Await` needs no frame — it only + /// registers interest in a terminal that arrives over the socket. Returns + /// `Err` only when the socket write fails, which ends the session. + async fn handle_command( + &self, + command: ClientCommand, + primary: &mut PrimaryState, + sink: &mut WsSink, + ) -> anyhow::Result<()> { + // A primary must be registered as such to initiate; a worker-only + // runtime rejects spawn/cancel/list locally so a misconfigured caller + // gets a clear answer instead of a CP policy denial round-trip. Await + // is allowed for either role (it only reads local tracking). + match command { + ClientCommand::Spawn { request, reply } => { + if self.cfg.agent_type != CpAgentType::Primary { + let _ = reply.send(Err(primary_only_error())); + return Ok(()); + } + let emission = primary.begin_spawn(self.next_id(), request, reply); + let frame = JsonRpcRequest::new( + emission.rpc_id, + methods::DELEGATE, + Some(serde_json::to_value(&emission.params)?), + ); + // A Spawn whose complete serialized frame exceeds the transport + // ceiling must NOT tear down the connection: an oversized prompt + // is a caller-side fault, not a socket fault. Detect it before + // writing, remove/fail the parked pending request with + // InvalidRequest, and return Ok so the serve loop keeps running. + if !frame_within_limit(&frame)? { + let bytes = serialized_len(&frame)?; + warn!( + rpc_id = emission.rpc_id, + frame_bytes = bytes, + limit = MAX_FRAME_BYTES, + "cp/delegate frame exceeds the transport limit; failing the spawn without disconnecting" + ); + primary.fail_pending_invalid_request( + emission.rpc_id, + format!( + "delegate frame ({bytes} bytes) exceeds the transport limit \ + ({MAX_FRAME_BYTES} bytes); reduce the prompt" + ), + ); + return Ok(()); + } + send(sink, &frame).await?; + } + ClientCommand::Cancel { + handle, + reason, + reply, + } => { + if self.cfg.agent_type != CpAgentType::Primary { + let _ = reply.send(Err(primary_only_error())); + return Ok(()); + } + let rpc_id = self.next_id(); + if let Some(params) = primary.begin_cancel(rpc_id, &handle, reason, reply) { + let frame = JsonRpcRequest::new( + rpc_id, + methods::CANCEL, + Some(serde_json::to_value(¶ms)?), + ); + send(sink, &frame).await?; + } + // begin_cancel already answered the caller if the handle was + // unknown; no frame in that case. + } + ClientCommand::ListAgents { reply } => { + let rpc_id = self.next_id(); + primary.begin_list_agents(rpc_id, reply); + let frame = + JsonRpcRequest::new(rpc_id, methods::LIST_AGENTS, Some(serde_json::json!({}))); + send(sink, &frame).await?; + } + ClientCommand::Await { handle, reply } => { + primary.begin_await(&handle, reply); + } + ClientCommand::Check { handle, reply } => { + primary.check(&handle, reply); + } + } + Ok(()) + } + + /// Classify one inbound frame, routing primary-side concerns (JSON-RPC + /// responses to our own requests, and the initiator-bound + /// `cp/delegate_result`) into `primary`, and everything else to the + /// worker-side [`Self::handle_frame`]. + fn classify_inbound(&self, text: &str, primary: &mut PrimaryState) -> FrameAction { + let msg: JsonRpcMessage = match serde_json::from_str(text) { + Ok(m) => m, + Err(e) => { + warn!(error = %e, "malformed frame from the control plane"); + return FrameAction::Ignore; + } + }; + // A frame with no method is a response to one of our own requests. + // Route it to the primary-side correlator first; if it does not own + // the id (e.g. a heartbeat/delegate_result ack), fall back to the + // worker-side "absorb and log errors" behavior. + if msg.method.is_none() { + let Some(id) = msg.id else { + if let Some(err) = msg.error { + warn!(code = err.code, message = %err.message, "control plane returned an error with no id"); + } + return FrameAction::Ignore; + }; + if primary.on_reply(id, msg.result, msg.error.clone()) { + return FrameAction::Ignore; + } + // Not a primary request: it is a response to a heartbeat or a + // delegate_result the worker side sent. Errors are worth a line. + if let Some(err) = msg.error { + warn!(code = err.code, message = %err.message, "control plane returned an error"); + } + return FrameAction::Ignore; + } + // A `cp/delegate_result` REQUEST is the CP delivering a terminal to us + // as the initiator. Route it to primary and ack it as ours. + if msg.method.as_deref() == Some(methods::DELEGATE_RESULT) { + let id = match msg.require_request_envelope() { + Ok(id) => id, + Err(err) => { + warn!(code = err.code, message = %err.message, "invalid cp/delegate_result envelope"); + return error_reply(msg.id.unwrap_or(0), err); + } + }; + let params: Option = + msg.params.and_then(|p| serde_json::from_value(p).ok()); + match params { + Some(params) => { + // Ack regardless of whether we track it: an untracked + // (id, admission) is a stale/foreign frame the CP cannot + // act on, and a JSON-RPC error would be misread as our + // failure rather than a routing miss. + let _known = primary.on_delegate_result(¶ms); + ok_reply(id) + } + None => error_reply( + id, + ErrorObject::new(codes::INVALID_PARAMS, "invalid cp/delegate_result params"), + ), + } + } else { + // Worker-side frames (cp/delegate, cp/cancel, unknown methods). + self.handle_frame(text) + } + } + /// Classify one inbound frame. Never spawns and never writes: the serve /// loop owns both, so this stays a pure-enough function to unit-test (it /// does signal cancellation, which has no other home). @@ -608,6 +873,66 @@ async fn send(sink: &mut WsSink, frame: &JsonRpcRequest) -> anyhow::Result<()> { Ok(()) } +/// Serialized byte length of an outbound frame. Pure and testable — the same +/// measurement [`send`] and the oversized-Spawn guard use, so the check and the +/// wire agree by construction. +fn serialized_len(frame: &JsonRpcRequest) -> anyhow::Result { + Ok(serde_json::to_string(frame)?.len()) +} + +/// Whether an outbound frame fits the transport ceiling. `Ok(false)` means the +/// caller must fail the request WITHOUT writing (an oversized cp/delegate is a +/// caller-side fault, not a socket fault, so it must not disconnect). +fn frame_within_limit(frame: &JsonRpcRequest) -> anyhow::Result { + Ok(serialized_len(frame)? <= MAX_FRAME_BYTES) +} + +/// Receive from an optional command channel. When the channel is `None` (never +/// created, or closed and cleared), this future is `pending` forever so the +/// serve loop's `select!` simply never wakes on this arm. +async fn recv_command( + rx: &mut Option>, +) -> Option { + match rx { + Some(rx) => rx.recv().await, + None => std::future::pending().await, + } +} + +/// Answer a command that arrived while the client is not connected. +fn reject_command(command: ClientCommand) { + use crate::control_plane::primary::CommandError; + match command { + ClientCommand::Spawn { reply, .. } => { + let _ = reply.send(Err(CommandError::NotConnected)); + } + ClientCommand::Cancel { reply, .. } => { + let _ = reply.send(Err(CommandError::NotConnected)); + } + ClientCommand::ListAgents { reply } => { + let _ = reply.send(Err(CommandError::NotConnected)); + } + ClientCommand::Await { reply, .. } => { + // Awaiting only reads local tracking, but a disconnected client + // holds none: there is nothing to wait on. + let _ = reply.send(Err(CommandError::NotConnected)); + } + ClientCommand::Check { reply, .. } => { + // Same rationale as Await: no tracking survives a disconnect. + let _ = reply.send(Err(CommandError::NotConnected)); + } + } +} + +/// Error answered to a spawn/cancel/list command on a runtime not registered +/// as `primary`. +fn primary_only_error() -> crate::control_plane::primary::CommandError { + crate::control_plane::primary::CommandError::Internal( + "this runtime is not registered as a control-plane primary and cannot initiate delegations" + .into(), + ) +} + fn validate_register_ack(result: serde_json::Value) -> anyhow::Result { let ack: RegisterAck = serde_json::from_value(result)?; anyhow::ensure!( @@ -662,6 +987,25 @@ enum Outcome { Disconnected, } +/// Test-only helpers for constructing a [`ControlPlaneHandle`] and its +/// receiver without a full client + socket, so sibling modules (the local IPC +/// server) can exercise their own plumbing. +#[cfg(test)] +pub(crate) mod test_support { + use super::*; + + /// A handle plus the receiver end of its command channel. The receiver is + /// returned so the caller can drop it (making the handle answer + /// `NotConnected`) or drain commands in a stub loop. + pub(crate) fn handle_and_rx() -> ( + ControlPlaneHandle, + tokio::sync::mpsc::Receiver, + ) { + let (tx, rx) = tokio::sync::mpsc::channel(COMMAND_QUEUE_DEPTH); + (ControlPlaneHandle { tx }, rx) + } +} + #[cfg(test)] mod tests { use super::*; @@ -974,4 +1318,61 @@ max_delegated_sessions = 3 .expect("shutdown is checked before dialling"); drop(tx); } + + #[test] + fn frame_within_limit_flags_an_oversized_delegate_frame() { + // A small frame fits; a frame whose serialized form exceeds the + // transport ceiling does not. This is the same measurement the + // oversized-Spawn guard uses before deciding to fail-without-writing. + let small = JsonRpcRequest::new( + 1, + methods::DELEGATE, + Some(serde_json::json!({"prompt": "hi"})), + ); + assert!(frame_within_limit(&small).unwrap()); + assert!(serialized_len(&small).unwrap() <= MAX_FRAME_BYTES); + + let huge = JsonRpcRequest::new( + 2, + methods::DELEGATE, + Some(serde_json::json!({ "prompt": "x".repeat(MAX_FRAME_BYTES + 1024) })), + ); + assert!( + !frame_within_limit(&huge).unwrap(), + "a frame larger than the ceiling must be flagged" + ); + assert!(serialized_len(&huge).unwrap() > MAX_FRAME_BYTES); + } + + #[tokio::test] + async fn an_oversized_spawn_fails_the_caller_and_does_not_disconnect() { + // Mirror the serve-loop path without a socket: begin_spawn parks the + // request, the frame is measured oversized, and the parked pending is + // failed with InvalidRequest — nothing tracked, no disconnect signal. + use crate::control_plane::primary::{ + CommandError, PrimaryState, SpawnRequest, SpawnTarget, + }; + let mut primary = PrimaryState::new(); + let (tx, rx) = tokio::sync::oneshot::channel(); + let request = SpawnRequest { + delegation_id: "d-huge".into(), + target: SpawnTarget::by_name("worker-1"), + prompt: "x".repeat(MAX_FRAME_BYTES + 4096), + deadline: chrono::Utc::now() + chrono::Duration::seconds(60), + parent: None, + }; + let emission = primary.begin_spawn(1, request, tx); + let frame = JsonRpcRequest::new( + emission.rpc_id, + methods::DELEGATE, + Some(serde_json::to_value(&emission.params).unwrap()), + ); + // The guard's decision: too large. + assert!(!frame_within_limit(&frame).unwrap()); + // The guard's action: fail the parked pending, track nothing. + assert!(primary.fail_pending_invalid_request(emission.rpc_id, "too large".into())); + assert_eq!(primary.tracked_len(), 0); + let outcome = rx.await.unwrap(); + assert!(matches!(outcome, Err(CommandError::InvalidRequest(_)))); + } } diff --git a/crates/openab-core/src/control_plane/local.rs b/crates/openab-core/src/control_plane/local.rs new file mode 100644 index 000000000..81a7151e2 --- /dev/null +++ b/crates/openab-core/src/control_plane/local.rs @@ -0,0 +1,1500 @@ +//! Local IPC surface for driving the primary side of the control-plane client. +//! +//! An out-of-process caller — a CLI subcommand, an MCP tool, an operator +//! script — needs to spawn/await/cancel delegations and list peers, but the +//! socket to the CP is owned by one in-process serve loop. This module bridges +//! the two with a Unix-domain socket that speaks newline-delimited JSON: one +//! [`LocalRequest`] per line in, one [`LocalResponse`] per line out. +//! +//! ## Security +//! +//! The socket is a local trust boundary: anything that can connect to it can +//! initiate delegations under this runtime's CP identity. So the socket is +//! created **owner-only** (`0600`) inside an **owner-only** parent directory +//! (`0700`), and the default path lives under the invoking user's home. There +//! is no authentication beyond filesystem permissions — the same posture as a +//! Docker or SSH agent socket — which is why the permissions are set +//! explicitly rather than left to the umask. +//! +//! The runtime only tightens (`chmod 0700`) a **final** socket directory it +//! creates itself (e.g. `$HOME/.openab/run`). An **existing** parent directory +//! is never re-permissioned: it must already be a directory, owned by the +//! current uid, with no group- or world-accessible bits — otherwise the server +//! refuses to bind rather than silently narrowing someone else's directory. +//! Likewise, a leftover socket node is removed before bind only when it is a +//! genuine Unix socket owned by the current uid; an arbitrary file or a symlink +//! is refused, so the server can never be tricked into unlinking an unrelated +//! path. +//! +//! ## Portability +//! +//! Unix-domain sockets are the whole mechanism, so the server/client are +//! `#[cfg(unix)]`. Non-Unix targets get stubs that compile and fail loudly at +//! runtime, keeping `cargo check --target x86_64-pc-windows-gnu` green without +//! a second transport. + +use serde::{Deserialize, Serialize}; + +use crate::control_plane::primary::DelegationStatusWire; + +/// Environment override for the local socket path. +pub const SOCKET_ENV: &str = "OPENAB_AGENT_SOCKET"; + +/// Default socket path: `$OPENAB_AGENT_SOCKET` if set, else +/// `$HOME/.openab/run/agent.sock`. The socket lives in a dedicated `run` +/// subdirectory so the runtime can create and lock down that final directory +/// (`0700`) without ever weakening or re-permissioning `$HOME/.openab`, which +/// may hold other state and be shared with other tools. Returns an error only +/// when neither the override nor `HOME` is available — a headless environment +/// with no home is a misconfiguration the caller must resolve, not something to +/// guess around. +pub fn default_socket_path() -> anyhow::Result { + if let Some(explicit) = std::env::var_os(SOCKET_ENV) { + return Ok(std::path::PathBuf::from(explicit)); + } + let home = std::env::var_os("HOME") + .ok_or_else(|| anyhow::anyhow!("neither {SOCKET_ENV} nor HOME is set"))?; + Ok(std::path::PathBuf::from(home) + .join(".openab") + .join("run") + .join("agent.sock")) +} + +/// A spawn target as expressed by a local caller. Exactly one field must be +/// set; the server rejects a request that sets both or neither. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default)] +pub struct LocalTarget { + #[serde(default, skip_serializing_if = "Option::is_none")] + pub name: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub labels: Option>, +} + +/// One request line on the local socket. +/// +/// `#[serde(tag = "op")]` makes each variant a self-describing object, which is +/// what a CLI or MCP tool marshals by hand: `{"op":"spawn", ...}`. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] +#[serde(tag = "op", rename_all = "snake_case")] +pub enum LocalRequest { + /// Initiate a delegation. Returns an opaque handle token once the CP acks + /// routing. `deadline_secs` is relative to receipt; the server converts it + /// to the absolute deadline the wire requires. + Spawn { + delegation_id: String, + target: LocalTarget, + prompt: String, + deadline_secs: u64, + /// Opaque parent handle token, if this spawn is a child of a + /// delegation this runtime is serving. + #[serde(default, skip_serializing_if = "Option::is_none")] + parent: Option, + }, + /// Block until the named delegation reaches a terminal state. + Await { handle: String }, + /// Nonblocking status query: answers with the delegation's current state + /// (pending / running / terminal) without waiting. + Check { handle: String }, + /// Cancel an in-flight delegation by its opaque handle token. + Cancel { handle: String, reason: String }, + /// Namespace-scoped roster snapshot. + ListAgents, +} + +/// A single agent as reported to a local caller. A projection of the wire +/// `AgentSummary`, kept separate so the local protocol does not re-export the +/// CP crate's types to a CLI. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +pub struct LocalAgent { + pub name: String, + pub agent_type: String, + pub instance_id: String, + pub labels: std::collections::BTreeMap, + pub active_sessions: u32, + pub max_delegated_sessions: u32, +} + +/// One response line on the local socket. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] +#[serde(tag = "kind", rename_all = "snake_case")] +pub enum LocalResponse { + /// A spawn was admitted; `handle` is the opaque token to await/cancel with. + Spawned { handle: String, assigned_to: String }, + /// A `check` found the delegation delegated but not yet acked by the CP. + Pending, + /// A `check` found the delegation admitted and being served by a peer. + Running { assigned_to: String }, + /// A delegation reached a terminal state. + Terminal { + status: DelegationStatusWire, + #[serde(default, skip_serializing_if = "Option::is_none")] + result: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + error: Option, + }, + /// A cancel was accepted by the CP. + Cancelled, + /// The roster snapshot. + Agents { agents: Vec }, + /// The request could not be satisfied. + Error { message: String }, +} + +impl LocalResponse { + fn error(message: impl Into) -> Self { + LocalResponse::Error { + message: message.into(), + } + } +} + +#[cfg(unix)] +pub use unix_impl::{serve_local, LocalClient}; + +#[cfg(unix)] +mod unix_impl { + use super::*; + use std::collections::{HashMap, VecDeque}; + use std::os::unix::fs::{FileTypeExt, MetadataExt, PermissionsExt}; + use std::path::{Path, PathBuf}; + use std::sync::{Arc, Mutex}; + + use rand::RngCore; + use tokio::io::{AsyncBufRead, AsyncBufReadExt, AsyncWriteExt, BufReader}; + use tokio::net::{UnixListener, UnixStream}; + use tokio::sync::watch; + use tracing::{debug, info, warn}; + + use crate::control_plane::primary::{ + ClientCommand, CommandError, DelegationHandle, DelegationSnapshot, SpawnRequest, + SpawnTarget, + }; + use crate::control_plane::ControlPlaneHandle; + + /// Maximum number of opaque handle tokens retained in the registry. Mirrors + /// the CP's `max_inflight_delegations` default (ADR §Resource caps): a + /// local surface can never track more live delegations than the hub admits, + /// and a bound keeps a long-lived server from leaking tokens for + /// delegations whose callers walked away. + const MAX_HANDLE_TOKENS: usize = 4096; + + /// Number of random bytes behind each opaque token (256 bits, rendered as + /// 64 lowercase hex chars). Large enough that a local caller cannot guess + /// or fabricate a valid token for a delegation it did not spawn. + const TOKEN_BYTES: usize = 32; + + /// Server-side map from opaque token → the real [`DelegationHandle`]. + /// + /// The token handed to a local caller is an unstructured random string; it + /// encodes nothing about the delegation (no `admission`, no id), so a + /// caller cannot forge one or infer another delegation's handle from its + /// own. The map is shared across every connection to this server (an + /// `Arc>`), so a token minted on one connection resolves on + /// another — the same delegation is reachable regardless of which + /// connection awaits or cancels it. + /// + /// Bounded at [`MAX_HANDLE_TOKENS`] with deterministic oldest-first + /// eviction: `order` records insertion order, and an insert over the cap + /// evicts the front entry. Terminal delegations are not evicted eagerly — + /// only age evicts — so a `check`/`await` after completion still resolves + /// until the token ages out. + #[derive(Default)] + pub(super) struct HandleRegistry { + by_token: HashMap, + order: VecDeque, + } + + impl HandleRegistry { + /// Mint a fresh opaque token for `handle`, insert it, and evict the + /// oldest token if the map is over capacity. Returns the token. + fn insert(&mut self, handle: DelegationHandle) -> String { + let token = mint_token(); + self.by_token.insert(token.clone(), handle); + self.order.push_back(token.clone()); + while self.order.len() > MAX_HANDLE_TOKENS { + if let Some(oldest) = self.order.pop_front() { + self.by_token.remove(&oldest); + } + } + token + } + + /// Resolve a token to its handle, or `None` if it was never minted here + /// (a fabricated token) or has aged out. + fn resolve(&self, token: &str) -> Option { + self.by_token.get(token).cloned() + } + + #[cfg(test)] + fn len(&self) -> usize { + self.by_token.len() + } + } + + /// A shared, connection-independent handle registry. + pub(super) type SharedRegistry = Arc>; + + /// Generate one opaque token: [`TOKEN_BYTES`] of CSPRNG output as hex. + fn mint_token() -> String { + let mut bytes = [0u8; TOKEN_BYTES]; + rand::thread_rng().fill_bytes(&mut bytes); + let mut s = String::with_capacity(TOKEN_BYTES * 2); + for b in bytes { + use std::fmt::Write as _; + let _ = write!(s, "{b:02x}"); + } + s + } + + /// Bind the local UDS server at `path`, serving requests by translating + /// them into [`ClientCommand`]s on `cp`. Runs until `shutdown` flips. + /// + /// The socket's final directory is created owner-only when the runtime + /// makes it; an existing parent is verified (directory, current uid, no + /// group/world bits) but never re-permissioned. A stale socket node from a + /// prior run is removed first — but only if it is a Unix socket owned by + /// the current uid, so the server never unlinks an arbitrary file or a + /// symlink someone planted at the path. + pub async fn serve_local( + path: PathBuf, + cp: ControlPlaneHandle, + mut shutdown: watch::Receiver, + ) -> anyhow::Result<()> { + prepare_socket_dir(&path)?; + remove_stale_socket(&path)?; + let listener = UnixListener::bind(&path) + .map_err(|e| anyhow::anyhow!("binding local socket {path:?}: {e}"))?; + // Owner-only on the socket node itself, not merely via the directory. + std::fs::set_permissions(&path, std::fs::Permissions::from_mode(0o600)) + .map_err(|e| anyhow::anyhow!("securing local socket {path:?}: {e}"))?; + info!(socket = %path.display(), "control-plane local IPC server listening"); + + // One registry for the whole server: tokens minted on any connection + // resolve on any other. + let registry: SharedRegistry = Arc::new(Mutex::new(HandleRegistry::default())); + + loop { + tokio::select! { + _ = shutdown.changed() => { + info!("control-plane local IPC server shutting down"); + break; + } + accepted = listener.accept() => { + match accepted { + Ok((stream, _addr)) => { + let cp = cp.clone(); + let registry = Arc::clone(®istry); + tokio::spawn(async move { + if let Err(e) = handle_conn(stream, cp, registry).await { + debug!(error = %format!("{e:#}"), "local IPC connection ended with error"); + } + }); + } + Err(e) => { + warn!(error = %e, "local IPC accept failed"); + } + } + } + } + } + // Best-effort cleanup so the next boot does not trip over the node. + let _ = std::fs::remove_file(&path); + Ok(()) + } + + /// Prepare the socket's parent directory. + /// + /// Two cases, deliberately asymmetric: + /// + /// * **The final directory does not exist.** The runtime creates *just that + /// directory* (not any ancestor) with mode `0700`, so a socket the + /// runtime owns lives in a directory the runtime owns. `create_dir` (not + /// `create_dir_all`) is used so a missing ancestor is an error, not a + /// silent chain of new directories — and the mode is applied by an + /// explicit `chmod` after creation so it is exact regardless of umask. + /// + /// * **The final directory already exists.** It is *verified*, never + /// modified: it must be a real directory (not a symlink to one — checked + /// via `symlink_metadata`), owned by the current uid, and carry no + /// group- or world-accessible bits. If any check fails the server + /// refuses to bind. The runtime never `chmod`s a directory it did not + /// create, because narrowing a pre-existing directory could strip access + /// another tool legitimately relies on, and a directory it does not own + /// is not its to tighten. + /// + /// Override-path ancestors are never created. For the built-in default on + /// a fresh volume, the runtime may create its own `$HOME/.openab` state + /// directory and then the final `run` directory, both owner-only. + fn prepare_socket_dir(path: &Path) -> anyhow::Result<()> { + let managed_state_dir = std::env::var_os("HOME") + .map(PathBuf::from) + .map(|home| home.join(".openab")); + prepare_socket_dir_with_managed_parent(path, managed_state_dir.as_deref()) + } + + fn prepare_socket_dir_with_managed_parent( + path: &Path, + managed_state_dir: Option<&Path>, + ) -> anyhow::Result<()> { + let Some(dir) = path.parent() else { + return Ok(()); + }; + if dir.as_os_str().is_empty() { + return Ok(()); + } + match std::fs::symlink_metadata(dir) { + Ok(meta) => { + // Existing entry: verify, never modify. + if !meta.file_type().is_dir() { + return Err(anyhow::anyhow!( + "socket parent {dir:?} exists but is not a directory (a symlink or file \ + at this path is refused rather than trusted)" + )); + } + let uid = unsafe { libc::getuid() }; + if meta.uid() != uid { + return Err(anyhow::anyhow!( + "socket parent {dir:?} is owned by uid {} not the current uid {uid}; \ + refusing to trust a directory this runtime does not own", + meta.uid() + )); + } + let mode = meta.permissions().mode() & 0o777; + if mode & 0o077 != 0 { + return Err(anyhow::anyhow!( + "socket parent {dir:?} has group/world bits set (mode {mode:#o}); \ + refusing to bind, and refusing to chmod a directory this runtime did \ + not create" + )); + } + Ok(()) + } + Err(e) if e.kind() == std::io::ErrorKind::NotFound => { + // The built-in default may be the first OpenAB component on a + // fresh volume. Create `$HOME/.openab` only when it is the + // explicitly identified managed state directory; arbitrary + // override ancestors remain the deployment's responsibility. + if let Some(parent) = dir.parent() { + if !parent.exists() { + anyhow::ensure!( + managed_state_dir.is_some_and(|managed| managed == parent), + "socket parent ancestor {parent:?} does not exist; refusing to create ancestors for an override path" + ); + std::fs::create_dir(parent).map_err(|e| { + anyhow::anyhow!("creating managed OpenAB state dir {parent:?}: {e}") + })?; + std::fs::set_permissions(parent, std::fs::Permissions::from_mode(0o700)) + .map_err(|e| { + anyhow::anyhow!("securing managed OpenAB state dir {parent:?}: {e}") + })?; + } + } + // Final dedicated `run` directory is ours to create and lock down. + std::fs::create_dir(dir) + .map_err(|e| anyhow::anyhow!("creating socket dir {dir:?}: {e}"))?; + std::fs::set_permissions(dir, std::fs::Permissions::from_mode(0o700)) + .map_err(|e| anyhow::anyhow!("securing socket dir {dir:?}: {e}"))?; + Ok(()) + } + Err(e) => Err(anyhow::anyhow!("inspecting socket dir {dir:?}: {e}")), + } + } + + /// Remove a leftover socket node so `bind` does not fail with `EADDRINUSE` + /// — but only when it is safe to do so. The node must exist as a genuine + /// Unix socket owned by the current uid; a symlink, a regular file, or a + /// node owned by someone else is refused rather than unlinked, so a hostile + /// or accidental entry planted at the socket path can never trick the + /// runtime into deleting an unrelated file. + fn remove_stale_socket(path: &Path) -> anyhow::Result<()> { + let meta = match std::fs::symlink_metadata(path) { + Ok(m) => m, + Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Ok(()), + Err(e) => return Err(anyhow::anyhow!("inspecting socket path {path:?}: {e}")), + }; + let ft = meta.file_type(); + if ft.is_symlink() { + return Err(anyhow::anyhow!( + "refusing to remove {path:?}: it is a symlink, not a socket this runtime created" + )); + } + if !ft.is_socket() { + return Err(anyhow::anyhow!( + "refusing to remove {path:?}: it is not a Unix socket (a stray regular file must \ + be cleared by hand rather than unlinked automatically)" + )); + } + let uid = unsafe { libc::getuid() }; + if meta.uid() != uid { + return Err(anyhow::anyhow!( + "refusing to remove socket {path:?}: it is owned by uid {} not the current uid \ + {uid}", + meta.uid() + )); + } + std::fs::remove_file(path) + .map_err(|e| anyhow::anyhow!("removing stale socket {path:?}: {e}")) + } + + /// Maximum length, in bytes, of a single newline-delimited protocol line + /// (request or response). A local caller (or a peer answering the client) + /// that never sends a newline must not be able to drive the reader's buffer + /// to unbounded growth, so the reader stops at this cap and reports an + /// error instead of accumulating. 1 MiB is far larger than any legitimate + /// request/response (handles are 64 hex chars, prompts are bounded by the + /// CP) while still bounding memory per connection. + const MAX_LINE_BYTES: usize = 1024 * 1024; + + /// Read one newline-terminated line from `reader`, bounded to + /// [`MAX_LINE_BYTES`]. + /// + /// Returns: + /// * `Ok(Some(line))` with the trailing `\n` (and any `\r`) trimmed, + /// * `Ok(None)` at clean EOF with no partial line buffered, + /// * `Err(..)` if the line would exceed the cap or the stream errors. + /// + /// Bytes are copied from the buffered reader only after checking the cap. + /// A newline consumes exactly one line and leaves subsequent buffered bytes + /// untouched for the next call, so pipelined requests are preserved. + async fn read_line_bounded(reader: &mut R) -> anyhow::Result> + where + R: AsyncBufRead + Unpin, + { + let mut acc: Vec = Vec::new(); + loop { + let available = reader.fill_buf().await?; + if available.is_empty() { + if acc.is_empty() { + return Ok(None); + } + let line = String::from_utf8(acc) + .map_err(|e| anyhow::anyhow!("line is not valid UTF-8: {e}"))?; + return Ok(Some(trim_line_ending(line))); + } + + if let Some(newline) = available.iter().position(|&b| b == b'\n') { + if acc.len() + newline > MAX_LINE_BYTES { + return Err(line_too_long_err()); + } + acc.extend_from_slice(&available[..newline]); + // Consume exactly this line, preserving every byte after the + // newline for the next call (pipelined requests are valid). + reader.consume(newline + 1); + let line = String::from_utf8(acc) + .map_err(|e| anyhow::anyhow!("line is not valid UTF-8: {e}"))?; + return Ok(Some(trim_line_ending(line))); + } + + let len = available.len(); + if acc.len() + len > MAX_LINE_BYTES { + return Err(line_too_long_err()); + } + acc.extend_from_slice(available); + reader.consume(len); + } + } + + fn line_too_long_err() -> anyhow::Error { + anyhow::anyhow!("protocol line exceeds the {MAX_LINE_BYTES}-byte limit") + } + + /// Strip a trailing line ending: a `\n` (optionally preceded by `\r`), or a + /// lone trailing `\r` when the `\n` was already consumed as the line + /// delimiter. Idempotent on already-trimmed input. + fn trim_line_ending(mut line: String) -> String { + if line.ends_with('\n') { + line.pop(); + } + if line.ends_with('\r') { + line.pop(); + } + line + } + + /// One client connection: read request lines, answer each with one + /// response line. A malformed line is answered with an `Error` rather than + /// dropping the connection, so a CLI gets a diagnostic. + async fn handle_conn( + stream: UnixStream, + cp: ControlPlaneHandle, + registry: SharedRegistry, + ) -> anyhow::Result<()> { + let (read_half, mut write_half) = stream.into_split(); + let mut reader = BufReader::new(read_half); + while let Some(line) = read_line_bounded(&mut reader).await? { + if line.trim().is_empty() { + continue; + } + let response = match serde_json::from_str::(&line) { + Ok(request) => dispatch(request, &cp, ®istry).await, + Err(e) => LocalResponse::error(format!("malformed request: {e}")), + }; + let mut encoded = serde_json::to_string(&response)?; + encoded.push('\n'); + write_half.write_all(encoded.as_bytes()).await?; + write_half.flush().await?; + } + Ok(()) + } + + /// Minimum accepted `deadline_secs`. Zero (or a negative-by-underflow + /// value) would spawn a delegation that is already past due, which the CP + /// would sweep on arrival — a caller error worth rejecting locally. + const MIN_DEADLINE_SECS: u64 = 1; + /// Maximum accepted `deadline_secs`. Caps a delegation's absolute lifetime + /// to the ADR's sensible ceiling (30 minutes); a larger value is *rejected* + /// — not clamped — so a caller learns its request was out of range rather + /// than silently getting a different deadline than it asked for. + const MAX_DEADLINE_SECS: u64 = 1800; + + /// Translate one request into a command, submit it, and map the reply. + async fn dispatch( + request: LocalRequest, + cp: &ControlPlaneHandle, + registry: &SharedRegistry, + ) -> LocalResponse { + match request { + LocalRequest::Spawn { + delegation_id, + target, + prompt, + deadline_secs, + parent, + } => { + let target = match into_spawn_target(target) { + Ok(t) => t, + Err(e) => return LocalResponse::error(e), + }; + // Reject a deadline outside the accepted range outright; both + // ends are validated before any frame is sent. An over-long + // value is refused rather than clamped, so the caller is never + // silently given a deadline it did not request. + if !(MIN_DEADLINE_SECS..=MAX_DEADLINE_SECS).contains(&deadline_secs) { + return LocalResponse::error(format!( + "deadline_secs must be within {MIN_DEADLINE_SECS}..={MAX_DEADLINE_SECS}" + )); + } + let parent = match parent { + Some(token) => match registry.lock().unwrap().resolve(&token) { + Some(h) => Some(h), + None => return LocalResponse::error("unknown parent handle token"), + }, + None => None, + }; + let deadline = chrono::Utc::now() + chrono::Duration::seconds(deadline_secs as i64); + let request = SpawnRequest { + delegation_id, + target, + prompt, + deadline, + parent, + }; + let (tx, rx) = tokio::sync::oneshot::channel(); + if cp + .submit(ClientCommand::Spawn { request, reply: tx }) + .await + .is_err() + { + return LocalResponse::error(CommandError::NotConnected.to_string()); + } + match rx.await { + Ok(Ok(ack)) => { + // Mint an opaque token for the handle only now that the + // spawn succeeded, and report the CP-assigned peer. + let token = registry.lock().unwrap().insert(ack.handle); + LocalResponse::Spawned { + handle: token, + assigned_to: ack.assigned_to, + } + } + Ok(Err(e)) => LocalResponse::error(e.to_string()), + Err(_) => LocalResponse::error("control-plane client dropped the reply"), + } + } + LocalRequest::Await { handle } => { + let Some(handle) = registry.lock().unwrap().resolve(&handle) else { + return LocalResponse::error("unknown handle token"); + }; + let (tx, rx) = tokio::sync::oneshot::channel(); + if cp + .submit(ClientCommand::Await { handle, reply: tx }) + .await + .is_err() + { + return LocalResponse::error(CommandError::NotConnected.to_string()); + } + match rx.await { + Ok(Ok(outcome)) => LocalResponse::Terminal { + status: outcome.status.into(), + result: outcome.result, + error: outcome.error, + }, + Ok(Err(e)) => LocalResponse::error(e.to_string()), + Err(_) => LocalResponse::error("control-plane client dropped the reply"), + } + } + LocalRequest::Check { handle } => { + let Some(handle) = registry.lock().unwrap().resolve(&handle) else { + return LocalResponse::error("unknown handle token"); + }; + let (tx, rx) = tokio::sync::oneshot::channel(); + if cp + .submit(ClientCommand::Check { handle, reply: tx }) + .await + .is_err() + { + return LocalResponse::error(CommandError::NotConnected.to_string()); + } + match rx.await { + Ok(Ok(DelegationSnapshot::Pending)) => LocalResponse::Pending, + Ok(Ok(DelegationSnapshot::Running { assigned_to })) => { + LocalResponse::Running { assigned_to } + } + Ok(Ok(DelegationSnapshot::Terminal(outcome))) => LocalResponse::Terminal { + status: outcome.status.into(), + result: outcome.result, + error: outcome.error, + }, + Ok(Err(e)) => LocalResponse::error(e.to_string()), + Err(_) => LocalResponse::error("control-plane client dropped the reply"), + } + } + LocalRequest::Cancel { handle, reason } => { + let Some(handle) = registry.lock().unwrap().resolve(&handle) else { + return LocalResponse::error("unknown handle token"); + }; + let (tx, rx) = tokio::sync::oneshot::channel(); + if cp + .submit(ClientCommand::Cancel { + handle, + reason, + reply: tx, + }) + .await + .is_err() + { + return LocalResponse::error(CommandError::NotConnected.to_string()); + } + match rx.await { + Ok(Ok(())) => LocalResponse::Cancelled, + Ok(Err(e)) => LocalResponse::error(e.to_string()), + Err(_) => LocalResponse::error("control-plane client dropped the reply"), + } + } + LocalRequest::ListAgents => { + let (tx, rx) = tokio::sync::oneshot::channel(); + if cp + .submit(ClientCommand::ListAgents { reply: tx }) + .await + .is_err() + { + return LocalResponse::error(CommandError::NotConnected.to_string()); + } + match rx.await { + Ok(Ok(agents)) => LocalResponse::Agents { + agents: agents + .into_iter() + .map(|a| LocalAgent { + name: a.name, + agent_type: a.agent_type.to_string(), + instance_id: a.instance_id, + labels: a.labels, + active_sessions: a.active_sessions, + max_delegated_sessions: a.max_delegated_sessions, + }) + .collect(), + }, + Ok(Err(e)) => LocalResponse::error(e.to_string()), + Err(_) => LocalResponse::error("control-plane client dropped the reply"), + } + } + } + } + + fn into_spawn_target(target: LocalTarget) -> Result { + match (target.name, target.labels) { + (Some(name), None) => { + if name.is_empty() { + Err("target.name is empty".into()) + } else { + Ok(SpawnTarget::by_name(name)) + } + } + (None, Some(labels)) => { + if labels.is_empty() { + Err("target.labels is empty".into()) + } else { + Ok(SpawnTarget::by_labels(labels)) + } + } + (Some(_), Some(_)) => { + Err("target must set exactly one of name or labels, not both".into()) + } + (None, None) => Err("target must set exactly one of name or labels".into()), + } + } + + /// A minimal client for the local socket, for a CLI or MCP tool linking + /// this crate. Opens one connection per call: the protocol is + /// request/response and the socket is cheap, so there is no pool. + pub struct LocalClient { + path: PathBuf, + } + + impl LocalClient { + /// Connect against an explicit path. + pub fn new(path: impl Into) -> Self { + Self { path: path.into() } + } + + /// Connect against [`default_socket_path`]. + pub fn from_default() -> anyhow::Result { + Ok(Self { + path: default_socket_path()?, + }) + } + + /// Send one request and read one response. + pub async fn request(&self, request: &LocalRequest) -> anyhow::Result { + let stream = UnixStream::connect(&self.path) + .await + .map_err(|e| anyhow::anyhow!("connecting to local socket {:?}: {e}", self.path))?; + let (read_half, mut write_half) = stream.into_split(); + let mut encoded = serde_json::to_string(request)?; + encoded.push('\n'); + write_half.write_all(encoded.as_bytes()).await?; + write_half.flush().await?; + let mut reader = BufReader::new(read_half); + match read_line_bounded(&mut reader).await? { + Some(line) => Ok(serde_json::from_str(&line)?), + None => Err(anyhow::anyhow!("local socket closed without a response")), + } + } + } + + #[cfg(test)] + mod tests { + use crate::control_plane::primary::{DelegationOutcome, SpawnAck}; + + #[test] + fn a_minted_token_is_opaque_hex_and_encodes_nothing_about_the_handle() { + let mut reg = HandleRegistry::default(); + let handle = DelegationHandle::new("discord:12345:thread", 42); + let token = reg.insert(handle.clone()); + // 32 random bytes → 64 lowercase hex chars. + assert_eq!(token.len(), TOKEN_BYTES * 2); + assert!(token + .chars() + .all(|c| c.is_ascii_hexdigit() && !c.is_ascii_uppercase())); + // The token is not derived from the handle: it is neither the old + // `admission:delegation_id` encoding nor does it embed the id + // verbatim (a random hex string may incidentally contain a short + // decimal like "42", so we test the structured forms, not digits). + assert_ne!( + token, + format!("{}:{}", handle.admission(), handle.delegation_id()) + ); + assert_ne!(token, handle.delegation_id()); + assert!(!token.contains("discord")); + assert!(!token.contains("thread")); + assert!(!token.contains(':')); + // It resolves back to the exact handle. + assert_eq!(reg.resolve(&token), Some(handle)); + } + + #[test] + fn two_tokens_for_the_same_handle_are_distinct_and_unguessable() { + let mut reg = HandleRegistry::default(); + let handle = DelegationHandle::new("d-1", 7); + let a = reg.insert(handle.clone()); + let b = reg.insert(handle); + assert_ne!(a, b, "each mint draws fresh randomness"); + } + + #[test] + fn a_fabricated_token_does_not_resolve() { + let mut reg = HandleRegistry::default(); + let real = reg.insert(DelegationHandle::new("d-1", 1)); + // A plausible-looking but never-minted token is rejected. + assert!(reg.resolve(&"0".repeat(TOKEN_BYTES * 2)).is_none()); + // The old fabricatable `admission:id` shape resolves to nothing. + assert!(reg.resolve("1:d-1").is_none()); + assert!(reg.resolve("").is_none()); + // The genuine one still resolves. + assert!(reg.resolve(&real).is_some()); + } + + #[test] + fn the_registry_is_bounded_and_evicts_oldest_first() { + let mut reg = HandleRegistry::default(); + // Fill to capacity, remembering the very first token. + let first = reg.insert(DelegationHandle::new("d-0", 0)); + for i in 1..MAX_HANDLE_TOKENS { + reg.insert(DelegationHandle::new(format!("d-{i}"), i as u64)); + } + assert_eq!(reg.len(), MAX_HANDLE_TOKENS); + assert!(reg.resolve(&first).is_some(), "still present at capacity"); + + // One more insert evicts the oldest (the first) token, not a newer one. + let newest = reg.insert(DelegationHandle::new("d-overflow", 9999)); + assert_eq!(reg.len(), MAX_HANDLE_TOKENS, "cap holds"); + assert!(reg.resolve(&first).is_none(), "oldest was evicted"); + assert!(reg.resolve(&newest).is_some(), "newest survives"); + } + + /// A stub CP loop that answers each command with a canned reply, so the + /// full local dispatch path can be exercised without a socket to a hub. + fn stub_cp() -> (ControlPlaneHandle, tokio::task::JoinHandle<()>) { + let (handle, mut rx) = crate::control_plane::client::test_support::handle_and_rx(); + let task = tokio::spawn(async move { + while let Some(command) = rx.recv().await { + match command { + ClientCommand::Spawn { request, reply } => { + let _ = reply.send(Ok(SpawnAck { + handle: DelegationHandle::new(request.delegation_id, 5), + assigned_to: "prod/worker-1".into(), + })); + } + ClientCommand::Check { reply, .. } => { + let _ = reply.send(Ok(DelegationSnapshot::Running { + assigned_to: "prod/worker-1".into(), + })); + } + ClientCommand::Await { reply, .. } => { + let _ = reply.send(Ok(DelegationOutcome { + status: openab_cp::proto::DelegationStatus::Completed, + result: Some("done".into()), + error: None, + })); + } + ClientCommand::Cancel { reply, .. } => { + let _ = reply.send(Ok(())); + } + ClientCommand::ListAgents { reply } => { + let _ = reply.send(Ok(vec![])); + } + } + } + }); + (handle, task) + } + + #[tokio::test] + async fn spawn_reports_the_assigned_peer_not_an_empty_string() { + let (cp, task) = stub_cp(); + let registry: SharedRegistry = Arc::new(Mutex::new(HandleRegistry::default())); + let resp = dispatch( + LocalRequest::Spawn { + delegation_id: "d-1".into(), + target: LocalTarget { + name: Some("worker-1".into()), + labels: None, + }, + prompt: "hi".into(), + deadline_secs: 60, + parent: None, + }, + &cp, + ®istry, + ) + .await; + match resp { + LocalResponse::Spawned { + handle, + assigned_to, + } => { + assert_eq!(assigned_to, "prod/worker-1", "assigned_to must be reported"); + // The handle is an opaque token, not `admission:id`. + assert_eq!(handle.len(), TOKEN_BYTES * 2); + assert!(!handle.contains(':')); + // And it resolves in the shared registry for a later check. + assert!(registry.lock().unwrap().resolve(&handle).is_some()); + } + other => panic!("expected Spawned, got {other:?}"), + } + drop(cp); + let _ = task.await; + } + + #[tokio::test] + async fn check_after_spawn_reports_running_without_blocking() { + let (cp, task) = stub_cp(); + let registry: SharedRegistry = Arc::new(Mutex::new(HandleRegistry::default())); + let spawned = dispatch( + LocalRequest::Spawn { + delegation_id: "d-1".into(), + target: LocalTarget { + name: Some("worker-1".into()), + labels: None, + }, + prompt: "hi".into(), + deadline_secs: 60, + parent: None, + }, + &cp, + ®istry, + ) + .await; + let LocalResponse::Spawned { handle, .. } = spawned else { + panic!("spawn failed"); + }; + let checked = dispatch(LocalRequest::Check { handle }, &cp, ®istry).await; + assert_eq!( + checked, + LocalResponse::Running { + assigned_to: "prod/worker-1".into() + } + ); + drop(cp); + let _ = task.await; + } + + #[tokio::test] + async fn check_of_a_fabricated_token_is_rejected_before_reaching_the_cp() { + let (cp, task) = stub_cp(); + let registry: SharedRegistry = Arc::new(Mutex::new(HandleRegistry::default())); + let resp = dispatch( + LocalRequest::Check { + handle: "deadbeef".into(), + }, + &cp, + ®istry, + ) + .await; + assert!(matches!(resp, LocalResponse::Error { .. })); + drop(cp); + let _ = task.await; + } + + #[tokio::test] + async fn a_zero_deadline_is_rejected_and_never_reaches_the_cp() { + let (cp, task) = stub_cp(); + let registry: SharedRegistry = Arc::new(Mutex::new(HandleRegistry::default())); + let resp = dispatch( + LocalRequest::Spawn { + delegation_id: "d-1".into(), + target: LocalTarget { + name: Some("worker-1".into()), + labels: None, + }, + prompt: "hi".into(), + deadline_secs: 0, + parent: None, + }, + &cp, + ®istry, + ) + .await; + match resp { + LocalResponse::Error { message } => assert!(message.contains("deadline_secs")), + other => panic!("expected an error, got {other:?}"), + } + // Nothing was tracked, since the spawn was rejected locally. + assert_eq!(registry.lock().unwrap().len(), 0); + drop(cp); + let _ = task.await; + } + + #[tokio::test] + async fn an_over_long_deadline_is_rejected_not_clamped() { + let (cp, task) = stub_cp(); + let registry: SharedRegistry = Arc::new(Mutex::new(HandleRegistry::default())); + // One second past the ceiling: a clamp would silently accept it; we + // require an explicit rejection instead. + let resp = dispatch( + LocalRequest::Spawn { + delegation_id: "d-1".into(), + target: LocalTarget { + name: Some("worker-1".into()), + labels: None, + }, + prompt: "hi".into(), + deadline_secs: MAX_DEADLINE_SECS + 1, + parent: None, + }, + &cp, + ®istry, + ) + .await; + match resp { + LocalResponse::Error { message } => assert!( + message.contains("deadline_secs"), + "over-long deadline must be rejected, got {message:?}" + ), + other => panic!("expected an error, got {other:?}"), + } + // Nothing was tracked, since the spawn was rejected locally. + assert_eq!(registry.lock().unwrap().len(), 0); + // The boundary value is still accepted. + let ok = dispatch( + LocalRequest::Spawn { + delegation_id: "d-2".into(), + target: LocalTarget { + name: Some("worker-1".into()), + labels: None, + }, + prompt: "hi".into(), + deadline_secs: MAX_DEADLINE_SECS, + parent: None, + }, + &cp, + ®istry, + ) + .await; + assert!(matches!(ok, LocalResponse::Spawned { .. })); + drop(cp); + let _ = task.await; + } + + use super::*; + + #[test] + fn spawn_target_conversion_enforces_exactly_one() { + assert!(into_spawn_target(LocalTarget { + name: Some("w1".into()), + labels: None + }) + .is_ok()); + let mut labels = std::collections::BTreeMap::new(); + labels.insert("k".into(), "v".into()); + assert!(into_spawn_target(LocalTarget { + name: None, + labels: Some(labels.clone()) + }) + .is_ok()); + assert!(into_spawn_target(LocalTarget { + name: Some("w1".into()), + labels: Some(labels) + }) + .is_err()); + assert!(into_spawn_target(LocalTarget::default()).is_err()); + assert!(into_spawn_target(LocalTarget { + name: Some(String::new()), + labels: None + }) + .is_err()); + } + + #[tokio::test] + async fn round_trips_a_request_and_response_over_a_real_socket() { + // Bind a bare listener that echoes a canned response, exercising the + // client's framing against a real UDS without a full CP. + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("agent.sock"); + let listener = UnixListener::bind(&path).unwrap(); + let server = tokio::spawn(async move { + let (stream, _) = listener.accept().await.unwrap(); + let (read_half, mut write_half) = stream.into_split(); + let mut reader = BufReader::new(read_half); + let line = read_line_bounded(&mut reader).await.unwrap().unwrap(); + let req: LocalRequest = serde_json::from_str(&line).unwrap(); + assert!(matches!(req, LocalRequest::ListAgents)); + let resp = LocalResponse::Agents { agents: vec![] }; + let mut encoded = serde_json::to_string(&resp).unwrap(); + encoded.push('\n'); + write_half.write_all(encoded.as_bytes()).await.unwrap(); + write_half.flush().await.unwrap(); + }); + let client = LocalClient::new(path); + let resp = client.request(&LocalRequest::ListAgents).await.unwrap(); + assert_eq!(resp, LocalResponse::Agents { agents: vec![] }); + server.await.unwrap(); + } + + #[tokio::test] + async fn the_server_secures_the_socket_and_its_directory_owner_only() { + use tokio::sync::watch; + let dir = tempfile::tempdir().unwrap(); + let sock = dir.path().join("nested").join("agent.sock"); + // A dummy CP handle: submissions will be rejected NotConnected, but + // the test only checks the socket's permissions, not a round trip. + let (client, _rx) = crate::control_plane::client::test_support::handle_and_rx(); + let (tx, rx) = watch::channel(false); + let server_path = sock.clone(); + let server = tokio::spawn(async move { serve_local(server_path, client, rx).await }); + // Wait for the socket to appear. + for _ in 0..100 { + if sock.exists() { + break; + } + tokio::time::sleep(std::time::Duration::from_millis(10)).await; + } + assert!(sock.exists(), "server did not create the socket"); + let sock_mode = std::fs::metadata(&sock).unwrap().permissions().mode() & 0o777; + assert_eq!(sock_mode, 0o600, "socket must be owner-only"); + let dir_mode = std::fs::metadata(sock.parent().unwrap()) + .unwrap() + .permissions() + .mode() + & 0o777; + assert_eq!(dir_mode, 0o700, "socket directory must be owner-only"); + let _ = tx.send(true); + let _ = server.await; + } + + // -- prepare_socket_dir / remove_stale_socket / bounded reader -- + + /// The default `/tmp`-style layout regression: pointing the socket at a + /// directory the runtime did NOT create (here, a pre-existing dir that + /// carries group/world bits, the way `/tmp` itself does at `01777`) is + /// refused, and the directory's mode is left exactly as it was — the + /// runtime never chmods a parent it did not create. + #[test] + fn an_existing_insecure_parent_is_refused_and_never_chmodded() { + let root = tempfile::tempdir().unwrap(); + let dir = root.path().join("shared"); + std::fs::create_dir(&dir).unwrap(); + // World/group accessible, like /tmp (minus the sticky bit, which is + // irrelevant to the group/world check). + std::fs::set_permissions(&dir, std::fs::Permissions::from_mode(0o777)).unwrap(); + let before = std::fs::metadata(&dir).unwrap().permissions().mode() & 0o777; + + let sock = dir.join("agent.sock"); + let err = prepare_socket_dir(&sock).expect_err("insecure parent must be refused"); + assert!( + format!("{err:#}").contains("group/world"), + "error should explain the group/world bits, got {err:#}" + ); + + let after = std::fs::metadata(&dir).unwrap().permissions().mode() & 0o777; + assert_eq!(before, after, "the directory must NOT have been chmodded"); + assert_eq!(after, 0o777, "the directory retains its original mode"); + } + + /// A parent that is a plain file (or a symlink) is refused rather than + /// trusted — `symlink_metadata` is used so a symlink-to-a-dir does not + /// slip through. + #[test] + fn an_existing_parent_that_is_not_a_directory_is_refused() { + let root = tempfile::tempdir().unwrap(); + let not_a_dir = root.path().join("parent"); + std::fs::write(¬_a_dir, b"i am a file").unwrap(); + let sock = not_a_dir.join("agent.sock"); + let err = prepare_socket_dir(&sock).expect_err("non-directory parent must be refused"); + assert!( + format!("{err:#}").contains("not a directory"), + "got {err:#}" + ); + } + + /// An owned, `0700`, real directory that already exists is accepted and + /// left untouched — the common warm-restart case where the runtime's + /// own `run` dir survives from a prior boot. + #[test] + fn an_owned_secure_existing_dir_is_accepted_untouched() { + let root = tempfile::tempdir().unwrap(); + let dir = root.path().join("run"); + std::fs::create_dir(&dir).unwrap(); + std::fs::set_permissions(&dir, std::fs::Permissions::from_mode(0o700)).unwrap(); + let sock = dir.join("agent.sock"); + prepare_socket_dir(&sock).expect("owned 0700 dir must be accepted"); + let mode = std::fs::metadata(&dir).unwrap().permissions().mode() & 0o777; + assert_eq!(mode, 0o700, "an already-secure dir is left exactly as-is"); + } + + /// A missing final directory is created by the runtime at exactly + /// `0700`, regardless of the umask in effect. + #[test] + fn a_missing_final_dir_is_created_owner_only() { + let root = tempfile::tempdir().unwrap(); + let dir = root.path().join("run"); // does not exist yet + let sock = dir.join("agent.sock"); + prepare_socket_dir(&sock).expect("missing final dir should be created"); + let meta = std::fs::metadata(&dir).unwrap(); + assert!(meta.is_dir()); + assert_eq!( + meta.permissions().mode() & 0o777, + 0o700, + "runtime-created dir must be 0700" + ); + } + + #[test] + fn fresh_default_home_creates_only_managed_state_and_run_dirs() { + let root = tempfile::tempdir().unwrap(); + let home = root.path().join("home"); + std::fs::create_dir(&home).unwrap(); + let state = home.join(".openab"); + let socket = state.join("run").join("agent.sock"); + + prepare_socket_dir_with_managed_parent(&socket, Some(&state)).unwrap(); + + for dir in [&state, &state.join("run")] { + let meta = std::fs::metadata(dir).unwrap(); + assert!(meta.is_dir()); + assert_eq!(meta.permissions().mode() & 0o777, 0o700); + } + } + + /// A stale entry at the socket path that is a regular file — not a Unix + /// socket — is refused rather than unlinked, so the runtime can never + /// be tricked into deleting an unrelated file someone planted there. + #[test] + fn a_stale_regular_file_at_the_socket_path_is_refused() { + let root = tempfile::tempdir().unwrap(); + let sock = root.path().join("agent.sock"); + std::fs::write(&sock, b"definitely not a socket").unwrap(); + let err = remove_stale_socket(&sock).expect_err("a regular file must not be unlinked"); + assert!( + format!("{err:#}").contains("not a Unix socket"), + "got {err:#}" + ); + // And it is still there — we refused, we did not delete. + assert!(sock.exists(), "the file must be left in place"); + } + + /// A symlink at the socket path is refused before any unlink, so a + /// planted symlink cannot redirect the removal at an arbitrary target. + #[test] + fn a_symlink_at_the_socket_path_is_refused() { + let root = tempfile::tempdir().unwrap(); + let target = root.path().join("victim"); + std::fs::write(&target, b"do not delete me").unwrap(); + let link = root.path().join("agent.sock"); + std::os::unix::fs::symlink(&target, &link).unwrap(); + let err = remove_stale_socket(&link).expect_err("a symlink must not be removed"); + assert!(format!("{err:#}").contains("symlink"), "got {err:#}"); + assert!(target.exists(), "the symlink's target must be untouched"); + assert!( + std::fs::symlink_metadata(&link).is_ok(), + "the symlink itself must be left in place" + ); + } + + /// A genuine leftover Unix socket owned by the current uid IS removed, + /// so a normal restart does not trip over `EADDRINUSE`. + #[tokio::test] + async fn a_genuine_owned_stale_socket_is_removed() { + let root = tempfile::tempdir().unwrap(); + let sock = root.path().join("agent.sock"); + let listener = UnixListener::bind(&sock).unwrap(); + drop(listener); // the socket node lingers on disk + assert!(std::fs::symlink_metadata(&sock) + .unwrap() + .file_type() + .is_socket()); + remove_stale_socket(&sock).expect("an owned stale socket should be removed"); + assert!( + std::fs::symlink_metadata(&sock).is_err(), + "the stale socket node should be gone" + ); + } + + /// A missing socket path is a no-op, not an error. + #[test] + fn remove_stale_socket_is_a_noop_when_absent() { + let root = tempfile::tempdir().unwrap(); + let sock = root.path().join("agent.sock"); + remove_stale_socket(&sock).expect("absent path is a clean no-op"); + } + + /// A line longer than the cap, with no newline, is rejected — and the + /// error surfaces rather than the reader growing without bound. + #[tokio::test] + async fn an_oversized_line_is_rejected() { + // MAX + 1 bytes of a single line with no newline anywhere. + let payload = vec![b'x'; MAX_LINE_BYTES + 1]; + let mut cursor = std::io::Cursor::new(payload); + let err = read_line_bounded(&mut cursor) + .await + .expect_err("an over-cap line must be rejected"); + assert!( + format!("{err:#}").contains("exceeds") && format!("{err:#}").contains("limit"), + "got {err:#}" + ); + } + + /// A line exactly at the cap, terminated by a newline, is accepted — + /// the boundary is inclusive of `MAX_LINE_BYTES` content bytes. + #[tokio::test] + async fn a_line_at_the_cap_is_accepted() { + let mut data = vec![b'y'; MAX_LINE_BYTES]; + data.push(b'\n'); + let mut cursor = std::io::Cursor::new(data); + let line = read_line_bounded(&mut cursor) + .await + .expect("a line at the cap must be accepted") + .expect("a full line, not EOF"); + assert_eq!(line.len(), MAX_LINE_BYTES); + } + + #[tokio::test] + async fn bounded_reader_preserves_pipelined_lines() { + let data = b"first\nsecond\n"; + let mut reader = BufReader::new(&data[..]); + assert_eq!( + read_line_bounded(&mut reader).await.unwrap().as_deref(), + Some("first") + ); + assert_eq!( + read_line_bounded(&mut reader).await.unwrap().as_deref(), + Some("second") + ); + assert!(read_line_bounded(&mut reader).await.unwrap().is_none()); + } + + /// The bounded reader trims CRLF and returns EOF cleanly, matching the + /// framing the previous `lines()` adapter provided. + #[tokio::test] + async fn bounded_reader_trims_crlf_and_reports_eof() { + let mut cursor = std::io::Cursor::new(b"hello\r\n".to_vec()); + let line = read_line_bounded(&mut cursor).await.unwrap().unwrap(); + assert_eq!(line, "hello"); + // Second read hits clean EOF. + assert!(read_line_bounded(&mut cursor).await.unwrap().is_none()); + } + } +} + +// --------------------------------------------------------------------------- +// Non-Unix stubs: compile everywhere, fail loudly at runtime. +// --------------------------------------------------------------------------- +#[cfg(not(unix))] +pub use non_unix_impl::{serve_local, LocalClient}; + +#[cfg(not(unix))] +mod non_unix_impl { + use super::*; + use std::path::PathBuf; + + use tokio::sync::watch; + + use crate::control_plane::ControlPlaneHandle; + + /// Not supported off Unix: there is no unix-domain socket transport. The + /// symbol exists so the crate compiles for `x86_64-pc-windows-gnu`. + pub async fn serve_local( + _path: PathBuf, + _cp: ControlPlaneHandle, + _shutdown: watch::Receiver, + ) -> anyhow::Result<()> { + Err(anyhow::anyhow!( + "the control-plane local IPC socket is only supported on Unix" + )) + } + + /// Non-Unix stub client. + pub struct LocalClient { + _path: PathBuf, + } + + impl LocalClient { + pub fn new(path: impl Into) -> Self { + Self { _path: path.into() } + } + + pub fn from_default() -> anyhow::Result { + Ok(Self { + _path: default_socket_path()?, + }) + } + + pub async fn request(&self, _request: &LocalRequest) -> anyhow::Result { + Err(anyhow::anyhow!( + "the control-plane local IPC socket is only supported on Unix" + )) + } + } +} + +#[cfg(test)] +mod protocol_tests { + use super::*; + + #[test] + fn default_socket_path_prefers_env_then_falls_back_to_home() { + // Env mutation is process-global; keep both cases in one test so they + // do not race a sibling test that also touches these vars, and restore + // on exit. + let prev_sock = std::env::var_os(SOCKET_ENV); + let prev_home = std::env::var_os("HOME"); + + std::env::set_var(SOCKET_ENV, "/tmp/custom-agent.sock"); + assert_eq!( + default_socket_path().unwrap(), + std::path::PathBuf::from("/tmp/custom-agent.sock") + ); + + std::env::remove_var(SOCKET_ENV); + std::env::set_var("HOME", "/home/tester"); + assert_eq!( + default_socket_path().unwrap(), + std::path::PathBuf::from("/home/tester/.openab/run/agent.sock") + ); + + match prev_sock { + Some(v) => std::env::set_var(SOCKET_ENV, v), + None => std::env::remove_var(SOCKET_ENV), + } + match prev_home { + Some(v) => std::env::set_var("HOME", v), + None => std::env::remove_var("HOME"), + } + } + + #[test] + fn requests_are_tagged_by_op() { + let spawn = LocalRequest::Spawn { + delegation_id: "d-1".into(), + target: LocalTarget { + name: Some("w1".into()), + labels: None, + }, + prompt: "hi".into(), + deadline_secs: 60, + parent: None, + }; + let v = serde_json::to_value(&spawn).unwrap(); + assert_eq!(v["op"], "spawn"); + assert_eq!(v["deadline_secs"], 60); + // Round-trips. + let back: LocalRequest = serde_json::from_value(v).unwrap(); + assert_eq!(back, spawn); + + let list = serde_json::to_value(LocalRequest::ListAgents).unwrap(); + assert_eq!(list["op"], "list_agents"); + + let check = serde_json::to_value(LocalRequest::Check { + handle: "tok".into(), + }) + .unwrap(); + assert_eq!(check["op"], "check"); + assert_eq!(check["handle"], "tok"); + } + + #[test] + fn responses_are_tagged_by_kind() { + let spawned = serde_json::to_value(LocalResponse::Spawned { + handle: "a1b2c3".into(), + assigned_to: "prod/w1".into(), + }) + .unwrap(); + assert_eq!(spawned["kind"], "spawned"); + assert_eq!(spawned["handle"], "a1b2c3"); + assert_eq!(spawned["assigned_to"], "prod/w1"); + + // The Check-family variants are each self-describing. + let pending = serde_json::to_value(LocalResponse::Pending).unwrap(); + assert_eq!(pending["kind"], "pending"); + let running = serde_json::to_value(LocalResponse::Running { + assigned_to: "prod/w1".into(), + }) + .unwrap(); + assert_eq!(running["kind"], "running"); + assert_eq!(running["assigned_to"], "prod/w1"); + + // The Terminal variant keeps its own `status` field alongside the tag. + let terminal = serde_json::to_value(LocalResponse::Terminal { + status: DelegationStatusWire::Completed, + result: Some("done".into()), + error: None, + }) + .unwrap(); + assert_eq!(terminal["kind"], "terminal"); + assert_eq!(terminal["status"], "completed"); + + let err = serde_json::to_value(LocalResponse::error("nope")).unwrap(); + assert_eq!(err["kind"], "error"); + assert_eq!(err["message"], "nope"); + } +} diff --git a/crates/openab-core/src/control_plane/mod.rs b/crates/openab-core/src/control_plane/mod.rs index c6c8d82ad..ceb877c98 100644 --- a/crates/openab-core/src/control_plane/mod.rs +++ b/crates/openab-core/src/control_plane/mod.rs @@ -17,11 +17,21 @@ pub mod client; pub mod executor; +pub mod local; +pub mod primary; -pub use client::ControlPlaneClient; +pub use client::{ControlPlaneClient, ControlPlaneHandle}; pub use executor::{ delegation_session_key, DelegationExecutor, PromptOutcome, PromptRunner, RouterPromptRunner, }; +pub use local::{ + default_socket_path, serve_local, LocalAgent, LocalClient, LocalRequest, LocalResponse, + LocalTarget, +}; +pub use primary::{ + ClientCommand, CommandError, DelegationHandle, DelegationLifecycle, DelegationOutcome, + DelegationSnapshot, DelegationStatusWire, SpawnAck, SpawnRequest, SpawnTarget, +}; use crate::config::CpAgentType; use openab_cp::proto::AgentType; diff --git a/crates/openab-core/src/control_plane/primary.rs b/crates/openab-core/src/control_plane/primary.rs new file mode 100644 index 000000000..7cfcf3b23 --- /dev/null +++ b/crates/openab-core/src/control_plane/primary.rs @@ -0,0 +1,1488 @@ +//! Primary-side (initiating) control-plane state. +//! +//! The worker side turns one inbound `cp/delegate` into one +//! `cp/delegate_result` ([`super::executor`]). The primary side is the mirror: +//! it *originates* `cp/delegate`, `cp/cancel`, and `cp/list_agents`, then +//! correlates the CP's replies and the later, initiator-bound +//! `cp/delegate_result` back to whoever asked for the work. +//! +//! Three deliberate seams keep this unit-testable without a socket: +//! +//! - **Commands are data.** External callers (the local UDS server, an MCP +//! tool, a CLI) submit a [`ClientCommand`] over an mpsc channel and await a +//! `oneshot`. The command carries no socket and no `&mut sink`; the serve +//! loop is the single owner that turns it into a frame. So this module can be +//! exercised by building commands and inspecting the frames/bookkeeping they +//! produce. +//! - **The delegation handle is opaque.** [`DelegationHandle`] encapsulates the +//! `(delegation_id, admission)` pair the wire uses to name one admission of a +//! reusable id. Callers hold it, pass it back to cancel or await, and never +//! pick it apart — a bare `delegation_id` is never a valid reference, exactly +//! as the protocol requires. +//! - **Correlation is a map, not a socket read.** [`PrimaryState`] owns the +//! `rpc_id → pending` and `(id, admission) → delegation` tables. The serve +//! loop feeds it inbound frames and outbound intents; every state transition +//! is a plain method with an assertable result. + +use std::collections::{HashMap, VecDeque}; + +use openab_cp::proto::{ + AdmissionToken, AgentSummary, DelegateAck, DelegateResultParams, DelegationStatus, ErrorObject, + TargetSelector, +}; +use tokio::sync::oneshot; +use tracing::{debug, warn}; + +/// Opaque reference to ONE admission of a delegation this runtime initiated. +/// +/// It wraps the coupled `(delegation_id, admission)` pair the protocol uses to +/// name a specific admission (see `openab_cp::proto::AdmissionToken`). It is +/// opaque on purpose: a caller cancels or awaits by handing the whole handle +/// back, never by reconstructing it from a bare `delegation_id` (which is +/// reusable and names only a *slot*). Cloneable and comparable so a caller can +/// hold one while awaiting and pass a copy to cancel. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub struct DelegationHandle { + delegation_id: String, + admission: AdmissionToken, +} + +impl DelegationHandle { + pub(crate) fn new(delegation_id: impl Into, admission: AdmissionToken) -> Self { + Self { + delegation_id: delegation_id.into(), + admission, + } + } + + pub fn delegation_id(&self) -> &str { + &self.delegation_id + } + + pub fn admission(&self) -> AdmissionToken { + self.admission + } + + /// The composite key the tracking table is indexed by. Kept private-ish: + /// callers compare handles, they do not build keys. + fn key(&self) -> (String, AdmissionToken) { + (self.delegation_id.clone(), self.admission) + } +} + +/// A spawn target: an exact logical name, or a set of labels that must all +/// match. Exactly one is set — the constructors enforce that so a caller +/// cannot submit an ambiguous or empty selector. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum SpawnTarget { + /// Route to the runtime with this exact logical `name`. + Name(String), + /// Route to any healthy runtime whose labels are a superset of these. + Labels(std::collections::BTreeMap), +} + +impl SpawnTarget { + pub fn by_name(name: impl Into) -> Self { + SpawnTarget::Name(name.into()) + } + + pub fn by_labels(labels: std::collections::BTreeMap) -> Self { + SpawnTarget::Labels(labels) + } + + pub(crate) fn into_selector(self) -> TargetSelector { + match self { + SpawnTarget::Name(name) => TargetSelector { + name: Some(name), + labels: None, + }, + SpawnTarget::Labels(labels) => TargetSelector { + name: None, + labels: Some(labels), + }, + } + } +} + +/// A request to spawn (delegate) work onto a peer runtime. +#[derive(Debug, Clone)] +pub struct SpawnRequest { + /// Caller-generated idempotency id. Reusable across cancel-then-retry. + pub delegation_id: String, + pub target: SpawnTarget, + pub prompt: String, + /// Absolute deadline by which the delegation must complete. + pub deadline: chrono::DateTime, + /// If this spawn happens while serving another delegation, its handle — + /// the CP derives ancestry, depth, cycle, and the child's deadline budget + /// from the admission this names. `None` for a root delegation. + pub parent: Option, +} + +/// Terminal outcome of an initiated delegation, as reported back to the caller +/// that spawned it. Mirrors the wire `DelegationStatus` plus the payload. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct DelegationOutcome { + pub status: DelegationStatus, + pub result: Option, + pub error: Option, +} +impl DelegationOutcome { + fn from_result(params: &DelegateResultParams) -> Self { + Self { + status: params.status.clone(), + result: params.result.clone(), + error: params.error.clone(), + } + } +} + +/// Serializable projection of [`DelegationStatus`] for the local IPC protocol. +/// +/// A distinct type so the local socket protocol (consumed by a CLI or MCP +/// tool) does not depend on `openab_cp::proto` types directly; the snake_case +/// wire words are identical, so a caller sees the same vocabulary the CP uses. +#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum DelegationStatusWire { + Completed, + Failed, + Timeout, + Cancelled, + TargetDisconnected, +} + +impl From for DelegationStatusWire { + fn from(s: DelegationStatus) -> Self { + match s { + DelegationStatus::Completed => DelegationStatusWire::Completed, + DelegationStatus::Failed => DelegationStatusWire::Failed, + DelegationStatus::Timeout => DelegationStatusWire::Timeout, + DelegationStatus::Cancelled => DelegationStatusWire::Cancelled, + DelegationStatus::TargetDisconnected => DelegationStatusWire::TargetDisconnected, + } + } +} + +/// The reply to a successful [`ClientCommand::Spawn`]: the opaque handle plus +/// the peer the CP routed the work to. `assigned_to` is only knowable from the +/// CP's `cp/delegate` ack, so it is carried here rather than reconstructed by +/// the caller (the previous shape returned a bare handle, forcing the local +/// layer to emit an empty `assigned_to`). +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct SpawnAck { + pub handle: DelegationHandle, + pub assigned_to: String, +} + +/// A nonblocking snapshot of a tracked delegation's lifecycle, returned by +/// [`ClientCommand::Check`]. Unlike [`ClientCommand::Await`], it never parks: +/// a still-running delegation answers `Running`, not a blocked reply. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum DelegationSnapshot { + /// Delegate frame sent, ack not yet received. + Pending, + /// Ack received; the peer runtime is serving it. + Running { assigned_to: String }, + /// A terminal `cp/delegate_result` has arrived. + Terminal(DelegationOutcome), +} + +/// A command submitted to the serve loop by an external initiating surface. +/// +/// Each carries its own `oneshot` reply sender: the serve loop performs the +/// frame write and (for spawn/cancel/list) the correlation, then answers the +/// caller. The loop is the only writer to the socket, so no command holds the +/// sink. +pub enum ClientCommand { + /// Initiate a delegation. Replies with the opaque handle and the assigned + /// peer once the CP acks routing (the terminal result is awaited + /// separately via [`Await`]). + Spawn { + request: SpawnRequest, + reply: oneshot::Sender>, + }, + /// Await the terminal result of a delegation this runtime initiated. + /// Blocks (parks) until a terminal frame arrives or the connection drops. + Await { + handle: DelegationHandle, + reply: oneshot::Sender>, + }, + /// Nonblocking status query for a delegation this runtime initiated. + /// Answers immediately with the current lifecycle snapshot. + Check { + handle: DelegationHandle, + reply: oneshot::Sender>, + }, + /// Cancel an in-flight delegation by opaque handle. + Cancel { + handle: DelegationHandle, + reason: String, + reply: oneshot::Sender>, + }, + /// Namespace-scoped roster snapshot. + ListAgents { + reply: oneshot::Sender, CommandError>>, + }, +} + +/// Why a primary-side command could not be satisfied. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum CommandError { + /// The client is not currently connected/registered, so no frame can be + /// originated. The caller may retry once the client reconnects. + NotConnected, + /// The CP answered the request with a JSON-RPC error. + Cp { code: i64, message: String }, + /// The named handle is not (or no longer) tracked here. + UnknownDelegation, + /// The connection dropped while this request/delegation was outstanding. + Disconnected, + /// The request could not be represented on the wire (e.g. its serialized + /// frame exceeds the transport limit). Distinct from `Internal`: it names + /// a caller-side fault the caller could fix (a smaller prompt), not a bug. + InvalidRequest(String), + /// The command channel or reply path was torn down (shutdown). + Internal(String), +} + +impl std::fmt::Display for CommandError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + CommandError::NotConnected => { + write!(f, "control-plane client is not connected") + } + CommandError::Cp { code, message } => { + write!( + f, + "control plane returned an error: {message} (code {code})" + ) + } + CommandError::UnknownDelegation => { + write!(f, "no such delegation is tracked by this runtime") + } + CommandError::Disconnected => { + write!( + f, + "control-plane connection dropped before the delegation finished" + ) + } + CommandError::InvalidRequest(m) => write!(f, "invalid request: {m}"), + CommandError::Internal(m) => write!(f, "internal error: {m}"), + } + } +} + +impl std::error::Error for CommandError {} + +impl CommandError { + fn from_cp(err: ErrorObject) -> Self { + CommandError::Cp { + code: err.code, + message: err.message, + } + } +} + +/// What the serve loop originated and is still awaiting a JSON-RPC reply for. +/// +/// Correlated by the JSON-RPC id we chose. `Spawn` is special: its reply is a +/// [`DelegateAck`], and success does NOT complete the caller — it registers the +/// admission so the later `cp/delegate_result` can be routed. Cancel and list +/// complete the caller directly. +enum PendingRequest { + Spawn { + delegation_id: String, + reply: oneshot::Sender>, + }, + Cancel { + reply: oneshot::Sender>, + }, + ListAgents { + reply: oneshot::Sender, CommandError>>, + }, +} + +/// Lifecycle of one initiated delegation, from ack to terminal. +struct TrackedDelegation { + state: DelegationLifecycle, + /// The peer the CP assigned this delegation to, learned at ack time and + /// retained so it can be accounted into the terminal-history byte bound + /// (the terminal `DelegationOutcome` does not itself carry it). + assigned_to: String, + /// Bytes this entry contributes to [`PrimaryState::terminal_bytes`] while + /// it is terminal; 0 while live. Stored per-entry so eviction subtracts + /// exactly what insertion added, with no re-measurement. + accounted_bytes: usize, +} + +/// The observable state of a tracked delegation. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum DelegationLifecycle { + /// Delegate frame sent, ack not yet received. + Pending, + /// Ack received; the peer runtime is (or will be) serving it. + Running { assigned_to: String }, + /// A terminal `cp/delegate_result` arrived. + Terminal(DelegationOutcome), +} + +/// The composite key naming one admission of a (reusable) delegation id. +type DelegationKey = (String, AdmissionToken); +/// A parked reply for a caller awaiting a delegation's terminal outcome. +type AwaitReply = oneshot::Sender>; +/// Retained terminal-history entry ceiling. Live delegations are never +/// evicted; this bounds only the count of retained TERMINAL outcomes plus the +/// live entries that share the map. +const MAX_TRACKED_DELEGATIONS: usize = 4096; +/// Retained terminal-history byte ceiling (16 MiB). A second, independent +/// bound alongside the count: a few very large terminal payloads can pin far +/// more memory than 4096 tiny ones, so history is capped on BOTH axes. +/// Only terminal entries contribute to (and are evicted to satisfy) this +/// bound; live delegations carry no accounted bytes until they terminate. +const MAX_TRACKED_BYTES: usize = 16 * 1024 * 1024; + +/// Primary-side correlation and delegation tracking for one runtime instance. +/// +/// Not `Sync`-shared across tasks: it lives inside the single-owner serve loop +/// and every method runs there. The serve loop hands it inbound frames and +/// outbound intents; nothing else touches it. +#[derive(Default)] +pub struct PrimaryState { + /// JSON-RPC id → the request the loop originated and awaits a reply for. + pending: HashMap, + /// (delegation_id, admission) → its lifecycle. + tracked: HashMap, + /// Insertion order for deterministic oldest-terminal history eviction. + tracked_order: VecDeque, + /// Callers awaiting a terminal result, keyed by the same composite key. + /// A caller can await before OR after the terminal arrives. + awaiters: HashMap>, + /// Running total of bytes accounted to TERMINAL entries currently retained + /// (delegation_id + result + error + assigned_to). Kept incrementally so + /// eviction does not re-scan the whole map; live entries contribute 0. + terminal_bytes: usize, +} + +/// Outcome of feeding a spawn command to the state: the params to emit, already +/// validated and parked under the rpc id the caller allocated. +pub(crate) struct SpawnEmission { + pub rpc_id: u64, + pub params: openab_cp::proto::DelegateParams, +} + +impl PrimaryState { + pub fn new() -> Self { + Self::default() + } + + /// Number of delegations currently tracked (any lifecycle state). + pub fn tracked_len(&self) -> usize { + self.tracked.len() + } + + /// Register a spawn: park its reply under `rpc_id` and produce the wire + /// params. The caller is answered later, when the ack lands. `rpc_id` is + /// allocated by the serve loop so every outbound frame — heartbeat, + /// delegate_result ack, and these primary requests — shares one id space. + pub(crate) fn begin_spawn( + &mut self, + rpc_id: u64, + request: SpawnRequest, + reply: oneshot::Sender>, + ) -> SpawnEmission { + let (parent_delegation_id, parent_admission) = match request.parent { + Some(h) => (Some(h.delegation_id), Some(h.admission)), + None => (None, None), + }; + let params = openab_cp::proto::DelegateParams { + delegation_id: request.delegation_id.clone(), + target: request.target.into_selector(), + prompt: request.prompt, + deadline: request.deadline, + parent_delegation_id, + parent_admission, + }; + self.pending.insert( + rpc_id, + PendingRequest::Spawn { + delegation_id: request.delegation_id, + reply, + }, + ); + SpawnEmission { rpc_id, params } + } + + /// Remove a parked pending request (a spawn whose frame the serve loop + /// could not put on the wire) and answer its caller + /// [`CommandError::InvalidRequest`]. Used for an outbound frame that is + /// well-formed but too large for the transport: the serve loop parks the + /// spawn in `begin_spawn`, discovers the serialized frame exceeds the + /// limit, and calls this to fail the caller WITHOUT tearing down the + /// connection. Returns true if a pending request was parked under + /// `rpc_id` (so this was ours to fail); false otherwise. Tracks nothing. + pub(crate) fn fail_pending_invalid_request(&mut self, rpc_id: u64, message: String) -> bool { + let Some(pending) = self.pending.remove(&rpc_id) else { + return false; + }; + match pending { + PendingRequest::Spawn { reply, .. } => { + let _ = reply.send(Err(CommandError::InvalidRequest(message))); + } + PendingRequest::Cancel { reply } => { + let _ = reply.send(Err(CommandError::InvalidRequest(message))); + } + PendingRequest::ListAgents { reply } => { + let _ = reply.send(Err(CommandError::InvalidRequest(message))); + } + } + true + } + + /// Register a cancel: park its reply under `rpc_id` and return the wire + /// params. Returns `None` (and answers the caller `UnknownDelegation`) if + /// the handle is not tracked here — a cancel for something we never spawned + /// is a caller error, not a frame worth sending. + pub(crate) fn begin_cancel( + &mut self, + rpc_id: u64, + handle: &DelegationHandle, + reason: String, + reply: oneshot::Sender>, + ) -> Option { + if !self.tracked.contains_key(&handle.key()) { + let _ = reply.send(Err(CommandError::UnknownDelegation)); + return None; + } + let params = openab_cp::proto::CancelParams { + delegation_id: handle.delegation_id.clone(), + admission: handle.admission, + reason, + }; + self.pending + .insert(rpc_id, PendingRequest::Cancel { reply }); + Some(params) + } + + /// Register a list_agents: park its reply under `rpc_id`. + pub(crate) fn begin_list_agents( + &mut self, + rpc_id: u64, + reply: oneshot::Sender, CommandError>>, + ) { + self.pending + .insert(rpc_id, PendingRequest::ListAgents { reply }); + } + + /// Await a delegation's terminal outcome. If it already ended, the caller + /// is answered immediately; otherwise the reply is parked until the + /// terminal frame arrives (or the connection drops). + pub fn begin_await( + &mut self, + handle: &DelegationHandle, + reply: oneshot::Sender>, + ) { + match self.tracked.get(&handle.key()) { + Some(TrackedDelegation { + state: DelegationLifecycle::Terminal(outcome), + .. + }) => { + let _ = reply.send(Ok(outcome.clone())); + } + Some(_) => { + self.awaiters.entry(handle.key()).or_default().push(reply); + } + None => { + let _ = reply.send(Err(CommandError::UnknownDelegation)); + } + } + } + + /// Nonblocking status query. Answers the caller immediately with the + /// current [`DelegationSnapshot`]; a running delegation returns `Running` + /// rather than parking, unlike [`begin_await`]. Unknown handles get + /// `UnknownDelegation`. + pub fn check( + &self, + handle: &DelegationHandle, + reply: oneshot::Sender>, + ) { + let snapshot = match self.tracked.get(&handle.key()) { + Some(TrackedDelegation { + state: DelegationLifecycle::Pending, + .. + }) => Some(DelegationSnapshot::Pending), + Some(TrackedDelegation { + state: DelegationLifecycle::Running { assigned_to }, + .. + }) => Some(DelegationSnapshot::Running { + assigned_to: assigned_to.clone(), + }), + Some(TrackedDelegation { + state: DelegationLifecycle::Terminal(outcome), + .. + }) => Some(DelegationSnapshot::Terminal(outcome.clone())), + None => None, + }; + match snapshot { + Some(s) => { + let _ = reply.send(Ok(s)); + } + None => { + let _ = reply.send(Err(CommandError::UnknownDelegation)); + } + } + } + /// pending request that originated it. Returns true if it was ours. + /// + /// `id` is the echoed request id; `result`/`error` are the two arms of a + /// JSON-RPC reply. + pub fn on_reply( + &mut self, + id: u64, + result: Option, + error: Option, + ) -> bool { + let Some(pending) = self.pending.remove(&id) else { + return false; + }; + match pending { + PendingRequest::Spawn { + delegation_id, + reply, + } => { + if let Some(err) = error { + let _ = reply.send(Err(CommandError::from_cp(err))); + return true; + } + let ack: Option = result.and_then(|v| serde_json::from_value(v).ok()); + match ack { + Some(ack) => { + // The ack MUST name the delegation we asked to spawn. + // A mismatched delegation_id is a protocol violation + // (a routing bug or a crossed reply): fail the caller + // and track NOTHING, so a foreign id can never seed the + // correlation table under this rpc id. + if ack.delegation_id != delegation_id { + warn!( + requested = %delegation_id, + acked = %ack.delegation_id, + "cp/delegate ack named a different delegation_id — rejecting" + ); + let _ = reply.send(Err(CommandError::Internal(format!( + "cp/delegate ack delegation_id {:?} does not match the request {:?}", + ack.delegation_id, delegation_id + )))); + return true; + } + let handle = + DelegationHandle::new(ack.delegation_id.clone(), ack.admission); + // Track the admission so the later result frame routes. + self.make_history_room(); + let key = handle.key(); + self.tracked.insert( + key.clone(), + TrackedDelegation { + state: DelegationLifecycle::Running { + assigned_to: ack.assigned_to.clone(), + }, + assigned_to: ack.assigned_to.clone(), + accounted_bytes: 0, + }, + ); + self.tracked_order.push_back(key); + debug!( + delegation_id = %delegation_id, + admission = ack.admission, + "delegation admitted by control plane" + ); + let _ = reply.send(Ok(SpawnAck { + handle, + assigned_to: ack.assigned_to, + })); + } + None => { + let _ = reply.send(Err(CommandError::Internal( + "cp/delegate ack carried no valid result".into(), + ))); + } + } + true + } + PendingRequest::Cancel { reply } => { + match error { + Some(err) => { + let _ = reply.send(Err(CommandError::from_cp(err))); + } + None => { + let _ = reply.send(Ok(())); + } + } + true + } + PendingRequest::ListAgents { reply } => { + if let Some(err) = error { + let _ = reply.send(Err(CommandError::from_cp(err))); + return true; + } + let parsed: Option = + result.and_then(|v| serde_json::from_value(v).ok()); + match parsed { + Some(list) => { + let _ = reply.send(Ok(list.agents)); + } + None => { + let _ = reply.send(Err(CommandError::Internal( + "cp/list_agents reply carried no valid result".into(), + ))); + } + } + true + } + } + } + + /// Evict oldest terminal history until BOTH bounds are satisfied for a new + /// tracked admission: at most [`MAX_TRACKED_DELEGATIONS`] entries AND at + /// most [`MAX_TRACKED_BYTES`] of accounted terminal payload. Running or + /// Pending work is never evicted; the CP/global capacity bounds it and the + /// count bound is expressed against the whole map so live work still + /// pushes out old history. Eviction always removes the OLDEST terminal + /// entry in insertion order, so it is deterministic and never touches live + /// delegations. + fn make_history_room(&mut self) { + while self.tracked.len() >= MAX_TRACKED_DELEGATIONS + || self.terminal_bytes > MAX_TRACKED_BYTES + { + let Some(index) = self.tracked_order.iter().position(|key| { + self.tracked.get(key).is_some_and(|tracked| { + matches!(tracked.state, DelegationLifecycle::Terminal(_)) + }) + }) else { + // No terminal entry left to evict: only live work remains, and + // live work is never evicted. Both bounds are best-effort + // against a flood of live delegations, which CP capacity caps. + break; + }; + let key = self + .tracked_order + .remove(index) + .expect("index came from position"); + if let Some(tracked) = self.tracked.remove(&key) { + self.terminal_bytes = self.terminal_bytes.saturating_sub(tracked.accounted_bytes); + } + self.awaiters.remove(&key); + } + } + + /// Bytes a terminal entry contributes to the history byte bound: + /// delegation_id + result + error + assigned_to. Counts the payload that + /// is actually retained; the admission (a `u64`) is negligible and fixed. + fn terminal_entry_bytes( + delegation_id: &str, + assigned_to: &str, + outcome: &DelegationOutcome, + ) -> usize { + delegation_id.len() + + assigned_to.len() + + outcome.result.as_deref().map_or(0, str::len) + + outcome.error.as_deref().map_or(0, str::len) + } + + /// Route an initiator-bound `cp/delegate_result` to its tracked delegation + /// and wake any awaiters. Returns true if the frame named a delegation we + /// track (and therefore should be acked as ours). + /// + /// Correlation is on the `(delegation_id, admission)` pair, never the + /// reusable id alone: a late result for a superseded admission is dropped. + pub fn on_delegate_result(&mut self, params: &DelegateResultParams) -> bool { + let key = (params.delegation_id.clone(), params.admission); + let Some(tracked) = self.tracked.get_mut(&key) else { + warn!( + delegation_id = %params.delegation_id, + admission = params.admission, + "cp/delegate_result for an untracked (id, admission) — ignoring" + ); + return false; + }; + // First terminal wins; a duplicate for the same admission is ignored. + if matches!(tracked.state, DelegationLifecycle::Terminal(_)) { + debug!( + delegation_id = %params.delegation_id, + admission = params.admission, + "duplicate terminal frame ignored" + ); + return true; + } + let outcome = DelegationOutcome::from_result(params); + let assigned_to = tracked.assigned_to.clone(); + let entry_bytes = Self::terminal_entry_bytes(¶ms.delegation_id, &assigned_to, &outcome); + tracked.state = DelegationLifecycle::Terminal(outcome.clone()); + tracked.accounted_bytes = entry_bytes; + self.terminal_bytes = self.terminal_bytes.saturating_add(entry_bytes); + if let Some(waiters) = self.awaiters.remove(&key) { + for w in waiters { + let _ = w.send(Ok(outcome.clone())); + } + } + // Move this entry to the back so terminal history is ordered by + // completion rather than admission. Enforce the byte bound now: older + // completions are evicted first, while this entry is evicted only if + // the retained set still cannot fit (for example, it alone is huge). + if let Some(index) = self.tracked_order.iter().position(|queued| queued == &key) { + self.tracked_order.remove(index); + } + self.tracked_order.push_back(key); + self.make_history_room(); + true + } + + /// Current lifecycle of a tracked delegation (for tests/introspection). + pub fn lifecycle(&self, handle: &DelegationHandle) -> Option { + self.tracked.get(&handle.key()).map(|t| t.state.clone()) + } + + /// Fail every outstanding primary-side interaction on connection loss. + /// + /// Pending requests (spawn/cancel/list awaiting a reply that will never + /// come), non-terminal tracked delegations, and their awaiters are all + /// answered `Disconnected`. Terminal delegations are dropped: their result + /// already reached the caller. + /// + /// This is the primary-side analogue of the executor's `cancel_all` on the + /// worker side — no local state may claim a delegation is live once the + /// only connection that could complete it is gone. + pub fn fail_all_live(&mut self) { + for (_id, pending) in self.pending.drain() { + match pending { + PendingRequest::Spawn { reply, .. } => { + let _ = reply.send(Err(CommandError::Disconnected)); + } + PendingRequest::Cancel { reply } => { + let _ = reply.send(Err(CommandError::Disconnected)); + } + PendingRequest::ListAgents { reply } => { + let _ = reply.send(Err(CommandError::Disconnected)); + } + } + } + // Mark non-terminal delegations failed and answer their awaiters. + let live_keys: Vec = self + .tracked + .iter() + .filter(|(_, t)| !matches!(t.state, DelegationLifecycle::Terminal(_))) + .map(|(k, _)| k.clone()) + .collect(); + for key in live_keys { + if let Some(waiters) = self.awaiters.remove(&key) { + for w in waiters { + let _ = w.send(Err(CommandError::Disconnected)); + } + } + self.tracked.remove(&key); + } + // Any awaiters that remain point at delegations that reached a terminal + // state (their result already reached the caller at `begin_await` time), + // so nothing is owed. Drop any that no longer name a tracked delegation + // to keep the map from leaking. Snapshot the keys first so the retain + // closure does not borrow `self.tracked` while `self.awaiters` is + // mutably borrowed. + let tracked_keys: std::collections::HashSet = + self.tracked.keys().cloned().collect(); + self.awaiters.retain(|k, _| tracked_keys.contains(k)); + self.tracked_order.retain(|key| tracked_keys.contains(key)); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use openab_cp::proto::DelegationStatus; + + fn spawn_request(id: &str) -> SpawnRequest { + SpawnRequest { + delegation_id: id.into(), + target: SpawnTarget::by_name("worker-1"), + prompt: "do the thing".into(), + deadline: chrono::Utc::now() + chrono::Duration::seconds(60), + parent: None, + } + } + + fn ack_value( + delegation_id: &str, + admission: AdmissionToken, + assigned_to: &str, + ) -> serde_json::Value { + serde_json::to_value(DelegateAck { + delegation_id: delegation_id.into(), + admission, + assigned_to: assigned_to.into(), + }) + .unwrap() + } + + #[test] + fn a_handle_is_opaque_and_carries_both_halves() { + let h = DelegationHandle::new("d-1", 7); + assert_eq!(h.delegation_id(), "d-1"); + assert_eq!(h.admission(), 7); + // Same id, different admission is a different handle. + assert_ne!(h, DelegationHandle::new("d-1", 8)); + } + + #[test] + fn spawn_target_is_exactly_one_of_name_or_labels() { + let by_name = SpawnTarget::by_name("w1").into_selector(); + assert_eq!(by_name.name.as_deref(), Some("w1")); + assert!(by_name.labels.is_none()); + + let mut labels = std::collections::BTreeMap::new(); + labels.insert("backend".to_string(), "kiro".to_string()); + let by_labels = SpawnTarget::by_labels(labels.clone()).into_selector(); + assert!(by_labels.name.is_none()); + assert_eq!(by_labels.labels, Some(labels)); + } + + #[test] + fn begin_spawn_parks_the_reply_and_produces_wire_params() { + let mut st = PrimaryState::new(); + let (tx, _rx) = oneshot::channel(); + let emission = st.begin_spawn(1, spawn_request("d-1"), tx); + assert_eq!(emission.rpc_id, 1); + assert_eq!(emission.params.delegation_id, "d-1"); + assert_eq!(emission.params.target.name.as_deref(), Some("worker-1")); + // A root spawn carries no parent reference on the wire. + assert!(emission.params.parent_delegation_id.is_none()); + assert!(emission.params.parent_admission.is_none()); + } + + #[test] + fn a_parented_spawn_carries_both_halves_of_the_parent_reference() { + let mut st = PrimaryState::new(); + let (tx, _rx) = oneshot::channel(); + let mut req = spawn_request("child"); + req.parent = Some(DelegationHandle::new("parent", 42)); + let emission = st.begin_spawn(7, req, tx); + assert_eq!( + emission.params.parent_delegation_id.as_deref(), + Some("parent") + ); + assert_eq!(emission.params.parent_admission, Some(42)); + } + + #[tokio::test] + async fn a_delegate_ack_resolves_the_spawn_and_starts_tracking() { + let mut st = PrimaryState::new(); + let (tx, rx) = oneshot::channel(); + let emission = st.begin_spawn(1, spawn_request("d-1"), tx); + assert!(st.on_reply( + emission.rpc_id, + Some(ack_value("d-1", 5, "prod/worker-1")), + None + )); + let ack = rx.await.unwrap().unwrap(); + assert_eq!(ack.assigned_to, "prod/worker-1"); + let handle = ack.handle; + assert_eq!(handle.delegation_id(), "d-1"); + assert_eq!(handle.admission(), 5); + assert_eq!( + st.lifecycle(&handle), + Some(DelegationLifecycle::Running { + assigned_to: "prod/worker-1".into() + }) + ); + } + + #[tokio::test] + async fn a_delegate_error_reply_fails_the_spawn_and_tracks_nothing() { + let mut st = PrimaryState::new(); + let (tx, rx) = oneshot::channel(); + let emission = st.begin_spawn(1, spawn_request("d-1"), tx); + let err = ErrorObject::new(openab_cp::proto::codes::NO_TARGET, "no such worker"); + assert!(st.on_reply(emission.rpc_id, None, Some(err))); + let outcome = rx.await.unwrap(); + assert!(matches!( + outcome, + Err(CommandError::Cp { + code: openab_cp::proto::codes::NO_TARGET, + .. + }) + )); + assert_eq!(st.tracked_len(), 0); + } + + #[tokio::test] + async fn a_terminal_result_wakes_an_earlier_awaiter() { + let mut st = PrimaryState::new(); + // Spawn + ack. + let (tx, rx) = oneshot::channel(); + let emission = st.begin_spawn(1, spawn_request("d-1"), tx); + st.on_reply(emission.rpc_id, Some(ack_value("d-1", 3, "prod/w1")), None); + let handle = rx.await.unwrap().unwrap().handle; + + // Await before the terminal frame arrives. + let (atx, arx) = oneshot::channel(); + st.begin_await(&handle, atx); + + let result = DelegateResultParams { + delegation_id: "d-1".into(), + admission: 3, + status: DelegationStatus::Completed, + result: Some("done".into()), + error: None, + }; + assert!(st.on_delegate_result(&result)); + let outcome = arx.await.unwrap().unwrap(); + assert_eq!(outcome.status, DelegationStatus::Completed); + assert_eq!(outcome.result.as_deref(), Some("done")); + } + + #[tokio::test] + async fn awaiting_an_already_terminal_delegation_answers_immediately() { + let mut st = PrimaryState::new(); + let (tx, rx) = oneshot::channel(); + let emission = st.begin_spawn(1, spawn_request("d-1"), tx); + st.on_reply(emission.rpc_id, Some(ack_value("d-1", 3, "prod/w1")), None); + let handle = rx.await.unwrap().unwrap().handle; + let result = DelegateResultParams { + delegation_id: "d-1".into(), + admission: 3, + status: DelegationStatus::Failed, + result: None, + error: Some("boom".into()), + }; + st.on_delegate_result(&result); + + let (atx, arx) = oneshot::channel(); + st.begin_await(&handle, atx); + let outcome = arx.await.unwrap().unwrap(); + assert_eq!(outcome.status, DelegationStatus::Failed); + assert_eq!(outcome.error.as_deref(), Some("boom")); + } + + #[tokio::test] + async fn a_stale_admission_result_is_dropped() { + let mut st = PrimaryState::new(); + let (tx, rx) = oneshot::channel(); + let emission = st.begin_spawn(1, spawn_request("d-1"), tx); + st.on_reply(emission.rpc_id, Some(ack_value("d-1", 10, "prod/w1")), None); + let _handle = rx.await.unwrap().unwrap().handle; + // A result for the same id but a different admission is not ours. + let stale = DelegateResultParams { + delegation_id: "d-1".into(), + admission: 9, + status: DelegationStatus::Completed, + result: Some("wrong".into()), + error: None, + }; + assert!( + !st.on_delegate_result(&stale), + "a stale admission is not acked as ours" + ); + } + + #[tokio::test] + async fn cancel_of_an_untracked_handle_never_emits_a_frame() { + let mut st = PrimaryState::new(); + let (tx, rx) = oneshot::channel(); + let emitted = st.begin_cancel(1, &DelegationHandle::new("never", 1), "stop".into(), tx); + assert!(emitted.is_none()); + assert!(matches!( + rx.await.unwrap(), + Err(CommandError::UnknownDelegation) + )); + } + + #[tokio::test] + async fn cancel_of_a_tracked_handle_emits_and_resolves_on_ok() { + let mut st = PrimaryState::new(); + let (tx, rx) = oneshot::channel(); + let emission = st.begin_spawn(1, spawn_request("d-1"), tx); + st.on_reply(emission.rpc_id, Some(ack_value("d-1", 3, "prod/w1")), None); + let handle = rx.await.unwrap().unwrap().handle; + + let (ctx, crx) = oneshot::channel(); + let params = st + .begin_cancel(2, &handle, "initiator gave up".into(), ctx) + .expect("a tracked handle emits a cancel frame"); + assert_eq!(params.delegation_id, "d-1"); + assert_eq!(params.admission, 3); + // CP acks the cancel. + assert!(st.on_reply(2, Some(serde_json::json!({"ok": true})), None)); + assert!(crx.await.unwrap().is_ok()); + } + + #[tokio::test] + async fn list_agents_resolves_with_the_roster() { + let mut st = PrimaryState::new(); + let (tx, rx) = oneshot::channel(); + st.begin_list_agents(1, tx); + let result = openab_cp::proto::ListAgentsResult { + namespace: "prod".into(), + agents: vec![AgentSummary { + name: "worker-1".into(), + agent_type: openab_cp::proto::AgentType::Worker, + instance_id: "i-1".into(), + labels: Default::default(), + active_sessions: 0, + max_delegated_sessions: 2, + }], + }; + assert!(st.on_reply(1, Some(serde_json::to_value(&result).unwrap()), None)); + let agents = rx.await.unwrap().unwrap(); + assert_eq!(agents.len(), 1); + assert_eq!(agents[0].name, "worker-1"); + } + #[tokio::test] + async fn check_reports_running_then_terminal_without_blocking() { + let mut st = PrimaryState::new(); + let (tx, rx) = oneshot::channel(); + let emission = st.begin_spawn(1, spawn_request("d-1"), tx); + st.on_reply(emission.rpc_id, Some(ack_value("d-1", 4, "prod/w1")), None); + let handle = rx.await.unwrap().unwrap().handle; + + // Running (ack received, no terminal yet) — never parks. + let (ctx, crx) = oneshot::channel(); + st.check(&handle, ctx); + assert_eq!( + crx.await.unwrap().unwrap(), + DelegationSnapshot::Running { + assigned_to: "prod/w1".into() + } + ); + + // Terminal after the result frame. + st.on_delegate_result(&DelegateResultParams { + delegation_id: "d-1".into(), + admission: 4, + status: DelegationStatus::Completed, + result: Some("done".into()), + error: None, + }); + let (ctx, crx) = oneshot::channel(); + st.check(&handle, ctx); + assert_eq!( + crx.await.unwrap().unwrap(), + DelegationSnapshot::Terminal(DelegationOutcome { + status: DelegationStatus::Completed, + result: Some("done".into()), + error: None, + }) + ); + } + + #[tokio::test] + async fn check_reports_pending_before_the_ack_lands() { + // A delegation whose lifecycle is Pending maps to a Pending snapshot. + // (In the running client a caller only obtains a handle at ack time, so + // this state is reached by seeding the tracking table directly — same + // module, private access — to prove the mapping is complete.) + let mut st = PrimaryState::new(); + let handle = DelegationHandle::new("d-pending", 1); + st.tracked.insert( + handle.key(), + TrackedDelegation { + state: DelegationLifecycle::Pending, + assigned_to: String::new(), + accounted_bytes: 0, + }, + ); + let (ctx, crx) = oneshot::channel(); + st.check(&handle, ctx); + assert_eq!(crx.await.unwrap().unwrap(), DelegationSnapshot::Pending); + } + + #[tokio::test] + async fn check_of_an_untracked_handle_is_unknown() { + let st = PrimaryState::new(); + let (ctx, crx) = oneshot::channel(); + st.check(&DelegationHandle::new("never", 1), ctx); + assert!(matches!( + crx.await.unwrap(), + Err(CommandError::UnknownDelegation) + )); + } + + #[tokio::test] + async fn a_foreign_reply_id_is_not_ours() { + let mut st = PrimaryState::new(); + assert!(!st.on_reply(999, Some(serde_json::json!({"ok": true})), None)); + } + + #[test] + fn terminal_history_is_bounded_by_evicting_the_oldest_terminal() { + let mut state = PrimaryState::new(); + for admission in 1..=MAX_TRACKED_DELEGATIONS as u64 { + let key = (format!("d-{admission}"), admission); + let outcome = DelegationOutcome { + status: DelegationStatus::Completed, + result: Some("done".into()), + error: None, + }; + let bytes = PrimaryState::terminal_entry_bytes(&key.0, "prod/w", &outcome); + state.tracked.insert( + key.clone(), + TrackedDelegation { + state: DelegationLifecycle::Terminal(outcome), + assigned_to: "prod/w".into(), + accounted_bytes: bytes, + }, + ); + state.terminal_bytes += bytes; + state.tracked_order.push_back(key); + } + state.make_history_room(); + let new_key = ("d-new".to_string(), 9000); + state.tracked.insert( + new_key.clone(), + TrackedDelegation { + state: DelegationLifecycle::Running { + assigned_to: "prod/w".into(), + }, + assigned_to: "prod/w".into(), + accounted_bytes: 0, + }, + ); + state.tracked_order.push_back(new_key); + assert_eq!(state.tracked.len(), MAX_TRACKED_DELEGATIONS); + assert!(!state.tracked.contains_key(&("d-1".to_string(), 1))); + } + + #[tokio::test] + async fn fail_all_live_disconnects_pending_and_running_but_not_terminal() { + let mut st = PrimaryState::new(); + + // A running delegation with a parked awaiter. + let (tx, rx) = oneshot::channel(); + let emission = st.begin_spawn(1, spawn_request("d-run"), tx); + st.on_reply( + emission.rpc_id, + Some(ack_value("d-run", 1, "prod/w1")), + None, + ); + let running = rx.await.unwrap().unwrap().handle; + let (arx_tx, arx) = oneshot::channel(); + st.begin_await(&running, arx_tx); + + // A terminal delegation: its awaiter was already answered. + let (tx2, rx2) = oneshot::channel(); + let emission2 = st.begin_spawn(2, spawn_request("d-done"), tx2); + st.on_reply( + emission2.rpc_id, + Some(ack_value("d-done", 2, "prod/w1")), + None, + ); + let done = rx2.await.unwrap().unwrap().handle; + st.on_delegate_result(&DelegateResultParams { + delegation_id: "d-done".into(), + admission: 2, + status: DelegationStatus::Completed, + result: Some("ok".into()), + error: None, + }); + + // A pending spawn awaiting its ack. + let (ptx, prx) = oneshot::channel(); + let _pending = st.begin_spawn(3, spawn_request("d-pending"), ptx); + + st.fail_all_live(); + + // The running delegation's awaiter is failed. + assert!(matches!( + arx.await.unwrap(), + Err(CommandError::Disconnected) + )); + // The pending spawn is failed. + assert!(matches!( + prx.await.unwrap(), + Err(CommandError::Disconnected) + )); + // The terminal delegation stays queryable. + assert_eq!( + st.lifecycle(&done), + Some(DelegationLifecycle::Terminal(DelegationOutcome { + status: DelegationStatus::Completed, + result: Some("ok".into()), + error: None, + })) + ); + // The running delegation is no longer tracked (it was live). + assert_eq!(st.lifecycle(&running), None); + } + + #[tokio::test] + async fn a_completed_delegation_survives_repeated_fail_all_live_reconnect_cycles() { + // Process-lifetime PrimaryState: the same instance is reused across + // reconnects. fail_all_live runs once per lost connection; a terminal + // outcome recorded before the first disconnect must remain queryable + // through every later cycle, while live/pending work is failed each + // time. This is the regression proving terminal history is preserved + // across the fail_all_live / reconnect-cycle semantics. + let mut st = PrimaryState::new(); + + // Complete a delegation while "connected". + let (tx, rx) = oneshot::channel(); + let emission = st.begin_spawn(1, spawn_request("d-done"), tx); + st.on_reply( + emission.rpc_id, + Some(ack_value("d-done", 1, "prod/w1")), + None, + ); + let done = rx.await.unwrap().unwrap().handle; + st.on_delegate_result(&DelegateResultParams { + delegation_id: "d-done".into(), + admission: 1, + status: DelegationStatus::Completed, + result: Some("done-payload".into()), + error: None, + }); + let terminal = DelegationLifecycle::Terminal(DelegationOutcome { + status: DelegationStatus::Completed, + result: Some("done-payload".into()), + error: None, + }); + assert_eq!(st.lifecycle(&done), Some(terminal.clone())); + let bytes_after_completion = st.terminal_bytes; + assert!( + bytes_after_completion > 0, + "a terminal payload is accounted" + ); + + // Simulate three reconnect cycles: each loses the connection (a live + // spawn is parked, then fail_all_live fires), then the connection is + // re-established (a new spawn is admitted and completes). + for cycle in 0..3u64 { + // A live spawn outstanding when the connection drops. + let (ptx, prx) = oneshot::channel(); + let _ = st.begin_spawn(100 + cycle, spawn_request("d-live"), ptx); + + st.fail_all_live(); + + // The live spawn is failed on every cycle. + assert!(matches!( + prx.await.unwrap(), + Err(CommandError::Disconnected) + )); + // The completed delegation survives, unchanged, with its payload + // still accounted; other cycles' terminals only add to the total. + assert_eq!(st.lifecycle(&done), Some(terminal.clone())); + assert!(st.terminal_bytes >= bytes_after_completion); + // An await against it still answers immediately from history. + let (atx, arx) = oneshot::channel(); + st.begin_await(&done, atx); + assert_eq!( + arx.await.unwrap().unwrap().result.as_deref(), + Some("done-payload") + ); + + // Reconnect: admit and complete another delegation this cycle. + let id = format!("d-cycle-{cycle}"); + let admission = 200 + cycle; + let (tx, rx) = oneshot::channel(); + let e = st.begin_spawn(300 + cycle, spawn_request(&id), tx); + st.on_reply(e.rpc_id, Some(ack_value(&id, admission, "prod/w1")), None); + let h = rx.await.unwrap().unwrap().handle; + st.on_delegate_result(&DelegateResultParams { + delegation_id: id.clone(), + admission, + status: DelegationStatus::Completed, + result: Some("ok".into()), + error: None, + }); + assert!(matches!( + st.lifecycle(&h), + Some(DelegationLifecycle::Terminal(_)) + )); + } + + // The original completed delegation is still there after all cycles. + assert_eq!(st.lifecycle(&done), Some(terminal)); + } + + #[tokio::test] + async fn an_ack_naming_a_different_delegation_id_tracks_nothing() { + // The CP acked with a delegation_id we never asked to spawn: a routing + // bug or crossed reply. The caller must be failed and NOTHING tracked, + // so a foreign id can never seed correlation under our rpc id. + let mut st = PrimaryState::new(); + let (tx, rx) = oneshot::channel(); + let emission = st.begin_spawn(1, spawn_request("d-requested"), tx); + assert!(st.on_reply( + emission.rpc_id, + Some(ack_value("d-foreign", 5, "prod/w1")), + None + )); + let outcome = rx.await.unwrap(); + assert!(matches!(outcome, Err(CommandError::Internal(_)))); + assert_eq!(st.tracked_len(), 0, "a mismatched ack tracks nothing"); + assert_eq!(st.terminal_bytes, 0); + } + + #[tokio::test] + async fn an_ack_with_the_matching_delegation_id_is_tracked() { + // The complement: an exact match is admitted normally. + let mut st = PrimaryState::new(); + let (tx, rx) = oneshot::channel(); + let emission = st.begin_spawn(1, spawn_request("d-1"), tx); + st.on_reply(emission.rpc_id, Some(ack_value("d-1", 5, "prod/w1")), None); + let handle = rx.await.unwrap().unwrap().handle; + assert_eq!( + st.lifecycle(&handle), + Some(DelegationLifecycle::Running { + assigned_to: "prod/w1".into() + }) + ); + } + + #[tokio::test] + async fn fail_pending_invalid_request_fails_the_caller_without_tracking() { + // The serve loop's oversized-outbound-Spawn path: the spawn was parked + // by begin_spawn, the serialized frame turned out too large, so the + // loop fails the parked pending request with InvalidRequest and tracks + // nothing — the connection is untouched. + let mut st = PrimaryState::new(); + let (tx, rx) = oneshot::channel(); + let emission = st.begin_spawn(7, spawn_request("d-huge"), tx); + assert!(st.fail_pending_invalid_request(emission.rpc_id, "frame too large".into())); + let outcome = rx.await.unwrap(); + assert!(matches!(outcome, Err(CommandError::InvalidRequest(_)))); + assert_eq!(st.tracked_len(), 0, "an oversized spawn tracks nothing"); + // A later reply for the same id is no longer ours (it was removed). + assert!(!st.on_reply(7, Some(ack_value("d-huge", 1, "prod/w1")), None)); + } + + #[tokio::test] + async fn fail_pending_invalid_request_is_a_noop_for_an_unknown_id() { + let mut st = PrimaryState::new(); + assert!(!st.fail_pending_invalid_request(999, "nope".into())); + } + + #[test] + fn terminal_history_is_bounded_by_retained_bytes_not_only_count() { + // A handful of very large terminal payloads must be evicted on the + // byte bound even though the count bound (4096 entries) is nowhere + // near reached. Oldest terminal goes first; live work is never touched. + let mut state = PrimaryState::new(); + // ~2 MiB result each; 16 MiB ceiling means ~8 fit, so inserting 12 + // must evict the oldest few to satisfy MAX_TRACKED_BYTES. + let big = "x".repeat(2 * 1024 * 1024); + for admission in 1..=12u64 { + let key = (format!("d-{admission}"), admission); + let outcome = DelegationOutcome { + status: DelegationStatus::Completed, + result: Some(big.clone()), + error: None, + }; + let bytes = PrimaryState::terminal_entry_bytes(&key.0, "prod/w", &outcome); + state.tracked.insert( + key.clone(), + TrackedDelegation { + state: DelegationLifecycle::Terminal(outcome), + assigned_to: "prod/w".into(), + accounted_bytes: bytes, + }, + ); + state.terminal_bytes += bytes; + state.tracked_order.push_back(key); + state.make_history_room(); + } + assert!( + state.terminal_bytes <= MAX_TRACKED_BYTES, + "retained terminal bytes ({}) must stay within the 16 MiB bound", + state.terminal_bytes + ); + assert!( + state.tracked.len() < 12, + "the byte bound must have evicted some entries (count bound was never hit)" + ); + // The oldest entries went first; the newest survives. + assert!(state.tracked.contains_key(&("d-12".to_string(), 12))); + assert!(!state.tracked.contains_key(&("d-1".to_string(), 1))); + } + + #[test] + fn out_of_order_completion_evicts_the_oldest_completion_not_oldest_admission() { + let mut state = PrimaryState::new(); + for admission in 1..=3u64 { + let key = (format!("d-{admission}"), admission); + state.tracked.insert( + key.clone(), + TrackedDelegation { + state: DelegationLifecycle::Running { + assigned_to: "prod/w".into(), + }, + assigned_to: "prod/w".into(), + accounted_bytes: 0, + }, + ); + state.tracked_order.push_back(key); + } + + // d-1 was admitted first but completes last. Three 6 MiB outcomes do + // not fit under the 16 MiB cap, so the oldest completion (d-2) must be + // evicted while the just-completed d-1 remains queryable. + let result = "x".repeat(6 * 1024 * 1024); + for admission in [2u64, 3, 1] { + assert!(state.on_delegate_result(&DelegateResultParams { + delegation_id: format!("d-{admission}"), + admission, + status: DelegationStatus::Completed, + result: Some(result.clone()), + error: None, + })); + } + + assert!(!state.tracked.contains_key(&("d-2".to_string(), 2))); + assert!(state.tracked.contains_key(&("d-3".to_string(), 3))); + assert!(state.tracked.contains_key(&("d-1".to_string(), 1))); + assert!(state.terminal_bytes <= MAX_TRACKED_BYTES); + } + + #[test] + fn a_live_delegation_is_never_evicted_by_the_byte_bound() { + // Even with the byte bound blown past, a live (Running) entry stays; + // only terminal history is evictable. + let mut state = PrimaryState::new(); + // One live entry first. + let live_key = ("d-live".to_string(), 1); + state.tracked.insert( + live_key.clone(), + TrackedDelegation { + state: DelegationLifecycle::Running { + assigned_to: "prod/w".into(), + }, + assigned_to: "prod/w".into(), + accounted_bytes: 0, + }, + ); + state.tracked_order.push_back(live_key.clone()); + // Then enough terminal bytes to exceed the ceiling. + let big = "x".repeat(9 * 1024 * 1024); + for admission in 2..=4u64 { + let key = (format!("d-{admission}"), admission); + let outcome = DelegationOutcome { + status: DelegationStatus::Completed, + result: Some(big.clone()), + error: None, + }; + let bytes = PrimaryState::terminal_entry_bytes(&key.0, "prod/w", &outcome); + state.tracked.insert( + key.clone(), + TrackedDelegation { + state: DelegationLifecycle::Terminal(outcome), + assigned_to: "prod/w".into(), + accounted_bytes: bytes, + }, + ); + state.terminal_bytes += bytes; + state.tracked_order.push_back(key); + state.make_history_room(); + } + assert!(state.terminal_bytes <= MAX_TRACKED_BYTES); + assert!( + state.tracked.contains_key(&live_key), + "the live delegation must survive byte-bound eviction" + ); + } +} diff --git a/src/main.rs b/src/main.rs index 381e748c5..c2d65e3b3 100644 --- a/src/main.rs +++ b/src/main.rs @@ -130,6 +130,14 @@ enum Commands { #[command(subcommand)] addon: McpAddon, }, + /// Agent Control Plane delegation operations over the local runtime socket. + Agent { + /// Override the local runtime socket (default: $OPENAB_AGENT_SOCKET or ~/.openab/run/agent.sock). + #[arg(long, global = true)] + socket: Option, + #[command(subcommand)] + action: AgentAction, + }, } #[derive(clap::Subcommand)] @@ -144,6 +152,39 @@ enum McpAddon { }, } +#[derive(clap::Subcommand)] +enum AgentAction { + /// Delegate a prompt by agent name or label selector. + Spawn { + /// Exact target agent name (mutually exclusive with --label). + #[arg(long)] + target: Option, + /// Label selector entry key=value; repeat for multiple labels. + #[arg(long = "label")] + labels: Vec, + #[arg(long)] + prompt: String, + #[arg(long, default_value_t = 300)] + deadline_secs: u64, + /// Return the opaque handle immediately instead of waiting for a result. + #[arg(long = "async")] + async_mode: bool, + /// Opaque parent handle when spawning a child delegation. + #[arg(long)] + parent: Option, + }, + /// List registered agents in this runtime's namespace. + List, + /// Return current state for an opaque delegation handle. + Status { handle: String }, + /// Cancel an in-flight delegation. + Cancel { + handle: String, + #[arg(long, default_value = "cancelled by operator")] + reason: String, + }, +} + #[derive(clap::Subcommand)] enum GmailNativeAction { /// Serve the adapter over loopback Streamable HTTP (non-loopback @@ -164,6 +205,166 @@ enum GmailNativeAction { }, } +async fn run_agent_command(socket: Option, action: AgentAction) -> anyhow::Result<()> { + use openab_core::control_plane::{ + default_socket_path, LocalClient, LocalRequest, LocalResponse, LocalTarget, + }; + + let client = LocalClient::new(socket.unwrap_or(default_socket_path()?)); + // A spawn is the one action bounded as a whole operation; the others are a + // single request/response. + if let AgentAction::Spawn { + target, + labels, + prompt, + deadline_secs, + async_mode, + parent, + } = action + { + let labels = parse_agent_labels(labels)?; + anyhow::ensure!( + target.is_some() ^ !labels.is_empty(), + "spawn requires exactly one of --target or one-or-more --label key=value" + ); + // Client-side deadline validation BEFORE any frame leaves the host, so + // an out-of-range value is a fast local error rather than a spawn the + // server would reject or silently clamp. Mirrors the MCP tool schema's + // `1..=1800` and the server's own bound. + validate_deadline_secs(deadline_secs)?; + let spawn = LocalRequest::Spawn { + delegation_id: new_delegation_id(), + target: LocalTarget { + name: target, + labels: (!labels.is_empty()).then_some(labels), + }, + prompt, + deadline_secs, + parent, + }; + // Same full-operation wrapper the MCP facade uses: admission + optional + // await under one `deadline_secs + 5` bound. + let client_ref = &client; + let response = + spawn_within_bound(deadline_secs, async_mode, spawn, move |req| async move { + client_ref.request(&req).await + }) + .await?; + if let LocalResponse::Error { message } = &response { + anyhow::bail!("{message}"); + } + println!("{}", serde_json::to_string_pretty(&response)?); + return Ok(()); + } + + let request = match action { + AgentAction::Spawn { .. } => unreachable!("handled above"), + AgentAction::List => LocalRequest::ListAgents, + AgentAction::Status { handle } => LocalRequest::Check { handle }, + AgentAction::Cancel { handle, reason } => LocalRequest::Cancel { handle, reason }, + }; + let response = client.request(&request).await?; + if let LocalResponse::Error { message } = &response { + anyhow::bail!("{message}"); + } + println!("{}", serde_json::to_string_pretty(&response)?); + Ok(()) +} + +/// Accepted `deadline_secs` window for `openab agent spawn`, enforced +/// client-side before any frame leaves the host. Matches the MCP tool schema +/// (`minimum: 1, maximum: 1800`) and the local API server's own bound, so all +/// three surfaces agree on the same 1-second..=30-minute window. +const MIN_DEADLINE_SECS: u64 = 1; +const MAX_DEADLINE_SECS: u64 = 1800; + +fn validate_deadline_secs(deadline_secs: u64) -> anyhow::Result<()> { + anyhow::ensure!( + (MIN_DEADLINE_SECS..=MAX_DEADLINE_SECS).contains(&deadline_secs), + "--deadline-secs must be between {MIN_DEADLINE_SECS} and {MAX_DEADLINE_SECS} (got {deadline_secs})" + ); + Ok(()) +} + +fn parse_agent_labels( + values: Vec, +) -> anyhow::Result> { + values + .into_iter() + .map(|entry| { + let (key, value) = entry + .split_once('=') + .ok_or_else(|| anyhow::anyhow!("invalid --label {entry:?}; expected key=value"))?; + anyhow::ensure!( + !key.is_empty() && !value.is_empty(), + "label key and value must be non-empty" + ); + Ok((key.to_string(), value.to_string())) + }) + .collect() +} + +use openab_core::control_plane::{LocalRequest, LocalResponse}; + +/// The one delegation-id generator shared by the MCP facade and the +/// `openab agent` CLI. A single generator keeps CLI- and MCP-initiated +/// delegations indistinguishable on the wire (no `d-cli-` vs `d-` fork) while +/// still being process-unique: `pid` separates concurrent processes, the +/// monotonic `seq` separates ids minted inside one process, and the nanosecond +/// clock separates process restarts that reuse a pid. +fn new_delegation_id() -> String { + static NEXT: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(1); + let seq = NEXT.fetch_add(1, std::sync::atomic::Ordering::Relaxed); + let nanos = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap_or_default() + .as_nanos(); + format!("d-{}-{nanos:032x}-{seq:016x}", std::process::id()) +} + +/// Run one full spawn operation under a single `deadline_secs + 5` timeout, +/// shared verbatim by the MCP facade and the `openab agent spawn` CLI. +/// +/// "Full operation" is the crux the two surfaces previously diverged on: it is +/// the admission round-trip *and*, unless `async_mode`, the blocking await of +/// the terminal result — both inside **one** [`tokio::time::timeout`]. Bounding +/// only the await (the old MCP behaviour) left a hung admission able to block +/// the caller's tool call forever; bounding neither (the old CLI behaviour) +/// left both unbounded. +/// +/// `send` performs one request/response against the local socket; it is a +/// closure so the same wrapper drives the facade's [`LocalClient`] and the +/// CLI's, and so tests can substitute an in-memory transport with no socket. +/// The `+ 5` is deliberate slack above the delegation deadline the server +/// itself enforces, so a healthy await running right up to its deadline is not +/// severed by this outer guard; only a genuinely stuck round-trip trips it. +async fn spawn_within_bound( + deadline_secs: u64, + async_mode: bool, + spawn: LocalRequest, + send: F, +) -> anyhow::Result +where + F: Fn(LocalRequest) -> Fut, + Fut: std::future::Future>, +{ + let bound = std::time::Duration::from_secs(deadline_secs.saturating_add(5)); + tokio::time::timeout(bound, async { + let spawned = send(spawn).await?; + // Async request, or the admission itself failed: hand back the first + // response verbatim without awaiting a terminal result. + if async_mode || matches!(spawned, LocalResponse::Error { .. }) { + return Ok(spawned); + } + let LocalResponse::Spawned { handle, .. } = spawned else { + anyhow::bail!("local API returned a non-spawn response"); + }; + send(LocalRequest::Await { handle }).await + }) + .await + .map_err(|_| anyhow::anyhow!("delegation exceeded the local wait bound of {bound:?}"))? +} + /// Returns true if any unified platform is enabled and its corresponding /// feature is compiled in. Google Chat uses its config-first resolver so /// `[googlechat].enabled` can activate the adapter without a duplicate env var, @@ -240,6 +441,11 @@ struct SupervisedControlPlaneTask { completion: tokio::sync::oneshot::Receiver>, } +struct SupervisedLocalApiTask { + abort: tokio::task::AbortHandle, + completion: tokio::sync::oneshot::Receiver, tokio::task::JoinError>>, +} + fn supervise_control_plane(handle: tokio::task::JoinHandle<()>) -> SupervisedControlPlaneTask { let abort = handle.abort_handle(); let (tx, completion) = tokio::sync::oneshot::channel(); @@ -249,6 +455,17 @@ fn supervise_control_plane(handle: tokio::task::JoinHandle<()>) -> SupervisedCon SupervisedControlPlaneTask { abort, completion } } +fn supervise_local_api( + handle: tokio::task::JoinHandle>, +) -> SupervisedLocalApiTask { + let abort = handle.abort_handle(); + let (tx, completion) = tokio::sync::oneshot::channel(); + tokio::spawn(async move { + let _ = tx.send(handle.await); + }); + SupervisedLocalApiTask { abort, completion } +} + async fn wait_for_control_plane_exit( task: &mut Option, ) -> anyhow::Error { @@ -262,6 +479,28 @@ async fn wait_for_control_plane_exit( } } +async fn wait_for_local_api_exit(task: &mut Option) -> anyhow::Error { + let Some(task) = task.as_mut() else { + return std::future::pending::().await; + }; + match (&mut task.completion).await { + Ok(Ok(Ok(()))) => anyhow::anyhow!("control-plane local API exited unexpectedly"), + Ok(Ok(Err(error))) => anyhow::anyhow!("control-plane local API failed: {error:#}"), + Ok(Err(error)) => anyhow::anyhow!("control-plane local API task failed: {error}"), + Err(_) => anyhow::anyhow!("control-plane local API supervisor stopped unexpectedly"), + } +} + +async fn wait_for_control_service_exit( + cp: &mut Option, + local: &mut Option, +) -> anyhow::Error { + tokio::select! { + error = wait_for_control_plane_exit(cp) => error, + error = wait_for_local_api_exit(local) => error, + } +} + /// Single LINE WORKS activation validator: the resolved (config → env → /// default) credentials must be complete and non-empty — the same rule the /// adapter constructor applies. Used by startup preflight AND cron platform @@ -463,6 +702,10 @@ async fn main() -> anyhow::Result<()> { } return Ok(()); } + Commands::Agent { socket, action } => { + run_agent_command(socket, action).await?; + return Ok(()); + } Commands::Run { config } => config, }; @@ -657,6 +900,10 @@ async fn main() -> anyhow::Result<()> { // the control-plane client is constructed much later (it needs the router), // and `cfg.agent` / `cfg.pool` are gone by then. let control_plane_cfg = cfg.control_plane.clone(); + let agent_socket_path = control_plane_cfg + .as_ref() + .map(|_| openab_core::control_plane::default_socket_path()) + .transpose()?; let prompt_hard_timeout_secs = cfg.pool.prompt_hard_timeout_secs; let pool_inner = acp::SessionPool::new( @@ -984,7 +1231,8 @@ async fn main() -> anyhow::Result<()> { // The auth key is used only for the `Authorization` header on the outbound // upgrade: it is never logged, and never reaches the agent subprocess — // `[agent].env` plumbing is untouched by this. - let mut cp_task = control_plane_cfg.map(|cp_cfg| { + let mut local_api_task: Option = None; + let mut cp_task = if let Some(cp_cfg) = control_plane_cfg { let runner: Arc = Arc::new( openab_core::control_plane::RouterPromptRunner::new(router.clone()), ); @@ -993,8 +1241,21 @@ async fn main() -> anyhow::Result<()> { runner, std::time::Duration::from_secs(prompt_hard_timeout_secs), )); - supervise_control_plane(tokio::spawn(client.run(shutdown_rx.clone()))) - }); + local_api_task = Some(supervise_local_api(tokio::spawn( + openab_core::control_plane::serve_local( + agent_socket_path + .clone() + .expect("control-plane runtime has a local socket path"), + client.handle(), + shutdown_rx.clone(), + ), + ))); + Some(supervise_control_plane(tokio::spawn( + client.run(shutdown_rx.clone()), + ))) + } else { + None + }; let dispatchers: Arc>>> = Arc::new(Mutex::new(Vec::new())); @@ -1861,7 +2122,7 @@ async fn main() -> anyhow::Result<()> { info!("discord bot running"); let discord_result = tokio::select! { result = client.start() => result, - error = wait_for_control_plane_exit(&mut cp_task) => { + error = wait_for_control_service_exit(&mut cp_task, &mut local_api_task) => { runtime_error = Some(error); Ok(()) } @@ -1889,7 +2150,7 @@ async fn main() -> anyhow::Result<()> { info!("running without discord, press ctrl+c to stop"); tokio::select! { _ = shutdown_signal() => info!("shutdown signal received"), - error = wait_for_control_plane_exit(&mut cp_task) => runtime_error = Some(error), + error = wait_for_control_service_exit(&mut cp_task, &mut local_api_task) => runtime_error = Some(error), } } // When discord feature is disabled at compile time, use this fallback block. @@ -1900,7 +2161,7 @@ async fn main() -> anyhow::Result<()> { info!("running without discord, press ctrl+c to stop"); tokio::select! { _ = shutdown_signal() => info!("shutdown signal received"), - error = wait_for_control_plane_exit(&mut cp_task) => runtime_error = Some(error), + error = wait_for_control_service_exit(&mut cp_task, &mut local_api_task) => runtime_error = Some(error), } } @@ -1932,6 +2193,25 @@ async fn main() -> anyhow::Result<()> { for d in dispatchers.lock().unwrap().iter() { d.shutdown(); } + // Stop accepting local CLI/MCP requests before the CP client and pool. + if let Some(mut task) = local_api_task { + match tokio::time::timeout(std::time::Duration::from_secs(5), &mut task.completion).await { + Ok(Ok(Ok(Ok(())))) => {} + Ok(Ok(Ok(Err(error)))) => { + tracing::warn!(error = %format!("{error:#}"), "control-plane local API failed during shutdown"); + } + Ok(Ok(Err(error))) => { + tracing::warn!(%error, "control-plane local API task failed during shutdown"); + } + Ok(Err(_)) => { + tracing::warn!("control-plane local API supervisor stopped during shutdown"); + } + Err(_) => { + tracing::warn!("control-plane local API missed the shutdown deadline — aborting"); + task.abort.abort(); + } + } + } // Stop the control-plane client BEFORE the pool: it cancels its in-flight // delegations (each stopping its agent and dropping its session) and closes // the socket, so the CP sees a clean disconnect instead of a lease timeout. @@ -2068,12 +2348,133 @@ mod tests { } } + #[test] + fn cli_agent_status_and_spawn_parse() { + let cli = Cli::try_parse_from(["openab", "agent", "status", "opaque-handle"]).unwrap(); + assert!(matches!( + cli.command, + Some(Commands::Agent { + action: AgentAction::Status { .. }, + .. + }) + )); + let cli = Cli::try_parse_from([ + "openab", + "agent", + "spawn", + "--target", + "worker-1", + "--prompt", + "review this", + "--async", + ]) + .unwrap(); + assert!(matches!( + cli.command, + Some(Commands::Agent { + action: AgentAction::Spawn { + async_mode: true, + .. + }, + .. + }) + )); + } + #[test] fn cli_setup_subcommand() { let cli = Cli::try_parse_from(["openab", "setup"]).unwrap(); assert!(matches!(cli.command.unwrap(), Commands::Setup { .. })); } + #[test] + fn cli_spawn_deadline_defaults_and_parses() { + // Default is 300 when the flag is absent. + let cli = Cli::try_parse_from([ + "openab", "agent", "spawn", "--target", "w1", "--prompt", "p", + ]) + .unwrap(); + match cli.command.unwrap() { + Commands::Agent { + action: AgentAction::Spawn { deadline_secs, .. }, + .. + } => assert_eq!(deadline_secs, 300), + _ => panic!("expected spawn"), + } + // `--deadline-secs` (clap kebab-cases the `deadline_secs` field) parses. + let cli = Cli::try_parse_from([ + "openab", + "agent", + "spawn", + "--target", + "w1", + "--prompt", + "p", + "--deadline-secs", + "42", + ]) + .unwrap(); + match cli.command.unwrap() { + Commands::Agent { + action: AgentAction::Spawn { deadline_secs, .. }, + .. + } => assert_eq!(deadline_secs, 42), + _ => panic!("expected spawn"), + } + } + + #[test] + fn cli_deadline_window_is_one_to_eighteen_hundred_inclusive() { + // Boundaries accepted. + assert!(validate_deadline_secs(1).is_ok()); + assert!(validate_deadline_secs(1800).is_ok()); + assert!(validate_deadline_secs(300).is_ok()); + // Out of range rejected on both ends, with an actionable message. + let zero = validate_deadline_secs(0).unwrap_err().to_string(); + assert!(zero.contains("between 1 and 1800"), "got {zero}"); + assert!(validate_deadline_secs(1801).is_err()); + assert!(validate_deadline_secs(u64::MAX).is_err()); + } + + #[test] + fn cli_spawn_requires_exactly_one_of_target_or_labels() { + // Exactly one of name / labels is the local API's XOR rule, enforced + // by run_agent_command before any frame leaves the host. Reproduce the + // exact predicate over the parsed pieces. + fn xor_ok(target: Option<&str>, labels: &[&str]) -> anyhow::Result<()> { + let target = target.map(str::to_string); + let labels = parse_agent_labels(labels.iter().map(|s| s.to_string()).collect())?; + anyhow::ensure!( + target.is_some() ^ !labels.is_empty(), + "spawn requires exactly one of --target or one-or-more --label key=value" + ); + Ok(()) + } + // name only → ok. + assert!(xor_ok(Some("w1"), &[]).is_ok()); + // labels only → ok. + assert!(xor_ok(None, &["backend=kiro"]).is_ok()); + assert!(xor_ok(None, &["backend=kiro", "tier=batch"]).is_ok()); + // both → rejected. + assert!(xor_ok(Some("w1"), &["backend=kiro"]).is_err()); + // neither → rejected. + assert!(xor_ok(None, &[]).is_err()); + } + + #[test] + fn cli_label_parsing_rejects_malformed_entries() { + assert!(parse_agent_labels(vec!["k=v".into()]).is_ok()); + // Missing '='. + assert!(parse_agent_labels(vec!["novalue".into()]).is_err()); + // Empty key or value. + assert!(parse_agent_labels(vec!["=v".into()]).is_err()); + assert!(parse_agent_labels(vec!["k=".into()]).is_err()); + // A parsed map keeps every entry. + let parsed = parse_agent_labels(vec!["a=1".into(), "b=2".into()]).unwrap(); + assert_eq!(parsed.len(), 2); + assert_eq!(parsed["a"], "1"); + } + #[test] fn has_unified_platform_checks_config_and_env() { // Run sequentially in one test to avoid env var race conditions @@ -2281,6 +2682,18 @@ agent_id = "1000002" assert!(error.to_string().contains("exited unexpectedly")); } + #[tokio::test] + async fn an_early_local_api_exit_is_supervised() { + let mut task = Some(supervise_local_api(tokio::spawn(async { Ok(()) }))); + let error = tokio::time::timeout( + std::time::Duration::from_secs(1), + wait_for_local_api_exit(&mut task), + ) + .await + .expect("the completed local API task is observed"); + assert!(error.to_string().contains("exited unexpectedly")); + } + #[test] fn a_control_plane_worker_boots_without_any_adapter() { let cfg = config::parse_config_str(&cp_section("worker"), "test").unwrap(); From 042bd8d67490237619be5f3dccaed5cede4ea18f Mon Sep 17 00:00:00 2001 From: Tautui Date: Sat, 31 Oct 2026 10:32:17 -0400 Subject: [PATCH 04/15] fix(cp): close review findings F1-F6 on the runtime client Port of reference c0611589, control_plane/* hunks only (cherry-pick -x not used; hunks applied from its diff). Not taken: src/main.rs (both hunks wire the F2 FacadeMode status, which needs acp_mcp.rs; skipped with acp_mcp.rs and pool.rs), acp/pool.rs, acp_mcp.rs, docs. Checks on macmini (cargo 1.99.1): - cargo check -p openab-core -p openab: pass, no warnings - cargo test -p openab-core --lib control_plane: 119 passed, 0 failed - Cargo.lock unchanged (no new dependencies) - git diff origin/main -- crates/openab-core/src/control_plane: 0 removed lines Linux run for this commit: not run (the plan assigns the Linux fix to commit 5). --- .../openab-core/src/control_plane/client.rs | 387 +++++++++++++-- .../openab-core/src/control_plane/executor.rs | 118 ++++- crates/openab-core/src/control_plane/local.rs | 381 +++++++++++++-- crates/openab-core/src/control_plane/mod.rs | 3 +- .../openab-core/src/control_plane/primary.rs | 443 +++++++++++++++--- 5 files changed, 1176 insertions(+), 156 deletions(-) diff --git a/crates/openab-core/src/control_plane/client.rs b/crates/openab-core/src/control_plane/client.rs index d4bc0404c..149278957 100644 --- a/crates/openab-core/src/control_plane/client.rs +++ b/crates/openab-core/src/control_plane/client.rs @@ -70,6 +70,15 @@ const CONNECT_TIMEOUT: Duration = Duration::from_secs(10); /// acks must land in backoff, not hang the client until shutdown. const REGISTER_TIMEOUT: Duration = Duration::from_secs(10); +/// Bound on ONE WebSocket write. The serve loop is the single sink owner and +/// awaits every write inline, so an unbounded `sink.send` under socket +/// backpressure would freeze heartbeats, inbound handling, completions, +/// command traffic and graceful shutdown together (F4). The CP server bounds +/// its writer with `write_timeout_secs` (default 30s); the runtime mirrors +/// that so a stalled peer is classified as disconnected and recovered by the +/// reconnect path instead of wedging the process. +const WRITE_TIMEOUT: Duration = Duration::from_secs(30); + type Ws = WebSocketStream>; /// Write half. Split from the read half because the serve loop must be able to /// write from a handler while the read future is still alive — one `select!` @@ -343,7 +352,9 @@ impl ControlPlaneClient { }; let frame = JsonRpcRequest::new(id, methods::REGISTER, Some(serde_json::to_value(¶ms)?)); - send(sink, &frame).await?; + // Bounded but not shutdown-raced: the caller's connect/register select + // already races the whole registration against shutdown. + write_bounded(sink, Message::Text(frame_text(&frame)?)).await?; // Anything other than the ack to this id is a protocol violation at // this point: registration is the first frame in both directions. @@ -424,6 +435,10 @@ impl ControlPlaneClient { tokio::select! { _ = shutdown.changed() => break Outcome::Shutdown, _ = heartbeat.tick() => { + // Age out parked requests/awaiters a withholding CP would + // otherwise let accumulate (F3). Piggybacks on the tick so + // the bound holds with no local caller involved. + primary.prune(std::time::Instant::now(), chrono::Utc::now()); let params = HeartbeatParams { instance_id: self.instance_id.clone(), active_delegated_sessions: self.executor.active(), @@ -433,8 +448,8 @@ impl ControlPlaneClient { methods::HEARTBEAT, Some(serde_json::to_value(¶ms)?), ); - if send(&mut sink, &frame).await.is_err() { - break Outcome::Disconnected; + if let Err(end) = send(&mut sink, &frame, shutdown).await { + break end.into(); } } // A primary-side command: the local initiating surface (UDS @@ -444,8 +459,8 @@ impl ControlPlaneClient { maybe_command = recv_command(command_rx) => { match maybe_command { Some(command) => { - if self.handle_command(command, primary, &mut sink).await.is_err() { - break Outcome::Disconnected; + if let Err(end) = self.handle_command(command, primary, &mut sink, shutdown).await { + break end.into(); } } // The command channel closed (no handles left). Stop @@ -459,8 +474,8 @@ impl ControlPlaneClient { // frame that closes the initiator's wait. Some(result) = result_rx.recv() => { let frame = delegate_result_request(self.next_id(), result)?; - if send(&mut sink, &frame).await.is_err() { - break Outcome::Disconnected; + if let Err(end) = send(&mut sink, &frame, shutdown).await { + break end.into(); } } inbound = stream.next() => { @@ -469,33 +484,53 @@ impl ControlPlaneClient { Ok(Message::Text(text)) => { match self.classify_inbound(&text, primary) { FrameAction::Serve { ack, forward } => { - let executor = Arc::clone(&self.executor); - let tx = result_tx.clone(); - serving.push(tokio::spawn(async move { - let result = executor.serve(forward).await; - // A closed channel means the connection - // that would carry this result is gone; - // the CP fails it as target_disconnected. - let _ = tx.send(result).await; - })); + // Admit HERE, on the loop's own task, before + // anything yields (F1). A `cp/cancel` that is + // the very next frame then finds the + // reservation instead of racing a spawned + // task that has not been polled yet. + let admitted = self.executor.admit_delegation(forward); // Prune finished tasks so a long-lived // connection does not accumulate handles. serving.retain(|h| !h.is_finished()); - if sink.send(Message::Text(ack)).await.is_err() { - break Outcome::Disconnected; + let refused = match admitted { + Ok(admitted) => { + let tx = result_tx.clone(); + serving.push(tokio::spawn(async move { + let result = admitted.run().await; + // A closed channel means the connection + // that would carry this result is gone; + // the CP fails it as target_disconnected. + let _ = tx.send(result).await; + })); + None + } + Err(refused) => Some(refused), + }; + if let Err(end) = write_or_shutdown(&mut sink, Message::Text(ack), shutdown).await { + break end.into(); + } + // A refusal is answered inline, after the ack, + // so the initiator is not left waiting on its + // deadline for a runtime that never started. + if let Some(refused) = refused { + let frame = delegate_result_request(self.next_id(), refused)?; + if let Err(end) = send(&mut sink, &frame, shutdown).await { + break end.into(); + } } } FrameAction::Reply(reply) => { - if sink.send(Message::Text(reply)).await.is_err() { - break Outcome::Disconnected; + if let Err(end) = write_or_shutdown(&mut sink, Message::Text(reply), shutdown).await { + break end.into(); } } FrameAction::Ignore => {} } } Ok(Message::Ping(p)) => { - if sink.send(Message::Pong(p)).await.is_err() { - break Outcome::Disconnected; + if let Err(end) = write_or_shutdown(&mut sink, Message::Pong(p), shutdown).await { + break end.into(); } } Ok(Message::Close(_)) => { @@ -546,21 +581,25 @@ impl ControlPlaneClient { handle.abort(); } } - let _ = sink.send(Message::Close(None)).await; - let _ = sink.close().await; + // Best-effort close handshake, bounded: a peer that stopped reading + // must not hold up reconnect or process shutdown. + let _ = write_bounded(&mut sink, Message::Close(None)).await; + let _ = tokio::time::timeout(WRITE_TIMEOUT, sink.close()).await; Ok(outcome) } /// Turn one primary-side command into an outbound frame, parking its reply /// in `primary` for later correlation. `Await` needs no frame — it only /// registers interest in a terminal that arrives over the socket. Returns - /// `Err` only when the socket write fails, which ends the session. + /// `Err` only when the socket write fails, stalls, or shutdown interrupts + /// it, each of which ends the session. async fn handle_command( &self, command: ClientCommand, primary: &mut PrimaryState, sink: &mut WsSink, - ) -> anyhow::Result<()> { + shutdown: &mut watch::Receiver, + ) -> Result<(), WriteEnd> { // A primary must be registered as such to initiate; a worker-only // runtime rejects spawn/cancel/list locally so a misconfigured caller // gets a clear answer instead of a CP policy denial round-trip. Await @@ -571,19 +610,27 @@ impl ControlPlaneClient { let _ = reply.send(Err(primary_only_error())); return Ok(()); } - let emission = primary.begin_spawn(self.next_id(), request, reply); + // `None`: the pending cap is reached and the caller was already + // answered `Overloaded`; nothing to send. + let Some(emission) = primary.begin_spawn(self.next_id(), request, reply) else { + return Ok(()); + }; let frame = JsonRpcRequest::new( emission.rpc_id, methods::DELEGATE, - Some(serde_json::to_value(&emission.params)?), + Some(encode(&emission.params)?), ); // A Spawn whose complete serialized frame exceeds the transport // ceiling must NOT tear down the connection: an oversized prompt // is a caller-side fault, not a socket fault. Detect it before // writing, remove/fail the parked pending request with // InvalidRequest, and return Ok so the serve loop keeps running. - if !frame_within_limit(&frame)? { - let bytes = serialized_len(&frame)?; + let fits = frame_within_limit(&frame).map_err(|e| { + error!(error = %e, "failed to measure a cp/delegate frame"); + WriteEnd::Disconnected + })?; + if !fits { + let bytes = serialized_len(&frame).unwrap_or(usize::MAX); warn!( rpc_id = emission.rpc_id, frame_bytes = bytes, @@ -599,7 +646,7 @@ impl ControlPlaneClient { ); return Ok(()); } - send(sink, &frame).await?; + send(sink, &frame, shutdown).await?; } ClientCommand::Cancel { handle, @@ -612,22 +659,21 @@ impl ControlPlaneClient { } let rpc_id = self.next_id(); if let Some(params) = primary.begin_cancel(rpc_id, &handle, reason, reply) { - let frame = JsonRpcRequest::new( - rpc_id, - methods::CANCEL, - Some(serde_json::to_value(¶ms)?), - ); - send(sink, &frame).await?; + let frame = + JsonRpcRequest::new(rpc_id, methods::CANCEL, Some(encode(¶ms)?)); + send(sink, &frame, shutdown).await?; } // begin_cancel already answered the caller if the handle was // unknown; no frame in that case. } ClientCommand::ListAgents { reply } => { let rpc_id = self.next_id(); - primary.begin_list_agents(rpc_id, reply); + if !primary.begin_list_agents(rpc_id, reply) { + return Ok(()); + } let frame = JsonRpcRequest::new(rpc_id, methods::LIST_AGENTS, Some(serde_json::json!({}))); - send(sink, &frame).await?; + send(sink, &frame, shutdown).await?; } ClientCommand::Await { handle, reply } => { primary.begin_await(&handle, reply); @@ -860,7 +906,10 @@ fn delegate_result_request( Ok(frame) } -async fn send(sink: &mut WsSink, frame: &JsonRpcRequest) -> anyhow::Result<()> { +/// Serialize an outbound frame and enforce the transport ceiling. The size +/// check is the same measurement [`serialized_len`] and the oversized-Spawn +/// guard use, so the check and the wire agree by construction. +fn frame_text(frame: &JsonRpcRequest) -> anyhow::Result { let text = serde_json::to_string(frame)?; if text.len() > MAX_FRAME_BYTES { anyhow::bail!( @@ -869,8 +918,93 @@ async fn send(sink: &mut WsSink, frame: &JsonRpcRequest) -> anyhow::Result<()> { MAX_FRAME_BYTES ); } - sink.send(Message::Text(text)).await?; - Ok(()) + Ok(text) +} + +/// One bounded WebSocket write. Generic over the sink so the stall path can be +/// exercised with a sink that never becomes ready. +async fn write_bounded(sink: &mut S, msg: Message) -> anyhow::Result<()> +where + S: futures_util::Sink + Unpin, + S::Error: std::fmt::Display, +{ + match tokio::time::timeout(WRITE_TIMEOUT, sink.send(msg)).await { + Ok(Ok(())) => Ok(()), + Ok(Err(e)) => Err(anyhow::anyhow!("control-plane socket write failed: {e}")), + Err(_) => Err(anyhow::anyhow!( + "control-plane socket write stalled for {}s; treating the connection as lost", + WRITE_TIMEOUT.as_secs() + )), + } +} + +/// Why a serve-loop write did not complete. Both end the session; they differ +/// in what the loop reports (and therefore whether it reconnects). +#[derive(Debug, PartialEq, Eq)] +enum WriteEnd { + /// Shutdown was requested while the write was pending. + Shutdown, + /// The write failed or exceeded [`WRITE_TIMEOUT`]. + Disconnected, +} + +impl From for Outcome { + fn from(end: WriteEnd) -> Self { + match end { + WriteEnd::Shutdown => Outcome::Shutdown, + WriteEnd::Disconnected => Outcome::Disconnected, + } + } +} + +/// Bounded write raced against shutdown: the serve loop's single-owner sink +/// must never make graceful shutdown wait on a stalled peer. +async fn write_or_shutdown( + sink: &mut S, + msg: Message, + shutdown: &mut watch::Receiver, +) -> Result<(), WriteEnd> +where + S: futures_util::Sink + Unpin, + S::Error: std::fmt::Display, +{ + tokio::select! { + biased; + _ = shutdown.changed() => Err(WriteEnd::Shutdown), + res = write_bounded(sink, msg) => res.map_err(|e| { + warn!(error = %e, "control-plane write ended the session"); + WriteEnd::Disconnected + }), + } +} + +/// Serve-loop frame write: size-checked, bounded, shutdown-racing. An +/// oversized frame is a programming error at this layer (payloads are capped +/// upstream), so it ends the session like a socket fault would. +async fn send( + sink: &mut S, + frame: &JsonRpcRequest, + shutdown: &mut watch::Receiver, +) -> Result<(), WriteEnd> +where + S: futures_util::Sink + Unpin, + S::Error: std::fmt::Display, +{ + let text = frame_text(frame).map_err(|e| { + error!(error = %e, "refusing to write an oversized control-plane frame"); + WriteEnd::Disconnected + })?; + write_or_shutdown(sink, Message::Text(text), shutdown).await +} + +/// Encode command params for the wire. Serializing our own wire structs cannot +/// realistically fail; if it does, the session ends rather than the frame +/// silently not being sent. +fn encode(params: &T) -> Result { + serde_json::to_value(params).map_err(|e| { + error!(error = %e, "failed to encode a control-plane frame"); + WriteEnd::Disconnected + }) } /// Serialized byte length of an outbound frame. Pure and testable — the same @@ -980,6 +1114,7 @@ fn error_reply(id: u64, error: ErrorObject) -> FrameAction { } /// Why a connection's serve loop ended. +#[derive(Debug, PartialEq, Eq)] enum Outcome { /// The process is shutting down; do not reconnect. Shutdown, @@ -1012,6 +1147,7 @@ mod tests { use crate::config::CpAgentType; use crate::control_plane::executor::PromptOutcome; use async_trait::async_trait; + use openab_cp::proto::DelegationStatus; struct NoopRunner; @@ -1361,7 +1497,9 @@ max_delegated_sessions = 3 deadline: chrono::Utc::now() + chrono::Duration::seconds(60), parent: None, }; - let emission = primary.begin_spawn(1, request, tx); + let emission = primary + .begin_spawn(1, request, tx) + .expect("under the pending cap"); let frame = JsonRpcRequest::new( emission.rpc_id, methods::DELEGATE, @@ -1375,4 +1513,167 @@ max_delegated_sessions = 3 let outcome = rx.await.unwrap(); assert!(matches!(outcome, Err(CommandError::InvalidRequest(_)))); } + + // ---- F1: cancel-before-admission ------------------------------------- + + /// A runner that records whether it was ever started. If the F1 race were + /// still open — admission on the spawned task's first poll — a cancel + /// classified before that poll would be acked as unknown and the runner + /// would then start anyway. + struct RecordingRunner { + started: std::sync::atomic::AtomicBool, + } + + #[async_trait] + impl PromptRunner for RecordingRunner { + async fn run( + &self, + _session_key: &str, + _forward: &DelegateForward, + ) -> anyhow::Result { + self.started + .store(true, std::sync::atomic::Ordering::SeqCst); + Ok(PromptOutcome::default()) + } + async fn cancel(&self, _session_key: &str) {} + async fn discard(&self, _session_key: &str) {} + } + + #[tokio::test] + async fn a_cancel_handled_before_the_serving_task_is_polled_never_starts_the_runner() { + // Reproduce the serve loop's ordering exactly: classify `cp/delegate` + // (which now admits synchronously), classify `cp/cancel` as the very + // next frame, and only THEN let the serving future run. `handle_frame` + // is the loop's classifier; `admit_delegation` + `Admitted::run` is + // what the `Serve` arm does around the spawn. + let runner = Arc::new(RecordingRunner { + started: std::sync::atomic::AtomicBool::new(false), + }); + let c = ControlPlaneClient::new(cfg(), Arc::clone(&runner) as _, Duration::from_secs(60)); + let delegate = serde_json::json!({"jsonrpc":"2.0","id":9,"method":"cp/delegate","params":{ + "delegation_id":"d-race","admission":7,"prompt":"hi", + "deadline":(chrono::Utc::now()+chrono::Duration::seconds(60)).to_rfc3339(), + "from":"prod/koudu","chain":["prod/koudu"]}}) + .to_string(); + let FrameAction::Serve { forward, .. } = c.handle_frame(&delegate) else { + panic!("expected a served delegation"); + }; + // The loop admits before it yields. + let admitted = c + .executor + .admit_delegation(forward) + .expect("capacity 3, nothing in flight"); + assert_eq!(c.executor.active(), 1, "the slot is reserved immediately"); + + // Next frame: the cancel. It must find the reservation. + let cancel = serde_json::json!({ + "jsonrpc": "2.0", "id": 10, "method": "cp/cancel", + "params": {"delegation_id": "d-race", "admission": 7, "reason": "initiator gave up"} + }) + .to_string(); + let FrameAction::Reply(reply) = c.handle_frame(&cancel) else { + panic!("expected a cancel ack"); + }; + assert!(reply.contains("\"ok\":true")); + + // Only now does the serving future get its first poll. + let result = admitted.run().await; + assert_eq!(result.status, DelegationStatus::Cancelled); + assert!( + !runner.started.load(std::sync::atomic::Ordering::SeqCst), + "a cancelled delegation must never reach the prompt runner" + ); + assert_eq!(c.executor.active(), 0, "the slot is released after the run"); + } + + #[tokio::test] + async fn an_admitted_delegation_dropped_without_running_releases_its_slot() { + let c = client(); + let forward = DelegateForward { + delegation_id: "d-drop".into(), + admission: 1, + prompt: "hi".into(), + deadline: chrono::Utc::now() + chrono::Duration::seconds(60), + from: "prod/koudu".into(), + chain: vec!["prod/koudu".into()], + }; + let admitted = c.executor.admit_delegation(forward).unwrap(); + assert_eq!(c.executor.active(), 1); + drop(admitted); + assert_eq!( + c.executor.active(), + 0, + "an ack write failure after admission must not leak the slot" + ); + } + + // ---- F4: bounded writes --------------------------------------------- + + /// A sink that is never ready: the shape of a peer that stopped reading + /// while the kernel buffers are full. + struct StalledSink; + + impl futures_util::Sink for StalledSink { + type Error = std::io::Error; + fn poll_ready( + self: std::pin::Pin<&mut Self>, + _cx: &mut std::task::Context<'_>, + ) -> std::task::Poll> { + std::task::Poll::Pending + } + fn start_send(self: std::pin::Pin<&mut Self>, _item: Message) -> Result<(), Self::Error> { + unreachable!("never ready") + } + fn poll_flush( + self: std::pin::Pin<&mut Self>, + _cx: &mut std::task::Context<'_>, + ) -> std::task::Poll> { + std::task::Poll::Pending + } + fn poll_close( + self: std::pin::Pin<&mut Self>, + _cx: &mut std::task::Context<'_>, + ) -> std::task::Poll> { + std::task::Poll::Pending + } + } + + #[tokio::test(start_paused = true)] + async fn a_stalled_write_is_classified_as_disconnected_after_the_write_timeout() { + let (_tx, mut shutdown) = watch::channel(false); + let frame = JsonRpcRequest::new(1, methods::HEARTBEAT, Some(serde_json::json!({}))); + let mut sink = StalledSink; + let writing = tokio::spawn(async move { send(&mut sink, &frame, &mut shutdown).await }); + tokio::task::yield_now().await; + // Just short of the bound: still pending. + tokio::time::advance(WRITE_TIMEOUT - Duration::from_millis(1)).await; + tokio::task::yield_now().await; + assert!( + !writing.is_finished(), + "the write is still within its bound" + ); + tokio::time::advance(Duration::from_millis(2)).await; + let end = writing.await.unwrap(); + assert_eq!(end, Err(WriteEnd::Disconnected)); + assert_eq!(Outcome::from(WriteEnd::Disconnected), Outcome::Disconnected); + } + + #[tokio::test(start_paused = true)] + async fn shutdown_wins_over_a_stalled_write() { + let (tx, mut shutdown) = watch::channel(false); + let mut sink = StalledSink; + let writing = tokio::spawn(async move { + write_or_shutdown(&mut sink, Message::Pong(vec![]), &mut shutdown).await + }); + tokio::task::yield_now().await; + tokio::time::advance(Duration::from_secs(1)).await; + assert!(!writing.is_finished()); + tx.send(true).unwrap(); + let end = tokio::time::timeout(Duration::from_secs(1), writing) + .await + .expect("shutdown must not wait on the stalled peer") + .unwrap(); + assert_eq!(end, Err(WriteEnd::Shutdown)); + assert_eq!(Outcome::from(WriteEnd::Shutdown), Outcome::Shutdown); + } } diff --git a/crates/openab-core/src/control_plane/executor.rs b/crates/openab-core/src/control_plane/executor.rs index 96fb04939..8793a2542 100644 --- a/crates/openab-core/src/control_plane/executor.rs +++ b/crates/openab-core/src/control_plane/executor.rs @@ -162,8 +162,10 @@ impl DelegationExecutor { delegation_id: &str, admission: AdmissionToken, ) -> std::result::Result, Refusal> { - let max = self.effective_max(); let mut g = self.lock_inflight(); + // Read the ceiling under the same lock as the count so a concurrent + // re-register ack cannot make one decision against a stale pair. + let max = self.effective_max(); if g.contains_key(delegation_id) { return Err(Refusal::Duplicate); } @@ -234,31 +236,54 @@ impl DelegationExecutor { } } - /// Admit, run, and classify one forwarded delegation. + /// Reserve the slot for one forwarded delegation — SYNCHRONOUSLY, on the + /// caller's task — or explain why not as a ready-to-send result payload. /// - /// Always resolves to a result frame payload — the refusal paths included, - /// so the initiator is never left waiting on its deadline for a runtime - /// that had already decided not to run. - pub async fn serve(self: Arc, forward: DelegateForward) -> DelegateResultParams { + /// Admission is deliberately split from execution ([`Admitted::run`]) so + /// the connection loop can reserve the slot *before* it yields. If the loop + /// instead spawned a task that admitted on its first poll, a `cp/cancel` + /// arriving as the very next frame would find nothing in flight, be acked + /// as "unknown", and the still-unpolled task would then admit and run work + /// the initiator had already cancelled — CP frame order does not imply + /// Tokio task scheduling order. With admission on the loop's own task, + /// every later frame observes the reservation. + /// + /// The refusal paths resolve to a result frame payload rather than an + /// error so the initiator is never left waiting on its deadline for a + /// runtime that had already decided not to run. + pub fn admit_delegation( + self: &Arc, + forward: DelegateForward, + ) -> std::result::Result { let id = forward.delegation_id.clone(); - let cancel = match self.admit(&id, forward.admission) { - Ok(signal) => signal, + match self.admit(&id, forward.admission) { + Ok(cancel) => Ok(Admitted { + reservation: SlotReservation { + executor: Arc::clone(self), + id, + }, + forward, + cancel, + }), Err(refusal) => { let error = refusal.message(); warn!(delegation_id = %id, from = %forward.from, %error, "delegation refused"); - return failed(&id, forward.admission, error); + Err(failed(&id, forward.admission, error)) } - }; - // RAII: the slot must free even if this task is ABORTED mid-await — - // the client aborts serving tasks that outlive the drain window on - // disconnect, and a plain post-await release would be skipped there, - // leaking the inflight entry forever (with the default cap of 1, the - // worker would refuse every delegation after reconnecting). - let _slot = SlotGuard { - executor: self.as_ref(), - id: id.clone(), - }; - self.execute(&forward, cancel).await + } + } + + /// Admit, run, and classify one forwarded delegation in one step. + /// + /// Convenience for callers that do not need to separate admission from + /// execution (tests, single-shot drivers). The connection loop uses + /// [`Self::admit_delegation`] + [`Admitted::run`] so admission happens + /// before it yields — see that method for why. + pub async fn serve(self: Arc, forward: DelegateForward) -> DelegateResultParams { + match self.admit_delegation(forward) { + Ok(admitted) => admitted.run().await, + Err(refused) => refused, + } } async fn execute( @@ -485,17 +510,60 @@ fn cap_result(text: String) -> String { /// Frees a delegation's inflight slot on drop — including the drop that /// happens when the serving task is aborted at an await point. -struct SlotGuard<'a> { - executor: &'a DelegationExecutor, +/// Owned RAII reservation of one inflight slot. Dropping it releases the +/// slot; it is created only by [`DelegationExecutor::admit_delegation`] after +/// a successful `admit`, so the release always pairs with a real insert. +struct SlotReservation { + executor: Arc, id: String, } -impl Drop for SlotGuard<'_> { +impl Drop for SlotReservation { fn drop(&mut self) { self.executor.release(&self.id); } } +/// A delegation whose slot is reserved but whose turn has not started. +/// +/// Produced by [`DelegationExecutor::admit_delegation`]; consumed by +/// [`Admitted::run`]. Owning this value IS holding the slot: dropping it +/// without running (or aborting the task that is running it) releases the +/// reservation — the same RAII guarantee the previous single-step `serve` +/// gave, now spanning the gap between admission on the connection loop and +/// execution on the spawned task. +#[must_use = "dropping an admitted delegation releases its slot without running it"] +pub struct Admitted { + reservation: SlotReservation, + forward: DelegateForward, + cancel: Arc, +} + +impl Admitted { + /// The id this reservation is for. + pub fn delegation_id(&self) -> &str { + &self.forward.delegation_id + } + + /// Run the admitted turn to a result payload. The slot is released when + /// this future completes or is dropped mid-await (task abort included — + /// the client aborts serving tasks that outlive the drain window on + /// disconnect, and a plain post-await release would be skipped there, + /// leaking the inflight entry forever; with the default cap of 1 the + /// worker would then refuse every delegation after reconnecting). + pub async fn run(self) -> DelegateResultParams { + let Admitted { + reservation, + forward, + cancel, + } = self; + let executor = Arc::clone(&reservation.executor); + // Held across the await so an abort releases the slot. + let _slot = reservation; + executor.execute(&forward, cancel).await + } +} + fn failed( delegation_id: &str, admission: AdmissionToken, @@ -878,8 +946,8 @@ mod tests { assert!(poisoned.is_err()); assert_eq!(ex.active(), 0); { - let _slot = SlotGuard { - executor: ex.as_ref(), + let _slot = SlotReservation { + executor: Arc::clone(&ex), id: "missing".into(), }; } diff --git a/crates/openab-core/src/control_plane/local.rs b/crates/openab-core/src/control_plane/local.rs index 81a7151e2..a75eb2c4b 100644 --- a/crates/openab-core/src/control_plane/local.rs +++ b/crates/openab-core/src/control_plane/local.rs @@ -158,6 +158,7 @@ mod unix_impl { use std::os::unix::fs::{FileTypeExt, MetadataExt, PermissionsExt}; use std::path::{Path, PathBuf}; use std::sync::{Arc, Mutex}; + use std::time::Duration; use rand::RngCore; use tokio::io::{AsyncBufRead, AsyncBufReadExt, AsyncWriteExt, BufReader}; @@ -267,11 +268,18 @@ mod unix_impl { // Owner-only on the socket node itself, not merely via the directory. std::fs::set_permissions(&path, std::fs::Permissions::from_mode(0o600)) .map_err(|e| anyhow::anyhow!("securing local socket {path:?}: {e}"))?; + // Remember WHICH node we bound so shutdown never unlinks a successor + // that took the pathname after us (F6). + let ours = node_identity(&path); info!(socket = %path.display(), "control-plane local IPC server listening"); // One registry for the whole server: tokens minted on any connection // resolve on any other. let registry: SharedRegistry = Arc::new(Mutex::new(HandleRegistry::default())); + // Every connection handler is tracked here so (a) their number is + // bounded by MAX_LOCAL_CONNECTIONS and (b) shutdown aborts them + // instead of leaving detached tasks blocked on replies (F3). + let mut handlers: tokio::task::JoinSet<()> = tokio::task::JoinSet::new(); loop { tokio::select! { @@ -279,12 +287,22 @@ mod unix_impl { info!("control-plane local IPC server shutting down"); break; } + // Reap finished handlers so the set's length is the live count. + Some(_) = handlers.join_next(), if !handlers.is_empty() => {} accepted = listener.accept() => { match accepted { Ok((stream, _addr)) => { + if handlers.len() >= MAX_LOCAL_CONNECTIONS { + warn!( + limit = MAX_LOCAL_CONNECTIONS, + "local IPC connection limit reached; refusing a new connection" + ); + handlers.spawn(refuse_conn(stream)); + continue; + } let cp = cp.clone(); let registry = Arc::clone(®istry); - tokio::spawn(async move { + handlers.spawn(async move { if let Err(e) = handle_conn(stream, cp, registry).await { debug!(error = %format!("{e:#}"), "local IPC connection ended with error"); } @@ -297,11 +315,39 @@ mod unix_impl { } } } - // Best-effort cleanup so the next boot does not trip over the node. - let _ = std::fs::remove_file(&path); + // Handlers still parked on a CP reply have no route to an answer once + // the runtime is leaving; end them now rather than at process exit. + handlers.abort_all(); + while handlers.join_next().await.is_some() {} + // Best-effort cleanup so the next boot does not trip over the node — + // but only if the path still names OUR node. + remove_if_same_node(&path, ours); Ok(()) } + /// Ceiling on concurrently served local connections. Each handler is one + /// task that may be parked on a CP reply for up to the request bound + /// below; without a cap a local caller (or a stuck script) could pin an + /// unbounded number of them. + const MAX_LOCAL_CONNECTIONS: usize = 64; + + /// Answer a connection refused at the cap with one error line so a CLI + /// gets a diagnostic instead of a bare EOF, then close it. + async fn refuse_conn(stream: UnixStream) { + let (_read_half, mut write_half) = stream.into_split(); + let response = LocalResponse::error(format!( + "too many concurrent local connections (limit {MAX_LOCAL_CONNECTIONS}); retry shortly" + )); + if let Ok(mut encoded) = serde_json::to_string(&response) { + encoded.push('\n'); + let _ = tokio::time::timeout(LOCAL_IO_TIMEOUT, async { + write_half.write_all(encoded.as_bytes()).await?; + write_half.flush().await + }) + .await; + } + } + /// Prepare the socket's parent directory. /// /// Two cases, deliberately asymmetric: @@ -406,6 +452,14 @@ mod unix_impl { /// node owned by someone else is refused rather than unlinked, so a hostile /// or accidental entry planted at the socket path can never trick the /// runtime into deleting an unrelated file. + /// + /// Ownership alone does not make a node stale: a second runtime under the + /// same uid (several bots on one host sharing `$HOME`) would otherwise + /// unlink its sibling's LIVE socket and bind over it (F6). So the node is + /// probed with a connect: a listener that answers is live and the bind is + /// refused with a pointer at `OPENAB_AGENT_SOCKET`; `ECONNREFUSED` is the + /// kernel telling us nobody is listening, which is the only state that is + /// unlinked. fn remove_stale_socket(path: &Path) -> anyhow::Result<()> { let meta = match std::fs::symlink_metadata(path) { Ok(m) => m, @@ -432,10 +486,67 @@ mod unix_impl { meta.uid() )); } + match std::os::unix::net::UnixStream::connect(path) { + Ok(_probe) => { + return Err(anyhow::anyhow!( + "another runtime is already serving {path:?}; refusing to replace a live \ + socket — run one runtime per socket path, or point OPENAB_AGENT_SOCKET at \ + a distinct path for this instance" + )); + } + Err(e) if e.kind() == std::io::ErrorKind::ConnectionRefused => { + // Nobody is listening: the classic leftover from an unclean exit. + } + Err(e) => { + // Anything else (permission, backlog full, ...) is not evidence + // of staleness, so do not unlink on it. + return Err(anyhow::anyhow!( + "cannot tell whether socket {path:?} is live ({e}); refusing to remove it" + )); + } + } std::fs::remove_file(path) .map_err(|e| anyhow::anyhow!("removing stale socket {path:?}: {e}")) } + /// Identity of the node this server bound: enough to tell OUR socket from a + /// successor that later took the same pathname. + #[derive(Debug, Clone, Copy, PartialEq, Eq)] + struct NodeIdentity { + dev: u64, + ino: u64, + } + + fn node_identity(path: &Path) -> Option { + let meta = std::fs::symlink_metadata(path).ok()?; + if !meta.file_type().is_socket() { + return None; + } + Some(NodeIdentity { + dev: meta.dev(), + ino: meta.ino(), + }) + } + + /// Unlink `path` only if it is still the node THIS server created. An older + /// runtime shutting down after a successor bound the same pathname must + /// not remove the successor's socket (F6); with the liveness probe above + /// that succession is refused, but the check keeps cleanup correct even if + /// the path was replaced by other means. + fn remove_if_same_node(path: &Path, ours: Option) { + let Some(ours) = ours else { + return; + }; + if node_identity(path) == Some(ours) { + let _ = std::fs::remove_file(path); + } else { + debug!( + socket = %path.display(), + "socket path no longer names the node this server bound; leaving it in place" + ); + } + } + /// Maximum length, in bytes, of a single newline-delimited protocol line /// (request or response). A local caller (or a peer answering the client) /// that never sends a newline must not be able to drive the reader's buffer @@ -514,6 +625,11 @@ mod unix_impl { /// One client connection: read request lines, answer each with one /// response line. A malformed line is answered with an `Error` rather than /// dropping the connection, so a CLI gets a diagnostic. + /// + /// Every wait is bounded (F3): the read for the next request line by + /// [`LOCAL_IDLE_TIMEOUT`], the response write by [`LOCAL_IO_TIMEOUT`], and + /// each CP round trip inside [`dispatch`] by its own bound — so a handler + /// task always ends, whatever the peer or the CP does. async fn handle_conn( stream: UnixStream, cp: ControlPlaneHandle, @@ -521,7 +637,21 @@ mod unix_impl { ) -> anyhow::Result<()> { let (read_half, mut write_half) = stream.into_split(); let mut reader = BufReader::new(read_half); - while let Some(line) = read_line_bounded(&mut reader).await? { + loop { + let line = match tokio::time::timeout( + LOCAL_IDLE_TIMEOUT, + read_line_bounded(&mut reader), + ) + .await + { + Ok(Ok(Some(line))) => line, + Ok(Ok(None)) => return Ok(()), + Ok(Err(e)) => return Err(e), + Err(_) => anyhow::bail!( + "local IPC connection idle for {}s; closing", + LOCAL_IDLE_TIMEOUT.as_secs() + ), + }; if line.trim().is_empty() { continue; } @@ -531,10 +661,54 @@ mod unix_impl { }; let mut encoded = serde_json::to_string(&response)?; encoded.push('\n'); - write_half.write_all(encoded.as_bytes()).await?; - write_half.flush().await?; + tokio::time::timeout(LOCAL_IO_TIMEOUT, async { + write_half.write_all(encoded.as_bytes()).await?; + write_half.flush().await + }) + .await + .map_err(|_| anyhow::anyhow!("local IPC response write stalled; closing"))??; } - Ok(()) + } + + /// How long a connection may sit between requests before it is closed. A + /// CLI or MCP call opens, asks, reads, and closes; an idle connection is a + /// client that forgot, and it must not pin a handler slot indefinitely. + const LOCAL_IDLE_TIMEOUT: Duration = Duration::from_secs(300); + /// Bound on writing one response line to a local peer. + const LOCAL_IO_TIMEOUT: Duration = Duration::from_secs(10); + /// Bound on a CP round trip for spawn/check/cancel/list. Mirrors the + /// primary state's own [`PENDING_REQUEST_TTL`] plus slack for the reply to + /// travel the command channel, so the state-side timeout normally fires + /// first and this is the backstop. + const LOCAL_REPLY_TIMEOUT: Duration = + Duration::from_secs(crate::control_plane::primary::PENDING_REQUEST_TTL.as_secs() + 15); + /// Bound on an `Await`: the longest deadline a caller may request plus the + /// grace the primary state gives a late terminal, plus slack. The state + /// normally answers `Timeout` before this; the constant guarantees the + /// handler task is finite regardless. + const LOCAL_AWAIT_TIMEOUT: Duration = Duration::from_secs(MAX_DEADLINE_SECS + 120); + + /// Wait for a parked reply, bounded. `None` means the bound elapsed: the + /// receiver is dropped here, which the primary state observes as + /// "caller gone" on its next prune and removes the parked entry without + /// sending anything. + async fn wait_reply( + rx: tokio::sync::oneshot::Receiver, + bound: Duration, + ) -> Option> { + match tokio::time::timeout(bound, rx).await { + Ok(Ok(v)) => Some(Ok(v)), + Ok(Err(_recv)) => Some(Err(())), + Err(_elapsed) => None, + } + } + + fn reply_timed_out() -> LocalResponse { + LocalResponse::error(CommandError::Timeout.to_string()) + } + + fn reply_dropped() -> LocalResponse { + LocalResponse::error("control-plane client dropped the reply") } /// Minimum accepted `deadline_secs`. Zero (or a negative-by-underflow @@ -597,8 +771,8 @@ mod unix_impl { { return LocalResponse::error(CommandError::NotConnected.to_string()); } - match rx.await { - Ok(Ok(ack)) => { + match wait_reply(rx, LOCAL_REPLY_TIMEOUT).await { + Some(Ok(Ok(ack))) => { // Mint an opaque token for the handle only now that the // spawn succeeded, and report the CP-assigned peer. let token = registry.lock().unwrap().insert(ack.handle); @@ -607,8 +781,9 @@ mod unix_impl { assigned_to: ack.assigned_to, } } - Ok(Err(e)) => LocalResponse::error(e.to_string()), - Err(_) => LocalResponse::error("control-plane client dropped the reply"), + Some(Ok(Err(e))) => LocalResponse::error(e.to_string()), + Some(Err(())) => reply_dropped(), + None => reply_timed_out(), } } LocalRequest::Await { handle } => { @@ -623,14 +798,15 @@ mod unix_impl { { return LocalResponse::error(CommandError::NotConnected.to_string()); } - match rx.await { - Ok(Ok(outcome)) => LocalResponse::Terminal { + match wait_reply(rx, LOCAL_AWAIT_TIMEOUT).await { + Some(Ok(Ok(outcome))) => LocalResponse::Terminal { status: outcome.status.into(), result: outcome.result, error: outcome.error, }, - Ok(Err(e)) => LocalResponse::error(e.to_string()), - Err(_) => LocalResponse::error("control-plane client dropped the reply"), + Some(Ok(Err(e))) => LocalResponse::error(e.to_string()), + Some(Err(())) => reply_dropped(), + None => reply_timed_out(), } } LocalRequest::Check { handle } => { @@ -645,18 +821,21 @@ mod unix_impl { { return LocalResponse::error(CommandError::NotConnected.to_string()); } - match rx.await { - Ok(Ok(DelegationSnapshot::Pending)) => LocalResponse::Pending, - Ok(Ok(DelegationSnapshot::Running { assigned_to })) => { + match wait_reply(rx, LOCAL_REPLY_TIMEOUT).await { + Some(Ok(Ok(DelegationSnapshot::Pending))) => LocalResponse::Pending, + Some(Ok(Ok(DelegationSnapshot::Running { assigned_to }))) => { LocalResponse::Running { assigned_to } } - Ok(Ok(DelegationSnapshot::Terminal(outcome))) => LocalResponse::Terminal { - status: outcome.status.into(), - result: outcome.result, - error: outcome.error, - }, - Ok(Err(e)) => LocalResponse::error(e.to_string()), - Err(_) => LocalResponse::error("control-plane client dropped the reply"), + Some(Ok(Ok(DelegationSnapshot::Terminal(outcome)))) => { + LocalResponse::Terminal { + status: outcome.status.into(), + result: outcome.result, + error: outcome.error, + } + } + Some(Ok(Err(e))) => LocalResponse::error(e.to_string()), + Some(Err(())) => reply_dropped(), + None => reply_timed_out(), } } LocalRequest::Cancel { handle, reason } => { @@ -675,10 +854,11 @@ mod unix_impl { { return LocalResponse::error(CommandError::NotConnected.to_string()); } - match rx.await { - Ok(Ok(())) => LocalResponse::Cancelled, - Ok(Err(e)) => LocalResponse::error(e.to_string()), - Err(_) => LocalResponse::error("control-plane client dropped the reply"), + match wait_reply(rx, LOCAL_REPLY_TIMEOUT).await { + Some(Ok(Ok(()))) => LocalResponse::Cancelled, + Some(Ok(Err(e))) => LocalResponse::error(e.to_string()), + Some(Err(())) => reply_dropped(), + None => reply_timed_out(), } } LocalRequest::ListAgents => { @@ -690,8 +870,8 @@ mod unix_impl { { return LocalResponse::error(CommandError::NotConnected.to_string()); } - match rx.await { - Ok(Ok(agents)) => LocalResponse::Agents { + match wait_reply(rx, LOCAL_REPLY_TIMEOUT).await { + Some(Ok(Ok(agents))) => LocalResponse::Agents { agents: agents .into_iter() .map(|a| LocalAgent { @@ -704,8 +884,9 @@ mod unix_impl { }) .collect(), }, - Ok(Err(e)) => LocalResponse::error(e.to_string()), - Err(_) => LocalResponse::error("control-plane client dropped the reply"), + Some(Ok(Err(e))) => LocalResponse::error(e.to_string()), + Some(Err(())) => reply_dropped(), + None => reply_timed_out(), } } } @@ -1288,6 +1469,140 @@ mod unix_impl { remove_stale_socket(&sock).expect("absent path is a clean no-op"); } + // -- F6: liveness and succession -- + + /// Two runtimes under one uid, same default path: the second must NOT + /// unlink the first's live socket. Ownership is not staleness; a + /// connect probe that succeeds is. + #[tokio::test] + async fn a_live_same_uid_socket_is_refused_not_unlinked() { + let root = tempfile::tempdir().unwrap(); + let sock = root.path().join("agent.sock"); + let _live = UnixListener::bind(&sock).unwrap(); + let err = remove_stale_socket(&sock).expect_err("a live listener must be refused"); + let msg = format!("{err:#}"); + assert!(msg.contains("already serving"), "got {msg}"); + assert!( + msg.contains("OPENAB_AGENT_SOCKET"), + "the error should point at the override: {msg}" + ); + assert!( + std::fs::symlink_metadata(&sock).is_ok(), + "the live socket node must still be there" + ); + } + + /// The full boot path refuses to start over a live sibling. + #[tokio::test] + async fn serve_local_refuses_to_replace_a_live_listener() { + let root = tempfile::tempdir().unwrap(); + let sock = root.path().join("nested").join("agent.sock"); + // The first runtime's boot: creates the 0700 dir and binds. + prepare_socket_dir(&sock).unwrap(); + let _live = UnixListener::bind(&sock).unwrap(); + let (cp, task) = stub_cp(); + let (_tx, rx) = watch::channel(false); + let err = serve_local(sock.clone(), cp, rx) + .await + .expect_err("second runtime must not bind over the first"); + assert!(format!("{err:#}").contains("already serving")); + assert!(std::fs::symlink_metadata(&sock).is_ok()); + task.abort(); + } + + /// An older runtime's shutdown cleanup must not remove a successor's + /// node that took the same pathname. The identity recorded at bind + /// time is what cleanup compares against. + #[tokio::test] + async fn cleanup_leaves_a_successors_node_alone() { + let root = tempfile::tempdir().unwrap(); + let sock = root.path().join("agent.sock"); + + // Predecessor binds and remembers its node. + let predecessor = UnixListener::bind(&sock).unwrap(); + let ours = node_identity(&sock); + assert!(ours.is_some()); + drop(predecessor); + std::fs::remove_file(&sock).unwrap(); + + // Successor takes the pathname (a different inode). + let _successor = UnixListener::bind(&sock).unwrap(); + let theirs = node_identity(&sock); + assert_ne!(ours, theirs, "a fresh bind is a fresh node"); + + // Predecessor's cleanup runs late: must be a no-op. + remove_if_same_node(&sock, ours); + assert!( + std::fs::symlink_metadata(&sock).is_ok(), + "the successor's socket must survive the predecessor's cleanup" + ); + + // The successor's own cleanup does remove it. + remove_if_same_node(&sock, theirs); + assert!(std::fs::symlink_metadata(&sock).is_err()); + } + + // -- F3: bounded handlers -- + + /// Connections past the cap are answered with one error line and + /// closed, rather than each pinning a handler task; the cap is the + /// number of LIVE handlers, so a slot frees when a connection ends. + #[tokio::test] + async fn connections_past_the_cap_are_refused_with_a_diagnostic() { + let root = tempfile::tempdir().unwrap(); + let sock = root.path().join("nested").join("agent.sock"); + let (cp, task) = stub_cp(); + let (tx, rx) = watch::channel(false); + let server_path = sock.clone(); + let server = tokio::spawn(async move { serve_local(server_path, cp, rx).await }); + for _ in 0..100 { + if sock.exists() { + break; + } + tokio::time::sleep(std::time::Duration::from_millis(10)).await; + } + + // Fill every slot with an idle connection (never sends a line). + let mut held = Vec::new(); + for _ in 0..MAX_LOCAL_CONNECTIONS { + held.push(UnixStream::connect(&sock).await.unwrap()); + } + // Let the accept loop register them all. + tokio::time::sleep(std::time::Duration::from_millis(50)).await; + + // The next one is refused with a diagnostic line. + let refused = UnixStream::connect(&sock).await.unwrap(); + let mut reader = BufReader::new(refused); + let mut line = String::new(); + tokio::time::timeout( + std::time::Duration::from_secs(5), + reader.read_line(&mut line), + ) + .await + .expect("a refused connection must be answered promptly") + .unwrap(); + let resp: LocalResponse = serde_json::from_str(line.trim()).unwrap(); + let LocalResponse::Error { message } = resp else { + panic!("expected an error line, got {line}"); + }; + assert!(message.contains("too many concurrent"), "got {message}"); + + // Free one slot; the next connection is served normally. + held.pop(); + tokio::time::sleep(std::time::Duration::from_millis(50)).await; + let client = LocalClient::new(sock.clone()); + let resp = client.request(&LocalRequest::ListAgents).await.unwrap(); + assert!(matches!(resp, LocalResponse::Agents { .. }), "got {resp:?}"); + + let _ = tx.send(true); + let _ = server.await; + task.abort(); + assert!( + std::fs::symlink_metadata(&sock).is_err(), + "our own node is removed on shutdown" + ); + } + /// A line longer than the cap, with no newline, is rejected — and the /// error surfaces rather than the reader growing without bound. #[tokio::test] diff --git a/crates/openab-core/src/control_plane/mod.rs b/crates/openab-core/src/control_plane/mod.rs index ceb877c98..a8f4cabe9 100644 --- a/crates/openab-core/src/control_plane/mod.rs +++ b/crates/openab-core/src/control_plane/mod.rs @@ -22,7 +22,8 @@ pub mod primary; pub use client::{ControlPlaneClient, ControlPlaneHandle}; pub use executor::{ - delegation_session_key, DelegationExecutor, PromptOutcome, PromptRunner, RouterPromptRunner, + delegation_session_key, Admitted, DelegationExecutor, PromptOutcome, PromptRunner, + RouterPromptRunner, }; pub use local::{ default_socket_path, serve_local, LocalAgent, LocalClient, LocalRequest, LocalResponse, diff --git a/crates/openab-core/src/control_plane/primary.rs b/crates/openab-core/src/control_plane/primary.rs index 7cfcf3b23..c440931fb 100644 --- a/crates/openab-core/src/control_plane/primary.rs +++ b/crates/openab-core/src/control_plane/primary.rs @@ -242,6 +242,14 @@ pub enum CommandError { /// frame exceeds the transport limit). Distinct from `Internal`: it names /// a caller-side fault the caller could fix (a smaller prompt), not a bug. InvalidRequest(String), + /// The CP did not answer this request within [`PENDING_REQUEST_TTL`], or a + /// delegation reached its deadline plus [`LATE_TERMINAL_GRACE`] without a + /// terminal frame. Local state for it has been dropped; a terminal that + /// arrives later is acked and ignored. + Timeout, + /// Too many requests are already outstanding to the CP + /// ([`MAX_PENDING_REQUESTS`]). Nothing was sent; the caller may retry. + Overloaded, /// The command channel or reply path was torn down (shutdown). Internal(String), } @@ -268,6 +276,15 @@ impl std::fmt::Display for CommandError { ) } CommandError::InvalidRequest(m) => write!(f, "invalid request: {m}"), + CommandError::Timeout => write!( + f, + "the control plane did not answer in time; local tracking for this request \ + was dropped" + ), + CommandError::Overloaded => write!( + f, + "too many control-plane requests are outstanding from this runtime; retry later" + ), CommandError::Internal(m) => write!(f, "internal error: {m}"), } } @@ -293,6 +310,9 @@ impl CommandError { enum PendingRequest { Spawn { delegation_id: String, + /// Carried through to the tracked entry so a delegation the CP never + /// terminates can be aged out locally (see [`PrimaryState::prune`]). + deadline: chrono::DateTime, reply: oneshot::Sender>, }, Cancel { @@ -303,9 +323,63 @@ enum PendingRequest { }, } +impl PendingRequest { + /// Whether the caller that parked this request has gone away (its receiver + /// dropped — a timed-out or disconnected local handler). Such an entry + /// has nobody to answer and is pruned without a reply. + fn caller_gone(&self) -> bool { + match self { + PendingRequest::Spawn { reply, .. } => reply.is_closed(), + PendingRequest::Cancel { reply } => reply.is_closed(), + PendingRequest::ListAgents { reply } => reply.is_closed(), + } + } + + fn fail(self, error: CommandError) { + match self { + PendingRequest::Spawn { reply, .. } => { + let _ = reply.send(Err(error)); + } + PendingRequest::Cancel { reply } => { + let _ = reply.send(Err(error)); + } + PendingRequest::ListAgents { reply } => { + let _ = reply.send(Err(error)); + } + } + } +} + +/// A parked request plus when it was parked, so [`PrimaryState::prune`] can +/// age it out if the CP keeps accepting writes but never answers (F3). +struct Parked { + since: std::time::Instant, + request: PendingRequest, +} + +/// Ceiling on requests outstanding to the CP at once. The command ingress +/// queue is bounded separately; this bounds what has already been consumed +/// from it and is waiting on a reply, so a withholding peer cannot grow +/// `pending` for the process lifetime. +pub const MAX_PENDING_REQUESTS: usize = 256; + +/// How long a parked request may wait for its JSON-RPC reply. The CP answers +/// spawn/cancel/list synchronously from its own state, so a minute is +/// generous; a longer silence means the reply is not coming. +pub const PENDING_REQUEST_TTL: std::time::Duration = std::time::Duration::from_secs(60); + +/// Grace past a delegation's own deadline before local tracking of a still +/// non-terminal entry is dropped and its awaiters answered `Timeout`. The CP +/// sweeps deadlines itself and sends `timed_out`; this is the fallback for a +/// CP that does not. +pub const LATE_TERMINAL_GRACE: chrono::Duration = chrono::Duration::seconds(60); + /// Lifecycle of one initiated delegation, from ack to terminal. struct TrackedDelegation { state: DelegationLifecycle, + /// The deadline the initiator asked for; live entries older than this + /// plus [`LATE_TERMINAL_GRACE`] are aged out by [`PrimaryState::prune`]. + deadline: chrono::DateTime, /// The peer the CP assigned this delegation to, learned at ack time and /// retained so it can be accounted into the terminal-history byte bound /// (the terminal `DelegationOutcome` does not itself carry it). @@ -350,7 +424,7 @@ const MAX_TRACKED_BYTES: usize = 16 * 1024 * 1024; #[derive(Default)] pub struct PrimaryState { /// JSON-RPC id → the request the loop originated and awaits a reply for. - pending: HashMap, + pending: HashMap, /// (delegation_id, admission) → its lifecycle. tracked: HashMap, /// Insertion order for deterministic oldest-terminal history eviction. @@ -385,12 +459,20 @@ impl PrimaryState { /// params. The caller is answered later, when the ack lands. `rpc_id` is /// allocated by the serve loop so every outbound frame — heartbeat, /// delegate_result ack, and these primary requests — shares one id space. + /// + /// Returns `None` — and answers the caller [`CommandError::Overloaded`] — + /// when [`MAX_PENDING_REQUESTS`] are already outstanding, so nothing is + /// sent and `pending` cannot grow past its cap. pub(crate) fn begin_spawn( &mut self, rpc_id: u64, request: SpawnRequest, reply: oneshot::Sender>, - ) -> SpawnEmission { + ) -> Option { + if self.pending_full() { + let _ = reply.send(Err(CommandError::Overloaded)); + return None; + } let (parent_delegation_id, parent_admission) = match request.parent { Some(h) => (Some(h.delegation_id), Some(h.admission)), None => (None, None), @@ -403,14 +485,106 @@ impl PrimaryState { parent_delegation_id, parent_admission, }; - self.pending.insert( + self.park( rpc_id, PendingRequest::Spawn { delegation_id: request.delegation_id, + deadline: request.deadline, reply, }, ); - SpawnEmission { rpc_id, params } + Some(SpawnEmission { rpc_id, params }) + } + + fn pending_full(&self) -> bool { + self.pending.len() >= MAX_PENDING_REQUESTS + } + + fn park(&mut self, rpc_id: u64, request: PendingRequest) { + self.pending.insert( + rpc_id, + Parked { + since: std::time::Instant::now(), + request, + }, + ); + } + + /// Number of requests parked awaiting a CP reply. + pub fn pending_len(&self) -> usize { + self.pending.len() + } + + /// Number of parked awaiters (callers blocked on a terminal). + pub fn awaiters_len(&self) -> usize { + self.awaiters.values().map(Vec::len).sum() + } + + /// Age out state a connected-but-withholding CP would otherwise let grow + /// for the process lifetime (F3). Called by the serve loop on every + /// heartbeat tick, so the bound holds without any local caller doing + /// anything. + /// + /// - A parked request whose caller has gone (receiver dropped) is removed + /// silently: there is nobody to answer. + /// - A parked request older than [`PENDING_REQUEST_TTL`] is failed + /// [`CommandError::Timeout`] and removed. + /// - An awaiter whose caller has gone is removed. + /// - A non-terminal delegation past `deadline + LATE_TERMINAL_GRACE` is + /// dropped from tracking and its awaiters answered `Timeout`. A terminal + /// frame that arrives for it afterwards is acked and ignored, exactly + /// like any other untracked terminal. + /// + /// Terminal history is untouched: it has its own count/byte bound. + pub fn prune(&mut self, now: std::time::Instant, wall: chrono::DateTime) { + let expired: Vec = self + .pending + .iter() + .filter(|(_, p)| { + p.request.caller_gone() || now.duration_since(p.since) > PENDING_REQUEST_TTL + }) + .map(|(id, _)| *id) + .collect(); + for id in expired { + if let Some(parked) = self.pending.remove(&id) { + if !parked.request.caller_gone() { + debug!( + rpc_id = id, + "pending control-plane request timed out locally" + ); + parked.request.fail(CommandError::Timeout); + } + } + } + + for waiters in self.awaiters.values_mut() { + waiters.retain(|w| !w.is_closed()); + } + self.awaiters.retain(|_, w| !w.is_empty()); + + let overdue: Vec = self + .tracked + .iter() + .filter(|(_, t)| { + !matches!(t.state, DelegationLifecycle::Terminal(_)) + && wall > t.deadline + LATE_TERMINAL_GRACE + }) + .map(|(k, _)| k.clone()) + .collect(); + for key in overdue { + debug!( + delegation_id = %key.0, + admission = key.1, + "delegation passed its deadline without a terminal frame; dropping local tracking" + ); + if let Some(waiters) = self.awaiters.remove(&key) { + for w in waiters { + let _ = w.send(Err(CommandError::Timeout)); + } + } + self.tracked.remove(&key); + self.tracked_order.retain(|k| k != &key); + } } /// Remove a parked pending request (a spawn whose frame the serve loop @@ -422,20 +596,10 @@ impl PrimaryState { /// connection. Returns true if a pending request was parked under /// `rpc_id` (so this was ours to fail); false otherwise. Tracks nothing. pub(crate) fn fail_pending_invalid_request(&mut self, rpc_id: u64, message: String) -> bool { - let Some(pending) = self.pending.remove(&rpc_id) else { + let Some(parked) = self.pending.remove(&rpc_id) else { return false; }; - match pending { - PendingRequest::Spawn { reply, .. } => { - let _ = reply.send(Err(CommandError::InvalidRequest(message))); - } - PendingRequest::Cancel { reply } => { - let _ = reply.send(Err(CommandError::InvalidRequest(message))); - } - PendingRequest::ListAgents { reply } => { - let _ = reply.send(Err(CommandError::InvalidRequest(message))); - } - } + parked.request.fail(CommandError::InvalidRequest(message)); true } @@ -454,24 +618,33 @@ impl PrimaryState { let _ = reply.send(Err(CommandError::UnknownDelegation)); return None; } + if self.pending_full() { + let _ = reply.send(Err(CommandError::Overloaded)); + return None; + } let params = openab_cp::proto::CancelParams { delegation_id: handle.delegation_id.clone(), admission: handle.admission, reason, }; - self.pending - .insert(rpc_id, PendingRequest::Cancel { reply }); + self.park(rpc_id, PendingRequest::Cancel { reply }); Some(params) } - /// Register a list_agents: park its reply under `rpc_id`. + /// Register a list_agents: park its reply under `rpc_id`. Returns `false` + /// (caller answered `Overloaded`, no frame to send) when the pending cap + /// is reached. pub(crate) fn begin_list_agents( &mut self, rpc_id: u64, reply: oneshot::Sender, CommandError>>, - ) { - self.pending - .insert(rpc_id, PendingRequest::ListAgents { reply }); + ) -> bool { + if self.pending_full() { + let _ = reply.send(Err(CommandError::Overloaded)); + return false; + } + self.park(rpc_id, PendingRequest::ListAgents { reply }); + true } /// Await a delegation's terminal outcome. If it already ended, the caller @@ -543,12 +716,13 @@ impl PrimaryState { result: Option, error: Option, ) -> bool { - let Some(pending) = self.pending.remove(&id) else { + let Some(parked) = self.pending.remove(&id) else { return false; }; - match pending { + match parked.request { PendingRequest::Spawn { delegation_id, + deadline, reply, } => { if let Some(err) = error { @@ -577,15 +751,32 @@ impl PrimaryState { } let handle = DelegationHandle::new(ack.delegation_id.clone(), ack.admission); + let key = handle.key(); + // A conforming CP mints one admission per ack; a repeat + // of an already-tracked (id, admission) would overwrite + // an entry (orphaning its accounted bytes) and push a + // permanent duplicate into the order queue. Refuse it. + if self.tracked.contains_key(&key) { + warn!( + delegation_id = %delegation_id, + admission = ack.admission, + "cp/delegate ack repeats an admission already tracked — rejecting" + ); + let _ = reply.send(Err(CommandError::Internal(format!( + "cp/delegate ack repeated admission {} of {:?}", + ack.admission, delegation_id + )))); + return true; + } // Track the admission so the later result frame routes. self.make_history_room(); - let key = handle.key(); self.tracked.insert( key.clone(), TrackedDelegation { state: DelegationLifecycle::Running { assigned_to: ack.assigned_to.clone(), }, + deadline, assigned_to: ack.assigned_to.clone(), accounted_bytes: 0, }, @@ -753,18 +944,8 @@ impl PrimaryState { /// worker side — no local state may claim a delegation is live once the /// only connection that could complete it is gone. pub fn fail_all_live(&mut self) { - for (_id, pending) in self.pending.drain() { - match pending { - PendingRequest::Spawn { reply, .. } => { - let _ = reply.send(Err(CommandError::Disconnected)); - } - PendingRequest::Cancel { reply } => { - let _ = reply.send(Err(CommandError::Disconnected)); - } - PendingRequest::ListAgents { reply } => { - let _ = reply.send(Err(CommandError::Disconnected)); - } - } + for (_id, parked) in self.pending.drain() { + parked.request.fail(CommandError::Disconnected); } // Mark non-terminal delegations failed and answer their awaiters. let live_keys: Vec = self @@ -848,7 +1029,7 @@ mod tests { fn begin_spawn_parks_the_reply_and_produces_wire_params() { let mut st = PrimaryState::new(); let (tx, _rx) = oneshot::channel(); - let emission = st.begin_spawn(1, spawn_request("d-1"), tx); + let emission = st.begin_spawn(1, spawn_request("d-1"), tx).unwrap(); assert_eq!(emission.rpc_id, 1); assert_eq!(emission.params.delegation_id, "d-1"); assert_eq!(emission.params.target.name.as_deref(), Some("worker-1")); @@ -863,7 +1044,7 @@ mod tests { let (tx, _rx) = oneshot::channel(); let mut req = spawn_request("child"); req.parent = Some(DelegationHandle::new("parent", 42)); - let emission = st.begin_spawn(7, req, tx); + let emission = st.begin_spawn(7, req, tx).unwrap(); assert_eq!( emission.params.parent_delegation_id.as_deref(), Some("parent") @@ -875,7 +1056,7 @@ mod tests { async fn a_delegate_ack_resolves_the_spawn_and_starts_tracking() { let mut st = PrimaryState::new(); let (tx, rx) = oneshot::channel(); - let emission = st.begin_spawn(1, spawn_request("d-1"), tx); + let emission = st.begin_spawn(1, spawn_request("d-1"), tx).unwrap(); assert!(st.on_reply( emission.rpc_id, Some(ack_value("d-1", 5, "prod/worker-1")), @@ -898,7 +1079,7 @@ mod tests { async fn a_delegate_error_reply_fails_the_spawn_and_tracks_nothing() { let mut st = PrimaryState::new(); let (tx, rx) = oneshot::channel(); - let emission = st.begin_spawn(1, spawn_request("d-1"), tx); + let emission = st.begin_spawn(1, spawn_request("d-1"), tx).unwrap(); let err = ErrorObject::new(openab_cp::proto::codes::NO_TARGET, "no such worker"); assert!(st.on_reply(emission.rpc_id, None, Some(err))); let outcome = rx.await.unwrap(); @@ -917,7 +1098,7 @@ mod tests { let mut st = PrimaryState::new(); // Spawn + ack. let (tx, rx) = oneshot::channel(); - let emission = st.begin_spawn(1, spawn_request("d-1"), tx); + let emission = st.begin_spawn(1, spawn_request("d-1"), tx).unwrap(); st.on_reply(emission.rpc_id, Some(ack_value("d-1", 3, "prod/w1")), None); let handle = rx.await.unwrap().unwrap().handle; @@ -942,7 +1123,7 @@ mod tests { async fn awaiting_an_already_terminal_delegation_answers_immediately() { let mut st = PrimaryState::new(); let (tx, rx) = oneshot::channel(); - let emission = st.begin_spawn(1, spawn_request("d-1"), tx); + let emission = st.begin_spawn(1, spawn_request("d-1"), tx).unwrap(); st.on_reply(emission.rpc_id, Some(ack_value("d-1", 3, "prod/w1")), None); let handle = rx.await.unwrap().unwrap().handle; let result = DelegateResultParams { @@ -965,7 +1146,7 @@ mod tests { async fn a_stale_admission_result_is_dropped() { let mut st = PrimaryState::new(); let (tx, rx) = oneshot::channel(); - let emission = st.begin_spawn(1, spawn_request("d-1"), tx); + let emission = st.begin_spawn(1, spawn_request("d-1"), tx).unwrap(); st.on_reply(emission.rpc_id, Some(ack_value("d-1", 10, "prod/w1")), None); let _handle = rx.await.unwrap().unwrap().handle; // A result for the same id but a different admission is not ours. @@ -998,7 +1179,7 @@ mod tests { async fn cancel_of_a_tracked_handle_emits_and_resolves_on_ok() { let mut st = PrimaryState::new(); let (tx, rx) = oneshot::channel(); - let emission = st.begin_spawn(1, spawn_request("d-1"), tx); + let emission = st.begin_spawn(1, spawn_request("d-1"), tx).unwrap(); st.on_reply(emission.rpc_id, Some(ack_value("d-1", 3, "prod/w1")), None); let handle = rx.await.unwrap().unwrap().handle; @@ -1017,7 +1198,7 @@ mod tests { async fn list_agents_resolves_with_the_roster() { let mut st = PrimaryState::new(); let (tx, rx) = oneshot::channel(); - st.begin_list_agents(1, tx); + assert!(st.begin_list_agents(1, tx)); let result = openab_cp::proto::ListAgentsResult { namespace: "prod".into(), agents: vec![AgentSummary { @@ -1038,7 +1219,7 @@ mod tests { async fn check_reports_running_then_terminal_without_blocking() { let mut st = PrimaryState::new(); let (tx, rx) = oneshot::channel(); - let emission = st.begin_spawn(1, spawn_request("d-1"), tx); + let emission = st.begin_spawn(1, spawn_request("d-1"), tx).unwrap(); st.on_reply(emission.rpc_id, Some(ack_value("d-1", 4, "prod/w1")), None); let handle = rx.await.unwrap().unwrap().handle; @@ -1084,6 +1265,7 @@ mod tests { handle.key(), TrackedDelegation { state: DelegationLifecycle::Pending, + deadline: chrono::Utc::now() + chrono::Duration::seconds(60), assigned_to: String::new(), accounted_bytes: 0, }, @@ -1125,6 +1307,7 @@ mod tests { key.clone(), TrackedDelegation { state: DelegationLifecycle::Terminal(outcome), + deadline: chrono::Utc::now() + chrono::Duration::seconds(60), assigned_to: "prod/w".into(), accounted_bytes: bytes, }, @@ -1140,6 +1323,7 @@ mod tests { state: DelegationLifecycle::Running { assigned_to: "prod/w".into(), }, + deadline: chrono::Utc::now() + chrono::Duration::seconds(60), assigned_to: "prod/w".into(), accounted_bytes: 0, }, @@ -1155,7 +1339,7 @@ mod tests { // A running delegation with a parked awaiter. let (tx, rx) = oneshot::channel(); - let emission = st.begin_spawn(1, spawn_request("d-run"), tx); + let emission = st.begin_spawn(1, spawn_request("d-run"), tx).unwrap(); st.on_reply( emission.rpc_id, Some(ack_value("d-run", 1, "prod/w1")), @@ -1167,7 +1351,7 @@ mod tests { // A terminal delegation: its awaiter was already answered. let (tx2, rx2) = oneshot::channel(); - let emission2 = st.begin_spawn(2, spawn_request("d-done"), tx2); + let emission2 = st.begin_spawn(2, spawn_request("d-done"), tx2).unwrap(); st.on_reply( emission2.rpc_id, Some(ack_value("d-done", 2, "prod/w1")), @@ -1223,7 +1407,7 @@ mod tests { // Complete a delegation while "connected". let (tx, rx) = oneshot::channel(); - let emission = st.begin_spawn(1, spawn_request("d-done"), tx); + let emission = st.begin_spawn(1, spawn_request("d-done"), tx).unwrap(); st.on_reply( emission.rpc_id, Some(ack_value("d-done", 1, "prod/w1")), @@ -1280,7 +1464,7 @@ mod tests { let id = format!("d-cycle-{cycle}"); let admission = 200 + cycle; let (tx, rx) = oneshot::channel(); - let e = st.begin_spawn(300 + cycle, spawn_request(&id), tx); + let e = st.begin_spawn(300 + cycle, spawn_request(&id), tx).unwrap(); st.on_reply(e.rpc_id, Some(ack_value(&id, admission, "prod/w1")), None); let h = rx.await.unwrap().unwrap().handle; st.on_delegate_result(&DelegateResultParams { @@ -1307,7 +1491,7 @@ mod tests { // so a foreign id can never seed correlation under our rpc id. let mut st = PrimaryState::new(); let (tx, rx) = oneshot::channel(); - let emission = st.begin_spawn(1, spawn_request("d-requested"), tx); + let emission = st.begin_spawn(1, spawn_request("d-requested"), tx).unwrap(); assert!(st.on_reply( emission.rpc_id, Some(ack_value("d-foreign", 5, "prod/w1")), @@ -1324,7 +1508,7 @@ mod tests { // The complement: an exact match is admitted normally. let mut st = PrimaryState::new(); let (tx, rx) = oneshot::channel(); - let emission = st.begin_spawn(1, spawn_request("d-1"), tx); + let emission = st.begin_spawn(1, spawn_request("d-1"), tx).unwrap(); st.on_reply(emission.rpc_id, Some(ack_value("d-1", 5, "prod/w1")), None); let handle = rx.await.unwrap().unwrap().handle; assert_eq!( @@ -1343,7 +1527,7 @@ mod tests { // nothing — the connection is untouched. let mut st = PrimaryState::new(); let (tx, rx) = oneshot::channel(); - let emission = st.begin_spawn(7, spawn_request("d-huge"), tx); + let emission = st.begin_spawn(7, spawn_request("d-huge"), tx).unwrap(); assert!(st.fail_pending_invalid_request(emission.rpc_id, "frame too large".into())); let outcome = rx.await.unwrap(); assert!(matches!(outcome, Err(CommandError::InvalidRequest(_)))); @@ -1379,6 +1563,7 @@ mod tests { key.clone(), TrackedDelegation { state: DelegationLifecycle::Terminal(outcome), + deadline: chrono::Utc::now() + chrono::Duration::seconds(60), assigned_to: "prod/w".into(), accounted_bytes: bytes, }, @@ -1412,6 +1597,7 @@ mod tests { state: DelegationLifecycle::Running { assigned_to: "prod/w".into(), }, + deadline: chrono::Utc::now() + chrono::Duration::seconds(60), assigned_to: "prod/w".into(), accounted_bytes: 0, }, @@ -1452,6 +1638,7 @@ mod tests { state: DelegationLifecycle::Running { assigned_to: "prod/w".into(), }, + deadline: chrono::Utc::now() + chrono::Duration::seconds(60), assigned_to: "prod/w".into(), accounted_bytes: 0, }, @@ -1471,6 +1658,7 @@ mod tests { key.clone(), TrackedDelegation { state: DelegationLifecycle::Terminal(outcome), + deadline: chrono::Utc::now() + chrono::Duration::seconds(60), assigned_to: "prod/w".into(), accounted_bytes: bytes, }, @@ -1485,4 +1673,151 @@ mod tests { "the live delegation must survive byte-bound eviction" ); } + + // ---- F3: bounded pending state ---------------------------------------- + + #[tokio::test] + async fn pending_requests_are_capped_and_the_overflow_is_answered_overloaded() { + let mut st = PrimaryState::new(); + let mut parked = Vec::new(); + for i in 0..MAX_PENDING_REQUESTS as u64 { + let (tx, rx) = oneshot::channel(); + assert!(st.begin_list_agents(i, tx), "under the cap"); + parked.push(rx); + } + assert_eq!(st.pending_len(), MAX_PENDING_REQUESTS); + + // One more of each kind: refused locally, nothing parked. + let (tx, rx) = oneshot::channel(); + assert!(!st.begin_list_agents(9_000, tx)); + assert!(matches!(rx.await, Ok(Err(CommandError::Overloaded)))); + + let (tx, rx) = oneshot::channel(); + assert!(st.begin_spawn(9_001, spawn_request("d-over"), tx).is_none()); + assert!(matches!(rx.await, Ok(Err(CommandError::Overloaded)))); + assert_eq!(st.pending_len(), MAX_PENDING_REQUESTS, "the cap holds"); + + // Replies drain the map and make room again. + st.on_reply(0, Some(serde_json::json!({"agents": []})), None); + assert_eq!(st.pending_len(), MAX_PENDING_REQUESTS - 1); + let (tx, _rx) = oneshot::channel(); + assert!(st.begin_list_agents(9_002, tx)); + } + + #[tokio::test] + async fn prune_times_out_stale_pending_and_forgets_callers_that_left() { + let mut st = PrimaryState::new(); + let t0 = std::time::Instant::now(); + let wall = chrono::Utc::now(); + + // A spawn the CP never answers. + let (tx, rx) = oneshot::channel(); + st.begin_spawn(1, spawn_request("d-silent"), tx).unwrap(); + // A list whose local caller already gave up (receiver dropped). + let (tx2, rx2) = oneshot::channel(); + assert!(st.begin_list_agents(2, tx2)); + drop(rx2); + assert_eq!(st.pending_len(), 2); + + // Young: the silent one stays; the orphan goes. + st.prune(t0 + std::time::Duration::from_secs(1), wall); + assert_eq!(st.pending_len(), 1); + + // Past the TTL: the silent one is failed Timeout and removed. + st.prune( + t0 + PENDING_REQUEST_TTL + std::time::Duration::from_secs(1), + wall, + ); + assert_eq!(st.pending_len(), 0); + assert!(matches!(rx.await, Ok(Err(CommandError::Timeout)))); + } + + #[tokio::test] + async fn prune_drops_orphaned_awaiters_and_ages_out_overdue_live_delegations() { + let mut st = PrimaryState::new(); + let now = std::time::Instant::now(); + + // Admit a delegation with a 60s deadline. + let (tx, rx) = oneshot::channel(); + let e = st.begin_spawn(1, spawn_request("d-late"), tx).unwrap(); + st.on_reply(e.rpc_id, Some(ack_value("d-late", 11, "prod/w1")), None); + let handle = rx.await.unwrap().unwrap().handle; + + // Two awaiters: one whose caller left, one still waiting. + let (gone_tx, gone_rx) = oneshot::channel(); + st.begin_await(&handle, gone_tx); + drop(gone_rx); + let (live_tx, live_rx) = oneshot::channel(); + st.begin_await(&handle, live_tx); + assert_eq!(st.awaiters_len(), 2); + + // Before deadline + grace: the orphan is dropped, the live one kept, + // and the delegation is still tracked. + st.prune(now, chrono::Utc::now()); + assert_eq!(st.awaiters_len(), 1); + assert_eq!(st.tracked_len(), 1); + + // Past deadline + grace with no terminal frame: tracking is dropped and + // the remaining awaiter is answered Timeout. + let overdue = chrono::Utc::now() + + chrono::Duration::seconds(60) + + LATE_TERMINAL_GRACE + + chrono::Duration::seconds(1); + st.prune(now, overdue); + assert_eq!(st.tracked_len(), 0); + assert_eq!(st.awaiters_len(), 0); + assert!(matches!(live_rx.await, Ok(Err(CommandError::Timeout)))); + + // A terminal that arrives afterwards names nothing we track: the + // serve loop acks it regardless (its contract for untracked results), + // and it must not resurrect tracking here. + let ours = st.on_delegate_result(&DelegateResultParams { + delegation_id: "d-late".into(), + admission: 11, + status: DelegationStatus::Completed, + result: Some("too late".into()), + error: None, + }); + assert!(!ours, "no longer ours"); + assert_eq!(st.tracked_len(), 0, "and it does not resurrect tracking"); + } + + #[tokio::test] + async fn prune_never_touches_terminal_history() { + let mut st = PrimaryState::new(); + let (tx, rx) = oneshot::channel(); + let e = st.begin_spawn(1, spawn_request("d-done"), tx).unwrap(); + st.on_reply(e.rpc_id, Some(ack_value("d-done", 1, "prod/w1")), None); + let _ = rx.await.unwrap().unwrap(); + st.on_delegate_result(&DelegateResultParams { + delegation_id: "d-done".into(), + admission: 1, + status: DelegationStatus::Completed, + result: Some("ok".into()), + error: None, + }); + // Far in the future: terminal history has its own bound and survives. + st.prune( + std::time::Instant::now(), + chrono::Utc::now() + chrono::Duration::days(365), + ); + assert_eq!(st.tracked_len(), 1); + } + + #[tokio::test] + async fn a_repeated_admission_in_a_delegate_ack_is_rejected_not_overwritten() { + let mut st = PrimaryState::new(); + let (tx, rx) = oneshot::channel(); + let e = st.begin_spawn(1, spawn_request("d-dup"), tx).unwrap(); + st.on_reply(e.rpc_id, Some(ack_value("d-dup", 3, "prod/w1")), None); + let _ = rx.await.unwrap().unwrap(); + + // A second spawn of the same id gets an ack that repeats admission 3. + let (tx2, rx2) = oneshot::channel(); + let e2 = st.begin_spawn(2, spawn_request("d-dup"), tx2).unwrap(); + st.on_reply(e2.rpc_id, Some(ack_value("d-dup", 3, "prod/w1")), None); + assert!(matches!(rx2.await, Ok(Err(CommandError::Internal(_))))); + assert_eq!(st.tracked_len(), 1, "the original entry stands"); + assert_eq!(st.tracked_order.len(), 1, "no duplicate in the order queue"); + } } From 52228926792e7aeb35ae5fcafc31fff2c7447101 Mon Sep 17 00:00:00 2001 From: chaodu-agent <274062505+chaodu-agent@users.noreply.github.com> Date: Wed, 9 Sep 2026 23:46:50 -0400 Subject: [PATCH 05/15] fix(cp): ext4-safe socket shutdown cleanup Take 91c0a89f7 (local.rs only) from the closed reference branch. Shutdown now drops our listener first, then unlinks the socket only if its (dev, ino) still matches the bound node AND a connect probe is refused. A successor that answers the probe is left alone, which covers ext4 reusing the freed inode number for the next create in the same directory. Reference CI (run 34434639420, head 91c0a89f7) reported the cap test `connections_past_the_cap_are_refused_with_a_diagnostic` failing on Linux. Root cause: NOT reproduced. The cap test passes on p1 (ext4) at this commit, so no change beyond the reference hunk was made. The CI cause remains UNVERIFIED. Commands and results: - macmini: cargo check -p openab-core -p openab: pass - macmini: cargo test -p openab-core --lib control_plane: 120 passed, 0 failed - p1 (ext4, Linux): cargo test -p openab-core --lib control_plane: 120 passed, 0 failed - p1: cargo test -p openab-core --lib control_plane::local, 25 repeated runs: 25/25 pass - p1: cap test alone under 8-way parallel load, 40 runs (--exact, full name): 40/40 "1 passed" - cp_client tests: not run, tests/ directory is added in a later commit Cherry-picked from 91c0a89f7 (see trailer). --- crates/openab-core/src/control_plane/local.rs | 100 +++++++++++++----- 1 file changed, 76 insertions(+), 24 deletions(-) diff --git a/crates/openab-core/src/control_plane/local.rs b/crates/openab-core/src/control_plane/local.rs index a75eb2c4b..e380db1e8 100644 --- a/crates/openab-core/src/control_plane/local.rs +++ b/crates/openab-core/src/control_plane/local.rs @@ -320,8 +320,10 @@ mod unix_impl { handlers.abort_all(); while handlers.join_next().await.is_some() {} // Best-effort cleanup so the next boot does not trip over the node — - // but only if the path still names OUR node. - remove_if_same_node(&path, ours); + // but never a successor's. Stop listening first so a connect probe + // cannot be answered by us, then unlink only if nobody else answers. + drop(listener); + remove_if_ours_and_dead(&path, ours); Ok(()) } @@ -528,22 +530,47 @@ mod unix_impl { }) } - /// Unlink `path` only if it is still the node THIS server created. An older - /// runtime shutting down after a successor bound the same pathname must - /// not remove the successor's socket (F6); with the liveness probe above - /// that succession is refused, but the check keeps cleanup correct even if - /// the path was replaced by other means. - fn remove_if_same_node(path: &Path, ours: Option) { + /// Shutdown cleanup: unlink `path` only if it still looks like the node + /// THIS server created AND nobody is listening on it. An older runtime + /// shutting down after a successor bound the same pathname must not + /// remove the successor's socket (F6); the liveness probe at startup + /// refuses that succession, but this keeps cleanup correct even if the + /// path was replaced by other means (an operator `rm`, a crash). + /// + /// Two checks, because neither is sufficient alone. `(dev, ino)` catches + /// a replaced node on filesystems that hand out fresh inode numbers, but + /// ext4 reuses a freed inode number for the very next create in the + /// directory, so a successor can carry OUR identity. The connect probe + /// catches that case: the caller has already closed its own listener, so + /// anyone who answers is a successor. `ECONNREFUSED` is the only state + /// that is unlinked. + fn remove_if_ours_and_dead(path: &Path, ours: Option) { let Some(ours) = ours else { return; }; - if node_identity(path) == Some(ours) { - let _ = std::fs::remove_file(path); - } else { + if node_identity(path) != Some(ours) { debug!( socket = %path.display(), "socket path no longer names the node this server bound; leaving it in place" ); + return; + } + match std::os::unix::net::UnixStream::connect(path) { + Err(e) if e.kind() == std::io::ErrorKind::ConnectionRefused => { + let _ = std::fs::remove_file(path); + } + Ok(_live) => { + debug!( + socket = %path.display(), + "a successor is listening on the socket path; leaving it in place" + ); + } + Err(e) => { + debug!( + socket = %path.display(), error = %e, + "cannot tell whether the socket path is live; leaving it in place" + ); + } } } @@ -1511,34 +1538,59 @@ mod unix_impl { } /// An older runtime's shutdown cleanup must not remove a successor's - /// node that took the same pathname. The identity recorded at bind - /// time is what cleanup compares against. + /// node that took the same pathname. On ext4 the successor is even + /// handed the predecessor's inode number back, so identity alone is + /// not enough: the connect probe is what protects the successor. #[tokio::test] - async fn cleanup_leaves_a_successors_node_alone() { + async fn cleanup_leaves_a_live_successors_node_alone() { let root = tempfile::tempdir().unwrap(); let sock = root.path().join("agent.sock"); - // Predecessor binds and remembers its node. + // Predecessor binds, remembers its node, and later closes its + // listener (the first step of its shutdown). let predecessor = UnixListener::bind(&sock).unwrap(); let ours = node_identity(&sock); assert!(ours.is_some()); drop(predecessor); - std::fs::remove_file(&sock).unwrap(); - // Successor takes the pathname (a different inode). - let _successor = UnixListener::bind(&sock).unwrap(); + // The pathname is replaced by other means and a successor binds it. + // (Deliberately no assertion that the identity differs: on ext4 it + // does not.) + std::fs::remove_file(&sock).unwrap(); + let successor = UnixListener::bind(&sock).unwrap(); let theirs = node_identity(&sock); - assert_ne!(ours, theirs, "a fresh bind is a fresh node"); - // Predecessor's cleanup runs late: must be a no-op. - remove_if_same_node(&sock, ours); + // Predecessor's cleanup runs late: the successor answers the probe, + // so it must be a no-op. + remove_if_ours_and_dead(&sock, ours); assert!( std::fs::symlink_metadata(&sock).is_ok(), - "the successor's socket must survive the predecessor's cleanup" + "the successor's live socket must survive the predecessor's cleanup" ); - // The successor's own cleanup does remove it. - remove_if_same_node(&sock, theirs); + // The successor's own shutdown: close, then clean up — removed. + drop(successor); + remove_if_ours_and_dead(&sock, theirs); + assert!(std::fs::symlink_metadata(&sock).is_err()); + } + + /// A node that no longer matches our identity is left alone even when + /// nothing is listening on it (a stranger's stale node is theirs to + /// clean up). + #[tokio::test] + async fn cleanup_leaves_a_foreign_dead_node_alone() { + let root = tempfile::tempdir().unwrap(); + let sock = root.path().join("agent.sock"); + let dead = UnixListener::bind(&sock).unwrap(); + drop(dead); + let not_ours = Some(NodeIdentity { + dev: u64::MAX, + ino: u64::MAX, + }); + remove_if_ours_and_dead(&sock, not_ours); + assert!(std::fs::symlink_metadata(&sock).is_ok()); + // And with the right identity, a dead node is removed. + remove_if_ours_and_dead(&sock, node_identity(&sock)); assert!(std::fs::symlink_metadata(&sock).is_err()); } From f2378adf8c0493b4a801a00047b53b05f35c456e Mon Sep 17 00:00:00 2001 From: Tautui Date: Sat, 31 Oct 2026 10:41:27 -0400 Subject: [PATCH 06/15] test(cp): client integration tests Adds crates/openab-core/tests/cp_client.rs, taken from the reference branch feat/cp-runtime-client-main (9455b504 and 3742e03c, the only commits that touch the file). The test boots openab-cp in-process on an ephemeral loopback port and drives the real ControlPlaneClient against it. Dev-dependencies added to crates/openab-core/Cargo.toml: openab-cp with the server feature, and axum 0.8 for axum::serve. Cargo.lock gains only the openab-core -> axum edge. Checks run on macmini (rustc 1.99.1): - cargo check -p openab-core -p openab: pass - cargo test -p openab-core --test cp_client: 11 passed, 0 failed - cargo test -p openab-core --lib control_plane: 120 passed, 0 failed Not run: Linux (not required for this commit). --- Cargo.lock | 1 + crates/openab-core/Cargo.toml | 7 + crates/openab-core/tests/cp_client.rs | 855 ++++++++++++++++++++++++++ 3 files changed, 863 insertions(+) create mode 100644 crates/openab-core/tests/cp_client.rs diff --git a/Cargo.lock b/Cargo.lock index af36789ac..625dc516c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2509,6 +2509,7 @@ dependencies = [ "aws-sdk-s3", "aws-sdk-secretsmanager", "aws-sigv4", + "axum", "base64", "bytes", "chrono", diff --git a/crates/openab-core/Cargo.toml b/crates/openab-core/Cargo.toml index 10d5811dc..317cfe285 100644 --- a/crates/openab-core/Cargo.toml +++ b/crates/openab-core/Cargo.toml @@ -64,6 +64,13 @@ libc = "0.2" [dev-dependencies] tokio = { version = "1", features = ["test-util"] } +# The control-plane client integration test drives the REAL CP server +# in-process (ephemeral loopback port), so the test build — and only the test +# build — needs the `server` feature. +openab-cp = { path = "../openab-cp", features = ["server"] } +# ... and `axum::serve` to actually bind it. Test-only: the runtime is a +# WebSocket *client* of the CP and links no HTTP server. +axum = "0.8" [features] default = ["discord", "slack", "secrets-aws", "agentcore", "config-s3"] diff --git a/crates/openab-core/tests/cp_client.rs b/crates/openab-core/tests/cp_client.rs new file mode 100644 index 000000000..38a99d2a2 --- /dev/null +++ b/crates/openab-core/tests/cp_client.rs @@ -0,0 +1,855 @@ +//! Control-plane client against the REAL control plane. +//! +//! The client's contract is a conversation, not a function: the first frame +//! must be `cp/register`, heartbeats must keep a lease alive, a delegation must +//! come back as a `cp/delegate_result` the *initiator* receives, and a +//! CP-initiated close must be recovered by re-registering. None of that is +//! observable from unit tests of the client alone — a frame the CP would reject +//! looks identical to one it accepts — so this boots `openab-cp` in-process on +//! an ephemeral loopback port and drives the real thing. +//! +//! Two participants: +//! - the **worker** is the real [`ControlPlaneClient`], with a scripted +//! [`PromptRunner`] in place of a coding agent (the ACP layer is not what is +//! under test here, and requiring a real agent would make this untestable in +//! CI); +//! - the **initiator** is a raw WebSocket client acting as the `primary`, so +//! the assertions are made where a real initiator would make them. + +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::{Arc, Mutex}; +use std::time::{Duration, Instant}; + +use async_trait::async_trait; +use futures_util::{SinkExt, StreamExt}; +use tokio::net::TcpStream; +use tokio_tungstenite::tungstenite::client::IntoClientRequest; +use tokio_tungstenite::tungstenite::Message; +use tokio_tungstenite::{MaybeTlsStream, WebSocketStream}; + +use openab_core::config::{ControlPlaneConfig, CpAgentType}; +use openab_core::control_plane::{ControlPlaneClient, PromptOutcome, PromptRunner}; +use openab_cp::config::CpConfig; +use openab_cp::proto::DelegateForward; +use openab_cp::server::{app, run_sweeper, sweep_leases, AppState}; + +const PRIMARY_KEY: &str = "k-primary"; +const WORKER_KEY: &str = "k-worker"; + +type Ws = WebSocketStream>; + +// --------------------------------------------------------------------------- +// Scripted prompt runner +// --------------------------------------------------------------------------- + +/// What the fake agent does with a delegated prompt. +#[derive(Clone, Copy)] +enum Script { + /// Answer immediately. + Answer, + /// Never answer on its own — the cancel/deadline paths need a turn that is + /// still running when they fire. + Hang, +} + +struct ScriptedRunner { + script: Script, + /// Prompts as the runner saw them, i.e. what the ACP layer would be given. + prompts: Mutex>, + started: AtomicUsize, + cancelled: Mutex>, + discarded: Mutex>, +} + +impl ScriptedRunner { + fn new(script: Script) -> Arc { + Arc::new(Self { + script, + prompts: Mutex::new(Vec::new()), + started: AtomicUsize::new(0), + cancelled: Mutex::new(Vec::new()), + discarded: Mutex::new(Vec::new()), + }) + } + fn started(&self) -> usize { + self.started.load(Ordering::Relaxed) + } + fn cancelled(&self) -> Vec { + self.cancelled.lock().unwrap().clone() + } + fn discarded(&self) -> Vec { + self.discarded.lock().unwrap().clone() + } + fn prompts(&self) -> Vec { + self.prompts.lock().unwrap().clone() + } +} + +#[async_trait] +impl PromptRunner for ScriptedRunner { + async fn run( + &self, + _session_key: &str, + forward: &DelegateForward, + ) -> anyhow::Result { + self.started.fetch_add(1, Ordering::Relaxed); + self.prompts.lock().unwrap().push(forward.prompt.clone()); + match self.script { + Script::Answer => Ok(PromptOutcome { + text: format!("done: {}", forward.prompt), + ..Default::default() + }), + Script::Hang => { + // Far longer than any deadline in this file: only cancellation + // or the local timeout may end it. + tokio::time::sleep(Duration::from_secs(3600)).await; + unreachable!("the hanging script must never answer") + } + } + } + + async fn cancel(&self, session_key: &str) { + self.cancelled.lock().unwrap().push(session_key.to_string()); + } + + async fn discard(&self, session_key: &str) { + self.discarded.lock().unwrap().push(session_key.to_string()); + } +} + +// --------------------------------------------------------------------------- +// Harness +// --------------------------------------------------------------------------- + +fn cp_config(extra: &str) -> CpConfig { + let raw = format!( + r#" +{extra} + +[[agents]] +key = "{PRIMARY_KEY}" +namespace = "prod" +name = "koudu" +type = "primary" + +[[agents]] +key = "{WORKER_KEY}" +namespace = "prod" +name = "worker-1" +type = "worker" +"# + ); + let cfg: CpConfig = toml::from_str(&raw).expect("CP test config parses"); + cfg.validate().expect("CP test config validates"); + cfg +} + +/// Boot a real CP on an ephemeral loopback port, with its sweeper running. +async fn spawn_cp(cfg: CpConfig) -> (Arc, String) { + let state = Arc::new(AppState::new(cfg)); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let router = app(state.clone()); + tokio::spawn(async move { + let _ = axum::serve(listener, router).await; + }); + // Lease expiry + delegation deadline sweeps, exactly as the binary runs them. + tokio::spawn(run_sweeper(state.clone())); + (state, format!("ws://{addr}/cp")) +} + +fn worker_cfg(url: &str) -> ControlPlaneConfig { + ControlPlaneConfig { + url: url.to_string(), + allow_insecure_transport: false, + auth_key: WORKER_KEY.to_string(), + namespace: "prod".into(), + name: "worker-1".into(), + agent_type: CpAgentType::Worker, + labels: [("backend".to_string(), "kiro".to_string())] + .into_iter() + .collect(), + max_delegated_sessions: 2, + } +} + +/// Start the real client. `prompt_hard_timeout` is the runtime's own per-turn +/// ceiling — the other clock bounding a delegation. +fn spawn_worker( + url: &str, + runner: Arc, + prompt_hard_timeout: Duration, +) -> ( + tokio::sync::watch::Sender, + tokio::task::JoinHandle<()>, + Arc, +) { + let (tx, rx) = tokio::sync::watch::channel(false); + let client = Arc::new(ControlPlaneClient::new( + worker_cfg(url), + runner, + prompt_hard_timeout, + )); + let handle = tokio::spawn(Arc::clone(&client).run(rx)); + (tx, handle, client) +} + +/// Raw initiator connection: a `primary` that registers by hand. +async fn connect_initiator(url: &str) -> Ws { + let mut req = url.into_client_request().unwrap(); + req.headers_mut().insert( + "authorization", + format!("Bearer {PRIMARY_KEY}").parse().unwrap(), + ); + let (mut ws, _) = tokio_tungstenite::connect_async(req).await.unwrap(); + let register = serde_json::json!({ + "jsonrpc": "2.0", "id": 1, "method": "cp/register", + "params": { + "protocol_version": 1, + "namespace": "prod", + "name": "koudu", + "type": "primary", + "instance_id": "i-initiator" + } + }) + .to_string(); + ws.send(Message::Text(register)).await.unwrap(); + let ack = next_json(&mut ws).await; + assert_eq!(ack["result"]["protocol_version"], 1, "initiator ack: {ack}"); + ws +} + +async fn next_json(ws: &mut Ws) -> serde_json::Value { + loop { + match tokio::time::timeout(Duration::from_secs(10), ws.next()) + .await + .expect("a frame within 10s") + .expect("the socket is open") + .expect("a valid frame") + { + Message::Text(text) => return serde_json::from_str(&text).expect("JSON frame"), + Message::Close(_) => panic!("the CP closed the initiator's socket"), + _ => continue, + } + } +} + +/// Read frames until one matches `method`, answering nothing else. +async fn next_request(ws: &mut Ws, method: &str) -> serde_json::Value { + let deadline = Instant::now() + Duration::from_secs(15); + while Instant::now() < deadline { + let v = next_json(ws).await; + if v["method"] == method { + return v; + } + } + panic!("no {method} frame arrived"); +} + +/// Wait until `predicate` holds, polling the CP's own registry/router state. +async fn wait_for(label: &str, mut predicate: impl FnMut() -> bool) { + let deadline = Instant::now() + Duration::from_secs(15); + while Instant::now() < deadline { + if predicate() { + return; + } + tokio::time::sleep(Duration::from_millis(25)).await; + } + panic!("timed out waiting for {label}"); +} + +fn worker_instances(state: &Arc) -> Vec { + state + .registry + .list("prod") + .into_iter() + .filter(|i| i.name == "worker-1") + .map(|i| i.instance_id) + .collect() +} + +async fn delegate(ws: &mut Ws, id: u64, delegation_id: &str, prompt: &str, deadline_secs: i64) { + let frame = serde_json::json!({ + "jsonrpc": "2.0", "id": id, "method": "cp/delegate", + "params": { + "delegation_id": delegation_id, + "target": {"name": "worker-1"}, + "prompt": prompt, + "deadline": (chrono::Utc::now() + chrono::Duration::seconds(deadline_secs)).to_rfc3339() + } + }) + .to_string(); + ws.send(Message::Text(frame)).await.unwrap(); +} + +// --------------------------------------------------------------------------- +// Tests +// --------------------------------------------------------------------------- + +/// Registration is the client's first frame, and the CP must accept the +/// identity it asserts — namespace, name, role, labels, and advertised budget +/// all come from `[control_plane]`. +#[tokio::test] +async fn the_client_registers_and_is_visible_in_the_registry() { + let (state, url) = spawn_cp(cp_config("")).await; + let runner = ScriptedRunner::new(Script::Answer); + let (shutdown, handle, client) = + spawn_worker(&url, Arc::clone(&runner), Duration::from_secs(60)); + + wait_for("the worker to register", || { + !worker_instances(&state).is_empty() + }) + .await; + + let worker = state + .registry + .list("prod") + .into_iter() + .find(|i| i.name == "worker-1") + .expect("the worker is registered"); + assert_eq!(worker.instance_id, client.instance_id()); + assert_eq!( + worker.labels.get("backend").map(String::as_str), + Some("kiro") + ); + assert_eq!(worker.max_delegated_sessions, 2, "the advertised budget"); + assert_eq!(worker.active_sessions, 0); + + let _ = shutdown.send(true); + let _ = tokio::time::timeout(Duration::from_secs(5), handle).await; +} + +/// Heartbeats are what hold the lease. Without them the CP deregisters the +/// instance and fails its in-flight delegations, so a client that registers but +/// never heartbeats is worse than one that never connected. +#[tokio::test] +async fn heartbeats_keep_the_lease_alive() { + // A 1s cadence with a 3s lease: two full sweeps' worth of misses would + // expire it. + let (state, url) = spawn_cp(cp_config( + "heartbeat_interval_secs = 1\nlease_expiry_secs = 3", + )) + .await; + let runner = ScriptedRunner::new(Script::Answer); + let (shutdown, handle, _client) = + spawn_worker(&url, Arc::clone(&runner), Duration::from_secs(60)); + wait_for("the worker to register", || { + !worker_instances(&state).is_empty() + }) + .await; + + tokio::time::sleep(Duration::from_millis(3500)).await; + assert!( + state.registry.expired(Duration::from_secs(2)).is_empty(), + "every lease is fresh, so heartbeats are landing" + ); + assert_eq!( + worker_instances(&state).len(), + 1, + "the CP's own sweeper left the registration alone" + ); + + let _ = shutdown.send(true); + let _ = tokio::time::timeout(Duration::from_secs(5), handle).await; +} + +/// The full round trip, asserted where a real initiator sees it: `cp/delegate` +/// is acked with the assignment, the prompt reaches the runner, and the reply +/// comes back as a `cp/delegate_result` with `status = completed`. +#[tokio::test] +async fn a_delegation_round_trips_from_initiator_to_worker_and_back() { + let (state, url) = spawn_cp(cp_config("")).await; + let runner = ScriptedRunner::new(Script::Answer); + let (shutdown, handle, _client) = + spawn_worker(&url, Arc::clone(&runner), Duration::from_secs(60)); + wait_for("the worker to register", || { + !worker_instances(&state).is_empty() + }) + .await; + + let mut initiator = connect_initiator(&url).await; + delegate(&mut initiator, 2, "d-round-trip", "ship it", 60).await; + + let ack = next_json(&mut initiator).await; + assert_eq!(ack["id"], 2, "the delegate ack: {ack}"); + assert_eq!(ack["result"]["assigned_to"], "prod/worker-1"); + + let result = next_request(&mut initiator, "cp/delegate_result").await; + assert_eq!(result["params"]["delegation_id"], "d-round-trip"); + assert_eq!( + result["params"]["status"], "completed", + "the worker reported completion: {result}" + ); + assert_eq!(result["params"]["result"], "done: ship it"); + assert_eq!( + runner.prompts(), + vec!["ship it".to_string()], + "the prompt reached the agent seam verbatim" + ); + + let _ = shutdown.send(true); + let _ = tokio::time::timeout(Duration::from_secs(5), handle).await; +} + +/// A turn that outruns the runtime's own ceiling is ended locally and reported +/// as `timeout`, rather than being left for the CP's deadline sweep. The +/// delegation deadline here is 60s and the local ceiling 1s, so the status the +/// initiator sees can only have come from the client. +#[tokio::test] +async fn the_local_deadline_reports_timeout_to_the_initiator() { + let (state, url) = spawn_cp(cp_config("")).await; + let runner = ScriptedRunner::new(Script::Hang); + let (shutdown, handle, _client) = + spawn_worker(&url, Arc::clone(&runner), Duration::from_secs(1)); + wait_for("the worker to register", || { + !worker_instances(&state).is_empty() + }) + .await; + + let mut initiator = connect_initiator(&url).await; + delegate(&mut initiator, 3, "d-timeout", "hang forever", 60).await; + assert_eq!(next_json(&mut initiator).await["id"], 3, "delegate ack"); + + let result = next_request(&mut initiator, "cp/delegate_result").await; + assert_eq!(result["params"]["delegation_id"], "d-timeout"); + assert_eq!(result["params"]["status"], "timeout", "{result}"); + assert_eq!(runner.started(), 1, "it really did start"); + // Local deadline cleanup: the agent is interrupted and the single-use + // session is dropped, so nothing is left behind for the next delegation. + wait_for("the session to be cleaned up", || { + !runner.cancelled().is_empty() && !runner.discarded().is_empty() + }) + .await; + + let _ = shutdown.send(true); + let _ = tokio::time::timeout(Duration::from_secs(5), handle).await; +} + +/// `cp/cancel` from the initiator must stop the turn *at the worker* and free +/// the slot. +/// +/// The CP removes the in-flight entry when it forwards the cancel and tells the +/// initiator nothing further (by design — the initiator asked), so the +/// observable effects are on the worker side: the agent is interrupted, the +/// single-use session is discarded, and the freed capacity accepts the next +/// delegation. The status mapping itself is pinned by the executor's unit tests. +#[tokio::test] +async fn a_cancel_stops_the_turn_and_frees_the_slot() { + let (state, url) = spawn_cp(cp_config("")).await; + let runner = ScriptedRunner::new(Script::Hang); + let (shutdown, handle, client) = + spawn_worker(&url, Arc::clone(&runner), Duration::from_secs(60)); + wait_for("the worker to register", || { + !worker_instances(&state).is_empty() + }) + .await; + + let mut initiator = connect_initiator(&url).await; + delegate(&mut initiator, 4, "d-cancel", "hang forever", 300).await; + let ack = next_json(&mut initiator).await; + assert_eq!(ack["id"], 4, "delegate ack"); + let admission = ack["result"]["admission"] + .as_u64() + .expect("the ack names the admission this cancel must target"); + wait_for("the turn to start", || runner.started() == 1).await; + + let cancel = serde_json::json!({ + "jsonrpc": "2.0", "id": 5, "method": "cp/cancel", + "params": {"delegation_id": "d-cancel", "admission": admission, "reason": "initiator changed its mind"} + }) + .to_string(); + initiator.send(Message::Text(cancel)).await.unwrap(); + + let expected_session = openab_core::control_plane::delegation_session_key( + client.instance_id(), + "d-cancel", + admission, + ); + wait_for("the worker to unwind the cancelled turn", || { + runner.cancelled().contains(&expected_session) + && runner.discarded().contains(&expected_session) + }) + .await; + + // The slot is free again: a second delegation is admitted and answered. + // (Same id would be refused as a duplicate, so use a new one.) + let runner2 = Arc::clone(&runner); + delegate(&mut initiator, 6, "d-after-cancel", "and now this", 60).await; + wait_for("the next delegation to start", || runner2.started() == 2).await; + + let _ = shutdown.send(true); + let _ = tokio::time::timeout(Duration::from_secs(5), handle).await; +} + +/// Lease expiry is the CP dropping a registration on its own initiative: it +/// closes the socket, because registration is first-frame-only and a connection +/// whose registry entry is gone can never recover. The client's job is to notice +/// and re-register — with the SAME instance id, since that identifies the +/// process, not the connection — and to be usable again afterwards. +#[tokio::test] +async fn a_cp_initiated_close_is_recovered_by_re_registering() { + let (state, url) = spawn_cp(cp_config("")).await; + let runner = ScriptedRunner::new(Script::Answer); + let (shutdown, handle, client) = + spawn_worker(&url, Arc::clone(&runner), Duration::from_secs(60)); + wait_for("the worker to register", || { + !worker_instances(&state).is_empty() + }) + .await; + let first_handle = state + .registry + .list("prod") + .into_iter() + .find(|i| i.name == "worker-1") + .map(|i| i.handle) + .expect("registered"); + + // Zero-window sweep: expire every lease, exactly as the sweeper would after + // a missed-heartbeat window. + sweep_leases(&state, Duration::ZERO); + assert!( + worker_instances(&state).is_empty(), + "the CP dropped the registration" + ); + + wait_for("the client to re-register", || { + state + .registry + .list("prod") + .into_iter() + .any(|i| i.name == "worker-1" && i.handle != first_handle) + }) + .await; + let second = state + .registry + .list("prod") + .into_iter() + .find(|i| i.name == "worker-1") + .expect("re-registered"); + assert_eq!( + second.instance_id, + client.instance_id(), + "one instance id per process, reused across reconnects" + ); + + // The new registration is not just present, it works. + let mut initiator = connect_initiator(&url).await; + delegate(&mut initiator, 7, "d-after-reconnect", "still there?", 60).await; + let ack = next_json(&mut initiator).await; + assert_eq!(ack["result"]["assigned_to"], "prod/worker-1", "{ack}"); + let result = next_request(&mut initiator, "cp/delegate_result").await; + assert_eq!(result["params"]["status"], "completed", "{result}"); + assert_eq!(result["params"]["result"], "done: still there?"); + + let _ = shutdown.send(true); + let _ = tokio::time::timeout(Duration::from_secs(5), handle).await; +} + +/// Shutdown must be a clean disconnect, not a lease timeout: the CP has to be +/// able to tell "this replica went away" from "this replica stopped answering". +#[tokio::test] +async fn shutdown_deregisters_the_instance() { + let (state, url) = spawn_cp(cp_config("")).await; + let runner = ScriptedRunner::new(Script::Answer); + let (shutdown, handle, _client) = + spawn_worker(&url, Arc::clone(&runner), Duration::from_secs(60)); + wait_for("the worker to register", || { + !worker_instances(&state).is_empty() + }) + .await; + + let _ = shutdown.send(true); + tokio::time::timeout(Duration::from_secs(5), handle) + .await + .expect("the client task ends on the shutdown signal") + .expect("without panicking"); + + wait_for("the CP to see the disconnect", || { + worker_instances(&state).is_empty() + }) + .await; +} + +/// Shutdown during an active delegation must be owned by the serve loop: it +/// cancels and discards the local session before returning, rather than an +/// outer select dropping the serve future and detaching its task. +#[tokio::test] +async fn shutdown_drains_an_inflight_delegation_before_exit() { + let (state, url) = spawn_cp(cp_config("")).await; + let runner = ScriptedRunner::new(Script::Hang); + let (shutdown, handle, _client) = + spawn_worker(&url, Arc::clone(&runner), Duration::from_secs(60)); + wait_for("the worker to register", || { + !worker_instances(&state).is_empty() + }) + .await; + + let mut initiator = connect_initiator(&url).await; + delegate(&mut initiator, 8, "d-shutdown", "hang until shutdown", 300).await; + assert_eq!(next_json(&mut initiator).await["id"], 8, "delegate ack"); + wait_for("the delegated turn to start", || runner.started() == 1).await; + + let _ = shutdown.send(true); + tokio::time::timeout(Duration::from_secs(7), handle) + .await + .expect("client drains within its bounded window") + .expect("client task does not panic"); + assert_eq!(runner.cancelled().len(), 1, "agent turn was interrupted"); + assert_eq!( + runner.discarded().len(), + 1, + "single-use session was discarded" + ); + wait_for("the CP to see the drained disconnect", || { + worker_instances(&state).is_empty() + }) + .await; +} + +/// The PR4 local API against two real runtime clients and the real CP server: +/// primary UDS -> primary client -> CP -> worker client -> scripted agent -> result. +#[tokio::test] +async fn primary_local_api_lists_spawns_and_awaits_over_the_real_cp() { + use openab_core::control_plane::{ + serve_local, ClientCommand, LocalClient, LocalRequest, LocalResponse, LocalTarget, + }; + + let (state, url) = spawn_cp(cp_config("")).await; + let worker_runner = ScriptedRunner::new(Script::Answer); + let (worker_shutdown, worker_task, _worker) = + spawn_worker(&url, Arc::clone(&worker_runner), Duration::from_secs(60)); + + let primary_runner = ScriptedRunner::new(Script::Answer); + let mut cfg = worker_cfg(&url); + cfg.auth_key = PRIMARY_KEY.to_string(); + cfg.name = "koudu".into(); + cfg.agent_type = CpAgentType::Primary; + let (primary_shutdown, primary_rx) = tokio::sync::watch::channel(false); + let primary = Arc::new(ControlPlaneClient::new( + cfg, + primary_runner, + Duration::from_secs(60), + )); + let primary_handle = primary.handle(); + let primary_task = tokio::spawn(primary.run(primary_rx)); + + wait_for("both runtimes to register", || { + state.registry.list("prod").len() >= 2 + }) + .await; + + let dir = tempfile::tempdir().unwrap(); + let socket = dir.path().join("run").join("agent.sock"); + let (local_shutdown, local_rx) = tokio::sync::watch::channel(false); + let local_task = tokio::spawn(serve_local(socket.clone(), primary_handle, local_rx)); + wait_for("local agent socket", || socket.exists()).await; + let local = LocalClient::new(socket); + + let list = local.request(&LocalRequest::ListAgents).await.unwrap(); + let LocalResponse::Agents { agents } = list else { + panic!("expected agent roster"); + }; + assert!(agents.iter().any(|a| a.name == "worker-1")); + + let spawned = local + .request(&LocalRequest::Spawn { + delegation_id: "d-local-e2e".into(), + target: LocalTarget { + name: Some("worker-1".into()), + labels: None, + }, + prompt: "ship it".into(), + deadline_secs: 60, + parent: None, + }) + .await + .unwrap(); + let LocalResponse::Spawned { + handle, + assigned_to, + } = spawned + else { + panic!("expected spawned handle"); + }; + assert_eq!(assigned_to, "prod/worker-1"); + assert_eq!(handle.len(), 64, "opaque random handle token"); + + let terminal = local + .request(&LocalRequest::Await { handle }) + .await + .unwrap(); + let LocalResponse::Terminal { + status, + result, + error, + } = terminal + else { + panic!("expected terminal result"); + }; + assert_eq!( + status, + openab_core::control_plane::DelegationStatusWire::Completed + ); + assert_eq!(result.as_deref(), Some("done: ship it")); + assert!(error.is_none()); + + let _ = local_shutdown.send(true); + let _ = primary_shutdown.send(true); + let _ = worker_shutdown.send(true); + let _ = local_task.await; + let _ = primary_task.await; + let _ = worker_task.await; + + // Compile-time assertion that the direct command type remains exported for + // in-process callers as well as the UDS frontends. + let _unused: Option = None; +} + +#[tokio::test] +async fn primary_local_api_cancels_a_running_delegation_over_the_real_cp() { + use openab_core::control_plane::{ + serve_local, LocalClient, LocalRequest, LocalResponse, LocalTarget, + }; + + let (state, url) = spawn_cp(cp_config("max_deadline_secs = 600")).await; + let worker_runner = ScriptedRunner::new(Script::Hang); + let (worker_shutdown, worker_task, _worker) = + spawn_worker(&url, Arc::clone(&worker_runner), Duration::from_secs(600)); + + let mut cfg = worker_cfg(&url); + cfg.auth_key = PRIMARY_KEY.to_string(); + cfg.name = "koudu".into(); + cfg.agent_type = CpAgentType::Primary; + let (primary_shutdown, primary_rx) = tokio::sync::watch::channel(false); + let primary = Arc::new(ControlPlaneClient::new( + cfg, + ScriptedRunner::new(Script::Answer), + Duration::from_secs(600), + )); + let handle = primary.handle(); + let primary_task = tokio::spawn(primary.run(primary_rx)); + wait_for("both runtimes to register", || { + state.registry.list("prod").len() >= 2 + }) + .await; + + let dir = tempfile::tempdir().unwrap(); + let socket = dir.path().join("run").join("agent.sock"); + let (local_shutdown, local_rx) = tokio::sync::watch::channel(false); + let local_task = tokio::spawn(serve_local(socket.clone(), handle, local_rx)); + wait_for("local agent socket", || socket.exists()).await; + let local = LocalClient::new(socket); + + let spawned = local + .request(&LocalRequest::Spawn { + delegation_id: "d-local-cancel".into(), + target: LocalTarget { + name: Some("worker-1".into()), + labels: None, + }, + prompt: "hang".into(), + deadline_secs: 300, + parent: None, + }) + .await + .unwrap(); + let LocalResponse::Spawned { handle, .. } = spawned else { + panic!("expected spawned handle"); + }; + wait_for("worker prompt to start", || worker_runner.started() == 1).await; + let cancelled = local + .request(&LocalRequest::Cancel { + handle, + reason: "operator cancelled".into(), + }) + .await + .unwrap(); + assert_eq!(cancelled, LocalResponse::Cancelled); + wait_for("worker cleanup", || { + !worker_runner.cancelled().is_empty() && !worker_runner.discarded().is_empty() + }) + .await; + + let _ = local_shutdown.send(true); + let _ = primary_shutdown.send(true); + let _ = worker_shutdown.send(true); + let _ = local_task.await; + let _ = primary_task.await; + let _ = worker_task.await; +} + +#[tokio::test] +async fn oversized_inprocess_spawn_is_rejected_without_disconnect() { + use openab_core::control_plane::{ClientCommand, CommandError, SpawnRequest, SpawnTarget}; + + let (state, url) = spawn_cp(cp_config("")).await; + let worker_runner = ScriptedRunner::new(Script::Answer); + let (worker_shutdown, worker_task, _worker) = + spawn_worker(&url, Arc::clone(&worker_runner), Duration::from_secs(60)); + + let mut cfg = worker_cfg(&url); + cfg.auth_key = PRIMARY_KEY.to_string(); + cfg.name = "koudu".into(); + cfg.agent_type = CpAgentType::Primary; + let (primary_shutdown, primary_rx) = tokio::sync::watch::channel(false); + let primary = Arc::new(ControlPlaneClient::new( + cfg, + ScriptedRunner::new(Script::Answer), + Duration::from_secs(60), + )); + let handle = primary.handle(); + let primary_task = tokio::spawn(primary.run(primary_rx)); + wait_for("both runtimes to register", || { + state.registry.list("prod").len() >= 2 + }) + .await; + + let (spawn_tx, spawn_rx) = tokio::sync::oneshot::channel(); + handle + .submit(ClientCommand::Spawn { + request: SpawnRequest { + delegation_id: "d-oversized".into(), + target: SpawnTarget::by_name("worker-1"), + prompt: "x".repeat(1024 * 1024 + 4096), + deadline: chrono::Utc::now() + chrono::Duration::seconds(60), + parent: None, + }, + reply: spawn_tx, + }) + .await + .unwrap(); + let result = tokio::time::timeout(Duration::from_secs(5), spawn_rx) + .await + .expect("oversized request is answered locally") + .unwrap(); + assert!(matches!(result, Err(CommandError::InvalidRequest(_)))); + + // The same connection remains usable: a subsequent list request reaches + // the real CP and returns the worker instead of reconnecting/flapping. + let (list_tx, list_rx) = tokio::sync::oneshot::channel(); + handle + .submit(ClientCommand::ListAgents { reply: list_tx }) + .await + .unwrap(); + let agents = tokio::time::timeout(Duration::from_secs(5), list_rx) + .await + .expect("connection remains usable") + .unwrap() + .unwrap(); + assert!(agents.iter().any(|agent| agent.name == "worker-1")); + assert_eq!( + state + .registry + .list("prod") + .into_iter() + .filter(|agent| agent.name == "koudu") + .count(), + 1, + "primary stayed registered on the same live connection" + ); + + let _ = primary_shutdown.send(true); + let _ = worker_shutdown.send(true); + let _ = primary_task.await; + let _ = worker_task.await; +} From 90a0100fb79014b6cee80a934d2efdc3346351a9 Mon Sep 17 00:00:00 2001 From: Tautui Date: Sat, 31 Oct 2026 10:44:02 -0400 Subject: [PATCH 07/15] test(cp): Phase 0 two-worker acceptance tests Add two in-process tests to crates/openab-core/tests/cp_client.rs, using the real CP server and the real ControlPlaneClient with a scripted runner: - phase0_parallel_spawns_land_on_different_workers: two workers with max_delegated_sessions = 1 and the same label. Two cp/delegate frames are written before either ack is read. Asserts the two acks name worker-1 and worker-2, and that each worker runner started exactly one delegation. - phase0_worker_identity_cannot_initiate: a raw worker-type connection sends cp/delegate. Asserts error code -32003 (POLICY_DENIED) and no in-flight entry. R9 (least-loaded tie-break): no change needed. The first test passed without modification. Router::delegate runs select and capacity reservation under the admission guard, so the first spawn is counted before the second selection. Commands run on macmini (cargo 1.99.1): - cargo test -p openab-core --test cp_client phase0_parallel_spawns_land_on_different_workers: ok - cargo test -p openab-core --test cp_client phase0_worker_identity_cannot_initiate: ok - cargo test -p openab-core --test cp_client: 13 passed, 0 failed - cargo test -p openab-core control_plane: 120 passed, 0 failed - cargo check -p openab-core -p openab: ok - rustfmt --check on cp_client.rs: ok (after rustfmt applied) --- crates/openab-core/tests/cp_client.rs | 160 ++++++++++++++++++++++++++ 1 file changed, 160 insertions(+) diff --git a/crates/openab-core/tests/cp_client.rs b/crates/openab-core/tests/cp_client.rs index 38a99d2a2..cd6b42d95 100644 --- a/crates/openab-core/tests/cp_client.rs +++ b/crates/openab-core/tests/cp_client.rs @@ -219,6 +219,30 @@ async fn connect_initiator(url: &str) -> Ws { ws } +/// Raw connection that registers as `agent_type` under `key` and `name`. +/// Lets a test send frames a real runtime would never send. +async fn connect_raw(url: &str, key: &str, name: &str, agent_type: &str) -> Ws { + let mut req = url.into_client_request().unwrap(); + req.headers_mut() + .insert("authorization", format!("Bearer {key}").parse().unwrap()); + let (mut ws, _) = tokio_tungstenite::connect_async(req).await.unwrap(); + let register = serde_json::json!({ + "jsonrpc": "2.0", "id": 1, "method": "cp/register", + "params": { + "protocol_version": 1, + "namespace": "prod", + "name": name, + "type": agent_type, + "instance_id": format!("i-raw-{name}") + } + }) + .to_string(); + ws.send(Message::Text(register)).await.unwrap(); + let ack = next_json(&mut ws).await; + assert_eq!(ack["result"]["protocol_version"], 1, "raw ack: {ack}"); + ws +} + async fn next_json(ws: &mut Ws) -> serde_json::Value { loop { match tokio::time::timeout(Duration::from_secs(10), ws.next()) @@ -853,3 +877,139 @@ async fn oversized_inprocess_spawn_is_rejected_without_disconnect() { let _ = primary_task.await; let _ = worker_task.await; } + +/// Phase 0 criterion 3: two parallel spawns land on different workers. +/// +/// Both workers advertise `max_delegated_sessions = 1` and carry the same +/// label, so a spawn selected by label has two candidates. The first spawn +/// fills its worker, so the second must pick the other one. If the first +/// spawn's capacity were not counted before the second selection, both would +/// land on the same worker. +#[tokio::test] +async fn phase0_parallel_spawns_land_on_different_workers() { + let raw = format!( + r#" +[[agents]] +key = "{PRIMARY_KEY}" +namespace = "prod" +name = "koudu" +type = "primary" + +[[agents]] +key = "{WORKER_KEY}" +namespace = "prod" +name = "worker-1" +type = "worker" + +[[agents]] +key = "k-worker-2" +namespace = "prod" +name = "worker-2" +type = "worker" +"# + ); + let cfg: CpConfig = toml::from_str(&raw).expect("two-worker config parses"); + cfg.validate().expect("two-worker config validates"); + let (state, url) = spawn_cp(cfg).await; + + let one_slot = |key: &str, name: &str| ControlPlaneConfig { + auth_key: key.to_string(), + name: name.into(), + max_delegated_sessions: 1, + ..worker_cfg(&url) + }; + let runner_1 = ScriptedRunner::new(Script::Hang); + let runner_2 = ScriptedRunner::new(Script::Hang); + let (tx_1, rx_1) = tokio::sync::watch::channel(false); + let (tx_2, rx_2) = tokio::sync::watch::channel(false); + let worker_1 = Arc::new(ControlPlaneClient::new( + one_slot(WORKER_KEY, "worker-1"), + runner_1.clone(), + Duration::from_secs(600), + )); + let worker_2 = Arc::new(ControlPlaneClient::new( + one_slot("k-worker-2", "worker-2"), + runner_2.clone(), + Duration::from_secs(600), + )); + let task_1 = tokio::spawn(Arc::clone(&worker_1).run(rx_1)); + let task_2 = tokio::spawn(Arc::clone(&worker_2).run(rx_2)); + wait_for("both workers to register", || { + state.registry.list("prod").len() >= 2 + }) + .await; + + // Two spawns, both written before either reply is read. + let mut initiator = connect_initiator(&url).await; + for (id, delegation_id) in [(2u64, "d-par-1"), (3u64, "d-par-2")] { + let frame = serde_json::json!({ + "jsonrpc": "2.0", "id": id, "method": "cp/delegate", + "params": { + "delegation_id": delegation_id, + "target": {"labels": {"backend": "kiro"}}, + "prompt": "sleep 30; hostname", + "deadline": (chrono::Utc::now() + chrono::Duration::seconds(60)).to_rfc3339() + } + }) + .to_string(); + initiator.send(Message::Text(frame)).await.unwrap(); + } + + let ack_1 = next_json(&mut initiator).await; + let ack_2 = next_json(&mut initiator).await; + assert_eq!(ack_1["id"], 2, "first ack: {ack_1}"); + assert_eq!(ack_2["id"], 3, "second ack: {ack_2}"); + let landed_1 = ack_1["result"]["assigned_to"].as_str().unwrap().to_string(); + let landed_2 = ack_2["result"]["assigned_to"].as_str().unwrap().to_string(); + assert_ne!( + landed_1, landed_2, + "two parallel spawns must not share a worker" + ); + let mut landed = vec![landed_1, landed_2]; + landed.sort(); + assert_eq!(landed, vec!["prod/worker-1", "prod/worker-2"]); + + // Each worker's runner saw exactly one of the two prompts. + wait_for("each worker to start its one delegation", || { + runner_1.started() == 1 && runner_2.started() == 1 + }) + .await; + + let _ = tx_1.send(true); + let _ = tx_2.send(true); + let _ = tokio::time::timeout(Duration::from_secs(5), task_1).await; + let _ = tokio::time::timeout(Duration::from_secs(5), task_2).await; +} + +/// Phase 0 criterion 5: a worker identity cannot initiate. The CP refuses a +/// `cp/delegate` from a worker with `POLICY_DENIED` (-32003) and admits nothing. +#[tokio::test] +async fn phase0_worker_identity_cannot_initiate() { + let (state, url) = spawn_cp(cp_config("")).await; + // A registered worker, so the refusal is about who is asking and not about + // a missing target. The target is the worker itself. + let mut worker = connect_raw(&url, WORKER_KEY, "worker-1", "worker").await; + wait_for("the worker to register", || { + !worker_instances(&state).is_empty() + }) + .await; + + delegate( + &mut worker, + 2, + "d-worker-initiates", + "should be refused", + 60, + ) + .await; + let reply = next_json(&mut worker).await; + assert_eq!(reply["id"], 2, "the refusal answers the request: {reply}"); + assert!(reply.get("error").is_some(), "no admission: {reply}"); + assert_eq!(reply["error"]["code"], -32003, "POLICY_DENIED: {reply}"); + + assert_eq!( + state.router.inflight_count(), + 0, + "a refused initiation leaves no in-flight entry" + ); +} From 9585c377bc51eb795c9a47523ec5ff1b2090b99e Mon Sep 17 00:00:00 2001 From: Tautui Date: Sat, 31 Oct 2026 10:45:57 -0400 Subject: [PATCH 08/15] docs(cp): [control_plane] config reference Adds the [control_plane] section to docs/config-reference.md, taken from the reference branch feat/cp-runtime-client-main (91c0a89f7) and placed before the [cron] section. Only the [control_plane] section is included; no MCP-tool docs. Edits to the reference text, so the docs match HEAD: - Headless table: a primary with no adapter and no [mcp] is a startup error on HEAD (headless_run_mode returns HeadlessMode::None), not "Full runtime" as the reference docs say. The reference code at 91c0a89f differs; reported separately. - Dropped the MCP-tool paragraphs and rows (facade tools, spawn_agent). - Delegation deadlines bullet now covers only `openab agent spawn`. - auth_key: removed the [secrets.refs] mention; no code path for it was found. - Kept the "One runtime per socket path" bullet from c0611589 (F6), which documents behavior shipped in commit 4 of this branch. Verification: - cargo check -p openab-core -p openab (macmini, rustc 1.99.1): Finished, exit 0. Cached result: this commit changes no Rust source. - CLI flags named in the docs (spawn --deadline-secs, --label, 1..=1800 validation) checked against src/main.rs: match. - Code claims (0700/0600 socket perms, 10 s connect timeout, allow_insecure_transport gate, deadline + 5 bound) checked against HEAD: match. - git diff origin/main -- docs/config-reference.md: additions only, no removed lines. - No automated docs test exists. --- docs/config-reference.md | 93 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 93 insertions(+) diff --git a/docs/config-reference.md b/docs/config-reference.md index 0d3362076..4ec1c3fb9 100644 --- a/docs/config-reference.md +++ b/docs/config-reference.md @@ -773,6 +773,99 @@ permissions scoped to the bucket. --- +## `[control_plane]` + +Enrol this runtime with an [Agent Control Plane](adr/agent-control-plane.md) so +it can delegate to, or serve delegations from, other OAB agents. + +Presence of the section is the opt-in — there is no cargo feature and no +env-var switch. Absent section = no outbound connection, no delegation serving, +and no behaviour change. Unknown keys are a hard startup failure. + +```toml +[control_plane] +url = "wss://cp.example.internal/cp" # the CP mounts the socket at /cp +# allow_insecure_transport = true # only for explicit non-loopback ws:// exceptions +auth_key = "${OPENAB_CP_KEY}" # per-agent credential, never shared +namespace = "prod" +name = "koudu" +type = "worker" # "primary" | "worker" +max_delegated_sessions = 2 # local concurrency budget + +[control_plane.labels] # optional selector labels +backend = "kiro" +tier = "batch" +``` + +| Field | Required | Default | Description | +|-------|----------|---------|-------------| +| `url` | ✅ | — | CP WebSocket endpoint, path `/cp`. Remote endpoints should use `wss://`; `ws://` is accepted only on loopback unless explicitly overridden | +| `allow_insecure_transport` | — | `false` | Permit non-loopback cleartext `ws://`. This sends the bearer key without connection-level encryption and logs a startup warning | +| `auth_key` | ✅ | — | Bearer key sent on the upgrade request. Use `${ENV}` — never a literal | +| `namespace` | ✅ | — | Asserted namespace; the CP verifies it against the key's claims | +| `name` | ✅ | — | Asserted logical agent name; likewise verified | +| `type` | ✅ | — | `primary` (initiates delegations) or `worker` (serves them) | +| `labels` | — | `{}` | Selector labels other agents can target by | +| `max_delegated_sessions` | — | `1` | Concurrency advertised at registration; the CP may clamp it, and the runtime enforces whatever the ack returns | + +`namespace`, `name`, and `type` are **assertions the CP verifies**, not +authorization inputs: each auth key is bound to immutable claims in CP config, +and a mismatch is rejected at registration. They live in config so a +misconfigured runtime fails loudly instead of being silently re-identified. +`observer` is not accepted here — it is a read-only lobby role, not a runtime +one. + +### Headless (no chat adapter) + +A worker may run without a chat adapter. It receives remote delegations. A +primary boots without a chat adapter only when `[mcp]` is also set. Either way, +a primary's initiating surface is the owner-only local socket and `openab agent`. + +| Config | Result | +|--------|--------| +| `[control_plane] type = "worker"`, no adapter | Full runtime — pool + control-plane client, no chat platform | +| `[control_plane] type = "primary"`, no adapter and no `[mcp]` | Startup error — a primary needs a chat adapter or `[mcp]` | +| `[mcp]` only, no adapter | Facade-only mode (unchanged) | +| `[mcp]` + worker, no adapter | Full runtime — facade listener + worker-side control-plane client | +| `[mcp]` + primary, no adapter | Full runtime — facade listener + primary registration and local CLI socket | + +### Operational notes + +- **Local agent API.** `$OPENAB_AGENT_SOCKET` overrides the default + `$HOME/.openab/run/agent.sock`. The runtime creates an owner-only directory + (`0700`) and socket (`0600`). The transport is a **Unix-domain socket and is + Unix-only**: non-Unix builds compile (the server and client are stubbed) but + fail loudly at runtime, and **no Windows support is planned** — CP + deployments are Linux containers or macOS. +- **One runtime per socket path.** Two runtimes under one user must set a + distinct `OPENAB_AGENT_SOCKET` for at least one of them: the runtime refuses + to start over a socket that answers a connect probe and only unlinks a node + nobody is listening on. +- **The key never reaches the agent.** Agent subprocesses start from + `env_clear()` with a fixed baseline plus explicit `[agent].env` keys; + `auth_key` is in neither, and it is never logged. +- **Reconnects are automatic** with equal jitter over 1/2/4/8/16/30s base + backoff and a 10-second connect deadline. One instance id is generated per + process and reused across reconnects, so the CP can tell a reconnecting + replica from a new one without a fleet-wide reconnect herd. +- **Transport is fail-closed.** Remote bearer credentials require `wss://` by + default. Loopback `ws://` is allowed for local development; a non-loopback + cleartext exception requires `allow_insecure_transport = true` and emits a + warning because the bearer key is not encrypted by that connection. +- **Per-turn ceiling.** A delegation is bounded by the nearer of its CP + deadline and `[pool].prompt_hard_timeout_secs`; exceeding the local one is + reported to the initiator as `timeout`. +- **Delegation deadlines are `1..=1800` seconds.** `openab agent spawn + --deadline-secs` rejects a value outside this window before any frame leaves + the host. A blocking spawn bounds admission plus the wait for the result under + one `deadline_secs + 5` timeout, so a hung admission cannot block the caller + past its deadline. +- **One session per delegation.** Each delegation runs in a fresh ACP session + that is discarded when it ends, so delegations never see each other's + conversation and none of them counts against the pool afterwards. + +--- + ## `[cron]` Everything cron-related lives under `[cron]`. From b1f3dc15d17011a8229ce1f9d5c577ecded57eb4 Mon Sep 17 00:00:00 2001 From: chaodu-agent <274062505+chaodu-agent@users.noreply.github.com> Date: Sat, 31 Oct 2026 11:40:53 -0400 Subject: [PATCH 09/15] fix(cp): drop spent supervision before shutdown; close review findings M1, L1-L3, L5, L9 M1: wait_for_control_plane_exit and wait_for_local_api_exit awaited a oneshot receiver inside tokio::select!. Once it yielded, the shutdown path awaited the same receiver again, which tokio panics on ("called after complete"), so pool.shutdown() and pre_shutdown were skipped. Each wait now drops its supervision after the receiver yields, and shutdown skips it. Test: tests::a_supervised_exit_is_not_awaited_again_at_shutdown. L1: openab agent --socket no longer evaluates default_socket_path() when an explicit socket is given, so HOME unset no longer fails an explicit path. L2: the headless-primary comment and the no-adapter error now state that a control-plane primary with no adapter also needs [mcp]. Boot behaviour is unchanged. L3: MIN_RUNTIME_FRAME_BYTES doc says the server does not enforce the floor. L5: client.rs comment matches behaviour: spawn and cancel are refused locally, list_agents is forwarded for either role. L9: spawn and delegation-id doc comments no longer claim the MCP facade calls them. No facade caller exists at this commit. Checks on the tree at this commit (macmini, cargo via ~/.cargo/bin): - cargo clippy -p openab-core -p openab -- -D warnings: pass (exit 0) - cargo test -p openab-core --lib control_plane: 120 passed, 0 failed - cargo test -p openab-core --test cp_client: 14 passed, 0 failed - cargo test -p openab: 33 passed, 0 failed (includes the M1 test) - cargo test --workspace: 895 passed, 1 failed: secrets::tests::resolve_exec_nonzero_exit (pre-existing; secrets.rs is not in this PR; the test expects the "exited with" text on this macOS host) - cargo fmt --all -- --check: fails across most of the repo at the base commit (pre-existing). The lines this commit adds are wrapped to rustfmt style. Commits in this pass were not built one by one. The checks above cover the final tree of the pass. --- .../openab-core/src/control_plane/client.rs | 10 +-- crates/openab-cp/src/proto.rs | 5 +- src/main.rs | 71 ++++++++++++++----- 3 files changed, 62 insertions(+), 24 deletions(-) diff --git a/crates/openab-core/src/control_plane/client.rs b/crates/openab-core/src/control_plane/client.rs index 149278957..1792f1d18 100644 --- a/crates/openab-core/src/control_plane/client.rs +++ b/crates/openab-core/src/control_plane/client.rs @@ -600,10 +600,12 @@ impl ControlPlaneClient { sink: &mut WsSink, shutdown: &mut watch::Receiver, ) -> Result<(), WriteEnd> { - // A primary must be registered as such to initiate; a worker-only - // runtime rejects spawn/cancel/list locally so a misconfigured caller - // gets a clear answer instead of a CP policy denial round-trip. Await - // is allowed for either role (it only reads local tracking). + // A primary must be registered as such to initiate. A worker-only runtime + // rejects spawn and cancel locally, so a misconfigured caller gets a clear + // answer instead of a CP policy denial round-trip. list_agents is forwarded + // for either role; whether the CP accepts a worker's list_agents is not + // verified here. Await is allowed for either role (it only reads local + // tracking). match command { ClientCommand::Spawn { request, reply } => { if self.cfg.agent_type != CpAgentType::Primary { diff --git a/crates/openab-cp/src/proto.rs b/crates/openab-cp/src/proto.rs index acf57b97b..6534c09b6 100644 --- a/crates/openab-cp/src/proto.rs +++ b/crates/openab-cp/src/proto.rs @@ -39,7 +39,10 @@ use serde_json::Value; pub const PROTOCOL_VERSION: u32 = 1; /// Minimum frame ceiling supported by the bundled runtime client. The server /// may configure a larger value, but a smaller one would let an otherwise -/// valid runtime result disconnect the whole worker connection. +/// valid runtime result disconnect the whole worker connection. The server +/// does not enforce this floor: `max_frame_bytes` is not checked against it, +/// so a hub configured below 1 MiB must keep this in mind. The default config +/// is above the floor. pub const MIN_RUNTIME_FRAME_BYTES: usize = 1024 * 1024; // --- JSON-RPC envelopes --- diff --git a/src/main.rs b/src/main.rs index c2d65e3b3..d51c0858e 100644 --- a/src/main.rs +++ b/src/main.rs @@ -210,7 +210,10 @@ async fn run_agent_command(socket: Option, action: AgentAction) -> anyh default_socket_path, LocalClient, LocalRequest, LocalResponse, LocalTarget, }; - let client = LocalClient::new(socket.unwrap_or(default_socket_path()?)); + let client = LocalClient::new(match socket { + Some(path) => path, + None => default_socket_path()?, + }); // A spawn is the one action bounded as a whole operation; the others are a // single request/response. if let AgentAction::Spawn { @@ -242,8 +245,9 @@ async fn run_agent_command(socket: Option, action: AgentAction) -> anyh deadline_secs, parent, }; - // Same full-operation wrapper the MCP facade uses: admission + optional - // await under one `deadline_secs + 5` bound. + // Full-operation wrapper: admission + optional await under one + // `deadline_secs + 5` bound. The MCP facade does not call it at this + // commit; the CLI is the only caller. let client_ref = &client; let response = spawn_within_bound(deadline_secs, async_mode, spawn, move |req| async move { @@ -272,9 +276,10 @@ async fn run_agent_command(socket: Option, action: AgentAction) -> anyh } /// Accepted `deadline_secs` window for `openab agent spawn`, enforced -/// client-side before any frame leaves the host. Matches the MCP tool schema -/// (`minimum: 1, maximum: 1800`) and the local API server's own bound, so all -/// three surfaces agree on the same 1-second..=30-minute window. +/// client-side before any frame leaves the host. Matches the bound the local +/// API server enforces (`minimum: 1, maximum: 1800`), so the CLI and the server +/// agree on the same 1-second..=30-minute window. No MCP tool schema consumes +/// these constants at this commit. const MIN_DEADLINE_SECS: u64 = 1; const MAX_DEADLINE_SECS: u64 = 1800; @@ -306,9 +311,10 @@ fn parse_agent_labels( use openab_core::control_plane::{LocalRequest, LocalResponse}; -/// The one delegation-id generator shared by the MCP facade and the -/// `openab agent` CLI. A single generator keeps CLI- and MCP-initiated -/// delegations indistinguishable on the wire (no `d-cli-` vs `d-` fork) while +/// The delegation-id generator for the `openab agent` CLI. The MCP facade is +/// not wired to it at this commit. Using one generator for every initiating +/// surface keeps the ids indistinguishable on the wire (no `d-cli-` vs `d-` +/// fork) while /// still being process-unique: `pid` separates concurrent processes, the /// monotonic `seq` separates ids minted inside one process, and the nanosecond /// clock separates process restarts that reuse a pid. @@ -322,8 +328,9 @@ fn new_delegation_id() -> String { format!("d-{}-{nanos:032x}-{seq:016x}", std::process::id()) } -/// Run one full spawn operation under a single `deadline_secs + 5` timeout, -/// shared verbatim by the MCP facade and the `openab agent spawn` CLI. +/// Run one full spawn operation under a single `deadline_secs + 5` timeout. +/// The `openab agent spawn` CLI is the only caller at this commit; the MCP +/// facade is not wired to it yet. /// /// "Full operation" is the crux the two surfaces previously diverged on: it is /// the admission round-trip *and*, unless `async_mode`, the blocking await of @@ -408,9 +415,9 @@ fn has_unified_platform(cfg: &config::Config) -> bool { /// - `[mcp]` alone stays exactly what it was: the facade in the foreground. /// /// `type = "primary"` alone (without `[mcp]`) deliberately does NOT unlock a -/// headless boot. A primary is the *initiating* side of a delegation; its -/// prompts come from a chat platform, so an adapter-less primary has no way to -/// be given work and is a misconfiguration worth failing on. +/// headless boot. The primary's initiating surface is the local socket and the +/// `openab agent` CLI. The headless boot serves that socket only through the +/// `[mcp]` path above, so a primary with no adapter and no `[mcp]` is refused. #[derive(Debug, PartialEq, Eq)] enum HeadlessMode { FullRuntime, @@ -469,10 +476,14 @@ fn supervise_local_api( async fn wait_for_control_plane_exit( task: &mut Option, ) -> anyhow::Error { - let Some(task) = task.as_mut() else { + let Some(supervised) = task.as_mut() else { return std::future::pending::().await; }; - match (&mut task.completion).await { + let outcome = (&mut supervised.completion).await; + // The receiver is spent once it has yielded. Drop the supervision here so + // the shutdown path skips it instead of polling it again, which panics. + *task = None; + match outcome { Ok(Ok(())) => anyhow::anyhow!("control-plane client exited unexpectedly"), Ok(Err(error)) => anyhow::anyhow!("control-plane client task failed: {error}"), Err(_) => anyhow::anyhow!("control-plane client supervisor stopped unexpectedly"), @@ -480,10 +491,13 @@ async fn wait_for_control_plane_exit( } async fn wait_for_local_api_exit(task: &mut Option) -> anyhow::Error { - let Some(task) = task.as_mut() else { + let Some(supervised) = task.as_mut() else { return std::future::pending::().await; }; - match (&mut task.completion).await { + let outcome = (&mut supervised.completion).await; + // Same as the control-plane wait: the spent receiver must not reach shutdown. + *task = None; + match outcome { Ok(Ok(Ok(()))) => anyhow::anyhow!("control-plane local API exited unexpectedly"), Ok(Ok(Err(error))) => anyhow::anyhow!("control-plane local API failed: {error:#}"), Ok(Err(error)) => anyhow::anyhow!("control-plane local API task failed: {error}"), @@ -774,7 +788,7 @@ async fn main() -> anyhow::Result<()> { } HeadlessMode::None => { anyhow::bail!( - "no adapter configured — add [discord], [slack], [telegram], [line], [lineworks], [feishu], [wecom], [googlechat], [teams], or [gateway] to config (or [mcp] for facade-only mode, or [control_plane] with type = \"worker\" for control-plane worker mode), or set the corresponding platform environment variables" + "no adapter configured — add [discord], [slack], [telegram], [line], [lineworks], [feishu], [wecom], [googlechat], [teams], or [gateway] to config (or [mcp] for facade-only mode, or [control_plane] with type = \"worker\" for control-plane worker mode; a control-plane primary with no adapter also needs [mcp]), or set the corresponding platform environment variables" ); } } @@ -2694,6 +2708,25 @@ agent_id = "1000002" assert!(error.to_string().contains("exited unexpectedly")); } + #[tokio::test] + async fn a_supervised_exit_is_not_awaited_again_at_shutdown() { + // Shutdown skips a task whose option is None. Reaching the await a + // second time would panic: tokio's oneshot receiver panics once spent. + let mut cp = Some(supervise_control_plane(tokio::spawn(async {}))); + let _ = wait_for_control_plane_exit(&mut cp).await; + assert!( + cp.is_none(), + "the spent control-plane supervision must be dropped" + ); + + let mut local = Some(supervise_local_api(tokio::spawn(async { Ok(()) }))); + let _ = wait_for_local_api_exit(&mut local).await; + assert!( + local.is_none(), + "the spent local-API supervision must be dropped" + ); + } + #[test] fn a_control_plane_worker_boots_without_any_adapter() { let cfg = config::parse_config_str(&cp_section("worker"), "test").unwrap(); From 228f714aa38c66ec64d0ab86964af4c46c4744d7 Mon Sep 17 00:00:00 2001 From: chaodu-agent <274062505+chaodu-agent@users.noreply.github.com> Date: Sat, 31 Oct 2026 11:40:53 -0400 Subject: [PATCH 10/15] test(cp): cover the runtime-side refusal of spawn on a worker (L11) Adds a_worker_runtime_refuses_spawn_locally. The worker's own client answers a Spawn with the primary-only error and sends nothing to the CP. Until now only the CP-side refusal (-32003) was tested. Checks (macmini): cargo test -p openab-core --test cp_client: 14 passed, 0 failed, including a_worker_runtime_refuses_spawn_locally and the two primary_local_api_* tests that cover primary-side spawn, list, await and cancel over the real CP (the AC4 evidence on a primary). cargo clippy -D warnings: pass on the same tree. --- crates/openab-core/tests/cp_client.rs | 51 +++++++++++++++++++++++++++ 1 file changed, 51 insertions(+) diff --git a/crates/openab-core/tests/cp_client.rs b/crates/openab-core/tests/cp_client.rs index cd6b42d95..a92c1b668 100644 --- a/crates/openab-core/tests/cp_client.rs +++ b/crates/openab-core/tests/cp_client.rs @@ -1013,3 +1013,54 @@ async fn phase0_worker_identity_cannot_initiate() { "a refused initiation leaves no in-flight entry" ); } + +/// Phase 0 criterion 5, runtime side: a worker runtime refuses `spawn` on its +/// own client, so the request never reaches the CP and no CP policy denial is +/// needed. +#[tokio::test] +async fn a_worker_runtime_refuses_spawn_locally() { + use openab_core::control_plane::{ClientCommand, CommandError, SpawnRequest, SpawnTarget}; + + let (state, url) = spawn_cp(cp_config("")).await; + let (worker_shutdown, worker_task, worker) = spawn_worker( + &url, + ScriptedRunner::new(Script::Answer), + Duration::from_secs(60), + ); + wait_for("the worker to register", || { + !worker_instances(&state).is_empty() + }) + .await; + + let (spawn_tx, spawn_rx) = tokio::sync::oneshot::channel(); + worker + .handle() + .submit(ClientCommand::Spawn { + request: SpawnRequest { + delegation_id: "d-worker-spawn".into(), + target: SpawnTarget::by_name("worker-1"), + prompt: "should be refused locally".into(), + deadline: chrono::Utc::now() + chrono::Duration::seconds(60), + parent: None, + }, + reply: spawn_tx, + }) + .await + .unwrap(); + let result = tokio::time::timeout(Duration::from_secs(5), spawn_rx) + .await + .expect("the refusal is answered locally") + .unwrap(); + assert!( + matches!(&result, Err(CommandError::Internal(message)) if message.contains("not registered as a control-plane primary")), + "a worker must refuse spawn with the primary-only error" + ); + assert_eq!( + state.router.inflight_count(), + 0, + "the local refusal sends nothing to the CP" + ); + + let _ = worker_shutdown.send(true); + let _ = worker_task.await; +} From bf26f4ab5f988846c2bd0b2f5b987a2132215c51 Mon Sep 17 00:00:00 2001 From: chaodu-agent <274062505+chaodu-agent@users.noreply.github.com> Date: Sat, 31 Oct 2026 11:40:53 -0400 Subject: [PATCH 11/15] docs(cp): name tool execution as a worker residual risk (M4) A cp/delegate from a primary enters the worker's normal agent turn with the agent's configured tool permissions. Nothing narrows tools per delegation. The config reference now says so, so the maintainer can accept it knowingly. Docs only. Code checks are the same as the preceding fix commit (clippy -D warnings pass, cp_client 14 passed, control_plane 120 passed, openab 33 passed on macmini). --- docs/config-reference.md | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/docs/config-reference.md b/docs/config-reference.md index 4ec1c3fb9..2cb1358b8 100644 --- a/docs/config-reference.md +++ b/docs/config-reference.md @@ -831,6 +831,11 @@ a primary's initiating surface is the owner-only local socket and `openab agent` ### Operational notes +- **Delegated prompts run with the worker's normal tool permissions.** A + `cp/delegate` from any primary the CP lets reach this worker enters the + worker's ordinary agent turn. The worker does not narrow tools per delegation, + so a delegated prompt can use any tool the configured `[agent]` can use. Give a + worker's key only to primaries you trust with that access. - **Local agent API.** `$OPENAB_AGENT_SOCKET` overrides the default `$HOME/.openab/run/agent.sock`. The runtime creates an owner-only directory (`0700`) and socket (`0600`). The transport is a **Unix-domain socket and is From 8354dc8ac8f927015506eb4f4097700936edc09c Mon Sep 17 00:00:00 2001 From: chaodu-agent <274062505+chaodu-agent@users.noreply.github.com> Date: Sat, 31 Oct 2026 11:52:28 -0400 Subject: [PATCH 12/15] test(cp): cover worker stop as target_disconnected (M6) Phase 0 criterion 4: stopping one worker fails its in-flight delegation as target_disconnected at the initiator, within the lease. No test covered it. The new test delegates to a hanging worker, stops that worker, and asserts the initiator's cp/delegate_result carries target_disconnected within 10s. The lease defaults to 30s and the delegation deadline is 300s, so a 10s bound rules out both as the source of the failure. Checks run on macmini at this commit: - cargo test -p openab-core --test cp_client a_stopped_worker: 1 passed, 0 failed - cargo test -p openab-core --test cp_client: 15 passed, 0 failed (after rustfmt layout fix) - rustfmt --check crates/openab-core/tests/cp_client.rs: clean (exit 0) --- crates/openab-core/tests/cp_client.rs | 43 +++++++++++++++++++++++++++ 1 file changed, 43 insertions(+) diff --git a/crates/openab-core/tests/cp_client.rs b/crates/openab-core/tests/cp_client.rs index a92c1b668..b7982a43f 100644 --- a/crates/openab-core/tests/cp_client.rs +++ b/crates/openab-core/tests/cp_client.rs @@ -630,6 +630,49 @@ async fn shutdown_drains_an_inflight_delegation_before_exit() { .await; } +/// Phase 0 criterion 4: stopping one worker fails its in-flight delegation as +/// `target_disconnected` at the initiator. The worker drops its own result on +/// the way out, so the CP's disconnect synthesis is what the initiator hears. +/// The lease defaults to 30s and the delegation deadline here is 300s; the +/// initiator must hear the failure within 10s, which rules out both. +#[tokio::test] +async fn a_stopped_worker_fails_its_inflight_delegation_as_target_disconnected() { + let (state, url) = spawn_cp(cp_config("")).await; + let runner = ScriptedRunner::new(Script::Hang); + let (shutdown, handle, _client) = + spawn_worker(&url, Arc::clone(&runner), Duration::from_secs(60)); + wait_for("the worker to register", || { + !worker_instances(&state).is_empty() + }) + .await; + + let mut initiator = connect_initiator(&url).await; + delegate( + &mut initiator, + 9, + "d-worker-stopped", + "hang until stopped", + 300, + ) + .await; + assert_eq!(next_json(&mut initiator).await["id"], 9, "delegate ack"); + wait_for("the delegated turn to start", || runner.started() == 1).await; + + let _ = shutdown.send(true); + let failure = tokio::time::timeout( + Duration::from_secs(10), + next_request(&mut initiator, "cp/delegate_result"), + ) + .await + .expect("the initiator hears the failure within 10s, well inside the lease"); + assert_eq!(failure["params"]["delegation_id"], "d-worker-stopped"); + assert_eq!( + failure["params"]["status"], "target_disconnected", + "the stopped worker fails its delegation: {failure}" + ); + let _ = tokio::time::timeout(Duration::from_secs(7), handle).await; +} + /// The PR4 local API against two real runtime clients and the real CP server: /// primary UDS -> primary client -> CP -> worker client -> scripted agent -> result. #[tokio::test] From 18955be3599269c28676b1644955e35561bc1e70 Mon Sep 17 00:00:00 2001 From: chaodu-agent <274062505+chaodu-agent@users.noreply.github.com> Date: Sat, 31 Oct 2026 11:52:28 -0400 Subject: [PATCH 13/15] fix(cp): correct adapter-less primary test comment and reflow id doc (L2, L14) L2: the comment on a_control_plane_primary_alone_does_not_unlock_a_headless_boot still gave the chat-only rationale. It now matches HeadlessMode: the primary's initiating surface is the local socket and `openab agent`, which the headless boot serves only with [mcp] or a chat adapter. The production wording already said this. L14: reflowed the new_delegation_id doc paragraph and removed the orphaned line break after "while". Comments only; no behaviour change. Checks run on macmini at this commit: - cargo clippy -- -D warnings: pass - cargo clippy --workspace -- -D warnings: pass - cargo test -p openab: 33 passed, 0 failed - cargo test -p openab-core --lib control_plane: 120 passed, 0 failed --- src/main.rs | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/src/main.rs b/src/main.rs index d51c0858e..f10c17950 100644 --- a/src/main.rs +++ b/src/main.rs @@ -314,10 +314,9 @@ use openab_core::control_plane::{LocalRequest, LocalResponse}; /// The delegation-id generator for the `openab agent` CLI. The MCP facade is /// not wired to it at this commit. Using one generator for every initiating /// surface keeps the ids indistinguishable on the wire (no `d-cli-` vs `d-` -/// fork) while -/// still being process-unique: `pid` separates concurrent processes, the -/// monotonic `seq` separates ids minted inside one process, and the nanosecond -/// clock separates process restarts that reuse a pid. +/// fork) while still being process-unique: `pid` separates concurrent +/// processes, the monotonic `seq` separates ids minted inside one process, and +/// the nanosecond clock separates process restarts that reuse a pid. fn new_delegation_id() -> String { static NEXT: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(1); let seq = NEXT.fetch_add(1, std::sync::atomic::Ordering::Relaxed); @@ -2735,8 +2734,9 @@ agent_id = "1000002" #[test] fn a_control_plane_primary_alone_does_not_unlock_a_headless_boot() { - // A primary initiates delegations; its prompts come from a chat - // platform, so an adapter-less primary has no way to be given work. + // A primary's initiating surface is the local socket and the `openab + // agent` CLI. The headless boot serves that socket only with `[mcp]` + // (or a chat adapter), so a primary with neither is refused. let cfg = config::parse_config_str(&cp_section("primary"), "test").unwrap(); assert_eq!(headless_run_mode(&cfg), HeadlessMode::None); } From bb3c631f034dbbbcd9a0d3cf7f08e3eaac206a30 Mon Sep 17 00:00:00 2001 From: chaodu-agent <274062505+chaodu-agent@users.noreply.github.com> Date: Sat, 31 Oct 2026 19:28:36 +0000 Subject: [PATCH 14/15] fix(cp): answer await/check from preserved history during backoff; bound close; warn on zero budget - F1: during reconnect backoff, Await and Check read the process-lifetime PrimaryState instead of answering NotConnected, so a delegation that completed before a blip stays readable. fail_all_live() now also runs on backoff entry, so an early Err out of connect_and_serve cannot leave a parked awaiter behind. - F2: warn when the register ack grants effective_max_delegated_sessions = 0. - F3: the end-of-session close handshake runs under one 3s CLOSE_TIMEOUT instead of two sequential 30s WRITE_TIMEOUT bounds. - F10: the Err arm applies the same stability-based backoff reset; an invalid cp/delegate_result logs its serde error; the ListAgents comment states that the CP serves it to every role. --- .../openab-core/src/control_plane/client.rs | 150 +++++++++++++++--- 1 file changed, 130 insertions(+), 20 deletions(-) diff --git a/crates/openab-core/src/control_plane/client.rs b/crates/openab-core/src/control_plane/client.rs index 1792f1d18..030e23933 100644 --- a/crates/openab-core/src/control_plane/client.rs +++ b/crates/openab-core/src/control_plane/client.rs @@ -79,6 +79,11 @@ const REGISTER_TIMEOUT: Duration = Duration::from_secs(10); /// reconnect path instead of wedging the process. const WRITE_TIMEOUT: Duration = Duration::from_secs(30); +/// Bound on the whole best-effort close handshake (Close frame plus flush) at +/// the end of a session. Short on purpose: the session is already over, and a +/// dead peer must not delay reconnect or process exit. +const CLOSE_TIMEOUT: Duration = Duration::from_secs(3); + type Ws = WebSocketStream>; /// Write half. Split from the read half because the serve loop must be able to /// write from a handler while the read future is still alive — one `select!` @@ -228,6 +233,10 @@ impl ControlPlaneClient { ); } Err(e) => { + // Same stability clock as a clean close: one transient + // failure after a long healthy session must not inherit an + // escalated backoff from an earlier flap. + backoff = backoff_after_session(backoff, session_started.elapsed()); error!(error = %format!("{e:#}"), backoff_secs = backoff, "control-plane connection failed"); } } @@ -237,14 +246,23 @@ impl ControlPlaneClient { retry_delay_ms = retry_delay.as_millis(), "control-plane reconnect delay selected" ); - // While disconnected, reject any submitted command promptly rather + // While disconnected, answer any submitted command promptly rather // than letting it sit in the buffer until the next session. The // channel stays open (the handle is process-lifetime), so this only // answers what is queued; new submissions during the sleep are - // answered too. A closed channel is replaced with `None` so - // `recv_command` becomes a pending future and the sleep proceeds. + // answered too. Commands that need the socket are rejected + // `NotConnected`; Await/Check read the preserved terminal history in + // `primary`, so a delegation that completed before the blip is + // still readable during backoff. A closed channel is replaced with + // `None` so `recv_command` becomes a pending future and the sleep + // proceeds. let sleep = tokio::time::sleep(retry_delay); tokio::pin!(sleep); + // Idempotent: the serve-loop teardown normally did this already, + // but an early `Err` out of `connect_and_serve` skips it. After + // this call `primary` holds terminal history only, so Await below + // answers immediately and never parks a reply no socket can serve. + primary.fail_all_live(); loop { tokio::select! { _ = &mut sleep => break, @@ -254,7 +272,7 @@ impl ControlPlaneClient { } maybe = recv_command(&mut command_rx) => { match maybe { - Some(command) => reject_command(command), + Some(command) => answer_while_disconnected(command, &mut primary), None => { command_rx = None; } } } @@ -286,6 +304,9 @@ impl ControlPlaneClient { }; self.executor .set_effective_max(ack.effective_max_delegated_sessions); + if let Some(warning) = register_ack_budget_warning(&ack) { + warn!(instance = %self.instance_id, "{warning}"); + } info!( instance = %self.instance_id, heartbeat_secs = ack.heartbeat_interval_secs, @@ -581,10 +602,14 @@ impl ControlPlaneClient { handle.abort(); } } - // Best-effort close handshake, bounded: a peer that stopped reading - // must not hold up reconnect or process shutdown. - let _ = write_bounded(&mut sink, Message::Close(None)).await; - let _ = tokio::time::timeout(WRITE_TIMEOUT, sink.close()).await; + // Best-effort close handshake under ONE short bound: a peer that + // stopped reading must not hold up reconnect or process shutdown. Using + // WRITE_TIMEOUT for each of the two steps could cost ~60s here. + let _ = tokio::time::timeout(CLOSE_TIMEOUT, async { + let _ = sink.send(Message::Close(None)).await; + let _ = sink.close().await; + }) + .await; Ok(outcome) } @@ -603,9 +628,9 @@ impl ControlPlaneClient { // A primary must be registered as such to initiate. A worker-only runtime // rejects spawn and cancel locally, so a misconfigured caller gets a clear // answer instead of a CP policy denial round-trip. list_agents is forwarded - // for either role; whether the CP accepts a worker's list_agents is not - // verified here. Await is allowed for either role (it only reads local - // tracking). + // for either role: the CP serves it to every registered type (pinned by + // `a_worker_runtime_refuses_spawn_locally`). Await is allowed for either + // role (it only reads local tracking). match command { ClientCommand::Spawn { request, reply } => { if self.cfg.agent_type != CpAgentType::Primary { @@ -731,7 +756,13 @@ impl ControlPlaneClient { } }; let params: Option = - msg.params.and_then(|p| serde_json::from_value(p).ok()); + msg.params.and_then(|p| match serde_json::from_value(p) { + Ok(params) => Some(params), + Err(e) => { + warn!(error = %e, rpc_id = id, "invalid cp/delegate_result params"); + None + } + }); match params { Some(params) => { // Ack regardless of whether we track it: an untracked @@ -1036,7 +1067,15 @@ async fn recv_command( } /// Answer a command that arrived while the client is not connected. -fn reject_command(command: ClientCommand) { +/// +/// Spawn, Cancel and ListAgents need the socket and are rejected +/// `NotConnected` (retryable). Await and Check only read local tracking, and +/// that tracking is process-lifetime: `fail_all_live` runs before backoff, so +/// every live delegation has already failed and what remains is terminal history. +/// Answering from it lets a caller read a result that landed before the blip +/// instead of getting a spurious `NotConnected`. An unknown or already-failed +/// handle gets `UnknownDelegation`, exactly as while connected. +fn answer_while_disconnected(command: ClientCommand, primary: &mut PrimaryState) { use crate::control_plane::primary::CommandError; match command { ClientCommand::Spawn { reply, .. } => { @@ -1048,14 +1087,13 @@ fn reject_command(command: ClientCommand) { ClientCommand::ListAgents { reply } => { let _ = reply.send(Err(CommandError::NotConnected)); } - ClientCommand::Await { reply, .. } => { - // Awaiting only reads local tracking, but a disconnected client - // holds none: there is nothing to wait on. - let _ = reply.send(Err(CommandError::NotConnected)); + ClientCommand::Await { handle, reply } => { + // Only terminal entries survive teardown, so this answers + // immediately; it never parks an awaiter that no socket can serve. + primary.begin_await(&handle, reply); } - ClientCommand::Check { reply, .. } => { - // Same rationale as Await: no tracking survives a disconnect. - let _ = reply.send(Err(CommandError::NotConnected)); + ClientCommand::Check { handle, reply } => { + primary.check(&handle, reply); } } } @@ -1080,6 +1118,18 @@ fn validate_register_ack(result: serde_json::Value) -> anyhow::Result Option<&'static str> { + (ack.effective_max_delegated_sessions == 0).then_some( + "control plane granted an effective_max_delegated_sessions of 0; \ + this runtime will refuse every delegation until it re-registers", + ) +} + fn jittered_backoff(base_secs: u64) -> Duration { let max_ms = base_secs.saturating_mul(1000).max(1); Duration::from_millis(rand::thread_rng().gen_range(max_ms / 2..=max_ms)) @@ -1261,6 +1311,66 @@ max_delegated_sessions = 3 .contains("requires")); } + #[test] + fn a_zero_budget_register_ack_is_accepted_but_warned() { + let ack = |max: u32| -> RegisterAck { + validate_register_ack(serde_json::json!({ + "protocol_version": PROTOCOL_VERSION, + "heartbeat_interval_secs": 15, + "lease_expiry_secs": 45, + "effective_max_delegated_sessions": max + })) + .expect("a zero budget is a legitimate drain, not a protocol error") + }; + let warning = register_ack_budget_warning(&ack(0)).expect("zero budget must not be silent"); + assert!(warning.contains("effective_max_delegated_sessions of 0")); + assert!(register_ack_budget_warning(&ack(1)).is_none()); + } + + #[test] + fn while_disconnected_await_and_check_read_local_tracking() { + use crate::control_plane::primary::{CommandError, DelegationHandle}; + let mut primary = PrimaryState::new(); + let handle = DelegationHandle::new("d-gone", 1); + + // Await/Check consult `primary`: an untracked handle is + // UnknownDelegation, not the socket-level NotConnected. + let (atx, mut arx) = tokio::sync::oneshot::channel(); + answer_while_disconnected( + ClientCommand::Await { + handle: handle.clone(), + reply: atx, + }, + &mut primary, + ); + assert_eq!(arx.try_recv().unwrap(), Err(CommandError::UnknownDelegation)); + let (ctx, mut crx) = tokio::sync::oneshot::channel(); + answer_while_disconnected( + ClientCommand::Check { + handle: handle.clone(), + reply: ctx, + }, + &mut primary, + ); + assert_eq!(crx.try_recv().unwrap(), Err(CommandError::UnknownDelegation)); + assert_eq!(primary.awaiters_len(), 0, "nothing is parked while disconnected"); + + // Commands that need the socket stay NotConnected (retryable). + let (ltx, mut lrx) = tokio::sync::oneshot::channel(); + answer_while_disconnected(ClientCommand::ListAgents { reply: ltx }, &mut primary); + assert!(matches!(lrx.try_recv().unwrap(), Err(CommandError::NotConnected))); + let (xtx, mut xrx) = tokio::sync::oneshot::channel(); + answer_while_disconnected( + ClientCommand::Cancel { + handle, + reason: "r".into(), + reply: xtx, + }, + &mut primary, + ); + assert_eq!(xrx.try_recv().unwrap(), Err(CommandError::NotConnected)); + } + #[test] fn reconnect_backoff_is_jittered_bounded_and_resets_only_after_stability() { for base in [1, 2, 4, 8, 16, 30] { From 1bd349029e227093db2bd5adcacf12e6ef68e886 Mon Sep 17 00:00:00 2001 From: chaodu-agent <274062505+chaodu-agent@users.noreply.github.com> Date: Sat, 31 Oct 2026 19:28:36 +0000 Subject: [PATCH 15/15] test(cp): register rejection backoff, malformed/unknown frames, await across reconnect - a_rejected_registration_backs_off_instead_of_spinning - unknown_methods_and_malformed_frames_do_not_end_the_session - a_completed_delegation_is_readable_during_reconnect_backoff - a_worker_runtime_refuses_spawn_locally now also pins that the CP serves cp/list_agents to a worker --- crates/openab-core/tests/cp_client.rs | 325 ++++++++++++++++++++++++++ 1 file changed, 325 insertions(+) diff --git a/crates/openab-core/tests/cp_client.rs b/crates/openab-core/tests/cp_client.rs index b7982a43f..6be01f97f 100644 --- a/crates/openab-core/tests/cp_client.rs +++ b/crates/openab-core/tests/cp_client.rs @@ -1104,6 +1104,331 @@ async fn a_worker_runtime_refuses_spawn_locally() { "the local refusal sends nothing to the CP" ); + // list_agents is NOT refused locally: it is forwarded for either role, and + // the CP serves it to a worker. Pin that so a policy change on either side + // shows up here. + let (list_tx, list_rx) = tokio::sync::oneshot::channel(); + worker + .handle() + .submit(ClientCommand::ListAgents { reply: list_tx }) + .await + .unwrap(); + let agents = tokio::time::timeout(Duration::from_secs(5), list_rx) + .await + .expect("the roster is answered") + .unwrap() + .expect("the CP serves cp/list_agents to a worker"); + assert!( + agents.iter().any(|a| a.name == "worker-1"), + "the roster includes the calling worker" + ); + let _ = worker_shutdown.send(true); let _ = worker_task.await; } + +// --------------------------------------------------------------------------- +// Failure paths +// --------------------------------------------------------------------------- + +/// A scripted stand-in for the CP, for failure modes the real server never +/// produces on demand. Counts accepted TCP connections; `script` drives each +/// upgraded socket. +async fn spawn_fake_cp(script: F) -> (String, Arc) +where + F: Fn(WebSocketStream) -> Fut + Send + Sync + 'static, + Fut: std::future::Future + Send + 'static, +{ + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let url = format!("ws://{}/cp", listener.local_addr().unwrap()); + let attempts = Arc::new(AtomicUsize::new(0)); + let counter = Arc::clone(&attempts); + let script = Arc::new(script); + tokio::spawn(async move { + while let Ok((tcp, _)) = listener.accept().await { + counter.fetch_add(1, Ordering::Relaxed); + let script = Arc::clone(&script); + tokio::spawn(async move { + if let Ok(ws) = tokio_tungstenite::accept_async(tcp).await { + (*script)(ws).await; + } + }); + } + }); + (url, attempts) +} + +/// Read the client's `cp/register` frame and return its JSON-RPC id. +async fn read_register(ws: &mut WebSocketStream) -> Option { + while let Some(Ok(msg)) = ws.next().await { + if let Message::Text(text) = msg { + let req: serde_json::Value = serde_json::from_str(&text).ok()?; + if req["method"] == "cp/register" { + return Some(req["id"].clone()); + } + } + } + None +} + +/// A CP that rejects `cp/register` (unknown key, wrong identity) must put the +/// client into its jittered exponential backoff, not a reconnect spin. With a +/// 1s base the attempts land at ~0s, 0.5-1s and 1.5-3s, and the fourth no +/// earlier than 3.5s, so a 3.2s window sees two or three attempts. +#[tokio::test] +async fn a_rejected_registration_backs_off_instead_of_spinning() { + let (url, attempts) = spawn_fake_cp(|mut ws| async move { + if let Some(id) = read_register(&mut ws).await { + let reply = serde_json::json!({ + "jsonrpc": "2.0", "id": id, + "error": {"code": -32001, "message": "unknown agent key"} + }); + let _ = ws.send(Message::Text(reply.to_string())).await; + } + // Hold the socket until the client drops it. + while let Some(Ok(_)) = ws.next().await {} + }) + .await; + + let (shutdown, handle, _client) = spawn_worker( + &url, + ScriptedRunner::new(Script::Answer), + Duration::from_secs(60), + ); + tokio::time::sleep(Duration::from_millis(3200)).await; + let seen = attempts.load(Ordering::Relaxed); + assert!( + (2..=3).contains(&seen), + "a rejected registration must retry with backoff (saw {seen} attempts in 3.2s)" + ); + + let _ = shutdown.send(true); + tokio::time::timeout(Duration::from_secs(5), handle) + .await + .expect("shutdown is honoured during backoff") + .unwrap(); +} + +/// An unknown method gets METHOD_NOT_FOUND, and a frame that is not JSON is +/// ignored. Neither may end the session: both later requests are answered on +/// the SAME connection, and the client never reconnects. +#[tokio::test] +async fn unknown_methods_and_malformed_frames_do_not_end_the_session() { + let (replies_tx, mut replies_rx) = tokio::sync::mpsc::unbounded_channel::(); + let replies_tx = Arc::new(replies_tx); + let (url, attempts) = spawn_fake_cp(move |mut ws| { + let replies_tx = Arc::clone(&replies_tx); + async move { + let Some(id) = read_register(&mut ws).await else { + return; + }; + let ack = serde_json::json!({ + "jsonrpc": "2.0", "id": id, + "result": { + "protocol_version": 1, + "heartbeat_interval_secs": 15, + "lease_expiry_secs": 45, + "effective_max_delegated_sessions": 2 + } + }); + ws.send(Message::Text(ack.to_string())).await.unwrap(); + for frame in [ + "{this is not json".to_string(), + serde_json::json!({"jsonrpc": "2.0", "id": 901, "method": "cp/bogus"}).to_string(), + "\u{0}\u{1}garbage".to_string(), + serde_json::json!({"jsonrpc": "2.0", "id": 902, "method": "cp/also_bogus"}) + .to_string(), + ] { + ws.send(Message::Text(frame)).await.unwrap(); + } + while let Some(Ok(msg)) = ws.next().await { + if let Message::Text(text) = msg { + if let Ok(v) = serde_json::from_str::(&text) { + let _ = replies_tx.send(v); + } + } + } + } + }) + .await; + + let (shutdown, handle, _client) = spawn_worker( + &url, + ScriptedRunner::new(Script::Answer), + Duration::from_secs(60), + ); + + let mut answered = std::collections::BTreeMap::new(); + let deadline = Instant::now() + Duration::from_secs(10); + while answered.len() < 2 && Instant::now() < deadline { + let Ok(Some(v)) = tokio::time::timeout(Duration::from_secs(10), replies_rx.recv()).await + else { + break; + }; + if let Some(id) = v["id"].as_u64().filter(|id| *id == 901 || *id == 902) { + answered.insert(id, v); + } + } + for id in [901, 902] { + let reply = answered + .get(&id) + .unwrap_or_else(|| panic!("request {id} was never answered")); + assert_eq!(reply["error"]["code"], -32601, "METHOD_NOT_FOUND: {reply}"); + } + // Give a wrongly-ended session time to reconnect before checking. + tokio::time::sleep(Duration::from_millis(300)).await; + assert_eq!( + attempts.load(Ordering::Relaxed), + 1, + "the session survived; the client never reconnected" + ); + + let _ = shutdown.send(true); + let _ = tokio::time::timeout(Duration::from_secs(5), handle).await; +} + +/// Terminal history is process-lifetime: a delegation that completed before a +/// connection loss can still be awaited and checked while the client is in +/// reconnect backoff. Socket-bound commands are still rejected NotConnected. +#[tokio::test] +async fn a_completed_delegation_is_readable_during_reconnect_backoff() { + use openab_core::control_plane::{ + ClientCommand, CommandError, DelegationSnapshot, DelegationStatusWire, SpawnRequest, + SpawnTarget, + }; + + // The real CP, with its serve task kept so it can be stopped: once the + // listener is gone every reconnect is refused, which pins the client in + // its backoff loop. + let state = Arc::new(AppState::new(cp_config(""))); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let url = format!("ws://{}/cp", listener.local_addr().unwrap()); + let router = app(state.clone()); + let server = tokio::spawn(async move { + let _ = axum::serve(listener, router).await; + }); + + let (worker_shutdown, worker_task, _worker) = spawn_worker( + &url, + ScriptedRunner::new(Script::Answer), + Duration::from_secs(60), + ); + let mut cfg = worker_cfg(&url); + cfg.auth_key = PRIMARY_KEY.to_string(); + cfg.name = "koudu".into(); + cfg.agent_type = CpAgentType::Primary; + let (primary_shutdown, primary_rx) = tokio::sync::watch::channel(false); + let primary = Arc::new(ControlPlaneClient::new( + cfg, + ScriptedRunner::new(Script::Answer), + Duration::from_secs(60), + )); + let cp = primary.handle(); + let primary_task = tokio::spawn(primary.run(primary_rx)); + wait_for("both runtimes to register", || { + state.registry.list("prod").len() >= 2 + }) + .await; + + // Spawn and await while connected. + let (spawn_tx, spawn_rx) = tokio::sync::oneshot::channel(); + cp.submit(ClientCommand::Spawn { + request: SpawnRequest { + delegation_id: "d-across-blip".into(), + target: SpawnTarget::by_name("worker-1"), + prompt: "finish before the blip".into(), + deadline: chrono::Utc::now() + chrono::Duration::seconds(60), + parent: None, + }, + reply: spawn_tx, + }) + .await + .unwrap(); + let ack = tokio::time::timeout(Duration::from_secs(10), spawn_rx) + .await + .expect("spawn is acked") + .unwrap() + .expect("spawn succeeds"); + let (await_tx, await_rx) = tokio::sync::oneshot::channel(); + cp.submit(ClientCommand::Await { + handle: ack.handle.clone(), + reply: await_tx, + }) + .await + .unwrap(); + let before = tokio::time::timeout(Duration::from_secs(10), await_rx) + .await + .expect("terminal arrives") + .unwrap() + .expect("the delegation completed"); + assert_eq!( + DelegationStatusWire::from(before.status.clone()), + DelegationStatusWire::Completed + ); + + // Lose the connection and keep it lost: stop accepting, then drop every + // registration so the CP closes the live sockets. + server.abort(); + let _ = server.await; + sweep_leases(&state, Duration::ZERO); + assert!(state.registry.list("prod").is_empty()); + // Let the primary notice the close and enter backoff. + tokio::time::sleep(Duration::from_millis(500)).await; + + let (await_tx, await_rx) = tokio::sync::oneshot::channel(); + cp.submit(ClientCommand::Await { + handle: ack.handle.clone(), + reply: await_tx, + }) + .await + .unwrap(); + let during = tokio::time::timeout(Duration::from_secs(5), await_rx) + .await + .expect("await is answered during backoff") + .unwrap(); + assert_eq!( + during, + Ok(before.clone()), + "the preserved terminal is readable while disconnected" + ); + + let (check_tx, check_rx) = tokio::sync::oneshot::channel(); + cp.submit(ClientCommand::Check { + handle: ack.handle.clone(), + reply: check_tx, + }) + .await + .unwrap(); + let snapshot = tokio::time::timeout(Duration::from_secs(5), check_rx) + .await + .expect("check is answered during backoff") + .unwrap(); + assert_eq!(snapshot, Ok(DelegationSnapshot::Terminal(before))); + + let (spawn_tx, spawn_rx) = tokio::sync::oneshot::channel(); + cp.submit(ClientCommand::Spawn { + request: SpawnRequest { + delegation_id: "d-while-down".into(), + target: SpawnTarget::by_name("worker-1"), + prompt: "no socket".into(), + deadline: chrono::Utc::now() + chrono::Duration::seconds(60), + parent: None, + }, + reply: spawn_tx, + }) + .await + .unwrap(); + let refused = tokio::time::timeout(Duration::from_secs(5), spawn_rx) + .await + .expect("spawn is answered during backoff") + .unwrap(); + assert!( + matches!(refused, Err(CommandError::NotConnected)), + "a socket-bound command is still NotConnected" + ); + + let _ = primary_shutdown.send(true); + let _ = worker_shutdown.send(true); + let _ = tokio::time::timeout(Duration::from_secs(5), primary_task).await; + let _ = tokio::time::timeout(Duration::from_secs(5), worker_task).await; +}