Merge pull request #97 from NianJiuZst/codex/cancel-session-lifecycle-rpcs

fix(session): cancel lifecycle RPCs cooperatively

crates/bsk-cli/src/daemon/ipc.rs

@@ -206,15 +206,15 @@ pub fn full_handler(status: DaemonStatus, state: Arc<DaemonState>) -> RpcHandler
                     Ok(v) => ResponseBody::Ok(v),
                     Err(e) => ResponseBody::Err(e),
                 },
-                Method::SessionStart => match handle_session_start(&state, params).await {
+                Method::SessionStart => match handle_session_start(&state, rpc_id, params).await {
                     Ok(v) => ResponseBody::Ok(v),
                     Err(e) => ResponseBody::Err(e),
                 },
-                Method::SessionStop => match handle_session_stop(&state, params).await {
+                Method::SessionStop => match handle_session_stop(&state, rpc_id, params).await {
                     Ok(v) => ResponseBody::Ok(v),
                     Err(e) => ResponseBody::Err(e),
                 },
-                Method::SessionStopAll => match handle_session_stop_all(&state).await {
+                Method::SessionStopAll => match handle_session_stop_all_rpc(&state, rpc_id).await {
                     Ok(v) => ResponseBody::Ok(v),
                     Err(e) => ResponseBody::Err(e),
                 },
@@ -434,8 +434,9 @@ async fn handle_wait_ms(
 /// cancellation surfaces (M10.2 + review C2):
 ///
 /// 1. [`AbortRegistry`] — answers daemon-local cancellable runners
-///    (currently `tool.wait_ms`). If a token is registered we trip it
-///    and stop.
+///    (`tool.wait_ms` and `session.*` lifecycle calls). If a token is
+///    registered we trip it and stop. Lifecycle handlers forward their
+///    own WS cancel after preserving request-before-cancel ordering.
 /// 2. [`super::inflight::ToolInflightRegistry`] — every IPC-tracked
 ///    `tool.*` RPC, registered the moment the IPC handler accepts the
 ///    request. Trips the entry's cancel token regardless of whether
@@ -634,7 +635,20 @@ fn clamp_browser_wait(wait_ms: Option<u64>) -> Option<Duration> {
     ))
 }
 
-async fn handle_session_start(state: &Arc<DaemonState>, params: Value) -> Result<Value, RpcError> {
+async fn handle_session_start(
+    state: &Arc<DaemonState>,
+    rpc_id: RpcId,
+    params: Value,
+) -> Result<Value, RpcError> {
+    let abort_guard = state
+        .abort_registry
+        .register(rpc_id)
+        .map_err(|err| RpcError {
+            code: ErrorCode::ProtocolError,
+            message: format!("session.start cancellation registration failed: {err:?}"),
+            data: None,
+        })?;
+    let cancel = abort_guard.token().clone();
     let params: CliSessionStartParams = if params.is_null() {
         CliSessionStartParams {
             browser_instance_id: None,
@@ -673,6 +687,7 @@ async fn handle_session_start(state: &Arc<DaemonState>, params: Value) -> Result
         },
         state.config.extension_connect_wait,
         DEFAULT_RPC_TIMEOUT,
+        Some(cancel),
     )
     .await
     {
@@ -696,6 +711,8 @@ fn map_start_error(err: StartSessionError) -> RpcError {
         StartSessionError::AmbiguousBrowserLabel { .. } => ErrorCode::InvalidParams,
         StartSessionError::IdExhausted => ErrorCode::ProtocolError,
         StartSessionError::Timeout => ErrorCode::Timeout,
+        StartSessionError::Cancelled => ErrorCode::Cancelled,
+        StartSessionError::CleanupFailed { .. } => ErrorCode::ProtocolError,
         StartSessionError::TransportClosed => ErrorCode::ProtocolError,
         StartSessionError::ExtensionError(inner) => inner.code,
     };
@@ -711,6 +728,16 @@ fn map_start_error(err: StartSessionError) -> RpcError {
             "label": label,
             "instance_ids": instance_ids,
         })),
+        StartSessionError::CleanupFailed {
+            session_id,
+            agent_window_id,
+            ..
+        } => Some(serde_json::json!({
+            "reason": "cleanup_failed",
+            "session_id": session_id.0,
+            "agent_window_id": agent_window_id,
+        })),
+        StartSessionError::ExtensionError(inner) => inner.data.clone(),
         _ => None,
     };
     RpcError {
@@ -720,7 +747,20 @@ fn map_start_error(err: StartSessionError) -> RpcError {
     }
 }
 
-async fn handle_session_stop(state: &Arc<DaemonState>, params: Value) -> Result<Value, RpcError> {
+async fn handle_session_stop(
+    state: &Arc<DaemonState>,
+    rpc_id: RpcId,
+    params: Value,
+) -> Result<Value, RpcError> {
+    let abort_guard = state
+        .abort_registry
+        .register(rpc_id)
+        .map_err(|err| RpcError {
+            code: ErrorCode::ProtocolError,
+            message: format!("session.stop cancellation registration failed: {err:?}"),
+            data: None,
+        })?;
+    let cancel = abort_guard.token().clone();
     let params: CliSessionStopParams = if params.is_null() {
         CliSessionStopParams {
             session_id: None,
@@ -734,7 +774,7 @@ async fn handle_session_stop(state: &Arc<DaemonState>, params: Value) -> Result<
         })?
     };
     if params.all {
-        return handle_session_stop_all(state).await;
+        return handle_session_stop_all(state, Some(cancel)).await;
     }
     let session_id = match params.session_id {
         Some(s) => SessionId(s),
@@ -753,6 +793,7 @@ async fn handle_session_stop(state: &Arc<DaemonState>, params: Value) -> Result<
         &state.session_interrupts,
         &session_id,
         DEFAULT_SESSION_STOP_TIMEOUT,
+        Some(cancel),
     )
     .await
     {
@@ -804,7 +845,25 @@ fn map_stop_error(err: StopSessionError) -> RpcError {
     }
 }
 
-async fn handle_session_stop_all(state: &Arc<DaemonState>) -> Result<Value, RpcError> {
+async fn handle_session_stop_all_rpc(
+    state: &Arc<DaemonState>,
+    rpc_id: RpcId,
+) -> Result<Value, RpcError> {
+    let abort_guard = state
+        .abort_registry
+        .register(rpc_id)
+        .map_err(|err| RpcError {
+            code: ErrorCode::ProtocolError,
+            message: format!("session.stop_all cancellation registration failed: {err:?}"),
+            data: None,
+        })?;
+    handle_session_stop_all(state, Some(abort_guard.token().clone())).await
+}
+
+async fn handle_session_stop_all(
+    state: &Arc<DaemonState>,
+    cancel: Option<super::abort::AbortToken>,
+) -> Result<Value, RpcError> {
     let ids: Vec<SessionId> = state
         .sessions
         .snapshot()
@@ -816,6 +875,16 @@ async fn handle_session_stop_all(state: &Arc<DaemonState>) -> Result<Value, RpcE
     let mut returned_tab_ids = Vec::new();
     let mut return_failures = Vec::new();
     for id in ids {
+        if cancel
+            .as_ref()
+            .is_some_and(super::abort::AbortToken::is_cancelled)
+        {
+            return Err(RpcError {
+                code: ErrorCode::Cancelled,
+                message: "session.stop_all was cancelled".into(),
+                data: None,
+            });
+        }
         match stop_session(
             &state.browsers,
             &state.sessions,
@@ -823,6 +892,7 @@ async fn handle_session_stop_all(state: &Arc<DaemonState>) -> Result<Value, RpcE
             &state.session_interrupts,
             &id,
             DEFAULT_SESSION_STOP_TIMEOUT,
+            cancel.clone(),
         )
         .await
         {
@@ -848,6 +918,15 @@ async fn handle_session_stop_all(state: &Arc<DaemonState>) -> Result<Value, RpcE
                 returned_tab_ids.extend(stop.returned_tab_ids);
                 return_failures.extend(stop.return_failures);
             }
+            Err(err)
+                if matches!(
+                    &err,
+                    StopSessionError::ExtensionError(inner)
+                        if inner.code == ErrorCode::Cancelled
+                ) =>
+            {
+                return Err(map_stop_error(err));
+            }
             Err(err) => {
                 debug!(session = %id, ?err, "session.stop_all: failure (continuing)");
                 let rpc = map_stop_error(err);

crates/bsk-cli/src/daemon/queue.rs

@@ -43,6 +43,7 @@ use serde_json::Value;
 use tokio::sync::{mpsc, oneshot};
 use tracing::{debug, warn};
 
+use super::abort::AbortToken;
 use super::browsers::BrowserRegistry;
 use super::inflight::{PromoteOutcome, ToolInflightEntry};
 use super::sessions::{SessionId, SessionRegistry};
@@ -83,6 +84,14 @@ pub struct ToolJob {
     /// IPC request id (e.g. `session.stop`'s queued teardown call —
     /// those already carry their own bespoke retry / abort path).
     pub inflight: Option<Arc<ToolInflightEntry>>,
+    /// Cancellation token for session-lifecycle work registered in the
+    /// daemon-local abort registry. Unlike `inflight`, lifecycle callers
+    /// forward their own WS cancel after the request has been queued.
+    pub lifecycle_cancel: Option<AbortToken>,
+    /// How long cancellation keeps the worker busy while the extension
+    /// finishes compensation. Session teardown uses its full RPC budget
+    /// so daemon state is reconciled even after the CLI exits.
+    pub cancel_cleanup_timeout: Duration,
     /// Owning session — used by `ToolInflightRegistry::cancel_session`
     /// to drain queued jobs that have not yet been promoted to
     /// "forwarded". Worker callers also have access to it via the
@@ -322,6 +331,8 @@ impl ToolQueueRegistry {
             timeout,
             false,
             inflight,
+            None,
+            CANCEL_CLEANUP_TIMEOUT,
         )
         .await;
         // A long-running request may outlive the idle threshold. Touching
@@ -360,6 +371,8 @@ impl ToolQueueRegistry {
             timeout,
             respond: respond_tx,
             inflight,
+            lifecycle_cancel: None,
+            cancel_cleanup_timeout: CANCEL_CLEANUP_TIMEOUT,
             session_id: sid.clone(),
         };
         match forward_one(sid, &self.browsers, &self.sessions, &job).await {
@@ -378,6 +391,7 @@ impl ToolQueueRegistry {
         method: Method,
         params: Value,
         timeout: Duration,
+        cancel: Option<AbortToken>,
     ) -> Result<Value, DispatchError> {
         let (sender, state) = {
             let mut guard = self.queues.lock().expect("tool queue registry poisoned");
@@ -404,6 +418,8 @@ impl ToolQueueRegistry {
             timeout,
             true,
             None,
+            cancel,
+            timeout,
         )
         .await
     }
@@ -419,6 +435,8 @@ async fn dispatch_with_sender(
     timeout: Duration,
     wait_for_capacity: bool,
     inflight: Option<Arc<ToolInflightEntry>>,
+    lifecycle_cancel: Option<AbortToken>,
+    cancel_cleanup_timeout: Duration,
 ) -> Result<Value, DispatchError> {
     let (respond_tx, respond_rx) = oneshot::channel();
     let job = ToolJob {
@@ -427,8 +445,11 @@ async fn dispatch_with_sender(
         timeout,
         respond: respond_tx,
         inflight,
+        lifecycle_cancel,
+        cancel_cleanup_timeout,
         session_id,
     };
+    let lifecycle_cancellable = job.lifecycle_cancel.is_some();
     if wait_for_capacity {
         let waited = tokio::time::timeout(
             timeout.saturating_add(Duration::from_secs(1)),
@@ -453,8 +474,15 @@ async fn dispatch_with_sender(
             mpsc::error::TrySendError::Closed(_) => DispatchError::QueueClosed,
         });
     }
-    let waited =
-        tokio::time::timeout(timeout.saturating_add(Duration::from_secs(1)), respond_rx).await;
+    let mut response_timeout = timeout.saturating_add(Duration::from_secs(1));
+    if lifecycle_cancellable {
+        // A lifecycle cancel can arrive at the original RPC deadline and
+        // intentionally keeps the worker alive while extension cleanup
+        // settles. Do not let this outer waiter abandon that reconciliation
+        // early, which could reopen a daemon session whose window is gone.
+        response_timeout = response_timeout.saturating_add(cancel_cleanup_timeout);
+    }
+    let waited = tokio::time::timeout(response_timeout, respond_rx).await;
     match waited {
         Ok(Ok(Ok(v))) => Ok(v),
         Ok(Ok(Err(rpc))) => Err(DispatchError::Rpc(rpc)),
@@ -492,11 +520,17 @@ async fn forward_one(
     // / browser resolution work, and before any WS frame leaves the
     // daemon (review C2). We keep the `tokio::select!` later so
     // cancels arriving mid-WS-hop also unblock the worker.
-    if let Some(entry) = job.inflight.as_ref()
-        && entry.is_cancelled()
+    if job
+        .inflight
+        .as_ref()
+        .is_some_and(|entry| entry.is_cancelled())
+        || job
+            .lifecycle_cancel
+            .as_ref()
+            .is_some_and(AbortToken::is_cancelled)
     {
         return Err(cancelled_error(
-            Some(entry),
+            job.inflight.as_deref(),
             "tool dispatch cancelled before forwarding",
         ));
     }
@@ -560,11 +594,20 @@ async fn forward_one(
             }
         }
         None => {
-            // Daemon-internal callers (e.g. queued session.stop drain)
-            // don't carry an inflight entry; they cannot be cancelled
-            // through the IPC `cancel` surface, so a plain sink send
-            // is sufficient and keeps the existing transport-error
-            // behaviour.
+            // Daemon-internal callers do not carry a tool-inflight entry.
+            // Session lifecycle jobs instead use the abort-registry token
+            // below; other internal callers keep the plain send path.
+            if job
+                .lifecycle_cancel
+                .as_ref()
+                .is_some_and(AbortToken::is_cancelled)
+            {
+                client.pending.lock().unwrap().cancel(&rpc_id);
+                return Err(cancelled_error(
+                    None,
+                    "session lifecycle cancelled before forwarding",
+                ));
+            }
             if client.sink.send(request).is_err() {
                 client.pending.lock().unwrap().cancel(&rpc_id);
                 return Err(RpcError {
@@ -573,19 +616,45 @@ async fn forward_one(
                     data: None,
                 });
             }
-            None
+            job.lifecycle_cancel.clone()
         }
     };
+    let forward_lifecycle_cancel = || {
+        if job.lifecycle_cancel.is_none() {
+            return;
+        }
+        let cancel = Frame::Request(RequestFrame {
+            id: format!("cancel-{rpc_id}"),
+            method: Method::Cancel,
+            params: Some(serde_json::json!({ "rpc_id": rpc_id })),
+        });
+        if let Err(err) = client.sink.send(cancel) {
+            warn!(%rpc_id, ?err, "failed to forward session lifecycle cancel");
+        }
+    };
+    let on_abort: Option<&(dyn Fn() + Sync)> = job
+        .lifecycle_cancel
+        .as_ref()
+        .map(|_| &forward_lifecycle_cancel as &(dyn Fn() + Sync));
     let waited = await_with_optional_cancel(
         job.timeout,
-        CANCEL_CLEANUP_TIMEOUT,
+        job.cancel_cleanup_timeout,
         waiter,
         cancel_token.as_ref(),
+        on_abort,
     )
     .await;
     let response = match waited {
         WaitOutcome::Response(resp) => resp,
         WaitOutcome::CancelledAfterResponse(resp) => {
+            if job.method == Method::ToolSessionStop
+                && let ResponseBody::Ok(value) = resp.body
+            {
+                // Teardown crossed its final irreversible boundary before
+                // the cancel landed. Commit the real extension result so the
+                // daemon does not retain a session whose window is gone.
+                return Ok(value);
+            }
             // Cancellation wins the external verdict, but only after the
             // extension's original RPC has settled. Preserve any non-cancel
             // error so compensation failures remain explicit instead of
@@ -600,13 +669,33 @@ async fn forward_one(
                 "tool dispatch cancelled after extension cleanup",
             ));
         }
+        WaitOutcome::TimedOutAfterResponse(resp) => match resp.body {
+            ResponseBody::Ok(value) if job.method == Method::ToolSessionStop => {
+                // The close crossed its irreversible boundary during the
+                // timeout cleanup grace. Reconcile daemon state instead of
+                // reopening a session whose window is already gone.
+                return Ok(value);
+            }
+            ResponseBody::Err(err)
+                if !matches!(err.code, ErrorCode::Cancelled | ErrorCode::UserAborted) =>
+            {
+                return Err(err);
+            }
+            _ => {
+                return Err(RpcError {
+                    code: ErrorCode::Timeout,
+                    message: format!("tool RPC timed out after {:?}", job.timeout),
+                    data: None,
+                });
+            }
+        },
         WaitOutcome::CleanupTimeout => {
             client.pending.lock().unwrap().cancel(&rpc_id);
             return Err(RpcError {
                 code: ErrorCode::Timeout,
                 message: format!(
                     "cancelled tool did not finish cleanup within {:?}",
-                    CANCEL_CLEANUP_TIMEOUT
+                    job.cancel_cleanup_timeout
                 ),
                 data: Some(serde_json::json!({ "reason": "cancel_cleanup_timeout" })),
             });
@@ -638,6 +727,7 @@ async fn forward_one(
 enum WaitOutcome {
     Response(bsk_protocol::ResponseFrame),
     CancelledAfterResponse(bsk_protocol::ResponseFrame),
+    TimedOutAfterResponse(bsk_protocol::ResponseFrame),
     CleanupTimeout,
     WaiterClosed,
     Timeout,
@@ -648,6 +738,7 @@ async fn await_with_optional_cancel(
     cleanup_timeout: Duration,
     mut waiter: oneshot::Receiver<bsk_protocol::ResponseFrame>,
     cancel: Option<&super::abort::AbortToken>,
+    on_abort: Option<&(dyn Fn() + Sync)>,
 ) -> WaitOutcome {
     match cancel {
         Some(token) => {
@@ -659,6 +750,9 @@ async fn await_with_optional_cancel(
                 // after compensation or the cleanup deadline expires.
                 biased;
                 _ = token.cancelled() => {
+                    if let Some(on_abort) = on_abort {
+                        on_abort();
+                    }
                     match tokio::time::timeout(cleanup_timeout, &mut waiter).await {
                         Ok(Ok(resp)) => WaitOutcome::CancelledAfterResponse(resp),
                         Ok(Err(_)) => WaitOutcome::WaiterClosed,
@@ -669,7 +763,17 @@ async fn await_with_optional_cancel(
                     Ok(resp) => WaitOutcome::Response(resp),
                     Err(_) => WaitOutcome::WaiterClosed,
                 },
-                _ = &mut deadline => WaitOutcome::Timeout,
+                _ = &mut deadline => {
+                    let Some(on_abort) = on_abort else {
+                        return WaitOutcome::Timeout;
+                    };
+                    on_abort();
+                    match tokio::time::timeout(cleanup_timeout, &mut waiter).await {
+                        Ok(Ok(resp)) => WaitOutcome::TimedOutAfterResponse(resp),
+                        Ok(Err(_)) => WaitOutcome::WaiterClosed,
+                        Err(_) => WaitOutcome::CleanupTimeout,
+                    }
+                },
             }
         }
         None => match tokio::time::timeout(timeout, &mut waiter).await {
@@ -759,6 +863,7 @@ mod await_with_optional_cancel_tests {
             Duration::from_secs(1),
             rx,
             Some(&token),
+            None,
         )
         .await;
         assert!(
@@ -781,6 +886,7 @@ mod await_with_optional_cancel_tests {
             Duration::from_secs(1),
             rx,
             Some(&token),
+            None,
         )
         .await;
         match outcome {
@@ -795,9 +901,14 @@ mod await_with_optional_cancel_tests {
     async fn no_cancel_token_still_returns_response() {
         let (tx, rx) = oneshot::channel();
         tx.send(dummy_response()).unwrap();
-        let outcome =
-            await_with_optional_cancel(Duration::from_secs(10), Duration::from_secs(1), rx, None)
-                .await;
+        let outcome = await_with_optional_cancel(
+            Duration::from_secs(10),
+            Duration::from_secs(1),
+            rx,
+            None,
+            None,
+        )
+        .await;
         assert!(matches!(outcome, WaitOutcome::Response(_)));
     }
 
@@ -812,6 +923,7 @@ mod await_with_optional_cancel_tests {
                 Duration::from_secs(1),
                 rx,
                 Some(&token),
+                None,
             )
             .await
         });
@@ -835,10 +947,37 @@ mod await_with_optional_cancel_tests {
             Duration::from_millis(10),
             rx,
             Some(&token),
+            None,
         )
         .await;
         assert!(matches!(outcome, WaitOutcome::CleanupTimeout));
     }
+
+    #[tokio::test]
+    async fn lifecycle_timeout_forwards_abort_and_waits_for_cleanup_response() {
+        use std::sync::atomic::{AtomicBool, Ordering};
+
+        let token = AbortToken::new();
+        let (tx, rx) = oneshot::channel();
+        let abort_forwarded = AtomicBool::new(false);
+        let forward_abort = || abort_forwarded.store(true, Ordering::SeqCst);
+        tokio::spawn(async move {
+            tokio::time::sleep(Duration::from_millis(20)).await;
+            tx.send(dummy_response()).unwrap();
+        });
+
+        let outcome = await_with_optional_cancel(
+            Duration::from_millis(5),
+            Duration::from_millis(100),
+            rx,
+            Some(&token),
+            Some(&forward_abort),
+        )
+        .await;
+
+        assert!(abort_forwarded.load(Ordering::SeqCst));
+        assert!(matches!(outcome, WaitOutcome::TimedOutAfterResponse(_)));
+    }
 }
 
 #[cfg(test)]
@@ -901,6 +1040,8 @@ mod tool_job_session_id_tests {
             timeout: Duration::from_secs(1),
             respond: tx,
             inflight: None,
+            lifecycle_cancel: None,
+            cancel_cleanup_timeout: CANCEL_CLEANUP_TIMEOUT,
             session_id: SessionId("sess-A".into()),
         };
         assert_eq!(job.session_id.0, "sess-A");

crates/bsk-cli/src/daemon/sessions.rs

@@ -12,13 +12,14 @@ use bsk_protocol::system::{BrowserStatusEntry, SessionStatusEntry};
 use bsk_protocol::tools::{
     SessionStartParams, SessionStartResult, SessionStopParams, SessionStopResult,
 };
-use bsk_protocol::{Frame, RequestFrame, ResponseBody, RpcError, RpcId};
+use bsk_protocol::{ErrorCode, Frame, Method, RequestFrame, ResponseBody, RpcError, RpcId};
 use rand::Rng;
 use serde::{Deserialize, Serialize};
 use tokio::time::{Duration, timeout};
 
+use super::abort::AbortToken;
 use super::browsers::{BrowserClient, BrowserId, BrowserRegistry, SelectError};
-use super::queue::{DispatchError, ToolQueueRegistry};
+use super::queue::{CANCEL_CLEANUP_TIMEOUT, DispatchError, ToolQueueRegistry};
 use super::session_interrupt::SessionInterruptRegistry;
 
 #[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
@@ -339,6 +340,14 @@ pub enum StartSessionError {
     IdExhausted,
     #[error("session creation timed out waiting for extension")]
     Timeout,
+    #[error("session creation was cancelled")]
+    Cancelled,
+    #[error("session creation cancellation cleanup failed for {session_id}: {message}")]
+    CleanupFailed {
+        session_id: SessionId,
+        agent_window_id: Option<i64>,
+        message: String,
+    },
     #[error("extension rejected tool.session_start: {0:?}")]
     ExtensionError(RpcError),
     #[error("transport closed while waiting for extension response")]
@@ -354,6 +363,8 @@ impl StartSessionError {
             StartSessionError::AmbiguousBrowserLabel { .. } => "invalid_params",
             StartSessionError::IdExhausted => "protocol_error",
             StartSessionError::Timeout => "timeout",
+            StartSessionError::Cancelled => "cancelled",
+            StartSessionError::CleanupFailed { .. } => "protocol_error",
             StartSessionError::TransportClosed => "protocol_error",
             StartSessionError::ExtensionError(err) => match err.code {
                 bsk_protocol::ErrorCode::Timeout => "timeout",
@@ -407,6 +418,7 @@ pub struct AgentWindowOptions {
 /// On success the matching per-session dispatch queue (M6.5,
 /// design §5) is also spawned so that subsequent `tool.*` RPCs serialise
 /// against this session's Agent Window / ref-store.
+#[allow(clippy::too_many_arguments)]
 pub async fn start_session(
     registry: &Arc<BrowserRegistry>,
     sessions: &Arc<SessionRegistry>,
@@ -415,24 +427,32 @@ pub async fn start_session(
     window: AgentWindowOptions,
     connect_wait: Duration,
     timeout_dur: Duration,
+    cancel: Option<AbortToken>,
 ) -> Result<Session, StartSessionError> {
-    let client: Arc<BrowserClient> = registry
-        .select_with_connect_wait(requested, connect_wait)
-        .await
-        .map_err(|e| match e {
-            SelectError::NoBrowserConnected => StartSessionError::NoBrowserConnected,
-            SelectError::MultipleBrowsersOnline => StartSessionError::MultipleBrowsersOnline {
-                browsers: snapshot_status_entries(registry, sessions),
-            },
-            SelectError::NotFound => StartSessionError::BrowserNotFound,
-            SelectError::AmbiguousLabel {
-                label,
-                instance_ids,
-            } => StartSessionError::AmbiguousBrowserLabel {
-                label,
-                instance_ids,
-            },
-        })?;
+    let selection = registry.select_with_connect_wait(requested, connect_wait);
+    tokio::pin!(selection);
+    let selected = match cancel.as_ref() {
+        Some(token) => tokio::select! {
+            biased;
+            _ = token.cancelled() => return Err(StartSessionError::Cancelled),
+            selected = &mut selection => selected,
+        },
+        None => selection.await,
+    };
+    let client: Arc<BrowserClient> = selected.map_err(|e| match e {
+        SelectError::NoBrowserConnected => StartSessionError::NoBrowserConnected,
+        SelectError::MultipleBrowsersOnline => StartSessionError::MultipleBrowsersOnline {
+            browsers: snapshot_status_entries(registry, sessions),
+        },
+        SelectError::NotFound => StartSessionError::BrowserNotFound,
+        SelectError::AmbiguousLabel {
+            label,
+            instance_ids,
+        } => StartSessionError::AmbiguousBrowserLabel {
+            label,
+            instance_ids,
+        },
+    })?;
     let session_id = sessions
         .reserve_id(client.id.clone(), SESSION_ID_MAX_RESERVE_ATTEMPTS, now_ms)
         .ok_or(StartSessionError::IdExhausted)?;
@@ -453,22 +473,71 @@ pub async fn start_session(
         let mut pending = client.pending.lock().unwrap();
         pending.register(rpc_id.clone())
     };
+    if cancel.as_ref().is_some_and(AbortToken::is_cancelled) {
+        sessions.cancel_reservation(&session_id);
+        client.pending.lock().unwrap().cancel(&rpc_id);
+        return Err(StartSessionError::Cancelled);
+    }
     if client.sink.send(Frame::Request(request)).is_err() {
         sessions.cancel_reservation(&session_id);
         client.pending.lock().unwrap().cancel(&rpc_id);
         return Err(StartSessionError::TransportClosed);
     }
-    let response = match timeout(timeout_dur, waiter).await {
-        Ok(Ok(resp)) => resp,
-        Ok(Err(_)) => {
+    let mut waiter = waiter;
+    enum StartWaitOutcome {
+        Response(bsk_protocol::ResponseFrame),
+        WaiterClosed,
+        Cancelled,
+        Timeout,
+    }
+    let wait_outcome = match cancel.as_ref() {
+        Some(token) => {
+            let deadline = tokio::time::sleep(timeout_dur);
+            tokio::pin!(deadline);
+            tokio::select! {
+                biased;
+                _ = token.cancelled() => StartWaitOutcome::Cancelled,
+                response = &mut waiter => match response {
+                    Ok(response) => StartWaitOutcome::Response(response),
+                    Err(_) => StartWaitOutcome::WaiterClosed,
+                },
+                _ = &mut deadline => StartWaitOutcome::Timeout,
+            }
+        }
+        None => match timeout(timeout_dur, &mut waiter).await {
+            Ok(Ok(response)) => StartWaitOutcome::Response(response),
+            Ok(Err(_)) => StartWaitOutcome::WaiterClosed,
+            Err(_) => StartWaitOutcome::Timeout,
+        },
+    };
+    let response = match wait_outcome {
+        StartWaitOutcome::Response(response) => response,
+        StartWaitOutcome::WaiterClosed => {
             sessions.cancel_reservation(&session_id);
             client.pending.lock().unwrap().cancel(&rpc_id);
             return Err(StartSessionError::TransportClosed);
         }
-        Err(_) => {
-            sessions.cancel_reservation(&session_id);
-            client.pending.lock().unwrap().cancel(&rpc_id);
-            return Err(StartSessionError::Timeout);
+        StartWaitOutcome::Cancelled => {
+            return Err(finish_aborted_start(
+                &client,
+                sessions,
+                &session_id,
+                &rpc_id,
+                waiter,
+                StartAbortReason::Cancelled,
+            )
+            .await);
+        }
+        StartWaitOutcome::Timeout => {
+            return Err(finish_aborted_start(
+                &client,
+                sessions,
+                &session_id,
+                &rpc_id,
+                waiter,
+                StartAbortReason::Timeout,
+            )
+            .await);
         }
     };
     let agent_window_id = match response.body {
@@ -504,6 +573,124 @@ pub async fn start_session(
     Ok(session)
 }
 
+#[derive(Clone, Copy)]
+enum StartAbortReason {
+    Cancelled,
+    Timeout,
+}
+
+impl StartAbortReason {
+    fn error(self) -> StartSessionError {
+        match self {
+            Self::Cancelled => StartSessionError::Cancelled,
+            Self::Timeout => StartSessionError::Timeout,
+        }
+    }
+}
+
+async fn finish_aborted_start(
+    client: &Arc<BrowserClient>,
+    sessions: &Arc<SessionRegistry>,
+    session_id: &SessionId,
+    rpc_id: &RpcId,
+    mut waiter: tokio::sync::oneshot::Receiver<bsk_protocol::ResponseFrame>,
+    reason: StartAbortReason,
+) -> StartSessionError {
+    sessions.cancel_reservation(session_id);
+    let cancel = Frame::Request(RequestFrame {
+        id: format!("cancel-{rpc_id}"),
+        method: Method::Cancel,
+        params: Some(serde_json::json!({ "rpc_id": rpc_id })),
+    });
+    if client.sink.send(cancel).is_err() {
+        client.pending.lock().unwrap().cancel(rpc_id);
+        return StartSessionError::TransportClosed;
+    }
+
+    match timeout(CANCEL_CLEANUP_TIMEOUT, &mut waiter).await {
+        Ok(Ok(response)) => match response.body {
+            ResponseBody::Err(err)
+                if matches!(err.code, ErrorCode::Cancelled | ErrorCode::UserAborted) =>
+            {
+                reason.error()
+            }
+            ResponseBody::Err(err) => StartSessionError::ExtensionError(err),
+            ResponseBody::Ok(value) => {
+                let agent_window_id = serde_json::from_value::<SessionStartResult>(value)
+                    .ok()
+                    .and_then(|result| result.agent_window_id);
+                match rollback_extension_session(client, session_id).await {
+                    Ok(()) => reason.error(),
+                    Err(message) => StartSessionError::CleanupFailed {
+                        session_id: session_id.clone(),
+                        agent_window_id,
+                        message,
+                    },
+                }
+            }
+        },
+        Ok(Err(_)) => {
+            client.pending.lock().unwrap().cancel(rpc_id);
+            StartSessionError::TransportClosed
+        }
+        Err(_) => {
+            // The cancel frame is already queued behind session_start. Even
+            // if chrome.windows.create is still blocked, the extension will
+            // observe the aborted signal and compensate before replying.
+            client.pending.lock().unwrap().cancel(rpc_id);
+            reason.error()
+        }
+    }
+}
+
+async fn rollback_extension_session(
+    client: &Arc<BrowserClient>,
+    session_id: &SessionId,
+) -> Result<(), String> {
+    let rollback_id = next_rpc_id("sess-start-rollback");
+    let waiter = {
+        let mut pending = client.pending.lock().unwrap();
+        pending.register(rollback_id.clone())
+    };
+    let request = RequestFrame {
+        id: rollback_id.clone(),
+        method: Method::ToolSessionStop,
+        params: Some(
+            serde_json::to_value(SessionStopParams {
+                session_id: session_id.0.clone(),
+            })
+            .expect("serialize session_start rollback params"),
+        ),
+    };
+    if client.sink.send(Frame::Request(request)).is_err() {
+        client.pending.lock().unwrap().cancel(&rollback_id);
+        return Err("browser disconnected before rollback was queued".into());
+    }
+    let response = match timeout(CANCEL_CLEANUP_TIMEOUT, waiter).await {
+        Ok(Ok(response)) => response,
+        Ok(Err(_)) => return Err("browser disconnected during rollback".into()),
+        Err(_) => {
+            client.pending.lock().unwrap().cancel(&rollback_id);
+            return Err(format!(
+                "extension did not finish rollback within {CANCEL_CLEANUP_TIMEOUT:?}"
+            ));
+        }
+    };
+    match response.body {
+        ResponseBody::Ok(value) => {
+            let result: SessionStopResult = serde_json::from_value(value)
+                .map_err(|err| format!("invalid rollback response: {err}"))?;
+            if result.return_failures.is_empty() {
+                Ok(())
+            } else {
+                Err("rollback reported borrowed-tab return failures".into())
+            }
+        }
+        ResponseBody::Err(err) if err.code == ErrorCode::NotFound => Ok(()),
+        ResponseBody::Err(err) => Err(format!("rollback failed: {} ({:?})", err.message, err.code)),
+    }
+}
+
 /// Tear down a single session id (round-trips `tool.session_stop` to the
 /// owning extension and unregisters on success). Also tears down the
 /// per-session dispatch queue so the worker task exits.
@@ -514,6 +701,7 @@ pub async fn stop_session(
     interrupts: &Arc<SessionInterruptRegistry>,
     session_id: &SessionId,
     timeout_dur: Duration,
+    cancel: Option<AbortToken>,
 ) -> Result<SessionStopResult, StopSessionError> {
     let session = sessions.get(session_id).ok_or(StopSessionError::NotFound)?;
     if registry.get(&session.browser_id).is_none() {
@@ -524,12 +712,17 @@ pub async fn stop_session(
         session_id: session_id.0.clone(),
     })
     .unwrap();
+    // Internal callers (idle reaping and shutdown) do not have a CLI-owned
+    // token, but still need timeout-driven WS cancellation so a late window
+    // close cannot leave daemon state pointing at a vanished session.
+    let lifecycle_cancel = Some(cancel.unwrap_or_default());
     match queues
         .dispatch_after_closing(
             session_id,
             bsk_protocol::Method::ToolSessionStop,
             params,
             timeout_dur,
+            lifecycle_cancel,
         )
         .await
     {
@@ -567,6 +760,10 @@ pub async fn stop_session(
                 drop_session_local(queues, interrupts, session_id);
                 return Ok(SessionStopResult::default());
             }
+            if err.code == bsk_protocol::ErrorCode::Timeout {
+                queues.reopen(session_id);
+                return Err(StopSessionError::Timeout);
+            }
             queues.reopen(session_id);
             Err(StopSessionError::ExtensionError(err))
         }

crates/bsk-cli/src/daemon/start.rs

@@ -491,6 +491,7 @@ pub(crate) fn spawn_session_idle_reaper(state: Arc<DaemonState>) -> tokio::task:
                     &state.session_interrupts,
                     &session_id,
                     Duration::from_secs(10),
+                    None,
                 )
                 .await
                 {

crates/bsk-cli/src/daemon/state.rs

@@ -34,8 +34,8 @@ pub struct DaemonState {
     /// `stop_session` / browser disconnect.
     pub tool_queues: Arc<ToolQueueRegistry>,
     /// Per-rpc-id cancellation tokens for daemon-side long-runners
-    /// (M9.3 — currently only `tool.wait_ms`). The CLI's `cancel
-    /// { rpc_id }` consults this registry first.
+    /// (`tool.wait_ms` plus `session.*` lifecycle calls). The CLI's
+    /// `cancel { rpc_id }` consults this registry first.
     pub abort_registry: Arc<AbortRegistry>,
     /// Tracks `tool.*` RPCs that have been forwarded to an extension
     /// over WS but have not yet received a response. Indexed by the