diff --git a/rust/src/session.rs b/rust/src/session.rs index 82e97b6606..6f1c068b79 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,27 @@ 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 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: 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 /// [`RecvErrorKind::Lagged`](crate::subscription::RecvErrorKind::Lagged) /// reporting the count of skipped events. Slow consumers do not block /// the session's event loop. @@ -438,7 +460,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). @@ -1155,11 +1180,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 } @@ -1431,6 +1461,7 @@ impl Client { capabilities.clone(), open_canvases.clone(), event_tx.clone(), + None, shutdown.clone(), external_tools_shutdown.clone(), ); @@ -1469,6 +1500,7 @@ impl Client { capabilities, open_canvases, event_tx, + resume_bootstrap: None, github_token_registration: ParkingLotMutex::new(github_token_registration), registration_token, }; @@ -1627,6 +1659,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(&event_tx)); let registration = self.register_session(&session_id); let registration_token = registration.token; let channels = registration.channels; @@ -1648,6 +1685,7 @@ impl Client { capabilities.clone(), open_canvases.clone(), event_tx.clone(), + resume_bootstrap.clone(), shutdown.clone(), external_tools_shutdown.clone(), ); @@ -1760,6 +1798,7 @@ impl Client { capabilities, open_canvases, event_tx, + resume_bootstrap, github_token_registration: ParkingLotMutex::new(github_token_registration), registration_token, }; @@ -1818,16 +1857,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 @@ -2056,6 +2099,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<()> { @@ -2097,7 +2141,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() => { @@ -2146,6 +2190,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() { @@ -2263,6 +2310,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, @@ -2364,10 +2412,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..50b02c085e 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, WeakSender}; use tokio_stream::wrappers::BroadcastStream; use tokio_stream::wrappers::errors::BroadcastStreamRecvError; use tokio_stream::{Stream, StreamExt as _}; @@ -135,6 +144,178 @@ 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, + live: WeakSender, +} + +impl ResumeBootstrap { + pub(crate) fn new(event_tx: &Sender) -> Arc { + Arc::new(Self { + state: Mutex::new(ResumeBootstrapState::Unclaimed(VecDeque::new())), + live: event_tx.downgrade(), + }) + } + + 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.clone()); + } + ResumeBootstrapState::Disabled => {} + } + // Other observers remain live even while the bootstrap owner catches up. + let _ = event_tx.send(event); + } + + pub(crate) fn subscribe( + self: &Arc, + event_tx: &Sender, + ) -> EventSubscription { + let mut state = self.state.lock(); + match &mut *state { + ResumeBootstrapState::Unclaimed(events) => { + let events = std::mem::take(events); + *state = ResumeBootstrapState::Claimed(events); + EventSubscription { + inner: None, + bootstrap: Some(self.clone()), + } + } + ResumeBootstrapState::Claimed(_) | ResumeBootstrapState::Disabled => { + EventSubscription::new(event_tx.subscribe()) + } + } + } + + fn pop(&self, live: &mut Option>) -> 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); + } + // 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(_)) { + *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. +/// 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>, +} + +impl EventSubscription { + pub(crate) fn new(rx: Receiver) -> Self { + Self { + inner: Some(BroadcastStream::new(rx)), + bootstrap: None, + } + } + + fn next_bootstrap_event(&mut self) -> Option { + let event = self + .bootstrap + .as_ref() + .and_then(|bootstrap| bootstrap.pop(&mut self.inner)); + 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 { + match self.next().await { + Some(Ok(event)) => Ok(event), + Some(Err(lagged)) => Err(lagged.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))); + } + 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)))) + } + 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 +381,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 +456,198 @@ 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(&tx); + for index in 0..600 { + bootstrap.publish(&tx, make_event(&format!("bootstrap-{index}"))); + } + + let mut sub = bootstrap.subscribe(&tx); + for index in 600..1200 { + bootstrap.publish(&tx, make_event(&format!("bootstrap-{index}"))); + } + + for index in 0..1200 { + assert_eq!(sub.recv().await.unwrap().id, format!("bootstrap-{index}")); + } + + 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(&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")); + + assert_eq!(first.recv().await.unwrap().id, "live"); + 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); + 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(&tx); + 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..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 @@ -350,6 +392,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 +635,133 @@ async fn cancelled_prepared_resume_cleans_up_and_allows_retry() { drop(session); } +#[tokio::test] +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-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) + .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"); + 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(); + 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); +} + +#[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 // --------------------------------------------------------------------------- @@ -884,14 +1079,75 @@ 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_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.expect_quiet().await; + server + .send_event(session_id.as_str(), "post-setup", "session.idle", true) + .await; + let mut first = 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" + ); + + // 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); }