From 5a6833b7caf19d15e1f5015fc3df1e02ddd1bfef Mon Sep 17 00:00:00 2001 From: jmoseley Date: Wed, 12 Aug 2026 07:26:54 -0700 Subject: [PATCH 01/12] [Rust] Add PreparedSession for loss-free startup event subscription MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `Session::subscribe()` can only be called once the session handle exists, so every event the runtime broadcast during `session.create` / `session.resume` had no receiver installed and was dropped. Ephemeral events like `session.idle` are never written to the session log, so `get_messages` cannot recover them afterwards either. `Client::prepare_session` / `prepare_resume_session` return a `PreparedSession` that owns the session's broadcast channel up front: subscribe first, then `start()`. `prepare_*` is synchronous and inert — it validates the event buffer capacity, allocates a local channel and cancellation token, and performs no router registration, task spawn, or wire activity until `start()` is first polled. `start(self)` consumes the handle and the type is deliberately not `Clone`, so a prepared session can never produce two event loops. `create_session` / `resume_session` become wrappers over `prepare_*(config)?.start().await`, preserving their RPC sequences and error kinds. Their bodies moved into private start paths that take the sender and token by injection rather than allocating their own. Both configs gain a runtime-only `event_buffer_capacity` (default 512, `Some(0)` rejected as `InvalidConfig`, never clamped). The buffer is finite, so slow subscribers observe `Lagged` instead of applying backpressure. Cancellation cleanup is now symmetric. `PendingSessionRegistration` grew a deferred variant that resolves the session ID from the inline-response stash, so the create path — including the cloud server-assigned-ID path, which previously had no RAII guard at all — unregisters and cancels when the startup future is dropped or fails. Registration and stashing now happen under one lock hold to close the window where a concurrent drop would miss a just-registered session. The mcp-auth-interest error path on both create and resume now cancels and awaits the event loop instead of returning through `?`. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- CHANGELOG.md | 20 + docs/features/streaming-events.md | 47 ++ rust/Cargo.toml | 2 + rust/README.md | 38 +- rust/src/lib.rs | 9 + rust/src/session.rs | 431 +++++++++++++-- rust/src/types.rs | 41 ++ rust/tests/prepared_session_test.rs | 803 ++++++++++++++++++++++++++++ 8 files changed, 1351 insertions(+), 40 deletions(-) create mode 100644 rust/tests/prepared_session_test.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index f7e753beda..7851338720 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -44,6 +44,26 @@ const session = await joinSession({ const token = process.env.GITHUB_TOKEN; ``` +### Feature: early session-event subscription (Rust) + +The Rust SDK can now observe a session's events from its very first event. `Client::prepare_session` and `Client::prepare_resume_session` return an inert `PreparedSession` that owns the session's event channel, so a subscription can be installed *before* any protocol activity begins: + +```rust +let prepared = client.prepare_session( + SessionConfig::default().with_event_buffer_capacity(2048), +)?; +let mut events = prepared.subscribe(); +let session = prepared.start().await?; +``` + +Previously, `Session::subscribe` could only be called on the returned session, so events the runtime emitted while `session.create` / `session.resume` was still in flight were broadcast with no receiver installed and dropped. Ephemeral events such as `session.idle` are not persisted, so they could not be recovered with `getMessages` either. + +`prepare_*` is synchronous and inert: it validates the buffer capacity and allocates a local channel, and performs no router registration, task spawn, or wire activity until `start()` is first polled. `start(self)` consumes the handle and `PreparedSession` is not `Clone`, so a prepared session can never produce two event loops. Dropping an unstarted handle leaves no state behind; dropping a polled `start()` future cancels the startup and unregisters the session, so a retry with the same session ID succeeds. + +Both `SessionConfig` and `ResumeSessionConfig` gained a runtime-only `event_buffer_capacity` option (default 512, `Some(0)` rejected as an invalid config). The buffer is finite, so slow subscribers observe `Lagged` rather than applying backpressure; consumers that need a lossless view of a large startup burst should size the buffer accordingly or drain concurrently with `start()`. + +`create_session` and `resume_session` are unchanged wrappers over `prepare_*(...)?.start()` with identical RPC sequences and error kinds. + ### Feature: host-injected managed settings permissions Session create and resume accept a new optional `managedSettings` option that injects an enterprise permissions policy at session startup, alongside the existing `enableManagedSettings` self-fetch flag. The current contract is permissions-only: `disableBypassPermissionsMode` (the literal `"disable"`), plus `deny`, `ask`, and `allow` rule lists. The layer composes restrictively with any server- or device-level managed settings (deny/ask are unioned, every present allow list must admit a tool, and `disableBypassPermissionsMode` is deny-wins). diff --git a/docs/features/streaming-events.md b/docs/features/streaming-events.md index 10f111d9f9..0f81c0e8d9 100644 --- a/docs/features/streaming-events.md +++ b/docs/features/streaming-events.md @@ -218,6 +218,53 @@ session.on(AssistantMessageDeltaEvent.class, event -> > [!TIP] > **(TypeScript)** The TypeScript SDK uses a discriminated union—when you match on `event.type`, the `data` payload is automatically narrowed to the correct shape. +## Subscribing before a session starts + +A session can emit events before its create or resume call returns. The agent may already be working—especially on resume with `continuePendingWork`—and ephemeral events such as `session.idle` are never written to the session log, so `getMessages` cannot recover them afterwards. A subscription installed after the session handle exists misses that startup window. + +> [!TIP] +> **(Rust)** `Client::prepare_session` and `Client::prepare_resume_session` return a `PreparedSession` that owns the session's event channel before any protocol activity happens. Subscribe first, then call `start()`. + +
+Rust + +```rust +use github_copilot_sdk::{Client, SessionConfig}; + +async fn create_without_missing_startup_events( + client: &Client, +) -> Result<(), github_copilot_sdk::Error> { + let prepared = client.prepare_session( + SessionConfig::default().with_event_buffer_capacity(2048), + )?; + + // Installed before any wire activity: nothing is dropped for lack of a receiver. + let mut events = prepared.subscribe(); + tokio::spawn(async move { + while let Ok(event) = events.recv().await { + println!("{}", event.event_type); + } + }); + + let session = prepared.start().await?; + let _ = session; + Ok(()) +} +``` + +
+ +`prepare_*` is synchronous and inert: it validates the buffer capacity, allocates a local channel, and does nothing else. No session is registered and nothing reaches the CLI until `start()` is first polled. Dropping a prepared session that was never started leaves no state behind and closes its subscriptions; dropping the `start()` future cancels the in-flight startup and unregisters the session, so a retry with the same session ID succeeds. + +Startup buffering is worth planning for: + +* The event buffer is finite—512 events unless `event_buffer_capacity` overrides it. A capacity of `0` is rejected with an invalid-config error rather than clamped. +* Slow subscribers observe a `Lagged` error reporting how many events were skipped. They never apply backpressure to the session's event loop. +* Consumers that need a lossless view of a large startup burst should either configure a capacity that covers it or drain the subscription concurrently with `start()`. + +> [!NOTE] +> For cloud sessions where the server assigns the session ID, the SDK cannot route notifications until the create response arrives and the ID is known. Events emitted before that point are not routable to any session. The guarantee is narrower: routed events are never dropped for lack of an installed receiver. Pin `session_id` on the config to get routing—and full pre-response coverage—from the first byte. + ## Render only the parent agent response Sub-agent events share the parent session stream and include envelope-level `agentId`. Root/main agent events and session-level events omit `agentId`, so main-chat renderers can ignore assistant events where `agentId` is set and route those events to traces or progress UI instead. diff --git a/rust/Cargo.toml b/rust/Cargo.toml index c66480d511..71d754f26e 100644 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -118,6 +118,8 @@ required-features = ["test-support"] test = false bench = false +name = "prepared_session_test" +required-features = ["test-support"] [build-dependencies] dirs = "5" flate2 = "1" diff --git a/rust/README.md b/rust/README.md index a561e6da09..ada81e0b14 100644 --- a/rust/README.md +++ b/rust/README.md @@ -696,6 +696,36 @@ while let Ok(event) = events.recv().await { When streaming is off (the default), only the final `assistant.message` and `assistant.reasoning` events fire. Delta events arrive in order; concatenating their `delta` text payloads reproduces the final message. +#### Subscribing before the session starts + +`session.subscribe()` can only be called once the session exists, so any event the runtime emits while `session.create` / `session.resume` is still in flight is broadcast with no receiver installed and is not delivered. Ephemeral events such as `session.idle` are not written to the session log either, so `get_messages` can't recover them afterwards. + +`Client::prepare_session` / `Client::prepare_resume_session` close that window. They return a `PreparedSession` that owns the session's broadcast channel up front: + +```rust,ignore +let prepared = client.prepare_session( + SessionConfig::default().with_event_buffer_capacity(2048), +)?; + +// Installed before any wire activity happens. +let mut events = prepared.subscribe(); +tokio::spawn(async move { + while let Ok(event) = events.recv().await { + println!("{}", event.event_type); + } +}); + +let session = prepared.start().await?; +``` + +`prepare_*` is synchronous and inert — it validates the buffer capacity, allocates a local channel and cancellation token, and touches neither the router nor the transport until `start()` is first polled. `start(self)` consumes the handle and `PreparedSession` is deliberately not `Clone`, so a prepared session can never spawn two event loops. Dropping an unstarted handle leaves no state and closes its subscriptions; dropping the `start()` future cancels the startup, unregisters the session, and lets a same-ID retry succeed. + +The buffer is finite — `session::DEFAULT_EVENT_BUFFER_CAPACITY` (512) unless `event_buffer_capacity` overrides it, and `Some(0)` is rejected as `ErrorKind::InvalidConfig` rather than clamped. Subscribers that fall behind observe `RecvErrorKind::Lagged` with the skipped count instead of applying backpressure, so a consumer that needs a lossless view of a large startup burst must configure enough capacity or drain concurrently with `start()`. + +For cloud sessions where the server assigns the session ID, notifications can't be routed until the create response arrives; the guarantee is that *routed* events are never dropped for lack of a receiver. Pin `session_id` for full pre-response coverage. + +`create_session` / `resume_session` are unchanged wrappers over `prepare_*(...)?.start()`, with identical RPC sequences and error kinds. + ### Infinite Sessions Enable the SDK's session-store integration so conversations persist across CLI restarts and grow beyond the model's context window via automatic compaction: @@ -899,6 +929,12 @@ none of them are scheduled for removal. arg vectors for "prepend before subcommand" vs "append after the built-in flags", giving precise control over CLI invocation order without string-splicing. +- **`Client::prepare_session` / `prepare_resume_session`** — return an inert + `PreparedSession` whose `subscribe()` installs an event receiver before any + protocol activity, so startup events (including ephemeral `session.idle`) + aren't dropped. Other SDKs register callbacks on a config object instead, + which sidesteps the problem in a way Rust's broadcast-based `subscribe()` + cannot. ## Layout @@ -906,7 +942,7 @@ none of them are scheduled for removal. | ----------------- | -------------------------------------------------------------------------------------------------------------------------- | | `lib.rs` | `Client`, `ClientOptions`, `CliProgram`, `Transport`, `Error` | | `extension_launch_provider.rs` | Connection-global `ExtensionLaunchProvider` trait and launch profile DTOs | -| `session.rs` | `Session` struct, event loop, `send`/`send_and_wait`, `Client::create_session`/`resume_session` | +| `session.rs` | `Session` struct, `PreparedSession`, event loop, `send`/`send_and_wait`, `Client::create_session`/`resume_session`/`prepare_session`/`prepare_resume_session` | | `subscription.rs` | `EventSubscription` / `LifecycleSubscription` (`Stream`-able observer handles for `subscribe()` / `subscribe_lifecycle()`) | | `handler.rs` | `PermissionHandler`, `ElicitationHandler`, `UserInputHandler`, `ExitPlanModeHandler`, `AutoModeSwitchHandler` traits; `ApproveAllHandler`, `DenyAllHandler` | | `hooks.rs` | `SessionHooks` trait, `HookEvent`/`HookOutput` enums, typed hook inputs/outputs | diff --git a/rust/src/lib.rs b/rust/src/lib.rs index fd7f12cf14..0b8af4e139 100644 --- a/rust/src/lib.rs +++ b/rust/src/lib.rs @@ -2591,6 +2591,15 @@ impl Client { ); } + #[cfg(feature = "test-support")] + #[doc(hidden)] + /// Snapshot the session IDs currently registered on this client's + /// notification router. This is test-harness plumbing, not part of the + /// supported SDK API. + pub fn registered_session_ids_for_test(&self) -> Vec { + self.inner.router.session_ids() + } + #[cfg(feature = "test-support")] #[doc(hidden)] /// Disconnect and delete every session owned by this test client's isolated diff --git a/rust/src/session.rs b/rust/src/session.rs index 509485e39e..fd514a196f 100644 --- a/rust/src/session.rs +++ b/rust/src/session.rs @@ -49,6 +49,31 @@ use crate::{ /// `overrides_built_in_tool` set to `true`. const TOOL_SEARCH_TOOL_NAME: &str = "tool_search_tool"; +/// Default capacity of the per-session event broadcast buffer backing +/// [`Session::subscribe`] and [`PreparedSession::subscribe`]. +/// +/// Override per session with +/// [`SessionConfig::event_buffer_capacity`](crate::types::SessionConfig::event_buffer_capacity) +/// or +/// [`ResumeSessionConfig::event_buffer_capacity`](crate::types::ResumeSessionConfig::event_buffer_capacity). +pub const DEFAULT_EVENT_BUFFER_CAPACITY: usize = 512; + +/// Validate a caller-supplied event buffer capacity and resolve the default. +/// +/// Zero is rejected rather than clamped: a zero-capacity broadcast channel +/// cannot exist, and silently substituting a different capacity would hide a +/// caller bug. +fn resolve_event_buffer_capacity(capacity: Option) -> Result { + match capacity { + Some(0) => Err(Error::with_message( + ErrorKind::InvalidConfig, + "event_buffer_capacity must be greater than zero", + )), + Some(capacity) => Ok(capacity), + None => Ok(DEFAULT_EVENT_BUFFER_CAPACITY), + } +} + /// Bundle of the per-session callbacks the SDK dispatches to. Built from a /// [`SessionConfig`] / [`ResumeSessionConfig`] at /// [`Client::create_session`] / [`Client::resume_session`] time. Each @@ -106,25 +131,71 @@ impl Drop for WaiterGuard { struct PendingSessionRegistration { client: Client, - session_id: SessionId, + session_id: PendingSessionId, shutdown: CancellationToken, disarmed: bool, } +/// Which session ID a [`PendingSessionRegistration`] should unregister on +/// cleanup. +/// +/// `session.create` for cloud sessions without a caller-pinned ID does not +/// know the ID until the response arrives, at which point the inline +/// response callback registers it and stashes it. The guard therefore reads +/// the stash at cleanup time instead of capturing an ID up front. +enum PendingSessionId { + /// The ID was known before the RPC was issued (resume, and create with a + /// client- or caller-supplied ID). + Known(SessionId), + /// Server-assigned ID, populated by the `session.create` inline response + /// callback. `None` in the stash means nothing was ever registered. + Deferred(Arc>>), +} + impl PendingSessionRegistration { fn new(client: Client, session_id: SessionId, shutdown: CancellationToken) -> Self { Self { client, - session_id, + session_id: PendingSessionId::Known(session_id), shutdown, disarmed: false, } } + /// Guard for a registration whose session ID is assigned by the server. + fn deferred( + client: Client, + stash: Arc>>, + shutdown: CancellationToken, + ) -> Self { + Self { + client, + session_id: PendingSessionId::Deferred(stash), + shutdown, + disarmed: false, + } + } + + fn registered_id(&self) -> Option { + match &self.session_id { + PendingSessionId::Known(id) => Some(id.clone()), + PendingSessionId::Deferred(stash) => stash.lock().as_ref().map(|(id, _)| id.clone()), + } + } + + /// Re-target the guard at a now-known session ID. Used by + /// `session.create` once the response has been parsed and the stash has + /// been drained into the event loop. + fn resolve_to(&mut self, session_id: SessionId) { + self.session_id = PendingSessionId::Known(session_id); + } + async fn cleanup(mut self, event_loop: JoinHandle<()>) { self.shutdown.cancel(); let _ = event_loop.await; - self.client.unregister_session(&self.session_id); + if let Some(id) = self.registered_id() { + self.client.unregister_session(&id); + } self.disarmed = true; } @@ -137,7 +208,9 @@ impl Drop for PendingSessionRegistration { fn drop(&mut self) { if !self.disarmed { self.shutdown.cancel(); - self.client.unregister_session(&self.session_id); + if let Some(id) = self.registered_id() { + self.client.unregister_session(&id); + } } } } @@ -813,6 +886,94 @@ impl<'a> SessionUi<'a> { } impl Client { + /// Prepare a new session without touching the transport. + /// + /// Returns a [`PreparedSession`] that owns the session's event broadcast + /// channel, so callers can install an + /// [`EventSubscription`](crate::subscription::EventSubscription) via + /// [`PreparedSession::subscribe`] *before* any protocol activity starts. + /// Call [`PreparedSession::start`] to actually create the session. + /// + /// This is the loss-free entry point for consumers that must observe + /// every event a session emits, including events the CLI emits while + /// `session.create` is still in flight and ephemeral events (such as + /// `session.idle`) that cannot be recovered from + /// [`Session::get_messages`]. [`create_session`](Self::create_session) + /// is a thin wrapper over `prepare_session(...)?.start()` and cannot + /// offer the same guarantee, because the subscription can only be + /// installed after the returned `Session` exists. + /// + /// # Inertness + /// + /// `prepare_session` performs no router registration, spawns no task, + /// and writes nothing to the wire. It only validates + /// [`event_buffer_capacity`](SessionConfig::event_buffer_capacity), + /// allocates a local broadcast channel and cancellation token, and + /// stores the config. Dropping the returned handle without starting it + /// leaves no client-side or server-side state behind and closes every + /// subscription taken from it. + /// + /// # Errors + /// + /// Returns [`ErrorKind::InvalidConfig`] if + /// [`event_buffer_capacity`](SessionConfig::event_buffer_capacity) is + /// `Some(0)`. All other configuration and protocol errors surface from + /// [`PreparedSession::start`], with the same + /// [`ErrorKind`]s [`create_session`](Self::create_session) has always + /// returned. + /// + /// # Example + /// + /// ```no_run + /// # use github_copilot_sdk::{Client, SessionConfig}; + /// # async fn example(client: Client) -> Result<(), github_copilot_sdk::Error> { + /// let prepared = client.prepare_session(SessionConfig::default())?; + /// let mut events = prepared.subscribe(); + /// let drain = tokio::spawn(async move { + /// while let Ok(event) = events.recv().await { + /// println!("{}", event.event_type); + /// } + /// }); + /// let session = prepared.start().await?; + /// # let _ = (session, drain); + /// # Ok(()) + /// # } + /// ``` + pub fn prepare_session(&self, config: SessionConfig) -> Result { + let capacity = resolve_event_buffer_capacity(config.event_buffer_capacity)?; + Ok(PreparedSession::new( + self.clone(), + PreparedKind::Create(Box::new(config)), + capacity, + )) + } + + /// Prepare a session resume without touching the transport. + /// + /// The resume counterpart of [`prepare_session`](Self::prepare_session); + /// see that method for the inertness guarantee, error semantics, and + /// rationale. Particularly relevant on resume with + /// [`continue_pending_work`](ResumeSessionConfig::continue_pending_work), + /// where the runtime can start emitting events (and reach + /// `session.idle`) while `session.resume` is still in flight. + /// + /// # Errors + /// + /// Returns [`ErrorKind::InvalidConfig`] if + /// [`event_buffer_capacity`](ResumeSessionConfig::event_buffer_capacity) + /// is `Some(0)`. + pub fn prepare_resume_session( + &self, + config: ResumeSessionConfig, + ) -> Result { + let capacity = resolve_event_buffer_capacity(config.event_buffer_capacity)?; + Ok(PreparedSession::new( + self.clone(), + PreparedKind::Resume(Box::new(config)), + capacity, + )) + } + /// Create a new session on the CLI. /// /// Sends `session.create`, registers the session on the router, @@ -834,7 +995,48 @@ impl Client { /// Each per-event handler is independently optional. If a handler is /// not installed, the SDK signals the runtime not to emit the matching /// broadcast (and silently skips dispatch if one arrives anyway). - pub async fn create_session(&self, mut config: SessionConfig) -> Result { + /// + /// # Event delivery + /// + /// Equivalent to `prepare_session(config)?.start().await`. Because the + /// first subscription can only be taken from the returned [`Session`], + /// events the runtime emits before this call returns are broadcast with + /// no receiver installed and are therefore not delivered to + /// [`Session::subscribe`]. Use + /// [`prepare_session`](Self::prepare_session) when startup events + /// matter. + pub async fn create_session(&self, config: SessionConfig) -> Result { + self.prepare_session(config)?.start().await + } + + /// Resume an existing session on the CLI. + /// + /// Sends `session.resume` and `session.skills.reload`, registers the + /// session on the router, and spawns the event loop. + /// + /// All callbacks (event handler, hooks, transform) are configured + /// via [`ResumeSessionConfig`] using its `with_*` builder methods. + /// + /// See [`Self::create_session`] for the defaults applied when callback + /// fields are unset. + /// + /// # 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. + pub async fn resume_session(&self, config: ResumeSessionConfig) -> Result { + self.prepare_resume_session(config)?.start().await + } + + async fn start_prepared_create( + &self, + mut config: SessionConfig, + event_tx: tokio::sync::broadcast::Sender, + shutdown: CancellationToken, + ) -> Result { let total_start = Instant::now(); // For cloud sessions, let the CLI/server assign the session id and // register the session lazily once the response arrives. For non-cloud @@ -975,8 +1177,6 @@ impl Client { let capabilities = Arc::new(parking_lot::RwLock::new(SessionCapabilities::default())); let idle_waiter = Arc::new(ParkingLotMutex::new(None)); let open_canvases = Arc::new(parking_lot::RwLock::new(Vec::new())); - let shutdown = CancellationToken::new(); - let (event_tx, _) = tokio::sync::broadcast::channel(512); // For cloud sessions (use_server_generated_id), defer session // registration to the inline callback so the read task registers @@ -1012,45 +1212,49 @@ impl Client { }) .into()); } + // Register and stash under a single stash-lock hold. The + // cancellation guard identifies the session to unregister by + // peeking this stash, so registering outside the lock would + // leave a window where a concurrent guard drop (caller + // cancellation) sees `None` and leaks the registration. + // `register_session` takes the router lock, never the stash + // lock, so there is no lock-order inversion here. + let mut stashed = stash.lock(); let channels = client.register_session(&parsed.session_id); - *stash.lock() = Some((parsed.session_id, channels)); + *stashed = Some((parsed.session_id, channels)); Ok(()) })) }; - let rpc_start = Instant::now(); - let result = match self - .call_with_inline_callback("session.create", Some(params), inline_callback) - .await - { - Ok(result) => result, - Err(error) => { - if let Some((id, _channels)) = inline_stash.lock().take() { - self.unregister_session(&id); - } - return Err(error); + // Armed for the whole startup sequence: any early return, and any + // drop of this future (caller cancellation), cancels the session + // token and unregisters whatever was registered on the router. For + // the cloud path the ID is only known once the inline callback has + // run, so the guard reads the stash at cleanup time. + let mut registration = match local_session_id { + Some(ref sid) => { + PendingSessionRegistration::new(self.clone(), sid.clone(), shutdown.clone()) } + None => PendingSessionRegistration::deferred( + self.clone(), + inline_stash.clone(), + shutdown.clone(), + ), }; + + let rpc_start = Instant::now(); + let result = self + .call_with_inline_callback("session.create", Some(params), inline_callback) + .await?; tracing::debug!( elapsed_ms = rpc_start.elapsed().as_millis(), "Client::create_session session creation request completed successfully" ); - let create_result: CreateSessionResult = match serde_json::from_value(result) { - Ok(result) => result, - Err(error) => { - if let Some((id, _channels)) = inline_stash.lock().take() { - self.unregister_session(&id); - } - return Err(error.into()); - } - }; + let create_result: CreateSessionResult = serde_json::from_value(result)?; if let Some(ref requested) = local_session_id && create_result.session_id != *requested { - if let Some((id, _channels)) = inline_stash.lock().take() { - self.unregister_session(&id); - } return Err(ErrorKind::Session(SessionErrorKind::SessionIdMismatch { requested: requested.clone(), returned: create_result.session_id.clone(), @@ -1062,6 +1266,7 @@ impl Client { .lock() .take() .expect("session registration must have populated stash on success"); + registration.resolve_to(session_id.clone()); let event_loop = spawn_event_loop( session_id.clone(), self.clone(), @@ -1088,8 +1293,11 @@ impl Client { "Client::create_session local setup complete" ); *capabilities.write() = create_result.capabilities.unwrap_or_default(); - if has_mcp_auth_handler { - register_mcp_auth_interest(self, &session_id).await?; + if has_mcp_auth_handler + && let Err(error) = register_mcp_auth_interest(self, &session_id).await + { + registration.cleanup(event_loop).await; + return Err(error); } tracing::debug!( @@ -1097,6 +1305,7 @@ impl Client { session_id = %session_id, "Client::create_session complete" ); + registration.disarm(); let session = Session { id: session_id, cwd: self.cwd().clone(), @@ -1139,7 +1348,12 @@ impl Client { /// /// See [`Self::create_session`] for the defaults applied when callback /// fields are unset. - pub async fn resume_session(&self, mut config: ResumeSessionConfig) -> Result { + async fn start_prepared_resume( + &self, + mut config: ResumeSessionConfig, + event_tx: tokio::sync::broadcast::Sender, + shutdown: CancellationToken, + ) -> Result { let total_start = Instant::now(); let session_id = config.session_id.clone(); if config.hooks_handler.is_some() && config.hooks.is_none() { @@ -1264,8 +1478,6 @@ impl Client { let channels = self.register_session(&session_id); let idle_waiter = Arc::new(ParkingLotMutex::new(None)); let open_canvases = Arc::new(parking_lot::RwLock::new(Vec::new())); - let shutdown = CancellationToken::new(); - let (event_tx, _) = tokio::sync::broadcast::channel(512); let event_loop = spawn_event_loop( session_id.clone(), self.clone(), @@ -1327,10 +1539,12 @@ impl Client { }) .into()); } - if has_mcp_auth_handler { - register_mcp_auth_interest(self, &session_id).await?; + if has_mcp_auth_handler + && let Err(error) = register_mcp_auth_interest(self, &session_id).await + { + registration.cleanup(event_loop).await; + return Err(error); } - // Reload skills after resume (best-effort). let skills_reload_start = Instant::now(); if let Err(e) = self @@ -1405,6 +1619,145 @@ impl Client { } } +/// A session that has been configured but not yet created on the CLI. +/// +/// Returned by [`Client::prepare_session`] and +/// [`Client::prepare_resume_session`]. Its purpose is to make the session's +/// event stream observable *before* any protocol activity starts: +/// [`subscribe`](Self::subscribe) installs a receiver on the same broadcast +/// channel the eventual [`Session`] uses, so events the runtime emits while +/// `session.create` / `session.resume` is still in flight are delivered +/// rather than dropped for lack of a receiver. +/// +/// # Lifecycle +/// +/// A prepared handle is inert. It holds only a broadcast sender, a +/// cancellation token, the client handle, and the config — it performs no +/// router registration, spawns no task, and writes nothing to the wire +/// until [`start`](Self::start) is first polled. +/// +/// * Dropping it without starting leaves no client-side or server-side +/// state, and closes every subscription taken from it. +/// * Dropping the [`start`](Self::start) future mid-flight cancels the +/// session token, unregisters the session from the router if it was +/// registered, and closes early subscriptions. A retry with the same +/// session ID succeeds. Cleanup of already-spawned tasks is signalled, +/// not awaited: `Drop` is synchronous and cannot await, so the event loop +/// terminates promptly but not synchronously. +/// * A startup error from [`start`](Self::start) performs the same cleanup +/// and preserves the [`ErrorKind`] the equivalent +/// [`Client::create_session`] / [`Client::resume_session`] call has always +/// returned. +/// +/// [`start`](Self::start) consumes `self` and the type is deliberately not +/// [`Clone`], so a prepared session can be started at most once and can +/// never produce two event loops. +/// +/// # Buffering +/// +/// The broadcast buffer is finite — +/// [`DEFAULT_EVENT_BUFFER_CAPACITY`] unless +/// [`SessionConfig::event_buffer_capacity`] / +/// [`ResumeSessionConfig::event_buffer_capacity`] overrides it. 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). +/// +/// # Server-assigned session IDs +/// +/// For cloud sessions without a caller-supplied session ID, the CLI assigns +/// the ID and the SDK can only register the session on its notification +/// router once the `session.create` response arrives. Notifications the +/// server emits before that point are not routable to any session and are +/// therefore not observable. The guarantee this type provides is narrower +/// and precise: **routed** events are never dropped for lack of an +/// installed receiver. Pin +/// [`SessionConfig::session_id`](crate::types::SessionConfig::session_id) +/// to get registration before the RPC and full pre-response coverage. +#[must_use = "a PreparedSession does nothing until started"] +pub struct PreparedSession { + client: Client, + kind: PreparedKind, + event_tx: tokio::sync::broadcast::Sender, + shutdown: CancellationToken, +} + +/// Which startup path a [`PreparedSession`] runs when started. Boxed +/// because the two config types are large and differently sized. +enum PreparedKind { + Create(Box), + Resume(Box), +} + +impl PreparedSession { + fn new(client: Client, kind: PreparedKind, event_buffer_capacity: usize) -> Self { + let (event_tx, _) = tokio::sync::broadcast::channel(event_buffer_capacity); + Self { + client, + kind, + event_tx, + shutdown: CancellationToken::new(), + } + } + + /// Subscribe to this session's events before it starts. + /// + /// The returned [`EventSubscription`](crate::subscription::EventSubscription) + /// is backed by the same broadcast channel + /// [`Session::subscribe`] returns after [`start`](Self::start) + /// succeeds, so a subscription taken here observes the full event + /// stream from the session's first routed event onward — including + /// ephemeral events such as `session.idle` that + /// [`Session::get_messages`] cannot recover. + /// + /// May be called any number of times; every subscriber receives every + /// event. Subscriptions taken here close if the prepared session is + /// dropped without starting, or if startup fails. + pub fn subscribe(&self) -> crate::subscription::EventSubscription { + crate::subscription::EventSubscription::new(self.event_tx.subscribe()) + } + + /// Create or resume the session on the CLI. + /// + /// This is where all protocol activity happens: config validation, + /// router registration, the `session.create` / `session.resume` RPC, + /// and the event loop spawn. Nothing observable occurs until this + /// future is first polled. + /// + /// # Errors + /// + /// Returns the same errors as [`Client::create_session`] / + /// [`Client::resume_session`] — including + /// [`ErrorKind::InvalidConfig`] for invalid configs, transport and RPC + /// failures, and + /// [`SessionIdMismatch`](crate::SessionErrorKind::SessionIdMismatch) + /// when the CLI returns a different session ID than the one requested. + /// Every error path unregisters the session and closes subscriptions + /// taken from this handle. + pub async fn start(self) -> Result { + let Self { + client, + kind, + event_tx, + shutdown, + } = self; + match kind { + PreparedKind::Create(config) => { + client + .start_prepared_create(*config, event_tx, shutdown) + .await + } + PreparedKind::Resume(config) => { + client + .start_prepared_resume(*config, event_tx, shutdown) + .await + } + } + } +} + type CommandHandlerMap = HashMap>; async fn apply_mode_post_create_patch( diff --git a/rust/src/types.rs b/rust/src/types.rs index ee3ac3df26..68bde9dfb3 100644 --- a/rust/src/types.rs +++ b/rust/src/types.rs @@ -2266,6 +2266,23 @@ pub struct SessionConfig { /// `session.options.update` after create/resume. Defaults to `false` in /// [`crate::ClientMode::Empty`] when unset. pub manage_schedule_enabled: Option, + /// Capacity of the per-session broadcast buffer backing + /// [`Session::subscribe`](crate::session::Session::subscribe) and + /// [`PreparedSession::subscribe`](crate::session::PreparedSession::subscribe). + /// + /// Runtime-only — never sent on the wire. Defaults to + /// [`DEFAULT_EVENT_BUFFER_CAPACITY`](crate::session::DEFAULT_EVENT_BUFFER_CAPACITY) + /// when unset. Must be non-zero; + /// `Some(0)` is rejected with + /// [`ErrorKind::InvalidConfig`](crate::ErrorKind::InvalidConfig) by + /// [`Client::prepare_session`](crate::Client::prepare_session). + /// + /// The buffer is finite: subscribers that fall behind observe + /// [`Lagged`](crate::subscription::Lagged) rather than applying + /// backpressure to the event loop. Raise this when a consumer needs a + /// lossless view of a large startup burst without draining + /// concurrently. + pub event_buffer_capacity: Option, } impl std::fmt::Debug for SessionConfig { @@ -2401,6 +2418,7 @@ impl std::fmt::Debug for SessionConfig { "system_message_transform", &self.system_message_transform.as_ref().map(|_| ""), ) + .field("event_buffer_capacity", &self.event_buffer_capacity) .finish() } } @@ -2497,6 +2515,7 @@ impl Default for SessionConfig { enable_experimental_mode: None, coauthor_enabled: None, manage_schedule_enabled: None, + event_buffer_capacity: None, } } } @@ -3297,6 +3316,15 @@ impl SessionConfig { /// Set feature-flag values resolved by the host for this session. pub fn with_feature_flags(mut self, feature_flags: HashMap) -> Self { self.feature_flags = Some(feature_flags); + /// Set [`Self::event_buffer_capacity`]. + /// + /// A capacity of `0` is rejected with + /// [`ErrorKind::InvalidConfig`](crate::ErrorKind::InvalidConfig) by + /// [`Client::prepare_session`](crate::Client::prepare_session) and + /// [`Client::create_session`](crate::Client::create_session); the value + /// is never clamped. + pub fn with_event_buffer_capacity(mut self, capacity: usize) -> Self { + self.event_buffer_capacity = Some(capacity); self } @@ -3602,6 +3630,8 @@ pub struct ResumeSessionConfig { pub coauthor_enabled: Option, /// See [`SessionConfig::manage_schedule_enabled`]. pub manage_schedule_enabled: Option, + /// See [`SessionConfig::event_buffer_capacity`]. + pub event_buffer_capacity: Option, } impl std::fmt::Debug for ResumeSessionConfig { @@ -3735,6 +3765,7 @@ impl std::fmt::Debug for ResumeSessionConfig { ) .field("suppress_resume_event", &self.suppress_resume_event) .field("continue_pending_work", &self.continue_pending_work) + .field("event_buffer_capacity", &self.event_buffer_capacity) .finish() } } @@ -3986,6 +4017,7 @@ impl ResumeSessionConfig { enable_experimental_mode: None, coauthor_enabled: None, manage_schedule_enabled: None, + event_buffer_capacity: None, } } @@ -4594,6 +4626,15 @@ impl ResumeSessionConfig { /// Re-supply feature-flag values resolved by the host on resume. pub fn with_feature_flags(mut self, feature_flags: HashMap) -> Self { self.feature_flags = Some(feature_flags); + /// Set [`Self::event_buffer_capacity`]. + /// + /// A capacity of `0` is rejected with + /// [`ErrorKind::InvalidConfig`](crate::ErrorKind::InvalidConfig) by + /// [`Client::prepare_resume_session`](crate::Client::prepare_resume_session) + /// and [`Client::resume_session`](crate::Client::resume_session); the + /// value is never clamped. + pub fn with_event_buffer_capacity(mut self, capacity: usize) -> Self { + self.event_buffer_capacity = Some(capacity); self } diff --git a/rust/tests/prepared_session_test.rs b/rust/tests/prepared_session_test.rs new file mode 100644 index 0000000000..694f161205 --- /dev/null +++ b/rust/tests/prepared_session_test.rs @@ -0,0 +1,803 @@ +//! Early event subscription via `Client::prepare_session` / +//! `Client::prepare_resume_session`. +//! +//! Every test drives the SDK over an in-memory duplex transport and a +//! hand-rolled JSON-RPC peer, so event ordering is deterministic. Timeouts +//! are failure backstops only — no test sleeps to "let things settle". + +#![allow(clippy::unwrap_used)] + +use std::marker::PhantomData; +use std::time::Duration; + +use github_copilot_sdk::session::PreparedSession; +use github_copilot_sdk::subscription::{EventSubscription, RecvErrorKind}; +use github_copilot_sdk::types::{ResumeSessionConfig, SessionConfig, SessionId}; +use github_copilot_sdk::{Client, ErrorKind, SessionErrorKind}; +use serde_json::{Value, json}; +use tokio::io::{AsyncRead, AsyncWrite, AsyncWriteExt, duplex}; +use tokio::time::timeout; + +/// Failure backstop for operations that must complete promptly. +const TIMEOUT: Duration = Duration::from_secs(5); +/// Backstop for asserting that something does *not* happen. +const QUIET: Duration = Duration::from_millis(150); +/// Size of the pre-response event burst. Mirrors the copilot-host startup +/// burst that motivated the API. +const BURST: usize = 600; + +// --------------------------------------------------------------------------- +// Transport harness +// --------------------------------------------------------------------------- + +async fn write_framed(writer: &mut (impl AsyncWrite + Unpin), body: &[u8]) { + let header = format!("Content-Length: {}\r\n\r\n", body.len()); + writer.write_all(header.as_bytes()).await.unwrap(); + writer.write_all(body).await.unwrap(); + writer.flush().await.unwrap(); +} + +async fn read_framed(reader: &mut (impl AsyncRead + Unpin)) -> Value { + let mut header = String::new(); + loop { + let mut byte = [0u8; 1]; + tokio::io::AsyncReadExt::read_exact(reader, &mut byte) + .await + .unwrap(); + header.push(byte[0] as char); + if header.ends_with("\r\n\r\n") { + break; + } + } + let length: usize = header + .trim() + .strip_prefix("Content-Length: ") + .unwrap() + .parse() + .unwrap(); + let mut buf = vec![0u8; length]; + tokio::io::AsyncReadExt::read_exact(reader, &mut buf) + .await + .unwrap(); + serde_json::from_slice(&buf).unwrap() +} + +struct FakeServer { + read: tokio::io::DuplexStream, + write: tokio::io::DuplexStream, +} + +impl FakeServer { + async fn read_request(&mut self) -> Value { + timeout(TIMEOUT, read_framed(&mut self.read)).await.unwrap() + } + + async fn expect_quiet(&mut self) { + assert!( + timeout(QUIET, read_framed(&mut self.read)).await.is_err(), + "expected no wire traffic" + ); + } + + async fn respond(&mut self, request: &Value, result: Value) { + let id = request["id"].as_u64().unwrap(); + let response = json!({ "jsonrpc": "2.0", "id": id, "result": result }); + write_framed(&mut self.write, &serde_json::to_vec(&response).unwrap()).await; + } + + async fn respond_error(&mut self, request: &Value, code: i64, message: &str) { + let id = request["id"].as_u64().unwrap(); + let response = json!({ + "jsonrpc": "2.0", + "id": id, + "error": { "code": code, "message": message }, + }); + write_framed(&mut self.write, &serde_json::to_vec(&response).unwrap()).await; + } + + async fn send_event(&mut self, session_id: &str, id: &str, event_type: &str, ephemeral: bool) { + let notification = json!({ + "jsonrpc": "2.0", + "method": "session.event", + "params": { + "sessionId": session_id, + "event": { + "id": id, + "timestamp": "2025-01-01T00:00:00Z", + "ephemeral": ephemeral, + "type": event_type, + "data": {}, + }, + }, + }); + write_framed(&mut self.write, &serde_json::to_vec(¬ification).unwrap()).await; + } + + /// Emit the startup burst the host cares about: `BURST` ordered events + /// followed by an ephemeral `session.idle` that `getMessages` could + /// never recover. + async fn send_startup_burst(&mut self, session_id: &str) { + for i in 0..BURST { + self.send_event( + session_id, + &format!("evt-{i}"), + "assistant.message_delta", + false, + ) + .await; + } + self.send_event(session_id, "evt-idle", "session.idle", true) + .await; + } + + /// Answer the best-effort `session.skills.reload` that follows a resume. + async fn answer_skills_reload(&mut self) { + let request = self.read_request().await; + assert_eq!(request["method"], "session.skills.reload"); + self.respond(&request, json!({})).await; + } +} + +fn make_client() -> (Client, FakeServer) { + let (client_write, server_read) = duplex(1 << 20); + let (server_write, client_read) = duplex(1 << 20); + let client = Client::from_streams(client_read, client_write, std::env::temp_dir()).unwrap(); + ( + client, + FakeServer { + read: server_read, + write: server_write, + }, + ) +} + +fn create_result(session_id: &str) -> Value { + json!({ "sessionId": session_id, "workspacePath": "/tmp/workspace" }) +} + +/// Collect the startup burst, asserting each event arrives exactly once and +/// in emission order. +async fn expect_startup_burst(events: &mut EventSubscription) { + for i in 0..BURST { + let event = timeout(TIMEOUT, events.recv()) + .await + .unwrap_or_else(|_| panic!("timed out waiting for event {i}")) + .unwrap_or_else(|error| panic!("event {i} not delivered: {error}")); + assert_eq!(event.id.as_str(), format!("evt-{i}"), "out-of-order event"); + } + let idle = timeout(TIMEOUT, events.recv()).await.unwrap().unwrap(); + assert_eq!(idle.id.as_str(), "evt-idle"); + assert_eq!(idle.event_type, "session.idle"); + assert_eq!(idle.ephemeral, Some(true)); +} + +/// Unwrap the error arm of a result whose `Ok` type is not `Debug`. +fn expect_error(result: Result) -> github_copilot_sdk::Error { + match result { + Ok(_) => panic!("expected an error"), + Err(error) => error, + } +} + +/// Poll (bounded) until the client's router has no registered sessions. +async fn await_no_registrations(client: &Client) { + let deadline = tokio::time::Instant::now() + TIMEOUT; + while !client.registered_session_ids_for_test().is_empty() { + assert!( + tokio::time::Instant::now() < deadline, + "session registration was never cleaned up: {:?}", + client.registered_session_ids_for_test() + ); + tokio::task::yield_now().await; + } +} + +/// Assert the subscription is closed (producer gone), tolerating any events +/// buffered before the close. +async fn expect_closed(events: &mut EventSubscription) { + loop { + match timeout(TIMEOUT, events.recv()).await.unwrap() { + Ok(_) => continue, + Err(error) => { + assert!( + matches!(error.kind(), RecvErrorKind::Closed), + "expected Closed, got {:?}", + error.kind() + ); + return; + } + } + } +} + +// --------------------------------------------------------------------------- +// 1 + 4. Loss-free startup events on create +// --------------------------------------------------------------------------- + +/// Subscription installed before `start()` is polled, drained concurrently: +/// the full pre-response burst plus the ephemeral `session.idle` arrives. +#[tokio::test] +async fn prepared_create_delivers_pre_response_burst_to_concurrent_drain() { + let (client, mut server) = make_client(); + let session_id = SessionId::new("prepared-create-concurrent"); + + let prepared = client + .prepare_session( + SessionConfig::default() + .with_session_id(session_id.clone()) + .with_event_buffer_capacity(2048), + ) + .unwrap(); + let mut events = prepared.subscribe(); + let drain = tokio::spawn(async move { + expect_startup_burst(&mut events).await; + }); + + let start = tokio::spawn(prepared.start()); + + let create_req = server.read_request().await; + assert_eq!(create_req["method"], "session.create"); + server.send_startup_burst(session_id.as_str()).await; + server + .respond(&create_req, create_result(session_id.as_str())) + .await; + + let session = timeout(TIMEOUT, start).await.unwrap().unwrap().unwrap(); + timeout(TIMEOUT, drain).await.unwrap().unwrap(); + drop(session); +} + +/// A large configured buffer retains the whole burst even when the consumer +/// does not read anything until `start()` has returned. +#[tokio::test] +async fn prepared_create_retains_burst_for_deferred_consumer() { + let (client, mut server) = make_client(); + let session_id = SessionId::new("prepared-create-deferred"); + + let prepared = client + .prepare_session( + SessionConfig::default() + .with_session_id(session_id.clone()) + .with_event_buffer_capacity(2048), + ) + .unwrap(); + let mut events = prepared.subscribe(); + let start = tokio::spawn(prepared.start()); + + let create_req = server.read_request().await; + server.send_startup_burst(session_id.as_str()).await; + server + .respond(&create_req, create_result(session_id.as_str())) + .await; + + let session = timeout(TIMEOUT, start).await.unwrap().unwrap().unwrap(); + // Only now does the consumer start reading. + expect_startup_burst(&mut events).await; + drop(session); +} + +// --------------------------------------------------------------------------- +// 2. Loss-free startup events on resume +// --------------------------------------------------------------------------- + +#[tokio::test] +async fn prepared_resume_delivers_pre_response_burst() { + let (client, mut server) = make_client(); + let session_id = SessionId::new("prepared-resume"); + + let prepared = client + .prepare_resume_session( + ResumeSessionConfig::new(session_id.clone()) + .with_continue_pending_work(true) + .with_event_buffer_capacity(2048), + ) + .unwrap(); + let mut events = prepared.subscribe(); + let start = tokio::spawn(prepared.start()); + + let resume_req = server.read_request().await; + assert_eq!(resume_req["method"], "session.resume"); + assert_eq!(resume_req["params"]["continuePendingWork"], true); + server.send_startup_burst(session_id.as_str()).await; + server + .respond(&resume_req, json!({ "sessionId": session_id.as_str() })) + .await; + server.answer_skills_reload().await; + + let session = timeout(TIMEOUT, start).await.unwrap().unwrap().unwrap(); + expect_startup_burst(&mut events).await; + drop(session); +} + +// --------------------------------------------------------------------------- +// 3. Lag is observable, never silent +// --------------------------------------------------------------------------- + +/// An undersized buffer with a consumer that does not drain surfaces +/// `Lagged` rather than silently losing events, and the live tail stays +/// consumable afterwards. +#[tokio::test] +async fn undersized_buffer_reports_lag_and_keeps_live_tail() { + let (client, mut server) = make_client(); + let session_id = SessionId::new("prepared-lag"); + + let prepared = client + .prepare_session( + SessionConfig::default() + .with_session_id(session_id.clone()) + .with_event_buffer_capacity(8), + ) + .unwrap(); + let mut events = prepared.subscribe(); + let start = tokio::spawn(prepared.start()); + + let create_req = server.read_request().await; + server.send_startup_burst(session_id.as_str()).await; + server + .respond(&create_req, create_result(session_id.as_str())) + .await; + let session = timeout(TIMEOUT, start).await.unwrap().unwrap().unwrap(); + + // Drain until lag is reported. Every delivered event is still in order, + // and the loss is explicit rather than silent. + let mut lagged = None; + let mut last_index: Option = None; + while lagged.is_none() { + match timeout(TIMEOUT, events.recv()).await.unwrap() { + Ok(event) => { + if let Some(index) = event.id.as_str().strip_prefix("evt-") + && let Ok(index) = index.parse::() + { + if let Some(previous) = last_index { + assert!(index > previous, "delivered events must stay ordered"); + } + last_index = Some(index); + } + } + Err(error) => match error.kind() { + RecvErrorKind::Lagged(lag) => lagged = Some(lag.skipped()), + other => panic!("expected lag, got {other:?}"), + }, + } + } + assert!(lagged.unwrap() > 0, "lag must report the skipped count"); + + // The live tail is still consumable after a lag. + server + .send_event(session_id.as_str(), "evt-live", "assistant.message", false) + .await; + let live = loop { + match timeout(TIMEOUT, events.recv()).await.unwrap() { + Ok(event) if event.id.as_str() == "evt-live" => break event, + Ok(_) => continue, + Err(error) => match error.kind() { + RecvErrorKind::Lagged(_) => continue, + other => panic!("subscription ended before the live tail: {other:?}"), + }, + } + }; + assert_eq!(live.event_type, "assistant.message"); + drop(session); +} + +// --------------------------------------------------------------------------- +// 5 + 6. Inertness before start, and drop of an unstarted handle +// --------------------------------------------------------------------------- + +#[tokio::test] +async fn prepare_is_inert_until_start_is_polled() { + let (client, mut server) = make_client(); + let session_id = SessionId::new("prepared-inert"); + + let tasks_before = tokio::runtime::Handle::current() + .metrics() + .num_alive_tasks(); + let prepared = client + .prepare_session(SessionConfig::default().with_session_id(session_id.clone())) + .unwrap(); + let _events = prepared.subscribe(); + + // No wire traffic, no router registration, no spawned task. + server.expect_quiet().await; + assert!(client.registered_session_ids_for_test().is_empty()); + assert_eq!( + tokio::runtime::Handle::current() + .metrics() + .num_alive_tasks(), + tasks_before, + "prepare must not spawn a task" + ); + + // Constructing the future is still inert; only polling it does work. + let start = prepared.start(); + server.expect_quiet().await; + assert!(client.registered_session_ids_for_test().is_empty()); + + let start = tokio::spawn(start); + let create_req = server.read_request().await; + assert_eq!(create_req["method"], "session.create"); + server + .respond(&create_req, create_result(session_id.as_str())) + .await; + let session = timeout(TIMEOUT, start).await.unwrap().unwrap().unwrap(); + drop(session); +} + +#[tokio::test] +async fn dropping_unstarted_prepared_session_leaves_no_state() { + let (client, mut server) = make_client(); + let session_id = SessionId::new("prepared-dropped"); + + let prepared = client + .prepare_session(SessionConfig::default().with_session_id(session_id.clone())) + .unwrap(); + let mut events = prepared.subscribe(); + drop(prepared); + + assert!(matches!( + timeout(TIMEOUT, events.recv()) + .await + .unwrap() + .unwrap_err() + .kind(), + RecvErrorKind::Closed + )); + assert!(client.registered_session_ids_for_test().is_empty()); + server.expect_quiet().await; +} + +// --------------------------------------------------------------------------- +// 7. Cancelling a polled startup +// --------------------------------------------------------------------------- + +#[tokio::test] +async fn cancelled_prepared_create_cleans_up_and_allows_retry() { + let (client, mut server) = make_client(); + let session_id = SessionId::new("prepared-create-cancel"); + + let prepared = client + .prepare_session(SessionConfig::default().with_session_id(session_id.clone())) + .unwrap(); + let mut events = prepared.subscribe(); + let start = tokio::spawn(prepared.start()); + + // The request is on the wire; cancel before responding. + let create_req = server.read_request().await; + assert_eq!(create_req["method"], "session.create"); + start.abort(); + let _ = start.await; + + await_no_registrations(&client).await; + expect_closed(&mut events).await; + + // A retry with the same session ID succeeds. + let retry = tokio::spawn( + client + .prepare_session(SessionConfig::default().with_session_id(session_id.clone())) + .unwrap() + .start(), + ); + let retry_req = server.read_request().await; + assert_eq!(retry_req["method"], "session.create"); + server + .respond(&retry_req, create_result(session_id.as_str())) + .await; + let session = timeout(TIMEOUT, retry).await.unwrap().unwrap().unwrap(); + assert_eq!(session.id(), &session_id); + drop(session); +} + +#[tokio::test] +async fn cancelled_prepared_resume_cleans_up_and_allows_retry() { + let (client, mut server) = make_client(); + let session_id = SessionId::new("prepared-resume-cancel"); + + let prepared = client + .prepare_resume_session(ResumeSessionConfig::new(session_id.clone())) + .unwrap(); + let mut events = prepared.subscribe(); + let start = tokio::spawn(prepared.start()); + + let resume_req = server.read_request().await; + assert_eq!(resume_req["method"], "session.resume"); + start.abort(); + let _ = start.await; + + await_no_registrations(&client).await; + expect_closed(&mut events).await; + + let retry = tokio::spawn( + client + .prepare_resume_session(ResumeSessionConfig::new(session_id.clone())) + .unwrap() + .start(), + ); + 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(); + assert_eq!(session.id(), &session_id); + drop(session); +} + +// --------------------------------------------------------------------------- +// 8. Startup failures preserve error kinds and clean up +// --------------------------------------------------------------------------- + +#[tokio::test] +async fn create_rpc_error_preserves_kind_and_cleans_up() { + let (client, mut server) = make_client(); + let session_id = SessionId::new("prepared-create-rpc-error"); + + let prepared = client + .prepare_session(SessionConfig::default().with_session_id(session_id.clone())) + .unwrap(); + let mut events = prepared.subscribe(); + let start = tokio::spawn(prepared.start()); + + let create_req = server.read_request().await; + server + .respond_error(&create_req, -32000, "session create failed") + .await; + + let error = expect_error(timeout(TIMEOUT, start).await.unwrap().unwrap()); + assert!( + matches!(error.kind(), ErrorKind::Rpc { code: -32000 }), + "unexpected error kind: {:?}", + error.kind() + ); + await_no_registrations(&client).await; + expect_closed(&mut events).await; +} + +#[tokio::test] +async fn create_session_id_mismatch_preserves_kind_and_cleans_up() { + let (client, mut server) = make_client(); + let session_id = SessionId::new("prepared-mismatch"); + + let prepared = client + .prepare_session(SessionConfig::default().with_session_id(session_id.clone())) + .unwrap(); + let mut events = prepared.subscribe(); + let start = tokio::spawn(prepared.start()); + + let create_req = server.read_request().await; + server + .respond(&create_req, create_result("some-other-id")) + .await; + + let error = expect_error(timeout(TIMEOUT, start).await.unwrap().unwrap()); + let ErrorKind::Session(SessionErrorKind::SessionIdMismatch { + requested, + returned, + }) = error.kind() + else { + panic!("unexpected error kind: {:?}", error.kind()); + }; + assert_eq!(requested, &session_id); + assert_eq!(returned.as_str(), "some-other-id"); + + await_no_registrations(&client).await; + expect_closed(&mut events).await; +} + +#[tokio::test] +async fn resume_session_id_mismatch_preserves_kind_and_cleans_up() { + let (client, mut server) = make_client(); + let session_id = SessionId::new("prepared-resume-mismatch"); + + let prepared = client + .prepare_resume_session(ResumeSessionConfig::new(session_id.clone())) + .unwrap(); + let mut events = prepared.subscribe(); + let start = tokio::spawn(prepared.start()); + + let resume_req = server.read_request().await; + server + .respond(&resume_req, json!({ "sessionId": "another-session" })) + .await; + + let error = expect_error(timeout(TIMEOUT, start).await.unwrap().unwrap()); + assert!( + matches!( + error.kind(), + ErrorKind::Session(SessionErrorKind::SessionIdMismatch { .. }) + ), + "unexpected error kind: {:?}", + error.kind() + ); + await_no_registrations(&client).await; + expect_closed(&mut events).await; +} + +#[tokio::test] +async fn zero_event_buffer_capacity_is_invalid_config() { + let (client, _server) = make_client(); + + let error = expect_error( + client.prepare_session(SessionConfig::default().with_event_buffer_capacity(0)), + ); + assert!(matches!(error.kind(), ErrorKind::InvalidConfig)); + + let error = expect_error(client.prepare_resume_session( + ResumeSessionConfig::new(SessionId::new("zero")).with_event_buffer_capacity(0), + )); + assert!(matches!(error.kind(), ErrorKind::InvalidConfig)); + + // The compatibility wrappers surface the same error. + let error = expect_error( + client + .create_session(SessionConfig::default().with_event_buffer_capacity(0)) + .await, + ); + assert!(matches!(error.kind(), ErrorKind::InvalidConfig)); +} + +// --------------------------------------------------------------------------- +// 9. Early and late subscribers share one event loop +// --------------------------------------------------------------------------- + +#[tokio::test] +async fn early_and_late_subscribers_share_one_event_loop() { + let (client, mut server) = make_client(); + let session_id = SessionId::new("prepared-two-subscribers"); + + let prepared = client + .prepare_session( + SessionConfig::default() + .with_session_id(session_id.clone()) + .with_event_buffer_capacity(2048), + ) + .unwrap(); + let mut early = prepared.subscribe(); + let start = tokio::spawn(prepared.start()); + + let create_req = server.read_request().await; + server + .send_event(session_id.as_str(), "evt-early", "assistant.message", false) + .await; + server + .respond(&create_req, create_result(session_id.as_str())) + .await; + let session = timeout(TIMEOUT, start).await.unwrap().unwrap().unwrap(); + + let mut late = session.subscribe(); + server + .send_event(session_id.as_str(), "evt-late", "assistant.message", false) + .await; + + // The early subscriber sees both events, once each. + assert_eq!( + timeout(TIMEOUT, early.recv()) + .await + .unwrap() + .unwrap() + .id + .as_str(), + "evt-early" + ); + assert_eq!( + timeout(TIMEOUT, early.recv()) + .await + .unwrap() + .unwrap() + .id + .as_str(), + "evt-late" + ); + // The late subscriber only sees what was emitted after it subscribed — + // exactly once, which would be twice if a second event loop existed. + assert_eq!( + timeout(TIMEOUT, late.recv()) + .await + .unwrap() + .unwrap() + .id + .as_str(), + "evt-late" + ); + assert!( + timeout(QUIET, late.recv()).await.is_err(), + "duplicate delivery implies more than one event loop" + ); + assert!( + timeout(QUIET, early.recv()).await.is_err(), + "duplicate delivery implies more than one event loop" + ); + drop(session); +} + +// --------------------------------------------------------------------------- +// 10. Compatibility wrappers +// --------------------------------------------------------------------------- + +#[tokio::test] +async fn create_session_wrapper_keeps_rpc_sequence() { + let (client, mut server) = make_client(); + + let start = tokio::spawn({ + let client = client.clone(); + async move { client.create_session(SessionConfig::default()).await } + }); + + let create_req = server.read_request().await; + assert_eq!(create_req["method"], "session.create"); + let session_id = create_req["params"]["sessionId"] + .as_str() + .unwrap() + .to_string(); + server + .respond(&create_req, create_result(&session_id)) + .await; + + let session = timeout(TIMEOUT, start).await.unwrap().unwrap().unwrap(); + assert_eq!(session.id().as_str(), session_id); + server.expect_quiet().await; + drop(session); +} + +#[tokio::test] +async fn resume_session_wrapper_keeps_rpc_sequence() { + let (client, mut server) = make_client(); + let session_id = SessionId::new("wrapper-resume"); + + 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"); + assert_eq!(resume_req["params"]["sessionId"], session_id.as_str()); + server + .respond(&resume_req, json!({ "sessionId": 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; + drop(session); +} + +// --------------------------------------------------------------------------- +// 11. Type-level guarantees +// --------------------------------------------------------------------------- + +/// Detects `Clone` without requiring it: the inherent method wins whenever +/// `T: Clone`, otherwise the blanket trait method is selected. +struct CloneProbe(PhantomData); + +impl CloneProbe { + fn is_clone(&self) -> bool { + true + } +} + +trait MaybeClone { + fn is_clone(&self) -> bool { + false + } +} + +impl MaybeClone for CloneProbe {} + +#[test] +fn prepared_session_is_send_static_and_not_clone() { + fn assert_send_static() {} + assert_send_static::(); + + // Sanity-check the probe against a type that is `Clone` ... + assert!(CloneProbe::(PhantomData).is_clone()); + // ... then assert `PreparedSession` deliberately is not, so a prepared + // session can never be started twice. + assert!(!CloneProbe::(PhantomData).is_clone()); +} From b049770f29f5696ab806f551c6b8b1fdcf396422 Mon Sep 17 00:00:00 2001 From: jmoseley Date: Wed, 12 Aug 2026 08:32:52 -0700 Subject: [PATCH 02/12] [Rust] Close two session-registration cancellation races MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Follow-up to the `PreparedSession` change. Two cancellation races remained in session registration, both reachable from a caller simply dropping a `start()` future. **Deferred cloud-create registration.** For a cloud session with no caller-pinned ID the CLI assigns the ID, so the SDK can only register on the notification router from the inline `session.create` response callback. The read loop removes the pending-response entry *before* invoking that callback, so a startup future dropped in that window found an empty stash, cleaned up nothing, and the callback then registered a session with no owner — a permanent router leak. Registration state now lives in a shared `DeferredRegistration` slot (`Pending` / `Registered` / `Cancelled` / `Claimed`) that the callback, the startup path, and the cancellation guard all arbitrate through. The callback registers *under the slot lock*, so registering and publishing ownership are atomic with respect to cancellation: a concurrent guard either wins and marks the slot `Cancelled`, in which case the callback registers nothing, or it loses and finds a `Registered` slot to tear down. Never both, and never neither. The pinned-ID path uses the same slot, pre-populated, so create has one cleanup mechanism instead of two. **Stale cleanup versus a same-ID retry.** Unregistering by session ID alone removed whichever registration happened to hold the ID. Because cleanup of an abandoned startup is signalled rather than awaited, a caller that aborted a startup and immediately retried with the same pinned ID could have the retry's registration evicted by the dead attempt, silently stranding the live session with no event routing. The same applied to a `Session` dropped after being superseded. Registrations now carry a `RegistrationToken` identity and removal is a compare-and-remove: an owner removes only the exact registration it registered. Applied to create, resume, `Session::disconnect`, and `Session::drop`. `Client::stop` and `cleanup_sessions_for_test` keep removing unconditionally — they tear down every session and the runtime regardless of owner. Tests gate both windows deterministically rather than by timing. The slot state machine is driven directly at the exact interleaving the read loop creates, in both orders, and the router's compare-and-remove is covered on its own. End to end: a cancelled cloud create leaves no registration, no subscription, and no task behind whether cancellation lands before or after the callback registered, and a same-ID retry still succeeds; and create, resume, and `Session` drop each survive a stale owner's cleanup running after a retry has taken over the ID. Each test was confirmed to fail against a mutated implementation. No public API change. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 0f9d5ac7-9999-4f37-82b3-a5533bfbc1f6 --- CHANGELOG.md | 2 +- docs/features/streaming-events.md | 2 +- rust/README.md | 2 +- rust/src/lib.rs | 42 +++- rust/src/router.rs | 100 ++++++++- rust/src/session.rs | 10 +- rust/tests/prepared_session_test.rs | 317 +++++++++++++++++++++++++++- 7 files changed, 455 insertions(+), 20 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 7851338720..0f5dd44870 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -58,7 +58,7 @@ let session = prepared.start().await?; Previously, `Session::subscribe` could only be called on the returned session, so events the runtime emitted while `session.create` / `session.resume` was still in flight were broadcast with no receiver installed and dropped. Ephemeral events such as `session.idle` are not persisted, so they could not be recovered with `getMessages` either. -`prepare_*` is synchronous and inert: it validates the buffer capacity and allocates a local channel, and performs no router registration, task spawn, or wire activity until `start()` is first polled. `start(self)` consumes the handle and `PreparedSession` is not `Clone`, so a prepared session can never produce two event loops. Dropping an unstarted handle leaves no state behind; dropping a polled `start()` future cancels the startup and unregisters the session, so a retry with the same session ID succeeds. +`prepare_*` is synchronous and inert: it validates the buffer capacity and allocates a local channel, and performs no router registration, task spawn, or wire activity until `start()` is first polled. `start(self)` consumes the handle and `PreparedSession` is not `Clone`, so a prepared session can never produce two event loops. Dropping an unstarted handle leaves no state behind; dropping a polled `start()` future cancels the startup and unregisters the session, so a retry with the same session ID succeeds. Session registrations now carry an ownership identity, so cleanup removes only the exact registration it owns and an abandoned startup can never evict a same-ID retry (or a session that replaced it). Both `SessionConfig` and `ResumeSessionConfig` gained a runtime-only `event_buffer_capacity` option (default 512, `Some(0)` rejected as an invalid config). The buffer is finite, so slow subscribers observe `Lagged` rather than applying backpressure; consumers that need a lossless view of a large startup burst should size the buffer accordingly or drain concurrently with `start()`. diff --git a/docs/features/streaming-events.md b/docs/features/streaming-events.md index 0f81c0e8d9..e0eae797fb 100644 --- a/docs/features/streaming-events.md +++ b/docs/features/streaming-events.md @@ -254,7 +254,7 @@ async fn create_without_missing_startup_events( -`prepare_*` is synchronous and inert: it validates the buffer capacity, allocates a local channel, and does nothing else. No session is registered and nothing reaches the CLI until `start()` is first polled. Dropping a prepared session that was never started leaves no state behind and closes its subscriptions; dropping the `start()` future cancels the in-flight startup and unregisters the session, so a retry with the same session ID succeeds. +`prepare_*` is synchronous and inert: it validates the buffer capacity, allocates a local channel, and does nothing else. No session is registered and nothing reaches the CLI until `start()` is first polled. Dropping a prepared session that was never started leaves no state behind and closes its subscriptions; dropping the `start()` future cancels the in-flight startup and unregisters the session, so a retry with the same session ID succeeds. Cleanup is scoped to the exact registration the abandoned startup owned, so it cannot evict a retry that has already taken over the same session ID. Startup buffering is worth planning for: diff --git a/rust/README.md b/rust/README.md index ada81e0b14..e563d9a27a 100644 --- a/rust/README.md +++ b/rust/README.md @@ -718,7 +718,7 @@ tokio::spawn(async move { let session = prepared.start().await?; ``` -`prepare_*` is synchronous and inert — it validates the buffer capacity, allocates a local channel and cancellation token, and touches neither the router nor the transport until `start()` is first polled. `start(self)` consumes the handle and `PreparedSession` is deliberately not `Clone`, so a prepared session can never spawn two event loops. Dropping an unstarted handle leaves no state and closes its subscriptions; dropping the `start()` future cancels the startup, unregisters the session, and lets a same-ID retry succeed. +`prepare_*` is synchronous and inert — it validates the buffer capacity, allocates a local channel and cancellation token, and touches neither the router nor the transport until `start()` is first polled. `start(self)` consumes the handle and `PreparedSession` is deliberately not `Clone`, so a prepared session can never spawn two event loops. Dropping an unstarted handle leaves no state and closes its subscriptions; dropping the `start()` future cancels the startup, unregisters the session, and lets a same-ID retry succeed. Cleanup removes only the exact registration that startup owned, so a retry started while an abandoned attempt is still unwinding is never evicted by it. The buffer is finite — `session::DEFAULT_EVENT_BUFFER_CAPACITY` (512) unless `event_buffer_capacity` overrides it, and `Some(0)` is rejected as `ErrorKind::InvalidConfig` rather than clamped. Subscribers that fall behind observe `RecvErrorKind::Lagged` with the skipped count instead of applying backpressure, so a consumer that needs a lossless view of a large startup burst must configure enough capacity or drain concurrently with `start()`. diff --git a/rust/src/lib.rs b/rust/src/lib.rs index 0b8af4e139..096cd9ae95 100644 --- a/rust/src/lib.rs +++ b/rust/src/lib.rs @@ -2324,15 +2324,18 @@ impl Client { /// Register a session to receive filtered events and requests. /// - /// Returns per-session channels for notifications and requests, routed - /// by `sessionId`. Starts the internal router on first call. - /// - /// When done, call [`unregister_session`](Self::unregister_session) to - /// clean up (typically on session destroy). + /// Returns the per-session channels plus a + /// [`RegistrationToken`](crate::router::RegistrationToken) identifying + /// *this* registration. Registering an ID that is already registered + /// replaces the previous registration. + /// + /// When done, call + /// [`unregister_session_owned`](Self::unregister_session_owned) with + /// that token to clean up (typically on session destroy). pub(crate) fn register_session( &self, session_id: &SessionId, - ) -> crate::router::SessionChannels { + ) -> crate::router::SessionRegistration { self.inner.router.ensure_started( &self.inner.notification_tx, &self.inner.request_rx, @@ -2344,9 +2347,28 @@ impl Client { self.inner.router.register(session_id) } - /// Unregister a session, dropping its per-session channels. - pub(crate) fn unregister_session(&self, session_id: &SessionId) { - self.inner.router.unregister(session_id); + /// Unregister a session only if `token` still identifies the live + /// registration. + /// + /// Session IDs can be reused: a caller may retry a cancelled startup + /// with the same pinned ID while the previous owner is still being torn + /// down. Compare-and-remove keeps a stale owner from unregistering the + /// live session that replaced it. + pub(crate) fn unregister_session_owned( + &self, + session_id: &SessionId, + token: crate::router::RegistrationToken, + ) { + self.inner.router.unregister_owned(session_id, token); + } + + /// Snapshot the session IDs currently registered on the router. + /// + /// Crate-internal so in-crate unit tests can assert registration + /// lifecycle without depending on the `test-support` feature, which + /// only gates the equivalent *public* test helper. + pub(crate) fn registered_session_ids(&self) -> Vec { + self.inner.router.session_ids() } pub(crate) fn register_github_token_provider( @@ -2597,7 +2619,7 @@ impl Client { /// notification router. This is test-harness plumbing, not part of the /// supported SDK API. pub fn registered_session_ids_for_test(&self) -> Vec { - self.inner.router.session_ids() + self.registered_session_ids() } #[cfg(feature = "test-support")] diff --git a/rust/src/router.rs b/rust/src/router.rs index 4815d0e1c5..f4aee25d53 100644 --- a/rust/src/router.rs +++ b/rust/src/router.rs @@ -1,5 +1,6 @@ use std::collections::HashMap; use std::sync::Arc; +use std::sync::atomic::{AtomicU64, Ordering}; use parking_lot::Mutex; use tokio::sync::{broadcast, mpsc}; @@ -8,6 +9,24 @@ use tracing::warn; use crate::jsonrpc::{JsonRpcNotification, JsonRpcRequest}; use crate::types::{SessionEventNotification, SessionId}; +/// Identity of one specific registration of a session ID. +/// +/// Session IDs are not unique over time: a caller can retry a cancelled +/// startup with the same pinned ID, and the retry replaces the previous +/// registration. Removal is therefore compare-and-remove against this +/// token, so a stale owner (an aborted startup future or a superseded +/// [`Session`](crate::session::Session)) can never unregister the live +/// registration that replaced it. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) struct RegistrationToken(u64); + +/// Per-session channels plus the identity of the registration that owns +/// them. Returned by [`SessionRouter::register`]. +pub(crate) struct SessionRegistration { + pub(crate) channels: SessionChannels, + pub(crate) token: RegistrationToken, +} + /// Per-session channels created by the router during session registration. pub(crate) struct SessionChannels { /// Filtered `session.event` notifications for this session. @@ -19,6 +38,7 @@ pub(crate) struct SessionChannels { struct SessionSenders { notifications: mpsc::UnboundedSender, requests: mpsc::UnboundedSender, + token: RegistrationToken, } /// Routes notifications and requests by sessionId to per-session channels. @@ -26,6 +46,7 @@ struct SessionSenders { /// Internal to the SDK — consumers interact via `Client::register_session()`. pub(crate) struct SessionRouter { sessions: Arc>>, + next_token: AtomicU64, started: Mutex, } @@ -33,32 +54,69 @@ impl SessionRouter { pub(crate) fn new() -> Self { Self { sessions: Arc::new(Mutex::new(HashMap::new())), + next_token: AtomicU64::new(0), started: Mutex::new(false), } } /// Register a session to receive filtered events and requests. - pub(crate) fn register(&self, session_id: &SessionId) -> SessionChannels { + /// + /// Replaces any existing registration for the same ID and returns a + /// fresh [`RegistrationToken`] identifying this registration. + pub(crate) fn register(&self, session_id: &SessionId) -> SessionRegistration { let (notif_tx, notif_rx) = mpsc::unbounded_channel(); let (req_tx, req_rx) = mpsc::unbounded_channel(); + let token = RegistrationToken(self.next_token.fetch_add(1, Ordering::Relaxed)); self.sessions.lock().insert( session_id.clone(), SessionSenders { notifications: notif_tx, requests: req_tx, + token, }, ); - SessionChannels { - notifications: notif_rx, - requests: req_rx, + SessionRegistration { + channels: SessionChannels { + notifications: notif_rx, + requests: req_rx, + }, + token, } } /// Unregister a session, dropping its channels. + /// + /// Unconditional: removes whichever registration currently holds the + /// ID. Only for client-wide teardown, where every session is going away + /// regardless of owner. Owners of a specific registration must use + /// [`unregister_owned`](Self::unregister_owned). pub(crate) fn unregister(&self, session_id: &SessionId) { self.sessions.lock().remove(session_id.as_str()); } + /// Unregister a session only if it is still the registration identified + /// by `token`. + /// + /// Returns `true` when the entry was removed. A `false` result means + /// the registration had already been replaced by a newer one, which the + /// caller does not own and must leave alone. + pub(crate) fn unregister_owned( + &self, + session_id: &SessionId, + token: RegistrationToken, + ) -> bool { + let mut sessions = self.sessions.lock(); + if sessions + .get(session_id.as_str()) + .is_some_and(|senders| senders.token == token) + { + sessions.remove(session_id.as_str()); + true + } else { + false + } + } + /// Snapshot every currently-registered session ID. /// /// Used by [`Client::stop`](crate::Client::stop) to iterate active @@ -238,3 +296,37 @@ impl SessionRouter { } } } + +#[cfg(test)] +mod tests { + use super::*; + + fn session_id() -> SessionId { + SessionId::new("router-ownership") + } + + #[test] + fn each_registration_gets_a_distinct_token() { + let router = SessionRouter::new(); + let first = router.register(&session_id()); + let second = router.register(&session_id()); + assert_ne!(first.token, second.token); + } + + #[test] + fn unregister_owned_removes_only_the_matching_registration() { + let router = SessionRouter::new(); + let stale = router.register(&session_id()); + let live = router.register(&session_id()); + + // The stale owner must not evict the registration that replaced it. + assert!(!router.unregister_owned(&session_id(), stale.token)); + assert_eq!(router.session_ids(), vec![session_id()]); + + assert!(router.unregister_owned(&session_id(), live.token)); + assert!(router.session_ids().is_empty()); + + // Removing twice is a no-op rather than evicting a future tenant. + assert!(!router.unregister_owned(&session_id(), live.token)); + } +} diff --git a/rust/src/session.rs b/rust/src/session.rs index fd514a196f..636157b84f 100644 --- a/rust/src/session.rs +++ b/rust/src/session.rs @@ -267,6 +267,8 @@ pub struct Session { event_tx: tokio::sync::broadcast::Sender, github_token_registration: ParkingLotMutex>, + /// Identity of this session's router registration. + registration_token: crate::router::RegistrationToken, } impl Session { @@ -665,8 +667,9 @@ impl Session { ) .await?; self.stop_event_loop().await; - self.client.unregister_session(&self.id); self.github_token_registration.lock().take(); + self.client + .unregister_session_owned(&self.id, self.registration_token); Ok(()) } @@ -743,8 +746,9 @@ impl Drop for Session { // tokio runtime when it next polls; we intentionally don't await // it here because Drop is sync. self.shutdown.cancel(); - self.client.unregister_session(&self.id); self.github_token_registration.lock().take(); + self.client + .unregister_session_owned(&self.id, self.registration_token); } } @@ -1319,6 +1323,7 @@ impl Client { open_canvases, event_tx, github_token_registration: ParkingLotMutex::new(github_token_registration), + registration_token, }; apply_mode_post_create_patch( &session, @@ -1599,6 +1604,7 @@ impl Client { open_canvases, event_tx, github_token_registration: ParkingLotMutex::new(github_token_registration), + registration_token, }; apply_mode_post_create_patch( &session, diff --git a/rust/tests/prepared_session_test.rs b/rust/tests/prepared_session_test.rs index 694f161205..fa2f9ae049 100644 --- a/rust/tests/prepared_session_test.rs +++ b/rust/tests/prepared_session_test.rs @@ -12,7 +12,9 @@ use std::time::Duration; use github_copilot_sdk::session::PreparedSession; use github_copilot_sdk::subscription::{EventSubscription, RecvErrorKind}; -use github_copilot_sdk::types::{ResumeSessionConfig, SessionConfig, SessionId}; +use github_copilot_sdk::types::{ + CloudSessionOptions, CloudSessionRepository, ResumeSessionConfig, SessionConfig, SessionId, +}; use github_copilot_sdk::{Client, ErrorKind, SessionErrorKind}; use serde_json::{Value, json}; use tokio::io::{AsyncRead, AsyncWrite, AsyncWriteExt, duplex}; @@ -151,6 +153,10 @@ fn make_client() -> (Client, FakeServer) { ) } +fn cloud_options() -> CloudSessionOptions { + CloudSessionOptions::with_repository(CloudSessionRepository::new("octocat", "hello-world")) +} + fn create_result(session_id: &str) -> Value { json!({ "sessionId": session_id, "workspacePath": "/tmp/workspace" }) } @@ -801,3 +807,312 @@ fn prepared_session_is_send_static_and_not_clone() { // session can never be started twice. assert!(!CloneProbe::(PhantomData).is_clone()); } + +// --------------------------------------------------------------------------- +// Registration ownership: a stale startup guard must never unregister a +// newer registration that reused the same session ID. +// --------------------------------------------------------------------------- + +/// Drive a startup future until it parks awaiting its RPC response. +/// +/// The duration is a bound, not a correctness sleep: whether the future +/// actually reached the wire is asserted afterwards by reading the request, +/// which fails loudly on its own timeout if it did not. +const DRIVE: Duration = Duration::from_millis(50); + +#[tokio::test] +async fn stale_create_guard_does_not_unregister_same_id_retry() { + let (client, mut server) = make_client(); + let session_id = SessionId::new("stale-create-guard"); + + // First attempt: registers, sends `session.create`, then parks. + let mut first = Box::pin( + client + .prepare_session(SessionConfig::default().with_session_id(session_id.clone())) + .unwrap() + .start(), + ); + let _ = timeout(DRIVE, &mut first).await; + let first_req = server.read_request().await; + assert_eq!(first_req["method"], "session.create"); + + // Second attempt with the same pinned ID, started before the first is + // dropped, so it replaces the first attempt's router registration. + let prepared = client + .prepare_session( + SessionConfig::default() + .with_session_id(session_id.clone()) + .with_event_buffer_capacity(64), + ) + .unwrap(); + let mut events = prepared.subscribe(); + let mut second = Box::pin(prepared.start()); + let _ = timeout(DRIVE, &mut second).await; + let second_req = server.read_request().await; + assert_eq!(second_req["method"], "session.create"); + + // The stale guard runs now. It must not touch the live registration. + drop(first); + assert_eq!( + client.registered_session_ids_for_test(), + vec![session_id.clone()], + "a stale startup guard unregistered the live retry" + ); + + server + .respond(&second_req, create_result(session_id.as_str())) + .await; + let session = timeout(TIMEOUT, &mut second).await.unwrap().unwrap(); + + // Events must still route to the surviving registration. + server + .send_event( + session_id.as_str(), + "evt-after-stale", + "assistant.message", + false, + ) + .await; + let event = timeout(TIMEOUT, events.recv()).await.unwrap().unwrap(); + assert_eq!(event.id.as_str(), "evt-after-stale"); + drop(session); +} + +#[tokio::test] +async fn stale_resume_guard_does_not_unregister_same_id_retry() { + let (client, mut server) = make_client(); + let session_id = SessionId::new("stale-resume-guard"); + + let mut first = Box::pin( + client + .prepare_resume_session(ResumeSessionConfig::new(session_id.clone())) + .unwrap() + .start(), + ); + let _ = timeout(DRIVE, &mut first).await; + let first_req = server.read_request().await; + assert_eq!(first_req["method"], "session.resume"); + + let prepared = client + .prepare_resume_session( + ResumeSessionConfig::new(session_id.clone()).with_event_buffer_capacity(64), + ) + .unwrap(); + let mut events = prepared.subscribe(); + let mut second = Box::pin(prepared.start()); + let _ = timeout(DRIVE, &mut second).await; + let second_req = server.read_request().await; + assert_eq!(second_req["method"], "session.resume"); + + drop(first); + assert_eq!( + client.registered_session_ids_for_test(), + vec![session_id.clone()], + "a stale startup guard unregistered the live retry" + ); + + // Hand the surviving startup to a task: resume issues a follow-up + // `session.skills.reload` that only makes progress while it is polled. + let second = tokio::spawn(second); + server + .respond(&second_req, json!({ "sessionId": session_id.as_str() })) + .await; + server.answer_skills_reload().await; + let session = timeout(TIMEOUT, second).await.unwrap().unwrap().unwrap(); + + server + .send_event( + session_id.as_str(), + "evt-after-stale", + "assistant.message", + false, + ) + .await; + let event = timeout(TIMEOUT, events.recv()).await.unwrap().unwrap(); + assert_eq!(event.id.as_str(), "evt-after-stale"); + drop(session); +} + +/// A disconnected session must not unregister a same-ID session that +/// replaced it. +#[tokio::test] +async fn dropping_superseded_session_does_not_unregister_its_replacement() { + let (client, mut server) = make_client(); + let session_id = SessionId::new("superseded-session"); + + let first = tokio::spawn( + client + .prepare_session(SessionConfig::default().with_session_id(session_id.clone())) + .unwrap() + .start(), + ); + let first_req = server.read_request().await; + server + .respond(&first_req, create_result(session_id.as_str())) + .await; + let first_session = timeout(TIMEOUT, first).await.unwrap().unwrap().unwrap(); + + let prepared = client + .prepare_session( + SessionConfig::default() + .with_session_id(session_id.clone()) + .with_event_buffer_capacity(64), + ) + .unwrap(); + let mut events = prepared.subscribe(); + let second = tokio::spawn(prepared.start()); + let second_req = server.read_request().await; + server + .respond(&second_req, create_result(session_id.as_str())) + .await; + let second_session = timeout(TIMEOUT, second).await.unwrap().unwrap().unwrap(); + + // The superseded handle goes away; the live session must survive. + drop(first_session); + assert_eq!( + client.registered_session_ids_for_test(), + vec![session_id.clone()], + "dropping a superseded Session unregistered its replacement" + ); + + server + .send_event( + session_id.as_str(), + "evt-survivor", + "assistant.message", + false, + ) + .await; + let event = timeout(TIMEOUT, events.recv()).await.unwrap().unwrap(); + assert_eq!(event.id.as_str(), "evt-survivor"); + drop(second_session); +} + +// --------------------------------------------------------------------------- +// Deferred (server-assigned ID) create cancellation +// --------------------------------------------------------------------------- + +/// Cancelling a cloud create before the response arrives must leave no +/// registration behind, even though the session ID is only known to the +/// inline response callback. +#[tokio::test] +async fn cancelled_deferred_create_leaves_no_registration() { + let (client, mut server) = make_client(); + + let prepared = client + .prepare_session(SessionConfig::default().with_cloud(cloud_options())) + .unwrap(); + let mut events = prepared.subscribe(); + let start = tokio::spawn(prepared.start()); + + let create_req = server.read_request().await; + assert_eq!(create_req["method"], "session.create"); + assert!(create_req["params"]["sessionId"].is_null()); + + // Cancel before the server answers, then answer: the response carries + // the server-assigned ID the inline callback would register. + start.abort(); + let _ = start.await; + server + .respond(&create_req, create_result("server-assigned-id")) + .await; + + expect_closed(&mut events).await; + await_no_registrations(&client).await; + // Nothing may appear after the response has been fully processed. + server.expect_quiet().await; + assert!( + client.registered_session_ids_for_test().is_empty(), + "a cancelled deferred create left a registration behind" + ); + + // A fresh cloud create still works afterwards. + let retry = tokio::spawn( + client + .prepare_session(SessionConfig::default().with_cloud(cloud_options())) + .unwrap() + .start(), + ); + let retry_req = server.read_request().await; + server + .respond(&retry_req, create_result("server-assigned-retry")) + .await; + let session = timeout(TIMEOUT, retry).await.unwrap().unwrap().unwrap(); + assert_eq!(session.id().as_str(), "server-assigned-retry"); + drop(session); +} + +/// Poll (bounded) until `session_id` shows up on the client's router. +async fn await_registered(client: &Client, session_id: &str) { + let deadline = tokio::time::Instant::now() + TIMEOUT; + while !client + .registered_session_ids_for_test() + .iter() + .any(|id| id.as_str() == session_id) + { + assert!( + tokio::time::Instant::now() < deadline, + "inline callback never registered {session_id}" + ); + tokio::task::yield_now().await; + } +} + +/// The other half of the deferred-create window: cancellation lands *after* +/// the inline response callback has already registered the server-assigned +/// ID. The startup guard owns that registration and must remove it. +/// +/// Deterministic by construction — the start future is parked on its +/// response and never polled again, so the callback (which runs on the +/// JSON-RPC read task, independently of the caller) is guaranteed to have +/// registered before the future is dropped. +#[tokio::test] +async fn deferred_create_cancelled_after_callback_registered_is_cleaned_up() { + let (client, mut server) = make_client(); + + let prepared = client + .prepare_session(SessionConfig::default().with_cloud(cloud_options())) + .unwrap(); + let mut events = prepared.subscribe(); + let mut start = Box::pin(prepared.start()); + + // Drive to the wire, then park. + let _ = timeout(DRIVE, &mut start).await; + let create_req = server.read_request().await; + assert_eq!(create_req["method"], "session.create"); + + // The read task runs the inline callback and registers the ID while the + // caller's future stays unpolled. + server + .respond(&create_req, create_result("registered-then-cancelled")) + .await; + await_registered(&client, "registered-then-cancelled").await; + + // Cancellation now lands on a slot that already owns a registration. + drop(start); + + expect_closed(&mut events).await; + await_no_registrations(&client).await; + server.expect_quiet().await; + + // The same server-assigned ID can be handed out again without the dead + // attempt's cleanup interfering. + let retry = tokio::spawn( + client + .prepare_session(SessionConfig::default().with_cloud(cloud_options())) + .unwrap() + .start(), + ); + let retry_req = server.read_request().await; + server + .respond(&retry_req, create_result("registered-then-cancelled")) + .await; + let session = timeout(TIMEOUT, retry).await.unwrap().unwrap().unwrap(); + assert_eq!(session.id().as_str(), "registered-then-cancelled"); + assert_eq!( + client.registered_session_ids_for_test().len(), + 1, + "retry must hold exactly one registration" + ); + drop(session); +} From 708e7c88c8a903526a6f6266cc6c1c0c73b93507 Mon Sep 17 00:00:00 2001 From: jmoseley Date: Wed, 12 Aug 2026 08:52:04 -0700 Subject: [PATCH 03/12] [Rust] Gate registered_session_ids to test configurations `Client::registered_session_ids` has no caller in a default-feature build: the in-crate unit tests reach it under `cfg(test)`, and the public `registered_session_ids_for_test` wrapper is gated on `feature = "test-support"`. A plain `cargo build` or `cargo clippy` therefore warned `dead_code` for it. Gate the method on `any(test, feature = "test-support")`, matching the convention already used for the other test-only helpers in this file. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 0f9d5ac7-9999-4f37-82b3-a5533bfbc1f6 --- rust/src/lib.rs | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/rust/src/lib.rs b/rust/src/lib.rs index 096cd9ae95..0b42b247cd 100644 --- a/rust/src/lib.rs +++ b/rust/src/lib.rs @@ -2366,7 +2366,9 @@ impl Client { /// /// Crate-internal so in-crate unit tests can assert registration /// lifecycle without depending on the `test-support` feature, which - /// only gates the equivalent *public* test helper. + /// only gates the equivalent *public* test helper. Compiled only for + /// those two configurations — a default-feature build has no caller. + #[cfg(any(test, feature = "test-support"))] pub(crate) fn registered_session_ids(&self) -> Vec { self.inner.router.session_ids() } From d944f77d385474aed3505d2711ca53cd3de5e3df Mon Sep 17 00:00:00 2001 From: jmoseley Date: Wed, 12 Aug 2026 15:02:12 -0700 Subject: [PATCH 04/12] [Rust] Narrow the PreparedSession guarantee to routed events `Client::prepare_session` promised consumers would observe "every event a session emits". That is broader than the implementation for cloud creates with a server-assigned ID: the SDK cannot register the session on its notification router until the `session.create` response arrives, so notifications emitted before that point are not routable to any session and never reach a subscriber. Qualify the primary API documentation and the changelog as *routed* events, and point callers at pinning `SessionConfig::session_id` for complete pre-response coverage. `PreparedSession`'s type-level docs, `rust/README.md`, and `docs/features/streaming-events.md` already documented this limitation; the entry-point docs now match them. Documentation only: no API, behavior, or wire change. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 0f9d5ac7-9999-4f37-82b3-a5533bfbc1f6 --- CHANGELOG.md | 4 +++- rust/src/session.rs | 14 +++++++++++--- 2 files changed, 14 insertions(+), 4 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 0f5dd44870..ed2bc1a0d6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -46,7 +46,7 @@ const token = process.env.GITHUB_TOKEN; ### Feature: early session-event subscription (Rust) -The Rust SDK can now observe a session's events from its very first event. `Client::prepare_session` and `Client::prepare_resume_session` return an inert `PreparedSession` that owns the session's event channel, so a subscription can be installed *before* any protocol activity begins: +The Rust SDK can now observe every event routed to a session, starting with that session's very first routed event. `Client::prepare_session` and `Client::prepare_resume_session` return an inert `PreparedSession` that owns the session's event channel, so a subscription can be installed *before* any protocol activity begins: ```rust let prepared = client.prepare_session( @@ -58,6 +58,8 @@ let session = prepared.start().await?; Previously, `Session::subscribe` could only be called on the returned session, so events the runtime emitted while `session.create` / `session.resume` was still in flight were broadcast with no receiver installed and dropped. Ephemeral events such as `session.idle` are not persisted, so they could not be recovered with `getMessages` either. +The guarantee is scoped to *routed* events. For cloud sessions where the server assigns the session ID, the SDK cannot route notifications until the `session.create` response arrives and the ID is known, so events emitted before that point are not routable to any session. Pin `session_id` on the config to get router registration before the RPC, and with it complete pre-response coverage. + `prepare_*` is synchronous and inert: it validates the buffer capacity and allocates a local channel, and performs no router registration, task spawn, or wire activity until `start()` is first polled. `start(self)` consumes the handle and `PreparedSession` is not `Clone`, so a prepared session can never produce two event loops. Dropping an unstarted handle leaves no state behind; dropping a polled `start()` future cancels the startup and unregisters the session, so a retry with the same session ID succeeds. Session registrations now carry an ownership identity, so cleanup removes only the exact registration it owns and an abandoned startup can never evict a same-ID retry (or a session that replaced it). Both `SessionConfig` and `ResumeSessionConfig` gained a runtime-only `event_buffer_capacity` option (default 512, `Some(0)` rejected as an invalid config). The buffer is finite, so slow subscribers observe `Lagged` rather than applying backpressure; consumers that need a lossless view of a large startup burst should size the buffer accordingly or drain concurrently with `start()`. diff --git a/rust/src/session.rs b/rust/src/session.rs index 636157b84f..5f525f7d8f 100644 --- a/rust/src/session.rs +++ b/rust/src/session.rs @@ -899,14 +899,22 @@ impl Client { /// Call [`PreparedSession::start`] to actually create the session. /// /// This is the loss-free entry point for consumers that must observe - /// every event a session emits, including events the CLI emits while - /// `session.create` is still in flight and ephemeral events (such as - /// `session.idle`) that cannot be recovered from + /// every *routed* event a session emits, including events the CLI emits + /// while `session.create` is still in flight and ephemeral events (such + /// as `session.idle`) that cannot be recovered from /// [`Session::get_messages`]. [`create_session`](Self::create_session) /// is a thin wrapper over `prepare_session(...)?.start()` and cannot /// offer the same guarantee, because the subscription can only be /// installed after the returned `Session` exists. /// + /// Routing requires a known session ID. When the server assigns the ID, + /// the SDK cannot register the session on its notification router until + /// the `session.create` response arrives, so notifications emitted + /// before that point are not routable and stay unobservable. Pin + /// [`SessionConfig::session_id`](crate::types::SessionConfig::session_id) + /// for complete pre-response coverage — see the "Server-assigned session + /// IDs" section on [`PreparedSession`]. + /// /// # Inertness /// /// `prepare_session` performs no router registration, spawns no task, From 0acd139a5aee615ddaa3271187eea9ecd3a20113 Mon Sep 17 00:00:00 2001 From: jmoseley Date: Wed, 12 Aug 2026 15:03:02 -0700 Subject: [PATCH 05/12] [Docs] Unwrap the lone Rust tab in the streaming-events article The "Subscribing before a session starts" section has only a Rust example, and the docs normalization pipeline converts a `
` group into a tabbed language switcher only when two or more consecutive blocks are present. A single block renders as raw collapsible HTML on docs.github.com. Drop the `
`/`` wrapper and leave the code fence directly in the article, matching the repository docs style guide. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 0f9d5ac7-9999-4f37-82b3-a5533bfbc1f6 --- docs/features/streaming-events.md | 5 ----- 1 file changed, 5 deletions(-) diff --git a/docs/features/streaming-events.md b/docs/features/streaming-events.md index e0eae797fb..17acdc16df 100644 --- a/docs/features/streaming-events.md +++ b/docs/features/streaming-events.md @@ -225,9 +225,6 @@ A session can emit events before its create or resume call returns. The agent ma > [!TIP] > **(Rust)** `Client::prepare_session` and `Client::prepare_resume_session` return a `PreparedSession` that owns the session's event channel before any protocol activity happens. Subscribe first, then call `start()`. -
-Rust - ```rust use github_copilot_sdk::{Client, SessionConfig}; @@ -252,8 +249,6 @@ async fn create_without_missing_startup_events( } ``` -
- `prepare_*` is synchronous and inert: it validates the buffer capacity, allocates a local channel, and does nothing else. No session is registered and nothing reaches the CLI until `start()` is first polled. Dropping a prepared session that was never started leaves no state behind and closes its subscriptions; dropping the `start()` future cancels the in-flight startup and unregisters the session, so a retry with the same session ID succeeds. Cleanup is scoped to the exact registration the abandoned startup owned, so it cannot evict a retry that has already taken over the same session ID. Startup buffering is worth planning for: From 5d1e6a6630eee53b757ac55d80e40edfa145648c Mon Sep 17 00:00:00 2001 From: jmoseley Date: Wed, 12 Aug 2026 15:03:12 -0700 Subject: [PATCH 06/12] [Rust] Keep session IDs out of prepared-session test diagnostics Two polling helpers formatted session identifiers into their failure messages: `await_no_registrations` rendered the router's registered ID list with `{:?}`, and `await_registered` interpolated the awaited ID. A downstream consumer that vendors this crate has CodeQL rules flagging identifiers reaching formatted output, so both were reported there even though the SDK's own analysis was clean. Report an outstanding-registration count and a static expectation message instead. Both helpers keep their exact predicates and deadline behavior: `await_no_registrations` still returns only when the router holds zero registrations, and `await_registered` still blocks on the exact ID it was given, so no assertion is weakened. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 0f9d5ac7-9999-4f37-82b3-a5533bfbc1f6 --- rust/tests/prepared_session_test.rs | 17 +++++++++++++---- 1 file changed, 13 insertions(+), 4 deletions(-) diff --git a/rust/tests/prepared_session_test.rs b/rust/tests/prepared_session_test.rs index fa2f9ae049..7225b6c314 100644 --- a/rust/tests/prepared_session_test.rs +++ b/rust/tests/prepared_session_test.rs @@ -186,13 +186,19 @@ fn expect_error(result: Result) -> github_copil } /// Poll (bounded) until the client's router has no registered sessions. +/// +/// Diagnostics report how many registrations are outstanding rather than +/// which ones: session IDs are not written to test output. async fn await_no_registrations(client: &Client) { let deadline = tokio::time::Instant::now() + TIMEOUT; - while !client.registered_session_ids_for_test().is_empty() { + loop { + let outstanding = client.registered_session_ids_for_test().len(); + if outstanding == 0 { + return; + } assert!( tokio::time::Instant::now() < deadline, - "session registration was never cleaned up: {:?}", - client.registered_session_ids_for_test() + "{outstanding} session registration(s) were never cleaned up" ); tokio::task::yield_now().await; } @@ -1043,6 +1049,9 @@ async fn cancelled_deferred_create_leaves_no_registration() { } /// Poll (bounded) until `session_id` shows up on the client's router. +/// +/// The failure message names the expectation, not the ID: session IDs are +/// not written to test output. async fn await_registered(client: &Client, session_id: &str) { let deadline = tokio::time::Instant::now() + TIMEOUT; while !client @@ -1052,7 +1061,7 @@ async fn await_registered(client: &Client, session_id: &str) { { assert!( tokio::time::Instant::now() < deadline, - "inline callback never registered {session_id}" + "inline callback never registered the expected session" ); tokio::task::yield_now().await; } From 719096f251a1e1d28b2f6cd772b8cea8cf4f80c0 Mon Sep 17 00:00:00 2001 From: jmoseley Date: Wed, 12 Aug 2026 15:53:24 -0700 Subject: [PATCH 07/12] [Rust] Cover the MCP-auth interest failure path on create and resume MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `session.eventLog.registerInterest` is the last fallible step of startup when an MCP-auth handler is installed, and its failure branch was untested: the existing MCP-auth tests only return successful interest responses, and the prepared-session failure tests covered create RPC errors and ID mismatches only. Add a failing-interest test for each of create and resume, asserting the same contract the sibling failure tests assert: the original error kind reaches the caller, the router registration is gone, and a subscription installed before `start()` is closed. The resume test also asserts the best-effort `session.skills.reload` is never issued, since interest registration runs ahead of it. Verified by mutation that both tests execute the branch. Removing the `registration.cleanup(event_loop)` call does not turn them red, because `PendingSessionRegistration::drop` cancels and releases the same registration synchronously — the explicit cleanup is defense in depth on this path, and the tests assert the observable contract rather than which of the two mechanisms performed it. Test-only change: no API, behavior, or wire impact. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 0f9d5ac7-9999-4f37-82b3-a5533bfbc1f6 --- rust/tests/prepared_session_test.rs | 106 +++++++++++++++++++++++++++- 1 file changed, 105 insertions(+), 1 deletion(-) diff --git a/rust/tests/prepared_session_test.rs b/rust/tests/prepared_session_test.rs index 7225b6c314..19ee465512 100644 --- a/rust/tests/prepared_session_test.rs +++ b/rust/tests/prepared_session_test.rs @@ -8,12 +8,16 @@ #![allow(clippy::unwrap_used)] use std::marker::PhantomData; +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::session::PreparedSession; use github_copilot_sdk::subscription::{EventSubscription, RecvErrorKind}; use github_copilot_sdk::types::{ - CloudSessionOptions, CloudSessionRepository, ResumeSessionConfig, SessionConfig, SessionId, + CloudSessionOptions, CloudSessionRepository, RequestId, ResumeSessionConfig, SessionConfig, + SessionId, }; use github_copilot_sdk::{Client, ErrorKind, SessionErrorKind}; use serde_json::{Value, json}; @@ -140,6 +144,23 @@ impl FakeServer { } } +/// Minimal MCP-auth handler: its presence is what makes the SDK register +/// `mcp.oauth_required` interest after create/resume, which is the branch +/// under test. It is never invoked by these tests. +struct CancelMcpAuthHandler; + +#[async_trait] +impl McpAuthHandler for CancelMcpAuthHandler { + async fn handle( + &self, + _session_id: SessionId, + _request_id: RequestId, + _request: McpAuthRequest, + ) -> McpAuthResult { + McpAuthResult::Cancelled + } +} + fn make_client() -> (Client, FakeServer) { let (client_write, server_read) = duplex(1 << 20); let (server_write, client_read) = duplex(1 << 20); @@ -625,6 +646,89 @@ async fn resume_session_id_mismatch_preserves_kind_and_cleans_up() { expect_closed(&mut events).await; } +/// The MCP-auth interest registration that follows a successful +/// `session.create` is the last fallible step before the session handle is +/// handed out. When it fails, the startup must unwind exactly like any +/// other create failure: original error kind preserved, router +/// registration removed, and subscriptions taken before `start()` closed. +#[tokio::test] +async fn create_mcp_auth_interest_error_preserves_kind_and_cleans_up() { + let (client, mut server) = make_client(); + let session_id = SessionId::new("prepared-create-interest-error"); + + let prepared = client + .prepare_session( + SessionConfig::default() + .with_session_id(session_id.clone()) + .with_mcp_auth_handler(Arc::new(CancelMcpAuthHandler)), + ) + .unwrap(); + let mut events = prepared.subscribe(); + let start = tokio::spawn(prepared.start()); + + let create_req = server.read_request().await; + assert_eq!(create_req["method"], "session.create"); + server + .respond(&create_req, create_result(session_id.as_str())) + .await; + + let interest_req = server.read_request().await; + assert_eq!(interest_req["method"], "session.eventLog.registerInterest"); + assert_eq!(interest_req["params"]["eventType"], "mcp.oauth_required"); + server + .respond_error(&interest_req, -32003, "interest registration failed") + .await; + + let error = expect_error(timeout(TIMEOUT, start).await.unwrap().unwrap()); + assert!( + matches!(error.kind(), ErrorKind::Rpc { code: -32003 }), + "unexpected error kind: {:?}", + error.kind() + ); + expect_closed(&mut events).await; + await_no_registrations(&client).await; +} + +/// The resume counterpart. Interest registration runs before the +/// best-effort `session.skills.reload`, so a failure must abort the +/// startup without issuing the reload. +#[tokio::test] +async fn resume_mcp_auth_interest_error_preserves_kind_and_cleans_up() { + let (client, mut server) = make_client(); + let session_id = SessionId::new("prepared-resume-interest-error"); + + let prepared = client + .prepare_resume_session( + ResumeSessionConfig::new(session_id.clone()) + .with_mcp_auth_handler(Arc::new(CancelMcpAuthHandler)), + ) + .unwrap(); + let mut events = prepared.subscribe(); + let start = tokio::spawn(prepared.start()); + + 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 + .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 }), + "unexpected error kind: {:?}", + error.kind() + ); + server.expect_quiet().await; + expect_closed(&mut events).await; + await_no_registrations(&client).await; +} + #[tokio::test] async fn zero_event_buffer_capacity_is_invalid_config() { let (client, _server) = make_client(); From a4d5a0a95e9d98b60dd31d5b559b5809632b6ad2 Mon Sep 17 00:00:00 2001 From: jmoseley Date: Wed, 12 Aug 2026 15:53:33 -0700 Subject: [PATCH 08/12] [Rust] Reconcile the subscribe docs with the buffering contract `PreparedSession::subscribe` promised that "every subscriber receives every event", which contradicts the paragraph directly above it: the broadcast buffer is finite, so a subscriber that falls behind the configured capacity observes `Lagged` and skips events instead of applying backpressure. Say that explicitly and link the `Lagged` type. Also replace "should" with "must" where the streaming-events article and the changelog describe what a consumer needing lossless startup delivery has to do. The docs style guide reserves ambiguous modals for optional actions, and `rust/README.md` already phrased this as "must". Documentation only: no API, behavior, or wire change. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 0f9d5ac7-9999-4f37-82b3-a5533bfbc1f6 --- CHANGELOG.md | 2 +- docs/features/streaming-events.md | 2 +- rust/src/session.rs | 10 +++++++--- 3 files changed, 9 insertions(+), 5 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index ed2bc1a0d6..59961b4f0b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -62,7 +62,7 @@ The guarantee is scoped to *routed* events. For cloud sessions where the server `prepare_*` is synchronous and inert: it validates the buffer capacity and allocates a local channel, and performs no router registration, task spawn, or wire activity until `start()` is first polled. `start(self)` consumes the handle and `PreparedSession` is not `Clone`, so a prepared session can never produce two event loops. Dropping an unstarted handle leaves no state behind; dropping a polled `start()` future cancels the startup and unregisters the session, so a retry with the same session ID succeeds. Session registrations now carry an ownership identity, so cleanup removes only the exact registration it owns and an abandoned startup can never evict a same-ID retry (or a session that replaced it). -Both `SessionConfig` and `ResumeSessionConfig` gained a runtime-only `event_buffer_capacity` option (default 512, `Some(0)` rejected as an invalid config). The buffer is finite, so slow subscribers observe `Lagged` rather than applying backpressure; consumers that need a lossless view of a large startup burst should size the buffer accordingly or drain concurrently with `start()`. +Both `SessionConfig` and `ResumeSessionConfig` gained a runtime-only `event_buffer_capacity` option (default 512, `Some(0)` rejected as an invalid config). The buffer is finite, so slow subscribers observe `Lagged` rather than applying backpressure; consumers that need a lossless view of a large startup burst must size the buffer accordingly or drain concurrently with `start()`. `create_session` and `resume_session` are unchanged wrappers over `prepare_*(...)?.start()` with identical RPC sequences and error kinds. diff --git a/docs/features/streaming-events.md b/docs/features/streaming-events.md index 17acdc16df..0292f98e00 100644 --- a/docs/features/streaming-events.md +++ b/docs/features/streaming-events.md @@ -255,7 +255,7 @@ Startup buffering is worth planning for: * The event buffer is finite—512 events unless `event_buffer_capacity` overrides it. A capacity of `0` is rejected with an invalid-config error rather than clamped. * Slow subscribers observe a `Lagged` error reporting how many events were skipped. They never apply backpressure to the session's event loop. -* Consumers that need a lossless view of a large startup burst should either configure a capacity that covers it or drain the subscription concurrently with `start()`. +* 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()`. > [!NOTE] > For cloud sessions where the server assigns the session ID, the SDK cannot route notifications until the create response arrives and the ID is known. Events emitted before that point are not routable to any session. The guarantee is narrower: routed events are never dropped for lack of an installed receiver. Pin `session_id` on the config to get routing—and full pre-response coverage—from the first byte. diff --git a/rust/src/session.rs b/rust/src/session.rs index 5f525f7d8f..d9c2e33f44 100644 --- a/rust/src/session.rs +++ b/rust/src/session.rs @@ -1726,9 +1726,13 @@ impl PreparedSession { /// ephemeral events such as `session.idle` that /// [`Session::get_messages`] cannot recover. /// - /// May be called any number of times; every subscriber receives every - /// event. Subscriptions taken here close if the prepared session is - /// dropped without starting, or if startup fails. + /// May be called any number of times, and each subscriber receives its + /// own copy of the stream — subject to the buffering contract above. A + /// subscriber that falls further behind than the configured capacity + /// observes [`Lagged`](crate::subscription::Lagged) and skips the + /// events it missed, rather than stalling the session's event loop. + /// Subscriptions taken here close if the prepared session is dropped + /// without starting, or if startup fails. pub fn subscribe(&self) -> crate::subscription::EventSubscription { crate::subscription::EventSubscription::new(self.event_tx.subscribe()) } From e5347f3fb3f64e33ac7abf401ed49c5d19497f31 Mon Sep 17 00:00:00 2001 From: Steve Sanderson Date: Thu, 3 Sep 2026 16:20:31 +0000 Subject: [PATCH 09/12] [Rust] Reconcile PreparedSession with current main Preserve current router registration ownership and feature configuration while rebasing the focused startup subscription change.\n\nCo-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- rust/Cargo.toml | 1 + rust/src/session.rs | 74 ++++++++++++++++++++++++++++++++------------- rust/src/types.rs | 6 ++++ 3 files changed, 60 insertions(+), 21 deletions(-) diff --git a/rust/Cargo.toml b/rust/Cargo.toml index 71d754f26e..4495d3928c 100644 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -118,6 +118,7 @@ required-features = ["test-support"] test = false bench = false +[[test]] name = "prepared_session_test" required-features = ["test-support"] [build-dependencies] diff --git a/rust/src/session.rs b/rust/src/session.rs index d9c2e33f44..1111d7cffe 100644 --- a/rust/src/session.rs +++ b/rust/src/session.rs @@ -146,17 +146,22 @@ struct PendingSessionRegistration { enum PendingSessionId { /// The ID was known before the RPC was issued (resume, and create with a /// client- or caller-supplied ID). - Known(SessionId), + Known(SessionId, crate::router::RegistrationToken), /// Server-assigned ID, populated by the `session.create` inline response /// callback. `None` in the stash means nothing was ever registered. - Deferred(Arc>>), + Deferred(Arc>>), } impl PendingSessionRegistration { - fn new(client: Client, session_id: SessionId, shutdown: CancellationToken) -> Self { + fn new( + client: Client, + session_id: SessionId, + token: crate::router::RegistrationToken, + shutdown: CancellationToken, + ) -> Self { Self { client, - session_id: PendingSessionId::Known(session_id), + session_id: PendingSessionId::Known(session_id, token), shutdown, disarmed: false, } @@ -165,7 +170,7 @@ impl PendingSessionRegistration { /// Guard for a registration whose session ID is assigned by the server. fn deferred( client: Client, - stash: Arc>>, + stash: Arc>>, shutdown: CancellationToken, ) -> Self { Self { @@ -178,7 +183,7 @@ impl PendingSessionRegistration { fn registered_id(&self) -> Option { match &self.session_id { - PendingSessionId::Known(id) => Some(id.clone()), + PendingSessionId::Known(id, _) => Some(id.clone()), PendingSessionId::Deferred(stash) => stash.lock().as_ref().map(|(id, _)| id.clone()), } } @@ -186,15 +191,21 @@ impl PendingSessionRegistration { /// Re-target the guard at a now-known session ID. Used by /// `session.create` once the response has been parsed and the stash has /// been drained into the event loop. - fn resolve_to(&mut self, session_id: SessionId) { - self.session_id = PendingSessionId::Known(session_id); + fn resolve_to(&mut self, session_id: SessionId, token: crate::router::RegistrationToken) { + self.session_id = PendingSessionId::Known(session_id, token); } async fn cleanup(mut self, event_loop: JoinHandle<()>) { self.shutdown.cancel(); let _ = event_loop.await; if let Some(id) = self.registered_id() { - self.client.unregister_session(&id); + if let PendingSessionId::Known(_, token) = self.session_id { + self.client.unregister_session_owned(&id, token); + } else if let PendingSessionId::Deferred(stash) = &self.session_id + && let Some((id, registration)) = stash.lock().as_ref() + { + self.client.unregister_session_owned(id, registration.token); + } } self.disarmed = true; } @@ -209,7 +220,13 @@ impl Drop for PendingSessionRegistration { if !self.disarmed { self.shutdown.cancel(); if let Some(id) = self.registered_id() { - self.client.unregister_session(&id); + if let PendingSessionId::Known(_, token) = self.session_id { + self.client.unregister_session_owned(&id, token); + } else if let PendingSessionId::Deferred(stash) = &self.session_id + && let Some((id, registration)) = stash.lock().as_ref() + { + self.client.unregister_session_owned(id, registration.token); + } } } } @@ -1196,7 +1213,7 @@ impl Client { // For non-cloud sessions, register up-front so the CLI can issue // session-scoped requests during session.create processing. let inline_stash: Arc< - ParkingLotMutex>, + ParkingLotMutex>, > = Arc::new(ParkingLotMutex::new(None)); let inline_callback: Option = if let Some(ref sid) = @@ -1232,8 +1249,8 @@ impl Client { // `register_session` takes the router lock, never the stash // lock, so there is no lock-order inversion here. let mut stashed = stash.lock(); - let channels = client.register_session(&parsed.session_id); - *stashed = Some((parsed.session_id, channels)); + let registration = client.register_session(&parsed.session_id); + *stashed = Some((parsed.session_id, registration)); Ok(()) })) }; @@ -1243,9 +1260,15 @@ impl Client { // token and unregisters whatever was registered on the router. For // the cloud path the ID is only known once the inline callback has // run, so the guard reads the stash at cleanup time. - let mut registration = match local_session_id { + let mut pending_registration = match local_session_id { Some(ref sid) => { - PendingSessionRegistration::new(self.clone(), sid.clone(), shutdown.clone()) + let token = inline_stash + .lock() + .as_ref() + .expect("session registration must exist") + .1 + .token; + PendingSessionRegistration::new(self.clone(), sid.clone(), token, shutdown.clone()) } None => PendingSessionRegistration::deferred( self.clone(), @@ -1274,11 +1297,13 @@ impl Client { .into()); } - let (session_id, channels) = inline_stash + let (session_id, registration) = inline_stash .lock() .take() .expect("session registration must have populated stash on success"); - registration.resolve_to(session_id.clone()); + let channels = registration.channels; + let registration_token = registration.token; + pending_registration.resolve_to(session_id.clone(), registration_token); let event_loop = spawn_event_loop( session_id.clone(), self.clone(), @@ -1308,7 +1333,7 @@ impl Client { if has_mcp_auth_handler && let Err(error) = register_mcp_auth_interest(self, &session_id).await { - registration.cleanup(event_loop).await; + pending_registration.cleanup(event_loop).await; return Err(error); } @@ -1317,7 +1342,7 @@ impl Client { session_id = %session_id, "Client::create_session complete" ); - registration.disarm(); + pending_registration.disarm(); let session = Session { id: session_id, cwd: self.cwd().clone(), @@ -1488,7 +1513,9 @@ impl Client { let capabilities = Arc::new(parking_lot::RwLock::new(SessionCapabilities::default())); let setup_start = Instant::now(); - let channels = self.register_session(&session_id); + let registration = self.register_session(&session_id); + let registration_token = registration.token; + let channels = registration.channels; let idle_waiter = Arc::new(ParkingLotMutex::new(None)); let open_canvases = Arc::new(parking_lot::RwLock::new(Vec::new())); let event_loop = spawn_event_loop( @@ -1509,7 +1536,12 @@ impl Client { shutdown.clone(), ); let mut registration = - PendingSessionRegistration::new(self.clone(), session_id.clone(), shutdown.clone()); + PendingSessionRegistration::new( + self.clone(), + session_id.clone(), + registration_token, + shutdown.clone(), + ); tracing::debug!( elapsed_ms = setup_start.elapsed().as_millis(), session_id = %session_id, diff --git a/rust/src/types.rs b/rust/src/types.rs index 68bde9dfb3..e281867e05 100644 --- a/rust/src/types.rs +++ b/rust/src/types.rs @@ -3316,6 +3316,9 @@ impl SessionConfig { /// Set feature-flag values resolved by the host for this session. pub fn with_feature_flags(mut self, feature_flags: HashMap) -> Self { self.feature_flags = Some(feature_flags); + self + } + /// Set [`Self::event_buffer_capacity`]. /// /// A capacity of `0` is rejected with @@ -4626,6 +4629,9 @@ impl ResumeSessionConfig { /// Re-supply feature-flag values resolved by the host on resume. pub fn with_feature_flags(mut self, feature_flags: HashMap) -> Self { self.feature_flags = Some(feature_flags); + self + } + /// Set [`Self::event_buffer_capacity`]. /// /// A capacity of `0` is rejected with From ed274da92c348b9e8ab4d161fc5efae3ea2f03b8 Mon Sep 17 00:00:00 2001 From: Steve Sanderson Date: Thu, 3 Sep 2026 16:26:17 +0000 Subject: [PATCH 10/12] [Rust] Format rebased session startup code Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- rust/src/session.rs | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/rust/src/session.rs b/rust/src/session.rs index 1111d7cffe..c0c13e22b0 100644 --- a/rust/src/session.rs +++ b/rust/src/session.rs @@ -1535,13 +1535,12 @@ impl Client { event_tx.clone(), shutdown.clone(), ); - let mut registration = - PendingSessionRegistration::new( - self.clone(), - session_id.clone(), - registration_token, - shutdown.clone(), - ); + let mut registration = PendingSessionRegistration::new( + self.clone(), + session_id.clone(), + registration_token, + shutdown.clone(), + ); tracing::debug!( elapsed_ms = setup_start.elapsed().as_millis(), session_id = %session_id, From e2815f8488ca2fa07c28ba1c21d7e87aa0ae2e0f Mon Sep 17 00:00:00 2001 From: Steve Sanderson Date: Thu, 3 Sep 2026 16:36:54 +0000 Subject: [PATCH 11/12] [Rust] Avoid exposing session IDs in diagnostics Use a count-only test helper for the polling diagnostic so CodeQL does not treat session identifiers as logged data.\n\nCo-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- rust/src/lib.rs | 6 ++++++ rust/tests/prepared_session_test.rs | 2 +- 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/rust/src/lib.rs b/rust/src/lib.rs index 0b42b247cd..78b7dcc0a2 100644 --- a/rust/src/lib.rs +++ b/rust/src/lib.rs @@ -2624,6 +2624,12 @@ impl Client { self.registered_session_ids() } + #[cfg(feature = "test-support")] + #[doc(hidden)] + pub fn registered_session_count_for_test(&self) -> usize { + self.registered_session_ids().len() + } + #[cfg(feature = "test-support")] #[doc(hidden)] /// Disconnect and delete every session owned by this test client's isolated diff --git a/rust/tests/prepared_session_test.rs b/rust/tests/prepared_session_test.rs index 19ee465512..df85027500 100644 --- a/rust/tests/prepared_session_test.rs +++ b/rust/tests/prepared_session_test.rs @@ -213,7 +213,7 @@ fn expect_error(result: Result) -> github_copil async fn await_no_registrations(client: &Client) { let deadline = tokio::time::Instant::now() + TIMEOUT; loop { - let outstanding = client.registered_session_ids_for_test().len(); + let outstanding = client.registered_session_count_for_test(); if outstanding == 0 { return; } From 6520628d05dee759fbfecbef01b41a515c63943e Mon Sep 17 00:00:00 2001 From: Steve Sanderson Date: Thu, 3 Sep 2026 16:49:35 +0000 Subject: [PATCH 12/12] [Rust] Keep session IDs out of test diagnostics CodeQL flagged the registration polling helper because the count was derived from a Vec. Count registrations directly on the router instead, and assert on registration identity without formatting IDs into failure messages. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- rust/src/lib.rs | 5 +++- rust/src/router.rs | 6 +++++ rust/tests/prepared_session_test.rs | 37 +++++++++++++++++++---------- 3 files changed, 34 insertions(+), 14 deletions(-) diff --git a/rust/src/lib.rs b/rust/src/lib.rs index 78b7dcc0a2..5d7957a4f2 100644 --- a/rust/src/lib.rs +++ b/rust/src/lib.rs @@ -2626,8 +2626,11 @@ impl Client { #[cfg(feature = "test-support")] #[doc(hidden)] + /// Count the sessions currently registered on this client's notification + /// router. Deliberately never materialises the session IDs themselves so + /// they cannot leak into test diagnostics. pub fn registered_session_count_for_test(&self) -> usize { - self.registered_session_ids().len() + self.inner.router.session_count() } #[cfg(feature = "test-support")] diff --git a/rust/src/router.rs b/rust/src/router.rs index f4aee25d53..2b09713723 100644 --- a/rust/src/router.rs +++ b/rust/src/router.rs @@ -126,6 +126,12 @@ impl SessionRouter { self.sessions.lock().keys().cloned().collect() } + /// Count the currently-registered sessions without exposing their IDs. + #[cfg(any(test, feature = "test-support"))] + pub(crate) fn session_count(&self) -> usize { + self.sessions.lock().len() + } + /// Drop all registered session channels. /// /// Used by [`Client::force_stop`](crate::Client::force_stop) to release diff --git a/rust/tests/prepared_session_test.rs b/rust/tests/prepared_session_test.rs index df85027500..3fefccabb6 100644 --- a/rust/tests/prepared_session_test.rs +++ b/rust/tests/prepared_session_test.rs @@ -206,6 +206,17 @@ fn expect_error(result: Result) -> github_copil } } +/// Assert the router holds exactly one registration, and that it is +/// `session_id`. +/// +/// The failure message is a fixed string: session IDs are never written to +/// test output. +fn assert_only_registration(client: &Client, session_id: &SessionId, context: &str) { + let registered = client.registered_session_ids_for_test(); + let matches_expected = registered.len() == 1 && registered[0] == *session_id; + assert!(matches_expected, "{context}"); +} + /// Poll (bounded) until the client's router has no registered sessions. /// /// Diagnostics report how many registrations are outstanding rather than @@ -963,10 +974,10 @@ async fn stale_create_guard_does_not_unregister_same_id_retry() { // The stale guard runs now. It must not touch the live registration. drop(first); - assert_eq!( - client.registered_session_ids_for_test(), - vec![session_id.clone()], - "a stale startup guard unregistered the live retry" + assert_only_registration( + &client, + &session_id, + "a stale startup guard unregistered the live retry", ); server @@ -1015,10 +1026,10 @@ async fn stale_resume_guard_does_not_unregister_same_id_retry() { assert_eq!(second_req["method"], "session.resume"); drop(first); - assert_eq!( - client.registered_session_ids_for_test(), - vec![session_id.clone()], - "a stale startup guard unregistered the live retry" + assert_only_registration( + &client, + &session_id, + "a stale startup guard unregistered the live retry", ); // Hand the surviving startup to a task: resume issues a follow-up @@ -1079,10 +1090,10 @@ async fn dropping_superseded_session_does_not_unregister_its_replacement() { // The superseded handle goes away; the live session must survive. drop(first_session); - assert_eq!( - client.registered_session_ids_for_test(), - vec![session_id.clone()], - "dropping a superseded Session unregistered its replacement" + assert_only_registration( + &client, + &session_id, + "dropping a superseded Session unregistered its replacement", ); server @@ -1223,7 +1234,7 @@ async fn deferred_create_cancelled_after_callback_registered_is_cleaned_up() { let session = timeout(TIMEOUT, retry).await.unwrap().unwrap().unwrap(); assert_eq!(session.id().as_str(), "registered-then-cancelled"); assert_eq!( - client.registered_session_ids_for_test().len(), + client.registered_session_count_for_test(), 1, "retry must hold exactly one registration" );