From d1a3fb815a429f3e5769b8e5313b2a74140b8c3f Mon Sep 17 00:00:00 2001 From: Hunter Sadler Date: Fri, 4 Sep 2026 12:40:02 -0600 Subject: [PATCH 1/4] Preserve resume bootstrap events Retain routed resume notifications until the first Session subscription catches up, then switch atomically to normal live broadcast delivery. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- rust/src/session.rs | 75 +++++++-- rust/src/subscription.rs | 245 ++++++++++++++++++++++++++-- rust/tests/prepared_session_test.rs | 128 ++++++++++++++- 3 files changed, 416 insertions(+), 32 deletions(-) diff --git a/rust/src/session.rs b/rust/src/session.rs index 0e64d6061c..2edfc5fb47 100644 --- a/rust/src/session.rs +++ b/rust/src/session.rs @@ -328,6 +328,9 @@ pub struct Session { open_canvases: Arc>>, /// Broadcast channel for runtime event subscribers — see [`Session::subscribe`]. event_tx: tokio::sync::broadcast::Sender, + /// Resume-only queue that retains routed events until the first + /// post-resume subscriber catches up and activates live broadcast delivery. + resume_bootstrap: Option>, github_token_registration: ParkingLotMutex>, /// Identity of this session's router registration. @@ -419,8 +422,19 @@ impl Session { /// loop or any combinator from `tokio_stream::StreamExt` / /// `futures::StreamExt`. /// - /// Each subscriber maintains its own queue. If a consumer cannot keep - /// up, the oldest events are dropped and `recv` returns + /// On a session returned by [`Client::resume_session`], the first + /// subscription also receives every routed event retained while resume + /// startup had no active [`PreparedSession`] subscriber. That bootstrap + /// prefix is lossless and ordered before live events. It is a one-shot + /// handoff: later subscriptions begin with live delivery, and dropping + /// the first subscription before draining it discards its remaining + /// bootstrap events. + /// + /// Bootstrap retention is unbounded until the first subscriber catches + /// up, so resume consumers that need events should subscribe promptly. + /// After the bootstrap handoff, each subscriber maintains its own finite + /// queue. If a consumer cannot keep + /// up, the oldest live events are dropped and `recv` returns /// [`RecvErrorKind::Lagged`](crate::subscription::RecvErrorKind::Lagged) /// reporting the count of skipped events. Slow consumers do not block /// the session's event loop. @@ -438,7 +452,10 @@ impl Session { /// # } /// ``` pub fn subscribe(&self) -> crate::subscription::EventSubscription { - crate::subscription::EventSubscription::new(self.event_tx.subscribe()) + match &self.resume_bootstrap { + Some(bootstrap) => bootstrap.subscribe(&self.event_tx), + None => crate::subscription::EventSubscription::new(self.event_tx.subscribe()), + } } /// The underlying Client (for advanced use cases). @@ -1152,11 +1169,16 @@ impl Client { /// /// # Event delivery /// - /// Equivalent to `prepare_resume_session(config)?.start().await`, and - /// carries the same startup-event caveat documented on - /// [`create_session`](Self::create_session). Use - /// [`prepare_resume_session`](Self::prepare_resume_session) when - /// startup events matter. + /// When no active [`PreparedSession`] subscription exists at startup, + /// routed events emitted during resume are retained in an ordered, + /// unbounded bootstrap queue. The first + /// [`Session::subscribe`] call receives that complete prefix before + /// switching atomically to normal live delivery. Later subscribers are + /// live-only. + /// + /// Use [`prepare_resume_session`](Self::prepare_resume_session) when the + /// event consumer must be installed before protocol activity begins or + /// when multiple startup observers are required. pub async fn resume_session(&self, config: ResumeSessionConfig) -> Result { self.prepare_resume_session(config)?.start().await } @@ -1428,6 +1450,7 @@ impl Client { capabilities.clone(), open_canvases.clone(), event_tx.clone(), + None, shutdown.clone(), external_tools_shutdown.clone(), ); @@ -1466,6 +1489,7 @@ impl Client { capabilities, open_canvases, event_tx, + resume_bootstrap: None, github_token_registration: ParkingLotMutex::new(github_token_registration), registration_token, }; @@ -1624,6 +1648,11 @@ impl Client { let capabilities = Arc::new(parking_lot::RwLock::new(SessionCapabilities::default())); let setup_start = Instant::now(); + // An active prepared subscription already owns startup delivery. The + // implicit queue is only needed by the compatibility resume wrapper, + // where Session::subscribe cannot be called until startup returns. + let resume_bootstrap = + (event_tx.receiver_count() == 0).then(crate::subscription::ResumeBootstrap::new); let registration = self.register_session(&session_id); let registration_token = registration.token; let channels = registration.channels; @@ -1645,6 +1674,7 @@ impl Client { capabilities.clone(), open_canvases.clone(), event_tx.clone(), + resume_bootstrap.clone(), shutdown.clone(), external_tools_shutdown.clone(), ); @@ -1757,6 +1787,7 @@ impl Client { capabilities, open_canvases, event_tx, + resume_bootstrap, github_token_registration: ParkingLotMutex::new(github_token_registration), registration_token, }; @@ -1815,16 +1846,20 @@ impl Client { /// /// # Buffering /// -/// The broadcast buffer is finite — -/// [`DEFAULT_EVENT_BUFFER_CAPACITY`] unless +/// Subscriptions taken directly from this prepared handle use the finite +/// broadcast buffer — [`DEFAULT_EVENT_BUFFER_CAPACITY`] unless /// [`SessionConfig::event_buffer_capacity`] / -/// [`ResumeSessionConfig::event_buffer_capacity`] overrides it. Subscribers -/// that fall behind observe +/// [`ResumeSessionConfig::event_buffer_capacity`] overrides it. Those +/// subscribers that fall behind observe /// [`Lagged`](crate::subscription::Lagged) instead of applying backpressure /// to the event loop. Consumers that need a lossless view of a large /// startup burst must either configure a capacity that covers it or drain /// the subscription concurrently with [`start`](Self::start). /// +/// If a resume starts with no active prepared subscription, the eventual +/// [`Session`] instead retains routed startup events in the resume bootstrap +/// queue documented on [`Session::subscribe`]. +/// /// # Server-assigned session IDs /// /// For cloud sessions without a caller-supplied session ID, the CLI assigns @@ -2053,6 +2088,7 @@ fn spawn_event_loop( capabilities: Arc>, open_canvases: Arc>>, event_tx: tokio::sync::broadcast::Sender, + resume_bootstrap: Option>, shutdown: CancellationToken, external_tools_shutdown: CancellationToken, ) -> JoinHandle<()> { @@ -2094,7 +2130,7 @@ fn spawn_event_loop( _ = shutdown.cancelled() => break, Some(notification) = notifications.recv() => { handle_notification( - &session_id, &client, &handlers, &command_handlers, notification, &idle_waiter, &capabilities, &open_canvases, &event_tx, &shutdown, &external_tools_shutdown, &pending_external_tools, + &session_id, &client, &handlers, &command_handlers, notification, &idle_waiter, &capabilities, &open_canvases, &event_tx, resume_bootstrap.as_ref(), &shutdown, &external_tools_shutdown, &pending_external_tools, ).await; } Some(request) = requests.recv() => { @@ -2260,6 +2296,7 @@ async fn handle_notification( capabilities: &Arc>, open_canvases: &Arc>>, event_tx: &tokio::sync::broadcast::Sender, + resume_bootstrap: Option<&Arc>, shutdown: &CancellationToken, external_tools_shutdown: &CancellationToken, pending_external_tools: &PendingExternalTools, @@ -2361,10 +2398,14 @@ async fn handle_notification( } } - // Fan out the event to runtime subscribers (`Session::subscribe`). `send` - // only errors when there are no receivers, which is the normal case - // before any consumer subscribes. - let _ = event_tx.send(event.clone()); + // Resume startup queues routed events until the first post-resume + // subscriber catches up. All other paths retain the existing bounded + // broadcast behavior. + if let Some(bootstrap) = resume_bootstrap { + bootstrap.publish(event_tx, event.clone()); + } else { + let _ = event_tx.send(event.clone()); + } tracing::debug!( elapsed_ms = dispatch_start.elapsed().as_millis(), diff --git a/rust/src/subscription.rs b/rust/src/subscription.rs index c3fc83b8b9..156f873130 100644 --- a/rust/src/subscription.rs +++ b/rust/src/subscription.rs @@ -19,18 +19,27 @@ //! also works for callers who don't need the [`Stream`](tokio_stream::Stream) //! surface. //! -//! # Lag policy +//! # Resume bootstrap and lag policy //! -//! Each subscriber maintains its own internal queue. If a consumer cannot -//! keep up, the oldest events are dropped and the next call yields +//! The first subscription on a session returned by +//! [`Client::resume_session`](crate::Client::resume_session) may begin with +//! a lossless, ordered bootstrap prefix retained during resume startup. Once +//! that subscriber catches up, delivery switches atomically to the normal +//! live broadcast stream. +//! +//! Each live subscriber maintains its own finite queue. If a consumer cannot +//! keep up, the oldest live events are dropped and the next call yields //! [`Lagged`](crate::subscription::Lagged) reporting how many events were skipped. //! Slow subscribers do not block the producer. +use std::collections::VecDeque; use std::fmt; use std::pin::Pin; +use std::sync::Arc; use std::task::{Context, Poll}; -use tokio::sync::broadcast::Receiver; +use parking_lot::Mutex; +use tokio::sync::broadcast::{Receiver, Sender}; use tokio_stream::wrappers::BroadcastStream; use tokio_stream::wrappers::errors::BroadcastStreamRecvError; use tokio_stream::{Stream, StreamExt as _}; @@ -135,6 +144,162 @@ impl From for RecvError { } } +enum ResumeBootstrapState { + Unclaimed(VecDeque), + Claimed(VecDeque), + Disabled, +} + +/// Lossless, one-shot queue for routed events emitted before the first +/// post-resume session subscription catches up. +pub(crate) struct ResumeBootstrap { + state: Mutex, +} + +impl ResumeBootstrap { + pub(crate) fn new() -> Arc { + Arc::new(Self { + state: Mutex::new(ResumeBootstrapState::Unclaimed(VecDeque::new())), + }) + } + + pub(crate) fn publish(&self, event_tx: &Sender, event: SessionEvent) { + let mut state = self.state.lock(); + match &mut *state { + ResumeBootstrapState::Unclaimed(events) | ResumeBootstrapState::Claimed(events) => { + events.push_back(event) + } + ResumeBootstrapState::Disabled => { + let _ = event_tx.send(event); + } + } + } + + pub(crate) fn subscribe( + self: &Arc, + event_tx: &Sender, + ) -> EventSubscription { + let mut state = self.state.lock(); + let receiver = event_tx.subscribe(); + let bootstrap = match &mut *state { + ResumeBootstrapState::Unclaimed(events) => { + let events = std::mem::take(events); + *state = ResumeBootstrapState::Claimed(events); + Some(self.clone()) + } + ResumeBootstrapState::Claimed(_) | ResumeBootstrapState::Disabled => None, + }; + EventSubscription::with_bootstrap(receiver, bootstrap) + } + + fn pop(&self) -> Option { + let mut state = self.state.lock(); + let ResumeBootstrapState::Claimed(events) = &mut *state else { + return None; + }; + if let Some(event) = events.pop_front() { + return Some(event); + } + *state = ResumeBootstrapState::Disabled; + None + } + + fn abandon(&self) { + let mut state = self.state.lock(); + if matches!(*state, ResumeBootstrapState::Claimed(_)) { + *state = ResumeBootstrapState::Disabled; + } + } +} + +/// Subscription to runtime events for a single +/// [`Session`](crate::session::Session). +/// +/// Created by [`Session::subscribe`](crate::session::Session::subscribe). +/// Implements [`Stream`] yielding `Result`. +/// Drop the value to unsubscribe; there is no separate cancel handle. +#[must_use = "subscriptions are inert until polled"] +pub struct EventSubscription { + inner: BroadcastStream, + bootstrap: Option>, +} + +impl EventSubscription { + pub(crate) fn new(rx: Receiver) -> Self { + Self::with_bootstrap(rx, None) + } + + fn with_bootstrap(rx: Receiver, bootstrap: Option>) -> Self { + Self { + inner: BroadcastStream::new(rx), + bootstrap, + } + } + + fn next_bootstrap_event(&mut self) -> Option { + let event = self + .bootstrap + .as_ref() + .and_then(|bootstrap| bootstrap.pop()); + if event.is_none() { + self.bootstrap = None; + } + event + } + + /// Receive the next event. + /// + /// Returns: + /// + /// - `Ok(event)` for the next delivered event. + /// - `Err(`[`RecvError`]`)` with [`RecvError::kind()`] [`RecvErrorKind::Lagged`] if the subscriber fell behind; + /// call `recv` again to continue from the next live event. + /// - `Err(`[`RecvError`]`)` with [`RecvError::kind()`] [`RecvErrorKind::Closed`] once the producer is gone. + /// + /// # Cancel safety + /// + /// **Cancel-safe.** Bootstrap events are removed before the future's + /// first suspension point. Once live delivery begins, this wraps a + /// `tokio::sync::broadcast::Receiver` via `BroadcastStream`, which is + /// cancel-safe by design. + pub async fn recv(&mut self) -> Result { + if let Some(event) = self.next_bootstrap_event() { + return Ok(event); + } + match self.inner.next().await { + Some(Ok(event)) => Ok(event), + Some(Err(BroadcastStreamRecvError::Lagged(n))) => Err(Lagged(n).into()), + None => Err(RecvErrorKind::Closed.into()), + } + } +} + +impl Stream for EventSubscription { + type Item = Result; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + if let Some(event) = self.next_bootstrap_event() { + return Poll::Ready(Some(Ok(event))); + } + match Pin::new(&mut self.inner).poll_next(cx) { + Poll::Ready(Some(Ok(event))) => Poll::Ready(Some(Ok(event))), + Poll::Ready(Some(Err(BroadcastStreamRecvError::Lagged(n)))) => { + Poll::Ready(Some(Err(Lagged(n)))) + } + Poll::Ready(None) => Poll::Ready(None), + Poll::Pending => Poll::Pending, + } + } +} + +impl Drop for EventSubscription { + fn drop(&mut self) { + if let Some(bootstrap) = self.bootstrap.take() { + bootstrap.abandon(); + } + } +} + macro_rules! define_subscription { ( $(#[$meta:meta])* @@ -200,16 +365,6 @@ macro_rules! define_subscription { }; } -define_subscription! { - /// Subscription to runtime events for a single - /// [`Session`](crate::session::Session). - /// - /// Created by [`Session::subscribe`](crate::session::Session::subscribe). - /// Implements [`Stream`] yielding `Result`. - /// Drop the value to unsubscribe; there is no separate cancel handle. - EventSubscription, SessionEvent -} - define_subscription! { /// Subscription to lifecycle events on a [`Client`](crate::Client). /// @@ -285,4 +440,66 @@ mod tests { assert_eq!(next.unwrap().unwrap().id, "a"); assert!(sub.next().await.is_none()); } + + #[tokio::test] + async fn resume_bootstrap_is_lossless_and_ordered_before_live_events() { + let (tx, _) = broadcast::channel(1); + let bootstrap = ResumeBootstrap::new(); + for index in 0..600 { + bootstrap.publish(&tx, make_event(&format!("bootstrap-{index}"))); + } + + let mut sub = bootstrap.subscribe(&tx); + bootstrap.publish(&tx, make_event("bootstrap-after-subscribe")); + + for index in 0..600 { + assert_eq!(sub.recv().await.unwrap().id, format!("bootstrap-{index}")); + } + assert_eq!(sub.recv().await.unwrap().id, "bootstrap-after-subscribe"); + + assert!(sub.next_bootstrap_event().is_none()); + bootstrap.publish(&tx, make_event("live")); + assert_eq!(sub.recv().await.unwrap().id, "live"); + } + + #[tokio::test] + async fn only_first_subscriber_claims_resume_bootstrap() { + let (tx, _) = broadcast::channel(8); + let bootstrap = ResumeBootstrap::new(); + bootstrap.publish(&tx, make_event("bootstrap")); + + let mut first = bootstrap.subscribe(&tx); + let mut second = bootstrap.subscribe(&tx); + + assert_eq!(first.recv().await.unwrap().id, "bootstrap"); + assert!(first.next_bootstrap_event().is_none()); + bootstrap.publish(&tx, make_event("live")); + + assert_eq!(first.recv().await.unwrap().id, "live"); + assert_eq!(second.recv().await.unwrap().id, "live"); + } + + #[tokio::test] + async fn dropping_bootstrap_owner_activates_live_delivery() { + let (tx, _) = broadcast::channel(8); + let bootstrap = ResumeBootstrap::new(); + bootstrap.publish(&tx, make_event("discarded-bootstrap")); + + let first = bootstrap.subscribe(&tx); + let mut second = bootstrap.subscribe(&tx); + drop(first); + + bootstrap.publish(&tx, make_event("live")); + assert_eq!(second.recv().await.unwrap().id, "live"); + } + + #[tokio::test] + async fn ordinary_subscription_does_not_replay_events_without_a_receiver() { + let (tx, _) = broadcast::channel(8); + assert!(tx.send(make_event("before-subscribe")).is_err()); + + let mut sub = EventSubscription::new(tx.subscribe()); + tx.send(make_event("live")).unwrap(); + assert_eq!(sub.recv().await.unwrap().id, "live"); + } } diff --git a/rust/tests/prepared_session_test.rs b/rust/tests/prepared_session_test.rs index 3fefccabb6..18b9c0ba1f 100644 --- a/rust/tests/prepared_session_test.rs +++ b/rust/tests/prepared_session_test.rs @@ -350,6 +350,32 @@ async fn prepared_resume_delivers_pre_response_burst() { let session = timeout(TIMEOUT, start).await.unwrap().unwrap().unwrap(); expect_startup_burst(&mut events).await; + let mut late = session.subscribe(); + assert!( + timeout(QUIET, late.recv()).await.is_err(), + "an active prepared subscriber must prevent implicit bootstrap replay" + ); + server + .send_event(session_id.as_str(), "evt-live", "assistant.message", false) + .await; + assert_eq!( + timeout(TIMEOUT, events.recv()) + .await + .unwrap() + .unwrap() + .id + .as_str(), + "evt-live" + ); + assert_eq!( + timeout(TIMEOUT, late.recv()) + .await + .unwrap() + .unwrap() + .id + .as_str(), + "evt-live" + ); drop(session); } @@ -567,6 +593,65 @@ async fn cancelled_prepared_resume_cleans_up_and_allows_retry() { drop(session); } +#[tokio::test] +async fn cancelled_resume_wrapper_cleans_implicit_bootstrap_and_allows_retry() { + let (client, mut server) = make_client(); + let session_id = SessionId::new("resume-wrapper-cancel"); + + let start = tokio::spawn({ + let client = client.clone(); + let session_id = session_id.clone(); + async move { + client + .resume_session(ResumeSessionConfig::new(session_id)) + .await + } + }); + + let resume_req = server.read_request().await; + assert_eq!(resume_req["method"], "session.resume"); + start.abort(); + let _ = start.await; + await_no_registrations(&client).await; + + let retry = tokio::spawn({ + let client = client.clone(); + let session_id = session_id.clone(); + async move { + client + .resume_session(ResumeSessionConfig::new(session_id)) + .await + } + }); + let retry_req = server.read_request().await; + assert_eq!(retry_req["method"], "session.resume"); + server + .respond(&retry_req, json!({ "sessionId": session_id.as_str() })) + .await; + server.answer_skills_reload().await; + let session = timeout(TIMEOUT, retry).await.unwrap().unwrap().unwrap(); + let mut events = session.subscribe(); + + server + .send_event( + session_id.as_str(), + "evt-after-retry", + "assistant.message", + false, + ) + .await; + assert_eq!( + timeout(TIMEOUT, events.recv()) + .await + .unwrap() + .unwrap() + .id + .as_str(), + "evt-after-retry" + ); + drop(session); +} + // --------------------------------------------------------------------------- // 8. Startup failures preserve error kinds and clean up // --------------------------------------------------------------------------- @@ -884,6 +969,7 @@ async fn resume_session_wrapper_keeps_rpc_sequence() { let resume_req = server.read_request().await; assert_eq!(resume_req["method"], "session.resume"); assert_eq!(resume_req["params"]["sessionId"], session_id.as_str()); + server.send_startup_burst(session_id.as_str()).await; server .respond(&resume_req, json!({ "sessionId": session_id.as_str() })) .await; @@ -891,7 +977,47 @@ async fn resume_session_wrapper_keeps_rpc_sequence() { let session = timeout(TIMEOUT, start).await.unwrap().unwrap().unwrap(); assert_eq!(session.id(), &session_id); - server.expect_quiet().await; + let mut first = session.subscribe(); + let mut second = session.subscribe(); + + expect_startup_burst(&mut first).await; + assert!( + timeout(QUIET, second.recv()).await.is_err(), + "only the first post-resume subscriber may claim the bootstrap" + ); + + // Polling past the retained prefix atomically activates the existing + // bounded broadcast stream. + assert!(timeout(QUIET, first.recv()).await.is_err()); + server + .send_event(session_id.as_str(), "evt-live", "assistant.message", false) + .await; + assert_eq!( + timeout(TIMEOUT, first.recv()) + .await + .unwrap() + .unwrap() + .id + .as_str(), + "evt-live" + ); + assert_eq!( + timeout(TIMEOUT, second.recv()) + .await + .unwrap() + .unwrap() + .id + .as_str(), + "evt-live" + ); + assert!( + timeout(QUIET, first.recv()).await.is_err(), + "first subscriber received a duplicate event" + ); + assert!( + timeout(QUIET, second.recv()).await.is_err(), + "second subscriber received a duplicate event" + ); drop(session); } From 4ac26236597e66ba64d412dbf879cfdc9314871c Mon Sep 17 00:00:00 2001 From: Hunter Sadler Date: Thu, 10 Sep 2026 21:15:47 -0600 Subject: [PATCH 2/4] Harden resume bootstrap catch-up and live observer handoff Keep later subscribers live during bootstrap catch-up, install the owner's broadcast receiver at the atomic empty-queue boundary, and release unclaimed queues on shutdown. Exercise response-gated startup bursts and concurrent handoff on current SDK main. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- rust/src/session.rs | 14 ++- rust/src/subscription.rs | 173 ++++++++++++++++++++++------ rust/tests/prepared_session_test.rs | 24 +++- 3 files changed, 172 insertions(+), 39 deletions(-) diff --git a/rust/src/session.rs b/rust/src/session.rs index 9ffb2bca02..a321a08e87 100644 --- a/rust/src/session.rs +++ b/rust/src/session.rs @@ -426,12 +426,15 @@ impl Session { /// subscription also receives every routed event retained while resume /// startup had no active [`PreparedSession`] subscriber. That bootstrap /// prefix is lossless and ordered before live events. It is a one-shot - /// handoff: later subscriptions begin with live delivery, and dropping + /// handoff: later subscriptions observe newly dispatched events even while + /// the first subscriber is catching up, and dropping /// the first subscription before draining it discards its remaining /// bootstrap events. /// /// Bootstrap retention is unbounded until the first subscriber catches - /// up, so resume consumers that need events should subscribe promptly. + /// up, so resume consumers should subscribe and drain promptly. Stopping + /// the event loop releases an unclaimed backlog; dropping a claimed + /// subscription releases its unread backlog. /// After the bootstrap handoff, each subscriber maintains its own finite /// queue. If a consumer cannot keep /// up, the oldest live events are dropped and `recv` returns @@ -1654,8 +1657,8 @@ impl Client { // An active prepared subscription already owns startup delivery. The // implicit queue is only needed by the compatibility resume wrapper, // where Session::subscribe cannot be called until startup returns. - let resume_bootstrap = - (event_tx.receiver_count() == 0).then(crate::subscription::ResumeBootstrap::new); + let resume_bootstrap = (event_tx.receiver_count() == 0) + .then(|| crate::subscription::ResumeBootstrap::new(&event_tx)); let registration = self.register_session(&session_id); let registration_token = registration.token; let channels = registration.channels; @@ -2182,6 +2185,9 @@ fn spawn_event_loop( else => break, } } + if let Some(bootstrap) = &resume_bootstrap { + bootstrap.release_unclaimed(); + } // Channels closed or shutdown signaled — fail any pending // send_and_wait so the caller observes a clean error. if let Some(waiter) = idle_waiter.lock().take() { diff --git a/rust/src/subscription.rs b/rust/src/subscription.rs index 156f873130..0084b74b89 100644 --- a/rust/src/subscription.rs +++ b/rust/src/subscription.rs @@ -39,7 +39,7 @@ use std::sync::Arc; use std::task::{Context, Poll}; use parking_lot::Mutex; -use tokio::sync::broadcast::{Receiver, Sender}; +use tokio::sync::broadcast::{Receiver, Sender, WeakSender}; use tokio_stream::wrappers::BroadcastStream; use tokio_stream::wrappers::errors::BroadcastStreamRecvError; use tokio_stream::{Stream, StreamExt as _}; @@ -154,12 +154,14 @@ enum ResumeBootstrapState { /// post-resume session subscription catches up. pub(crate) struct ResumeBootstrap { state: Mutex, + live: WeakSender, } impl ResumeBootstrap { - pub(crate) fn new() -> Arc { + pub(crate) fn new(event_tx: &Sender) -> Arc { Arc::new(Self { state: Mutex::new(ResumeBootstrapState::Unclaimed(VecDeque::new())), + live: event_tx.downgrade(), }) } @@ -167,12 +169,12 @@ impl ResumeBootstrap { let mut state = self.state.lock(); match &mut *state { ResumeBootstrapState::Unclaimed(events) | ResumeBootstrapState::Claimed(events) => { - events.push_back(event) - } - ResumeBootstrapState::Disabled => { - let _ = event_tx.send(event); + events.push_back(event.clone()); } + ResumeBootstrapState::Disabled => {} } + // Other observers remain live even while the bootstrap owner catches up. + let _ = event_tx.send(event); } pub(crate) fn subscribe( @@ -180,19 +182,22 @@ impl ResumeBootstrap { event_tx: &Sender, ) -> EventSubscription { let mut state = self.state.lock(); - let receiver = event_tx.subscribe(); - let bootstrap = match &mut *state { + match &mut *state { ResumeBootstrapState::Unclaimed(events) => { let events = std::mem::take(events); *state = ResumeBootstrapState::Claimed(events); - Some(self.clone()) + EventSubscription { + inner: None, + bootstrap: Some(self.clone()), + } } - ResumeBootstrapState::Claimed(_) | ResumeBootstrapState::Disabled => None, - }; - EventSubscription::with_bootstrap(receiver, bootstrap) + ResumeBootstrapState::Claimed(_) | ResumeBootstrapState::Disabled => { + EventSubscription::new(event_tx.subscribe()) + } + } } - fn pop(&self) -> Option { + fn pop(&self, live: &mut Option>) -> Option { let mut state = self.state.lock(); let ResumeBootstrapState::Claimed(events) = &mut *state else { return None; @@ -200,10 +205,23 @@ impl ResumeBootstrap { if let Some(event) = events.pop_front() { return Some(event); } + // Installing the live receiver under the publication lock makes the + // empty-queue boundary gap-free without replaying broadcast duplicates. + *live = self + .live + .upgrade() + .map(|sender| BroadcastStream::new(sender.subscribe())); *state = ResumeBootstrapState::Disabled; None } + pub(crate) fn release_unclaimed(&self) { + let mut state = self.state.lock(); + if matches!(*state, ResumeBootstrapState::Unclaimed(_)) { + *state = ResumeBootstrapState::Disabled; + } + } + fn abandon(&self) { let mut state = self.state.lock(); if matches!(*state, ResumeBootstrapState::Claimed(_)) { @@ -220,19 +238,15 @@ impl ResumeBootstrap { /// Drop the value to unsubscribe; there is no separate cancel handle. #[must_use = "subscriptions are inert until polled"] pub struct EventSubscription { - inner: BroadcastStream, + inner: Option>, bootstrap: Option>, } impl EventSubscription { pub(crate) fn new(rx: Receiver) -> Self { - Self::with_bootstrap(rx, None) - } - - fn with_bootstrap(rx: Receiver, bootstrap: Option>) -> Self { Self { - inner: BroadcastStream::new(rx), - bootstrap, + inner: Some(BroadcastStream::new(rx)), + bootstrap: None, } } @@ -240,7 +254,7 @@ impl EventSubscription { let event = self .bootstrap .as_ref() - .and_then(|bootstrap| bootstrap.pop()); + .and_then(|bootstrap| bootstrap.pop(&mut self.inner)); if event.is_none() { self.bootstrap = None; } @@ -263,12 +277,9 @@ impl EventSubscription { /// `tokio::sync::broadcast::Receiver` via `BroadcastStream`, which is /// cancel-safe by design. pub async fn recv(&mut self) -> Result { - if let Some(event) = self.next_bootstrap_event() { - return Ok(event); - } - match self.inner.next().await { + match self.next().await { Some(Ok(event)) => Ok(event), - Some(Err(BroadcastStreamRecvError::Lagged(n))) => Err(Lagged(n).into()), + Some(Err(lagged)) => Err(lagged.into()), None => Err(RecvErrorKind::Closed.into()), } } @@ -281,7 +292,10 @@ impl Stream for EventSubscription { if let Some(event) = self.next_bootstrap_event() { return Poll::Ready(Some(Ok(event))); } - match Pin::new(&mut self.inner).poll_next(cx) { + let Some(inner) = self.inner.as_mut() else { + return Poll::Ready(None); + }; + match Pin::new(inner).poll_next(cx) { Poll::Ready(Some(Ok(event))) => Poll::Ready(Some(Ok(event))), Poll::Ready(Some(Err(BroadcastStreamRecvError::Lagged(n)))) => { Poll::Ready(Some(Err(Lagged(n)))) @@ -444,18 +458,19 @@ mod tests { #[tokio::test] async fn resume_bootstrap_is_lossless_and_ordered_before_live_events() { let (tx, _) = broadcast::channel(1); - let bootstrap = ResumeBootstrap::new(); + let bootstrap = ResumeBootstrap::new(&tx); for index in 0..600 { bootstrap.publish(&tx, make_event(&format!("bootstrap-{index}"))); } let mut sub = bootstrap.subscribe(&tx); - bootstrap.publish(&tx, make_event("bootstrap-after-subscribe")); + for index in 600..1200 { + bootstrap.publish(&tx, make_event(&format!("bootstrap-{index}"))); + } - for index in 0..600 { + for index in 0..1200 { assert_eq!(sub.recv().await.unwrap().id, format!("bootstrap-{index}")); } - assert_eq!(sub.recv().await.unwrap().id, "bootstrap-after-subscribe"); assert!(sub.next_bootstrap_event().is_none()); bootstrap.publish(&tx, make_event("live")); @@ -465,13 +480,16 @@ mod tests { #[tokio::test] async fn only_first_subscriber_claims_resume_bootstrap() { let (tx, _) = broadcast::channel(8); - let bootstrap = ResumeBootstrap::new(); + let bootstrap = ResumeBootstrap::new(&tx); bootstrap.publish(&tx, make_event("bootstrap")); let mut first = bootstrap.subscribe(&tx); let mut second = bootstrap.subscribe(&tx); + bootstrap.publish(&tx, make_event("during-catchup")); + assert_eq!(second.recv().await.unwrap().id, "during-catchup"); assert_eq!(first.recv().await.unwrap().id, "bootstrap"); + assert_eq!(first.recv().await.unwrap().id, "during-catchup"); assert!(first.next_bootstrap_event().is_none()); bootstrap.publish(&tx, make_event("live")); @@ -479,10 +497,99 @@ mod tests { assert_eq!(second.recv().await.unwrap().id, "live"); } + #[tokio::test] + async fn bootstrap_stream_drains_before_closing_without_retaining_the_sender() { + let (tx, _) = broadcast::channel(1); + let bootstrap = ResumeBootstrap::new(&tx); + bootstrap.publish(&tx, make_event("first")); + let mut events = bootstrap.subscribe(&tx); + bootstrap.publish(&tx, make_event("second")); + drop(tx); + + assert_eq!(events.next().await.unwrap().unwrap().id, "first"); + assert_eq!(events.next().await.unwrap().unwrap().id, "second"); + assert!(events.next().await.is_none()); + assert!(matches!( + events.recv().await.unwrap_err().kind(), + RecvErrorKind::Closed + )); + } + + #[tokio::test] + async fn bootstrap_handoff_is_cancel_safe_and_preserves_live_lag() { + use futures_util::FutureExt; + + let (tx, _) = broadcast::channel(1); + let bootstrap = ResumeBootstrap::new(&tx); + let mut events = bootstrap.subscribe(&tx); + // Poll through the empty bootstrap and cancel the pending live receive. + assert!(events.recv().now_or_never().is_none()); + bootstrap.publish(&tx, make_event("overwritten")); + bootstrap.publish(&tx, make_event("live")); + assert!(matches!( + events.recv().await.unwrap_err().kind(), + RecvErrorKind::Lagged(_) + )); + assert_eq!(events.recv().await.unwrap().id, "live"); + } + + #[tokio::test] + async fn publication_racing_catchup_has_no_gap_or_duplicate() { + for _ in 0..32 { + let (tx, _) = broadcast::channel(1024); + let bootstrap = ResumeBootstrap::new(&tx); + bootstrap.publish(&tx, make_event("prefix")); + let mut events = bootstrap.subscribe(&tx); + assert_eq!(events.recv().await.unwrap().id, "prefix"); + + let barrier = Arc::new(std::sync::Barrier::new(2)); + let producer = std::thread::spawn({ + let barrier = barrier.clone(); + let bootstrap = bootstrap.clone(); + move || { + barrier.wait(); + for index in 0..600 { + bootstrap.publish(&tx, make_event(&format!("event-{index}"))); + } + } + }); + barrier.wait(); + for index in 0..600 { + let event = tokio::time::timeout(std::time::Duration::from_secs(5), events.recv()) + .await + .unwrap() + .unwrap(); + assert_eq!(event.id, format!("event-{index}")); + } + producer.join().unwrap(); + assert!(events.next().await.is_none()); + } + } + + #[tokio::test] + async fn shutdown_releases_only_unclaimed_bootstrap() { + let (tx, _) = broadcast::channel(1); + let bootstrap = ResumeBootstrap::new(&tx); + bootstrap.publish(&tx, make_event("unclaimed")); + bootstrap.release_unclaimed(); + assert!(matches!( + *bootstrap.state.lock(), + ResumeBootstrapState::Disabled + )); + + let claimed = ResumeBootstrap::new(&tx); + claimed.publish(&tx, make_event("claimed")); + let mut events = claimed.subscribe(&tx); + claimed.release_unclaimed(); + drop(tx); + assert_eq!(events.recv().await.unwrap().id, "claimed"); + assert!(events.next().await.is_none()); + } + #[tokio::test] async fn dropping_bootstrap_owner_activates_live_delivery() { let (tx, _) = broadcast::channel(8); - let bootstrap = ResumeBootstrap::new(); + let bootstrap = ResumeBootstrap::new(&tx); bootstrap.publish(&tx, make_event("discarded-bootstrap")); let first = bootstrap.subscribe(&tx); diff --git a/rust/tests/prepared_session_test.rs b/rust/tests/prepared_session_test.rs index 18b9c0ba1f..7a4c5d4077 100644 --- a/rust/tests/prepared_session_test.rs +++ b/rust/tests/prepared_session_test.rs @@ -969,18 +969,38 @@ async fn resume_session_wrapper_keeps_rpc_sequence() { let resume_req = server.read_request().await; assert_eq!(resume_req["method"], "session.resume"); assert_eq!(resume_req["params"]["sessionId"], session_id.as_str()); - server.send_startup_burst(session_id.as_str()).await; + server + .send_event( + session_id.as_str(), + "pre-response", + "session.model_change", + false, + ) + .await; server .respond(&resume_req, json!({ "sessionId": session_id.as_str() })) .await; + // Runtime response-gated notifications can precede the remaining setup RPCs. + server.send_startup_burst(session_id.as_str()).await; server.answer_skills_reload().await; let session = timeout(TIMEOUT, start).await.unwrap().unwrap().unwrap(); assert_eq!(session.id(), &session_id); + server + .send_event(session_id.as_str(), "post-setup", "session.idle", true) + .await; let mut first = session.subscribe(); - let mut second = session.subscribe(); + assert_eq!( + timeout(TIMEOUT, first.recv()).await.unwrap().unwrap().id, + "pre-response" + ); expect_startup_burst(&mut first).await; + assert_eq!( + timeout(TIMEOUT, first.recv()).await.unwrap().unwrap().id, + "post-setup" + ); + let mut second = session.subscribe(); assert!( timeout(QUIET, second.recv()).await.is_err(), "only the first post-resume subscriber may claim the bootstrap" From 7350e99aad25d7cbf124951b0ac54a2cf5d15f80 Mon Sep 17 00:00:00 2001 From: Hunter Sadler Date: Thu, 10 Sep 2026 22:35:00 -0600 Subject: [PATCH 3/4] Cover resume bootstrap ownership and populated cleanup Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- rust/src/session.rs | 7 +- rust/src/subscription.rs | 39 +++++++++ rust/tests/prepared_session_test.rs | 126 ++++++++++++++++++++++++++-- 3 files changed, 163 insertions(+), 9 deletions(-) diff --git a/rust/src/session.rs b/rust/src/session.rs index a321a08e87..6f1c068b79 100644 --- a/rust/src/session.rs +++ b/rust/src/session.rs @@ -432,9 +432,14 @@ impl Session { /// bootstrap events. /// /// Bootstrap retention is unbounded until the first subscriber catches - /// up, so resume consumers should subscribe and drain promptly. Stopping + /// up: callers that never subscribe or cannot catch up can retain an + /// arbitrarily large backlog. Resume consumers should subscribe and drain + /// promptly. Ownership is assigned by the first `subscribe()` call, not + /// by the first poll. Stopping /// the event loop releases an unclaimed backlog; dropping a claimed /// subscription releases its unread backlog. + /// This retention covers events routed to this session, not events lost + /// to overflow in the client-global notification router. /// After the bootstrap handoff, each subscriber maintains its own finite /// queue. If a consumer cannot keep /// up, the oldest live events are dropped and `recv` returns diff --git a/rust/src/subscription.rs b/rust/src/subscription.rs index 0084b74b89..0b3a4df6f5 100644 --- a/rust/src/subscription.rs +++ b/rust/src/subscription.rs @@ -497,6 +497,45 @@ mod tests { assert_eq!(second.recv().await.unwrap().id, "live"); } + #[tokio::test] + async fn concurrent_subscribers_claim_bootstrap_exactly_once() { + use futures_util::FutureExt; + + let (tx, _) = broadcast::channel(8); + let bootstrap = ResumeBootstrap::new(&tx); + bootstrap.publish(&tx, make_event("bootstrap")); + let barrier = std::sync::Barrier::new(2); + let mut subscriptions = std::thread::scope(|scope| { + let subscribe = || { + barrier.wait(); + bootstrap.subscribe(&tx) + }; + let first = scope.spawn(subscribe); + let second = scope.spawn(subscribe); + [first.join().unwrap(), second.join().unwrap()] + }); + + let mut owners = 0; + for sub in &mut subscriptions { + if let Some(event) = sub.recv().now_or_never() { + assert_eq!(event.unwrap().id, "bootstrap"); + owners += 1; + } + } + assert_eq!(owners, 1); + + bootstrap.publish(&tx, make_event("during-catchup")); + for sub in &mut subscriptions { + assert_eq!(sub.recv().await.unwrap().id, "during-catchup"); + assert!(sub.recv().now_or_never().is_none()); + } + bootstrap.publish(&tx, make_event("live")); + for sub in &mut subscriptions { + assert_eq!(sub.recv().await.unwrap().id, "live"); + assert!(sub.recv().now_or_never().is_none()); + } + } + #[tokio::test] async fn bootstrap_stream_drains_before_closing_without_retaining_the_sender() { let (tx, _) = broadcast::channel(1); diff --git a/rust/tests/prepared_session_test.rs b/rust/tests/prepared_session_test.rs index 7a4c5d4077..f4767dbd04 100644 --- a/rust/tests/prepared_session_test.rs +++ b/rust/tests/prepared_session_test.rs @@ -12,12 +12,14 @@ use std::sync::Arc; use std::time::Duration; use async_trait::async_trait; -use github_copilot_sdk::handler::{McpAuthHandler, McpAuthRequest, McpAuthResult}; +use github_copilot_sdk::handler::{ + McpAuthHandler, McpAuthRequest, McpAuthResult, PermissionHandler, PermissionResult, +}; use github_copilot_sdk::session::PreparedSession; use github_copilot_sdk::subscription::{EventSubscription, RecvErrorKind}; use github_copilot_sdk::types::{ - CloudSessionOptions, CloudSessionRepository, RequestId, ResumeSessionConfig, SessionConfig, - SessionId, + CloudSessionOptions, CloudSessionRepository, PermissionRequestData, RequestId, + ResumeSessionConfig, SessionConfig, SessionId, }; use github_copilot_sdk::{Client, ErrorKind, SessionErrorKind}; use serde_json::{Value, json}; @@ -142,6 +144,46 @@ impl FakeServer { assert_eq!(request["method"], "session.skills.reload"); self.respond(&request, json!({})).await; } + + /// The permission handler runs after publication on the same session loop, + /// proving the preceding burst is queued without claiming a subscription. + async fn await_bootstrap_publication( + &mut self, + session_id: &str, + published: &tokio::sync::Notify, + ) { + self.send_startup_burst(session_id).await; + let notification = json!({ + "jsonrpc": "2.0", + "method": "session.event", + "params": { + "sessionId": session_id, + "event": { + "id": "publication-fence", + "timestamp": "2025-01-01T00:00:00Z", + "type": "permission.requested", + "data": { "requestId": "publication-fence", "kind": "read" }, + }, + }, + }); + write_framed(&mut self.write, &serde_json::to_vec(¬ification).unwrap()).await; + timeout(TIMEOUT, published.notified()).await.unwrap(); + } +} + +struct PublicationFence(Arc); + +#[async_trait] +impl PermissionHandler for PublicationFence { + async fn handle( + &self, + _session_id: SessionId, + _request_id: RequestId, + _request: PermissionRequestData, + ) -> PermissionResult { + self.0.notify_one(); + PermissionResult::no_result() + } } /// Minimal MCP-auth handler: its presence is what makes the SDK register @@ -594,25 +636,54 @@ async fn cancelled_prepared_resume_cleans_up_and_allows_retry() { } #[tokio::test] -async fn cancelled_resume_wrapper_cleans_implicit_bootstrap_and_allows_retry() { +async fn populated_resume_bootstrap_cleans_up_after_setup_failure_or_cancellation() { + for cancel in [false, true] { + check_populated_resume_cleanup(cancel).await; + } +} + +async fn check_populated_resume_cleanup(cancel: bool) { let (client, mut server) = make_client(); - let session_id = SessionId::new("resume-wrapper-cancel"); + let session_id = SessionId::new("resume-wrapper-cleanup"); + let published = Arc::new(tokio::sync::Notify::new()); let start = tokio::spawn({ let client = client.clone(); let session_id = session_id.clone(); + let published = published.clone(); async move { client - .resume_session(ResumeSessionConfig::new(session_id)) + .resume_session( + ResumeSessionConfig::new(session_id) + .with_permission_handler(Arc::new(PublicationFence(published))) + .with_mcp_auth_handler(Arc::new(CancelMcpAuthHandler)), + ) .await } }); let resume_req = server.read_request().await; assert_eq!(resume_req["method"], "session.resume"); - start.abort(); - let _ = start.await; + server + .respond(&resume_req, json!({ "sessionId": session_id.as_str() })) + .await; + let interest_req = server.read_request().await; + assert_eq!(interest_req["method"], "session.eventLog.registerInterest"); + server + .await_bootstrap_publication(session_id.as_str(), &published) + .await; + if cancel { + start.abort(); + assert!(start.await.err().unwrap().is_cancelled()); + } else { + server + .respond_error(&interest_req, -32004, "interest registration failed") + .await; + let error = expect_error(timeout(TIMEOUT, start).await.unwrap().unwrap()); + assert!(matches!(error.kind(), ErrorKind::Rpc { code: -32004 })); + } await_no_registrations(&client).await; + server.expect_quiet().await; let retry = tokio::spawn({ let client = client.clone(); @@ -652,6 +723,45 @@ async fn cancelled_resume_wrapper_cleans_implicit_bootstrap_and_allows_retry() { drop(session); } +#[tokio::test] +async fn stopping_session_releases_populated_unclaimed_bootstrap() { + let (client, mut server) = make_client(); + let session_id = SessionId::new("resume-wrapper-stop"); + let published = Arc::new(tokio::sync::Notify::new()); + let start = tokio::spawn({ + let client = client.clone(); + let session_id = session_id.clone(); + let published = published.clone(); + async move { + client + .resume_session( + ResumeSessionConfig::new(session_id) + .with_permission_handler(Arc::new(PublicationFence(published))), + ) + .await + } + }); + let resume_req = server.read_request().await; + server + .respond(&resume_req, json!({ "sessionId": session_id.as_str() })) + .await; + server + .await_bootstrap_publication(session_id.as_str(), &published) + .await; + server.answer_skills_reload().await; + let session = timeout(TIMEOUT, start).await.unwrap().unwrap().unwrap(); + + timeout(TIMEOUT, session.stop_event_loop()).await.unwrap(); + let mut events = session.subscribe(); + assert!( + timeout(QUIET, events.recv()).await.is_err(), + "stopping must discard the unclaimed bootstrap, not replay it" + ); + drop(session); + expect_closed(&mut events).await; + await_no_registrations(&client).await; +} + // --------------------------------------------------------------------------- // 8. Startup failures preserve error kinds and clean up // --------------------------------------------------------------------------- From 02546ce40a605dc0b1f937039e70a558ec0bcdd9 Mon Sep 17 00:00:00 2001 From: Hunter Sadler Date: Fri, 11 Sep 2026 10:05:59 -0600 Subject: [PATCH 4/4] Clarify immediate resume subscription ownership Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- rust/src/subscription.rs | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/rust/src/subscription.rs b/rust/src/subscription.rs index 0b3a4df6f5..50b02c085e 100644 --- a/rust/src/subscription.rs +++ b/rust/src/subscription.rs @@ -236,7 +236,9 @@ impl ResumeBootstrap { /// Created by [`Session::subscribe`](crate::session::Session::subscribe). /// Implements [`Stream`] yielding `Result`. /// Drop the value to unsubscribe; there is no separate cancel handle. -#[must_use = "subscriptions are inert until polled"] +/// A resume bootstrap is claimed when this subscription is created, not when +/// it is first polled. Dropping its owner discards any unread bootstrap events. +#[must_use = "dropping the subscription unsubscribes and discards any owned resume bootstrap backlog"] pub struct EventSubscription { inner: Option>, bootstrap: Option>,