feat(acp-nats): add prompt handler (#20)

- Add prompt handler with session validation, cancel pre-flight, backpressure
- Expand Bridge with CancelledSessions, PendingSessionPromptResponseWaiters
- Update cancel to mark sessions cancelled and resolve pending prompt waiters
- Add prompt_timeout and max_concurrent_client_tasks to config
- Add session_prompt NATS subject

Signed-off-by: Yordis Prieto <yordis.prieto@gmail.com>

Author: yordis

rsworkspace/crates/AGENTS.md

@@ -1,6 +1,6 @@
 Prefer domain-specific value objects over primitives (e.g. `AcpPrefix` not `String`). Each type's factory must guarantee correctness at construction—invalid instances should be unrepresentable. Validate per-type, not per-aggregate: avoid validating unrelated fields together in a single constructor.
 
-Every value object lives in its own file named after the type (e.g. `acp_prefix.rs`, `ext_method_name.rs`, `session_id.rs`). Never inline a value object into a config, aggregate, or service file.
+Every value object lives in its own file named after the type (e.g. `acp_prefix.rs`, `ext_method_name.rs`, `session_id.rs`). Never inline a value object into a config, aggregate, or service file. File layout: `src/{type_snake_case}.rs`; export in `lib.rs` as `pub use {module}::{Type, TypeError}`.
 
 You must use the `test-support` feature to share test helpers between crates.
 Prefer one trait per operation over a single trait with multiple operations.

rsworkspace/crates/acp-nats/Cargo.toml

@@ -13,7 +13,7 @@ async-nats = "0.45.0"
 async-trait = "0.1.89"
 serde = { version = "1.0.228", features = ["derive"] }
 serde_json = "1.0.149"
-tokio = { version = "1.49.0", features = ["rt", "macros", "sync", "time"] }
+tokio = { version = "1.49.0", features = ["rt", "macros", "sync", "time", "test-util"] }
 tracing = "0.1.44"
 
 trogon-nats = { path = "../trogon-nats" }

rsworkspace/crates/acp-nats/src/agent/cancel.rs

@@ -1,13 +1,16 @@
 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, PromptResponse, Result, StopReason};
+use agent_client_protocol::{Error, ErrorCode};
 use tracing::{info, instrument, warn};
 use trogon_std::time::GetElapsed;
 
-/// Publishes the cancel notification to the backend via NATS (fire-and-forget).
-/// The publish failure is logged and recorded as a metric but does not propagate
-/// to the caller, so the client always receives `Ok(())`.
+/// Handles cancel notification requests.
+///
+/// Marks the session as cancelled, resolves any pending prompt waiters, and publishes
+/// the cancellation to the backend. The backend publish is fire-and-forget - if it fails,
+/// the error is logged and recorded in metrics, but the method still returns `Ok(())`.
 #[instrument(
     name = "acp.session.cancel",
     skip(bridge, args),
@@ -25,15 +28,22 @@ pub async fn handle<N: RequestClient + PublishClient + FlushClient, C: GetElapse
         bridge
             .metrics
             .record_request("cancel", bridge.clock.elapsed(start).as_secs_f64(), false);
-        bridge
-            .metrics
-            .record_error("session_validate", "invalid_session_id");
+        bridge.metrics.record_error("cancel", "invalid_session_id");
         Error::new(
             ErrorCode::InvalidParams.into(),
             format!("Invalid session ID: {}", e),
         )
     })?;
 
+    bridge
+        .cancelled_sessions
+        .mark_cancelled(args.session_id.clone(), &bridge.clock);
+
+    bridge.pending_session_prompt_responses.resolve_waiter(
+        &args.session_id,
+        Ok(PromptResponse::new(StopReason::Cancelled)),
+    );
+
     let subject = agent::session_cancel(bridge.config.acp_prefix(), &args.session_id.to_string());
 
     let publish_result = nats::publish(
@@ -46,20 +56,26 @@ pub async fn handle<N: RequestClient + PublishClient + FlushClient, C: GetElapse
     )
     .await;
 
-    if let Err(error) = publish_result {
-        warn!(
-            session_id = %args.session_id,
-            error = %error,
-            "Failed to publish cancel notification to backend"
-        );
-        bridge
-            .metrics
-            .record_error("cancel", "cancel_publish_failed");
-    }
+    let publish_ok = match publish_result {
+        Ok(()) => true,
+        Err(error) => {
+            warn!(
+                session_id = %args.session_id,
+                error = %error,
+                "Failed to publish cancel notification to backend"
+            );
+            bridge
+                .metrics
+                .record_error("cancel", "cancel_publish_failed");
+            false
+        }
+    };
 
-    bridge
-        .metrics
-        .record_request("cancel", bridge.clock.elapsed(start).as_secs_f64(), true);
+    bridge.metrics.record_request(
+        "cancel",
+        bridge.clock.elapsed(start).as_secs_f64(),
+        publish_ok,
+    );
 
     Ok(())
 }
@@ -183,6 +199,23 @@ mod tests {
             .is_some()
     }
 
+    #[tokio::test]
+    async fn cancel_resolves_pending_prompt_waiter_with_cancelled() {
+        let (_mock, bridge) = mock_bridge();
+        let rx = bridge
+            .pending_session_prompt_responses
+            .register_waiter(agent_client_protocol::SessionId::from("s1"))
+            .unwrap();
+
+        let _ = bridge.cancel(CancelNotification::new("s1")).await;
+
+        let result = rx.await.unwrap().unwrap();
+        assert_eq!(
+            result.stop_reason,
+            agent_client_protocol::StopReason::Cancelled
+        );
+    }
+
     #[tokio::test]
     async fn cancel_publishes_to_correct_subject() {
         let (mock, bridge) = mock_bridge();
@@ -223,8 +256,8 @@ mod tests {
             "expected acp.request.count with method=cancel, success=false on validation failure"
         );
         assert!(
-            has_error_metric(&finished_metrics, "session_validate", "invalid_session_id"),
-            "expected acp.errors.total with operation=session_validate, reason=invalid_session_id"
+            has_error_metric(&finished_metrics, "cancel", "invalid_session_id"),
+            "expected acp.errors.total with operation=cancel, reason=invalid_session_id"
         );
         provider.shutdown().unwrap();
     }
@@ -258,8 +291,8 @@ mod tests {
             "expected acp.errors.total with operation=cancel, reason=cancel_publish_failed"
         );
         assert!(
-            has_request_metric(&finished_metrics, "cancel", true),
-            "publish failure is fire-and-forget; caller still gets Ok, so success=true"
+            has_request_metric(&finished_metrics, "cancel", false),
+            "request metric records publish outcome; success=false when publish fails"
         );
         provider.shutdown().unwrap();
     }

rsworkspace/crates/acp-nats/src/agent/cancelled_sessions.rs

@@ -0,0 +1,66 @@
+//! Pre-flight checks to avoid sending prompt requests for cancelled sessions.
+//!
+//! Entries are evicted after `CANCELLED_SESSION_TTL` to prevent unbounded growth
+//! when sessions are cancelled but never receive a subsequent prompt.
+
+use std::cell::{Cell, RefCell};
+use std::collections::HashMap;
+use std::time::Duration;
+
+use agent_client_protocol::SessionId;
+use trogon_std::time::GetElapsed;
+
+const CANCELLED_SESSION_TTL: Duration = Duration::from_secs(300); // 5 minutes
+const CLEANUP_EVERY: usize = 16;
+
+/// Pre-flight checks to avoid sending prompt requests for cancelled sessions.
+///
+/// Entries are evicted after `CANCELLED_SESSION_TTL` to prevent unbounded growth
+/// when sessions are cancelled but never receive a subsequent prompt.
+pub(crate) struct CancelledSessions<I: Copy> {
+    map: RefCell<HashMap<SessionId, I>>,
+    cleanup_counter: Cell<usize>,
+}
+
+impl<I: Copy> CancelledSessions<I> {
+    pub fn new() -> Self {
+        Self {
+            map: RefCell::new(HashMap::new()),
+            cleanup_counter: Cell::new(0),
+        }
+    }
+
+    pub fn mark_cancelled<C: GetElapsed<Instant = I>>(&self, session_id: SessionId, clock: &C) {
+        let mut map = self.map.borrow_mut();
+        map.insert(session_id, clock.now());
+        let cleanup_due_count = self.cleanup_counter.get().wrapping_add(1);
+        self.cleanup_counter.set(cleanup_due_count);
+        if cleanup_due_count.is_multiple_of(CLEANUP_EVERY) {
+            map.retain(|_, ts| clock.elapsed(*ts) < CANCELLED_SESSION_TTL);
+        }
+    }
+
+    /// Atomically checks if a session is cancelled and clears it if so.
+    ///
+    /// Returns `Some(())` if the session was cancelled (and has now been cleared),
+    /// or `None` if the session was not found or has expired.
+    pub fn take_if_cancelled<C: GetElapsed<Instant = I>>(
+        &self,
+        session_id: &SessionId,
+        clock: &C,
+    ) -> Option<()> {
+        let mut map = self.map.borrow_mut();
+
+        let is_valid = map
+            .get(session_id)
+            .is_some_and(|ts| clock.elapsed(*ts) < CANCELLED_SESSION_TTL);
+
+        if is_valid {
+            map.remove(session_id);
+            Some(())
+        } else {
+            map.remove(session_id);
+            None
+        }
+    }
+}

rsworkspace/crates/acp-nats/src/agent/mod.rs

@@ -1,42 +1,79 @@
+// Safety: `CancelledSessions` and `PendingSessionPromptResponseWaiters` use `RefCell`
+// for interior mutability. This is sound because the `Bridge` must be driven from a
+// single task (or a single-threaded `LocalSet`). The fire-and-forget publish in
+// `new_session` uses `tokio::spawn` and captures only cloned, `Send` values —
+// it never touches the `RefCell` fields.
+
 mod authenticate;
 mod cancel;
+mod cancelled_sessions;
 mod ext_method;
 mod ext_notification;
 mod initialize;
 mod load_session;
 mod new_session;
+mod pending_prompt_waiters;
+mod prompt;
 mod set_session_mode;
 
+use cancelled_sessions::CancelledSessions;
+use pending_prompt_waiters::PendingSessionPromptResponseWaiters;
+
 use crate::config::Config;
 use crate::nats::{FlushClient, PublishClient, RequestClient};
 use crate::telemetry::metrics::Metrics;
-use agent_client_protocol::ErrorCode;
 use agent_client_protocol::{
-    Agent, AuthenticateRequest, AuthenticateResponse, CancelNotification, Error, ExtNotification,
+    Agent, AuthenticateRequest, AuthenticateResponse, CancelNotification, ExtNotification,
     ExtRequest, ExtResponse, InitializeRequest, InitializeResponse, LoadSessionRequest,
     LoadSessionResponse, NewSessionRequest, NewSessionResponse, PromptRequest, PromptResponse,
     Result, SetSessionModeRequest, SetSessionModeResponse,
 };
 use opentelemetry::metrics::Meter;
+use std::cell::Cell;
+use std::marker::PhantomData;
 use trogon_std::time::GetElapsed;
 
 pub struct Bridge<N: RequestClient + PublishClient + FlushClient, C: GetElapsed> {
     pub(crate) nats: N,
     pub(crate) clock: C,
-    pub(crate) config: Config,
     pub(crate) metrics: Metrics,
+    pub(crate) cancelled_sessions: CancelledSessions<C::Instant>,
+    pub(crate) pending_session_prompt_responses: PendingSessionPromptResponseWaiters<C::Instant>,
+    pub(crate) config: Config,
+    _not_send_sync: PhantomData<std::rc::Rc<()>>,
+    agent_prompt_requests_in_flight: Cell<usize>,
 }
 
 impl<N: RequestClient + PublishClient + FlushClient, C: GetElapsed> Bridge<N, C> {
     pub fn new(nats: N, clock: C, meter: &Meter, config: Config) -> Self {
         Self {
             nats,
             clock,
-            config,
             metrics: Metrics::new(meter),
+            cancelled_sessions: CancelledSessions::new(),
+            pending_session_prompt_responses: PendingSessionPromptResponseWaiters::new(),
+            config,
+            _not_send_sync: PhantomData,
+            agent_prompt_requests_in_flight: Cell::new(0),
+        }
+    }
+
+    pub(crate) fn try_acquire_prompt_slot(&self) -> bool {
+        let max = self.config.max_concurrent_client_tasks();
+        if self.agent_prompt_requests_in_flight.get() >= max {
+            false
+        } else {
+            self.agent_prompt_requests_in_flight
+                .set(self.agent_prompt_requests_in_flight.get() + 1);
+            true
         }
     }
 
+    pub(crate) fn release_prompt_slot(&self) {
+        self.agent_prompt_requests_in_flight
+            .set(self.agent_prompt_requests_in_flight.get().saturating_sub(1));
+    }
+
     pub(crate) fn nats(&self) -> &N {
         &self.nats
     }
@@ -67,11 +104,8 @@ impl<N: RequestClient + PublishClient + FlushClient, C: GetElapsed> Agent for Br
         set_session_mode::handle(self, args).await
     }
 
-    async fn prompt(&self, _args: PromptRequest) -> Result<PromptResponse> {
-        Err(Error::new(
-            ErrorCode::InternalError.into(),
-            "not yet implemented",
-        ))
+    async fn prompt(&self, args: PromptRequest) -> Result<PromptResponse> {
+        prompt::handle(self, args).await
     }
 
     async fn cancel(&self, args: CancelNotification) -> Result<()> {
@@ -86,39 +120,3 @@ impl<N: RequestClient + PublishClient + FlushClient, C: GetElapsed> Agent for Br
         ext_notification::handle(self, args).await
     }
 }
-
-#[cfg(test)]
-mod tests {
-    use super::Bridge;
-    use crate::config::Config;
-    use agent_client_protocol::{Agent, PromptRequest};
-    use trogon_nats::AdvancedMockNatsClient;
-
-    fn mock_bridge() -> Bridge<AdvancedMockNatsClient, trogon_std::time::SystemClock> {
-        Bridge::new(
-            AdvancedMockNatsClient::new(),
-            trogon_std::time::SystemClock,
-            &opentelemetry::global::meter("acp-nats-test"),
-            Config::for_test("acp"),
-        )
-    }
-
-    #[tokio::test]
-    async fn stub_methods_return_not_implemented() {
-        let bridge = mock_bridge();
-        let msg = "not yet implemented";
-
-        assert!(
-            bridge
-                .prompt(PromptRequest::new("s1", vec![]))
-                .await
-                .is_err()
-        );
-
-        let err = bridge
-            .prompt(PromptRequest::new("s1", vec![]))
-            .await
-            .unwrap_err();
-        assert!(err.to_string().contains(msg));
-    }
-}