feat(cp): graceful shutdown — SIGTERM/SIGHUP handling and connection drain

crates/openab-core/src/acp/agentcore.rs

@@ -271,17 +271,22 @@ 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(());
                 }
             }
         }
 
         // 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<ShellHandle> {
-        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::<Value>(&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::<Value>(&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() {

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<Vec<ConfigOption>> {
         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<crate::acp::protocol::UsageReport> {
         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(&registrar));
-        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(&registrar));
+        super::install_facade_token(
+            &mut state,
+            "discord:acp_x",
+            "T_pred".into(),
+            Some(&registrar),
+        );
+        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(&registrar),
+        );
 
-        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<dyn crate::acp_mcp::SessionTokenRegistrar> = 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(&registrar));
 
-        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]

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<UsageReport> {
-    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<UsageReport> {
                     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()?,