Skip to main content

github_copilot_sdk/
session.rs

1use std::collections::HashMap;
2use std::panic::AssertUnwindSafe;
3use std::path::{Path, PathBuf};
4use std::sync::Arc;
5use std::time::{Duration, Instant};
6
7use futures_util::FutureExt;
8use parking_lot::Mutex as ParkingLotMutex;
9use serde_json::Value;
10use tokio::sync::oneshot;
11use tokio::task::JoinHandle;
12use tokio_util::sync::CancellationToken;
13use tracing::{Instrument, error, warn};
14
15use crate::canvas::CanvasHandler;
16use crate::generated::api_types::{
17    LogRequest, ModelSwitchAutoTierRequest, ModelSwitchAutoTierResult, ModelSwitchToRequest,
18    OpenCanvasInstance, PermissionDecisionRequest, RegisterEventInterestParams,
19    ToolsGetCurrentMetadataResult, rpc_methods,
20};
21use crate::generated::session_events::{
22    CommandExecuteData, ElicitationRequestedData, ExternalToolRequestedData, McpOauthRequiredData,
23    SessionCanvasClosedData, SessionErrorData, SessionEventType, SessionIdleData, SessionMode,
24};
25use crate::handler::{
26    AutoModeSwitchHandler, AutoModeSwitchResponse, ElicitationHandler, ExitPlanModeHandler,
27    McpAuthHandler, McpAuthRequest, McpAuthResult, PermissionHandler, PermissionResult,
28    UserInputHandler, UserInputResponse,
29};
30use crate::hooks::SessionHooks;
31use crate::provider_token::BearerTokenProvider;
32use crate::session_fs::SessionFsProvider;
33use crate::trace_context::inject_trace_context;
34use crate::transforms::SystemMessageTransform;
35use crate::types::{
36    AutoTier, AutoTierPreference, CommandContext, CommandDefinition, CommandHandler,
37    CreateSessionResult, ElicitationRequest, ElicitationResult, ExitPlanModeData,
38    GetMessagesResponse, MessageOptions, PermissionRequestData, RequestId, ResumeSessionConfig,
39    ResumeSessionResult, SectionOverride, SessionCapabilities, SessionConfig, SessionEvent,
40    SessionId, SetModelOptions, SystemMessageConfig, ToolInvocation, ToolResult,
41    ToolResultExpanded, TraceContext, UiInputOptions, ensure_attachment_display_names,
42};
43use crate::{
44    Client, Error, ErrorKind, JsonRpcResponse, SessionErrorKind, SessionEventNotification,
45    error_codes,
46};
47
48/// Fixed name of the runtime's built-in tool-search tool. A client can replace
49/// its behavior by registering a tool with this exact name and
50/// `overrides_built_in_tool` set to `true`.
51const TOOL_SEARCH_TOOL_NAME: &str = "tool_search_tool";
52
53/// Default capacity of the per-session event broadcast buffer backing
54/// [`Session::subscribe`] and [`PreparedSession::subscribe`].
55///
56/// Override per session with
57/// [`SessionConfig::event_buffer_capacity`](crate::types::SessionConfig::event_buffer_capacity)
58/// or
59/// [`ResumeSessionConfig::event_buffer_capacity`](crate::types::ResumeSessionConfig::event_buffer_capacity).
60pub const DEFAULT_EVENT_BUFFER_CAPACITY: usize = 512;
61
62/// Validate a caller-supplied event buffer capacity and resolve the default.
63///
64/// Zero is rejected rather than clamped: a zero-capacity broadcast channel
65/// cannot exist, and silently substituting a different capacity would hide a
66/// caller bug.
67fn resolve_event_buffer_capacity(capacity: Option<usize>) -> Result<usize, Error> {
68    match capacity {
69        Some(0) => Err(Error::with_message(
70            ErrorKind::InvalidConfig,
71            "event_buffer_capacity must be greater than zero",
72        )),
73        Some(capacity) => Ok(capacity),
74        None => Ok(DEFAULT_EVENT_BUFFER_CAPACITY),
75    }
76}
77
78/// Bundle of the per-session callbacks the SDK dispatches to. Built from a
79/// [`SessionConfig`] / [`ResumeSessionConfig`] at
80/// [`Client::create_session`] / [`Client::resume_session`] time. Each
81/// field is `None` (or an empty map for tools) when the caller didn't
82/// install a handler -- in that case the SDK skips dispatch for that
83/// event type. The wire flags on `session.create` / `session.resume`
84/// are derived from these fields.
85#[derive(Clone)]
86pub(crate) struct SessionHandlers {
87    pub permission: Option<Arc<dyn PermissionHandler>>,
88    pub managed_settings_enabled: bool,
89    pub elicitation: Option<Arc<dyn ElicitationHandler>>,
90    pub mcp_auth: Option<Arc<dyn McpAuthHandler>>,
91    pub user_input: Option<Arc<dyn UserInputHandler>>,
92    pub exit_plan_mode: Option<Arc<dyn ExitPlanModeHandler>>,
93    pub auto_mode_switch: Option<Arc<dyn AutoModeSwitchHandler>>,
94    pub tools: Arc<HashMap<String, Arc<dyn crate::tool::ToolHandler>>>,
95}
96
97type PendingExternalTools = Arc<ParkingLotMutex<HashMap<RequestId, Arc<CancellationToken>>>>;
98
99struct PendingExternalToolGuard {
100    request_id: RequestId,
101    token: Arc<CancellationToken>,
102    pending: PendingExternalTools,
103}
104
105impl Drop for PendingExternalToolGuard {
106    fn drop(&mut self) {
107        let mut pending = self.pending.lock();
108        if pending
109            .get(&self.request_id)
110            .is_some_and(|token| Arc::ptr_eq(token, &self.token))
111        {
112            pending.remove(&self.request_id);
113        }
114    }
115}
116
117impl PendingExternalToolGuard {
118    fn claim(&self) -> bool {
119        let mut pending = self.pending.lock();
120        if pending
121            .get(&self.request_id)
122            .is_some_and(|token| Arc::ptr_eq(token, &self.token))
123        {
124            pending.remove(&self.request_id);
125            true
126        } else {
127            false
128        }
129    }
130}
131
132fn has_managed_settings(
133    enable_managed_settings: Option<bool>,
134    managed_settings: Option<&crate::types::ManagedSettings>,
135) -> bool {
136    enable_managed_settings == Some(true) || managed_settings.is_some()
137}
138
139/// Shared state between a [`Session`] and its event loop, used by [`Session::send_and_wait`].
140struct IdleWaiter {
141    tx: oneshot::Sender<Result<Option<SessionEvent>, Error>>,
142    last_assistant_message: Option<SessionEvent>,
143    started_at: Instant,
144    first_assistant_message_seen: bool,
145}
146
147fn structured_output_error(message: impl Into<String>) -> Error {
148    Error::with_message(
149        ErrorKind::Session(SessionErrorKind::AgentError),
150        message.into(),
151    )
152}
153
154fn is_structured_output_event(event: &SessionEvent) -> bool {
155    matches!(
156        event.parsed_type(),
157        SessionEventType::UserMessage
158            | SessionEventType::AssistantMessage
159            | SessionEventType::SessionIdle
160            | SessionEventType::SessionError
161    )
162}
163
164struct StructuredOutputState {
165    message_id: String,
166    started: bool,
167    final_message: Option<SessionEvent>,
168}
169
170impl StructuredOutputState {
171    fn observe(&mut self, event: SessionEvent) -> Result<Option<SessionEvent>, Error> {
172        if event.agent_id.as_deref().is_some_and(|id| !id.is_empty()) {
173            return Ok(None);
174        }
175        match event.parsed_type() {
176            SessionEventType::UserMessage => {
177                let data: crate::session_events::UserMessageData =
178                    serde_json::from_value(event.data)?;
179                if data.message_id.as_deref() == Some(self.message_id.as_str()) {
180                    self.started = true;
181                }
182            }
183            SessionEventType::AssistantMessage => {
184                let data: crate::session_events::AssistantMessageData =
185                    serde_json::from_value(event.data.clone())?;
186                if data.originating_message_id.as_deref() == Some(self.message_id.as_str()) {
187                    self.started = true;
188                    self.final_message =
189                        if data.tool_requests.is_some_and(|tools| !tools.is_empty()) {
190                            None
191                        } else {
192                            Some(event)
193                        };
194                }
195            }
196            SessionEventType::SessionIdle if self.started => {
197                let data: SessionIdleData = serde_json::from_value(event.data)?;
198                if data.mode == Some(SessionMode::Autopilot) {
199                    return Ok(None);
200                }
201                if data.aborted == Some(true) {
202                    return Err(structured_output_error(
203                        "session aborted before structured output completed",
204                    ));
205                }
206                let result = self.final_message.take().ok_or_else(|| {
207                    structured_output_error("run completed without a structured assistant response")
208                })?;
209                let data: crate::session_events::AssistantMessageData =
210                    serde_json::from_value(result.data.clone())?;
211                if data.content.trim().is_empty() {
212                    return Err(structured_output_error(
213                        "run completed without a structured assistant response",
214                    ));
215                }
216                return Ok(Some(result));
217            }
218            SessionEventType::SessionError if self.started => {
219                let data: SessionErrorData = serde_json::from_value(event.data)?;
220                return Err(structured_output_error(format!(
221                    "session error: {}",
222                    data.message
223                )));
224            }
225            _ => {}
226        }
227        Ok(None)
228    }
229}
230
231/// RAII guard that clears the [`Session::idle_waiter`] slot on drop. Used
232/// by [`Session::send_and_wait`] to ensure the slot doesn't leak if the
233/// caller's future is cancelled (outer `tokio::time::timeout` / `select!`
234/// / dropped JoinHandle). Synchronous clear via `parking_lot::Mutex` —
235/// no async drop needed.
236///
237/// Without this, an outer cancellation between "install waiter" and
238/// "drain channel" would leave the slot occupied, causing all subsequent
239/// `send` and `send_and_wait` calls on the session to return
240/// [`SendWhileWaiting`](SessionErrorKind::SendWhileWaiting). Closes RFD-400
241/// review finding #2.
242struct WaiterGuard {
243    slot: Arc<ParkingLotMutex<Option<IdleWaiter>>>,
244}
245
246impl Drop for WaiterGuard {
247    fn drop(&mut self) {
248        self.slot.lock().take();
249    }
250}
251
252struct PendingSessionRegistration {
253    client: Client,
254    session_id: PendingSessionId,
255    shutdown: CancellationToken,
256    external_tools_shutdown: CancellationToken,
257    disarmed: bool,
258}
259
260/// Which session ID a [`PendingSessionRegistration`] should unregister on
261/// cleanup.
262///
263/// `session.create` for cloud sessions without a caller-pinned ID does not
264/// know the ID until the response arrives, at which point the inline
265/// response callback registers it and stashes it. The guard therefore reads
266/// the stash at cleanup time instead of capturing an ID up front.
267enum PendingSessionId {
268    /// The ID was known before the RPC was issued (resume, and create with a
269    /// client- or caller-supplied ID).
270    Known(SessionId, crate::router::RegistrationToken),
271    /// Server-assigned ID, populated by the `session.create` inline response
272    /// callback. `None` in the stash means nothing was ever registered.
273    Deferred(Arc<ParkingLotMutex<Option<(SessionId, crate::router::SessionRegistration)>>>),
274}
275
276impl PendingSessionRegistration {
277    fn new(
278        client: Client,
279        session_id: SessionId,
280        token: crate::router::RegistrationToken,
281        shutdown: CancellationToken,
282        external_tools_shutdown: CancellationToken,
283    ) -> Self {
284        Self {
285            client,
286            session_id: PendingSessionId::Known(session_id, token),
287            shutdown,
288            external_tools_shutdown,
289            disarmed: false,
290        }
291    }
292
293    /// Guard for a registration whose session ID is assigned by the server.
294    fn deferred(
295        client: Client,
296        stash: Arc<ParkingLotMutex<Option<(SessionId, crate::router::SessionRegistration)>>>,
297        shutdown: CancellationToken,
298        external_tools_shutdown: CancellationToken,
299    ) -> Self {
300        Self {
301            client,
302            session_id: PendingSessionId::Deferred(stash),
303            shutdown,
304            external_tools_shutdown,
305            disarmed: false,
306        }
307    }
308
309    fn registered_id(&self) -> Option<SessionId> {
310        match &self.session_id {
311            PendingSessionId::Known(id, _) => Some(id.clone()),
312            PendingSessionId::Deferred(stash) => stash.lock().as_ref().map(|(id, _)| id.clone()),
313        }
314    }
315
316    /// Re-target the guard at a now-known session ID. Used by
317    /// `session.create` once the response has been parsed and the stash has
318    /// been drained into the event loop.
319    fn resolve_to(&mut self, session_id: SessionId, token: crate::router::RegistrationToken) {
320        self.session_id = PendingSessionId::Known(session_id, token);
321    }
322
323    async fn cleanup(mut self, event_loop: JoinHandle<()>) {
324        self.external_tools_shutdown.cancel();
325        self.shutdown.cancel();
326        let _ = event_loop.await;
327        if let Some(id) = self.registered_id() {
328            if let PendingSessionId::Known(_, token) = self.session_id {
329                self.client.unregister_session_owned(&id, token);
330            } else if let PendingSessionId::Deferred(stash) = &self.session_id
331                && let Some((id, registration)) = stash.lock().as_ref()
332            {
333                self.client.unregister_session_owned(id, registration.token);
334            }
335        }
336        self.disarmed = true;
337    }
338
339    fn disarm(&mut self) {
340        self.disarmed = true;
341    }
342}
343
344impl Drop for PendingSessionRegistration {
345    fn drop(&mut self) {
346        if !self.disarmed {
347            self.external_tools_shutdown.cancel();
348            self.shutdown.cancel();
349            if let Some(id) = self.registered_id() {
350                if let PendingSessionId::Known(_, token) = self.session_id {
351                    self.client.unregister_session_owned(&id, token);
352                } else if let PendingSessionId::Deferred(stash) = &self.session_id
353                    && let Some((id, registration)) = stash.lock().as_ref()
354                {
355                    self.client.unregister_session_owned(id, registration.token);
356                }
357            }
358        }
359    }
360}
361
362/// A session on a GitHub Copilot CLI server.
363///
364/// Created via [`Client::create_session`] or [`Client::resume_session`].
365/// Owns an internal event loop that dispatches events to the per-callback
366/// handlers installed on the session config.
367///
368/// Protocol methods (`send`, `get_events`, `abort`, etc.) automatically
369/// inject the session ID into RPC params.
370///
371/// Call [`destroy`](Self::destroy) for graceful cleanup (RPC + local). If dropped
372/// without calling `destroy`, the `Drop` impl aborts the event loop and
373/// unregisters from the router as a best-effort safety net.
374pub struct Session {
375    id: SessionId,
376    cwd: PathBuf,
377    workspace_path: Option<PathBuf>,
378    remote_url: Option<String>,
379    client: Client,
380    /// Handle to the spawned event-loop task. Sync `parking_lot::Mutex`
381    /// because the lock is never held across an `.await` and the `Drop`
382    /// impl needs to take the handle synchronously without `try_lock`
383    /// fallibility.
384    event_loop: ParkingLotMutex<Option<JoinHandle<()>>>,
385    /// Cooperative shutdown signal for the event loop. The loop selects
386    /// on [`shutdown.cancelled()`](CancellationToken::cancelled) alongside
387    /// its inbound channels; [`Session::stop_event_loop`] and [`Drop`]
388    /// both call [`cancel()`](CancellationToken::cancel) to ask the loop
389    /// to exit between iterations rather than aborting the task (which
390    /// can land at any await point and leave the session mid-protocol).
391    /// See RFD-400 review finding #3.
392    ///
393    /// `CancellationToken` is the canonical signalling primitive in
394    /// `tokio_util`; it is what `tonic` uses for the equivalent task-
395    /// coordination case. Advanced consumers can obtain a child token
396    /// via [`Session::cancellation_token`] to bind their own work to
397    /// the session lifetime.
398    shutdown: CancellationToken,
399    /// Cancels only host-owned external tool callbacks. Disconnect signals this
400    /// before the destroy RPC without stopping unrelated event delivery.
401    external_tools_shutdown: CancellationToken,
402    /// Only populated while a `send_and_wait` call is in flight.
403    ///
404    /// Sync `parking_lot::Mutex` because the lock is never held across an
405    /// `.await`, and synchronous access lets the `WaiterGuard` RAII helper
406    /// in `send_and_wait` clear the slot from a `Drop` impl on caller-side
407    /// cancellation. See RFD-400 review (cancel-safety hardening).
408    idle_waiter: Arc<ParkingLotMutex<Option<IdleWaiter>>>,
409    /// Capabilities negotiated with the CLI, updated on `capabilities.changed` events.
410    capabilities: Arc<parking_lot::RwLock<SessionCapabilities>>,
411    /// Canvas instances currently known to be open for this session.
412    open_canvases: Arc<parking_lot::RwLock<Vec<OpenCanvasInstance>>>,
413    /// Broadcast channel for runtime event subscribers — see [`Session::subscribe`].
414    event_tx: tokio::sync::broadcast::Sender<SessionEvent>,
415    github_token_registration:
416        ParkingLotMutex<Option<crate::github_token::GitHubTokenRegistration>>,
417    /// Identity of this session's router registration.
418    registration_token: crate::router::RegistrationToken,
419}
420
421impl Session {
422    /// Session ID assigned by the CLI.
423    pub fn id(&self) -> &SessionId {
424        &self.id
425    }
426
427    /// Working directory of the CLI process.
428    pub fn cwd(&self) -> &PathBuf {
429        &self.cwd
430    }
431
432    /// Workspace directory for the session (if using infinite sessions).
433    pub fn workspace_path(&self) -> Option<&Path> {
434        self.workspace_path.as_deref()
435    }
436
437    /// Remote session URL, if the session is running remotely.
438    pub fn remote_url(&self) -> Option<&str> {
439        self.remote_url.as_deref()
440    }
441
442    /// Session capabilities negotiated with the CLI.
443    ///
444    /// Capabilities are set during session creation and updated at runtime
445    /// via `capabilities.changed` events.
446    pub fn capabilities(&self) -> SessionCapabilities {
447        self.capabilities.read().clone()
448    }
449
450    /// Open canvas instances reported by the most recent `session.resume`
451    /// response or surfaced by inbound `canvas.opened` events.
452    pub fn open_canvases(&self) -> Vec<OpenCanvasInstance> {
453        self.open_canvases.read().clone()
454    }
455
456    /// Returns a [`CancellationToken`] that fires when this session shuts
457    /// down (via [`Session::stop_event_loop`], [`Session::destroy`], or
458    /// [`Drop`]).
459    ///
460    /// Use this to bind an external task's lifetime to the session — when
461    /// the session shuts down, awaiting [`cancelled()`](CancellationToken::cancelled)
462    /// resolves so cooperative consumers can stop cleanly.
463    ///
464    /// The returned handle is a *child* token: calling
465    /// [`cancel()`](CancellationToken::cancel) on it cancels only the
466    /// caller's child, not the session itself. To cancel the session, call
467    /// [`Session::stop_event_loop`].
468    ///
469    /// # Example
470    ///
471    /// ```no_run
472    /// # async fn example(session: github_copilot_sdk::session::Session) {
473    /// let token = session.cancellation_token();
474    /// tokio::select! {
475    ///     _ = token.cancelled() => println!("session shut down"),
476    ///     _ = tokio::time::sleep(std::time::Duration::from_secs(60)) => {
477    ///         println!("60s elapsed, session still alive");
478    ///     }
479    /// }
480    /// # }
481    /// ```
482    pub fn cancellation_token(&self) -> CancellationToken {
483        self.shutdown.child_token()
484    }
485
486    /// Subscribe to events for this session.
487    ///
488    /// Returns an [`EventSubscription`](crate::subscription::EventSubscription)
489    /// that yields every [`SessionEvent`] dispatched on this session's
490    /// event loop. Drop the value to unsubscribe; there is no separate
491    /// cancel handle.
492    ///
493    /// **Observe-only.** Subscribers receive a clone of every
494    /// [`SessionEvent`] but cannot influence permission decisions, tool
495    /// results, or anything else that requires returning a value. Those
496    /// remain the responsibility of the per-callback handlers passed via
497    /// [`SessionConfig`]'s `with_*_handler`
498    /// builder methods.
499    ///
500    /// The returned handle implements both an inherent
501    /// [`recv`](crate::subscription::EventSubscription::recv) method and
502    /// [`Stream`](tokio_stream::Stream), so callers can use a `while let`
503    /// loop or any combinator from `tokio_stream::StreamExt` /
504    /// `futures::StreamExt`.
505    ///
506    /// Each subscriber maintains its own queue. If a consumer cannot keep
507    /// up, the oldest events are dropped and `recv` returns
508    /// [`RecvErrorKind::Lagged`](crate::subscription::RecvErrorKind::Lagged)
509    /// reporting the count of skipped events. Slow consumers do not block
510    /// the session's event loop.
511    ///
512    /// # Example
513    ///
514    /// ```no_run
515    /// # async fn example(session: github_copilot_sdk::session::Session) {
516    /// let mut events = session.subscribe();
517    /// tokio::spawn(async move {
518    ///     while let Ok(event) = events.recv().await {
519    ///         println!("[{}] event {}", event.id, event.event_type);
520    ///     }
521    /// });
522    /// # }
523    /// ```
524    pub fn subscribe(&self) -> crate::subscription::EventSubscription {
525        crate::subscription::EventSubscription::new(self.event_tx.subscribe())
526    }
527
528    /// The underlying Client (for advanced use cases).
529    pub fn client(&self) -> &Client {
530        &self.client
531    }
532
533    /// Typed RPC namespace for this session.
534    ///
535    /// Every protocol method lives here under its schema-aligned path —
536    /// e.g. `session.rpc().workspaces().list_files()`. Wire method names
537    /// and request/response types are generated from the protocol schema,
538    /// so the typed namespace can't drift from the wire contract.
539    ///
540    /// The hand-authored helpers on [`Session`] delegate to this namespace
541    /// and remain the recommended entry point for everyday use; reach for
542    /// `rpc()` when you want a method without a hand-written wrapper.
543    pub fn rpc(&self) -> crate::generated::rpc::SessionRpc<'_> {
544        crate::generated::rpc::SessionRpc { session: self }
545    }
546
547    /// Stop the internal event loop. Called automatically on [`destroy`](Self::destroy).
548    ///
549    /// Cooperative: signals shutdown via the session's [`CancellationToken`]
550    /// and awaits the loop's natural exit rather than aborting the task, so
551    /// the loop always stops between iterations instead of at an arbitrary
552    /// await point. See RFD-400 review finding #3.
553    ///
554    /// Inbound requests are dispatched to their own spawned tasks, which this
555    /// call does not await. A handler (permission callback, tool call,
556    /// elicitation response) still running at teardown may therefore outlive
557    /// the loop, and its response can be lost if the connection closes first.
558    /// Await your own handler work before calling this if it must complete.
559    pub async fn stop_event_loop(&self) {
560        self.shutdown.cancel();
561        let handle = self.event_loop.lock().take();
562        if let Some(handle) = handle {
563            let _ = handle.await;
564        }
565        // Fail any pending send_and_wait so it returns immediately.
566        if let Some(waiter) = self.idle_waiter.lock().take() {
567            let _ = waiter.tx.send(Err(
568                ErrorKind::Session(SessionErrorKind::EventLoopClosed).into()
569            ));
570        }
571    }
572
573    /// Send a user message to the agent.
574    ///
575    /// Accepts anything convertible to [`MessageOptions`] — pass a `&str` for the
576    /// trivial case, or build a `MessageOptions` for mode/attachments. The
577    /// `wait_timeout` field on `MessageOptions` is ignored here (use
578    /// [`send_and_wait`](Self::send_and_wait) if you need to wait).
579    ///
580    /// Returns the assigned message ID, which can be used to correlate the
581    /// send with later [`SessionEvent`]s emitted in
582    /// response (assistant messages, tool requests, etc.).
583    ///
584    /// Returns an error if a [`send_and_wait`](Self::send_and_wait) call is
585    /// currently in flight, since the plain send would race with the waiter.
586    ///
587    /// # Cancel safety
588    ///
589    /// **Cancel-safe.** The underlying `session.send` RPC is dispatched
590    /// through the writer-actor (see [`Client::call`](crate::Client::call)),
591    /// so dropping this future after the actor has committed to writing
592    /// will not produce a partial frame on the wire. If the caller's
593    /// future is dropped between "frame enqueued" and "response received",
594    /// the message has already landed on the wire — the agent will process
595    /// it and emit events normally; the caller just won't see the returned
596    /// message ID.
597    pub async fn send(&self, opts: impl Into<MessageOptions>) -> Result<String, Error> {
598        if self.idle_waiter.lock().is_some() {
599            return Err(ErrorKind::Session(SessionErrorKind::SendWhileWaiting).into());
600        }
601        self.send_inner(opts.into()).await
602    }
603
604    async fn send_inner(&self, opts: MessageOptions) -> Result<String, Error> {
605        let mut params = serde_json::json!({
606            "sessionId": self.id,
607            "prompt": opts.prompt,
608        });
609        if let Some(source) = opts.source {
610            params["source"] = serde_json::to_value(source)?;
611        }
612        if let Some(m) = opts.mode {
613            params["mode"] = serde_json::to_value(m)?;
614        }
615        if let Some(am) = opts.agent_mode {
616            params["agentMode"] = serde_json::to_value(am)?;
617        }
618        if let Some(mut a) = opts.attachments {
619            ensure_attachment_display_names(&mut a);
620            params["attachments"] = serde_json::to_value(a)?;
621        }
622        if let Some(headers) = opts.request_headers
623            && !headers.is_empty()
624        {
625            params["requestHeaders"] = serde_json::to_value(headers)?;
626        }
627        if let Some(display_prompt) = opts.display_prompt {
628            params["displayPrompt"] = serde_json::to_value(display_prompt)?;
629        }
630        if let Some(schema) = opts.response_schema {
631            params["responseFormat"] = serde_json::json!({
632                "type": "json_schema",
633                "jsonSchema": { "name": "response", "strict": true, "schema": schema }
634            });
635        }
636        let trace_ctx = if opts.traceparent.is_some() || opts.tracestate.is_some() {
637            TraceContext {
638                traceparent: opts.traceparent,
639                tracestate: opts.tracestate,
640            }
641        } else {
642            self.client.resolve_trace_context().await
643        };
644        inject_trace_context(&mut params, &trace_ctx);
645        let rpc_start = Instant::now();
646        let result = self.client.call("session.send", Some(params)).await?;
647        let message_id = result
648            .get("messageId")
649            .and_then(|v| v.as_str())
650            .map(|s| s.to_string())
651            .unwrap_or_default();
652        tracing::debug!(
653            elapsed_ms = rpc_start.elapsed().as_millis(),
654            session_id = %self.id,
655            message_id = %message_id,
656            "Session::send completed successfully"
657        );
658        Ok(message_id)
659    }
660
661    /// Send a user message and wait for the agent to finish processing.
662    ///
663    /// Accepts anything convertible to [`MessageOptions`] — pass a `&str` for the
664    /// trivial case, or build a `MessageOptions` for mode/attachments/timeout.
665    /// Blocks until `session.idle` (success) or `session.error` (failure),
666    /// returning the last `assistant.message` event captured during streaming.
667    /// Times out after `MessageOptions::wait_timeout` (default 60 seconds).
668    ///
669    /// Only one unformatted `send_and_wait` may be active per session. Calling
670    /// [`send`](Self::send) during that wait also returns an error. Schema-bearing
671    /// waits instead correlate by originating message ID and support concurrency.
672    /// They select the last root message without tool requests at non-autopilot
673    /// idle, failing on aborted idle, session errors after starting, or no result.
674    ///
675    /// # Cancel safety
676    ///
677    /// **Cancel-safe.** A `WaiterGuard` clears the in-flight slot on every
678    /// exit path (success, internal failure, internal timeout, *and*
679    /// external cancellation via `tokio::time::timeout` / `select!` /
680    /// dropped JoinHandle). Subsequent `send` and `send_and_wait` calls on
681    /// this session will succeed normally — the slot is never leaked.
682    pub async fn send_and_wait(
683        &self,
684        opts: impl Into<MessageOptions>,
685    ) -> Result<Option<SessionEvent>, Error> {
686        let total_start = Instant::now();
687        let opts = opts.into();
688        if opts.response_schema.is_some() {
689            return self.send_and_wait_structured(opts).await.map(Some);
690        }
691        let timeout_duration = opts.wait_timeout.unwrap_or(Duration::from_secs(60));
692        let (tx, rx) = oneshot::channel();
693
694        {
695            let mut guard = self.idle_waiter.lock();
696            if guard.is_some() {
697                return Err(ErrorKind::Session(SessionErrorKind::SendWhileWaiting).into());
698            }
699            *guard = Some(IdleWaiter {
700                tx,
701                last_assistant_message: None,
702                started_at: total_start,
703                first_assistant_message_seen: false,
704            });
705        }
706
707        // RAII: clears the idle_waiter slot on every exit path, including
708        // external cancellation (caller's outer `select!` / `timeout` /
709        // dropped future). Without this, an outer cancellation would leak
710        // the slot and brick subsequent `send`/`send_and_wait` calls.
711        let _waiter_guard = WaiterGuard {
712            slot: self.idle_waiter.clone(),
713        };
714
715        let result = tokio::time::timeout(timeout_duration, async {
716            self.send_inner(opts).await?;
717            match rx.await {
718                Ok(result) => result,
719                Err(_) => Err(ErrorKind::Session(SessionErrorKind::EventLoopClosed).into()),
720            }
721        })
722        .await;
723
724        match result {
725            Ok(inner) => {
726                tracing::debug!(
727                    elapsed_ms = total_start.elapsed().as_millis(),
728                    session_id = %self.id,
729                    completed_by = if inner.is_ok() { "idle" } else { "error" },
730                    "Session::send_and_wait complete"
731                );
732                inner
733            }
734            Err(_) => {
735                tracing::warn!(
736                    elapsed_ms = total_start.elapsed().as_millis(),
737                    session_id = %self.id,
738                    completed_by = "timeout",
739                    "Session::send_and_wait failed"
740                );
741                Err(ErrorKind::Session(SessionErrorKind::Timeout(timeout_duration)).into())
742            }
743        }
744    }
745
746    /// Infer an output schema with the same `schemars` integration as custom tools,
747    /// then deserialize the final correlated root response at non-autopilot idle.
748    ///
749    /// Requires the `derive` feature. Provider schema restrictions apply. Serde
750    /// validates JSON/type compatibility, not every JSON Schema constraint.
751    /// Options must not specify a schema or immediate delivery. Dropping this
752    /// future or timing out unsubscribes the wait without aborting agent work.
753    #[cfg(feature = "derive")]
754    pub async fn send_and_wait_typed<T>(&self, opts: impl Into<MessageOptions>) -> Result<T, Error>
755    where
756        T: schemars::JsonSchema + serde::de::DeserializeOwned,
757    {
758        let mut opts = opts.into();
759        if opts.response_schema.is_some()
760            || opts.mode == Some(crate::types::DeliveryMode::Immediate)
761        {
762            return Err(Error::with_message(
763                ErrorKind::InvalidConfig,
764                "typed structured output cannot specify a response schema or immediate delivery",
765            ));
766        }
767        opts.response_schema = Some(crate::tool::schema_for::<T>());
768        let event = self.send_and_wait_structured(opts).await?;
769        let data: crate::session_events::AssistantMessageData = serde_json::from_value(event.data)?;
770        serde_json::from_str::<Option<T>>(&data.content)?.ok_or_else(|| {
771            structured_output_error("structured response was JSON null, not a result")
772        })
773    }
774
775    async fn send_and_wait_structured(&self, opts: MessageOptions) -> Result<SessionEvent, Error> {
776        let duration = opts.wait_timeout.unwrap_or(Duration::from_secs(60));
777        let mut events = self.subscribe();
778        let wait = async {
779            let mut admission = Box::pin(self.send(opts));
780            let mut pending = Vec::new();
781            let message_id = loop {
782                tokio::select! {
783                    result = &mut admission => break result?,
784                    event = events.recv() => {
785                        let event = event.map_err(|err| structured_output_error(err.to_string()))?;
786                        if is_structured_output_event(&event) {
787                            pending.push(event);
788                        }
789                    }
790                    _ = self.shutdown.cancelled() =>
791                        return Err(structured_output_error("session closed before structured output completed")),
792                }
793            };
794            let mut state = StructuredOutputState {
795                message_id,
796                started: false,
797                final_message: None,
798            };
799            for event in pending {
800                if let Some(result) = state.observe(event)? {
801                    return Ok(result);
802                }
803            }
804            loop {
805                tokio::select! {
806                    event = events.recv() => {
807                        let event = event.map_err(|err| structured_output_error(err.to_string()))?;
808                        if let Some(result) = state.observe(event)? {
809                            return Ok(result);
810                        }
811                    }
812                    _ = self.shutdown.cancelled() =>
813                        return Err(structured_output_error("session closed before structured output completed")),
814                }
815            }
816        };
817        tokio::time::timeout(duration, wait)
818            .await
819            .map_err(|_| Error::from(ErrorKind::Session(SessionErrorKind::Timeout(duration))))?
820    }
821
822    /// Retrieve the session's timeline events.
823    pub async fn get_events(&self) -> Result<Vec<SessionEvent>, Error> {
824        let result = self
825            .client
826            .call(
827                "session.getMessages",
828                Some(serde_json::json!({ "sessionId": self.id })),
829            )
830            .await?;
831        let response: GetMessagesResponse = serde_json::from_value(result)?;
832        Ok(response.events)
833    }
834
835    /// Deprecated alias for [`get_events`](Self::get_events).
836    #[deprecated(since = "0.1.0", note = "Use `get_events()` instead")]
837    pub async fn get_messages(&self) -> Result<Vec<SessionEvent>, Error> {
838        self.get_events().await
839    }
840
841    /// Abort the current agent turn.
842    ///
843    /// # Cancel safety
844    ///
845    /// **Cancel-safe.** Single `session.abort` RPC; the underlying
846    /// [`Client::call`](crate::Client::call) is cancel-safe via the
847    /// writer-actor.
848    pub async fn abort(&self) -> Result<(), Error> {
849        self.client
850            .call(
851                "session.abort",
852                Some(serde_json::json!({ "sessionId": self.id })),
853            )
854            .await?;
855        Ok(())
856    }
857
858    /// Switch to a different model.
859    ///
860    /// Pass `None` for `opts` if no extra configuration is needed.
861    pub async fn set_model(&self, model: &str, opts: Option<SetModelOptions>) -> Result<(), Error> {
862        let opts = opts.unwrap_or_default();
863        let auto_tier = opts.auto_tier.clone();
864        let request = ModelSwitchToRequest {
865            auto_tier: match &auto_tier {
866                Some(AutoTierPreference::Tier(tier)) => Some(tier.clone()),
867                _ => None,
868            },
869            compaction_decision: None,
870            context_tier: opts.context_tier,
871            defer_if_model_change_queued: None,
872            model_capabilities: opts.model_capabilities,
873            model_change_scope: None,
874            model_id: model.to_string(),
875            picker_persistence: None,
876            reasoning_effort: opts.reasoning_effort,
877            reasoning_summary: opts.reasoning_summary,
878            repo_scope: None,
879            require_available: None,
880            run_compaction_preflight: None,
881            source: None,
882            verbosity: None,
883        };
884
885        if matches!(auto_tier, Some(AutoTierPreference::Reset)) {
886            // The generated request skips a `None` tier, which the runtime reads
887            // as "leave the preference alone" rather than "use provider-default
888            // routing", so send an explicit null instead.
889            let mut wire_params = serde_json::to_value(request)?;
890            wire_params["sessionId"] = serde_json::Value::String(self.id.to_string());
891            wire_params["autoTier"] = serde_json::Value::Null;
892            self.client
893                .call("session.model.switchTo", Some(wire_params))
894                .await?;
895            return Ok(());
896        }
897
898        self.rpc().model().switch_to(request).await?;
899        Ok(())
900    }
901
902    /// Change the Auto routing preference without changing the selected model.
903    ///
904    /// The runtime does not apply the preference immediately. It records the
905    /// request and commits it only when a later user turn using the `auto`
906    /// model successfully obtains a usable model from the provider. A
907    /// [`ModelSwitchAutoTierStatus::Pending`] status therefore confirms that the
908    /// request was accepted, not that it took effect.
909    ///
910    /// Watch for the outcome through the `session.model_change` event on
911    /// success, or the ephemeral `session.auto_tier_switch_failed` event on
912    /// failure. You can also read the current committed and in-flight state at
913    /// any time through `session.rpc().model().get_current()`.
914    ///
915    /// Only the most recent request survives: issuing a new request replaces any
916    /// earlier one that has not yet been claimed by a turn.
917    ///
918    /// Pass `None` to return to the provider's default Auto routing.
919    ///
920    /// **Experimental.** Part of an experimental Auto routing surface and may
921    /// change or be removed in a future release.
922    ///
923    /// # Cancel safety
924    ///
925    /// **Cancel-safe.** Single `session.model.switchAutoTier` RPC; the
926    /// underlying [`Client::call`](crate::Client::call) is cancel-safe via the
927    /// writer-actor.
928    ///
929    /// [`ModelSwitchAutoTierStatus::Pending`]: crate::generated::api_types::ModelSwitchAutoTierStatus::Pending
930    pub async fn set_auto_tier(
931        &self,
932        auto_tier: Option<AutoTier>,
933    ) -> Result<ModelSwitchAutoTierResult, Error> {
934        self.rpc()
935            .model()
936            .switch_auto_tier(ModelSwitchAutoTierRequest {
937                auto_tier,
938                source: None,
939            })
940            .await
941    }
942
943    /// Disconnect this session from the CLI.
944    ///
945    /// Sends the `session.detach` RPC, stops the event loop, and unregisters
946    /// the session from the client. **Session state on disk** (conversation
947    /// history, planning state, artifacts) is **preserved**, so the
948    /// conversation can be resumed later via [`Client::resume_session`]
949    /// using this session's ID. To permanently remove all on-disk session
950    /// data, use [`Client::delete_session`] instead.
951    ///
952    /// The caller should ensure the session is idle (e.g. [`send_and_wait`]
953    /// has returned) before disconnecting; in-flight tool or event handlers
954    /// may otherwise observe failures.
955    ///
956    /// [`Client::resume_session`]: crate::Client::resume_session
957    /// [`Client::delete_session`]: crate::Client::delete_session
958    /// [`send_and_wait`]: Self::send_and_wait
959    pub async fn disconnect(&self) -> Result<(), Error> {
960        self.client.detach_session(&self.id).await?;
961        self.external_tools_shutdown.cancel();
962        self.stop_event_loop().await;
963        self.github_token_registration.lock().take();
964        self.client
965            .unregister_session_owned(&self.id, self.registration_token);
966        Ok(())
967    }
968
969    /// Deprecated alias for [`disconnect`](Self::disconnect).
970    /// Prefer `disconnect` in new code.
971    #[deprecated(since = "0.1.0", note = "Use `disconnect()` instead")]
972    pub async fn destroy(&self) -> Result<(), Error> {
973        self.disconnect().await
974    }
975
976    /// Write a log message to the session.
977    ///
978    /// Pass `None` for `opts` to use defaults (info level, persisted).
979    pub async fn log(
980        &self,
981        message: &str,
982        opts: Option<crate::types::LogOptions>,
983    ) -> Result<(), Error> {
984        let opts = opts.unwrap_or_default();
985        let level = match opts.level {
986            Some(level) => Some(serde_json::from_value(serde_json::to_value(level)?)?),
987            None => None,
988        };
989        let request = LogRequest {
990            message: message.to_string(),
991            level,
992            ephemeral: opts.ephemeral,
993            r#type: None,
994            tip: None,
995            url: None,
996        };
997        self.rpc().log(request).await?;
998        Ok(())
999    }
1000
1001    /// Returns the UI sub-API for elicitation, confirmation, selection, and
1002    /// free-form input.
1003    ///
1004    /// All UI methods route through `session.ui.*` RPCs and require host
1005    /// support — check `session.capabilities().ui.elicitation` before use.
1006    pub fn ui(&self) -> SessionUi<'_> {
1007        SessionUi { session: self }
1008    }
1009
1010    /// Returns an error if the host doesn't support elicitation.
1011    fn assert_elicitation(&self) -> Result<(), Error> {
1012        if self
1013            .capabilities
1014            .read()
1015            .ui
1016            .as_ref()
1017            .and_then(|u| u.elicitation)
1018            != Some(true)
1019        {
1020            return Err(ErrorKind::Session(SessionErrorKind::ElicitationNotSupported).into());
1021        }
1022        Ok(())
1023    }
1024}
1025
1026impl Drop for Session {
1027    fn drop(&mut self) {
1028        // Cooperative shutdown: cancel the event loop's token to signal
1029        // exit between iterations. The loop will see the cancellation on
1030        // its next select poll and break cleanly. We do NOT abort the
1031        // JoinHandle — that would land at any await point in the loop body,
1032        // potentially leaving the CLI with an unanswered request id.
1033        // RFD-400 review finding #3. Requests already dispatched to their
1034        // own tasks are not tracked here and may outlive the session.
1035        //
1036        // The handle itself is left in `event_loop` to be reaped by the
1037        // tokio runtime when it next polls; we intentionally don't await
1038        // it here because Drop is sync.
1039        self.shutdown.cancel();
1040        self.external_tools_shutdown.cancel();
1041        self.github_token_registration.lock().take();
1042        self.client
1043            .unregister_session_owned(&self.id, self.registration_token);
1044    }
1045}
1046
1047/// UI sub-API for a [`Session`] — elicitation, confirmation, selection,
1048/// and free-form input.
1049///
1050/// Acquired via [`Session::ui`]. Methods route to `session.ui.*` RPCs and
1051/// require host elicitation support — check
1052/// `session.capabilities().ui.elicitation` before use.
1053pub struct SessionUi<'a> {
1054    session: &'a Session,
1055}
1056
1057impl<'a> SessionUi<'a> {
1058    /// Request user input via an interactive UI form (elicitation).
1059    ///
1060    /// Sends a JSON Schema describing form fields to the CLI host. The host
1061    /// renders a form dialog and returns the user's response.
1062    ///
1063    /// Prefer the typed convenience methods [`confirm`](Self::confirm),
1064    /// [`select`](Self::select), and [`input`](Self::input) for common cases.
1065    pub async fn elicitation(
1066        &self,
1067        message: &str,
1068        schema: Value,
1069    ) -> Result<ElicitationResult, Error> {
1070        self.session.assert_elicitation()?;
1071        let result = self
1072            .session
1073            .client
1074            .call(
1075                "session.ui.elicitation",
1076                Some(serde_json::json!({
1077                    "sessionId": self.session.id,
1078                    "message": message,
1079                    "requestedSchema": schema,
1080                })),
1081            )
1082            .await?;
1083        let elicitation: ElicitationResult = serde_json::from_value(result)?;
1084        Ok(elicitation)
1085    }
1086
1087    /// Ask the user a yes/no confirmation question.
1088    ///
1089    /// Returns `true` if the user accepted and confirmed, `false` otherwise.
1090    pub async fn confirm(&self, message: &str) -> Result<bool, Error> {
1091        self.session.assert_elicitation()?;
1092        let schema = serde_json::json!({
1093            "type": "object",
1094            "properties": {
1095                "confirmed": {
1096                    "type": "boolean",
1097                    "default": true,
1098                }
1099            },
1100            "required": ["confirmed"]
1101        });
1102        let result = self.elicitation(message, schema).await?;
1103        Ok(result.action == "accept"
1104            && result
1105                .content
1106                .and_then(|c| c.get("confirmed").and_then(|v| v.as_bool()))
1107                == Some(true))
1108    }
1109
1110    /// Ask the user to select from a list of options.
1111    ///
1112    /// Returns the selected option string on accept, or `None` on decline/cancel.
1113    pub async fn select(&self, message: &str, options: &[&str]) -> Result<Option<String>, Error> {
1114        self.session.assert_elicitation()?;
1115        let schema = serde_json::json!({
1116            "type": "object",
1117            "properties": {
1118                "selection": {
1119                    "type": "string",
1120                    "enum": options,
1121                }
1122            },
1123            "required": ["selection"]
1124        });
1125        let result = self.elicitation(message, schema).await?;
1126        if result.action != "accept" {
1127            return Ok(None);
1128        }
1129        let selection = result.content.and_then(|c| {
1130            c.get("selection")
1131                .and_then(|v| v.as_str())
1132                .map(String::from)
1133        });
1134        Ok(selection)
1135    }
1136
1137    /// Ask the user for free-form text input.
1138    ///
1139    /// Returns the input string on accept, or `None` on decline/cancel.
1140    /// Use [`UiInputOptions`] to set validation constraints and field metadata.
1141    pub async fn input(
1142        &self,
1143        message: &str,
1144        options: Option<&UiInputOptions<'_>>,
1145    ) -> Result<Option<String>, Error> {
1146        self.session.assert_elicitation()?;
1147        let mut field = serde_json::json!({ "type": "string" });
1148        if let Some(opts) = options {
1149            if let Some(title) = opts.title {
1150                field["title"] = Value::String(title.to_string());
1151            }
1152            if let Some(desc) = opts.description {
1153                field["description"] = Value::String(desc.to_string());
1154            }
1155            if let Some(min) = opts.min_length {
1156                field["minLength"] = Value::Number(min.into());
1157            }
1158            if let Some(max) = opts.max_length {
1159                field["maxLength"] = Value::Number(max.into());
1160            }
1161            if let Some(fmt) = &opts.format {
1162                field["format"] = Value::String(fmt.as_str().to_string());
1163            }
1164            if let Some(default) = opts.default {
1165                field["default"] = Value::String(default.to_string());
1166            }
1167        }
1168        let schema = serde_json::json!({
1169            "type": "object",
1170            "properties": { "value": field },
1171            "required": ["value"]
1172        });
1173        let result = self.elicitation(message, schema).await?;
1174        if result.action != "accept" {
1175            return Ok(None);
1176        }
1177        let value = result
1178            .content
1179            .and_then(|c| c.get("value").and_then(|v| v.as_str()).map(String::from));
1180        Ok(value)
1181    }
1182}
1183
1184impl Client {
1185    /// Prepare a new session without touching the transport.
1186    ///
1187    /// Returns a [`PreparedSession`] that owns the session's event broadcast
1188    /// channel, so callers can install an
1189    /// [`EventSubscription`](crate::subscription::EventSubscription) via
1190    /// [`PreparedSession::subscribe`] *before* any protocol activity starts.
1191    /// Call [`PreparedSession::start`] to actually create the session.
1192    ///
1193    /// This is the loss-free entry point for consumers that must observe
1194    /// every *routed* event a session emits, including events the CLI emits
1195    /// while `session.create` is still in flight and ephemeral events (such
1196    /// as `session.idle`) that cannot be recovered from
1197    /// [`Session::get_messages`]. [`create_session`](Self::create_session)
1198    /// is a thin wrapper over `prepare_session(...)?.start()` and cannot
1199    /// offer the same guarantee, because the subscription can only be
1200    /// installed after the returned `Session` exists.
1201    ///
1202    /// Routing requires a known session ID. When the server assigns the ID,
1203    /// the SDK cannot register the session on its notification router until
1204    /// the `session.create` response arrives, so notifications emitted
1205    /// before that point are not routable and stay unobservable. Pin
1206    /// [`SessionConfig::session_id`](crate::types::SessionConfig::session_id)
1207    /// for complete pre-response coverage — see the "Server-assigned session
1208    /// IDs" section on [`PreparedSession`].
1209    ///
1210    /// # Inertness
1211    ///
1212    /// `prepare_session` performs no router registration, spawns no task,
1213    /// and writes nothing to the wire. It only validates
1214    /// [`event_buffer_capacity`](SessionConfig::event_buffer_capacity),
1215    /// allocates a local broadcast channel and cancellation token, and
1216    /// stores the config. Dropping the returned handle without starting it
1217    /// leaves no client-side or server-side state behind and closes every
1218    /// subscription taken from it.
1219    ///
1220    /// # Errors
1221    ///
1222    /// Returns [`ErrorKind::InvalidConfig`] if
1223    /// [`event_buffer_capacity`](SessionConfig::event_buffer_capacity) is
1224    /// `Some(0)`. All other configuration and protocol errors surface from
1225    /// [`PreparedSession::start`], with the same
1226    /// [`ErrorKind`]s [`create_session`](Self::create_session) has always
1227    /// returned.
1228    ///
1229    /// # Example
1230    ///
1231    /// ```no_run
1232    /// # use github_copilot_sdk::{Client, SessionConfig};
1233    /// # async fn example(client: Client) -> Result<(), github_copilot_sdk::Error> {
1234    /// let prepared = client.prepare_session(SessionConfig::default())?;
1235    /// let mut events = prepared.subscribe();
1236    /// let drain = tokio::spawn(async move {
1237    ///     while let Ok(event) = events.recv().await {
1238    ///         println!("{}", event.event_type);
1239    ///     }
1240    /// });
1241    /// let session = prepared.start().await?;
1242    /// # let _ = (session, drain);
1243    /// # Ok(())
1244    /// # }
1245    /// ```
1246    pub fn prepare_session(&self, config: SessionConfig) -> Result<PreparedSession, Error> {
1247        let capacity = resolve_event_buffer_capacity(config.event_buffer_capacity)?;
1248        Ok(PreparedSession::new(
1249            self.clone(),
1250            PreparedKind::Create(Box::new(config)),
1251            capacity,
1252        ))
1253    }
1254
1255    /// Prepare a session resume without touching the transport.
1256    ///
1257    /// The resume counterpart of [`prepare_session`](Self::prepare_session);
1258    /// see that method for the inertness guarantee, error semantics, and
1259    /// rationale. Particularly relevant on resume with
1260    /// [`continue_pending_work`](ResumeSessionConfig::continue_pending_work),
1261    /// where the runtime can start emitting events (and reach
1262    /// `session.idle`) while `session.resume` is still in flight.
1263    ///
1264    /// # Errors
1265    ///
1266    /// Returns [`ErrorKind::InvalidConfig`] if
1267    /// [`event_buffer_capacity`](ResumeSessionConfig::event_buffer_capacity)
1268    /// is `Some(0)`.
1269    pub fn prepare_resume_session(
1270        &self,
1271        config: ResumeSessionConfig,
1272    ) -> Result<PreparedSession, Error> {
1273        let capacity = resolve_event_buffer_capacity(config.event_buffer_capacity)?;
1274        Ok(PreparedSession::new(
1275            self.clone(),
1276            PreparedKind::Resume(Box::new(config)),
1277            capacity,
1278        ))
1279    }
1280
1281    /// Create a new session on the CLI.
1282    ///
1283    /// Sends `session.create`, registers the session on the router,
1284    /// and spawns an internal event loop that dispatches to the handler.
1285    ///
1286    /// All callbacks (per-event handlers, tool handlers, hooks, transform)
1287    /// are configured via [`SessionConfig`] using its `with_*_handler` /
1288    /// `with_tools` / `with_hooks` / `with_system_message_transform` builder
1289    /// methods.
1290    ///
1291    /// If [`hooks_handler`](SessionConfig::hooks_handler) is set, the
1292    /// wire-level `hooks` flag is automatically enabled.
1293    ///
1294    /// If [`system_message_transform`](SessionConfig::system_message_transform) is set, the SDK injects
1295    /// `action: "transform"` sections into the [`SystemMessageConfig`] wire
1296    /// format and handles `systemMessage.transform` RPC callbacks during
1297    /// the session.
1298    ///
1299    /// Each per-event handler is independently optional. If a handler is
1300    /// not installed, the SDK signals the runtime not to emit the matching
1301    /// broadcast (and silently skips dispatch if one arrives anyway).
1302    ///
1303    /// # Event delivery
1304    ///
1305    /// Equivalent to `prepare_session(config)?.start().await`. Because the
1306    /// first subscription can only be taken from the returned [`Session`],
1307    /// events the runtime emits before this call returns are broadcast with
1308    /// no receiver installed and are therefore not delivered to
1309    /// [`Session::subscribe`]. Use
1310    /// [`prepare_session`](Self::prepare_session) when startup events
1311    /// matter.
1312    pub async fn create_session(&self, config: SessionConfig) -> Result<Session, Error> {
1313        self.prepare_session(config)?.start().await
1314    }
1315
1316    /// Resume an existing session on the CLI.
1317    ///
1318    /// Sends `session.resume` and `session.skills.reload`, registers the
1319    /// session on the router, and spawns the event loop.
1320    ///
1321    /// All callbacks (event handler, hooks, transform) are configured
1322    /// via [`ResumeSessionConfig`] using its `with_*` builder methods.
1323    ///
1324    /// See [`Self::create_session`] for the defaults applied when callback
1325    /// fields are unset.
1326    ///
1327    /// # Event delivery
1328    ///
1329    /// Equivalent to `prepare_resume_session(config)?.start().await`, and
1330    /// carries the same startup-event caveat documented on
1331    /// [`create_session`](Self::create_session). Use
1332    /// [`prepare_resume_session`](Self::prepare_resume_session) when
1333    /// startup events matter.
1334    pub async fn resume_session(&self, config: ResumeSessionConfig) -> Result<Session, Error> {
1335        self.prepare_resume_session(config)?.start().await
1336    }
1337
1338    async fn start_prepared_create(
1339        &self,
1340        mut config: SessionConfig,
1341        event_tx: tokio::sync::broadcast::Sender<SessionEvent>,
1342        shutdown: CancellationToken,
1343    ) -> Result<Session, Error> {
1344        let total_start = Instant::now();
1345        // For cloud sessions, let the CLI/server assign the session id and
1346        // register the session lazily once the response arrives. For non-cloud
1347        // sessions we generate the id client-side (when the caller didn't
1348        // supply one) so the session can be registered BEFORE the RPC — the
1349        // CLI may issue session-scoped requests (e.g. sessionFs.writeFile for
1350        // workspace metadata) during session.create processing, before it has
1351        // sent the response.
1352        let caller_session_id = config.session_id.clone();
1353        let use_server_generated_id = config.cloud.is_some() && caller_session_id.is_none();
1354        let local_session_id: Option<SessionId> = if use_server_generated_id {
1355            None
1356        } else {
1357            Some(
1358                caller_session_id
1359                    .clone()
1360                    .unwrap_or_else(|| SessionId::new(uuid::Uuid::new_v4().to_string())),
1361            )
1362        };
1363        if config.hooks_handler.is_some() && config.hooks.is_none() {
1364            config.hooks = Some(true);
1365        }
1366        if let Some(transforms) = config.system_message_transform.clone() {
1367            inject_transform_sections(&mut config, transforms.as_ref());
1368        }
1369        let mode = self.inner.mode;
1370        if mode == crate::ClientMode::Empty && config.available_tools.is_none() {
1371            return Err(Error::with_message(
1372                ErrorKind::InvalidConfig,
1373                "ClientMode::Empty requires available_tools to be set on the session config. \
1374                 Use ToolSet to specify which tools the session may use (e.g. \
1375                 ToolSet::new().add_builtin_many(BUILTIN_TOOLS_ISOLATED)).",
1376            ));
1377        }
1378        crate::mode::validate_tool_filter_list(
1379            "available_tools",
1380            config.available_tools.as_deref(),
1381        )?;
1382        crate::mode::validate_tool_filter_list("excluded_tools", config.excluded_tools.as_deref())?;
1383        config.system_message =
1384            crate::mode::system_message_for_mode(mode, config.system_message.take());
1385        config.memory = crate::mode::memory_for_mode(mode, config.memory.take());
1386        config.enable_experimental_mode =
1387            crate::mode::experimental_mode_for_mode(mode, config.enable_experimental_mode);
1388        if mode == crate::ClientMode::Empty {
1389            if config.enable_session_telemetry.is_none() {
1390                config.enable_session_telemetry = Some(false);
1391            }
1392            if config.skip_embedding_retrieval.is_none() {
1393                config.skip_embedding_retrieval = Some(true);
1394            }
1395            if config.enable_on_demand_instruction_discovery.is_none() {
1396                config.enable_on_demand_instruction_discovery = Some(false);
1397            }
1398            if config.enable_file_hooks.is_none() {
1399                config.enable_file_hooks = Some(false);
1400            }
1401            if config.enable_host_git_operations.is_none() {
1402                config.enable_host_git_operations = Some(false);
1403            }
1404            if config.enable_session_store.is_none() {
1405                config.enable_session_store = Some(false);
1406            }
1407            if config.enable_skills.is_none() {
1408                config.enable_skills = Some(false);
1409            }
1410        }
1411        if mode == crate::ClientMode::Empty && config.mcp_oauth_token_storage.is_none() {
1412            config.mcp_oauth_token_storage = Some("in-memory".into());
1413        }
1414        if mode == crate::ClientMode::Empty && config.embedding_cache_storage.is_none() {
1415            config.embedding_cache_storage = Some("in-memory".into());
1416        }
1417        config.custom_agents_local_only =
1418            crate::mode::resolve_custom_agents_local_only(mode, config.custom_agents_local_only);
1419        let opt_skip_custom_instructions = config.skip_custom_instructions;
1420        let opt_custom_agents_local_only = config.custom_agents_local_only;
1421        let opt_coauthor_enabled = config.coauthor_enabled;
1422        let opt_manage_schedule_enabled = config.manage_schedule_enabled;
1423        let opt_included_builtin_skills = config.included_builtin_skills.take();
1424        let (mut wire, mut runtime) = config.into_wire(local_session_id.clone())?;
1425        wire.enable_github_telemetry_forwarding =
1426            self.inner.on_github_telemetry.is_some().then_some(true);
1427
1428        let permission_handler = crate::permission::resolve_handler(
1429            runtime.permission_handler.take(),
1430            runtime.permission_policy.take(),
1431        );
1432        let handlers = SessionHandlers {
1433            permission: permission_handler,
1434            managed_settings_enabled: has_managed_settings(
1435                wire.enable_managed_settings,
1436                wire.managed_settings.as_ref(),
1437            ),
1438            elicitation: runtime.elicitation_handler.take(),
1439            mcp_auth: runtime.mcp_auth_handler.take(),
1440            user_input: runtime.user_input_handler.take(),
1441            exit_plan_mode: runtime.exit_plan_mode_handler.take(),
1442            auto_mode_switch: runtime.auto_mode_switch_handler.take(),
1443            tools: Arc::new(std::mem::take(&mut runtime.tool_handlers)),
1444        };
1445        let hooks = runtime.hooks_handler.take();
1446        let transforms = runtime.system_message_transform.take();
1447        let tools_count = wire.tools.as_ref().map_or(0, Vec::len);
1448        let commands_count = runtime.commands.as_ref().map_or(0, Vec::len);
1449        let has_hooks = hooks.is_some();
1450        let command_handlers = build_command_handler_map(runtime.commands.as_deref());
1451        let canvas_handler = runtime.canvas_handler.take();
1452        let session_fs_provider = runtime.session_fs_provider.take();
1453        let bearer_token_providers = std::mem::take(&mut runtime.bearer_token_providers);
1454        let github_token_registration = runtime
1455            .github_token_provider
1456            .take()
1457            .map(|provider| self.register_github_token_provider(provider));
1458        wire.github_token_provider_registration_id = github_token_registration
1459            .as_ref()
1460            .map(|registration| registration.id().to_string());
1461        let has_mcp_auth_handler = handlers.mcp_auth.is_some();
1462        if self.inner.session_fs_configured && session_fs_provider.is_none() {
1463            return Err(ErrorKind::Session(SessionErrorKind::SessionFsProviderRequired).into());
1464        }
1465        if self.inner.session_fs_sqlite_declared
1466            && let Some(ref provider) = session_fs_provider
1467            && provider.sqlite().is_none()
1468        {
1469            return Err(Error::with_message(
1470                ErrorKind::InvalidConfig,
1471                "SessionFs capabilities declare SQLite support but the provider \
1472                 does not implement SessionFsSqliteProvider",
1473            ));
1474        }
1475
1476        let mut params = serde_json::to_value(&wire)?;
1477        let trace_ctx = self.resolve_trace_context().await;
1478        inject_trace_context(&mut params, &trace_ctx);
1479
1480        let setup_start = Instant::now();
1481        let capabilities = Arc::new(parking_lot::RwLock::new(SessionCapabilities::default()));
1482        let idle_waiter = Arc::new(ParkingLotMutex::new(None));
1483        let open_canvases = Arc::new(parking_lot::RwLock::new(Vec::new()));
1484        let external_tools_shutdown = self.inner.rpc.connection_closed_token();
1485
1486        // For cloud sessions (use_server_generated_id), defer session
1487        // registration to the inline callback so the read task registers
1488        // the session synchronously the instant the response arrives.
1489        // For non-cloud sessions, register up-front so the CLI can issue
1490        // session-scoped requests during session.create processing.
1491        let inline_stash: Arc<
1492            ParkingLotMutex<Option<(SessionId, crate::router::SessionRegistration)>>,
1493        > = Arc::new(ParkingLotMutex::new(None));
1494
1495        let inline_callback: Option<crate::jsonrpc::InlineResponseCallback> = if let Some(ref sid) =
1496            local_session_id
1497        {
1498            let channels = self.register_session(sid);
1499            *inline_stash.lock() = Some((sid.clone(), channels));
1500            None
1501        } else {
1502            let client = self.clone();
1503            let stash = inline_stash.clone();
1504            let expected = caller_session_id.clone();
1505            Some(Box::new(move |response| {
1506                let result = response.result.as_ref().ok_or_else(|| {
1507                    Error::with_message(ErrorKind::Json, "session.create response had no result")
1508                })?;
1509                let parsed: CreateSessionResult =
1510                    serde_json::from_value(result.clone()).map_err(Error::from)?;
1511                if let Some(requested) = expected.as_ref()
1512                    && parsed.session_id != *requested
1513                {
1514                    return Err(ErrorKind::Session(SessionErrorKind::SessionIdMismatch {
1515                        requested: requested.clone(),
1516                        returned: parsed.session_id,
1517                    })
1518                    .into());
1519                }
1520                // Register and stash under a single stash-lock hold. The
1521                // cancellation guard identifies the session to unregister by
1522                // peeking this stash, so registering outside the lock would
1523                // leave a window where a concurrent guard drop (caller
1524                // cancellation) sees `None` and leaks the registration.
1525                // `register_session` takes the router lock, never the stash
1526                // lock, so there is no lock-order inversion here.
1527                let mut stashed = stash.lock();
1528                let registration = client.register_session(&parsed.session_id);
1529                *stashed = Some((parsed.session_id, registration));
1530                Ok(())
1531            }))
1532        };
1533
1534        // Armed for the whole startup sequence: any early return, and any
1535        // drop of this future (caller cancellation), cancels the session
1536        // token and unregisters whatever was registered on the router. For
1537        // the cloud path the ID is only known once the inline callback has
1538        // run, so the guard reads the stash at cleanup time.
1539        let mut pending_registration = match local_session_id {
1540            Some(ref sid) => {
1541                let token = inline_stash
1542                    .lock()
1543                    .as_ref()
1544                    .expect("session registration must exist")
1545                    .1
1546                    .token;
1547                PendingSessionRegistration::new(
1548                    self.clone(),
1549                    sid.clone(),
1550                    token,
1551                    shutdown.clone(),
1552                    external_tools_shutdown.clone(),
1553                )
1554            }
1555            None => PendingSessionRegistration::deferred(
1556                self.clone(),
1557                inline_stash.clone(),
1558                shutdown.clone(),
1559                external_tools_shutdown.clone(),
1560            ),
1561        };
1562
1563        let rpc_start = Instant::now();
1564        let result = self
1565            .call_with_inline_callback("session.create", Some(params), inline_callback)
1566            .await?;
1567        tracing::debug!(
1568            elapsed_ms = rpc_start.elapsed().as_millis(),
1569            "Client::create_session session creation request completed successfully"
1570        );
1571        let create_result: CreateSessionResult = serde_json::from_value(result)?;
1572
1573        if let Some(ref requested) = local_session_id
1574            && create_result.session_id != *requested
1575        {
1576            return Err(ErrorKind::Session(SessionErrorKind::SessionIdMismatch {
1577                requested: requested.clone(),
1578                returned: create_result.session_id.clone(),
1579            })
1580            .into());
1581        }
1582
1583        let (session_id, registration) = inline_stash
1584            .lock()
1585            .take()
1586            .expect("session registration must have populated stash on success");
1587        let channels = registration.channels;
1588        let registration_token = registration.token;
1589        pending_registration.resolve_to(session_id.clone(), registration_token);
1590        let event_loop = spawn_event_loop(
1591            session_id.clone(),
1592            self.clone(),
1593            handlers,
1594            hooks,
1595            transforms,
1596            command_handlers,
1597            canvas_handler,
1598            session_fs_provider,
1599            bearer_token_providers,
1600            channels,
1601            idle_waiter.clone(),
1602            capabilities.clone(),
1603            open_canvases.clone(),
1604            event_tx.clone(),
1605            shutdown.clone(),
1606            external_tools_shutdown.clone(),
1607        );
1608        tracing::debug!(
1609            elapsed_ms = setup_start.elapsed().as_millis(),
1610            session_id = %session_id,
1611            tools_count,
1612            commands_count,
1613            has_hooks,
1614            "Client::create_session local setup complete"
1615        );
1616        *capabilities.write() = create_result.capabilities.unwrap_or_default();
1617        if has_mcp_auth_handler
1618            && let Err(error) = register_mcp_auth_interest(self, &session_id).await
1619        {
1620            pending_registration.cleanup(event_loop).await;
1621            return Err(error);
1622        }
1623
1624        tracing::debug!(
1625            elapsed_ms = total_start.elapsed().as_millis(),
1626            session_id = %session_id,
1627            "Client::create_session complete"
1628        );
1629        pending_registration.disarm();
1630        let session = Session {
1631            id: session_id,
1632            cwd: self.cwd().clone(),
1633            workspace_path: create_result.workspace_path,
1634            remote_url: create_result.remote_url,
1635            client: self.clone(),
1636            event_loop: ParkingLotMutex::new(Some(event_loop)),
1637            shutdown,
1638            external_tools_shutdown,
1639            idle_waiter,
1640            capabilities,
1641            open_canvases,
1642            event_tx,
1643            github_token_registration: ParkingLotMutex::new(github_token_registration),
1644            registration_token,
1645        };
1646        apply_mode_post_create_patch(
1647            &session,
1648            mode,
1649            opt_skip_custom_instructions,
1650            opt_custom_agents_local_only,
1651            opt_coauthor_enabled,
1652            opt_manage_schedule_enabled,
1653            opt_included_builtin_skills,
1654        )
1655        .await?;
1656        if let Some(registration) = session.github_token_registration.lock().as_ref() {
1657            registration.claim(session.id.clone());
1658        } else {
1659            self.retire_github_token_provider(&session.id);
1660        }
1661        Ok(session)
1662    }
1663
1664    /// Resume an existing session on the CLI.
1665    ///
1666    /// Sends `session.resume` and `session.skills.reload`, registers the
1667    /// session on the router, and spawns the event loop.
1668    ///
1669    /// All callbacks (event handler, hooks, transform) are configured
1670    /// via [`ResumeSessionConfig`] using its `with_*` builder methods.
1671    ///
1672    /// See [`Self::create_session`] for the defaults applied when callback
1673    /// fields are unset.
1674    async fn start_prepared_resume(
1675        &self,
1676        mut config: ResumeSessionConfig,
1677        event_tx: tokio::sync::broadcast::Sender<SessionEvent>,
1678        shutdown: CancellationToken,
1679    ) -> Result<Session, Error> {
1680        let total_start = Instant::now();
1681        let session_id = config.session_id.clone();
1682        if config.hooks_handler.is_some() && config.hooks.is_none() {
1683            config.hooks = Some(true);
1684        }
1685        if let Some(transforms) = config.system_message_transform.clone() {
1686            inject_transform_sections_resume(&mut config, transforms.as_ref());
1687        }
1688        let mode = self.inner.mode;
1689        if mode == crate::ClientMode::Empty && config.available_tools.is_none() {
1690            return Err(Error::with_message(
1691                ErrorKind::InvalidConfig,
1692                "ClientMode::Empty requires available_tools to be set on the session config. \
1693                 Use ToolSet to specify which tools the session may use (e.g. \
1694                 ToolSet::new().add_builtin_many(BUILTIN_TOOLS_ISOLATED)).",
1695            ));
1696        }
1697        crate::mode::validate_tool_filter_list(
1698            "available_tools",
1699            config.available_tools.as_deref(),
1700        )?;
1701        crate::mode::validate_tool_filter_list("excluded_tools", config.excluded_tools.as_deref())?;
1702        config.system_message =
1703            crate::mode::system_message_for_mode(mode, config.system_message.take());
1704        config.memory = crate::mode::memory_for_mode(mode, config.memory.take());
1705        config.enable_experimental_mode =
1706            crate::mode::experimental_mode_for_mode(mode, config.enable_experimental_mode);
1707        if mode == crate::ClientMode::Empty {
1708            if config.enable_session_telemetry.is_none() {
1709                config.enable_session_telemetry = Some(false);
1710            }
1711            if config.skip_embedding_retrieval.is_none() {
1712                config.skip_embedding_retrieval = Some(true);
1713            }
1714            if config.enable_on_demand_instruction_discovery.is_none() {
1715                config.enable_on_demand_instruction_discovery = Some(false);
1716            }
1717            if config.enable_file_hooks.is_none() {
1718                config.enable_file_hooks = Some(false);
1719            }
1720            if config.enable_host_git_operations.is_none() {
1721                config.enable_host_git_operations = Some(false);
1722            }
1723            if config.enable_session_store.is_none() {
1724                config.enable_session_store = Some(false);
1725            }
1726            if config.enable_skills.is_none() {
1727                config.enable_skills = Some(false);
1728            }
1729        }
1730        if mode == crate::ClientMode::Empty && config.mcp_oauth_token_storage.is_none() {
1731            config.mcp_oauth_token_storage = Some("in-memory".into());
1732        }
1733        if mode == crate::ClientMode::Empty && config.embedding_cache_storage.is_none() {
1734            config.embedding_cache_storage = Some("in-memory".into());
1735        }
1736        config.custom_agents_local_only =
1737            crate::mode::resolve_custom_agents_local_only(mode, config.custom_agents_local_only);
1738        let opt_skip_custom_instructions = config.skip_custom_instructions;
1739        let opt_custom_agents_local_only = config.custom_agents_local_only;
1740        let opt_coauthor_enabled = config.coauthor_enabled;
1741        let opt_manage_schedule_enabled = config.manage_schedule_enabled;
1742        let opt_included_builtin_skills = config.included_builtin_skills.take();
1743        let (mut wire, mut runtime) = config.into_wire()?;
1744        wire.enable_github_telemetry_forwarding =
1745            self.inner.on_github_telemetry.is_some().then_some(true);
1746
1747        let permission_handler = crate::permission::resolve_handler(
1748            runtime.permission_handler.take(),
1749            runtime.permission_policy.take(),
1750        );
1751        let handlers = SessionHandlers {
1752            permission: permission_handler,
1753            managed_settings_enabled: has_managed_settings(
1754                wire.enable_managed_settings,
1755                wire.managed_settings.as_ref(),
1756            ),
1757            elicitation: runtime.elicitation_handler.take(),
1758            mcp_auth: runtime.mcp_auth_handler.take(),
1759            user_input: runtime.user_input_handler.take(),
1760            exit_plan_mode: runtime.exit_plan_mode_handler.take(),
1761            auto_mode_switch: runtime.auto_mode_switch_handler.take(),
1762            tools: Arc::new(std::mem::take(&mut runtime.tool_handlers)),
1763        };
1764        let hooks = runtime.hooks_handler.take();
1765        let transforms = runtime.system_message_transform.take();
1766        let tools_count = wire.tools.as_ref().map_or(0, Vec::len);
1767        let commands_count = runtime.commands.as_ref().map_or(0, Vec::len);
1768        let has_hooks = hooks.is_some();
1769        let command_handlers = build_command_handler_map(runtime.commands.as_deref());
1770        let canvas_handler = runtime.canvas_handler.take();
1771        let session_fs_provider = runtime.session_fs_provider.take();
1772        let bearer_token_providers = std::mem::take(&mut runtime.bearer_token_providers);
1773        let github_token_registration = runtime
1774            .github_token_provider
1775            .take()
1776            .map(|provider| self.register_github_token_provider(provider));
1777        wire.github_token_provider_registration_id = github_token_registration
1778            .as_ref()
1779            .map(|registration| registration.id().to_string());
1780        let has_mcp_auth_handler = handlers.mcp_auth.is_some();
1781        if self.inner.session_fs_configured && session_fs_provider.is_none() {
1782            return Err(ErrorKind::Session(SessionErrorKind::SessionFsProviderRequired).into());
1783        }
1784        if self.inner.session_fs_sqlite_declared
1785            && let Some(ref provider) = session_fs_provider
1786            && provider.sqlite().is_none()
1787        {
1788            return Err(Error::with_message(
1789                ErrorKind::InvalidConfig,
1790                "SessionFs capabilities declare SQLite support but the provider \
1791                 does not implement SessionFsSqliteProvider",
1792            ));
1793        }
1794
1795        let mut params = serde_json::to_value(&wire)?;
1796        let trace_ctx = self.resolve_trace_context().await;
1797        inject_trace_context(&mut params, &trace_ctx);
1798
1799        let capabilities = Arc::new(parking_lot::RwLock::new(SessionCapabilities::default()));
1800        let setup_start = Instant::now();
1801        let registration = self.register_session(&session_id);
1802        let registration_token = registration.token;
1803        let channels = registration.channels;
1804        let idle_waiter = Arc::new(ParkingLotMutex::new(None));
1805        let open_canvases = Arc::new(parking_lot::RwLock::new(Vec::new()));
1806        let external_tools_shutdown = self.inner.rpc.connection_closed_token();
1807        let event_loop = spawn_event_loop(
1808            session_id.clone(),
1809            self.clone(),
1810            handlers,
1811            hooks,
1812            transforms,
1813            command_handlers,
1814            canvas_handler,
1815            session_fs_provider,
1816            bearer_token_providers,
1817            channels,
1818            idle_waiter.clone(),
1819            capabilities.clone(),
1820            open_canvases.clone(),
1821            event_tx.clone(),
1822            shutdown.clone(),
1823            external_tools_shutdown.clone(),
1824        );
1825        let mut registration = PendingSessionRegistration::new(
1826            self.clone(),
1827            session_id.clone(),
1828            registration_token,
1829            shutdown.clone(),
1830            external_tools_shutdown.clone(),
1831        );
1832        tracing::debug!(
1833            elapsed_ms = setup_start.elapsed().as_millis(),
1834            session_id = %session_id,
1835            tools_count,
1836            commands_count,
1837            has_hooks,
1838            "Client::resume_session local setup complete"
1839        );
1840
1841        let rpc_start = Instant::now();
1842        let result = match self.call("session.resume", Some(params)).await {
1843            Ok(result) => result,
1844            Err(error) => {
1845                registration.cleanup(event_loop).await;
1846                return Err(error);
1847            }
1848        };
1849        tracing::debug!(
1850            elapsed_ms = rpc_start.elapsed().as_millis(),
1851            session_id = %session_id,
1852            "Client::resume_session session resume request completed successfully"
1853        );
1854
1855        let resume_result: ResumeSessionResult = match serde_json::from_value(result) {
1856            Ok(result) => result,
1857            Err(error) => {
1858                registration.cleanup(event_loop).await;
1859                return Err(error.into());
1860            }
1861        };
1862        let cli_session_id = resume_result
1863            .session_id
1864            .clone()
1865            .unwrap_or_else(|| session_id.clone());
1866        if cli_session_id != session_id {
1867            registration.cleanup(event_loop).await;
1868            return Err(ErrorKind::Session(SessionErrorKind::SessionIdMismatch {
1869                requested: session_id,
1870                returned: cli_session_id,
1871            })
1872            .into());
1873        }
1874        if has_mcp_auth_handler
1875            && let Err(error) = register_mcp_auth_interest(self, &session_id).await
1876        {
1877            registration.cleanup(event_loop).await;
1878            return Err(error);
1879        }
1880        // Reload skills after resume (best-effort).
1881        let skills_reload_start = Instant::now();
1882        if let Err(e) = self
1883            .call(
1884                "session.skills.reload",
1885                Some(serde_json::json!({ "sessionId": session_id })),
1886            )
1887            .await
1888        {
1889            warn!(
1890                elapsed_ms = skills_reload_start.elapsed().as_millis(),
1891                session_id = %session_id,
1892                error = %e,
1893                "Client::resume_session skills reload request failed"
1894            );
1895        } else {
1896            tracing::debug!(
1897                elapsed_ms = skills_reload_start.elapsed().as_millis(),
1898                session_id = %session_id,
1899                "Client::resume_session skills reload request completed successfully"
1900            );
1901        }
1902
1903        *capabilities.write() = resume_result.capabilities.unwrap_or_default();
1904        // Upsert resume snapshots rather than replacing wholesale. Live
1905        // `session.canvas.opened` notifications can arrive on the event loop
1906        // while `session.resume` is in flight; a wholesale replace would
1907        // discard those updates.
1908        {
1909            let mut snapshots = open_canvases.write();
1910            for snapshot in resume_result.open_canvases.unwrap_or_default() {
1911                upsert_open_canvas_snapshot(&mut snapshots, snapshot);
1912            }
1913        }
1914
1915        tracing::debug!(
1916            elapsed_ms = total_start.elapsed().as_millis(),
1917            session_id = %session_id,
1918            "Client::resume_session complete"
1919        );
1920        registration.disarm();
1921        let session = Session {
1922            id: session_id,
1923            cwd: self.cwd().clone(),
1924            workspace_path: resume_result.workspace_path,
1925            remote_url: resume_result.remote_url,
1926            client: self.clone(),
1927            event_loop: ParkingLotMutex::new(Some(event_loop)),
1928            shutdown,
1929            external_tools_shutdown,
1930            idle_waiter,
1931            capabilities,
1932            open_canvases,
1933            event_tx,
1934            github_token_registration: ParkingLotMutex::new(github_token_registration),
1935            registration_token,
1936        };
1937        apply_mode_post_create_patch(
1938            &session,
1939            mode,
1940            opt_skip_custom_instructions,
1941            opt_custom_agents_local_only,
1942            opt_coauthor_enabled,
1943            opt_manage_schedule_enabled,
1944            opt_included_builtin_skills,
1945        )
1946        .await?;
1947        if let Some(registration) = session.github_token_registration.lock().as_ref() {
1948            registration.claim(session.id.clone());
1949        } else {
1950            self.retire_github_token_provider(&session.id);
1951        }
1952        Ok(session)
1953    }
1954}
1955
1956/// A session that has been configured but not yet created on the CLI.
1957///
1958/// Returned by [`Client::prepare_session`] and
1959/// [`Client::prepare_resume_session`]. Its purpose is to make the session's
1960/// event stream observable *before* any protocol activity starts:
1961/// [`subscribe`](Self::subscribe) installs a receiver on the same broadcast
1962/// channel the eventual [`Session`] uses, so events the runtime emits while
1963/// `session.create` / `session.resume` is still in flight are delivered
1964/// rather than dropped for lack of a receiver.
1965///
1966/// # Lifecycle
1967///
1968/// A prepared handle is inert. It holds only a broadcast sender, a
1969/// cancellation token, the client handle, and the config — it performs no
1970/// router registration, spawns no task, and writes nothing to the wire
1971/// until [`start`](Self::start) is first polled.
1972///
1973/// * Dropping it without starting leaves no client-side or server-side
1974///   state, and closes every subscription taken from it.
1975/// * Dropping the [`start`](Self::start) future mid-flight cancels the
1976///   session token, unregisters the session from the router if it was
1977///   registered, and closes early subscriptions. A retry with the same
1978///   session ID succeeds. Cleanup of already-spawned tasks is signalled,
1979///   not awaited: `Drop` is synchronous and cannot await, so the event loop
1980///   terminates promptly but not synchronously.
1981/// * A startup error from [`start`](Self::start) performs the same cleanup
1982///   and preserves the [`ErrorKind`] the equivalent
1983///   [`Client::create_session`] / [`Client::resume_session`] call has always
1984///   returned.
1985///
1986/// [`start`](Self::start) consumes `self` and the type is deliberately not
1987/// [`Clone`], so a prepared session can be started at most once and can
1988/// never produce two event loops.
1989///
1990/// # Buffering
1991///
1992/// The broadcast buffer is finite —
1993/// [`DEFAULT_EVENT_BUFFER_CAPACITY`] unless
1994/// [`SessionConfig::event_buffer_capacity`] /
1995/// [`ResumeSessionConfig::event_buffer_capacity`] overrides it. Subscribers
1996/// that fall behind observe
1997/// [`Lagged`](crate::subscription::Lagged) instead of applying backpressure
1998/// to the event loop. Consumers that need a lossless view of a large
1999/// startup burst must either configure a capacity that covers it or drain
2000/// the subscription concurrently with [`start`](Self::start).
2001///
2002/// # Server-assigned session IDs
2003///
2004/// For cloud sessions without a caller-supplied session ID, the CLI assigns
2005/// the ID and the SDK can only register the session on its notification
2006/// router once the `session.create` response arrives. Notifications the
2007/// server emits before that point are not routable to any session and are
2008/// therefore not observable. The guarantee this type provides is narrower
2009/// and precise: **routed** events are never dropped for lack of an
2010/// installed receiver. Pin
2011/// [`SessionConfig::session_id`](crate::types::SessionConfig::session_id)
2012/// to get registration before the RPC and full pre-response coverage.
2013#[must_use = "a PreparedSession does nothing until started"]
2014pub struct PreparedSession {
2015    client: Client,
2016    kind: PreparedKind,
2017    event_tx: tokio::sync::broadcast::Sender<SessionEvent>,
2018    shutdown: CancellationToken,
2019}
2020
2021/// Which startup path a [`PreparedSession`] runs when started. Boxed
2022/// because the two config types are large and differently sized.
2023enum PreparedKind {
2024    Create(Box<SessionConfig>),
2025    Resume(Box<ResumeSessionConfig>),
2026}
2027
2028impl PreparedSession {
2029    fn new(client: Client, kind: PreparedKind, event_buffer_capacity: usize) -> Self {
2030        let (event_tx, _) = tokio::sync::broadcast::channel(event_buffer_capacity);
2031        Self {
2032            client,
2033            kind,
2034            event_tx,
2035            shutdown: CancellationToken::new(),
2036        }
2037    }
2038
2039    /// Subscribe to this session's events before it starts.
2040    ///
2041    /// The returned [`EventSubscription`](crate::subscription::EventSubscription)
2042    /// is backed by the same broadcast channel
2043    /// [`Session::subscribe`] returns after [`start`](Self::start)
2044    /// succeeds, so a subscription taken here observes the full event
2045    /// stream from the session's first routed event onward — including
2046    /// ephemeral events such as `session.idle` that
2047    /// [`Session::get_messages`] cannot recover.
2048    ///
2049    /// May be called any number of times, and each subscriber receives its
2050    /// own copy of the stream — subject to the buffering contract above. A
2051    /// subscriber that falls further behind than the configured capacity
2052    /// observes [`Lagged`](crate::subscription::Lagged) and skips the
2053    /// events it missed, rather than stalling the session's event loop.
2054    /// Subscriptions taken here close if the prepared session is dropped
2055    /// without starting, or if startup fails.
2056    pub fn subscribe(&self) -> crate::subscription::EventSubscription {
2057        crate::subscription::EventSubscription::new(self.event_tx.subscribe())
2058    }
2059
2060    /// Create or resume the session on the CLI.
2061    ///
2062    /// This is where all protocol activity happens: config validation,
2063    /// router registration, the `session.create` / `session.resume` RPC,
2064    /// and the event loop spawn. Nothing observable occurs until this
2065    /// future is first polled.
2066    ///
2067    /// # Errors
2068    ///
2069    /// Returns the same errors as [`Client::create_session`] /
2070    /// [`Client::resume_session`] — including
2071    /// [`ErrorKind::InvalidConfig`] for invalid configs, transport and RPC
2072    /// failures, and
2073    /// [`SessionIdMismatch`](crate::SessionErrorKind::SessionIdMismatch)
2074    /// when the CLI returns a different session ID than the one requested.
2075    /// Every error path unregisters the session and closes subscriptions
2076    /// taken from this handle.
2077    pub async fn start(self) -> Result<Session, Error> {
2078        let Self {
2079            client,
2080            kind,
2081            event_tx,
2082            shutdown,
2083        } = self;
2084        match kind {
2085            PreparedKind::Create(config) => {
2086                client
2087                    .start_prepared_create(*config, event_tx, shutdown)
2088                    .await
2089            }
2090            PreparedKind::Resume(config) => {
2091                client
2092                    .start_prepared_resume(*config, event_tx, shutdown)
2093                    .await
2094            }
2095        }
2096    }
2097}
2098
2099type CommandHandlerMap = HashMap<String, Arc<dyn CommandHandler>>;
2100
2101async fn apply_mode_post_create_patch(
2102    session: &Session,
2103    mode: crate::ClientMode,
2104    opt_skip_custom_instructions: Option<bool>,
2105    opt_custom_agents_local_only: Option<bool>,
2106    opt_coauthor_enabled: Option<bool>,
2107    opt_manage_schedule_enabled: Option<bool>,
2108    opt_included_builtin_skills: Option<Vec<String>>,
2109) -> Result<(), Error> {
2110    let Some(patch) = build_mode_post_create_patch(
2111        mode,
2112        opt_skip_custom_instructions,
2113        opt_custom_agents_local_only,
2114        opt_coauthor_enabled,
2115        opt_manage_schedule_enabled,
2116        opt_included_builtin_skills,
2117    ) else {
2118        return Ok(());
2119    };
2120    if let Err(error) = session.rpc().options().update(patch).await {
2121        let _ = session.disconnect().await;
2122        return Err(error);
2123    }
2124    Ok(())
2125}
2126
2127/// Builds the `session.options.update` patch applied immediately after a session
2128/// is created or resumed, or returns `None` when no patch should be sent.
2129///
2130/// Under [`ClientMode::Empty`](crate::ClientMode::Empty) the overridable feature
2131/// flags fall back to safe defaults (caller values win), while
2132/// `installed_plugins` is unconditionally empty. `included_builtin_skills`
2133/// defaults to an empty list, but callers can explicitly allow selected
2134/// runtime-bundled skills. Under other modes only explicitly-set fields are
2135/// forwarded.
2136fn build_mode_post_create_patch(
2137    mode: crate::ClientMode,
2138    opt_skip_custom_instructions: Option<bool>,
2139    opt_custom_agents_local_only: Option<bool>,
2140    opt_coauthor_enabled: Option<bool>,
2141    opt_manage_schedule_enabled: Option<bool>,
2142    opt_included_builtin_skills: Option<Vec<String>>,
2143) -> Option<crate::generated::api_types::SessionUpdateOptionsParams> {
2144    use crate::generated::api_types::SessionUpdateOptionsParams;
2145    let mut patch = SessionUpdateOptionsParams::default();
2146    let should_send = if mode == crate::ClientMode::Empty {
2147        patch.skip_custom_instructions = Some(opt_skip_custom_instructions.unwrap_or(true));
2148        patch.custom_agents_local_only = Some(opt_custom_agents_local_only.unwrap_or(true));
2149        patch.coauthor_enabled = Some(opt_coauthor_enabled.unwrap_or(false));
2150        patch.manage_schedule_enabled = Some(opt_manage_schedule_enabled.unwrap_or(false));
2151        patch.installed_plugins = Some(Vec::new());
2152        patch.included_builtin_skills = Some(opt_included_builtin_skills.unwrap_or_default());
2153        true
2154    } else {
2155        let mut any = false;
2156        if let Some(v) = opt_skip_custom_instructions {
2157            patch.skip_custom_instructions = Some(v);
2158            any = true;
2159        }
2160        if let Some(v) = opt_custom_agents_local_only {
2161            patch.custom_agents_local_only = Some(v);
2162            any = true;
2163        }
2164        if let Some(v) = opt_coauthor_enabled {
2165            patch.coauthor_enabled = Some(v);
2166            any = true;
2167        }
2168        if let Some(v) = opt_manage_schedule_enabled {
2169            patch.manage_schedule_enabled = Some(v);
2170            any = true;
2171        }
2172        if let Some(v) = opt_included_builtin_skills {
2173            patch.included_builtin_skills = Some(v);
2174            any = true;
2175        }
2176        any
2177    };
2178    if !should_send {
2179        return None;
2180    }
2181    Some(patch)
2182}
2183
2184fn build_command_handler_map(commands: Option<&[CommandDefinition]>) -> Arc<CommandHandlerMap> {
2185    let map = match commands {
2186        Some(commands) => commands
2187            .iter()
2188            .filter(|cmd| !cmd.name.is_empty())
2189            .map(|cmd| (cmd.name.clone(), cmd.handler.clone()))
2190            .collect(),
2191        None => HashMap::new(),
2192    };
2193    Arc::new(map)
2194}
2195
2196fn upsert_open_canvas_snapshot(
2197    snapshots: &mut Vec<OpenCanvasInstance>,
2198    snapshot: OpenCanvasInstance,
2199) {
2200    if let Some(existing) = snapshots
2201        .iter_mut()
2202        .find(|open| open.instance_id == snapshot.instance_id)
2203    {
2204        *existing = snapshot;
2205    } else {
2206        snapshots.push(snapshot);
2207    }
2208}
2209
2210fn remove_open_canvas_snapshot(snapshots: &mut Vec<OpenCanvasInstance>, instance_id: &str) {
2211    snapshots.retain(|open| open.instance_id != instance_id);
2212}
2213
2214#[allow(clippy::too_many_arguments)]
2215fn spawn_event_loop(
2216    session_id: SessionId,
2217    client: Client,
2218    handlers: SessionHandlers,
2219    hooks: Option<Arc<dyn SessionHooks>>,
2220    transforms: Option<Arc<dyn SystemMessageTransform>>,
2221    command_handlers: Arc<CommandHandlerMap>,
2222    canvas_handler: Option<Arc<dyn CanvasHandler>>,
2223    session_fs_provider: Option<Arc<dyn SessionFsProvider>>,
2224    bearer_token_providers: HashMap<String, Arc<dyn BearerTokenProvider>>,
2225    channels: crate::router::SessionChannels,
2226    idle_waiter: Arc<ParkingLotMutex<Option<IdleWaiter>>>,
2227    capabilities: Arc<parking_lot::RwLock<SessionCapabilities>>,
2228    open_canvases: Arc<parking_lot::RwLock<Vec<OpenCanvasInstance>>>,
2229    event_tx: tokio::sync::broadcast::Sender<SessionEvent>,
2230    shutdown: CancellationToken,
2231    external_tools_shutdown: CancellationToken,
2232) -> JoinHandle<()> {
2233    let crate::router::SessionChannels {
2234        mut notifications,
2235        mut requests,
2236    } = channels;
2237    let pending_external_tools: PendingExternalTools =
2238        Arc::new(ParkingLotMutex::new(HashMap::new()));
2239
2240    let span = tracing::error_span!("session_event_loop", session_id = %session_id);
2241    tokio::spawn(
2242        async move {
2243            loop {
2244                // `mpsc::UnboundedReceiver::recv` and
2245                // `CancellationToken::cancelled` are both cancel-safe per
2246                // RFD 400.
2247                //
2248                // Inbound JSON-RPC *requests* are dispatched fire-and-forget:
2249                // each `handle_request` runs in its own spawned task that
2250                // awaits the handler and sends that request's response. This
2251                // mirrors the other Copilot SDKs and moves concurrency to the
2252                // request-dispatch boundary, so any slow handler — not just
2253                // `userInput.request` (which can stay pending for the full
2254                // input backstop of several minutes), but also `exitPlanMode`,
2255                // `autoModeSwitch`, hooks, transforms, or canvas/session-FS
2256                // providers — cannot park the reader loop and starve sibling
2257                // requests or co-emitted notifications. JSON-RPC permits
2258                // concurrent requests and out-of-order responses, so the SDK
2259                // does not serialize them.
2260                //
2261                // `handle_notification` is awaited inline because it only
2262                // performs fast dispatch work; its slow interactive callbacks
2263                // (permission/tool/elicitation) are themselves spawned as child
2264                // tasks. All of these spawned tasks intentionally outlive the
2265                // parent loop and own their own cleanup — RFD 400's "spawn
2266                // background tasks to perform cancel-unsafe operations" pattern.
2267                tokio::select! {
2268                    _ = shutdown.cancelled() => break,
2269                    Some(notification) = notifications.recv() => {
2270                        handle_notification(
2271                            &session_id, &client, &handlers, &command_handlers, notification, &idle_waiter, &capabilities, &open_canvases, &event_tx, &shutdown, &external_tools_shutdown, &pending_external_tools,
2272                        ).await;
2273                    }
2274                    Some(request) = requests.recv() => {
2275                        // Clone the Arc-backed dispatch context into the task so
2276                        // the spawned `handle_request` future is `'static`. All
2277                        // clones are cheap (Arc refcount bumps / small maps).
2278                        let span = tracing::error_span!("session_request_handler", session_id = %session_id);
2279                        let session_id = session_id.clone();
2280                        let client = client.clone();
2281                        let handlers = handlers.clone();
2282                        let hooks = hooks.clone();
2283                        let transforms = transforms.clone();
2284                        let canvas_handler = canvas_handler.clone();
2285                        let session_fs_provider = session_fs_provider.clone();
2286                        let bearer_token_providers = bearer_token_providers.clone();
2287                        let request_id = request.id;
2288                        let method = request.method.clone();
2289                        tokio::spawn(
2290                            async move {
2291                                let ctx = RequestDispatchContext {
2292                                    client: &client,
2293                                    handlers: &handlers,
2294                                    hooks: hooks.as_deref(),
2295                                    transforms: transforms.as_deref(),
2296                                    canvas_handler: canvas_handler.as_ref(),
2297                                    session_fs_provider: session_fs_provider.as_ref(),
2298                                    bearer_token_providers: &bearer_token_providers,
2299                                };
2300                                let dispatch = handle_request(&session_id, ctx, request);
2301                                if AssertUnwindSafe(dispatch).catch_unwind().await.is_err() {
2302                                    // Tokio isolates the panic to this task, so without a
2303                                    // reply the CLI waits out its own timeout on this id.
2304                                    error!(method = %method, "request handler panicked");
2305                                    let _ = send_error_response(
2306                                        &client,
2307                                        request_id,
2308                                        error_codes::INTERNAL_ERROR,
2309                                        "request handler panicked",
2310                                    )
2311                                    .await;
2312                                }
2313                            }
2314                            .instrument(span),
2315                        );
2316                    }
2317                    else => break,
2318                }
2319            }
2320            // Channels closed or shutdown signaled — fail any pending
2321            // send_and_wait so the caller observes a clean error.
2322            if let Some(waiter) = idle_waiter.lock().take() {
2323                let _ = waiter
2324                    .tx
2325                    .send(Err(ErrorKind::Session(SessionErrorKind::EventLoopClosed).into()));
2326            }
2327        }
2328        .instrument(span),
2329    )
2330}
2331
2332fn extract_request_id(data: &Value) -> Option<RequestId> {
2333    data.get("requestId")
2334        .and_then(|v| v.as_str())
2335        .filter(|s| !s.is_empty())
2336        .map(RequestId::new)
2337}
2338
2339fn permission_request_data(
2340    event_data: &Value,
2341    managed_settings_enabled: bool,
2342) -> PermissionRequestData {
2343    let request_data = event_data
2344        .get("permissionRequest")
2345        .cloned()
2346        .unwrap_or_else(|| event_data.clone());
2347    let managed_approval_required = match request_data.get("managedApprovalRequired") {
2348        None => None,
2349        Some(Value::Bool(value)) => Some(*value),
2350        Some(_) => Some(true),
2351    };
2352    match serde_json::from_value::<PermissionRequestData>(request_data) {
2353        Ok(mut data) => {
2354            data.extra = event_data.clone();
2355            data.managed_settings_enabled = managed_settings_enabled;
2356            data
2357        }
2358        Err(_) => PermissionRequestData {
2359            kind: None,
2360            tool_call_id: None,
2361            managed_approval_required,
2362            managed_settings_enabled,
2363            extra: event_data.clone(),
2364        },
2365    }
2366}
2367
2368/// Build the full `session.permissions.handlePendingPermissionRequest`
2369/// params for a permission result.
2370///
2371/// `decisionContext` is a sibling of `result` and is only present when the
2372/// handler attributed the decision — omitting it preserves legacy behavior.
2373///
2374/// Returns `None` when the SDK must not send a response.
2375fn permission_response_params(
2376    session_id: &SessionId,
2377    request_id: &RequestId,
2378    result: &PermissionResult,
2379) -> Option<Value> {
2380    let (decision, decision_context) = match result {
2381        PermissionResult::Decision { decision, context } => (decision, context.clone()),
2382        PermissionResult::NoResult => return None,
2383    };
2384    let mut params = serde_json::to_value(PermissionDecisionRequest {
2385        decision_context,
2386        request_id: request_id.clone(),
2387        result: decision.clone(),
2388    })
2389    .expect("serializing permission response should succeed");
2390    params["sessionId"] =
2391        serde_json::to_value(session_id).expect("serializing session ID should succeed");
2392    Some(params)
2393}
2394
2395async fn register_mcp_auth_interest(client: &Client, session_id: &SessionId) -> Result<(), Error> {
2396    let mut params = serde_json::to_value(RegisterEventInterestParams {
2397        event_type: "mcp.oauth_required".to_string(),
2398    })?;
2399    params["sessionId"] = Value::String(session_id.to_string());
2400    client
2401        .call(rpc_methods::SESSION_EVENTLOG_REGISTERINTEREST, Some(params))
2402        .await?;
2403    Ok(())
2404}
2405
2406fn tool_failure_result(message: impl Into<String>) -> ToolResult {
2407    let message = message.into();
2408    ToolResult::Expanded(ToolResultExpanded {
2409        text_result_for_llm: message.clone(),
2410        result_type: "failure".to_string(),
2411        binary_results_for_llm: None,
2412        session_log: None,
2413        error: Some(message),
2414        tool_telemetry: None,
2415        tool_references: None,
2416    })
2417}
2418
2419fn is_autopilot_continuation_idle(event: &SessionEvent) -> bool {
2420    event
2421        .typed_data::<SessionIdleData>()
2422        .is_some_and(|data| data.mode == Some(SessionMode::Autopilot))
2423}
2424
2425/// Process a notification from the CLI's broadcast channel.
2426#[allow(clippy::too_many_arguments)]
2427async fn handle_notification(
2428    session_id: &SessionId,
2429    client: &Client,
2430    handlers: &SessionHandlers,
2431    command_handlers: &Arc<CommandHandlerMap>,
2432    notification: SessionEventNotification,
2433    idle_waiter: &Arc<ParkingLotMutex<Option<IdleWaiter>>>,
2434    capabilities: &Arc<parking_lot::RwLock<SessionCapabilities>>,
2435    open_canvases: &Arc<parking_lot::RwLock<Vec<OpenCanvasInstance>>>,
2436    event_tx: &tokio::sync::broadcast::Sender<SessionEvent>,
2437    shutdown: &CancellationToken,
2438    external_tools_shutdown: &CancellationToken,
2439    pending_external_tools: &PendingExternalTools,
2440) {
2441    let dispatch_start = Instant::now();
2442    let event = &notification.event;
2443    let event_type = event.parsed_type();
2444    if event_type == SessionEventType::PermissionRequested {
2445        tracing::debug!(
2446            session_id = %session_id,
2447            event_type = %event.event_type,
2448            "Session::handle_notification permission request received"
2449        );
2450    }
2451
2452    // Signal send_and_wait if active. The lock is only contended when
2453    // a send_and_wait call is in flight (idle_waiter is Some).
2454    match event_type {
2455        SessionEventType::AssistantMessage
2456        | SessionEventType::SessionIdle
2457        | SessionEventType::SessionError => {
2458            let mut guard = idle_waiter.lock();
2459            if let Some(waiter) = guard.as_mut() {
2460                match event_type {
2461                    SessionEventType::AssistantMessage => {
2462                        if !waiter.first_assistant_message_seen {
2463                            waiter.first_assistant_message_seen = true;
2464                            tracing::debug!(
2465                                elapsed_ms = waiter.started_at.elapsed().as_millis(),
2466                                session_id = %session_id,
2467                                "Session::send_and_wait first assistant message"
2468                            );
2469                        }
2470                        waiter.last_assistant_message = Some(event.clone());
2471                    }
2472                    SessionEventType::SessionIdle if is_autopilot_continuation_idle(event) => {}
2473                    SessionEventType::SessionIdle | SessionEventType::SessionError => {
2474                        if let Some(waiter) = guard.take() {
2475                            if event_type == SessionEventType::SessionIdle {
2476                                tracing::debug!(
2477                                    elapsed_ms = waiter.started_at.elapsed().as_millis(),
2478                                    session_id = %session_id,
2479                                    "Session::send_and_wait idle received"
2480                                );
2481                                let _ = waiter.tx.send(Ok(waiter.last_assistant_message));
2482                            } else {
2483                                let error_msg = event
2484                                    .typed_data::<SessionErrorData>()
2485                                    .map(|d| d.message)
2486                                    .or_else(|| {
2487                                        event
2488                                            .data
2489                                            .get("message")
2490                                            .and_then(|v| v.as_str())
2491                                            .map(|s| s.to_string())
2492                                    })
2493                                    .unwrap_or_else(|| "session error".to_string());
2494                                let _ = waiter.tx.send(Err(Error::with_message(
2495                                    ErrorKind::Session(SessionErrorKind::AgentError),
2496                                    error_msg,
2497                                )));
2498                            }
2499                        }
2500                    }
2501                    _ => {}
2502                }
2503            }
2504        }
2505        _ => {}
2506    }
2507
2508    // Update the snapshot caches BEFORE broadcasting so subscribers that
2509    // call `Session::capabilities()` / `Session::open_canvases()` in
2510    // response to the event observe the new state.
2511    if event_type == SessionEventType::CapabilitiesChanged {
2512        match serde_json::from_value::<SessionCapabilities>(notification.event.data.clone()) {
2513            Ok(changed) => *capabilities.write() = changed,
2514            Err(e) => warn!(error = %e, "failed to deserialize capabilities.changed payload"),
2515        }
2516    }
2517    if event_type == SessionEventType::SessionCanvasOpened {
2518        match serde_json::from_value::<OpenCanvasInstance>(notification.event.data.clone()) {
2519            Ok(open_canvas) => {
2520                upsert_open_canvas_snapshot(&mut open_canvases.write(), open_canvas);
2521            }
2522            Err(e) => warn!(error = %e, "failed to deserialize session.canvas.opened payload"),
2523        }
2524    }
2525    if event_type == SessionEventType::SessionCanvasClosed {
2526        match serde_json::from_value::<SessionCanvasClosedData>(notification.event.data.clone()) {
2527            Ok(closed) => {
2528                if closed.instance_id.is_empty() {
2529                    warn!("failed to deserialize session.canvas.closed payload");
2530                } else {
2531                    remove_open_canvas_snapshot(&mut open_canvases.write(), &closed.instance_id);
2532                }
2533            }
2534            Err(e) => warn!(error = %e, "failed to deserialize session.canvas.closed payload"),
2535        }
2536    }
2537
2538    // Fan out the event to runtime subscribers (`Session::subscribe`). `send`
2539    // only errors when there are no receivers, which is the normal case
2540    // before any consumer subscribes.
2541    let _ = event_tx.send(event.clone());
2542
2543    tracing::debug!(
2544        elapsed_ms = dispatch_start.elapsed().as_millis(),
2545        session_id = %session_id,
2546        event_type = %notification.event.event_type,
2547        "Session::handle_notification dispatch"
2548    );
2549
2550    // Notification-based permission/tool/elicitation requests require a
2551    // separate RPC callback. Spawn concurrently since the CLI doesn't block.
2552    match event_type {
2553        SessionEventType::ExternalToolCompleted => {
2554            if let Some(request_id) = extract_request_id(&notification.event.data)
2555                && let Some(token) = pending_external_tools.lock().remove(&request_id)
2556            {
2557                token.cancel();
2558            }
2559        }
2560        SessionEventType::PermissionRequested => {
2561            let Some(request_id) = extract_request_id(&notification.event.data) else {
2562                return;
2563            };
2564            // Honor the runtime's `resolvedByHook` signal — when the
2565            // server has already resolved the permission via a hook,
2566            // clients must not send a second response.
2567            if notification
2568                .event
2569                .data
2570                .get("resolvedByHook")
2571                .and_then(|v| v.as_bool())
2572                .unwrap_or(false)
2573            {
2574                return;
2575            }
2576            // Multi-client safety: if this client has no permission
2577            // handler installed, don't respond — another client on the
2578            // same CLI may handle it.
2579            let Some(permission_handler) = handlers.permission.clone() else {
2580                return;
2581            };
2582            let client = client.clone();
2583            let sid = session_id.clone();
2584            let shutdown = shutdown.clone();
2585            let data = permission_request_data(
2586                &notification.event.data,
2587                handlers.managed_settings_enabled,
2588            );
2589            let span = tracing::error_span!(
2590                "permission_request_handler",
2591                session_id = %sid,
2592                request_id = %request_id
2593            );
2594            tokio::spawn(
2595                async move {
2596                    let handler_start = Instant::now();
2597                    let result = permission_handler
2598                        .handle(sid.clone(), request_id.clone(), data)
2599                        .await;
2600                    tracing::debug!(
2601                        elapsed_ms = handler_start.elapsed().as_millis(),
2602                        session_id = %sid,
2603                        request_id = %request_id,
2604                        "PermissionHandler::handle dispatch"
2605                    );
2606                    let Some(params) = permission_response_params(&sid, &request_id, &result)
2607                    else {
2608                        // Handler returned Deferred / NoResult — it will
2609                        // call handlePendingPermissionRequest itself (or
2610                        // leave the request unanswered).
2611                        return;
2612                    };
2613                    let rpc_start = Instant::now();
2614                    let method =
2615                        rpc_methods::SESSION_PERMISSIONS_HANDLEPENDINGPERMISSIONREQUEST;
2616                    tokio::select! {
2617                        biased;
2618                        response = client.call(method, Some(params)) => {
2619                            match response {
2620                                Ok(_) => tracing::debug!(
2621                                    elapsed_ms = rpc_start.elapsed().as_millis(),
2622                                    session_id = %sid,
2623                                    request_id = %request_id,
2624                                    method,
2625                                    "Session::handle_notification response sent successfully"
2626                                ),
2627                                Err(error) => warn!(
2628                                    error = %error,
2629                                    session_id = %sid,
2630                                    request_id = %request_id,
2631                                    method,
2632                                    "failed to deliver permission decision back to the runtime"
2633                                ),
2634                            }
2635                        }
2636                        _ = shutdown.cancelled() => {
2637                            warn!(
2638                                elapsed_ms = rpc_start.elapsed().as_millis(),
2639                                session_id = %sid,
2640                                request_id = %request_id,
2641                                method,
2642                                delivery_outcome = "unknown",
2643                                "permission confirmation acknowledgement wait cancelled during session shutdown"
2644                            );
2645                        }
2646                    }
2647                }
2648                .instrument(span),
2649            );
2650        }
2651        SessionEventType::ExternalToolRequested => {
2652            let Some(request_id) = extract_request_id(&notification.event.data) else {
2653                return;
2654            };
2655            let data: ExternalToolRequestedData =
2656                match serde_json::from_value(notification.event.data.clone()) {
2657                    Ok(d) => d,
2658                    Err(e) => {
2659                        warn!(error = %e, "failed to deserialize external_tool.requested");
2660                        let client = client.clone();
2661                        let sid = session_id.clone();
2662                        let span = tracing::error_span!(
2663                            "external_tool_deserialize_error",
2664                            session_id = %sid,
2665                            request_id = %request_id
2666                        );
2667                        tokio::spawn(
2668                            async move {
2669                                let rpc_start = Instant::now();
2670                                let _ = client
2671                                .call(
2672                                    "session.tools.handlePendingToolCall",
2673                                    Some(serde_json::json!({
2674                                        "sessionId": sid,
2675                                        "requestId": request_id,
2676                                        "error": format!("Failed to deserialize tool request: {e}"),
2677                                    })),
2678                                )
2679                                .await;
2680                                tracing::debug!(
2681                                    elapsed_ms = rpc_start.elapsed().as_millis(),
2682                                    session_id = %sid,
2683                                    request_id = %request_id,
2684                                    "Session::handle_notification response sent successfully"
2685                                );
2686                            }
2687                            .instrument(span),
2688                        );
2689                        return;
2690                    }
2691                };
2692            // Multi-client safety: look up a handler for the requested
2693            // tool name. If this client has no handler installed for that
2694            // tool, don't respond — another connected client may have one.
2695            let tool_handler = if data.tool_name.is_empty() {
2696                None
2697            } else {
2698                handlers.tools.get(&data.tool_name).cloned()
2699            };
2700            let Some(tool_handler) = tool_handler else {
2701                return;
2702            };
2703            let cancellation = Arc::new(external_tools_shutdown.child_token());
2704            {
2705                let mut pending = pending_external_tools.lock();
2706                if external_tools_shutdown.is_cancelled() || pending.contains_key(&request_id) {
2707                    return;
2708                }
2709                pending.insert(request_id.clone(), cancellation.clone());
2710            }
2711            let client = client.clone();
2712            let sid = session_id.clone();
2713            let pending_external_tools = pending_external_tools.clone();
2714            let guard_request_id = request_id.clone();
2715            let guard_cancellation = cancellation.clone();
2716            let span = tracing::error_span!(
2717                "external_tool_handler",
2718                session_id = %sid,
2719                request_id = %request_id
2720            );
2721            tokio::spawn(
2722                async move {
2723                    let guard = PendingExternalToolGuard {
2724                        request_id: guard_request_id,
2725                        token: guard_cancellation,
2726                        pending: pending_external_tools,
2727                    };
2728                    if cancellation.is_cancelled() {
2729                        return;
2730                    }
2731                    // `tool_name.is_empty()` would have produced a `None`
2732                    // lookup in `handlers.tools` and short-circuited at the
2733                    // outer guard above, so only the tool_call_id check is
2734                    // reachable here.
2735                    if data.tool_call_id.is_empty() {
2736                        if !guard.claim() {
2737                            return;
2738                        }
2739                        let error_msg = "Missing toolCallId";
2740                        let rpc_start = Instant::now();
2741                        let _ = client
2742                            .call(
2743                                "session.tools.handlePendingToolCall",
2744                                Some(serde_json::json!({
2745                                    "sessionId": sid,
2746                                    "requestId": request_id,
2747                                    "error": error_msg,
2748                                })),
2749                            )
2750                            .await;
2751                        tracing::debug!(
2752                            elapsed_ms = rpc_start.elapsed().as_millis(),
2753                            session_id = %sid,
2754                            request_id = %request_id,
2755                            "Session::handle_notification response sent successfully"
2756                        );
2757                        return;
2758                    }
2759                    let tool_call_id = data.tool_call_id.clone();
2760                    let tool_name = data.tool_name.clone();
2761                    // The built-in tool-search tool receives a snapshot of the
2762                    // session's currently initialized tools so an override can
2763                    // filter the live catalog without issuing its own RPC. Fetch
2764                    // it only for that tool to avoid a round-trip on every tool
2765                    // call; a failed fetch leaves the snapshot `None` rather than
2766                    // failing the tool.
2767                    let available_tools = if tool_name == TOOL_SEARCH_TOOL_NAME {
2768                        let metadata_result = tokio::select! {
2769                            biased;
2770                            _ = cancellation.cancelled() => return,
2771                            result = client.call(
2772                                rpc_methods::SESSION_TOOLS_GETCURRENTMETADATA,
2773                                Some(serde_json::json!({ "sessionId": sid })),
2774                            ) => result,
2775                        };
2776                        match metadata_result {
2777                            Ok(value) => {
2778                                serde_json::from_value::<ToolsGetCurrentMetadataResult>(value)
2779                                    .ok()
2780                                    .and_then(|result| result.tools)
2781                            }
2782                            Err(_) => None,
2783                        }
2784                    } else {
2785                        None
2786                    };
2787                    let invocation = ToolInvocation {
2788                        session_id: sid.clone(),
2789                        tool_call_id: data.tool_call_id,
2790                        tool_name: data.tool_name,
2791                        arguments: data
2792                            .arguments
2793                            .unwrap_or(Value::Object(serde_json::Map::new())),
2794                        available_tools,
2795                        traceparent: data.traceparent,
2796                        tracestate: data.tracestate,
2797                    };
2798                    let handler_start = Instant::now();
2799                    let tool_result = tokio::select! {
2800                        biased;
2801                        _ = cancellation.cancelled() => return,
2802                        result = tool_handler.call(invocation) => match result {
2803                            Ok(r) => r,
2804                            Err(e) => tool_failure_result(e.to_string()),
2805                        },
2806                    };
2807                    tracing::debug!(
2808                        elapsed_ms = handler_start.elapsed().as_millis(),
2809                        session_id = %sid,
2810                        request_id = %request_id,
2811                        tool_call_id = %tool_call_id,
2812                        tool_name = %tool_name,
2813                        "ToolHandler::call dispatch"
2814                    );
2815                    if !guard.claim() {
2816                        return;
2817                    }
2818                    let result_value = serde_json::to_value(tool_result).unwrap_or(Value::Null);
2819                    let rpc_start = Instant::now();
2820                    let _ = client
2821                        .call(
2822                            "session.tools.handlePendingToolCall",
2823                            Some(serde_json::json!({
2824                                "sessionId": sid,
2825                                "requestId": request_id,
2826                                "result": result_value,
2827                            })),
2828                        )
2829                        .await;
2830                    tracing::debug!(
2831                        elapsed_ms = rpc_start.elapsed().as_millis(),
2832                        session_id = %sid,
2833                        request_id = %request_id,
2834                        tool_call_id = %tool_call_id,
2835                        tool_name = %tool_name,
2836                        "Session::handle_notification response sent successfully"
2837                    );
2838                }
2839                .instrument(span),
2840            );
2841        }
2842        SessionEventType::UserInputRequested => {
2843            // Notification-only signal for observers (UI, telemetry).
2844            // The CLI follows up with a `userInput.request` JSON-RPC call
2845            // that drives the `UserInputHandler` dispatch — handling
2846            // the notification here too would double-fire the handler
2847            // and produce duplicate prompts on the consumer side. See
2848            // github/github-app#4249.
2849        }
2850        SessionEventType::ElicitationRequested => {
2851            let Some(request_id) = extract_request_id(&notification.event.data) else {
2852                return;
2853            };
2854            // Multi-client safety: if this client has no elicitation
2855            // handler installed, don't respond — another client on the
2856            // same CLI may handle it.
2857            let Some(elicitation_handler) = handlers.elicitation.clone() else {
2858                return;
2859            };
2860            let elicitation_data: ElicitationRequestedData =
2861                match serde_json::from_value(notification.event.data.clone()) {
2862                    Ok(d) => d,
2863                    Err(e) => {
2864                        warn!(error = %e, "failed to deserialize elicitation request");
2865                        return;
2866                    }
2867                };
2868            let request = ElicitationRequest {
2869                message: elicitation_data.message,
2870                requested_schema: elicitation_data
2871                    .requested_schema
2872                    .map(|s| serde_json::to_value(s).unwrap_or(Value::Null)),
2873                mode: elicitation_data.mode.map(|m| match m {
2874                    crate::generated::session_events::ElicitationRequestedMode::Form => {
2875                        crate::types::ElicitationMode::Form
2876                    }
2877                    crate::generated::session_events::ElicitationRequestedMode::Url => {
2878                        crate::types::ElicitationMode::Url
2879                    }
2880                    _ => crate::types::ElicitationMode::Unknown,
2881                }),
2882                elicitation_source: elicitation_data.elicitation_source,
2883                url: elicitation_data.url,
2884            };
2885            let client = client.clone();
2886            let sid = session_id.clone();
2887            let span = tracing::error_span!(
2888                "elicitation_request_handler",
2889                session_id = %sid,
2890                request_id = %request_id
2891            );
2892            tokio::spawn(
2893                async move {
2894                    let cancel = ElicitationResult {
2895                        action: "cancel".to_string(),
2896                        content: None,
2897                    };
2898                    // Dispatch to a nested task so panics are caught as JoinErrors.
2899                    let handler_task = tokio::spawn({
2900                        let sid = sid.clone();
2901                        let request_id = request_id.clone();
2902                        let span = tracing::error_span!(
2903                            "elicitation_callback",
2904                            session_id = %sid,
2905                            request_id = %request_id
2906                        );
2907                        async move {
2908                            let handler_start = Instant::now();
2909                            let response = elicitation_handler
2910                                .handle(sid.clone(), request_id.clone(), request)
2911                                .await;
2912                            tracing::debug!(
2913                                elapsed_ms = handler_start.elapsed().as_millis(),
2914                                session_id = %sid,
2915                                request_id = %request_id,
2916                                "ElicitationHandler::handle dispatch"
2917                            );
2918                            response
2919                        }
2920                        .instrument(span)
2921                    });
2922                    let result = match handler_task.await {
2923                        Ok(r) => r,
2924                        Err(_) => cancel.clone(),
2925                    };
2926                    let rpc_start = Instant::now();
2927                    if let Err(e) = client
2928                        .call(
2929                            "session.ui.handlePendingElicitation",
2930                            Some(serde_json::json!({
2931                                "sessionId": sid,
2932                                "requestId": request_id,
2933                                "result": result,
2934                            })),
2935                        )
2936                        .await
2937                    {
2938                        // RPC failed — attempt cancel as last resort
2939                        warn!(error = %e, "handlePendingElicitation failed, sending cancel");
2940                        let _ = client
2941                            .call(
2942                                "session.ui.handlePendingElicitation",
2943                                Some(serde_json::json!({
2944                                    "sessionId": sid,
2945                                    "requestId": request_id,
2946                                    "result": cancel,
2947                                })),
2948                            )
2949                            .await;
2950                    } else {
2951                        tracing::debug!(
2952                            elapsed_ms = rpc_start.elapsed().as_millis(),
2953                            session_id = %sid,
2954                            request_id = %request_id,
2955                            "Session::handle_notification response sent successfully"
2956                        );
2957                    }
2958                }
2959                .instrument(span),
2960            );
2961        }
2962        SessionEventType::McpOauthRequired => {
2963            let Some(request_id) = extract_request_id(&notification.event.data) else {
2964                return;
2965            };
2966            let Some(mcp_auth_handler) = handlers.mcp_auth.clone() else {
2967                warn!(
2968                    session_id = %session_id,
2969                    request_id = %request_id,
2970                    "received MCP OAuth request without a registered MCP auth handler"
2971                );
2972                return;
2973            };
2974            let data: McpOauthRequiredData =
2975                match serde_json::from_value(notification.event.data.clone()) {
2976                    Ok(d) => d,
2977                    Err(e) => {
2978                        warn!(error = %e, "failed to deserialize MCP OAuth request");
2979                        return;
2980                    }
2981                };
2982            let request = McpAuthRequest {
2983                request_id: request_id.clone(),
2984                server_name: data.server_name,
2985                server_url: data.server_url,
2986                reason: data.reason,
2987                www_authenticate_params: data.www_authenticate_params,
2988                resource_metadata: data.resource_metadata,
2989                static_client_config: data.static_client_config,
2990            };
2991            let client = client.clone();
2992            let sid = session_id.clone();
2993            let span = tracing::error_span!(
2994                "mcp_auth_request_handler",
2995                session_id = %sid,
2996                request_id = %request_id
2997            );
2998            tokio::spawn(
2999                async move {
3000                    let cancel = McpAuthResult::Cancelled;
3001                    let handler_task = tokio::spawn({
3002                        let sid = sid.clone();
3003                        let request_id = request_id.clone();
3004                        let span = tracing::error_span!(
3005                            "mcp_auth_callback",
3006                            session_id = %sid,
3007                            request_id = %request_id
3008                        );
3009                        async move {
3010                            let handler_start = Instant::now();
3011                            let response = mcp_auth_handler
3012                                .handle(sid.clone(), request_id.clone(), request)
3013                                .await;
3014                            tracing::debug!(
3015                                elapsed_ms = handler_start.elapsed().as_millis(),
3016                                session_id = %sid,
3017                                request_id = %request_id,
3018                                "McpAuthHandler::handle dispatch"
3019                            );
3020                            response
3021                        }
3022                        .instrument(span)
3023                    });
3024                    let result = match handler_task.await {
3025                        Ok(result) => result,
3026                        Err(_) => cancel,
3027                    };
3028                    let rpc_start = Instant::now();
3029                    let _ = client
3030                        .call(
3031                            "session.mcp.oauth.handlePendingRequest",
3032                            Some(serde_json::json!({
3033                                "sessionId": sid,
3034                                "requestId": request_id,
3035                                "result": result.into_wire(),
3036                            })),
3037                        )
3038                        .await;
3039                    tracing::debug!(
3040                        elapsed_ms = rpc_start.elapsed().as_millis(),
3041                        "Session::handle_notification MCP auth response sent"
3042                    );
3043                }
3044                .instrument(span),
3045            );
3046        }
3047        SessionEventType::CommandExecute => {
3048            let data: CommandExecuteData =
3049                match serde_json::from_value(notification.event.data.clone()) {
3050                    Ok(d) => d,
3051                    Err(e) => {
3052                        warn!(error = %e, "failed to deserialize command.execute");
3053                        return;
3054                    }
3055                };
3056            let client = client.clone();
3057            let command_handlers = command_handlers.clone();
3058            let sid = session_id.clone();
3059            let span = tracing::error_span!("command_handler", session_id = %sid);
3060            tokio::spawn(
3061                async move {
3062                    let request_id = data.request_id;
3063                    let ack_error = match command_handlers.get(&data.command_name).cloned() {
3064                        None => Some(format!("Unknown command: {}", data.command_name)),
3065                        Some(handler) => {
3066                            let command_name = data.command_name.clone();
3067                            let ctx = CommandContext {
3068                                session_id: sid.clone(),
3069                                command: data.command,
3070                                command_name: data.command_name,
3071                                args: data.args,
3072                            };
3073                            let handler_start = Instant::now();
3074                            let result = handler.on_command(ctx).await;
3075                            tracing::debug!(
3076                                elapsed_ms = handler_start.elapsed().as_millis(),
3077                                session_id = %sid,
3078                                request_id = %request_id,
3079                                command_name = %command_name,
3080                                "CommandHandler::call dispatch"
3081                            );
3082                            match result {
3083                                Ok(()) => None,
3084                                Err(e) => Some(e.to_string()),
3085                            }
3086                        }
3087                    };
3088                    let mut params = serde_json::json!({
3089                        "sessionId": sid,
3090                        "requestId": request_id,
3091                    });
3092                    if let Some(error_msg) = ack_error {
3093                        params["error"] = serde_json::Value::String(error_msg);
3094                    }
3095                    let rpc_start = Instant::now();
3096                    let _ = client
3097                        .call("session.commands.handlePendingCommand", Some(params))
3098                        .await;
3099                    tracing::debug!(
3100                        elapsed_ms = rpc_start.elapsed().as_millis(),
3101                        session_id = %sid,
3102                        request_id = %request_id,
3103                        "Session::handle_notification response sent successfully"
3104                    );
3105                }
3106                .instrument(span),
3107            );
3108        }
3109        _ => {}
3110    }
3111}
3112
3113struct RequestDispatchContext<'a> {
3114    client: &'a Client,
3115    handlers: &'a SessionHandlers,
3116    hooks: Option<&'a dyn SessionHooks>,
3117    transforms: Option<&'a dyn SystemMessageTransform>,
3118    canvas_handler: Option<&'a Arc<dyn CanvasHandler>>,
3119    session_fs_provider: Option<&'a Arc<dyn SessionFsProvider>>,
3120    bearer_token_providers: &'a HashMap<String, Arc<dyn BearerTokenProvider>>,
3121}
3122
3123/// Process a JSON-RPC request from the CLI.
3124async fn handle_request(
3125    session_id: &SessionId,
3126    ctx: RequestDispatchContext<'_>,
3127    request: crate::JsonRpcRequest,
3128) {
3129    let sid = session_id.clone();
3130    let client = ctx.client;
3131    let handlers = ctx.handlers;
3132    let hooks = ctx.hooks;
3133    let transforms = ctx.transforms;
3134    let canvas_handler = ctx.canvas_handler;
3135    let session_fs_provider = ctx.session_fs_provider;
3136    let bearer_token_providers = ctx.bearer_token_providers;
3137
3138    if request.method.starts_with("sessionFs.") {
3139        crate::session_fs_dispatch::dispatch(client, session_fs_provider, request).await;
3140        return;
3141    }
3142
3143    if request.method.starts_with("canvas.") {
3144        crate::canvas_dispatch::dispatch(client, canvas_handler, request).await;
3145        return;
3146    }
3147
3148    if request.method == crate::generated::api_types::rpc_methods::PROVIDERTOKEN_GETTOKEN {
3149        crate::provider_token_dispatch::dispatch(client, bearer_token_providers, request).await;
3150        return;
3151    }
3152
3153    match request.method.as_str() {
3154        "hooks.invoke" => {
3155            let params = request.params.as_ref();
3156            let hook_type = params
3157                .and_then(|p| p.get("hookType"))
3158                .and_then(|v| v.as_str())
3159                .unwrap_or("");
3160            let input = params
3161                .and_then(|p| p.get("input"))
3162                .cloned()
3163                .unwrap_or(Value::Object(Default::default()));
3164
3165            let rpc_result = if let Some(hooks) = hooks {
3166                match crate::hooks::dispatch_hook(hooks, &sid, hook_type, input).await {
3167                    Ok(output) => output,
3168                    Err(e) => {
3169                        warn!(error = %e, hook_type = hook_type, "hook dispatch failed");
3170                        serde_json::json!({ "output": {} })
3171                    }
3172                }
3173            } else {
3174                serde_json::json!({ "output": {} })
3175            };
3176
3177            let rpc_response = JsonRpcResponse {
3178                jsonrpc: "2.0".to_string(),
3179                id: request.id,
3180                result: Some(rpc_result),
3181                error: None,
3182            };
3183            let _ = client.send_response(&rpc_response).await;
3184        }
3185
3186        "userInput.request" => {
3187            let params = request.params.as_ref();
3188            let Some(question) = params
3189                .and_then(|p| p.get("question"))
3190                .and_then(|v| v.as_str())
3191            else {
3192                warn!("userInput.request missing 'question' field");
3193                let rpc_response = JsonRpcResponse {
3194                    jsonrpc: "2.0".to_string(),
3195                    id: request.id,
3196                    result: None,
3197                    error: Some(crate::JsonRpcError {
3198                        code: error_codes::INVALID_PARAMS,
3199                        message: "missing required field: question".to_string(),
3200                        data: None,
3201                    }),
3202                };
3203                let _ = client.send_response(&rpc_response).await;
3204                return;
3205            };
3206            let question = question.to_string();
3207            let choices = params
3208                .and_then(|p| p.get("choices"))
3209                .and_then(|v| v.as_array())
3210                .map(|arr| {
3211                    arr.iter()
3212                        .filter_map(|v| v.as_str().map(|s| s.to_string()))
3213                        .collect()
3214                });
3215            let allow_freeform = params
3216                .and_then(|p| p.get("allowFreeform"))
3217                .and_then(|v| v.as_bool());
3218
3219            let handler_start = Instant::now();
3220            let response = if let Some(user_input_handler) = handlers.user_input.as_ref() {
3221                user_input_handler
3222                    .handle(sid.clone(), question, choices, allow_freeform)
3223                    .await
3224            } else {
3225                None
3226            };
3227            tracing::debug!(
3228                elapsed_ms = handler_start.elapsed().as_millis(),
3229                session_id = %sid,
3230                "UserInputHandler::handle dispatch"
3231            );
3232
3233            let rpc_result = match response {
3234                Some(UserInputResponse {
3235                    answer,
3236                    was_freeform,
3237                }) => serde_json::json!({
3238                    "answer": answer,
3239                    "wasFreeform": was_freeform,
3240                }),
3241                None => serde_json::json!({ "noResponse": true }),
3242            };
3243            let rpc_response = JsonRpcResponse {
3244                jsonrpc: "2.0".to_string(),
3245                id: request.id,
3246                result: Some(rpc_result),
3247                error: None,
3248            };
3249            let _ = client.send_response(&rpc_response).await;
3250        }
3251
3252        "exitPlanMode.request" => {
3253            let params = request
3254                .params
3255                .as_ref()
3256                .cloned()
3257                .unwrap_or(Value::Object(serde_json::Map::new()));
3258            let data: ExitPlanModeData = match serde_json::from_value(params) {
3259                Ok(d) => d,
3260                Err(e) => {
3261                    warn!(error = %e, "failed to deserialize exitPlanMode.request params, using defaults");
3262                    ExitPlanModeData::default()
3263                }
3264            };
3265
3266            let rpc_result = if let Some(exit_plan_handler) = handlers.exit_plan_mode.as_ref() {
3267                let result = exit_plan_handler.handle(sid, data).await;
3268                serde_json::to_value(result).expect("ExitPlanModeResult serialization cannot fail")
3269            } else {
3270                serde_json::json!({ "approved": true })
3271            };
3272            let rpc_response = JsonRpcResponse {
3273                jsonrpc: "2.0".to_string(),
3274                id: request.id,
3275                result: Some(rpc_result),
3276                error: None,
3277            };
3278            let _ = client.send_response(&rpc_response).await;
3279        }
3280
3281        "autoModeSwitch.request" => {
3282            let error_code = request
3283                .params
3284                .as_ref()
3285                .and_then(|p| p.get("errorCode"))
3286                .and_then(|v| v.as_str())
3287                .map(|s| s.to_string());
3288            let retry_after_seconds = request
3289                .params
3290                .as_ref()
3291                .and_then(|p| p.get("retryAfterSeconds"))
3292                .and_then(|v| v.as_f64());
3293
3294            let answer = if let Some(auto_mode_handler) = handlers.auto_mode_switch.as_ref() {
3295                auto_mode_handler
3296                    .handle(sid, error_code, retry_after_seconds)
3297                    .await
3298            } else {
3299                AutoModeSwitchResponse::No
3300            };
3301            let rpc_response = JsonRpcResponse {
3302                jsonrpc: "2.0".to_string(),
3303                id: request.id,
3304                result: Some(serde_json::json!({ "response": answer })),
3305                error: None,
3306            };
3307            let _ = client.send_response(&rpc_response).await;
3308        }
3309
3310        "systemMessage.transform" => {
3311            let params = request.params.as_ref();
3312            let sections: HashMap<String, crate::transforms::TransformSection> =
3313                match params.and_then(|p| p.get("sections")) {
3314                    Some(v) => match serde_json::from_value(v.clone()) {
3315                        Ok(s) => s,
3316                        Err(e) => {
3317                            let _ = send_error_response(
3318                                client,
3319                                request.id,
3320                                error_codes::INVALID_PARAMS,
3321                                &format!("invalid sections: {e}"),
3322                            )
3323                            .await;
3324                            return;
3325                        }
3326                    },
3327                    None => {
3328                        let _ = send_error_response(
3329                            client,
3330                            request.id,
3331                            error_codes::INVALID_PARAMS,
3332                            "missing sections parameter",
3333                        )
3334                        .await;
3335                        return;
3336                    }
3337                };
3338
3339            let rpc_result = if let Some(transforms) = transforms {
3340                let transform_start = Instant::now();
3341                let response =
3342                    crate::transforms::dispatch_transform(transforms, &sid, sections).await;
3343                tracing::debug!(
3344                    elapsed_ms = transform_start.elapsed().as_millis(),
3345                    session_id = %sid,
3346                    "SystemMessageTransform::transform_section dispatch"
3347                );
3348                match serde_json::to_value(response) {
3349                    Ok(v) => v,
3350                    Err(e) => {
3351                        warn!(error = %e, "failed to serialize transform response");
3352                        serde_json::json!({ "sections": {} })
3353                    }
3354                }
3355            } else {
3356                // No transforms registered — pass through all sections unchanged.
3357                let passthrough: HashMap<String, crate::transforms::TransformSection> = sections;
3358                serde_json::json!({ "sections": passthrough })
3359            };
3360
3361            let rpc_response = JsonRpcResponse {
3362                jsonrpc: "2.0".to_string(),
3363                id: request.id,
3364                result: Some(rpc_result),
3365                error: None,
3366            };
3367            let _ = client.send_response(&rpc_response).await;
3368        }
3369
3370        method => {
3371            warn!(
3372                method = method,
3373                "unhandled request method in session event loop"
3374            );
3375            let _ = send_error_response(
3376                client,
3377                request.id,
3378                error_codes::METHOD_NOT_FOUND,
3379                &format!("unknown method: {method}"),
3380            )
3381            .await;
3382        }
3383    }
3384}
3385
3386async fn send_error_response(
3387    client: &Client,
3388    id: u64,
3389    code: i32,
3390    message: &str,
3391) -> Result<(), Error> {
3392    let response = JsonRpcResponse {
3393        jsonrpc: "2.0".to_string(),
3394        id,
3395        result: None,
3396        error: Some(crate::JsonRpcError {
3397            code,
3398            message: message.to_string(),
3399            data: None,
3400        }),
3401    };
3402    client.send_response(&response).await
3403}
3404
3405/// Inject `action: "transform"` sections into a `SystemMessageConfig`,
3406/// forcing `mode: "customize"` (required by the CLI for transforms to fire).
3407/// Preserves any existing caller-provided section overrides.
3408fn apply_transform_sections(
3409    sys_msg: &mut SystemMessageConfig,
3410    transforms: &dyn SystemMessageTransform,
3411) {
3412    sys_msg.mode = Some("customize".to_string());
3413    let sections = sys_msg.sections.get_or_insert_with(HashMap::new);
3414    for id in transforms.section_ids() {
3415        sections.entry(id).or_insert_with(|| SectionOverride {
3416            action: Some("transform".to_string()),
3417            content: None,
3418        });
3419    }
3420}
3421
3422fn inject_transform_sections(config: &mut SessionConfig, transforms: &dyn SystemMessageTransform) {
3423    let sys_msg = config.system_message.get_or_insert_with(Default::default);
3424    apply_transform_sections(sys_msg, transforms);
3425}
3426
3427fn inject_transform_sections_resume(
3428    config: &mut ResumeSessionConfig,
3429    transforms: &dyn SystemMessageTransform,
3430) {
3431    let sys_msg = config.system_message.get_or_insert_with(Default::default);
3432    apply_transform_sections(sys_msg, transforms);
3433}
3434
3435#[cfg(test)]
3436mod tests {
3437    use serde_json::json;
3438
3439    use super::{
3440        build_mode_post_create_patch, has_managed_settings, is_autopilot_continuation_idle,
3441        permission_request_data, permission_response_params,
3442    };
3443    use crate::handler::PermissionResult;
3444    use crate::types::{
3445        PermissionDecisionContext, PermissionDecisionOutcome, PermissionDecisionSource,
3446        PermissionDecisionSurface, RequestId, SessionEvent, SessionId,
3447    };
3448
3449    #[test]
3450    fn identifies_only_autopilot_continuation_idles() {
3451        let mut event = SessionEvent {
3452            id: "event-1".to_string(),
3453            timestamp: "2026-01-01T00:00:00Z".to_string(),
3454            parent_id: None,
3455            ephemeral: None,
3456            agent_id: None,
3457            debug_cli_received_at_ms: None,
3458            debug_ws_forwarded_at_ms: None,
3459            event_type: "session.idle".to_string(),
3460            data: json!({ "mode": "autopilot" }),
3461        };
3462
3463        assert!(is_autopilot_continuation_idle(&event));
3464
3465        event.data = json!({ "mode": "interactive" });
3466        assert!(!is_autopilot_continuation_idle(&event));
3467
3468        event.data = json!({});
3469        assert!(!is_autopilot_continuation_idle(&event));
3470    }
3471
3472    #[test]
3473    fn empty_mode_post_patch_sets_empty_included_builtin_skills() {
3474        let patch =
3475            build_mode_post_create_patch(crate::ClientMode::Empty, None, None, None, None, None)
3476                .expect("empty mode always sends a patch");
3477        assert_eq!(
3478            patch.included_builtin_skills,
3479            Some(Vec::new()),
3480            "empty mode must fail closed with an empty includedBuiltinSkills list"
3481        );
3482        assert_eq!(patch.installed_plugins.as_ref().map(|p| p.len()), Some(0));
3483        // Serializes as an explicit empty array (not omitted).
3484        let value = serde_json::to_value(&patch).expect("serialize patch");
3485        assert_eq!(value["includedBuiltinSkills"], serde_json::json!([]));
3486    }
3487
3488    #[test]
3489    fn empty_mode_post_patch_preserves_explicit_builtin_skill_allowlist() {
3490        let patch = build_mode_post_create_patch(
3491            crate::ClientMode::Empty,
3492            Some(false),
3493            Some(false),
3494            Some(true),
3495            Some(true),
3496            Some(vec!["code-review".to_string()]),
3497        )
3498        .expect("empty mode always sends a patch");
3499        assert_eq!(
3500            patch.included_builtin_skills,
3501            Some(vec!["code-review".to_string()])
3502        );
3503    }
3504
3505    #[test]
3506    fn copilot_cli_mode_does_not_inject_included_builtin_skills() {
3507        // No fields set -> no patch at all.
3508        assert!(
3509            build_mode_post_create_patch(
3510                crate::ClientMode::CopilotCli,
3511                None,
3512                None,
3513                None,
3514                None,
3515                None
3516            )
3517            .is_none()
3518        );
3519        // A field set -> patch sent, but skills field stays absent.
3520        let patch = build_mode_post_create_patch(
3521            crate::ClientMode::CopilotCli,
3522            Some(true),
3523            None,
3524            None,
3525            None,
3526            None,
3527        )
3528        .expect("a set field triggers a patch");
3529        assert_eq!(patch.included_builtin_skills, None);
3530        assert!(patch.installed_plugins.is_none());
3531        let value = serde_json::to_value(&patch).expect("serialize patch");
3532        assert!(value.get("includedBuiltinSkills").is_none());
3533
3534        let patch = build_mode_post_create_patch(
3535            crate::ClientMode::CopilotCli,
3536            None,
3537            None,
3538            None,
3539            None,
3540            Some(vec!["code-review".to_string()]),
3541        )
3542        .expect("an explicit allowlist triggers a patch");
3543        assert_eq!(
3544            patch.included_builtin_skills,
3545            Some(vec!["code-review".to_string()])
3546        );
3547    }
3548
3549    #[test]
3550    fn direct_injection_enables_managed_safeguards() {
3551        let settings = crate::types::ManagedSettings::default();
3552        assert!(has_managed_settings(None, Some(&settings)));
3553        assert!(!has_managed_settings(None, None));
3554    }
3555
3556    fn attribution_context() -> PermissionDecisionContext {
3557        PermissionDecisionContext {
3558            outcome: PermissionDecisionOutcome::AutoApproved,
3559            response_capability: None,
3560            source: PermissionDecisionSource::AssistedApproval,
3561            surface: PermissionDecisionSurface::CopilotApp,
3562        }
3563    }
3564
3565    #[test]
3566    fn response_params_omit_decision_context_without_attribution() {
3567        for (result, expected) in [
3568            (
3569                PermissionResult::approve_once(),
3570                json!({ "kind": "approve-once" }),
3571            ),
3572            (PermissionResult::reject(None), json!({ "kind": "reject" })),
3573            (
3574                PermissionResult::reject(Some("bad".to_string())),
3575                json!({ "kind": "reject", "feedback": "bad" }),
3576            ),
3577            (
3578                PermissionResult::user_not_available(),
3579                json!({ "kind": "user-not-available" }),
3580            ),
3581        ] {
3582            let params = permission_response_params(
3583                &SessionId::from("session-1"),
3584                &RequestId::from("permission-1"),
3585                &result,
3586            )
3587            .unwrap();
3588            assert_eq!(
3589                params,
3590                json!({
3591                    "sessionId": "session-1",
3592                    "requestId": "permission-1",
3593                    "result": expected,
3594                })
3595            );
3596        }
3597    }
3598
3599    #[test]
3600    fn response_params_forward_decision_context_alongside_result() {
3601        let params = permission_response_params(
3602            &SessionId::from("session-1"),
3603            &RequestId::from("permission-1"),
3604            &PermissionResult::approve_once().with_context(attribution_context()),
3605        )
3606        .unwrap();
3607        assert_eq!(
3608            params,
3609            json!({
3610                "sessionId": "session-1",
3611                "requestId": "permission-1",
3612                "result": { "kind": "approve-once" },
3613                "decisionContext": {
3614                    "outcome": "auto_approved",
3615                    "source": "assisted_approval",
3616                    "surface": "copilot_app",
3617                },
3618            })
3619        );
3620        // The context is a sibling of `result`, never nested inside it.
3621        assert!(params["result"].get("decisionContext").is_none());
3622    }
3623
3624    #[test]
3625    fn response_params_suppressed_for_no_result() {
3626        assert!(
3627            permission_response_params(
3628                &SessionId::from("session-1"),
3629                &RequestId::from("permission-1"),
3630                &PermissionResult::NoResult,
3631            )
3632            .is_none()
3633        );
3634    }
3635
3636    #[test]
3637    fn with_context_is_a_no_op_on_no_result() {
3638        let result = PermissionResult::no_result().with_context(attribution_context());
3639        assert!(matches!(result, PermissionResult::NoResult));
3640    }
3641
3642    #[test]
3643    fn with_context_replaces_rather_than_nests() {
3644        let result = PermissionResult::approve_once()
3645            .with_context(attribution_context())
3646            .with_context(PermissionDecisionContext {
3647                outcome: PermissionDecisionOutcome::PromptedUser,
3648                response_capability: None,
3649                source: PermissionDecisionSource::HumanResponse,
3650                surface: PermissionDecisionSurface::Sdk,
3651            });
3652        let params = permission_response_params(
3653            &SessionId::from("session-1"),
3654            &RequestId::from("permission-1"),
3655            &result,
3656        )
3657        .unwrap();
3658        assert_eq!(
3659            params["decisionContext"],
3660            json!({
3661                "outcome": "prompted_user",
3662                "source": "human_response",
3663                "surface": "sdk",
3664            })
3665        );
3666    }
3667
3668    #[test]
3669    fn permission_request_data_reads_nested_managed_approval_metadata() {
3670        let data = permission_request_data(
3671            &json!({
3672                "requestId": "permission-1",
3673                "permissionRequest": {
3674                    "kind": "read",
3675                    "managedApprovalRequired": true,
3676                    "path": "/workspace/file.txt"
3677                }
3678            }),
3679            false,
3680        );
3681
3682        assert_eq!(data.managed_approval_required, Some(true));
3683        assert_eq!(
3684            data.extra["permissionRequest"]["path"],
3685            "/workspace/file.txt"
3686        );
3687    }
3688
3689    #[test]
3690    fn permission_request_data_preserves_managed_flag_when_other_fields_are_malformed() {
3691        let data = permission_request_data(
3692            &json!({
3693                "requestId": "permission-1",
3694                "permissionRequest": {
3695                    "kind": "read",
3696                    "managedApprovalRequired": true,
3697                    "toolCallId": 42
3698                }
3699            }),
3700            false,
3701        );
3702
3703        assert_eq!(data.managed_approval_required, Some(true));
3704        assert_eq!(data.extra["requestId"], "permission-1");
3705    }
3706
3707    #[test]
3708    fn permission_request_data_fails_closed_for_malformed_managed_flag() {
3709        let data = permission_request_data(
3710            &json!({
3711                "requestId": "permission-1",
3712                "permissionRequest": {
3713                    "kind": "read",
3714                    "managedApprovalRequired": "yes",
3715                    "path": "/workspace/file.txt"
3716                }
3717            }),
3718            false,
3719        );
3720
3721        assert_eq!(data.managed_approval_required, Some(true));
3722    }
3723
3724    #[test]
3725    fn permission_request_data_preserves_valid_false_managed_flag() {
3726        let data = permission_request_data(
3727            &json!({
3728                "requestId": "permission-1",
3729                "permissionRequest": {
3730                    "kind": "read",
3731                    "managedApprovalRequired": false,
3732                    "path": "/workspace/file.txt"
3733                }
3734            }),
3735            false,
3736        );
3737
3738        assert_eq!(data.managed_approval_required, Some(false));
3739    }
3740}