feat(acp-nats): add ext_session_prompt_response client handler

yordis · yordis · commit 7e90a3cc27d4 · 2026-03-10T05:19:59.000-04:00
diff --git a/rsworkspace/crates/acp-nats/src/client/ext_session_prompt_response.rs b/rsworkspace/crates/acp-nats/src/client/ext_session_prompt_response.rs
@@ -0,0 +1,187 @@
+use super::Bridge;
+use crate::session_id::AcpSessionId;
+use crate::nats::{FlushClient, PublishClient, RequestClient};
+use agent_client_protocol::{PromptResponse, SessionId};
+use tracing::{instrument, warn};
+use trogon_std::time::GetElapsed;
+
+#[instrument(
+    name = "acp.client.ext.session.prompt_response",
+    skip(payload, bridge),
+    fields(session_id = %session_id)
+)]
+pub async fn handle<
+    N: RequestClient + PublishClient + FlushClient,
+    C: GetElapsed,
+>(
+    session_id: &str,
+    payload: &[u8],
+    bridge: &Bridge<N, C>,
+) {
+    let Ok(validated) = AcpSessionId::new(session_id) else {
+        warn!(
+            session_id = %session_id,
+            "Invalid session_id in prompt response notification"
+        );
+        bridge
+            .metrics
+            .record_error("client.ext.session.prompt_response", "invalid_session_id");
+        return;
+    };
+
+    let session_id_typed: SessionId = validated.as_str().to_string().into();
+
+    match serde_json::from_slice::<PromptResponse>(payload) {
+        Ok(response) => {
+            bridge
+                .pending_session_prompt_responses
+                .purge_expired_timed_out_waiters(&bridge.clock);
+            let suppress_missing_waiter_warning = bridge
+                .pending_session_prompt_responses
+                .should_suppress_missing_waiter_warning(&session_id_typed, &bridge.clock);
+
+            if !bridge
+                .pending_session_prompt_responses
+                .resolve_waiter(&session_id_typed, Ok(response))
+            {
+                if !suppress_missing_waiter_warning {
+                    warn!(
+                        session_id = %session_id,
+                        "No pending prompt response waiter found for session"
+                    );
+                }
+            }
+        }
+        Err(e) => {
+            warn!(error = %e, session_id = %session_id, "Failed to parse prompt response");
+            bridge
+                .pending_session_prompt_responses
+                .purge_expired_timed_out_waiters(&bridge.clock);
+            let suppress_missing_waiter_warning = bridge
+                .pending_session_prompt_responses
+                .should_suppress_missing_waiter_warning(&session_id_typed, &bridge.clock);
+
+            if !bridge
+                .pending_session_prompt_responses
+                .resolve_waiter(&session_id_typed, Err(e.to_string()))
+            {
+                if !suppress_missing_waiter_warning {
+                    warn!(
+                        session_id = %session_id,
+                        "No pending prompt response waiter found for session"
+                    );
+                }
+            }
+            bridge
+                .metrics
+                .record_error("client.ext.session.prompt_response", "prompt_response_parse_failed");
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::agent::Bridge;
+    use crate::config::Config;
+    use agent_client_protocol::StopReason;
+    use trogon_nats::MockNatsClient;
+    use trogon_std::time::MockClock;
+
+    fn make_bridge() -> Bridge<MockNatsClient, MockClock> {
+        Bridge::new(
+            MockNatsClient::new(),
+            MockClock::new(),
+            &opentelemetry::global::meter("acp-nats-test"),
+            Config::for_test("acp"),
+        )
+    }
+
+    #[tokio::test]
+    async fn resolves_waiter() {
+        let bridge = make_bridge();
+        let session_id: SessionId = "prompt-resp-001".into();
+
+        let (rx, _guard) = bridge
+            .pending_session_prompt_responses
+            .register_waiter(session_id.clone())
+            .unwrap();
+
+        let response = PromptResponse::new(StopReason::EndTurn);
+        let payload = serde_json::to_vec(&response).unwrap();
+
+        super::handle("prompt-resp-001", &payload, &bridge).await;
+
+        let result = rx
+            .await
+            .expect("Should receive response")
+            .expect("Prompt response should not include error");
+        assert_eq!(result.stop_reason, StopReason::EndTurn);
+    }
+
+    #[tokio::test]
+    async fn no_waiter_does_not_panic() {
+        let bridge = make_bridge();
+
+        let response = PromptResponse::new(StopReason::EndTurn);
+        let payload = serde_json::to_vec(&response).unwrap();
+
+        super::handle("no-waiter-session", &payload, &bridge).await;
+    }
+
+    #[tokio::test]
+    async fn invalid_payload_removes_waiter() {
+        let bridge = make_bridge();
+        let session_id: SessionId = "bad-payload-001".into();
+
+        let (rx, _guard) = bridge
+            .pending_session_prompt_responses
+            .register_waiter(session_id.clone())
+            .unwrap();
+
+        super::handle("bad-payload-001", b"not json", &bridge).await;
+
+        let result = rx
+            .await
+            .expect("Should receive resolved parse error")
+            .err()
+            .expect("Parse failure should be forwarded to waiter");
+        assert!(!result.is_empty(), "Expected parse error to be forwarded");
+    }
+
+    #[tokio::test]
+    async fn invalid_session_id_is_rejected() {
+        let bridge = make_bridge();
+        let session_id: SessionId = "valid-session".into();
+
+        let (rx, _guard) = bridge
+            .pending_session_prompt_responses
+            .register_waiter(session_id.clone())
+            .unwrap();
+
+        let response = PromptResponse::new(StopReason::EndTurn);
+        let payload = serde_json::to_vec(&response).unwrap();
+
+        super::handle("session.with.dots", &payload, &bridge).await;
+        super::handle("session*wild", &payload, &bridge).await;
+        super::handle("session id", &payload, &bridge).await;
+
+        assert!(
+            bridge
+                .pending_session_prompt_responses
+                .has_waiter(&session_id),
+            "invalid session IDs should not resolve valid waiter",
+        );
+
+        bridge
+            .pending_session_prompt_responses
+            .remove_waiter_for_test(&session_id);
+        assert!(
+            !bridge
+                .pending_session_prompt_responses
+                .has_waiter(&session_id),
+            "waiter should be removed"
+        );
+        drop(rx);
+    }
+}
diff --git a/rsworkspace/crates/acp-nats/src/client/mod.rs b/rsworkspace/crates/acp-nats/src/client/mod.rs
@@ -1,3 +1,4 @@
+pub(crate) mod ext_session_prompt_response;
 pub(crate) mod fs_read_text_file;
 pub(crate) mod fs_write_text_file;
 pub(crate) mod request_permission;
@@ -222,6 +223,21 @@ async fn dispatch_client_method<
         ClientMethod::SessionUpdate => {
             session_update::handle(&payload, ctx.client, &parsed.session_id).await;
         }
+        ClientMethod::ExtSessionPromptResponse => {
+            if reply.is_some() {
+                warn!(
+                    session_id = %parsed.session_id,
+                    method = ?parsed.method,
+                    "Unexpected reply subject on prompt response notification"
+                );
+            }
+            ext_session_prompt_response::handle(
+                parsed.session_id.as_str(),
+                &payload,
+                ctx.bridge,
+            )
+            .await;
+        }
         ClientMethod::TerminalCreate => {
             terminal_create::handle(
                 &payload,
diff --git a/rsworkspace/crates/acp-nats/src/nats/client_method.rs b/rsworkspace/crates/acp-nats/src/nats/client_method.rs
@@ -9,6 +9,7 @@ pub enum ClientMethod {
     TerminalOutput,
     TerminalRelease,
     TerminalWaitForExit,
+    ExtSessionPromptResponse,
 }
 
 impl ClientMethod {
@@ -23,6 +24,7 @@ impl ClientMethod {
             "client.terminal.output" => Some(Self::TerminalOutput),
             "client.terminal.release" => Some(Self::TerminalRelease),
             "client.terminal.wait_for_exit" => Some(Self::TerminalWaitForExit),
+            "client.ext.session.prompt_response" => Some(Self::ExtSessionPromptResponse),
             _ => None,
         }
     }
diff --git a/rsworkspace/crates/acp-nats/src/nats/parsing.rs b/rsworkspace/crates/acp-nats/src/nats/parsing.rs
@@ -92,6 +92,14 @@ mod tests {
         assert_eq!(parsed.method, ClientMethod::TerminalWaitForExit);
     }
 
+    #[test]
+    fn test_parse_ext_session_prompt_response() {
+        let subject = "acp.sess999.client.ext.session.prompt_response";
+        let parsed = parse_client_subject(subject).unwrap();
+        assert_eq!(parsed.session_id.as_str(), "sess999");
+        assert_eq!(parsed.method, ClientMethod::ExtSessionPromptResponse);
+    }
+
     #[test]
     fn test_parse_with_custom_prefix() {
         let subject = "myapp.sess123.client.session.update";
diff --git a/rsworkspace/crates/acp-nats/src/nats/subjects.rs b/rsworkspace/crates/acp-nats/src/nats/subjects.rs
@@ -78,6 +78,14 @@ pub mod client {
         format!("{}.{}.client.terminal.wait_for_exit", prefix, session_id)
     }
 
+    pub fn ext(prefix: &str, session_id: &str, method: &str) -> String {
+        format!("{}.{}.client.ext.{}", prefix, session_id, method)
+    }
+
+    pub fn ext_session_prompt_response(prefix: &str, session_id: &str) -> String {
+        ext(prefix, session_id, "session.prompt_response")
+    }
+
     pub mod wildcards {
         pub fn all(prefix: &str) -> String {
             format!("{}.*.client.>", prefix)
@@ -162,6 +170,14 @@ mod tests {
         );
     }
 
+    #[test]
+    fn client_ext_session_prompt_response_subject() {
+        assert_eq!(
+            client::ext_session_prompt_response("acp", "s1"),
+            "acp.s1.client.ext.session.prompt_response"
+        );
+    }
+
     #[test]
     fn client_wildcards_all() {
         assert_eq!(client::wildcards::all("foo"), "foo.*.client.>");
diff --git a/rsworkspace/crates/acp-nats/src/pending_prompt_waiters.rs b/rsworkspace/crates/acp-nats/src/pending_prompt_waiters.rs
@@ -126,8 +126,16 @@ impl<I: Copy> PendingSessionPromptResponseWaiters<I> {
         });
     }
 
+    /// Returns true if a late prompt response for this session should not emit a missing-waiter warning.
+    pub(crate) fn should_suppress_missing_waiter_warning<C: GetElapsed<Instant = I>>(
+        &self,
+        session_id: &SessionId,
+        _clock: &C,
+    ) -> bool {
+        self.timed_out.lock().unwrap().contains_key(session_id)
+    }
+
     /// Delivers a backend prompt result to the currently waiting caller for `session_id`.
-    #[allow(dead_code)]
     pub fn resolve_waiter(
         &self,
         session_id: &SessionId,
@@ -152,6 +160,11 @@ impl<I: Copy> PendingSessionPromptResponseWaiters<I> {
         }
     }
 
+    #[cfg(test)]
+    pub(crate) fn has_waiter(&self, session_id: &SessionId) -> bool {
+        self.waiters.lock().unwrap().contains_key(session_id)
+    }
+
     #[cfg(test)]
     pub(crate) fn remove_waiter_for_test(&self, session_id: &SessionId) {
         self.waiters.lock().unwrap().remove(session_id);

Original file line number	Diff line number	Diff line change
`@@ -9,6 +9,7 @@ pub enum ClientMethod {`
`9`	`9`	`TerminalOutput,`
`10`	`10`	`TerminalRelease,`
`11`	`11`	`TerminalWaitForExit,`
	`12`	`+ ExtSessionPromptResponse,`
`12`	`13`	`}`
`13`	`14`
`14`	`15`	`impl ClientMethod {`
`@@ -23,6 +24,7 @@ impl ClientMethod {`
`23`	`24`	`"client.terminal.output" => Some(Self::TerminalOutput),`
`24`	`25`	`"client.terminal.release" => Some(Self::TerminalRelease),`
`25`	`26`	`"client.terminal.wait_for_exit" => Some(Self::TerminalWaitForExit),`
	`27`	`+ "client.ext.session.prompt_response" => Some(Self::ExtSessionPromptResponse),`
`26`	`28`	`_ => None,`
`27`	`29`	`}`
`28`	`30`	`}`