Refactor session shutdown from Notify to CancellationToken

Replaces the per-session `Arc<Notify>` shutdown signal with
`tokio_util::sync::CancellationToken`, the canonical primitive for
SDK-internal task coordination in async Rust. Same cooperative-shutdown
semantics, materially better fit for the use case:

- `tonic` uses `CancellationToken` for the equivalent task coordination
  case (request cancel propagated into spawned server handlers).
- Parent/child token tree gives us a clean way to expose a public
  `Session::cancellation_token()` accessor that returns a child token —
  consumers can `select!` on `child.cancelled()` to bind their own work
  to the session lifetime, but cancelling the child does NOT shut down
  the session itself (child cancel is isolated).
- `Drop` impl just calls `cancel()`, no buffering quirks (Notify's
  `notify_one` was already documented as buffer-on-no-listeners; token
  cancellation is sticky, which is the simpler mental model).

Plumbing:

- `tokio-util = { version = "0.7", default-features = false }` added as
  a runtime dep. No new transitive deps the SDK didn't already pull
  through tokio.
- `Session.shutdown` field type changes from `Arc<Notify>` to
  `CancellationToken` (which is itself internally `Arc`-backed and
  `Clone`).
- `stop_event_loop()` calls `cancel()` instead of `notify_one()`.
- `Drop` impl calls `cancel()` instead of `notify_one()`.
- Event-loop `select!` arm becomes `_ = shutdown.cancelled() => break`.
- New public `Session::cancellation_token()` returning a `child_token()`
  — power-user API for binding external tasks to session lifetime.

Tests:

- `cancellation_token_fires_on_session_drop`: dropping the session
  fires the child token within 2s.
- `cancellation_token_child_cancel_does_not_kill_session`: cancelling
  a child does not propagate up to the parent.

Closes the cancellation idiom gap stephentoub flagged on the SKILL doc.
The framing "drop the future or call abort()" was too absolute — drop-
to-cancel is the right primitive for caller-owned futures, but for
SDK-internal task coordination CancellationToken is the canonical
answer (citations: tokio-util docs, tonic cancellation example,
withoutboats "Asynchronous clean-up", `tokio-graceful-shutdown` crate).

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Author: tclem

rust/CHANGELOG.md

@@ -71,6 +71,13 @@ public surface.
 - `Session::disconnect` (canonical) and `Session::destroy` (alias)
   preserve on-disk session state for later resume.
 - `Session::stop_event_loop` for shutting down the per-session loop.
+- `Session::cancellation_token()` returns a [`tokio_util::sync::CancellationToken`]
+  child token that fires when the session shuts down (via
+  `stop_event_loop`, `destroy`, or `Drop`). Lets external tasks bind their
+  lifetime to a session via `tokio::select!` without taking a strong
+  reference to the session. Cancelling the returned child token does not
+  shut the session down — only `stop_event_loop` (or dropping the session)
+  does.
 
 #### Handlers + helpers
 - `SessionHandler` trait with default fallback impls for each event

rust/Cargo.lock

@@ -45,6 +45,7 @@ serde_json = "1"
 thiserror = "2"
 tokio = { version = "1", features = ["io-util", "sync", "rt", "process", "net", "time", "macros"] }
 tokio-stream = { version = "0.1", features = ["sync"] }
+tokio-util = { version = "0.7", default-features = false }
 tracing = "0.1"
 dirs = "5"
 parking_lot = "0.12"

rust/Cargo.toml

@@ -5,8 +5,9 @@ use std::time::Duration;
 
 use parking_lot::Mutex as ParkingLotMutex;
 use serde_json::Value;
-use tokio::sync::{Notify, oneshot};
+use tokio::sync::oneshot;
 use tokio::task::JoinHandle;
+use tokio_util::sync::CancellationToken;
 use tracing::{Instrument, warn};
 
 use crate::generated::api_types::{
@@ -87,14 +88,19 @@ pub struct Session {
     /// fallibility.
     event_loop: ParkingLotMutex<Option<JoinHandle<()>>>,
     /// Cooperative shutdown signal for the event loop. The loop selects
-    /// on `shutdown.notified()` alongside its inbound channels;
-    /// [`Session::stop_event_loop`] and [`Drop`] both call `notify_one()`
-    /// (which buffers a single signal so it is not lost while the loop
-    /// is inside an awaiting branch) to ask the loop to exit between
-    /// iterations rather than `JoinHandle::abort` (which can land at any
-    /// await point and leave the session in mid-protocol state). See
-    /// RFD-400 review finding #3.
-    shutdown: Arc<Notify>,
+    /// on [`shutdown.cancelled()`](CancellationToken::cancelled) alongside
+    /// its inbound channels; [`Session::stop_event_loop`] and [`Drop`]
+    /// both call [`cancel()`](CancellationToken::cancel) to ask the loop
+    /// to exit between iterations rather than aborting the task (which
+    /// can land at any await point and leave the session mid-protocol).
+    /// See RFD-400 review finding #3.
+    ///
+    /// `CancellationToken` is the canonical signalling primitive in
+    /// `tokio_util`; it is what `tonic` uses for the equivalent task-
+    /// coordination case. Advanced consumers can obtain a child token
+    /// via [`Session::cancellation_token`] to bind their own work to
+    /// the session lifetime.
+    shutdown: CancellationToken,
     /// Only populated while a `send_and_wait` call is in flight.
     ///
     /// Sync `parking_lot::Mutex` because the lock is never held across an
@@ -137,6 +143,36 @@ impl Session {
         self.capabilities.read().clone()
     }
 
+    /// Returns a [`CancellationToken`] that fires when this session shuts
+    /// down (via [`Session::stop_event_loop`], [`Session::destroy`], or
+    /// [`Drop`]).
+    ///
+    /// Use this to bind an external task's lifetime to the session — when
+    /// the session shuts down, awaiting [`cancelled()`](CancellationToken::cancelled)
+    /// resolves so cooperative consumers can stop cleanly.
+    ///
+    /// The returned handle is a *child* token: calling
+    /// [`cancel()`](CancellationToken::cancel) on it cancels only the
+    /// caller's child, not the session itself. To cancel the session, call
+    /// [`Session::stop_event_loop`].
+    ///
+    /// # Example
+    ///
+    /// ```no_run
+    /// # async fn example(session: github_copilot_sdk::session::Session) {
+    /// let token = session.cancellation_token();
+    /// tokio::select! {
+    ///     _ = token.cancelled() => println!("session shut down"),
+    ///     _ = tokio::time::sleep(std::time::Duration::from_secs(60)) => {
+    ///         println!("60s elapsed, session still alive");
+    ///     }
+    /// }
+    /// # }
+    /// ```
+    pub fn cancellation_token(&self) -> CancellationToken {
+        self.shutdown.child_token()
+    }
+
     /// Subscribe to events for this session.
     ///
     /// Returns an [`EventSubscription`](crate::subscription::EventSubscription)
@@ -200,13 +236,13 @@ impl Session {
 
     /// Stop the internal event loop. Called automatically on [`destroy`](Self::destroy).
     ///
-    /// Cooperative: signals shutdown via the session's `Notify` 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
+    /// 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.
     pub async fn stop_event_loop(&self) {
-        self.shutdown.notify_one();
+        self.shutdown.cancel();
         let handle = self.event_loop.lock().take();
         if let Some(handle) = handle {
             let _ = handle.await;
@@ -618,17 +654,18 @@ impl Session {
 
 impl Drop for Session {
     fn drop(&mut self) {
-        // Cooperative shutdown: notify the event loop to exit between
-        // iterations. The loop will see the signal 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.
+        // 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.
         //
         // The handle itself is left in `event_loop` to be reaped by the
         // tokio runtime when it next polls; we intentionally don't await
         // it here because Drop is sync.
-        self.shutdown.notify_one();
+        self.shutdown.cancel();
         self.client.unregister_session(&self.id);
     }
 }
@@ -824,7 +861,7 @@ impl Client {
         let channels = self.register_session(&session_id);
 
         let idle_waiter = Arc::new(ParkingLotMutex::new(None));
-        let shutdown = Arc::new(Notify::new());
+        let shutdown = CancellationToken::new();
         let (event_tx, _) = tokio::sync::broadcast::channel(512);
         let event_loop = spawn_event_loop(
             session_id.clone(),
@@ -927,7 +964,7 @@ impl Client {
         let channels = self.register_session(&cli_session_id);
 
         let idle_waiter = Arc::new(ParkingLotMutex::new(None));
-        let shutdown = Arc::new(Notify::new());
+        let shutdown = CancellationToken::new();
         let (event_tx, _) = tokio::sync::broadcast::channel(512);
         let event_loop = spawn_event_loop(
             cli_session_id.clone(),
@@ -986,7 +1023,7 @@ fn spawn_event_loop(
     idle_waiter: Arc<ParkingLotMutex<Option<IdleWaiter>>>,
     capabilities: Arc<parking_lot::RwLock<SessionCapabilities>>,
     event_tx: tokio::sync::broadcast::Sender<SessionEvent>,
-    shutdown: Arc<Notify>,
+    shutdown: CancellationToken,
 ) -> JoinHandle<()> {
     let crate::router::SessionChannels {
         mut notifications,
@@ -997,18 +1034,18 @@ fn spawn_event_loop(
     tokio::spawn(
         async move {
             loop {
-                // `mpsc::UnboundedReceiver::recv` and `Notify::notified()`
-                // are both cancel-safe per RFD 400. The selected branch's
-                // `await`'d handler is *not* mid-cancelled by the select
-                // — once a branch fires it runs to completion within the
-                // loop's iteration. Spawned child tasks inside
-                // `handle_notification` (permission/tool/elicitation
-                // callbacks) intentionally outlive the parent loop and
-                // own their own cleanup; this is RFD 400's "spawn
-                // background tasks to perform cancel-unsafe operations"
-                // pattern and is correct as-is.
+                // `mpsc::UnboundedReceiver::recv` and
+                // `CancellationToken::cancelled` are both cancel-safe per
+                // RFD 400. The selected branch's `await`'d handler is
+                // *not* mid-cancelled by the select — once a branch fires
+                // it runs to completion within the loop's iteration.
+                // Spawned child tasks inside `handle_notification`
+                // (permission/tool/elicitation callbacks) intentionally
+                // outlive the parent loop and own their own cleanup;
+                // this is RFD 400's "spawn background tasks to perform
+                // cancel-unsafe operations" pattern and is correct as-is.
                 tokio::select! {
-                    _ = shutdown.notified() => break,
+                    _ = shutdown.cancelled() => break,
                     Some(notification) = notifications.recv() => {
                         handle_notification(
                             &session_id, &client, &handler, &command_handlers, notification, &idle_waiter, &capabilities, &event_tx,

rust/src/session.rs

@@ -2052,6 +2052,45 @@ async fn drop_session_does_not_abort_handler() {
     );
 }
 
+/// `Session::cancellation_token()` returns a child token that fires when
+/// the session shuts down. Lets external tasks bind their lifetime to the
+/// session via `tokio::select!` without taking a strong reference to the
+/// session itself.
+#[tokio::test]
+async fn cancellation_token_fires_on_session_drop() {
+    let handler = Arc::new(ApproveAllHandler);
+    let (session, _server) = create_session_pair(handler).await;
+
+    let token = session.cancellation_token();
+    assert!(!token.is_cancelled());
+
+    drop(session);
+
+    // The session's Drop impl cancels the parent token, which propagates
+    // to all child tokens.
+    timeout(Duration::from_secs(2), token.cancelled())
+        .await
+        .expect("child token must observe cancellation after session drop");
+    assert!(token.is_cancelled());
+}
+
+/// Cancelling a child token returned by `cancellation_token()` does NOT
+/// shut the session down — child tokens isolate consumer-side cancel
+/// logic from the session's own lifecycle.
+#[tokio::test]
+async fn cancellation_token_child_cancel_does_not_kill_session() {
+    let handler = Arc::new(ApproveAllHandler);
+    let (session, _server) = create_session_pair(handler).await;
+
+    let child = session.cancellation_token();
+    child.cancel();
+
+    // Session's own token (and event loop) are untouched. Issue a cheap
+    // RPC and confirm it still works.
+    let parent = session.cancellation_token();
+    assert!(!parent.is_cancelled());
+}
+
 #[tokio::test]
 async fn elicitation_requested_dispatches_to_handler_and_responds() {
     use github_copilot_sdk::types::ElicitationResult;