From 910b0c19415a2fc17ca2433de5880e64f654acee Mon Sep 17 00:00:00 2001 From: chaodu-agent Date: Wed, 12 Aug 2026 09:14:34 -0400 Subject: [PATCH 1/5] =?UTF-8?q?feat(cp):=20runtime=20control-plane=20clien?= =?UTF-8?q?t=20=E2=80=94=20[control=5Fplane]=20config,=20worker=20delegati?= =?UTF-8?q?on=20serving,=20headless=20mode=20(PR=203/4)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit PR 3/4 of the Agent Control Plane stack: the OAB runtime side. The runtime dials out to the CP, registers, serves `cp/delegate` through the same ACP turn driver every chat platform uses, and reports each turn as `cp/delegate_result`. Config surface (`[control_plane]`, `crates/openab-core/src/config.rs`) - `url`, `auth_key`, `namespace`, `name`, `type` (`primary` | `worker`), optional `labels`, `max_delegated_sessions` (default 1). - Presence is the opt-in — no cargo feature — following the `[mcp]` precedent. `deny_unknown_fields`, so a mistyped key fails startup instead of defaulting silently. `${ENV}` expansion comes from the shared loader. - `CpAgentType` has only Primary and Worker: `observer` is a read-only lobby role, so a runtime claiming it is unrepresentable rather than merely validated against. - Validation: non-empty url/auth_key/namespace/name, and `max_delegated_sessions > 0`. Each of these would otherwise fail at registration and reconnect forever. State machine (`control_plane/client.rs`) - One process-lifetime instance id (uuid), reused across reconnects: it identifies the replica, not the connection. - Connect ws/wss with `Authorization: Bearer` (never in the URL), send `cp/register` as the first frame, adopt `effective_max_delegated_sessions` from the ack, then one loop over inbound frames / heartbeat tick / delegation completions / shutdown watch. - On close or error: cancel every local in-flight delegation (each stops its agent and drops its session; the results are deliberately dropped, since the CP synthesizes `target_disconnected`), then 1/2/4/8/16/30s shutdown-aware backoff, reconnect, re-register. - CP-issued requests get JSON-RPC result acks; a delegation's outcome never rides in that ack. Delegation serving (`control_plane/executor.rs`) - Admission never executes: over-cap or duplicate `delegation_id` returns `status = failed` without touching the session pool. - Bounded by the nearer of the CP deadline and `[pool].prompt_hard_timeout_secs`. Outcomes map to completed / failed (agent error, broker error, silent failure) / timeout / cancelled. - One fresh session per delegation, keyed `control-plane:`, discarded on every terminal outcome. Adds `SessionPool::discard_session` for the non-cancel teardown (`reset_session` would emit a spurious `session/cancel` and log an error for the benign already-gone case). Prompt seam refactor safety - `AdapterRouter::stream_prompt_blocks` now returns `Result`. Nothing else changed inside it: two values are captured before the existing composition consumes them, and the delivery paths are byte-for-byte the same. `Err` still means what it meant (incomplete user-visible delivery). - All platform callers map back with `.map(|_| ())`, and `DispatchTarget` stays at `Result<()>`, so Discord/Slack/gateway/ambient — and the existing dispatch mock — are untouched. Every pre-existing adapter test passes unmodified. - The executor reuses that same driver with a private sink adapter (send-once, no placeholder, no reactions) and reads the reply from the return value. Feature split (Dockerfile-neutral) - `openab-cp` gains a default `server` feature gating config/registry/policy/ router/events/server and the binary (`required-features`); `proto` — the wire contract — stays unconditional, as do serde/serde_json/chrono. - `openab-core` depends on `openab-cp` with `default-features = false` for wire types only. Verified: `cargo tree -p openab-core -e normal` shows only chrono/serde/serde_json under openab-cp, an inverted axum query finds nothing, and `cargo build -p openab-cp` still produces the binary. - No new crate, no new runtime cargo feature. Headless mode (`src/main.rs`) - `[agent]` + `[control_plane] type = "worker"` with no adapter is a valid boot; a `primary` alone is not (its prompts come from a chat platform). `[mcp]`-only is unchanged, and mcp + worker runs both. - Extracted as `headless_run_mode` so the matrix is unit-tested. - Client spawned after router construction, wired to the existing shutdown watch, and stopped before `pool.shutdown()` so the CP sees a clean disconnect rather than a lease timeout. Security - `auth_key` is used only for the upgrade header (marked sensitive), is never logged, and never enters the agent child env — `[agent].env` plumbing is untouched. Tests - Config: absent section, full section incl. `type` rename and `${ENV}`, defaults, observer refusal, unknown-key refusal, every validation failure. - Run modes: the full headless matrix. - Executor: cap enforcement, duplicate id, all outcome mappings, and cancel/deadline cleanup asserting the session is gone. - `crates/openab-core/tests/cp_client.rs`: seven tests against the REAL in-process CP (ephemeral loopback port + `run_sweeper`), with a raw-WS primary as initiator and a scripted `PromptRunner` at the trait seam — register/ack, heartbeats holding the lease, delegate→result round trip, local deadline → `timeout`, cancel mid-flight, zero-lease sweep → reconnect + re-register (same instance id, and the new registration serves), and clean deregistration on shutdown. Deviations from the design 1. `openab-cp`'s `config.rs` is behind `server` too. The design said "proto.rs (+ config types needed by proto) stays unconditional"; proto needs nothing from config (the dependency runs the other way), so gating it keeps anyhow and toml out of the runtime build. 2. Completed delegations report through an mpsc channel rather than a `JoinSet` polled in the `select!`. The inbound branch must spawn while the completion branch is armed, and one `select!` cannot lend the same `JoinSet` to both — the borrow checker rejects it. Same for the socket, which is split into sink/stream halves. Behaviour is unchanged; handles are still tracked for the drain-then-abort teardown. 3. `axum` added as a dev-dependency of `openab-core`: `openab_cp::server::app` returns an `axum::Router`, and the integration test needs `axum::serve` to bind it. Test-only — the runtime links no HTTP server. Audit round-1 fixes (all four blocking findings): 1. [mcp] + [control_plane type=primary] now takes the full boot path so the facade and the CP client both run; facade-only had foreclosed the client. Run-mode test updated to pin both-run. 2. Reconnect backoff no longer resets on every clean close: a session must live STABLE_SESSION_SECS (60s) before a close resets backoff to 1s, so a CP that accepts-then-closes (lease misconfig, crash loop) faces 1/2/4/8/16/30 escalation instead of a 1s reconnect storm. 3. The executor's best-effort session/cancel before discard is now bounded at 5s (the pool's own cleanup bound): a wedged agent stdin can no longer block discard and leak the delegation slot. Both sites (cancel, deadline). 4. Shutdown is storm-proof: the dial/register phase races the shutdown signal via select (a hung connect cannot stall shutdown), and main aborts the client task if it misses the 10s deadline instead of detaching it over pool teardown. --- Cargo.lock | 2 + crates/openab-core/Cargo.toml | 11 + crates/openab-core/src/acp/pool.rs | 29 + crates/openab-core/src/adapter.rs | 45 +- crates/openab-core/src/config.rs | 238 ++++- .../openab-core/src/control_plane/client.rs | 738 ++++++++++++++++ .../openab-core/src/control_plane/executor.rs | 816 ++++++++++++++++++ crates/openab-core/src/control_plane/mod.rs | 51 ++ crates/openab-core/src/dispatch.rs | 4 + crates/openab-core/src/lib.rs | 3 + crates/openab-core/tests/cp_client.rs | 563 ++++++++++++ crates/openab-cp/Cargo.toml | 53 +- crates/openab-cp/src/lib.rs | 19 + docs/adr/agent-control-plane.md | 2 +- docs/config-reference.md | 70 ++ src/main.rs | 221 ++++- 16 files changed, 2832 insertions(+), 33 deletions(-) create mode 100644 crates/openab-core/src/control_plane/client.rs create mode 100644 crates/openab-core/src/control_plane/executor.rs create mode 100644 crates/openab-core/src/control_plane/mod.rs create mode 100644 crates/openab-core/tests/cp_client.rs 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..17261aaa4 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,74 @@ 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 + /// runtime enforces whatever the ack returns, and its own local cap on top. + #[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 +316,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 +2399,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 +2637,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..f8cf73a8b --- /dev/null +++ b/crates/openab-core/src/control_plane/client.rs @@ -0,0 +1,738 @@ +//! 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::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); + +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 = self.register(&mut sink, &mut stream).await?; + 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); + let (ws, _resp) = tokio_tungstenite::connect_async(request) + .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); + 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", + "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.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", "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..f2b0b7fd6 --- /dev/null +++ b/crates/openab-core/src/control_plane/executor.rs @@ -0,0 +1,816 @@ +//! 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::{DelegateForward, DelegateResultParams, DelegationStatus}; +use sha2::{Digest, Sha256}; +use tokio::sync::Notify; +use tracing::{info, warn}; + +/// Session-pool key for 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` is mixed in so the same id seen after a +/// reconnect maps to a different session. +pub fn delegation_session_key(instance_id: &str, delegation_id: &str) -> String { + let mut hasher = Sha256::new(); + hasher.update(instance_id.as_bytes()); + hasher.update(delegation_id.as_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) -> 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(), 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) -> bool { + let signal = self + .inflight + .lock() + .expect("inflight mutex") + .get(delegation_id) + .map(Arc::clone); + 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(Arc::clone) + .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) { + Ok(signal) => signal, + Err(refusal) => { + let error = refusal.message(); + warn!(delegation_id = %id, from = %forward.from, %error, "delegation refused"); + return failed(&id, error); + } + }; + let outcome = self.execute(&forward, cancel).await; + self.release(&id); + outcome + } + + async fn execute( + &self, + forward: &DelegateForward, + cancel: Arc, + ) -> DelegateResultParams { + let id = &forward.delegation_id; + let session_key = delegation_session_key(&self.instance_id, id); + + // 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); + }; + 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"); + // Best-effort, BOUNDED: session/cancel writes to the agent's + // stdin, which can wedge (dead child, full pipe). The pool's + // own cleanup uses the same 5s bound. Discard must run either + // way or the slot leaks. + let _ = tokio::time::timeout( + std::time::Duration::from_secs(5), + self.runner.cancel(&session_key), + ) + .await; + self.runner.discard(&session_key).await; + return DelegateResultParams { + delegation_id: id.clone(), + 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"); + // Best-effort, BOUNDED: session/cancel writes to the agent's + // stdin, which can wedge (dead child, full pipe). The pool's + // own cleanup uses the same 5s bound. Discard must run either + // way or the slot leaks. + let _ = tokio::time::timeout( + std::time::Duration::from_secs(5), + self.runner.cancel(&session_key), + ) + .await; + self.runner.discard(&session_key).await; + timed_out(id) + } + 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.runner.discard(&session_key).await; + failed(id, error) + } + Ok(Ok(outcome)) => { + self.runner.discard(&session_key).await; + if let Some(error) = outcome.error { + warn!(delegation_id = %id, %error, "delegation ended in an agent error"); + return failed(id, error); + } + if outcome.silent_failure { + warn!(delegation_id = %id, "delegation produced an empty turn (silent failure)"); + return failed( + id, + "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(), + status: DelegationStatus::Completed, + result: Some(outcome.text), + error: None, + } + } + } + } +} + +fn failed(delegation_id: &str, error: impl Into) -> DelegateResultParams { + DelegateResultParams { + delegation_id: delegation_id.to_string(), + status: DelegationStatus::Failed, + result: None, + error: Some(error.into()), + } +} + +fn timed_out(delegation_id: &str) -> DelegateResultParams { + DelegateResultParams { + delegation_id: delegation_id.to_string(), + 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(), + 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")], + "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").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").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").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"); + 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"), "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"); + 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")); + } + + #[test] + fn session_keys_are_namespaced_bounded_and_instance_scoped() { + let a = delegation_session_key("i-1", "d-1"); + let b = delegation_session_key("i-2", "d-1"); + assert!(a.starts_with("control-plane:")); + assert_ne!(a, b, "a replayed id after reconnect gets a fresh session"); + assert_eq!(a.len(), "control-plane:".len() + 64); + // A hostile id cannot forge another platform's key shape. + let hostile = delegation_session_key("i-1", "discord:12345"); + 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()); + } +} 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..a5ff06938 --- /dev/null +++ b/crates/openab-core/tests/cp_client.rs @@ -0,0 +1,563 @@ +//! 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; + assert_eq!(next_json(&mut initiator).await["id"], 4, "delegate ack"); + 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", "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"); + 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..ebe156e6a 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" diff --git a/docs/config-reference.md b/docs/config-reference.md index af30e0940..35f6241a6 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 the lower of the two | + +`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()); + } } From 8db2649e5ac5110d334b9179d7653a6bcbbb1c10 Mon Sep 17 00:00:00 2001 From: chaodu-agent Date: Wed, 12 Aug 2026 12:12:10 -0400 Subject: [PATCH 2/5] fix(cp): cap delegate_result bodies below the CP transport frame limit MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Self-review finding. The worker sent the agent's full final text in cp/delegate_result (SinkAdapter deliberately does no chunking), but the CP enforces max_frame_bytes (default 1 MiB) at the WS transport BEFORE parsing — its router-level max_result_bytes truncation can never save an oversized frame. A single >1 MiB result therefore dropped the whole connection, and every in-flight delegation on that worker died as target_disconnected. The executor now caps Completed results at 512 KiB (half the default frame budget, leaving JSON-envelope headroom) with a UTF-8-boundary-safe truncation marker; the CP still applies its own, typically smaller, max_result_bytes on arrival. Regression test covers the oversized, multibyte-straddling, and under-cap cases. --- .../openab-core/src/control_plane/executor.rs | 52 ++++++++++++++++++- 1 file changed, 51 insertions(+), 1 deletion(-) diff --git a/crates/openab-core/src/control_plane/executor.rs b/crates/openab-core/src/control_plane/executor.rs index f2b0b7fd6..54474c7ce 100644 --- a/crates/openab-core/src/control_plane/executor.rs +++ b/crates/openab-core/src/control_plane/executor.rs @@ -316,7 +316,7 @@ impl DelegationExecutor { DelegateResultParams { delegation_id: id.clone(), status: DelegationStatus::Completed, - result: Some(outcome.text), + result: Some(cap_result(outcome.text)), error: None, } } @@ -324,6 +324,35 @@ impl DelegationExecutor { } } +/// 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; + +fn cap_result(text: String) -> String { + if text.len() <= MAX_RESULT_BYTES { + return text; + } + let marker = format!( + "\n…[truncated by worker: {} bytes total exceeded the transport budget]", + text.len() + ); + let budget = MAX_RESULT_BYTES.saturating_sub(marker.len()); + let mut cut = budget.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 failed(delegation_id: &str, error: impl Into) -> DelegateResultParams { DelegateResultParams { delegation_id: delegation_id.to_string(), @@ -813,4 +842,25 @@ mod tests { 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"); + } } From e2619e9d8d4b238afe3fe684866cfdff07684d21 Mon Sep 17 00:00:00 2001 From: chaodu-agent Date: Wed, 12 Aug 2026 22:26:23 -0400 Subject: [PATCH 3/5] fix(cp): free the delegation slot via RAII so task abort cannot leak it MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Self-review round-2 finding. serve() released its inflight slot with a plain call after the execute().await — but the client ABORTS serving tasks that outlive the 5s drain window on disconnect, and an aborted task never runs code after its await point. The leaked entry then survives reconnection (executor state is process-lifetime), permanently consuming capacity: with the default max_delegated_sessions = 1, one abort turns the worker into a zombie that refuses every delegation while heartbeating happily. The window is realistic, not theoretical: the bounded (5s) session/cancel added in the previous fix means a wedged agent stdin makes the task overrun the equally-sized drain window by construction. release() now lives in a Drop guard taken at admission, so the abort's unwind frees the slot. Regression test aborts a serving task mid-run and asserts active() == 0; mutation-verified (plain-release pattern fails it 5/5, the guard passes). --- .../openab-core/src/control_plane/executor.rs | 56 ++++++++++++++++++- 1 file changed, 53 insertions(+), 3 deletions(-) diff --git a/crates/openab-core/src/control_plane/executor.rs b/crates/openab-core/src/control_plane/executor.rs index 54474c7ce..e6473a22f 100644 --- a/crates/openab-core/src/control_plane/executor.rs +++ b/crates/openab-core/src/control_plane/executor.rs @@ -219,9 +219,16 @@ impl DelegationExecutor { return failed(&id, error); } }; - let outcome = self.execute(&forward, cancel).await; - self.release(&id); - outcome + // 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( @@ -353,6 +360,19 @@ fn cap_result(text: String) -> String { out } +/// 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, error: impl Into) -> DelegateResultParams { DelegateResultParams { delegation_id: delegation_id.to_string(), @@ -863,4 +883,34 @@ mod tests { // Under the cap: untouched. assert_eq!(cap_result("small".into()), "small"); } + + #[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); + } } From 655b42487007dcab54d0fd3a302f4f6633e1c346 Mon Sep 17 00:00:00 2001 From: chaodu-agent Date: Fri, 14 Aug 2026 00:07:38 -0400 Subject: [PATCH 4/5] fix(cp): echo the CP admission token in every delegate_result and match it on cancel MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Round-8 made the admission token protocol-visible on the result path; round-9 extended it to cancellation. The runtime side of both: - The executor carries forward.admission into all five terminal frame kinds (completed, failed, refused-at-admission, timeout, cancelled). - The worker-side cancel handler matches (delegation_id, admission): a stale cancel naming a superseded admission of a reused id is ignored instead of aborting the live delegation — the worker-side half of the misdelivery the wire token closes. Regression pins both directions. - The initiator-side integration test echoes the ack's admission in its cp/cancel, exercising the CP's required-field validation end-to-end. The real-server integration suite is the first live client of the token-enforcing contract on both the result and cancel paths. --- .../openab-core/src/control_plane/client.rs | 11 +- .../openab-core/src/control_plane/executor.rs | 114 ++++++++++++++---- crates/openab-core/tests/cp_client.rs | 8 +- 3 files changed, 104 insertions(+), 29 deletions(-) diff --git a/crates/openab-core/src/control_plane/client.rs b/crates/openab-core/src/control_plane/client.rs index f8cf73a8b..a8a64e213 100644 --- a/crates/openab-core/src/control_plane/client.rs +++ b/crates/openab-core/src/control_plane/client.rs @@ -442,7 +442,9 @@ impl ControlPlaneClient { msg.params.and_then(|p| serde_json::from_value(p).ok()); match params { Some(params) => { - let known = self.executor.cancel(¶ms.delegation_id); + let known = self + .executor + .cancel(¶ms.delegation_id, params.admission); info!( delegation_id = %params.delegation_id, reason = %params.reason, @@ -622,6 +624,7 @@ max_delegated_sessions = 3 "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", @@ -639,6 +642,10 @@ max_delegated_sessions = 3 "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()]); } @@ -668,7 +675,7 @@ max_delegated_sessions = 3 let c = client(); let frame = serde_json::json!({ "jsonrpc": "2.0", "id": 4, "method": "cp/cancel", - "params": {"delegation_id": "gone", "reason": "initiator gave up"} + "params": {"delegation_id": "gone", "admission": 3, "reason": "initiator gave up"} }) .to_string(); let FrameAction::Reply(reply) = c.handle_frame(&frame) else { diff --git a/crates/openab-core/src/control_plane/executor.rs b/crates/openab-core/src/control_plane/executor.rs index e6473a22f..e62433380 100644 --- a/crates/openab-core/src/control_plane/executor.rs +++ b/crates/openab-core/src/control_plane/executor.rs @@ -24,7 +24,7 @@ use std::time::Duration; use anyhow::Result; use async_trait::async_trait; -use openab_cp::proto::{DelegateForward, DelegateResultParams, DelegationStatus}; +use openab_cp::proto::{AdmissionToken, DelegateForward, DelegateResultParams, DelegationStatus}; use sha2::{Digest, Sha256}; use tokio::sync::Notify; use tracing::{info, warn}; @@ -86,7 +86,7 @@ pub struct DelegationExecutor { /// Admitted delegations → their cancel signal. Also the capacity counter: /// its length is the number of active delegated sessions reported in /// `cp/heartbeat`. - inflight: Mutex>>, + inflight: Mutex)>>, } /// Why admission refused, as the message sent back in `status = failed`. @@ -141,7 +141,11 @@ impl DelegationExecutor { } /// Reserve a slot for `delegation_id`, or explain why not. - fn admit(&self, delegation_id: &str) -> std::result::Result, Refusal> { + 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) { @@ -152,7 +156,7 @@ impl DelegationExecutor { return Err(Refusal::OverCapacity { active, max }); } let signal = Arc::new(Notify::new()); - g.insert(delegation_id.to_string(), Arc::clone(&signal)); + g.insert(delegation_id.to_string(), (admission, Arc::clone(&signal))); Ok(signal) } @@ -170,13 +174,28 @@ impl DelegationExecutor { /// `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) -> bool { - let signal = self - .inflight - .lock() - .expect("inflight mutex") - .get(delegation_id) - .map(Arc::clone); + 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(); @@ -197,7 +216,7 @@ impl DelegationExecutor { .lock() .expect("inflight mutex") .values() - .map(Arc::clone) + .map(|(_, s)| Arc::clone(s)) .collect(); for s in signals { s.notify_one(); @@ -211,12 +230,12 @@ impl DelegationExecutor { /// 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) { + 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, error); + return failed(&id, forward.admission, error); } }; // RAII: the slot must free even if this task is ABORTED mid-await — @@ -244,7 +263,7 @@ impl DelegationExecutor { // 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); + return timed_out(id, forward.admission); }; let budget = remaining.min(self.prompt_hard_timeout); info!( @@ -275,6 +294,7 @@ impl DelegationExecutor { self.runner.discard(&session_key).await; return DelegateResultParams { delegation_id: id.clone(), + admission: forward.admission, status: DelegationStatus::Cancelled, result: None, error: None, @@ -296,25 +316,26 @@ impl DelegationExecutor { ) .await; self.runner.discard(&session_key).await; - timed_out(id) + 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.runner.discard(&session_key).await; - failed(id, error) + failed(id, forward.admission, error) } Ok(Ok(outcome)) => { self.runner.discard(&session_key).await; if let Some(error) = outcome.error { warn!(delegation_id = %id, %error, "delegation ended in an agent error"); - return failed(id, 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", ); @@ -322,6 +343,7 @@ impl DelegationExecutor { 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, @@ -373,18 +395,27 @@ impl Drop for SlotGuard<'_> { } } -fn failed(delegation_id: &str, error: impl Into) -> DelegateResultParams { +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, error: Some(error.into()), } } -fn timed_out(delegation_id: &str) -> DelegateResultParams { +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()), @@ -569,6 +600,7 @@ mod tests { 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(), @@ -672,7 +704,7 @@ mod tests { 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").expect("first slot"); + 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")); @@ -684,7 +716,7 @@ mod tests { 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").expect("slot"); + 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")); @@ -696,7 +728,7 @@ mod tests { 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").expect("slot"); + 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!( @@ -788,7 +820,7 @@ mod tests { while ex.active() == 0 { tokio::task::yield_now().await; } - assert!(ex.cancel("d-9"), "the id is in flight"); + 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()); @@ -826,7 +858,7 @@ mod tests { #[test] fn cancel_of_an_unknown_id_is_a_no_op() { let ex = executor(FakeRunner::completing("ok"), 1); - assert!(!ex.cancel("never-seen")); + assert!(!ex.cancel("never-seen", 1)); } #[test] @@ -913,4 +945,36 @@ mod tests { 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/tests/cp_client.rs b/crates/openab-core/tests/cp_client.rs index a5ff06938..4b81b18c3 100644 --- a/crates/openab-core/tests/cp_client.rs +++ b/crates/openab-core/tests/cp_client.rs @@ -445,12 +445,16 @@ async fn a_cancel_stops_the_turn_and_frees_the_slot() { let mut initiator = connect_initiator(&url).await; delegate(&mut initiator, 4, "d-cancel", "hang forever", 300).await; - assert_eq!(next_json(&mut initiator).await["id"], 4, "delegate ack"); + 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", "reason": "initiator changed its mind"} + "params": {"delegation_id": "d-cancel", "admission": admission, "reason": "initiator changed its mind"} }) .to_string(); initiator.send(Message::Text(cancel)).await.unwrap(); From d3d130b0094d9bed7d719e452c9e94cee67feb7e Mon Sep 17 00:00:00 2001 From: chaodu-agent Date: Fri, 14 Aug 2026 22:01:23 +0000 Subject: [PATCH 5/5] fix(review): cap error payloads, admission-scope session keys, bound teardown MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Round-7 review fixes: - F1: failed() now caps every error string (64 KiB, UTF-8-safe) inside the constructor so no call site can bypass it; oversized-error regression test mirrors the success-path one - F34: delegation_session_key mixes in the CP admission token, so a re-admission of a reusable id can never resume an earlier admission's session (e.g. one orphaned by a drain-timeout abort); false reconnect claim in the doc corrected; key-property test pins admission scoping - F22/F31: duplicated cancel+discard blocks extracted into bounded helpers; discard now shares the 5s bound instead of being unbounded on every terminal path - F15: client WS now sets max_message_size/max_frame_size (1 MiB), mirroring the CP's accept-side max_frame_bytes instead of tungstenite's 64 MiB default - F5: the cp/register ack wait is bounded (10s, mirroring the CP's register_timeout_secs) so a stalled CP enters backoff instead of hanging - F27: config/docs no longer claim a local min() the code does not implement — the runtime enforces the ack value - F37: ADR section 7 facade label corrected from PR 3/4 to PR 4/4 --- crates/openab-core/src/config.rs | 5 +- .../openab-core/src/control_plane/client.rs | 36 ++++- .../openab-core/src/control_plane/executor.rs | 138 +++++++++++++----- crates/openab-core/tests/cp_client.rs | 7 +- docs/adr/agent-control-plane.md | 2 +- docs/config-reference.md | 2 +- 6 files changed, 141 insertions(+), 49 deletions(-) diff --git a/crates/openab-core/src/config.rs b/crates/openab-core/src/config.rs index 17261aaa4..7b5a37d0f 100644 --- a/crates/openab-core/src/config.rs +++ b/crates/openab-core/src/config.rs @@ -168,8 +168,9 @@ pub struct ControlPlaneConfig { /// `cp/delegate` label targeting. #[serde(default)] pub labels: BTreeMap, - /// Concurrency budget advertised at registration. The CP may clamp it; the - /// runtime enforces whatever the ack returns, and its own local cap on top. + /// 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, } diff --git a/crates/openab-core/src/control_plane/client.rs b/crates/openab-core/src/control_plane/client.rs index a8a64e213..b3f67409d 100644 --- a/crates/openab-core/src/control_plane/client.rs +++ b/crates/openab-core/src/control_plane/client.rs @@ -31,6 +31,7 @@ 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}; @@ -48,6 +49,16 @@ const STABLE_SESSION_SECS: u64 = 60; /// 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!` @@ -161,7 +172,14 @@ impl ControlPlaneClient { ) -> anyhow::Result { let ws = self.connect().await?; let (mut sink, mut stream) = ws.split(); - let ack = self.register(&mut sink, &mut stream).await?; + 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!( @@ -184,9 +202,19 @@ impl ControlPlaneClient { // Belt and braces: the key must not surface in a `{:?}` of the request. value.set_sensitive(true); request.headers_mut().insert("Authorization", value); - let (ws, _resp) = tokio_tungstenite::connect_async(request) - .await - .map_err(|e| anyhow::anyhow!("control-plane handshake failed: {e}"))?; + // 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) } diff --git a/crates/openab-core/src/control_plane/executor.rs b/crates/openab-core/src/control_plane/executor.rs index e62433380..c49868b41 100644 --- a/crates/openab-core/src/control_plane/executor.rs +++ b/crates/openab-core/src/control_plane/executor.rs @@ -29,17 +29,26 @@ use sha2::{Digest, Sha256}; use tokio::sync::Notify; use tracing::{info, warn}; -/// Session-pool key for one delegation. +/// 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` is mixed in so the same id seen after a -/// reconnect maps to a different session. -pub fn delegation_session_key(instance_id: &str, delegation_id: &str) -> String { +/// 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()) } @@ -256,7 +265,7 @@ impl DelegationExecutor { cancel: Arc, ) -> DelegateResultParams { let id = &forward.delegation_id; - let session_key = delegation_session_key(&self.instance_id, 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 @@ -282,16 +291,7 @@ impl DelegationExecutor { biased; _ = cancelled => { info!(delegation_id = %id, "delegation cancelled"); - // Best-effort, BOUNDED: session/cancel writes to the agent's - // stdin, which can wedge (dead child, full pipe). The pool's - // own cleanup uses the same 5s bound. Discard must run either - // way or the slot leaks. - let _ = tokio::time::timeout( - std::time::Duration::from_secs(5), - self.runner.cancel(&session_key), - ) - .await; - self.runner.discard(&session_key).await; + self.cancel_and_discard(&session_key).await; return DelegateResultParams { delegation_id: id.clone(), admission: forward.admission, @@ -306,27 +306,18 @@ impl DelegationExecutor { match outcome { Err(_elapsed) => { warn!(delegation_id = %id, budget_secs = budget.as_secs(), "delegation exceeded its local deadline"); - // Best-effort, BOUNDED: session/cancel writes to the agent's - // stdin, which can wedge (dead child, full pipe). The pool's - // own cleanup uses the same 5s bound. Discard must run either - // way or the slot leaks. - let _ = tokio::time::timeout( - std::time::Duration::from_secs(5), - self.runner.cancel(&session_key), - ) - .await; - self.runner.discard(&session_key).await; + 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.runner.discard(&session_key).await; + self.bounded_discard(&session_key).await; failed(id, forward.admission, error) } Ok(Ok(outcome)) => { - self.runner.discard(&session_key).await; + 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); @@ -351,8 +342,32 @@ impl DelegationExecutor { } } } + + /// 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, @@ -363,16 +378,24 @@ impl DelegationExecutor { /// its (typically smaller) `max_result_bytes` on what arrives. const MAX_RESULT_BYTES: usize = 512 * 1024; -fn cap_result(text: String) -> String { - if text.len() <= MAX_RESULT_BYTES { +/// 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 budget = MAX_RESULT_BYTES.saturating_sub(marker.len()); - let mut cut = budget.min(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; } @@ -382,6 +405,10 @@ fn cap_result(text: String) -> String { 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> { @@ -408,7 +435,10 @@ fn failed( admission, status: DelegationStatus::Failed, result: None, - error: Some(error.into()), + // 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)), } } @@ -692,7 +722,7 @@ mod tests { assert!(res.error.is_none()); assert_eq!( runner.discarded(), - vec![delegation_session_key("i-test", "d-1")], + 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"); @@ -799,7 +829,7 @@ mod tests { 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"); + 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); @@ -824,7 +854,7 @@ mod tests { 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"); + 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); @@ -863,13 +893,17 @@ mod tests { #[test] fn session_keys_are_namespaced_bounded_and_instance_scoped() { - let a = delegation_session_key("i-1", "d-1"); - let b = delegation_session_key("i-2", "d-1"); + 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, "a replayed id after reconnect gets a fresh session"); + 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"); + let hostile = delegation_session_key("i-1", "discord:12345", 1); assert!(hostile.starts_with("control-plane:")); assert_eq!(hostile.len(), a.len()); } @@ -916,6 +950,32 @@ mod tests { 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 diff --git a/crates/openab-core/tests/cp_client.rs b/crates/openab-core/tests/cp_client.rs index 4b81b18c3..c3122280f 100644 --- a/crates/openab-core/tests/cp_client.rs +++ b/crates/openab-core/tests/cp_client.rs @@ -459,8 +459,11 @@ async fn a_cancel_stops_the_turn_and_frees_the_slot() { .to_string(); initiator.send(Message::Text(cancel)).await.unwrap(); - let expected_session = - openab_core::control_plane::delegation_session_key(client.instance_id(), "d-cancel"); + 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) diff --git a/docs/adr/agent-control-plane.md b/docs/adr/agent-control-plane.md index ebe156e6a..6ee55977b 100644 --- a/docs/adr/agent-control-plane.md +++ b/docs/adr/agent-control-plane.md @@ -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 35f6241a6..ff2a5426c 100644 --- a/docs/config-reference.md +++ b/docs/config-reference.md @@ -796,7 +796,7 @@ tier = "batch" | `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 the lower of the two | +| `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,