feat(acp-nats): wire cancel/prompt/session handlers to Bridge infrastructure

yordis · yordis · commit 99c1e6ad6b23 · 2026-03-14T07:31:23.000-04:00
- Wire cancel handler to mark sessions and resolve pending prompt waiters
- Add pre-flight cancellation check in prompt handler
- Replace inline publish_session_ready with bridge.spawn_session_ready()
  in new_session and load_session handlers
- Remove #[allow(dead_code)] from Bridge infrastructure now that handlers use it

Signed-off-by: Yordis Prieto &lt;yordis.prieto@gmail.com&gt;
diff --git a/rsworkspace/crates/acp-nats/src/agent/cancel.rs b/rsworkspace/crates/acp-nats/src/agent/cancel.rs
@@ -1,7 +1,9 @@
 use super::Bridge;
 use crate::nats::{self, FlushClient, PublishClient, RequestClient, agent};
 use crate::session_id::AcpSessionId;
-use agent_client_protocol::{CancelNotification, Error, ErrorCode, Result};
+use agent_client_protocol::{
+    CancelNotification, Error, ErrorCode, PromptResponse, Result, StopReason,
+};
 use tracing::{info, instrument, warn};
 use trogon_std::time::GetElapsed;
 
@@ -57,6 +59,16 @@ pub async fn handle<N: RequestClient + PublishClient + FlushClient, C: GetElapse
             .record_error("cancel", "cancel_publish_failed");
     }
 
+    bridge
+        .cancelled_sessions
+        .mark_cancelled(args.session_id.clone(), &bridge.clock);
+    bridge
+        .pending_session_prompt_responses
+        .cancel_waiter_for_session(
+            &args.session_id,
+            Ok(PromptResponse::new(StopReason::Cancelled)),
+        );
+
     bridge.metrics.record_request(
         "cancel",
         bridge.clock.elapsed(start).as_secs_f64(),
diff --git a/rsworkspace/crates/acp-nats/src/agent/load_session.rs b/rsworkspace/crates/acp-nats/src/agent/load_session.rs
@@ -1,13 +1,7 @@
 use super::Bridge;
-use crate::acp_prefix::AcpPrefix;
-use crate::config::SESSION_READY_DELAY;
 use crate::error::AGENT_UNAVAILABLE;
-use crate::nats::{
-    self, ExtSessionReady, FlushClient, FlushPolicy, PublishClient, PublishOptions, RequestClient,
-    RetryPolicy, agent,
-};
+use crate::nats::{self, FlushClient, PublishClient, RequestClient, agent};
 use crate::session_id::AcpSessionId;
-use crate::telemetry::metrics::Metrics;
 use agent_client_protocol::{Error, ErrorCode, LoadSessionRequest, LoadSessionResponse, Result};
 use tracing::{info, instrument, warn};
 use trogon_nats::NatsError;
@@ -88,14 +82,7 @@ pub async fn handle<N: RequestClient + PublishClient + FlushClient, C: GetElapse
     .map_err(map_load_session_error);
 
     if result.is_ok() {
-        let nats = bridge.nats.clone();
-        let prefix = bridge.config.acp_prefix.clone();
-        let session_id = session_id.clone();
-        let metrics = bridge.metrics.clone();
-        // TODO: track the JoinHandle so we can drain in-flight publishes on graceful shutdown.
-        tokio::spawn(async move {
-            publish_session_ready(&nats, &prefix, &session_id, &metrics).await;
-        });
+        bridge.spawn_session_ready(&args.session_id);
     }
 
     bridge.metrics.record_request(
@@ -107,38 +94,6 @@ pub async fn handle<N: RequestClient + PublishClient + FlushClient, C: GetElapse
     result
 }
 
-async fn publish_session_ready<N: PublishClient + FlushClient>(
-    nats: &N,
-    prefix: &AcpPrefix,
-    session_id: &AcpSessionId,
-    metrics: &Metrics,
-) {
-    tokio::time::sleep(SESSION_READY_DELAY).await;
-
-    let subject = agent::ext_session_ready(prefix.as_str(), session_id.as_str());
-    info!(session_id = %session_id, subject = %subject, "Publishing session.ready");
-
-    let message = ExtSessionReady::new(agent_client_protocol::SessionId::from(
-        session_id.to_string(),
-    ));
-
-    let options = PublishOptions::builder()
-        .publish_retry_policy(RetryPolicy::standard())
-        .flush_policy(FlushPolicy::standard())
-        .build();
-
-    if let Err(e) = nats::publish(nats, &subject, &message, options).await {
-        warn!(
-            error = %e,
-            session_id = %session_id,
-            "Failed to publish session.ready"
-        );
-        metrics.record_error("session_ready", "session_ready_publish_failed");
-    } else {
-        info!(session_id = %session_id, "Published session.ready");
-    }
-}
-
 #[cfg(test)]
 mod tests {
     use super::{Bridge, map_load_session_error};
diff --git a/rsworkspace/crates/acp-nats/src/agent/mod.rs b/rsworkspace/crates/acp-nats/src/agent/mod.rs
@@ -30,18 +30,14 @@ use std::time::Duration;
 use tracing::{info, warn};
 use trogon_std::time::GetElapsed;
 
-#[allow(dead_code)]
 const CANCELLED_SESSION_TTL: Duration = Duration::from_secs(300);
-#[allow(dead_code)]
 const CLEANUP_EVERY: usize = 16;
 
-#[allow(dead_code)]
 pub(crate) struct CancelledSessions<I: Copy> {
     map: Mutex<HashMap<SessionId, I>>,
     cleanup_counter: std::sync::atomic::AtomicUsize,
 }
 
-#[allow(dead_code)]
 impl<I: Copy> CancelledSessions<I> {
     pub fn new() -> Self {
         Self {
@@ -85,7 +81,6 @@ pub struct Bridge<N: RequestClient + PublishClient + FlushClient, C: GetElapsed>
     pub(crate) nats: N,
     pub(crate) clock: C,
     pub(crate) metrics: Metrics,
-    #[allow(dead_code)]
     pub(crate) cancelled_sessions: CancelledSessions<C::Instant>,
     pub(crate) pending_session_prompt_responses: PendingSessionPromptResponseWaiters<C::Instant>,
     pub(crate) prompt_slot_counter: PromptSlotCounter,
@@ -112,7 +107,6 @@ impl<N: RequestClient + PublishClient + FlushClient, C: GetElapsed> Bridge<N, C>
         &self.nats
     }
 
-    #[allow(dead_code)]
     pub(crate) fn register_session_ready_task(&self, task: tokio::task::JoinHandle<()>) {
         let mut tasks = self.session_ready_publish_tasks.lock().unwrap();
         tasks.retain(|task| !task.is_finished());
@@ -134,7 +128,6 @@ impl<N: RequestClient + PublishClient + FlushClient, C: GetElapsed> Bridge<N, C>
         }
     }
 
-    #[allow(dead_code)]
     pub(crate) fn spawn_session_ready(&self, session_id: &SessionId) {
         let nats_clone = self.nats.clone();
         let prefix = self.config.acp_prefix().to_string();
@@ -224,63 +217,3 @@ mod send_sync_tests {
         assert_send_sync::<Bridge<AdvancedMockNatsClient, SystemClock>>();
     }
 }
-
-#[cfg(test)]
-mod cancelled_sessions_tests {
-    use super::*;
-    use agent_client_protocol::SessionId;
-    use trogon_std::time::MockClock;
-
-    fn session(id: &str) -> SessionId {
-        SessionId::new(id.to_string())
-    }
-
-    #[test]
-    fn mark_and_take_within_ttl() {
-        let clock = MockClock::new();
-        let cs = CancelledSessions::new();
-        cs.mark_cancelled(session("s1"), &clock);
-        assert!(cs.take_if_cancelled(&session("s1"), &clock).is_some());
-    }
-
-    #[test]
-    fn take_removes_entry() {
-        let clock = MockClock::new();
-        let cs = CancelledSessions::new();
-        cs.mark_cancelled(session("s1"), &clock);
-        cs.take_if_cancelled(&session("s1"), &clock);
-        assert!(cs.take_if_cancelled(&session("s1"), &clock).is_none());
-    }
-
-    #[test]
-    fn take_returns_none_for_unknown_session() {
-        let clock = MockClock::new();
-        let cs = CancelledSessions::new();
-        assert!(cs.take_if_cancelled(&session("nope"), &clock).is_none());
-    }
-
-    #[test]
-    fn expired_entry_returns_none() {
-        let clock = MockClock::new();
-        let cs = CancelledSessions::new();
-        cs.mark_cancelled(session("s1"), &clock);
-        clock.advance(CANCELLED_SESSION_TTL + Duration::from_secs(1));
-        assert!(cs.take_if_cancelled(&session("s1"), &clock).is_none());
-    }
-
-    #[test]
-    fn cleanup_evicts_expired_entries() {
-        let clock = MockClock::new();
-        let cs = CancelledSessions::new();
-
-        cs.mark_cancelled(session("old"), &clock);
-        clock.advance(CANCELLED_SESSION_TTL + Duration::from_secs(1));
-
-        for i in 0..CLEANUP_EVERY {
-            cs.mark_cancelled(session(&format!("s{i}")), &clock);
-        }
-
-        let map = cs.map.lock().unwrap();
-        assert!(!map.contains_key(&session("old")));
-    }
-}
diff --git a/rsworkspace/crates/acp-nats/src/agent/new_session.rs b/rsworkspace/crates/acp-nats/src/agent/new_session.rs
@@ -1,14 +1,7 @@
 use super::Bridge;
-use crate::config::SESSION_READY_DELAY;
 use crate::error::AGENT_UNAVAILABLE;
-use crate::nats::{
-    self, ExtSessionReady, FlushClient, FlushPolicy, PublishClient, PublishOptions, RequestClient,
-    RetryPolicy, agent,
-};
-use crate::telemetry::metrics::Metrics;
-use agent_client_protocol::{
-    Error, ErrorCode, NewSessionRequest, NewSessionResponse, Result, SessionId,
-};
+use crate::nats::{self, FlushClient, PublishClient, RequestClient, agent};
+use agent_client_protocol::{Error, ErrorCode, NewSessionRequest, NewSessionResponse, Result};
 use tracing::{Span, info, instrument, warn};
 use trogon_nats::NatsError;
 use trogon_std::time::GetElapsed;
@@ -83,15 +76,7 @@ pub async fn handle<N: RequestClient + PublishClient + FlushClient, C: GetElapse
         info!(session_id = %response.session_id, "Session created");
 
         bridge.metrics.record_session_created();
-
-        let nats = bridge.nats.clone();
-        let prefix = bridge.config.acp_prefix.clone();
-        let session_id = response.session_id.clone();
-        let metrics = bridge.metrics.clone();
-        // TODO: track the JoinHandle so we can drain in-flight publishes on graceful shutdown.
-        tokio::spawn(async move {
-            publish_session_ready(&nats, prefix.as_str(), &session_id, &metrics).await;
-        });
+        bridge.spawn_session_ready(&response.session_id);
     }
 
     bridge.metrics.record_request(
@@ -103,36 +88,6 @@ pub async fn handle<N: RequestClient + PublishClient + FlushClient, C: GetElapse
     result
 }
 
-async fn publish_session_ready<N: PublishClient + FlushClient>(
-    nats: &N,
-    prefix: &str,
-    session_id: &SessionId,
-    metrics: &Metrics,
-) {
-    tokio::time::sleep(SESSION_READY_DELAY).await;
-
-    let subject = agent::ext_session_ready(prefix, &session_id.to_string());
-    info!(session_id = %session_id, subject = %subject, "Publishing session.ready");
-
-    let message = ExtSessionReady::new(session_id.clone());
-
-    let options = PublishOptions::builder()
-        .publish_retry_policy(RetryPolicy::standard())
-        .flush_policy(FlushPolicy::standard())
-        .build();
-
-    if let Err(e) = nats::publish(nats, &subject, &message, options).await {
-        warn!(
-            error = %e,
-            session_id = %session_id,
-            "Failed to publish session.ready"
-        );
-        metrics.record_error("session_ready", "session_ready_publish_failed");
-    } else {
-        info!(session_id = %session_id, "Published session.ready");
-    }
-}
-
 #[cfg(test)]
 mod tests {
     use super::{Bridge, map_new_session_error};
diff --git a/rsworkspace/crates/acp-nats/src/agent/prompt.rs b/rsworkspace/crates/acp-nats/src/agent/prompt.rs
@@ -84,6 +84,17 @@ pub async fn handle<N: RequestClient + PublishClient + FlushClient, C: GetElapse
             Err(error) => return Err(invalid_session_id_error(bridge, error)),
         };
 
+        if bridge
+            .cancelled_sessions
+            .take_if_cancelled(&args.session_id, &bridge.clock)
+            .is_some()
+        {
+            info!(session_id = %args.session_id, "Prompt pre-flight: session already cancelled");
+            return Ok(PromptResponse::new(
+                agent_client_protocol::StopReason::Cancelled,
+            ));
+        }
+
         let nats = bridge.nats();
         let subject = agent::session_prompt(bridge.config.acp_prefix(), session_id.as_str());
 
diff --git a/rsworkspace/crates/acp-nats/src/pending_prompt_waiters.rs b/rsworkspace/crates/acp-nats/src/pending_prompt_waiters.rs
@@ -299,28 +299,4 @@ mod tests {
             .expect("new waiter should receive a valid response");
         assert_eq!(response.stop_reason, StopReason::EndTurn);
     }
-
-    #[tokio::test]
-    async fn cancel_waiter_for_session_resolves_pending() {
-        let waiters = PendingSessionPromptResponseWaiters::<MockInstant>::new();
-        let session_id = SessionId::from("s1");
-
-        let (rx, _guard, _token) = waiters.register_waiter(session_id.clone()).unwrap();
-        let cancelled = waiters
-            .cancel_waiter_for_session(&session_id, Ok(PromptResponse::new(StopReason::Cancelled)));
-        assert!(cancelled);
-
-        let response = rx.await.unwrap().unwrap();
-        assert_eq!(response.stop_reason, StopReason::Cancelled);
-    }
-
-    #[test]
-    fn cancel_waiter_for_session_returns_false_when_no_waiter() {
-        let waiters = PendingSessionPromptResponseWaiters::<MockInstant>::new();
-        let session_id = SessionId::from("s1");
-        assert!(!waiters.cancel_waiter_for_session(
-            &session_id,
-            Ok(PromptResponse::new(StopReason::Cancelled)),
-        ));
-    }
 }
