From cbc98aa3252a509ce9f762a561a2b7e568611e7d Mon Sep 17 00:00:00 2001 From: jmoseley Date: Wed, 12 Aug 2026 07:26:54 -0700 Subject: [PATCH 01/13] [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 | 4 + rust/README.md | 38 +- rust/src/lib.rs | 9 + rust/src/session.rs | 431 +++++++++++++-- rust/src/types.rs | 47 ++ rust/tests/prepared_session_test.rs | 803 ++++++++++++++++++++++++++++ 8 files changed, 1359 insertions(+), 40 deletions(-) create mode 100644 rust/tests/prepared_session_test.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index e9f22a3df..c6d2b6c20 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,26 @@ See [GitHub Releases](https://github.com/github/copilot-sdk/releases) for the fu ## [Unreleased] +### 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 10f111d9f..0f81c0e8d 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 0f18a9b15..e0c60bec7 100644 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -90,6 +90,10 @@ required-features = ["test-support"] name = "protocol_version_test" required-features = ["test-support"] +[[test]] +name = "prepared_session_test" +required-features = ["test-support"] + [build-dependencies] base64 = "0.22" dirs = "5" diff --git a/rust/README.md b/rust/README.md index 29fe67355..c84c66134 100644 --- a/rust/README.md +++ b/rust/README.md @@ -583,6 +583,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: @@ -786,13 +816,19 @@ 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 | File | Description | | ----------------- | -------------------------------------------------------------------------------------------------------------------------- | | `lib.rs` | `Client`, `ClientOptions`, `CliProgram`, `Transport`, `Error` | -| `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 cafa3c596..d7c7d770b 100644 --- a/rust/src/lib.rs +++ b/rust/src/lib.rs @@ -2209,6 +2209,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 c6c806b1c..5c4e34fbc 100644 --- a/rust/src/session.rs +++ b/rust/src/session.rs @@ -47,6 +47,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 @@ -104,25 +129,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; } @@ -135,7 +206,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); + } } } } @@ -795,6 +868,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, @@ -816,7 +977,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 @@ -949,8 +1151,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 @@ -986,45 +1186,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(), @@ -1036,6 +1240,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(), @@ -1062,8 +1267,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!( @@ -1071,6 +1279,7 @@ impl Client { session_id = %session_id, "Client::create_session complete" ); + registration.disarm(); let session = Session { id: session_id, cwd: self.cwd().clone(), @@ -1106,7 +1315,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() { @@ -1223,8 +1437,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(), @@ -1286,10 +1498,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 @@ -1357,6 +1571,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 d3c4faa16..0d4c12ac8 100644 --- a/rust/src/types.rs +++ b/rust/src/types.rs @@ -2216,6 +2216,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 { @@ -2340,6 +2357,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() } } @@ -2431,6 +2449,7 @@ impl Default for SessionConfig { enable_experimental_mode: None, coauthor_enabled: None, manage_schedule_enabled: None, + event_buffer_capacity: None, } } } @@ -3183,6 +3202,18 @@ impl SessionConfig { self } + /// 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 + } + /// Inject ExP assignment ("flight") data for this session, in the same /// JSON shape the Copilot CLI fetches from the experimentation service /// (`CopilotExpAssignmentResponse`). The runtime feeds it into the same @@ -3464,6 +3495,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 { @@ -3586,6 +3619,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() } } @@ -3821,6 +3855,7 @@ impl ResumeSessionConfig { enable_experimental_mode: None, coauthor_enabled: None, manage_schedule_enabled: None, + event_buffer_capacity: None, } } @@ -4393,6 +4428,18 @@ impl ResumeSessionConfig { self } + /// 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 + } + /// Inject ExP assignment ("flight") data on resume. See /// [`SessionConfig::with_exp_assignments`]. Re-supply the assignments on /// resume so the runtime re-applies them after a CLI process restart. diff --git a/rust/tests/prepared_session_test.rs b/rust/tests/prepared_session_test.rs new file mode 100644 index 000000000..694f16120 --- /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 bb8ca4327d82124840e52b392f18efb194d8e341 Mon Sep 17 00:00:00 2001 From: jmoseley Date: Wed, 12 Aug 2026 08:32:52 -0700 Subject: [PATCH 02/13] [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 | 404 ++++++++++++++++++++++------ rust/tests/prepared_session_test.rs | 317 +++++++++++++++++++++- 7 files changed, 773 insertions(+), 96 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index c6d2b6c20..cce83113c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -21,7 +21,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 0f81c0e8d..e0eae797f 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 c84c66134..18e03b72e 100644 --- a/rust/README.md +++ b/rust/README.md @@ -605,7 +605,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 d7c7d770b..b965e6085 100644 --- a/rust/src/lib.rs +++ b/rust/src/lib.rs @@ -1976,15 +1976,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, @@ -1994,9 +1997,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() } /// Returns the protocol version negotiated with the CLI server, if any. @@ -2215,7 +2237,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 adc192382..4aee33180 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 @@ -226,3 +284,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 5c4e34fbc..3664854e2 100644 --- a/rust/src/session.rs +++ b/rust/src/session.rs @@ -134,66 +134,206 @@ struct PendingSessionRegistration { disarmed: bool, } -/// Which session ID a [`PendingSessionRegistration`] should unregister on -/// cleanup. +/// Which registration a [`PendingSessionRegistration`] owns and may remove +/// 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. +/// Removal is always by *identity*, never by session ID alone: a caller can +/// abort a startup and immediately retry with the same pinned ID, so a +/// stale guard must not unregister the retry that replaced it. 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>>), + /// The registration was made before the RPC and its identity is known. + Known { + id: SessionId, + token: crate::router::RegistrationToken, + }, + /// `session.create` where registration happens (or has yet to happen) + /// inside the inline response callback. The shared slot arbitrates + /// between the callback and this guard. + Deferred(DeferredRegistrationSlot), +} + +/// State of a `session.create` registration that the inline response +/// callback owns until the startup path claims it. +/// +/// The callback runs on the JSON-RPC read task, which removes the pending +/// response entry *before* invoking the callback. A startup future dropped +/// in that window would otherwise see nothing to clean up and the callback +/// would then register a session nobody owns. The state machine closes that +/// window: every transition happens under one lock, so the callback and the +/// guard always agree on who owns the registration. +enum DeferredRegistration { + /// No registration exists yet. The callback may still create one. + Pending, + /// A registration exists and this slot owns it. + Registered { + id: SessionId, + channels: crate::router::SessionChannels, + token: crate::router::RegistrationToken, + }, + /// Startup was cancelled or failed. The callback must not register. + Cancelled, + /// The startup path took ownership; the guard tracks it as + /// [`PendingSessionId::Known`] from here on. + Claimed, +} + +/// Handle to a [`DeferredRegistration`], shared between the inline response +/// callback and the startup cancellation guard. +#[derive(Clone)] +struct DeferredRegistrationSlot(Arc>); + +impl DeferredRegistrationSlot { + /// Slot for a session whose ID the server assigns: nothing is + /// registered until the response arrives. + fn pending() -> Self { + Self(Arc::new(ParkingLotMutex::new( + DeferredRegistration::Pending, + ))) + } + + /// Slot for a session registered up front, before the RPC was issued. + fn registered( + id: SessionId, + channels: crate::router::SessionChannels, + token: crate::router::RegistrationToken, + ) -> Self { + Self(Arc::new(ParkingLotMutex::new( + DeferredRegistration::Registered { + id, + channels, + token, + }, + ))) + } + + /// Register `id` on behalf of the inline response callback. + /// + /// Registration happens *under the slot lock* so it is atomic with + /// publishing the result: a guard running concurrently either wins and + /// marks the slot cancelled (in which case nothing is registered at + /// all), or loses and finds a `Registered` slot to clean up. Never both. + fn register(&self, client: &Client, id: SessionId) { + let mut state = self.0.lock(); + if matches!(*state, DeferredRegistration::Pending) { + let registration = client.register_session(&id); + *state = DeferredRegistration::Registered { + id, + channels: registration.channels, + token: registration.token, + }; + } + // Cancelled: startup is gone, so deliberately register nothing — + // there is no owner left to tear it down. Registered/Claimed: + // already resolved; a second registration would orphan the first. + } + + /// Take ownership of the registration on the successful startup path. + fn claim( + &self, + ) -> Option<( + SessionId, + crate::router::SessionChannels, + crate::router::RegistrationToken, + )> { + let mut state = self.0.lock(); + match std::mem::replace(&mut *state, DeferredRegistration::Claimed) { + DeferredRegistration::Registered { + id, + channels, + token, + } => Some((id, channels, token)), + other => { + *state = other; + None + } + } + } + + /// Cancel the slot on behalf of a startup guard. + /// + /// Returns the registration to remove, if one exists. After this call + /// the callback will not register anything. + fn cancel(&self) -> Option<(SessionId, crate::router::RegistrationToken)> { + let mut state = self.0.lock(); + match std::mem::replace(&mut *state, DeferredRegistration::Cancelled) { + DeferredRegistration::Registered { + id, + channels, + token, + } => { + drop(channels); + Some((id, token)) + } + DeferredRegistration::Pending | DeferredRegistration::Cancelled => None, + // The startup path owns it now; leave it alone. + DeferredRegistration::Claimed => { + *state = DeferredRegistration::Claimed; + None + } + } + } } 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 { + id: session_id, + token, + }, shutdown, disarmed: false, } } - /// Guard for a registration whose session ID is assigned by the server. + /// Guard for a `session.create` registration arbitrated through a + /// shared slot. fn deferred( client: Client, - stash: Arc>>, + slot: DeferredRegistrationSlot, shutdown: CancellationToken, ) -> Self { Self { client, - session_id: PendingSessionId::Deferred(stash), + session_id: PendingSessionId::Deferred(slot), 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 registration. Used by + /// `session.create` once the slot has been claimed by the startup path. + fn resolve_to(&mut self, session_id: SessionId, token: crate::router::RegistrationToken) { + self.session_id = PendingSessionId::Known { + id: session_id, + token, + }; } - /// 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); + /// Remove the registration this guard owns, if it still owns one. + fn release(&mut self) { + match &self.session_id { + PendingSessionId::Known { id, token } => { + self.client.unregister_session_owned(id, *token); + } + PendingSessionId::Deferred(slot) => { + if let Some((id, token)) = slot.cancel() { + self.client.unregister_session_owned(&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); - } + self.release(); self.disarmed = true; } @@ -206,9 +346,7 @@ impl Drop for PendingSessionRegistration { fn drop(&mut self) { if !self.disarmed { self.shutdown.cancel(); - if let Some(id) = self.registered_id() { - self.client.unregister_session(&id); - } + self.release(); } } } @@ -263,6 +401,10 @@ pub struct Session { open_canvases: Arc>>, /// Broadcast channel for runtime event subscribers — see [`Session::subscribe`]. event_tx: tokio::sync::broadcast::Sender, + /// Identity of this session's router registration. Unregistering is a + /// compare-and-remove against this token so a superseded handle for a + /// reused session ID cannot evict the live registration. + registration_token: crate::router::RegistrationToken, } impl Session { @@ -649,7 +791,8 @@ impl Session { ) .await?; self.stop_event_loop().await; - self.client.unregister_session(&self.id); + self.client + .unregister_session_owned(&self.id, self.registration_token); Ok(()) } @@ -726,7 +869,8 @@ 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.client + .unregister_session_owned(&self.id, self.registration_token); } } @@ -1157,19 +1301,28 @@ impl Client { // the session synchronously the instant the response arrives. // For non-cloud sessions, register up-front so the CLI can issue // session-scoped requests during session.create processing. - let inline_stash: Arc< - ParkingLotMutex>, - > = Arc::new(ParkingLotMutex::new(None)); + // Either way the registration lives in a shared slot that arbitrates + // between the callback, this startup path, and the cancellation + // guard below. + let slot = match local_session_id { + Some(ref sid) => { + let registration = self.register_session(sid); + DeferredRegistrationSlot::registered( + sid.clone(), + registration.channels, + registration.token, + ) + } + None => DeferredRegistrationSlot::pending(), + }; - let inline_callback: Option = if let Some(ref sid) = - local_session_id + let inline_callback: Option = if local_session_id + .is_some() { - let channels = self.register_session(sid); - *inline_stash.lock() = Some((sid.clone(), channels)); None } else { let client = self.clone(); - let stash = inline_stash.clone(); + let slot = slot.clone(); let expected = caller_session_id.clone(); Some(Box::new(move |response| { let result = response.result.as_ref().ok_or_else(|| { @@ -1186,35 +1339,22 @@ 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); - *stashed = Some((parsed.session_id, channels)); + // Registers only if the slot is still `Pending`. The read + // task removes the pending-response entry before calling + // this, so a startup future dropped in that window has + // already marked the slot `Cancelled` and nothing is + // registered for a caller that no longer exists. + slot.register(&client, parsed.session_id); Ok(()) })) }; // 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(), - ), - }; + // token and removes the registration this startup owns — by + // identity, so a same-ID retry started in the meantime survives. + let mut registration = + PendingSessionRegistration::deferred(self.clone(), slot.clone(), shutdown.clone()); let rpc_start = Instant::now(); let result = self @@ -1236,11 +1376,10 @@ impl Client { .into()); } - let (session_id, channels) = inline_stash - .lock() - .take() - .expect("session registration must have populated stash on success"); - registration.resolve_to(session_id.clone()); + let (session_id, channels, registration_token) = slot + .claim() + .expect("session registration must have populated the slot on success"); + registration.resolve_to(session_id.clone(), registration_token); let event_loop = spawn_event_loop( session_id.clone(), self.clone(), @@ -1292,6 +1431,7 @@ impl Client { capabilities, open_canvases, event_tx, + registration_token, }; apply_mode_post_create_patch( &session, @@ -1434,7 +1574,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 session_registration = self.register_session(&session_id); + let registration_token = session_registration.token; + let channels = session_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( @@ -1454,8 +1596,12 @@ impl Client { event_tx.clone(), shutdown.clone(), ); - let mut registration = - PendingSessionRegistration::new(self.clone(), session_id.clone(), 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, @@ -1557,6 +1703,7 @@ impl Client { capabilities, open_canvases, event_tx, + registration_token, }; apply_mode_post_create_patch( &session, @@ -1593,9 +1740,11 @@ impl Client { /// * 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. +/// session ID succeeds — cleanup removes only the exact registration that +/// startup owned, so a retry started before the abandoned attempt has +/// finished unwinding is never evicted by it. 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 @@ -2916,8 +3065,107 @@ fn inject_transform_sections_resume( mod tests { use serde_json::json; - use super::{has_managed_settings, notification_permission_payload, permission_request_data}; + use super::{ + DeferredRegistrationSlot, has_managed_settings, notification_permission_payload, + permission_request_data, + }; use crate::handler::PermissionResult; + use crate::types::SessionId; + + fn test_client() -> crate::Client { + let (client_write, _server_read) = tokio::io::duplex(1024); + let (_server_write, client_read) = tokio::io::duplex(1024); + crate::Client::from_streams(client_read, client_write, std::env::temp_dir()) + .expect("from_streams") + } + + /// The window the inline response callback runs in: the JSON-RPC read + /// task removes the pending-response entry *before* invoking the + /// callback, so a startup future dropped in between finds nothing to + /// clean up. Driving the two sides in that exact order asserts the + /// callback declines to register for a caller that is already gone. + #[tokio::test] + async fn deferred_slot_cancelled_before_callback_registers_nothing() { + let client = test_client(); + let slot = DeferredRegistrationSlot::pending(); + + // Gate: the guard lands first, while the slot is still `Pending`. + assert!(slot.cancel().is_none(), "nothing was registered yet"); + + // The inline callback now runs with a server-assigned ID. + slot.register(&client, SessionId::new("server-assigned")); + + assert!( + client.registered_session_ids().is_empty(), + "a cancelled startup left a registration behind" + ); + assert!( + slot.claim().is_none(), + "a cancelled slot must stay unclaimable" + ); + client.force_stop(); + } + + /// The other interleaving: the callback registers first, so the guard + /// finds the registration and owns its removal. + #[tokio::test] + async fn deferred_slot_registered_before_cancel_is_cleaned_up() { + let client = test_client(); + let slot = DeferredRegistrationSlot::pending(); + let id = SessionId::new("server-assigned"); + + slot.register(&client, id.clone()); + assert_eq!(client.registered_session_ids(), vec![id.clone()]); + + let (cancelled_id, token) = slot.cancel().expect("guard must own the registration"); + assert_eq!(cancelled_id, id); + client.unregister_session_owned(&cancelled_id, token); + assert!(client.registered_session_ids().is_empty()); + client.force_stop(); + } + + /// Once the startup path claims the registration, the guard tracks it + /// by identity instead and must not tear it down through the slot. + #[tokio::test] + async fn claimed_slot_is_not_cancelled_by_a_later_guard_drop() { + let client = test_client(); + let slot = DeferredRegistrationSlot::pending(); + let id = SessionId::new("server-assigned"); + + slot.register(&client, id.clone()); + let (claimed_id, _channels, _token) = slot.claim().expect("startup path claims once"); + assert_eq!(claimed_id, id); + + assert!( + slot.cancel().is_none(), + "a claimed slot must not be cancellable" + ); + assert!(slot.claim().is_none(), "a slot can only be claimed once"); + assert_eq!(client.registered_session_ids(), vec![id]); + client.force_stop(); + } + + /// A duplicate callback invocation must not orphan the first + /// registration by silently replacing it. + #[tokio::test] + async fn deferred_slot_registers_at_most_once() { + let client = test_client(); + let slot = DeferredRegistrationSlot::pending(); + let id = SessionId::new("server-assigned"); + + slot.register(&client, id.clone()); + let first_token = match &*slot.0.lock() { + super::DeferredRegistration::Registered { token, .. } => *token, + _ => panic!("expected a registered slot"), + }; + slot.register(&client, id.clone()); + let second_token = match &*slot.0.lock() { + super::DeferredRegistration::Registered { token, .. } => *token, + _ => panic!("expected a registered slot"), + }; + assert_eq!(first_token, second_token, "slot re-registered the session"); + client.force_stop(); + } #[test] fn direct_injection_enables_managed_safeguards() { diff --git a/rust/tests/prepared_session_test.rs b/rust/tests/prepared_session_test.rs index 694f16120..fa2f9ae04 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 a7471a051699b7a9fbeaede2c5339b3c72d8f3a1 Mon Sep 17 00:00:00 2001 From: jmoseley Date: Wed, 12 Aug 2026 08:52:04 -0700 Subject: [PATCH 03/13] [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 b965e6085..1b616dae2 100644 --- a/rust/src/lib.rs +++ b/rust/src/lib.rs @@ -2016,7 +2016,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 69da770371b7463ec23420b7ef2e9806081ac50c Mon Sep 17 00:00:00 2001 From: jmoseley Date: Wed, 12 Aug 2026 15:02:12 -0700 Subject: [PATCH 04/13] [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 cce83113c..389d0c210 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,7 +9,7 @@ See [GitHub Releases](https://github.com/github/copilot-sdk/releases) for the fu ### 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( @@ -21,6 +21,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 3664854e2..7f295643d 100644 --- a/rust/src/session.rs +++ b/rust/src/session.rs @@ -1021,14 +1021,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 de2ef2597db401cd558437295df610d9f3d95693 Mon Sep 17 00:00:00 2001 From: jmoseley Date: Wed, 12 Aug 2026 15:03:02 -0700 Subject: [PATCH 05/13] [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 e0eae797f..17acdc16d 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 d8b48c79622db583963c2ee8c40a9521f0b59d77 Mon Sep 17 00:00:00 2001 From: jmoseley Date: Wed, 12 Aug 2026 15:03:12 -0700 Subject: [PATCH 06/13] [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 fa2f9ae04..7225b6c31 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 dd4e56f4c5ba993db49f7931efedaf4594f88b80 Mon Sep 17 00:00:00 2001 From: jmoseley Date: Wed, 12 Aug 2026 15:53:24 -0700 Subject: [PATCH 07/13] [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 7225b6c31..19ee46551 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 c98cbee6fa83b114b14ca6b62d14b126e5d6fd6e Mon Sep 17 00:00:00 2001 From: jmoseley Date: Wed, 12 Aug 2026 15:53:33 -0700 Subject: [PATCH 08/13] [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 389d0c210..7425ec075 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -25,7 +25,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 17acdc16d..0292f98e0 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 7f295643d..bc7a2de34 100644 --- a/rust/src/session.rs +++ b/rust/src/session.rs @@ -1821,9 +1821,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 e02ccab343d0c7902387ddeed9d2a840829720fa Mon Sep 17 00:00:00 2001 From: jmoseley Date: Wed, 12 Aug 2026 17:34:45 -0700 Subject: [PATCH 09/13] [Rust] Own CLI process termination through to the OS reaper MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `force_stop` sent a kill and returned. Nothing waited on the child, so on Unix it stayed a zombie for as long as the parent lived — fine for a short-lived consumer, but an embedded host is long-lived and accumulates them every time a shutdown times out. The deeper problem was ownership. `stop()` took the child out of its slot and then awaited the kill inline, so an outer timeout cancelling it in that window dropped the handle with the future: the process was neither reaped nor recoverable, and a following `force_stop` found nothing to do. A plain async sibling of `force_stop` would not have fixed that. Introduce `ChildLifecycle`, which owns the child rather than lending it to whoever asks. The kill is delivered synchronously under the slot lock — `start_kill` needs no reactor, so signalling never depends on a task being polled — and the claimed child is then reaped on a dedicated thread with its own runtime. A `tokio::spawn`ed reaper would be cancelled when its runtime drops, which for an embedded host is precisely when termination matters; a thread is independent of every caller runtime, so a handle can be awaited even from a different one. The `enable_all()` on that runtime is load-bearing: once `try_wait` misses, readiness comes from the signal driver. The outcome is published on a `watch` channel, so concurrent and repeat callers observe the same terminal result instead of racing for the handle, and `stop()` claims through the same path — cancelling it can no longer strand the process. Public API (additive; `force_stop` keeps its signature and semantics): Client::force_stop_and_wait() -> Result> Client::start_force_stop() -> ForcedShutdown ForcedShutdown::wait() / ::pid() `Option` is `None` only when the client never spawned a child (stream-backed and in-process transports); failure is always `Err`. `ReapGuard` publishes a definitive error if the reaper cannot run at all, so no waiter is left pending — a guarantee about the state machine, not a time bound, since the underlying wait is unbounded by nature. 15 tests, three of them mutation-verified: reverting the detached ownership hangs the cancel-after-claim recovery; moving the kill back into the reaper leaves the process in state S; dropping `enable_all` hangs the still-running-child reap. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 0f9d5ac7-9999-4f37-82b3-a5533bfbc1f6 --- rust/src/child.rs | 677 ++++++++++++++++++++++++++++++++++++++++++++++ rust/src/lib.rs | 329 ++++++++++++++++++---- 2 files changed, 956 insertions(+), 50 deletions(-) create mode 100644 rust/src/child.rs diff --git a/rust/src/child.rs b/rust/src/child.rs new file mode 100644 index 000000000..97691dba5 --- /dev/null +++ b/rust/src/child.rs @@ -0,0 +1,677 @@ +//! Lifecycle of the CLI child process: termination that is guaranteed to +//! reach the OS reaper. +//! +//! Terminating a child is two steps — deliver the signal, then wait for the +//! kernel to release the process entry — and the second step is the one +//! that is easy to lose. A caller that owns the [`Child`] across an +//! `.await` loses it if its future is cancelled: the handle drops, no one +//! waits, and on Unix the process stays a zombie for as long as the parent +//! lives. That is fatal for an embedded host, which is long-lived by +//! definition and shuts sessions down under a timeout. +//! +//! [`ChildLifecycle`] closes that hole by making the *lifecycle* own the +//! reap rather than the caller. The kill is delivered synchronously, so +//! signalling never depends on a task being polled; the claimed child is +//! then reaped on a dedicated thread with its own runtime, so neither a +//! cancelled future nor a caller's runtime being torn down can strand it. +//! Callers observe the outcome through a [`watch`] channel, which means +//! concurrent and repeat callers all await the same terminal result +//! instead of racing for the handle. + +use std::process::ExitStatus; +use std::sync::Arc; + +use parking_lot::Mutex; +use tokio::process::Child; +use tokio::sync::watch; +use tracing::{info, warn}; + +use crate::{Error, ErrorKind}; + +/// Terminal state of the CLI child process. +/// +/// Every state other than `Pending` is definitive, so a waiter always has +/// something to return and never blocks forever. +#[derive(Clone, Debug)] +enum ReapState { + /// No reaper has finished yet. + Pending, + /// The child was waited on and the OS released it. `None` means this + /// client never owned a child (stream-backed or in-process transport). + Reaped(Option), + /// Termination could not be confirmed. The process may still exist. + Failed { + kind: std::io::ErrorKind, + message: String, + }, +} + +impl ReapState { + fn from_wait(result: std::io::Result) -> Self { + match result { + Ok(status) => Self::Reaped(Some(status)), + Err(error) => Self::Failed { + kind: error.kind(), + message: error.to_string(), + }, + } + } + + /// Resolve a definitive state into the public result, or `None` while + /// the outcome is still pending. + fn resolve(&self) -> Option, Error>> { + match self { + Self::Pending => None, + Self::Reaped(status) => Some(Ok(*status)), + Self::Failed { kind, message } => Some(Err(Error::new( + ErrorKind::Io, + std::io::Error::new(*kind, message.clone()), + ))), + } + } +} + +/// Owns the CLI child process and guarantees it is reaped exactly once. +pub(crate) struct ChildLifecycle { + /// Holds the child until a caller claims it for termination. Empty + /// from that moment on: the detached reaper owns it instead. + child: Mutex>, + /// Last known process ID. Fixed at construction, so diagnostics can + /// still name the process after the child has been claimed. + pid: Option, + /// Terminal outcome, shared by every waiter. + state: Arc>, +} + +impl ChildLifecycle { + pub(crate) fn new(child: Option) -> Self { + let pid = child.as_ref().and_then(Child::id); + // A client with no child has nothing to reap, so its outcome is + // already final and waiters resolve immediately. + let initial = match child { + Some(_) => ReapState::Pending, + None => ReapState::Reaped(None), + }; + let (state, _) = watch::channel(initial); + Self { + child: Mutex::new(child), + pid, + state: Arc::new(state), + } + } + + /// Process ID of the live child, or `None` once it has been claimed + /// for termination. + pub(crate) fn pid(&self) -> Option { + self.child.lock().as_ref().and_then(Child::id) + } + + /// Whether a live child is still held. False once termination starts. + pub(crate) fn has_child(&self) -> bool { + self.child.lock().is_some() + } + + /// Begin terminating the child and return a handle to its completion. + /// + /// Synchronous and idempotent. The first caller signals the child and + /// hands it to a dedicated reaper thread that waits to completion; + /// later callers — and callers racing on another thread — get a handle + /// to that same reaper's outcome. Dropping the returned handle never + /// cancels the termination, which is the whole point: the reap + /// outlives whatever future asked for it. + pub(crate) fn begin_termination(&self) -> ForcedShutdown { + // Claim, signal, and reset under one lock. + // + // The kill is delivered here rather than inside the reaper because + // `start_kill` needs no reactor: signalling must not depend on a + // task ever being polled, or a caller tearing its runtime down + // would leave the CLI process alive. Holding one guard across the + // take and the reset also stops a concurrent caller subscribing in + // between and binding to a previous attempt's outcome. + let (claimed, handle) = { + let mut slot = self.child.lock(); + match slot.take() { + Some(mut child) => { + let pid = child.id(); + if let Err(error) = child.start_kill() { + // Usually just an already-exited child; the wait + // below still reports its real status. + warn!(pid = ?pid, error = %error, "kill signal not delivered to CLI process"); + } + self.state.send_replace(ReapState::Pending); + (Some(child), self.handle()) + } + // Either termination already started (the running reaper + // will publish) or there was never a child (already + // `Reaped(None)`). + None => (None, self.handle()), + } + }; + let Some(child) = claimed else { + return handle; + }; + + let pid = child.id(); + info!(pid = ?pid, "terminating CLI process"); + // Reap on a dedicated thread with its own runtime rather than on a + // caller's. A `tokio::spawn`ed task is cancelled when its runtime + // drops, which for an embedded host is precisely the moment + // termination matters — the host tears its runtime down as it + // shuts down. Owning a thread makes the reap independent of every + // caller runtime, so a handle really can be awaited anywhere. + // Termination happens once per client, so the thread is not a hot + // path. + let guard = ReapGuard { + state: Arc::clone(&self.state), + completed: false, + }; + if let Err(error) = std::thread::Builder::new() + .name("copilot-cli-reaper".to_string()) + .spawn(move || { + // Moved in, so it drops with the thread even if the body + // panics before publishing. + let mut guard = guard; + // `enable_all` is load-bearing on Unix, not boilerplate: + // once `try_wait` misses, readiness for `Child::wait` + // comes from the signal driver. Without it, a child that + // outlives the reaper's first poll — a process wedged in + // uninterruptible I/O, exactly what `force_stop` exists + // for — would never be reaped. + match tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + { + Ok(runtime) => { + let outcome = runtime.block_on(reap(child)); + guard.completed = true; + guard.state.send_replace(outcome); + } + Err(error) => { + warn!(error = %error, "could not build a runtime to reap the CLI process"); + } + } + }) + { + // The guard was moved into the closure and dropped with it, so + // waiters have already been resolved with a definitive error. + warn!(pid = ?pid, error = %error, "could not start the CLI reaper thread"); + } + handle + } + + fn handle(&self) -> ForcedShutdown { + ForcedShutdown { + pid: self.pid, + state: self.state.subscribe(), + } + } + + /// Best-effort signal used on `Drop`, where there is no opportunity to + /// await anything. Leaves the child in place so a reaper that is + /// already running keeps ownership. + pub(crate) fn signal_on_drop(&self) { + if let Some(child) = self.child.lock().as_mut() { + let pid = child.id(); + match child.start_kill() { + Ok(()) => info!(pid = ?pid, "kill signal sent for CLI process on drop"), + Err(error) => { + warn!(pid = ?pid, error = %error, "failed to kill CLI process on drop"); + } + } + } + } +} + +/// Publishes a terminal state if the reaper never publishes one of its +/// own — its thread fails to start, its runtime cannot be built, or it +/// panics — so no waiter is left pending. +struct ReapGuard { + state: Arc>, + completed: bool, +} + +impl Drop for ReapGuard { + fn drop(&mut self) { + if !self.completed { + self.state.send_replace(ReapState::Failed { + kind: std::io::ErrorKind::Interrupted, + message: "the CLI child was signalled but its reaper could not run to completion \ + (the reaper thread or its runtime failed to start, or it panicked); the \ + OS may not have released the process" + .to_string(), + }); + } + } +} + +/// Wait for the OS to release a child that has already been signalled. +async fn reap(mut child: Child) -> ReapState { + let pid = child.id(); + let state = ReapState::from_wait(child.wait().await); + info!(pid = ?pid, outcome = ?state, "CLI process reaped"); + state +} +/// Owned handle to a CLI process being terminated. +/// +/// Returned by [`Client::start_force_stop`](crate::Client::start_force_stop). +/// Await [`wait`](Self::wait) to observe termination completing — that is, +/// the child killed *and* reaped, with no zombie left behind. +/// +/// The handle borrows nothing from the client and reaping is owned by a +/// dedicated thread, so it can be moved to another task — or another +/// runtime — and awaited there, including after the runtime that started +/// termination has been torn down. Dropping it without awaiting does not +/// cancel the termination; it only gives up watching. Any number of +/// handles may exist for one child, and they all resolve to the same +/// outcome. +#[derive(Debug)] +#[must_use = "termination continues either way; await this handle to observe it completing"] +pub struct ForcedShutdown { + pid: Option, + state: watch::Receiver, +} + +impl ForcedShutdown { + /// Process ID of the terminating child, if this client spawned one. + /// + /// Reports which process is being terminated; it carries no outcome. + pub fn pid(&self) -> Option { + self.pid + } + + /// Wait until the child has been killed and reaped. + /// + /// Resolves to the child's [`ExitStatus`], or `None` when the client + /// never owned a child (stream-backed and in-process transports). + /// `None` therefore means "nothing to terminate", never "termination + /// failed" — a failure is always an `Err`. The status itself is the + /// child's, and a killed process reports an unsuccessful one; it says + /// nothing about whether teardown worked. + /// + /// Repeat and concurrent waiters all observe the same outcome, and a + /// handle may be awaited on any runtime — reaping is owned by a + /// dedicated thread, not by the runtime that started it. + /// + /// # Cancel safety + /// + /// Cancel-safe, and cancelling gives nothing up: termination is owned + /// by the SDK, not by this future, so dropping it leaves the reap + /// running. Acquire another handle to await the same completion. + /// + /// # Errors + /// + /// Returns [`ErrorKind::Io`] if waiting on the child failed, or if the + /// reaper could not run to completion — its thread failed to start, + /// its runtime could not be built, or it panicked. The child has + /// always been signalled by then; only confirmation is lost. + /// + /// No outcome is ever lost: every failure resolves to an `Err` rather + /// than a pending wait. That is a guarantee about the state machine, + /// not about elapsed time — the underlying wait is unbounded, so a + /// process the kernel will not release (uninterruptible I/O) parks + /// this future for as long as that lasts. Wrap it in + /// [`tokio::time::timeout`] if the caller needs a bounded shutdown. + pub async fn wait(mut self) -> Result, Error> { + loop { + if let Some(outcome) = self.state.borrow_and_update().resolve() { + return outcome; + } + if self.state.changed().await.is_err() { + // The client and the reaper are both gone without + // publishing. Report it rather than wait on a sender that + // no longer exists. + return Err(Error::new( + ErrorKind::Io, + std::io::Error::other( + "the client was dropped before CLI process termination completed", + ), + )); + } + } + } +} + +#[cfg(test)] +mod tests { + use std::process::Stdio; + use std::time::Duration; + + use tokio::process::Command; + use tokio::time::timeout; + + use super::*; + + /// Failure backstop. Every wait in these tests must resolve well + /// inside it; exceeding it means something hung. + const TIMEOUT: Duration = Duration::from_secs(10); + + /// A child that stays alive until it is killed. + fn spawn_sleeper() -> Child { + let mut command = if cfg!(windows) { + let mut command = Command::new("cmd"); + command.args(["/C", "ping -n 120 127.0.0.1 > nul"]); + command + } else { + let mut command = Command::new("sleep"); + command.arg("120"); + command + }; + command + .stdin(Stdio::null()) + .stdout(Stdio::null()) + .stderr(Stdio::null()) + .spawn() + .expect("spawn sleeper child") + } + + /// A child that exits on its own, immediately. + fn spawn_exiting() -> Child { + let mut command = if cfg!(windows) { + let mut command = Command::new("cmd"); + command.args(["/C", "exit 0"]); + command + } else { + Command::new("true") + }; + command + .stdin(Stdio::null()) + .stdout(Stdio::null()) + .stderr(Stdio::null()) + .spawn() + .expect("spawn exiting child") + } + + /// The regression this type exists for. A caller that claims the child + /// and then disappears — an outer timeout firing between the claim and + /// the reap — must not strand the process. The recovery caller has to + /// find the termination still owned and running, and observe it + /// completing. + /// + /// Deterministic by construction: the first handle is dropped + /// synchronously, before it is ever polled, which is the exact window + /// that used to lose the child. + #[tokio::test] + async fn handle_dropped_before_polling_still_reaps_and_stays_observable() { + let lifecycle = ChildLifecycle::new(Some(spawn_sleeper())); + let pid = lifecycle.pid().expect("sleeper should report a pid"); + + let abandoned = lifecycle.begin_termination(); + assert_eq!(abandoned.pid(), Some(pid)); + drop(abandoned); + + let recovery = lifecycle.begin_termination(); + assert_eq!(recovery.pid(), Some(pid), "the handle must not be lost"); + let status = timeout(TIMEOUT, recovery.wait()) + .await + .expect("recovery waiter hung") + .expect("recovery waiter failed") + .expect("a spawned child must report an exit status"); + assert!( + !status.success(), + "a killed child must not report success: {status:?}" + ); + assert!(lifecycle.pid().is_none()); + assert!(!lifecycle.has_child()); + } + + /// Every waiter — those that arrive before the reap finishes and those + /// that arrive long after — resolves to the same terminal status. + #[tokio::test] + async fn concurrent_and_repeat_waiters_observe_one_outcome() { + let lifecycle = ChildLifecycle::new(Some(spawn_sleeper())); + + let waiters: Vec<_> = (0..8) + .map(|_| tokio::spawn(lifecycle.begin_termination().wait())) + .collect(); + + let mut statuses = Vec::new(); + for waiter in waiters { + let status = timeout(TIMEOUT, waiter) + .await + .expect("waiter hung") + .expect("waiter panicked") + .expect("waiter failed"); + statuses.push(status); + } + let first = statuses[0]; + assert!(first.is_some(), "expected an exit status"); + assert!( + statuses.iter().all(|status| *status == first), + "waiters disagreed about the outcome: {statuses:?}" + ); + + // A waiter created after the fact resolves to the same value. + let late = timeout(TIMEOUT, lifecycle.begin_termination().wait()) + .await + .expect("late waiter hung") + .expect("late waiter failed"); + assert_eq!(late, first); + } + + /// A child that has already exited is reaped without hanging, and its + /// real status survives the failed kill. + #[tokio::test] + async fn already_exited_child_is_reaped_without_hanging() { + let mut child = spawn_exiting(); + let expected = child.wait().await.expect("child should exit"); + // Re-wrap a child that is already gone: `wait` is idempotent and + // keeps reporting the same status. + let lifecycle = ChildLifecycle::new(Some(child)); + + let status = timeout(TIMEOUT, lifecycle.begin_termination().wait()) + .await + .expect("waiting on an exited child hung") + .expect("waiting on an exited child failed"); + assert_eq!(status, Some(expected)); + } + + /// Clients with no child of their own (stream-backed, in-process) + /// resolve immediately instead of waiting for a process that does not + /// exist. + #[tokio::test] + async fn childless_lifecycle_resolves_immediately() { + let lifecycle = ChildLifecycle::new(None); + assert!(!lifecycle.has_child()); + assert!(lifecycle.pid().is_none()); + + for _ in 0..3 { + let outcome = timeout(TIMEOUT, lifecycle.begin_termination().wait()) + .await + .expect("childless wait hung") + .expect("childless wait failed"); + assert_eq!(outcome, None); + } + } + + /// Gate: a handle really can be awaited on a *different* runtime than + /// the one that started termination. Runtime A is dropped immediately + /// after the claim; the reap is owned by a dedicated thread, so it + /// completes regardless, and runtime B observes the real exit status + /// rather than an error or a hang. + #[test] + fn termination_survives_its_originating_runtime() { + let runtime_a = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .expect("build runtime A"); + let lifecycle = runtime_a.block_on(async { ChildLifecycle::new(Some(spawn_sleeper())) }); + + let handle = runtime_a.block_on(async { lifecycle.begin_termination() }); + drop(runtime_a); + + let runtime_b = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .expect("build runtime B"); + let status = runtime_b + .block_on(async { timeout(TIMEOUT, handle.wait()).await }) + .expect("waiter hung after its originating runtime was dropped") + .expect("termination must not fail when its originating runtime goes away") + .expect("a spawned child must report an exit status"); + assert!(!status.success(), "a killed child must not report success"); + } + + /// Termination needs no runtime in context at all: the kill is + /// synchronous and the reap owns its own. This is what makes the + /// synchronous `force_stop` usable while a caller is tearing its + /// runtime down. + #[test] + fn termination_without_any_runtime_context_still_kills_and_reaps() { + let setup = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .expect("build runtime"); + let lifecycle = setup.block_on(async { ChildLifecycle::new(Some(spawn_sleeper())) }); + drop(setup); + + // No runtime context here whatsoever. + assert!(tokio::runtime::Handle::try_current().is_err()); + let handle = lifecycle.begin_termination(); + assert!(!lifecycle.has_child(), "the child must have been claimed"); + + let observer = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .expect("build observing runtime"); + let status = observer + .block_on(async { timeout(TIMEOUT, handle.wait()).await }) + .expect("wait hung") + .expect("termination without a runtime must still succeed") + .expect("a spawned child must report an exit status"); + assert!(!status.success()); + } + + /// Gate: dropping every observer — the lifecycle itself and every + /// handle — must not cancel the kill or the reap. Nothing is left + /// running, and nothing keeps the lifecycle alive. + #[cfg(unix)] + #[test] + fn dropping_every_observer_does_not_cancel_the_reap() { + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .expect("build runtime"); + let lifecycle = runtime.block_on(async { ChildLifecycle::new(Some(spawn_sleeper())) }); + let pid = lifecycle.pid().expect("sleeper should report a pid"); + + let handle = lifecycle.begin_termination(); + drop(handle); + drop(lifecycle); + drop(runtime); + + // The reaper thread owns the child; give it a moment to finish. + let deadline = std::time::Instant::now() + std::time::Duration::from_secs(5); + loop { + let output = std::process::Command::new("ps") + .args(["-o", "stat=", "-p", &pid.to_string()]) + .output() + .expect("run ps"); + let state = String::from_utf8_lossy(&output.stdout).trim().to_string(); + if state.is_empty() { + return; + } + assert!( + std::time::Instant::now() < deadline, + "the CLI process was never reaped (ps state {state:?})" + ); + std::thread::sleep(std::time::Duration::from_millis(50)); + } + } + + /// The critical guarantee behind `force_stop`'s synchronous contract: + /// the kill must not depend on a task ever being polled. A runtime + /// torn down immediately after the claim discards the reaper, and the + /// process must still have been signalled — a live orphaned CLI is a + /// worse outcome than the zombie this work set out to remove. + /// + /// Unix-only because it reads the process state directly; the + /// behaviour it guards is platform-independent. + #[cfg(unix)] + #[test] + fn signal_is_delivered_even_when_the_reaper_never_runs() { + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .expect("build runtime"); + let lifecycle = runtime.block_on(async { ChildLifecycle::new(Some(spawn_sleeper())) }); + let pid = lifecycle.pid().expect("sleeper should report a pid"); + + // Claim inside the runtime, then discard the runtime before the + // reaper can be polled. + let _handle = runtime.block_on(async { lifecycle.begin_termination() }); + drop(runtime); + std::thread::sleep(std::time::Duration::from_millis(300)); + + let output = std::process::Command::new("ps") + .args(["-o", "stat=", "-p", &pid.to_string()]) + .output() + .expect("run ps"); + let state = String::from_utf8_lossy(&output.stdout).trim().to_string(); + assert!( + state.is_empty() || state.starts_with('Z'), + "the CLI process was left running instead of signalled (ps state {state:?})" + ); + } + + /// A child that exits on its own after a beat, so a reaper polling it + /// misses on the first `try_wait` and must fall through to the + /// signal-driven path. + fn spawn_slow_exiting() -> Child { + let mut command = if cfg!(windows) { + let mut command = Command::new("cmd"); + command.args(["/C", "ping -n 3 127.0.0.1 > nul"]); + command + } else { + let mut command = Command::new("sleep"); + command.arg("1"); + command + }; + command + .stdin(Stdio::null()) + .stdout(Stdio::null()) + .stderr(Stdio::null()) + .spawn() + .expect("spawn slow-exiting child") + } + + /// Guards the `enable_all()` on the reaper's runtime, which is + /// load-bearing rather than boilerplate. + /// + /// Every other test kills its child before the reaper runs, so the + /// child is already a zombie and the first `try_wait` succeeds — the + /// signal path is never touched, and the suite would pass even without + /// a signal driver. This case waits on a child that is still running + /// at the first poll, on a runtime other than the one that spawned it: + /// the shape of the wedged process `force_stop` exists for. + #[test] + fn reap_completes_for_a_child_still_running_at_the_first_poll() { + let spawner = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .expect("build spawning runtime"); + let child = spawner.block_on(async { spawn_slow_exiting() }); + drop(spawner); + + let reaper = std::thread::spawn(move || { + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .expect("build reaper runtime"); + runtime.block_on(reap(child)) + }); + + match reaper.join().expect("reaper thread panicked") { + ReapState::Reaped(Some(status)) => assert!( + status.success(), + "a child left to exit on its own should report success: {status:?}" + ), + other => panic!("expected a reaped child, got {other:?}"), + } + } + + #[test] + fn forced_shutdown_handle_is_send_and_static() { + fn assert_send_static() {} + assert_send_static::(); + } +} diff --git a/rust/src/lib.rs b/rust/src/lib.rs index 1b616dae2..ad758153b 100644 --- a/rust/src/lib.rs +++ b/rust/src/lib.rs @@ -6,6 +6,7 @@ /// Canvas declarations, provider callbacks, and host-side canvas RPC types. pub mod canvas; mod canvas_dispatch; +mod child; /// Bundled CLI binary extraction and caching. #[cfg(feature = "bundled-cli")] pub(crate) mod embeddedcli; @@ -13,6 +14,7 @@ mod errors; /// In-process FFI transport hosting the runtime cdylib (`Transport::InProcess`). #[cfg(feature = "bundled-in-process")] pub(crate) mod ffi; +pub use child::ForcedShutdown; pub use errors::*; /// Connection-level Copilot request handler — intercept and replace the /// model-layer HTTP and WebSocket traffic the runtime issues for both CAPI and @@ -103,7 +105,7 @@ use tokio::io::{AsyncBufReadExt, AsyncRead, AsyncWrite, BufReader}; use tokio::net::TcpStream; use tokio::process::{Child, Command}; use tokio::sync::{broadcast, mpsc, oneshot}; -use tracing::{Instrument, debug, error, info, warn}; +use tracing::{Instrument, debug, info, warn}; pub use types::*; mod sdk_protocol_version; @@ -987,7 +989,7 @@ impl std::fmt::Debug for Client { } struct ClientInner { - child: parking_lot::Mutex>, + child: child::ChildLifecycle, #[cfg(feature = "bundled-in-process")] /// In-process FFI runtime host, set only for [`Transport::InProcess`]. /// Closing it tears down the native runtime connection. @@ -1550,7 +1552,7 @@ impl Client { let client = Self { inner: Arc::new(ClientInner { - child: parking_lot::Mutex::new(child), + child: child::ChildLifecycle::new(child), #[cfg(feature = "bundled-in-process")] ffi_host: parking_lot::Mutex::new(None), rpc, @@ -2374,8 +2376,11 @@ impl Client { } /// Return the OS process ID of the CLI child process, if one was spawned. + /// + /// Returns `None` once termination has begun — the child is owned by + /// the SDK's reaper from that point on. pub fn pid(&self) -> Option { - self.inner.child.lock().as_ref().and_then(|c| c.id()) + self.inner.child.pid() } /// Cooperatively shut down the client and the CLI child process. @@ -2395,15 +2400,20 @@ impl Client { /// /// # Cancel safety /// - /// **Cancel-unsafe but recoverable.** The body sequentially destroys - /// every registered session (each via [`Client::call`](Self::call), - /// individually cancel-safe) before killing the child. Cancelling - /// `stop()` mid-loop leaves some sessions still in the router map - /// and the child still running. Recovery: call [`force_stop`](Self::force_stop) - /// (sync, kills the child unconditionally and clears router state) - /// or call `stop()` again with a fresh future. The documented - /// `tokio::time::timeout(..., client.stop())` pattern in the example - /// below uses `force_stop` as the fallback for exactly this case. + /// **Cancel-unsafe but recoverable, and the child is never stranded.** + /// The body sequentially destroys every registered session (each via + /// [`Client::call`](Self::call), individually cancel-safe) before + /// terminating the child. Cancelling `stop()` mid-loop leaves some + /// sessions still in the router map and, if the cancellation lands + /// before termination begins, the child still running. + /// + /// Termination itself is not lost either way: the child is claimed + /// synchronously and handed to a reaper the SDK owns, so cancelling + /// this future cannot leave a signalled-but-unreaped process behind. + /// Recovery: call [`force_stop_and_wait`](Self::force_stop_and_wait) + /// to drive teardown to completion and observe the same terminal + /// outcome, [`force_stop`](Self::force_stop) for the synchronous + /// fallback, or `stop()` again with a fresh future. pub async fn stop(&self) -> std::result::Result<(), StopErrors> { let pid = self.pid(); info!(pid = ?pid, "stopping CLI process"); @@ -2432,7 +2442,7 @@ impl Client { self.inner.router.unregister(&session_id); } - let should_shutdown_runtime = self.inner.child.lock().is_some(); + let should_shutdown_runtime = self.inner.child.has_child(); #[cfg(feature = "bundled-in-process")] let should_shutdown_runtime = should_shutdown_runtime || self.inner.ffi_host.lock().is_some(); @@ -2471,25 +2481,20 @@ impl Client { } } - let child = self.inner.child.lock().take(); + // Claim the child *before* awaiting anything else. The claim is + // synchronous and hands ownership to the SDK's reaper, so if this + // future is cancelled from here on, termination still runs to + // completion and a later caller can await the same outcome. + let termination = self.inner.child.begin_termination(); *self.inner.state.lock() = ConnectionState::Disconnected; *self.inner.models_cache.lock() = Arc::new(tokio::sync::OnceCell::new()); - if let Some(mut child) = child { - match child.try_wait() { - Ok(Some(_status)) => {} - Ok(None) => { - // The runtime completes all cleanup before responding to - // runtime.shutdown and then leaves termination to us; it - // deliberately keeps its JSON-RPC server alive to send the - // response and never self-exits. Waiting for a self-exit - // that will never come just wastes time, so terminate the - // child immediately. - if let Err(e) = child.kill().await { - errors.push(e.into()); - } - } - Err(e) => errors.push(e.into()), - } + // The runtime completes all cleanup before responding to + // runtime.shutdown and then leaves termination to us; it + // deliberately keeps its JSON-RPC server alive to send the + // response and never self-exits. Waiting for a self-exit that will + // never come just wastes time, so terminate the child immediately. + if let Err(e) = termination.wait().await { + errors.push(e); } // The runtime.shutdown RPC above already asked the runtime to clean up; @@ -2514,10 +2519,17 @@ impl Client { /// /// Synchronous fallback when [`stop`](Self::stop) is unsuitable — for /// example when the awaiting tokio runtime is shutting down or the - /// process is wedged on I/O. Sends a kill signal without awaiting - /// reaper completion and immediately drops all per-session router - /// state so dependent tasks observe a closed channel rather than a - /// hang. + /// process is wedged on I/O. Sends a kill signal and immediately drops + /// all per-session router state so dependent tasks observe a closed + /// channel rather than a hang. + /// + /// This returns before the child has been reaped. Termination itself + /// is owned by the SDK and continues in the background, but a caller + /// that must *know* the process is gone — an embedded host tearing + /// down under a timeout, where a surviving zombie accumulates for the + /// lifetime of the host — should use + /// [`force_stop_and_wait`](Self::force_stop_and_wait) or + /// [`start_force_stop`](Self::start_force_stop) instead. /// /// # Cancel safety /// @@ -2540,13 +2552,34 @@ impl Client { /// # } /// ``` pub fn force_stop(&self) { + drop(self.start_force_stop()); + } + + /// Forcibly stop the CLI process and return a handle to its + /// termination. + /// + /// Performs the same synchronous teardown as + /// [`force_stop`](Self::force_stop) — kill signal, transport close, + /// router cleanup — and additionally hands back an owned + /// [`ForcedShutdown`] that resolves when the child has been killed + /// *and* reaped. + /// + /// Use this when the awaiting code lives somewhere else: the handle + /// borrows nothing from the client, so it can be moved into another + /// task or onto a runtime that will outlive this one. Dropping it does + /// not cancel termination, and any number of handles may await the + /// same child. + /// + /// # Cancel safety + /// + /// **Synchronous.** Termination is owned by the SDK from the moment + /// this returns, so neither dropping the handle nor cancelling a + /// future awaiting it can strand the process. + #[must_use = "the child is terminating either way; await the handle to observe it finishing"] + pub fn start_force_stop(&self) -> ForcedShutdown { let pid = self.pid(); info!(pid = ?pid, "force-stopping CLI process"); - if let Some(mut child) = self.inner.child.lock().take() - && let Err(e) = child.start_kill() - { - error!(pid = ?pid, error = %e, "failed to send kill signal"); - } + let termination = self.inner.child.begin_termination(); self.inner.rpc.force_close(); #[cfg(feature = "bundled-in-process")] { @@ -2559,6 +2592,55 @@ impl Client { self.inner.router.clear(); *self.inner.state.lock() = ConnectionState::Disconnected; *self.inner.models_cache.lock() = Arc::new(tokio::sync::OnceCell::new()); + termination + } + + /// Forcibly stop the CLI process and wait until the OS has reaped it. + /// + /// The awaited counterpart of [`force_stop`](Self::force_stop): it + /// tears down the client identically, then waits for the child to be + /// killed *and* released by the OS, so no zombie survives the call. + /// Resolves to the child's exit status, or `None` when this client + /// never spawned one (stream-backed and in-process transports). + /// + /// Repeat and concurrent calls are safe: the first begins + /// termination, and every caller observes the same terminal outcome + /// rather than racing for the process handle. + /// + /// # Cancel safety + /// + /// **Cancel-safe, and cancelling costs nothing.** Termination is owned + /// by the SDK rather than by this future, so a timeout that elapses + /// here leaves the reap running; call this again — or await a + /// [`start_force_stop`](Self::start_force_stop) handle — to observe + /// the same completion. + /// + /// # Errors + /// + /// Returns [`ErrorKind::Io`] if waiting on the child failed, or if the + /// reaper could not run to completion. The child is signalled in every + /// case; only confirmation of the reap is lost. No outcome is ever + /// lost either — a failure resolves to an `Err` rather than a pending + /// wait — but the wait itself is unbounded, so wrap this in + /// [`tokio::time::timeout`] if the caller needs a bounded shutdown. + /// + /// # Example + /// + /// ```no_run + /// # async fn example(client: github_copilot_sdk::Client) { + /// // Graceful first; on timeout, force-stop and confirm the process + /// // is really gone before the host finishes shutting down. + /// if tokio::time::timeout( + /// std::time::Duration::from_secs(5), + /// client.stop(), + /// ).await.is_err() + /// { + /// let _ = client.force_stop_and_wait().await; + /// } + /// # } + /// ``` + pub async fn force_stop_and_wait(&self) -> Result> { + self.start_force_stop().wait().await } /// Subscribe to lifecycle events. @@ -2602,14 +2684,7 @@ impl Client { impl Drop for ClientInner { fn drop(&mut self) { - if let Some(ref mut child) = *self.child.lock() { - let pid = child.id(); - if let Err(e) = child.start_kill() { - error!(pid = ?pid, error = %e, "failed to kill CLI process on drop"); - } else { - info!(pid = ?pid, "kill signal sent for CLI process on drop"); - } - } + self.child.signal_on_drop(); #[cfg(feature = "bundled-in-process")] { if let Some(host) = self.ffi_host.lock().take() { @@ -3248,10 +3323,164 @@ mod tests { client.force_stop(); } + /// A child that stays alive until it is killed, for client-level + /// lifecycle tests that must not depend on a CLI being installed. + fn spawn_test_child() -> tokio::process::Child { + let mut command = if cfg!(windows) { + let mut command = tokio::process::Command::new("cmd"); + command.args(["/C", "ping -n 120 127.0.0.1 > nul"]); + command + } else { + let mut command = tokio::process::Command::new("sleep"); + command.arg("120"); + command + }; + command + .stdin(std::process::Stdio::null()) + .stdout(std::process::Stdio::null()) + .stderr(std::process::Stdio::null()) + .spawn() + .expect("spawn test child") + } + + const LIFECYCLE_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(10); + + /// The host-facing guarantee: after `force_stop_and_wait` resolves, + /// the child has been killed *and* reaped — no zombie is left for a + /// long-lived embedded host to accumulate. + #[tokio::test] + async fn force_stop_and_wait_reaps_the_child() { + let client = client_with_child_and_handler(Some(spawn_test_child()), None); + assert!(client.pid().is_some()); + + let status = tokio::time::timeout(LIFECYCLE_TIMEOUT, client.force_stop_and_wait()) + .await + .expect("force_stop_and_wait hung") + .expect("force_stop_and_wait failed") + .expect("a spawned child must report an exit status"); + assert!(!status.success(), "a killed child must not report success"); + assert!(client.pid().is_none()); + assert!(matches!( + *client.inner.state.lock(), + ConnectionState::Disconnected + )); + } + + /// Repeat and concurrent callers converge on the same terminal + /// outcome instead of racing for the process handle, and none of them + /// hangs. + #[tokio::test] + async fn force_stop_and_wait_is_idempotent_and_concurrent_safe() { + let client = client_with_child_and_handler(Some(spawn_test_child()), None); + + let concurrent: Vec<_> = (0..6) + .map(|_| { + let client = client.clone(); + tokio::spawn(async move { client.force_stop_and_wait().await }) + }) + .collect(); + let mut outcomes = Vec::new(); + for task in concurrent { + outcomes.push( + tokio::time::timeout(LIFECYCLE_TIMEOUT, task) + .await + .expect("concurrent waiter hung") + .expect("concurrent waiter panicked") + .expect("concurrent waiter failed"), + ); + } + let first = outcomes[0]; + assert!(first.is_some()); + assert!( + outcomes.iter().all(|outcome| *outcome == first), + "concurrent callers disagreed: {outcomes:?}" + ); + + // And again, long after termination completed. + let repeat = tokio::time::timeout(LIFECYCLE_TIMEOUT, client.force_stop_and_wait()) + .await + .expect("repeat call hung") + .expect("repeat call failed"); + assert_eq!(repeat, first); + } + + /// A `stop()` whose future goes away must never strand the child. + /// + /// On the default current-thread runtime the abort lands before the + /// spawned `stop()` is ever polled, so this covers the + /// cancelled-*before*-the-claim branch: the child is still held, and + /// the recovery call terminates it and reports its status. The + /// cancelled-*after*-the-claim branch — the one the old ownership + /// model lost outright — is covered deterministically by + /// `child::tests::handle_dropped_before_polling_still_reaps_and_stays_observable`. + #[tokio::test] + async fn cancelled_stop_leaves_the_child_recoverable() { + let client = client_with_child_and_handler(Some(spawn_test_child()), None); + + let stopping = tokio::spawn({ + let client = client.clone(); + async move { client.stop().await } + }); + stopping.abort(); + let _ = stopping.await; + + let status = tokio::time::timeout(LIFECYCLE_TIMEOUT, client.force_stop_and_wait()) + .await + .expect("recovery hung") + .expect("recovery failed") + .expect("the child must still be accounted for after a cancelled stop"); + assert!(!status.success()); + assert!(client.pid().is_none()); + } + + /// A handle taken before the client is torn down can be awaited + /// afterwards, and on another task — the shape an embedded host uses + /// when its own shutdown outlives the client. + #[tokio::test] + async fn start_force_stop_handle_is_awaitable_elsewhere() { + let client = client_with_child_and_handler(Some(spawn_test_child()), None); + let handle = client.start_force_stop(); + assert!(handle.pid().is_some()); + + let awaited = tokio::spawn(async move { handle.wait().await }); + let status = tokio::time::timeout(LIFECYCLE_TIMEOUT, awaited) + .await + .expect("detached waiter hung") + .expect("detached waiter panicked") + .expect("detached waiter failed"); + assert!(status.is_some()); + } + + /// Clients with no child of their own resolve immediately rather than + /// waiting on a process that does not exist. + #[tokio::test] + async fn force_stop_and_wait_without_a_child_resolves_immediately() { + let (client_write, _server_read) = tokio::io::duplex(8192); + let (_server_write, client_read) = tokio::io::duplex(8192); + let client = Client::from_streams(client_read, client_write, std::env::temp_dir()).unwrap(); + + for _ in 0..3 { + let outcome = tokio::time::timeout(LIFECYCLE_TIMEOUT, client.force_stop_and_wait()) + .await + .expect("childless force_stop_and_wait hung") + .expect("childless force_stop_and_wait failed"); + assert_eq!(outcome, None); + } + } + fn client_with_list_models_handler(handler: Arc) -> Client { + client_with_child_and_handler(None, Some(handler)) + } + + /// Build a client around an arbitrary child process, so lifecycle + /// behaviour can be exercised without a real CLI on the machine. + fn client_with_child_and_handler( + child: Option, + handler: Option>, + ) -> Client { Client { inner: Arc::new(ClientInner { - child: parking_lot::Mutex::new(None), + child: child::ChildLifecycle::new(child), #[cfg(feature = "bundled-in-process")] ffi_host: parking_lot::Mutex::new(None), rpc: { @@ -3268,7 +3497,7 @@ mod tests { negotiated_protocol_version: OnceLock::new(), state: parking_lot::Mutex::new(ConnectionState::Connected), lifecycle_tx: broadcast::channel(16).0, - on_list_models: Some(handler), + on_list_models: handler, models_cache: parking_lot::Mutex::new(Arc::new(tokio::sync::OnceCell::new())), session_fs_configured: false, session_fs_sqlite_declared: false, From 036bd233a7a9f816ef9a3d00d5623550e8f347ae Mon Sep 17 00:00:00 2001 From: jmoseley Date: Wed, 12 Aug 2026 17:34:53 -0700 Subject: [PATCH 10/13] [Docs] Document awaited CLI process termination MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Cover the three forced-shutdown shapes and when each resolves, the ownership guarantee that makes cancellation safe, and the cross-runtime promise for `ForcedShutdown` — each claim matching what the tests prove rather than what the design aspires to. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 0f9d5ac7-9999-4f37-82b3-a5533bfbc1f6 --- CHANGELOG.md | 21 +++++++++++++++++++++ rust/README.md | 42 ++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 63 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 7425ec075..1bec18126 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,27 @@ See [GitHub Releases](https://github.com/github/copilot-sdk/releases) for the fu ## [Unreleased] +### Feature: confirmed CLI process termination (Rust) + +`Client::force_stop` sends a kill signal and returns without waiting for the OS to release the process. That is fine for a short-lived CLI consumer, but an embedded host that shuts sessions down under a timeout keeps running, and on Unix every unreaped child stays a zombie for as long as its parent lives. + +The Rust SDK now exposes termination you can await: + +```rust +// Await the reap: resolves once the child is killed *and* released. +let exit_status = client.force_stop_and_wait().await?; + +// Or take an owned handle and await it somewhere else. +let shutdown = client.start_force_stop(); +tokio::spawn(async move { shutdown.wait().await }); +``` + +`force_stop_and_wait` resolves to the child's `ExitStatus`, or `None` for clients that never spawned one (stream-backed and in-process transports) — `None` means "nothing to terminate", never "termination failed", which is always an `Err`. `start_force_stop` returns a `ForcedShutdown` that borrows nothing from the client and is not bound to the runtime that started termination, so it can be awaited from another task or another runtime. + +Termination is now owned by the SDK rather than by the future that requests it. Previously `stop()` took the child out of its slot and then awaited the kill inline: if an outer timeout cancelled it in that window, the handle went with the cancelled future — the child was neither reaped nor recoverable, and a following `force_stop` found nothing to do. The child is now claimed synchronously and handed to a detached reaper, so cancelling `stop()`, cancelling `force_stop_and_wait()`, or dropping a `ForcedShutdown` cannot strand a signalled-but-unreaped process. Repeat and concurrent callers observe the same terminal outcome instead of racing for the handle. + +`force_stop` keeps its existing synchronous, infallible signature and semantics — the kill signal is still delivered synchronously, before anything is scheduled, so it works with no tokio runtime in context at all. Reaping runs on a dedicated thread with its own runtime rather than on a caller's, because a `tokio::spawn`ed reaper is cancelled when its runtime drops, which for an embedded host is exactly the moment termination matters. If the reaper cannot run to completion anyway — its thread fails to start, or it panics — waiters get a definitive `ErrorKind::Io` rather than sitting pending. + ### Feature: early session-event subscription (Rust) 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: diff --git a/rust/README.md b/rust/README.md index 18e03b72e..dbb245c94 100644 --- a/rust/README.md +++ b/rust/README.md @@ -74,6 +74,48 @@ let pong = client.ping("hello").await?; // Shutdown client.stop().await?; + +// Forced shutdown that confirms the OS reaped the CLI process, for hosts +// that must not leak a zombie when graceful shutdown times out. +if tokio::time::timeout(Duration::from_secs(5), client.stop()).await.is_err() { + let exit_status = client.force_stop_and_wait().await?; +} +``` + +### Terminating the CLI process + +`stop()` is the cooperative path. When it is unsuitable — a wedged +process, or an outer timeout that has already elapsed — three forced +variants differ only in how much of the teardown you wait for: + +| Method | Shape | Resolves when | +| ------ | ----- | ------------- | +| `force_stop()` | sync, infallible | the kill signal has been sent | +| `start_force_stop()` | sync, returns `ForcedShutdown` | immediately; await the handle for the reap | +| `force_stop_and_wait()` | async | the child is killed **and** reaped | + +Use `force_stop_and_wait()` (or await a `ForcedShutdown`) when the caller +must know the process is really gone. Sending a kill only starts +termination; until someone waits on the child, a Unix process stays a +zombie for as long as its parent lives — which for a long-lived embedded +host means they accumulate. + +Termination is owned by the SDK, not by the future that asks for it. The +kill is delivered synchronously — it never depends on a task being +polled — and the claimed child is then reaped on a dedicated thread with +its own runtime. So cancelling `stop()` or `force_stop_and_wait()`, +dropping a `ForcedShutdown`, or tearing down the runtime that started +termination cannot strand a signalled-but-unreaped process. Repeat and +concurrent callers all observe the same terminal outcome instead of +racing for the handle, and because the reaper is not bound to a caller's +runtime, a handle can be awaited on a different one: + +```rust,ignore +let shutdown = client.start_force_stop(); // teardown starts now +tokio::spawn(async move { + let exit_status = shutdown.wait().await?; // observed elsewhere + Ok::<_, github_copilot_sdk::Error>(exit_status) +}); ``` After `Client::start` succeeds, inspect its startup cost without parsing logs: From 17b1a7f36aa42dd792e01d40c62d3cb283e08034 Mon Sep 17 00:00:00 2001 From: jmoseley Date: Wed, 12 Aug 2026 17:55:37 -0700 Subject: [PATCH 11/13] [Rust] Reap the CLI process when a client is simply dropped MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `ClientInner::drop` signalled the child and dropped it without waiting, so a client that was never stopped left behind exactly the zombie this work set out to remove. Signalling is not termination. Give `ChildLifecycle` its own `Drop` that hands the child to the reaper instead, so every path out of a client — graceful stop, forced stop, or plain drop — ends with the process released by the OS. `ClientInner` no longer needs to do anything: the field's drop covers it. Mutation-verified: reverting drop-time termination to a bare `start_kill` leaves the process in `ps` state `Z`, which the new test catches. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 0f9d5ac7-9999-4f37-82b3-a5533bfbc1f6 --- rust/src/child.rs | 55 ++++++++++++++++++++++++++++++++++++----------- rust/src/lib.rs | 3 ++- 2 files changed, 45 insertions(+), 13 deletions(-) diff --git a/rust/src/child.rs b/rust/src/child.rs index 97691dba5..605c84a2b 100644 --- a/rust/src/child.rs +++ b/rust/src/child.rs @@ -205,19 +205,16 @@ impl ChildLifecycle { state: self.state.subscribe(), } } +} - /// Best-effort signal used on `Drop`, where there is no opportunity to - /// await anything. Leaves the child in place so a reaper that is - /// already running keeps ownership. - pub(crate) fn signal_on_drop(&self) { - if let Some(child) = self.child.lock().as_mut() { - let pid = child.id(); - match child.start_kill() { - Ok(()) => info!(pid = ?pid, "kill signal sent for CLI process on drop"), - Err(error) => { - warn!(pid = ?pid, error = %error, "failed to kill CLI process on drop"); - } - } +/// Terminating on drop is what keeps the reaping guarantee honest for a +/// client that is simply dropped: signalling alone would leave the same +/// zombie this type exists to prevent, so hand the child to the reaper. +impl Drop for ChildLifecycle { + fn drop(&mut self) { + if self.has_child() { + info!("client dropped with a live CLI process; terminating it"); + drop(self.begin_termination()); } } } @@ -669,6 +666,40 @@ mod tests { } } + /// Dropping a client that was never stopped must still reap its CLI + /// process. Signalling alone would leave exactly the zombie this type + /// exists to prevent. + #[cfg(unix)] + #[test] + fn dropping_an_unstopped_lifecycle_reaps_the_child() { + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .expect("build runtime"); + let lifecycle = runtime.block_on(async { ChildLifecycle::new(Some(spawn_sleeper())) }); + let pid = lifecycle.pid().expect("sleeper should report a pid"); + + // No stop, no force stop, no handle — just drop it. + drop(lifecycle); + + let deadline = std::time::Instant::now() + std::time::Duration::from_secs(5); + loop { + let output = std::process::Command::new("ps") + .args(["-o", "stat=", "-p", &pid.to_string()]) + .output() + .expect("run ps"); + let state = String::from_utf8_lossy(&output.stdout).trim().to_string(); + if state.is_empty() { + return; + } + assert!( + std::time::Instant::now() < deadline, + "the CLI process was left unreaped after drop (ps state {state:?})" + ); + std::thread::sleep(std::time::Duration::from_millis(50)); + } + } + #[test] fn forced_shutdown_handle_is_send_and_static() { fn assert_send_static() {} diff --git a/rust/src/lib.rs b/rust/src/lib.rs index ad758153b..36a003eef 100644 --- a/rust/src/lib.rs +++ b/rust/src/lib.rs @@ -2684,7 +2684,8 @@ impl Client { impl Drop for ClientInner { fn drop(&mut self) { - self.child.signal_on_drop(); + // The CLI child is terminated and reaped by `ChildLifecycle`'s own + // `Drop`, which runs when this struct's fields drop. #[cfg(feature = "bundled-in-process")] { if let Some(host) = self.ffi_host.lock().take() { From 0601c9459305cfef19848df55b6a4d041260728d Mon Sep 17 00:00:00 2001 From: jmoseley Date: Wed, 12 Aug 2026 18:12:05 -0700 Subject: [PATCH 12/13] [Rust] Reap by polling instead of SIGCHLD wakeup MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit CI caught this on Ubuntu: dropping a client left its CLI process in `Z` for the full 5s the new test allows. The reaper awaited `Child::wait`, which on Unix falls back to the SIGCHLD driver once `try_wait` misses — and the signal registration belongs to the runtime that spawned the child, not to the reaper's own. The reaper could therefore sleep waiting for a wakeup that was delivered elsewhere. macOS happened to win the race locally; Linux did not. Poll `try_wait` on a short backoff (2ms doubling to 50ms) instead. That makes the reap independent of which runtime owns the signal registration, which is the whole premise of running it on its own thread. It costs a handful of wakeups once per client teardown. `enable_all()` on the reaper runtime stays load-bearing, now for the timer rather than the signal driver, and the test that guards it is unchanged: it reaps a child that is still running at the first check. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 0f9d5ac7-9999-4f37-82b3-a5533bfbc1f6 --- rust/src/child.rs | 47 ++++++++++++++++++++++++++++++++++++----------- 1 file changed, 36 insertions(+), 11 deletions(-) diff --git a/rust/src/child.rs b/rust/src/child.rs index 605c84a2b..a0298bafc 100644 --- a/rust/src/child.rs +++ b/rust/src/child.rs @@ -20,6 +20,7 @@ use std::process::ExitStatus; use std::sync::Arc; +use std::time::Duration; use parking_lot::Mutex; use tokio::process::Child; @@ -171,10 +172,9 @@ impl ChildLifecycle { // Moved in, so it drops with the thread even if the body // panics before publishing. let mut guard = guard; - // `enable_all` is load-bearing on Unix, not boilerplate: - // once `try_wait` misses, readiness for `Child::wait` - // comes from the signal driver. Without it, a child that - // outlives the reaper's first poll — a process wedged in + // `enable_all` is load-bearing, not boilerplate: `reap` + // sleeps between polls, so without the timer a child that + // outlives the first check — a process wedged in // uninterruptible I/O, exactly what `force_stop` exists // for — would never be reaped. match tokio::runtime::Builder::new_current_thread() @@ -242,11 +242,36 @@ impl Drop for ReapGuard { } /// Wait for the OS to release a child that has already been signalled. +/// +/// Polls rather than awaiting `Child::wait`. The child is usually a zombie +/// by the first check, but not always, and once `try_wait` misses, +/// `wait` depends on the SIGCHLD driver waking *this* runtime — which is +/// not guaranteed when the child was spawned under a different one, as it +/// always is here. Polling keeps the reap independent of which runtime +/// owns the signal registration. It runs once per client teardown, so the +/// wakeups are immaterial. async fn reap(mut child: Child) -> ReapState { + const FIRST_INTERVAL: Duration = Duration::from_millis(2); + const MAX_INTERVAL: Duration = Duration::from_millis(50); + let pid = child.id(); - let state = ReapState::from_wait(child.wait().await); - info!(pid = ?pid, outcome = ?state, "CLI process reaped"); - state + let mut interval = FIRST_INTERVAL; + loop { + match child.try_wait() { + Ok(Some(status)) => { + info!(pid = ?pid, ?status, "CLI process reaped"); + return ReapState::Reaped(Some(status)); + } + Ok(None) => { + tokio::time::sleep(interval).await; + interval = (interval * 2).min(MAX_INTERVAL); + } + Err(error) => { + warn!(pid = ?pid, error = %error, "could not reap the CLI process"); + return ReapState::from_wait(Err(error)); + } + } + } } /// Owned handle to a CLI process being terminated. /// @@ -635,10 +660,10 @@ mod tests { /// load-bearing rather than boilerplate. /// /// Every other test kills its child before the reaper runs, so the - /// child is already a zombie and the first `try_wait` succeeds — the - /// signal path is never touched, and the suite would pass even without - /// a signal driver. This case waits on a child that is still running - /// at the first poll, on a runtime other than the one that spawned it: + /// child is already a zombie and the very first `try_wait` succeeds — + /// the retry loop is never entered, and the suite would pass even + /// without a timer. This case reaps a child that is still running at + /// the first check, on a runtime other than the one that spawned it: /// the shape of the wedged process `force_stop` exists for. #[test] fn reap_completes_for_a_child_still_running_at_the_first_poll() { From e9e4fc6165ce6947ba6891629f0ee98f7f4f73f2 Mon Sep 17 00:00:00 2001 From: jmoseley Date: Wed, 12 Aug 2026 18:32:09 -0700 Subject: [PATCH 13/13] [Rust] Report a refused kill instead of waiting on a live process MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `start_kill` failures were logged and swallowed. That is right when the child is already gone — the reap then reports its real status — but wrong when the OS refuses the signal: nothing will make that process exit, so the reaper would poll forever, every `ForcedShutdown::wait()` would park indefinitely, and `force_stop()` would report success. Distinguish the two with `signal_refusal`, which treats a failed kill as benign only when the child has actually exited. A genuine refusal is published as a definitive error *before* the reaper starts, so waiters observe it rather than blocking. Ownership is unchanged: the reaper still takes the child and still overwrites the error with the truth if the process does go away, so a refusal degrades to an observable error rather than to a silent success or a hang. The decision is a pure function so it can be tested directly; fabricating a real refusal would mean signalling a process the test does not own. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 0f9d5ac7-9999-4f37-82b3-a5533bfbc1f6 --- rust/src/child.rs | 64 ++++++++++++++++++++++++++++++++++++++++++----- 1 file changed, 58 insertions(+), 6 deletions(-) diff --git a/rust/src/child.rs b/rust/src/child.rs index a0298bafc..bcf7fb7b8 100644 --- a/rust/src/child.rs +++ b/rust/src/child.rs @@ -134,13 +134,13 @@ impl ChildLifecycle { match slot.take() { Some(mut child) => { let pid = child.id(); - if let Err(error) = child.start_kill() { - // Usually just an already-exited child; the wait - // below still reports its real status. - warn!(pid = ?pid, error = %error, "kill signal not delivered to CLI process"); + let kill = child.start_kill(); + let refusal = signal_refusal(kill, || matches!(child.try_wait(), Ok(Some(_)))); + if let Some(error) = &refusal { + warn!(pid = ?pid, error = %error, "kill signal refused by the OS"); } self.state.send_replace(ReapState::Pending); - (Some(child), self.handle()) + (Some((child, refusal)), self.handle()) } // Either termination already started (the running reaper // will publish) or there was never a child (already @@ -148,11 +148,22 @@ impl ChildLifecycle { None => (None, self.handle()), } }; - let Some(child) = claimed else { + let Some((child, refusal)) = claimed else { return handle; }; let pid = child.id(); + // Publish the refusal *before* the reaper starts, so waiters see it + // rather than blocking on a process that was never signalled. The + // reaper still takes ownership and still overwrites this with the + // truth if the process does go away, so a refusal degrades to an + // observable error rather than to a silent success or a hang. + if let Some(error) = refusal { + self.state.send_replace(ReapState::Failed { + kind: error.kind(), + message: format!("the CLI child could not be signalled: {error}"), + }); + } info!(pid = ?pid, "terminating CLI process"); // Reap on a dedicated thread with its own runtime rather than on a // caller's. A `tokio::spawn`ed task is cancelled when its runtime @@ -241,6 +252,24 @@ impl Drop for ReapGuard { } } +/// Classify a `start_kill` result. +/// +/// A failed kill is benign when the child is already gone — the reap +/// reports its real status. Anything else means the signal was refused, so +/// nothing will make the process exit and it must not be reported as a +/// successful stop. `exited` is evaluated only when it matters, since it +/// costs a `waitpid`. +fn signal_refusal( + kill: std::io::Result<()>, + exited: impl FnOnce() -> bool, +) -> Option { + match kill { + Ok(()) => None, + Err(_) if exited() => None, + Err(error) => Some(error), + } +} + /// Wait for the OS to release a child that has already been signalled. /// /// Polls rather than awaiting `Child::wait`. The child is usually a zombie @@ -725,6 +754,29 @@ mod tests { } } + /// A kill that fails because the child is already gone is benign — + /// the reap reports its real status. A kill the OS refuses is not: + /// nothing will make that process exit, so it must surface as an + /// error rather than as a successful stop or an endless wait. + #[test] + fn signal_refusal_distinguishes_a_dead_child_from_a_refused_kill() { + assert!(signal_refusal(Ok(()), || panic!("must not probe on success")).is_none()); + assert!( + signal_refusal( + Err(std::io::Error::from(std::io::ErrorKind::InvalidInput)), + || true + ) + .is_none(), + "an already-exited child is not a refusal" + ); + let refusal = signal_refusal( + Err(std::io::Error::from(std::io::ErrorKind::PermissionDenied)), + || false, + ) + .expect("a refused kill on a live child must be reported"); + assert_eq!(refusal.kind(), std::io::ErrorKind::PermissionDenied); + } + #[test] fn forced_shutdown_handle_is_send_and_static() { fn assert_send_static() {}