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

fix(session): cancel lifecycle RPCs cooperatively

apps/extension/src/tools/__tests__/dispatcher.test.ts

@@ -203,6 +203,55 @@ describe("ToolDispatcher", () => {
     expect(sent[0]).toEqual({ id: "r-1", result: {} });
   });
 
+  it("forwards daemon cancellation into an in-flight session stop", async () => {
+    const { transport, sent, deliver } = fakeTransport();
+    const remove = vi.fn(async () => {});
+    const sessions = new SessionManager({
+      agentWindow: {
+        create: vi.fn(async () => 4242),
+        remove,
+        ensureActiveTab: vi.fn(async () => {}),
+      },
+    });
+    await sessions.start("aa11");
+    let finishDetach: () => void = () => {};
+    const detach = new Promise<void>((resolve) => {
+      finishDetach = resolve;
+    });
+    const cdp = {
+      send: vi.fn(),
+      detachSession: vi.fn(() => detach),
+      ensureNetworkCapture: vi.fn(async () => {}),
+      networkEntriesSince: vi.fn(() => ({
+        tab_id: 7,
+        entries: [],
+        next_since: 0,
+        truncated: false,
+      })),
+      setDeviceMetricsOverride: vi.fn(async () => {}),
+      clearDeviceMetricsOverride: vi.fn(async () => {}),
+      setUserAgentOverride: vi.fn(async () => {}),
+      setTouchEmulationEnabled: vi.fn(async () => {}),
+    };
+    const dispatcher = new ToolDispatcher({ transport, sessions, cdp: cdp as TestDispatcherCdp });
+    dispatcher.start();
+
+    deliver(makeRequest("tool.session_stop", { session_id: "aa11" }));
+    await vi.waitFor(() => expect(cdp.detachSession).toHaveBeenCalledWith("aa11"));
+    deliver({ id: "cancel-stop", method: "cancel", params: { rpc_id: "r-1" } });
+    await flushMicrotasks();
+    finishDetach();
+    await flushMicrotasks();
+
+    expect(sent).toContainEqual({ id: "cancel-stop", result: { cancelled: true } });
+    expect(sent).toContainEqual({
+      id: "r-1",
+      error: expect.objectContaining({ code: "cancelled" }),
+    });
+    expect(sessions.has("aa11")).toBe(true);
+    expect(remove).not.toHaveBeenCalled();
+  });
+
   it("returns not_found when stopping an unknown session", async () => {
     const { transport, sent, deliver } = fakeTransport();
     const sessions = new SessionManager({

apps/extension/src/tools/__tests__/session.test.ts

@@ -66,6 +66,52 @@ function makeApis(
 }
 
 describe("handleSessionStop with auto-return", () => {
+  it("leaves the session untouched when cancellation arrived before teardown", async () => {
+    const aw = fakeAgentWindow([100]);
+    const sm = new SessionManager({ agentWindow: aw });
+    await sm.start("aa11");
+    const cdp = { detachSession: vi.fn(async () => {}) };
+    const controller = new AbortController();
+    controller.abort();
+
+    const res = await handleSessionStop(
+      sm,
+      { session_id: "aa11" },
+      { cdp, signal: controller.signal },
+    );
+
+    expect(res).toMatchObject({ code: "cancelled" });
+    expect(sm.has("aa11")).toBe(true);
+    expect(cdp.detachSession).not.toHaveBeenCalled();
+    expect(aw.remove).not.toHaveBeenCalled();
+  });
+
+  it("does not close the Agent Window when cancellation lands during CDP detach", async () => {
+    const aw = fakeAgentWindow([100]);
+    const sm = new SessionManager({ agentWindow: aw });
+    await sm.start("aa11");
+    const controller = new AbortController();
+    let finishDetach: () => void = () => {};
+    const detach = new Promise<void>((resolve) => {
+      finishDetach = resolve;
+    });
+    const cdp = { detachSession: vi.fn(() => detach) };
+
+    const stopping = handleSessionStop(
+      sm,
+      { session_id: "aa11" },
+      { cdp, signal: controller.signal },
+    );
+    await vi.waitFor(() => expect(cdp.detachSession).toHaveBeenCalledWith("aa11"));
+    controller.abort();
+    finishDetach();
+
+    const res = await stopping;
+    expect(res).toMatchObject({ code: "cancelled" });
+    expect(sm.has("aa11")).toBe(true);
+    expect(aw.remove).not.toHaveBeenCalled();
+  });
+
   it("returns every borrowed tab and closes the Agent Window in the right order", async () => {
     const aw = fakeAgentWindow([100]);
     const sm = new SessionManager({ agentWindow: aw });

apps/extension/src/tools/dispatcher.ts

@@ -290,6 +290,7 @@ export class ToolDispatcher {
         await this.releaseHoverLatch((req.params as SessionStopParams).session_id);
         return handleSessionStop(this.sessions, req.params as SessionStopParams, {
           cdp: this.cdp,
+          signal,
         });
       }
       case "tool.tab_list":

apps/extension/src/tools/session.ts

@@ -76,6 +76,7 @@ export interface SessionStopResult {
 }
 
 export interface SessionStopDeps {
+  signal?: AbortSignal;
   cdp?: {
     detachSession(sessionId: string): Promise<void>;
   };
@@ -177,6 +178,9 @@ export async function handleSessionStop(
       message: `session ${params.session_id} unknown`,
     };
   }
+  if (deps.signal?.aborted) {
+    return { code: "cancelled", message: "session_stop aborted before teardown" };
+  }
 
   // Step 1: auto-return borrowed tabs. Iterate over a snapshot of the
   // ids so deletions during iteration do not break the Map iterator.
@@ -187,6 +191,7 @@ export async function handleSessionStop(
     try {
       const tabManagement = {
         ...(deps.tabManagement ?? {}),
+        signal: deps.signal,
         isAgentWindowId:
           deps.tabManagement?.isAgentWindowId ??
           ((windowId: number) => manager.findByWindowId(windowId) !== null),
@@ -214,6 +219,17 @@ export async function handleSessionStop(
     }
   }
 
+  if (deps.signal?.aborted) {
+    const nonCancelFailures = returnFailures?.filter((failure) => failure.code !== "cancelled");
+    if (nonCancelFailures && nonCancelFailures.length > 0) {
+      return {
+        ...(returnedTabIds.length > 0 ? { returned_tab_ids: returnedTabIds } : {}),
+        return_failures: nonCancelFailures,
+      };
+    }
+    return { code: "cancelled", message: "session_stop aborted during tab return" };
+  }
+
   const result: SessionStopResult = {};
   if (returnedTabIds.length > 0) result.returned_tab_ids = returnedTabIds;
   if (returnFailures && returnFailures.length > 0) {
@@ -233,6 +249,10 @@ export async function handleSessionStop(
   // Step 3: detach CDP sessions this session opened (no-op if none).
   await deps.cdp?.detachSession(params.session_id);
 
+  if (deps.signal?.aborted) {
+    return { code: "cancelled", message: "session_stop aborted before window close" };
+  }
+
   // Step 4: close the Agent Window and drop the context.
   await manager.stop(params.session_id);
 

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

crates/bsk-cli/tests/per_session_queue.rs

@@ -325,6 +325,7 @@ async fn session_stop_fast_fails_while_tool_is_in_flight() {
         &state.session_interrupts,
         &sid,
         Duration::from_secs(5),
+        None,
     )
     .await;
     assert!(
@@ -366,6 +367,7 @@ async fn session_stop_fast_fails_while_tool_is_in_flight() {
         &state.session_interrupts,
         &sid,
         Duration::from_secs(5),
+        None,
     )
     .await
     .expect("session.stop ok after tool completes");

crates/bsk-cli/tests/sessions_ipc.rs

@@ -10,9 +10,14 @@ use std::time::Duration;
 
 use bsk::daemon::{self, DaemonConfig};
 use bsk::ipc_client::IpcClient;
+use bsk_protocol::cancel::{CancelParams, CancelResult};
 use bsk_protocol::system::{HandshakeParams, HandshakeResult, StatusResult};
-use bsk_protocol::tools::{SessionStartParams, SessionStartResult, SessionStopParams};
-use bsk_protocol::{BrowserPeerInfo, Frame, Method, RequestFrame, ResponseBody, ResponseFrame};
+use bsk_protocol::tools::{
+    SessionStartParams, SessionStartResult, SessionStopParams, SessionStopResult,
+};
+use bsk_protocol::{
+    BrowserPeerInfo, ErrorCode, Frame, Method, RequestFrame, ResponseBody, ResponseFrame, RpcError,
+};
 use futures_util::{SinkExt, StreamExt};
 use rand::Rng;
 use tokio_tungstenite::tungstenite::handshake::client::generate_key;
@@ -23,6 +28,9 @@ use support::{wait_for_browser_count, wait_for_no_sessions};
 
 const TEST_EXT_ID: &str = "abcdefghijklmnopabcdefghijklmnop";
 
+type TestWs =
+    tokio_tungstenite::WebSocketStream<tokio_tungstenite::MaybeTlsStream<tokio::net::TcpStream>>;
+
 fn tempfile_path(prefix: &str) -> PathBuf {
     let mut p = std::env::temp_dir();
     let mut rng = rand::thread_rng();
@@ -56,9 +64,7 @@ async fn spawn_daemon_with_session_idle(session_idle: Duration) -> (daemon::Daem
     (handle, sock)
 }
 
-async fn connect_ext(
-    addr: std::net::SocketAddr,
-) -> tokio_tungstenite::WebSocketStream<tokio_tungstenite::MaybeTlsStream<tokio::net::TcpStream>> {
+async fn connect_ext(addr: std::net::SocketAddr) -> TestWs {
     let origin = format!("chrome-extension://{TEST_EXT_ID}");
     let url = format!("ws://{addr}/");
     let req = Request::builder()
@@ -76,11 +82,7 @@ async fn connect_ext(
     ws
 }
 
-async fn handshake_as_ext(
-    ws: &mut tokio_tungstenite::WebSocketStream<
-        tokio_tungstenite::MaybeTlsStream<tokio::net::TcpStream>,
-    >,
-) -> HandshakeResult {
+async fn handshake_as_ext(ws: &mut TestWs) -> HandshakeResult {
     let params = HandshakeParams {
         client: "browser-skill-extension".into(),
         version: "0.1.0-dev.0".parse().unwrap(),
@@ -114,6 +116,128 @@ async fn handshake_as_ext(
     }
 }
 
+async fn next_extension_request(ws: &mut TestWs) -> RequestFrame {
+    loop {
+        let message = ws.next().await.expect("daemon closed websocket").unwrap();
+        match message {
+            Message::Text(text) => match serde_json::from_str::<Frame>(&text).unwrap() {
+                Frame::Request(request) => return request,
+                Frame::Response(_) | Frame::Event(_) => continue,
+            },
+            Message::Ping(payload) => ws.send(Message::Pong(payload)).await.unwrap(),
+            Message::Close(_) => panic!("daemon closed websocket before sending request"),
+            _ => {}
+        }
+    }
+}
+
+async fn send_extension_response(ws: &mut TestWs, response: ResponseFrame) {
+    ws.send(Message::Text(serde_json::to_string(&response).unwrap()))
+        .await
+        .unwrap();
+}
+
+async fn acknowledge_extension_cancel(ws: &mut TestWs, cancel: RequestFrame, target_rpc_id: &str) {
+    assert_eq!(cancel.method, Method::Cancel);
+    let params: CancelParams = serde_json::from_value(cancel.params.unwrap()).unwrap();
+    assert_eq!(params.rpc_id, target_rpc_id);
+    send_extension_response(
+        ws,
+        ResponseFrame {
+            id: cancel.id,
+            body: ResponseBody::Ok(serde_json::to_value(CancelResult { cancelled: true }).unwrap()),
+        },
+    )
+    .await;
+}
+
+async fn respond_to_aborted_start(
+    ws: &mut TestWs,
+    start_seen: Option<tokio::sync::oneshot::Sender<()>>,
+    agent_window_id: i64,
+) {
+    let start = next_extension_request(ws).await;
+    assert_eq!(start.method, Method::ToolSessionStart);
+    let params: SessionStartParams = serde_json::from_value(start.params.clone().unwrap()).unwrap();
+    if let Some(start_seen) = start_seen {
+        start_seen.send(()).unwrap();
+    }
+
+    let cancel = next_extension_request(ws).await;
+    acknowledge_extension_cancel(ws, cancel, &start.id).await;
+
+    // Model the hardest race: chrome.windows.create completed just as
+    // cancellation arrived, so the extension's original call succeeds.
+    send_extension_response(
+        ws,
+        ResponseFrame {
+            id: start.id,
+            body: ResponseBody::Ok(
+                serde_json::to_value(SessionStartResult {
+                    agent_window_id: Some(agent_window_id),
+                })
+                .unwrap(),
+            ),
+        },
+    )
+    .await;
+
+    let rollback = next_extension_request(ws).await;
+    assert_eq!(rollback.method, Method::ToolSessionStop);
+    let rollback_params: SessionStopParams =
+        serde_json::from_value(rollback.params.unwrap()).unwrap();
+    assert_eq!(rollback_params.session_id, params.session_id);
+    send_extension_response(
+        ws,
+        ResponseFrame {
+            id: rollback.id,
+            body: ResponseBody::Ok(serde_json::to_value(SessionStopResult::default()).unwrap()),
+        },
+    )
+    .await;
+}
+
+async fn respond_to_aborted_stop(
+    ws: &mut TestWs,
+    stop_seen: Option<tokio::sync::oneshot::Sender<()>>,
+) {
+    let start = next_extension_request(ws).await;
+    assert_eq!(start.method, Method::ToolSessionStart);
+    send_extension_response(
+        ws,
+        ResponseFrame {
+            id: start.id,
+            body: ResponseBody::Ok(
+                serde_json::to_value(SessionStartResult {
+                    agent_window_id: Some(4242),
+                })
+                .unwrap(),
+            ),
+        },
+    )
+    .await;
+
+    let stop = next_extension_request(ws).await;
+    assert_eq!(stop.method, Method::ToolSessionStop);
+    if let Some(stop_seen) = stop_seen {
+        stop_seen.send(()).unwrap();
+    }
+    let cancel = next_extension_request(ws).await;
+    acknowledge_extension_cancel(ws, cancel, &stop.id).await;
+    send_extension_response(
+        ws,
+        ResponseFrame {
+            id: stop.id,
+            body: ResponseBody::Err(RpcError {
+                code: ErrorCode::Cancelled,
+                message: "session_stop aborted before window close".into(),
+                data: None,
+            }),
+        },
+    )
+    .await;
+}
+
 #[tokio::test]
 async fn session_start_stop_round_trip_via_ipc() {
     let (handle, sock) = spawn_daemon().await;
@@ -250,6 +374,240 @@ async fn session_start_stop_round_trip_via_ipc() {
     handle.shutdown().await;
 }
 
+#[tokio::test]
+async fn cancelling_session_start_rolls_back_a_late_extension_success() {
+    const START_RPC_ID: &str = "session-start-cancel";
+
+    let (handle, sock) = spawn_daemon().await;
+    let mut ws = connect_ext(handle.ws_addr()).await;
+    let _ = handshake_as_ext(&mut ws).await;
+    let (start_seen_tx, start_seen_rx) = tokio::sync::oneshot::channel();
+    let (release_tx, release_rx) = tokio::sync::oneshot::channel();
+
+    let responder = tokio::spawn(async move {
+        respond_to_aborted_start(&mut ws, Some(start_seen_tx), 4242).await;
+        let _ = release_rx.await;
+    });
+
+    let mut start_ipc = IpcClient::connect(&sock).await.unwrap();
+    let start_call = tokio::spawn(async move {
+        let outcome: Result<serde_json::Value, RpcError> = start_ipc
+            .call_with_id(
+                START_RPC_ID.to_string(),
+                Method::SessionStart,
+                Some(serde_json::json!({ "focused": false })),
+                Duration::from_secs(5),
+            )
+            .await
+            .unwrap();
+        outcome
+    });
+
+    start_seen_rx.await.unwrap();
+    let mut cancel_ipc = IpcClient::connect(&sock).await.unwrap();
+    let cancel: CancelResult = cancel_ipc
+        .call(
+            "cancel-start",
+            Method::Cancel,
+            Some(CancelParams {
+                rpc_id: START_RPC_ID.to_string(),
+            }),
+            Duration::from_secs(2),
+        )
+        .await
+        .unwrap()
+        .unwrap();
+    assert!(cancel.cancelled);
+
+    let start_error = start_call.await.unwrap().unwrap_err();
+    assert_eq!(start_error.code, ErrorCode::Cancelled);
+
+    let mut status_ipc = IpcClient::connect(&sock).await.unwrap();
+    let status: StatusResult = status_ipc
+        .call::<(), _>(
+            "status-after-start-cancel",
+            Method::SystemStatus,
+            None,
+            Duration::from_secs(2),
+        )
+        .await
+        .unwrap()
+        .unwrap();
+    assert!(status.sessions.is_empty());
+    assert_eq!(status.browsers[0].session_count, 0);
+
+    let _ = release_tx.send(());
+    responder.await.unwrap();
+    handle.shutdown().await;
+}
+
+#[tokio::test]
+async fn timing_out_session_start_rolls_back_a_late_extension_success() {
+    let (handle, _sock) = spawn_daemon().await;
+    let state = handle.state();
+    let mut ws = connect_ext(handle.ws_addr()).await;
+    let _ = handshake_as_ext(&mut ws).await;
+    let (release_tx, release_rx) = tokio::sync::oneshot::channel();
+
+    let responder = tokio::spawn(async move {
+        respond_to_aborted_start(&mut ws, None, 4343).await;
+        let _ = release_rx.await;
+    });
+
+    let start = bsk::daemon::sessions::start_session(
+        &state.browsers,
+        &state.sessions,
+        &state.tool_queues,
+        None,
+        bsk::daemon::sessions::AgentWindowOptions::default(),
+        Duration::ZERO,
+        Duration::from_millis(20),
+        None,
+    )
+    .await;
+    assert!(matches!(
+        start,
+        Err(bsk::daemon::sessions::StartSessionError::Timeout)
+    ));
+    assert!(state.sessions.is_empty());
+
+    let _ = release_tx.send(());
+    responder.await.unwrap();
+    handle.shutdown().await;
+}
+
+#[tokio::test]
+async fn cancelling_session_stop_keeps_the_session_retryable() {
+    const STOP_RPC_ID: &str = "session-stop-cancel";
+
+    let (handle, sock) = spawn_daemon().await;
+    let mut ws = connect_ext(handle.ws_addr()).await;
+    let _ = handshake_as_ext(&mut ws).await;
+    let (stop_seen_tx, stop_seen_rx) = tokio::sync::oneshot::channel();
+    let (release_tx, release_rx) = tokio::sync::oneshot::channel();
+
+    let responder = tokio::spawn(async move {
+        respond_to_aborted_stop(&mut ws, Some(stop_seen_tx)).await;
+        let _ = release_rx.await;
+    });
+
+    let mut start_ipc = IpcClient::connect(&sock).await.unwrap();
+    let start: serde_json::Value = start_ipc
+        .call(
+            "start-before-stop-cancel",
+            Method::SessionStart,
+            Some(serde_json::json!({ "focused": false })),
+            Duration::from_secs(5),
+        )
+        .await
+        .unwrap()
+        .unwrap();
+    let session_id = start["session_id"].as_str().unwrap().to_string();
+
+    let mut stop_ipc = IpcClient::connect(&sock).await.unwrap();
+    let stop_session_id = session_id.clone();
+    let stop_call = tokio::spawn(async move {
+        let outcome: Result<serde_json::Value, RpcError> = stop_ipc
+            .call_with_id(
+                STOP_RPC_ID.to_string(),
+                Method::SessionStop,
+                Some(serde_json::json!({
+                    "session_id": stop_session_id,
+                    "all": false,
+                })),
+                Duration::from_secs(5),
+            )
+            .await
+            .unwrap();
+        outcome
+    });
+
+    stop_seen_rx.await.unwrap();
+    let mut cancel_ipc = IpcClient::connect(&sock).await.unwrap();
+    let cancel: CancelResult = cancel_ipc
+        .call(
+            "cancel-stop",
+            Method::Cancel,
+            Some(CancelParams {
+                rpc_id: STOP_RPC_ID.to_string(),
+            }),
+            Duration::from_secs(2),
+        )
+        .await
+        .unwrap()
+        .unwrap();
+    assert!(cancel.cancelled);
+
+    let stop_error = stop_call.await.unwrap().unwrap_err();
+    assert_eq!(stop_error.code, ErrorCode::Cancelled);
+
+    let mut status_ipc = IpcClient::connect(&sock).await.unwrap();
+    let status: StatusResult = status_ipc
+        .call::<(), _>(
+            "status-after-stop-cancel",
+            Method::SystemStatus,
+            None,
+            Duration::from_secs(2),
+        )
+        .await
+        .unwrap()
+        .unwrap();
+    assert_eq!(status.sessions.len(), 1);
+    assert_eq!(status.sessions[0].session_id, session_id);
+
+    let _ = release_tx.send(());
+    responder.await.unwrap();
+    handle.shutdown().await;
+}
+
+#[tokio::test]
+async fn timing_out_session_stop_cancels_extension_teardown_and_keeps_session() {
+    let (handle, sock) = spawn_daemon().await;
+    let state = handle.state();
+    let mut ws = connect_ext(handle.ws_addr()).await;
+    let _ = handshake_as_ext(&mut ws).await;
+    let (release_tx, release_rx) = tokio::sync::oneshot::channel();
+
+    let responder = tokio::spawn(async move {
+        respond_to_aborted_stop(&mut ws, None).await;
+        let _ = release_rx.await;
+    });
+
+    let mut start_ipc = IpcClient::connect(&sock).await.unwrap();
+    let start: serde_json::Value = start_ipc
+        .call(
+            "start-before-stop-timeout",
+            Method::SessionStart,
+            Some(serde_json::json!({ "focused": false })),
+            Duration::from_secs(5),
+        )
+        .await
+        .unwrap()
+        .unwrap();
+    let session_id =
+        bsk::daemon::sessions::SessionId(start["session_id"].as_str().unwrap().to_string());
+
+    let stop = bsk::daemon::sessions::stop_session(
+        &state.browsers,
+        &state.sessions,
+        &state.tool_queues,
+        &state.session_interrupts,
+        &session_id,
+        Duration::from_millis(20),
+        None,
+    )
+    .await;
+    assert!(matches!(
+        stop,
+        Err(bsk::daemon::sessions::StopSessionError::Timeout)
+    ));
+    assert!(state.sessions.get(&session_id).is_some());
+
+    let _ = release_tx.send(());
+    responder.await.unwrap();
+    handle.shutdown().await;
+}
+
 #[tokio::test]
 async fn session_idle_timeout_stops_and_unregisters_session() {
     let (handle, _sock) = spawn_daemon_with_session_idle(Duration::from_millis(50)).await;