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 40e9a67ca..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) { @@ -1250,7 +1251,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 +1273,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] @@ -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 2ea8b28c0..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. @@ -3644,14 +3739,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 +3761,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", ); @@ -4357,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 { @@ -4392,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 @@ -4400,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 @@ -4408,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 @@ -4416,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) @@ -4424,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) --- @@ -4514,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] @@ -4543,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, )); } @@ -4555,7 +4685,8 @@ mod tests { AllowUsers::Involved, false, // is_thread false, // bot_involved (irrelevant for non-thread) - false, false, + false, + false, )); } @@ -4567,7 +4698,8 @@ mod tests { AllowUsers::Involved, true, // is_thread false, // bot_involved - false, false, + false, + false, )); } @@ -4579,7 +4711,8 @@ mod tests { AllowUsers::Involved, true, // is_thread true, // bot_involved - false, false, + false, + false, )); } @@ -4629,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 4fa1ce9e1..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::>(), @@ -434,7 +438,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-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..4d63a473a 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 { @@ -377,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] 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-cp/cp.toml.example b/crates/openab-cp/cp.toml.example index f4ea04ddc..8121a0913 100644 --- a/crates/openab-cp/cp.toml.example +++ b/crates/openab-cp/cp.toml.example @@ -60,6 +60,19 @@ 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 — +# 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 + # 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..b2a04e6cb 100644 --- a/crates/openab-cp/src/config.rs +++ b/crates/openab-cp/src/config.rs @@ -90,6 +90,32 @@ 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). + /// + /// 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 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: 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 + /// 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 /// `cp/event` frames. The lobby is an audit surface, not a second /// delivery path: excerpts are truncated with a marker, never rejected. @@ -184,6 +210,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 } @@ -322,6 +351,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 7a34753c3..1f6462cf9 100644 --- a/crates/openab-cp/src/main.rs +++ b/crates/openab-cp/src/main.rs @@ -1,13 +1,14 @@ //! 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, 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 +18,143 @@ struct Cli { config: String, } +/// 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)] + interrupt: Option, + #[cfg(unix)] + 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}; + // 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")) + .ok(); + let hangup = signal(SignalKind::hangup()) + .inspect_err(|e| warn!(error = %e, "failed to install SIGHUP handler")) + .ok(); + Self { + interrupt, + terminate, + hangup, + } + } + #[cfg(not(unix))] + { + Self {} + } + } + + /// `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)] + { + // 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 => { + // 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 — ctrl-c will no longer stop the CP \ + gracefully (the OS default still kills it)" + ); + std::future::pending::<()>().await; + } + } + } + }; + tokio::select! { + _ = interrupt => "SIGINT", + _ = Self::arm(self.terminate.as_mut()) => "SIGTERM", + _ = Self::arm(self.hangup.as_mut()) => "SIGHUP", + } + } + #[cfg(not(unix))] + { + let _ = tokio::signal::ctrl_c().await; + "SIGINT" + } + } +} + +/// 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() @@ -44,13 +182,89 @@ async fn main() -> Result<()> { let listener = tokio::net::TcpListener::bind(&listen) .await .with_context(|| format!("binding {listen}"))?; - axum::serve(listener, app(state)) - .with_graceful_shutdown(async { - let _ = tokio::signal::ctrl_c().await; - info!("shutdown signal received"); - }) - .await?; + // 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::<()>(); + + // 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 { + axum::serve(listener, app(state)) + .with_graceful_shutdown(async move { + let _ = stop_accepting_rx.await; + }) + .await + } + }); + + 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 + // honest, and during a shutdown its findings would race the closes above. 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/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..5795a3b57 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,14 @@ 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, + /// 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 { @@ -70,6 +78,8 @@ impl AppState { router: Router::new(), rpc_id: AtomicU64::new(1), conns: Mutex::new(BTreeMap::new()), + shutdown: shutdown_signal(), + shutdown_deadline: Mutex::new(None), } } @@ -77,6 +87,48 @@ 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. 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. + 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)); + } + + /// 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 + /// 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 +203,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 +281,62 @@ 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())), + } +} + +/// 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; +} + +/// 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, } } @@ -249,35 +395,67 @@ 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() { + send_final_frame( + &mut sink, + restart_close(reason), + final_frame_budget(&state, write_timeout), + ) + .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)), + }; + send_final_frame(&mut sink, frame, final_frame_budget(&state, write_timeout)).await; + return; + } + } { Ok(Some(text)) => text, Ok(None) => { warn!(agent = %identity.name, "connection closed before registration"); @@ -289,11 +467,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 shutdown_rx, + final_frame_budget(&state, write_timeout), ) .await; return; @@ -303,11 +480,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 shutdown_rx, + final_frame_budget(&state, write_timeout), ) .await; return; @@ -367,11 +543,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 shutdown_rx, + final_frame_budget(&state, write_timeout), ) .await; return; @@ -413,21 +588,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 { + send_final_frame(&mut sink, frame, final_frame_budget(&state, write_timeout)).await; + } } // `_registered` runs teardown on the way out. return; @@ -451,7 +627,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 +647,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 +683,61 @@ 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. + // + // 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 deadline = state.shutdown_deadline().unwrap_or_else(Instant::now); + while let Some(text) = rx.recv().await { + let budget = 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 — 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. + 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), handle, @@ -516,14 +746,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 shutdown_rx, - ) - .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 @@ -861,9 +1087,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, 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. +/// +/// 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 +/// 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(); + let mut undeliverable = 0usize; + for (inst, frame) in + state + .router + .fail_all(&state.registry, &state.events, REASON_SHUTDOWN, &mut next) + { + 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); + let outcome = tokio::time::timeout(drain, state.connections_drained()).await; + if outcome.is_err() { + warn!( + drain_secs = state.cfg.shutdown_drain_secs, + undeliverable, "shutdown drain budget elapsed — exiting with connections still open" + ); + } else { + info!(undeliverable, "all connections drained"); + } +} + #[cfg(test)] mod tests { use super::*; @@ -1243,6 +1545,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 new file mode 100644 index 000000000..c50bca22c --- /dev/null +++ b/crates/openab-cp/tests/cp_shutdown.rs @@ -0,0 +1,444 @@ +//! 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 [`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(); + drop(probe); + port +} + +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", + std::process::id() + )); + let raw = format!( + r#" +listen = "127.0.0.1:{port}" +register_timeout_secs = 30 +{extra} + +[[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); + } +} + +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 +} + +/// 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, extra); + 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) +} + +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 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 = cp.connect(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 +/// 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 = cp.connect(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 = cp.connect(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: + // 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 = cp.connect(KEY).await; + assert_eq!( + register(&mut initiator, "koudu", "i-1").await["result"]["protocol_version"], + 1 + ); + let mut worker = cp.connect(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:?}" + ); +} diff --git a/crates/openab-cp/tests/ws_lifecycle.rs b/crates/openab-cp/tests/ws_lifecycle.rs index d2ebabe74..4738c0c0c 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,275 @@ 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:?}" + ); +} + +/// 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`. + 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) + // 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"); + ( + child, + format!("ws://127.0.0.1:{port}/cp"), + CfgGuard(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_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:?}" + ); +} + +#[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() { + // 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) = 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}" + ); + } +} 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/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..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] @@ -3309,6 +3318,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. diff --git a/docs/adr/agent-control-plane.md b/docs/adr/agent-control-plane.md index f5a5195d0..e7ba5adde 100644 --- a/docs/adr/agent-control-plane.md +++ b/docs/adr/agent-control-plane.md @@ -546,10 +546,30 @@ 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. `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 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 + operator's "stop waiting" and exits immediately with the signal's + 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 + 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..1dd2b4696 100644 --- a/docs/control-plane.md +++ b/docs/control-plane.md @@ -56,7 +56,12 @@ 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. 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 @@ -121,7 +126,23 @@ 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` + (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 + 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) 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); }