From 91432ca4b516682e5913fe08096ec70dcf544157 Mon Sep 17 00:00:00 2001 From: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 1 Oct 2026 01:36:06 +0800 Subject: [PATCH 01/13] test: resolve clippy --all-targets lint warnings under rust 1.98 Fixes pre-existing test-code lints that `cargo clippy --workspace --all-targets -- -D warnings` reports on current stable: literal-bool assert_eq, manual str::repeat, unwrap_or on literal Some/None, get().is_some()/is_none() on maps, an unused cfg(test) helper, and one intentional ENV_LOCK hold across .await (env-serialization in tests). (cherry picked from commit ae0e563f74036c5f16bff1cf28d602d813e542d3) Refs #1475 Generated with [Devin](https://devin.ai) Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- crates/openab-core/src/cron.rs | 4 ++-- crates/openab-core/src/discord.rs | 7 +++---- crates/openab-core/src/format.rs | 2 +- crates/openab-mcp/src/auth.rs | 4 ++-- crates/openab-mcp/src/mcp/runtime.rs | 1 + crates/openab-mcp/src/native/gmail.rs | 6 ------ 6 files changed, 9 insertions(+), 15 deletions(-) diff --git a/crates/openab-core/src/cron.rs b/crates/openab-core/src/cron.rs index 40e9a67ca..7f9d59846 100644 --- a/crates/openab-core/src/cron.rs +++ b/crates/openab-core/src/cron.rs @@ -1250,7 +1250,7 @@ message = "hello" "#; let cfg: UsercronFile = toml::from_str(toml_str).unwrap(); let job = &cfg.jobs[0]; - assert_eq!(job.enabled, true); + assert!(job.enabled); assert_eq!(job.platform, "discord"); assert_eq!(job.sender_name, "openab-cron"); assert_eq!(job.timezone, "UTC"); @@ -1272,7 +1272,7 @@ channel = "123" message = "hello" "#; let cfg: UsercronFile = toml::from_str(toml_str).unwrap(); - assert_eq!(cfg.jobs[0].enabled, false); + assert!(!cfg.jobs[0].enabled); } #[test] diff --git a/crates/openab-core/src/discord.rs b/crates/openab-core/src/discord.rs index 2ea8b28c0..c34c68a7d 100644 --- a/crates/openab-core/src/discord.rs +++ b/crates/openab-core/src/discord.rs @@ -3644,14 +3644,14 @@ mod tests { fn image_attachment_block_includes_url_and_metadata() { // Simulates the format string used in the image attachment handler. let filename = "screenshot.png"; - let content_type = Some("image/png"); + let content_type = "image/png"; let size: u32 = 142048; let url = "https://cdn.discordapp.com/attachments/123/456/screenshot.png"; let text = format!( "[Image attachment]\nfilename: {}\ncontent_type: {}\nsize_bytes: {}\nurl: {} (expires ~24h)", filename, - content_type.unwrap_or("unknown"), + content_type, size, url, ); @@ -3666,11 +3666,10 @@ mod tests { #[test] fn image_attachment_block_missing_content_type_falls_back() { - let content_type: Option<&str> = None; let text = format!( "[Image attachment]\nfilename: {}\ncontent_type: {}\nsize_bytes: {}\nurl: {} (expires ~24h)", "photo.jpg", - content_type.unwrap_or("unknown"), + "unknown", 99999, "https://cdn.discordapp.com/attachments/1/2/photo.jpg", ); diff --git a/crates/openab-core/src/format.rs b/crates/openab-core/src/format.rs index 4fa1ce9e1..41a26d0df 100644 --- a/crates/openab-core/src/format.rs +++ b/crates/openab-core/src/format.rs @@ -434,7 +434,7 @@ mod tests { #[test] fn fenced_hard_split_grapheme_safe() { // Emoji inside a code fence: grapheme-safe, all content preserved. - let content: String = std::iter::repeat("🎉❤️你").take(20).collect(); + let content: String = "🎉❤️你".repeat(20); let text = format!("```\n{content}\n```"); let chunks = split_message(&text, 20); assert_length_invariant(&chunks, 20); diff --git a/crates/openab-mcp/src/auth.rs b/crates/openab-mcp/src/auth.rs index 7a8f2fd38..3c21b09c6 100644 --- a/crates/openab-mcp/src/auth.rs +++ b/crates/openab-mcp/src/auth.rs @@ -2355,14 +2355,14 @@ mod tests { let map = read_auth_file(&path).unwrap(); assert!( - map.get("mcp-pending:stale").is_none(), + !map.contains_key("mcp-pending:stale"), "stale pending swept" ); assert!( matches!(map.get("mcp-pending:fresh"), Some(AuthEntry::Pending(_))), "fresh pending kept" ); - assert!(map.get("codex").is_some(), "real tenant untouched"); + assert!(map.contains_key("codex"), "real tenant untouched"); } #[cfg(unix)] diff --git a/crates/openab-mcp/src/mcp/runtime.rs b/crates/openab-mcp/src/mcp/runtime.rs index 29f91b29f..0251287d3 100644 --- a/crates/openab-mcp/src/mcp/runtime.rs +++ b/crates/openab-mcp/src/mcp/runtime.rs @@ -3309,6 +3309,7 @@ mod tests { } #[tokio::test] + #[allow(clippy::await_holding_lock)] async fn start_paste_login_builtin_without_env_var_errors_loud() { let _guard = ENV_LOCK.lock().unwrap_or_else(|e| e.into_inner()); unsafe { diff --git a/crates/openab-mcp/src/native/gmail.rs b/crates/openab-mcp/src/native/gmail.rs index 0549b0935..26288b5a5 100644 --- a/crates/openab-mcp/src/native/gmail.rs +++ b/crates/openab-mcp/src/native/gmail.rs @@ -121,12 +121,6 @@ impl GmailNative { } } - #[cfg(test)] - fn with_api_base(mut self, base: &str) -> Self { - self.api_base = base.to_string(); - self - } - /// Valid bearer token: cached when fresh, refreshed via the stored /// refresh token when near expiry (60s buffer, same policy as the MCP /// runtime), otherwise an actionable re-login error. From afc77510e84d21dfc4abe072de907b336b8b0875 Mon Sep 17 00:00:00 2001 From: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 1 Oct 2026 01:36:36 +0800 Subject: [PATCH 02/13] =?UTF-8?q?feat(cp):=20graceful=20shutdown=20?= =?UTF-8?q?=E2=80=94=20SIGTERM/SIGHUP=20handling=20and=20connection=20drai?= =?UTF-8?q?n?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit docker stop, ECS and k8s send SIGTERM before SIGKILL, but openab-cp only handled ctrl-c: the default disposition killed the process, so connections saw TCP resets, in-flight delegations vanished without terminal frames, and the sweeper was aborted while WS tasks were still mid-write. Now all three signals run one drain path: every in-flight delegation is resolved first (initiators get a synthesized target_disconnected terminal, serving runtimes a cp/cancel, both stamped "control plane shutting down"), then every connection — including pre-registration sockets the registry cannot reach — flushes its outbound queue and closes with code 1012 (service restart) inside shutdown_drain_secs, and only then is the sweeper stopped and the process exits. A shutting_down latch on Router, checked under the same inflight lock that fail_all drains, refuses admissions racing the drain with SATURATED — no delegation can be accepted after the last terminal synthesis. Refs #1475 Generated with [Devin](https://devin.ai) Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- crates/openab-cp/cp.toml.example | 10 + crates/openab-cp/src/config.rs | 18 ++ crates/openab-cp/src/main.rs | 60 ++++- crates/openab-cp/src/router.rs | 244 ++++++++++++++++++++- crates/openab-cp/src/server.rs | 292 ++++++++++++++++++++----- crates/openab-cp/tests/ws_lifecycle.rs | 200 ++++++++++++++++- docs/adr/agent-control-plane.md | 15 +- docs/control-plane.md | 14 +- 8 files changed, 782 insertions(+), 71 deletions(-) diff --git a/crates/openab-cp/cp.toml.example b/crates/openab-cp/cp.toml.example index f4ea04ddc..5fb04c678 100644 --- a/crates/openab-cp/cp.toml.example +++ b/crates/openab-cp/cp.toml.example @@ -60,6 +60,16 @@ write_timeout_secs = 30 # max_frame_bytes so one maximum-size frame always fits. max_outbound_queue_bytes = 16777216 +# Grace period for live connections after SIGTERM/SIGHUP/SIGINT, seconds. +# `docker stop` and k8s/ECS send SIGTERM first and SIGKILL after their own +# grace period, so keep this comfortably below it. On the signal the CP +# resolves every in-flight delegation (initiators get a synthesized +# target_disconnected terminal, serving runtimes a cp/cancel), then each +# connection flushes its queued outbound frames and closes with code 1012 — +# all inside this budget; a stalled peer can hold it no longer than one +# write_timeout_secs more. 0 exits as soon as the close signal is raised. +shutdown_drain_secs = 5 + # Process-wide ceiling on simultaneously in-flight delegations, enforced at # admission (before any capacity is reserved) with SATURATED. Per-target # max_delegated_sessions is advertised by the runtime and bounds one target's diff --git a/crates/openab-cp/src/config.rs b/crates/openab-cp/src/config.rs index a7fd7b6f9..40b0cb7d8 100644 --- a/crates/openab-cp/src/config.rs +++ b/crates/openab-cp/src/config.rs @@ -90,6 +90,21 @@ pub struct CpConfig { /// short enough that a dead one frees its quota promptly. #[serde(default = "default_write_timeout_secs")] pub write_timeout_secs: u64, + + /// How long the CP waits for live connections to drain after a shutdown + /// signal (SIGTERM/SIGHUP/SIGINT) before the process exits. + /// + /// On shutdown every in-flight delegation is resolved first — the + /// initiator gets a synthesized `target_disconnected` terminal and the + /// serving runtime a `cp/cancel` — then each connection flushes its + /// queued outbound frames and sends its close frame, all inside this + /// budget. A stalled peer's close write is additionally bounded by + /// `write_timeout_secs`, so it can never hold the CP past this budget + /// plus one write timeout. 0 disables the wait: connections are + /// signalled and the process exits as soon as the listener is closed + /// (sockets still get their close frames if they flush in time). + #[serde(default = "default_shutdown_drain_secs")] + pub shutdown_drain_secs: u64, /// Maximum size of prompt/result excerpts mirrored to observers in /// `cp/event` frames. The lobby is an audit surface, not a second /// delivery path: excerpts are truncated with a marker, never rejected. @@ -184,6 +199,9 @@ fn default_max_connections_per_identity() -> u32 { fn default_write_timeout_secs() -> u64 { 30 } +fn default_shutdown_drain_secs() -> u64 { + 5 +} fn default_max_outbound_queue_bytes() -> usize { 16 * 1024 * 1024 } diff --git a/crates/openab-cp/src/main.rs b/crates/openab-cp/src/main.rs index 7a34753c3..f8e268bd3 100644 --- a/crates/openab-cp/src/main.rs +++ b/crates/openab-cp/src/main.rs @@ -7,7 +7,7 @@ use clap::Parser; use tracing::info; use openab_cp::config::CpConfig; -use openab_cp::server::{app, run_sweeper, AppState}; +use openab_cp::server::{app, graceful_shutdown, run_sweeper, AppState}; #[derive(Parser)] #[command(name = "openab-cp", about = "OpenAB Agent Control Plane")] @@ -17,6 +17,49 @@ struct Cli { config: String, } +/// Wait for a shutdown signal: SIGINT (ctrl-c) on every platform, plus +/// SIGTERM and SIGHUP on unix — the signals `docker stop`, ECS and k8s send +/// before escalating to SIGKILL. Without handlers the default disposition +/// kills the process outright: no close frames, no terminal results, and +/// runtimes see TCP resets. +async fn shutdown_signal() -> &'static str { + #[cfg(unix)] + { + use tokio::signal::unix::{signal, SignalKind}; + let mut term = match signal(SignalKind::terminate()) { + Ok(s) => s, + Err(e) => { + tracing::warn!( + error = %e, + "failed to install SIGTERM handler — falling back to ctrl-c only" + ); + let _ = tokio::signal::ctrl_c().await; + return "SIGINT"; + } + }; + let mut hup = match signal(SignalKind::hangup()) { + Ok(s) => s, + Err(e) => { + tracing::warn!(error = %e, "failed to install SIGHUP handler"); + return tokio::select! { + _ = tokio::signal::ctrl_c() => "SIGINT", + _ = term.recv() => "SIGTERM", + }; + } + }; + tokio::select! { + _ = tokio::signal::ctrl_c() => "SIGINT", + _ = term.recv() => "SIGTERM", + _ = hup.recv() => "SIGHUP", + } + } + #[cfg(not(unix))] + { + let _ = tokio::signal::ctrl_c().await; + "SIGINT" + } +} + #[tokio::main] async fn main() -> Result<()> { tracing_subscriber::fmt() @@ -44,13 +87,22 @@ async fn main() -> Result<()> { let listener = tokio::net::TcpListener::bind(&listen) .await .with_context(|| format!("binding {listen}"))?; - axum::serve(listener, app(state)) + axum::serve(listener, app(state.clone())) .with_graceful_shutdown(async { - let _ = tokio::signal::ctrl_c().await; - info!("shutdown signal received"); + let signal = shutdown_signal().await; + info!(signal, "shutdown signal received — draining connections"); }) .await?; + // The listener has stopped accepting and in-flight HTTP work is done; + // the WS connection tasks are detached spawns still running. Give them a + // real shutdown — synthesized terminals, close frames, bounded drain — + // instead of letting process exit reset their sockets. The sweeper + // stops last: a lease/deadline expiry during the drain still + // synthesizes normally. + graceful_shutdown(&state).await; sweeper.abort(); + let _ = sweeper.await; + info!("openab-cp stopped"); Ok(()) } diff --git a/crates/openab-cp/src/router.rs b/crates/openab-cp/src/router.rs index ff9d75d7a..36a19a572 100644 --- a/crates/openab-cp/src/router.rs +++ b/crates/openab-cp/src/router.rs @@ -9,11 +9,16 @@ //! instance fail immediately with `target_disconnected`. //! - **Initiator disconnect** — its in-flight delegations are cancelled //! downstream (best effort); nobody is left to receive the result. -//! - **CP restart** — the table dies with the process, and the connections -//! die with it, so the CP cannot synthesize anything: in-flight delegations -//! end as initiator-side timeouts against the already-propagated deadline. -//! Late `cp/delegate_result` frames for unknown ids are acknowledged and -//! dropped (logged), so reconnecting runtimes do not error-loop. +//! - **CP shutdown** — a *graceful* stop (SIGTERM/SIGHUP/SIGINT) drains the +//! table through [`Router::fail_all`]: initiators get a synthesized +//! `target_disconnected` terminal, serving runtimes a `cp/cancel`, and the +//! `shutting_down` latch refuses any later admission with `SATURATED`. A +//! *hard* restart (SIGKILL, crash) is different — the table dies with the +//! process and the connections die with it, so the CP cannot synthesize +//! anything: in-flight delegations end as initiator-side timeouts against +//! the already-propagated deadline. Late `cp/delegate_result` frames for +//! unknown ids are acknowledged and dropped (logged), so reconnecting +//! runtimes do not error-loop. //! - **Saturation** — routing never queues; `SATURATED` is returned //! immediately (fast-fail, no hidden buffer). //! @@ -76,6 +81,7 @@ //! total and acyclic. use std::collections::BTreeMap; +use std::sync::atomic::{AtomicBool, Ordering}; use chrono::{DateTime, Utc}; use parking_lot::Mutex; @@ -198,6 +204,12 @@ pub struct Router { /// would disclose other namespaces' delegation volume. Commit matching /// only requires never-reuse per `(namespace, delegation_id)` key. admission: Mutex>, + /// Latched by [`Router::fail_all`] before the table is drained; checked + /// under the `inflight` lock at admission insert so no delegation can be + /// admitted after the drain snapshot — without it a late admission would + /// be accepted but never terminated (no entry left for `fail_all`, a + /// dangling `requested` for observers, a leaked reservation). + shutting_down: AtomicBool, } pub enum DelegateOutcome { @@ -385,6 +397,7 @@ impl Router { Self { inflight: Mutex::new(BTreeMap::new()), admission: Mutex::new(BTreeMap::new()), + shutting_down: AtomicBool::new(false), } } @@ -659,7 +672,26 @@ impl Router { // never-announced (no observer terminal, no synthesized frames). announced: false, }; - self.inflight.lock().insert(key.clone(), entry.clone()); + // The insert and the shutdown check share one `inflight` critical + // section with `fail_all`'s drain: `fail_all` latches + // `shutting_down` before taking the lock, so an insert that reads a + // clear latch also committed before the take — `fail_all` removes + // the entry and the `Forward::Gone` path below reports the loss — + // while one that reads the latched flag is refused here, before the + // admission can outlive the in-flight table itself. + { + let mut g = self.inflight.lock(); + if self.shutting_down.load(Ordering::SeqCst) { + // Roll back the capacity reservation: nothing was inserted, + // nothing was announced, and no terminal is owed. + registry.adjust_sessions(target.handle, -1); + return DelegateOutcome::Rejected(ErrorObject::new( + codes::SATURATED, + "control plane shutting down — not accepting new delegations", + )); + } + g.insert(key.clone(), entry.clone()); + } // Admission is committed: the token is minted, capacity is reserved, // and the entry is inserted — the duplicate check rides the in-flight @@ -1379,6 +1411,94 @@ impl Router { affected } + /// Graceful-shutdown drain: end EVERY in-flight delegation at once. + /// Both peers are still connected while the process drains, so unlike + /// the per-instance failure paths each initiator gets a synthesized + /// `target_disconnected` terminal AND each serving runtime a best-effort + /// `cp/cancel` — the same dual emission `sweep_deadlines` uses for the + /// case where both peers live. `reason` marks the frames as a CP + /// shutdown diagnostic rather than an ordinary peer disconnect. + /// + /// Same table discipline as `fail_instance`: an entry never announced + /// gets no terminal and no frames (its reservation still releases). A + /// delegation inserted between the drain below and its own announce + /// section finds no entry on re-check, emits nothing, and reports the + /// loss to its caller — the `Forward::Gone` path. + /// + /// The `shutting_down` latch is set BEFORE the drain so `delegate`'s + /// insert — checked under the same lock — can never land an entry after + /// this take: either the insert runs first and is drained here, or it + /// reads the latch and rejects with `SATURATED`. No third interleaving + /// exists, so no admission can be accepted-and-never-terminated. + pub fn fail_all( + &self, + registry: &Registry, + events: &EventHub, + reason: &str, + rpc_id: &mut impl FnMut() -> u64, + ) -> Vec<(Instance, String)> { + self.shutting_down.store(true, Ordering::SeqCst); + let entries: Vec = { + let mut g = self.inflight.lock(); + std::mem::take(&mut *g).into_values().collect() + }; + let mut frames = Vec::new(); + for e in entries { + // The serving instance's slot was reserved at admission and is + // always released — announced or not. + registry.adjust_sessions(e.to_handle, -1); + warn!( + delegation = %e.delegation_id, + announced = e.announced, + "in-flight delegation ended by control-plane shutdown" + ); + if !e.announced { + continue; + } + // Completion-shaped, and it must agree with the initiator-bound + // terminal below: one announced admission gets exactly one + // terminal on each surface. + events.emit( + registry, + &e.namespace, + CpEvent::DelegationCompleted { + delegation_id: e.delegation_id.clone(), + admission: e.generation, + from: e.from_logical.clone(), + to: e.to_logical.clone(), + status: DelegationStatus::TargetDisconnected, + result_excerpt: None, + error: Some(events.cp_diagnostic(reason)), + }, + ); + if let Some(init) = registry.get(e.from_handle) { + let params = DelegateResultParams { + delegation_id: e.delegation_id.clone(), + admission: e.generation, + status: DelegationStatus::TargetDisconnected, + result: None, + error: Some(reason.to_string()), + }; + if let Some(text) = synthesized_frame(rpc_id(), methods::DELEGATE_RESULT, ¶ms) { + frames.push((init, text)); + } + } + if let Some(target) = registry.get(e.to_handle) { + let params = CancelParams { + delegation_id: e.delegation_id.clone(), + // Names the admission that is over, so the frame can + // overtake nothing — the connection is closing anyway. + admission: e.generation, + reason: reason.to_string(), + }; + if let Some(text) = synthesized_frame(rpc_id(), methods::CANCEL, ¶ms) { + frames.push((target, text)); + } + } + } + frames + } + /// Deadline sweep: expire overdue delegations. Returns frames to deliver /// (timeout result to the initiator, best-effort cancel to the server). pub fn sweep_deadlines( @@ -5103,6 +5223,118 @@ metadata_only = true .contains("initiator prod/koudu disconnected")); } + #[test] + fn fail_all_tells_both_sides_and_releases_capacity() { + // Graceful shutdown differs from a single disconnect: both peers are + // still alive while the CP drains, so the initiator gets a + // `target_disconnected` terminal AND the serving runtime a cancel — + // the dual emission the deadline sweep performs. + let mut w = world(); + let mut lobby = observe(&w, "prod"); + let a = accept(do_delegate(&w, delegate_params("d-1", "worker-1", 60))); + w.worker_rx.try_recv().unwrap(); // the forward itself + assert_eq!( + w.registry.get(w.h_worker).unwrap().active_sessions, + 1, + "the admission reserved a serving slot" + ); + + let mut id = 0u64; + let mut next = || { + id += 1; + id + }; + let frames = w.router.fail_all( + &w.registry, + &w.events, + "control plane shutting down", + &mut next, + ); + + assert_eq!(w.router.inflight_count(), 0, "the table is empty"); + assert_eq!( + w.registry.get(w.h_worker).unwrap().active_sessions, + 0, + "the reserved slot is released even though both peers still exist" + ); + + let init_frames: Vec<&String> = frames + .iter() + .filter(|(i, _)| i.handle == w.h_primary) + .map(|(_, t)| t) + .collect(); + let worker_frames: Vec<&String> = frames + .iter() + .filter(|(i, _)| i.handle == w.h_worker) + .map(|(_, t)| t) + .collect(); + assert_eq!(init_frames.len(), 1, "one terminal for the initiator"); + assert_eq!(worker_frames.len(), 1, "one cancel for the worker"); + + let terminal: serde_json::Value = serde_json::from_str(init_frames[0]).unwrap(); + assert_eq!(terminal["method"], "cp/delegate_result"); + assert_eq!(terminal["params"]["status"], "target_disconnected"); + assert_eq!(terminal["params"]["admission"].as_u64(), Some(a.admission)); + assert_eq!(terminal["params"]["error"], "control plane shutting down"); + + let cancel: serde_json::Value = serde_json::from_str(worker_frames[0]).unwrap(); + assert_eq!(cancel["method"], "cp/cancel"); + assert_eq!(cancel["params"]["delegation_id"], "d-1"); + assert_eq!(cancel["params"]["admission"].as_u64(), Some(a.admission)); + + // The observer sees the terminal, completion-shaped and agreeing + // with the initiator's frame (the terminal contract's "one terminal + // per announced admission, agreeing on both surfaces"). + let ev = events_of(&mut lobby); + assert_eq!(ev.len(), 2, "requested + completed"); + assert_eq!(ev[1]["event"], "delegation_completed"); + assert_eq!(ev[1]["status"], "target_disconnected"); + assert_eq!(ev[1]["admission"].as_u64(), Some(a.admission)); + } + + #[test] + fn delegate_after_fail_all_is_refused_instead_of_orphaned() { + // The shutdown latch closes the admit-after-drain race: a delegation + // arriving after `fail_all` must be refused (SATURATED — back off and + // retry), never accepted into a table nobody will drain again. + let mut w = world(); + let mut lobby = observe(&w, "prod"); + let mut id = 0u64; + let mut next = || { + id += 1; + id + }; + w.router.fail_all( + &w.registry, + &w.events, + "control plane shutting down", + &mut next, + ); + + assert!(matches!( + do_delegate(&w, delegate_params("d-late", "worker-1", 60)), + DelegateOutcome::Rejected(e) if e.code == codes::SATURATED + )); + assert_eq!( + w.router.inflight_count(), + 0, + "no entry may be inserted once the drain latch is set" + ); + assert_eq!( + w.registry.get(w.h_worker).unwrap().active_sessions, + 0, + "a refused admission must not leak the capacity reservation" + ); + assert!( + w.worker_rx.try_recv().is_err(), + "the refused admission must not forward anything downstream" + ); + assert!( + events_of(&mut lobby).is_empty(), + "no requested/terminal may be emitted for a refused admission" + ); + } + #[test] fn metadata_only_namespace_omits_prompt_and_result_excerpts() { let w = world_with_cfg( diff --git a/crates/openab-cp/src/server.rs b/crates/openab-cp/src/server.rs index c442d6d39..1e5977e04 100644 --- a/crates/openab-cp/src/server.rs +++ b/crates/openab-cp/src/server.rs @@ -46,7 +46,7 @@ use crate::proto::{ DelegateResultParams, DeregisterReason, ErrorObject, JsonRpcErrorResponse, JsonRpcMessage, JsonRpcResponse, ListAgentsResult, RegisterAck, RegisterParams, PROTOCOL_VERSION, }; -use crate::registry::{outbound_channel, shutdown_signal, Instance, Registry}; +use crate::registry::{outbound_channel, shutdown_signal, Instance, Registry, ShutdownTx}; use crate::router::{CompleteOutcome, DelegateOutcome, Router}; pub struct AppState { @@ -59,6 +59,11 @@ pub struct AppState { /// Live connections per identity (`namespace/name`), counted from the /// upgrade so pre-registration sockets are bounded too. conns: Mutex>, + /// Process-wide shutdown broadcast, latched once by + /// [`graceful_shutdown`] on SIGTERM/SIGHUP/SIGINT. Every connection task + /// subscribes — including pre-registration sockets, which the + /// per-connection signal held by the registry cannot reach. + shutdown: ShutdownTx, } impl AppState { @@ -70,6 +75,7 @@ impl AppState { router: Router::new(), rpc_id: AtomicU64::new(1), conns: Mutex::new(BTreeMap::new()), + shutdown: shutdown_signal(), } } @@ -77,6 +83,29 @@ impl AppState { self.rpc_id.fetch_add(1, Ordering::Relaxed) } + /// Subscribe to the process-wide shutdown signal: resolves to the + /// shutdown reason once [`AppState::begin_shutdown`] latches it. + pub fn shutdown_rx(&self) -> watch::Receiver> { + self.shutdown.subscribe() + } + + /// Latch the shutdown reason: every live connection task closes, and a + /// connection upgraded mid-drain sees the latched reason at subscribe + /// time. + fn begin_shutdown(&self, reason: &'static str) { + self.shutdown.send_replace(Some(reason)); + } + + /// Resolves when no live connection permits remain — every connection + /// task (registered or still pre-registration) has exited. Polled: + /// drain progress is a wait on an ever-shrinking set, not a precise + /// signal. + pub async fn connections_drained(&self) { + while !self.conns.lock().is_empty() { + tokio::time::sleep(Duration::from_millis(10)).await; + } + } + /// Take a connection slot for `identity`, or `None` when the identity is /// already at `max_connections_per_identity`. The returned guard releases /// the slot on drop — including on every early return and on an upgrade @@ -151,6 +180,57 @@ fn policy_close(reason: &'static str) -> Message { })) } +/// Reason every connection gets on its close frame when the CP shuts down +/// (SIGTERM/SIGHUP/SIGINT), and carried on the terminal/cancel frames +/// synthesized for in-flight delegations so a runtime can tell "the CP is +/// going away" from "my peer disconnected". +pub const REASON_SHUTDOWN: &str = "control plane shutting down"; + +/// A CP-initiated close for process shutdown: code 1012 (service restart — +/// the client should reconnect) rather than the 1008 the misbehaviour +/// closes above use, because nothing about a deploy is a policy violation. +fn restart_close(reason: &'static str) -> Message { + Message::Close(Some(CloseFrame { + code: close_code::RESTART, + reason: reason.into(), + })) +} + +/// The two signals that can terminate a connection: the per-connection one +/// held by the registry entry (lease expiry, terminal-frame backpressure) +/// and the process-wide one latched on [`AppState`] when the CP begins +/// draining. The process signal is subscribed at upgrade time so +/// pre-registration sockets — which the registry cannot see yet — close +/// with a reason too. +struct CloseWatch { + conn: watch::Receiver>, + process: watch::Receiver>, +} + +impl CloseWatch { + /// Wait for the next transition on either signal. `watch` semantics: a + /// value already seen is not reported again, so a caller that has + /// consumed a reason keeps writing its close frame undisturbed. + async fn changed(&mut self) { + tokio::select! { + _ = self.conn.changed() => {} + _ = self.process.changed() => {} + } + } + + /// The process-wide shutdown reason, when the CP is draining. + fn process_reason(&self) -> Option<&'static str> { + *self.process.borrow() + } + + /// The close reason set on either signal, if any; the process reason + /// wins because "the CP is going away" is the actionable fact when both + /// are latched (e.g. a lease sweep racing the drain). + fn reason(&self) -> Option<&'static str> { + (*self.process.borrow()).or(*self.conn.borrow()) + } +} + /// Why a bounded write did not complete. Either way the connection ends. enum WriteStop { /// Transport error, or the peer did not accept the frame within @@ -178,19 +258,27 @@ enum WriteStop { /// A cancelled or timed-out write can leave a partially written frame on the /// wire; that is acceptable precisely because both outcomes end the /// connection (the caller breaks to teardown, dropping the socket). +/// +/// `biased` polling order matters for the drain path: the caller has +/// already dequeued `msg` from the outbound channel, so a close signal that +/// wins a fair race would destroy a frame that was ready to write. The send +/// is tried first — an immediately-writable socket completes it and the +/// close signal is observed on the next call — while a write that cannot +/// finish now still pends and is interrupted by the signal or the timeout. async fn send_bounded( sink: &mut futures_util::stream::SplitSink, msg: Message, write_timeout: Duration, - shutdown_rx: &mut watch::Receiver>, + close: &mut CloseWatch, ) -> Result<(), WriteStop> { tokio::select! { - _ = shutdown_rx.changed() => Err(WriteStop::Shutdown(*shutdown_rx.borrow_and_update())), + biased; sent = tokio::time::timeout(write_timeout, sink.send(msg)) => match sent { Ok(Ok(())) => Ok(()), Ok(Err(_)) => Err(WriteStop::Disconnected), Err(_) => Err(WriteStop::Disconnected), }, + _ = close.changed() => Err(WriteStop::Shutdown(close.reason())), } } @@ -249,35 +337,62 @@ async fn handle_connection( let (mut sink, mut stream) = socket.split(); let write_timeout = Duration::from_secs(state.cfg.write_timeout_secs); - // Shutdown signal so the CP can close this socket when it drops the - // registration on its own initiative (lease expiry) or must terminate the - // connection (terminal-frame backpressure). + // Shutdown signals so the CP can close this socket when it drops the + // registration on its own initiative (lease expiry), must terminate the + // connection (terminal-frame backpressure), or the whole process is + // draining (SIGTERM/SIGHUP/SIGINT via `graceful_shutdown`). // - // Created BEFORE the registration read — and therefore before the registry - // ever holds a clone — so no signal can be missed, and every write in this - // task, registration-phase writes included, can be raced against it. Kept + // The per-connection sender is created BEFORE the registration read — + // and therefore before the registry ever holds a clone — so no signal + // can be missed, and every write in this task, registration-phase writes + // included, can be raced against either signal. The senders are kept // alive here for the whole connection: closing is driven by an explicit // signal, never by the registry happening to drop its side. let shutdown = shutdown_signal(); - let mut shutdown_rx = shutdown.subscribe(); + let mut close = CloseWatch { + conn: shutdown.subscribe(), + process: state.shutdown_rx(), + }; + + // A socket that finished its upgrade while the CP was already draining + // must not register into a dying process: close it immediately. `watch` + // only reports transitions that postdate the subscription, so the + // already-latched reason is checked once explicitly. + if let Some(reason) = close.process_reason() { + let _ = send_bounded(&mut sink, restart_close(reason), write_timeout, &mut close).await; + return; + } // --- Registration: mandatory first frame, within a deadline --- // An authenticated peer must not be able to park idle sockets: pings keep // the transport alive but do not extend this deadline. - let register = match tokio::time::timeout( - Duration::from_secs(state.cfg.register_timeout_secs), - async { - loop { - match stream.next().await { - Some(Ok(Message::Text(text))) => return Some(text), - Some(Ok(Message::Ping(_) | Message::Pong(_))) => continue, - _ => return None, + let register = match tokio::select! { + res = tokio::time::timeout( + Duration::from_secs(state.cfg.register_timeout_secs), + async { + loop { + match stream.next().await { + Some(Ok(Message::Text(text))) => return Some(text), + Some(Ok(Message::Ping(_) | Message::Pong(_))) => continue, + _ => return None, + } } - } - }, - ) - .await - { + }, + ) => res, + _ = close.changed() => { + // Shutdown while still pre-registration: no registry entry or + // in-flight delegation exists to tear down — just tell the peer + // why its socket is closing. Only the process signal can latch + // before registration (the per-connection sender has no owner + // yet), but the frame is still picked by which signal fired. + let frame = match close.process_reason() { + Some(reason) => restart_close(reason), + None => policy_close(close.reason().unwrap_or(REASON_SHUTDOWN)), + }; + let _ = send_bounded(&mut sink, frame, write_timeout, &mut close).await; + return; + } + } { Ok(Some(text)) => text, Ok(None) => { warn!(agent = %identity.name, "connection closed before registration"); @@ -293,7 +408,7 @@ async fn handle_connection( &mut sink, policy_close(REASON_REGISTER_TIMEOUT), write_timeout, - &mut shutdown_rx, + &mut close, ) .await; return; @@ -307,7 +422,7 @@ async fn handle_connection( &mut sink, Message::Text(serde_json::to_string(&resp).expect("serializable").into()), write_timeout, - &mut shutdown_rx, + &mut close, ) .await; return; @@ -371,7 +486,7 @@ async fn handle_connection( &mut sink, Message::Text(serde_json::to_string(&resp).expect("serializable").into()), write_timeout, - &mut shutdown_rx, + &mut close, ) .await; return; @@ -413,21 +528,22 @@ async fn handle_connection( &mut sink, Message::Text(serde_json::to_string(&resp).expect("serializable").into()), write_timeout, - &mut shutdown_rx, + &mut close, ) .await { // CP-initiated closes carry meaning even here: if the CP signalled // this connection while the ack was in flight, still tell the client - // why before tearing down. - if let WriteStop::Shutdown(Some(reason)) = stop { - let _ = send_bounded( - &mut sink, - policy_close(reason), - write_timeout, - &mut shutdown_rx, - ) - .await; + // why before tearing down — a restart frame when the whole process + // is draining, a policy frame when this connection alone is ending. + if let WriteStop::Shutdown(_) = stop { + let frame = match close.process_reason() { + Some(reason) => Some(restart_close(reason)), + None => close.reason().map(policy_close), + }; + if let Some(frame) = frame { + let _ = send_bounded(&mut sink, frame, write_timeout, &mut close).await; + } } // `_registered` runs teardown on the way out. return; @@ -451,7 +567,7 @@ async fn handle_connection( // duration of its own arm body, and `break` acts on the loop below. macro_rules! write_or_break { ($msg:expr) => { - match send_bounded(&mut sink, $msg, write_timeout, &mut shutdown_rx).await { + match send_bounded(&mut sink, $msg, write_timeout, &mut close).await { Ok(()) => {} Err(WriteStop::Shutdown(reason)) => { cp_close_reason = reason; @@ -471,15 +587,15 @@ async fn handle_connection( } loop { tokio::select! { - // The CP dropped this registration (lease expiry) or must - // terminate the connection (terminal-frame backpressure): the - // socket must go too. Keeping it open would leave a - // connection whose every frame hits an absent registry entry and - // which can never re-register, since registration is - // first-frame-only. Closing lets the client reconnect, - // re-authenticate, and register again. - _ = shutdown_rx.changed() => { - cp_close_reason = *shutdown_rx.borrow_and_update(); + // The CP dropped this registration (lease expiry), must + // terminate the connection (terminal-frame backpressure), or the + // process is shutting down: the socket must go too. Keeping it + // open would leave a connection whose every frame hits an absent + // registry entry and which can never re-register, since + // registration is first-frame-only. Closing lets the client + // reconnect, re-authenticate, and register again. + _ = close.changed() => { + cp_close_reason = close.reason(); break; } outbound = rx.recv() => { @@ -507,7 +623,39 @@ async fn handle_connection( } } - if let Some(reason) = cp_close_reason { + // Process shutdown — SIGTERM/SIGHUP/SIGINT — is distinguished from a + // per-connection termination by which signal latched. On the process + // path the outbound queue is flushed FIRST: the synthesized terminals + // and cancels that `graceful_shutdown` queued must reach the wire ahead + // of the close frame. `rx.close()` refuses any further enqueue, so the + // loop ends on its own once the backlog is flushed; the drain deadline + // bounds the backlog and `write_timeout` bounds each stalled write. + if let Some(reason) = close.process_reason() { + rx.close(); + let drain_deadline = Instant::now() + Duration::from_secs(state.cfg.shutdown_drain_secs); + while let Some(text) = rx.recv().await { + let budget = drain_deadline.saturating_duration_since(Instant::now()); + if budget.is_zero() { + break; + } + if !matches!( + tokio::time::timeout(budget.min(write_timeout), sink.send(text.into())).await, + Ok(Ok(())) + ) { + return; // the peer is gone — teardown runs via the guard + } + } + info!( + agent = %format!("{}/{}", identity.namespace, identity.name), + handle, + reason, + "closing connection: control plane shutting down" + ); + // Bounded like every other write: a peer that has stopped reading + // must not be able to hold teardown (and its quota slot) by refusing + // to accept the close frame. + let _ = send_bounded(&mut sink, restart_close(reason), write_timeout, &mut close).await; + } else if let Some(reason) = cp_close_reason { info!( agent = %format!("{}/{}", identity.namespace, identity.name), handle, @@ -517,13 +665,7 @@ async fn handle_connection( // Bounded like every other write: a peer that has stopped reading must // not be able to delay teardown (and its quota slot) by refusing to // accept the close frame. - let _ = send_bounded( - &mut sink, - policy_close(reason), - write_timeout, - &mut shutdown_rx, - ) - .await; + let _ = send_bounded(&mut sink, policy_close(reason), write_timeout, &mut close).await; } // Teardown runs here, when `_registered` drops — on this path and on an @@ -1008,6 +1150,48 @@ pub async fn run_sweeper(state: Arc) { } } +/// Graceful shutdown (SIGTERM/SIGHUP/SIGINT): resolve every in-flight +/// delegation with synthesized terminal frames, signal every connection to +/// close, then wait — bounded by `shutdown_drain_secs` — for the connection +/// tasks to flush their queues and exit. +/// +/// This runs AFTER the HTTP listener has stopped accepting (axum's graceful +/// shutdown already resolved); the WebSocket connection tasks are detached +/// spawns that keep running until told to close. The caller stops the +/// sweeper and exits after this returns — the sweeper deliberately outlives +/// the drain so a mid-drain lease/deadline expiry still synthesizes +/// normally. +/// +/// Order is deliberate: terminals are synthesized BEFORE the close signal +/// so they are already on each connection's outbound queue when it starts +/// draining. Signalling first would let tasks close ahead of the frames +/// they were about to deliver, and "synthesize where possible" would +/// degrade to "almost never". +pub async fn graceful_shutdown(state: &Arc) { + info!("draining connections and resolving in-flight delegations"); + let mut next = || state.next_rpc_id(); + for (inst, frame) in + state + .router + .fail_all(&state.registry, &state.events, REASON_SHUTDOWN, &mut next) + { + let _ = inst.tx.try_send(frame); + } + state.begin_shutdown(REASON_SHUTDOWN); + match tokio::time::timeout( + Duration::from_secs(state.cfg.shutdown_drain_secs), + state.connections_drained(), + ) + .await + { + Ok(()) => info!("all connections drained"), + Err(_) => warn!( + drain_secs = state.cfg.shutdown_drain_secs, + "shutdown drain budget elapsed — exiting with connections still open" + ), + } +} + #[cfg(test)] mod tests { use super::*; diff --git a/crates/openab-cp/tests/ws_lifecycle.rs b/crates/openab-cp/tests/ws_lifecycle.rs index d2ebabe74..baac339cf 100644 --- a/crates/openab-cp/tests/ws_lifecycle.rs +++ b/crates/openab-cp/tests/ws_lifecycle.rs @@ -17,7 +17,7 @@ use tokio_tungstenite::tungstenite::{Error as WsError, Message}; use tokio_tungstenite::{MaybeTlsStream, WebSocketStream}; use openab_cp::config::CpConfig; -use openab_cp::server::{app, sweep_leases, AppState}; +use openab_cp::server::{app, graceful_shutdown, sweep_leases, AppState, REASON_SHUTDOWN}; const KEY: &str = "k-primary"; const KEY_WORKER: &str = "k-worker"; @@ -744,3 +744,201 @@ max_outbound_queue_bytes = 67108864" "the released slot must be reusable by the same identity" ); } + +// --- Graceful shutdown: SIGTERM/SIGHUP/SIGINT drain --- + +/// The close every connection must see on CP shutdown: 1012 (service +/// restart — the client should reconnect) carrying the shutdown reason, +/// distinct from the 1008 policy closes the CP sends for misbehaviour. +fn shutdown_close() -> Closed { + Closed::Frame { + code: 1012, + reason: REASON_SHUTDOWN.to_string(), + } +} + +#[tokio::test] +async fn shutdown_synthesizes_terminals_then_closes_with_restart() { + // The container-stop path: every in-flight delegation gets a synthesized + // terminal — the initiator sees `target_disconnected` and the serving + // runtime a `cp/cancel` — riding out BEFORE the sockets close, and each + // socket closes with a 1012 frame naming the reason instead of a bare + // TCP reset. + let (state, url) = spawn_cp(cfg( + "register_timeout_secs = 30\nmax_connections_per_identity = 2\nshutdown_drain_secs = 5", + )) + .await; + + let mut initiator = connect(&url).await.expect("initiator accepted"); + assert_eq!( + register(&mut initiator, "i-1").await["result"]["protocol_version"], + 1 + ); + let mut worker = connect_as(&url, KEY_WORKER).await.expect("worker accepted"); + assert_eq!( + register_worker(&mut worker, "i-w", 2).await["result"]["protocol_version"], + 1 + ); + + let deadline = (chrono::Utc::now() + chrono::Duration::seconds(600)).to_rfc3339(); + let delegate = serde_json::json!({ + "jsonrpc": "2.0", "id": 10, "method": "cp/delegate", + "params": { + "delegation_id": "d-1", + "target": {"name": "worker-1"}, + "prompt": "do it", + "deadline": deadline + } + }) + .to_string(); + initiator + .send(Message::Text(delegate.into())) + .await + .unwrap(); + let ack = await_frame(&mut initiator, "delegate ack", |v| v["id"] == 10).await; + let admission = ack["result"]["admission"] + .as_u64() + .expect("ack must carry the admission token"); + await_frame(&mut worker, "forwarded delegate", |v| { + v["method"] == "cp/delegate" + }) + .await; + assert_eq!(state.router.inflight_count(), 1, "delegation in flight"); + + graceful_shutdown(&state).await; + + // The synthesized terminal reaches the initiator ahead of the close. + let terminal = await_frame(&mut initiator, "synthesized terminal", |v| { + v["method"] == "cp/delegate_result" + }) + .await; + assert_eq!(terminal["params"]["delegation_id"], "d-1"); + assert_eq!( + terminal["params"]["admission"].as_u64(), + Some(admission), + "the terminal must name the admission it ends" + ); + assert_eq!(terminal["params"]["status"], "target_disconnected"); + assert_eq!(terminal["params"]["error"], REASON_SHUTDOWN); + + // The serving runtime is told to stop working on the same admission. + let cancel = await_frame(&mut worker, "synthesized cancel", |v| { + v["method"] == "cp/cancel" + }) + .await; + assert_eq!(cancel["params"]["delegation_id"], "d-1"); + assert_eq!(cancel["params"]["admission"].as_u64(), Some(admission)); + assert_eq!(cancel["params"]["reason"], REASON_SHUTDOWN); + + // Then every connection gets the restart close, not a TCP reset. + for (ws, who) in [(&mut initiator, "initiator"), (&mut worker, "worker")] { + let closed = wait_closed(ws, Duration::from_secs(5)) + .await + .unwrap_or_else(|| panic!("{who} must observe a close")); + assert_eq!( + closed, + shutdown_close(), + "{who}: shutdown must close with 1012 + reason, got {closed:?}" + ); + } + + assert_eq!(state.router.inflight_count(), 0); + assert!( + state.registry.list("prod").is_empty(), + "every registration is torn down" + ); +} + +#[tokio::test] +async fn shutdown_closes_sockets_that_never_registered() { + // The per-connection close signal lives in the registry, so a socket + // that has not completed `cp/register` is invisible to it. Shutdown must + // still reach it: the process-wide signal is what every connection task + // subscribes to from the upgrade on. + let (state, url) = spawn_cp(cfg("register_timeout_secs = 30")).await; + let mut ws = connect(&url).await.expect("connection accepted"); + + graceful_shutdown(&state).await; + + let closed = wait_closed(&mut ws, Duration::from_secs(5)).await; + assert_eq!( + closed, + Some(shutdown_close()), + "a parked pre-registration socket must get the shutdown close, got {closed:?}" + ); +} + +/// Spawn a real `openab-cp` process bound to an ephemeral port; returns the +/// child, its WS URL, and the config path (for cleanup). This is the only +/// proof the OS-level path works: `docker stop`, ECS, and k8s all deliver +/// SIGTERM before SIGKILL, and the binary used to die by default disposition +/// without ever running a shutdown path. +#[cfg(unix)] +async fn spawn_cp_process() -> (tokio::process::Child, String, std::path::PathBuf) { + // Grab an ephemeral port, release it, hand it to the child — the usual + // tiny race, absorbed by `connect_retry`. + let port = std::net::TcpListener::bind("127.0.0.1:0") + .expect("bind ephemeral") + .local_addr() + .unwrap() + .port(); + let cfg_path = std::env::temp_dir().join(format!( + "openab-cp-shutdown-{}-{port}.toml", + std::process::id() + )); + std::fs::write( + &cfg_path, + format!( + "listen = \"127.0.0.1:{port}\"\n\n\ + [[agents]]\nkey = \"{KEY}\"\nnamespace = \"prod\"\nname = \"koudu\"\ntype = \"primary\"\n\n\ + [[agents]]\nkey = \"{KEY_WORKER}\"\nnamespace = \"prod\"\nname = \"worker-1\"\ntype = \"worker\"\n" + ), + ) + .expect("write cp config"); + let child = tokio::process::Command::new(env!("CARGO_BIN_EXE_openab-cp")) + .arg("--config") + .arg(&cfg_path) + .kill_on_drop(true) + .spawn() + .expect("spawn openab-cp"); + (child, format!("ws://127.0.0.1:{port}/cp"), cfg_path) +} + +#[cfg(unix)] +#[tokio::test] +async fn sigterm_and_sighup_run_the_drain_and_exit_cleanly() { + // SIGTERM is what every containerized deployment actually sends; + // SIGHUP shares the path. Each must produce the graceful close on live + // sockets and a clean exit code — never the old default-disposition kill. + for sig in ["TERM", "HUP"] { + let (mut child, url, cfg_path) = spawn_cp_process().await; + let mut ws = connect_retry(&url).await; + assert_eq!( + register(&mut ws, "i-1").await["result"]["protocol_version"], + 1 + ); + + let status = std::process::Command::new("kill") + .arg(format!("-{sig}")) + .arg(child.id().expect("child pid").to_string()) + .status() + .expect("kill must run"); + assert!(status.success(), "kill -{sig} failed: {status}"); + + let closed = wait_closed(&mut ws, Duration::from_secs(10)).await; + assert_eq!( + closed, + Some(shutdown_close()), + "SIG{sig}: the socket must see a graceful close, got {closed:?}" + ); + let status = tokio::time::timeout(Duration::from_secs(10), child.wait()) + .await + .unwrap_or_else(|_| panic!("SIG{sig}: the CP must exit")) + .expect("wait must succeed"); + assert!( + status.success(), + "SIG{sig}: the CP must exit cleanly, got {status}" + ); + let _ = std::fs::remove_file(&cfg_path); + } +} diff --git a/docs/adr/agent-control-plane.md b/docs/adr/agent-control-plane.md index f5a5195d0..ae6a46525 100644 --- a/docs/adr/agent-control-plane.md +++ b/docs/adr/agent-control-plane.md @@ -546,10 +546,17 @@ recovery semantics: instance's live id return the same `POLICY_DENIED` error object, byte for byte, so cancel cannot be used as an existence oracle for other tenants' delegation ids. The CP's own logs keep the distinction. -- **CP restart semantics.** A CP restart is equivalent to every lease - expiring at once *with* the connection closure that implies — except that - the CP is not there to send it: the in-flight table and the sockets die - together with the process, so no synthesized `timeout` or +- **CP restart semantics.** A *graceful* CP stop (SIGTERM/SIGHUP/SIGINT) + drains first: every in-flight delegation is resolved — initiators get a + synthesized `target_disconnected` terminal and serving runtimes a + `cp/cancel`, both stamped `control plane shutting down`, and admissions + racing the drain are refused `SATURATED` — then every connection closes + with code 1012, all inside `shutdown_drain_secs` plus at most one + `write_timeout_secs` for a stalled close write. A + *hard* restart (SIGKILL, crash, drain budget exhausted) is equivalent to + every lease expiring at once *with* the connection closure that implies — + except that the CP is not there to send it: the in-flight table and the + sockets die together with the process, so no synthesized `timeout` or `target_disconnected` frame can be emitted for delegations that were in flight. Runtimes observe the transport drop, reconnect with backoff, and re-register (new handles, empty in-flight table). Initiators reconcile diff --git a/docs/control-plane.md b/docs/control-plane.md index e29564dc9..d1aaaeda1 100644 --- a/docs/control-plane.md +++ b/docs/control-plane.md @@ -56,7 +56,10 @@ issue #1474). - CP-initiated closes use WS code 1008 with a reason: `registration timeout`, `lease expired`, or `outbound queue overflow`. On any of these, - reconnect, re-authenticate, and re-register. + reconnect, re-authenticate, and re-register. Process shutdown (SIGTERM, + SIGHUP, SIGINT — e.g. `docker stop`, ECS, k8s) uses code **1012** + (service restart) with reason `control plane shutting down` instead: same + recovery, different signal — your client did nothing wrong. - A peer that stops reading is disconnected: any single outbound write that blocks longer than `write_timeout_secs` is treated as a dead peer, so keep draining the socket even while busy. The same rule applies to the queue @@ -121,7 +124,14 @@ issue #1474). the CP (`default_max_delegated_sessions_cap`, or a per-identity override). The ack's `effective_max_delegated_sessions` is the value that counts. - After a lease expires or the CP restarts, in-flight delegations are gone: - initiators reconcile against their own deadlines and re-delegate. + initiators reconcile against their own deadlines and re-delegate. On a + *graceful* stop the CP resolves them first: within `shutdown_drain_secs` + (plus at most one `write_timeout_secs` for a stalled close write) every + initiator gets a synthesized `target_disconnected` terminal (error + `control plane shutting down`) and every serving runtime a `cp/cancel`, + ahead of the 1012 close; a `cp/delegate` that reaches the CP mid-drain is + refused `SATURATED` — back off and retry after reconnecting. A SIGKILL + still leaves reconciliation to your deadlines. ## Observer surface (lobby) From 6ca7417f79c88f8d0468343c8cdfd91d87e821a4 Mon Sep 17 00:00:00 2001 From: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 1 Oct 2026 01:36:36 +0800 Subject: [PATCH 03/13] style: normalize workspace formatting under rustfmt 1.98 Whitespace-only `cargo fmt --all` pass over files outside the feature change: the workspace predates the current stable rustfmt, so `cargo fmt --all -- --check` fails on ~36 unrelated files and blocks the fmt gate for any change. Refs #1475 Generated with [Devin](https://devin.ai) Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- crates/openab-core/src/acp/agentcore.rs | 34 +- crates/openab-core/src/acp/pool.rs | 123 +- crates/openab-core/src/acp/protocol.rs | 22 +- crates/openab-core/src/acp_mcp.rs | 77 +- crates/openab-core/src/adapter.rs | 39 +- crates/openab-core/src/ambient.rs | 33 +- crates/openab-core/src/config.rs | 7 +- crates/openab-core/src/cron.rs | 20 +- crates/openab-core/src/discord.rs | 492 ++-- crates/openab-core/src/filestore.rs | 132 +- crates/openab-core/src/format.rs | 6 +- crates/openab-core/src/gateway.rs | 211 +- crates/openab-core/src/hooks.rs | 11 +- crates/openab-core/src/lib.rs | 12 +- crates/openab-core/src/media.rs | 81 +- crates/openab-core/src/multibot_cache.rs | 4 +- crates/openab-core/src/redact.rs | 34 +- crates/openab-core/src/remind.rs | 51 +- crates/openab-core/src/secrets.rs | 57 +- crates/openab-core/src/slack.rs | 245 +- crates/openab-core/src/trust.rs | 55 +- .../openab-gateway/src/adapters/acp_server.rs | 2012 ++++++++++++----- crates/openab-gateway/src/adapters/feishu.rs | 685 ++++-- .../src/adapters/feishu_card.rs | 32 +- .../openab-gateway/src/adapters/googlechat.rs | 295 ++- .../openab-gateway/src/adapters/lineworks.rs | 3 +- crates/openab-gateway/src/adapters/mod.rs | 26 +- crates/openab-gateway/src/adapters/teams.rs | 16 +- .../openab-gateway/src/adapters/telegram.rs | 303 ++- crates/openab-gateway/src/adapters/wecom.rs | 178 +- crates/openab-gateway/src/lib.rs | 86 +- crates/openab-gateway/src/store.rs | 6 +- crates/openab-mcp/src/mcp/runtime.rs | 15 +- src/acp_tunnel_source.rs | 131 +- src/ctl.rs | 67 +- src/main.rs | 131 +- 36 files changed, 3960 insertions(+), 1772 deletions(-) diff --git a/crates/openab-core/src/acp/agentcore.rs b/crates/openab-core/src/acp/agentcore.rs index f7a923e09..d0f847644 100644 --- a/crates/openab-core/src/acp/agentcore.rs +++ b/crates/openab-core/src/acp/agentcore.rs @@ -271,9 +271,12 @@ where .as_millis(); let runtime_sid = format!("oab-reconnect-{ts:020}-{ts:013x}"); match self.open_shell(&runtime_sid).await { - Ok(handle) => { self.sessions.insert(acp_sid.clone(), handle); } + Ok(handle) => { + self.sessions.insert(acp_sid.clone(), handle); + } Err(e) => { - self.write_error(id, -32000, &format!("reconnect failed: {e}")).await?; + self.write_error(id, -32000, &format!("reconnect failed: {e}")) + .await?; return Ok(()); } } @@ -281,7 +284,9 @@ where // Allocate ID before borrowing sessions let kiro_id = self.alloc_id(); - let kiro_sid = self.sessions.get(&acp_sid) + let kiro_sid = self + .sessions + .get(&acp_sid) .map(|s| s.kiro_session_id.clone()) .unwrap_or_default(); let mut fwd_params = params.clone(); @@ -425,7 +430,8 @@ where } async fn open_shell(&self, session_id: &str) -> Result { - let (request, host) = build_signed_request(&self.runtime_arn, session_id, &self.region).await?; + let (request, host) = + build_signed_request(&self.runtime_arn, session_id, &self.region).await?; // Manual TLS connection — gives us full control, avoids connect_async host override let tcp = tokio::net::TcpStream::connect(format!("{host}:443")) @@ -564,7 +570,9 @@ where Ok(Some(line)) => { // Check if this is the initialize response (has "id":0 or "id": 0) if let Ok(v) = serde_json::from_str::(&line) { - if v.get("id").and_then(|i| i.as_u64()) == Some(0) && v.get("result").is_some() { + if v.get("id").and_then(|i| i.as_u64()) == Some(0) + && v.get("result").is_some() + { info!(attempt, "agent initialized"); initialized = true; break; @@ -580,10 +588,14 @@ where } } } - if initialized { break; } + if initialized { + break; + } } if !initialized { - return Err(anyhow!("agent failed to respond to initialize after 5 attempts")); + return Err(anyhow!( + "agent failed to respond to initialize after 5 attempts" + )); } // Send session/new to kiro-cli to create a session @@ -616,7 +628,8 @@ where Ok(Some(line)) => { if let Ok(v) = serde_json::from_str::(&line) { if v.get("id").and_then(|i| i.as_u64()) == Some(1) { - sid = v.pointer("/result/sessionId") + sid = v + .pointer("/result/sessionId") .and_then(|s| s.as_str()) .unwrap_or("default") .to_string(); @@ -710,7 +723,10 @@ async fn build_signed_request( .header("connection", "Upgrade") .header("upgrade", "websocket") .header("sec-websocket-version", "13") - .header("sec-websocket-key", tokio_tungstenite::tungstenite::handshake::client::generate_key()); + .header( + "sec-websocket-key", + tokio_tungstenite::tungstenite::handshake::client::generate_key(), + ); // Add SigV4 auth headers (x-amz-date, authorization) for (name, value) in instructions.headers() { diff --git a/crates/openab-core/src/acp/pool.rs b/crates/openab-core/src/acp/pool.rs index 86b2ee989..7784b4e36 100644 --- a/crates/openab-core/src/acp/pool.rs +++ b/crates/openab-core/src/acp/pool.rs @@ -112,7 +112,12 @@ fn classify_hung( /// either resumes the session, so both are credentials. Extracted from the loop in `cleanup_idle` /// so the redaction can be exercised by a test for real — R1 redacted the sites it enumerated and /// this force-evict site was outside that list, logging both ids raw. -fn warn_force_evicting_hung(key: &str, session_id: Option<&str>, age_secs: u64, threshold_secs: u64) { +fn warn_force_evicting_hung( + key: &str, + session_id: Option<&str>, + age_secs: u64, + threshold_secs: u64, +) { warn!( thread_id = %crate::redact::redact_session_ids(key), session_id = %session_id.map(crate::redact::redact_session_ids).unwrap_or_default(), @@ -692,7 +697,12 @@ impl SessionPool { // supersedes under the same key (its guard cannot fire if that predecessor is hung). F3. #[cfg(feature = "acp-mcp")] if let Some(token) = session_token { - install_facade_token(&mut state, thread_id, token, self.session_registrar.as_ref()); + install_facade_token( + &mut state, + thread_id, + token, + self.session_registrar.as_ref(), + ); } self.save_mapping(&state.persisted); @@ -727,11 +737,12 @@ impl SessionPool { { let conn = { let state = self.state.read().await; - state - .active - .get(thread_id) - .cloned() - .ok_or_else(|| anyhow!("no connection for thread {}", crate::redact::redact_session_ids(thread_id)))? + state.active.get(thread_id).cloned().ok_or_else(|| { + anyhow!( + "no connection for thread {}", + crate::redact::redact_session_ids(thread_id) + ) + })? }; let mut conn = conn.lock().await; @@ -759,11 +770,12 @@ impl SessionPool { ) -> Result> { let conn = { let state = self.state.read().await; - state - .active - .get(thread_id) - .cloned() - .ok_or_else(|| anyhow!("no connection for thread {}", crate::redact::redact_session_ids(thread_id)))? + state.active.get(thread_id).cloned().ok_or_else(|| { + anyhow!( + "no connection for thread {}", + crate::redact::redact_session_ids(thread_id) + ) + })? }; let mut conn = conn.lock().await; conn.set_config_option(config_id, value).await @@ -775,11 +787,12 @@ impl SessionPool { pub async fn get_usage(&self, thread_id: &str) -> Result { let conn = { let state = self.state.read().await; - state - .active - .get(thread_id) - .cloned() - .ok_or_else(|| anyhow!("no connection for thread {}", crate::redact::redact_session_ids(thread_id)))? + state.active.get(thread_id).cloned().ok_or_else(|| { + anyhow!( + "no connection for thread {}", + crate::redact::redact_session_ids(thread_id) + ) + })? }; let mut conn = conn.lock().await; conn.get_usage().await @@ -794,7 +807,12 @@ impl SessionPool { .cancel_handles .get(thread_id) .cloned() - .ok_or_else(|| anyhow!("no session for thread {}", crate::redact::redact_session_ids(thread_id)))? + .ok_or_else(|| { + anyhow!( + "no session for thread {}", + crate::redact::redact_session_ids(thread_id) + ) + })? }; let data = serde_json::to_string(&serde_json::json!({ "jsonrpc": "2.0", @@ -852,7 +870,10 @@ impl SessionPool { info!(thread_id = %crate::redact::redact_session_ids(thread_id), "session reset"); Ok(()) } else { - Err(anyhow!("no session for thread {}", crate::redact::redact_session_ids(thread_id))) + Err(anyhow!( + "no session for thread {}", + crate::redact::redact_session_ids(thread_id) + )) } } @@ -1084,11 +1105,28 @@ mod tests { let mut state = empty_pool_state(); // Predecessor registers, then a successor takes over the SAME key. - super::install_facade_token(&mut state, "discord:acp_x", "T_pred".into(), Some(®istrar)); - assert!(reg.revoked().is_empty(), "nothing to revoke on the first install"); - super::install_facade_token(&mut state, "discord:acp_x", "T_succ".into(), Some(®istrar)); + super::install_facade_token( + &mut state, + "discord:acp_x", + "T_pred".into(), + Some(®istrar), + ); + assert!( + reg.revoked().is_empty(), + "nothing to revoke on the first install" + ); + super::install_facade_token( + &mut state, + "discord:acp_x", + "T_succ".into(), + Some(®istrar), + ); - assert_eq!(reg.revoked(), vec!["T_pred"], "the predecessor token must be revoked"); + assert_eq!( + reg.revoked(), + vec!["T_pred"], + "the predecessor token must be revoked" + ); assert_eq!( state.facade_tokens.get("discord:acp_x").map(String::as_str), Some("T_succ"), @@ -1105,14 +1143,25 @@ mod tests { let reg = Arc::new(CountingRegistrar::default()); let registrar: Arc = reg.clone(); let mut state = empty_pool_state(); - state.facade_tokens.insert("discord:acp_x".into(), "T_hung".into()); + state + .facade_tokens + .insert("discord:acp_x".into(), "T_hung".into()); // A different session's token must be untouched. - state.facade_tokens.insert("discord:acp_y".into(), "T_other".into()); + state + .facade_tokens + .insert("discord:acp_y".into(), "T_other".into()); super::revoke_facade_token_for_key(&mut state, "discord:acp_x", Some(®istrar)); - assert_eq!(reg.revoked(), vec!["T_hung"], "only the evicted session's token is revoked"); - assert!(!state.facade_tokens.contains_key("discord:acp_x"), "and it is forgotten"); + assert_eq!( + reg.revoked(), + vec!["T_hung"], + "only the evicted session's token is revoked" + ); + assert!( + !state.facade_tokens.contains_key("discord:acp_x"), + "and it is forgotten" + ); assert_eq!( state.facade_tokens.get("discord:acp_y").map(String::as_str), Some("T_other"), @@ -1319,11 +1368,23 @@ mod tests { }); let out = String::from_utf8(buf.lock().unwrap().clone()).unwrap(); - assert!(out.contains("force-evicting hung session"), "the warning must fire: {out}"); + assert!( + out.contains("force-evicting hung session"), + "the warning must fire: {out}" + ); assert!(!out.contains(uuid), "no raw uuid may reach the log: {out}"); - assert!(!out.contains("acp_") && !out.contains("sess_"), "no raw id prefix either: {out}"); - assert!(out.contains('#'), "the redaction tag must be present: {out}"); - assert!(out.contains("discord"), "the readable platform half must survive: {out}"); + assert!( + !out.contains("acp_") && !out.contains("sess_"), + "no raw id prefix either: {out}" + ); + assert!( + out.contains('#'), + "the redaction tag must be present: {out}" + ); + assert!( + out.contains("discord"), + "the readable platform half must survive: {out}" + ); } #[test] diff --git a/crates/openab-core/src/acp/protocol.rs b/crates/openab-core/src/acp/protocol.rs index 456245cae..217e774ed 100644 --- a/crates/openab-core/src/acp/protocol.rs +++ b/crates/openab-core/src/acp/protocol.rs @@ -284,9 +284,15 @@ pub fn parse_turn_result(result: &Value) -> TurnResult { .and_then(|v| v.as_str()) .map(String::from); let usage = result.get("usage"); - let input_tokens = usage.and_then(|u| u.get("inputTokens")).and_then(|v| v.as_u64()); - let output_tokens = usage.and_then(|u| u.get("outputTokens")).and_then(|v| v.as_u64()); - let total_tokens = usage.and_then(|u| u.get("totalTokens")).and_then(|v| v.as_u64()); + let input_tokens = usage + .and_then(|u| u.get("inputTokens")) + .and_then(|v| v.as_u64()); + let output_tokens = usage + .and_then(|u| u.get("outputTokens")) + .and_then(|v| v.as_u64()); + let total_tokens = usage + .and_then(|u| u.get("totalTokens")) + .and_then(|v| v.as_u64()); TurnResult { stop_reason, input_tokens, @@ -335,7 +341,11 @@ pub struct UsageReport { /// Returns `None` when `success` is not true or the data shape is missing — /// callers should treat that as "usage not supported by this agent". pub fn parse_usage_report(result: &Value) -> Option { - if !result.get("success").and_then(|v| v.as_bool()).unwrap_or(false) { + if !result + .get("success") + .and_then(|v| v.as_bool()) + .unwrap_or(false) + { return None; } let data = result.get("data")?; @@ -361,9 +371,7 @@ pub fn parse_usage_report(result: &Value) -> Option { let has_limit = b .get("hasLimit") .and_then(|v| v.as_bool()) - .unwrap_or_else(|| { - b.get("limit").and_then(|v| v.as_f64()).is_some() - }); + .unwrap_or_else(|| b.get("limit").and_then(|v| v.as_f64()).is_some()); Some(UsageBreakdown { display_name: b.get("displayName")?.as_str()?.to_string(), used: b.get("used")?.as_f64()?, diff --git a/crates/openab-core/src/acp_mcp.rs b/crates/openab-core/src/acp_mcp.rs index ab82c3091..0921eae8c 100644 --- a/crates/openab-core/src/acp_mcp.rs +++ b/crates/openab-core/src/acp_mcp.rs @@ -35,7 +35,6 @@ pub trait AcpMcpTunnel: Send + Sync { params: Option, ) -> Result; - /// Resolve a declared server NAME to the `server_id` the registry keys that tunnel by, for one /// channel. /// @@ -70,10 +69,6 @@ pub trait AcpMcpTunnel: Send + Sync { fn attached_server_names(&self, channel_id: &str) -> Vec; } - - - - /// Write `value` to `path` atomically, owner-only. /// /// Same-directory temp file created `0600` *before* any bytes reach it, then `rename`. Writing in @@ -96,7 +91,9 @@ async fn write_json_atomic(path: &std::path::Path, value: &Value) -> std::io::Re static TMP_SEQ: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0); let tmp = dir.join(format!( ".{}.openab-tmp.{}.{}", - path.file_name().and_then(|n| n.to_str()).unwrap_or("mcp.json"), + path.file_name() + .and_then(|n| n.to_str()) + .unwrap_or("mcp.json"), std::process::id(), TMP_SEQ.fetch_add(1, std::sync::atomic::Ordering::Relaxed) )); @@ -150,9 +147,6 @@ async fn write_json_atomic(path: &std::path::Path, value: &Value) -> std::io::Re } } - - - /// Author the static `openab` facade entry into the one file openab owns /// (Facade mode). The entry is byte-identical for every session — the bridge it used to be /// compared against here was removed with the rest of Option C, so the comparison is gone @@ -191,10 +185,11 @@ pub async fn write_facade_mcp_config(workdir: &str, facade_url: &str) -> std::io /// pass it and identifies no vendor at spawn time (D-21). openab never puts the entry in /// place itself either (D-15). pub fn facade_config_path(workdir: &str) -> std::path::PathBuf { - std::path::Path::new(workdir).join(".openab").join("mcp-facade.json") + std::path::Path::new(workdir) + .join(".openab") + .join("mcp-facade.json") } - /// Broker-side session credential hook (Facade mode). Implemented by the root /// (closing over the facade's `SessionTokens` registry — core stays free of /// the openab-mcp dependency); the pool calls it at session spawn/evict. @@ -263,8 +258,6 @@ pub fn report_facade_status(mcp_configured: bool, workdir: &str) { } } - - /// The destructive cases for the facade config writer — the one file openab owns, /// `.openab/mcp-facade.json`. It touches a user's `mcp.json` nowhere at all any more; the /// boundary test below is what asserts that. @@ -293,9 +286,11 @@ mod facade_config_writer { let path = facade_config_path(wd.to_str().unwrap()); assert_eq!(path, wd.join(".openab").join("mcp-facade.json")); - let v: Value = - serde_json::from_slice(&tokio::fs::read(&path).await.unwrap()).unwrap(); - assert_eq!(v["mcpServers"]["openab"]["url"], json!("http://127.0.0.1:8848/mcp")); + let v: Value = serde_json::from_slice(&tokio::fs::read(&path).await.unwrap()).unwrap(); + assert_eq!( + v["mcpServers"]["openab"]["url"], + json!("http://127.0.0.1:8848/mcp") + ); // The token is never in the file — the literal is, and the agent's env supplies the value. assert_eq!( v["mcpServers"]["openab"]["headers"]["Authorization"], @@ -320,7 +315,9 @@ mod facade_config_writer { let kiro = wd.join(".kiro").join("settings").join("mcp.json"); let agent = wd.join(".kiro").join("agents").join("terra.json"); for f in [&cursor, &kiro, &agent] { - tokio::fs::create_dir_all(f.parent().unwrap()).await.unwrap(); + tokio::fs::create_dir_all(f.parent().unwrap()) + .await + .unwrap(); } // Deliberately including a `//` comment: the shape that the old merge path destroyed. let cursor_body = "{\n // mine\n \"mcpServers\": {\"mine\": {\"command\": \"x\"}}\n}"; @@ -338,7 +335,10 @@ mod facade_config_writer { // longer retires those — see the PR body: that cleanup became operator-performed, and it // is a policy bypass rather than a convenience, so it is stated rather than silently // dropped. - assert_eq!(tokio::fs::read_to_string(&cursor).await.unwrap(), cursor_body); + assert_eq!( + tokio::fs::read_to_string(&cursor).await.unwrap(), + cursor_body + ); assert_eq!(tokio::fs::read_to_string(&kiro).await.unwrap(), kiro_body); assert_eq!(tokio::fs::read_to_string(&agent).await.unwrap(), agent_body); let _ = tokio::fs::remove_dir_all(&wd).await; @@ -379,7 +379,10 @@ mod facade_config_writer { let v: Value = serde_json::from_slice(&tokio::fs::read(facade_config_path(&w)).await.unwrap()) .unwrap(); - assert_eq!(v["mcpServers"]["openab"]["url"], json!("http://127.0.0.1:8848/mcp")); + assert_eq!( + v["mcpServers"]["openab"]["url"], + json!("http://127.0.0.1:8848/mcp") + ); let _ = tokio::fs::remove_dir_all(&wd).await; } @@ -393,24 +396,22 @@ mod facade_config_writer { .await .unwrap(); let path = facade_config_path(wd.to_str().unwrap()); - let mode = tokio::fs::metadata(&path).await.unwrap().permissions().mode() & 0o777; - assert_eq!(mode, 0o600, "0600 must be set at creation, not chmod'd after"); + let mode = tokio::fs::metadata(&path) + .await + .unwrap() + .permissions() + .mode() + & 0o777; + assert_eq!( + mode, 0o600, + "0600 must be set at creation, not chmod'd after" + ); let _ = tokio::fs::remove_dir_all(&wd).await; } } #[cfg(test)] mod tests { - use super::{ - }; - - - - - - - - /// Unique throwaway workdir with a `.kiro/agents/` tree. async fn tmp_workdir(tag: &str) -> std::path::PathBuf { @@ -425,19 +426,5 @@ mod tests { dir } - - - - - - - - - - - - const INIT_BODY: &str = r#"{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2025-03-26","capabilities":{},"clientInfo":{"name":"t","version":"1"}}}"#; - - } diff --git a/crates/openab-core/src/adapter.rs b/crates/openab-core/src/adapter.rs index e899a586e..de8f83cb7 100644 --- a/crates/openab-core/src/adapter.rs +++ b/crates/openab-core/src/adapter.rs @@ -4,7 +4,9 @@ use serde::Serialize; use std::sync::Arc; use tracing::{error, warn}; -use crate::acp::{classify_notification, parse_turn_result, AcpEvent, ContentBlock, SessionPool, TurnResult}; +use crate::acp::{ + classify_notification, parse_turn_result, AcpEvent, ContentBlock, SessionPool, TurnResult, +}; use crate::config::{ReactionsConfig, ToolDisplay}; use crate::error_display::{format_coded_error, format_user_error}; use crate::format; @@ -72,7 +74,12 @@ pub fn parse_output_directives(content: &str) -> (OutputDirectives, String) { "reply_to" => { let v = value.trim(); // Validate: non-empty, reasonable length, no whitespace/control chars - if !v.is_empty() && v.len() <= 64 && v.chars().all(|c| c.is_ascii_alphanumeric() || c == '.' || c == '-' || c == '_') { + if !v.is_empty() + && v.len() <= 64 + && v.chars().all(|c| { + c.is_ascii_alphanumeric() || c == '.' || c == '-' || c == '_' + }) + { directives.reply_to = Some(v.to_string()); } } @@ -1388,16 +1395,17 @@ fn contains_bot_mention(content: &str) -> bool { while i + 2 < bytes.len() { if bytes[i] == b'<' && bytes[i + 1] == b'@' { // Skip optional '!' (nickname mention) or '&' (role mention) - let start = if i + 2 < bytes.len() - && (bytes[i + 2] == b'!' || bytes[i + 2] == b'&') - { + let start = if i + 2 < bytes.len() && (bytes[i + 2] == b'!' || bytes[i + 2] == b'&') { i + 3 } else { i + 2 }; if start < bytes.len() && bytes[start].is_ascii_digit() { if let Some(end) = content[start..].find('>') { - if content[start..start + end].chars().all(|c| c.is_ascii_digit()) { + if content[start..start + end] + .chars() + .all(|c| c.is_ascii_digit()) + { return true; } } @@ -1626,8 +1634,7 @@ fn compose_display( // matches the sibling finished-fallback summary and // the pre-PR raw-count behaviour that users are used // to. A hidden group of `a×2` contributes 2, not 1. - let hidden_groups = - running_groups.len() - TOOL_COLLAPSE_THRESHOLD; + let hidden_groups = running_groups.len() - TOOL_COLLAPSE_THRESHOLD; let hidden_calls: usize = running_groups .iter() .take(hidden_groups) @@ -1736,7 +1743,11 @@ fn propagate_mentions_to_chunks( } else { let footer = format!( "\n{}", - missing.iter().map(|m| m.as_str()).collect::>().join(" ") + missing + .iter() + .map(|m| m.as_str()) + .collect::>() + .join(" ") ); if chunk.chars().count() + footer.chars().count() <= limit { format!("{chunk}{footer}") @@ -1761,7 +1772,10 @@ mod tests { assert_eq!(reply_message_limit("slack", 4096), 4096); // and a long reply under the ACP limit is a single chunk (delivered whole) let long = "x".repeat(50_000); - assert_eq!(crate::format::split_message(&long, reply_message_limit("acp", 4096)).len(), 1); + assert_eq!( + crate::format::split_message(&long, reply_message_limit("acp", 4096)).len(), + 1 + ); } #[test] @@ -2101,7 +2115,10 @@ mod tests { tool("3", "grep", ToolState::Completed), ]; let out = compose_display(&tools, "done", false, ToolDisplay::Full); - assert!(!out.contains("(×"), "should not collapse across order: {out}"); + assert!( + !out.contains("(×"), + "should not collapse across order: {out}" + ); assert_eq!(out.matches("`grep`").count(), 2, "output: {out}"); assert_eq!(out.matches("`curl`").count(), 1, "output: {out}"); } diff --git a/crates/openab-core/src/ambient.rs b/crates/openab-core/src/ambient.rs index 57f9e3213..8daca550b 100644 --- a/crates/openab-core/src/ambient.rs +++ b/crates/openab-core/src/ambient.rs @@ -346,7 +346,9 @@ impl AmbientDispatcher { /// Check if ambient mode is active and this channel is in the allowlist. pub fn is_ambient_channel(&self, channel_id: u64) -> bool { - self.config.enabled && !self.enabled_channels.is_empty() && self.enabled_channels.contains(&channel_id) + self.config.enabled + && !self.enabled_channels.is_empty() + && self.enabled_channels.contains(&channel_id) } /// Decide whether a message should be ambient-buffered. @@ -431,11 +433,7 @@ impl AmbientDispatcher { let state = channels.get(channel_id).unwrap(); // Non-blocking try_send — if buffer is full (hard_cap), drop the message. if let Err(e) = state.tx.try_send(msg) { - debug!( - channel_id, - "ambient buffer full, dropping message: {}", - e - ); + debug!(channel_id, "ambient buffer full, dropping message: {}", e); } } @@ -448,7 +446,10 @@ impl AmbientDispatcher { if let Some(state) = channels.get(channel_id) { // Cancel any in-flight post — consumer discards current batch on next check. state.post_guard.cancel(); - debug!(channel_id, "ambient buffer discard requested (mention arrived)"); + debug!( + channel_id, + "ambient buffer discard requested (mention arrived)" + ); } } @@ -504,7 +505,8 @@ async fn ambient_consumer_loop( let base = Duration::from_secs(base_secs); let jitter_range = base.as_millis() as f64 * 0.2; let jitter_ms = rand::thread_rng().gen_range(-jitter_range..jitter_range) as i64; - let interval = Duration::from_millis((base.as_millis() as i64 + jitter_ms).max(1000) as u64); + let interval = + Duration::from_millis((base.as_millis() as i64 + jitter_ms).max(1000) as u64); let deadline = tokio::time::Instant::now() + interval; let mut batch = vec![first]; @@ -569,8 +571,16 @@ async fn ambient_consumer_loop( // Debug mode: notify channel that a flush is happening. if config.debug { let senders: Vec<&str> = batch.iter().map(|m| m.sender_name.as_str()).collect(); - let transcript = batch.iter() - .map(|m| format!("{} (<@{}>): {}", m.sender_name, m.sender_id, m.prompt.replace('`', "'"))) + let transcript = batch + .iter() + .map(|m| { + format!( + "{} (<@{}>): {}", + m.sender_name, + m.sender_id, + m.prompt.replace('`', "'") + ) + }) .collect::>() .join("\n"); let mut debug_msg = format!( @@ -750,7 +760,8 @@ mod tests { #[test] fn payload_transcript_includes_sender_uid() { - let blocks = build_ambient_payload(&[msg("hello")], super::DEFAULT_AMBIENT_SYSTEM_INSTRUCTION); + let blocks = + build_ambient_payload(&[msg("hello")], super::DEFAULT_AMBIENT_SYSTEM_INSTRUCTION); let transcript = match &blocks[1] { ContentBlock::Text { text } => text.as_str(), #[allow(unreachable_patterns)] diff --git a/crates/openab-core/src/config.rs b/crates/openab-core/src/config.rs index fdd89dcb3..191788afd 100644 --- a/crates/openab-core/src/config.rs +++ b/crates/openab-core/src/config.rs @@ -133,7 +133,8 @@ fn default_mcp_listen() -> String { } #[derive(Debug, Clone, Deserialize)] -pub struct AgentCoreConfig { /// AgentCore Runtime ARN (required) +pub struct AgentCoreConfig { + /// AgentCore Runtime ARN (required) pub runtime_arn: String, /// ACP agent command to run in the PTY shell (default: kiro-cli acp --trust-all-tools) #[serde(default = "default_agentcore_shell_command")] @@ -3016,9 +3017,7 @@ allowed_users = ["U1234567890abcdef0123456789abcdef"] // --- use_adc: config value resolves without touching env --- let r = GoogleChatConfig { use_adc: Some(true), - adc_target_service_account: Some( - "chat-bot@project.iam.gserviceaccount.com".into(), - ), + adc_target_service_account: Some("chat-bot@project.iam.gserviceaccount.com".into()), ..Default::default() } .resolve(); diff --git a/crates/openab-core/src/cron.rs b/crates/openab-core/src/cron.rs index 7f9d59846..73edd0b7b 100644 --- a/crates/openab-core/src/cron.rs +++ b/crates/openab-core/src/cron.rs @@ -251,10 +251,12 @@ fn should_create_cron_thread(job: &CronJobConfig) -> bool { } fn cron_sender_thread_id(channel: &ChannelRef) -> Option { - channel - .thread_id - .clone() - .or_else(|| channel.parent_id.as_ref().map(|_| channel.channel_id.clone())) + channel.thread_id.clone().or_else(|| { + channel + .parent_id + .as_ref() + .map(|_| channel.channel_id.clone()) + }) } /// Validate all cronjob configs (fail-fast on bad cron expressions or timezones). @@ -647,8 +649,7 @@ async fn fire_cronjob( DisableOnSuccessResult::NotAchieved(reason) => { info!( id = job.id.as_deref().unwrap_or(""), - reason, - "disable_on_success not achieved, firing cronjob normally" + reason, "disable_on_success not achieved, firing cronjob normally" ); } } @@ -725,7 +726,7 @@ async fn fire_cronjob( thread_id: cron_sender_thread_id(&reply_channel), is_bot: true, timestamp: Some(Utc::now().to_rfc3339()), - message_id: None, // cron jobs don't originate from a message + message_id: None, // cron jobs don't originate from a message receiver_id: None, // cron jobs are self-triggered, no external receiver }; let sender_json = match serde_json::to_string(&sender) { @@ -1630,7 +1631,10 @@ message = "a" origin_event_id: None, }; - assert_eq!(cron_sender_thread_id(&channel).as_deref(), Some("thread-456")); + assert_eq!( + cron_sender_thread_id(&channel).as_deref(), + Some("thread-456") + ); } // --- validate_cronjobs tests --- diff --git a/crates/openab-core/src/discord.rs b/crates/openab-core/src/discord.rs index c34c68a7d..7c2817e8a 100644 --- a/crates/openab-core/src/discord.rs +++ b/crates/openab-core/src/discord.rs @@ -17,7 +17,9 @@ use serenity::builder::{ }; use serenity::http::Http; use serenity::model::application::ButtonStyle; -use serenity::model::application::{Command, CommandOptionType, ComponentInteractionDataKind, Interaction}; +use serenity::model::application::{ + Command, CommandOptionType, ComponentInteractionDataKind, Interaction, +}; use serenity::model::channel::{AutoArchiveDuration, Message, MessageType, Reaction, ReactionType}; use serenity::model::gateway::Ready; use serenity::model::id::{ChannelId, MessageId, UserId}; @@ -212,7 +214,11 @@ impl ChatAdapter for DiscordAdapter { let ch_id: u64 = Self::resolve_channel(channel).parse()?; // Truncate at char boundary to avoid panic on multi-byte chars (中文/Emoji). let truncated: &str = if title.chars().count() > 100 { - let end = title.char_indices().nth(100).map(|(i, _)| i).unwrap_or(title.len()); + let end = title + .char_indices() + .nth(100) + .map(|(i, _)| i) + .unwrap_or(title.len()); &title[..end] } else { title @@ -356,7 +362,9 @@ impl EventHandler for Handler { let key = msg.channel_id.to_string(); { let mut cache = self.multibot_threads.lock().await; - cache.entry(key.clone()).or_insert_with(tokio::time::Instant::now); + cache + .entry(key.clone()) + .or_insert_with(tokio::time::Instant::now); } // Persist to disk — multibot is irreversible self.multibot_cache.mark_multibot(&key).await; @@ -499,11 +507,14 @@ impl EventHandler for Handler { // non-allowed channels. Moved before bot gating so ambient context // can be resolved early — bot messages in ambient contexts must bypass // discord-level bot gating (#1197). - let (in_thread, bot_owns_thread, thread_parent_id, is_dm, is_structural_thread, structural_parent_id) = match msg - .channel_id - .to_channel(&ctx.http) - .await - { + let ( + in_thread, + bot_owns_thread, + thread_parent_id, + is_dm, + is_structural_thread, + structural_parent_id, + ) = match msg.channel_id.to_channel(&ctx.http).await { Ok(serenity::model::channel::Channel::Guild(gc)) => { let parent = gc.parent_id.map(|id| id.get().to_string()); let has_thread_metadata = gc.thread_metadata.is_some(); @@ -532,7 +543,11 @@ impl EventHandler for Handler { if has_thread_metadata { parent } else { None }, false, has_thread_metadata, - if has_thread_metadata { parent_u64 } else { None }, + if has_thread_metadata { + parent_u64 + } else { + None + }, ) } Ok(serenity::model::channel::Channel::Private(_)) => { @@ -552,7 +567,12 @@ impl EventHandler for Handler { // Check if message is in an ambient context (resolved early so bot // messages destined for ambient can bypass discord-level bot gating). let in_ambient_context = self.ambient.as_ref().is_some_and(|ambient| { - ambient.should_buffer(channel_id, is_structural_thread, bot_owns_thread, structural_parent_id) + ambient.should_buffer( + channel_id, + is_structural_thread, + bot_owns_thread, + structural_parent_id, + ) }); // --- Ambient early-route for bot messages --- @@ -599,13 +619,15 @@ impl EventHandler for Handler { let target = Arc::clone(&self.router) as Arc; debug!(channel_id = %msg.channel_id, bot_id = %msg.author.id, "ambient early-route: bot msg buffered"); - ambient.submit( - &channel_id.to_string(), - channel_ref, - adapter.clone(), - target, - ambient_msg, - ).await; + ambient + .submit( + &channel_id.to_string(), + channel_ref, + adapter.clone(), + target, + ambient_msg, + ) + .await; } } return; @@ -761,13 +783,15 @@ impl EventHandler for Handler { }; let target = Arc::clone(&self.router) as Arc; - ambient.submit( - &channel_id.to_string(), - channel_ref, - adapter.clone(), - target, - ambient_msg, - ).await; + ambient + .submit( + &channel_id.to_string(), + channel_ref, + adapter.clone(), + target, + ambient_msg, + ) + .await; return; } } @@ -918,7 +942,8 @@ impl EventHandler for Handler { // be uploaded to S3, not inlined). let attachment_size = u64::from(attachment.size); #[cfg(feature = "filestore")] - let skip_cap = self.filestore.is_some() && attachment_size > crate::media::TEXT_INLINE_LIMIT; + let skip_cap = + self.filestore.is_some() && attachment_size > crate::media::TEXT_INLINE_LIMIT; #[cfg(not(feature = "filestore"))] let skip_cap = false; if !skip_cap && text_file_bytes + attachment_size > TEXT_TOTAL_CAP { @@ -999,7 +1024,9 @@ impl EventHandler for Handler { attachment.content_type.as_deref(), None, fs, - ).await { + ) + .await + { extra_blocks.push(block); } } @@ -1064,10 +1091,11 @@ impl EventHandler for Handler { let trigger_msg = discord_msg_ref(&msg); // Per-thread streaming: check if another bot is present in this thread - let other_bot_present_flag = { - let cache = self.multibot_threads.lock().await; - cache.contains_key(&msg.channel_id.to_string()) - } || self.multibot_cache.is_multibot(&msg.channel_id.to_string()); + let other_bot_present_flag = + { + let cache = self.multibot_threads.lock().await; + cache.contains_key(&msg.channel_id.to_string()) + } || self.multibot_cache.is_multibot(&msg.channel_id.to_string()); // Backfill thread_id: when OAB just created a new thread, the sender // was built before the thread existed. Patch it so the agent sees @@ -1236,24 +1264,30 @@ impl EventHandler for Handler { if !in_allowed_thread { return; } - (ChannelRef { - platform: "discord".into(), - channel_id: channel_id.get().to_string(), - thread_id: None, - parent_id: parent.map(|p| p.to_string()), - origin_event_id: None, - }, true) + ( + ChannelRef { + platform: "discord".into(), + channel_id: channel_id.get().to_string(), + thread_id: None, + parent_id: parent.map(|p| p.to_string()), + origin_event_id: None, + }, + true, + ) } else { if !in_allowed_channel { return; } - (ChannelRef { - platform: "discord".into(), - channel_id: channel_id.get().to_string(), - thread_id: None, - parent_id: None, - origin_event_id: None, - }, false) + ( + ChannelRef { + platform: "discord".into(), + channel_id: channel_id.get().to_string(), + thread_id: None, + parent_id: None, + origin_event_id: None, + }, + false, + ) } } _ => return, @@ -1267,7 +1301,8 @@ impl EventHandler for Handler { self.allow_user_messages, AllowUsers::Involved | AllowUsers::MultibotMentions ) { - self.bot_participated_in_thread(&ctx.http, channel_id, bot_id).await + self.bot_participated_in_thread(&ctx.http, channel_id, bot_id) + .await } else { // For non-thread: still check multibot cache for dispatch info. let mb = { @@ -1301,17 +1336,16 @@ impl EventHandler for Handler { tokio::spawn(async move { // F2 fix: Fetch user info first, then apply user gating with confirmed bot status. - let (sender_name, display_name, is_bot_confirmed) = - match user_id.to_user(&http).await { - Ok(user) => { - let display = user.global_name.as_ref().unwrap_or(&user.name).clone(); - (user.name.clone(), display, user.bot) - } - Err(_) => { - let fallback = user_id.to_string(); - (fallback.clone(), fallback, is_reactor_bot) - } - }; + let (sender_name, display_name, is_bot_confirmed) = match user_id.to_user(&http).await { + Ok(user) => { + let display = user.global_name.as_ref().unwrap_or(&user.name).clone(); + (user.name.clone(), display, user.bot) + } + Err(_) => { + let fallback = user_id.to_string(); + (fallback.clone(), fallback, is_reactor_bot) + } + }; // Defense-in-depth: if to_user() reveals this is a bot but member was // None (rare edge case), re-apply bot gating retroactively. @@ -1319,8 +1353,7 @@ impl EventHandler for Handler { match allow_bot_messages { AllowBots::Off | AllowBots::Mentions => return, AllowBots::All => { - if !trusted_bot_ids.is_empty() - && !trusted_bot_ids.contains(&user_id.get()) + if !trusted_bot_ids.is_empty() && !trusted_bot_ids.contains(&user_id.get()) { return; } @@ -1403,23 +1436,31 @@ impl EventHandler for Handler { CreateCommand::new("reset").description("Reset the conversation session"), CreateCommand::new("remind") .description("Set a one-shot reminder to mention users/roles after a delay") - .add_option(CreateCommandOption::new( - CommandOptionType::String, - "targets", - "Users/roles to mention (e.g. @user1 @role1)", - ).required(true)) - .add_option(CreateCommandOption::new( - CommandOptionType::String, - "message", - "Reminder message", - ).required(true)) - .add_option(CreateCommandOption::new( - CommandOptionType::String, - "delay", - "Delay before firing (e.g. 30m, 2h, 1d)", - ).required(true)), - CreateCommand::new("auth") - .description("Authenticate the backend agent (device flow)"), + .add_option( + CreateCommandOption::new( + CommandOptionType::String, + "targets", + "Users/roles to mention (e.g. @user1 @role1)", + ) + .required(true), + ) + .add_option( + CreateCommandOption::new( + CommandOptionType::String, + "message", + "Reminder message", + ) + .required(true), + ) + .add_option( + CreateCommandOption::new( + CommandOptionType::String, + "delay", + "Delay before firing (e.g. 30m, 2h, 1d)", + ) + .required(true), + ), + CreateCommand::new("auth").description("Authenticate the backend agent (device flow)"), CreateCommand::new("usage") .description("Show backend account usage and billing information"), CreateCommand::new("export-thread") @@ -1491,7 +1532,8 @@ impl EventHandler for Handler { self.handle_config_command(&ctx, &cmd, "model").await; } Interaction::Command(cmd) if cmd.data.name == "effort" => { - self.handle_config_command(&ctx, &cmd, "thought_level").await; + self.handle_config_command(&ctx, &cmd, "thought_level") + .await; } Interaction::Command(cmd) if cmd.data.name == "agents" => { self.handle_config_command(&ctx, &cmd, "agent").await; @@ -1708,7 +1750,9 @@ impl Handler { if !self.router.pool().has_active_session(&thread_key).await { let response = CreateInteractionResponse::Message( CreateInteractionResponseMessage::new() - .content("⚠️ No active session. Start a conversation first by @mentioning the bot.") + .content( + "⚠️ No active session. Start a conversation first by @mentioning the bot.", + ) .ephemeral(true), ); if let Err(e) = cmd.create_response(&ctx.http, response).await { @@ -1719,8 +1763,9 @@ impl Handler { // The ACP round-trip can exceed Discord's 3-second interaction // deadline — acknowledge with a deferred ephemeral response first. - let defer = - CreateInteractionResponse::Defer(CreateInteractionResponseMessage::new().ephemeral(true)); + let defer = CreateInteractionResponse::Defer( + CreateInteractionResponseMessage::new().ephemeral(true), + ); if let Err(e) = cmd.create_response(&ctx.http, defer).await { tracing::error!(error = %e, "failed to defer /usage response"); return; @@ -1857,15 +1902,18 @@ impl Handler { // Extract options let opts = &cmd.data.options; - let targets_raw = opts.iter() + let targets_raw = opts + .iter() .find(|o| o.name == "targets") .and_then(|o| o.value.as_str()) .unwrap_or(""); - let message = opts.iter() + let message = opts + .iter() .find(|o| o.name == "message") .and_then(|o| o.value.as_str()) .unwrap_or(""); - let delay_raw = opts.iter() + let delay_raw = opts + .iter() .find(|o| o.name == "delay") .and_then(|o| o.value.as_str()) .unwrap_or(""); @@ -1927,7 +1975,10 @@ impl Handler { if targets.len() > remind::MAX_TARGETS { let response = CreateInteractionResponse::Message( CreateInteractionResponseMessage::new() - .content(format!("⚠️ Too many targets (max {}). Use a @role instead.", remind::MAX_TARGETS)) + .content(format!( + "⚠️ Too many targets (max {}). Use a @role instead.", + remind::MAX_TARGETS + )) .ephemeral(true), ); let _ = cmd.create_response(&ctx.http, response).await; @@ -2027,7 +2078,9 @@ impl Handler { if AUTH_IN_PROGRESS.swap(true, std::sync::atomic::Ordering::Acquire) { let response = CreateInteractionResponse::Message( CreateInteractionResponseMessage::new() - .content("⚠️ Authentication already in progress. Please wait for it to complete.") + .content( + "⚠️ Authentication already in progress. Please wait for it to complete.", + ) .ephemeral(true), ); let _ = cmd.create_response(&ctx.http, response).await; @@ -2040,7 +2093,9 @@ impl Handler { AUTH_IN_PROGRESS.store(false, std::sync::atomic::Ordering::Release); let response = CreateInteractionResponse::Message( CreateInteractionResponseMessage::new() - .content("⚠️ No auth command configured (`OPENAB_AGENT_AUTH_COMMAND` not set).") + .content( + "⚠️ No auth command configured (`OPENAB_AGENT_AUTH_COMMAND` not set).", + ) .ephemeral(true), ); let _ = cmd.create_response(&ctx.http, response).await; @@ -2063,9 +2118,9 @@ impl Handler { let user_id = cmd.user.id.get(); tokio::spawn(async move { + use std::sync::Arc; use tokio::io::AsyncBufReadExt; use tokio::process::Command as TokioCommand; - use std::sync::Arc; // Drop guard ensures AUTH_IN_PROGRESS is cleared even on panic. struct AuthGuard; @@ -2089,13 +2144,15 @@ impl Handler { Ok(c) => c, Err(e) => { tracing::error!(error = %e, "/auth: failed to spawn auth command"); - let _ = http.create_followup_message( - &token, - &CreateInteractionResponseFollowup::new() - .content(format!("❌ Failed to start auth command: {e}")) - .ephemeral(true), - Vec::new(), - ).await; + let _ = http + .create_followup_message( + &token, + &CreateInteractionResponseFollowup::new() + .content(format!("❌ Failed to start auth command: {e}")) + .ephemeral(true), + Vec::new(), + ) + .await; return; } }; @@ -2114,7 +2171,10 @@ impl Handler { let mut reader = tokio::io::BufReader::new(stdout).lines(); while let Ok(Some(line)) = reader.next_line().await { let has_url = line.contains("http://") || line.contains("https://"); - lines_out.lock().unwrap_or_else(|e| e.into_inner()).push(line); + lines_out + .lock() + .unwrap_or_else(|e| e.into_inner()) + .push(line); if has_url { url_found_out.notify_one(); } @@ -2129,7 +2189,10 @@ impl Handler { let mut reader = tokio::io::BufReader::new(stderr).lines(); while let Ok(Some(line)) = reader.next_line().await { let has_url = line.contains("http://") || line.contains("https://"); - lines_err.lock().unwrap_or_else(|e| e.into_inner()).push(line); + lines_err + .lock() + .unwrap_or_else(|e| e.into_inner()) + .push(line); if has_url { url_found_err.notify_one(); } @@ -2159,12 +2222,8 @@ impl Handler { // Handle an early exit (the command terminated during the URL window). if let Some(res) = early_exit { let _ = tokio::join!(stdout_task, stderr_task); - let collected = strip_ansi_codes( - &lines - .lock() - .unwrap_or_else(|e| e.into_inner()) - .join("\n"), - ); + let collected = + strip_ansi_codes(&lines.lock().unwrap_or_else(|e| e.into_inner()).join("\n")); let detail = if collected.trim().is_empty() { String::new() } else { @@ -2184,13 +2243,15 @@ impl Handler { } Err(e) => format!("❌ Error waiting for auth command: {e}"), }; - let _ = http.create_followup_message( - &token, - &CreateInteractionResponseFollowup::new() - .content(content) - .ephemeral(true), - Vec::new(), - ).await; + let _ = http + .create_followup_message( + &token, + &CreateInteractionResponseFollowup::new() + .content(content) + .ephemeral(true), + Vec::new(), + ) + .await; return; } @@ -2221,13 +2282,15 @@ impl Handler { // `truncate_to_utf16_budget` for the testable implementation. let truncated = truncate_to_utf16_budget(&output, prefix, suffix, 2000); let msg = format!("{prefix}{truncated}{suffix}"); - let _ = http.create_followup_message( - &token, - &CreateInteractionResponseFollowup::new() - .content(msg) - .ephemeral(true), - Vec::new(), - ).await; + let _ = http + .create_followup_message( + &token, + &CreateInteractionResponseFollowup::new() + .content(msg) + .ephemeral(true), + Vec::new(), + ) + .await; // Wait for the process to complete (user authorizes in browser). // Use 14min (not 15) to leave headroom for the Discord interaction token TTL. @@ -2235,44 +2298,55 @@ impl Handler { match tokio::time::timeout(timeout, child.wait()).await { Ok(Ok(status)) if status.success() => { info!("/auth: authentication successful"); - let _ = http.create_followup_message( - &token, - &CreateInteractionResponseFollowup::new() - .content("✅ Authentication successful!") - .ephemeral(true), - Vec::new(), - ).await; + let _ = http + .create_followup_message( + &token, + &CreateInteractionResponseFollowup::new() + .content("✅ Authentication successful!") + .ephemeral(true), + Vec::new(), + ) + .await; } Ok(Ok(status)) => { warn!(%status, "/auth: authentication failed"); - let _ = http.create_followup_message( - &token, - &CreateInteractionResponseFollowup::new() - .content(format!("❌ Authentication failed (exit code: {}).", status)) - .ephemeral(true), - Vec::new(), - ).await; + let _ = http + .create_followup_message( + &token, + &CreateInteractionResponseFollowup::new() + .content(format!( + "❌ Authentication failed (exit code: {}).", + status + )) + .ephemeral(true), + Vec::new(), + ) + .await; } Ok(Err(e)) => { tracing::error!(error = %e, "/auth: error waiting for auth process"); - let _ = http.create_followup_message( - &token, - &CreateInteractionResponseFollowup::new() - .content(format!("❌ Auth process error: {e}")) - .ephemeral(true), - Vec::new(), - ).await; + let _ = http + .create_followup_message( + &token, + &CreateInteractionResponseFollowup::new() + .content(format!("❌ Auth process error: {e}")) + .ephemeral(true), + Vec::new(), + ) + .await; } Err(_) => { warn!("/auth: timed out waiting for authorization"); let _ = child.kill().await; - let _ = http.create_followup_message( - &token, - &CreateInteractionResponseFollowup::new() - .content("⏰ Authentication timed out. Run `/auth` again to retry.") - .ephemeral(true), - Vec::new(), - ).await; + let _ = http + .create_followup_message( + &token, + &CreateInteractionResponseFollowup::new() + .content("⏰ Authentication timed out. Run `/auth` again to retry.") + .ephemeral(true), + Vec::new(), + ) + .await; } } @@ -2319,9 +2393,7 @@ impl Handler { ); (in_thread, gc.name.clone()) } - Ok(serenity::model::channel::Channel::Private(_)) => { - (self.allow_dm, "dm".to_string()) - } + Ok(serenity::model::channel::Channel::Private(_)) => (self.allow_dm, "dm".to_string()), Ok(_) => (false, "channel".to_string()), Err(e) => { tracing::warn!(channel_id = %channel_id, error = %e, "failed to inspect channel for export"); @@ -2343,16 +2415,34 @@ impl Handler { // --- Parse and validate filter params (mutual exclusion) --- let opts = &cmd.data.options; - let limit_opt = opts.iter().find(|o| o.name == "limit").and_then(|o| o.value.as_i64()); - let since_opt = opts.iter().find(|o| o.name == "since").and_then(|o| o.value.as_str()); - let days_opt = opts.iter().find(|o| o.name == "days").and_then(|o| o.value.as_i64()); - let all_opt = opts.iter().find(|o| o.name == "all").and_then(|o| o.value.as_bool()).unwrap_or(false); + let limit_opt = opts + .iter() + .find(|o| o.name == "limit") + .and_then(|o| o.value.as_i64()); + let since_opt = opts + .iter() + .find(|o| o.name == "since") + .and_then(|o| o.value.as_str()); + let days_opt = opts + .iter() + .find(|o| o.name == "days") + .and_then(|o| o.value.as_i64()); + let all_opt = opts + .iter() + .find(|o| o.name == "all") + .and_then(|o| o.value.as_bool()) + .unwrap_or(false); - let filter_count = limit_opt.is_some() as u8 + since_opt.is_some() as u8 + days_opt.is_some() as u8 + all_opt as u8; + let filter_count = limit_opt.is_some() as u8 + + since_opt.is_some() as u8 + + days_opt.is_some() as u8 + + all_opt as u8; if filter_count > 1 { let response = CreateInteractionResponse::Message( CreateInteractionResponseMessage::new() - .content("⚠️ Please specify only one filter: `limit`, `since`, `days`, or `all`.") + .content( + "⚠️ Please specify only one filter: `limit`, `since`, `days`, or `all`.", + ) .ephemeral(true), ); let _ = cmd.create_response(&ctx.http, response).await; @@ -2620,8 +2710,7 @@ fn format_usage_body(report: &UsageReport) -> (String, bool) { fn build_usage_reply(report: &UsageReport) -> (String, CreateEmbed) { let (body, over_limit) = format_usage_body(report); let content = format!("📊 **Usage — {}**\n{}", report.plan_name, body); - let mut embed = - CreateEmbed::new().colour(if over_limit { 0xE74C3C } else { 0x2ECC71 }); + let mut embed = CreateEmbed::new().colour(if over_limit { 0xE74C3C } else { 0x2ECC71 }); if let Some(reset) = &report.billing_cycle_reset { embed = embed.footer(CreateEmbedFooter::new(format!( "Billing cycle resets {reset}" @@ -2791,7 +2880,10 @@ async fn export_channel_messages( let filename = export_filename(channel_id, channel_name); if attachment_size_limit < 2048 { - tracing::warn!(attachment_size_limit, "attachment_size_limit is very small; export will likely be truncated"); + tracing::warn!( + attachment_size_limit, + "attachment_size_limit is very small; export will likely be truncated" + ); } let max_bytes = usize::try_from(attachment_size_limit) .unwrap_or(8 * 1024 * 1024) @@ -2861,10 +2953,7 @@ fn format_export_message(msg: &Message) -> String { let bot_marker = if msg.author.bot { " [bot]" } else { "" }; let mut out = format!( "[{}] {}{} ({})\n", - msg.timestamp, - msg.author.name, - bot_marker, - msg.author.id + msg.timestamp, msg.author.name, bot_marker, msg.author.id ); if msg.content.is_empty() { @@ -3284,7 +3373,7 @@ fn truncate_to_utf16_budget(body: &str, prefix: &str, suffix: &str, limit: usize mod tests { use super::*; use crate::acp::protocol::ConfigOptionValue; - use crate::bot_turns::{TurnResult, HARD_BOT_TURN_LIMIT, BOT_TURN_LIMIT_WARNING_PREFIX}; + use crate::bot_turns::{TurnResult, BOT_TURN_LIMIT_WARNING_PREFIX, HARD_BOT_TURN_LIMIT}; // --- truncate_for_discord (select menu option 100-char cap) --- @@ -3402,7 +3491,10 @@ mod tests { assert!(content.contains("12781.64 / 10000")); assert!(content.contains("Overage charges: 111.27 USD")); let json = serde_json::to_value(&embed).expect("embed serializes"); - assert!(json.get("description").is_none(), "body must not be in embed"); + assert!( + json.get("description").is_none(), + "body must not be in embed" + ); assert!(json.get("title").is_none(), "title must not be in embed"); assert_eq!(json["color"], 0xE74C3C, "over limit → red strip"); assert_eq!(json["footer"]["text"], "Billing cycle resets 2026-08-01"); @@ -3464,7 +3556,10 @@ mod tests { #[test] fn truncate_utf16_respects_prefix_suffix_budget() { // limit 10, prefix "pre" (3) + suffix "su" (2) = 5 → 5 ASCII units left. - assert_eq!(truncate_to_utf16_budget("abcdefghij", "pre", "su", 10), "abcde"); + assert_eq!( + truncate_to_utf16_budget("abcdefghij", "pre", "su", 10), + "abcde" + ); } /// A supplementary-plane scalar counts as TWO UTF-16 code units, not one. @@ -4356,9 +4451,8 @@ mod tests { trusted_bot_ids: &HashSet, author_id: u64, ) -> bool { - let trusted_mention = is_mentioned - && !trusted_bot_ids.is_empty() - && trusted_bot_ids.contains(&author_id); + let trusted_mention = + is_mentioned && !trusted_bot_ids.is_empty() && trusted_bot_ids.contains(&author_id); if !trusted_mention { match allow_bot_messages { @@ -4391,7 +4485,12 @@ mod tests { #[test] fn bot_admission_untrusted_mention_blocked_by_off() { let trusted = HashSet::from([42]); - assert!(!should_admit_bot_message(AllowBots::Off, true, &trusted, 99)); + assert!(!should_admit_bot_message( + AllowBots::Off, + true, + &trusted, + 99 + )); } /// GIVEN: allow_bot_messages=Off, trusted bot without @mention @@ -4399,7 +4498,12 @@ mod tests { #[test] fn bot_admission_trusted_no_mention_blocked_by_off() { let trusted = HashSet::from([42]); - assert!(!should_admit_bot_message(AllowBots::Off, false, &trusted, 42)); + assert!(!should_admit_bot_message( + AllowBots::Off, + false, + &trusted, + 42 + )); } /// GIVEN: allow_bot_messages=Off, empty trusted_bot_ids, bot @mentions @@ -4407,7 +4511,12 @@ mod tests { #[test] fn bot_admission_empty_trusted_ids_off_mode() { let trusted: HashSet = HashSet::new(); - assert!(!should_admit_bot_message(AllowBots::Off, true, &trusted, 42)); + assert!(!should_admit_bot_message( + AllowBots::Off, + true, + &trusted, + 42 + )); } /// GIVEN: allow_bot_messages=Mentions, trusted bot @mentions @@ -4415,7 +4524,12 @@ mod tests { #[test] fn bot_admission_mentions_mode_trusted_mention() { let trusted = HashSet::from([42]); - assert!(should_admit_bot_message(AllowBots::Mentions, true, &trusted, 42)); + assert!(should_admit_bot_message( + AllowBots::Mentions, + true, + &trusted, + 42 + )); } /// GIVEN: allow_bot_messages=All, untrusted bot (not in trusted_bot_ids) @@ -4423,7 +4537,12 @@ mod tests { #[test] fn bot_admission_all_mode_untrusted_bot_rejected() { let trusted = HashSet::from([42]); - assert!(!should_admit_bot_message(AllowBots::All, false, &trusted, 99)); + assert!(!should_admit_bot_message( + AllowBots::All, + false, + &trusted, + 99 + )); } // --- DM gating tests (#656) --- @@ -4513,19 +4632,28 @@ mod tests { #[test] fn dedup_detects_existing_bot_warning() { - let msg = format!("{} (20/20). A human must reply.", BOT_TURN_LIMIT_WARNING_PREFIX); + let msg = format!( + "{} (20/20). A human must reply.", + BOT_TURN_LIMIT_WARNING_PREFIX + ); assert!(turn_limit_warning_present(&[(true, &msg)])); } #[test] fn dedup_ignores_human_warning_text() { - let msg = format!("{} (20/20). A human must reply.", BOT_TURN_LIMIT_WARNING_PREFIX); + let msg = format!( + "{} (20/20). A human must reply.", + BOT_TURN_LIMIT_WARNING_PREFIX + ); assert!(!turn_limit_warning_present(&[(false, &msg)])); } #[test] fn dedup_returns_false_when_no_warning() { - assert!(!turn_limit_warning_present(&[(true, "hello"), (false, "world")])); + assert!(!turn_limit_warning_present(&[ + (true, "hello"), + (false, "world") + ])); } #[test] @@ -4542,7 +4670,10 @@ mod tests { fn reaction_mentions_mode_always_rejected() { assert!(!should_process_reaction( AllowUsers::Mentions, - true, true, false, false, + true, + true, + false, + false, )); } @@ -4554,7 +4685,8 @@ mod tests { AllowUsers::Involved, false, // is_thread false, // bot_involved (irrelevant for non-thread) - false, false, + false, + false, )); } @@ -4566,7 +4698,8 @@ mod tests { AllowUsers::Involved, true, // is_thread false, // bot_involved - false, false, + false, + false, )); } @@ -4578,7 +4711,8 @@ mod tests { AllowUsers::Involved, true, // is_thread true, // bot_involved - false, false, + false, + false, )); } @@ -4628,7 +4762,9 @@ mod tests { assert!(!should_process_reaction( AllowUsers::MultibotMentions, false, // is_thread - false, false, false, + false, + false, + false, )); } } diff --git a/crates/openab-core/src/filestore.rs b/crates/openab-core/src/filestore.rs index 0ba54ded1..92139b43c 100644 --- a/crates/openab-core/src/filestore.rs +++ b/crates/openab-core/src/filestore.rs @@ -92,11 +92,7 @@ impl Filestore { /// /// The object key is `{prefix}{uuid}_{filename}`. On success returns the /// presigned URL as a String. On failure logs the error and returns Err. - pub async fn upload_and_presign( - &self, - filename: &str, - data: &[u8], - ) -> anyhow::Result { + pub async fn upload_and_presign(&self, filename: &str, data: &[u8]) -> anyhow::Result { // Sanitize filename: strip path separators, traversal sequences, // double quotes (breaks Content-Disposition header parsing), and // non-ASCII chars. Limit length to prevent excessively long S3 keys. @@ -107,13 +103,12 @@ impl Filestore { .filter(|c| c.is_ascii_graphic() || *c == ' ') .take(200) .collect(); - let safe_name = if safe_name.is_empty() { "unnamed" } else { &safe_name }; - let key = format!( - "{}{}_{}", - self.prefix, - uuid::Uuid::new_v4(), - safe_name - ); + let safe_name = if safe_name.is_empty() { + "unnamed" + } else { + &safe_name + }; + let key = format!("{}{}_{}", self.prefix, uuid::Uuid::new_v4(), safe_name); // Upload the object (3-minute timeout to prevent indefinite hangs) const UPLOAD_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(180); @@ -188,7 +183,7 @@ impl Filestore { } else { &safe_name }; - let key = format!("{}{}_{}",self.prefix, uuid::Uuid::new_v4(), safe_name); + let key = format!("{}{}_{}", self.prefix, uuid::Uuid::new_v4(), safe_name); // Pre-check reported size if reported_size > self.max_file_size { @@ -260,25 +255,26 @@ impl Filestore { .send() .await { - Ok(resp) => { - match resp.e_tag { - Some(tag) => { - parts.push( - CompletedPart::builder() - .part_number(part_number) - .e_tag(tag) - .build(), - ); - part_number += 1; - } - None => { - upload_error = Some(anyhow::anyhow!("upload_part {part_number} returned no ETag")); - break; - } + Ok(resp) => match resp.e_tag { + Some(tag) => { + parts.push( + CompletedPart::builder() + .part_number(part_number) + .e_tag(tag) + .build(), + ); + part_number += 1; } - } + None => { + upload_error = Some(anyhow::anyhow!( + "upload_part {part_number} returned no ETag" + )); + break; + } + }, Err(e) => { - upload_error = Some(anyhow::anyhow!("upload_part {part_number} failed: {e}")); + upload_error = + Some(anyhow::anyhow!("upload_part {part_number} failed: {e}")); break; } } @@ -294,27 +290,31 @@ impl Filestore { .key(&key) .upload_id(&upload_id) .part_number(part_number) - .body(aws_sdk_s3::primitives::ByteStream::from(std::mem::take(&mut buffer))) + .body(aws_sdk_s3::primitives::ByteStream::from(std::mem::take( + &mut buffer, + ))) .send() .await { - Ok(resp) => { - match resp.e_tag { - Some(tag) => { - parts.push( - CompletedPart::builder() - .part_number(part_number) - .e_tag(tag) - .build(), - ); - } - None => { - upload_error = Some(anyhow::anyhow!("upload_part {part_number} (final) returned no ETag")); - } + Ok(resp) => match resp.e_tag { + Some(tag) => { + parts.push( + CompletedPart::builder() + .part_number(part_number) + .e_tag(tag) + .build(), + ); } - } + None => { + upload_error = Some(anyhow::anyhow!( + "upload_part {part_number} (final) returned no ETag" + )); + } + }, Err(e) => { - upload_error = Some(anyhow::anyhow!("upload_part {part_number} (final) failed: {e}")); + upload_error = Some(anyhow::anyhow!( + "upload_part {part_number} (final) failed: {e}" + )); } } } @@ -361,7 +361,9 @@ impl Filestore { .upload_id(&upload_id) .send() .await; - return Err(anyhow::anyhow!("stream produced no data — file may be empty or download failed")); + return Err(anyhow::anyhow!( + "stream produced no data — file may be empty or download failed" + )); } // If buffer has remaining data but no parts yet (file < 16 MB), upload as single part @@ -383,7 +385,8 @@ impl Filestore { Some(tag) => tag, None => { error!(bucket = %self.bucket, key = %key, "upload single part returned no ETag, aborting"); - let _ = self.client + let _ = self + .client .abort_multipart_upload() .bucket(&self.bucket) .key(&key) @@ -393,16 +396,12 @@ impl Filestore { return Err(anyhow::anyhow!("upload single part returned no ETag")); } }; - parts.push( - CompletedPart::builder() - .part_number(1) - .e_tag(e_tag) - .build(), - ); + parts.push(CompletedPart::builder().part_number(1).e_tag(e_tag).build()); } Err(e) => { error!(bucket = %self.bucket, key = %key, error = %e, "upload single part failed, aborting"); - let _ = self.client + let _ = self + .client .abort_multipart_upload() .bucket(&self.bucket) .key(&key) @@ -419,7 +418,8 @@ impl Filestore { .set_parts(Some(parts)) .build(); - if let Err(e) = self.client + if let Err(e) = self + .client .complete_multipart_upload() .bucket(&self.bucket) .key(&key) @@ -430,7 +430,8 @@ impl Filestore { { error!(bucket = %self.bucket, key = %key, error = %e, "complete_multipart_upload failed, aborting"); // Abort the multipart upload to clean up orphaned parts - let _ = self.client + let _ = self + .client .abort_multipart_upload() .bucket(&self.bucket) .key(&key) @@ -480,11 +481,20 @@ impl Filestore { /// Format the hint block returned to the agent when a large file is uploaded /// to the filestore instead of being inlined. -pub fn format_filestore_hint(filename: &str, size_bytes: u64, presigned_url: &str, ttl_secs: u64) -> String { +pub fn format_filestore_hint( + filename: &str, + size_bytes: u64, + presigned_url: &str, + ttl_secs: u64, +) -> String { let size_kb = size_bytes / 1024; let ttl_minutes = ttl_secs / 60; // Sanitize filename for prompt safety — strip control characters - let safe_filename: String = filename.chars().filter(|c| !c.is_control()).take(200).collect(); + let safe_filename: String = filename + .chars() + .filter(|c| !c.is_control()) + .take(200) + .collect(); format!( "[File: {safe_filename}]\n\ This file ({size_kb} KB) exceeds the 512 KB inline limit. \ @@ -550,7 +560,9 @@ secret_access_key = "SECRET" assert!(hint.contains("[File: big-log.txt]")); assert!(hint.contains("1024 KB")); assert!(hint.contains("exceeds the 512 KB inline limit")); - assert!(hint.contains("https://bucket.s3.amazonaws.com/incoming/uuid_big-log.txt?X-Amz-...")); + assert!( + hint.contains("https://bucket.s3.amazonaws.com/incoming/uuid_big-log.txt?X-Amz-...") + ); assert!(hint.contains("expires in 60 minutes")); } diff --git a/crates/openab-core/src/format.rs b/crates/openab-core/src/format.rs index 41a26d0df..91887c40f 100644 --- a/crates/openab-core/src/format.rs +++ b/crates/openab-core/src/format.rs @@ -413,7 +413,11 @@ mod tests { let chunks = split_message(line, 20); assert_length_invariant(&chunks, 20); assert!(chunks.len() > 1); - let rejoined = chunks.iter().map(|c| c.trim()).collect::>().join(" "); + let rejoined = chunks + .iter() + .map(|c| c.trim()) + .collect::>() + .join(" "); assert_eq!( rejoined.split_whitespace().collect::>(), line.split_whitespace().collect::>(), diff --git a/crates/openab-core/src/gateway.rs b/crates/openab-core/src/gateway.rs index e63416b56..f8dc253a7 100644 --- a/crates/openab-core/src/gateway.rs +++ b/crates/openab-core/src/gateway.rs @@ -105,7 +105,10 @@ struct EventFilterParams<'a> { /// Returns `true` if the event should be skipped (filtered out). fn should_skip_event(event: &GatewayEvent, filter: &EventFilterParams) -> bool { // Bot filter - if event.sender.is_bot && !filter.allow_bot_messages && !filter.trusted_bot_ids.contains(&event.sender.id) { + if event.sender.is_bot + && !filter.allow_bot_messages + && !filter.trusted_bot_ids.contains(&event.sender.id) + { tracing::info!(sender = %event.sender.id, "gateway: bot not in trusted_bot_ids, skipping"); return true; } @@ -120,7 +123,8 @@ fn should_skip_event(event: &GatewayEvent, filter: &EventFilterParams) -> bool { return true; } // @mention gating: in groups, only respond if bot is mentioned - let is_group = event.channel.channel_type == "group" || event.channel.channel_type == "supergroup"; + let is_group = + event.channel.channel_type == "group" || event.channel.channel_type == "supergroup"; let in_thread = event.channel.thread_id.is_some(); if is_group && !in_thread { if let Some(bot_name) = filter.bot_username { @@ -244,7 +248,8 @@ fn gateway_delivery_result(resp: GatewayResponse) -> Result { } else { Err(anyhow::anyhow!( "gateway reported failure: {}", - resp.error.unwrap_or_else(|| "unspecified error".to_string()) + resp.error + .unwrap_or_else(|| "unspecified error".to_string()) )) } } @@ -557,7 +562,8 @@ impl ChatAdapter for GatewayAdapter { content: &str, reply_to_message_id: &str, ) -> Result { - self.send_gateway_reply(channel, content, Some(reply_to_message_id)).await + self.send_gateway_reply(channel, content, Some(reply_to_message_id)) + .await } async fn create_thread( @@ -708,10 +714,14 @@ impl ChatAdapter for GatewayAdapter { match tokio::time::timeout( std::time::Duration::from_millis(EDIT_RESPONSE_TIMEOUT_MS), rx, - ).await { + ) + .await + { Ok(Ok(resp)) if resp.success => Ok(()), Ok(Ok(resp)) => { - let err_msg = resp.error.clone() + let err_msg = resp + .error + .clone() .unwrap_or_else(|| "gateway reported edit failure".to_string()); tracing::warn!(request_id = %id, error = %err_msg, "edit_message gateway replied failure"); Err(anyhow::anyhow!("edit failure: {err_msg}")) @@ -1464,7 +1474,11 @@ pub async fn process_gateway_event( } else { event.timestamp.clone() }), - message_id: if event.message_id.is_empty() { None } else { Some(event.message_id.clone()) }, + message_id: if event.message_id.is_empty() { + None + } else { + Some(event.message_id.clone()) + }, receiver_id: None, }; let sender_json = serde_json::to_string(&sender_ctx).unwrap_or_default(); @@ -1505,25 +1519,24 @@ pub async fn process_gateway_event( }; match att.attachment_type.as_str() { - "image" => { - match bytes_result { - Ok(bytes) => { - use base64::Engine; - let b64 = base64::engine::general_purpose::STANDARD.encode(&bytes); - extra_blocks.push(ContentBlock::Image { - media_type: att.mime_type.clone(), - data: b64, - }); - } - Err(e) => { - tracing::warn!(filename = %att.filename, error = %e, "gateway image read failed"); - } + "image" => match bytes_result { + Ok(bytes) => { + use base64::Engine; + let b64 = base64::engine::general_purpose::STANDARD.encode(&bytes); + extra_blocks.push(ContentBlock::Image { + media_type: att.mime_type.clone(), + data: b64, + }); } - } + Err(e) => { + tracing::warn!(filename = %att.filename, error = %e, "gateway image read failed"); + } + }, "text_file" => { match bytes_result { Ok(bytes) => { - let safe_filename: String = att.filename + let safe_filename: String = att + .filename .chars() .filter(|c| !c.is_control()) .take(200) @@ -1538,7 +1551,14 @@ pub async fn process_gateway_event( // Large file — upload to filestore if available #[cfg(feature = "filestore")] if let Some(ref fs) = ctx.filestore { - if let Some((block, _)) = crate::media::upload_bytes_to_filestore_public(&att.filename, &bytes, fs).await { + if let Some((block, _)) = + crate::media::upload_bytes_to_filestore_public( + &att.filename, + &bytes, + fs, + ) + .await + { extra_blocks.push(block); } else { // Upload refused (size cap) — emit degraded hint, don't inline oversized body @@ -1572,35 +1592,38 @@ pub async fn process_gateway_event( } } } - "audio" if ctx.stt_config.enabled => { - match bytes_result { - Ok(bytes) => { - match crate::stt::transcribe( - &crate::media::HTTP_CLIENT, - &ctx.stt_config, - bytes, - att.filename.clone(), - &att.mime_type, - ).await { - Some(transcript) => { - extra_blocks.push(ContentBlock::Text { - text: format!("[Voice message transcript]: {transcript}"), - }); - } - None => { - extra_blocks.push(ContentBlock::Text { - text: format!("[Voice message — transcription failed for {}]", att.filename), - }); - } + "audio" if ctx.stt_config.enabled => match bytes_result { + Ok(bytes) => { + match crate::stt::transcribe( + &crate::media::HTTP_CLIENT, + &ctx.stt_config, + bytes, + att.filename.clone(), + &att.mime_type, + ) + .await + { + Some(transcript) => { + extra_blocks.push(ContentBlock::Text { + text: format!("[Voice message transcript]: {transcript}"), + }); + } + None => { + extra_blocks.push(ContentBlock::Text { + text: format!( + "[Voice message — transcription failed for {}]", + att.filename + ), + }); } - } - Err(_) => { - extra_blocks.push(ContentBlock::Text { - text: format!("[Voice message — read failed for {}]", att.filename), - }); } } - } + Err(_) => { + extra_blocks.push(ContentBlock::Text { + text: format!("[Voice message — read failed for {}]", att.filename), + }); + } + }, _ => {} } } @@ -1609,20 +1632,38 @@ pub async fn process_gateway_event( let prompt = event.content.text.clone(); let trimmed = prompt.trim(); if trimmed == "/reset" { - let thread_id_str = event.channel.thread_id.as_deref().unwrap_or(&event.channel.id); + let thread_id_str = event + .channel + .thread_id + .as_deref() + .unwrap_or(&event.channel.id); let thread_key = format!("{}:{}", event.platform, thread_id_str); - let dropped = ctx.dispatcher.cancel_buffered_thread(event.platform.as_str(), thread_id_str); + let dropped = ctx + .dispatcher + .cancel_buffered_thread(event.platform.as_str(), thread_id_str); let msg = match (ctx.router.pool().reset_session(&thread_key).await, dropped) { (Ok(()), 0) => "🔄 Session reset. Start a new conversation!".to_string(), - (Ok(()), n) => format!("🔄 Session reset. Dropped {n} buffered message(s). Start a new conversation!"), + (Ok(()), n) => format!( + "🔄 Session reset. Dropped {n} buffered message(s). Start a new conversation!" + ), (Err(_), 0) => "⚠️ No active session to reset.".to_string(), - (Err(_), n) => format!("🔄 Dropped {n} buffered message(s). No active session to reset."), + (Err(_), n) => { + format!("🔄 Dropped {n} buffered message(s). No active session to reset.") + } }; let _ = ctx.adapter.send_message(&channel, &msg).await; return Ok(false); } if trimmed == "/cancel" { - let thread_key = format!("{}:{}", event.platform, event.channel.thread_id.as_deref().unwrap_or(&event.channel.id)); + let thread_key = format!( + "{}:{}", + event.platform, + event + .channel + .thread_id + .as_deref() + .unwrap_or(&event.channel.id) + ); let msg = match ctx.router.pool().cancel_session(&thread_key).await { Ok(()) => "🛑 Cancel signal sent.".to_string(), Err(e) => format!("⚠️ {e}"), @@ -1631,7 +1672,15 @@ pub async fn process_gateway_event( return Ok(false); } { - let thread_key = format!("{}:{}", event.platform, event.channel.thread_id.as_deref().unwrap_or(&event.channel.id)); + let thread_key = format!( + "{}:{}", + event.platform, + event + .channel + .thread_id + .as_deref() + .unwrap_or(&event.channel.id) + ); if let Some(msg) = handle_config_command(trimmed, &ctx.router, &thread_key).await { let _ = ctx.adapter.send_message(&channel, &msg).await; return Ok(false); @@ -1645,32 +1694,26 @@ pub async fn process_gateway_event( let sender_id = event.sender.id.clone(); tokio::spawn(async move { - let thread_channel = if event.channel.channel_type == "supergroup" - && channel.thread_id.is_none() - { - let title = crate::format::shorten_thread_name(&prompt); - match adapter.create_thread(&channel, &trigger_msg, &title).await { - Ok(tc) => tc, - Err(e) => { - tracing::warn!("create_thread failed, replying in channel: {e}"); - channel.clone() + let thread_channel = + if event.channel.channel_type == "supergroup" && channel.thread_id.is_none() { + let title = crate::format::shorten_thread_name(&prompt); + match adapter.create_thread(&channel, &trigger_msg, &title).await { + Ok(tc) => tc, + Err(e) => { + tracing::warn!("create_thread failed, replying in channel: {e}"); + channel.clone() + } } - } - } else { - channel.clone() - }; + } else { + channel.clone() + }; let thread_id = thread_channel .thread_id .as_deref() .unwrap_or(&thread_channel.channel_id); - let thread_key = dispatcher.key( - &thread_channel.platform, - thread_id, - &sender_id, - ); - let estimated_tokens = - crate::dispatch::estimate_tokens(&prompt, &extra_blocks); + let thread_key = dispatcher.key(&thread_channel.platform, thread_id, &sender_id); + let estimated_tokens = crate::dispatch::estimate_tokens(&prompt, &extra_blocks); let buf_msg = crate::dispatch::BufferedMessage { sender_json, sender_name, @@ -1822,7 +1865,14 @@ mod tests { assert!(!echo_allowed(key), "still throttled within the window"); } - fn make_event(is_bot: bool, sender_id: &str, channel_id: &str, channel_type: &str, thread_id: Option<&str>, mentions: Vec<&str>) -> GatewayEvent { + fn make_event( + is_bot: bool, + sender_id: &str, + channel_id: &str, + channel_type: &str, + thread_id: Option<&str>, + mentions: Vec<&str>, + ) -> GatewayEvent { serde_json::from_value(serde_json::json!({ "schema": "openab.gateway.event.v1", "event_id": "evt1", @@ -1833,10 +1883,15 @@ mod tests { "content": { "type": "text", "text": "hello" }, "mentions": mentions, "message_id": "msg1" - })).unwrap() + })) + .unwrap() } - fn default_filter<'a>(allowed_channels: &'a HashSet, allowed_users: &'a HashSet, trusted_bot_ids: &'a HashSet) -> EventFilterParams<'a> { + fn default_filter<'a>( + allowed_channels: &'a HashSet, + allowed_users: &'a HashSet, + trusted_bot_ids: &'a HashSet, + ) -> EventFilterParams<'a> { EventFilterParams { allow_all_channels: true, allowed_channels, diff --git a/crates/openab-core/src/hooks.rs b/crates/openab-core/src/hooks.rs index 164ce39d0..51a63ab7d 100644 --- a/crates/openab-core/src/hooks.rs +++ b/crates/openab-core/src/hooks.rs @@ -111,9 +111,10 @@ fn write_temp_script(name: &str, content: &str) -> anyhow::Result { let mut f = builder.tempfile()?; f.write_all(content.as_bytes())?; - let path = f.into_temp_path().keep().map_err(|e| { - anyhow::anyhow!("failed to persist temp script: {}", e.error) - })?; + let path = f + .into_temp_path() + .keep() + .map_err(|e| anyhow::anyhow!("failed to persist temp script: {}", e.error))?; Ok(path) } @@ -128,9 +129,7 @@ async fn fetch_and_verify(url: &str, expected_hex: &str) -> anyhow::Result MAX_SCRIPT_SIZE { - anyhow::bail!( - "hook script too large: {content_length} bytes (max {MAX_SCRIPT_SIZE})" - ); + anyhow::bail!("hook script too large: {content_length} bytes (max {MAX_SCRIPT_SIZE})"); } let body = resp.bytes().await?; if body.len() > MAX_SCRIPT_SIZE { diff --git a/crates/openab-core/src/lib.rs b/crates/openab-core/src/lib.rs index 0e61e7cb2..52c3c3925 100644 --- a/crates/openab-core/src/lib.rs +++ b/crates/openab-core/src/lib.rs @@ -1,14 +1,15 @@ pub mod acp; -pub mod adapter; #[cfg(feature = "acp-mcp")] pub mod acp_mcp; -pub mod redact; +pub mod adapter; pub mod bot_turns; pub mod config; pub mod cron; pub mod directives; pub mod dispatch; pub mod error_display; +#[cfg(feature = "filestore")] +pub mod filestore; pub mod format; pub mod gateway; pub mod hooks; @@ -17,9 +18,8 @@ pub mod media; pub mod multibot_cache; #[cfg(feature = "pre-seed")] pub mod pre_seed; -#[cfg(feature = "filestore")] -pub mod filestore; pub mod reactions; +pub mod redact; #[cfg(feature = "discord")] pub mod remind; pub mod secrets; @@ -28,9 +28,9 @@ pub mod stt; pub mod timestamp; pub mod trust; -#[cfg(feature = "discord")] -pub mod discord; #[cfg(feature = "discord")] pub mod ambient; +#[cfg(feature = "discord")] +pub mod discord; #[cfg(feature = "slack")] pub mod slack; diff --git a/crates/openab-core/src/media.rs b/crates/openab-core/src/media.rs index cadc56445..42b1da6d1 100644 --- a/crates/openab-core/src/media.rs +++ b/crates/openab-core/src/media.rs @@ -98,10 +98,7 @@ fn hex_prefix(body: &[u8]) -> String { /// (`image/*`) because CDNs commonly serve any file type as `application/octet-stream`; /// rejecting that header would silently break real downloads. The magic-byte check /// examines the actual bytes regardless of what the server claims. -fn validate_image_response( - content_type: Option<&str>, - body: &[u8], -) -> Result<(), MediaFetchError> { +fn validate_image_response(content_type: Option<&str>, body: &[u8]) -> Result<(), MediaFetchError> { // Reject explicitly-text responses early (e.g. Slack HTML login page at HTTP 200). // application/octet-stream and other generic types pass through to magic-byte check. if let Some(ct) = content_type { @@ -344,7 +341,11 @@ pub async fn download_and_transcribe( }; if bytes.len() as u64 > MAX_SIZE { - error!(filename, size = bytes.len(), "downloaded audio exceeds 25MB limit"); + error!( + filename, + size = bytes.len(), + "downloaded audio exceeds 25MB limit" + ); return None; } @@ -574,15 +575,27 @@ async fn download_text_file_inner( let stream = Box::pin(resp.bytes_stream()); let upload_result = tokio::time::timeout( STREAM_TIMEOUT, - fs.stream_upload_and_presign(filename, stream, content_length, Some("text/plain; charset=utf-8")), + fs.stream_upload_and_presign( + filename, + stream, + content_length, + Some("text/plain; charset=utf-8"), + ), ) .await; return match upload_result { Ok(Ok((presigned_url, actual_bytes))) => { let hint = crate::filestore::format_filestore_hint( - filename, actual_bytes, &presigned_url, fs.presigned_ttl_secs(), + filename, + actual_bytes, + &presigned_url, + fs.presigned_ttl_secs(), + ); + tracing::info!( + filename, + size = actual_bytes, + "text file streamed to filestore (inline fallback)" ); - tracing::info!(filename, size = actual_bytes, "text file streamed to filestore (inline fallback)"); Some((ContentBlock::Text { text: hint }, 0)) } Ok(Err(e)) => { @@ -596,7 +609,10 @@ async fn download_text_file_inner( Some((ContentBlock::Text { text: hint }, 0)) } Err(_) => { - tracing::error!(filename, "filestore stream upload timed out (inline fallback)"); + tracing::error!( + filename, + "filestore stream upload timed out (inline fallback)" + ); let size_kb = content_length / 1024; let hint = format!( "[File: {safe_filename}]\n\ @@ -739,7 +755,12 @@ async fn download_and_upload_to_filestore( // Cap file size to prevent abuse (configurable, default 250 MB, max 500 MB). let max_size = filestore.max_file_size(); if size > max_size { - tracing::warn!(filename, size, max = max_size, "text file exceeds filestore size limit, skipping"); + tracing::warn!( + filename, + size, + max = max_size, + "text file exceeds filestore size limit, skipping" + ); return None; } @@ -789,7 +810,12 @@ async fn download_and_upload_to_filestore( let upload_result = tokio::time::timeout( STREAM_TIMEOUT, - filestore.stream_upload_and_presign(filename, stream, size, Some("text/plain; charset=utf-8")), + filestore.stream_upload_and_presign( + filename, + stream, + size, + Some("text/plain; charset=utf-8"), + ), ) .await; @@ -801,7 +827,11 @@ async fn download_and_upload_to_filestore( &presigned_url, filestore.presigned_ttl_secs(), ); - tracing::info!(filename, size = actual_bytes, "text file streamed to filestore"); + tracing::info!( + filename, + size = actual_bytes, + "text file streamed to filestore" + ); Some((ContentBlock::Text { text: hint }, 0)) } Ok(Err(e)) => { @@ -856,7 +886,12 @@ pub async fn download_and_upload_any_file( ) -> Option<(ContentBlock, u64)> { let max_size = filestore.max_file_size(); if size > max_size { - tracing::warn!(filename, size, max = max_size, "file exceeds filestore size limit, skipping"); + tracing::warn!( + filename, + size, + max = max_size, + "file exceeds filestore size limit, skipping" + ); return None; } @@ -881,7 +916,12 @@ pub async fn download_and_upload_any_file( // Content-Length pre-check if let Some(content_length) = resp.content_length() { if content_length > max_size { - tracing::warn!(filename, content_length, max = max_size, "Content-Length exceeds filestore limit"); + tracing::warn!( + filename, + content_length, + max = max_size, + "Content-Length exceeds filestore limit" + ); return None; } } @@ -922,7 +962,12 @@ pub async fn download_and_upload_any_file( Note: this URL expires in {} minutes.", filestore.presigned_ttl_secs() / 60 ); - tracing::info!(filename, mime, size = actual_bytes, "file uploaded to filestore (any-file path)"); + tracing::info!( + filename, + mime, + size = actual_bytes, + "file uploaded to filestore (any-file path)" + ); Some((ContentBlock::Text { text: hint }, 0)) } Ok(Err(e)) => { @@ -984,7 +1029,11 @@ async fn upload_bytes_to_filestore( &presigned_url, filestore.presigned_ttl_secs(), ); - tracing::info!(filename, size = actual_size, "text file uploaded to filestore"); + tracing::info!( + filename, + size = actual_size, + "text file uploaded to filestore" + ); // Return 0 inline bytes — the file is stored externally, not inlined // into the prompt. This prevents it from counting against the // caller's aggregate size cap. diff --git a/crates/openab-core/src/multibot_cache.rs b/crates/openab-core/src/multibot_cache.rs index a9bd8c82e..05d01e027 100644 --- a/crates/openab-core/src/multibot_cache.rs +++ b/crates/openab-core/src/multibot_cache.rs @@ -61,7 +61,9 @@ impl MultibotCache { threads.clone() }; let path = self.path.clone(); - tokio::task::spawn_blocking(move || persist(&path, &snapshot)).await.ok(); + tokio::task::spawn_blocking(move || persist(&path, &snapshot)) + .await + .ok(); } } diff --git a/crates/openab-core/src/redact.rs b/crates/openab-core/src/redact.rs index 85954da12..7a23a7510 100644 --- a/crates/openab-core/src/redact.rs +++ b/crates/openab-core/src/redact.rs @@ -24,9 +24,14 @@ /// identifier in a log at all. pub fn redact_session_ids(s: &str) -> String { s.split(':') - .map(|seg| match seg.strip_prefix("acp_").or_else(|| seg.strip_prefix("sess_")) { - Some(uuid) if !uuid.is_empty() => hash_tag(uuid), - _ => seg.to_string(), + .map(|seg| { + match seg + .strip_prefix("acp_") + .or_else(|| seg.strip_prefix("sess_")) + { + Some(uuid) if !uuid.is_empty() => hash_tag(uuid), + _ => seg.to_string(), + } }) .collect::>() .join(":") @@ -53,7 +58,10 @@ mod tests { fn both_encodings_hash_alike_and_everything_else_passes_through() { let u = "00000000-0000-0000-0000-000000000000"; let tag = redact_session_ids(&format!("acp_{u}")); - assert!(tag.starts_with('#') && tag.len() == 9, "expected #<8hex>, got {tag}"); + assert!( + tag.starts_with('#') && tag.len() == 9, + "expected #<8hex>, got {tag}" + ); assert_eq!( redact_session_ids(&format!("sess_{u}")), tag, @@ -65,10 +73,22 @@ mod tests { format!("acp:{tag}"), "a : pool key must redact the id half and keep the platform greppable" ); - assert_eq!(redact_session_ids("1234567890"), "1234567890", "public ids stay greppable"); - assert_eq!(redact_session_ids("-"), "-", "the no-session sentinel is not a session"); + assert_eq!( + redact_session_ids("1234567890"), + "1234567890", + "public ids stay greppable" + ); + assert_eq!( + redact_session_ids("-"), + "-", + "the no-session sentinel is not a session" + ); assert_eq!(redact_session_ids(""), "", "empty in, empty out"); - assert_eq!(redact_session_ids("acp_"), "acp_", "a bare prefix carries no uuid to hide"); + assert_eq!( + redact_session_ids("acp_"), + "acp_", + "a bare prefix carries no uuid to hide" + ); assert_eq!( redact_session_ids("discord:1234567890"), "discord:1234567890", diff --git a/crates/openab-core/src/remind.rs b/crates/openab-core/src/remind.rs index 9472c53d8..b4c43495b 100644 --- a/crates/openab-core/src/remind.rs +++ b/crates/openab-core/src/remind.rs @@ -108,7 +108,9 @@ pub fn sanitize_message(msg: &str) -> String { /// Validate reminder message length. pub fn validate_message(msg: &str) -> Result<(), String> { if msg.len() > MAX_MESSAGE_LEN { - Err(format!("message too long (max {MAX_MESSAGE_LEN} characters)")) + Err(format!( + "message too long (max {MAX_MESSAGE_LEN} characters)" + )) } else { Ok(()) } @@ -130,7 +132,9 @@ pub fn parse_delay(input: &str) -> Result { if ch.is_ascii_digit() { num_buf.push(ch); } else { - let n: u64 = num_buf.parse().map_err(|_| format!("invalid number in delay: {input}"))?; + let n: u64 = num_buf + .parse() + .map_err(|_| format!("invalid number in delay: {input}"))?; num_buf.clear(); let multiplier = match ch { 'm' => 60, @@ -144,7 +148,9 @@ pub fn parse_delay(input: &str) -> Result { // Handle bare number (default to minutes) if !num_buf.is_empty() { - let n: u64 = num_buf.parse().map_err(|_| format!("invalid number in delay: {input}"))?; + let n: u64 = num_buf + .parse() + .map_err(|_| format!("invalid number in delay: {input}"))?; total_secs += n * 60; // default unit = minutes } @@ -164,18 +170,24 @@ pub fn format_delay(secs: u64) -> String { let h = (secs % 86400) / 3600; let m = (secs % 3600) / 60; let mut parts = Vec::new(); - if d > 0 { parts.push(format!("{d}d")); } - if h > 0 { parts.push(format!("{h}h")); } - if m > 0 { parts.push(format!("{m}m")); } - if parts.is_empty() { "< 1m".into() } else { parts.join(" ") } + if d > 0 { + parts.push(format!("{d}d")); + } + if h > 0 { + parts.push(format!("{h}h")); + } + if m > 0 { + parts.push(format!("{m}m")); + } + if parts.is_empty() { + "< 1m".into() + } else { + parts.join(" ") + } } /// Spawn a tokio task that fires the reminder after the delay. -pub fn schedule_reminder( - http: Arc, - store: ReminderStore, - reminder: Reminder, -) { +pub fn schedule_reminder(http: Arc, store: ReminderStore, reminder: Reminder) { let now = Utc::now(); let delay = if reminder.fire_at > now { (reminder.fire_at - now).to_std().unwrap_or_default() @@ -370,9 +382,18 @@ mod tests { #[test] fn sanitize_message_strips_everyone_here() { - assert_eq!(sanitize_message("hello @everyone"), "hello @\u{200b}everyone"); - assert_eq!(sanitize_message("hey @here check"), "hey @\u{200b}here check"); - assert_eq!(sanitize_message("@everyone @here"), "@\u{200b}everyone @\u{200b}here"); + assert_eq!( + sanitize_message("hello @everyone"), + "hello @\u{200b}everyone" + ); + assert_eq!( + sanitize_message("hey @here check"), + "hey @\u{200b}here check" + ); + assert_eq!( + sanitize_message("@everyone @here"), + "@\u{200b}everyone @\u{200b}here" + ); } #[test] diff --git a/crates/openab-core/src/secrets.rs b/crates/openab-core/src/secrets.rs index e6a7967fd..bb5b06226 100644 --- a/crates/openab-core/src/secrets.rs +++ b/crates/openab-core/src/secrets.rs @@ -23,9 +23,9 @@ pub async fn resolve(cfg: &SecretsConfig) -> anyhow::Result { let value = if uri.starts_with("aws-sm://") { #[cfg(feature = "secrets-aws")] { - let client = aws_client.as_ref().ok_or_else(|| { - anyhow::anyhow!("secret '{key}': AWS client not initialized") - })?; + let client = aws_client + .as_ref() + .ok_or_else(|| anyhow::anyhow!("secret '{key}': AWS client not initialized"))?; resolve_aws_sm(key, uri, client).await? } #[cfg(not(feature = "secrets-aws"))] @@ -81,12 +81,16 @@ async fn resolve_aws_sm( .await .map_err(|e| { error!(secret = key, secret_id = %secret_id, "AWS Secrets Manager error"); - anyhow::anyhow!("secret '{key}': failed to fetch '{secret_id}' from AWS Secrets Manager: {e}") + anyhow::anyhow!( + "secret '{key}': failed to fetch '{secret_id}' from AWS Secrets Manager: {e}" + ) })?; - let secret_string = resp - .secret_string() - .ok_or_else(|| anyhow::anyhow!("secret '{key}': '{secret_id}' has no string value (binary secrets not supported)"))?; + let secret_string = resp.secret_string().ok_or_else(|| { + anyhow::anyhow!( + "secret '{key}': '{secret_id}' has no string value (binary secrets not supported)" + ) + })?; // Parse as JSON and extract the key let json: serde_json::Value = serde_json::from_str(secret_string) @@ -95,7 +99,9 @@ async fn resolve_aws_sm( let value = json .get(&json_key) .and_then(|v| v.as_str()) - .ok_or_else(|| anyhow::anyhow!("secret '{key}': JSON key '{json_key}' not found in '{secret_id}'"))?; + .ok_or_else(|| { + anyhow::anyhow!("secret '{key}': JSON key '{json_key}' not found in '{secret_id}'") + })?; Ok(value.to_owned()) } @@ -117,9 +123,9 @@ fn parse_aws_sm_uri(uri: &str) -> Option<(String, String)> { async fn resolve_exec(key: &str, uri: &str, cfg: &SecretsConfig) -> anyhow::Result { let rest = uri.strip_prefix("exec://").unwrap(); let mut parts_iter = rest.splitn(3, ' '); - let script = parts_iter.next().ok_or_else(|| { - anyhow::anyhow!("secret '{key}': exec:// URI missing script path") - })?; + let script = parts_iter + .next() + .ok_or_else(|| anyhow::anyhow!("secret '{key}': exec:// URI missing script path"))?; if script.is_empty() { anyhow::bail!("secret '{key}': exec:// URI has empty script path"); } @@ -240,9 +246,13 @@ mod tests { #[cfg(feature = "secrets-aws")] #[test] fn parse_aws_sm_uri_with_arn() { - let uri = "aws-sm://arn:aws:secretsmanager:us-east-1:123456789:secret:my-secret-abc123#api_key"; + let uri = + "aws-sm://arn:aws:secretsmanager:us-east-1:123456789:secret:my-secret-abc123#api_key"; let (id, key) = parse_aws_sm_uri(uri).unwrap(); - assert_eq!(id, "arn:aws:secretsmanager:us-east-1:123456789:secret:my-secret-abc123"); + assert_eq!( + id, + "arn:aws:secretsmanager:us-east-1:123456789:secret:my-secret-abc123" + ); assert_eq!(key, "api_key"); } @@ -349,8 +359,7 @@ mod tests { exec: ExecSecretsConfig { timeout_seconds: 5 }, refs: HashMap::new(), }; - let result = - resolve_exec("test", "exec:///nonexistent/script arg1 arg2", &cfg).await; + let result = resolve_exec("test", "exec:///nonexistent/script arg1 arg2", &cfg).await; let err = result.unwrap_err().to_string(); assert!(err.contains("not found")); assert!(err.contains("pre_boot")); @@ -406,7 +415,10 @@ mod tests { }; let result = resolve(&cfg).await; assert!(result.is_err()); - assert!(result.unwrap_err().to_string().contains("unrecognized URI scheme")); + assert!(result + .unwrap_err() + .to_string() + .contains("unrecognized URI scheme")); } #[cfg(feature = "secrets-aws")] @@ -430,7 +442,9 @@ mod tests { refs: HashMap::new(), }; // /usr/bin/env prints all env vars; grep for our dummy var - let result = resolve_exec("test", "exec:///usr/bin/env", &cfg).await.unwrap(); + let result = resolve_exec("test", "exec:///usr/bin/env", &cfg) + .await + .unwrap(); assert!( !result.contains("OPENAB_TEST_LEAKED_SECRET"), "exec script should not see unrelated env vars" @@ -452,7 +466,10 @@ mod tests { // sleep 10 will be killed after 1s timeout let result = resolve_exec("test", "exec:///bin/sleep 10", &cfg).await; let err = result.unwrap_err().to_string(); - assert!(err.contains("timed out"), "expected timeout error, got: {err}"); + assert!( + err.contains("timed out"), + "expected timeout error, got: {err}" + ); } #[test] @@ -471,8 +488,8 @@ args = ["--key", "${secrets.key}"] "#; let substituted = substitute(raw, &secrets); // Verify the substituted text is valid TOML that parses correctly - let cfg: crate::config::Config = toml::from_str(&substituted) - .expect("substituted config should be valid TOML"); + let cfg: crate::config::Config = + toml::from_str(&substituted).expect("substituted config should be valid TOML"); assert_eq!(cfg.discord.unwrap().bot_token, "xoxb-secret-value"); assert_eq!(cfg.agent.args[1], "sk-with\"special\\chars"); } diff --git a/crates/openab-core/src/slack.rs b/crates/openab-core/src/slack.rs index 04af2d9bc..012b85caa 100644 --- a/crates/openab-core/src/slack.rs +++ b/crates/openab-core/src/slack.rs @@ -144,7 +144,6 @@ impl SlackAdapter { self.multibot_cache.mark_multibot(thread_ts).await; } - /// Insert a stream entry, bounding the map so aborted turns (begin without a /// matching finish) can't leak unboundedly. Normal lifecycle: stream_begin /// inserts, stream_finish removes. @@ -451,7 +450,8 @@ impl ChatAdapter for SlackAdapter { // instead of failing outright. Err(e) if is_block_payload_rejected(&e) => { warn!(error = %e, "markdown block rejected; retrying chat.postMessage text-only"); - let fallback = build_post_message_text_only(&channel.channel_id, thread_ts, content); + let fallback = + build_post_message_text_only(&channel.channel_id, thread_ts, content); self.api_post("chat.postMessage", fallback).await? } Err(e) => return Err(e), @@ -591,7 +591,10 @@ impl ChatAdapter for SlackAdapter { if let Some(ts) = resp["ts"].as_str() { self.insert_stream( ts.to_string(), - StreamEntry { active: true, degraded_buf: String::new() }, + StreamEntry { + active: true, + degraded_buf: String::new(), + }, ) .await; return Ok(make_ref(ts.to_string())); @@ -605,14 +608,20 @@ impl ChatAdapter for SlackAdapter { } else { // Expected for bot-authored turns (no recipient bound) and non-user // triggers, so warn! rather than error! to avoid on-call noise. - warn!(thread_ts, "no recipient for turn; falling back to post+edit"); + warn!( + thread_ts, + "no recipient for turn; falling back to post+edit" + ); } // Degraded fallback: plain placeholder via send_message; mark inactive. let msg = self.send_message(channel, "…").await?; self.insert_stream( msg.message_id.clone(), - StreamEntry { active: false, degraded_buf: String::new() }, + StreamEntry { + active: false, + degraded_buf: String::new(), + }, ) .await; Ok(msg) @@ -1596,14 +1605,11 @@ async fn handle_message( size, Some(bot_token), filestore, - ).await; + ) + .await; #[cfg(not(feature = "filestore"))] - let text_file_result = media::download_and_read_text_file( - url, - filename, - size, - Some(bot_token), - ).await; + let text_file_result = + media::download_and_read_text_file(url, filename, size, Some(bot_token)).await; if let Some((block, actual_bytes)) = text_file_result { if text_file_bytes + actual_bytes > TEXT_TOTAL_CAP { debug!( @@ -1652,7 +1658,9 @@ async fn handle_message( Some(mimetype), Some(bot_token), fs, - ).await { + ) + .await + { extra_blocks.push(block); } } @@ -1676,9 +1684,7 @@ async fn handle_message( warn!(filename, error = %e, "image post-processing failed"); failed_image_files.push(filename.to_string()); } - Err(media::MediaFetchError::HttpStatus(status)) - if status.is_client_error() => - { + Err(media::MediaFetchError::HttpStatus(status)) if status.is_client_error() => { warn!(filename, %status, "image download denied"); failed_image_files.push(filename.to_string()); } @@ -1880,7 +1886,10 @@ fn strip_mime_params(mimetype: &str) -> &str { /// mrkdwn delimiters; without escaping, `` or `` `<@U123>` `` would render /// as mentions or @-here pings. pub(crate) fn sanitize_slack_filename(s: &str) -> String { - s.replace('&', "&").replace('`', "'").replace('<', "(").replace('>', ")") + s.replace('&', "&") + .replace('`', "'") + .replace('<', "(") + .replace('>', ")") } /// Returns `true` if `text` contains a Slack user mention for `uid`. @@ -1891,8 +1900,12 @@ pub(crate) fn sanitize_slack_filename(s: &str) -> String { /// misses it. fn text_mentions_uid(text: &str, uid: &str) -> bool { let prefix = format!("<@{uid}"); - text.match_indices(&prefix) - .any(|(i, _)| matches!(text.as_bytes().get(i + prefix.len()), Some(b'>') | Some(b'|'))) + text.match_indices(&prefix).any(|(i, _)| { + matches!( + text.as_bytes().get(i + prefix.len()), + Some(b'>') | Some(b'|') + ) + }) } fn bot_id_matches_trusted( @@ -2061,7 +2074,12 @@ fn markdown_to_mrkdwn(text: &str) -> String { text.into_owned() } -fn build_start_stream_body(channel: &str, thread_ts: &str, user_id: &str, team_id: &str) -> serde_json::Value { +fn build_start_stream_body( + channel: &str, + thread_ts: &str, + user_id: &str, + team_id: &str, +) -> serde_json::Value { serde_json::json!({ "channel": channel, "thread_ts": thread_ts, @@ -2141,13 +2159,29 @@ mod tests { #[tokio::test] async fn degraded_stream_append_accumulates() { - let adapter = SlackAdapter::new("xoxb-test".into(), std::time::Duration::from_secs(60), AllowBots::Off, true, crate::multibot_cache::MultibotCache::load("/dev/null".into()), true); + let adapter = SlackAdapter::new( + "xoxb-test".into(), + std::time::Duration::from_secs(60), + AllowBots::Off, + true, + crate::multibot_cache::MultibotCache::load("/dev/null".into()), + true, + ); adapter.streams.lock().await.insert( "TS".into(), - StreamEntry { active: false, degraded_buf: String::new() }, + StreamEntry { + active: false, + degraded_buf: String::new(), + }, + ); + assert_eq!( + adapter.accumulate_degraded("TS", "a").await.as_deref(), + Some("a") + ); + assert_eq!( + adapter.accumulate_degraded("TS", "b").await.as_deref(), + Some("ab") ); - assert_eq!(adapter.accumulate_degraded("TS", "a").await.as_deref(), Some("a")); - assert_eq!(adapter.accumulate_degraded("TS", "b").await.as_deref(), Some("ab")); // missing stream is not resurrected: assert_eq!(adapter.accumulate_degraded("MISSING", "x").await, None); } @@ -2238,7 +2272,10 @@ mod tests { #[test] fn mentions_uid_labelled_form_mid_sentence() { - assert!(text_mentions_uid("please ask <@U123BOT|handle> to run", "U123BOT")); + assert!(text_mentions_uid( + "please ask <@U123BOT|handle> to run", + "U123BOT" + )); } #[test] @@ -2356,7 +2393,10 @@ mod tests { #[test] fn sanitize_leaves_normal_filename_unchanged() { assert_eq!(sanitize_slack_filename("photo.png"), "photo.png"); - assert_eq!(sanitize_slack_filename("my file (1).jpg"), "my file (1).jpg"); + assert_eq!( + sanitize_slack_filename("my file (1).jpg"), + "my file (1).jpg" + ); } #[test] @@ -2374,10 +2414,7 @@ mod tests { #[test] fn sanitize_combined_injection_attempt() { // A filename constructed to inject a Slack @here ping. - assert_eq!( - sanitize_slack_filename("``"), - "'(!here)'" - ); + assert_eq!(sanitize_slack_filename("``"), "'(!here)'"); } #[test] @@ -2386,8 +2423,14 @@ mod tests { // "<@here>" would round-trip back to "<@here>" and trigger a mention // ping if & is not escaped. The & must be escaped first so downstream // Slack entity decoding cannot reconstruct a mrkdwn delimiter. - assert_eq!(sanitize_slack_filename("<@here>"), "&lt;@here&gt;"); - assert_eq!(sanitize_slack_filename("file&name.png"), "file&name.png"); + assert_eq!( + sanitize_slack_filename("<@here>"), + "&lt;@here&gt;" + ); + assert_eq!( + sanitize_slack_filename("file&name.png"), + "file&name.png" + ); } // --- strip_mime_params tests --- @@ -2427,11 +2470,7 @@ mod tests { #[test] fn trusted_bot_ids_accepts_resolved_bot_user_id() { let trusted = HashSet::from(["U123BOT".to_string()]); - assert!(bot_id_matches_trusted( - &trusted, - "B123BOT", - Some("U123BOT") - )); + assert!(bot_id_matches_trusted(&trusted, "B123BOT", Some("U123BOT"))); } #[test] @@ -2450,7 +2489,14 @@ mod tests { #[test] fn streaming_per_thread() { let ttl = std::time::Duration::from_secs(300); - let adapter = SlackAdapter::new("xoxb-test".into(), ttl, AllowBots::Mentions, false, crate::multibot_cache::MultibotCache::load("/dev/null".into()), true); + let adapter = SlackAdapter::new( + "xoxb-test".into(), + ttl, + AllowBots::Mentions, + false, + crate::multibot_cache::MultibotCache::load("/dev/null".into()), + true, + ); assert!( adapter.use_streaming(false), @@ -2467,22 +2513,70 @@ mod tests { let ttl = std::time::Duration::from_secs(60); // assistant_mode=true → status API on; native streaming on (no other bot), // off when another bot is present; post+edit streaming on regardless. - let adapter = SlackAdapter::new("xoxb-test".into(), ttl, AllowBots::Off, true, crate::multibot_cache::MultibotCache::load("/dev/null".into()), true); - assert!(adapter.uses_assistant_status(), "assistant_mode enables status API"); - assert!(adapter.use_streaming(false), "post+edit streaming on when no other bot"); - assert!(adapter.uses_native_streaming(false), "native streaming on when no other bot"); - assert!(!adapter.uses_native_streaming(true), "other bot present disables native"); + let adapter = SlackAdapter::new( + "xoxb-test".into(), + ttl, + AllowBots::Off, + true, + crate::multibot_cache::MultibotCache::load("/dev/null".into()), + true, + ); + assert!( + adapter.uses_assistant_status(), + "assistant_mode enables status API" + ); + assert!( + adapter.use_streaming(false), + "post+edit streaming on when no other bot" + ); + assert!( + adapter.uses_native_streaming(false), + "native streaming on when no other bot" + ); + assert!( + !adapter.uses_native_streaming(true), + "other bot present disables native" + ); // assistant_mode=false → no status API, no native streaming; post+edit still streams. - let adapter2 = SlackAdapter::new("xoxb-test".into(), ttl, AllowBots::Off, false, crate::multibot_cache::MultibotCache::load("/dev/null".into()), true); + let adapter2 = SlackAdapter::new( + "xoxb-test".into(), + ttl, + AllowBots::Off, + false, + crate::multibot_cache::MultibotCache::load("/dev/null".into()), + true, + ); assert!(!adapter2.uses_assistant_status()); - assert!(adapter2.use_streaming(false), "post+edit streaming independent of assistant_mode"); - assert!(!adapter2.uses_native_streaming(false), "native streaming requires assistant_mode"); + assert!( + adapter2.use_streaming(false), + "post+edit streaming independent of assistant_mode" + ); + assert!( + !adapter2.uses_native_streaming(false), + "native streaming requires assistant_mode" + ); // streaming=false → send-once: neither post+edit nor native, even alone. - let adapter3 = SlackAdapter::new("xoxb-test".into(), ttl, AllowBots::Off, true, crate::multibot_cache::MultibotCache::load("/dev/null".into()), false); - assert!(!adapter3.use_streaming(false), "streaming=false forces send-once (no post+edit)"); - assert!(!adapter3.uses_native_streaming(false), "streaming=false disables native even with assistant_mode"); - assert!(adapter3.uses_assistant_status(), "streaming switch does not affect assistant status API"); + let adapter3 = SlackAdapter::new( + "xoxb-test".into(), + ttl, + AllowBots::Off, + true, + crate::multibot_cache::MultibotCache::load("/dev/null".into()), + false, + ); + assert!( + !adapter3.use_streaming(false), + "streaming=false forces send-once (no post+edit)" + ); + assert!( + !adapter3.uses_native_streaming(false), + "streaming=false disables native even with assistant_mode" + ); + assert!( + adapter3.uses_assistant_status(), + "streaming switch does not affect assistant status API" + ); } /// chat.postMessage body carries Block Kit `markdown` blocks with the raw @@ -2495,7 +2589,10 @@ mod tests { assert_eq!(b["blocks"][0]["type"], "markdown"); // Raw markdown preserved — heading is NOT flattened to `*Heading*`. assert_eq!(b["blocks"][0]["text"], "## Heading\n- item"); - assert!(b["text"].is_string(), "text fallback present for a11y/notifs"); + assert!( + b["text"].is_string(), + "text fallback present for a11y/notifs" + ); } /// thread_ts is omitted (top-level post) when the channel has no thread. @@ -2536,7 +2633,14 @@ mod tests { #[test] fn typical_long_table_stays_in_one_chunk() { let ttl = std::time::Duration::from_secs(300); - let adapter = SlackAdapter::new("xoxb-test".into(), ttl, AllowBots::Mentions, true, crate::multibot_cache::MultibotCache::load("/dev/null".into()), true); + let adapter = SlackAdapter::new( + "xoxb-test".into(), + ttl, + AllowBots::Mentions, + true, + crate::multibot_cache::MultibotCache::load("/dev/null".into()), + true, + ); let limit = adapter.message_limit(); assert_eq!(limit, MARKDOWN_BLOCK_LIMIT); let mut table = String::from("| col a | col b | col c |\n|---|---|---|\n"); @@ -2588,7 +2692,9 @@ mod tests { // Exact error-code match, not substring: a future code that merely // contains `invalid_blocks` must NOT trigger a text-only retry. assert!( - !is_block_payload_rejected(&anyhow!("Slack API chat.postMessage: invalid_blocks_field")), + !is_block_payload_rejected(&anyhow!( + "Slack API chat.postMessage: invalid_blocks_field" + )), "must match the error code exactly, not as a substring" ); } @@ -2598,7 +2704,14 @@ mod tests { #[test] fn slack_renders_native_tables() { let ttl = std::time::Duration::from_secs(300); - let adapter = SlackAdapter::new("xoxb-test".into(), ttl, AllowBots::Mentions, true, crate::multibot_cache::MultibotCache::load("/dev/null".into()), true); + let adapter = SlackAdapter::new( + "xoxb-test".into(), + ttl, + AllowBots::Mentions, + true, + crate::multibot_cache::MultibotCache::load("/dev/null".into()), + true, + ); assert!(adapter.renders_native_tables("slack")); } } @@ -2622,7 +2735,19 @@ mod socket_keepalive_tests { cur }) .collect(); - assert_eq!(seq, vec![1, 2, 4, 8, 16, MAX_BACKOFF_SECS, MAX_BACKOFF_SECS, MAX_BACKOFF_SECS]); + assert_eq!( + seq, + vec![ + 1, + 2, + 4, + 8, + 16, + MAX_BACKOFF_SECS, + MAX_BACKOFF_SECS, + MAX_BACKOFF_SECS + ] + ); assert_eq!(next_backoff(MAX_BACKOFF_SECS), MAX_BACKOFF_SECS); } @@ -2632,9 +2757,15 @@ mod socket_keepalive_tests { fn idle_detects_half_open_at_boundary() { let timeout = Duration::from_secs(IDLE_TIMEOUT_SECS); assert!(!socket_idle(Duration::from_secs(0), timeout)); - assert!(!socket_idle(Duration::from_secs(IDLE_TIMEOUT_SECS - 1), timeout)); + assert!(!socket_idle( + Duration::from_secs(IDLE_TIMEOUT_SECS - 1), + timeout + )); assert!(socket_idle(Duration::from_secs(IDLE_TIMEOUT_SECS), timeout)); - assert!(socket_idle(Duration::from_secs(IDLE_TIMEOUT_SECS + 10), timeout)); + assert!(socket_idle( + Duration::from_secs(IDLE_TIMEOUT_SECS + 10), + timeout + )); } /// Regression test for #1279: a stalled WebSocket handshake during reconnect diff --git a/crates/openab-core/src/trust.rs b/crates/openab-core/src/trust.rs index 5d3ab343b..e34a842e0 100644 --- a/crates/openab-core/src/trust.rs +++ b/crates/openab-core/src/trust.rs @@ -206,10 +206,10 @@ mod tests { fn cfg() -> TrustConfig { // L2 open, explicit allowed channel; L3 with one allowed user. TrustConfig::new( - None, // allow_all_channels → true + None, // allow_all_channels → true ["chan-1".to_string()], - None, // allow_dm → true - None, // allow_all_users → false (deny) + None, // allow_dm → true + None, // allow_all_users → false (deny) ["user-1".to_string()], ) } @@ -225,7 +225,10 @@ mod tests { #[test] fn allowed_user_in_scope_channel_is_allowed() { - assert_eq!(cfg().decide("any-channel", false, "user-1"), Decision::Allow); + assert_eq!( + cfg().decide("any-channel", false, "user-1"), + Decision::Allow + ); } #[test] @@ -239,7 +242,10 @@ mod tests { #[test] fn untrusted_user_in_dm_denied_identity_not_scope() { // DM surface open by default → reaches L3 → identity deny (echo path). - assert_eq!(cfg().decide("dm-chan", true, "stranger"), Decision::DenyIdentity); + assert_eq!( + cfg().decide("dm-chan", true, "stranger"), + Decision::DenyIdentity + ); } #[test] @@ -266,13 +272,25 @@ mod tests { #[test] fn allow_all_users_opens_l3() { - let c = TrustConfig::new(None, std::iter::empty(), None, Some(true), std::iter::empty()); + let c = TrustConfig::new( + None, + std::iter::empty(), + None, + Some(true), + std::iter::empty(), + ); assert_eq!(c.decide("c", false, "anyone"), Decision::Allow); } #[test] fn dm_closed_denies_scope_even_for_allowed_user() { - let c = TrustConfig::new(None, std::iter::empty(), Some(false), None, ["user-1".to_string()]); + let c = TrustConfig::new( + None, + std::iter::empty(), + Some(false), + None, + ["user-1".to_string()], + ); // allowed user, but DM surface disabled → DenyScope (no echo). assert_eq!(c.decide("dm", true, "user-1"), Decision::DenyScope); // same user in a channel (L2 open) → Allow. @@ -291,7 +309,10 @@ mod tests { fn registry_returns_default_for_unknown_platform() { let reg = PlatformTrustConfigs::new(); // unknown platform → default (L3 deny-all) → stranger denied identity. - assert_eq!(reg.decide("mars", "c", false, "stranger"), Decision::DenyIdentity); + assert_eq!( + reg.decide("mars", "c", false, "stranger"), + Decision::DenyIdentity + ); } #[test] @@ -302,15 +323,27 @@ mod tests { TrustConfig::new(None, std::iter::empty(), None, None, ["123".to_string()]), ); assert_eq!(reg.decide("telegram", "c", false, "123"), Decision::Allow); - assert_eq!(reg.decide("telegram", "c", false, "999"), Decision::DenyIdentity); + assert_eq!( + reg.decide("telegram", "c", false, "999"), + Decision::DenyIdentity + ); // unregistered platform still gets deny-all default. - assert_eq!(reg.decide("discord", "c", false, "123"), Decision::DenyIdentity); + assert_eq!( + reg.decide("discord", "c", false, "123"), + Decision::DenyIdentity + ); } #[test] fn empty_sender_is_never_identity_allowed() { // Even with allow_all_users = true, an empty sender_id fails closed. - let open = TrustConfig::new(None, std::iter::empty(), None, Some(true), std::iter::empty()); + let open = TrustConfig::new( + None, + std::iter::empty(), + None, + Some(true), + std::iter::empty(), + ); assert!(!open.identity_allowed("")); assert_eq!(open.decide("c", false, ""), Decision::DenyIdentity); // non-empty still allowed under allow_all_users. diff --git a/crates/openab-gateway/src/adapters/acp_server.rs b/crates/openab-gateway/src/adapters/acp_server.rs index 747132e28..d702e50dd 100644 --- a/crates/openab-gateway/src/adapters/acp_server.rs +++ b/crates/openab-gateway/src/adapters/acp_server.rs @@ -745,10 +745,10 @@ async fn route_client_response( } // Accept a numeric id (what we mint) or a stringified number ("1") from a spec-loose // client, so its responses still correlate to the pending request instead of being dropped. - let Some(id) = raw - .get("id") - .and_then(|v| v.as_u64().or_else(|| v.as_str().and_then(|s| s.parse().ok()))) - else { + let Some(id) = raw.get("id").and_then(|v| { + v.as_u64() + .or_else(|| v.as_str().and_then(|s| s.parse().ok())) + }) else { return false; }; if let Some(tx) = pending.lock().await.remove(&id) { @@ -757,7 +757,10 @@ async fn route_client_response( // `debug!`, not `warn!`: since the tunnel cancels on expiry, a reply arriving after we // gave up is EXPECTED traffic from a well-behaved peer, not a client defect. Logging it at // warn made a correct peer look broken. - debug!(id, "acp: client response arrived after its request was abandoned; discarding"); + debug!( + id, + "acp: client response arrived after its request was abandoned; discarding" + ); } true } @@ -831,7 +834,15 @@ async fn mcp_connect( acp_id: acp_id.to_string(), }) .unwrap(); - let frame = send_request(out_tx, pending, next_id, "mcp/connect", params, timeout_secs).await?; + let frame = send_request( + out_tx, + pending, + next_id, + "mcp/connect", + params, + timeout_secs, + ) + .await?; let result: McpConnectResult = serde_json::from_value(frame_result(frame)?) .map_err(|e| format!("mcp/connect: malformed result: {e}"))?; Ok(result.connection_id) @@ -871,9 +882,16 @@ async fn mcp_disconnect( connection_id: connection_id.to_string(), }) .unwrap(); - send_request(out_tx, pending, next_id, "mcp/disconnect", params, timeout_secs) - .await - .map(|_| ()) + send_request( + out_tx, + pending, + next_id, + "mcp/disconnect", + params, + timeout_secs, + ) + .await + .map(|_| ()) } /// Monotonic attach ordering for last-attach-wins (ADR §6.1). @@ -1007,8 +1025,7 @@ const INNER_MCP_PROTOCOL_VERSION: &str = "2025-06-18"; /// in the abstract, so a transport or framing change can invalidate the set while the method list /// stays identical. A bare list of version strings with no stated reason is how this becomes /// silently wrong later. -const SUPPORTED_INNER_MCP_PROTOCOL_VERSIONS: [&str; 3] = - ["2025-06-18", "2025-03-26", "2024-11-05"]; +const SUPPORTED_INNER_MCP_PROTOCOL_VERSIONS: [&str; 3] = ["2025-06-18", "2025-03-26", "2024-11-05"]; /// Perform the inner MCP handshake on a freshly connected tunnel. /// @@ -1332,8 +1349,17 @@ fn spawn_acp_tunnels( } establish_tasks.push(tokio::spawn(async move { if let Err(e) = establish_and_register_tunnel( - out_tx, pending, next_id, srv.id, srv.name, channel_id, registry, 30, owner, - connection_generation, connection_closed, + out_tx, + pending, + next_id, + srv.id, + srv.name, + channel_id, + registry, + 30, + owner, + connection_generation, + connection_closed, ) .await { @@ -1493,8 +1519,11 @@ async fn handle_acp_connection(state: Arc, socket: WebSocket) { Ok(r) => r, Err(e) => { if !is_notification { - let err_resp = - JsonRpcResponse::error(Value::Null, -32600, format!("Invalid Request: {e}")); + let err_resp = JsonRpcResponse::error( + Value::Null, + -32600, + format!("Invalid Request: {e}"), + ); let _ = out_tx.send(serde_json::to_string(&err_resp).unwrap()); } continue; @@ -1545,9 +1574,9 @@ async fn handle_acp_connection(state: Arc, socket: WebSocket) { continue; } // Required params per schema: { cwd, mcpServers }. - if let Err(msg) = - validate_params::(req.params.as_ref()) - { + if let Err(msg) = validate_params::( + req.params.as_ref(), + ) { let resp = JsonRpcResponse::error(id, -32602, msg); let _ = out_tx.send(serde_json::to_string(&resp).unwrap()); continue; @@ -1564,18 +1593,16 @@ async fn handle_acp_connection(state: Arc, socket: WebSocket) { } // Bound the declaration fan-out BEFORE the session exists or any task is // spawned (R3-F1): an over-declaring request must cost nothing. - let acp_mcp_servers = match accept_acp_servers(parse_acp_mcp_servers( - req.params.as_ref(), - )) { - Ok(list) => list, - Err(msg) => { - let resp = JsonRpcResponse::error(id, ACP_OVERLOADED, msg); - let _ = out_tx.send(serde_json::to_string(&resp).unwrap()); - continue; - } - }; - let (resp, channel_id) = - handle_session_new(&sessions, id.clone()).await; + let acp_mcp_servers = + match accept_acp_servers(parse_acp_mcp_servers(req.params.as_ref())) { + Ok(list) => list, + Err(msg) => { + let resp = JsonRpcResponse::error(id, ACP_OVERLOADED, msg); + let _ = out_tx.send(serde_json::to_string(&resp).unwrap()); + continue; + } + }; + let (resp, channel_id) = handle_session_new(&sessions, id.clone()).await; let _ = out_tx.send(serde_json::to_string(&resp).unwrap()); // If the client declared "type":"acp" MCP servers, open + register a tunnel to @@ -1605,9 +1632,9 @@ async fn handle_acp_connection(state: Arc, socket: WebSocket) { } // Required params per schema: { sessionId, cwd, mcpServers? }. The // sessionId's `sess_` shape is checked further in the handler. - if let Err(msg) = - validate_params::(req.params.as_ref()) - { + if let Err(msg) = validate_params::( + req.params.as_ref(), + ) { let resp = JsonRpcResponse::error(id, -32602, msg); let _ = out_tx.send(serde_json::to_string(&resp).unwrap()); continue; @@ -1918,7 +1945,11 @@ fn handle_initialize(req: &JsonRpcRequest) -> JsonRpcResponse { match serde_json::from_value(req.params.clone().unwrap_or(Value::Null)) { Ok(r) => r, Err(e) => { - return JsonRpcResponse::error(id, -32602, format!("Invalid initialize params: {e}")); + return JsonRpcResponse::error( + id, + -32602, + format!("Invalid initialize params: {e}"), + ); } }; // Negotiate: respond with the version we will use = the lower of the client's and @@ -2054,7 +2085,10 @@ async fn handle_session_resume( id: Value, params: Option<&Value>, ) -> (JsonRpcResponse, Option) { - let session_id = match params.and_then(|p| p.get("sessionId")).and_then(|v| v.as_str()) { + let session_id = match params + .and_then(|p| p.get("sessionId")) + .and_then(|v| v.as_str()) + { Some(s) => s.to_string(), None => { return ( @@ -2147,7 +2181,9 @@ async fn handle_session_cancel( "session/cancel is a notification and must not carry an id", )); } - let sess_key = params.and_then(|p| p.get("sessionId")).and_then(|v| v.as_str()); + let sess_key = params + .and_then(|p| p.get("sessionId")) + .and_then(|v| v.as_str()); if let Some(k) = sess_key { let notify = sessions.lock().await.get(k).and_then(|s| s.cancel.clone()); if let Some(n) = notify { @@ -2178,7 +2214,10 @@ pub(crate) fn redact_id(id: &str) -> String { // and `prompt dispatched` prints two of them on a single line. Correlating a session across // logs is the entire reason the tag exists, so producing several defeats the purpose more // completely than not redacting would. - let uuid = id.strip_prefix("acp_").or_else(|| id.strip_prefix("sess_")).unwrap_or(id); + let uuid = id + .strip_prefix("acp_") + .or_else(|| id.strip_prefix("sess_")) + .unwrap_or(id); use sha2::{Digest as _, Sha256}; let digest = Sha256::digest(uuid.as_bytes()); let short: String = digest.iter().take(4).map(|b| format!("{b:02x}")).collect(); @@ -2212,7 +2251,13 @@ mod redact_id_cross_encoding { fn openab_core_tag(uuid: &str) -> String { use sha2::{Digest as _, Sha256}; let d = Sha256::digest(uuid.as_bytes()); - format!("#{}", d.iter().take(4).map(|b| format!("{b:02x}")).collect::()) + format!( + "#{}", + d.iter() + .take(4) + .map(|b| format!("{b:02x}")) + .collect::() + ) } } @@ -2270,12 +2315,11 @@ async fn handle_session_prompt( let channel_id = match sessions.lock().await.get(&session_id) { Some(s) => s.channel_id.clone(), None => { - let resp = - JsonRpcResponse::error( - id, - -32602, - format!("Unknown session: {}", redact_id(&session_id)), - ); + let resp = JsonRpcResponse::error( + id, + -32602, + format!("Unknown session: {}", redact_id(&session_id)), + ); let _ = out_tx.send(serde_json::to_string(&resp).unwrap()); release_prompt(sessions, &session_id).await; return; @@ -2325,11 +2369,7 @@ async fn handle_session_prompt( if state.event_tx.send(json).is_err() { // No receivers — agent/core not connected warn!("ACP: event_tx send failed — no agent connected"); - let resp = JsonRpcResponse::error( - id, - -32603, - "No agent backend connected", - ); + let resp = JsonRpcResponse::error(id, -32603, "No agent backend connected"); let _ = out_tx.send(serde_json::to_string(&resp).unwrap()); release_prompt(sessions, &session_id).await; // Cleanup registry — only this turn's own sink (F4). @@ -2527,7 +2567,10 @@ pub async fn handle_reply(reply: &GatewayReply, registry: &AcpReplyRegistry) { sink.tx.clone() } Some(_) => { - debug!(channel = key, "ACP dropping stale reply from a superseded turn"); + debug!( + channel = key, + "ACP dropping stale reply from a superseded turn" + ); return; } None => return, @@ -2538,15 +2581,24 @@ pub async fn handle_reply(reply: &GatewayReply, registry: &AcpReplyRegistry) { Some("edit_message") => { // Streaming update — send as text snapshot if tx.send(ReplyChunk::Text(full_text)).is_err() { - debug!(channel = key, "ACP reply send failed (client likely disconnected)"); - registry.lock().unwrap_or_else(|e| e.into_inner()).remove(key); + debug!( + channel = key, + "ACP reply send failed (client likely disconnected)" + ); + registry + .lock() + .unwrap_or_else(|e| e.into_inner()) + .remove(key); } } None | Some("send_message") => { // Final message let _ = tx.send(ReplyChunk::Text(full_text)); let _ = tx.send(ReplyChunk::Done); - registry.lock().unwrap_or_else(|e| e.into_inner()).remove(key); + registry + .lock() + .unwrap_or_else(|e| e.into_inner()) + .remove(key); } Some("add_reaction") | Some("remove_reaction") => { // Reactions are agent state indicators — could map to notifications later @@ -2576,12 +2628,21 @@ mod acp_conformance { where T: serde::Serialize + serde::de::DeserializeOwned, { - let a: T = serde_json::from_value(wire.clone()) - .unwrap_or_else(|e| panic!("emitted wire is not valid ACP {}: {e}\n wire={wire}", std::any::type_name::())); + let a: T = serde_json::from_value(wire.clone()).unwrap_or_else(|e| { + panic!( + "emitted wire is not valid ACP {}: {e}\n wire={wire}", + std::any::type_name::() + ) + }); let v1 = serde_json::to_value(&a).unwrap(); let b: T = serde_json::from_value(v1.clone()).expect("re-parse of generated form"); let v2 = serde_json::to_value(&b).unwrap(); - assert_eq!(v1, v2, "ACP serde is not a stable fixed point for {}", std::any::type_name::()); + assert_eq!( + v1, + v2, + "ACP serde is not a stable fixed point for {}", + std::any::type_name::() + ); } // --- outbound responses (exact shapes handle_* emit) --- @@ -2608,7 +2669,9 @@ mod acp_conformance { #[test] fn new_session_response() { - conforms::(json!({ "sessionId": "sess_00000000-0000-0000-0000-000000000000" })); + conforms::( + json!({ "sessionId": "sess_00000000-0000-0000-0000-000000000000" }), + ); } #[test] @@ -2652,12 +2715,12 @@ mod acp_conformance { // astral-plane emoji, ZWJ sequence, regional-indicator flag, VS16 emoji, // astral-plane CJK, and a mixed run. const EDGE_TEXT: &[&str] = &[ - "🎉", // U+1F389, 4-byte astral emoji - "👨‍👩‍👧‍👦", // ZWJ family (7 codepoints joined by ZWJ) - "🇹🇼", // regional-indicator pair (flag) - "❤️", // U+2764 + U+FE0F (VS16) - "𠀀", // U+20000, astral-plane CJK - "🎉 你好 (๑•̀ㅂ•́)و ❤️", // mixed emoji + CJK + kaomoji + VS16 + "🎉", // U+1F389, 4-byte astral emoji + "👨‍👩‍👧‍👦", // ZWJ family (7 codepoints joined by ZWJ) + "🇹🇼", // regional-indicator pair (flag) + "❤️", // U+2764 + U+FE0F (VS16) + "𠀀", // U+20000, astral-plane CJK + "🎉 你好 (๑•̀ㅂ•́)و ❤️", // mixed emoji + CJK + kaomoji + VS16 ]; #[test] @@ -2698,7 +2761,13 @@ mod acp_conformance { #[test] fn prompt_response_all_stop_reasons() { - for sr in ["end_turn", "max_tokens", "max_turn_requests", "refusal", "cancelled"] { + for sr in [ + "end_turn", + "max_tokens", + "max_turn_requests", + "refusal", + "cancelled", + ] { conforms::(json!({ "stopReason": sr })); } } @@ -2725,8 +2794,12 @@ mod acp_conformance { serde_json::from_str(r#"{"jsonrpc":"2.0","method":"session/cancel"}"#).unwrap(); assert!(notif.get("id").is_none(), "no id member → notification"); let req_null: Value = - serde_json::from_str(r#"{"jsonrpc":"2.0","method":"session/cancel","id":null}"#).unwrap(); - assert!(req_null.get("id").is_some(), "explicit id:null → request (id member present)"); + serde_json::from_str(r#"{"jsonrpc":"2.0","method":"session/cancel","id":null}"#) + .unwrap(); + assert!( + req_null.get("id").is_some(), + "explicit id:null → request (id member present)" + ); let req_num: Value = serde_json::from_str(r#"{"jsonrpc":"2.0","method":"initialize","id":7}"#).unwrap(); assert_eq!(req_num.get("id"), Some(&json!(7))); @@ -2738,14 +2811,36 @@ mod acp_conformance { fn session_param_validation() { use super::validate_params; // session/new requires { cwd, mcpServers } - assert!(validate_params::(Some(&json!({"cwd": "/w", "mcpServers": []}))).is_ok()); - assert!(validate_params::(Some(&json!({"mcpServers": []}))).is_err(), "missing cwd"); - assert!(validate_params::(Some(&json!({"cwd": "/w"}))).is_err(), "missing mcpServers"); - assert!(validate_params::(None).is_err(), "missing params"); + assert!(validate_params::(Some( + &json!({"cwd": "/w", "mcpServers": []}) + )) + .is_ok()); + assert!( + validate_params::(Some(&json!({"mcpServers": []}))).is_err(), + "missing cwd" + ); + assert!( + validate_params::(Some(&json!({"cwd": "/w"}))).is_err(), + "missing mcpServers" + ); + assert!( + validate_params::(None).is_err(), + "missing params" + ); // session/resume requires { sessionId, cwd } - assert!(validate_params::(Some(&json!({"sessionId": "sess_x", "cwd": "/w", "mcpServers": []}))).is_ok()); - assert!(validate_params::(Some(&json!({"cwd": "/w"}))).is_err(), "missing sessionId"); - assert!(validate_params::(Some(&json!({"sessionId": "sess_x"}))).is_err(), "missing cwd"); + assert!(validate_params::(Some( + &json!({"sessionId": "sess_x", "cwd": "/w", "mcpServers": []}) + )) + .is_ok()); + assert!( + validate_params::(Some(&json!({"cwd": "/w"}))).is_err(), + "missing sessionId" + ); + assert!( + validate_params::(Some(&json!({"sessionId": "sess_x"}))) + .is_err(), + "missing cwd" + ); } // --- prompt content blocks (F10): unsupported block types rejected, not dropped --- @@ -2796,13 +2891,16 @@ mod acp_conformance { // R17-F3a — a plain-string prompt is non-conformant (schema requires // `prompt: [ContentBlock]`) → rejected, surfaced as -32602 at the call site. assert!( - extract_prompt_params(Some(&json!({"sessionId": "sess_x", "prompt": "hello"}))).is_err(), + extract_prompt_params(Some(&json!({"sessionId": "sess_x", "prompt": "hello"}))) + .is_err(), "a bare string prompt must be rejected, not coerced" ); // an object (non-array, non-string) prompt is likewise rejected. assert!( - extract_prompt_params(Some(&json!({"sessionId": "sess_x", "prompt": {"type": "text"}}))) - .is_err(), + extract_prompt_params(Some( + &json!({"sessionId": "sess_x", "prompt": {"type": "text"}}) + )) + .is_err(), "a non-array prompt must be rejected" ); } @@ -2833,8 +2931,11 @@ mod acp_conformance { use axum::http::HeaderMap; let mut h = HeaderMap::new(); assert_eq!(subprotocol_token(&h), None); // no header - // the browser offers "openab.bearer., acp.v1" → extract the token - h.insert("sec-websocket-protocol", "openab.bearer.abc123, acp.v1".parse().unwrap()); + // the browser offers "openab.bearer., acp.v1" → extract the token + h.insert( + "sec-websocket-protocol", + "openab.bearer.abc123, acp.v1".parse().unwrap(), + ); assert_eq!(subprotocol_token(&h), Some("abc123")); // only the real protocol, no bearer entry → None h.insert("sec-websocket-protocol", "acp.v1".parse().unwrap()); @@ -2873,14 +2974,7 @@ mod acp_streaming { #[test] fn multibyte_codepoints_never_split() { // each snapshot appends a whole multi-byte grapheme; reconstruction is exact - let snaps = [ - "a", - "a🎉", - "a🎉你", - "a🎉你👨‍👩‍👧‍👦", - "a🎉你👨‍👩‍👧‍👦🇹🇼", - "a🎉你👨‍👩‍👧‍👦🇹🇼❤️", - ]; + let snaps = ["a", "a🎉", "a🎉你", "a🎉你👨‍👩‍👧‍👦", "a🎉你👨‍👩‍👧‍👦🇹🇼", "a🎉你👨‍👩‍👧‍👦🇹🇼❤️"]; assert_eq!(replay(&snaps), *snaps.last().unwrap()); } @@ -2936,7 +3030,11 @@ mod acp_requests { pending: &Arc>>>, ) { let f: serde_json::Value = serde_json::from_str(&out_rx.recv().await.unwrap()).unwrap(); - assert_eq!(f["params"]["method"], json!("initialize"), "expected the inner MCP initialize"); + assert_eq!( + f["params"]["method"], + json!("initialize"), + "expected the inner MCP initialize" + ); route_client_response( pending, &json!({"jsonrpc":"2.0","id":f["id"],"result":{ @@ -2951,8 +3049,7 @@ mod acp_requests { assert_eq!(n["params"]["method"], json!("notifications/initialized")); } - fn new_pending( - ) -> Arc>>> { + fn new_pending() -> Arc>>> { Arc::new(tokio::sync::Mutex::new(HashMap::new())) } @@ -2967,8 +3064,7 @@ mod acp_requests { let p2 = pending.clone(); let version = version.map(str::to_string); let responder = tokio::spawn(async move { - let f: serde_json::Value = - serde_json::from_str(&out_rx.recv().await.unwrap()).unwrap(); + let f: serde_json::Value = serde_json::from_str(&out_rx.recv().await.unwrap()).unwrap(); let result = match version { Some(v) => json!({ "protocolVersion": v, "capabilities": {"tools": {}}, @@ -2984,7 +3080,8 @@ mod acp_requests { route_client_response(&p2, &json!({"jsonrpc":"2.0","id":f["id"],"result":result})) .await; // Drain `notifications/initialized` if the handshake got far enough to send it. - let _ = tokio::time::timeout(std::time::Duration::from_millis(200), out_rx.recv()).await; + let _ = + tokio::time::timeout(std::time::Duration::from_millis(200), out_rx.recv()).await; }); let r = super::inner_mcp_handshake(&out_tx, &pending, &next_id, "conn-1", 5).await; let _ = responder.await; @@ -3013,7 +3110,10 @@ mod acp_requests { let err = handshake_answering(Some("1999-01-01")) .await .expect_err("an unknown revision must not be negotiated into"); - assert!(err.contains("1999-01-01"), "the error must name what the peer answered: {err}"); + assert!( + err.contains("1999-01-01"), + "the error must name what the peer answered: {err}" + ); assert!( err.contains(super::INNER_MCP_PROTOCOL_VERSION), "the error must name what we requested: {err}" @@ -3031,7 +3131,10 @@ mod acp_requests { let err = handshake_answering(None) .await .expect_err("no protocolVersion is not a compliant initialize result"); - assert!(err.contains("no `protocolVersion`"), "unexpected error: {err}"); + assert!( + err.contains("no `protocolVersion`"), + "unexpected error: {err}" + ); } #[tokio::test] @@ -3039,9 +3142,11 @@ mod acp_requests { let pending = new_pending(); let (tx, rx) = oneshot::channel(); pending.lock().await.insert(5, tx); - let consumed = - route_client_response(&pending, &json!({"jsonrpc":"2.0","id":5,"result":{"ok":true}})) - .await; + let consumed = route_client_response( + &pending, + &json!({"jsonrpc":"2.0","id":5,"result":{"ok":true}}), + ) + .await; assert!(consumed, "an id+result frame is a response we consume"); assert_eq!(rx.await.unwrap()["result"]["ok"], json!(true)); assert!( @@ -3054,7 +3159,9 @@ mod acp_requests { async fn route_client_response_ignores_requests_and_notifications() { let pending = new_pending(); // has `method` → a request, not a response - assert!(!route_client_response(&pending, &json!({"jsonrpc":"2.0","id":1,"method":"foo"})).await); + assert!( + !route_client_response(&pending, &json!({"jsonrpc":"2.0","id":1,"method":"foo"})).await + ); // notification-shaped, no result/error → not a response assert!(!route_client_response(&pending, &json!({"jsonrpc":"2.0","method":"bar"})).await); // id present but neither result nor error → not a response @@ -3069,7 +3176,10 @@ mod acp_requests { &json!({"jsonrpc":"2.0","id":99,"error":{"code":-1,"message":"x"}}), ) .await; - assert!(consumed, "an unmatched response is still consumed (logged, no panic)"); + assert!( + consumed, + "an unmatched response is still consumed (logged, no panic)" + ); } #[tokio::test] @@ -3088,8 +3198,11 @@ mod acp_requests { assert_eq!(v["method"], json!("mcp/message")); assert_eq!(v["params"]["connectionId"], json!("conn-1")); let id = v["id"].as_u64().unwrap(); - route_client_response(&pending2, &json!({"jsonrpc":"2.0","id":id,"result":{"pong":true}})) - .await; + route_client_response( + &pending2, + &json!({"jsonrpc":"2.0","id":id,"result":{"pong":true}}), + ) + .await; }); let resp = send_request( @@ -3104,7 +3217,11 @@ mod acp_requests { .unwrap(); assert_eq!(resp["result"]["pong"], json!(true)); driver.await.unwrap(); - assert_eq!(next_id.load(Ordering::Relaxed), 2, "the id counter advanced"); + assert_eq!( + next_id.load(Ordering::Relaxed), + 2, + "the id counter advanced" + ); } #[tokio::test] @@ -3117,7 +3234,8 @@ mod acp_requests { // tools/list into an inner result, routing each reply by the frame's own outer id. let pending2 = pending.clone(); let ext = tokio::spawn(async move { - let f1: serde_json::Value = serde_json::from_str(&out_rx.recv().await.unwrap()).unwrap(); + let f1: serde_json::Value = + serde_json::from_str(&out_rx.recv().await.unwrap()).unwrap(); assert_eq!(f1["method"], json!("mcp/connect")); assert_eq!(f1["params"]["acpId"], json!("srv-1")); route_client_response( @@ -3126,7 +3244,8 @@ mod acp_requests { ) .await; - let f2: serde_json::Value = serde_json::from_str(&out_rx.recv().await.unwrap()).unwrap(); + let f2: serde_json::Value = + serde_json::from_str(&out_rx.recv().await.unwrap()).unwrap(); assert_eq!(f2["method"], json!("mcp/message")); assert_eq!(f2["params"]["connectionId"], json!("conn-9")); assert_eq!(f2["params"]["method"], json!("tools/list")); @@ -3170,8 +3289,11 @@ mod acp_requests { assert_eq!(f["method"], json!("mcp/message")); assert_eq!(f["params"]["connectionId"], json!("conn-9")); assert_eq!(f["params"]["method"], json!("tools/call")); - route_client_response(&pending2, &json!({"jsonrpc":"2.0","id":f["id"],"result":{"ok":true}})) - .await; + route_client_response( + &pending2, + &json!({"jsonrpc":"2.0","id":f["id"],"result":{"ok":true}}), + ) + .await; }); let result = handle @@ -3224,14 +3346,23 @@ mod acp_requests { let mut reg = registry.lock().unwrap(); // Deliberately out of insertion order relative to rank, so a "first match" answer and a // "newest" answer differ. - reg.insert(("acp_1".into(), "srv-new".into()), tunnel_ranked("conn-b", 7, 2)); - reg.insert(("acp_1".into(), "srv-old".into()), tunnel_ranked("conn-a", 3, 9)); + reg.insert( + ("acp_1".into(), "srv-new".into()), + tunnel_ranked("conn-b", 7, 2), + ); + reg.insert( + ("acp_1".into(), "srv-old".into()), + tunnel_ranked("conn-a", 3, 9), + ); reg.insert(("acp_1".into(), "other".into()), { let mut h = tunnel_ranked("conn-c", 9, 9); h.server_name = "notes".into(); h }); - reg.insert(("acp_2".into(), "elsewhere".into()), tunnel_ranked("conn-d", 99, 99)); + reg.insert( + ("acp_2".into(), "elsewhere".into()), + tunnel_ranked("conn-d", 99, 99), + ); } assert_eq!( super::resolve_by_name(®istry, "acp_1", "katashiro").as_deref(), @@ -3370,8 +3501,11 @@ mod acp_requests { let f: serde_json::Value = serde_json::from_str(&out_rx.recv().await.unwrap()).unwrap(); assert_eq!(f["method"], json!("mcp/connect")); assert_eq!(f["params"]["acpId"], json!("srv-1")); - route_client_response(&pending2, &json!({"jsonrpc":"2.0","id":f["id"],"result":{"connectionId":"conn-9"}})) - .await; + route_client_response( + &pending2, + &json!({"jsonrpc":"2.0","id":f["id"],"result":{"connectionId":"conn-9"}}), + ) + .await; answer_inner_handshake(&mut out_rx, &pending2).await; }); @@ -3563,7 +3697,10 @@ mod acp_handlers { #[test] fn initialize_returns_conformant_capabilities() { - let v = serde_json::to_value(handle_initialize(&init_req(Some(json!({"protocolVersion": 1}))))).unwrap(); + let v = serde_json::to_value(handle_initialize(&init_req(Some( + json!({"protocolVersion": 1}), + )))) + .unwrap(); assert_eq!(v["id"], json!(1)); let result = &v["result"]; assert_eq!(result["protocolVersion"], json!(1)); @@ -3575,26 +3712,44 @@ mod acp_handlers { // the serde default {http:false, sse:false} — the same VALUES, but arrived at by silence // while the gateway ships MCP-over-ACP and says so nowhere. let mcp = &result["agentCapabilities"]["mcpCapabilities"]; - assert_eq!(mcp["http"], json!(false), "no code forwards an http declaration anywhere"); - assert_eq!(mcp["sse"], json!(false), "same for sse — parse_acp_mcp_servers drops both"); assert_eq!( - mcp["_meta"]["dev.openab/acp"], json!(true), + mcp["http"], + json!(false), + "no code forwards an http declaration anywhere" + ); + assert_eq!( + mcp["sse"], + json!(false), + "same for sse — parse_acp_mcp_servers drops both" + ); + assert_eq!( + mcp["_meta"]["dev.openab/acp"], + json!(true), "the ACP capability is a reverse-DNS-namespaced extension under _meta (F1(b))" ); assert!( mcp["_meta"]["acp"].is_null(), "the bare `_meta.acp` key is gone — the informal convention the framework supersedes" ); - assert!(mcp.get("acp").is_none(), "no core mcpCapabilities.acp field (would fork the schema)"); + assert!( + mcp.get("acp").is_none(), + "no core mcpCapabilities.acp field (would fork the schema)" + ); } #[test] fn initialize_negotiates_version_and_rejects_bad() { // a higher client version negotiates down to ours (1) - let v = serde_json::to_value(handle_initialize(&init_req(Some(json!({"protocolVersion": 5}))))).unwrap(); + let v = serde_json::to_value(handle_initialize(&init_req(Some( + json!({"protocolVersion": 5}), + )))) + .unwrap(); assert_eq!(v["result"]["protocolVersion"], json!(1)); // version 0 is below our minimum → -32602 - let v = serde_json::to_value(handle_initialize(&init_req(Some(json!({"protocolVersion": 0}))))).unwrap(); + let v = serde_json::to_value(handle_initialize(&init_req(Some( + json!({"protocolVersion": 0}), + )))) + .unwrap(); assert_eq!(v["error"]["code"], json!(-32602)); // missing protocolVersion → -32602 let v = serde_json::to_value(handle_initialize(&init_req(Some(json!({}))))).unwrap(); @@ -3607,11 +3762,16 @@ mod acp_handlers { #[tokio::test] async fn session_new_mints_and_stores_a_session() { let sessions = new_sessions(); - let v = - serde_json::to_value(handle_session_new(&sessions, json!(2)).await.0).unwrap(); + let v = serde_json::to_value(handle_session_new(&sessions, json!(2)).await.0).unwrap(); let sid = v["result"]["sessionId"].as_str().unwrap(); - assert!(sid.starts_with("sess_"), "sessionId must be sess_: {sid}"); - assert!(sessions.lock().await.contains_key(sid), "session must be stored"); + assert!( + sid.starts_with("sess_"), + "sessionId must be sess_: {sid}" + ); + assert!( + sessions.lock().await.contains_key(sid), + "session must be stored" + ); } #[test] @@ -3627,8 +3787,14 @@ mod acp_handlers { assert_eq!( parse_acp_mcp_servers(Some(¶ms)), vec![ - AcpMcpServer { id: "srv-1".into(), name: "browser".into() }, - AcpMcpServer { id: "srv-2".into(), name: "other".into() }, + AcpMcpServer { + id: "srv-1".into(), + name: "browser".into() + }, + AcpMcpServer { + id: "srv-2".into(), + name: "other".into() + }, ] ); // no mcpServers -> empty @@ -3636,25 +3802,34 @@ mod acp_handlers { assert!(parse_acp_mcp_servers(None).is_empty()); } - #[tokio::test] async fn session_resume_valid_stores_and_invalid_errors() { let sessions = new_sessions(); // valid sess_ → {} and the session is (re)stored let sid = format!("sess_{}", Uuid::new_v4()); let params = json!({"sessionId": sid, "cwd": "/w", "mcpServers": []}); - let v = serde_json::to_value(handle_session_resume(&sessions, json!(3), Some(¶ms)).await.0) - .unwrap(); + let v = serde_json::to_value( + handle_session_resume(&sessions, json!(3), Some(¶ms)) + .await + .0, + ) + .unwrap(); assert_eq!(v["result"], json!({})); assert!(sessions.lock().await.contains_key(&sid)); // malformed sessionId shape → -32602 let bad = json!({"sessionId": "not-a-session", "cwd": "/w", "mcpServers": []}); - let v = serde_json::to_value(handle_session_resume(&sessions, json!(4), Some(&bad)).await.0) - .unwrap(); + let v = serde_json::to_value( + handle_session_resume(&sessions, json!(4), Some(&bad)) + .await + .0, + ) + .unwrap(); assert_eq!(v["error"]["code"], json!(-32602)); // missing sessionId → -32602 let v = serde_json::to_value( - handle_session_resume(&sessions, json!(5), Some(&json!({"cwd": "/w"}))).await.0, + handle_session_resume(&sessions, json!(5), Some(&json!({"cwd": "/w"}))) + .await + .0, ) .unwrap(); assert_eq!(v["error"]["code"], json!(-32602)); @@ -3685,23 +3860,39 @@ mod acp_review_fixes { for _ in 0..MAX_SESSIONS_PER_CONNECTION { let sid = format!("sess_{}", Uuid::new_v4()); let p = json!({ "sessionId": sid }); - let v = serde_json::to_value(handle_session_resume(&sessions, json!(1), Some(&p)).await.0) - .unwrap(); + let v = + serde_json::to_value(handle_session_resume(&sessions, json!(1), Some(&p)).await.0) + .unwrap(); assert_eq!(v["result"], json!({}), "resume under cap should succeed"); ids.push(sid); } assert_eq!(sessions.lock().await.len(), MAX_SESSIONS_PER_CONNECTION); // A new distinct session over the cap is refused with ACP_OVERLOADED. let over = json!({ "sessionId": format!("sess_{}", Uuid::new_v4()) }); - let v = serde_json::to_value(handle_session_resume(&sessions, json!(2), Some(&over)).await.0) - .unwrap(); - assert_eq!(v["error"]["code"], json!(ACP_OVERLOADED), "over-cap resume must be refused"); + let v = serde_json::to_value( + handle_session_resume(&sessions, json!(2), Some(&over)) + .await + .0, + ) + .unwrap(); + assert_eq!( + v["error"]["code"], + json!(ACP_OVERLOADED), + "over-cap resume must be refused" + ); // Re-resuming an already-present session is exempt (idempotent). let existing = json!({ "sessionId": ids[0] }); - let v = - serde_json::to_value(handle_session_resume(&sessions, json!(3), Some(&existing)).await.0) - .unwrap(); - assert_eq!(v["result"], json!({}), "re-resume of existing session must bypass the cap"); + let v = serde_json::to_value( + handle_session_resume(&sessions, json!(3), Some(&existing)) + .await + .0, + ) + .unwrap(); + assert_eq!( + v["result"], + json!({}), + "re-resume of existing session must bypass the cap" + ); } // --- R3-F1: declaration fan-out is bounded before it costs anything --- @@ -3855,14 +4046,23 @@ mod acp_review_fixes { // 1. missing sessionId let (resp, chan) = handle_session_resume(&sessions, json!(1), Some(&json!({"cwd": "/w"}))).await; - assert_eq!(serde_json::to_value(resp).unwrap()["error"]["code"], json!(-32602)); + assert_eq!( + serde_json::to_value(resp).unwrap()["error"]["code"], + json!(-32602) + ); assert!(chan.is_none(), "missing sessionId must not yield a channel"); // 2. malformed sessionId let bad = json!({"sessionId": "not-a-session", "cwd": "/w"}); let (resp, chan) = handle_session_resume(&sessions, json!(2), Some(&bad)).await; - assert_eq!(serde_json::to_value(resp).unwrap()["error"]["code"], json!(-32602)); - assert!(chan.is_none(), "malformed sessionId must not yield a channel"); + assert_eq!( + serde_json::to_value(resp).unwrap()["error"]["code"], + json!(-32602) + ); + assert!( + chan.is_none(), + "malformed sessionId must not yield a channel" + ); // 3. over the per-connection cap — note the id IS well formed, so the old // derive-from-params guard would have happily produced a channel here. @@ -3877,7 +4077,10 @@ mod acp_review_fixes { serde_json::to_value(resp).unwrap()["error"]["code"], json!(ACP_OVERLOADED) ); - assert!(chan.is_none(), "an over-cap resume must not yield a channel"); + assert!( + chan.is_none(), + "an over-cap resume must not yield a channel" + ); // 4. busy — likewise a well-formed id on a session that really exists. let busy_sid = format!("sess_{}", Uuid::new_v4()); @@ -3891,8 +4094,14 @@ mod acp_review_fixes { ); let (resp, chan) = handle_session_resume(&sessions, json!(5), Some(&json!({"sessionId": busy_sid}))).await; - assert_eq!(serde_json::to_value(resp).unwrap()["error"]["code"], json!(-32001)); - assert!(chan.is_none(), "a busy-rejected resume must not yield a channel"); + assert_eq!( + serde_json::to_value(resp).unwrap()["error"]["code"], + json!(-32001) + ); + assert!( + chan.is_none(), + "a busy-rejected resume must not yield a channel" + ); } fn reply(channel_id: &str, reply_to: &str, text: &str, command: Option<&str>) -> GatewayReply { @@ -3900,7 +4109,10 @@ mod acp_review_fixes { schema: "openab.gateway.reply.v1".into(), reply_to: reply_to.into(), platform: "acp".into(), - channel: crate::schema::ReplyChannel { id: channel_id.into(), thread_id: None }, + channel: crate::schema::ReplyChannel { + id: channel_id.into(), + thread_id: None, + }, content: crate::schema::Content { content_type: "text".into(), text: text.into(), @@ -3918,22 +4130,33 @@ mod acp_review_fixes { async fn handle_reply_fences_stale_turn() { let registry = new_reply_registry(); let (tx, mut rx) = mpsc::unbounded_channel::(); - registry - .lock() - .unwrap() - .insert( - "acp_chan".into(), - ReplySink { turn_id: "evt_current".into(), tx, owner: "conn-test".into(), generation: 0 }, - ); + registry.lock().unwrap().insert( + "acp_chan".into(), + ReplySink { + turn_id: "evt_current".into(), + tx, + owner: "conn-test".into(), + generation: 0, + }, + ); // Stale reply (previous turn's event id) → dropped. - handle_reply(&reply("acp_chan", "evt_stale", "leaked", Some("edit_message")), ®istry) - .await; - assert!(rx.try_recv().is_err(), "stale reply must not reach the active turn"); + handle_reply( + &reply("acp_chan", "evt_stale", "leaked", Some("edit_message")), + ®istry, + ) + .await; + assert!( + rx.try_recv().is_err(), + "stale reply must not reach the active turn" + ); // Matching reply → delivered. - handle_reply(&reply("acp_chan", "evt_current", "hello", Some("edit_message")), ®istry) - .await; + handle_reply( + &reply("acp_chan", "evt_current", "hello", Some("edit_message")), + ®istry, + ) + .await; match rx.try_recv() { Ok(ReplyChunk::Text(t)) => assert_eq!(t, "hello"), _ => panic!("expected the matching reply to be delivered"), @@ -3948,30 +4171,64 @@ mod acp_review_fixes { let registry = new_reply_registry(); let sink = |turn: &str, owner: &str, generation: u64| { let (tx, _rx) = mpsc::unbounded_channel::(); - super::ReplySink { turn_id: turn.into(), tx, owner: owner.into(), generation } + super::ReplySink { + turn_id: turn.into(), + tx, + owner: owner.into(), + generation, + } + }; + let current_turn = || { + registry + .lock() + .unwrap() + .get("acp_x") + .map(|s| s.turn_id.clone()) }; - let current_turn = || registry.lock().unwrap().get("acp_x").map(|s| s.turn_id.clone()); // Connection A (gen 1) installs its sink. - assert!(super::install_reply_sink(®istry, "acp_x", sink("evt_a", "conn-A", 1))); + assert!(super::install_reply_sink( + ®istry, + "acp_x", + sink("evt_a", "conn-A", 1) + )); // Newer connection B (gen 2) resumes the same session and takes over (reconnect semantics). - assert!(super::install_reply_sink(®istry, "acp_x", sink("evt_b", "conn-B", 2))); - assert_eq!(current_turn().as_deref(), Some("evt_b"), "the newer connection owns the sink"); + assert!(super::install_reply_sink( + ®istry, + "acp_x", + sink("evt_b", "conn-B", 2) + )); + assert_eq!( + current_turn().as_deref(), + Some("evt_b"), + "the newer connection owns the sink" + ); // The older connection A, its prompt processed late, must NOT clobber B's live sink. assert!( !super::install_reply_sink(®istry, "acp_x", sink("evt_a2", "conn-A", 1)), "an older connection cannot install over a newer one" ); - assert_eq!(current_turn().as_deref(), Some("evt_b"), "B's sink survives A's late install"); + assert_eq!( + current_turn().as_deref(), + Some("evt_b"), + "B's sink survives A's late install" + ); // A's turn completing must not remove B's sink (turn-scoped removal). super::remove_reply_sink_if_owner(®istry, "acp_x", "evt_a"); - assert_eq!(current_turn().as_deref(), Some("evt_b"), "A's completion must not remove B's sink"); + assert_eq!( + current_turn().as_deref(), + Some("evt_b"), + "A's completion must not remove B's sink" + ); // B's own completion removes B's sink. super::remove_reply_sink_if_owner(®istry, "acp_x", "evt_b"); - assert!(current_turn().is_none(), "the owner's completion removes its own sink"); + assert!( + current_turn().is_none(), + "the owner's completion removes its own sink" + ); } // F4 — the SAME connection starting its next turn replaces its own sink; the stale prior turn's @@ -3981,14 +4238,24 @@ mod acp_review_fixes { let registry = new_reply_registry(); let sink = |turn: &str| { let (tx, _rx) = mpsc::unbounded_channel::(); - super::ReplySink { turn_id: turn.into(), tx, owner: "conn-A".into(), generation: 5 } + super::ReplySink { + turn_id: turn.into(), + tx, + owner: "conn-A".into(), + generation: 5, + } }; assert!(super::install_reply_sink(®istry, "acp_x", sink("evt_1"))); assert!(super::install_reply_sink(®istry, "acp_x", sink("evt_2"))); // Stale turn 1's completion must not remove turn 2's sink. super::remove_reply_sink_if_owner(®istry, "acp_x", "evt_1"); assert_eq!( - registry.lock().unwrap().get("acp_x").map(|s| s.turn_id.clone()).as_deref(), + registry + .lock() + .unwrap() + .get("acp_x") + .map(|s| s.turn_id.clone()) + .as_deref(), Some("evt_2"), "the stale turn must not remove the same connection's newer sink" ); @@ -4001,10 +4268,19 @@ mod acp_review_fixes { // the `else` of the bearer branch), so a keyed bind is unaffected by the allowlist. #[test] fn acp_origin_ok_keyless_gating() { - let allow = vec!["https://app.example".to_string(), "http://localhost:5173".to_string()]; + let allow = vec![ + "https://app.example".to_string(), + "http://localhost:5173".to_string(), + ]; // Absent Origin (non-browser client) → accept, regardless of allowlist. - assert!(acp_origin_ok(None, &allow), "no Origin (non-browser) must be admitted"); - assert!(acp_origin_ok(None, &[]), "no Origin must be admitted even with empty allowlist"); + assert!( + acp_origin_ok(None, &allow), + "no Origin (non-browser) must be admitted" + ); + assert!( + acp_origin_ok(None, &[]), + "no Origin must be admitted even with empty allowlist" + ); // Allowlisted browser Origin → accept (exact match, both entries). assert!(acp_origin_ok(Some("https://app.example"), &allow)); assert!(acp_origin_ok(Some("http://localhost:5173"), &allow)); @@ -4034,7 +4310,10 @@ mod acp_review_fixes { assert_eq!(ws_bearer_token(&h), Some("sekret")); // The subprotocol path still carries the key. let mut h = HeaderMap::new(); - h.insert("sec-websocket-protocol", "openab.bearer.sekret, acp.v1".parse().unwrap()); + h.insert( + "sec-websocket-protocol", + "openab.bearer.sekret, acp.v1".parse().unwrap(), + ); assert_eq!(ws_bearer_token(&h), Some("sekret")); } @@ -4043,10 +4322,18 @@ mod acp_review_fixes { #[test] fn ws_subprotocol_token_charset() { for &b in b"AZaz09._~+-!#$%&'*^`|" { - assert!(is_ws_subprotocol_token_char(b), "{} should be token-safe", b as char); + assert!( + is_ws_subprotocol_token_char(b), + "{} should be token-safe", + b as char + ); } for &b in b"=/,; @\"" { - assert!(!is_ws_subprotocol_token_char(b), "{} should be rejected", b as char); + assert!( + !is_ws_subprotocol_token_char(b), + "{} should be rejected", + b as char + ); } } @@ -4079,8 +4366,18 @@ mod acp_review_fixes { let (out_tx, mut out_rx) = mpsc::unbounded_channel::(); let params = json!({"sessionId": sid, "prompt": [{"type": "text", "text": "hi"}]}); - handle_session_prompt(&state, &sessions, json!(7), Some(¶ms), &out_tx, sid.clone(), cancel, "conn-test", 0) - .await; + handle_session_prompt( + &state, + &sessions, + json!(7), + Some(¶ms), + &out_tx, + sid.clone(), + cancel, + "conn-test", + 0, + ) + .await; // The final response (matching our request id) must carry stopReason "cancelled". let mut final_resp = None; @@ -4099,7 +4396,10 @@ mod acp_review_fixes { // And the reservation is released. let g = sessions.lock().await; let s = g.get(&sid).unwrap(); - assert!(!s.busy && s.cancel.is_none(), "cancel must release busy + cancel handle"); + assert!( + !s.busy && s.cancel.is_none(), + "cancel must release busy + cancel handle" + ); } // R16-F2 — session/resume on a session with a prompt in flight is rejected (busy), so the @@ -4121,7 +4421,11 @@ mod acp_review_fixes { let params = json!({"sessionId": sid, "cwd": "/w", "mcpServers": []}); let (resp, resumed) = handle_session_resume(&sessions, json!(9), Some(¶ms)).await; let v = serde_json::to_value(resp).unwrap(); - assert_eq!(v["error"]["code"], json!(-32001), "resume while busy must be rejected"); + assert_eq!( + v["error"]["code"], + json!(-32001), + "resume while busy must be rejected" + ); assert!( resumed.is_none(), "a rejected resume must not hand back a channel — that value is the caller's \ @@ -4132,7 +4436,10 @@ mod acp_review_fixes { let g = sessions.lock().await; let s = g.get(&sid).unwrap(); assert!(s.busy, "busy must remain set after a rejected resume"); - assert!(s.cancel.is_some(), "the active prompt's cancel handle must survive resume"); + assert!( + s.cancel.is_some(), + "the active prompt's cancel handle must survive resume" + ); } // R16-F3(A) — Phase-1 send-once: the ACP path streams the whole reply as a SINGLE terminal @@ -4152,7 +4459,11 @@ mod acp_review_fixes { let cancel = Arc::new(tokio::sync::Notify::new()); sessions.lock().await.insert( sid.clone(), - AcpSession { channel_id: channel_id.clone(), busy: true, cancel: Some(cancel.clone()) }, + AcpSession { + channel_id: channel_id.clone(), + busy: true, + cancel: Some(cancel.clone()), + }, ); let (out_tx, mut out_rx) = mpsc::unbounded_channel::(); @@ -4161,21 +4472,40 @@ mod acp_review_fixes { let sid2 = sid.clone(); let handle = tokio::spawn(async move { let params = json!({"sessionId": sid2, "prompt": [{"type": "text", "text": "hi"}]}); - handle_session_prompt(&st2, &sessions2, json!(11), Some(¶ms), &out_tx, sid2.clone(), cancel, "conn-test", 0) - .await; + handle_session_prompt( + &st2, + &sessions2, + json!(11), + Some(¶ms), + &out_tx, + sid2.clone(), + cancel, + "conn-test", + 0, + ) + .await; }); // Wait for the handler to register its reply sink, then feed one final reply. let mut turn_id = None; for _ in 0..10_000 { - if let Some(t) = registry.lock().unwrap().get(&channel_id).map(|s| s.turn_id.clone()) { + if let Some(t) = registry + .lock() + .unwrap() + .get(&channel_id) + .map(|s| s.turn_id.clone()) + { turn_id = Some(t); break; } tokio::task::yield_now().await; } let turn_id = turn_id.expect("handler must register a reply sink"); - handle_reply(&reply(&channel_id, &turn_id, "hello world", Some("send_message")), ®istry).await; + handle_reply( + &reply(&channel_id, &turn_id, "hello world", Some("send_message")), + ®istry, + ) + .await; handle.await.unwrap(); let mut chunks = Vec::new(); @@ -4185,15 +4515,28 @@ mod acp_review_fixes { if v["method"] == json!("session/update") && v["params"]["update"]["sessionUpdate"] == json!("agent_message_chunk") { - chunks.push(v["params"]["update"]["content"]["text"].as_str().unwrap_or("").to_string()); + chunks.push( + v["params"]["update"]["content"]["text"] + .as_str() + .unwrap_or("") + .to_string(), + ); } if v.get("id") == Some(&json!(11)) { final_stop = v["result"]["stopReason"].as_str().map(str::to_string); } } - assert_eq!(chunks.len(), 1, "Phase-1 must stream exactly one terminal chunk, got {chunks:?}"); + assert_eq!( + chunks.len(), + 1, + "Phase-1 must stream exactly one terminal chunk, got {chunks:?}" + ); assert_eq!(chunks[0], "hello world"); - assert_eq!(final_stop.as_deref(), Some("end_turn"), "a completed turn ends end_turn"); + assert_eq!( + final_stop.as_deref(), + Some("end_turn"), + "a completed turn ends end_turn" + ); } // R17-F3c — a request-shaped `session/cancel` (id present) must NOT be acknowledged with @@ -4207,7 +4550,11 @@ mod acp_review_fixes { .await .expect("a request-shaped cancel must produce a response, not silence"); let v = serde_json::to_value(&resp).unwrap(); - assert_eq!(v["error"]["code"], json!(-32600), "request-shaped cancel must be -32600"); + assert_eq!( + v["error"]["code"], + json!(-32600), + "request-shaped cancel must be -32600" + ); assert_eq!(v["id"], json!(42)); assert!( v.get("result").is_none(), @@ -4232,7 +4579,10 @@ mod acp_review_fixes { ); let params = json!({"sessionId": sid}); let resp = handle_session_cancel(&sessions, Value::Null, Some(¶ms), true).await; - assert!(resp.is_none(), "a notification cancel must produce no response frame"); + assert!( + resp.is_none(), + "a notification cancel must produce no response frame" + ); // notify_one stored a permit, so notified() resolves immediately — the signal fired. assert!( tokio::time::timeout(std::time::Duration::from_millis(200), cancel.notified()) @@ -4328,19 +4678,23 @@ mod acp_ws_integration { } match frame["params"]["method"].as_str() { Some("initialize") => { - send(ws, json!({ - "jsonrpc": "2.0", "id": frame["id"].clone(), - "result": { - "protocolVersion": "2025-06-18", - "capabilities": { "tools": {} }, - "serverInfo": { "name": "test-ext", "version": "0" } - } - })).await; - Some("initialize") - } - // A notification: no reply is owed, but it must still be taken off the socket or the - // next `recv` in a test mistakes it for the frame it was waiting for. - Some("notifications/initialized") => Some("initialized"), + send( + ws, + json!({ + "jsonrpc": "2.0", "id": frame["id"].clone(), + "result": { + "protocolVersion": "2025-06-18", + "capabilities": { "tools": {} }, + "serverInfo": { "name": "test-ext", "version": "0" } + } + }), + ) + .await; + Some("initialize") + } + // A notification: no reply is owed, but it must still be taken off the socket or the + // next `recv` in a test mistakes it for the frame it was waiting for. + Some("notifications/initialized") => Some("initialized"), _ => None, } } @@ -4348,20 +4702,28 @@ mod acp_ws_integration { /// Drive `initialize` + `session/new` declaring one `type:acp` server, then answer the /// `mcp/connect` the gateway sends back. Returns the session id. async fn handshake(ws: &mut Ws, acp_id: &str, name: &str, connection_id: &str) -> String { - send(ws, json!({ - "jsonrpc": "2.0", "id": 1, "method": "initialize", - "params": {"protocolVersion": 1, "clientCapabilities": {}} - })).await; + send( + ws, + json!({ + "jsonrpc": "2.0", "id": 1, "method": "initialize", + "params": {"protocolVersion": 1, "clientCapabilities": {}} + }), + ) + .await; let init = recv(ws).await; assert!(init.get("result").is_some(), "initialize failed: {init}"); - send(ws, json!({ - "jsonrpc": "2.0", "id": 2, "method": "session/new", - "params": { - "cwd": "/w", - "mcpServers": [{"type": "acp", "id": acp_id, "name": name}] - } - })).await; + send( + ws, + json!({ + "jsonrpc": "2.0", "id": 2, "method": "session/new", + "params": { + "cwd": "/w", + "mcpServers": [{"type": "acp", "id": acp_id, "name": name}] + } + }), + ) + .await; // The gateway now does two things concurrently: answer session/new, and open the tunnel // by sending mcp/connect. Order is not guaranteed, so accept either first. @@ -4381,17 +4743,25 @@ mod acp_ws_integration { } if frame.get("method").and_then(Value::as_str) == Some("mcp/connect") { assert_eq!( - frame["params"]["acpId"], json!(acp_id), + frame["params"]["acpId"], + json!(acp_id), "mcp/connect must name the declared id" ); - send(ws, json!({ - "jsonrpc": "2.0", "id": frame["id"].clone(), - "result": {"connectionId": connection_id} - })).await; + send( + ws, + json!({ + "jsonrpc": "2.0", "id": frame["id"].clone(), + "result": {"connectionId": connection_id} + }), + ) + .await; connected = true; } else if frame.get("id") == Some(&json!(2)) { session_id = Some( - frame["result"]["sessionId"].as_str().expect("sessionId").to_string(), + frame["result"]["sessionId"] + .as_str() + .expect("sessionId") + .to_string(), ); } } @@ -4429,10 +4799,14 @@ mod acp_ws_integration { async fn the_inner_mcp_lifecycle_completes_before_the_tunnel_is_registered() { let (url, registry) = serve().await; let (mut ws, _) = tokio_tungstenite::connect_async(&url).await.unwrap(); - send(&mut ws, json!({ - "jsonrpc": "2.0", "id": 1, "method": "initialize", - "params": {"protocolVersion": 1, "clientCapabilities": {}} - })).await; + send( + &mut ws, + json!({ + "jsonrpc": "2.0", "id": 1, "method": "initialize", + "params": {"protocolVersion": 1, "clientCapabilities": {}} + }), + ) + .await; let _ = recv(&mut ws).await; send(&mut ws, json!({ "jsonrpc": "2.0", "id": 2, "method": "session/new", @@ -4445,40 +4819,48 @@ mod acp_ws_integration { // should also say which bug. let collect = async { let mut order: Vec = Vec::new(); - let mut connected = false; - let mut lifecycle = 0; - while !connected || lifecycle < 2 { - let f = recv(&mut ws).await; - if f.get("method").and_then(Value::as_str) == Some("mcp/connect") { - order.push("mcp/connect".into()); - send(&mut ws, json!({ - "jsonrpc": "2.0", "id": f["id"].clone(), - "result": {"connectionId": "conn-1"} - })).await; - connected = true; - } else if f.get("method").and_then(Value::as_str) == Some("mcp/message") { - let inner = f["params"]["method"].as_str().unwrap_or("").to_string(); - order.push(inner.clone()); - if inner == "initialize" { - // The registry must still be empty: a server that has not answered - // `initialize` has not agreed to serve anything yet. - assert!( - registry.lock().unwrap().is_empty(), - "the tunnel was registered before the MCP handshake completed" - ); - send(&mut ws, json!({ - "jsonrpc": "2.0", "id": f["id"].clone(), - "result": { - "protocolVersion": "2025-06-18", - "capabilities": { "tools": {} }, - "serverInfo": { "name": "test-ext", "version": "0" } - } - })).await; + let mut connected = false; + let mut lifecycle = 0; + while !connected || lifecycle < 2 { + let f = recv(&mut ws).await; + if f.get("method").and_then(Value::as_str) == Some("mcp/connect") { + order.push("mcp/connect".into()); + send( + &mut ws, + json!({ + "jsonrpc": "2.0", "id": f["id"].clone(), + "result": {"connectionId": "conn-1"} + }), + ) + .await; + connected = true; + } else if f.get("method").and_then(Value::as_str) == Some("mcp/message") { + let inner = f["params"]["method"].as_str().unwrap_or("").to_string(); + order.push(inner.clone()); + if inner == "initialize" { + // The registry must still be empty: a server that has not answered + // `initialize` has not agreed to serve anything yet. + assert!( + registry.lock().unwrap().is_empty(), + "the tunnel was registered before the MCP handshake completed" + ); + send( + &mut ws, + json!({ + "jsonrpc": "2.0", "id": f["id"].clone(), + "result": { + "protocolVersion": "2025-06-18", + "capabilities": { "tools": {} }, + "serverInfo": { "name": "test-ext", "version": "0" } + } + }), + ) + .await; + } + lifecycle += 1; } - lifecycle += 1; } - } - order + order }; let order = tokio::time::timeout(std::time::Duration::from_secs(10), collect) .await @@ -4511,10 +4893,14 @@ mod acp_ws_integration { async fn a_server_that_refuses_initialize_is_not_registered() { let (url, registry) = serve().await; let (mut ws, _) = tokio_tungstenite::connect_async(&url).await.unwrap(); - send(&mut ws, json!({ - "jsonrpc": "2.0", "id": 1, "method": "initialize", - "params": {"protocolVersion": 1, "clientCapabilities": {}} - })).await; + send( + &mut ws, + json!({ + "jsonrpc": "2.0", "id": 1, "method": "initialize", + "params": {"protocolVersion": 1, "clientCapabilities": {}} + }), + ) + .await; let _ = recv(&mut ws).await; send(&mut ws, json!({ "jsonrpc": "2.0", "id": 2, "method": "session/new", @@ -4529,16 +4915,24 @@ mod acp_ws_integration { let f = recv(&mut ws).await; match f.get("method").and_then(Value::as_str) { Some("mcp/connect") => { - send(&mut ws, json!({ - "jsonrpc": "2.0", "id": f["id"].clone(), - "result": {"connectionId": "conn-1"} - })).await; + send( + &mut ws, + json!({ + "jsonrpc": "2.0", "id": f["id"].clone(), + "result": {"connectionId": "conn-1"} + }), + ) + .await; } Some("mcp/message") if f["params"]["method"] == json!("initialize") => { - send(&mut ws, json!({ - "jsonrpc": "2.0", "id": f["id"].clone(), - "error": {"code": -32603, "message": "not accepting connections"} - })).await; + send( + &mut ws, + json!({ + "jsonrpc": "2.0", "id": f["id"].clone(), + "error": {"code": -32603, "message": "not accepting connections"} + }), + ) + .await; refused = true; } _ => {} @@ -4614,15 +5008,23 @@ mod acp_ws_integration { // B opens FIRST — older connection — and declares nothing yet. let (mut b, _) = tokio_tungstenite::connect_async(&url).await.unwrap(); - send(&mut b, json!({ - "jsonrpc": "2.0", "id": 1, "method": "initialize", - "params": {"protocolVersion": 1, "clientCapabilities": {}} - })).await; + send( + &mut b, + json!({ + "jsonrpc": "2.0", "id": 1, "method": "initialize", + "params": {"protocolVersion": 1, "clientCapabilities": {}} + }), + ) + .await; let _ = recv(&mut b).await; - send(&mut b, json!({ - "jsonrpc": "2.0", "id": 2, "method": "session/new", - "params": {"cwd": "/w", "mcpServers": []} - })).await; + send( + &mut b, + json!({ + "jsonrpc": "2.0", "id": 2, "method": "session/new", + "params": {"cwd": "/w", "mcpServers": []} + }), + ) + .await; let session_id = loop { let f = recv(&mut b).await; if f.get("id") == Some(&json!(2)) { @@ -4634,16 +5036,24 @@ mod acp_ws_integration { // `mcp/connect` is captured and deliberately left unanswered, so it takes the LOWER attach // number while making no progress. let (mut c, _) = tokio_tungstenite::connect_async(&url).await.unwrap(); - send(&mut c, json!({ - "jsonrpc": "2.0", "id": 1, "method": "initialize", - "params": {"protocolVersion": 1, "clientCapabilities": {}} - })).await; + send( + &mut c, + json!({ + "jsonrpc": "2.0", "id": 1, "method": "initialize", + "params": {"protocolVersion": 1, "clientCapabilities": {}} + }), + ) + .await; let _ = recv(&mut c).await; - send(&mut c, json!({ - "jsonrpc": "2.0", "id": 2, "method": "session/resume", - "params": {"sessionId": session_id.clone(), "cwd": "/w", - "mcpServers": [{"type": "acp", "id": "srv-c", "name": "katashiro"}]} - })).await; + send( + &mut c, + json!({ + "jsonrpc": "2.0", "id": 2, "method": "session/resume", + "params": {"sessionId": session_id.clone(), "cwd": "/w", + "mcpServers": [{"type": "acp", "id": "srv-c", "name": "katashiro"}]} + }), + ) + .await; let c_connect_id = loop { let f = recv(&mut c).await; if f.get("method").and_then(Value::as_str) == Some("mcp/connect") { @@ -4653,19 +5063,27 @@ mod acp_ws_integration { // Now the OLDER connection declares the same name under a different id and completes, so it // registers with the HIGHER attach number. - send(&mut b, json!({ - "jsonrpc": "2.0", "id": 9, "method": "session/resume", - "params": {"sessionId": session_id, "cwd": "/w", - "mcpServers": [{"type": "acp", "id": "srv-b", "name": "katashiro"}]} - })).await; + send( + &mut b, + json!({ + "jsonrpc": "2.0", "id": 9, "method": "session/resume", + "params": {"sessionId": session_id, "cwd": "/w", + "mcpServers": [{"type": "acp", "id": "srv-b", "name": "katashiro"}]} + }), + ) + .await; let mut done = false; while !done { let f = recv(&mut b).await; if f.get("method").and_then(Value::as_str) == Some("mcp/connect") { - send(&mut b, json!({ - "jsonrpc": "2.0", "id": f["id"].clone(), - "result": {"connectionId": "conn-b"} - })).await; + send( + &mut b, + json!({ + "jsonrpc": "2.0", "id": f["id"].clone(), + "result": {"connectionId": "conn-b"} + }), + ) + .await; } else if handled_inner_lifecycle(&mut b, &f).await == Some("initialize") { done = true; } @@ -4673,10 +5091,14 @@ mod acp_ws_integration { wait_for_tunnels(®istry, 1).await; // Finally let the stalled, newer-connection establish finish. - send(&mut c, json!({ - "jsonrpc": "2.0", "id": c_connect_id, - "result": {"connectionId": "conn-c"} - })).await; + send( + &mut c, + json!({ + "jsonrpc": "2.0", "id": c_connect_id, + "result": {"connectionId": "conn-c"} + }), + ) + .await; loop { let f = recv(&mut c).await; if handled_inner_lifecycle(&mut c, &f).await == Some("initialize") { @@ -4725,15 +5147,23 @@ mod acp_ws_integration { async fn within_one_connection_a_late_finishing_older_establish_loses_to_its_successor() { let (url, registry) = serve().await; let (mut ws, _) = tokio_tungstenite::connect_async(&url).await.unwrap(); - send(&mut ws, json!({ - "jsonrpc": "2.0", "id": 1, "method": "initialize", - "params": {"protocolVersion": 1, "clientCapabilities": {}} - })).await; + send( + &mut ws, + json!({ + "jsonrpc": "2.0", "id": 1, "method": "initialize", + "params": {"protocolVersion": 1, "clientCapabilities": {}} + }), + ) + .await; let _ = recv(&mut ws).await; - send(&mut ws, json!({ - "jsonrpc": "2.0", "id": 2, "method": "session/new", - "params": {"cwd": "/w", "mcpServers": []} - })).await; + send( + &mut ws, + json!({ + "jsonrpc": "2.0", "id": 2, "method": "session/new", + "params": {"cwd": "/w", "mcpServers": []} + }), + ) + .await; let session_id = loop { let f = recv(&mut ws).await; if f.get("id") == Some(&json!(2)) { @@ -4743,11 +5173,15 @@ mod acp_ws_integration { // First resume: declare srv-1 and PARK it — its `mcp/connect` is captured, not answered, so // it holds the lower attach number while making no progress. - send(&mut ws, json!({ - "jsonrpc": "2.0", "id": 3, "method": "session/resume", - "params": {"sessionId": session_id.clone(), "cwd": "/w", - "mcpServers": [{"type": "acp", "id": "srv-1", "name": "katashiro"}]} - })).await; + send( + &mut ws, + json!({ + "jsonrpc": "2.0", "id": 3, "method": "session/resume", + "params": {"sessionId": session_id.clone(), "cwd": "/w", + "mcpServers": [{"type": "acp", "id": "srv-1", "name": "katashiro"}]} + }), + ) + .await; let parked_connect_id = loop { let f = recv(&mut ws).await; if f.get("method").and_then(Value::as_str) == Some("mcp/connect") { @@ -4757,21 +5191,29 @@ mod acp_ws_integration { }; // Second resume on the SAME connection: same name, new id, driven to completion. - send(&mut ws, json!({ - "jsonrpc": "2.0", "id": 4, "method": "session/resume", - "params": {"sessionId": session_id, "cwd": "/w", - "mcpServers": [{"type": "acp", "id": "srv-2", "name": "katashiro"}]} - })).await; + send( + &mut ws, + json!({ + "jsonrpc": "2.0", "id": 4, "method": "session/resume", + "params": {"sessionId": session_id, "cwd": "/w", + "mcpServers": [{"type": "acp", "id": "srv-2", "name": "katashiro"}]} + }), + ) + .await; let mut registered = false; while !registered { let f = recv(&mut ws).await; if f.get("method").and_then(Value::as_str) == Some("mcp/connect") && f["params"]["acpId"] == json!("srv-2") { - send(&mut ws, json!({ - "jsonrpc": "2.0", "id": f["id"].clone(), - "result": {"connectionId": "conn-2"} - })).await; + send( + &mut ws, + json!({ + "jsonrpc": "2.0", "id": f["id"].clone(), + "result": {"connectionId": "conn-2"} + }), + ) + .await; } else if handled_inner_lifecycle(&mut ws, &f).await == Some("initialize") { registered = true; } @@ -4779,10 +5221,14 @@ mod acp_ws_integration { wait_for_tunnels(®istry, 1).await; // Now release the parked, EARLIER establish. - send(&mut ws, json!({ - "jsonrpc": "2.0", "id": parked_connect_id, - "result": {"connectionId": "conn-1"} - })).await; + send( + &mut ws, + json!({ + "jsonrpc": "2.0", "id": parked_connect_id, + "result": {"connectionId": "conn-1"} + }), + ) + .await; loop { let f = recv(&mut ws).await; if handled_inner_lifecycle(&mut ws, &f).await == Some("initialize") { @@ -4825,15 +5271,23 @@ mod acp_ws_integration { // B opens first, so its connection is the OLDER one. It declares nothing yet. let (mut b, _) = tokio_tungstenite::connect_async(&url).await.unwrap(); - send(&mut b, json!({ - "jsonrpc": "2.0", "id": 1, "method": "initialize", - "params": {"protocolVersion": 1, "clientCapabilities": {}} - })).await; + send( + &mut b, + json!({ + "jsonrpc": "2.0", "id": 1, "method": "initialize", + "params": {"protocolVersion": 1, "clientCapabilities": {}} + }), + ) + .await; let _ = recv(&mut b).await; - send(&mut b, json!({ - "jsonrpc": "2.0", "id": 2, "method": "session/new", - "params": {"cwd": "/w", "mcpServers": []} - })).await; + send( + &mut b, + json!({ + "jsonrpc": "2.0", "id": 2, "method": "session/new", + "params": {"cwd": "/w", "mcpServers": []} + }), + ) + .await; let session_id = loop { let f = recv(&mut b).await; if f.get("id") == Some(&json!(2)) { @@ -4843,23 +5297,35 @@ mod acp_ws_integration { // C opens second (NEWER connection) and establishes srv-3 under the shared name. let (mut c, _) = tokio_tungstenite::connect_async(&url).await.unwrap(); - send(&mut c, json!({ - "jsonrpc": "2.0", "id": 1, "method": "initialize", - "params": {"protocolVersion": 1, "clientCapabilities": {}} - })).await; + send( + &mut c, + json!({ + "jsonrpc": "2.0", "id": 1, "method": "initialize", + "params": {"protocolVersion": 1, "clientCapabilities": {}} + }), + ) + .await; let _ = recv(&mut c).await; - send(&mut c, json!({ - "jsonrpc": "2.0", "id": 2, "method": "session/resume", - "params": {"sessionId": session_id.clone(), "cwd": "/w", - "mcpServers": [{"type": "acp", "id": "srv-3", "name": "katashiro"}]} - })).await; + send( + &mut c, + json!({ + "jsonrpc": "2.0", "id": 2, "method": "session/resume", + "params": {"sessionId": session_id.clone(), "cwd": "/w", + "mcpServers": [{"type": "acp", "id": "srv-3", "name": "katashiro"}]} + }), + ) + .await; loop { let f = recv(&mut c).await; if f.get("method").and_then(Value::as_str) == Some("mcp/connect") { - send(&mut c, json!({ - "jsonrpc": "2.0", "id": f["id"].clone(), - "result": {"connectionId": "conn-c"} - })).await; + send( + &mut c, + json!({ + "jsonrpc": "2.0", "id": f["id"].clone(), + "result": {"connectionId": "conn-c"} + }), + ) + .await; } else if handled_inner_lifecycle(&mut c, &f).await == Some("initialize") { break; } @@ -4867,19 +5333,27 @@ mod acp_ws_integration { wait_for_tunnels(®istry, 1).await; // The OLDER connection now declares the same NAME under a different id, and completes. - send(&mut b, json!({ - "jsonrpc": "2.0", "id": 9, "method": "session/resume", - "params": {"sessionId": session_id, "cwd": "/w", - "mcpServers": [{"type": "acp", "id": "srv-4", "name": "katashiro"}]} - })).await; + send( + &mut b, + json!({ + "jsonrpc": "2.0", "id": 9, "method": "session/resume", + "params": {"sessionId": session_id, "cwd": "/w", + "mcpServers": [{"type": "acp", "id": "srv-4", "name": "katashiro"}]} + }), + ) + .await; let mut done = false; while !done { let f = recv(&mut b).await; if f.get("method").and_then(Value::as_str) == Some("mcp/connect") { - send(&mut b, json!({ - "jsonrpc": "2.0", "id": f["id"].clone(), - "result": {"connectionId": "conn-b"} - })).await; + send( + &mut b, + json!({ + "jsonrpc": "2.0", "id": f["id"].clone(), + "result": {"connectionId": "conn-b"} + }), + ) + .await; } else if handled_inner_lifecycle(&mut b, &f).await == Some("initialize") { done = true; } @@ -4906,7 +5380,11 @@ mod acp_ws_integration { let reg = registry.lock().unwrap(); reg.keys().map(|(_, id)| id.clone()).collect() }; - assert_eq!(ids, vec!["srv-3".to_string()], "the newer connection's tunnel must hold the name"); + assert_eq!( + ids, + vec!["srv-3".to_string()], + "the newer connection's tunnel must hold the name" + ); } /// An older connection's late resume must not TAKE OVER a newer connection's tunnel either. @@ -4929,15 +5407,23 @@ mod acp_ws_integration { // Connection B (older) opens the session but establishes nothing yet. let (mut b, _) = tokio_tungstenite::connect_async(&url).await.unwrap(); - send(&mut b, json!({ - "jsonrpc": "2.0", "id": 1, "method": "initialize", - "params": {"protocolVersion": 1, "clientCapabilities": {}} - })).await; + send( + &mut b, + json!({ + "jsonrpc": "2.0", "id": 1, "method": "initialize", + "params": {"protocolVersion": 1, "clientCapabilities": {}} + }), + ) + .await; let _ = recv(&mut b).await; - send(&mut b, json!({ - "jsonrpc": "2.0", "id": 2, "method": "session/new", - "params": {"cwd": "/w", "mcpServers": []} - })).await; + send( + &mut b, + json!({ + "jsonrpc": "2.0", "id": 2, "method": "session/new", + "params": {"cwd": "/w", "mcpServers": []} + }), + ) + .await; let session_id = loop { let f = recv(&mut b).await; if f.get("id") == Some(&json!(2)) { @@ -4947,23 +5433,35 @@ mod acp_ws_integration { // Connection C (newer) resumes and establishes srv-1. let (mut c, _) = tokio_tungstenite::connect_async(&url).await.unwrap(); - send(&mut c, json!({ - "jsonrpc": "2.0", "id": 1, "method": "initialize", - "params": {"protocolVersion": 1, "clientCapabilities": {}} - })).await; + send( + &mut c, + json!({ + "jsonrpc": "2.0", "id": 1, "method": "initialize", + "params": {"protocolVersion": 1, "clientCapabilities": {}} + }), + ) + .await; let _ = recv(&mut c).await; - send(&mut c, json!({ - "jsonrpc": "2.0", "id": 2, "method": "session/resume", - "params": {"sessionId": session_id.clone(), "cwd": "/w", - "mcpServers": [{"type": "acp", "id": "srv-1", "name": "katashiro"}]} - })).await; + send( + &mut c, + json!({ + "jsonrpc": "2.0", "id": 2, "method": "session/resume", + "params": {"sessionId": session_id.clone(), "cwd": "/w", + "mcpServers": [{"type": "acp", "id": "srv-1", "name": "katashiro"}]} + }), + ) + .await; loop { let f = recv(&mut c).await; if f.get("method").and_then(Value::as_str) == Some("mcp/connect") { - send(&mut c, json!({ - "jsonrpc": "2.0", "id": f["id"].clone(), - "result": {"connectionId": "conn-c"} - })).await; + send( + &mut c, + json!({ + "jsonrpc": "2.0", "id": f["id"].clone(), + "result": {"connectionId": "conn-c"} + }), + ) + .await; } else if handled_inner_lifecycle(&mut c, &f).await == Some("initialize") { break; } @@ -4971,19 +5469,27 @@ mod acp_ws_integration { wait_for_tunnels(®istry, 1).await; // The OLDER connection now resumes declaring THE SAME id, and answers its handshake fully. - send(&mut b, json!({ - "jsonrpc": "2.0", "id": 9, "method": "session/resume", - "params": {"sessionId": session_id, "cwd": "/w", - "mcpServers": [{"type": "acp", "id": "srv-1", "name": "katashiro"}]} - })).await; + send( + &mut b, + json!({ + "jsonrpc": "2.0", "id": 9, "method": "session/resume", + "params": {"sessionId": session_id, "cwd": "/w", + "mcpServers": [{"type": "acp", "id": "srv-1", "name": "katashiro"}]} + }), + ) + .await; let mut answered = false; while !answered { let f = recv(&mut b).await; if f.get("method").and_then(Value::as_str) == Some("mcp/connect") { - send(&mut b, json!({ - "jsonrpc": "2.0", "id": f["id"].clone(), - "result": {"connectionId": "conn-b"} - })).await; + send( + &mut b, + json!({ + "jsonrpc": "2.0", "id": f["id"].clone(), + "result": {"connectionId": "conn-b"} + }), + ) + .await; } else if handled_inner_lifecycle(&mut b, &f).await == Some("initialize") { answered = true; } @@ -4995,7 +5501,10 @@ mod acp_ws_integration { loop { let f = recv(&mut c).await; if f.get("method").and_then(Value::as_str) == Some("mcp/disconnect") { - return f["params"]["connectionId"].as_str().unwrap_or("?").to_string(); + return f["params"]["connectionId"] + .as_str() + .unwrap_or("?") + .to_string(); } } }) @@ -5007,7 +5516,11 @@ mod acp_ws_integration { late resume wins precisely because it ran late", stolen.ok() ); - assert_eq!(registry.lock().unwrap().len(), 1, "exactly one tunnel must hold the slot"); + assert_eq!( + registry.lock().unwrap().len(), + 1, + "exactly one tunnel must hold the slot" + ); } /// An older connection's late resume must not retire a NEWER connection's tunnel. @@ -5032,26 +5545,38 @@ mod acp_ws_integration { // Connection C (newer) resumes the same session and adds srv-3 under a different name. let (mut c, _) = tokio_tungstenite::connect_async(&url).await.unwrap(); - send(&mut c, json!({ - "jsonrpc": "2.0", "id": 1, "method": "initialize", - "params": {"protocolVersion": 1, "clientCapabilities": {}} - })).await; + send( + &mut c, + json!({ + "jsonrpc": "2.0", "id": 1, "method": "initialize", + "params": {"protocolVersion": 1, "clientCapabilities": {}} + }), + ) + .await; let _ = recv(&mut c).await; - send(&mut c, json!({ - "jsonrpc": "2.0", "id": 2, "method": "session/resume", - "params": {"sessionId": session_id.clone(), "cwd": "/w", - "mcpServers": [{"type": "acp", "id": "srv-2", "name": "katashiro"}, - {"type": "acp", "id": "srv-3", "name": "notes"}]} - })).await; + send( + &mut c, + json!({ + "jsonrpc": "2.0", "id": 2, "method": "session/resume", + "params": {"sessionId": session_id.clone(), "cwd": "/w", + "mcpServers": [{"type": "acp", "id": "srv-2", "name": "katashiro"}, + {"type": "acp", "id": "srv-3", "name": "notes"}]} + }), + ) + .await; loop { let f = recv(&mut c).await; if f.get("method").and_then(Value::as_str) == Some("mcp/connect") && f["params"]["acpId"] == json!("srv-3") { - send(&mut c, json!({ - "jsonrpc": "2.0", "id": f["id"].clone(), - "result": {"connectionId": "conn-c"} - })).await; + send( + &mut c, + json!({ + "jsonrpc": "2.0", "id": f["id"].clone(), + "result": {"connectionId": "conn-c"} + }), + ) + .await; } else if handled_inner_lifecycle(&mut c, &f).await == Some("initialize") { break; } @@ -5059,11 +5584,15 @@ mod acp_ws_integration { wait_for_tunnels(®istry, 2).await; // Now the OLDER connection resumes, still declaring only what it knew about. - send(&mut b, json!({ - "jsonrpc": "2.0", "id": 9, "method": "session/resume", - "params": {"sessionId": session_id, "cwd": "/w", - "mcpServers": [{"type": "acp", "id": "srv-2", "name": "katashiro"}]} - })).await; + send( + &mut b, + json!({ + "jsonrpc": "2.0", "id": 9, "method": "session/resume", + "params": {"sessionId": session_id, "cwd": "/w", + "mcpServers": [{"type": "acp", "id": "srv-2", "name": "katashiro"}]} + }), + ) + .await; // srv-3 belongs to a newer connection and must survive. Poll: a wrongful retire is async. let deadline = std::time::Instant::now() + std::time::Duration::from_secs(2); @@ -5099,16 +5628,24 @@ mod acp_ws_integration { wait_for_tunnels(®istry, 1).await; let (mut b, _) = tokio_tungstenite::connect_async(&url).await.unwrap(); - send(&mut b, json!({ - "jsonrpc": "2.0", "id": 1, "method": "initialize", - "params": {"protocolVersion": 1, "clientCapabilities": {}} - })).await; + send( + &mut b, + json!({ + "jsonrpc": "2.0", "id": 1, "method": "initialize", + "params": {"protocolVersion": 1, "clientCapabilities": {}} + }), + ) + .await; let _ = recv(&mut b).await; // No `mcpServers` key at all. - send(&mut b, json!({ - "jsonrpc": "2.0", "id": 2, "method": "session/resume", - "params": {"sessionId": session_id, "cwd": "/w"} - })).await; + send( + &mut b, + json!({ + "jsonrpc": "2.0", "id": 2, "method": "session/resume", + "params": {"sessionId": session_id, "cwd": "/w"} + }), + ) + .await; let resumed = recv(&mut b).await; assert!(resumed.get("result").is_some(), "resume failed: {resumed}"); @@ -5120,12 +5657,19 @@ mod acp_ws_integration { // reads as if it were the only thing standing between `null` and a full sweep, and the next // person to relax the schema needs the two facts in one place. for shape in [json!(null), json!({}), json!("nonsense")] { - send(&mut b, json!({ - "jsonrpc": "2.0", "id": 3, "method": "session/resume", - "params": {"sessionId": session_id, "cwd": "/w", "mcpServers": shape} - })).await; + send( + &mut b, + json!({ + "jsonrpc": "2.0", "id": 3, "method": "session/resume", + "params": {"sessionId": session_id, "cwd": "/w", "mcpServers": shape} + }), + ) + .await; let r = recv(&mut b).await; - assert!(r.get("result").is_some() || r.get("error").is_some(), "no reply: {r}"); + assert!( + r.get("result").is_some() || r.get("error").is_some(), + "no reply: {r}" + ); } // The tunnel must stay. Poll, because a wrongful teardown is asynchronous. @@ -5169,11 +5713,15 @@ mod acp_ws_integration { wait_for_tunnels(®istry, 1).await; // Same connection, same id, re-declared. - send(&mut a, json!({ - "jsonrpc": "2.0", "id": 3, "method": "session/resume", - "params": {"sessionId": session_id, "cwd": "/w", - "mcpServers": [{"type": "acp", "id": "srv-1", "name": "katashiro"}]} - })).await; + send( + &mut a, + json!({ + "jsonrpc": "2.0", "id": 3, "method": "session/resume", + "params": {"sessionId": session_id, "cwd": "/w", + "mcpServers": [{"type": "acp", "id": "srv-1", "name": "katashiro"}]} + }), + ) + .await; // Nothing may be disconnected. A sweep-everything implementation retires conn-1 here and // then re-establishes it, which this catches; the correct implementation sends no @@ -5182,7 +5730,10 @@ mod acp_ws_integration { loop { let f = recv(&mut a).await; if f.get("method").and_then(Value::as_str) == Some("mcp/disconnect") { - return f["params"]["connectionId"].as_str().unwrap_or("?").to_string(); + return f["params"]["connectionId"] + .as_str() + .unwrap_or("?") + .to_string(); } } }) @@ -5193,7 +5744,11 @@ mod acp_ws_integration { set is not 'registered minus declared', it is 'everything'", quiet.ok() ); - assert_eq!(registry.lock().unwrap().len(), 1, "the tunnel must still be registered"); + assert_eq!( + registry.lock().unwrap().len(), + 1, + "the tunnel must still be registered" + ); } /// A resume on a NEW connection that stops declaring a server retires its tunnel. @@ -5220,15 +5775,23 @@ mod acp_ws_integration { // Connection 2 — a genuine reconnect — resumes the same session declaring NOTHING. let (mut b, _) = tokio_tungstenite::connect_async(&url).await.unwrap(); - send(&mut b, json!({ - "jsonrpc": "2.0", "id": 1, "method": "initialize", - "params": {"protocolVersion": 1, "clientCapabilities": {}} - })).await; + send( + &mut b, + json!({ + "jsonrpc": "2.0", "id": 1, "method": "initialize", + "params": {"protocolVersion": 1, "clientCapabilities": {}} + }), + ) + .await; let _ = recv(&mut b).await; - send(&mut b, json!({ - "jsonrpc": "2.0", "id": 2, "method": "session/resume", - "params": {"sessionId": session_id, "cwd": "/w", "mcpServers": []} - })).await; + send( + &mut b, + json!({ + "jsonrpc": "2.0", "id": 2, "method": "session/resume", + "params": {"sessionId": session_id, "cwd": "/w", "mcpServers": []} + }), + ) + .await; let resumed = recv(&mut b).await; assert!(resumed.get("result").is_some(), "resume failed: {resumed}"); @@ -5285,10 +5848,14 @@ mod acp_ws_integration { // Socket A: declare `katashiro` as srv-1, then STALL — do not answer its `mcp/connect`. let (mut a, _) = tokio_tungstenite::connect_async(&url).await.unwrap(); - send(&mut a, json!({ - "jsonrpc": "2.0", "id": 1, "method": "initialize", - "params": {"protocolVersion": 1, "clientCapabilities": {}} - })).await; + send( + &mut a, + json!({ + "jsonrpc": "2.0", "id": 1, "method": "initialize", + "params": {"protocolVersion": 1, "clientCapabilities": {}} + }), + ) + .await; let _ = recv(&mut a).await; send(&mut a, json!({ "jsonrpc": "2.0", "id": 2, "method": "session/new", @@ -5311,23 +5878,35 @@ mod acp_ws_integration { // Socket B: resume the SAME session — same channel — declaring the same name as srv-2, and // carry it all the way to registered. This is the newer attach. let (mut b, _) = tokio_tungstenite::connect_async(&url).await.unwrap(); - send(&mut b, json!({ - "jsonrpc": "2.0", "id": 1, "method": "initialize", - "params": {"protocolVersion": 1, "clientCapabilities": {}} - })).await; + send( + &mut b, + json!({ + "jsonrpc": "2.0", "id": 1, "method": "initialize", + "params": {"protocolVersion": 1, "clientCapabilities": {}} + }), + ) + .await; let _ = recv(&mut b).await; - send(&mut b, json!({ - "jsonrpc": "2.0", "id": 2, "method": "session/resume", - "params": {"sessionId": session_id, "cwd": "/w", - "mcpServers": [{"type": "acp", "id": "srv-2", "name": "katashiro"}]} - })).await; + send( + &mut b, + json!({ + "jsonrpc": "2.0", "id": 2, "method": "session/resume", + "params": {"sessionId": session_id, "cwd": "/w", + "mcpServers": [{"type": "acp", "id": "srv-2", "name": "katashiro"}]} + }), + ) + .await; loop { let f = recv(&mut b).await; if f.get("method").and_then(Value::as_str) == Some("mcp/connect") { - send(&mut b, json!({ - "jsonrpc": "2.0", "id": f["id"].clone(), - "result": {"connectionId": "conn-new"} - })).await; + send( + &mut b, + json!({ + "jsonrpc": "2.0", "id": f["id"].clone(), + "result": {"connectionId": "conn-new"} + }), + ) + .await; } else if handled_inner_lifecycle(&mut b, &f).await == Some("initialize") { break; } @@ -5335,10 +5914,14 @@ mod acp_ws_integration { wait_for_tunnels(®istry, 1).await; // Now let the OLDER establish finish. - send(&mut a, json!({ - "jsonrpc": "2.0", "id": a_connect_id.unwrap(), - "result": {"connectionId": "conn-old"} - })).await; + send( + &mut a, + json!({ + "jsonrpc": "2.0", "id": a_connect_id.unwrap(), + "result": {"connectionId": "conn-old"} + }), + ) + .await; loop { let f = recv(&mut a).await; if handled_inner_lifecycle(&mut a, &f).await == Some("initialize") { @@ -5404,10 +5987,14 @@ mod acp_ws_integration { // Socket A declares srv-1 and stalls with its `mcp/connect` unanswered. let (mut a, _) = tokio_tungstenite::connect_async(&url).await.unwrap(); - send(&mut a, json!({ - "jsonrpc": "2.0", "id": 1, "method": "initialize", - "params": {"protocolVersion": 1, "clientCapabilities": {}} - })).await; + send( + &mut a, + json!({ + "jsonrpc": "2.0", "id": 1, "method": "initialize", + "params": {"protocolVersion": 1, "clientCapabilities": {}} + }), + ) + .await; let _ = recv(&mut a).await; send(&mut a, json!({ "jsonrpc": "2.0", "id": 2, "method": "session/new", @@ -5426,23 +6013,35 @@ mod acp_ws_integration { // Socket B reconnects and re-declares THE SAME id, completing fully. let (mut b, _) = tokio_tungstenite::connect_async(&url).await.unwrap(); - send(&mut b, json!({ - "jsonrpc": "2.0", "id": 1, "method": "initialize", - "params": {"protocolVersion": 1, "clientCapabilities": {}} - })).await; + send( + &mut b, + json!({ + "jsonrpc": "2.0", "id": 1, "method": "initialize", + "params": {"protocolVersion": 1, "clientCapabilities": {}} + }), + ) + .await; let _ = recv(&mut b).await; - send(&mut b, json!({ - "jsonrpc": "2.0", "id": 2, "method": "session/resume", - "params": {"sessionId": session_id.unwrap(), "cwd": "/w", - "mcpServers": [{"type": "acp", "id": "srv-1", "name": "katashiro"}]} - })).await; + send( + &mut b, + json!({ + "jsonrpc": "2.0", "id": 2, "method": "session/resume", + "params": {"sessionId": session_id.unwrap(), "cwd": "/w", + "mcpServers": [{"type": "acp", "id": "srv-1", "name": "katashiro"}]} + }), + ) + .await; loop { let f = recv(&mut b).await; if f.get("method").and_then(Value::as_str) == Some("mcp/connect") { - send(&mut b, json!({ - "jsonrpc": "2.0", "id": f["id"].clone(), - "result": {"connectionId": "conn-new"} - })).await; + send( + &mut b, + json!({ + "jsonrpc": "2.0", "id": f["id"].clone(), + "result": {"connectionId": "conn-new"} + }), + ) + .await; } else if handled_inner_lifecycle(&mut b, &f).await == Some("initialize") { break; } @@ -5450,10 +6049,14 @@ mod acp_ws_integration { wait_for_tunnels(®istry, 1).await; // Now let the OLDER establish finish, on the same key. - send(&mut a, json!({ - "jsonrpc": "2.0", "id": a_connect_id.unwrap(), - "result": {"connectionId": "conn-old"} - })).await; + send( + &mut a, + json!({ + "jsonrpc": "2.0", "id": a_connect_id.unwrap(), + "result": {"connectionId": "conn-old"} + }), + ) + .await; loop { let f = recv(&mut a).await; if handled_inner_lifecycle(&mut a, &f).await == Some("initialize") { @@ -5472,15 +6075,21 @@ mod acp_ws_integration { } } }; - let disconnected = - tokio::time::timeout(std::time::Duration::from_secs(10), wait_disconnect) - .await - .expect( - "an older establish reusing the same server_id overwrote the newer live handle \ + let disconnected = tokio::time::timeout( + std::time::Duration::from_secs(10), + wait_disconnect, + ) + .await + .expect( + "an older establish reusing the same server_id overwrote the newer live handle \ instead of standing down — same-key ordering is not enforced", - ); + ); assert_eq!(disconnected, "conn-old"); - assert_eq!(registry.lock().unwrap().len(), 1, "exactly one tunnel must remain"); + assert_eq!( + registry.lock().unwrap().len(), + 1, + "exactly one tunnel must remain" + ); } /// A server that *succeeds* at `initialize` but answers a protocol version we do not speak @@ -5501,10 +6110,14 @@ mod acp_ws_integration { async fn a_server_answering_an_unsupported_protocol_version_is_not_registered() { let (url, registry) = serve().await; let (mut ws, _) = tokio_tungstenite::connect_async(&url).await.unwrap(); - send(&mut ws, json!({ - "jsonrpc": "2.0", "id": 1, "method": "initialize", - "params": {"protocolVersion": 1, "clientCapabilities": {}} - })).await; + send( + &mut ws, + json!({ + "jsonrpc": "2.0", "id": 1, "method": "initialize", + "params": {"protocolVersion": 1, "clientCapabilities": {}} + }), + ) + .await; let _ = recv(&mut ws).await; send(&mut ws, json!({ "jsonrpc": "2.0", "id": 2, "method": "session/new", @@ -5528,20 +6141,28 @@ mod acp_ws_integration { let f = recv(&mut ws).await; match f.get("method").and_then(Value::as_str) { Some("mcp/connect") => { - send(&mut ws, json!({ - "jsonrpc": "2.0", "id": f["id"].clone(), - "result": {"connectionId": "conn-1"} - })).await; + send( + &mut ws, + json!({ + "jsonrpc": "2.0", "id": f["id"].clone(), + "result": {"connectionId": "conn-1"} + }), + ) + .await; } Some("mcp/message") if f["params"]["method"] == json!("initialize") => { - send(&mut ws, json!({ - "jsonrpc": "2.0", "id": f["id"].clone(), - "result": { - "protocolVersion": "2019-01-01", - "capabilities": { "tools": {} }, - "serverInfo": { "name": "old-ext", "version": "0" } - } - })).await; + send( + &mut ws, + json!({ + "jsonrpc": "2.0", "id": f["id"].clone(), + "result": { + "protocolVersion": "2019-01-01", + "capabilities": { "tools": {} }, + "serverInfo": { "name": "old-ext", "version": "0" } + } + }), + ) + .await; answered = true; } _ => {} @@ -5600,10 +6221,15 @@ mod acp_ws_integration { // Server side, exactly as core reaches a session's tunnel. let handle = { let reg = registry.lock().unwrap(); - reg.values().next().expect("a tunnel must be registered").clone() + reg.values() + .next() + .expect("a tunnel must be registered") + .clone() }; let call = tokio::spawn(async move { - handle.mcp_message("tools/call", Some(json!({"name": "katashiro.click"})), 5).await + handle + .mcp_message("tools/call", Some(json!({"name": "katashiro.click"})), 5) + .await }); let framed = recv(&mut ws).await; @@ -5612,10 +6238,14 @@ mod acp_ws_integration { assert_eq!(framed["params"]["method"], json!("tools/call")); assert_eq!(framed["params"]["params"]["name"], json!("katashiro.click")); - send(&mut ws, json!({ - "jsonrpc": "2.0", "id": framed["id"].clone(), - "result": {"content": [{"type": "text", "text": "clicked"}]} - })).await; + send( + &mut ws, + json!({ + "jsonrpc": "2.0", "id": framed["id"].clone(), + "result": {"content": [{"type": "text", "text": "clicked"}]} + }), + ) + .await; let got = call.await.unwrap().expect("the call must succeed"); assert_eq!(got["content"][0]["text"], json!("clicked")); @@ -5637,12 +6267,19 @@ mod acp_ws_integration { let call = tokio::spawn(async move { handle.mcp_message("tools/call", None, 5).await }); let framed = recv(&mut ws).await; - send(&mut ws, json!({ - "jsonrpc": "2.0", "id": framed["id"].clone(), - "error": {"code": -32603, "message": "no active tab"} - })).await; + send( + &mut ws, + json!({ + "jsonrpc": "2.0", "id": framed["id"].clone(), + "error": {"code": -32603, "message": "no active tab"} + }), + ) + .await; - let err = call.await.unwrap().expect_err("a remote error must not read as success"); + let err = call + .await + .unwrap() + .expect_err("a remote error must not read as success"); assert!( err.contains("no active tab"), "the client's message must reach the caller, got: {err}" @@ -5690,7 +6327,10 @@ mod acp_ws_integration { "a cancellation is a notification: giving it an `id` would oblige a reply nobody reads" ); - let err = call.await.unwrap().expect_err("a timed-out call must not read as success"); + let err = call + .await + .unwrap() + .expect_err("a timed-out call must not read as success"); assert!(err.contains("timed out"), "got: {err}"); } @@ -5708,10 +6348,14 @@ mod acp_ws_integration { async fn pending_tunnel_establishes_do_not_consume_the_prompt_budget() { let (url, _registry) = serve().await; let (mut ws, _) = tokio_tungstenite::connect_async(&url).await.unwrap(); - send(&mut ws, json!({ - "jsonrpc": "2.0", "id": 1, "method": "initialize", - "params": {"protocolVersion": 1, "clientCapabilities": {}} - })).await; + send( + &mut ws, + json!({ + "jsonrpc": "2.0", "id": 1, "method": "initialize", + "params": {"protocolVersion": 1, "clientCapabilities": {}} + }), + ) + .await; let _ = recv(&mut ws).await; // Enough sessions to park strictly more than MAX_INFLIGHT_PROMPTS establishes. @@ -5734,10 +6378,14 @@ mod acp_ws_integration { let servers: Vec = (0..MAX_ACP_SERVERS_PER_SESSION) .map(|i| json!({"type": "acp", "id": format!("s{n}-{i}"), "name": format!("n{n}-{i}")})) .collect(); - send(&mut ws, json!({ - "jsonrpc": "2.0", "id": req_id, "method": "session/new", - "params": {"cwd": "/w", "mcpServers": servers} - })).await; + send( + &mut ws, + json!({ + "jsonrpc": "2.0", "id": req_id, "method": "session/new", + "params": {"cwd": "/w", "mcpServers": servers} + }), + ) + .await; // Collect this session's response; mcp/connect frames are observed and NEVER answered, // which is what keeps each establish task in flight. let mut got_resp = false; @@ -5810,11 +6458,12 @@ mod acp_ws_integration { // substring check is satisfied by every OTHER error too — a session-not-found, a params // error, or the budget message itself once someone rewords it — so it could only ever fail // for one spelling of one regression, and would pass silently through the rest. - let err = f - .get("error") - .unwrap_or_else(|| panic!("the prompt must reach the backend and fail there, got: {f}")); + let err = f.get("error").unwrap_or_else(|| { + panic!("the prompt must reach the backend and fail there, got: {f}") + }); assert_eq!( - err["message"], json!("No agent backend connected"), + err["message"], + json!("No agent backend connected"), "{connects} parked mcp/connects must not spend the prompt budget; the prompt must \ reach the backend and fail only for the missing backend, got: {f}" ); @@ -5832,18 +6481,26 @@ mod acp_ws_integration { let (mut ws, _) = tokio_tungstenite::connect_async(&url).await.unwrap(); // Declare TWO servers, answer both mcp/connect calls. - send(&mut ws, json!({ - "jsonrpc": "2.0", "id": 1, "method": "initialize", - "params": {"protocolVersion": 1, "clientCapabilities": {}} - })).await; + send( + &mut ws, + json!({ + "jsonrpc": "2.0", "id": 1, "method": "initialize", + "params": {"protocolVersion": 1, "clientCapabilities": {}} + }), + ) + .await; let _ = recv(&mut ws).await; - send(&mut ws, json!({ - "jsonrpc": "2.0", "id": 2, "method": "session/new", - "params": {"cwd": "/w", "mcpServers": [ - {"type": "acp", "id": "keep-1", "name": "katashiro"}, - {"type": "acp", "id": "drop-1", "name": "other"} - ]} - })).await; + send( + &mut ws, + json!({ + "jsonrpc": "2.0", "id": 2, "method": "session/new", + "params": {"cwd": "/w", "mcpServers": [ + {"type": "acp", "id": "keep-1", "name": "katashiro"}, + {"type": "acp", "id": "drop-1", "name": "other"} + ]} + }), + ) + .await; // Two servers: two `mcp/connect`s AND two lifecycle pairs (initialize + notification). // Exiting on the connects alone leaves the handshakes unread and both establishes fail. let mut session_id = None; @@ -5857,10 +6514,14 @@ mod acp_ws_integration { } if f.get("method").and_then(Value::as_str) == Some("mcp/connect") { let acp_id = f["params"]["acpId"].as_str().unwrap().to_string(); - send(&mut ws, json!({ - "jsonrpc": "2.0", "id": f["id"].clone(), - "result": {"connectionId": format!("conn-{acp_id}")} - })).await; + send( + &mut ws, + json!({ + "jsonrpc": "2.0", "id": f["id"].clone(), + "result": {"connectionId": format!("conn-{acp_id}")} + }), + ) + .await; connects += 1; } else if f.get("id") == Some(&json!(2)) { session_id = Some(f["result"]["sessionId"].as_str().unwrap().to_string()); @@ -5869,14 +6530,18 @@ mod acp_ws_integration { wait_for_tunnels(®istry, 2).await; // Resume declaring ONLY the first — "other" is withdrawn. - send(&mut ws, json!({ - "jsonrpc": "2.0", "id": 3, "method": "session/resume", - "params": { - "sessionId": session_id.unwrap(), - "cwd": "/w", - "mcpServers": [{"type": "acp", "id": "keep-2", "name": "katashiro"}] - } - })).await; + send( + &mut ws, + json!({ + "jsonrpc": "2.0", "id": 3, "method": "session/resume", + "params": { + "sessionId": session_id.unwrap(), + "cwd": "/w", + "mcpServers": [{"type": "acp", "id": "keep-2", "name": "katashiro"}] + } + }), + ) + .await; // TWO disconnects are owed here, for different reasons, and both are correct: // - `conn-drop-1` because the resume WITHDREW that declaration (this item), and @@ -5897,10 +6562,14 @@ mod acp_ws_integration { disconnected.push(f["params"]["connectionId"].as_str().unwrap().to_string()); } Some("mcp/connect") => { - send(&mut ws, json!({ - "jsonrpc": "2.0", "id": f["id"].clone(), - "result": {"connectionId": "conn-keep-2"} - })).await; + send( + &mut ws, + json!({ + "jsonrpc": "2.0", "id": f["id"].clone(), + "result": {"connectionId": "conn-keep-2"} + }), + ) + .await; reconnected = true; } _ => {} @@ -5944,8 +6613,11 @@ mod acp_ws_integration { let (url, _registry) = serve().await; let (mut ws, _) = tokio_tungstenite::connect_async(&url).await.unwrap(); let oversized = "x".repeat(super::MAX_FRAME_BYTES + 64); - send(&mut ws, json!({"jsonrpc": "2.0", "id": 1, "method": "initialize", "pad": oversized})) - .await; + send( + &mut ws, + json!({"jsonrpc": "2.0", "id": 1, "method": "initialize", "pad": oversized}), + ) + .await; // The connection must go away rather than answer. Bounded so a regression that keeps it // open fails here instead of hanging. @@ -5979,22 +6651,36 @@ mod acp_ws_integration { // Over 1 MiB, comfortably under 8 MiB — so it reaches the per-kind check, not the // transport one. Asserting the gap between the two ceilings is the point. let pad = "y".repeat(super::MAX_NON_TUNNEL_FRAME_BYTES + 4096); - assert!(pad.len() < super::MAX_FRAME_BYTES, "must not trip the transport ceiling"); - send(&mut ws, json!({ - "jsonrpc": "2.0", "id": 7, "method": "initialize", - "params": {"protocolVersion": 1, "clientCapabilities": {}, "pad": pad} - })).await; + assert!( + pad.len() < super::MAX_FRAME_BYTES, + "must not trip the transport ceiling" + ); + send( + &mut ws, + json!({ + "jsonrpc": "2.0", "id": 7, "method": "initialize", + "params": {"protocolVersion": 1, "clientCapabilities": {}, "pad": pad} + }), + ) + .await; let resp = recv(&mut ws).await; assert_eq!(resp["id"], json!(7)); - assert!(resp.get("error").is_some(), "an oversized request must be answered with an error: {resp}"); + assert!( + resp.get("error").is_some(), + "an oversized request must be answered with an error: {resp}" + ); // And the connection still works — the discriminating half. A regression that closed here // would satisfy every assertion above. - send(&mut ws, json!({ - "jsonrpc": "2.0", "id": 8, "method": "initialize", - "params": {"protocolVersion": 1, "clientCapabilities": {}} - })).await; + send( + &mut ws, + json!({ + "jsonrpc": "2.0", "id": 8, "method": "initialize", + "params": {"protocolVersion": 1, "clientCapabilities": {}} + }), + ) + .await; let after = recv(&mut ws).await; assert_eq!(after["id"], json!(8)); assert!( @@ -6048,24 +6734,32 @@ mod acp_ws_integration { async fn two_same_name_servers_in_one_session_evict_down_to_one() { let (url, registry) = serve().await; let (mut ws, _) = tokio_tungstenite::connect_async(&url).await.unwrap(); - send(&mut ws, json!({ - "jsonrpc": "2.0", "id": 1, "method": "initialize", - "params": {"protocolVersion": 1, "clientCapabilities": {}} - })).await; + send( + &mut ws, + json!({ + "jsonrpc": "2.0", "id": 1, "method": "initialize", + "params": {"protocolVersion": 1, "clientCapabilities": {}} + }), + ) + .await; let init = recv(&mut ws).await; assert!(init.get("result").is_some(), "initialize failed: {init}"); // Same name, different ids — accepted, because the dedup key is the id. - send(&mut ws, json!({ - "jsonrpc": "2.0", "id": 2, "method": "session/new", - "params": { - "cwd": "/w", - "mcpServers": [ - {"type": "acp", "id": "uuid-a", "name": "katashiro"}, - {"type": "acp", "id": "uuid-b", "name": "katashiro"} - ] - } - })).await; + send( + &mut ws, + json!({ + "jsonrpc": "2.0", "id": 2, "method": "session/new", + "params": { + "cwd": "/w", + "mcpServers": [ + {"type": "acp", "id": "uuid-a", "name": "katashiro"}, + {"type": "acp", "id": "uuid-b", "name": "katashiro"} + ] + } + }), + ) + .await; // Answer both connects, absorb both inner lifecycles, and wait for the disconnect the // loser is owed. Bounded so a regression fails naming what was missing rather than @@ -6093,28 +6787,45 @@ mod acp_ws_integration { } match frame.get("method").and_then(Value::as_str) { Some("mcp/connect") => { - let acp_id = frame["params"]["acpId"].as_str().expect("acpId").to_string(); + let acp_id = frame["params"]["acpId"] + .as_str() + .expect("acpId") + .to_string(); let conn = format!("conn-{}", acp_id.trim_start_matches("uuid-")); connected.insert(acp_id, conn.clone()); - send(&mut ws, json!({ - "jsonrpc": "2.0", "id": frame["id"].clone(), - "result": {"connectionId": conn} - })).await; + send( + &mut ws, + json!({ + "jsonrpc": "2.0", "id": frame["id"].clone(), + "result": {"connectionId": conn} + }), + ) + .await; } Some("mcp/disconnect") => { disconnected = Some( - frame["params"]["connectionId"].as_str().expect("connectionId").to_string(), + frame["params"]["connectionId"] + .as_str() + .expect("connectionId") + .to_string(), ); } _ => { if frame.get("id") == Some(&json!(2)) { - assert!(frame.get("result").is_some(), "session/new refused: {frame}"); + assert!( + frame.get("result").is_some(), + "session/new refused: {frame}" + ); session_seen = true; } } } } - assert_eq!(connected.len(), 2, "both declared servers must be asked to connect"); + assert_eq!( + connected.len(), + 2, + "both declared servers must be asked to connect" + ); let evicted = disconnected.expect( "the evicted tunnel is owed an mcp/disconnect — if this is missing, the eviction \ branch did not run and this scenario no longer covers it", @@ -6125,7 +6836,11 @@ mod acp_ws_integration { let reg = registry.lock().unwrap(); reg.values().map(|h| h.connection_id.clone()).collect() }; - assert_eq!(survivors.len(), 1, "last-attach-wins must collapse the name to one tunnel"); + assert_eq!( + survivors.len(), + 1, + "last-attach-wins must collapse the name to one tunnel" + ); assert_ne!( survivors[0], evicted, "the disconnect must name the tunnel that LOST, not the one still registered" @@ -6154,19 +6869,27 @@ mod acp_ws_integration { // The client comes back on a fresh socket and resumes, re-declaring under the same name // with the fresh id its runtime mints per connection. let (mut second, _) = tokio_tungstenite::connect_async(&url).await.unwrap(); - send(&mut second, json!({ - "jsonrpc": "2.0", "id": 1, "method": "initialize", - "params": {"protocolVersion": 1, "clientCapabilities": {}} - })).await; + send( + &mut second, + json!({ + "jsonrpc": "2.0", "id": 1, "method": "initialize", + "params": {"protocolVersion": 1, "clientCapabilities": {}} + }), + ) + .await; let _ = recv(&mut second).await; - send(&mut second, json!({ - "jsonrpc": "2.0", "id": 2, "method": "session/resume", - "params": { - "sessionId": session_id, - "cwd": "/w", - "mcpServers": [{"type": "acp", "id": "uuid-new", "name": "katashiro"}] - } - })).await; + send( + &mut second, + json!({ + "jsonrpc": "2.0", "id": 2, "method": "session/resume", + "params": { + "sessionId": session_id, + "cwd": "/w", + "mcpServers": [{"type": "acp", "id": "uuid-new", "name": "katashiro"}] + } + }), + ) + .await; let mut connected_new = false; let mut new_lifecycle = 0; while !connected_new || new_lifecycle < 2 { @@ -6176,10 +6899,14 @@ mod acp_ws_integration { continue; } if frame.get("method").and_then(Value::as_str) == Some("mcp/connect") { - send(&mut second, json!({ - "jsonrpc": "2.0", "id": frame["id"].clone(), - "result": {"connectionId": "conn-new"} - })).await; + send( + &mut second, + json!({ + "jsonrpc": "2.0", "id": frame["id"].clone(), + "result": {"connectionId": "conn-new"} + }), + ) + .await; connected_new = true; } if frame.get("id") == Some(&json!(2)) { @@ -6214,7 +6941,8 @@ mod acp_ws_integration { let framed = recv(&mut first).await; assert_eq!(framed["method"], json!("mcp/disconnect")); assert_eq!( - framed["params"]["connectionId"], json!("conn-old"), + framed["params"]["connectionId"], + json!("conn-old"), "the evicted connection is the one owed a disconnect, not its replacement" ); @@ -6223,7 +6951,11 @@ mod acp_ws_integration { let reg = registry.lock().unwrap(); reg.values().map(|h| h.connection_id.clone()).collect() }; - assert_eq!(names, vec!["conn-new".to_string()], "only the newest tunnel may remain"); + assert_eq!( + names, + vec!["conn-new".to_string()], + "only the newest tunnel may remain" + ); // The survivor must WORK, not merely be present in the registry: a replacement that // registers a handle whose socket is gone satisfies every assertion above. This is the @@ -6236,23 +6968,36 @@ mod acp_ws_integration { // the reason, where a missing-frame check could only time out. let handle = { let reg = registry.lock().unwrap(); - reg.values().next().expect("the survivor must be registered").clone() + reg.values() + .next() + .expect("the survivor must be registered") + .clone() }; let call = tokio::spawn(async move { - handle.mcp_message("tools/call", Some(json!({"name": "katashiro.click"})), 5).await + handle + .mcp_message("tools/call", Some(json!({"name": "katashiro.click"})), 5) + .await }); let request = recv(&mut second).await; assert_eq!( - request["method"], json!("mcp/message"), + request["method"], + json!("mcp/message"), "the next frame owed to the surviving connection is the tool call, not a disconnect" ); assert_eq!(request["params"]["connectionId"], json!("conn-new")); - send(&mut second, json!({ - "jsonrpc": "2.0", "id": request["id"].clone(), - "result": {"content": [{"type": "text", "text": "clicked"}]} - })).await; - let got = call.await.unwrap().expect("the surviving tunnel must carry a call"); + send( + &mut second, + json!({ + "jsonrpc": "2.0", "id": request["id"].clone(), + "result": {"content": [{"type": "text", "text": "clicked"}]} + }), + ) + .await; + let got = call + .await + .unwrap() + .expect("the surviving tunnel must carry a call"); assert_eq!(got["content"][0]["text"], json!("clicked")); } } @@ -6293,14 +7038,25 @@ mod acp_log_redaction { let channel_id = format!("acp_{uuid}"); let tag = redact_id(&channel_id); - assert_eq!(tag, redact_id(&channel_id), "same id must tag identically, or logs stop \ - being correlatable across lines"); - assert!(!tag.contains(&uuid.to_string()), "the uuid must not survive into the tag"); + assert_eq!( + tag, + redact_id(&channel_id), + "same id must tag identically, or logs stop \ + being correlatable across lines" + ); + assert!( + !tag.contains(&uuid.to_string()), + "the uuid must not survive into the tag" + ); assert!( !tag.contains(&channel_id) && !channel_id.contains(&tag[1..]), "the tag must not be a substring of the id or vice versa" ); - assert_ne!(tag, redact_id(&format!("acp_{}", Uuid::new_v4())), "different ids differ"); + assert_ne!( + tag, + redact_id(&format!("acp_{}", Uuid::new_v4())), + "different ids differ" + ); } /// No `info!`/`warn!`/`error!` in this file may carry a raw channel or session id. @@ -6343,7 +7099,10 @@ mod acp_log_redaction { let line = src[..start].matches('\n').count() + 1; offenders.push(format!( "{line}: {}", - body.split_whitespace().take(8).collect::>().join(" ") + body.split_whitespace() + .take(8) + .collect::>() + .join(" ") )); } } @@ -6391,7 +7150,10 @@ mod acp_teardown_ownership { { let mut r = reg.lock().unwrap(); // The successor: same channel, fresh server id, different connection. - r.insert(("acp_x".into(), "srv-new".into()), tunnel("conn-B", "cid-new")); + r.insert( + ("acp_x".into(), "srv-new".into()), + tunnel("conn-B", "cid-new"), + ); } teardown(®, &["acp_x"], "conn-A"); @@ -6414,8 +7176,14 @@ mod acp_teardown_ownership { teardown(®, &["acp_x"], "conn-A"); let r = reg.lock().unwrap(); - assert!(!r.contains_key(&("acp_x".to_string(), "srv-a".to_string())), "conn-A's own tunnel must go"); - assert!(r.contains_key(&("acp_x".to_string(), "srv-b".to_string())), "conn-B's must stay"); + assert!( + !r.contains_key(&("acp_x".to_string(), "srv-a".to_string())), + "conn-A's own tunnel must go" + ); + assert!( + r.contains_key(&("acp_x".to_string(), "srv-b".to_string())), + "conn-B's must stay" + ); assert_eq!(r.len(), 1); } @@ -6430,7 +7198,12 @@ mod acp_teardown_ownership { let (tx, _rx) = mpsc::unbounded_channel::(); reg.lock().unwrap().insert( "acp_x".into(), - ReplySink { turn_id: "evt_new".into(), tx, owner: "conn-B".into(), generation: 0 }, + ReplySink { + turn_id: "evt_new".into(), + tx, + owner: "conn-B".into(), + generation: 0, + }, ); } { @@ -6445,5 +7218,4 @@ mod acp_teardown_ownership { "conn-A's cleanup must leave conn-B's sink in place, or the live session goes mute" ); } - } diff --git a/crates/openab-gateway/src/adapters/feishu.rs b/crates/openab-gateway/src/adapters/feishu.rs index c3df29e99..135c28b15 100644 --- a/crates/openab-gateway/src/adapters/feishu.rs +++ b/crates/openab-gateway/src/adapters/feishu.rs @@ -453,7 +453,9 @@ mod event_types { return None; } - let content_json: serde_json::Value = msg.content.as_deref() + let content_json: serde_json::Value = msg + .content + .as_deref() .and_then(|s| serde_json::from_str(s).ok())?; let message_id = msg.message_id.as_deref()?; @@ -462,9 +464,8 @@ mod event_types { let (clean_text, mention_ids, media_refs) = match msg_type { "image" => { let image_key = content_json.get("image_key")?.as_str()?; - let mentions = extract_mentions( - "", msg.mentions.as_deref().unwrap_or(&[]), bot_open_id, - ); + let mentions = + extract_mentions("", msg.mentions.as_deref().unwrap_or(&[]), bot_open_id); let refs = vec![MediaRef::Image { message_id: message_id.to_string(), image_key: image_key.to_string(), @@ -473,12 +474,12 @@ mod event_types { } "file" => { let file_key = content_json.get("file_key")?.as_str()?; - let file_name = content_json.get("file_name") + let file_name = content_json + .get("file_name") .and_then(|v| v.as_str()) .unwrap_or("unknown"); - let mentions = extract_mentions( - "", msg.mentions.as_deref().unwrap_or(&[]), bot_open_id, - ); + let mentions = + extract_mentions("", msg.mentions.as_deref().unwrap_or(&[]), bot_open_id); let refs = vec![MediaRef::File { message_id: message_id.to_string(), file_key: file_key.to_string(), @@ -488,9 +489,8 @@ mod event_types { } "audio" => { let file_key = content_json.get("file_key")?.as_str()?; - let mentions = extract_mentions( - "", msg.mentions.as_deref().unwrap_or(&[]), bot_open_id, - ); + let mentions = + extract_mentions("", msg.mentions.as_deref().unwrap_or(&[]), bot_open_id); let refs = vec![MediaRef::Audio { message_id: message_id.to_string(), file_key: file_key.to_string(), @@ -512,7 +512,9 @@ mod event_types { } } Some("img") => { - if let Some(key) = el.get("image_key").and_then(|v| v.as_str()) { + if let Some(key) = + el.get("image_key").and_then(|v| v.as_str()) + { refs.push(MediaRef::Image { message_id: message_id.to_string(), image_key: key.to_string(), @@ -543,7 +545,10 @@ mod event_types { } _ => { // text - let raw_text = content_json.get("text").and_then(|v| v.as_str()).unwrap_or(""); + let raw_text = content_json + .get("text") + .and_then(|v| v.as_str()) + .unwrap_or(""); if raw_text.trim().is_empty() { return None; } @@ -768,7 +773,8 @@ impl FeishuTokenCache { let expire = body.get("expire").and_then(|v| v.as_u64()).unwrap_or(7200); - let token = body.get("tenant_access_token") + let token = body + .get("tenant_access_token") .and_then(|v| v.as_str()) .map(|s| s.to_string()) .ok_or_else(|| anyhow::anyhow!("feishu token refresh: missing tenant_access_token"))?; @@ -902,7 +908,10 @@ async fn get_ws_endpoint( let body: serde_json::Value = resp.json().await?; let code = body.get("code").and_then(|v| v.as_i64()).unwrap_or(-1); if code != 0 { - let msg = body.get("msg").and_then(|v| v.as_str()).unwrap_or("unknown"); + let msg = body + .get("msg") + .and_then(|v| v.as_str()) + .unwrap_or("unknown"); anyhow::bail!("feishu ws endpoint error: code={code} msg={msg}"); } body.get("data") @@ -1135,10 +1144,16 @@ async fn handle_ws_message( // - MultibotMentions: bypass only if participated AND no other bot in thread // - Mentions: never bypass let bypass_mention = detect_and_mark_multibot( - &envelope, bot_id_ref, config, participated_threads, multibot_threads, + &envelope, + bot_id_ref, + config, + participated_threads, + multibot_threads, ); - if let Some((mut gateway_event, media_refs)) = parse_message_event(&envelope, bot_id_ref, config, bypass_mention) { + if let Some((mut gateway_event, media_refs)) = + parse_message_event(&envelope, bot_id_ref, config, bypass_mention) + { // Also dedupe by message_id if dedupe.is_duplicate(&gateway_event.message_id) { return; @@ -1170,8 +1185,13 @@ async fn handle_ws_message( // Resolve sender display name (lazy, cached) let name = resolve_user_name( - &gateway_event.sender.id, name_cache, token_cache, client, &config.api_base(), - ).await; + &gateway_event.sender.id, + name_cache, + token_cache, + client, + &config.api_base(), + ) + .await; gateway_event.sender.name = name.clone(); gateway_event.sender.display_name = name; @@ -1181,14 +1201,29 @@ async fn handle_ws_message( let api_base = config.api_base(); for media_ref in &media_refs { let attachment = match media_ref { - MediaRef::Image { message_id, image_key } => { - download_feishu_image(client, &api_base, &token, message_id, image_key).await + MediaRef::Image { + message_id, + image_key, + } => { + download_feishu_image(client, &api_base, &token, message_id, image_key) + .await } - MediaRef::File { message_id, file_key, file_name } => { - download_feishu_file(client, &api_base, &token, message_id, file_key, file_name).await + MediaRef::File { + message_id, + file_key, + file_name, + } => { + download_feishu_file( + client, &api_base, &token, message_id, file_key, file_name, + ) + .await } - MediaRef::Audio { message_id, file_key } => { - download_feishu_audio(client, &api_base, &token, message_id, file_key).await + MediaRef::Audio { + message_id, + file_key, + } => { + download_feishu_audio(client, &api_base, &token, message_id, file_key) + .await } }; gateway_event.content.attachments.push(attachment); @@ -1197,7 +1232,9 @@ async fn handle_ws_message( } // Skip if no text and no attachments (e.g. unsupported file type) - if gateway_event.content.text.trim().is_empty() && gateway_event.content.attachments.is_empty() { + if gateway_event.content.text.trim().is_empty() + && gateway_event.content.attachments.is_empty() + { return; } @@ -1293,10 +1330,7 @@ fn is_feishu_cap_reached_body(body: &str) -> bool { .get("code") .and_then(|c| c.as_i64()) .is_some_and(|code| code == 230072), - Err(_) => { - body.contains("230072") - || body.contains("number of times it can be edited") - } + Err(_) => body.contains("230072") || body.contains("number of times it can be edited"), } } @@ -1321,10 +1355,7 @@ pub enum EditOutcome { /// over `EDIT_COUNTS_CACHE_MAX`, the oldest *insertions* are evicted (not the /// lowest-count entries) so active streams are not bumped out from under /// themselves. -fn increment_edit_count( - cache: &Arc>, - message_id: &str, -) { +fn increment_edit_count(cache: &Arc>, message_id: &str) { let mut c = cache.lock(); let was_new = !c.counts.contains_key(message_id); let entry = c.counts.entry(message_id.to_string()).or_insert(0); @@ -1340,10 +1371,7 @@ fn increment_edit_count( /// Mark a message_id as cap-reached; subsequent edit attempts skip the API /// call and signal `EditOutcome::CapReached` directly so the core finalize /// path can take over. -fn mark_edit_cap( - cache: &Arc>, - message_id: &str, -) { +fn mark_edit_cap(cache: &Arc>, message_id: &str) { let mut c = cache.lock(); let was_new = !c.counts.contains_key(message_id); c.counts.insert(message_id.to_string(), u32::MAX); @@ -1371,10 +1399,7 @@ fn evict_if_overcap(c: &mut EditCountsCache) { /// Return true if this message_id has already reached the edit cap (either /// tracked locally or marked via 230072 sentinel). -fn is_edit_cap_reached( - cache: &Arc>, - message_id: &str, -) -> bool { +fn is_edit_cap_reached(cache: &Arc>, message_id: &str) -> bool { let c = cache.lock(); c.counts .get(message_id) @@ -1387,11 +1412,7 @@ fn is_edit_cap_reached( /// and generic failure. Performs a preemptive local cap check (`FEISHU_EDIT_CAP`) /// before hitting the network, and detects the server-side errcode 230072 via /// body-code-first parsing if the local count drifts from reality. -async fn edit_feishu_message( - adapter: &FeishuAdapter, - message_id: &str, - text: &str, -) -> EditOutcome { +async fn edit_feishu_message(adapter: &FeishuAdapter, message_id: &str, text: &str) -> EditOutcome { // Pre-check: if we've already tracked >= FEISHU_EDIT_CAP edits (or the sentinel // u32::MAX from a 230072 response), skip the API call and signal CapReached so // the caller can stop attempting edits and let the core finalize path recover. @@ -1413,9 +1434,14 @@ async fn edit_feishu_message( "msg_type": "post", "content": post_content.to_string(), }); - match adapter.client.put(&url).bearer_auth(&token) + match adapter + .client + .put(&url) + .bearer_auth(&token) .header("Content-Type", "application/json; charset=utf-8") - .json(&body).send().await + .json(&body) + .send() + .await { Ok(resp) => { let status = resp.status(); @@ -1498,10 +1524,7 @@ async fn edit_feishu_message( /// `message_id` shape is validated by the caller (`handle_reply` dispatch seam, /// via `is_valid_feishu_message_id`) before this is reached, so it is safe to /// interpolate into the URL path here. -async fn delete_feishu_message( - adapter: &FeishuAdapter, - message_id: &str, -) -> Result<(), String> { +async fn delete_feishu_message(adapter: &FeishuAdapter, message_id: &str) -> Result<(), String> { let token = adapter .token_cache .get_token(&adapter.client) @@ -1509,13 +1532,7 @@ async fn delete_feishu_message( .map_err(|e| format!("token error: {e}"))?; let api_base = adapter.config.api_base(); let url = format!("{}/open-apis/im/v1/messages/{}", api_base, message_id); - match adapter - .client - .delete(&url) - .bearer_auth(&token) - .send() - .await - { + match adapter.client.delete(&url).bearer_auth(&token).send().await { Ok(resp) if resp.status().is_success() => { tracing::info!(message_id = %message_id, "feishu message deleted"); Ok(()) @@ -1582,9 +1599,7 @@ fn has_code_fence(text: &str) -> bool { fn has_markdown_table(text: &str) -> bool { text.split('\n').any(|line| { let t = line.trim(); - t.contains('|') - && t.contains('-') - && t.chars().all(|c| matches!(c, '|' | '-' | ':' | ' ')) + t.contains('|') && t.contains('-') && t.chars().all(|c| matches!(c, '|' | '-' | ':' | ' ')) }) } @@ -1783,9 +1798,19 @@ fn try_parse_link(chars: &[char], start: usize) -> Option<(String, String, usize /// Reference to a media resource that needs async download after parse_message_event. pub enum MediaRef { - Image { message_id: String, image_key: String }, - File { message_id: String, file_key: String, file_name: String }, - Audio { message_id: String, file_key: String }, + Image { + message_id: String, + image_key: String, + }, + File { + message_id: String, + file_key: String, + file_name: String, + }, + Audio { + message_id: String, + file_key: String, + }, } const IMAGE_MAX_DIMENSION_PX: u32 = 1200; @@ -1861,13 +1886,21 @@ pub async fn download_feishu_image( if let Some(cl) = resp.headers().get(reqwest::header::CONTENT_LENGTH) { if let Ok(size) = cl.to_str().unwrap_or("0").parse::() { if size > IMAGE_MAX_DOWNLOAD { - tracing::warn!(image_key, size, "feishu image Content-Length exceeds 10MB limit, skipping download"); + tracing::warn!( + image_key, + size, + "feishu image Content-Length exceeds 10MB limit, skipping download" + ); return crate::schema::Attachment::rejected( "image", format!("{}.jpg", image_key), "application/octet-stream", size, - format!("size exceeded: {} exceeds {}", format_bytes(size), format_bytes(IMAGE_MAX_DOWNLOAD)), + format!( + "size exceeded: {} exceeds {}", + format_bytes(size), + format_bytes(IMAGE_MAX_DOWNLOAD) + ), ); } } @@ -1887,13 +1920,21 @@ pub async fn download_feishu_image( }; // Fallback check (Content-Length may be absent or misreported) if bytes.len() as u64 > IMAGE_MAX_DOWNLOAD { - tracing::warn!(image_key, size = bytes.len(), "feishu image exceeds 10MB limit"); + tracing::warn!( + image_key, + size = bytes.len(), + "feishu image exceeds 10MB limit" + ); return crate::schema::Attachment::rejected( "image", format!("{}.jpg", image_key), "application/octet-stream", bytes.len() as u64, - format!("size exceeded: {} exceeds {}", format_bytes(bytes.len() as u64), format_bytes(IMAGE_MAX_DOWNLOAD)), + format!( + "size exceeded: {} exceeds {}", + format_bytes(bytes.len() as u64), + format_bytes(IMAGE_MAX_DOWNLOAD) + ), ); } let (compressed, mime) = match resize_and_compress(&bytes) { @@ -1946,9 +1987,9 @@ pub async fn download_feishu_file( // Only download text-like files let ext = file_name.rsplit('.').next().unwrap_or("").to_lowercase(); const TEXT_EXTS: &[&str] = &[ - "txt", "csv", "log", "md", "json", "jsonl", "yaml", "yml", "toml", "xml", - "rs", "py", "js", "ts", "jsx", "tsx", "go", "java", "c", "cpp", "h", "hpp", - "rb", "sh", "bash", "sql", "html", "css", "ini", "cfg", "conf", "env", + "txt", "csv", "log", "md", "json", "jsonl", "yaml", "yml", "toml", "xml", "rs", "py", "js", + "ts", "jsx", "tsx", "go", "java", "c", "cpp", "h", "hpp", "rb", "sh", "bash", "sql", + "html", "css", "ini", "cfg", "conf", "env", ]; if !TEXT_EXTS.contains(&ext.as_str()) { tracing::debug!(file_name, "skipping non-text file attachment"); @@ -1992,13 +2033,21 @@ pub async fn download_feishu_file( if let Some(cl) = resp.headers().get(reqwest::header::CONTENT_LENGTH) { if let Ok(size) = cl.to_str().unwrap_or("0").parse::() { if size > FILE_MAX_DOWNLOAD { - tracing::warn!(file_name, size, "feishu file Content-Length exceeds 512KB limit, skipping download"); + tracing::warn!( + file_name, + size, + "feishu file Content-Length exceeds 512KB limit, skipping download" + ); return crate::schema::Attachment::rejected( "text_file", file_name.to_string(), "application/octet-stream", size, - format!("size exceeded: {} exceeds {}", format_bytes(size), format_bytes(FILE_MAX_DOWNLOAD)), + format!( + "size exceeded: {} exceeds {}", + format_bytes(size), + format_bytes(FILE_MAX_DOWNLOAD) + ), ); } } @@ -2018,13 +2067,21 @@ pub async fn download_feishu_file( }; // Fallback check (Content-Length may be absent or misreported) if bytes.len() as u64 > FILE_MAX_DOWNLOAD { - tracing::warn!(file_name, size = bytes.len(), "feishu file exceeds 512KB limit"); + tracing::warn!( + file_name, + size = bytes.len(), + "feishu file exceeds 512KB limit" + ); return crate::schema::Attachment::rejected( "text_file", file_name.to_string(), "application/octet-stream", bytes.len() as u64, - format!("size exceeded: {} exceeds {}", format_bytes(bytes.len() as u64), format_bytes(FILE_MAX_DOWNLOAD)), + format!( + "size exceeded: {} exceeds {}", + format_bytes(bytes.len() as u64), + format_bytes(FILE_MAX_DOWNLOAD) + ), ); } let path = match crate::store::store_media(&bytes).await { @@ -2064,7 +2121,9 @@ pub async fn download_feishu_audio( use urlencoding::encode; let url = format!( "{}/open-apis/im/v1/messages/{}/resources/{}?type=file", - api_base, encode(message_id), encode(file_key) + api_base, + encode(message_id), + encode(file_key) ); let resp = match client.get(&url).bearer_auth(token).send().await { Ok(r) => r, @@ -2105,7 +2164,11 @@ pub async fn download_feishu_audio( format!("{}.ogg", file_key), "audio/ogg", size, - format!("size exceeded: {} exceeds {}", format_bytes(size), format_bytes(AUDIO_MAX_DOWNLOAD)), + format!( + "size exceeded: {} exceeds {}", + format_bytes(size), + format_bytes(AUDIO_MAX_DOWNLOAD) + ), ); } } @@ -2124,13 +2187,21 @@ pub async fn download_feishu_audio( } }; if bytes.len() as u64 > AUDIO_MAX_DOWNLOAD { - tracing::warn!(file_key, size = bytes.len(), "feishu audio exceeds 25MB limit"); + tracing::warn!( + file_key, + size = bytes.len(), + "feishu audio exceeds 25MB limit" + ); return crate::schema::Attachment::rejected( "audio", format!("{}.ogg", file_key), "audio/ogg", bytes.len() as u64, - format!("size exceeded: {} exceeds {}", format_bytes(bytes.len() as u64), format_bytes(AUDIO_MAX_DOWNLOAD)), + format!( + "size exceeded: {} exceeds {}", + format_bytes(bytes.len() as u64), + format_bytes(AUDIO_MAX_DOWNLOAD) + ), ); } tracing::debug!(file_key, size = bytes.len(), "feishu audio downloaded"); @@ -2180,7 +2251,10 @@ pub async fn send_post_message( ) } else { ( - format!("{}/open-apis/im/v1/messages?receive_id_type=chat_id", api_base), + format!( + "{}/open-apis/im/v1/messages?receive_id_type=chat_id", + api_base + ), serde_json::json!({ "receive_id": chat_id, "msg_type": "post", @@ -2315,13 +2389,18 @@ async fn add_reaction(adapter: &FeishuAdapter, message_id: &str, emoji: &str) { }; let token = match adapter.token_cache.get_token(&adapter.client).await { Ok(t) => t, - Err(e) => { tracing::error!(err = %e, "feishu: cannot get token for reaction"); return; } + Err(e) => { + tracing::error!(err = %e, "feishu: cannot get token for reaction"); + return; + } }; let url = format!( "{}/open-apis/im/v1/messages/{}/reactions", - adapter.config.api_base(), message_id + adapter.config.api_base(), + message_id ); - let _ = adapter.client + let _ = adapter + .client .post(&url) .bearer_auth(&token) .json(&serde_json::json!({"reaction_type": {"emoji_type": reaction_type}})) @@ -2337,15 +2416,26 @@ async fn remove_reaction(adapter: &FeishuAdapter, message_id: &str, emoji: &str) }; let token = match adapter.token_cache.get_token(&adapter.client).await { Ok(t) => t, - Err(e) => { tracing::error!(err = %e, "feishu: cannot get token for reaction"); return; } + Err(e) => { + tracing::error!(err = %e, "feishu: cannot get token for reaction"); + return; + } }; // Feishu remove reaction needs reaction_id. Simpler approach: delete by type. // GET reactions, find matching, DELETE by id. let list_url = format!( "{}/open-apis/im/v1/messages/{}/reactions?reaction_type={}", - adapter.config.api_base(), message_id, reaction_type + adapter.config.api_base(), + message_id, + reaction_type ); - let resp = match adapter.client.get(&list_url).bearer_auth(&token).send().await { + let resp = match adapter + .client + .get(&list_url) + .bearer_auth(&token) + .send() + .await + { Ok(r) => r, Err(_) => return, }; @@ -2357,15 +2447,24 @@ async fn remove_reaction(adapter: &FeishuAdapter, message_id: &str, emoji: &str) if let Some(items) = body.pointer("/data/items").and_then(|v| v.as_array()) { let bot_id = adapter.bot_open_id.read().await; for item in items { - let is_ours = item.pointer("/operator/operator_id/open_id") - .and_then(|v| v.as_str()) == bot_id.as_deref(); + let is_ours = item + .pointer("/operator/operator_id/open_id") + .and_then(|v| v.as_str()) + == bot_id.as_deref(); if is_ours { if let Some(reaction_id) = item.get("reaction_id").and_then(|v| v.as_str()) { let del_url = format!( "{}/open-apis/im/v1/messages/{}/reactions/{}", - adapter.config.api_base(), message_id, reaction_id + adapter.config.api_base(), + message_id, + reaction_id ); - let _ = adapter.client.delete(&del_url).bearer_auth(&token).send().await; + let _ = adapter + .client + .delete(&del_url) + .bearer_auth(&token) + .send() + .await; return; } } @@ -2394,7 +2493,8 @@ fn check_thread_participated( // Intentionally recover from poisoned mutex — cache data loss is acceptable // and preferable to panicking the gateway. let c = cache.lock(); - c.get(tid).is_some_and(|ts| ts.elapsed().as_secs() < session_ttl_secs) + c.get(tid) + .is_some_and(|ts| ts.elapsed().as_secs() < session_ttl_secs) }) .unwrap_or(false) } @@ -2417,9 +2517,8 @@ fn detect_and_mark_multibot( participated_threads: &Arc>>, multibot_threads: &Arc>>, ) -> bool { - let self_participated = check_thread_participated( - envelope, participated_threads, config.session_ttl_secs, - ); + let self_participated = + check_thread_participated(envelope, participated_threads, config.session_ttl_secs); let thread_id_for_check = envelope .event @@ -2438,14 +2537,15 @@ fn detect_and_mark_multibot( .and_then(|m| m.mentions.as_ref()); if let Some(mention_list) = mentions { let bot_self_id = bot_open_id.unwrap_or(""); - let mention_ids: Vec<_> = mention_list.iter().filter_map(|m| { - m.id.as_ref().and_then(|id| id.open_id.as_deref()) - }).collect(); + let mention_ids: Vec<_> = mention_list + .iter() + .filter_map(|m| m.id.as_ref().and_then(|id| id.open_id.as_deref())) + .collect(); let mentions_other_bot = if !config.trusted_bot_ids.is_empty() { - mention_ids.iter().any(|oid| { - config.trusted_bot_ids.iter().any(|bid| bid == oid) - }) + mention_ids + .iter() + .any(|oid| config.trusted_bot_ids.iter().any(|bid| bid == oid)) } else if !config.allowed_users.is_empty() { mention_ids.iter().any(|oid| { *oid != bot_self_id && !config.allowed_users.iter().any(|u| u == oid) @@ -2591,9 +2691,7 @@ pub async fn handle_reply( // behavior). The card_message_id comes from our own session state // (gateway-internal, trusted) so it bypasses the shape check. let card_msg = { - let mut reg = adapter - .stream_sessions - .lock(); + let mut reg = adapter.stream_sessions.lock(); reg.remove(&reply.reply_to).map(|s| s.card_message_id) }; let target = card_msg.as_deref().unwrap_or(&reply.reply_to); @@ -2647,7 +2745,9 @@ pub async fn handle_reply( } // quote_message_id (agent-controlled reply-to) takes priority over thread_id - let reply_target = reply.quote_message_id.as_deref() + let reply_target = reply + .quote_message_id + .as_deref() .or(reply.channel.thread_id.as_deref()); let thread_id = reply.channel.thread_id.as_deref(); @@ -2656,11 +2756,27 @@ pub async fn handle_reply( // Use post (rich text) format for markdown rendering. // When in a thread (thread_id present), use reply API to stay in the same thread. if text.len() <= limit { - let result = send_post_message(&adapter.client, &api_base, &token, &reply.channel.id, reply_target, text).await; + let result = send_post_message( + &adapter.client, + &api_base, + &token, + &reply.channel.id, + reply_target, + text, + ) + .await; // Fallback: if quote_message_id caused failure, retry without it let result = if result.is_none() && reply.quote_message_id.is_some() { tracing::warn!(quote_message_id = ?reply.quote_message_id, channel_id = %reply.channel.id, "reply-to failed, falling back to plain send"); - send_post_message(&adapter.client, &api_base, &token, &reply.channel.id, thread_id, text).await + send_post_message( + &adapter.client, + &api_base, + &token, + &reply.channel.id, + thread_id, + text, + ) + .await } else { result }; @@ -2669,7 +2785,11 @@ pub async fn handle_reply( adapter.dedupe.is_duplicate(&msg_id); // Record thread participation for mention bypass if let Some(tid) = thread_id { - record_participation(&adapter.participated_threads, tid, adapter.config.session_ttl_secs); + record_participation( + &adapter.participated_threads, + tid, + adapter.config.session_ttl_secs, + ); } // Send response with message_id back to OAB core (for streaming edit) if let Some(ref req_id) = reply.request_id { @@ -2714,7 +2834,16 @@ pub async fn handle_reply( let mut succeeded = 0usize; let mut last_msg_id: Option = None; for chunk in &chunks { - if let Some(msg_id) = send_post_message(&adapter.client, &api_base, &token, &reply.channel.id, reply_target, chunk).await { + if let Some(msg_id) = send_post_message( + &adapter.client, + &api_base, + &token, + &reply.channel.id, + reply_target, + chunk, + ) + .await + { adapter.dedupe.is_duplicate(&msg_id); succeeded += 1; last_msg_id = Some(msg_id); @@ -2724,7 +2853,16 @@ pub async fn handle_reply( if succeeded == 0 && reply.quote_message_id.is_some() { tracing::warn!(quote_message_id = ?reply.quote_message_id, channel_id = %reply.channel.id, "chunked reply-to failed, falling back to plain send"); for chunk in &chunks { - if let Some(msg_id) = send_post_message(&adapter.client, &api_base, &token, &reply.channel.id, thread_id, chunk).await { + if let Some(msg_id) = send_post_message( + &adapter.client, + &api_base, + &token, + &reply.channel.id, + thread_id, + chunk, + ) + .await + { adapter.dedupe.is_duplicate(&msg_id); succeeded += 1; last_msg_id = Some(msg_id); @@ -2733,7 +2871,11 @@ pub async fn handle_reply( } if succeeded > 0 { if let Some(tid) = thread_id { - record_participation(&adapter.participated_threads, tid, adapter.config.session_ttl_secs); + record_participation( + &adapter.participated_threads, + tid, + adapter.config.session_ttl_secs, + ); } } // Report back to core. Success requires every chunk delivered — partial @@ -2818,9 +2960,7 @@ async fn handle_card_edit( Active { card_id: String, seq: i64 }, } let existing = { - let mut reg = adapter - .stream_sessions - .lock(); + let mut reg = adapter.stream_sessions.lock(); match reg.get_mut(&om_post) { None => Existing::None, Some(s) if s.finalized => Existing::Finalized, @@ -2882,9 +3022,7 @@ async fn handle_card_edit( // the full reply as a fresh message (reusing #1122 recovery). feishu_card::CardOutcome::Failed { code, message } => { let card_msg = { - let mut reg = adapter - .stream_sessions - .lock(); + let mut reg = adapter.stream_sessions.lock(); reg.remove(&om_post).map(|s| s.card_message_id) }; if adapter.config.card_fallback_to_post { @@ -2963,15 +3101,17 @@ async fn promote_and_respond( }; // 1. Create the card entity seeded with the current full text. - let card_id = match feishu_card::create_streaming_card(&adapter.client, &api_base, &token, text) - .await - { - Ok(id) => id, - Err(outcome) => { - tracing::warn!(?outcome, "feishu card promote: create failed; post-edit fallback"); - return promote_fallback_post(reply, adapter, event_tx).await; - } - }; + let card_id = + match feishu_card::create_streaming_card(&adapter.client, &api_base, &token, text).await { + Ok(id) => id, + Err(outcome) => { + tracing::warn!( + ?outcome, + "feishu card promote: create failed; post-edit fallback" + ); + return promote_fallback_post(reply, adapter, event_tx).await; + } + }; // 2. Send the interactive card message into the same thread. let reply_target = reply @@ -3001,9 +3141,7 @@ async fn promote_and_respond( // 4. Register the session keyed by om_post (core stays oblivious). Seed // last_text with the current full text for the finalize rebuild. { - let mut reg = adapter - .stream_sessions - .lock(); + let mut reg = adapter.stream_sessions.lock(); reg.promote(&om_post, card_id, card_msg_id, text.to_string()); } tracing::info!(om_post = %om_post, "feishu reply promoted post → card"); @@ -3081,9 +3219,7 @@ async fn try_send_initial_card( // Session keyed by the card's own message_id — core edits this id directly, // so subsequent edit_message calls land on the Active path with no swap. { - let mut reg = adapter - .stream_sessions - .lock(); + let mut reg = adapter.stream_sessions.lock(); reg.promote(&card_msg_id, card_id, card_msg_id.clone(), text.to_string()); } // Mirror the post path's self-echo dedupe + thread participation tracking. @@ -3255,11 +3391,9 @@ fn verify_signature( fn decrypt_event(encrypt_key: &str, encrypted: &str) -> anyhow::Result { use sha2::{Digest, Sha256}; let key = Sha256::digest(encrypt_key.as_bytes()); - let cipher_bytes = base64::Engine::decode( - &base64::engine::general_purpose::STANDARD, - encrypted, - ) - .map_err(|e| anyhow::anyhow!("base64 decode failed: {e}"))?; + let cipher_bytes = + base64::Engine::decode(&base64::engine::general_purpose::STANDARD, encrypted) + .map_err(|e| anyhow::anyhow!("base64 decode failed: {e}"))?; if cipher_bytes.len() < 16 { anyhow::bail!("encrypted data too short"); @@ -3308,7 +3442,10 @@ pub async fn webhook( .and_then(|v| v.to_str().ok()) .unwrap_or("unknown"); if feishu.rate_limiter.check(ip) { - return (axum::http::StatusCode::TOO_MANY_REQUESTS, "rate limit exceeded") + return ( + axum::http::StatusCode::TOO_MANY_REQUESTS, + "rate limit exceeded", + ) .into_response(); } @@ -3432,16 +3569,25 @@ pub async fn webhook( // Check participated threads and multibot detection for mention bypass let bypass_mention = detect_and_mark_multibot( - &envelope, bot_id_ref, &feishu.config, - &feishu.participated_threads, &feishu.multibot_threads, + &envelope, + bot_id_ref, + &feishu.config, + &feishu.participated_threads, + &feishu.multibot_threads, ); - if let Some((mut gateway_event, media_refs)) = parse_message_event(&envelope, bot_id_ref, &feishu.config, bypass_mention) { + if let Some((mut gateway_event, media_refs)) = + parse_message_event(&envelope, bot_id_ref, &feishu.config, bypass_mention) + { if !feishu.dedupe.is_duplicate(&gateway_event.message_id) { let name = resolve_user_name( - &gateway_event.sender.id, &feishu.name_cache, &feishu.token_cache, - &feishu.client, &feishu.config.api_base(), - ).await; + &gateway_event.sender.id, + &feishu.name_cache, + &feishu.token_cache, + &feishu.client, + &feishu.config.api_base(), + ) + .await; gateway_event.sender.name = name.clone(); gateway_event.sender.display_name = name; @@ -3451,14 +3597,46 @@ pub async fn webhook( let api_base = feishu.config.api_base(); for media_ref in &media_refs { let attachment = match media_ref { - MediaRef::Image { message_id, image_key } => { - download_feishu_image(&feishu.client, &api_base, &token, message_id, image_key).await + MediaRef::Image { + message_id, + image_key, + } => { + download_feishu_image( + &feishu.client, + &api_base, + &token, + message_id, + image_key, + ) + .await } - MediaRef::File { message_id, file_key, file_name } => { - download_feishu_file(&feishu.client, &api_base, &token, message_id, file_key, file_name).await + MediaRef::File { + message_id, + file_key, + file_name, + } => { + download_feishu_file( + &feishu.client, + &api_base, + &token, + message_id, + file_key, + file_name, + ) + .await } - MediaRef::Audio { message_id, file_key } => { - download_feishu_audio(&feishu.client, &api_base, &token, message_id, file_key).await + MediaRef::Audio { + message_id, + file_key, + } => { + download_feishu_audio( + &feishu.client, + &api_base, + &token, + message_id, + file_key, + ) + .await } }; gateway_event.content.attachments.push(attachment); @@ -3467,7 +3645,9 @@ pub async fn webhook( } // Skip if no text and no attachments (e.g. unsupported file type) - if gateway_event.content.text.trim().is_empty() && gateway_event.content.attachments.is_empty() { + if gateway_event.content.text.trim().is_empty() + && gateway_event.content.attachments.is_empty() + { return axum::http::StatusCode::OK.into_response(); } @@ -3591,12 +3771,24 @@ mod tests { #[test] fn should_use_card_auto_promotes_on_code_fence() { - assert!(should_use_card("here:\n```\ncode\n```", StreamingMode::Auto, 4000)); - assert!(should_use_card("```python\nprint(1)\n```", StreamingMode::Auto, 4000)); + assert!(should_use_card( + "here:\n```\ncode\n```", + StreamingMode::Auto, + 4000 + )); + assert!(should_use_card( + "```python\nprint(1)\n```", + StreamingMode::Auto, + 4000 + )); // Indented fence (markdown_to_post trims leading whitespace too). assert!(should_use_card(" ```\ncode", StreamingMode::Auto, 4000)); // Inline single backticks are NOT a fence. - assert!(!should_use_card("use `cargo` to build", StreamingMode::Auto, 4000)); + assert!(!should_use_card( + "use `cargo` to build", + StreamingMode::Auto, + 4000 + )); } #[test] @@ -3612,13 +3804,29 @@ mod tests { #[test] fn should_use_card_auto_plain_short_text_stays_post() { - assert!(!should_use_card("just a short reply", StreamingMode::Auto, 4000)); + assert!(!should_use_card( + "just a short reply", + StreamingMode::Auto, + 4000 + )); // A line with a pipe but no delimiter row is not a table. - assert!(!should_use_card("a | b but not a table", StreamingMode::Auto, 4000)); + assert!(!should_use_card( + "a | b but not a table", + StreamingMode::Auto, + 4000 + )); // Bullet list with dashes is not a table (no pipe). - assert!(!should_use_card("- item one\n- item two", StreamingMode::Auto, 4000)); + assert!(!should_use_card( + "- item one\n- item two", + StreamingMode::Auto, + 4000 + )); // Thematic break (---) is not a table (no pipe). - assert!(!should_use_card("above\n\n---\n\nbelow", StreamingMode::Auto, 4000)); + assert!(!should_use_card( + "above\n\n---\n\nbelow", + StreamingMode::Auto, + 4000 + )); } // --- Token tests --- @@ -3866,7 +4074,7 @@ mod tests { fn parse_group_without_mention_filtered() { let env = make_envelope("group", "just chatting", "ou_user1", None); let cfg = test_config(); // require_mention = true - // Gateway-side mention gating: group message without bot mention is filtered + // Gateway-side mention gating: group message without bot mention is filtered assert!(parse_message_event(&env, Some("ou_bot"), &cfg, false).is_none()); } @@ -3882,7 +4090,13 @@ mod tests { #[test] fn parse_skips_bot_sender() { let mut env = make_envelope("p2p", "hello", "ou_bot", None); - env.event.as_mut().unwrap().sender.as_mut().unwrap().sender_type = Some("bot".into()); + env.event + .as_mut() + .unwrap() + .sender + .as_mut() + .unwrap() + .sender_type = Some("bot".into()); let cfg = test_config(); assert!(parse_message_event(&env, Some("ou_bot"), &cfg, false).is_none()); } @@ -3897,7 +4111,13 @@ mod tests { #[test] fn parse_skips_non_text_message() { let mut env = make_envelope("p2p", "hello", "ou_user1", None); - env.event.as_mut().unwrap().message.as_mut().unwrap().message_type = Some("sticker".into()); + env.event + .as_mut() + .unwrap() + .message + .as_mut() + .unwrap() + .message_type = Some("sticker".into()); let cfg = test_config(); assert!(parse_message_event(&env, Some("ou_bot"), &cfg, false).is_none()); } @@ -3991,11 +4211,25 @@ mod tests { let name_cache = Arc::new(parking_lot::Mutex::new(HashMap::new())); let client = reqwest::Client::new(); - let name = resolve_user_name("ou_user1", &name_cache, &token_cache, &client, &server.uri()).await; + let name = resolve_user_name( + "ou_user1", + &name_cache, + &token_cache, + &client, + &server.uri(), + ) + .await; assert_eq!(name, "Alice"); // Second call should use cache (expect(1) above ensures no second API call) - let name2 = resolve_user_name("ou_user1", &name_cache, &token_cache, &client, &server.uri()).await; + let name2 = resolve_user_name( + "ou_user1", + &name_cache, + &token_cache, + &client, + &server.uri(), + ) + .await; assert_eq!(name2, "Alice"); } @@ -4022,7 +4256,14 @@ mod tests { let name_cache = Arc::new(parking_lot::Mutex::new(HashMap::new())); let client = reqwest::Client::new(); - let name = resolve_user_name("ou_unknown", &name_cache, &token_cache, &client, &server.uri()).await; + let name = resolve_user_name( + "ou_unknown", + &name_cache, + &token_cache, + &client, + &server.uri(), + ) + .await; assert_eq!(name, "ou_unknown"); } @@ -4033,10 +4274,17 @@ mod tests { // If mention key appears in normal text too, only the first occurrence is removed let mentions = vec![FeishuMention { key: Some("@_user_1".into()), - id: Some(FeishuMentionId { open_id: Some("ou_bot".into()) }), + id: Some(FeishuMentionId { + open_id: Some("ou_bot".into()), + }), name: Some("Bot".into()), }]; - let env = make_envelope("group", "@_user_1 tell me about @_user_1 patterns", "ou_user1", Some(mentions)); + let env = make_envelope( + "group", + "@_user_1 tell me about @_user_1 patterns", + "ou_user1", + Some(mentions), + ); let cfg = test_config(); let (evt, _media) = parse_message_event(&env, Some("ou_bot"), &cfg, false).unwrap(); // Only first @_user_1 removed, second preserved @@ -4067,7 +4315,9 @@ mod tests { fn parse_allowed_groups_blocks_unlisted() { let mentions = vec![FeishuMention { key: Some("@_user_1".into()), - id: Some(FeishuMentionId { open_id: Some("ou_bot".into()) }), + id: Some(FeishuMentionId { + open_id: Some("ou_bot".into()), + }), name: Some("Bot".into()), }]; let env = make_envelope("group", "@_user_1 hello", "ou_user1", Some(mentions)); @@ -4080,7 +4330,9 @@ mod tests { fn parse_allowed_groups_permits_listed() { let mentions = vec![FeishuMention { key: Some("@_user_1".into()), - id: Some(FeishuMentionId { open_id: Some("ou_bot".into()) }), + id: Some(FeishuMentionId { + open_id: Some("ou_bot".into()), + }), name: Some("Bot".into()), }]; let env = make_envelope("group", "@_user_1 hello", "ou_user1", Some(mentions)); @@ -4141,7 +4393,13 @@ mod tests { #[test] fn parse_thread_id_from_root_id() { let mut env = make_envelope("p2p", "reply", "ou_user1", None); - env.event.as_mut().unwrap().message.as_mut().unwrap().root_id = Some("om_root".into()); + env.event + .as_mut() + .unwrap() + .message + .as_mut() + .unwrap() + .root_id = Some("om_root".into()); let cfg = test_config(); let (evt, _media) = parse_message_event(&env, Some("ou_bot"), &cfg, false).unwrap(); assert_eq!(evt.channel.thread_id, Some("om_root".into())); @@ -4150,7 +4408,13 @@ mod tests { #[test] fn parse_thread_id_from_parent_id() { let mut env = make_envelope("p2p", "reply", "ou_user1", None); - env.event.as_mut().unwrap().message.as_mut().unwrap().parent_id = Some("om_parent".into()); + env.event + .as_mut() + .unwrap() + .message + .as_mut() + .unwrap() + .parent_id = Some("om_parent".into()); let cfg = test_config(); let (evt, _media) = parse_message_event(&env, Some("ou_bot"), &cfg, false).unwrap(); assert_eq!(evt.channel.thread_id, Some("om_parent".into())); @@ -4175,9 +4439,15 @@ mod tests { #[test] fn participated_thread_bypasses_mention_gating() { let cfg = test_config(); // require_mention = true - // Build envelope with root_id (in a thread) + // Build envelope with root_id (in a thread) let mut env = make_envelope("group", "Hello", "ou_user1", None); - env.event.as_mut().unwrap().message.as_mut().unwrap().root_id = Some("root_123".into()); + env.event + .as_mut() + .unwrap() + .message + .as_mut() + .unwrap() + .root_id = Some("root_123".into()); // Without participation: no @mention → None assert!(parse_message_event(&env, Some("ou_bot"), &cfg, false).is_none()); // With participation: no @mention → Some (bypass) @@ -4190,7 +4460,7 @@ mod tests { #[test] fn participated_no_effect_without_thread() { let cfg = test_config(); // require_mention = true - // Message in main channel (no thread_id) — participated flag doesn't help + // Message in main channel (no thread_id) — participated flag doesn't help let env = make_envelope("group", "Hello", "ou_user1", None); assert!(parse_message_event(&env, Some("ou_bot"), &cfg, true).is_none()); } @@ -4216,7 +4486,13 @@ mod tests { let mut cfg = test_config(); cfg.allow_user_messages = AllowUsers::MultibotMentions; let mut env = make_envelope("group", "Hello", "ou_user1", None); - env.event.as_mut().unwrap().message.as_mut().unwrap().root_id = Some("root_456".into()); + env.event + .as_mut() + .unwrap() + .message + .as_mut() + .unwrap() + .root_id = Some("root_456".into()); // participated + no other bot → bypass_mention_gating=true let result = parse_message_event(&env, Some("ou_bot"), &cfg, true); assert!(result.is_some()); @@ -4227,7 +4503,13 @@ mod tests { let mut cfg = test_config(); cfg.allow_user_messages = AllowUsers::MultibotMentions; let mut env = make_envelope("group", "Hello", "ou_user1", None); - env.event.as_mut().unwrap().message.as_mut().unwrap().root_id = Some("root_456".into()); + env.event + .as_mut() + .unwrap() + .message + .as_mut() + .unwrap() + .root_id = Some("root_456".into()); // not participated → bypass_mention_gating=false assert!(parse_message_event(&env, Some("ou_bot"), &cfg, false).is_none()); } @@ -4237,7 +4519,13 @@ mod tests { let mut cfg = test_config(); cfg.allow_user_messages = AllowUsers::Mentions; let mut env = make_envelope("group", "Hello", "ou_user1", None); - env.event.as_mut().unwrap().message.as_mut().unwrap().root_id = Some("root_789".into()); + env.event + .as_mut() + .unwrap() + .message + .as_mut() + .unwrap() + .root_id = Some("root_789".into()); // Even with bypass_mention_gating=true, Mentions mode never bypasses // (caller would pass false because Mentions mode always returns false) assert!(parse_message_event(&env, Some("ou_bot"), &cfg, false).is_none()); @@ -4245,7 +4533,7 @@ mod tests { #[test] fn quote_message_id_takes_priority_over_thread_id() { - use crate::schema::{GatewayReply, ReplyChannel, Content}; + use crate::schema::{Content, GatewayReply, ReplyChannel}; let reply = GatewayReply { schema: "openab.gateway.reply.v1".into(), reply_to: "evt_123".into(), @@ -4264,14 +4552,16 @@ mod tests { quote_message_id: Some("om_specific".into()), }; // quote_message_id should take priority - let reply_target = reply.quote_message_id.as_deref() + let reply_target = reply + .quote_message_id + .as_deref() .or(reply.channel.thread_id.as_deref()); assert_eq!(reply_target, Some("om_specific")); } #[test] fn reply_target_falls_back_to_thread_id_when_no_quote() { - use crate::schema::{GatewayReply, ReplyChannel, Content}; + use crate::schema::{Content, GatewayReply, ReplyChannel}; let reply = GatewayReply { schema: "openab.gateway.reply.v1".into(), reply_to: "evt_123".into(), @@ -4289,14 +4579,16 @@ mod tests { request_id: None, quote_message_id: None, }; - let reply_target = reply.quote_message_id.as_deref() + let reply_target = reply + .quote_message_id + .as_deref() .or(reply.channel.thread_id.as_deref()); assert_eq!(reply_target, Some("om_root")); } #[test] fn reply_target_is_none_when_both_absent() { - use crate::schema::{GatewayReply, ReplyChannel, Content}; + use crate::schema::{Content, GatewayReply, ReplyChannel}; let reply = GatewayReply { schema: "openab.gateway.reply.v1".into(), reply_to: "evt_123".into(), @@ -4314,7 +4606,9 @@ mod tests { request_id: None, quote_message_id: None, }; - let reply_target = reply.quote_message_id.as_deref() + let reply_target = reply + .quote_message_id + .as_deref() .or(reply.channel.thread_id.as_deref()); assert_eq!(reply_target, None); } @@ -4506,7 +4800,9 @@ mod tests { #[test] fn message_id_validation_accepts_valid_shapes() { - assert!(is_valid_feishu_message_id("om_dc13264520392907fcq2e6kpngacls")); + assert!(is_valid_feishu_message_id( + "om_dc13264520392907fcq2e6kpngacls" + )); assert!(is_valid_feishu_message_id("om_abc123")); assert!(is_valid_feishu_message_id("om_A_B_c_1_2_3")); } @@ -4709,7 +5005,10 @@ mod tests { "report.pdf", ) .await; - assert!(att.status.is_some(), "non-text extension must have status set"); + assert!( + att.status.is_some(), + "non-text extension must have status set" + ); let reason = att.status.unwrap(); assert!(reason.contains("unsupported format"), "got: {reason}"); } @@ -4767,7 +5066,12 @@ mod tests { let adapter = FeishuAdapter::new(config); let (tx, mut rx) = tokio::sync::broadcast::channel(16); - handle_reply(&edit_reply("om_ph1", "short reply", None, Some("r1")), &adapter, &tx).await; + handle_reply( + &edit_reply("om_ph1", "short reply", None, Some("r1")), + &adapter, + &tx, + ) + .await; assert!( adapter.stream_sessions.lock().is_empty(), @@ -4921,7 +5225,10 @@ mod tests { ); let resp: crate::schema::GatewayResponse = serde_json::from_str(&rx.try_recv().unwrap()).unwrap(); - assert!(!resp.success, "failure must surface so core finalize takes over"); + assert!( + !resp.success, + "failure must surface so core finalize takes over" + ); } /// Card mode: the FIRST reply (command=None) goes straight to a card with no @@ -4974,7 +5281,9 @@ mod tests { // Session keyed by the card's own message_id (no post→card swap). { let reg = adapter.stream_sessions.lock(); - let s = reg.get("om_cardZ").expect("session keyed by card message id"); + let s = reg + .get("om_cardZ") + .expect("session keyed by card message id"); assert_eq!(s.card_id, "cardZ"); } let resp: crate::schema::GatewayResponse = diff --git a/crates/openab-gateway/src/adapters/feishu_card.rs b/crates/openab-gateway/src/adapters/feishu_card.rs index 9c5577644..f2ff5a39c 100644 --- a/crates/openab-gateway/src/adapters/feishu_card.rs +++ b/crates/openab-gateway/src/adapters/feishu_card.rs @@ -850,7 +850,9 @@ mod tests { .and(path("/open-apis/cardkit/v1/cards")) // Only assert the envelope; the escaped `data` string is covered by // the pure card_v2 test above. - .and(body_partial_json(serde_json::json!({ "type": "card_json" }))) + .and(body_partial_json( + serde_json::json!({ "type": "card_json" }), + )) .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ "code": 0, "msg": "success", @@ -998,8 +1000,7 @@ mod tests { .mount(&server) .await; let client = reqwest::Client::new(); - let out = - finish_card_stream(&client, &server.uri(), "tok", card_id, "final text", 9).await; + let out = finish_card_stream(&client, &server.uri(), "tok", card_id, "final text", 9).await; assert_eq!(out, CardOutcome::Updated); } @@ -1030,8 +1031,15 @@ mod tests { .mount(&server) .await; let client = reqwest::Client::new(); - let mid = - send_card_message(&client, &server.uri(), "tok", "oc_chat", Some(root), "cardabc").await; + let mid = send_card_message( + &client, + &server.uri(), + "tok", + "oc_chat", + Some(root), + "cardabc", + ) + .await; assert_eq!(mid.as_deref(), Some("om_sent789")); } @@ -1052,7 +1060,8 @@ mod tests { .mount(&server) .await; let client = reqwest::Client::new(); - let mid = send_card_message(&client, &server.uri(), "tok", "oc_chat", None, "cardabc").await; + let mid = + send_card_message(&client, &server.uri(), "tok", "oc_chat", None, "cardabc").await; assert_eq!(mid.as_deref(), Some("om_new456")); } @@ -1088,7 +1097,12 @@ mod tests { fn registry_promote_get_remove() { let mut reg = FeishuStreamRegistry::default(); assert!(!reg.contains("om_post1")); - reg.promote("om_post1", "card_a".into(), "om_card_msg".into(), "seed".into()); + reg.promote( + "om_post1", + "card_a".into(), + "om_card_msg".into(), + "seed".into(), + ); assert!(reg.contains("om_post1")); let s = reg.get("om_post1").expect("session present"); assert_eq!(s.card_id, "card_a"); @@ -1109,8 +1123,8 @@ mod tests { let mut reg = FeishuStreamRegistry::default(); reg.promote("om_post1", "card_a".into(), "m1".into(), "t1".into()); reg.get_mut("om_post1").unwrap().next_sequence(); // seq = 1 - // Re-promoting the same key replaces the session and keeps a single - // order entry (a fresh session, sequence reset). + // Re-promoting the same key replaces the session and keeps a single + // order entry (a fresh session, sequence reset). reg.promote("om_post1", "card_b".into(), "m2".into(), "t2".into()); assert_eq!(reg.len(), 1); let s = reg.get("om_post1").unwrap(); diff --git a/crates/openab-gateway/src/adapters/googlechat.rs b/crates/openab-gateway/src/adapters/googlechat.rs index 03866cdb0..0ddc9d026 100644 --- a/crates/openab-gateway/src/adapters/googlechat.rs +++ b/crates/openab-gateway/src/adapters/googlechat.rs @@ -34,8 +34,7 @@ const CHAT_API_REQUEST_TIMEOUT: std::time::Duration = std::time::Duration::from_ /// End-to-end budget for one outbound Google Chat reply. Core waits 35 seconds /// for an acknowledged gateway reply; finishing within 30 seconds leaves room /// for response serialization and WebSocket/broadcast scheduling. -const GOOGLE_CHAT_DELIVERY_DEADLINE: std::time::Duration = - std::time::Duration::from_secs(30); +const GOOGLE_CHAT_DELIVERY_DEADLINE: std::time::Duration = std::time::Duration::from_secs(30); /// Cap on text file attachments per message (matches Discord/Slack). const TEXT_FILE_COUNT_CAP: usize = 5; /// Cap on aggregate text file bytes per message (matches Discord/Slack 1 MB). @@ -300,9 +299,9 @@ fn normalize_adc_target_service_account(value: Option) -> Result {} @@ -763,10 +760,7 @@ pub async fn webhook( if let Some(ref adapter) = state.google_chat { if let Some(ref verifier) = adapter.jwt_verifier { - let auth_header = match headers - .get("authorization") - .and_then(|v| v.to_str().ok()) - { + let auth_header = match headers.get("authorization").and_then(|v| v.to_str().ok()) { Some(h) => h, None => { warn!("googlechat webhook: missing authorization header"); @@ -843,12 +837,7 @@ pub async fn webhook( let thread_id = msg.thread.as_ref().map(|t| t.name.clone()); - let message_id = msg - .name - .rsplit('/') - .next() - .unwrap_or(&msg.name) - .to_string(); + let message_id = msg.name.rsplit('/').next().unwrap_or(&msg.name).to_string(); // No attachments → emit event synchronously and respond 200 if media_refs.is_empty() { @@ -1106,7 +1095,9 @@ impl GoogleChatTokenCache { } async fn refresh(&self, client: &reqwest::Client) -> Result<(String, u64), String> { - let jwt = self.build_jwt().map_err(|e| format!("JWT build error: {e}"))?; + let jwt = self + .build_jwt() + .map_err(|e| format!("JWT build error: {e}"))?; let resp = client .post("https://oauth2.googleapis.com/token") .timeout(TOKEN_REQUEST_TIMEOUT) @@ -1163,8 +1154,7 @@ impl GoogleChatTokenCache { let key = jsonwebtoken::EncodingKey::from_rsa_pem(self.private_key.as_bytes()) .map_err(|e| format!("RSA key parse error: {e}"))?; let header = jsonwebtoken::Header::new(jsonwebtoken::Algorithm::RS256); - jsonwebtoken::encode(&header, &claims, &key) - .map_err(|e| format!("JWT encode error: {e}")) + jsonwebtoken::encode(&header, &claims, &key).map_err(|e| format!("JWT encode error: {e}")) } } @@ -1366,9 +1356,7 @@ impl MetadataTokenSource { if elapsed < ttl { *self.refresh_retry_after.write().await = Some( Instant::now() - + std::time::Duration::from_secs( - ADC_REFRESH_RETRY_COOLDOWN_SECS, - ), + + std::time::Duration::from_secs(ADC_REFRESH_RETRY_COOLDOWN_SECS), ); warn!( "googlechat ADC refresh failed ({e}); serving cached token \ @@ -1827,9 +1815,9 @@ fn split_text(text: &str, limit: usize) -> Vec<&str> { /// Whitelist of text-like file extensions for `download_googlechat_file`. const TEXT_EXTS: &[&str] = &[ - "txt", "csv", "log", "md", "json", "jsonl", "yaml", "yml", "toml", "xml", - "rs", "py", "js", "ts", "jsx", "tsx", "go", "java", "c", "cpp", "h", "hpp", - "rb", "sh", "bash", "sql", "html", "css", "ini", "cfg", "conf", + "txt", "csv", "log", "md", "json", "jsonl", "yaml", "yml", "toml", "xml", "rs", "py", "js", + "ts", "jsx", "tsx", "go", "java", "c", "cpp", "h", "hpp", "rb", "sh", "bash", "sql", "html", + "css", "ini", "cfg", "conf", ]; /// Parse Google Chat attachment array into media references for async download. @@ -1931,7 +1919,13 @@ pub async fn download_googlechat_image( content_name: &str, ) -> crate::schema::Attachment { let url = media_url(api_base, resource_name); - let resp = match client.get(&url).bearer_auth(token).timeout(MEDIA_REQUEST_TIMEOUT).send().await { + let resp = match client + .get(&url) + .bearer_auth(token) + .timeout(MEDIA_REQUEST_TIMEOUT) + .send() + .await + { Ok(r) => r, Err(e) => { warn!(content_name, error = %e, "googlechat image download failed"); @@ -1958,13 +1952,20 @@ pub async fn download_googlechat_image( if let Some(cl) = resp.headers().get(reqwest::header::CONTENT_LENGTH) { if let Ok(size) = cl.to_str().unwrap_or("0").parse::() { if size > IMAGE_MAX_DOWNLOAD { - warn!(content_name, size, "googlechat image Content-Length exceeds 10MB limit"); + warn!( + content_name, + size, "googlechat image Content-Length exceeds 10MB limit" + ); return crate::schema::Attachment::rejected( "image", content_name.to_string(), "image/jpeg", size, - format!("size exceeded: {} exceeds {}", format_bytes(size), format_bytes(IMAGE_MAX_DOWNLOAD)), + format!( + "size exceeded: {} exceeds {}", + format_bytes(size), + format_bytes(IMAGE_MAX_DOWNLOAD) + ), ); } } @@ -1983,13 +1984,21 @@ pub async fn download_googlechat_image( } }; if bytes.len() as u64 > IMAGE_MAX_DOWNLOAD { - warn!(content_name, size = bytes.len(), "googlechat image exceeds 10MB limit"); + warn!( + content_name, + size = bytes.len(), + "googlechat image exceeds 10MB limit" + ); return crate::schema::Attachment::rejected( "image", content_name.to_string(), "image/jpeg", bytes.len() as u64, - format!("size exceeded: {} exceeds {}", format_bytes(bytes.len() as u64), format_bytes(IMAGE_MAX_DOWNLOAD)), + format!( + "size exceeded: {} exceeds {}", + format_bytes(bytes.len() as u64), + format_bytes(IMAGE_MAX_DOWNLOAD) + ), ); } let (compressed, mime) = match resize_and_compress(&bytes) { @@ -2052,7 +2061,13 @@ pub async fn download_googlechat_file( } let max_size = FILE_MAX_DOWNLOAD.min(remaining_budget); let url = media_url(api_base, resource_name); - let resp = match client.get(&url).bearer_auth(token).timeout(MEDIA_REQUEST_TIMEOUT).send().await { + let resp = match client + .get(&url) + .bearer_auth(token) + .timeout(MEDIA_REQUEST_TIMEOUT) + .send() + .await + { Ok(r) => r, Err(e) => { warn!(content_name, error = %e, "googlechat file download failed"); @@ -2079,13 +2094,22 @@ pub async fn download_googlechat_file( if let Some(cl) = resp.headers().get(reqwest::header::CONTENT_LENGTH) { if let Ok(size) = cl.to_str().unwrap_or("0").parse::() { if size > max_size { - warn!(content_name, size, limit = max_size, "googlechat file Content-Length exceeds limit"); + warn!( + content_name, + size, + limit = max_size, + "googlechat file Content-Length exceeds limit" + ); return crate::schema::Attachment::rejected( "text_file", content_name.to_string(), "text/plain", size, - format!("size exceeded: {} exceeds {}", format_bytes(size), format_bytes(max_size)), + format!( + "size exceeded: {} exceeds {}", + format_bytes(size), + format_bytes(max_size) + ), ); } } @@ -2104,13 +2128,22 @@ pub async fn download_googlechat_file( } }; if bytes.len() as u64 > max_size { - warn!(content_name, size = bytes.len(), limit = max_size, "googlechat file exceeds size limit"); + warn!( + content_name, + size = bytes.len(), + limit = max_size, + "googlechat file exceeds size limit" + ); return crate::schema::Attachment::rejected( "text_file", content_name.to_string(), "text/plain", bytes.len() as u64, - format!("size exceeded: {} exceeds {}", format_bytes(bytes.len() as u64), format_bytes(max_size)), + format!( + "size exceeded: {} exceeds {}", + format_bytes(bytes.len() as u64), + format_bytes(max_size) + ), ); } let path = match crate::store::store_media(&bytes).await { @@ -2148,7 +2181,13 @@ pub async fn download_googlechat_audio( content_type: &str, ) -> crate::schema::Attachment { let url = media_url(api_base, resource_name); - let resp = match client.get(&url).bearer_auth(token).timeout(MEDIA_REQUEST_TIMEOUT).send().await { + let resp = match client + .get(&url) + .bearer_auth(token) + .timeout(MEDIA_REQUEST_TIMEOUT) + .send() + .await + { Ok(r) => r, Err(e) => { warn!(content_name, error = %e, "googlechat audio download failed"); @@ -2175,13 +2214,20 @@ pub async fn download_googlechat_audio( if let Some(cl) = resp.headers().get(reqwest::header::CONTENT_LENGTH) { if let Ok(size) = cl.to_str().unwrap_or("0").parse::() { if size > AUDIO_MAX_DOWNLOAD { - warn!(content_name, size, "googlechat audio Content-Length exceeds 25MB limit"); + warn!( + content_name, + size, "googlechat audio Content-Length exceeds 25MB limit" + ); return crate::schema::Attachment::rejected( "audio", content_name.to_string(), "audio/ogg", size, - format!("size exceeded: {} exceeds {}", format_bytes(size), format_bytes(AUDIO_MAX_DOWNLOAD)), + format!( + "size exceeded: {} exceeds {}", + format_bytes(size), + format_bytes(AUDIO_MAX_DOWNLOAD) + ), ); } } @@ -2200,13 +2246,21 @@ pub async fn download_googlechat_audio( } }; if bytes.len() as u64 > AUDIO_MAX_DOWNLOAD { - warn!(content_name, size = bytes.len(), "googlechat audio exceeds 25MB limit"); + warn!( + content_name, + size = bytes.len(), + "googlechat audio exceeds 25MB limit" + ); return crate::schema::Attachment::rejected( "audio", content_name.to_string(), "audio/ogg", bytes.len() as u64, - format!("size exceeded: {} exceeds {}", format_bytes(bytes.len() as u64), format_bytes(AUDIO_MAX_DOWNLOAD)), + format!( + "size exceeded: {} exceeds {}", + format_bytes(bytes.len() as u64), + format_bytes(AUDIO_MAX_DOWNLOAD) + ), ); } let path = match crate::store::store_media(&bytes).await { @@ -2311,7 +2365,10 @@ mod tests { let msg = payload.message.as_ref().unwrap(); assert_eq!(msg.argument_text.as_deref(), Some("hi")); assert_eq!(msg.thread.as_ref().unwrap().name, "spaces/SP/threads/t1"); - assert_eq!(payload.space.as_ref().unwrap().space_type.as_deref(), Some("ROOM")); + assert_eq!( + payload.space.as_ref().unwrap().space_type.as_deref(), + Some("ROOM") + ); } #[test] @@ -2546,10 +2603,7 @@ mod tests { server.uri(), ); - let token = src - .get_token() - .await - .expect("should mint a chat.bot token"); + let token = src.get_token().await.expect("should mint a chat.bot token"); assert_eq!(token, "chat-bot-tok"); } @@ -2578,11 +2632,7 @@ mod tests { .mount(&server) .await; - let src = MetadataTokenSource::with_bases( - same_sa.into(), - server.uri(), - server.uri(), - ); + let src = MetadataTokenSource::with_bases(same_sa.into(), server.uri(), server.uri()); let err = src .get_token() .await @@ -2636,12 +2686,10 @@ mod tests { // Start both production policy helpers from a builder carrying the // same explicit proxy. Metadata must clear it; IAM must retain it. let metadata_client = MetadataTokenSource::build_metadata_client( - reqwest::Client::builder() - .proxy(reqwest::Proxy::all(iam_proxy.uri()).unwrap()), + reqwest::Client::builder().proxy(reqwest::Proxy::all(iam_proxy.uri()).unwrap()), ); let iam_client = MetadataTokenSource::build_iam_client( - reqwest::Client::builder() - .proxy(reqwest::Proxy::all(iam_proxy.uri()).unwrap()), + reqwest::Client::builder().proxy(reqwest::Proxy::all(iam_proxy.uri()).unwrap()), ); let src = MetadataTokenSource::with_bases_and_clients( "chat-bot@project.iam.gserviceaccount.com".into(), @@ -2821,7 +2869,10 @@ mod tests { adc_target_service_account: Some("chat-bot@project.iam.gserviceaccount.com".into()), ..Default::default() }); - assert!(adapter.token_cache.is_none(), "unreadable key file → no cache"); + assert!( + adapter.token_cache.is_none(), + "unreadable key file → no cache" + ); assert!( adapter.metadata_source.is_some(), "use_adc=true → ADC source installed when the key file could not be read" @@ -2880,7 +2931,10 @@ mod tests { "api_not_enabled" ); // Anything else stays unclassified rather than guessing. - assert_eq!(classify_generate_access_token_error(500, "boom"), "unclassified"); + assert_eq!( + classify_generate_access_token_error(500, "boom"), + "unclassified" + ); } #[test] @@ -2990,9 +3044,13 @@ mod tests { assert!(!is_google_chat_message_name("spaces/")); assert!(!is_google_chat_message_name("spaces/SP")); assert!(!is_google_chat_message_name("spaces/SP/messages/")); - assert!(!is_google_chat_message_name("spaces/SP/messages/msg1/extra")); + assert!(!is_google_chat_message_name( + "spaces/SP/messages/msg1/extra" + )); assert!(!is_google_chat_message_name("spaces/SP/messages/msg1?x=1")); - assert!(!is_google_chat_message_name("spaces/SP/messages/msg1#fragment")); + assert!(!is_google_chat_message_name( + "spaces/SP/messages/msg1#fragment" + )); assert!(!is_google_chat_message_name("spaces/SP/messages/msg%2F1")); } @@ -3093,7 +3151,10 @@ mod tests { server.uri(), server.uri(), ); - let err = src.get_token().await.expect_err("blank token must be rejected"); + let err = src + .get_token() + .await + .expect_err("blank token must be rejected"); assert!( err.contains("missing or empty accessToken"), "unexpected error: {err}" @@ -3402,15 +3463,16 @@ mod tests { #[tokio::test] async fn handle_reply_sends_gateway_response_success() { - use wiremock::{Mock, MockServer, ResponseTemplate}; use wiremock::matchers::{method, path_regex}; + use wiremock::{Mock, MockServer, ResponseTemplate}; let mock_server = MockServer::start().await; Mock::given(method("POST")) .and(path_regex("/spaces/.*/messages")) - .respond_with(ResponseTemplate::new(200).set_body_json( - serde_json::json!({"name": "spaces/TEST/messages/msg_abc"}), - )) + .respond_with( + ResponseTemplate::new(200) + .set_body_json(serde_json::json!({"name": "spaces/TEST/messages/msg_abc"})), + ) .mount(&mock_server) .await; @@ -3448,8 +3510,8 @@ mod tests { #[tokio::test] async fn handle_reply_sends_failure_response_on_api_error() { - use wiremock::{Mock, MockServer, ResponseTemplate}; use wiremock::matchers::{method, path_regex}; + use wiremock::{Mock, MockServer, ResponseTemplate}; let mock_server = MockServer::start().await; Mock::given(method("POST")) @@ -3489,13 +3551,17 @@ mod tests { assert!(!resp.success); assert!(resp.message_id.is_none()); let err = resp.error.expect("error should be set on send failure"); - assert!(err.contains("500"), "error should include status code, got: {}", err); + assert!( + err.contains("500"), + "error should include status code, got: {}", + err + ); } #[tokio::test] async fn handle_reply_empty_message_short_circuits() { - use wiremock::{Mock, MockServer, ResponseTemplate}; use wiremock::matchers::{method, path_regex}; + use wiremock::{Mock, MockServer, ResponseTemplate}; let mock_server = MockServer::start().await; // Mount a mock that would fail the test if called @@ -3531,7 +3597,10 @@ mod tests { adapter.handle_reply(&reply, &event_tx).await; let received = event_rx.try_recv(); - assert!(received.is_ok(), "expected failure GatewayResponse for empty message"); + assert!( + received.is_ok(), + "expected failure GatewayResponse for empty message" + ); let resp: GatewayResponse = serde_json::from_str(&received.unwrap()).unwrap(); assert_eq!(resp.request_id, "req_empty"); assert!(!resp.success); @@ -3540,8 +3609,8 @@ mod tests { #[tokio::test] async fn handle_reply_multi_chunk_failure_includes_error() { - use wiremock::{Mock, MockServer, ResponseTemplate}; use wiremock::matchers::{method, path_regex}; + use wiremock::{Mock, MockServer, ResponseTemplate}; let mock_server = MockServer::start().await; Mock::given(method("POST")) @@ -3620,15 +3689,16 @@ mod tests { #[tokio::test] async fn handle_reply_edit_message_does_not_send_response() { - use wiremock::{Mock, MockServer, ResponseTemplate}; use wiremock::matchers::{method, path_regex}; + use wiremock::{Mock, MockServer, ResponseTemplate}; let mock_server = MockServer::start().await; Mock::given(method("PATCH")) .and(path_regex("/spaces/.*/messages/.*")) - .respond_with(ResponseTemplate::new(200).set_body_json( - serde_json::json!({"name": "spaces/SP/messages/msg1"}), - )) + .respond_with( + ResponseTemplate::new(200) + .set_body_json(serde_json::json!({"name": "spaces/SP/messages/msg1"})), + ) .expect(1) .mount(&mock_server) .await; @@ -3663,15 +3733,16 @@ mod tests { #[tokio::test] async fn handle_reply_multi_chunk_sends_gateway_response() { - use wiremock::{Mock, MockServer, ResponseTemplate}; use wiremock::matchers::{method, path_regex}; + use wiremock::{Mock, MockServer, ResponseTemplate}; let mock_server = MockServer::start().await; Mock::given(method("POST")) .and(path_regex("/spaces/.*/messages")) - .respond_with(ResponseTemplate::new(200).set_body_json( - serde_json::json!({"name": "spaces/TEST/messages/first_chunk"}), - )) + .respond_with( + ResponseTemplate::new(200) + .set_body_json(serde_json::json!({"name": "spaces/TEST/messages/first_chunk"})), + ) .mount(&mock_server) .await; @@ -3705,7 +3776,10 @@ mod tests { let resp: GatewayResponse = serde_json::from_str(&received.unwrap()).unwrap(); assert_eq!(resp.request_id, "req_multi"); assert!(resp.success); - assert_eq!(resp.message_id, Some("spaces/TEST/messages/first_chunk".into())); + assert_eq!( + resp.message_id, + Some("spaces/TEST/messages/first_chunk".into()) + ); } #[tokio::test] @@ -3714,16 +3788,17 @@ mod tests { // Result contract reports no overall message receipt, but the error // retains explicit partial-delivery context so operators know one chunk // already reached the space and should not blindly retry the whole turn. - use wiremock::{Mock, MockServer, ResponseTemplate}; use wiremock::matchers::{method, path_regex}; + use wiremock::{Mock, MockServer, ResponseTemplate}; let mock_server = MockServer::start().await; // First request: 200 OK with message name Mock::given(method("POST")) .and(path_regex("/spaces/.*/messages")) - .respond_with(ResponseTemplate::new(200).set_body_json( - serde_json::json!({"name": "spaces/TEST/messages/first_chunk"}), - )) + .respond_with( + ResponseTemplate::new(200) + .set_body_json(serde_json::json!({"name": "spaces/TEST/messages/first_chunk"})), + ) .up_to_n_times(1) .mount(&mock_server) .await; @@ -3764,7 +3839,10 @@ mod tests { let resp: GatewayResponse = serde_json::from_str(&received.unwrap()).unwrap(); assert_eq!(resp.request_id, "req_partial"); assert!(!resp.success, "partial failure must report success=false"); - assert!(resp.message_id.is_none(), "partial delivery is not an overall receipt"); + assert!( + resp.message_id.is_none(), + "partial delivery is not an overall receipt" + ); let err = resp.error.expect("partial failure should set error"); assert!(err.contains("chunk 2/2"), "{err}"); assert!(err.contains("after 1 successful chunk(s)"), "{err}"); @@ -3780,9 +3858,10 @@ mod tests { let server = MockServer::start().await; Mock::given(method("POST")) .and(path_regex("/spaces/.*/messages")) - .respond_with(ResponseTemplate::new(200).set_body_json( - serde_json::json!({"name": "spaces/TEST/messages/first_chunk"}), - )) + .respond_with( + ResponseTemplate::new(200) + .set_body_json(serde_json::json!({"name": "spaces/TEST/messages/first_chunk"})), + ) .up_to_n_times(1) .mount(&server) .await; @@ -3791,9 +3870,7 @@ mod tests { .respond_with( ResponseTemplate::new(200) .set_delay(std::time::Duration::from_millis(250)) - .set_body_json( - serde_json::json!({"name": "spaces/TEST/messages/late_chunk"}), - ), + .set_body_json(serde_json::json!({"name": "spaces/TEST/messages/late_chunk"})), ) .mount(&server) .await; @@ -3923,7 +4000,10 @@ mod tests { #[test] fn media_url_preserves_slashes_and_encodes_specials() { - let url = media_url("https://chat.googleapis.com/v1", "spaces/SP/messages/MSG/attachments/ATT"); + let url = media_url( + "https://chat.googleapis.com/v1", + "spaces/SP/messages/MSG/attachments/ATT", + ); assert_eq!( url, "https://chat.googleapis.com/v1/media/spaces/SP/messages/MSG/attachments/ATT?alt=media" @@ -3937,8 +4017,8 @@ mod tests { #[tokio::test] async fn download_googlechat_image_resizes_and_returns_attachment() { - use wiremock::{Mock, MockServer, ResponseTemplate}; use wiremock::matchers::{method, path_regex}; + use wiremock::{Mock, MockServer, ResponseTemplate}; // Generate a small valid PNG let img = image::RgbImage::from_pixel(10, 10, image::Rgb([255, 0, 0])); @@ -3989,22 +4069,23 @@ mod tests { ) .await; let att = result; - assert!(att.status.is_some(), "non-text extension must have status set"); + assert!( + att.status.is_some(), + "non-text extension must have status set" + ); let reason = att.status.unwrap(); assert!(reason.contains("unsupported format"), "got: {reason}"); } #[tokio::test] async fn download_googlechat_file_text_extension_succeeds() { - use wiremock::{Mock, MockServer, ResponseTemplate}; use wiremock::matchers::{method, path_regex}; + use wiremock::{Mock, MockServer, ResponseTemplate}; let mock_server = MockServer::start().await; Mock::given(method("GET")) .and(path_regex("/media/.*")) - .respond_with( - ResponseTemplate::new(200).set_body_bytes(b"hello world".to_vec()), - ) + .respond_with(ResponseTemplate::new(200).set_body_bytes(b"hello world".to_vec())) .mount(&mock_server) .await; @@ -4026,8 +4107,8 @@ mod tests { #[tokio::test] async fn download_googlechat_audio_returns_attachment() { - use wiremock::{Mock, MockServer, ResponseTemplate}; use wiremock::matchers::{method, path_regex}; + use wiremock::{Mock, MockServer, ResponseTemplate}; let mock_server = MockServer::start().await; let audio_bytes = vec![0u8; 1024]; @@ -4056,8 +4137,8 @@ mod tests { #[tokio::test] async fn download_googlechat_image_rejects_oversized_content_length() { - use wiremock::{Mock, MockServer, ResponseTemplate}; use wiremock::matchers::{method, path_regex}; + use wiremock::{Mock, MockServer, ResponseTemplate}; let mock_server = MockServer::start().await; Mock::given(method("GET")) diff --git a/crates/openab-gateway/src/adapters/lineworks.rs b/crates/openab-gateway/src/adapters/lineworks.rs index 036542c36..9cd53ab92 100644 --- a/crates/openab-gateway/src/adapters/lineworks.rs +++ b/crates/openab-gateway/src/adapters/lineworks.rs @@ -1251,8 +1251,7 @@ mod tests { tokio::spawn(async move { // Concurrent task refreshes the token (as a successful // get_token re-exchange would). - *cache.token.write().await = - Some(("tok_new".to_string(), Instant::now(), 3600)); + *cache.token.write().await = Some(("tok_new".to_string(), Instant::now(), 3600)); barrier.wait().await; }) }; diff --git a/crates/openab-gateway/src/adapters/mod.rs b/crates/openab-gateway/src/adapters/mod.rs index 105f51081..e35c2bcf1 100644 --- a/crates/openab-gateway/src/adapters/mod.rs +++ b/crates/openab-gateway/src/adapters/mod.rs @@ -1,23 +1,23 @@ -#[cfg(feature = "telegram")] -pub mod telegram; -#[cfg(feature = "line")] -pub mod line; +#[cfg(feature = "acp")] +#[allow(clippy::all, dead_code, unused)] +pub mod acp_schema; +#[cfg(feature = "acp")] +pub mod acp_server; #[cfg(feature = "feishu")] pub mod feishu; #[cfg(feature = "feishu")] pub mod feishu_card; #[cfg(feature = "googlechat")] pub mod googlechat; -#[cfg(feature = "wecom")] -pub mod wecom; -#[cfg(feature = "teams")] -pub mod teams; -#[cfg(feature = "acp")] -pub mod acp_server; -#[cfg(feature = "acp")] -#[allow(clippy::all, dead_code, unused)] -pub mod acp_schema; +#[cfg(feature = "line")] +pub mod line; #[cfg(feature = "lineworks")] pub mod lineworks; #[cfg(feature = "lineworks")] pub mod lineworks_flex; +#[cfg(feature = "teams")] +pub mod teams; +#[cfg(feature = "telegram")] +pub mod telegram; +#[cfg(feature = "wecom")] +pub mod wecom; diff --git a/crates/openab-gateway/src/adapters/teams.rs b/crates/openab-gateway/src/adapters/teams.rs index 0ee438d11..a097df9a0 100644 --- a/crates/openab-gateway/src/adapters/teams.rs +++ b/crates/openab-gateway/src/adapters/teams.rs @@ -282,7 +282,9 @@ impl TeamsAdapter { } // B2: Validate channel endorsements — key must endorse the activity's channelId - let channel_id = activity.channel_id.as_deref() + let channel_id = activity + .channel_id + .as_deref() .ok_or_else(|| anyhow::anyhow!("activity missing channelId"))?; if key.endorsements.is_empty() { anyhow::bail!("JWK has no endorsements — cannot verify channelId={channel_id}"); @@ -308,9 +310,13 @@ impl TeamsAdapter { let token_data = decode::(token, &decoding_key, &validation)?; // B1: Validate serviceUrl claim matches activity's serviceUrl - let activity_service_url = activity.service_url.as_deref() + let activity_service_url = activity + .service_url + .as_deref() .ok_or_else(|| anyhow::anyhow!("activity missing serviceUrl"))?; - let token_service_url = token_data.claims.get("serviceurl") + let token_service_url = token_data + .claims + .get("serviceurl") .and_then(|v| v.as_str()) .ok_or_else(|| anyhow::anyhow!("JWT missing serviceurl claim"))?; if token_service_url != activity_service_url { @@ -806,7 +812,9 @@ mod tests { async fn jwt_rejects_garbage_token() { let adapter = TeamsAdapter::new(make_config(vec![])); let activity = make_activity_with_tenant(Some("t1")); - let result = adapter.validate_jwt("Bearer not.a.valid.jwt", &activity).await; + let result = adapter + .validate_jwt("Bearer not.a.valid.jwt", &activity) + .await; assert!(result.is_err()); } diff --git a/crates/openab-gateway/src/adapters/telegram.rs b/crates/openab-gateway/src/adapters/telegram.rs index f4c381c1a..53ee66b6e 100644 --- a/crates/openab-gateway/src/adapters/telegram.rs +++ b/crates/openab-gateway/src/adapters/telegram.rs @@ -1,4 +1,7 @@ -use crate::media::{format_bytes, resize_and_compress, MediaKind, AUDIO_MAX_DOWNLOAD, FILE_MAX_DOWNLOAD, IMAGE_MAX_DOWNLOAD}; +use crate::media::{ + format_bytes, resize_and_compress, MediaKind, AUDIO_MAX_DOWNLOAD, FILE_MAX_DOWNLOAD, + IMAGE_MAX_DOWNLOAD, +}; use crate::schema::*; use crate::store; use axum::extract::State; @@ -106,12 +109,19 @@ pub async fn webhook( // Priority: CF-Connecting-IP (edge proxy, non-spoofable) > X-Real-IP > X-Forwarded-For // NOTE: XFF leftmost entry is client-supplied and spoofable. Phase 2 must implement // trusted proxy configuration for accurate IP extraction when enforcement is enabled. - let source_ip = headers.get("cf-connecting-ip") + let source_ip = headers + .get("cf-connecting-ip") .and_then(|v| v.to_str().ok()) .map(|s| s.trim().to_string()) - .or_else(|| headers.get("x-real-ip").and_then(|v| v.to_str().ok()).map(|s| s.trim().to_string())) .or_else(|| { - headers.get("x-forwarded-for") + headers + .get("x-real-ip") + .and_then(|v| v.to_str().ok()) + .map(|s| s.trim().to_string()) + }) + .or_else(|| { + headers + .get("x-forwarded-for") .and_then(|v| v.to_str().ok()) .and_then(|v| v.split(',').next()) .map(|s| s.trim().to_string()) @@ -159,19 +169,25 @@ pub async fn webhook( let client = &state.client; if is_photo { if let Some(largest) = msg.photo.iter().max_by_key(|p| p.width * p.height) { - let att = download_telegram_media(client, token, &largest.file_id, MediaKind::Image).await; + let att = + download_telegram_media(client, token, &largest.file_id, MediaKind::Image) + .await; attachments.push(att); } } else if let Some(doc) = msg.document { let file_name = doc.file_name.unwrap_or_else(|| "unknown.txt".to_string()); let mime_type = doc.mime_type.unwrap_or_else(|| "text/plain".to_string()); - let att = download_telegram_document(client, token, &doc.file_id, &file_name, &mime_type).await; + let att = + download_telegram_document(client, token, &doc.file_id, &file_name, &mime_type) + .await; attachments.push(att); } else if let Some(voice) = msg.voice { - let att = download_telegram_media(client, token, &voice.file_id, MediaKind::Audio).await; + let att = + download_telegram_media(client, token, &voice.file_id, MediaKind::Audio).await; attachments.push(att); } else if let Some(audio) = msg.audio { - let att = download_telegram_media(client, token, &audio.file_id, MediaKind::Audio).await; + let att = + download_telegram_media(client, token, &audio.file_id, MediaKind::Audio).await; attachments.push(att); } } @@ -277,7 +293,8 @@ fn check_telegram_subnet(ip_str: &str) -> bool { std::net::IpAddr::V4(v4) => { let octets = v4.octets(); // 149.154.160.0/20 → 149.154.160.0 - 149.154.175.255 - let in_range1 = octets[0] == 149 && octets[1] == 154 && (160..=175).contains(&octets[2]); + let in_range1 = + octets[0] == 149 && octets[1] == 154 && (160..=175).contains(&octets[2]); // 91.108.4.0/22 → 91.108.4.0 - 91.108.7.255 let in_range2 = octets[0] == 91 && octets[1] == 108 && (4..=7).contains(&octets[2]); in_range1 || in_range2 @@ -338,7 +355,11 @@ fn is_complex_markdown(text: &str) -> bool { /// Compute a stable draft_id from channel + thread to avoid collisions in forum topics. fn compute_draft_id(chat_id: &str, thread_id: &Option) -> i64 { let chan: i64 = chat_id.parse::().unwrap_or(1).abs(); - let tid: i64 = thread_id.as_deref().and_then(|t| t.parse::().ok()).unwrap_or(0).abs(); + let tid: i64 = thread_id + .as_deref() + .and_then(|t| t.parse::().ok()) + .unwrap_or(0) + .abs(); (chan.wrapping_add(tid)) % 1_000_000 + 1 } @@ -360,12 +381,20 @@ async fn send_rich_message( "message_thread_id": thread_id, "rich_message": { "markdown": text }, }); - let resp = client.post(&url).json(&body).send().await.map_err(|e| e.to_string())?; + let resp = client + .post(&url) + .json(&body) + .send() + .await + .map_err(|e| e.to_string())?; let json: serde_json::Value = resp.json().await.map_err(|e| e.to_string())?; if json["ok"].as_bool() == Some(true) { Ok(json) } else { - Err(json["description"].as_str().unwrap_or("unknown error").to_string()) + Err(json["description"] + .as_str() + .unwrap_or("unknown error") + .to_string()) } } @@ -391,12 +420,20 @@ async fn send_rich_message_draft( "draft_id": draft_id, "rich_message": if text.contains(" return, Err(e) => warn!("sendRichMessage failed ({e}), falling back to sendMessage"), } @@ -680,25 +733,48 @@ async fn download_telegram_media( let fallback_filename = format!("{}.{}", file_id, ext); let get_file_url = format!("{TELEGRAM_API_BASE}/bot{}/getFile", bot_token); - let resp = match client.get(&get_file_url).query(&[("file_id", file_id)]).send().await { + let resp = match client + .get(&get_file_url) + .query(&[("file_id", file_id)]) + .send() + .await + { Ok(r) => r, Err(e) => { warn!(file_id, error = %e, kind = ?kind, "Telegram getFile request failed"); - return Attachment::rejected(att_type, fallback_filename, default_mime, 0, "download failed: network error"); + return Attachment::rejected( + att_type, + fallback_filename, + default_mime, + 0, + "download failed: network error", + ); } }; let body: serde_json::Value = match resp.json().await { Ok(v) => v, Err(e) => { warn!(file_id, error = %e, kind = ?kind, "Telegram getFile JSON parse failed"); - return Attachment::rejected(att_type, fallback_filename, default_mime, 0, "download failed: invalid API response"); + return Attachment::rejected( + att_type, + fallback_filename, + default_mime, + 0, + "download failed: invalid API response", + ); } }; let file_path = match body["result"]["file_path"].as_str() { Some(p) => p, None => { warn!(file_id, kind = ?kind, "Telegram getFile response missing file_path"); - return Attachment::rejected(att_type, fallback_filename, default_mime, 0, "download failed: invalid API response"); + return Attachment::rejected( + att_type, + fallback_filename, + default_mime, + 0, + "download failed: invalid API response", + ); } }; @@ -707,7 +783,13 @@ async fn download_telegram_media( Ok(r) => r, Err(e) => { warn!(file_id, error = %e, kind = ?kind, "Telegram media download request failed"); - return Attachment::rejected(att_type, fallback_filename, default_mime, 0, "download failed: network error"); + return Attachment::rejected( + att_type, + fallback_filename, + default_mime, + 0, + "download failed: network error", + ); } }; @@ -732,7 +814,11 @@ async fn download_telegram_media( fallback_filename, default_mime, size, - format!("size exceeded: {} exceeds {}", format_bytes(size), format_bytes(max_size)), + format!( + "size exceeded: {} exceeds {}", + format_bytes(size), + format_bytes(max_size) + ), ); } } @@ -749,7 +835,13 @@ async fn download_telegram_media( Ok(b) => b, Err(e) => { warn!(file_id, error = %e, kind = ?kind, "Telegram media body read failed"); - return Attachment::rejected(att_type, fallback_filename, default_mime, 0, "download failed: body read error"); + return Attachment::rejected( + att_type, + fallback_filename, + default_mime, + 0, + "download failed: body read error", + ); } }; if bytes.len() as u64 > max_size { @@ -759,7 +851,11 @@ async fn download_telegram_media( fallback_filename, default_mime, bytes.len() as u64, - format!("size exceeded: {} exceeds {}", format_bytes(bytes.len() as u64), format_bytes(max_size)), + format!( + "size exceeded: {} exceeds {}", + format_bytes(bytes.len() as u64), + format_bytes(max_size) + ), ); } @@ -768,7 +864,13 @@ async fn download_telegram_media( Ok((c, m)) => (c, m), Err(e) => { error!(err = %e, "Telegram image processing failed"); - return Attachment::rejected(att_type, fallback_filename, default_mime, bytes.len() as u64, "processing failed: image encoding error"); + return Attachment::rejected( + att_type, + fallback_filename, + default_mime, + bytes.len() as u64, + "processing failed: image encoding error", + ); } }, MediaKind::Audio => (bytes.to_vec(), content_type), @@ -779,17 +881,27 @@ async fn download_telegram_media( Some(p) => p, None => { warn!(file_id, kind = ?kind, "Telegram media store failed"); - return Attachment::rejected(att_type, fallback_filename, default_mime, data_bytes.len() as u64, "processing failed: storage error"); + return Attachment::rejected( + att_type, + fallback_filename, + default_mime, + data_bytes.len() as u64, + "processing failed: storage error", + ); } }; info!(file_id, size = data_bytes.len(), kind = ?kind, "Telegram media stored"); Attachment { attachment_type: att_type.into(), - filename: format!("{}.{}", file_id, match kind { - MediaKind::Image => "jpg", - MediaKind::Audio => crate::media::audio_extension(&mime), - }), + filename: format!( + "{}.{}", + file_id, + match kind { + MediaKind::Image => "jpg", + MediaKind::Audio => crate::media::audio_extension(&mime), + } + ), mime_type: mime, data: String::new(), // No base64 — using file path size: data_bytes.len() as u64, @@ -819,25 +931,51 @@ async fn download_telegram_document( } let get_file_url = format!("{TELEGRAM_API_BASE}/bot{}/getFile", bot_token); - let resp = match client.get(&get_file_url).query(&[("file_id", file_id)]).send().await { + let resp = match client + .get(&get_file_url) + .query(&[("file_id", file_id)]) + .send() + .await + { Ok(r) => r, Err(e) => { warn!(file_id, error = %e, "Telegram document getFile request failed"); - return Attachment::rejected("text_file", file_name.to_string(), mime_type, 0, "download failed: network error"); + return Attachment::rejected( + "text_file", + file_name.to_string(), + mime_type, + 0, + "download failed: network error", + ); } }; let body: serde_json::Value = match resp.json().await { Ok(v) => v, Err(e) => { warn!(file_id, error = %e, "Telegram document getFile JSON parse failed"); - return Attachment::rejected("text_file", file_name.to_string(), mime_type, 0, "download failed: invalid API response"); + return Attachment::rejected( + "text_file", + file_name.to_string(), + mime_type, + 0, + "download failed: invalid API response", + ); } }; let file_path = match body["result"]["file_path"].as_str() { Some(p) => p, None => { - warn!(file_id, "Telegram document getFile response missing file_path"); - return Attachment::rejected("text_file", file_name.to_string(), mime_type, 0, "download failed: invalid API response"); + warn!( + file_id, + "Telegram document getFile response missing file_path" + ); + return Attachment::rejected( + "text_file", + file_name.to_string(), + mime_type, + 0, + "download failed: invalid API response", + ); } }; @@ -846,12 +984,22 @@ async fn download_telegram_document( Ok(r) => r, Err(e) => { warn!(file_id, err = %e, "Telegram document network error"); - return Attachment::rejected("text_file", file_name.to_string(), mime_type, 0, "download failed: network error"); + return Attachment::rejected( + "text_file", + file_name.to_string(), + mime_type, + 0, + "download failed: network error", + ); } }; if !resp.status().is_success() { let status = resp.status(); - warn!(file_id, status = status.as_u16(), "Telegram document HTTP error"); + warn!( + file_id, + status = status.as_u16(), + "Telegram document HTTP error" + ); return Attachment::rejected( "text_file", file_name.to_string(), @@ -864,13 +1012,20 @@ async fn download_telegram_document( if let Some(cl) = resp.headers().get(reqwest::header::CONTENT_LENGTH) { if let Ok(size) = cl.to_str().unwrap_or("0").parse::() { if size > FILE_MAX_DOWNLOAD { - warn!(file_id, size, "Telegram document Content-Length exceeds limit"); + warn!( + file_id, + size, "Telegram document Content-Length exceeds limit" + ); return Attachment::rejected( "text_file", file_name.to_string(), mime_type, size, - format!("size exceeded: {} exceeds {}", format_bytes(size), format_bytes(FILE_MAX_DOWNLOAD)), + format!( + "size exceeded: {} exceeds {}", + format_bytes(size), + format_bytes(FILE_MAX_DOWNLOAD) + ), ); } } @@ -880,23 +1035,40 @@ async fn download_telegram_document( Ok(b) => b, Err(e) => { warn!(file_id, error = %e, "Telegram document body read failed"); - return Attachment::rejected("text_file", file_name.to_string(), mime_type, 0, "download failed: body read error"); + return Attachment::rejected( + "text_file", + file_name.to_string(), + mime_type, + 0, + "download failed: body read error", + ); } }; if bytes.len() as u64 > FILE_MAX_DOWNLOAD { - warn!(file_id, size = bytes.len(), "Telegram document exceeds limit"); + warn!( + file_id, + size = bytes.len(), + "Telegram document exceeds limit" + ); return Attachment::rejected( "text_file", file_name.to_string(), mime_type, bytes.len() as u64, - format!("size exceeded: {} exceeds {}", format_bytes(bytes.len() as u64), format_bytes(FILE_MAX_DOWNLOAD)), + format!( + "size exceeded: {} exceeds {}", + format_bytes(bytes.len() as u64), + format_bytes(FILE_MAX_DOWNLOAD) + ), ); } // Validate UTF-8 — reject binary files if String::from_utf8(bytes.to_vec()).is_err() { - warn!(file_id, file_name, "Telegram document is not valid UTF-8, skipping"); + warn!( + file_id, + file_name, "Telegram document is not valid UTF-8, skipping" + ); return Attachment::rejected( "text_file", file_name.to_string(), @@ -910,10 +1082,21 @@ async fn download_telegram_document( Some(p) => p, None => { warn!(file_id, file_name, "Telegram document store failed"); - return Attachment::rejected("text_file", file_name.to_string(), mime_type, bytes.len() as u64, "processing failed: storage error"); + return Attachment::rejected( + "text_file", + file_name.to_string(), + mime_type, + bytes.len() as u64, + "processing failed: storage error", + ); } }; - info!(file_id, file_name, size = bytes.len(), "Telegram document stored"); + info!( + file_id, + file_name, + size = bytes.len(), + "Telegram document stored" + ); Attachment { attachment_type: "text_file".into(), @@ -941,17 +1124,29 @@ mod tests { "application/pdf", ) .await; - assert!(att.status.is_some(), "non-text extension must have status set"); + assert!( + att.status.is_some(), + "non-text extension must have status set" + ); let reason = att.status.unwrap(); assert!(reason.contains("unsupported format"), "got: {reason}"); - assert!(reason.contains("pdf"), "expected file extension in reason, got: {reason}"); + assert!( + reason.contains("pdf"), + "expected file extension in reason, got: {reason}" + ); } #[test] fn test_is_markdown_parse_error() { - assert!(is_markdown_parse_error("Bad Request: can't find end of italic entity at byte offset 37")); - assert!(is_markdown_parse_error("Bad Request: can't parse entities: Can't find end of bold entity")); - assert!(is_markdown_parse_error("can't parse entities in message text")); + assert!(is_markdown_parse_error( + "Bad Request: can't find end of italic entity at byte offset 37" + )); + assert!(is_markdown_parse_error( + "Bad Request: can't parse entities: Can't find end of bold entity" + )); + assert!(is_markdown_parse_error( + "can't parse entities in message text" + )); assert!(!is_markdown_parse_error("Unauthorized")); assert!(!is_markdown_parse_error("Bad Request: chat not found")); } @@ -982,8 +1177,12 @@ mod tests { fn test_is_complex_markdown() { // Tables assert!(is_complex_markdown("| Col1 | Col2 |\n|---|---|\n| a | b |")); - assert!(is_complex_markdown("| Col1 | Col2 |\n| :--- | ---: |\n| a | b |")); - assert!(is_complex_markdown("| A | B |\n| :---: | :---: |\n| x | y |")); + assert!(is_complex_markdown( + "| Col1 | Col2 |\n| :--- | ---: |\n| a | b |" + )); + assert!(is_complex_markdown( + "| A | B |\n| :---: | :---: |\n| x | y |" + )); // Code blocks — intentionally NOT complex (preserves syntax highlighting on legacy path) assert!(!is_complex_markdown("```rust\nfn main() {}\n```")); assert!(!is_complex_markdown("~~~\ncode\n~~~")); diff --git a/crates/openab-gateway/src/adapters/wecom.rs b/crates/openab-gateway/src/adapters/wecom.rs index 97aa84551..9bac4ab82 100644 --- a/crates/openab-gateway/src/adapters/wecom.rs +++ b/crates/openab-gateway/src/adapters/wecom.rs @@ -1,6 +1,6 @@ +use crate::media::format_bytes; use anyhow::Result; use axum::extract::State; -use crate::media::format_bytes; use std::sync::Arc; use tokio::sync::RwLock; use tracing::{info, warn}; @@ -47,7 +47,10 @@ impl WecomConfig { .unwrap_or(3); if encoding_aes_key.len() != 43 { - warn!("WECOM_ENCODING_AES_KEY must be 43 characters, got {}", encoding_aes_key.len()); + warn!( + "WECOM_ENCODING_AES_KEY must be 43 characters, got {}", + encoding_aes_key.len() + ); return None; } @@ -138,7 +141,10 @@ fn decrypt_message( .map_err(|e| anyhow::anyhow!("base64 decode failed: {e}"))?; if cipher_bytes.is_empty() || !cipher_bytes.len().is_multiple_of(16) { - anyhow::bail!("ciphertext length {} not a multiple of 16", cipher_bytes.len()); + anyhow::bail!( + "ciphertext length {} not a multiple of 16", + cipher_bytes.len() + ); } type Aes256CbcDec = cbc::Decryptor; @@ -153,12 +159,17 @@ fn decrypt_message( .map_err(|e| anyhow::anyhow!("aes decrypt failed: {e}"))?; // Strip WeCom PKCS7 padding (block_size=32): last byte indicates pad length (1-32) - let pad_byte = *plaintext.last().ok_or_else(|| anyhow::anyhow!("empty plaintext"))? as usize; + let pad_byte = *plaintext + .last() + .ok_or_else(|| anyhow::anyhow!("empty plaintext"))? as usize; if pad_byte == 0 || pad_byte > 32 || pad_byte > plaintext.len() { anyhow::bail!("invalid wecom padding value: {pad_byte}"); } let pad_start = plaintext.len() - pad_byte; - if !plaintext[pad_start..].iter().all(|&b| b as usize == pad_byte) { + if !plaintext[pad_start..] + .iter() + .all(|&b| b as usize == pad_byte) + { anyhow::bail!("invalid PKCS#7 padding: not all padding bytes match"); } let plaintext = &plaintext[..pad_start]; @@ -342,7 +353,6 @@ impl WecomAdapter { } } - pub async fn handle_reply( &self, reply: &crate::schema::GatewayReply, @@ -375,7 +385,10 @@ impl WecomAdapter { .unwrap_or(&reply.channel.id); let has_pending = { - let pending = self.pending_streams.lock().unwrap_or_else(|e| e.into_inner()); + let pending = self + .pending_streams + .lock() + .unwrap_or_else(|e| e.into_inner()); pending.contains_key(&reply.channel.id) }; let is_streaming_placeholder = reply.request_id.is_some() && !has_pending; @@ -398,10 +411,16 @@ impl WecomAdapter { let (text_tx, text_rx) = tokio::sync::watch::channel(String::new()); { - let mut pending = self.pending_streams.lock().unwrap_or_else(|e| e.into_inner()); - pending.insert(reply.channel.id.clone(), PendingStream { - text_watch: text_tx, - }); + let mut pending = self + .pending_streams + .lock() + .unwrap_or_else(|e| e.into_inner()); + pending.insert( + reply.channel.id.clone(), + PendingStream { + text_watch: text_tx, + }, + ); } let client = self.client.clone(); let token_cache = self.token_cache.clone(); @@ -453,9 +472,16 @@ impl WecomAdapter { return; } flush_thinking( - &client, &token_cache, &corp_id, &secret, &agent_id, - thinking_id.as_deref(), &flush_to_user, &last_text, - ).await; + &client, + &token_cache, + &corp_id, + &secret, + &agent_id, + thinking_id.as_deref(), + &flush_to_user, + &last_text, + ) + .await; }); if let Some(ref req_id) = reply.request_id { @@ -479,7 +505,10 @@ impl WecomAdapter { // between our earlier read of `has_pending` and now. If it did, // fall through to the direct-send path so the chunk isn't lost. let appended = { - let pending = self.pending_streams.lock().unwrap_or_else(|e| e.into_inner()); + let pending = self + .pending_streams + .lock() + .unwrap_or_else(|e| e.into_inner()); if let Some(stream) = pending.get(&reply.channel.id) { let current = stream.text_watch.borrow().clone(); let combined = if current.is_empty() { @@ -548,15 +577,21 @@ impl WecomAdapter { if text.is_empty() { return; } - let pending = self.pending_streams.lock().unwrap_or_else(|e| e.into_inner()); + let pending = self + .pending_streams + .lock() + .unwrap_or_else(|e| e.into_inner()); if let Some(stream) = pending.get(&reply.channel.id) { let _ = stream.text_watch.send(text.to_string()); } } - async fn send_text(&self, to_user: &str, text: &str) -> Result { - let agent_id: u64 = self.config.agent_id.parse().expect("agent_id validated at startup"); + let agent_id: u64 = self + .config + .agent_id + .parse() + .expect("agent_id validated at startup"); let body = serde_json::json!({ "touser": to_user, "msgtype": "text", @@ -590,16 +625,30 @@ async fn post_with_token_retry( body: &serde_json::Value, ) -> Result { let token = token_cache.get_token(client, corp_id, secret).await?; - let url = format!("{}{}?access_token={}", token_cache.base_url, api_path, token); + let url = format!( + "{}{}?access_token={}", + token_cache.base_url, api_path, token + ); let resp: serde_json::Value = client.post(&url).json(body).send().await?.json().await?; let errcode = resp["errcode"].as_i64().unwrap_or(-1); if errcode == 42001 { - warn!(api_path, "wecom: access_token expired, refreshing and retrying"); + warn!( + api_path, + "wecom: access_token expired, refreshing and retrying" + ); let new_token = token_cache.force_refresh(client, corp_id, secret).await?; - let retry_url = format!("{}{}?access_token={}", token_cache.base_url, api_path, new_token); - let retry_resp: serde_json::Value = - client.post(&retry_url).json(body).send().await?.json().await?; + let retry_url = format!( + "{}{}?access_token={}", + token_cache.base_url, api_path, new_token + ); + let retry_resp: serde_json::Value = client + .post(&retry_url) + .json(body) + .send() + .await? + .json() + .await?; let retry_code = retry_resp["errcode"].as_i64().unwrap_or(-1); if retry_code != 0 { anyhow::bail!( @@ -780,7 +829,11 @@ async fn flush_thinking( to_user: &str, text: &str, ) { - info!(?thinking_msg_id, text_len = text.len(), "wecom: flush_thinking starting"); + info!( + ?thinking_msg_id, + text_len = text.len(), + "wecom: flush_thinking starting" + ); // Recall thinking placeholder (only when streaming was enabled) if let Some(id) = thinking_msg_id { @@ -947,7 +1000,10 @@ pub async fn webhook( if let Ok(ts) = timestamp.parse::() { let now = chrono::Utc::now().timestamp(); if (now - ts).abs() > 300 { - warn!(timestamp_age_secs = now - ts, "wecom webhook: rejecting stale callback"); + warn!( + timestamp_age_secs = now - ts, + "wecom webhook: rejecting stale callback" + ); return axum::http::StatusCode::FORBIDDEN.into_response(); } } @@ -988,7 +1044,10 @@ pub async fn webhook( return axum::http::StatusCode::FORBIDDEN.into_response(); } - info!(encrypt_len = envelope.encrypt.len(), "wecom: decrypting callback"); + info!( + encrypt_len = envelope.encrypt.len(), + "wecom: decrypting callback" + ); let decrypted = match decrypt_message( &wecom.config.encoding_aes_key, &envelope.encrypt, @@ -999,7 +1058,10 @@ pub async fn webhook( d } Err(e) => { - warn!(encrypt_len = envelope.encrypt.len(), "wecom decrypt failed: {e}"); + warn!( + encrypt_len = envelope.encrypt.len(), + "wecom decrypt failed: {e}" + ); return "success".into_response(); } }; @@ -1076,7 +1138,12 @@ pub async fn webhook( ); event.content.attachments = attachments; - let att_sizes: Vec = event.content.attachments.iter().map(|a| a.data.len()).collect(); + let att_sizes: Vec = event + .content + .attachments + .iter() + .map(|a| a.data.len()) + .collect(); info!( attachments = event.content.attachments.len(), text_len = event.content.text.len(), @@ -1151,7 +1218,11 @@ async fn download_wecom_image( "wecom_image.jpg", "image/jpeg", size, - format!("size exceeded: {} exceeds {}", format_bytes(size), format_bytes(IMAGE_MAX_DOWNLOAD)), + format!( + "size exceeded: {} exceeds {}", + format_bytes(size), + format_bytes(IMAGE_MAX_DOWNLOAD) + ), ); } } @@ -1176,7 +1247,11 @@ async fn download_wecom_image( "wecom_image.jpg", "image/jpeg", bytes.len() as u64, - format!("size exceeded: {} exceeds {}", format_bytes(bytes.len() as u64), format_bytes(IMAGE_MAX_DOWNLOAD)), + format!( + "size exceeded: {} exceeds {}", + format_bytes(bytes.len() as u64), + format_bytes(IMAGE_MAX_DOWNLOAD) + ), ); } let (compressed, mime) = match resize_and_compress(&bytes) { @@ -1222,13 +1297,21 @@ const FILE_MAX_DOWNLOAD: u64 = 20 * 1024 * 1024; const TEXT_EXTENSIONS: &[&str] = &[ "txt", "csv", "log", "md", "json", "jsonl", "yaml", "yml", "toml", "xml", "rs", "py", "js", "ts", "jsx", "tsx", "go", "java", "c", "cpp", "h", "hpp", "rb", "sh", "bash", "zsh", "fish", - "ps1", "bat", "sql", "html", "css", "scss", "less", "ini", "cfg", "conf", "env", - "swift", "kt", "scala", "r", "pl", "lua", "graphql", "tsv", + "ps1", "bat", "sql", "html", "css", "scss", "less", "ini", "cfg", "conf", "env", "swift", "kt", + "scala", "r", "pl", "lua", "graphql", "tsv", ]; const TEXT_FILENAMES: &[&str] = &[ - "dockerfile", "makefile", "justfile", "rakefile", "gemfile", - "procfile", "vagrantfile", ".gitignore", ".dockerignore", ".editorconfig", + "dockerfile", + "makefile", + "justfile", + "rakefile", + "gemfile", + "procfile", + "vagrantfile", + ".gitignore", + ".dockerignore", + ".editorconfig", ]; fn is_text_file(filename: &str) -> bool { @@ -1326,7 +1409,11 @@ async fn download_wecom_file( filename.to_string(), "application/octet-stream", size, - format!("size exceeded: {} exceeds {}", format_bytes(size), format_bytes(FILE_MAX_DOWNLOAD)), + format!( + "size exceeded: {} exceeds {}", + format_bytes(size), + format_bytes(FILE_MAX_DOWNLOAD) + ), ); } } @@ -1351,7 +1438,11 @@ async fn download_wecom_file( filename.to_string(), "application/octet-stream", bytes.len() as u64, - format!("size exceeded: {} exceeds {}", format_bytes(bytes.len() as u64), format_bytes(FILE_MAX_DOWNLOAD)), + format!( + "size exceeded: {} exceeds {}", + format_bytes(bytes.len() as u64), + format_bytes(FILE_MAX_DOWNLOAD) + ), ); } @@ -1451,7 +1542,10 @@ mod tests { ("WECOM_CORP_ID", "ww_test_corp"), ("WECOM_SECRET", "test_secret"), ("WECOM_TOKEN", "test_token"), - ("WECOM_ENCODING_AES_KEY", "abcdefghijklmnopqrstuvwxyz0123456789ABCDEFG"), + ( + "WECOM_ENCODING_AES_KEY", + "abcdefghijklmnopqrstuvwxyz0123456789ABCDEFG", + ), ("WECOM_AGENT_ID", "1000002"), ]); let config = WecomConfig::from_reader(env).unwrap(); @@ -1613,11 +1707,17 @@ mod tests { let cache = WecomTokenCache::with_base_url(server.uri()); let client = reqwest::Client::new(); - let token = cache.get_token(&client, "ww_test_corp", "test_secret").await.unwrap(); + let token = cache + .get_token(&client, "ww_test_corp", "test_secret") + .await + .unwrap(); assert_eq!(token, "test_token_abc"); // Second call uses cache (mock expects exactly 1 call) - let token2 = cache.get_token(&client, "ww_test_corp", "test_secret").await.unwrap(); + let token2 = cache + .get_token(&client, "ww_test_corp", "test_secret") + .await + .unwrap(); assert_eq!(token2, "test_token_abc"); } diff --git a/crates/openab-gateway/src/lib.rs b/crates/openab-gateway/src/lib.rs index 587cc15fe..523207740 100644 --- a/crates/openab-gateway/src/lib.rs +++ b/crates/openab-gateway/src/lib.rs @@ -97,7 +97,6 @@ pub struct AppState { pub client: reqwest::Client, } - impl AppState { /// Create a minimal AppState for testing. Only requires an `event_tx` sender; /// all adapter fields default to `None`/empty. This decouples adapter tests @@ -141,9 +140,11 @@ impl AppState { event_tx, reply_token_cache: Arc::new(std::sync::Mutex::new(HashMap::new())), line_webhook_semaphore: Arc::new(Semaphore::new(LINE_WEBHOOK_CONCURRENCY_MAX)), - lineworks_webhook_semaphore: Arc::new(Semaphore::new(LINEWORKS_WEBHOOK_CONCURRENCY_MAX)), - lineworks_ingress_queue: Arc::new(Semaphore::new(LINEWORKS_INGRESS_QUEUE_MAX)), - trust_probe: None, + lineworks_webhook_semaphore: Arc::new(Semaphore::new( + LINEWORKS_WEBHOOK_CONCURRENCY_MAX, + )), + lineworks_ingress_queue: Arc::new(Semaphore::new(LINEWORKS_INGRESS_QUEUE_MAX)), + trust_probe: None, client: reqwest::Client::new(), } } @@ -184,8 +185,8 @@ impl AppState { // Feishu #[cfg(feature = "feishu")] - let feishu = adapters::feishu::FeishuConfig::from_env() - .map(adapters::feishu::FeishuAdapter::new); + let feishu = + adapters::feishu::FeishuConfig::from_env().map(adapters::feishu::FeishuAdapter::new); // Google Chat #[cfg(feature = "googlechat")] @@ -206,16 +207,20 @@ impl AppState { // WeCom #[cfg(feature = "wecom")] - let wecom = adapters::wecom::WecomConfig::from_env() - .map(adapters::wecom::WecomAdapter::new); + let wecom = + adapters::wecom::WecomConfig::from_env().map(adapters::wecom::WecomAdapter::new); // ACP Server #[cfg(feature = "acp")] let acp = adapters::acp_server::AcpConfig::from_env(); #[cfg(feature = "acp")] - let acp_reply_registry = acp.as_ref().map(|_| adapters::acp_server::new_reply_registry()); + let acp_reply_registry = acp + .as_ref() + .map(|_| adapters::acp_server::new_reply_registry()); #[cfg(feature = "acp")] - let acp_tunnel_registry = acp.as_ref().map(|_| adapters::acp_server::new_tunnel_registry()); + let acp_tunnel_registry = acp + .as_ref() + .map(|_| adapters::acp_server::new_tunnel_registry()); // LINE WORKS #[cfg(feature = "lineworks")] let lineworks = adapters::lineworks::LineWorksConfig::from_env().map(|config| { @@ -260,9 +265,11 @@ impl AppState { event_tx, reply_token_cache: Arc::new(std::sync::Mutex::new(HashMap::new())), line_webhook_semaphore: Arc::new(Semaphore::new(LINE_WEBHOOK_CONCURRENCY_MAX)), - lineworks_webhook_semaphore: Arc::new(Semaphore::new(LINEWORKS_WEBHOOK_CONCURRENCY_MAX)), - lineworks_ingress_queue: Arc::new(Semaphore::new(LINEWORKS_INGRESS_QUEUE_MAX)), - trust_probe: None, + lineworks_webhook_semaphore: Arc::new(Semaphore::new( + LINEWORKS_WEBHOOK_CONCURRENCY_MAX, + )), + lineworks_ingress_queue: Arc::new(Semaphore::new(LINEWORKS_INGRESS_QUEUE_MAX)), + trust_probe: None, client, } } @@ -422,7 +429,12 @@ impl AppState { /// no adapter, matching env-only semantics. #[cfg(feature = "wecom")] pub fn apply_wecom_config(&mut self, cfg: GatewayWecomConfig) { - let streaming = if cfg.streaming_enabled { "true" } else { "false" }.to_string(); + let streaming = if cfg.streaming_enabled { + "true" + } else { + "false" + } + .to_string(); let debounce = cfg.debounce_secs.to_string(); self.wecom = adapters::wecom::WecomConfig::from_reader(|k| match k { "WECOM_CORP_ID" => cfg.corp_id.clone(), @@ -601,10 +613,16 @@ impl Default for ServeConfig { /// Start the standalone gateway server. This is the main entry point extracted /// from the gateway binary — the binary becomes a thin wrapper around this. pub async fn serve(config: ServeConfig) -> anyhow::Result<()> { - use axum::{routing::{get, post}, Router}; + use axum::{ + routing::{get, post}, + Router, + }; use tracing::{info, warn}; - let ServeConfig { listen_addr, ws_token } = config; + let ServeConfig { + listen_addr, + ws_token, + } = config; if ws_token.is_none() { warn!("GATEWAY_WS_TOKEN not set — WebSocket connections are NOT authenticated (insecure)"); @@ -620,7 +638,10 @@ pub async fn serve(config: ServeConfig) -> anyhow::Result<()> { // ACP Server adapter. Fail-open (no transport key) is only allowed on a loopback // bind; a non-loopback bind without OPENAB_ACP_AUTH_KEY refuses to mount /acp. #[cfg(feature = "acp")] - if std::env::var("OPENAB_ACP_ENABLED").map(|v| v == "true" || v == "1").unwrap_or(false) { + if std::env::var("OPENAB_ACP_ENABLED") + .map(|v| v == "true" || v == "1") + .unwrap_or(false) + { let acp_key = std::env::var("OPENAB_ACP_AUTH_KEY").ok(); match adapters::acp_server::acp_auth_ok_for_bind(acp_key.as_deref(), &listen_addr) { Ok(()) => { @@ -732,7 +753,11 @@ pub async fn serve(config: ServeConfig) -> anyhow::Result<()> { let api_base = f.config.api_base(); let idle_ms = f.config.card_idle_finalize_ms; tokio::spawn(adapters::feishu::run_idle_reaper( - sessions, token_cache, client, api_base, idle_ms, + sessions, + token_cache, + client, + api_base, + idle_ms, )); info!(idle_ms, "feishu card-streaming idle reaper started"); } @@ -740,8 +765,8 @@ pub async fn serve(config: ServeConfig) -> anyhow::Result<()> { // Google Chat adapter #[cfg(feature = "googlechat")] - let googlechat_webhook_path = std::env::var("GOOGLE_CHAT_WEBHOOK_PATH") - .unwrap_or_else(|_| "/webhook/googlechat".into()); + let googlechat_webhook_path = + std::env::var("GOOGLE_CHAT_WEBHOOK_PATH").unwrap_or_else(|_| "/webhook/googlechat".into()); #[cfg(feature = "googlechat")] let google_chat = { let enabled = std::env::var("GOOGLE_CHAT_ENABLED") @@ -749,7 +774,10 @@ pub async fn serve(config: ServeConfig) -> anyhow::Result<()> { .unwrap_or(false); if enabled { info!(path = %googlechat_webhook_path, "googlechat adapter enabled"); - app = app.route(&googlechat_webhook_path, post(adapters::googlechat::webhook)); + app = app.route( + &googlechat_webhook_path, + post(adapters::googlechat::webhook), + ); Some(adapters::googlechat::GoogleChatAdapter::from_parts( adapters::googlechat::GoogleChatParts::from_env(), )) @@ -879,7 +907,11 @@ pub async fn serve(config: ServeConfig) -> anyhow::Result<()> { cache.retain(|_, (_, t)| t.elapsed().as_secs() < REPLY_TOKEN_TTL_SECS); let after = cache.len(); if before != after { - info!(removed = before - after, remaining = after, "reply token cache sweep"); + info!( + removed = before - after, + remaining = after, + "reply token cache sweep" + ); } } }); @@ -896,7 +928,11 @@ pub async fn serve(config: ServeConfig) -> anyhow::Result<()> { urls.retain(|_, (_, t)| t.elapsed().as_secs() < 4 * 3600); let after = urls.len(); if before != after { - info!(removed = before - after, remaining = after, "teams service_url cache cleanup"); + info!( + removed = before - after, + remaining = after, + "teams service_url cache cleanup" + ); } } }); @@ -1097,9 +1133,7 @@ async fn handle_oab_connection(state: Arc, socket: axum::extract::ws:: "lineworks" => { if let Some(ref lineworks) = state_for_recv.lineworks { let ok = adapters::lineworks::dispatch_lineworks_reply( - &client, - lineworks, - &reply, + &client, lineworks, &reply, ) .await; if !ok { diff --git a/crates/openab-gateway/src/store.rs b/crates/openab-gateway/src/store.rs index b08e69903..8927a8c5a 100644 --- a/crates/openab-gateway/src/store.rs +++ b/crates/openab-gateway/src/store.rs @@ -54,7 +54,11 @@ const MAX_STORE_SIZE: usize = 20 * 1024 * 1024; // 20 MB (matches AUDIO_MAX_DOWN /// Rejects files exceeding MAX_STORE_SIZE as a defense-in-depth measure. pub async fn store_media(bytes: &[u8]) -> Option { if bytes.len() > MAX_STORE_SIZE { - error!(size = bytes.len(), max = MAX_STORE_SIZE, "store_media rejected: exceeds size limit"); + error!( + size = bytes.len(), + max = MAX_STORE_SIZE, + "store_media rejected: exceeds size limit" + ); return None; } let dir = media_dir().await; diff --git a/crates/openab-mcp/src/mcp/runtime.rs b/crates/openab-mcp/src/mcp/runtime.rs index 0251287d3..25ba71294 100644 --- a/crates/openab-mcp/src/mcp/runtime.rs +++ b/crates/openab-mcp/src/mcp/runtime.rs @@ -2927,7 +2927,10 @@ mod tests { custom_headers, ) .await; - assert!(result.is_err(), "redirect response must not complete MCP POST"); + assert!( + result.is_err(), + "redirect response must not complete MCP POST" + ); assert_eq!( *original_headers.lock().await, ( @@ -3069,7 +3072,10 @@ mod tests { assert!(err.contains(header_name), "got: {err}"); assert!(!err.contains("secret-must-not-leak"), "value leaked: {err}"); assert!(!err.contains("circuit-breaker open"), "got: {err}"); - assert!(matches!(&mgr.statuses().await[0].1, ServerStatus::Failed(_))); + assert!(matches!( + &mgr.statuses().await[0].1, + ServerStatus::Failed(_) + )); } } } @@ -3096,7 +3102,10 @@ mod tests { !err.contains("custom-secret-must-not-leak"), "header value leaked: {err}" ); - assert!(matches!(&mgr.statuses().await[0].1, ServerStatus::Failed(_))); + assert!(matches!( + &mgr.statuses().await[0].1, + ServerStatus::Failed(_) + )); } #[tokio::test] diff --git a/src/acp_tunnel_source.rs b/src/acp_tunnel_source.rs index 25440f6bb..efd0bb4d0 100644 --- a/src/acp_tunnel_source.rs +++ b/src/acp_tunnel_source.rs @@ -484,8 +484,8 @@ impl openab_core::acp_mcp::SessionTokenRegistrar for FacadeRegistrar { mod tests { use super::{AcpTunnelSource, CapabilitySource, SessionCtx, Tool}; use openab_core::acp_mcp::AcpMcpTunnel; - use std::collections::HashSet; use serde_json::{json, Map, Value}; + use std::collections::HashSet; use std::sync::Arc; /// Tunnel double: reports declared servers, records forwarded `tools/call`s, @@ -518,17 +518,16 @@ mod tests { /// Make `tools/list` answer with these tool names. fn set_tools_list(&self, names: &[&str]) { - *self.tools_list.lock().unwrap() = - Some(names.iter().map(|n| n.to_string()).collect()); + *self.tools_list.lock().unwrap() = Some(names.iter().map(|n| n.to_string()).collect()); } /// Make a specific server's `tools/list` return exactly `names`, verbatim — the tool names /// need not be prefixed with the server's name. fn set_server_tools(&self, server_name: &str, names: &[&str]) { - self.server_tools - .lock() - .unwrap() - .insert(server_name.to_string(), names.iter().map(|n| n.to_string()).collect()); + self.server_tools.lock().unwrap().insert( + server_name.to_string(), + names.iter().map(|n| n.to_string()).collect(), + ); } /// Make `tools/list` fail (no answer configured). @@ -615,9 +614,9 @@ mod tests { shared .into_iter() .filter(|n| match &server_name { - Some(name) => n - .split_once('.') - .is_some_and(|(prefix, _)| prefix == name), + Some(name) => { + n.split_once('.').is_some_and(|(prefix, _)| prefix == name) + } None => true, }) .collect() @@ -717,7 +716,10 @@ mod tests { .call(Some(&ctx()), "notes.anything", &Map::new()) .await .unwrap(); - assert!(!is_err, "no allowlist gate: a connected server's declared tool dispatches"); + assert!( + !is_err, + "no allowlist gate: a connected server's declared tool dispatches" + ); let fwd = tunnel.forwarded.lock().unwrap(); assert_eq!(fwd[0].1, "uuid-n"); assert_eq!(fwd[0].2["name"], "notes.anything"); @@ -819,18 +821,36 @@ mod tests { // `attached_server_names`, but a name already in the cache keeps its tools even while it is // momentarily not attached — the UNION in `tools()` is what preserves this. let tunnel = FakeTunnel::with(&[("katashiro", "uuid-abc")]); - tunnel.set_tools_list(&["katashiro.click", "katashiro.navigate", "katashiro.read_dom", "katashiro.screenshot", "katashiro.type"]); + tunnel.set_tools_list(&[ + "katashiro.click", + "katashiro.navigate", + "katashiro.read_dom", + "katashiro.screenshot", + "katashiro.type", + ]); let src = AcpTunnelSource::new(tunnel.clone()); let _ = src.tools(Some(&ctx())); settle().await; - let before: Vec = - src.tools(Some(&ctx())).iter().map(|t| t.name.to_string()).collect(); - assert!(!before.is_empty(), "precondition: discovery must have populated the catalog"); + let before: Vec = src + .tools(Some(&ctx())) + .iter() + .map(|t| t.name.to_string()) + .collect(); + assert!( + !before.is_empty(), + "precondition: discovery must have populated the catalog" + ); tunnel.detach("katashiro"); - let after: Vec = - src.tools(Some(&ctx())).iter().map(|t| t.name.to_string()).collect(); - assert_eq!(after, before, "a detached tunnel must not shrink the catalog"); + let after: Vec = src + .tools(Some(&ctx())) + .iter() + .map(|t| t.name.to_string()) + .collect(); + assert_eq!( + after, before, + "a detached tunnel must not shrink the catalog" + ); } #[tokio::test] @@ -838,13 +858,22 @@ mod tests { // Once a catalog has been discovered, a LATER failed fetch must not empty it — the // invariant that protects a running deployment through a flap. let tunnel = FakeTunnel::with(&[("katashiro", "uuid-abc")]); - tunnel.set_tools_list(&["katashiro.click", "katashiro.navigate", "katashiro.read_dom", "katashiro.screenshot", "katashiro.type"]); + tunnel.set_tools_list(&[ + "katashiro.click", + "katashiro.navigate", + "katashiro.read_dom", + "katashiro.screenshot", + "katashiro.type", + ]); let src = AcpTunnelSource::new(tunnel.clone()); let _ = src.tools(Some(&ctx())); settle().await; let discovered = src.tools(Some(&ctx())).len(); - assert!(discovered > 0, "precondition: discovery must have populated the catalog"); + assert!( + discovered > 0, + "precondition: discovery must have populated the catalog" + ); tunnel.fail_tools_list(); let _ = src.tools(Some(&ctx())); @@ -967,10 +996,16 @@ mod tests { .call(Some(&ctx()), "notes.click", &Map::new()) .await .unwrap(); - assert!(!err, "prefix selects the tunnel; the server decides what it implements"); + assert!( + !err, + "prefix selects the tunnel; the server decides what it implements" + ); let fwd = tunnel.forwarded.lock().unwrap(); assert_eq!(fwd.len(), 1); - assert_eq!(fwd[0].1, "uuid-n", "routed to notes' tunnel, not katashiro's"); + assert_eq!( + fwd[0].1, "uuid-n", + "routed to notes' tunnel, not katashiro's" + ); assert_eq!(fwd[0].2["name"], "notes.click"); } @@ -1006,15 +1041,27 @@ mod tests { let _ = src.tools(Some(&ctx())); settle().await; - let names: Vec = - src.tools(Some(&ctx())).iter().map(|t| t.name.to_string()).collect(); + let names: Vec = src + .tools(Some(&ctx())) + .iter() + .map(|t| t.name.to_string()) + .collect(); assert_eq!(names, ["build"], "the bare tool is advertised"); let (_v, is_err) = src.call(Some(&ctx()), "build", &Map::new()).await.unwrap(); - assert!(!is_err, "a bare tool must dispatch, not be rejected as malformed"); + assert!( + !is_err, + "a bare tool must dispatch, not be rejected as malformed" + ); let fwd = tunnel.forwarded.lock().unwrap(); - assert_eq!(fwd[0].1, "uuid-p", "routed to the publishing server's tunnel"); - assert_eq!(fwd[0].2["name"], "build", "the original tool name is forwarded unchanged"); + assert_eq!( + fwd[0].1, "uuid-p", + "routed to the publishing server's tunnel" + ); + assert_eq!( + fwd[0].2["name"], "build", + "the original tool name is forwarded unchanged" + ); } #[tokio::test] @@ -1031,7 +1078,10 @@ mod tests { let _ = src.tools(Some(&ctx())); settle().await; - let (_v, is_err) = src.call(Some(&ctx()), "katashiro.click", &Map::new()).await.unwrap(); + let (_v, is_err) = src + .call(Some(&ctx()), "katashiro.click", &Map::new()) + .await + .unwrap(); assert!(!is_err); let fwd = tunnel.forwarded.lock().unwrap(); assert_eq!( @@ -1087,7 +1137,10 @@ mod tests { for tool in ["screenshot", "beta.screenshot"] { let (_v, is_err) = src.call(Some(&ctx()), tool, &Map::new()).await.unwrap(); - assert!(!is_err, "{tool} must dispatch — being advertised is the promise"); + assert!( + !is_err, + "{tool} must dispatch — being advertised is the promise" + ); } let fwd = tunnel.forwarded.lock().unwrap(); let routed: Vec<(String, String)> = fwd @@ -1256,10 +1309,22 @@ mod tests { let _ = src.tools(Some(&ctx())); settle().await; - let (_v, is_err) = src.call(Some(&ctx()), "browser.click", &Map::new()).await.unwrap(); - assert!(!is_err, "a mismatched-prefix tool must route to its publisher, not a phantom server"); + let (_v, is_err) = src + .call(Some(&ctx()), "browser.click", &Map::new()) + .await + .unwrap(); + assert!( + !is_err, + "a mismatched-prefix tool must route to its publisher, not a phantom server" + ); let fwd = tunnel.forwarded.lock().unwrap(); - assert_eq!(fwd[0].1, "uuid-k", "routed to katashiro, not the 'browser' prefix"); - assert_eq!(fwd[0].2["name"], "browser.click", "the original tool name is forwarded unchanged"); + assert_eq!( + fwd[0].1, "uuid-k", + "routed to katashiro, not the 'browser' prefix" + ); + assert_eq!( + fwd[0].2["name"], "browser.click", + "the original tool name is forwarded unchanged" + ); } } diff --git a/src/ctl.rs b/src/ctl.rs index e40fb7fc9..bce7cb9f9 100644 --- a/src/ctl.rs +++ b/src/ctl.rs @@ -68,9 +68,7 @@ pub trait CtlHandler: Send + Sync + 'static { /// Start the control socket server. Call this from `openab run` startup. /// Returns a JoinHandle; abort it on shutdown. #[cfg(unix)] -pub fn spawn_server( - handler: std::sync::Arc, -) -> tokio::task::JoinHandle<()> { +pub fn spawn_server(handler: std::sync::Arc) -> tokio::task::JoinHandle<()> { spawn_server_at(socket_path(), handler) } @@ -116,10 +114,7 @@ pub fn spawn_server_at( } #[cfg(unix)] -async fn handle_conn( - stream: UnixStream, - handler: &dyn CtlHandler, -) -> anyhow::Result<()> { +async fn handle_conn(stream: UnixStream, handler: &dyn CtlHandler) -> anyhow::Result<()> { let (reader, mut writer) = stream.into_split(); let mut lines = BufReader::new(reader).lines(); if let Some(line) = lines.next_line().await? { @@ -127,7 +122,9 @@ async fn handle_conn( let resp = match req.action { Action::Set => { let val = req.value.as_deref().unwrap_or(""); - handler.handle_set(req.thread_id.as_deref(), &req.key, val).await + handler + .handle_set(req.thread_id.as_deref(), &req.key, val) + .await } Action::Get => handler.handle_get(req.thread_id.as_deref(), &req.key).await, }; @@ -155,7 +152,10 @@ pub fn new_registry() -> ThreadRegistry { #[cfg(unix)] #[allow(dead_code)] pub async fn register_thread(registry: &ThreadRegistry, thread_id: &str, platform: &str) { - registry.write().await.insert(thread_id.to_string(), platform.to_string()); + registry + .write() + .await + .insert(thread_id.to_string(), platform.to_string()); } /// Type-alias for the Discord shard slot. When the discord feature is disabled, @@ -182,7 +182,11 @@ impl RuntimeHandler { registry: Arc>>, shard: ShardSlot, ) -> Self { - Self { adapters, registry, shard } + Self { + adapters, + registry, + shard, + } } /// Resolve which adapter to use for a given thread_id. @@ -234,7 +238,8 @@ impl CtlHandler for RuntimeHandler { let Some((adapter, tid)) = self.resolve(thread_id).await else { return Response { ok: false, - message: "unknown thread (use --thread or register via message dispatch)".into(), + message: "unknown thread (use --thread or register via message dispatch)" + .into(), value: None, }; }; @@ -262,7 +267,8 @@ impl CtlHandler for RuntimeHandler { let Some((_adapter, tid)) = self.resolve(thread_id).await else { return Response { ok: false, - message: "unknown thread (use --thread or register via message dispatch)".into(), + message: "unknown thread (use --thread or register via message dispatch)" + .into(), value: None, }; }; @@ -477,7 +483,12 @@ mod tests { struct MockHandler; #[async_trait::async_trait] impl CtlHandler for MockHandler { - async fn handle_set(&self, thread_id: Option<&str>, key: &str, value: &str) -> Response { + async fn handle_set( + &self, + thread_id: Option<&str>, + key: &str, + value: &str, + ) -> Response { Response { ok: true, message: format!("{key} = {value} (thread: {})", thread_id.unwrap_or("none")), @@ -503,24 +514,30 @@ mod tests { tokio::time::sleep(std::time::Duration::from_millis(50)).await; // Test set - let resp = send_request_to(&sock, &Request { - action: Action::Set, - key: "thread.name".into(), - value: Some("hello world".into()), - thread_id: Some("999".into()), - }) + let resp = send_request_to( + &sock, + &Request { + action: Action::Set, + key: "thread.name".into(), + value: Some("hello world".into()), + thread_id: Some("999".into()), + }, + ) .await .unwrap(); assert!(resp.ok); assert_eq!(resp.message, "thread.name = hello world (thread: 999)"); // Test get - let resp = send_request_to(&sock, &Request { - action: Action::Get, - key: "thread.name".into(), - value: None, - thread_id: None, - }) + let resp = send_request_to( + &sock, + &Request { + action: Action::Get, + key: "thread.name".into(), + value: None, + thread_id: None, + }, + ) .await .unwrap(); assert!(resp.ok); diff --git a/src/main.rs b/src/main.rs index f49879778..c3aa77326 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,3 +1,7 @@ +#[cfg(feature = "acp")] +mod acp_tunnel; +#[cfg(feature = "acp")] +mod acp_tunnel_source; mod ctl; #[cfg(any( feature = "telegram", @@ -10,10 +14,6 @@ mod ctl; feature = "lineworks", ))] mod unified_adapter; -#[cfg(feature = "acp")] -mod acp_tunnel; -#[cfg(feature = "acp")] -mod acp_tunnel_source; use openab_core::acp; use openab_core::adapter::{self, AdapterRouter}; use openab_core::bot_turns; @@ -495,8 +495,8 @@ async fn main() -> anyhow::Result<()> { #[cfg(feature = "acp")] let acp_tunnel_registry = openab_gateway::adapters::acp_server::new_tunnel_registry(); #[cfg(feature = "acp")] - let acp_tunnel: Arc = Arc::new( - acp_tunnel::RootAcpTunnel::new( + let acp_tunnel: Arc = + Arc::new(acp_tunnel::RootAcpTunnel::new( acp_tunnel_registry.clone(), // Browser control requires `[mcp]`, so the absent case is unreachable in practice; // fall back through the SAME function serde uses rather than repeating the literal. @@ -511,8 +511,7 @@ async fn main() -> anyhow::Result<()> { openab_gateway::adapters::acp_server::warn_if_tunnel_timeout_is_ineffective(t); t }, - ), - ); + )); // OAB MCP Facade (`[mcp]` in config.toml — OAB MCP Adapter ADR §6.2): // serve the loopback Streamable HTTP MCP server in-process so any coding @@ -543,15 +542,13 @@ async fn main() -> anyhow::Result<()> { // be skipped in bridge mode; with the bridge gone there is no mode in which the facade // runs without it. #[cfg(feature = "acp")] - let sources: Vec> = - vec![Arc::new(acp_tunnel_source::AcpTunnelSource::new( - acp_tunnel.clone(), - ))]; + let sources: Vec> = vec![Arc::new( + acp_tunnel_source::AcpTunnelSource::new(acp_tunnel.clone()), + )]; #[cfg(not(feature = "acp"))] let sources: Vec> = Vec::new(); tokio::spawn(async move { - if let Err(e) = - openab_mcp::mcp::facade::serve_http_with(&listen, sources, tokens).await + if let Err(e) = openab_mcp::mcp::facade::serve_http_with(&listen, sources, tokens).await { tracing::error!(error = %format!("{e:#}"), listen, "OAB MCP facade exited"); std::process::exit(1); @@ -579,7 +576,10 @@ async fn main() -> anyhow::Result<()> { facade_serving.then(|| { format!( "http://{}/mcp", - cfg.mcp.as_ref().map(|m| m.listen.as_str()).unwrap_or("127.0.0.1:8848") + cfg.mcp + .as_ref() + .map(|m| m.listen.as_str()) + .unwrap_or("127.0.0.1:8848") ) }), ); @@ -915,22 +915,21 @@ async fn main() -> anyhow::Result<()> { // Initialize filestore (for uploading file attachments to S3/R2). #[cfg(feature = "filestore")] - let filestore: Option> = if let Some(ref fs_cfg) = - cfg.filestore - { - info!( - bucket = %fs_cfg.bucket, - region = %fs_cfg.region, - prefix = %fs_cfg.prefix, - presigned_ttl = fs_cfg.presigned_ttl, - "filestore enabled" - ); - Some(Arc::new( - openab_core::filestore::Filestore::new(fs_cfg).await, - )) - } else { - None - }; + let filestore: Option> = + if let Some(ref fs_cfg) = cfg.filestore { + info!( + bucket = %fs_cfg.bucket, + region = %fs_cfg.region, + prefix = %fs_cfg.prefix, + presigned_ttl = fs_cfg.presigned_ttl, + "filestore enabled" + ); + Some(Arc::new( + openab_core::filestore::Filestore::new(fs_cfg).await, + )) + } else { + None + }; #[cfg(feature = "slack")] let shared_slack_adapter: Option> = cfg.slack.as_ref().map(|s| { @@ -990,13 +989,7 @@ async fn main() -> anyhow::Result<()> { configured_platforms.push("telegram"); } #[cfg(feature = "googlechat")] - if cfg - .googlechat - .clone() - .unwrap_or_default() - .resolve() - .enabled - { + if cfg.googlechat.clone().unwrap_or_default().resolve().enabled { configured_platforms.push("googlechat"); } #[cfg(feature = "lineworks")] @@ -1111,15 +1104,15 @@ async fn main() -> anyhow::Result<()> { #[cfg(feature = "filestore")] let gw_filestore = filestore.clone(); Some(tokio::spawn(async move { - if let Err(e) = - gateway::run_gateway_adapter( - params, - shutdown_rx, - gw_dispatcher, - gw_router, - #[cfg(feature = "filestore")] - gw_filestore, - ).await + if let Err(e) = gateway::run_gateway_adapter( + params, + shutdown_rx, + gw_dispatcher, + gw_router, + #[cfg(feature = "filestore")] + gw_filestore, + ) + .await { error!("gateway adapter error: {e}"); } @@ -1199,7 +1192,6 @@ async fn main() -> anyhow::Result<()> { )); } - // First-class `[telegram]` config overrides env-derived values // (config-authoritative + ${} expansion + TELEGRAM_* env fallback). #[cfg_attr(not(feature = "telegram"), allow(unused_variables))] @@ -1399,7 +1391,10 @@ async fn main() -> anyhow::Result<()> { f.config.api_base(), idle_ms, )); - info!(idle_ms, "unified: feishu card-streaming idle reaper started"); + info!( + idle_ms, + "unified: feishu card-streaming idle reaper started" + ); } if f.config.connection_mode == feishu::ConnectionMode::Websocket { let (shutdown_tx, shutdown_rx) = tokio::sync::watch::channel(false); @@ -1541,19 +1536,22 @@ async fn main() -> anyhow::Result<()> { info!(addr = %listen_addr, "unified webhook server starting"); - (Some(tokio::spawn(async move { - let listener = match tokio::net::TcpListener::bind(&listen_addr).await { - Ok(l) => l, - Err(e) => { - error!(addr = %listen_addr, error = %e, "unified webhook server bind failed"); - return; + ( + Some(tokio::spawn(async move { + let listener = match tokio::net::TcpListener::bind(&listen_addr).await { + Ok(l) => l, + Err(e) => { + error!(addr = %listen_addr, error = %e, "unified webhook server bind failed"); + return; + } + }; + info!(addr = %listen_addr, "unified webhook server listening"); + if let Err(e) = axum::serve(listener, app).await { + error!(error = %e, "unified webhook server error"); } - }; - info!(addr = %listen_addr, "unified webhook server listening"); - if let Err(e) = axum::serve(listener, app).await { - error!(error = %e, "unified webhook server error"); - } - })), Some(cron_unified_adapter)) + })), + Some(cron_unified_adapter), + ) } else { (None, None) } @@ -2005,13 +2003,10 @@ allowed_users = ["u1"] assert_ne!(trust.decide("c2", false, "u1"), Decision::Allow); // Empty lists → allow-all (matching the old inline filter default). - let gw_open = config::parse_config_str( - "[gateway]\nurl = \"ws://gw:8080/ws\"\n", - "test", - ) - .unwrap() - .gateway - .unwrap(); + let gw_open = config::parse_config_str("[gateway]\nurl = \"ws://gw:8080/ws\"\n", "test") + .unwrap() + .gateway + .unwrap(); let trust_open = gateway_section_trust(&gw_open); assert_eq!(trust_open.decide("any", false, "anyone"), Decision::Allow); } From 92c367b57d278c9700fb49438dd250f395015c1d Mon Sep 17 00:00:00 2001 From: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 30 Sep 2026 09:44:04 +0800 Subject: [PATCH 04/13] test: make run_inline_script_timeout resilient to transient spawn errors MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Under heavy parallel load (e.g. full workspace `cargo test` on a busy box), child spawn/wait can transiently fail with a non-timeout error before the 1s timeout path is reached — the previous exact-message assertion then flaked (observed in a detached replay run). Retry up to 3 attempts and require at least one to hit the timeout path; a real regression that never reaches the timeout arm still fails. --- crates/openab-core/src/hooks.rs | 23 ++++++++++++++++++++--- 1 file changed, 20 insertions(+), 3 deletions(-) diff --git a/crates/openab-core/src/hooks.rs b/crates/openab-core/src/hooks.rs index 51a63ab7d..4d63a473a 100644 --- a/crates/openab-core/src/hooks.rs +++ b/crates/openab-core/src/hooks.rs @@ -376,9 +376,26 @@ mod tests { timeout_seconds: 1, on_failure: OnFailure::Abort, }; - let result = run_hook("test", &hook).await; - assert!(result.is_err()); - assert!(result.unwrap_err().to_string().contains("timed out")); + // The timeout path itself is deterministic once the child is spawned, + // but under heavy parallel load spawn/wait can transiently fail with a + // non-timeout error (EAGAIN, reaper races). Retry a few attempts and + // require at least one to hit the timeout path. + let mut last_err = String::new(); + let mut timed_out = false; + for _ in 0..3 { + let result = run_hook("test", &hook).await; + assert!(result.is_err()); + let msg = result.unwrap_err().to_string(); + if msg.contains("timed out") { + timed_out = true; + break; + } + last_err = msg; + } + assert!( + timed_out, + "no attempt hit the timeout path; last error: {last_err}" + ); } #[tokio::test] From 132d42efce4fb1b125a9c85c8240699a832f4ad6 Mon Sep 17 00:00:00 2001 From: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 1 Oct 2026 04:57:02 +0800 Subject: [PATCH 05/13] fix(cp): close-frame writes must not be interrupted by a second signal MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pre-landing review of the shutdown drain found the last frame raced the close watch: a second latch landing mid-write (a lease sweep during the drain, or the process signal during a per-connection close) aborted the close frame entirely — the peer saw a bare TCP reset on precisely the path that exists to prevent one. Final frames now go through send_final_frame, bounded only by write_timeout since the connection is already committed to ending. A second shutdown signal during the drain is also now honored as a force-quit: exit immediately with the signal's conventional status (128 + signum), instead of waiting out the drain budget. Refs #1475 Generated with [Devin](https://devin.ai) Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- crates/openab-cp/src/main.rs | 37 ++++++++++++++++++++++ crates/openab-cp/src/server.rs | 44 ++++++++++++++++++-------- crates/openab-cp/tests/ws_lifecycle.rs | 30 ++++++++++++++++++ docs/adr/agent-control-plane.md | 4 ++- 4 files changed, 101 insertions(+), 14 deletions(-) diff --git a/crates/openab-cp/src/main.rs b/crates/openab-cp/src/main.rs index f8e268bd3..e7d1a19f5 100644 --- a/crates/openab-cp/src/main.rs +++ b/crates/openab-cp/src/main.rs @@ -60,6 +60,16 @@ async fn shutdown_signal() -> &'static str { } } +/// Conventional exit status for a signal-driven exit: 128 + signal number. +fn signal_exit_code(signal: &str) -> i32 { + match signal { + "SIGHUP" => 129, + "SIGINT" => 130, + "SIGTERM" => 143, + _ => 1, + } +} + #[tokio::main] async fn main() -> Result<()> { tracing_subscriber::fmt() @@ -94,6 +104,17 @@ async fn main() -> Result<()> { }) .await?; + // A second signal while the drain is still running is the operator's + // "stop waiting": exit immediately with the signal's conventional + // status. `shutdown_signal` registers fresh streams each call, so the + // repeat SIGINT/SIGTERM/SIGHUP still resolves it — the drain itself is + // bounded, so this only shortens a wait, never causes one. + let force_quit = tokio::spawn(async { + let signal = shutdown_signal().await; + tracing::warn!(signal, "second shutdown signal — forcing immediate exit"); + std::process::exit(signal_exit_code(signal)); + }); + // The listener has stopped accepting and in-flight HTTP work is done; // the WS connection tasks are detached spawns still running. Give them a // real shutdown — synthesized terminals, close frames, bounded drain — @@ -101,8 +122,24 @@ async fn main() -> Result<()> { // stops last: a lease/deadline expiry during the drain still // synthesizes normally. graceful_shutdown(&state).await; + force_quit.abort(); sweeper.abort(); let _ = sweeper.await; info!("openab-cp stopped"); Ok(()) } + +#[cfg(test)] +mod tests { + use super::signal_exit_code; + + #[test] + fn signal_exit_codes_follow_the_128_plus_signum_convention() { + // The force-quit path must report WHICH signal ended the process so + // an operator (or a supervisor's exit-code log) can tell a forced + // exit from a clean drain — and from a crash. + assert_eq!(signal_exit_code("SIGHUP"), 129); + assert_eq!(signal_exit_code("SIGINT"), 130); + assert_eq!(signal_exit_code("SIGTERM"), 143); + } +} diff --git a/crates/openab-cp/src/server.rs b/crates/openab-cp/src/server.rs index 1e5977e04..c42d15bef 100644 --- a/crates/openab-cp/src/server.rs +++ b/crates/openab-cp/src/server.rs @@ -282,6 +282,23 @@ async fn send_bounded( } } +/// Send the LAST frame a connection will ever get — a close frame, or a +/// terminal error response written just before teardown. Unlike +/// [`send_bounded`] this is raced only against `write_timeout`, never the +/// close watch: the connection is already committed to ending, and a second +/// signal latching mid-write (e.g. a lease sweep landing during the +/// shutdown drain, or the process signal during a per-connection close) +/// must not eat the frame the peer was promised. A peer that has stopped +/// reading is still bounded by `write_timeout`, so it cannot hold teardown +/// longer than that. +async fn send_final_frame( + sink: &mut futures_util::stream::SplitSink, + msg: Message, + write_timeout: Duration, +) { + let _ = tokio::time::timeout(write_timeout, sink.send(msg)).await; +} + async fn ws_handler( State(state): State>, headers: HeaderMap, @@ -359,7 +376,7 @@ async fn handle_connection( // only reports transitions that postdate the subscription, so the // already-latched reason is checked once explicitly. if let Some(reason) = close.process_reason() { - let _ = send_bounded(&mut sink, restart_close(reason), write_timeout, &mut close).await; + send_final_frame(&mut sink, restart_close(reason), write_timeout).await; return; } @@ -389,7 +406,7 @@ async fn handle_connection( Some(reason) => restart_close(reason), None => policy_close(close.reason().unwrap_or(REASON_SHUTDOWN)), }; - let _ = send_bounded(&mut sink, frame, write_timeout, &mut close).await; + send_final_frame(&mut sink, frame, write_timeout).await; return; } } { @@ -404,11 +421,10 @@ async fn handle_connection( timeout_secs = state.cfg.register_timeout_secs, "no cp/register within the registration deadline — closing" ); - let _ = send_bounded( + send_final_frame( &mut sink, policy_close(REASON_REGISTER_TIMEOUT), write_timeout, - &mut close, ) .await; return; @@ -418,11 +434,10 @@ async fn handle_connection( Ok(ok) => ok, Err((id, err)) => { let resp = JsonRpcErrorResponse::new(id, err); - let _ = send_bounded( + send_final_frame( &mut sink, Message::Text(serde_json::to_string(&resp).expect("serializable").into()), write_timeout, - &mut close, ) .await; return; @@ -482,11 +497,10 @@ async fn handle_connection( ), ), ); - let _ = send_bounded( + send_final_frame( &mut sink, Message::Text(serde_json::to_string(&resp).expect("serializable").into()), write_timeout, - &mut close, ) .await; return; @@ -542,7 +556,7 @@ async fn handle_connection( None => close.reason().map(policy_close), }; if let Some(frame) = frame { - let _ = send_bounded(&mut sink, frame, write_timeout, &mut close).await; + send_final_frame(&mut sink, frame, write_timeout).await; } } // `_registered` runs teardown on the way out. @@ -653,8 +667,10 @@ async fn handle_connection( ); // Bounded like every other write: a peer that has stopped reading // must not be able to hold teardown (and its quota slot) by refusing - // to accept the close frame. - let _ = send_bounded(&mut sink, restart_close(reason), write_timeout, &mut close).await; + // to accept the close frame. The write is NOT raced against the + // close watch: a second signal latching mid-write (a lease sweep + // landing mid-drain) must not eat the promised close frame. + send_final_frame(&mut sink, restart_close(reason), write_timeout).await; } else if let Some(reason) = cp_close_reason { info!( agent = %format!("{}/{}", identity.namespace, identity.name), @@ -664,8 +680,10 @@ async fn handle_connection( ); // Bounded like every other write: a peer that has stopped reading must // not be able to delay teardown (and its quota slot) by refusing to - // accept the close frame. - let _ = send_bounded(&mut sink, policy_close(reason), write_timeout, &mut close).await; + // accept the close frame. Not raced against the close watch either: + // the process signal may latch mid-write while this frame is still + // the only explanation the peer will get. + send_final_frame(&mut sink, policy_close(reason), write_timeout).await; } // Teardown runs here, when `_registered` drops — on this path and on an diff --git a/crates/openab-cp/tests/ws_lifecycle.rs b/crates/openab-cp/tests/ws_lifecycle.rs index baac339cf..e926bcea5 100644 --- a/crates/openab-cp/tests/ws_lifecycle.rs +++ b/crates/openab-cp/tests/ws_lifecycle.rs @@ -904,6 +904,36 @@ async fn spawn_cp_process() -> (tokio::process::Child, String, std::path::PathBu (child, format!("ws://127.0.0.1:{port}/cp"), cfg_path) } +#[tokio::test] +async fn a_second_signal_latching_mid_close_still_sends_a_close_frame() { + // The close a connection is promised must not be eaten by a second + // signal latching while its frame is still on the wire: a lease sweep + // firing during the shutdown drain (or the process signal landing + // during a per-connection close) previously aborted the close write — + // the peer saw a bare TCP reset instead of the Close frame. Whichever + // shutdown path wins, some Close frame must always be observed. + let (state, url) = spawn_cp(cfg("register_timeout_secs = 30")).await; + let mut ws = connect(&url).await.expect("connection accepted"); + assert_eq!( + register(&mut ws, "i-1").await["result"]["protocol_version"], + 1 + ); + let handle = state.registry.list("prod")[0].handle; + + // Latch the per-connection signal the way the sweeper does, then the + // process-wide one immediately behind it — the second latch can land + // while the first close frame is still being written. + assert!(state.registry.signal_shutdown(handle, "lease expired")); + graceful_shutdown(&state).await; + + let closed = wait_closed(&mut ws, Duration::from_secs(5)).await; + assert!( + matches!(closed, Some(Closed::Frame { .. })), + "a second signal mid-close must not reduce the peer's close frame \ + to a bare reset — got {closed:?}" + ); +} + #[cfg(unix)] #[tokio::test] async fn sigterm_and_sighup_run_the_drain_and_exit_cleanly() { diff --git a/docs/adr/agent-control-plane.md b/docs/adr/agent-control-plane.md index ae6a46525..ff1a93485 100644 --- a/docs/adr/agent-control-plane.md +++ b/docs/adr/agent-control-plane.md @@ -552,7 +552,9 @@ recovery semantics: `cp/cancel`, both stamped `control plane shutting down`, and admissions racing the drain are refused `SATURATED` — then every connection closes with code 1012, all inside `shutdown_drain_secs` plus at most one - `write_timeout_secs` for a stalled close write. A + `write_timeout_secs` for a stalled close write. A second shutdown signal + during the drain is the operator's "stop waiting": the process exits + immediately with the signal's conventional status (128 + signum). A *hard* restart (SIGKILL, crash, drain budget exhausted) is equivalent to every lease expiring at once *with* the connection closure that implies — except that the CP is not there to send it: the in-flight table and the From e42a7bfe1e116071655de6a35771a5a49d633dbc Mon Sep 17 00:00:00 2001 From: worker-w15 Date: Thu, 1 Oct 2026 21:40:11 +0800 Subject: [PATCH 06/13] test(cp): prove SIGTERM resolves in-flight work on a real process MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The in-process regressions cover the drain's individual steps, but the part that actually breaks in a container is a race between the drain and process exit: the synthesized terminals and the 1012 close frames must be on the wire before `main` returns, or the runtime sees the same TCP reset the shutdown path exists to prevent. Only a real process sent a real SIGTERM reproduces that ordering. The test therefore spawns the binary on an ephemeral port, holds one delegation in flight, and sends SIGTERM: the initiator must receive its `target_disconnected` terminal naming the admission, the serving runtime its `cp/cancel`, both sockets a 1012 close carrying the reason, and the process exit 0 rather than dying by signal. A second test races fresh upgrades against the drain and asserts none is ever acked — a CP that is leaving must not admit a registration it cannot serve. Refs #1475 --- crates/openab-cp/tests/cp_shutdown.rs | 380 ++++++++++++++++++++++++++ 1 file changed, 380 insertions(+) create mode 100644 crates/openab-cp/tests/cp_shutdown.rs diff --git a/crates/openab-cp/tests/cp_shutdown.rs b/crates/openab-cp/tests/cp_shutdown.rs new file mode 100644 index 000000000..f9f24a025 --- /dev/null +++ b/crates/openab-cp/tests/cp_shutdown.rs @@ -0,0 +1,380 @@ +//! Container-grade shutdown of a REAL `openab-cp` process, with work in +//! flight. +//! +//! `docker stop`, ECS and k8s all stop a task with SIGTERM and escalate to +//! SIGKILL when the grace period ends. Handling only ctrl-c/SIGINT meant every +//! deploy and scale-in killed the CP mid-flight: connected runtimes saw a TCP +//! reset instead of a Close frame, and in-flight delegations were resolved +//! only by each client's own deadline reconciliation against a vanished CP. +//! +//! Signal handling lives in `main`, so unit tests cannot reach it — but neither +//! can an in-process test of `graceful_shutdown` prove the part that actually +//! breaks in a container: that the synthesized terminals and the close frames +//! are written to the wire BEFORE the process exits. That ordering is a race +//! between the drain and process exit, and only a real process reproduces it. + +#![cfg(unix)] + +use std::path::PathBuf; +use std::process::Stdio; +use std::sync::atomic::{AtomicU32, Ordering}; +use std::time::{Duration, Instant}; + +use futures_util::{SinkExt, StreamExt}; +use tokio::net::TcpStream; +use tokio::process::{Child, Command}; +use tokio_tungstenite::tungstenite::client::IntoClientRequest; +use tokio_tungstenite::tungstenite::Message; +use tokio_tungstenite::{MaybeTlsStream, WebSocketStream}; + +const KEY: &str = "k-primary"; +const KEY_WORKER: &str = "k-worker"; + +/// Reason the CP must put on every close frame when it shuts down, and on the +/// terminal/cancel frames it synthesizes for in-flight delegations. +/// +/// Spelled out as a literal rather than imported from the crate: these are the +/// values a runtime sees on the wire, and a test that imported the constant +/// would pass even if the constant itself were wrong. +const REASON_SHUTDOWN: &str = "control plane shutting down"; + +type Ws = WebSocketStream>; + +/// Distinguishes concurrent tests' temp configs without a random dependency. +static SEQ: AtomicU32 = AtomicU32::new(0); + +/// A loopback port that was free a moment ago. The CP is started immediately +/// after, and [`connect_retry`] absorbs the race either way. +fn free_port() -> u16 { + let probe = std::net::TcpListener::bind("127.0.0.1:0").expect("a free loopback port"); + let port = probe.local_addr().expect("bound address").port(); + drop(probe); + port +} + +fn config_file(port: u16) -> PathBuf { + let seq = SEQ.fetch_add(1, Ordering::Relaxed); + let path = std::env::temp_dir().join(format!( + "openab-cp-sigterm-{}-{seq}.toml", + std::process::id() + )); + let raw = format!( + r#" +listen = "127.0.0.1:{port}" +register_timeout_secs = 30 +shutdown_drain_secs = 5 + +[[agents]] +key = "{KEY}" +namespace = "prod" +name = "koudu" +type = "primary" + +[[agents]] +key = "{KEY_WORKER}" +namespace = "prod" +name = "worker-1" +type = "worker" +"# + ); + std::fs::write(&path, raw).expect("writing the CP config"); + path +} + +/// A spawned CP process plus the temp config it reads (removed on drop). +struct Cp { + child: Child, + url: String, + config: PathBuf, +} + +impl Drop for Cp { + fn drop(&mut self) { + // `kill_on_drop` reaps the child if the test panics; the config file is + // a temp artifact either way. + let _ = std::fs::remove_file(&self.config); + } +} + +/// Start the real `openab-cp` binary on a free loopback port. +async fn spawn_cp() -> Cp { + let port = free_port(); + let config = config_file(port); + let child = Command::new(env!("CARGO_BIN_EXE_openab-cp")) + .arg("--config") + .arg(&config) + // Not piped: a full pipe buffer would block the child mid-write. + .stdout(Stdio::null()) + .stderr(Stdio::null()) + .kill_on_drop(true) + .spawn() + .expect("spawning openab-cp"); + Cp { + child, + url: format!("ws://127.0.0.1:{port}/cp"), + config, + } +} + +async fn connect_as(url: &str, key: &str) -> Result { + let mut req = url.into_client_request().expect("client request"); + req.headers_mut().insert( + "authorization", + format!("Bearer {key}").parse().expect("header value"), + ); + tokio_tungstenite::connect_async(req) + .await + .map(|(ws, _)| ws) +} + +/// Connect as `key`, retrying while the CP is still binding its port. +async fn connect_retry(url: &str, key: &str) -> Ws { + for _ in 0..100 { + if let Ok(ws) = connect_as(url, key).await { + return ws; + } + tokio::time::sleep(Duration::from_millis(50)).await; + } + panic!("the CP never accepted a connection"); +} + +fn register_frame(name: &str, instance_id: &str) -> String { + let ty = if name == "koudu" { "primary" } else { "worker" }; + serde_json::json!({ + "jsonrpc": "2.0", "id": 1, "method": "cp/register", + "params": { + "protocol_version": 1, + "namespace": "prod", + "name": name, + "type": ty, + "instance_id": instance_id + } + }) + .to_string() +} + +/// Send `cp/register` and return the parsed ack. +async fn register(ws: &mut Ws, name: &str, instance_id: &str) -> serde_json::Value { + ws.send(Message::Text(register_frame(name, instance_id).into())) + .await + .unwrap(); + let msg = tokio::time::timeout(Duration::from_secs(10), ws.next()) + .await + .expect("register must be answered") + .expect("stream open") + .expect("no ws error"); + serde_json::from_str(msg.to_text().expect("text frame")).unwrap() +} + +/// Read JSON frames until one satisfies `want`; panics on timeout or a stream +/// that ends first (a reset is exactly the failure these tests exist to catch). +async fn await_frame( + ws: &mut Ws, + what: &str, + want: impl Fn(&serde_json::Value) -> bool, +) -> serde_json::Value { + let deadline = Instant::now() + Duration::from_secs(20); + while Instant::now() < deadline { + match tokio::time::timeout(Duration::from_millis(200), ws.next()).await { + Ok(Some(Ok(Message::Text(t)))) => { + let v: serde_json::Value = serde_json::from_str(&t).unwrap(); + if want(&v) { + return v; + } + } + Ok(Some(Ok(_))) => continue, + Ok(None) | Ok(Some(Err(_))) => panic!("stream ended while waiting for {what}"), + Err(_) => continue, + } + } + panic!("never received {what}"); +} + +/// The Close frame the CP must send, read as `(code, reason)`. +async fn await_close(ws: &mut Ws) -> (u16, String) { + let deadline = Instant::now() + Duration::from_secs(20); + while Instant::now() < deadline { + match tokio::time::timeout(Duration::from_millis(200), ws.next()).await { + Ok(Some(Ok(Message::Close(Some(cf))))) => { + return (cf.code.into(), cf.reason.to_string()) + } + Ok(Some(Ok(_))) => continue, + Ok(None) | Ok(Some(Err(_))) => panic!("the socket dropped instead of a Close frame"), + Err(_) => continue, + } + } + panic!("never received a Close frame"); +} + +/// Send SIGTERM to `child`, the way every container runtime stops a task. +/// Synchronous on purpose: it must happen the instant the test asks for it, +/// with no await point in between where the CP could reach its own exit. +fn send_sigterm(pid: u32) { + let status = std::process::Command::new("kill") + .arg("-TERM") + .arg(pid.to_string()) + .status() + .expect("running kill(1)"); + assert!(status.success(), "kill -TERM failed: {status}"); +} + +fn is_shutdown_close(code: u16, reason: &str) -> bool { + // 1012 (service restart), not the 1008 the CP sends for misbehaviour: + // nothing about a deploy is a policy violation, and the client should + // reconnect rather than fix anything. + code == 1012 && reason == REASON_SHUTDOWN +} + +#[tokio::test] +async fn sigterm_resolves_in_flight_delegations_before_the_process_exits() { + // The whole point of the issue, over a real socket and a real process: + // under SIGTERM the CP must resolve the delegation it was serving, close + // both sockets with a reason, and only then exit 0 — all inside the + // orchestrator's grace period, before it escalates to SIGKILL. + let mut cp = spawn_cp().await; + + let mut initiator = connect_retry(&cp.url, KEY).await; + assert_eq!( + register(&mut initiator, "koudu", "i-1").await["result"]["protocol_version"], + 1 + ); + let mut worker = connect_retry(&cp.url, KEY_WORKER).await; + assert_eq!( + register(&mut worker, "worker-1", "i-w").await["result"]["protocol_version"], + 1 + ); + + // One delegation in flight and never completed. + let deadline = (chrono::Utc::now() + chrono::Duration::seconds(600)).to_rfc3339(); + initiator + .send(Message::Text( + serde_json::json!({ + "jsonrpc": "2.0", "id": 10, "method": "cp/delegate", + "params": { + "delegation_id": "d-1", + "target": {"name": "worker-1"}, + "prompt": "do it", + "deadline": deadline + } + }) + .to_string() + .into(), + )) + .await + .unwrap(); + let ack = await_frame(&mut initiator, "the delegation ack", |v| v["id"] == 10).await; + assert!(ack.get("error").is_none(), "delegation accepted: {ack}"); + let admission = ack["result"]["admission"] + .as_u64() + .expect("the ack carries the admission token"); + await_frame(&mut worker, "the forwarded delegation", |v| { + v["method"] == "cp/delegate" + }) + .await; + + send_sigterm(cp.child.id().expect("the CP is still running")); + + // The initiator is told how the delegation ended, naming the admission it + // ends — not left to reconcile its own deadline against a vanished CP. + let terminal = await_frame(&mut initiator, "the shutdown terminal", |v| { + v["method"] == "cp/delegate_result" + }) + .await; + assert_eq!(terminal["params"]["delegation_id"], "d-1"); + assert_eq!( + terminal["params"]["status"], "target_disconnected", + "shutdown resolves the delegation like a lost target, not a timeout: {terminal}" + ); + assert_eq!( + terminal["params"]["admission"].as_u64(), + Some(admission), + "the synthesized terminal names the admission it ends" + ); + assert_eq!(terminal["params"]["error"], REASON_SHUTDOWN); + + // The serving runtime is told to stop working before its socket closes. + let cancel = await_frame(&mut worker, "the shutdown cancel", |v| { + v["method"] == "cp/cancel" + }) + .await; + assert_eq!(cancel["params"]["delegation_id"], "d-1"); + assert_eq!(cancel["params"]["admission"].as_u64(), Some(admission)); + assert_eq!(cancel["params"]["reason"], REASON_SHUTDOWN); + + // Both sockets get a Close frame that says why — a client can tell a + // graceful CP shutdown from a network drop. + let (i_code, i_reason) = await_close(&mut initiator).await; + assert!( + is_shutdown_close(i_code, &i_reason), + "initiator close was {i_code} {i_reason:?}" + ); + let (w_code, w_reason) = await_close(&mut worker).await; + assert!( + is_shutdown_close(w_code, &w_reason), + "worker close was {w_code} {w_reason:?}" + ); + + // And the process leaves on its own, cleanly: an orchestrator must not + // have to escalate to SIGKILL. + let status = tokio::time::timeout(Duration::from_secs(30), cp.child.wait()) + .await + .expect("the CP must exit after SIGTERM without a SIGKILL") + .expect("reaping the CP"); + assert_eq!( + status.code(), + Some(0), + "a drained shutdown exits 0, not by signal: {status:?}" + ); +} + +#[tokio::test] +async fn a_connection_opened_during_the_drain_never_registers() { + // The drain only means something if the CP stops admitting work while it + // runs: a runtime that reconnects into a closing CP must be turned away + // (or closed politely), never handed a socket that dies mid-registration. + let mut cp = spawn_cp().await; + let mut existing = connect_retry(&cp.url, KEY).await; + register(&mut existing, "koudu", "i-1").await; + + send_sigterm(cp.child.id().expect("the CP is still running")); + + // Race the drain with fresh upgrades. Every outcome is acceptable except + // a successful registration: a refused connection, or a socket closed with + // the shutdown reason before it can register. + let mut registered = false; + for _ in 0..40 { + match connect_as(&cp.url, KEY_WORKER).await { + // The listener is already gone, or the CP is not far enough into + // its shutdown for the refusal to be observable yet. + Err(_) => {} + Ok(mut ws) => match tokio::time::timeout(Duration::from_secs(5), ws.next()).await { + Ok(Some(Ok(Message::Close(Some(cf))))) => assert!( + is_shutdown_close(cf.code.into(), cf.reason.as_ref()), + "a socket closed during the drain must carry the shutdown reason: {cf:?}" + ), + Ok(Some(Ok(Message::Text(t)))) => { + let v: serde_json::Value = serde_json::from_str(&t).unwrap(); + assert!( + v.get("result").is_none(), + "a socket admitted during the drain must not be acked: {v}" + ); + registered = true; + } + Ok(Some(Ok(_))) => continue, + Ok(None) | Ok(Some(Err(_))) => {} + Err(_) => panic!("a socket accepted during the drain never ended"), + }, + } + tokio::time::sleep(Duration::from_millis(50)).await; + } + assert!( + !registered, + "the CP must not have acked a registration once it started draining" + ); + + let status = tokio::time::timeout(Duration::from_secs(30), cp.child.wait()) + .await + .expect("the CP must exit after SIGTERM") + .expect("reaping the CP"); + assert_eq!(status.code(), Some(0), "{status:?}"); +} From a41d5190b928a29b9e771b3902e1d243d7e1c558 Mon Sep 17 00:00:00 2001 From: worker-w15 Date: Thu, 1 Oct 2026 22:16:36 +0800 Subject: [PATCH 07/13] fix(cp): a stalled peer must not be able to pin or truncate the drain MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pre-landing review reproduced two defects in the shutdown sequence, both of which end with the orchestrator's SIGKILL doing the closing the CP was supposed to do: - The drain ran AFTER `axum::serve`'s graceful shutdown resolved. That wait ends when the last in-flight HTTP request finishes, so a peer that opened a socket and then stopped sending — a half-sent request, a stalled upgrade — pinned it indefinitely: no `fail_all`, no synthesized terminals, no close frames, just a process that sat there until it was killed. The shutdown sequence is now spawned at the signal, and the listener's own wait is bounded by the same budget as the drain. - The second-signal force-quit was armed only after that wait, so the operator's "stop waiting" never fired in exactly the case it exists for. It is armed at the signal and left armed until the process is gone, so it covers the drain and the listener's bookkeeping alike. Also tightened, from the same review: - Every budget in the drain is now a remainder of ONE deadline, latched with the close signal. Each connection previously started a fresh `shutdown_drain_secs` when it left its main loop, so a late one could be granted a longer life than the drain it is part of — and the close write, bounded by `write_timeout_secs`, could push the process exit past the documented ceiling. `shutdown_drain_secs` is now the single bound, and the docs say so. - `shutdown_drain_secs = 0` is refused: it emitted the observer-side `delegation_completed` events and then exited before any frame reached a wire, so the two surfaces disagreed by construction. - A terminal that the drain cannot queue (full peer queue) is now logged at warn level naming the agent, instead of being dropped silently — the drain is that delegation's last chance, and the surviving outcome is the initiator's own deadline reconciliation. - A `cp/delegate_result` that arrives during the drain no longer blames backpressure: the queue was closed deliberately, not overflowed, and the connection is already closing with the shutdown reason. Regression tests: `a_stalled_http_request_cannot_pin_the_shutdown` and `a_second_signal_exits_immediately_while_the_drain_is_pending` drive the real binary with a half-sent HTTP request and a SIGTERM pair (both fail against the previous `main`), and `shutdown_survives_an_initiator_whose_queue_is_already_full` pins the refusal path. The mid-close test now asserts the actual 1012 close instead of accepting any close frame, and the spawned CP's stdio is no longer inherited into the harness log. Refs #1475 --- crates/openab-cp/cp.toml.example | 6 +- crates/openab-cp/src/config.rs | 25 ++- crates/openab-cp/src/main.rs | 97 ++++++++--- crates/openab-cp/src/server.rs | 212 +++++++++++++++++++++---- crates/openab-cp/tests/cp_shutdown.rs | 74 +++++++++ crates/openab-cp/tests/ws_lifecycle.rs | 14 +- docs/adr/agent-control-plane.md | 13 +- docs/control-plane.md | 21 ++- 8 files changed, 385 insertions(+), 77 deletions(-) diff --git a/crates/openab-cp/cp.toml.example b/crates/openab-cp/cp.toml.example index 5fb04c678..82a3a7a5d 100644 --- a/crates/openab-cp/cp.toml.example +++ b/crates/openab-cp/cp.toml.example @@ -66,8 +66,10 @@ max_outbound_queue_bytes = 16777216 # resolves every in-flight delegation (initiators get a synthesized # target_disconnected terminal, serving runtimes a cp/cancel), then each # connection flushes its queued outbound frames and closes with code 1012 — -# all inside this budget; a stalled peer can hold it no longer than one -# write_timeout_secs more. 0 exits as soon as the close signal is raised. +# and this budget is the ceiling for all of it, close write and process exit +# included, so a peer that stops reading cannot push the exit past it. Must +# be greater than 0: a zero budget would emit the observer-side terminal +# events and then exit before any of the frames reached a wire. shutdown_drain_secs = 5 # Process-wide ceiling on simultaneously in-flight delegations, enforced at diff --git a/crates/openab-cp/src/config.rs b/crates/openab-cp/src/config.rs index 40b0cb7d8..964b37fc3 100644 --- a/crates/openab-cp/src/config.rs +++ b/crates/openab-cp/src/config.rs @@ -97,12 +97,18 @@ pub struct CpConfig { /// On shutdown every in-flight delegation is resolved first — the /// initiator gets a synthesized `target_disconnected` terminal and the /// serving runtime a `cp/cancel` — then each connection flushes its - /// queued outbound frames and sends its close frame, all inside this - /// budget. A stalled peer's close write is additionally bounded by - /// `write_timeout_secs`, so it can never hold the CP past this budget - /// plus one write timeout. 0 disables the wait: connections are - /// signalled and the process exits as soon as the listener is closed - /// (sockets still get their close frames if they flush in time). + /// queued outbound frames and sends its close frame. This budget is the + /// ceiling for ALL of it: the flush, the close write and the process exit + /// are remainders of one deadline, so a peer that stops reading cannot + /// push the exit past it (that is the point — an orchestrator escalates to + /// SIGKILL when ITS grace period ends, and a CP still waiting then gets + /// the TCP reset this whole path exists to avoid). + /// + /// Must be greater than 0. A zero budget would exit before any + /// synthesized terminal or close frame reached a wire, while still having + /// emitted the observer-side `delegation_completed` events — the two + /// surfaces would disagree by construction, and the shutdown would be + /// indistinguishable from being killed. #[serde(default = "default_shutdown_drain_secs")] pub shutdown_drain_secs: u64, /// Maximum size of prompt/result excerpts mirrored to observers in @@ -340,6 +346,13 @@ impl CpConfig { if self.max_inflight_delegations == 0 { bail!("max_inflight_delegations must be at least 1"); } + // Zero would exit the process before any synthesized terminal or close frame + // reached a wire — while the observer-side `delegation_completed` + // events had already been emitted, so the two surfaces would disagree + // and the shutdown would be indistinguishable from a SIGKILL. + if self.shutdown_drain_secs == 0 { + bail!("shutdown_drain_secs must be greater than 0"); + } // Zero would clamp every uncapped identity to no capacity at all, so // every delegation to it would be SATURATED. if self.default_max_delegated_sessions_cap == 0 { diff --git a/crates/openab-cp/src/main.rs b/crates/openab-cp/src/main.rs index e7d1a19f5..a0ca2c292 100644 --- a/crates/openab-cp/src/main.rs +++ b/crates/openab-cp/src/main.rs @@ -1,10 +1,11 @@ //! Standalone control-plane binary. use std::sync::Arc; +use std::time::Duration; use anyhow::{Context, Result}; use clap::Parser; -use tracing::info; +use tracing::{info, warn}; use openab_cp::config::CpConfig; use openab_cp::server::{app, graceful_shutdown, run_sweeper, AppState}; @@ -92,37 +93,87 @@ async fn main() -> Result<()> { ); let state = Arc::new(AppState::new(cfg)); + let drain_budget = Duration::from_secs(state.cfg.shutdown_drain_secs); + // The process's exit is bounded by the drain budget plus slack for the + // listener's own bookkeeping. One knob, one bound: an orchestrator + // escalates to SIGKILL at the end of its grace period, so the CP's own + // deadline has to be the shorter one. + let hard_bound = drain_budget + Duration::from_secs(1); let sweeper = tokio::spawn(run_sweeper(state.clone())); let listener = tokio::net::TcpListener::bind(&listen) .await .with_context(|| format!("binding {listen}"))?; - axum::serve(listener, app(state.clone())) - .with_graceful_shutdown(async { + + // The shutdown sequence is SPAWNED rather than awaited after the listener. + // Ordering it after `serve` was a real defect: axum's graceful shutdown + // ends when the last in-flight HTTP request finishes, and a peer that + // opened a socket and then said nothing (a half-sent request, a stalled + // upgrade) pins that wait for as long as the orchestrator allows the + // process to live — so the drain would never start, and the CP would be + // SIGKILLed with its delegations unresolved. Spawning it makes the drain + // begin the instant the signal lands, whatever hyper is doing. + let (stop_accepting_tx, stop_accepting_rx) = tokio::sync::oneshot::channel::<()>(); + let (drained_tx, drained_rx) = tokio::sync::oneshot::channel::<()>(); + tokio::spawn({ + let state = Arc::clone(&state); + async move { let signal = shutdown_signal().await; info!(signal, "shutdown signal received — draining connections"); - }) - .await?; - - // A second signal while the drain is still running is the operator's - // "stop waiting": exit immediately with the signal's conventional - // status. `shutdown_signal` registers fresh streams each call, so the - // repeat SIGINT/SIGTERM/SIGHUP still resolves it — the drain itself is - // bounded, so this only shortens a wait, never causes one. - let force_quit = tokio::spawn(async { - let signal = shutdown_signal().await; - tracing::warn!(signal, "second shutdown signal — forcing immediate exit"); - std::process::exit(signal_exit_code(signal)); + // A second signal at ANY point before the process exits is the + // operator's "stop waiting": exit immediately with the signal's + // conventional status (128 + signum). Armed before the drain and + // deliberately left armed — the drain and the listener's wait are + // separate waits, and the escape hatch has to cover both. The task + // ends with the process if it never fires. + tokio::spawn(async { + let signal = shutdown_signal().await; + warn!(signal, "second shutdown signal — forcing immediate exit"); + std::process::exit(signal_exit_code(signal)); + }); + // Nothing new may enter a CP that is leaving. + let _ = stop_accepting_tx.send(()); + graceful_shutdown(&state).await; + let _ = drained_tx.send(()); + } }); - // The listener has stopped accepting and in-flight HTTP work is done; - // the WS connection tasks are detached spawns still running. Give them a - // real shutdown — synthesized terminals, close frames, bounded drain — - // instead of letting process exit reset their sockets. The sweeper - // stops last: a lease/deadline expiry during the drain still - // synthesizes normally. - graceful_shutdown(&state).await; - force_quit.abort(); + // The listener stops accepting as soon as the signal lands, and the wait + // for it is bounded by the same hard bound — hyper's bookkeeping must not + // be able to outlive the drain it is supposed to follow. + let serving = + axum::serve(listener, app(Arc::clone(&state))).with_graceful_shutdown(async move { + let _ = stop_accepting_rx.await; + }); + let finished = tokio::time::timeout(hard_bound, async { + tokio::join!( + async { + // The drain is the shutdown: the synthesized terminals and the + // close frames get their chance to reach the wire here. + let _ = drained_rx.await; + }, + async { + match tokio::time::timeout(hard_bound, serving).await { + Ok(Ok(())) => {} + Ok(Err(e)) => warn!(error = %e, "listener stopped with an error"), + Err(_) => warn!( + bound_secs = hard_bound.as_secs(), + "the listener's graceful wait outlived the shutdown bound — dropping it" + ), + } + }, + ); + }) + .await; + if finished.is_err() { + warn!( + bound_secs = hard_bound.as_secs(), + "shutdown did not finish within its bound — exiting with the drain unfinished" + ); + } + + // Only now: the sweeper exists to keep registrations and delegations + // honest, and during a shutdown its findings would race the closes above. sweeper.abort(); let _ = sweeper.await; info!("openab-cp stopped"); diff --git a/crates/openab-cp/src/server.rs b/crates/openab-cp/src/server.rs index c42d15bef..58bef9dce 100644 --- a/crates/openab-cp/src/server.rs +++ b/crates/openab-cp/src/server.rs @@ -64,6 +64,9 @@ pub struct AppState { /// subscribes — including pre-registration sockets, which the /// per-connection signal held by the registry cannot reach. shutdown: ShutdownTx, + /// When the current drain expires, latched together with `shutdown`. + /// Every budget in the drain is a remainder of this one instant. + shutdown_deadline: Mutex>, } impl AppState { @@ -76,6 +79,7 @@ impl AppState { rpc_id: AtomicU64::new(1), conns: Mutex::new(BTreeMap::new()), shutdown: shutdown_signal(), + shutdown_deadline: Mutex::new(None), } } @@ -91,11 +95,30 @@ impl AppState { /// Latch the shutdown reason: every live connection task closes, and a /// connection upgraded mid-drain sees the latched reason at subscribe - /// time. + /// time. The drain deadline is latched with it, so every connection + /// budgets against the SAME instant rather than each starting a fresh + /// `shutdown_drain_secs` when it happens to leave its main loop — a + /// connection that noticed the signal late must not be granted a longer + /// life than the drain it is part of. fn begin_shutdown(&self, reason: &'static str) { + let deadline = Instant::now() + Duration::from_secs(self.cfg.shutdown_drain_secs); + *self.shutdown_deadline.lock() = Some(deadline); self.shutdown.send_replace(Some(reason)); } + /// The instant the current drain expires. `None` until the CP begins + /// shutting down. + fn shutdown_deadline(&self) -> Option { + *self.shutdown_deadline.lock() + } + + /// Whether the CP has begun draining. Distinguishes "this peer stopped + /// reading" from "the CP is leaving" when a frame cannot be delivered — + /// the two deserve different explanations. + pub fn shutting_down(&self) -> bool { + self.shutdown_deadline.lock().is_some() + } + /// Resolves when no live connection permits remain — every connection /// task (registered or still pre-registration) has exited. Polled: /// drain progress is a wait on an ever-shrinking set, not a precise @@ -642,13 +665,18 @@ async fn handle_connection( // path the outbound queue is flushed FIRST: the synthesized terminals // and cancels that `graceful_shutdown` queued must reach the wire ahead // of the close frame. `rx.close()` refuses any further enqueue, so the - // loop ends on its own once the backlog is flushed; the drain deadline - // bounds the backlog and `write_timeout` bounds each stalled write. + // loop ends on its own once the backlog is flushed. + // + // Every budget below is a remainder of ONE instant — the deadline latched + // when the CP began draining — not a fresh `shutdown_drain_secs` per + // connection and not a fresh one per write. That is what makes the + // documented bound true: `shutdown_drain_secs` covers the whole drain, + // the flush and the final close write included. if let Some(reason) = close.process_reason() { rx.close(); - let drain_deadline = Instant::now() + Duration::from_secs(state.cfg.shutdown_drain_secs); + let deadline = state.shutdown_deadline().unwrap_or_else(Instant::now); while let Some(text) = rx.recv().await { - let budget = drain_deadline.saturating_duration_since(Instant::now()); + let budget = deadline.saturating_duration_since(Instant::now()); if budget.is_zero() { break; } @@ -665,12 +693,27 @@ async fn handle_connection( reason, "closing connection: control plane shutting down" ); - // Bounded like every other write: a peer that has stopped reading + // Bounded like every other write — a peer that has stopped reading // must not be able to hold teardown (and its quota slot) by refusing - // to accept the close frame. The write is NOT raced against the - // close watch: a second signal latching mid-write (a lease sweep + // to accept the close frame — but by what is LEFT of the drain, not by + // a fresh timeout: a stalled peer cannot push the process's exit past + // the budget the operator configured. The write is NOT raced against + // the close watch: a second signal latching mid-write (a lease sweep // landing mid-drain) must not eat the promised close frame. - send_final_frame(&mut sink, restart_close(reason), write_timeout).await; + let remaining = deadline.saturating_duration_since(Instant::now()); + if !remaining.is_zero() { + send_final_frame( + &mut sink, + restart_close(reason), + remaining.min(write_timeout), + ) + .await; + } else { + warn!( + handle, + "drain deadline elapsed before this connection's close frame could be written" + ); + } } else if let Some(reason) = cp_close_reason { info!( agent = %format!("{}/{}", identity.namespace, identity.name), @@ -1021,9 +1064,23 @@ fn handle_frame(state: &Arc, handle: u64, text: &str) -> Option) { } /// Graceful shutdown (SIGTERM/SIGHUP/SIGINT): resolve every in-flight -/// delegation with synthesized terminal frames, signal every connection to -/// close, then wait — bounded by `shutdown_drain_secs` — for the connection -/// tasks to flush their queues and exit. +/// delegation with synthesized terminal frames, latch the process-wide close +/// signal, then wait — bounded by the SAME `shutdown_drain_secs` deadline the +/// connection tasks are working against — for the connection tasks to flush +/// their queues and exit. /// -/// This runs AFTER the HTTP listener has stopped accepting (axum's graceful -/// shutdown already resolved); the WebSocket connection tasks are detached -/// spawns that keep running until told to close. The caller stops the -/// sweeper and exits after this returns — the sweeper deliberately outlives -/// the drain so a mid-drain lease/deadline expiry still synthesizes +/// Runs concurrently with the listener's own graceful shutdown, NOT after it: +/// hyper's wait ends when the last in-flight HTTP request finishes, and a peer +/// that opened a socket and then said nothing would pin that wait for as long +/// as the orchestrator lets the process live — taking the drain down with it. +/// The caller stops the sweeper once this returns; the sweeper deliberately +/// outlives the drain so a mid-drain lease/deadline expiry still synthesizes /// normally. /// /// Order is deliberate: terminals are synthesized BEFORE the close signal @@ -1188,25 +1247,43 @@ pub async fn run_sweeper(state: Arc) { pub async fn graceful_shutdown(state: &Arc) { info!("draining connections and resolving in-flight delegations"); let mut next = || state.next_rpc_id(); + let mut undeliverable = 0usize; for (inst, frame) in state .router .fail_all(&state.registry, &state.events, REASON_SHUTDOWN, &mut next) { - let _ = inst.tx.try_send(frame); + if inst.tx.try_send(frame).is_err() { + // The peer's bounded queue is already full (or was closed by a + // connection leaving the drain): the terminal cannot be delivered + // and the observer event is already emitted, so the two surfaces + // disagree — unavoidably, since a full queue is exactly a peer + // that stopped reading. Named at warn level because the drain is + // the last chance this delegation could ever be resolved for that + // initiator, and its only remaining outcome is deadline + // reconciliation on the client. + undeliverable += 1; + warn!( + agent = %inst.logical_id(), + handle = inst.handle, + "shutdown terminal could not be queued — the peer's outbound \ + queue refused it; the delegation will be reconciled by the \ + initiator's deadline" + ); + } } + // The deadline is latched with the close signal below, and covers + // everything after it. + let drain = Duration::from_secs(state.cfg.shutdown_drain_secs); state.begin_shutdown(REASON_SHUTDOWN); - match tokio::time::timeout( - Duration::from_secs(state.cfg.shutdown_drain_secs), - state.connections_drained(), - ) - .await - { - Ok(()) => info!("all connections drained"), - Err(_) => warn!( + let outcome = tokio::time::timeout(drain, state.connections_drained()).await; + if outcome.is_err() { + warn!( drain_secs = state.cfg.shutdown_drain_secs, - "shutdown drain budget elapsed — exiting with connections still open" - ), + undeliverable, "shutdown drain budget elapsed — exiting with connections still open" + ); + } else { + info!(undeliverable, "all connections drained"); } } @@ -1445,6 +1522,79 @@ mod tests { .expect("token on the ack") } + #[tokio::test] + async fn shutdown_survives_an_initiator_whose_queue_is_already_full() { + // The one outcome the drain cannot repair: an initiator whose bounded + // queue is full when the terminal is synthesized. The frame cannot be + // delivered — that is what "bounded" means — so the two surfaces + // disagree by construction: the observer event is emitted, the wire + // frame is refused. What must NOT happen is a panic, a delegation row + // left behind, a leaked serving reservation, or a drain that never + // ends. The refusal is logged at warn level, because the drain is the + // last chance that delegation had to be resolved for that initiator. + let state = state_with(""); + let (h_i, mut rx_i) = register_test_instance(&state, "koudu", AgentType::Primary, 4); + let (h_w, mut rx_w) = register_test_instance(&state, "worker-1", AgentType::Worker, 1); + let (_h_o, mut rx_o) = register_test_instance(&state, "lobby", AgentType::Observer, 0); + delegate_through_handler(&state, h_i, "d-1", "worker-1"); + rx_w.try_recv().expect("worker received the forward"); + while rx_o.try_recv().is_ok() {} + + // Nobody drains this queue — exactly the situation a peer that stopped + // reading creates, and the only way to make `try_send` refuse. + let filler = "x".repeat(64 * 1024); + let mut enqueued = 0; + while state + .registry + .get(h_i) + .unwrap() + .tx + .try_send(filler.clone()) + .is_ok() + { + enqueued += 1; + assert!(enqueued < 1000, "the bounded queue must refuse eventually"); + } + assert!(enqueued > 0, "the first frame must fit"); + + graceful_shutdown(&state).await; + + assert_eq!( + state.router.inflight_count(), + 0, + "the delegation must still be ended" + ); + assert_eq!( + state.registry.get(h_w).unwrap().active_sessions, + 0, + "the serving reservation must still be released" + ); + assert!( + state.shutting_down(), + "the CP must stay latched as draining afterwards" + ); + let mut terminal_queued = false; + while let Ok(frame) = rx_i.try_recv() { + if frame.contains("cp/delegate_result") { + terminal_queued = true; + } + } + assert!( + !terminal_queued, + "a full queue cannot take the terminal — the documented exception" + ); + let events: Vec = std::iter::from_fn(|| rx_o.try_recv().ok()).collect(); + assert!( + events.iter().any(|e| e.contains("delegation_completed")), + "the observer-side terminal is still emitted: that is why the two \ + surfaces disagree, and why the refusal must be logged: {events:?}" + ); + assert!( + rx_w.try_recv().unwrap().contains("cp/cancel"), + "the serving runtime's queue was healthy, so its cancel must land" + ); + } + #[test] fn registered_teardown_survives_a_panic() { // Teardown used to run only on the connection task's normal return diff --git a/crates/openab-cp/tests/cp_shutdown.rs b/crates/openab-cp/tests/cp_shutdown.rs index f9f24a025..68e9b58a6 100644 --- a/crates/openab-cp/tests/cp_shutdown.rs +++ b/crates/openab-cp/tests/cp_shutdown.rs @@ -225,6 +225,80 @@ fn is_shutdown_close(code: u16, reason: &str) -> bool { code == 1012 && reason == REASON_SHUTDOWN } +/// Open a TCP connection to the CP and start an HTTP request that never +/// finishes: the headers stop one CRLF short of complete, so the server keeps +/// the connection waiting for the rest. Nothing in the CP's own shutdown can +/// complete it — it exists to prove the shutdown does not wait on it. +async fn stall_an_http_request(url: &str) -> tokio::net::TcpStream { + use tokio::io::AsyncWriteExt; + let host = url.trim_start_matches("ws://").trim_end_matches("/cp"); + let mut stream = TcpStream::connect(host).await.expect("connecting"); + stream + .write_all(b"GET /health HTTP/1.1\r\nHost: x\r\n") + .await + .expect("sending a partial request"); + stream +} + +#[tokio::test] +async fn a_stalled_http_request_cannot_pin_the_shutdown() { + // A CP whose drain waits for the listener's own graceful shutdown is pinned + // by a peer that opened a socket and then said nothing: axum ends that wait + // when the last in-flight request finishes, so the drain would never start + // and the process would sit until the orchestrator's SIGKILL — the reset + // this whole path exists to prevent. + let mut cp = spawn_cp().await; + let mut initiator = connect_retry(&cp.url, KEY).await; + register(&mut initiator, "koudu", "i-1").await; + let _stalled = stall_an_http_request(&cp.url).await; + + send_sigterm(cp.child.id().expect("the CP is still running")); + + // The drain runs regardless of the stalled request: the live socket still + // gets its close frame. + let (code, reason) = await_close(&mut initiator).await; + assert!( + is_shutdown_close(code, &reason), + "the drain must run past a stalled request — close was {code} {reason:?}" + ); + // And the process leaves on its own instead of waiting to be killed. + let status = tokio::time::timeout(Duration::from_secs(30), cp.child.wait()) + .await + .expect("a stalled request must not pin the process past the drain budget") + .expect("reaping the CP"); + assert_eq!(status.code(), Some(0), "{status:?}"); +} + +#[tokio::test] +async fn a_second_signal_exits_immediately_while_the_drain_is_pending() { + // The operator's escape hatch has to cover the whole window between the + // first signal and the exit — including a listener that will not finish its + // own wait. Armed only after the drain, it would never fire there. + let mut cp = spawn_cp().await; + let mut initiator = connect_retry(&cp.url, KEY).await; + register(&mut initiator, "koudu", "i-1").await; + let _stalled = stall_an_http_request(&cp.url).await; + + let pid = cp.child.id().expect("the CP is still running"); + send_sigterm(pid); + // Observed close frame: the drain has started and the listener is all that + // is left keeping the process alive. + let (code, reason) = await_close(&mut initiator).await; + assert!(is_shutdown_close(code, &reason), "{code} {reason:?}"); + + send_sigterm(pid); + + let status = tokio::time::timeout(Duration::from_secs(20), cp.child.wait()) + .await + .expect("a second signal must force the exit, not wait out the budget") + .expect("reaping the CP"); + assert_eq!( + status.code(), + Some(143), + "a forced exit reports 128 + SIGTERM, not a signal death: {status:?}" + ); +} + #[tokio::test] async fn sigterm_resolves_in_flight_delegations_before_the_process_exits() { // The whole point of the issue, over a real socket and a real process: diff --git a/crates/openab-cp/tests/ws_lifecycle.rs b/crates/openab-cp/tests/ws_lifecycle.rs index e926bcea5..b3f63e2fb 100644 --- a/crates/openab-cp/tests/ws_lifecycle.rs +++ b/crates/openab-cp/tests/ws_lifecycle.rs @@ -898,6 +898,11 @@ async fn spawn_cp_process() -> (tokio::process::Child, String, std::path::PathBu let child = tokio::process::Command::new(env!("CARGO_BIN_EXE_openab-cp")) .arg("--config") .arg(&cfg_path) + // Not inherited and not piped: the child's `tracing` output would + // otherwise interleave with the harness log, and a pipe nobody drains + // can block the child mid-write. + .stdout(std::process::Stdio::null()) + .stderr(std::process::Stdio::null()) .kill_on_drop(true) .spawn() .expect("spawn openab-cp"); @@ -927,10 +932,11 @@ async fn a_second_signal_latching_mid_close_still_sends_a_close_frame() { graceful_shutdown(&state).await; let closed = wait_closed(&mut ws, Duration::from_secs(5)).await; - assert!( - matches!(closed, Some(Closed::Frame { .. })), - "a second signal mid-close must not reduce the peer's close frame \ - to a bare reset — got {closed:?}" + assert_eq!( + closed, + Some(shutdown_close()), + "a second signal mid-close must not reduce the peer's close frame to \ + a bare reset — nor downgrade the reason it is being told: got {closed:?}" ); } diff --git a/docs/adr/agent-control-plane.md b/docs/adr/agent-control-plane.md index ff1a93485..7734cef82 100644 --- a/docs/adr/agent-control-plane.md +++ b/docs/adr/agent-control-plane.md @@ -551,10 +551,15 @@ recovery semantics: synthesized `target_disconnected` terminal and serving runtimes a `cp/cancel`, both stamped `control plane shutting down`, and admissions racing the drain are refused `SATURATED` — then every connection closes - with code 1012, all inside `shutdown_drain_secs` plus at most one - `write_timeout_secs` for a stalled close write. A second shutdown signal - during the drain is the operator's "stop waiting": the process exits - immediately with the signal's conventional status (128 + signum). A + with code 1012. `shutdown_drain_secs` is the single ceiling for all of it + (the flush, the close write, and the process exit), so a peer that stops + reading cannot push the exit past it; the process exits 0 once the drain + finishes. A second shutdown signal before the process is gone is the + operator's "stop waiting" and exits immediately with the signal's + conventional status (128 + signum — 143 for SIGTERM). Nothing waits on the + HTTP listener's own graceful shutdown: the drain is spawned at the signal, + because a peer that opened a socket and then stopped sending would + otherwise pin that wait indefinitely and take the drain down with it. A *hard* restart (SIGKILL, crash, drain budget exhausted) is equivalent to every lease expiring at once *with* the connection closure that implies — except that the CP is not there to send it: the in-flight table and the diff --git a/docs/control-plane.md b/docs/control-plane.md index d1aaaeda1..6c04f3cfc 100644 --- a/docs/control-plane.md +++ b/docs/control-plane.md @@ -59,7 +59,9 @@ issue #1474). reconnect, re-authenticate, and re-register. Process shutdown (SIGTERM, SIGHUP, SIGINT — e.g. `docker stop`, ECS, k8s) uses code **1012** (service restart) with reason `control plane shutting down` instead: same - recovery, different signal — your client did nothing wrong. + recovery, different signal — your client did nothing wrong. Shutdown does + not wait for the HTTP listener to finish: the drain starts at the signal, + so a peer that connected and then stopped sending cannot delay your close. - A peer that stops reading is disconnected: any single outbound write that blocks longer than `write_timeout_secs` is treated as a dead peer, so keep draining the socket even while busy. The same rule applies to the queue @@ -126,12 +128,17 @@ issue #1474). - After a lease expires or the CP restarts, in-flight delegations are gone: initiators reconcile against their own deadlines and re-delegate. On a *graceful* stop the CP resolves them first: within `shutdown_drain_secs` - (plus at most one `write_timeout_secs` for a stalled close write) every - initiator gets a synthesized `target_disconnected` terminal (error - `control plane shutting down`) and every serving runtime a `cp/cancel`, - ahead of the 1012 close; a `cp/delegate` that reaches the CP mid-drain is - refused `SATURATED` — back off and retry after reconnecting. A SIGKILL - still leaves reconciliation to your deadlines. + (the single ceiling — the flush, the close write and the process exit all + fit inside it) every initiator gets a synthesized `target_disconnected` + terminal (error `control plane shutting down`) and every serving runtime a + `cp/cancel`, ahead of the 1012 close; a `cp/delegate` that reaches the CP + mid-drain is refused `SATURATED` — back off and retry after reconnecting. + The process then exits 0. A second SIGTERM/SIGINT/SIGHUP before it is gone + is your "stop waiting": it exits immediately with 128 + signum (143 for + SIGTERM). A SIGKILL still leaves reconciliation to your deadlines. An + initiator whose outbound queue was already full cannot receive its + synthesized terminal — the CP logs that refusal; the delegation is then + reconciled by the initiator's own deadline, like any lost result. ## Observer surface (lobby) From 0a1601ed12f720fa2d4a949aeb32553154b74213 Mon Sep 17 00:00:00 2001 From: worker-w15 Date: Thu, 1 Oct 2026 22:30:14 +0800 Subject: [PATCH 08/13] fix(cp): anchor the shutdown bound at the signal, not at startup MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Second review round. The previous commit armed `timeout(hard_bound, …)` in `main` before the signal ever arrived, and `axum::serve` is lazy — so the "bound" started ticking at startup and became the process's LIFETIME: with the default `shutdown_drain_secs = 5` the CP exited 0 about six seconds after boot, having served nobody, and no signal path ever ran. Trading "the CP dies mid-delegation" for "the CP vanishes on its own" is not a fix. The bound now applies where it belongs: the listener runs as its own task, the signal is awaited in `main` and the drain happens inline (bounded, as before), and only the listener's remaining graceful wait gets a one-second slack — a deliberate allowance, not a reasoned budget, since `/health` answers immediately and `/cp` finishes its response as soon as the handshake is done: anything still pending is a peer that connected and stopped talking. Also from the same round: - The second-signal escape hatch reused freshly registered streams, so a repeat arriving between the first signal resolving and that registration was discarded by tokio's registry — no listener, no delivery. `ShutdownSignals` now registers once and is awaited twice: the first wait consumes the signal that starts the drain, the second (in the force-quit task) waits for the next one, with no gap and no risk of the first signal being consumed twice. - The early-return close paths — a socket caught mid-upgrade by the drain, a pre-registration close, a refused register frame, the observer cap — still wrote with the full `write_timeout_secs`, so a task could outlive the drain by 30s while the docs claimed `shutdown_drain_secs` was the ceiling. They now spend what is left of the drain (`final_frame_budget`). - A drain that exhausts its budget still exits 0 (it always did); the docs now say so, since a non-zero status would read as a crash. - `spawn_cp_process` removes its temp config through an RAII guard, so a panicking assert no longer leaves one per run in /tmp. Regression test: `the_drain_budget_bounds_the_shutdown_not_the_process_lifetime` starts a CP with a one-second drain, waits three, and asserts it still registers a runtime — it fails against the previous `main`, which had already exited. Refs #1475 --- crates/openab-cp/src/main.rs | 247 +++++++++++++++---------- crates/openab-cp/src/server.rs | 35 +++- crates/openab-cp/tests/cp_shutdown.rs | 39 +++- crates/openab-cp/tests/ws_lifecycle.rs | 22 ++- docs/adr/agent-control-plane.md | 9 +- docs/control-plane.md | 6 +- 6 files changed, 239 insertions(+), 119 deletions(-) diff --git a/crates/openab-cp/src/main.rs b/crates/openab-cp/src/main.rs index a0ca2c292..8cb22dd24 100644 --- a/crates/openab-cp/src/main.rs +++ b/crates/openab-cp/src/main.rs @@ -18,46 +18,99 @@ struct Cli { config: String, } -/// Wait for a shutdown signal: SIGINT (ctrl-c) on every platform, plus -/// SIGTERM and SIGHUP on unix — the signals `docker stop`, ECS and k8s send -/// before escalating to SIGKILL. Without handlers the default disposition -/// kills the process outright: no close frames, no terminal results, and -/// runtimes see TCP resets. -async fn shutdown_signal() -> &'static str { +/// Slack the listener gets, after the drain has finished, to complete its own +/// graceful shutdown. +/// +/// The CP has no long-lived HTTP route: `/health` answers immediately and +/// `/cp` completes its response as soon as the handshake is done (the WebSocket +/// runs detached). So anything still pending once the drain ends is a peer that +/// connected and then stopped talking — not work worth waiting for. The +/// orchestrator's grace period is the real deadline; this only decides how much +/// of it the CP spends on a socket that will never answer. +const LISTENER_SHUTDOWN_SLACK: Duration = Duration::from_secs(1); + +/// The signals that mean "shut down": SIGINT (ctrl-c) on every platform, plus +/// SIGTERM and SIGHUP on unix — what `docker stop`, ECS and k8s send before +/// escalating to SIGKILL. Without handlers the default disposition kills the +/// process outright: no close frames, no terminal results, runtimes see TCP +/// resets, and in-flight delegations are left to their clients' deadlines. +/// +/// Registration is separated from waiting on purpose. tokio's signal registry +/// discards an event that has no listener registered, so a watcher built +/// *inside* the handler would leave a window in which a repeat signal — the +/// operator's "stop waiting" — is dropped on the floor. [`ShutdownSignals`] +/// registers once and can then be awaited any number of times, which is what +/// makes "signal it twice" reliable. +struct ShutdownSignals { #[cfg(unix)] - { - use tokio::signal::unix::{signal, SignalKind}; - let mut term = match signal(SignalKind::terminate()) { - Ok(s) => s, - Err(e) => { - tracing::warn!( - error = %e, - "failed to install SIGTERM handler — falling back to ctrl-c only" - ); - let _ = tokio::signal::ctrl_c().await; - return "SIGINT"; - } - }; - let mut hup = match signal(SignalKind::hangup()) { - Ok(s) => s, - Err(e) => { - tracing::warn!(error = %e, "failed to install SIGHUP handler"); - return tokio::select! { - _ = tokio::signal::ctrl_c() => "SIGINT", - _ = term.recv() => "SIGTERM", - }; - } - }; - tokio::select! { - _ = tokio::signal::ctrl_c() => "SIGINT", - _ = term.recv() => "SIGTERM", - _ = hup.recv() => "SIGHUP", + terminate: Option, + #[cfg(unix)] + hangup: Option, +} + +impl ShutdownSignals { + /// Register the signal streams. Synchronous: the listeners exist from this + /// moment on, whatever the caller awaits next. + fn register() -> Self { + #[cfg(unix)] + { + use tokio::signal::unix::{signal, SignalKind}; + // A stream that cannot be registered is simply absent: ctrl-c and + // whatever else did register still work, and the failure is logged + // rather than fatal — a CP that cannot be stopped politely should + // still run. + let terminate = signal(SignalKind::terminate()) + .inspect_err(|e| { + warn!(error = %e, "failed to install SIGTERM handler"); + }) + .ok(); + let hangup = signal(SignalKind::hangup()) + .inspect_err(|e| { + warn!(error = %e, "failed to install SIGHUP handler"); + }) + .ok(); + Self { terminate, hangup } + } + #[cfg(not(unix))] + { + Self {} } } - #[cfg(not(unix))] - { - let _ = tokio::signal::ctrl_c().await; - "SIGINT" + + /// Wait for the next shutdown signal and name it. + async fn recv(&mut self) -> &'static str { + #[cfg(unix)] + { + // `pending()` for an unregistered stream: it can never fire, and + // selecting on it keeps the arm shape uniform. + let never = || std::future::pending::<()>(); + let term = async { + match self.terminate.as_mut() { + Some(s) => { + s.recv().await; + } + None => never().await, + } + }; + let hup = async { + match self.hangup.as_mut() { + Some(s) => { + s.recv().await; + } + None => never().await, + } + }; + tokio::select! { + _ = tokio::signal::ctrl_c() => "SIGINT", + _ = term => "SIGTERM", + _ = hup => "SIGHUP", + } + } + #[cfg(not(unix))] + { + let _ = tokio::signal::ctrl_c().await; + "SIGINT" + } } } @@ -93,83 +146,73 @@ async fn main() -> Result<()> { ); let state = Arc::new(AppState::new(cfg)); - let drain_budget = Duration::from_secs(state.cfg.shutdown_drain_secs); - // The process's exit is bounded by the drain budget plus slack for the - // listener's own bookkeeping. One knob, one bound: an orchestrator - // escalates to SIGKILL at the end of its grace period, so the CP's own - // deadline has to be the shorter one. - let hard_bound = drain_budget + Duration::from_secs(1); let sweeper = tokio::spawn(run_sweeper(state.clone())); let listener = tokio::net::TcpListener::bind(&listen) .await .with_context(|| format!("binding {listen}"))?; - // The shutdown sequence is SPAWNED rather than awaited after the listener. - // Ordering it after `serve` was a real defect: axum's graceful shutdown - // ends when the last in-flight HTTP request finishes, and a peer that - // opened a socket and then said nothing (a half-sent request, a stalled - // upgrade) pins that wait for as long as the orchestrator allows the - // process to live — so the drain would never start, and the CP would be - // SIGKILLed with its delegations unresolved. Spawning it makes the drain - // begin the instant the signal lands, whatever hyper is doing. + // Registered once, before anything is awaited: see `ShutdownSignals`. The + // first wait consumes the signal that starts the drain; the second (in the + // task below) is the operator's "stop waiting". + let mut signals = ShutdownSignals::register(); let (stop_accepting_tx, stop_accepting_rx) = tokio::sync::oneshot::channel::<()>(); - let (drained_tx, drained_rx) = tokio::sync::oneshot::channel::<()>(); - tokio::spawn({ + + // The listener runs as its own task. Awaiting `axum::serve` here and then + // draining would be the obvious shape and it is wrong: axum's graceful + // shutdown ends when the last in-flight HTTP request finishes, so a peer + // that opened a socket and then said nothing (a half-sent request, a stalled + // upgrade) would pin that wait — and with it the drain — for as long as the + // orchestrator lets the process live. Serving from a task lets the signal + // be handled here, immediately, whatever hyper is doing. + let mut serve_done = tokio::spawn({ let state = Arc::clone(&state); async move { - let signal = shutdown_signal().await; - info!(signal, "shutdown signal received — draining connections"); - // A second signal at ANY point before the process exits is the - // operator's "stop waiting": exit immediately with the signal's - // conventional status (128 + signum). Armed before the drain and - // deliberately left armed — the drain and the listener's wait are - // separate waits, and the escape hatch has to cover both. The task - // ends with the process if it never fires. - tokio::spawn(async { - let signal = shutdown_signal().await; - warn!(signal, "second shutdown signal — forcing immediate exit"); - std::process::exit(signal_exit_code(signal)); - }); - // Nothing new may enter a CP that is leaving. - let _ = stop_accepting_tx.send(()); - graceful_shutdown(&state).await; - let _ = drained_tx.send(()); + axum::serve(listener, app(state)) + .with_graceful_shutdown(async move { + let _ = stop_accepting_rx.await; + }) + .await } }); - // The listener stops accepting as soon as the signal lands, and the wait - // for it is bounded by the same hard bound — hyper's bookkeeping must not - // be able to outlive the drain it is supposed to follow. - let serving = - axum::serve(listener, app(Arc::clone(&state))).with_graceful_shutdown(async move { - let _ = stop_accepting_rx.await; - }); - let finished = tokio::time::timeout(hard_bound, async { - tokio::join!( - async { - // The drain is the shutdown: the synthesized terminals and the - // close frames get their chance to reach the wire here. - let _ = drained_rx.await; - }, - async { - match tokio::time::timeout(hard_bound, serving).await { - Ok(Ok(())) => {} - Ok(Err(e)) => warn!(error = %e, "listener stopped with an error"), - Err(_) => warn!( - bound_secs = hard_bound.as_secs(), - "the listener's graceful wait outlived the shutdown bound — dropping it" - ), - } - }, - ); - }) - .await; - if finished.is_err() { - warn!( - bound_secs = hard_bound.as_secs(), - "shutdown did not finish within its bound — exiting with the drain unfinished" - ); + let signal = signals.recv().await; + info!(signal, "shutdown signal received — draining connections"); + + // A second signal at ANY point before the process exits is "stop waiting": + // exit immediately with the signal's conventional status (128 + signum). + // Armed from the SAME already-registered streams — the first `recv` above + // consumed its event, so this one waits for the next — which closes the + // window in which a repeat would be dropped by an unregistered listener. + // Deliberately left armed: the drain and the listener's wait are separate + // waits and the escape hatch must cover both. The task ends with the + // process if it never fires. + tokio::spawn(async move { + let signal = signals.recv().await; + warn!(signal, "second shutdown signal — forcing immediate exit"); + std::process::exit(signal_exit_code(signal)); + }); + + // Nothing new may enter a CP that is leaving. Then drain: bounded by + // `shutdown_drain_secs`, which resolves the in-flight delegations, lets + // every connection flush and close, and returns. + let _ = stop_accepting_tx.send(()); + graceful_shutdown(&state).await; + + // The drain is done; give the listener a bounded moment to finish its own + // bookkeeping, then stop regardless. The bound starts HERE, not at boot: + // a timer armed at startup would become the process's lifetime. + match tokio::time::timeout(LISTENER_SHUTDOWN_SLACK, &mut serve_done).await { + Ok(Ok(Ok(()))) => {} + Ok(Ok(Err(e))) => warn!(error = %e, "listener stopped with an error"), + Ok(Err(join)) => warn!("listener task ended abnormally: {join:?}"), + Err(_) => { + warn!( + slack_secs = LISTENER_SHUTDOWN_SLACK.as_secs(), + "the listener's graceful wait outlived the shutdown slack — dropping it" + ); + serve_done.abort(); + } } // Only now: the sweeper exists to keep registrations and delegations diff --git a/crates/openab-cp/src/server.rs b/crates/openab-cp/src/server.rs index 58bef9dce..6b6724e64 100644 --- a/crates/openab-cp/src/server.rs +++ b/crates/openab-cp/src/server.rs @@ -322,6 +322,24 @@ async fn send_final_frame( let _ = tokio::time::timeout(write_timeout, sink.send(msg)).await; } +/// Budget for the last frame of a connection that is ending: the write timeout, +/// clipped to what is left of the drain once the CP is shutting down. +/// +/// The main-loop teardown already spends the remaining drain directly; this +/// covers the paths that never reach the main loop — a socket caught mid-upgrade +/// by the drain, a pre-registration close, a refused register frame. Without it +/// those writes could still hold their task for a full `write_timeout_secs` +/// (30s by default) after a 5s drain expired, so `shutdown_drain_secs` would +/// not really be the ceiling the docs promise. +fn final_frame_budget(state: &AppState, write_timeout: Duration) -> Duration { + match state.shutdown_deadline() { + Some(deadline) => deadline + .saturating_duration_since(Instant::now()) + .min(write_timeout), + None => write_timeout, + } +} + async fn ws_handler( State(state): State>, headers: HeaderMap, @@ -399,7 +417,12 @@ async fn handle_connection( // only reports transitions that postdate the subscription, so the // already-latched reason is checked once explicitly. if let Some(reason) = close.process_reason() { - send_final_frame(&mut sink, restart_close(reason), write_timeout).await; + send_final_frame( + &mut sink, + restart_close(reason), + final_frame_budget(&state, write_timeout), + ) + .await; return; } @@ -429,7 +452,7 @@ async fn handle_connection( Some(reason) => restart_close(reason), None => policy_close(close.reason().unwrap_or(REASON_SHUTDOWN)), }; - send_final_frame(&mut sink, frame, write_timeout).await; + send_final_frame(&mut sink, frame, final_frame_budget(&state, write_timeout)).await; return; } } { @@ -447,7 +470,7 @@ async fn handle_connection( send_final_frame( &mut sink, policy_close(REASON_REGISTER_TIMEOUT), - write_timeout, + final_frame_budget(&state, write_timeout), ) .await; return; @@ -460,7 +483,7 @@ async fn handle_connection( send_final_frame( &mut sink, Message::Text(serde_json::to_string(&resp).expect("serializable").into()), - write_timeout, + final_frame_budget(&state, write_timeout), ) .await; return; @@ -523,7 +546,7 @@ async fn handle_connection( send_final_frame( &mut sink, Message::Text(serde_json::to_string(&resp).expect("serializable").into()), - write_timeout, + final_frame_budget(&state, write_timeout), ) .await; return; @@ -579,7 +602,7 @@ async fn handle_connection( None => close.reason().map(policy_close), }; if let Some(frame) = frame { - send_final_frame(&mut sink, frame, write_timeout).await; + send_final_frame(&mut sink, frame, final_frame_budget(&state, write_timeout)).await; } } // `_registered` runs teardown on the way out. diff --git a/crates/openab-cp/tests/cp_shutdown.rs b/crates/openab-cp/tests/cp_shutdown.rs index 68e9b58a6..96474ea47 100644 --- a/crates/openab-cp/tests/cp_shutdown.rs +++ b/crates/openab-cp/tests/cp_shutdown.rs @@ -52,7 +52,7 @@ fn free_port() -> u16 { port } -fn config_file(port: u16) -> PathBuf { +fn config_file(port: u16, extra: &str) -> PathBuf { let seq = SEQ.fetch_add(1, Ordering::Relaxed); let path = std::env::temp_dir().join(format!( "openab-cp-sigterm-{}-{seq}.toml", @@ -62,7 +62,7 @@ fn config_file(port: u16) -> PathBuf { r#" listen = "127.0.0.1:{port}" register_timeout_secs = 30 -shutdown_drain_secs = 5 +{extra} [[agents]] key = "{KEY}" @@ -98,8 +98,13 @@ impl Drop for Cp { /// Start the real `openab-cp` binary on a free loopback port. async fn spawn_cp() -> Cp { + spawn_cp_with("shutdown_drain_secs = 5").await +} + +/// Start the real binary with extra config lines injected. +async fn spawn_cp_with(extra: &str) -> Cp { let port = free_port(); - let config = config_file(port); + let config = config_file(port, extra); let child = Command::new(env!("CARGO_BIN_EXE_openab-cp")) .arg("--config") .arg(&config) @@ -225,6 +230,34 @@ fn is_shutdown_close(code: u16, reason: &str) -> bool { code == 1012 && reason == REASON_SHUTDOWN } +#[tokio::test] +async fn the_drain_budget_bounds_the_shutdown_not_the_process_lifetime() { + // `shutdown_drain_secs` is a shutdown budget, not a lifetime: a CP with a + // one-second drain must still be serving three seconds later, when no + // signal has been delivered. A bound armed at startup instead of at the + // signal would end the process on its own — turning a stop-mid-delegation + // bug into "the CP vanishes", with no signal ever involved. + let mut cp = spawn_cp_with("shutdown_drain_secs = 1").await; + + tokio::time::sleep(Duration::from_secs(3)).await; + let mut ws = connect_retry(&cp.url, KEY).await; + let ack = register(&mut ws, "koudu", "i-1").await; + assert_eq!( + ack["result"]["protocol_version"], 1, + "the CP must still be serving after 3x its drain budget with no signal: {ack}" + ); + + // And the budget still applies to the shutdown itself. + send_sigterm(cp.child.id().expect("the CP is still running")); + let (code, reason) = await_close(&mut ws).await; + assert!(is_shutdown_close(code, &reason), "{code} {reason:?}"); + let status = tokio::time::timeout(Duration::from_secs(20), cp.child.wait()) + .await + .expect("the CP must exit after SIGTERM") + .expect("reaping the CP"); + assert_eq!(status.code(), Some(0), "{status:?}"); +} + /// Open a TCP connection to the CP and start an HTTP request that never /// finishes: the headers stop one CRLF short of complete, so the server keeps /// the connection waiting for the rest. Nothing in the CP's own shutdown can diff --git a/crates/openab-cp/tests/ws_lifecycle.rs b/crates/openab-cp/tests/ws_lifecycle.rs index b3f63e2fb..2c6c3747f 100644 --- a/crates/openab-cp/tests/ws_lifecycle.rs +++ b/crates/openab-cp/tests/ws_lifecycle.rs @@ -874,7 +874,18 @@ async fn shutdown_closes_sockets_that_never_registered() { /// SIGTERM before SIGKILL, and the binary used to die by default disposition /// without ever running a shutdown path. #[cfg(unix)] -async fn spawn_cp_process() -> (tokio::process::Child, String, std::path::PathBuf) { +/// The temp config a spawned CP reads, removed when the guard drops — including +/// on a panicking assert, which otherwise leaves one file per failed run in +/// `/tmp`. +struct CfgGuard(std::path::PathBuf); + +impl Drop for CfgGuard { + fn drop(&mut self) { + let _ = std::fs::remove_file(&self.0); + } +} + +async fn spawn_cp_process() -> (tokio::process::Child, String, CfgGuard) { // Grab an ephemeral port, release it, hand it to the child — the usual // tiny race, absorbed by `connect_retry`. let port = std::net::TcpListener::bind("127.0.0.1:0") @@ -906,7 +917,11 @@ async fn spawn_cp_process() -> (tokio::process::Child, String, std::path::PathBu .kill_on_drop(true) .spawn() .expect("spawn openab-cp"); - (child, format!("ws://127.0.0.1:{port}/cp"), cfg_path) + ( + child, + format!("ws://127.0.0.1:{port}/cp"), + CfgGuard(cfg_path), + ) } #[tokio::test] @@ -947,7 +962,7 @@ async fn sigterm_and_sighup_run_the_drain_and_exit_cleanly() { // SIGHUP shares the path. Each must produce the graceful close on live // sockets and a clean exit code — never the old default-disposition kill. for sig in ["TERM", "HUP"] { - let (mut child, url, cfg_path) = spawn_cp_process().await; + let (mut child, url, _cfg) = spawn_cp_process().await; let mut ws = connect_retry(&url).await; assert_eq!( register(&mut ws, "i-1").await["result"]["protocol_version"], @@ -975,6 +990,5 @@ async fn sigterm_and_sighup_run_the_drain_and_exit_cleanly() { status.success(), "SIG{sig}: the CP must exit cleanly, got {status}" ); - let _ = std::fs::remove_file(&cfg_path); } } diff --git a/docs/adr/agent-control-plane.md b/docs/adr/agent-control-plane.md index 7734cef82..42a9a31a2 100644 --- a/docs/adr/agent-control-plane.md +++ b/docs/adr/agent-control-plane.md @@ -554,12 +554,17 @@ recovery semantics: with code 1012. `shutdown_drain_secs` is the single ceiling for all of it (the flush, the close write, and the process exit), so a peer that stops reading cannot push the exit past it; the process exits 0 once the drain - finishes. A second shutdown signal before the process is gone is the + finishes. A drain that exhausts its budget also exits 0 — with a warning + naming it — because a non-zero status would be indistinguishable from a + crash to most supervisors, and the frames that did not make it are lost + either way. A second shutdown signal before the process is gone is the operator's "stop waiting" and exits immediately with the signal's conventional status (128 + signum — 143 for SIGTERM). Nothing waits on the HTTP listener's own graceful shutdown: the drain is spawned at the signal, because a peer that opened a socket and then stopped sending would - otherwise pin that wait indefinitely and take the drain down with it. A + otherwise pin that wait indefinitely and take the drain down with it. That + wait is bounded by one second of slack *after* the drain ends (and is + dropped if it overruns), so no client can delay the exit. A *hard* restart (SIGKILL, crash, drain budget exhausted) is equivalent to every lease expiring at once *with* the connection closure that implies — except that the CP is not there to send it: the in-flight table and the diff --git a/docs/control-plane.md b/docs/control-plane.md index 6c04f3cfc..d973fe6d4 100644 --- a/docs/control-plane.md +++ b/docs/control-plane.md @@ -133,8 +133,10 @@ issue #1474). terminal (error `control plane shutting down`) and every serving runtime a `cp/cancel`, ahead of the 1012 close; a `cp/delegate` that reaches the CP mid-drain is refused `SATURATED` — back off and retry after reconnecting. - The process then exits 0. A second SIGTERM/SIGINT/SIGHUP before it is gone - is your "stop waiting": it exits immediately with 128 + signum (143 for + The process then exits 0 — including when the budget ran out with + connections still open, which is logged as a warning (a non-zero status + would read as a crash). A second SIGTERM/SIGINT/SIGHUP before it is gone is + your "stop waiting": it exits immediately with 128 + signum (143 for SIGTERM). A SIGKILL still leaves reconciliation to your deadlines. An initiator whose outbound queue was already full cannot receive its synthesized terminal — the CP logs that refusal; the delegation is then From 06ef6dc0fec055a9cf3e9ddc747feb7c1d7f35c1 Mon Sep 17 00:00:00 2001 From: worker-w15 Date: Thu, 1 Oct 2026 22:52:12 +0800 Subject: [PATCH 09/13] fix(cp): close the last signal gap and make the process tests diagnosable Third review round (verdict: pass, nits addressed). - SIGINT now goes through the same boot-registered `Signal` machinery as SIGTERM and SIGHUP instead of `tokio::signal::ctrl_c()`, which re-registers its handler on every call: a second ctrl-c arriving between one wait returning and the next being created was dropped by tokio's registry, so the escape hatch could miss exactly the signal it exists for. ctrl-c stays as an arm for hosts where the SIGINT stream cannot be registered. - `CfgGuard` and `spawn_cp_process` in the lifecycle tests are now both `#[cfg(unix)]` again (a doc comment had captured the attribute), so a Windows `--all-targets` check compiles. - `Cp::connect` reports the child's status when its retries run out. A CP that exited during startup and a CP that is slow are indistinguishable from the client side, and the port-grab-then-release race is exactly where a one-in-a-hundred flake lands; the status makes the next occurrence actionable instead of guesswork. - The mid-drain registration race waits for the shutdown latch before racing it, since a `cp/register` that completes before the latch is legitimately acked and asserting otherwise would be asserting something untrue. - Docs: the listener's one second of slack after the drain is now stated, and the ADR no longer describes a spawn ordering the code stopped using. Refs #1475 --- crates/openab-cp/src/main.rs | 56 ++++++++++++++------------ crates/openab-cp/tests/cp_shutdown.rs | 49 +++++++++++++--------- crates/openab-cp/tests/ws_lifecycle.rs | 14 ++++--- docs/adr/agent-control-plane.md | 12 +++--- docs/control-plane.md | 12 +++--- 5 files changed, 82 insertions(+), 61 deletions(-) diff --git a/crates/openab-cp/src/main.rs b/crates/openab-cp/src/main.rs index 8cb22dd24..6cc30e098 100644 --- a/crates/openab-cp/src/main.rs +++ b/crates/openab-cp/src/main.rs @@ -42,6 +42,8 @@ const LISTENER_SHUTDOWN_SLACK: Duration = Duration::from_secs(1); /// registers once and can then be awaited any number of times, which is what /// makes "signal it twice" reliable. struct ShutdownSignals { + #[cfg(unix)] + interrupt: Option, #[cfg(unix)] terminate: Option, #[cfg(unix)] @@ -55,21 +57,28 @@ impl ShutdownSignals { #[cfg(unix)] { use tokio::signal::unix::{signal, SignalKind}; - // A stream that cannot be registered is simply absent: ctrl-c and - // whatever else did register still work, and the failure is logged - // rather than fatal — a CP that cannot be stopped politely should - // still run. + // SIGINT goes through the same `Signal` machinery as SIGTERM and + // SIGHUP rather than through `tokio::signal::ctrl_c()`, which + // re-registers its handler on every call: a repeat arriving between + // one wait returning and the next being created would be dropped by + // the registry — the gap the double-signal escape hatch cannot + // have. A stream that cannot be registered is simply absent: the + // others still work, and the failure is logged rather than fatal, + // because a CP that cannot be stopped politely should still run. + let interrupt = signal(SignalKind::interrupt()) + .inspect_err(|e| warn!(error = %e, "failed to install SIGINT handler")) + .ok(); let terminate = signal(SignalKind::terminate()) - .inspect_err(|e| { - warn!(error = %e, "failed to install SIGTERM handler"); - }) + .inspect_err(|e| warn!(error = %e, "failed to install SIGTERM handler")) .ok(); let hangup = signal(SignalKind::hangup()) - .inspect_err(|e| { - warn!(error = %e, "failed to install SIGHUP handler"); - }) + .inspect_err(|e| warn!(error = %e, "failed to install SIGHUP handler")) .ok(); - Self { terminate, hangup } + Self { + interrupt, + terminate, + hangup, + } } #[cfg(not(unix))] { @@ -83,27 +92,22 @@ impl ShutdownSignals { { // `pending()` for an unregistered stream: it can never fire, and // selecting on it keeps the arm shape uniform. - let never = || std::future::pending::<()>(); - let term = async { - match self.terminate.as_mut() { + async fn arm(signal: Option<&mut tokio::signal::unix::Signal>) { + match signal { Some(s) => { s.recv().await; } - None => never().await, + None => std::future::pending::<()>().await, } - }; - let hup = async { - match self.hangup.as_mut() { - Some(s) => { - s.recv().await; - } - None => never().await, - } - }; + } tokio::select! { + // Kept as the fallback for a host where the SIGINT stream + // could not be registered; whichever arm fires first names the + // same signal. _ = tokio::signal::ctrl_c() => "SIGINT", - _ = term => "SIGTERM", - _ = hup => "SIGHUP", + _ = arm(self.interrupt.as_mut()) => "SIGINT", + _ = arm(self.terminate.as_mut()) => "SIGTERM", + _ = arm(self.hangup.as_mut()) => "SIGHUP", } } #[cfg(not(unix))] diff --git a/crates/openab-cp/tests/cp_shutdown.rs b/crates/openab-cp/tests/cp_shutdown.rs index 96474ea47..e5336472b 100644 --- a/crates/openab-cp/tests/cp_shutdown.rs +++ b/crates/openab-cp/tests/cp_shutdown.rs @@ -44,7 +44,7 @@ type Ws = WebSocketStream>; static SEQ: AtomicU32 = AtomicU32::new(0); /// A loopback port that was free a moment ago. The CP is started immediately -/// after, and [`connect_retry`] absorbs the race either way. +/// after, and [`Cp::connect`] absorbs the race either way. fn free_port() -> u16 { let probe = std::net::TcpListener::bind("127.0.0.1:0").expect("a free loopback port"); let port = probe.local_addr().expect("bound address").port(); @@ -96,6 +96,26 @@ impl Drop for Cp { } } +impl Cp { + /// Connect as `key`, retrying while the CP is still binding its port. + /// + /// The retry exists because the port is grabbed and released before the + /// child binds it. When the retries run out the child's fate is part of the + /// panic: from the client side "the CP exited during startup" and "the CP + /// is slow" look identical, and the exit status is the only thing that + /// tells them apart. + async fn connect(&mut self, key: &str) -> Ws { + for _ in 0..100 { + if let Ok(ws) = connect_as(&self.url, key).await { + return ws; + } + tokio::time::sleep(Duration::from_millis(50)).await; + } + let status = self.child.try_wait().expect("querying the child"); + panic!("the CP never accepted a connection (child status: {status:?})"); + } +} + /// Start the real `openab-cp` binary on a free loopback port. async fn spawn_cp() -> Cp { spawn_cp_with("shutdown_drain_secs = 5").await @@ -132,17 +152,6 @@ async fn connect_as(url: &str, key: &str) -> Result Ws { - for _ in 0..100 { - if let Ok(ws) = connect_as(url, key).await { - return ws; - } - tokio::time::sleep(Duration::from_millis(50)).await; - } - panic!("the CP never accepted a connection"); -} - fn register_frame(name: &str, instance_id: &str) -> String { let ty = if name == "koudu" { "primary" } else { "worker" }; serde_json::json!({ @@ -240,7 +249,7 @@ async fn the_drain_budget_bounds_the_shutdown_not_the_process_lifetime() { let mut cp = spawn_cp_with("shutdown_drain_secs = 1").await; tokio::time::sleep(Duration::from_secs(3)).await; - let mut ws = connect_retry(&cp.url, KEY).await; + let mut ws = cp.connect(KEY).await; let ack = register(&mut ws, "koudu", "i-1").await; assert_eq!( ack["result"]["protocol_version"], 1, @@ -281,7 +290,7 @@ async fn a_stalled_http_request_cannot_pin_the_shutdown() { // and the process would sit until the orchestrator's SIGKILL — the reset // this whole path exists to prevent. let mut cp = spawn_cp().await; - let mut initiator = connect_retry(&cp.url, KEY).await; + let mut initiator = cp.connect(KEY).await; register(&mut initiator, "koudu", "i-1").await; let _stalled = stall_an_http_request(&cp.url).await; @@ -308,7 +317,7 @@ async fn a_second_signal_exits_immediately_while_the_drain_is_pending() { // first signal and the exit — including a listener that will not finish its // own wait. Armed only after the drain, it would never fire there. let mut cp = spawn_cp().await; - let mut initiator = connect_retry(&cp.url, KEY).await; + let mut initiator = cp.connect(KEY).await; register(&mut initiator, "koudu", "i-1").await; let _stalled = stall_an_http_request(&cp.url).await; @@ -340,12 +349,12 @@ async fn sigterm_resolves_in_flight_delegations_before_the_process_exits() { // orchestrator's grace period, before it escalates to SIGKILL. let mut cp = spawn_cp().await; - let mut initiator = connect_retry(&cp.url, KEY).await; + let mut initiator = cp.connect(KEY).await; assert_eq!( register(&mut initiator, "koudu", "i-1").await["result"]["protocol_version"], 1 ); - let mut worker = connect_retry(&cp.url, KEY_WORKER).await; + let mut worker = cp.connect(KEY_WORKER).await; assert_eq!( register(&mut worker, "worker-1", "i-w").await["result"]["protocol_version"], 1 @@ -440,10 +449,14 @@ async fn a_connection_opened_during_the_drain_never_registers() { // runs: a runtime that reconnects into a closing CP must be turned away // (or closed politely), never handed a socket that dies mid-registration. let mut cp = spawn_cp().await; - let mut existing = connect_retry(&cp.url, KEY).await; + let mut existing = cp.connect(KEY).await; register(&mut existing, "koudu", "i-1").await; send_sigterm(cp.child.id().expect("the CP is still running")); + // Let the shutdown actually latch: a `cp/register` that completes BEFORE + // it is legitimately acked, so racing the latch itself would assert + // something untrue. What matters is that nothing registers after it. + tokio::time::sleep(Duration::from_millis(250)).await; // Race the drain with fresh upgrades. Every outcome is acceptable except // a successful registration: a refused connection, or a socket closed with diff --git a/crates/openab-cp/tests/ws_lifecycle.rs b/crates/openab-cp/tests/ws_lifecycle.rs index 2c6c3747f..bfe3e751e 100644 --- a/crates/openab-cp/tests/ws_lifecycle.rs +++ b/crates/openab-cp/tests/ws_lifecycle.rs @@ -868,23 +868,25 @@ async fn shutdown_closes_sockets_that_never_registered() { ); } -/// Spawn a real `openab-cp` process bound to an ephemeral port; returns the -/// child, its WS URL, and the config path (for cleanup). This is the only -/// proof the OS-level path works: `docker stop`, ECS, and k8s all deliver -/// SIGTERM before SIGKILL, and the binary used to die by default disposition -/// without ever running a shutdown path. -#[cfg(unix)] /// The temp config a spawned CP reads, removed when the guard drops — including /// on a panicking assert, which otherwise leaves one file per failed run in /// `/tmp`. +#[cfg(unix)] struct CfgGuard(std::path::PathBuf); +#[cfg(unix)] impl Drop for CfgGuard { fn drop(&mut self) { let _ = std::fs::remove_file(&self.0); } } +/// Spawn a real `openab-cp` process bound to an ephemeral port; returns the +/// child, its WS URL, and a guard that cleans up its config. This is the only +/// proof the OS-level path works: `docker stop`, ECS, and k8s all deliver +/// SIGTERM before SIGKILL, and the binary used to die by default disposition +/// without ever running a shutdown path. +#[cfg(unix)] async fn spawn_cp_process() -> (tokio::process::Child, String, CfgGuard) { // Grab an ephemeral port, release it, hand it to the child — the usual // tiny race, absorbed by `connect_retry`. diff --git a/docs/adr/agent-control-plane.md b/docs/adr/agent-control-plane.md index 42a9a31a2..cd02b870e 100644 --- a/docs/adr/agent-control-plane.md +++ b/docs/adr/agent-control-plane.md @@ -559,12 +559,12 @@ recovery semantics: crash to most supervisors, and the frames that did not make it are lost either way. A second shutdown signal before the process is gone is the operator's "stop waiting" and exits immediately with the signal's - conventional status (128 + signum — 143 for SIGTERM). Nothing waits on the - HTTP listener's own graceful shutdown: the drain is spawned at the signal, - because a peer that opened a socket and then stopped sending would - otherwise pin that wait indefinitely and take the drain down with it. That - wait is bounded by one second of slack *after* the drain ends (and is - dropped if it overruns), so no client can delay the exit. A + conventional status (128 + signum — 143 for SIGTERM). The listener runs as + its own task so nothing waits on the HTTP listener's own graceful shutdown: + a peer that opened a socket and then stopped sending would otherwise pin + that wait indefinitely and take the drain down with it. After the drain + ends, that wait gets one second of slack and is dropped if it overruns, so + no client can delay the exit. A *hard* restart (SIGKILL, crash, drain budget exhausted) is equivalent to every lease expiring at once *with* the connection closure that implies — except that the CP is not there to send it: the in-flight table and the diff --git a/docs/control-plane.md b/docs/control-plane.md index d973fe6d4..1dd2b4696 100644 --- a/docs/control-plane.md +++ b/docs/control-plane.md @@ -128,11 +128,13 @@ issue #1474). - After a lease expires or the CP restarts, in-flight delegations are gone: initiators reconcile against their own deadlines and re-delegate. On a *graceful* stop the CP resolves them first: within `shutdown_drain_secs` - (the single ceiling — the flush, the close write and the process exit all - fit inside it) every initiator gets a synthesized `target_disconnected` - terminal (error `control plane shutting down`) and every serving runtime a - `cp/cancel`, ahead of the 1012 close; a `cp/delegate` that reaches the CP - mid-drain is refused `SATURATED` — back off and retry after reconnecting. + (the single ceiling — the flush, the close write and the drain itself all + fit inside it; the process then allows the HTTP listener one further second + of slack to finish its own bookkeeping) every initiator gets a synthesized + `target_disconnected` terminal (error `control plane shutting down`) and + every serving runtime a `cp/cancel`, ahead of the 1012 close; a + `cp/delegate` that reaches the CP mid-drain is refused `SATURATED` — back + off and retry after reconnecting. The process then exits 0 — including when the budget ran out with connections still open, which is logged as a warning (a non-zero status would read as a crash). A second SIGTERM/SIGINT/SIGHUP before it is gone is From d2a00dbd7b770f7bdb4af26cf6fab84df1ea780e Mon Sep 17 00:00:00 2001 From: worker-w15 Date: Thu, 1 Oct 2026 23:18:10 +0800 Subject: [PATCH 10/13] =?UTF-8?q?fix(cp):=20one=20listener=20per=20signal?= =?UTF-8?q?=20=E2=80=94=20a=20single=20ctrl-c=20must=20not=20force-exit?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review round four caught a regression I introduced in the previous commit, and measured it: on unix `tokio::signal::ctrl_c()` IS `signal(SignalKind::interrupt())`, so selecting on a registered SIGINT stream AND ctrl-c handed the SAME ctrl-c to two listeners. The force-quit task then read the first signal's notification as a second one and called exit(130) immediately — one ctrl-c skipped the drain in 8 runs out of 10, which is the TCP-reset failure this whole issue exists to remove. Exactly one listener per signal now: the registered stream when it exists, ctrl-c only where it could not be registered (where it is also the only thing that can work). Verified against the built binary: a single SIGINT now drains and exits 0 with the 1012 close in 8/8 runs, and a double SIGINT still force-exits 130 in 3/3. Also from the same round: the mid-drain registration test slept 250 ms before racing the latch, which made it vacuous — the CP exits well inside that window, so every attempt took the "connection refused" arm and the test asserted nothing. Replaced with a deterministic in-process test that latches the drain directly (`AppState::begin_shutdown`, now public) and asserts a socket that upgrades afterwards gets the shutdown close instead of a registration. Refs #1475 --- crates/openab-cp/src/main.rs | 43 ++++++++++++++------ crates/openab-cp/src/server.rs | 2 +- crates/openab-cp/tests/cp_shutdown.rs | 56 -------------------------- crates/openab-cp/tests/ws_lifecycle.rs | 22 ++++++++++ 4 files changed, 53 insertions(+), 70 deletions(-) diff --git a/crates/openab-cp/src/main.rs b/crates/openab-cp/src/main.rs index 6cc30e098..2f0401443 100644 --- a/crates/openab-cp/src/main.rs +++ b/crates/openab-cp/src/main.rs @@ -86,28 +86,45 @@ impl ShutdownSignals { } } + /// `pending()` for an unregistered stream: it can never fire, and + /// selecting on it keeps the arm shape uniform. + #[cfg(unix)] + async fn arm(signal: Option<&mut tokio::signal::unix::Signal>) { + match signal { + Some(s) => { + s.recv().await; + } + None => std::future::pending::<()>().await, + } + } + /// Wait for the next shutdown signal and name it. async fn recv(&mut self) -> &'static str { #[cfg(unix)] { - // `pending()` for an unregistered stream: it can never fire, and - // selecting on it keeps the arm shape uniform. - async fn arm(signal: Option<&mut tokio::signal::unix::Signal>) { - match signal { + // Exactly ONE listener per signal, always. On unix + // `tokio::signal::ctrl_c()` IS `signal(SignalKind::interrupt())`, so + // selecting on both a registered SIGINT stream and ctrl-c would + // hand the SAME ctrl-c to two listeners — and the force-quit task + // would read the first signal's notification as a second one and + // exit immediately, skipping the drain entirely. (Measured: one + // ctrl-c force-exited 8 times out of 10.) ctrl-c is therefore the + // arm only where the stream could not be registered, where it is + // also the only thing that can work. + let interrupt = async { + match self.interrupt.as_mut() { Some(s) => { s.recv().await; } - None => std::future::pending::<()>().await, + None => { + let _ = tokio::signal::ctrl_c().await; + } } - } + }; tokio::select! { - // Kept as the fallback for a host where the SIGINT stream - // could not be registered; whichever arm fires first names the - // same signal. - _ = tokio::signal::ctrl_c() => "SIGINT", - _ = arm(self.interrupt.as_mut()) => "SIGINT", - _ = arm(self.terminate.as_mut()) => "SIGTERM", - _ = arm(self.hangup.as_mut()) => "SIGHUP", + _ = interrupt => "SIGINT", + _ = Self::arm(self.terminate.as_mut()) => "SIGTERM", + _ = Self::arm(self.hangup.as_mut()) => "SIGHUP", } } #[cfg(not(unix))] diff --git a/crates/openab-cp/src/server.rs b/crates/openab-cp/src/server.rs index 6b6724e64..5795a3b57 100644 --- a/crates/openab-cp/src/server.rs +++ b/crates/openab-cp/src/server.rs @@ -100,7 +100,7 @@ impl AppState { /// `shutdown_drain_secs` when it happens to leave its main loop — a /// connection that noticed the signal late must not be granted a longer /// life than the drain it is part of. - fn begin_shutdown(&self, reason: &'static str) { + pub fn begin_shutdown(&self, reason: &'static str) { let deadline = Instant::now() + Duration::from_secs(self.cfg.shutdown_drain_secs); *self.shutdown_deadline.lock() = Some(deadline); self.shutdown.send_replace(Some(reason)); diff --git a/crates/openab-cp/tests/cp_shutdown.rs b/crates/openab-cp/tests/cp_shutdown.rs index e5336472b..c50bca22c 100644 --- a/crates/openab-cp/tests/cp_shutdown.rs +++ b/crates/openab-cp/tests/cp_shutdown.rs @@ -442,59 +442,3 @@ async fn sigterm_resolves_in_flight_delegations_before_the_process_exits() { "a drained shutdown exits 0, not by signal: {status:?}" ); } - -#[tokio::test] -async fn a_connection_opened_during_the_drain_never_registers() { - // The drain only means something if the CP stops admitting work while it - // runs: a runtime that reconnects into a closing CP must be turned away - // (or closed politely), never handed a socket that dies mid-registration. - let mut cp = spawn_cp().await; - let mut existing = cp.connect(KEY).await; - register(&mut existing, "koudu", "i-1").await; - - send_sigterm(cp.child.id().expect("the CP is still running")); - // Let the shutdown actually latch: a `cp/register` that completes BEFORE - // it is legitimately acked, so racing the latch itself would assert - // something untrue. What matters is that nothing registers after it. - tokio::time::sleep(Duration::from_millis(250)).await; - - // Race the drain with fresh upgrades. Every outcome is acceptable except - // a successful registration: a refused connection, or a socket closed with - // the shutdown reason before it can register. - let mut registered = false; - for _ in 0..40 { - match connect_as(&cp.url, KEY_WORKER).await { - // The listener is already gone, or the CP is not far enough into - // its shutdown for the refusal to be observable yet. - Err(_) => {} - Ok(mut ws) => match tokio::time::timeout(Duration::from_secs(5), ws.next()).await { - Ok(Some(Ok(Message::Close(Some(cf))))) => assert!( - is_shutdown_close(cf.code.into(), cf.reason.as_ref()), - "a socket closed during the drain must carry the shutdown reason: {cf:?}" - ), - Ok(Some(Ok(Message::Text(t)))) => { - let v: serde_json::Value = serde_json::from_str(&t).unwrap(); - assert!( - v.get("result").is_none(), - "a socket admitted during the drain must not be acked: {v}" - ); - registered = true; - } - Ok(Some(Ok(_))) => continue, - Ok(None) | Ok(Some(Err(_))) => {} - Err(_) => panic!("a socket accepted during the drain never ended"), - }, - } - tokio::time::sleep(Duration::from_millis(50)).await; - } - assert!( - !registered, - "the CP must not have acked a registration once it started draining" - ); - - let status = tokio::time::timeout(Duration::from_secs(30), cp.child.wait()) - .await - .expect("the CP must exit after SIGTERM") - .expect("reaping the CP"); - assert_eq!(status.code(), Some(0), "{status:?}"); -} diff --git a/crates/openab-cp/tests/ws_lifecycle.rs b/crates/openab-cp/tests/ws_lifecycle.rs index bfe3e751e..4738c0c0c 100644 --- a/crates/openab-cp/tests/ws_lifecycle.rs +++ b/crates/openab-cp/tests/ws_lifecycle.rs @@ -957,6 +957,28 @@ async fn a_second_signal_latching_mid_close_still_sends_a_close_frame() { ); } +#[tokio::test] +async fn a_socket_upgraded_after_the_drain_latched_is_never_registered() { + // The drain's first act is to latch the process-wide reason; a socket that + // completes its upgrade after that must be closed with the reason instead + // of registering into a CP that is leaving. Deterministic here on purpose: + // racing a real process's drain would only prove that the process was + // usually already gone. + let (state, url) = spawn_cp(cfg("register_timeout_secs = 30")).await; + state.begin_shutdown(REASON_SHUTDOWN); + + let mut ws = connect(&url) + .await + .expect("the upgrade itself still completes"); + let closed = wait_closed(&mut ws, Duration::from_secs(5)).await; + assert_eq!( + closed, + Some(shutdown_close()), + "a socket arriving after the latch must get the shutdown close, never a \ + registration: got {closed:?}" + ); +} + #[cfg(unix)] #[tokio::test] async fn sigterm_and_sighup_run_the_drain_and_exit_cleanly() { From c641f6a3e72812e1ed36b154060533bdb67d2449 Mon Sep 17 00:00:00 2001 From: worker-w15 Date: Thu, 1 Oct 2026 23:48:15 +0800 Subject: [PATCH 11/13] docs(cp): state the listener slack where every shutdown doc is read MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Final review round (verdict: pass, nits closed). `shutdown_drain_secs` bounds the flush, the close write and the drain; after it, the CP allows the HTTP listener one further second for its own bookkeeping before dropping it. `docs/control-plane.md` already said so, the ADR and `cp.toml.example` did not — the example claimed the process exit itself fit inside the budget. Also: where the SIGINT stream could not be registered, `ctrl_c()` resolves immediately with an error, which `recv` was naming as a SIGINT — a shutdown signal that never arrived. It now waits instead, and says the CP will not stop on ctrl-c. Refs #1475 --- crates/openab-cp/cp.toml.example | 7 ++++--- crates/openab-cp/src/main.rs | 8 +++++++- docs/adr/agent-control-plane.md | 6 +++--- 3 files changed, 14 insertions(+), 7 deletions(-) diff --git a/crates/openab-cp/cp.toml.example b/crates/openab-cp/cp.toml.example index 82a3a7a5d..8121a0913 100644 --- a/crates/openab-cp/cp.toml.example +++ b/crates/openab-cp/cp.toml.example @@ -66,9 +66,10 @@ max_outbound_queue_bytes = 16777216 # resolves every in-flight delegation (initiators get a synthesized # target_disconnected terminal, serving runtimes a cp/cancel), then each # connection flushes its queued outbound frames and closes with code 1012 — -# and this budget is the ceiling for all of it, close write and process exit -# included, so a peer that stops reading cannot push the exit past it. Must -# be greater than 0: a zero budget would emit the observer-side terminal +# and this budget is the ceiling for all of it, close write included, so a +# peer that stops reading cannot push the exit past it; once the drain is done +# the CP allows the HTTP listener one further second for its own bookkeeping. +# Must be greater than 0: a zero budget would emit the observer-side terminal # events and then exit before any of the frames reached a wire. shutdown_drain_secs = 5 diff --git a/crates/openab-cp/src/main.rs b/crates/openab-cp/src/main.rs index 2f0401443..69bffc8cb 100644 --- a/crates/openab-cp/src/main.rs +++ b/crates/openab-cp/src/main.rs @@ -117,7 +117,13 @@ impl ShutdownSignals { s.recv().await; } None => { - let _ = tokio::signal::ctrl_c().await; + // A failed registration resolves immediately with an + // error — reporting that as a SIGINT would look like a + // shutdown signal that never arrived. Wait instead. + if let Err(e) = tokio::signal::ctrl_c().await { + warn!(error = %e, "cannot wait for SIGINT — the CP will not stop on ctrl-c"); + std::future::pending::<()>().await; + } } } }; diff --git a/docs/adr/agent-control-plane.md b/docs/adr/agent-control-plane.md index cd02b870e..d18d509a7 100644 --- a/docs/adr/agent-control-plane.md +++ b/docs/adr/agent-control-plane.md @@ -552,9 +552,9 @@ recovery semantics: `cp/cancel`, both stamped `control plane shutting down`, and admissions racing the drain are refused `SATURATED` — then every connection closes with code 1012. `shutdown_drain_secs` is the single ceiling for all of it - (the flush, the close write, and the process exit), so a peer that stops - reading cannot push the exit past it; the process exits 0 once the drain - finishes. A drain that exhausts its budget also exits 0 — with a warning + (the flush, the close write, and the drain), so a peer that stops reading + cannot push it past its budget; the process exits 0 once the drain finishes + (plus the listener slack below). A drain that exhausts its budget also exits 0 — with a warning naming it — because a non-zero status would be indistinguishable from a crash to most supervisors, and the frames that did not make it are lost either way. A second shutdown signal before the process is gone is the From 8620c2a7faa937d3f0667b1f50927b8a0b2ab12f Mon Sep 17 00:00:00 2001 From: worker-w15 Date: Fri, 2 Oct 2026 00:00:10 +0800 Subject: [PATCH 12/13] docs(cp): the field's own doc must state the listener slack too MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `shutdown_drain_secs` bounds the flush, the close write and the drain; the process then allows the HTTP listener one further second for its own graceful wait. The ADR, `docs/control-plane.md` and `cp.toml.example` now say so — the rustdoc on the field itself did not, and a reader of `config.rs` would conclude the exit is hard-bounded at the budget. Also: the run-on ADR clause that referred forward to "the listener slack below", and the warning text for a failed ctrl-c registration, which claimed the CP "will not stop on ctrl-c" — the OS default still kills it; only the graceful path is gone. Refs #1475 --- crates/openab-cp/src/config.rs | 5 +++++ crates/openab-cp/src/main.rs | 6 +++++- docs/adr/agent-control-plane.md | 5 +++-- 3 files changed, 13 insertions(+), 3 deletions(-) diff --git a/crates/openab-cp/src/config.rs b/crates/openab-cp/src/config.rs index 964b37fc3..ee14b526c 100644 --- a/crates/openab-cp/src/config.rs +++ b/crates/openab-cp/src/config.rs @@ -104,6 +104,11 @@ pub struct CpConfig { /// SIGKILL when ITS grace period ends, and a CP still waiting then gets /// the TCP reset this whole path exists to avoid). /// + /// It is not the process's whole exit budget: after the drain, the CP + /// allows the HTTP listener a further second (its own graceful wait) before + /// dropping it, so a peer that connected and stopped talking costs at most + /// that much. + /// /// Must be greater than 0. A zero budget would exit before any /// synthesized terminal or close frame reached a wire, while still having /// emitted the observer-side `delegation_completed` events — the two diff --git a/crates/openab-cp/src/main.rs b/crates/openab-cp/src/main.rs index 69bffc8cb..1f6462cf9 100644 --- a/crates/openab-cp/src/main.rs +++ b/crates/openab-cp/src/main.rs @@ -121,7 +121,11 @@ impl ShutdownSignals { // error — reporting that as a SIGINT would look like a // shutdown signal that never arrived. Wait instead. if let Err(e) = tokio::signal::ctrl_c().await { - warn!(error = %e, "cannot wait for SIGINT — the CP will not stop on ctrl-c"); + warn!( + error = %e, + "cannot wait for SIGINT — ctrl-c will no longer stop the CP \ + gracefully (the OS default still kills it)" + ); std::future::pending::<()>().await; } } diff --git a/docs/adr/agent-control-plane.md b/docs/adr/agent-control-plane.md index d18d509a7..e7ba5adde 100644 --- a/docs/adr/agent-control-plane.md +++ b/docs/adr/agent-control-plane.md @@ -553,8 +553,9 @@ recovery semantics: racing the drain are refused `SATURATED` — then every connection closes with code 1012. `shutdown_drain_secs` is the single ceiling for all of it (the flush, the close write, and the drain), so a peer that stops reading - cannot push it past its budget; the process exits 0 once the drain finishes - (plus the listener slack below). A drain that exhausts its budget also exits 0 — with a warning + cannot push it past its budget. The process exits 0 once the drain + finishes, plus up to one further second for the listener's own graceful wait + (below). A drain that exhausts its budget also exits 0 — with a warning naming it — because a non-zero status would be indistinguishable from a crash to most supervisors, and the frames that did not make it are lost either way. A second shutdown signal before the process is gone is the From b67103654934af874b5ef11d06440f5a5af9a1cf Mon Sep 17 00:00:00 2001 From: worker-w15 Date: Fri, 2 Oct 2026 00:06:24 +0800 Subject: [PATCH 13/13] docs(cp): drop the self-contradicting "and the process exit" clause MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The previous commit appended the listener-slack paragraph to this rustdoc but left the earlier sentence claiming the budget bounds "the flush, the close write and the process exit" — one paragraph contradicting the next, and the last remaining place in the repo asserting a bound the code does not keep. The field bounds the drain; the exit allows the listener one further second after it, and an operator sizing a grace period needs to know that. Refs #1475 --- crates/openab-cp/src/config.rs | 20 ++++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/crates/openab-cp/src/config.rs b/crates/openab-cp/src/config.rs index ee14b526c..b2a04e6cb 100644 --- a/crates/openab-cp/src/config.rs +++ b/crates/openab-cp/src/config.rs @@ -92,22 +92,22 @@ pub struct CpConfig { pub write_timeout_secs: u64, /// How long the CP waits for live connections to drain after a shutdown - /// signal (SIGTERM/SIGHUP/SIGINT) before the process exits. + /// signal (SIGTERM/SIGHUP/SIGINT). /// /// On shutdown every in-flight delegation is resolved first — the /// initiator gets a synthesized `target_disconnected` terminal and the /// serving runtime a `cp/cancel` — then each connection flushes its /// queued outbound frames and sends its close frame. This budget is the - /// ceiling for ALL of it: the flush, the close write and the process exit - /// are remainders of one deadline, so a peer that stops reading cannot - /// push the exit past it (that is the point — an orchestrator escalates to - /// SIGKILL when ITS grace period ends, and a CP still waiting then gets - /// the TCP reset this whole path exists to avoid). + /// ceiling for all of THAT: the flush and the close write are remainders + /// of one deadline, so a peer that stops reading cannot push the drain past + /// it (that is the point — an orchestrator escalates to SIGKILL when ITS + /// grace period ends, and a CP still waiting then gets the TCP reset this + /// whole path exists to avoid). /// - /// It is not the process's whole exit budget: after the drain, the CP - /// allows the HTTP listener a further second (its own graceful wait) before - /// dropping it, so a peer that connected and stopped talking costs at most - /// that much. + /// It is not the process's whole exit budget: once the drain is done the CP + /// allows the HTTP listener one further second for its own graceful wait + /// before dropping it, so size your orchestrator's grace period with that + /// second of headroom. /// /// Must be greater than 0. A zero budget would exit before any /// synthesized terminal or close frame reached a wire, while still having