diff --git a/Cargo.lock b/Cargo.lock index fc11ec7e5..4709cad63 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2533,6 +2533,7 @@ dependencies = [ "aws-sdk-s3", "aws-sdk-secretsmanager", "aws-sigv4", + "axum", "base64", "bytes", "chrono", @@ -2546,6 +2547,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 1ede398f1..94b2503bb 100644 --- a/crates/openab-core/Cargo.toml +++ b/crates/openab-core/Cargo.toml @@ -54,12 +54,23 @@ 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" [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/src/acp/pool.rs b/crates/openab-core/src/acp/pool.rs index 86b2ee989..b92dc53a8 100644 --- a/crates/openab-core/src/acp/pool.rs +++ b/crates/openab-core/src/acp/pool.rs @@ -856,6 +856,35 @@ impl SessionPool { } } + /// 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_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 fa7e95dba..df8eacf2f 100644 --- a/crates/openab-core/src/adapter.rs +++ b/crates/openab-core/src/adapter.rs @@ -447,6 +447,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 { @@ -682,11 +702,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, @@ -697,7 +726,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()); @@ -1116,6 +1145,14 @@ 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 silent_failure = + final_content.is_empty() && turn_result.is_silent_failure(); let final_content = if final_content.is_empty() { if turn_result.is_silent_failure() { warn!( @@ -1352,7 +1389,11 @@ impl AdapterRouter { "streaming finalization had delivery failures; user view is incomplete" )) } else { - Ok(()) + Ok(PromptExecution { + final_text: final_content, + terminal_error, + silent_failure, + }) } }) }) diff --git a/crates/openab-core/src/config.rs b/crates/openab-core/src/config.rs index a9bc26abd..7b5a37d0f 100644 --- a/crates/openab-core/src/config.rs +++ b/crates/openab-core/src/config.rs @@ -1,7 +1,7 @@ use crate::markdown::TableMode; use regex::Regex; use serde::Deserialize; -use std::collections::HashMap; +use std::collections::{BTreeMap, HashMap}; use std::path::Path; /// Controls how incoming messages are dispatched to ACP turns. @@ -132,6 +132,75 @@ 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(Debug, 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, + /// 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 +} + #[derive(Debug, Clone, Deserialize)] pub struct AgentCoreConfig { /// AgentCore Runtime ARN (required) pub runtime_arn: String, @@ -248,6 +317,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)] @@ -2327,6 +2400,30 @@ 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" + ); + } + 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) } @@ -2541,6 +2638,146 @@ 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 = \"ws://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 = \"ws://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 = \"ws://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 = \"ws://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")); + } 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..b3f67409d --- /dev/null +++ b/crates/openab-core/src/control_plane/client.rs @@ -0,0 +1,773 @@ +//! 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, PROTOCOL_VERSION, +}; +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; +use crate::control_plane::executor::{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); + +/// Inbound WS message/frame ceiling, mirroring the CP server's own +/// `max_frame_bytes` default (1 MiB). Outbound frames are already capped at +/// the executor; this closes the other direction. +const MAX_INBOUND_FRAME_BYTES: usize = 1024 * 1024; + +/// 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 { + 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; + } + let mut shutdown_signal = shutdown.clone(); + 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(); + let served = tokio::select! { + r = self.connect_and_serve(&mut shutdown) => r, + // connect()/register() are not themselves shutdown-aware; this + // select is what keeps a hung dial from stalling shutdown. + _ = shutdown_signal.changed() => { + info!("control-plane client shutting down"); + return; + } + }; + 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. + if session_started.elapsed() >= Duration::from_secs(STABLE_SESSION_SECS) { + backoff = 1; + } + warn!( + backoff_secs = backoff, + "control-plane connection closed — reconnecting" + ); + } + Err(e) => { + error!(error = %format!("{e:#}"), backoff_secs = backoff, "control-plane connection failed"); + } + } + tokio::select! { + _ = tokio::time::sleep(Duration::from_secs(backoff)) => {} + _ = shutdown.changed() => { + info!("control-plane client shutting down"); + return; + } + } + backoff = (backoff * 2).min(MAX_BACKOFF_SECS); + } + } + + async fn connect_and_serve( + &self, + shutdown: &mut watch::Receiver, + ) -> anyhow::Result { + let ws = self.connect().await?; + let (mut sink, mut stream) = ws.split(); + let ack = tokio::time::timeout(REGISTER_TIMEOUT, self.register(&mut sink, &mut stream)) + .await + .map_err(|_| { + anyhow::anyhow!( + "control plane did not ack registration within {}s", + REGISTER_TIMEOUT.as_secs() + ) + })??; + 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 + } + + /// 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_INBOUND_FRAME_BYTES), + max_frame_size: Some(MAX_INBOUND_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)?)); + sink.send(Message::Text(serde_json::to_string(&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 Ok(serde_json::from_value(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 = JsonRpcRequest::new( + self.next_id(), + methods::DELEGATE_RESULT, + Some(serde_json::to_value(&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 => { + 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 = %params.reason, + 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, +} + +async fn send(sink: &mut WsSink, frame: &JsonRpcRequest) -> anyhow::Result<()> { + let text = serde_json::to_string(frame)?; + sink.send(Message::Text(text)).await?; + Ok(()) +} + +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 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] + 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..c49868b41 --- /dev/null +++ b/crates/openab-core/src/control_plane/executor.rs @@ -0,0 +1,1040 @@ +//! 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)`, so no delegation can observe another's +//! conversation, and a replayed id after a reconnect cannot resume a stale +//! one. 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}; +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) + } + + /// Number of admitted, not-yet-finished delegations. + pub fn active(&self) -> u32 { + self.inflight.lock().expect("inflight mutex").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.inflight.lock().expect("inflight mutex"); + 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.inflight + .lock() + .expect("inflight mutex") + .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.inflight.lock().expect("inflight mutex"); + 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 + .inflight + .lock() + .expect("inflight mutex") + .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)) => { + self.bounded_discard(&session_key).await; + if let Some(error) = outcome.error { + warn!(delegation_id = %id, %error, "delegation ended in an agent error"); + return failed(id, forward.admission, error); + } + if outcome.silent_failure { + 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", + ); + } + 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 are bounded so + /// a wedged teardown cannot burn the client's disconnect drain window. + async fn cancel_and_discard(&self, session_key: &str) { + let _ = tokio::time::timeout(TEARDOWN_BOUND, self.runner.cancel(session_key)).await; + self.bounded_discard(session_key).await; + } + + /// 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"); + } + } +} + +/// 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(5); + +/// 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; + +fn cap_text(text: String, budget: usize) -> String { + if text.len() <= budget { + return text; + } + let marker = format!( + "\n…[truncated by worker: {} bytes total exceeded the transport budget]", + text.len() + ); + let keep = budget.saturating_sub(marker.len()); + let mut cut = keep.min(text.len()); + while cut > 0 && !text.is_char_boundary(cut) { + cut -= 1; + } + 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_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, + 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() + } + 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) { + 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()); + } + } + + 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.discarded(), + 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); + } + + #[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 { + 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"); + } +} 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/dispatch.rs b/crates/openab-core/src/dispatch.rs index 64ba68917..3e848e8ad 100644 --- a/crates/openab-core/src/dispatch.rs +++ b/crates/openab-core/src/dispatch.rs @@ -194,6 +194,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-core/tests/cp_client.rs b/crates/openab-core/tests/cp_client.rs new file mode 100644 index 000000000..c3122280f --- /dev/null +++ b/crates/openab-core/tests/cp_client.rs @@ -0,0 +1,570 @@ +//! 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(), + 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; +} diff --git a/crates/openab-cp/Cargo.toml b/crates/openab-cp/Cargo.toml index 86ca5e30c..8597fda08 100644 --- a/crates/openab-cp/Cargo.toml +++ b/crates/openab-cp/Cargo.toml @@ -6,19 +6,52 @@ license = "MIT" description = "OpenAB Agent Control Plane — registry, router, and policy for direct inter-agent delegation" [dependencies] -tokio = { version = "1", features = ["full"] } -axum = { version = "0.8", features = ["ws"] } -futures-util = "0.3" +# --- wire types (always compiled) --- +# `proto` is the shared contract between the CP and every runtime that speaks +# to it, so it must be usable WITHOUT the server: `openab-core` depends on this +# crate with `default-features = false` purely for these types, and must not +# inherit axum/tokio/hyper through that edge (the runtime image would otherwise +# grow a whole HTTP server it never binds). serde = { version = "1", features = ["derive"] } serde_json = "1" -toml = "0.8" -tracing = "0.1" -tracing-subscriber = { version = "0.3", features = ["env-filter"] } -anyhow = "1" chrono = { version = "0.4", features = ["serde"] } -parking_lot = "0.12" -clap = { version = "4", features = ["derive"] } -subtle = "2" + +# --- server-only (feature `server`, on by default) --- +tokio = { version = "1", features = ["full"], optional = true } +axum = { version = "0.8", features = ["ws"], optional = true } +futures-util = { version = "0.3", optional = true } +toml = { version = "0.8", optional = true } +tracing = { version = "0.1", optional = true } +tracing-subscriber = { version = "0.3", features = ["env-filter"], optional = true } +anyhow = { version = "1", optional = true } +parking_lot = { version = "0.12", optional = true } +clap = { version = "4", features = ["derive"], optional = true } +subtle = { version = "2", optional = true } + +[features] +# Default: the full control plane (config, registry, policy, router, events, +# WebSocket server) plus the `openab-cp` binary. Turning it off leaves only +# `proto` — the wire contract — which is what runtime clients need. +default = ["server"] +server = [ + "dep:tokio", + "dep:axum", + "dep:futures-util", + "dep:toml", + "dep:tracing", + "dep:tracing-subscriber", + "dep:anyhow", + "dep:parking_lot", + "dep:clap", + "dep:subtle", +] + +# The binary IS the server; without the feature there is nothing for it to +# start, so it is skipped instead of failing to compile. +[[bin]] +name = "openab-cp" +path = "src/main.rs" +required-features = ["server"] [dev-dependencies] # WebSocket client for the end-to-end admission/lifecycle tests. Same version diff --git a/crates/openab-cp/src/lib.rs b/crates/openab-cp/src/lib.rs index 2ffefece0..48baade90 100644 --- a/crates/openab-cp/src/lib.rs +++ b/crates/openab-cp/src/lib.rs @@ -5,11 +5,30 @@ //! register over WebSocket; the CP authenticates them against config-bound //! identities, enforces delegation policy authoritatively, and routes //! `cp/delegate` / `cp/delegate_result` frames between them. +//! +//! ## Crate layout: `proto` vs the `server` feature +//! +//! [`proto`] is the wire contract and is ALWAYS compiled. Everything else — +//! config, registry, policy, router, observer fan-out, and the WebSocket +//! server — sits behind the default `server` feature, together with the axum +//! and tokio dependencies it needs. +//! +//! That split exists so the OAB runtime (`openab-core`) can depend on this +//! crate with `default-features = false` to speak the protocol without +//! linking an HTTP server it never binds. Anything a runtime client needs +//! belongs in `proto`; anything that only the hub does must stay behind the +//! feature. +#[cfg(feature = "server")] pub mod config; +#[cfg(feature = "server")] pub mod events; +#[cfg(feature = "server")] pub mod policy; pub mod proto; +#[cfg(feature = "server")] pub mod registry; +#[cfg(feature = "server")] pub mod router; +#[cfg(feature = "server")] pub mod server; diff --git a/docs/adr/agent-control-plane.md b/docs/adr/agent-control-plane.md index 6efbbf146..6ee55977b 100644 --- a/docs/adr/agent-control-plane.md +++ b/docs/adr/agent-control-plane.md @@ -125,7 +125,7 @@ The config section and subsystem are named `control_plane`, not ```toml [control_plane] -url = "wss://cp.example.internal/acp" +url = "wss://cp.example.internal/cp" auth_key = "${OPENAB_CP_KEY}" # per-agent credential, never shared namespace = "prod" name = "koudu" @@ -682,7 +682,7 @@ Two distinct auth boundaries exist, and they must not be conflated: private network) in front — bearer keys must never cross untrusted cleartext TCP. See the "v1 contract amendments" in §4 for the enforced registration semantics. -2. **Agent subprocess ↔ local facade (PR 3/4, not yet shipped):** the UDS +2. **Agent subprocess ↔ local facade (PR 4/4, not yet shipped):** the UDS path is the only thing the child needs; filesystem permissions on the socket are the local auth boundary. The *local facade* is never exposed on TCP — this claim is about the UDS facade, not about the CP itself, diff --git a/docs/config-reference.md b/docs/config-reference.md index af30e0940..ff2a5426c 100644 --- a/docs/config-reference.md +++ b/docs/config-reference.md @@ -765,6 +765,76 @@ 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 +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 (`ws://` or `wss://`), path `/cp` | +| `auth_key` | ✅ | — | Bearer key sent on the upgrade request. Use `${ENV}` or `[secrets.refs]` — 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 config with `[agent]` and `[control_plane] type = "worker"` and **no** chat +adapter is a valid deployment: the runtime's work arrives as delegations rather +than chat messages. + +| Config | Result | +|--------|--------| +| `[control_plane] type = "worker"`, no adapter | Worker mode — pool + control-plane client, no chat platform | +| `[control_plane] type = "primary"`, no adapter | Startup error — a primary's prompts come from a chat platform | +| `[mcp]` only, no adapter | Facade-only mode (unchanged) | +| `[mcp]` + worker, no adapter | Both — the facade listener and the control-plane client | + +### Operational notes + +- **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 1/2/4/8/16/30s backoff. One instance id is + generated per process and reused across reconnects, so the CP can tell a + reconnecting replica from a new one. +- **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`. +- **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]`. diff --git a/src/main.rs b/src/main.rs index a2ee786ac..bdffcd0e3 100644 --- a/src/main.rs +++ b/src/main.rs @@ -195,6 +195,49 @@ 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: +/// +/// - `[control_plane] type = "worker"` wins, because the runtime it needs is a +/// superset of the facade's: the normal boot path builds the session pool and +/// router, starts the CP client, and — when `[mcp]` is also present — spawns +/// the facade alongside, so `mcp + CP` runs both rather than only one. +/// - `[mcp]` alone stays exactly what it was: the facade in the foreground. +/// +/// `type = "primary"` deliberately does NOT unlock a headless boot. A primary is +/// the *initiating* side of a delegation; its prompts come from a chat platform, +/// so a primary with no adapter has no way to be given work and is a +/// misconfiguration worth failing on. +#[derive(Debug, PartialEq, Eq)] +enum HeadlessMode { + ControlPlaneWorker, + 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::ControlPlaneWorker + } else if cfg.mcp.is_some() { + // [mcp] + [control_plane type=primary]: full boot, not facade-only — + // the facade serves in the background AND the CP client registers + // (visible in the roster, ready for primary-side initiation in the + // next slice). Facade-only forecloses the client entirely. + if cfg.control_plane.is_some() { + HeadlessMode::ControlPlaneWorker + } else { + HeadlessMode::FacadeOnly + } + } else { + HeadlessMode::None + } +} + /// 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 +464,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::ControlPlaneWorker => { + 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], [wecom], [googlechat], or [gateway] to config (or [mcp] for facade-only mode, or [control_plane] with type = \"worker\" for control-plane worker mode), or set platform env vars (TELEGRAM_BOT_TOKEN, etc.)" + ); + } } - 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 +629,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 +953,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 cp_handle = 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), + )); + tokio::spawn(client.run(shutdown_rx.clone())) + }); + let dispatchers: Arc>>> = Arc::new(Mutex::new(Vec::new())); // Spawn cleanup task @@ -1795,6 +1890,24 @@ 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(handle) = cp_handle { + let abort = handle.abort_handle(); + if tokio::time::timeout(std::time::Duration::from_secs(10), handle) + .await + .is_err() + { + // Do not let a wedged dial/serve loop outlive the pool teardown + // below — a detached task writing into dead sessions is worse + // than an aborted socket (the CP synthesizes target_disconnected). + tracing::warn!("control-plane client missed the shutdown deadline — aborting"); + abort.abort(); + } + } let shutdown_pool = pool; shutdown_pool.shutdown().await; if let Some(ref hook) = shutdown_hook { @@ -2084,4 +2197,74 @@ 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 = \"ws://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); + } + + #[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::ControlPlaneWorker); + } + + #[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::ControlPlaneWorker); + } + + #[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::ControlPlaneWorker); + } + + /// 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()); + } }