diff --git a/CHANGELOG.md b/CHANGELOG.md index e9f22a3df..1bec18126 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,49 @@ 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: + +```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. + +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 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. + ### 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..0292f98e0 100644 --- a/docs/features/streaming-events.md +++ b/docs/features/streaming-events.md @@ -218,6 +218,48 @@ 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 +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. 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: + +* 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 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. + ## 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..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: @@ -583,6 +625,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. 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()`. + +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 +858,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/child.rs b/rust/src/child.rs new file mode 100644 index 000000000..bcf7fb7b8 --- /dev/null +++ b/rust/src/child.rs @@ -0,0 +1,785 @@ +//! 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 std::time::Duration; + +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(); + 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, refusal)), 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, 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 + // 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, 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() + .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(), + } + } +} + +/// 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()); + } + } +} + +/// 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(), + }); + } + } +} + +/// 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 +/// 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 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. +/// +/// 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 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() { + 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:?}"), + } + } + + /// 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)); + } + } + + /// 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() {} + assert_send_static::(); + } +} diff --git a/rust/src/lib.rs b/rust/src/lib.rs index cafa3c596..36a003eef 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, @@ -1976,15 +1978,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. + /// 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`](Self::unregister_session) to - /// clean up (typically on session destroy). + /// 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 +1999,30 @@ 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. 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() } /// Returns the protocol version negotiated with the CLI server, if any. @@ -2209,6 +2235,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.registered_session_ids() + } + #[cfg(feature = "test-support")] #[doc(hidden)] /// Disconnect and delete every session owned by this test client's isolated @@ -2341,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. @@ -2362,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"); @@ -2399,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(); @@ -2438,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; @@ -2481,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 /// @@ -2507,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")] { @@ -2526,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. @@ -2569,14 +2684,8 @@ 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"); - } - } + // 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() { @@ -3215,10 +3324,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: { @@ -3235,7 +3498,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, 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 c6c806b1c..bc7a2de34 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,211 @@ impl Drop for WaiterGuard { struct PendingSessionRegistration { client: Client, - session_id: SessionId, + session_id: PendingSessionId, shutdown: CancellationToken, disarmed: bool, } +/// Which registration a [`PendingSessionRegistration`] owns and may remove +/// on cleanup. +/// +/// 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 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 { + id: session_id, + token, + }, + shutdown, + disarmed: false, + } + } + + /// Guard for a `session.create` registration arbitrated through a + /// shared slot. + fn deferred( + client: Client, + slot: DeferredRegistrationSlot, + shutdown: CancellationToken, + ) -> Self { Self { client, - session_id, + session_id: PendingSessionId::Deferred(slot), shutdown, disarmed: false, } } + /// 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, + }; + } + + /// 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; - self.client.unregister_session(&self.session_id); + self.release(); self.disarmed = true; } @@ -135,7 +346,7 @@ impl Drop for PendingSessionRegistration { fn drop(&mut self) { if !self.disarmed { self.shutdown.cancel(); - self.client.unregister_session(&self.session_id); + self.release(); } } } @@ -190,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 { @@ -576,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(()) } @@ -653,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); } } @@ -795,6 +1012,102 @@ 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 *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, + /// 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 +1129,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,27 +1303,34 @@ 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 // 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(|| { @@ -986,45 +1347,36 @@ impl Client { }) .into()); } - let channels = client.register_session(&parsed.session_id); - *stash.lock() = 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 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 = match self + let result = 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); - } - }; + .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(), @@ -1032,10 +1384,10 @@ impl Client { .into()); } - let (session_id, channels) = inline_stash - .lock() - .take() - .expect("session registration must have populated stash on success"); + 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(), @@ -1062,8 +1414,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 +1426,7 @@ impl Client { session_id = %session_id, "Client::create_session complete" ); + registration.disarm(); let session = Session { id: session_id, cwd: self.cwd().clone(), @@ -1083,6 +1439,7 @@ impl Client { capabilities, open_canvases, event_tx, + registration_token, }; apply_mode_post_create_patch( &session, @@ -1106,7 +1463,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() { @@ -1220,11 +1582,11 @@ 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 shutdown = CancellationToken::new(); - let (event_tx, _) = tokio::sync::broadcast::channel(512); let event_loop = spawn_event_loop( session_id.clone(), self.clone(), @@ -1242,8 +1604,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, @@ -1286,10 +1652,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 @@ -1343,6 +1711,7 @@ impl Client { capabilities, open_canvases, event_tx, + registration_token, }; apply_mode_post_create_patch( &session, @@ -1357,6 +1726,151 @@ 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 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 +/// 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, 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()) + } + + /// 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( @@ -2563,8 +3077,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/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..19ee46551 --- /dev/null +++ b/rust/tests/prepared_session_test.rs @@ -0,0 +1,1231 @@ +//! 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::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, RequestId, 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; + } +} + +/// 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); + let client = Client::from_streams(client_read, client_write, std::env::temp_dir()).unwrap(); + ( + client, + FakeServer { + read: server_read, + write: server_write, + }, + ) +} + +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" }) +} + +/// 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. +/// +/// 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; + loop { + let outstanding = client.registered_session_ids_for_test().len(); + if outstanding == 0 { + return; + } + assert!( + tokio::time::Instant::now() < deadline, + "{outstanding} session registration(s) were never cleaned up" + ); + 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; +} + +/// 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(); + + 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()); +} + +// --------------------------------------------------------------------------- +// 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. +/// +/// 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 + .registered_session_ids_for_test() + .iter() + .any(|id| id.as_str() == session_id) + { + assert!( + tokio::time::Instant::now() < deadline, + "inline callback never registered the expected session" + ); + 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); +}