fix(rust): answer the request id when a handler panics

Requests are dispatched to their own spawned tasks, so a panic inside a handler is isolated by tokio and the request is simply never answered - the caller waits out its own timeout. Catch the unwind at the dispatch boundary and reply with a JSON-RPC internal error (-32603) instead; the panic payload is not exposed.

Also corrects the stop_event_loop, Drop and lifecycle-test documentation, which still claimed in-flight handlers complete before the loop exits. That stopped being true when request dispatch moved to spawned tasks.

Fixes #2053.

Author: Zhangyi Yuan

rust/src/session.rs

@@ -1,14 +1,16 @@
 use std::collections::HashMap;
+use std::panic::AssertUnwindSafe;
 use std::path::{Path, PathBuf};
 use std::sync::Arc;
 use std::time::{Duration, Instant};
 
+use futures_util::FutureExt;
 use parking_lot::Mutex as ParkingLotMutex;
 use serde_json::Value;
 use tokio::sync::oneshot;
 use tokio::task::JoinHandle;
 use tokio_util::sync::CancellationToken;
-use tracing::{Instrument, warn};
+use tracing::{Instrument, error, warn};
 
 use crate::canvas::CanvasHandler;
 use crate::generated::api_types::{
@@ -321,10 +323,15 @@ impl Session {
     /// Stop the internal event loop. Called automatically on [`destroy`](Self::destroy).
     ///
     /// Cooperative: signals shutdown via the session's [`CancellationToken`]
-    /// and awaits the loop's natural exit rather than aborting the task.
-    /// Any in-flight handler (permission callback, tool call, elicitation
-    /// response) completes before the loop exits, so the CLI never sees a
-    /// half-handled request. See RFD-400 review finding #3.
+    /// and awaits the loop's natural exit rather than aborting the task, so
+    /// the loop always stops between iterations instead of at an arbitrary
+    /// await point. See RFD-400 review finding #3.
+    ///
+    /// Inbound requests are dispatched to their own spawned tasks, which this
+    /// call does not await. A handler (permission callback, tool call,
+    /// elicitation response) still running at teardown may therefore outlive
+    /// the loop, and its response can be lost if the connection closes first.
+    /// Await your own handler work before calling this if it must complete.
     pub async fn stop_event_loop(&self) {
         self.shutdown.cancel();
         let handle = self.event_loop.lock().take();
@@ -643,11 +650,11 @@ impl Drop for Session {
     fn drop(&mut self) {
         // Cooperative shutdown: cancel the event loop's token to signal
         // exit between iterations. The loop will see the cancellation on
-        // its next select poll and break cleanly without interrupting an
-        // in-flight handler. We do NOT abort the JoinHandle — that would
-        // land at any await point in the loop body, potentially leaving
-        // the CLI with an unanswered request id. RFD-400 review finding
-        // #3.
+        // its next select poll and break cleanly. We do NOT abort the
+        // JoinHandle — that would land at any await point in the loop body,
+        // potentially leaving the CLI with an unanswered request id.
+        // RFD-400 review finding #3. Requests already dispatched to their
+        // own tasks are not tracked here and may outlive the session.
         //
         // The handle itself is left in `event_loop` to be reaped by the
         // tokio runtime when it next polls; we intentionally don't await
@@ -1506,6 +1513,8 @@ fn spawn_event_loop(
                         let canvas_handler = canvas_handler.clone();
                         let session_fs_provider = session_fs_provider.clone();
                         let bearer_token_providers = bearer_token_providers.clone();
+                        let request_id = request.id;
+                        let method = request.method.clone();
                         tokio::spawn(
                             async move {
                                 let ctx = RequestDispatchContext {
@@ -1517,7 +1526,19 @@ fn spawn_event_loop(
                                     session_fs_provider: session_fs_provider.as_ref(),
                                     bearer_token_providers: &bearer_token_providers,
                                 };
-                                handle_request(&session_id, ctx, request).await;
+                                let dispatch = handle_request(&session_id, ctx, request);
+                                if AssertUnwindSafe(dispatch).catch_unwind().await.is_err() {
+                                    // Tokio isolates the panic to this task, so without a
+                                    // reply the CLI waits out its own timeout on this id.
+                                    error!(method = %method, "request handler panicked");
+                                    let _ = send_error_response(
+                                        &client,
+                                        request_id,
+                                        error_codes::INTERNAL_ERROR,
+                                        "request handler panicked",
+                                    )
+                                    .await;
+                                }
                             }
                             .instrument(span),
                         );

rust/tests/session_test.rs

@@ -2803,10 +2803,10 @@ async fn send_and_wait_drop_clears_waiter() {
 }
 
 /// Cancel-safety regression: `Session::stop_event_loop` must NOT abort
-/// the event-loop task mid-handler. An in-flight handler (here a slow
-/// `userInput.request` callback) must run to completion before the loop
-/// exits — the CLI receives the response on the wire before the session
-/// tears down.
+/// the event-loop task at an arbitrary await point. Requests are
+/// dispatched to their own tasks, so a handler that is still running when
+/// shutdown is signalled keeps going and its response still reaches the
+/// wire rather than being lost mid-protocol.
 ///
 /// Closes RFD-400 review finding #3.
 #[tokio::test]
@@ -2850,31 +2850,84 @@ async fn stop_event_loop_completes_in_flight_handler() {
     // Give the loop a moment to dispatch into the handler.
     tokio::time::sleep(Duration::from_millis(20)).await;
 
-    // Now request shutdown. The loop is parked in handle_request awaiting
-    // the slow handler. `notify_one()` buffers the signal until the loop
-    // re-enters its select, which can only happen after the handler
-    // returns and the response is sent on the wire.
+    // Now request shutdown while the spawned handler is still sleeping.
     let stop_handle = tokio::spawn({
         let session = session.clone();
         async move { session.stop_event_loop().await }
     });
 
-    // Verify the handler's response lands on the wire BEFORE the loop
-    // exits — i.e. stop_event_loop did not abort mid-handler.
+    // The handler task is independent of the loop, so its response still
+    // lands on the wire instead of being lost to an aborted task.
     let response = timeout(Duration::from_secs(2), server.read_response())
         .await
         .unwrap();
     assert_eq!(response["id"], 900);
     assert_eq!(response["result"]["answer"], "completed");
 
-    // stop_event_loop completes after the handler returns and the loop
-    // observes the buffered shutdown signal on its next select iteration.
     timeout(Duration::from_secs(2), stop_handle)
         .await
         .unwrap()
         .unwrap();
 }
 
+/// A panicking request handler must still answer its request id. Tokio
+/// isolates the panic to the spawned handler task, so without an explicit
+/// reply the caller would wait out its own timeout on a request that can
+/// never complete.
+#[tokio::test]
+async fn panicking_request_handler_responds_with_internal_error() {
+    struct PanickingHandler;
+    #[async_trait]
+    impl UserInputHandler for PanickingHandler {
+        async fn handle(
+            &self,
+            _session_id: SessionId,
+            _question: String,
+            _choices: Option<Vec<String>>,
+            _allow_freeform: Option<bool>,
+        ) -> Option<UserInputResponse> {
+            panic!("handler blew up");
+        }
+    }
+
+    let (session, mut server) = create_session_pair_with_config(|cfg| {
+        cfg.with_user_input_handler(Arc::new(PanickingHandler))
+    })
+    .await;
+
+    server
+        .send_request(
+            901,
+            "userInput.request",
+            serde_json::json!({
+                "sessionId": server.session_id,
+                "question": "boom",
+                "choices": null,
+                "allowFreeform": true,
+            }),
+        )
+        .await;
+
+    let response = timeout(TIMEOUT, server.read_response()).await.unwrap();
+    assert_eq!(response["id"], 901);
+    assert_eq!(response["error"]["code"], -32603);
+    assert!(response.get("result").is_none());
+
+    // The loop survives the panicking handler and keeps serving requests.
+    server
+        .send_request(
+            902,
+            "unknown.method",
+            serde_json::json!({ "sessionId": server.session_id }),
+        )
+        .await;
+    let response = timeout(TIMEOUT, server.read_response()).await.unwrap();
+    assert_eq!(response["id"], 902);
+    assert_eq!(response["error"]["code"], -32601);
+
+    session.stop_event_loop().await;
+}
+
 /// Cancel-safety regression: dropping a Session does NOT abort the event
 /// loop mid-handler. The loop sees the buffered shutdown signal on its
 /// next select iteration and exits cleanly. This is the Drop equivalent