Skip to main content

github_copilot_sdk/
session.rs

1use std::collections::HashMap;
2use std::path::{Path, PathBuf};
3use std::sync::Arc;
4use std::time::{Duration, Instant};
5
6use parking_lot::Mutex as ParkingLotMutex;
7use serde_json::Value;
8use tokio::sync::oneshot;
9use tokio::task::JoinHandle;
10use tokio_util::sync::CancellationToken;
11use tracing::{Instrument, warn};
12
13use crate::canvas::CanvasHandler;
14use crate::generated::api_types::{
15    LogRequest, ModelSwitchToRequest, OpenCanvasInstance, RegisterEventInterestParams, rpc_methods,
16};
17use crate::generated::session_events::{
18    CommandExecuteData, ElicitationRequestedData, ExternalToolRequestedData, McpOauthRequiredData,
19    SessionCanvasClosedData, SessionErrorData, SessionEventType,
20};
21use crate::handler::{
22    AutoModeSwitchHandler, AutoModeSwitchResponse, ElicitationHandler, ExitPlanModeHandler,
23    McpAuthHandler, McpAuthRequest, McpAuthResult, PermissionHandler, PermissionResult,
24    UserInputHandler, UserInputResponse,
25};
26use crate::hooks::SessionHooks;
27use crate::provider_token::BearerTokenProvider;
28use crate::session_fs::SessionFsProvider;
29use crate::trace_context::inject_trace_context;
30use crate::transforms::SystemMessageTransform;
31use crate::types::{
32    CommandContext, CommandDefinition, CommandHandler, CreateSessionResult, ElicitationRequest,
33    ElicitationResult, ExitPlanModeData, GetMessagesResponse, MessageOptions,
34    PermissionRequestData, RequestId, ResumeSessionConfig, ResumeSessionResult, SectionOverride,
35    SessionCapabilities, SessionConfig, SessionEvent, SessionId, SetModelOptions,
36    SystemMessageConfig, ToolInvocation, ToolResult, ToolResultExpanded, TraceContext,
37    UiInputOptions, ensure_attachment_display_names,
38};
39use crate::{
40    Client, Error, ErrorKind, JsonRpcResponse, SessionErrorKind, SessionEventNotification,
41    error_codes,
42};
43
44/// Bundle of the per-session callbacks the SDK dispatches to. Built from a
45/// [`SessionConfig`] / [`ResumeSessionConfig`] at
46/// [`Client::create_session`] / [`Client::resume_session`] time. Each
47/// field is `None` (or an empty map for tools) when the caller didn't
48/// install a handler -- in that case the SDK skips dispatch for that
49/// event type. The wire flags on `session.create` / `session.resume`
50/// are derived from these fields.
51#[derive(Clone)]
52pub(crate) struct SessionHandlers {
53    pub permission: Option<Arc<dyn PermissionHandler>>,
54    pub elicitation: Option<Arc<dyn ElicitationHandler>>,
55    pub mcp_auth: Option<Arc<dyn McpAuthHandler>>,
56    pub user_input: Option<Arc<dyn UserInputHandler>>,
57    pub exit_plan_mode: Option<Arc<dyn ExitPlanModeHandler>>,
58    pub auto_mode_switch: Option<Arc<dyn AutoModeSwitchHandler>>,
59    pub tools: Arc<HashMap<String, Arc<dyn crate::tool::ToolHandler>>>,
60}
61
62/// Shared state between a [`Session`] and its event loop, used by [`Session::send_and_wait`].
63struct IdleWaiter {
64    tx: oneshot::Sender<Result<Option<SessionEvent>, Error>>,
65    last_assistant_message: Option<SessionEvent>,
66    started_at: Instant,
67    first_assistant_message_seen: bool,
68}
69
70/// RAII guard that clears the [`Session::idle_waiter`] slot on drop. Used
71/// by [`Session::send_and_wait`] to ensure the slot doesn't leak if the
72/// caller's future is cancelled (outer `tokio::time::timeout` / `select!`
73/// / dropped JoinHandle). Synchronous clear via `parking_lot::Mutex` —
74/// no async drop needed.
75///
76/// Without this, an outer cancellation between "install waiter" and
77/// "drain channel" would leave the slot occupied, causing all subsequent
78/// `send` and `send_and_wait` calls on the session to return
79/// [`SendWhileWaiting`](SessionErrorKind::SendWhileWaiting). Closes RFD-400
80/// review finding #2.
81struct WaiterGuard {
82    slot: Arc<ParkingLotMutex<Option<IdleWaiter>>>,
83}
84
85impl Drop for WaiterGuard {
86    fn drop(&mut self) {
87        self.slot.lock().take();
88    }
89}
90
91struct PendingSessionRegistration {
92    client: Client,
93    session_id: SessionId,
94    shutdown: CancellationToken,
95    disarmed: bool,
96}
97
98impl PendingSessionRegistration {
99    fn new(client: Client, session_id: SessionId, shutdown: CancellationToken) -> Self {
100        Self {
101            client,
102            session_id,
103            shutdown,
104            disarmed: false,
105        }
106    }
107
108    async fn cleanup(mut self, event_loop: JoinHandle<()>) {
109        self.shutdown.cancel();
110        let _ = event_loop.await;
111        self.client.unregister_session(&self.session_id);
112        self.disarmed = true;
113    }
114
115    fn disarm(&mut self) {
116        self.disarmed = true;
117    }
118}
119
120impl Drop for PendingSessionRegistration {
121    fn drop(&mut self) {
122        if !self.disarmed {
123            self.shutdown.cancel();
124            self.client.unregister_session(&self.session_id);
125        }
126    }
127}
128
129/// A session on a GitHub Copilot CLI server.
130///
131/// Created via [`Client::create_session`] or [`Client::resume_session`].
132/// Owns an internal event loop that dispatches events to the per-callback
133/// handlers installed on the session config.
134///
135/// Protocol methods (`send`, `get_events`, `abort`, etc.) automatically
136/// inject the session ID into RPC params.
137///
138/// Call [`destroy`](Self::destroy) for graceful cleanup (RPC + local). If dropped
139/// without calling `destroy`, the `Drop` impl aborts the event loop and
140/// unregisters from the router as a best-effort safety net.
141pub struct Session {
142    id: SessionId,
143    cwd: PathBuf,
144    workspace_path: Option<PathBuf>,
145    remote_url: Option<String>,
146    client: Client,
147    /// Handle to the spawned event-loop task. Sync `parking_lot::Mutex`
148    /// because the lock is never held across an `.await` and the `Drop`
149    /// impl needs to take the handle synchronously without `try_lock`
150    /// fallibility.
151    event_loop: ParkingLotMutex<Option<JoinHandle<()>>>,
152    /// Cooperative shutdown signal for the event loop. The loop selects
153    /// on [`shutdown.cancelled()`](CancellationToken::cancelled) alongside
154    /// its inbound channels; [`Session::stop_event_loop`] and [`Drop`]
155    /// both call [`cancel()`](CancellationToken::cancel) to ask the loop
156    /// to exit between iterations rather than aborting the task (which
157    /// can land at any await point and leave the session mid-protocol).
158    /// See RFD-400 review finding #3.
159    ///
160    /// `CancellationToken` is the canonical signalling primitive in
161    /// `tokio_util`; it is what `tonic` uses for the equivalent task-
162    /// coordination case. Advanced consumers can obtain a child token
163    /// via [`Session::cancellation_token`] to bind their own work to
164    /// the session lifetime.
165    shutdown: CancellationToken,
166    /// Only populated while a `send_and_wait` call is in flight.
167    ///
168    /// Sync `parking_lot::Mutex` because the lock is never held across an
169    /// `.await`, and synchronous access lets the `WaiterGuard` RAII helper
170    /// in `send_and_wait` clear the slot from a `Drop` impl on caller-side
171    /// cancellation. See RFD-400 review (cancel-safety hardening).
172    idle_waiter: Arc<ParkingLotMutex<Option<IdleWaiter>>>,
173    /// Capabilities negotiated with the CLI, updated on `capabilities.changed` events.
174    capabilities: Arc<parking_lot::RwLock<SessionCapabilities>>,
175    /// Canvas instances currently known to be open for this session.
176    open_canvases: Arc<parking_lot::RwLock<Vec<OpenCanvasInstance>>>,
177    /// Broadcast channel for runtime event subscribers — see [`Session::subscribe`].
178    event_tx: tokio::sync::broadcast::Sender<SessionEvent>,
179}
180
181impl Session {
182    /// Session ID assigned by the CLI.
183    pub fn id(&self) -> &SessionId {
184        &self.id
185    }
186
187    /// Working directory of the CLI process.
188    pub fn cwd(&self) -> &PathBuf {
189        &self.cwd
190    }
191
192    /// Workspace directory for the session (if using infinite sessions).
193    pub fn workspace_path(&self) -> Option<&Path> {
194        self.workspace_path.as_deref()
195    }
196
197    /// Remote session URL, if the session is running remotely.
198    pub fn remote_url(&self) -> Option<&str> {
199        self.remote_url.as_deref()
200    }
201
202    /// Session capabilities negotiated with the CLI.
203    ///
204    /// Capabilities are set during session creation and updated at runtime
205    /// via `capabilities.changed` events.
206    pub fn capabilities(&self) -> SessionCapabilities {
207        self.capabilities.read().clone()
208    }
209
210    /// Open canvas instances reported by the most recent `session.resume`
211    /// response or surfaced by inbound `canvas.opened` events.
212    pub fn open_canvases(&self) -> Vec<OpenCanvasInstance> {
213        self.open_canvases.read().clone()
214    }
215
216    /// Returns a [`CancellationToken`] that fires when this session shuts
217    /// down (via [`Session::stop_event_loop`], [`Session::destroy`], or
218    /// [`Drop`]).
219    ///
220    /// Use this to bind an external task's lifetime to the session — when
221    /// the session shuts down, awaiting [`cancelled()`](CancellationToken::cancelled)
222    /// resolves so cooperative consumers can stop cleanly.
223    ///
224    /// The returned handle is a *child* token: calling
225    /// [`cancel()`](CancellationToken::cancel) on it cancels only the
226    /// caller's child, not the session itself. To cancel the session, call
227    /// [`Session::stop_event_loop`].
228    ///
229    /// # Example
230    ///
231    /// ```no_run
232    /// # async fn example(session: github_copilot_sdk::session::Session) {
233    /// let token = session.cancellation_token();
234    /// tokio::select! {
235    ///     _ = token.cancelled() => println!("session shut down"),
236    ///     _ = tokio::time::sleep(std::time::Duration::from_secs(60)) => {
237    ///         println!("60s elapsed, session still alive");
238    ///     }
239    /// }
240    /// # }
241    /// ```
242    pub fn cancellation_token(&self) -> CancellationToken {
243        self.shutdown.child_token()
244    }
245
246    /// Subscribe to events for this session.
247    ///
248    /// Returns an [`EventSubscription`](crate::subscription::EventSubscription)
249    /// that yields every [`SessionEvent`] dispatched on this session's
250    /// event loop. Drop the value to unsubscribe; there is no separate
251    /// cancel handle.
252    ///
253    /// **Observe-only.** Subscribers receive a clone of every
254    /// [`SessionEvent`] but cannot influence permission decisions, tool
255    /// results, or anything else that requires returning a value. Those
256    /// remain the responsibility of the per-callback handlers passed via
257    /// [`SessionConfig`]'s `with_*_handler`
258    /// builder methods.
259    ///
260    /// The returned handle implements both an inherent
261    /// [`recv`](crate::subscription::EventSubscription::recv) method and
262    /// [`Stream`](tokio_stream::Stream), so callers can use a `while let`
263    /// loop or any combinator from `tokio_stream::StreamExt` /
264    /// `futures::StreamExt`.
265    ///
266    /// Each subscriber maintains its own queue. If a consumer cannot keep
267    /// up, the oldest events are dropped and `recv` returns
268    /// [`RecvErrorKind::Lagged`](crate::subscription::RecvErrorKind::Lagged)
269    /// reporting the count of skipped events. Slow consumers do not block
270    /// the session's event loop.
271    ///
272    /// # Example
273    ///
274    /// ```no_run
275    /// # async fn example(session: github_copilot_sdk::session::Session) {
276    /// let mut events = session.subscribe();
277    /// tokio::spawn(async move {
278    ///     while let Ok(event) = events.recv().await {
279    ///         println!("[{}] event {}", event.id, event.event_type);
280    ///     }
281    /// });
282    /// # }
283    /// ```
284    pub fn subscribe(&self) -> crate::subscription::EventSubscription {
285        crate::subscription::EventSubscription::new(self.event_tx.subscribe())
286    }
287
288    /// The underlying Client (for advanced use cases).
289    pub fn client(&self) -> &Client {
290        &self.client
291    }
292
293    /// Typed RPC namespace for this session.
294    ///
295    /// Every protocol method lives here under its schema-aligned path —
296    /// e.g. `session.rpc().workspaces().list_files()`. Wire method names
297    /// and request/response types are generated from the protocol schema,
298    /// so the typed namespace can't drift from the wire contract.
299    ///
300    /// The hand-authored helpers on [`Session`] delegate to this namespace
301    /// and remain the recommended entry point for everyday use; reach for
302    /// `rpc()` when you want a method without a hand-written wrapper.
303    pub fn rpc(&self) -> crate::generated::rpc::SessionRpc<'_> {
304        crate::generated::rpc::SessionRpc { session: self }
305    }
306
307    /// Stop the internal event loop. Called automatically on [`destroy`](Self::destroy).
308    ///
309    /// Cooperative: signals shutdown via the session's [`CancellationToken`]
310    /// and awaits the loop's natural exit rather than aborting the task.
311    /// Any in-flight handler (permission callback, tool call, elicitation
312    /// response) completes before the loop exits, so the CLI never sees a
313    /// half-handled request. See RFD-400 review finding #3.
314    pub async fn stop_event_loop(&self) {
315        self.shutdown.cancel();
316        let handle = self.event_loop.lock().take();
317        if let Some(handle) = handle {
318            let _ = handle.await;
319        }
320        // Fail any pending send_and_wait so it returns immediately.
321        if let Some(waiter) = self.idle_waiter.lock().take() {
322            let _ = waiter.tx.send(Err(
323                ErrorKind::Session(SessionErrorKind::EventLoopClosed).into()
324            ));
325        }
326    }
327
328    /// Send a user message to the agent.
329    ///
330    /// Accepts anything convertible to [`MessageOptions`] — pass a `&str` for the
331    /// trivial case, or build a `MessageOptions` for mode/attachments. The
332    /// `wait_timeout` field on `MessageOptions` is ignored here (use
333    /// [`send_and_wait`](Self::send_and_wait) if you need to wait).
334    ///
335    /// Returns the assigned message ID, which can be used to correlate the
336    /// send with later [`SessionEvent`]s emitted in
337    /// response (assistant messages, tool requests, etc.).
338    ///
339    /// Returns an error if a [`send_and_wait`](Self::send_and_wait) call is
340    /// currently in flight, since the plain send would race with the waiter.
341    ///
342    /// # Cancel safety
343    ///
344    /// **Cancel-safe.** The underlying `session.send` RPC is dispatched
345    /// through the writer-actor (see [`Client::call`](crate::Client::call)),
346    /// so dropping this future after the actor has committed to writing
347    /// will not produce a partial frame on the wire. If the caller's
348    /// future is dropped between "frame enqueued" and "response received",
349    /// the message has already landed on the wire — the agent will process
350    /// it and emit events normally; the caller just won't see the returned
351    /// message ID.
352    pub async fn send(&self, opts: impl Into<MessageOptions>) -> Result<String, Error> {
353        if self.idle_waiter.lock().is_some() {
354            return Err(ErrorKind::Session(SessionErrorKind::SendWhileWaiting).into());
355        }
356        self.send_inner(opts.into()).await
357    }
358
359    async fn send_inner(&self, opts: MessageOptions) -> Result<String, Error> {
360        let mut params = serde_json::json!({
361            "sessionId": self.id,
362            "prompt": opts.prompt,
363        });
364        if let Some(m) = opts.mode {
365            params["mode"] = serde_json::to_value(m)?;
366        }
367        if let Some(am) = opts.agent_mode {
368            params["agentMode"] = serde_json::to_value(am)?;
369        }
370        if let Some(mut a) = opts.attachments {
371            ensure_attachment_display_names(&mut a);
372            params["attachments"] = serde_json::to_value(a)?;
373        }
374        if let Some(headers) = opts.request_headers
375            && !headers.is_empty()
376        {
377            params["requestHeaders"] = serde_json::to_value(headers)?;
378        }
379        if let Some(display_prompt) = opts.display_prompt {
380            params["displayPrompt"] = serde_json::to_value(display_prompt)?;
381        }
382        let trace_ctx = if opts.traceparent.is_some() || opts.tracestate.is_some() {
383            TraceContext {
384                traceparent: opts.traceparent,
385                tracestate: opts.tracestate,
386            }
387        } else {
388            self.client.resolve_trace_context().await
389        };
390        inject_trace_context(&mut params, &trace_ctx);
391        let rpc_start = Instant::now();
392        let result = self.client.call("session.send", Some(params)).await?;
393        let message_id = result
394            .get("messageId")
395            .and_then(|v| v.as_str())
396            .map(|s| s.to_string())
397            .unwrap_or_default();
398        tracing::debug!(
399            elapsed_ms = rpc_start.elapsed().as_millis(),
400            session_id = %self.id,
401            message_id = %message_id,
402            "Session::send completed successfully"
403        );
404        Ok(message_id)
405    }
406
407    /// Send a user message and wait for the agent to finish processing.
408    ///
409    /// Accepts anything convertible to [`MessageOptions`] — pass a `&str` for the
410    /// trivial case, or build a `MessageOptions` for mode/attachments/timeout.
411    /// Blocks until `session.idle` (success) or `session.error` (failure),
412    /// returning the last `assistant.message` event captured during streaming.
413    /// Times out after `MessageOptions::wait_timeout` (default 60 seconds).
414    ///
415    /// Only one `send_and_wait` call may be active per session at a time.
416    /// Calling [`send`](Self::send) while a `send_and_wait`
417    /// is in flight will also return an error.
418    ///
419    /// # Cancel safety
420    ///
421    /// **Cancel-safe.** A `WaiterGuard` clears the in-flight slot on every
422    /// exit path (success, internal failure, internal timeout, *and*
423    /// external cancellation via `tokio::time::timeout` / `select!` /
424    /// dropped JoinHandle). Subsequent `send` and `send_and_wait` calls on
425    /// this session will succeed normally — the slot is never leaked.
426    pub async fn send_and_wait(
427        &self,
428        opts: impl Into<MessageOptions>,
429    ) -> Result<Option<SessionEvent>, Error> {
430        let total_start = Instant::now();
431        let opts = opts.into();
432        let timeout_duration = opts.wait_timeout.unwrap_or(Duration::from_secs(60));
433        let (tx, rx) = oneshot::channel();
434
435        {
436            let mut guard = self.idle_waiter.lock();
437            if guard.is_some() {
438                return Err(ErrorKind::Session(SessionErrorKind::SendWhileWaiting).into());
439            }
440            *guard = Some(IdleWaiter {
441                tx,
442                last_assistant_message: None,
443                started_at: total_start,
444                first_assistant_message_seen: false,
445            });
446        }
447
448        // RAII: clears the idle_waiter slot on every exit path, including
449        // external cancellation (caller's outer `select!` / `timeout` /
450        // dropped future). Without this, an outer cancellation would leak
451        // the slot and brick subsequent `send`/`send_and_wait` calls.
452        let _waiter_guard = WaiterGuard {
453            slot: self.idle_waiter.clone(),
454        };
455
456        let result = tokio::time::timeout(timeout_duration, async {
457            self.send_inner(opts).await?;
458            match rx.await {
459                Ok(result) => result,
460                Err(_) => Err(ErrorKind::Session(SessionErrorKind::EventLoopClosed).into()),
461            }
462        })
463        .await;
464
465        match result {
466            Ok(inner) => {
467                tracing::debug!(
468                    elapsed_ms = total_start.elapsed().as_millis(),
469                    session_id = %self.id,
470                    completed_by = if inner.is_ok() { "idle" } else { "error" },
471                    "Session::send_and_wait complete"
472                );
473                inner
474            }
475            Err(_) => {
476                tracing::warn!(
477                    elapsed_ms = total_start.elapsed().as_millis(),
478                    session_id = %self.id,
479                    completed_by = "timeout",
480                    "Session::send_and_wait failed"
481                );
482                Err(ErrorKind::Session(SessionErrorKind::Timeout(timeout_duration)).into())
483            }
484        }
485    }
486
487    /// Retrieve the session's timeline events.
488    pub async fn get_events(&self) -> Result<Vec<SessionEvent>, Error> {
489        let result = self
490            .client
491            .call(
492                "session.getMessages",
493                Some(serde_json::json!({ "sessionId": self.id })),
494            )
495            .await?;
496        let response: GetMessagesResponse = serde_json::from_value(result)?;
497        Ok(response.events)
498    }
499
500    /// Deprecated alias for [`get_events`](Self::get_events).
501    #[deprecated(since = "0.1.0", note = "Use `get_events()` instead")]
502    pub async fn get_messages(&self) -> Result<Vec<SessionEvent>, Error> {
503        self.get_events().await
504    }
505
506    /// Abort the current agent turn.
507    ///
508    /// # Cancel safety
509    ///
510    /// **Cancel-safe.** Single `session.abort` RPC; the underlying
511    /// [`Client::call`](crate::Client::call) is cancel-safe via the
512    /// writer-actor.
513    pub async fn abort(&self) -> Result<(), Error> {
514        self.client
515            .call(
516                "session.abort",
517                Some(serde_json::json!({ "sessionId": self.id })),
518            )
519            .await?;
520        Ok(())
521    }
522
523    /// Switch to a different model.
524    ///
525    /// Pass `None` for `opts` if no extra configuration is needed.
526    pub async fn set_model(&self, model: &str, opts: Option<SetModelOptions>) -> Result<(), Error> {
527        let opts = opts.unwrap_or_default();
528        let request = ModelSwitchToRequest {
529            model_id: model.to_string(),
530            reasoning_effort: opts.reasoning_effort,
531            reasoning_summary: opts.reasoning_summary,
532            verbosity: None,
533            context_tier: opts.context_tier,
534            model_capabilities: opts.model_capabilities,
535        };
536        self.rpc().model().switch_to(request).await?;
537        Ok(())
538    }
539
540    /// Disconnect this session from the CLI.
541    ///
542    /// Sends the `session.destroy` RPC, stops the event loop, and unregisters
543    /// the session from the client. **Session state on disk** (conversation
544    /// history, planning state, artifacts) is **preserved**, so the
545    /// conversation can be resumed later via [`Client::resume_session`]
546    /// using this session's ID. To permanently remove all on-disk session
547    /// data, use [`Client::delete_session`] instead.
548    ///
549    /// The caller should ensure the session is idle (e.g. [`send_and_wait`]
550    /// has returned) before disconnecting; in-flight tool or event handlers
551    /// may otherwise observe failures.
552    ///
553    /// [`Client::resume_session`]: crate::Client::resume_session
554    /// [`Client::delete_session`]: crate::Client::delete_session
555    /// [`send_and_wait`]: Self::send_and_wait
556    pub async fn disconnect(&self) -> Result<(), Error> {
557        self.client
558            .call(
559                "session.destroy",
560                Some(serde_json::json!({ "sessionId": self.id })),
561            )
562            .await?;
563        self.stop_event_loop().await;
564        self.client.unregister_session(&self.id);
565        Ok(())
566    }
567
568    /// Deprecated alias for [`disconnect`](Self::disconnect). The
569    /// underlying wire RPC happens to be named `session.destroy`, but it
570    /// only severs the connection — on-disk session state is preserved.
571    /// Prefer `disconnect` in new code.
572    #[deprecated(since = "0.1.0", note = "Use `disconnect()` instead")]
573    pub async fn destroy(&self) -> Result<(), Error> {
574        self.disconnect().await
575    }
576
577    /// Write a log message to the session.
578    ///
579    /// Pass `None` for `opts` to use defaults (info level, persisted).
580    pub async fn log(
581        &self,
582        message: &str,
583        opts: Option<crate::types::LogOptions>,
584    ) -> Result<(), Error> {
585        let opts = opts.unwrap_or_default();
586        let level = match opts.level {
587            Some(level) => Some(serde_json::from_value(serde_json::to_value(level)?)?),
588            None => None,
589        };
590        let request = LogRequest {
591            message: message.to_string(),
592            level,
593            ephemeral: opts.ephemeral,
594            r#type: None,
595            tip: None,
596            url: None,
597        };
598        self.rpc().log(request).await?;
599        Ok(())
600    }
601
602    /// Returns the UI sub-API for elicitation, confirmation, selection, and
603    /// free-form input.
604    ///
605    /// All UI methods route through `session.ui.*` RPCs and require host
606    /// support — check `session.capabilities().ui.elicitation` before use.
607    pub fn ui(&self) -> SessionUi<'_> {
608        SessionUi { session: self }
609    }
610
611    /// Returns an error if the host doesn't support elicitation.
612    fn assert_elicitation(&self) -> Result<(), Error> {
613        if self
614            .capabilities
615            .read()
616            .ui
617            .as_ref()
618            .and_then(|u| u.elicitation)
619            != Some(true)
620        {
621            return Err(ErrorKind::Session(SessionErrorKind::ElicitationNotSupported).into());
622        }
623        Ok(())
624    }
625}
626
627impl Drop for Session {
628    fn drop(&mut self) {
629        // Cooperative shutdown: cancel the event loop's token to signal
630        // exit between iterations. The loop will see the cancellation on
631        // its next select poll and break cleanly without interrupting an
632        // in-flight handler. We do NOT abort the JoinHandle — that would
633        // land at any await point in the loop body, potentially leaving
634        // the CLI with an unanswered request id. RFD-400 review finding
635        // #3.
636        //
637        // The handle itself is left in `event_loop` to be reaped by the
638        // tokio runtime when it next polls; we intentionally don't await
639        // it here because Drop is sync.
640        self.shutdown.cancel();
641        self.client.unregister_session(&self.id);
642    }
643}
644
645/// UI sub-API for a [`Session`] — elicitation, confirmation, selection,
646/// and free-form input.
647///
648/// Acquired via [`Session::ui`]. Methods route to `session.ui.*` RPCs and
649/// require host elicitation support — check
650/// `session.capabilities().ui.elicitation` before use.
651pub struct SessionUi<'a> {
652    session: &'a Session,
653}
654
655impl<'a> SessionUi<'a> {
656    /// Request user input via an interactive UI form (elicitation).
657    ///
658    /// Sends a JSON Schema describing form fields to the CLI host. The host
659    /// renders a form dialog and returns the user's response.
660    ///
661    /// Prefer the typed convenience methods [`confirm`](Self::confirm),
662    /// [`select`](Self::select), and [`input`](Self::input) for common cases.
663    pub async fn elicitation(
664        &self,
665        message: &str,
666        schema: Value,
667    ) -> Result<ElicitationResult, Error> {
668        self.session.assert_elicitation()?;
669        let result = self
670            .session
671            .client
672            .call(
673                "session.ui.elicitation",
674                Some(serde_json::json!({
675                    "sessionId": self.session.id,
676                    "message": message,
677                    "requestedSchema": schema,
678                })),
679            )
680            .await?;
681        let elicitation: ElicitationResult = serde_json::from_value(result)?;
682        Ok(elicitation)
683    }
684
685    /// Ask the user a yes/no confirmation question.
686    ///
687    /// Returns `true` if the user accepted and confirmed, `false` otherwise.
688    pub async fn confirm(&self, message: &str) -> Result<bool, Error> {
689        self.session.assert_elicitation()?;
690        let schema = serde_json::json!({
691            "type": "object",
692            "properties": {
693                "confirmed": {
694                    "type": "boolean",
695                    "default": true,
696                }
697            },
698            "required": ["confirmed"]
699        });
700        let result = self.elicitation(message, schema).await?;
701        Ok(result.action == "accept"
702            && result
703                .content
704                .and_then(|c| c.get("confirmed").and_then(|v| v.as_bool()))
705                == Some(true))
706    }
707
708    /// Ask the user to select from a list of options.
709    ///
710    /// Returns the selected option string on accept, or `None` on decline/cancel.
711    pub async fn select(&self, message: &str, options: &[&str]) -> Result<Option<String>, Error> {
712        self.session.assert_elicitation()?;
713        let schema = serde_json::json!({
714            "type": "object",
715            "properties": {
716                "selection": {
717                    "type": "string",
718                    "enum": options,
719                }
720            },
721            "required": ["selection"]
722        });
723        let result = self.elicitation(message, schema).await?;
724        if result.action != "accept" {
725            return Ok(None);
726        }
727        let selection = result.content.and_then(|c| {
728            c.get("selection")
729                .and_then(|v| v.as_str())
730                .map(String::from)
731        });
732        Ok(selection)
733    }
734
735    /// Ask the user for free-form text input.
736    ///
737    /// Returns the input string on accept, or `None` on decline/cancel.
738    /// Use [`UiInputOptions`] to set validation constraints and field metadata.
739    pub async fn input(
740        &self,
741        message: &str,
742        options: Option<&UiInputOptions<'_>>,
743    ) -> Result<Option<String>, Error> {
744        self.session.assert_elicitation()?;
745        let mut field = serde_json::json!({ "type": "string" });
746        if let Some(opts) = options {
747            if let Some(title) = opts.title {
748                field["title"] = Value::String(title.to_string());
749            }
750            if let Some(desc) = opts.description {
751                field["description"] = Value::String(desc.to_string());
752            }
753            if let Some(min) = opts.min_length {
754                field["minLength"] = Value::Number(min.into());
755            }
756            if let Some(max) = opts.max_length {
757                field["maxLength"] = Value::Number(max.into());
758            }
759            if let Some(fmt) = &opts.format {
760                field["format"] = Value::String(fmt.as_str().to_string());
761            }
762            if let Some(default) = opts.default {
763                field["default"] = Value::String(default.to_string());
764            }
765        }
766        let schema = serde_json::json!({
767            "type": "object",
768            "properties": { "value": field },
769            "required": ["value"]
770        });
771        let result = self.elicitation(message, schema).await?;
772        if result.action != "accept" {
773            return Ok(None);
774        }
775        let value = result
776            .content
777            .and_then(|c| c.get("value").and_then(|v| v.as_str()).map(String::from));
778        Ok(value)
779    }
780}
781
782impl Client {
783    /// Create a new session on the CLI.
784    ///
785    /// Sends `session.create`, registers the session on the router,
786    /// and spawns an internal event loop that dispatches to the handler.
787    ///
788    /// All callbacks (per-event handlers, tool handlers, hooks, transform)
789    /// are configured via [`SessionConfig`] using its `with_*_handler` /
790    /// `with_tools` / `with_hooks` / `with_system_message_transform` builder
791    /// methods.
792    ///
793    /// If [`hooks_handler`](SessionConfig::hooks_handler) is set, the
794    /// wire-level `hooks` flag is automatically enabled.
795    ///
796    /// If [`system_message_transform`](SessionConfig::system_message_transform) is set, the SDK injects
797    /// `action: "transform"` sections into the [`SystemMessageConfig`] wire
798    /// format and handles `systemMessage.transform` RPC callbacks during
799    /// the session.
800    ///
801    /// Each per-event handler is independently optional. If a handler is
802    /// not installed, the SDK signals the runtime not to emit the matching
803    /// broadcast (and silently skips dispatch if one arrives anyway).
804    pub async fn create_session(&self, mut config: SessionConfig) -> Result<Session, Error> {
805        let total_start = Instant::now();
806        // For cloud sessions, let the CLI/server assign the session id and
807        // register the session lazily once the response arrives. For non-cloud
808        // sessions we generate the id client-side (when the caller didn't
809        // supply one) so the session can be registered BEFORE the RPC — the
810        // CLI may issue session-scoped requests (e.g. sessionFs.writeFile for
811        // workspace metadata) during session.create processing, before it has
812        // sent the response.
813        let caller_session_id = config.session_id.clone();
814        let use_server_generated_id = config.cloud.is_some() && caller_session_id.is_none();
815        let local_session_id: Option<SessionId> = if use_server_generated_id {
816            None
817        } else {
818            Some(
819                caller_session_id
820                    .clone()
821                    .unwrap_or_else(|| SessionId::new(uuid::Uuid::new_v4().to_string())),
822            )
823        };
824        if config.hooks_handler.is_some() && config.hooks.is_none() {
825            config.hooks = Some(true);
826        }
827        if let Some(transforms) = config.system_message_transform.clone() {
828            inject_transform_sections(&mut config, transforms.as_ref());
829        }
830        let mode = self.inner.mode;
831        if mode == crate::ClientMode::Empty && config.available_tools.is_none() {
832            return Err(Error::with_message(
833                ErrorKind::InvalidConfig,
834                "ClientMode::Empty requires available_tools to be set on the session config. \
835                 Use ToolSet to specify which tools the session may use (e.g. \
836                 ToolSet::new().add_builtin_many(BUILTIN_TOOLS_ISOLATED)).",
837            ));
838        }
839        crate::mode::validate_tool_filter_list(
840            "available_tools",
841            config.available_tools.as_deref(),
842        )?;
843        crate::mode::validate_tool_filter_list("excluded_tools", config.excluded_tools.as_deref())?;
844        config.system_message =
845            crate::mode::system_message_for_mode(mode, config.system_message.take());
846        config.memory = crate::mode::memory_for_mode(mode, config.memory.take());
847        if mode == crate::ClientMode::Empty {
848            if config.enable_session_telemetry.is_none() {
849                config.enable_session_telemetry = Some(false);
850            }
851            if config.skip_embedding_retrieval.is_none() {
852                config.skip_embedding_retrieval = Some(true);
853            }
854            if config.enable_on_demand_instruction_discovery.is_none() {
855                config.enable_on_demand_instruction_discovery = Some(false);
856            }
857            if config.enable_file_hooks.is_none() {
858                config.enable_file_hooks = Some(false);
859            }
860            if config.enable_host_git_operations.is_none() {
861                config.enable_host_git_operations = Some(false);
862            }
863            if config.enable_session_store.is_none() {
864                config.enable_session_store = Some(false);
865            }
866            if config.enable_skills.is_none() {
867                config.enable_skills = Some(false);
868            }
869        }
870        if mode == crate::ClientMode::Empty && config.mcp_oauth_token_storage.is_none() {
871            config.mcp_oauth_token_storage = Some("in-memory".into());
872        }
873        if mode == crate::ClientMode::Empty && config.embedding_cache_storage.is_none() {
874            config.embedding_cache_storage = Some("in-memory".into());
875        }
876        let opt_skip_custom_instructions = config.skip_custom_instructions;
877        let opt_custom_agents_local_only = config.custom_agents_local_only;
878        let opt_coauthor_enabled = config.coauthor_enabled;
879        let opt_manage_schedule_enabled = config.manage_schedule_enabled;
880        let (mut wire, mut runtime) = config.into_wire(local_session_id.clone())?;
881        wire.enable_github_telemetry_forwarding =
882            self.inner.on_github_telemetry.is_some().then_some(true);
883
884        let permission_handler = crate::permission::resolve_handler(
885            runtime.permission_handler.take(),
886            runtime.permission_policy.take(),
887        );
888        let handlers = SessionHandlers {
889            permission: permission_handler,
890            elicitation: runtime.elicitation_handler.take(),
891            mcp_auth: runtime.mcp_auth_handler.take(),
892            user_input: runtime.user_input_handler.take(),
893            exit_plan_mode: runtime.exit_plan_mode_handler.take(),
894            auto_mode_switch: runtime.auto_mode_switch_handler.take(),
895            tools: Arc::new(std::mem::take(&mut runtime.tool_handlers)),
896        };
897        let hooks = runtime.hooks_handler.take();
898        let transforms = runtime.system_message_transform.take();
899        let tools_count = wire.tools.as_ref().map_or(0, Vec::len);
900        let commands_count = runtime.commands.as_ref().map_or(0, Vec::len);
901        let has_hooks = hooks.is_some();
902        let command_handlers = build_command_handler_map(runtime.commands.as_deref());
903        let canvas_handler = runtime.canvas_handler.take();
904        let session_fs_provider = runtime.session_fs_provider.take();
905        let bearer_token_providers = std::mem::take(&mut runtime.bearer_token_providers);
906        let has_mcp_auth_handler = handlers.mcp_auth.is_some();
907        if self.inner.session_fs_configured && session_fs_provider.is_none() {
908            return Err(ErrorKind::Session(SessionErrorKind::SessionFsProviderRequired).into());
909        }
910        if self.inner.session_fs_sqlite_declared
911            && let Some(ref provider) = session_fs_provider
912            && provider.sqlite().is_none()
913        {
914            return Err(Error::with_message(
915                ErrorKind::InvalidConfig,
916                "SessionFs capabilities declare SQLite support but the provider \
917                 does not implement SessionFsSqliteProvider",
918            ));
919        }
920
921        let mut params = serde_json::to_value(&wire)?;
922        let trace_ctx = self.resolve_trace_context().await;
923        inject_trace_context(&mut params, &trace_ctx);
924
925        let setup_start = Instant::now();
926        let capabilities = Arc::new(parking_lot::RwLock::new(SessionCapabilities::default()));
927        let idle_waiter = Arc::new(ParkingLotMutex::new(None));
928        let open_canvases = Arc::new(parking_lot::RwLock::new(Vec::new()));
929        let shutdown = CancellationToken::new();
930        let (event_tx, _) = tokio::sync::broadcast::channel(512);
931
932        // For cloud sessions (use_server_generated_id), defer session
933        // registration to the inline callback so the read task registers
934        // the session synchronously the instant the response arrives.
935        // For non-cloud sessions, register up-front so the CLI can issue
936        // session-scoped requests during session.create processing.
937        let inline_stash: Arc<
938            ParkingLotMutex<Option<(SessionId, crate::router::SessionChannels)>>,
939        > = Arc::new(ParkingLotMutex::new(None));
940
941        let inline_callback: Option<crate::jsonrpc::InlineResponseCallback> = if let Some(ref sid) =
942            local_session_id
943        {
944            let channels = self.register_session(sid);
945            *inline_stash.lock() = Some((sid.clone(), channels));
946            None
947        } else {
948            let client = self.clone();
949            let stash = inline_stash.clone();
950            let expected = caller_session_id.clone();
951            Some(Box::new(move |response| {
952                let result = response.result.as_ref().ok_or_else(|| {
953                    Error::with_message(ErrorKind::Json, "session.create response had no result")
954                })?;
955                let parsed: CreateSessionResult =
956                    serde_json::from_value(result.clone()).map_err(Error::from)?;
957                if let Some(requested) = expected.as_ref()
958                    && parsed.session_id != *requested
959                {
960                    return Err(ErrorKind::Session(SessionErrorKind::SessionIdMismatch {
961                        requested: requested.clone(),
962                        returned: parsed.session_id,
963                    })
964                    .into());
965                }
966                let channels = client.register_session(&parsed.session_id);
967                *stash.lock() = Some((parsed.session_id, channels));
968                Ok(())
969            }))
970        };
971
972        let rpc_start = Instant::now();
973        let result = match self
974            .call_with_inline_callback("session.create", Some(params), inline_callback)
975            .await
976        {
977            Ok(result) => result,
978            Err(error) => {
979                if let Some((id, _channels)) = inline_stash.lock().take() {
980                    self.unregister_session(&id);
981                }
982                return Err(error);
983            }
984        };
985        tracing::debug!(
986            elapsed_ms = rpc_start.elapsed().as_millis(),
987            "Client::create_session session creation request completed successfully"
988        );
989        let create_result: CreateSessionResult = match serde_json::from_value(result) {
990            Ok(result) => result,
991            Err(error) => {
992                if let Some((id, _channels)) = inline_stash.lock().take() {
993                    self.unregister_session(&id);
994                }
995                return Err(error.into());
996            }
997        };
998
999        if let Some(ref requested) = local_session_id
1000            && create_result.session_id != *requested
1001        {
1002            if let Some((id, _channels)) = inline_stash.lock().take() {
1003                self.unregister_session(&id);
1004            }
1005            return Err(ErrorKind::Session(SessionErrorKind::SessionIdMismatch {
1006                requested: requested.clone(),
1007                returned: create_result.session_id.clone(),
1008            })
1009            .into());
1010        }
1011
1012        let (session_id, channels) = inline_stash
1013            .lock()
1014            .take()
1015            .expect("session registration must have populated stash on success");
1016        let event_loop = spawn_event_loop(
1017            session_id.clone(),
1018            self.clone(),
1019            handlers,
1020            hooks,
1021            transforms,
1022            command_handlers,
1023            canvas_handler,
1024            session_fs_provider,
1025            bearer_token_providers,
1026            channels,
1027            idle_waiter.clone(),
1028            capabilities.clone(),
1029            open_canvases.clone(),
1030            event_tx.clone(),
1031            shutdown.clone(),
1032        );
1033        tracing::debug!(
1034            elapsed_ms = setup_start.elapsed().as_millis(),
1035            session_id = %session_id,
1036            tools_count,
1037            commands_count,
1038            has_hooks,
1039            "Client::create_session local setup complete"
1040        );
1041        *capabilities.write() = create_result.capabilities.unwrap_or_default();
1042        if has_mcp_auth_handler {
1043            register_mcp_auth_interest(self, &session_id).await?;
1044        }
1045
1046        tracing::debug!(
1047            elapsed_ms = total_start.elapsed().as_millis(),
1048            session_id = %session_id,
1049            "Client::create_session complete"
1050        );
1051        let session = Session {
1052            id: session_id,
1053            cwd: self.cwd().clone(),
1054            workspace_path: create_result.workspace_path,
1055            remote_url: create_result.remote_url,
1056            client: self.clone(),
1057            event_loop: ParkingLotMutex::new(Some(event_loop)),
1058            shutdown,
1059            idle_waiter,
1060            capabilities,
1061            open_canvases,
1062            event_tx,
1063        };
1064        apply_mode_post_create_patch(
1065            &session,
1066            mode,
1067            opt_skip_custom_instructions,
1068            opt_custom_agents_local_only,
1069            opt_coauthor_enabled,
1070            opt_manage_schedule_enabled,
1071        )
1072        .await?;
1073        Ok(session)
1074    }
1075
1076    /// Resume an existing session on the CLI.
1077    ///
1078    /// Sends `session.resume` and `session.skills.reload`, registers the
1079    /// session on the router, and spawns the event loop.
1080    ///
1081    /// All callbacks (event handler, hooks, transform) are configured
1082    /// via [`ResumeSessionConfig`] using its `with_*` builder methods.
1083    ///
1084    /// See [`Self::create_session`] for the defaults applied when callback
1085    /// fields are unset.
1086    pub async fn resume_session(&self, mut config: ResumeSessionConfig) -> Result<Session, Error> {
1087        let total_start = Instant::now();
1088        let session_id = config.session_id.clone();
1089        if config.hooks_handler.is_some() && config.hooks.is_none() {
1090            config.hooks = Some(true);
1091        }
1092        if let Some(transforms) = config.system_message_transform.clone() {
1093            inject_transform_sections_resume(&mut config, transforms.as_ref());
1094        }
1095        let mode = self.inner.mode;
1096        if mode == crate::ClientMode::Empty && config.available_tools.is_none() {
1097            return Err(Error::with_message(
1098                ErrorKind::InvalidConfig,
1099                "ClientMode::Empty requires available_tools to be set on the session config. \
1100                 Use ToolSet to specify which tools the session may use (e.g. \
1101                 ToolSet::new().add_builtin_many(BUILTIN_TOOLS_ISOLATED)).",
1102            ));
1103        }
1104        crate::mode::validate_tool_filter_list(
1105            "available_tools",
1106            config.available_tools.as_deref(),
1107        )?;
1108        crate::mode::validate_tool_filter_list("excluded_tools", config.excluded_tools.as_deref())?;
1109        config.system_message =
1110            crate::mode::system_message_for_mode(mode, config.system_message.take());
1111        config.memory = crate::mode::memory_for_mode(mode, config.memory.take());
1112        if mode == crate::ClientMode::Empty {
1113            if config.enable_session_telemetry.is_none() {
1114                config.enable_session_telemetry = Some(false);
1115            }
1116            if config.skip_embedding_retrieval.is_none() {
1117                config.skip_embedding_retrieval = Some(true);
1118            }
1119            if config.enable_on_demand_instruction_discovery.is_none() {
1120                config.enable_on_demand_instruction_discovery = Some(false);
1121            }
1122            if config.enable_file_hooks.is_none() {
1123                config.enable_file_hooks = Some(false);
1124            }
1125            if config.enable_host_git_operations.is_none() {
1126                config.enable_host_git_operations = Some(false);
1127            }
1128            if config.enable_session_store.is_none() {
1129                config.enable_session_store = Some(false);
1130            }
1131            if config.enable_skills.is_none() {
1132                config.enable_skills = Some(false);
1133            }
1134        }
1135        if mode == crate::ClientMode::Empty && config.mcp_oauth_token_storage.is_none() {
1136            config.mcp_oauth_token_storage = Some("in-memory".into());
1137        }
1138        if mode == crate::ClientMode::Empty && config.embedding_cache_storage.is_none() {
1139            config.embedding_cache_storage = Some("in-memory".into());
1140        }
1141        let opt_skip_custom_instructions = config.skip_custom_instructions;
1142        let opt_custom_agents_local_only = config.custom_agents_local_only;
1143        let opt_coauthor_enabled = config.coauthor_enabled;
1144        let opt_manage_schedule_enabled = config.manage_schedule_enabled;
1145        let (mut wire, mut runtime) = config.into_wire()?;
1146        wire.enable_github_telemetry_forwarding =
1147            self.inner.on_github_telemetry.is_some().then_some(true);
1148
1149        let permission_handler = crate::permission::resolve_handler(
1150            runtime.permission_handler.take(),
1151            runtime.permission_policy.take(),
1152        );
1153        let handlers = SessionHandlers {
1154            permission: permission_handler,
1155            elicitation: runtime.elicitation_handler.take(),
1156            mcp_auth: runtime.mcp_auth_handler.take(),
1157            user_input: runtime.user_input_handler.take(),
1158            exit_plan_mode: runtime.exit_plan_mode_handler.take(),
1159            auto_mode_switch: runtime.auto_mode_switch_handler.take(),
1160            tools: Arc::new(std::mem::take(&mut runtime.tool_handlers)),
1161        };
1162        let hooks = runtime.hooks_handler.take();
1163        let transforms = runtime.system_message_transform.take();
1164        let tools_count = wire.tools.as_ref().map_or(0, Vec::len);
1165        let commands_count = runtime.commands.as_ref().map_or(0, Vec::len);
1166        let has_hooks = hooks.is_some();
1167        let command_handlers = build_command_handler_map(runtime.commands.as_deref());
1168        let canvas_handler = runtime.canvas_handler.take();
1169        let session_fs_provider = runtime.session_fs_provider.take();
1170        let bearer_token_providers = std::mem::take(&mut runtime.bearer_token_providers);
1171        let has_mcp_auth_handler = handlers.mcp_auth.is_some();
1172        if self.inner.session_fs_configured && session_fs_provider.is_none() {
1173            return Err(ErrorKind::Session(SessionErrorKind::SessionFsProviderRequired).into());
1174        }
1175        if self.inner.session_fs_sqlite_declared
1176            && let Some(ref provider) = session_fs_provider
1177            && provider.sqlite().is_none()
1178        {
1179            return Err(Error::with_message(
1180                ErrorKind::InvalidConfig,
1181                "SessionFs capabilities declare SQLite support but the provider \
1182                 does not implement SessionFsSqliteProvider",
1183            ));
1184        }
1185
1186        let mut params = serde_json::to_value(&wire)?;
1187        let trace_ctx = self.resolve_trace_context().await;
1188        inject_trace_context(&mut params, &trace_ctx);
1189
1190        let capabilities = Arc::new(parking_lot::RwLock::new(SessionCapabilities::default()));
1191        let setup_start = Instant::now();
1192        let channels = self.register_session(&session_id);
1193        let idle_waiter = Arc::new(ParkingLotMutex::new(None));
1194        let open_canvases = Arc::new(parking_lot::RwLock::new(Vec::new()));
1195        let shutdown = CancellationToken::new();
1196        let (event_tx, _) = tokio::sync::broadcast::channel(512);
1197        let event_loop = spawn_event_loop(
1198            session_id.clone(),
1199            self.clone(),
1200            handlers,
1201            hooks,
1202            transforms,
1203            command_handlers,
1204            canvas_handler,
1205            session_fs_provider,
1206            bearer_token_providers,
1207            channels,
1208            idle_waiter.clone(),
1209            capabilities.clone(),
1210            open_canvases.clone(),
1211            event_tx.clone(),
1212            shutdown.clone(),
1213        );
1214        let mut registration =
1215            PendingSessionRegistration::new(self.clone(), session_id.clone(), shutdown.clone());
1216        tracing::debug!(
1217            elapsed_ms = setup_start.elapsed().as_millis(),
1218            session_id = %session_id,
1219            tools_count,
1220            commands_count,
1221            has_hooks,
1222            "Client::resume_session local setup complete"
1223        );
1224
1225        let rpc_start = Instant::now();
1226        let result = match self.call("session.resume", Some(params)).await {
1227            Ok(result) => result,
1228            Err(error) => {
1229                registration.cleanup(event_loop).await;
1230                return Err(error);
1231            }
1232        };
1233        tracing::debug!(
1234            elapsed_ms = rpc_start.elapsed().as_millis(),
1235            session_id = %session_id,
1236            "Client::resume_session session resume request completed successfully"
1237        );
1238
1239        let resume_result: ResumeSessionResult = match serde_json::from_value(result) {
1240            Ok(result) => result,
1241            Err(error) => {
1242                registration.cleanup(event_loop).await;
1243                return Err(error.into());
1244            }
1245        };
1246        let cli_session_id = resume_result
1247            .session_id
1248            .clone()
1249            .unwrap_or_else(|| session_id.clone());
1250        if cli_session_id != session_id {
1251            registration.cleanup(event_loop).await;
1252            return Err(ErrorKind::Session(SessionErrorKind::SessionIdMismatch {
1253                requested: session_id,
1254                returned: cli_session_id,
1255            })
1256            .into());
1257        }
1258        if has_mcp_auth_handler {
1259            register_mcp_auth_interest(self, &session_id).await?;
1260        }
1261
1262        // Reload skills after resume (best-effort).
1263        let skills_reload_start = Instant::now();
1264        if let Err(e) = self
1265            .call(
1266                "session.skills.reload",
1267                Some(serde_json::json!({ "sessionId": session_id })),
1268            )
1269            .await
1270        {
1271            warn!(
1272                elapsed_ms = skills_reload_start.elapsed().as_millis(),
1273                session_id = %session_id,
1274                error = %e,
1275                "Client::resume_session skills reload request failed"
1276            );
1277        } else {
1278            tracing::debug!(
1279                elapsed_ms = skills_reload_start.elapsed().as_millis(),
1280                session_id = %session_id,
1281                "Client::resume_session skills reload request completed successfully"
1282            );
1283        }
1284
1285        *capabilities.write() = resume_result.capabilities.unwrap_or_default();
1286        // Upsert resume snapshots rather than replacing wholesale. Live
1287        // `session.canvas.opened` notifications can arrive on the event loop
1288        // while `session.resume` is in flight; a wholesale replace would
1289        // discard those updates.
1290        {
1291            let mut snapshots = open_canvases.write();
1292            for snapshot in resume_result.open_canvases.unwrap_or_default() {
1293                upsert_open_canvas_snapshot(&mut snapshots, snapshot);
1294            }
1295        }
1296
1297        tracing::debug!(
1298            elapsed_ms = total_start.elapsed().as_millis(),
1299            session_id = %session_id,
1300            "Client::resume_session complete"
1301        );
1302        registration.disarm();
1303        let session = Session {
1304            id: session_id,
1305            cwd: self.cwd().clone(),
1306            workspace_path: resume_result.workspace_path,
1307            remote_url: resume_result.remote_url,
1308            client: self.clone(),
1309            event_loop: ParkingLotMutex::new(Some(event_loop)),
1310            shutdown,
1311            idle_waiter,
1312            capabilities,
1313            open_canvases,
1314            event_tx,
1315        };
1316        apply_mode_post_create_patch(
1317            &session,
1318            mode,
1319            opt_skip_custom_instructions,
1320            opt_custom_agents_local_only,
1321            opt_coauthor_enabled,
1322            opt_manage_schedule_enabled,
1323        )
1324        .await?;
1325        Ok(session)
1326    }
1327}
1328
1329type CommandHandlerMap = HashMap<String, Arc<dyn CommandHandler>>;
1330
1331async fn apply_mode_post_create_patch(
1332    session: &Session,
1333    mode: crate::ClientMode,
1334    opt_skip_custom_instructions: Option<bool>,
1335    opt_custom_agents_local_only: Option<bool>,
1336    opt_coauthor_enabled: Option<bool>,
1337    opt_manage_schedule_enabled: Option<bool>,
1338) -> Result<(), Error> {
1339    use crate::generated::api_types::SessionUpdateOptionsParams;
1340    let mut patch = SessionUpdateOptionsParams::default();
1341    let should_send = if mode == crate::ClientMode::Empty {
1342        patch.skip_custom_instructions = Some(opt_skip_custom_instructions.unwrap_or(true));
1343        patch.custom_agents_local_only = Some(opt_custom_agents_local_only.unwrap_or(true));
1344        patch.coauthor_enabled = Some(opt_coauthor_enabled.unwrap_or(false));
1345        patch.manage_schedule_enabled = Some(opt_manage_schedule_enabled.unwrap_or(false));
1346        patch.installed_plugins = Some(Vec::new());
1347        true
1348    } else {
1349        let mut any = false;
1350        if let Some(v) = opt_skip_custom_instructions {
1351            patch.skip_custom_instructions = Some(v);
1352            any = true;
1353        }
1354        if let Some(v) = opt_custom_agents_local_only {
1355            patch.custom_agents_local_only = Some(v);
1356            any = true;
1357        }
1358        if let Some(v) = opt_coauthor_enabled {
1359            patch.coauthor_enabled = Some(v);
1360            any = true;
1361        }
1362        if let Some(v) = opt_manage_schedule_enabled {
1363            patch.manage_schedule_enabled = Some(v);
1364            any = true;
1365        }
1366        any
1367    };
1368    if !should_send {
1369        return Ok(());
1370    }
1371    if let Err(error) = session.rpc().options().update(patch).await {
1372        let _ = session.disconnect().await;
1373        return Err(error);
1374    }
1375    Ok(())
1376}
1377
1378fn build_command_handler_map(commands: Option<&[CommandDefinition]>) -> Arc<CommandHandlerMap> {
1379    let map = match commands {
1380        Some(commands) => commands
1381            .iter()
1382            .filter(|cmd| !cmd.name.is_empty())
1383            .map(|cmd| (cmd.name.clone(), cmd.handler.clone()))
1384            .collect(),
1385        None => HashMap::new(),
1386    };
1387    Arc::new(map)
1388}
1389
1390fn upsert_open_canvas_snapshot(
1391    snapshots: &mut Vec<OpenCanvasInstance>,
1392    snapshot: OpenCanvasInstance,
1393) {
1394    if let Some(existing) = snapshots
1395        .iter_mut()
1396        .find(|open| open.instance_id == snapshot.instance_id)
1397    {
1398        *existing = snapshot;
1399    } else {
1400        snapshots.push(snapshot);
1401    }
1402}
1403
1404fn remove_open_canvas_snapshot(snapshots: &mut Vec<OpenCanvasInstance>, instance_id: &str) {
1405    snapshots.retain(|open| open.instance_id != instance_id);
1406}
1407
1408#[allow(clippy::too_many_arguments)]
1409fn spawn_event_loop(
1410    session_id: SessionId,
1411    client: Client,
1412    handlers: SessionHandlers,
1413    hooks: Option<Arc<dyn SessionHooks>>,
1414    transforms: Option<Arc<dyn SystemMessageTransform>>,
1415    command_handlers: Arc<CommandHandlerMap>,
1416    canvas_handler: Option<Arc<dyn CanvasHandler>>,
1417    session_fs_provider: Option<Arc<dyn SessionFsProvider>>,
1418    bearer_token_providers: HashMap<String, Arc<dyn BearerTokenProvider>>,
1419    channels: crate::router::SessionChannels,
1420    idle_waiter: Arc<ParkingLotMutex<Option<IdleWaiter>>>,
1421    capabilities: Arc<parking_lot::RwLock<SessionCapabilities>>,
1422    open_canvases: Arc<parking_lot::RwLock<Vec<OpenCanvasInstance>>>,
1423    event_tx: tokio::sync::broadcast::Sender<SessionEvent>,
1424    shutdown: CancellationToken,
1425) -> JoinHandle<()> {
1426    let crate::router::SessionChannels {
1427        mut notifications,
1428        mut requests,
1429    } = channels;
1430
1431    let span = tracing::error_span!("session_event_loop", session_id = %session_id);
1432    tokio::spawn(
1433        async move {
1434            loop {
1435                // `mpsc::UnboundedReceiver::recv` and
1436                // `CancellationToken::cancelled` are both cancel-safe per
1437                // RFD 400. The selected branch's `await`'d handler is
1438                // *not* mid-cancelled by the select — once a branch fires
1439                // it runs to completion within the loop's iteration.
1440                // Spawned child tasks inside `handle_notification`
1441                // (permission/tool/elicitation callbacks) intentionally
1442                // outlive the parent loop and own their own cleanup;
1443                // this is RFD 400's "spawn background tasks to perform
1444                // cancel-unsafe operations" pattern and is correct as-is.
1445                tokio::select! {
1446                    _ = shutdown.cancelled() => break,
1447                    Some(notification) = notifications.recv() => {
1448                        handle_notification(
1449                            &session_id, &client, &handlers, &command_handlers, notification, &idle_waiter, &capabilities, &open_canvases, &event_tx,
1450                        ).await;
1451                    }
1452                    Some(request) = requests.recv() => {
1453                        let ctx = RequestDispatchContext {
1454                            client: &client,
1455                            handlers: &handlers,
1456                            hooks: hooks.as_deref(),
1457                            transforms: transforms.as_deref(),
1458                            canvas_handler: canvas_handler.as_ref(),
1459                            session_fs_provider: session_fs_provider.as_ref(),
1460                            bearer_token_providers: &bearer_token_providers,
1461                        };
1462                        handle_request(&session_id, ctx, request).await;
1463                    }
1464                    else => break,
1465                }
1466            }
1467            // Channels closed or shutdown signaled — fail any pending
1468            // send_and_wait so the caller observes a clean error.
1469            if let Some(waiter) = idle_waiter.lock().take() {
1470                let _ = waiter
1471                    .tx
1472                    .send(Err(ErrorKind::Session(SessionErrorKind::EventLoopClosed).into()));
1473            }
1474        }
1475        .instrument(span),
1476    )
1477}
1478
1479fn extract_request_id(data: &Value) -> Option<RequestId> {
1480    data.get("requestId")
1481        .and_then(|v| v.as_str())
1482        .filter(|s| !s.is_empty())
1483        .map(RequestId::new)
1484}
1485
1486/// Map a [`PermissionResult`] to the `result` payload sent back to the
1487/// server via `session.permissions.handlePendingPermissionRequest`.
1488///
1489/// Returns `None` when the SDK must not send a response.
1490fn notification_permission_payload(result: &PermissionResult) -> Option<Value> {
1491    match result {
1492        PermissionResult::NoResult => None,
1493        PermissionResult::Decision(decision) => Some(
1494            serde_json::to_value(decision).expect("serializing permission decision should succeed"),
1495        ),
1496    }
1497}
1498
1499async fn register_mcp_auth_interest(client: &Client, session_id: &SessionId) -> Result<(), Error> {
1500    let mut params = serde_json::to_value(RegisterEventInterestParams {
1501        event_type: "mcp.oauth_required".to_string(),
1502    })?;
1503    params["sessionId"] = Value::String(session_id.to_string());
1504    client
1505        .call(rpc_methods::SESSION_EVENTLOG_REGISTERINTEREST, Some(params))
1506        .await?;
1507    Ok(())
1508}
1509
1510fn tool_failure_result(message: impl Into<String>) -> ToolResult {
1511    let message = message.into();
1512    ToolResult::Expanded(ToolResultExpanded {
1513        text_result_for_llm: message.clone(),
1514        result_type: "failure".to_string(),
1515        binary_results_for_llm: None,
1516        session_log: None,
1517        error: Some(message),
1518        tool_telemetry: None,
1519    })
1520}
1521
1522/// Process a notification from the CLI's broadcast channel.
1523#[allow(clippy::too_many_arguments)]
1524async fn handle_notification(
1525    session_id: &SessionId,
1526    client: &Client,
1527    handlers: &SessionHandlers,
1528    command_handlers: &Arc<CommandHandlerMap>,
1529    notification: SessionEventNotification,
1530    idle_waiter: &Arc<ParkingLotMutex<Option<IdleWaiter>>>,
1531    capabilities: &Arc<parking_lot::RwLock<SessionCapabilities>>,
1532    open_canvases: &Arc<parking_lot::RwLock<Vec<OpenCanvasInstance>>>,
1533    event_tx: &tokio::sync::broadcast::Sender<SessionEvent>,
1534) {
1535    let dispatch_start = Instant::now();
1536    let event = notification.event.clone();
1537    let event_type = event.parsed_type();
1538    if event_type == SessionEventType::PermissionRequested {
1539        tracing::debug!(
1540            session_id = %session_id,
1541            event_type = %event.event_type,
1542            "Session::handle_notification permission request received"
1543        );
1544    }
1545
1546    // Signal send_and_wait if active. The lock is only contended when
1547    // a send_and_wait call is in flight (idle_waiter is Some).
1548    match event_type {
1549        SessionEventType::AssistantMessage
1550        | SessionEventType::SessionIdle
1551        | SessionEventType::SessionError => {
1552            let mut guard = idle_waiter.lock();
1553            if let Some(waiter) = guard.as_mut() {
1554                match event_type {
1555                    SessionEventType::AssistantMessage => {
1556                        if !waiter.first_assistant_message_seen {
1557                            waiter.first_assistant_message_seen = true;
1558                            tracing::debug!(
1559                                elapsed_ms = waiter.started_at.elapsed().as_millis(),
1560                                session_id = %session_id,
1561                                "Session::send_and_wait first assistant message"
1562                            );
1563                        }
1564                        waiter.last_assistant_message = Some(event.clone());
1565                    }
1566                    SessionEventType::SessionIdle | SessionEventType::SessionError => {
1567                        if let Some(waiter) = guard.take() {
1568                            if event_type == SessionEventType::SessionIdle {
1569                                tracing::debug!(
1570                                    elapsed_ms = waiter.started_at.elapsed().as_millis(),
1571                                    session_id = %session_id,
1572                                    "Session::send_and_wait idle received"
1573                                );
1574                                let _ = waiter.tx.send(Ok(waiter.last_assistant_message));
1575                            } else {
1576                                let error_msg = event
1577                                    .typed_data::<SessionErrorData>()
1578                                    .map(|d| d.message)
1579                                    .or_else(|| {
1580                                        event
1581                                            .data
1582                                            .get("message")
1583                                            .and_then(|v| v.as_str())
1584                                            .map(|s| s.to_string())
1585                                    })
1586                                    .unwrap_or_else(|| "session error".to_string());
1587                                let _ = waiter.tx.send(Err(Error::with_message(
1588                                    ErrorKind::Session(SessionErrorKind::AgentError),
1589                                    error_msg,
1590                                )));
1591                            }
1592                        }
1593                    }
1594                    _ => {}
1595                }
1596            }
1597        }
1598        _ => {}
1599    }
1600
1601    // Update the snapshot caches BEFORE broadcasting so subscribers that
1602    // call `Session::capabilities()` / `Session::open_canvases()` in
1603    // response to the event observe the new state.
1604    if event_type == SessionEventType::CapabilitiesChanged {
1605        match serde_json::from_value::<SessionCapabilities>(notification.event.data.clone()) {
1606            Ok(changed) => *capabilities.write() = changed,
1607            Err(e) => warn!(error = %e, "failed to deserialize capabilities.changed payload"),
1608        }
1609    }
1610    if event_type == SessionEventType::SessionCanvasOpened {
1611        match serde_json::from_value::<OpenCanvasInstance>(notification.event.data.clone()) {
1612            Ok(open_canvas) => {
1613                upsert_open_canvas_snapshot(&mut open_canvases.write(), open_canvas);
1614            }
1615            Err(e) => warn!(error = %e, "failed to deserialize session.canvas.opened payload"),
1616        }
1617    }
1618    if event_type == SessionEventType::SessionCanvasClosed {
1619        match serde_json::from_value::<SessionCanvasClosedData>(notification.event.data.clone()) {
1620            Ok(closed) => {
1621                if closed.instance_id.is_empty() {
1622                    warn!("failed to deserialize session.canvas.closed payload");
1623                } else {
1624                    remove_open_canvas_snapshot(&mut open_canvases.write(), &closed.instance_id);
1625                }
1626            }
1627            Err(e) => warn!(error = %e, "failed to deserialize session.canvas.closed payload"),
1628        }
1629    }
1630
1631    // Fan out the event to runtime subscribers (`Session::subscribe`). `send`
1632    // only errors when there are no receivers, which is the normal case
1633    // before any consumer subscribes.
1634    let _ = event_tx.send(event.clone());
1635
1636    tracing::debug!(
1637        elapsed_ms = dispatch_start.elapsed().as_millis(),
1638        session_id = %session_id,
1639        event_type = %notification.event.event_type,
1640        "Session::handle_notification dispatch"
1641    );
1642
1643    // Notification-based permission/tool/elicitation requests require a
1644    // separate RPC callback. Spawn concurrently since the CLI doesn't block.
1645    match event_type {
1646        SessionEventType::PermissionRequested => {
1647            let Some(request_id) = extract_request_id(&notification.event.data) else {
1648                return;
1649            };
1650            // Honor the runtime's `resolvedByHook` signal — when the
1651            // server has already resolved the permission via a hook,
1652            // clients must not send a second response.
1653            if notification
1654                .event
1655                .data
1656                .get("resolvedByHook")
1657                .and_then(|v| v.as_bool())
1658                .unwrap_or(false)
1659            {
1660                return;
1661            }
1662            // Multi-client safety: if this client has no permission
1663            // handler installed, don't respond — another client on the
1664            // same CLI may handle it.
1665            let Some(permission_handler) = handlers.permission.clone() else {
1666                return;
1667            };
1668            let client = client.clone();
1669            let sid = session_id.clone();
1670            let data: PermissionRequestData =
1671                serde_json::from_value(notification.event.data.clone()).unwrap_or_else(|_| {
1672                    PermissionRequestData {
1673                        kind: None,
1674                        tool_call_id: None,
1675                        extra: notification.event.data.clone(),
1676                    }
1677                });
1678            let span = tracing::error_span!(
1679                "permission_request_handler",
1680                session_id = %sid,
1681                request_id = %request_id
1682            );
1683            tokio::spawn(
1684                async move {
1685                    let handler_start = Instant::now();
1686                    let result = permission_handler
1687                        .handle(sid.clone(), request_id.clone(), data)
1688                        .await;
1689                    tracing::debug!(
1690                        elapsed_ms = handler_start.elapsed().as_millis(),
1691                        session_id = %sid,
1692                        request_id = %request_id,
1693                        "PermissionHandler::handle dispatch"
1694                    );
1695                    let Some(result_value) = notification_permission_payload(&result) else {
1696                        // Handler returned Deferred / NoResult — it will
1697                        // call handlePendingPermissionRequest itself (or
1698                        // leave the request unanswered).
1699                        return;
1700                    };
1701                    let rpc_start = Instant::now();
1702                    let _ = client
1703                        .call(
1704                            "session.permissions.handlePendingPermissionRequest",
1705                            Some(serde_json::json!({
1706                                "sessionId": sid,
1707                                "requestId": request_id,
1708                                "result": result_value,
1709                            })),
1710                        )
1711                        .await;
1712                    tracing::debug!(
1713                        elapsed_ms = rpc_start.elapsed().as_millis(),
1714                        session_id = %sid,
1715                        request_id = %request_id,
1716                        "Session::handle_notification response sent successfully"
1717                    );
1718                }
1719                .instrument(span),
1720            );
1721        }
1722        SessionEventType::ExternalToolRequested => {
1723            let Some(request_id) = extract_request_id(&notification.event.data) else {
1724                return;
1725            };
1726            let data: ExternalToolRequestedData =
1727                match serde_json::from_value(notification.event.data.clone()) {
1728                    Ok(d) => d,
1729                    Err(e) => {
1730                        warn!(error = %e, "failed to deserialize external_tool.requested");
1731                        let client = client.clone();
1732                        let sid = session_id.clone();
1733                        let span = tracing::error_span!(
1734                            "external_tool_deserialize_error",
1735                            session_id = %sid,
1736                            request_id = %request_id
1737                        );
1738                        tokio::spawn(
1739                            async move {
1740                                let rpc_start = Instant::now();
1741                                let _ = client
1742                                .call(
1743                                    "session.tools.handlePendingToolCall",
1744                                    Some(serde_json::json!({
1745                                        "sessionId": sid,
1746                                        "requestId": request_id,
1747                                        "error": format!("Failed to deserialize tool request: {e}"),
1748                                    })),
1749                                )
1750                                .await;
1751                                tracing::debug!(
1752                                    elapsed_ms = rpc_start.elapsed().as_millis(),
1753                                    session_id = %sid,
1754                                    request_id = %request_id,
1755                                    "Session::handle_notification response sent successfully"
1756                                );
1757                            }
1758                            .instrument(span),
1759                        );
1760                        return;
1761                    }
1762                };
1763            // Multi-client safety: look up a handler for the requested
1764            // tool name. If this client has no handler installed for that
1765            // tool, don't respond — another connected client may have one.
1766            let tool_handler = if data.tool_name.is_empty() {
1767                None
1768            } else {
1769                handlers.tools.get(&data.tool_name).cloned()
1770            };
1771            let Some(tool_handler) = tool_handler else {
1772                return;
1773            };
1774            let client = client.clone();
1775            let sid = session_id.clone();
1776            let span = tracing::error_span!(
1777                "external_tool_handler",
1778                session_id = %sid,
1779                request_id = %request_id
1780            );
1781            tokio::spawn(
1782                async move {
1783                    // `tool_name.is_empty()` would have produced a `None`
1784                    // lookup in `handlers.tools` and short-circuited at the
1785                    // outer guard above, so only the tool_call_id check is
1786                    // reachable here.
1787                    if data.tool_call_id.is_empty() {
1788                        let error_msg = "Missing toolCallId";
1789                        let rpc_start = Instant::now();
1790                        let _ = client
1791                            .call(
1792                                "session.tools.handlePendingToolCall",
1793                                Some(serde_json::json!({
1794                                    "sessionId": sid,
1795                                    "requestId": request_id,
1796                                    "error": error_msg,
1797                                })),
1798                            )
1799                            .await;
1800                        tracing::debug!(
1801                            elapsed_ms = rpc_start.elapsed().as_millis(),
1802                            session_id = %sid,
1803                            request_id = %request_id,
1804                            "Session::handle_notification response sent successfully"
1805                        );
1806                        return;
1807                    }
1808                    let tool_call_id = data.tool_call_id.clone();
1809                    let tool_name = data.tool_name.clone();
1810                    let invocation = ToolInvocation {
1811                        session_id: sid.clone(),
1812                        tool_call_id: data.tool_call_id,
1813                        tool_name: data.tool_name,
1814                        arguments: data
1815                            .arguments
1816                            .unwrap_or(Value::Object(serde_json::Map::new())),
1817                        traceparent: data.traceparent,
1818                        tracestate: data.tracestate,
1819                    };
1820                    let handler_start = Instant::now();
1821                    let tool_result = match tool_handler.call(invocation).await {
1822                        Ok(r) => r,
1823                        Err(e) => tool_failure_result(e.to_string()),
1824                    };
1825                    tracing::debug!(
1826                        elapsed_ms = handler_start.elapsed().as_millis(),
1827                        session_id = %sid,
1828                        request_id = %request_id,
1829                        tool_call_id = %tool_call_id,
1830                        tool_name = %tool_name,
1831                        "ToolHandler::call dispatch"
1832                    );
1833                    let result_value = serde_json::to_value(tool_result).unwrap_or(Value::Null);
1834                    let rpc_start = Instant::now();
1835                    let _ = client
1836                        .call(
1837                            "session.tools.handlePendingToolCall",
1838                            Some(serde_json::json!({
1839                                "sessionId": sid,
1840                                "requestId": request_id,
1841                                "result": result_value,
1842                            })),
1843                        )
1844                        .await;
1845                    tracing::debug!(
1846                        elapsed_ms = rpc_start.elapsed().as_millis(),
1847                        session_id = %sid,
1848                        request_id = %request_id,
1849                        tool_call_id = %tool_call_id,
1850                        tool_name = %tool_name,
1851                        "Session::handle_notification response sent successfully"
1852                    );
1853                }
1854                .instrument(span),
1855            );
1856        }
1857        SessionEventType::UserInputRequested => {
1858            // Notification-only signal for observers (UI, telemetry).
1859            // The CLI follows up with a `userInput.request` JSON-RPC call
1860            // that drives the `UserInputHandler` dispatch — handling
1861            // the notification here too would double-fire the handler
1862            // and produce duplicate prompts on the consumer side. See
1863            // github/github-app#4249.
1864        }
1865        SessionEventType::ElicitationRequested => {
1866            let Some(request_id) = extract_request_id(&notification.event.data) else {
1867                return;
1868            };
1869            // Multi-client safety: if this client has no elicitation
1870            // handler installed, don't respond — another client on the
1871            // same CLI may handle it.
1872            let Some(elicitation_handler) = handlers.elicitation.clone() else {
1873                return;
1874            };
1875            let elicitation_data: ElicitationRequestedData =
1876                match serde_json::from_value(notification.event.data.clone()) {
1877                    Ok(d) => d,
1878                    Err(e) => {
1879                        warn!(error = %e, "failed to deserialize elicitation request");
1880                        return;
1881                    }
1882                };
1883            let request = ElicitationRequest {
1884                message: elicitation_data.message,
1885                requested_schema: elicitation_data
1886                    .requested_schema
1887                    .map(|s| serde_json::to_value(s).unwrap_or(Value::Null)),
1888                mode: elicitation_data.mode.map(|m| match m {
1889                    crate::generated::session_events::ElicitationRequestedMode::Form => {
1890                        crate::types::ElicitationMode::Form
1891                    }
1892                    crate::generated::session_events::ElicitationRequestedMode::Url => {
1893                        crate::types::ElicitationMode::Url
1894                    }
1895                    _ => crate::types::ElicitationMode::Unknown,
1896                }),
1897                elicitation_source: elicitation_data.elicitation_source,
1898                url: elicitation_data.url,
1899            };
1900            let client = client.clone();
1901            let sid = session_id.clone();
1902            let span = tracing::error_span!(
1903                "elicitation_request_handler",
1904                session_id = %sid,
1905                request_id = %request_id
1906            );
1907            tokio::spawn(
1908                async move {
1909                    let cancel = ElicitationResult {
1910                        action: "cancel".to_string(),
1911                        content: None,
1912                    };
1913                    // Dispatch to a nested task so panics are caught as JoinErrors.
1914                    let handler_task = tokio::spawn({
1915                        let sid = sid.clone();
1916                        let request_id = request_id.clone();
1917                        let span = tracing::error_span!(
1918                            "elicitation_callback",
1919                            session_id = %sid,
1920                            request_id = %request_id
1921                        );
1922                        async move {
1923                            let handler_start = Instant::now();
1924                            let response = elicitation_handler
1925                                .handle(sid.clone(), request_id.clone(), request)
1926                                .await;
1927                            tracing::debug!(
1928                                elapsed_ms = handler_start.elapsed().as_millis(),
1929                                session_id = %sid,
1930                                request_id = %request_id,
1931                                "ElicitationHandler::handle dispatch"
1932                            );
1933                            response
1934                        }
1935                        .instrument(span)
1936                    });
1937                    let result = match handler_task.await {
1938                        Ok(r) => r,
1939                        Err(_) => cancel.clone(),
1940                    };
1941                    let rpc_start = Instant::now();
1942                    if let Err(e) = client
1943                        .call(
1944                            "session.ui.handlePendingElicitation",
1945                            Some(serde_json::json!({
1946                                "sessionId": sid,
1947                                "requestId": request_id,
1948                                "result": result,
1949                            })),
1950                        )
1951                        .await
1952                    {
1953                        // RPC failed — attempt cancel as last resort
1954                        warn!(error = %e, "handlePendingElicitation failed, sending cancel");
1955                        let _ = client
1956                            .call(
1957                                "session.ui.handlePendingElicitation",
1958                                Some(serde_json::json!({
1959                                    "sessionId": sid,
1960                                    "requestId": request_id,
1961                                    "result": cancel,
1962                                })),
1963                            )
1964                            .await;
1965                    } else {
1966                        tracing::debug!(
1967                            elapsed_ms = rpc_start.elapsed().as_millis(),
1968                            session_id = %sid,
1969                            request_id = %request_id,
1970                            "Session::handle_notification response sent successfully"
1971                        );
1972                    }
1973                }
1974                .instrument(span),
1975            );
1976        }
1977        SessionEventType::McpOauthRequired => {
1978            let Some(request_id) = extract_request_id(&notification.event.data) else {
1979                return;
1980            };
1981            let Some(mcp_auth_handler) = handlers.mcp_auth.clone() else {
1982                warn!(
1983                    session_id = %session_id,
1984                    request_id = %request_id,
1985                    "received MCP OAuth request without a registered MCP auth handler"
1986                );
1987                return;
1988            };
1989            let data: McpOauthRequiredData =
1990                match serde_json::from_value(notification.event.data.clone()) {
1991                    Ok(d) => d,
1992                    Err(e) => {
1993                        warn!(error = %e, "failed to deserialize MCP OAuth request");
1994                        return;
1995                    }
1996                };
1997            let request = McpAuthRequest {
1998                request_id: request_id.clone(),
1999                server_name: data.server_name,
2000                server_url: data.server_url,
2001                reason: data.reason,
2002                www_authenticate_params: data.www_authenticate_params,
2003                resource_metadata: data.resource_metadata,
2004                static_client_config: data.static_client_config,
2005            };
2006            let client = client.clone();
2007            let sid = session_id.clone();
2008            let span = tracing::error_span!(
2009                "mcp_auth_request_handler",
2010                session_id = %sid,
2011                request_id = %request_id
2012            );
2013            tokio::spawn(
2014                async move {
2015                    let cancel = McpAuthResult::Cancelled;
2016                    let handler_task = tokio::spawn({
2017                        let sid = sid.clone();
2018                        let request_id = request_id.clone();
2019                        let span = tracing::error_span!(
2020                            "mcp_auth_callback",
2021                            session_id = %sid,
2022                            request_id = %request_id
2023                        );
2024                        async move {
2025                            let handler_start = Instant::now();
2026                            let response = mcp_auth_handler
2027                                .handle(sid.clone(), request_id.clone(), request)
2028                                .await;
2029                            tracing::debug!(
2030                                elapsed_ms = handler_start.elapsed().as_millis(),
2031                                session_id = %sid,
2032                                request_id = %request_id,
2033                                "McpAuthHandler::handle dispatch"
2034                            );
2035                            response
2036                        }
2037                        .instrument(span)
2038                    });
2039                    let result = match handler_task.await {
2040                        Ok(result) => result,
2041                        Err(_) => cancel,
2042                    };
2043                    let rpc_start = Instant::now();
2044                    let _ = client
2045                        .call(
2046                            "session.mcp.oauth.handlePendingRequest",
2047                            Some(serde_json::json!({
2048                                "sessionId": sid,
2049                                "requestId": request_id,
2050                                "result": result.into_wire(),
2051                            })),
2052                        )
2053                        .await;
2054                    tracing::debug!(
2055                        elapsed_ms = rpc_start.elapsed().as_millis(),
2056                        "Session::handle_notification MCP auth response sent"
2057                    );
2058                }
2059                .instrument(span),
2060            );
2061        }
2062        SessionEventType::CommandExecute => {
2063            let data: CommandExecuteData =
2064                match serde_json::from_value(notification.event.data.clone()) {
2065                    Ok(d) => d,
2066                    Err(e) => {
2067                        warn!(error = %e, "failed to deserialize command.execute");
2068                        return;
2069                    }
2070                };
2071            let client = client.clone();
2072            let command_handlers = command_handlers.clone();
2073            let sid = session_id.clone();
2074            let span = tracing::error_span!("command_handler", session_id = %sid);
2075            tokio::spawn(
2076                async move {
2077                    let request_id = data.request_id;
2078                    let ack_error = match command_handlers.get(&data.command_name).cloned() {
2079                        None => Some(format!("Unknown command: {}", data.command_name)),
2080                        Some(handler) => {
2081                            let command_name = data.command_name.clone();
2082                            let ctx = CommandContext {
2083                                session_id: sid.clone(),
2084                                command: data.command,
2085                                command_name: data.command_name,
2086                                args: data.args,
2087                            };
2088                            let handler_start = Instant::now();
2089                            let result = handler.on_command(ctx).await;
2090                            tracing::debug!(
2091                                elapsed_ms = handler_start.elapsed().as_millis(),
2092                                session_id = %sid,
2093                                request_id = %request_id,
2094                                command_name = %command_name,
2095                                "CommandHandler::call dispatch"
2096                            );
2097                            match result {
2098                                Ok(()) => None,
2099                                Err(e) => Some(e.to_string()),
2100                            }
2101                        }
2102                    };
2103                    let mut params = serde_json::json!({
2104                        "sessionId": sid,
2105                        "requestId": request_id,
2106                    });
2107                    if let Some(error_msg) = ack_error {
2108                        params["error"] = serde_json::Value::String(error_msg);
2109                    }
2110                    let rpc_start = Instant::now();
2111                    let _ = client
2112                        .call("session.commands.handlePendingCommand", Some(params))
2113                        .await;
2114                    tracing::debug!(
2115                        elapsed_ms = rpc_start.elapsed().as_millis(),
2116                        session_id = %sid,
2117                        request_id = %request_id,
2118                        "Session::handle_notification response sent successfully"
2119                    );
2120                }
2121                .instrument(span),
2122            );
2123        }
2124        _ => {}
2125    }
2126}
2127
2128struct RequestDispatchContext<'a> {
2129    client: &'a Client,
2130    handlers: &'a SessionHandlers,
2131    hooks: Option<&'a dyn SessionHooks>,
2132    transforms: Option<&'a dyn SystemMessageTransform>,
2133    canvas_handler: Option<&'a Arc<dyn CanvasHandler>>,
2134    session_fs_provider: Option<&'a Arc<dyn SessionFsProvider>>,
2135    bearer_token_providers: &'a HashMap<String, Arc<dyn BearerTokenProvider>>,
2136}
2137
2138/// Process a JSON-RPC request from the CLI.
2139async fn handle_request(
2140    session_id: &SessionId,
2141    ctx: RequestDispatchContext<'_>,
2142    request: crate::JsonRpcRequest,
2143) {
2144    let sid = session_id.clone();
2145    let client = ctx.client;
2146    let handlers = ctx.handlers;
2147    let hooks = ctx.hooks;
2148    let transforms = ctx.transforms;
2149    let canvas_handler = ctx.canvas_handler;
2150    let session_fs_provider = ctx.session_fs_provider;
2151    let bearer_token_providers = ctx.bearer_token_providers;
2152
2153    if request.method.starts_with("sessionFs.") {
2154        crate::session_fs_dispatch::dispatch(client, session_fs_provider, request).await;
2155        return;
2156    }
2157
2158    if request.method.starts_with("canvas.") {
2159        crate::canvas_dispatch::dispatch(client, canvas_handler, request).await;
2160        return;
2161    }
2162
2163    if request.method == crate::generated::api_types::rpc_methods::PROVIDERTOKEN_GETTOKEN {
2164        crate::provider_token_dispatch::dispatch(client, bearer_token_providers, request).await;
2165        return;
2166    }
2167
2168    match request.method.as_str() {
2169        "hooks.invoke" => {
2170            let params = request.params.as_ref();
2171            let hook_type = params
2172                .and_then(|p| p.get("hookType"))
2173                .and_then(|v| v.as_str())
2174                .unwrap_or("");
2175            let input = params
2176                .and_then(|p| p.get("input"))
2177                .cloned()
2178                .unwrap_or(Value::Object(Default::default()));
2179
2180            let rpc_result = if let Some(hooks) = hooks {
2181                match crate::hooks::dispatch_hook(hooks, &sid, hook_type, input).await {
2182                    Ok(output) => output,
2183                    Err(e) => {
2184                        warn!(error = %e, hook_type = hook_type, "hook dispatch failed");
2185                        serde_json::json!({ "output": {} })
2186                    }
2187                }
2188            } else {
2189                serde_json::json!({ "output": {} })
2190            };
2191
2192            let rpc_response = JsonRpcResponse {
2193                jsonrpc: "2.0".to_string(),
2194                id: request.id,
2195                result: Some(rpc_result),
2196                error: None,
2197            };
2198            let _ = client.send_response(&rpc_response).await;
2199        }
2200
2201        "userInput.request" => {
2202            let params = request.params.as_ref();
2203            let Some(question) = params
2204                .and_then(|p| p.get("question"))
2205                .and_then(|v| v.as_str())
2206            else {
2207                warn!("userInput.request missing 'question' field");
2208                let rpc_response = JsonRpcResponse {
2209                    jsonrpc: "2.0".to_string(),
2210                    id: request.id,
2211                    result: None,
2212                    error: Some(crate::JsonRpcError {
2213                        code: error_codes::INVALID_PARAMS,
2214                        message: "missing required field: question".to_string(),
2215                        data: None,
2216                    }),
2217                };
2218                let _ = client.send_response(&rpc_response).await;
2219                return;
2220            };
2221            let question = question.to_string();
2222            let choices = params
2223                .and_then(|p| p.get("choices"))
2224                .and_then(|v| v.as_array())
2225                .map(|arr| {
2226                    arr.iter()
2227                        .filter_map(|v| v.as_str().map(|s| s.to_string()))
2228                        .collect()
2229                });
2230            let allow_freeform = params
2231                .and_then(|p| p.get("allowFreeform"))
2232                .and_then(|v| v.as_bool());
2233
2234            let handler_start = Instant::now();
2235            let response = if let Some(user_input_handler) = handlers.user_input.as_ref() {
2236                user_input_handler
2237                    .handle(sid.clone(), question, choices, allow_freeform)
2238                    .await
2239            } else {
2240                None
2241            };
2242            tracing::debug!(
2243                elapsed_ms = handler_start.elapsed().as_millis(),
2244                session_id = %sid,
2245                "UserInputHandler::handle dispatch"
2246            );
2247
2248            let rpc_result = match response {
2249                Some(UserInputResponse {
2250                    answer,
2251                    was_freeform,
2252                }) => serde_json::json!({
2253                    "answer": answer,
2254                    "wasFreeform": was_freeform,
2255                }),
2256                None => serde_json::json!({ "noResponse": true }),
2257            };
2258            let rpc_response = JsonRpcResponse {
2259                jsonrpc: "2.0".to_string(),
2260                id: request.id,
2261                result: Some(rpc_result),
2262                error: None,
2263            };
2264            let _ = client.send_response(&rpc_response).await;
2265        }
2266
2267        "exitPlanMode.request" => {
2268            let params = request
2269                .params
2270                .as_ref()
2271                .cloned()
2272                .unwrap_or(Value::Object(serde_json::Map::new()));
2273            let data: ExitPlanModeData = match serde_json::from_value(params) {
2274                Ok(d) => d,
2275                Err(e) => {
2276                    warn!(error = %e, "failed to deserialize exitPlanMode.request params, using defaults");
2277                    ExitPlanModeData::default()
2278                }
2279            };
2280
2281            let rpc_result = if let Some(exit_plan_handler) = handlers.exit_plan_mode.as_ref() {
2282                let result = exit_plan_handler.handle(sid, data).await;
2283                serde_json::to_value(result).expect("ExitPlanModeResult serialization cannot fail")
2284            } else {
2285                serde_json::json!({ "approved": true })
2286            };
2287            let rpc_response = JsonRpcResponse {
2288                jsonrpc: "2.0".to_string(),
2289                id: request.id,
2290                result: Some(rpc_result),
2291                error: None,
2292            };
2293            let _ = client.send_response(&rpc_response).await;
2294        }
2295
2296        "autoModeSwitch.request" => {
2297            let error_code = request
2298                .params
2299                .as_ref()
2300                .and_then(|p| p.get("errorCode"))
2301                .and_then(|v| v.as_str())
2302                .map(|s| s.to_string());
2303            let retry_after_seconds = request
2304                .params
2305                .as_ref()
2306                .and_then(|p| p.get("retryAfterSeconds"))
2307                .and_then(|v| v.as_f64());
2308
2309            let answer = if let Some(auto_mode_handler) = handlers.auto_mode_switch.as_ref() {
2310                auto_mode_handler
2311                    .handle(sid, error_code, retry_after_seconds)
2312                    .await
2313            } else {
2314                AutoModeSwitchResponse::No
2315            };
2316            let rpc_response = JsonRpcResponse {
2317                jsonrpc: "2.0".to_string(),
2318                id: request.id,
2319                result: Some(serde_json::json!({ "response": answer })),
2320                error: None,
2321            };
2322            let _ = client.send_response(&rpc_response).await;
2323        }
2324
2325        "systemMessage.transform" => {
2326            let params = request.params.as_ref();
2327            let sections: HashMap<String, crate::transforms::TransformSection> =
2328                match params.and_then(|p| p.get("sections")) {
2329                    Some(v) => match serde_json::from_value(v.clone()) {
2330                        Ok(s) => s,
2331                        Err(e) => {
2332                            let _ = send_error_response(
2333                                client,
2334                                request.id,
2335                                error_codes::INVALID_PARAMS,
2336                                &format!("invalid sections: {e}"),
2337                            )
2338                            .await;
2339                            return;
2340                        }
2341                    },
2342                    None => {
2343                        let _ = send_error_response(
2344                            client,
2345                            request.id,
2346                            error_codes::INVALID_PARAMS,
2347                            "missing sections parameter",
2348                        )
2349                        .await;
2350                        return;
2351                    }
2352                };
2353
2354            let rpc_result = if let Some(transforms) = transforms {
2355                let transform_start = Instant::now();
2356                let response =
2357                    crate::transforms::dispatch_transform(transforms, &sid, sections).await;
2358                tracing::debug!(
2359                    elapsed_ms = transform_start.elapsed().as_millis(),
2360                    session_id = %sid,
2361                    "SystemMessageTransform::transform_section dispatch"
2362                );
2363                match serde_json::to_value(response) {
2364                    Ok(v) => v,
2365                    Err(e) => {
2366                        warn!(error = %e, "failed to serialize transform response");
2367                        serde_json::json!({ "sections": {} })
2368                    }
2369                }
2370            } else {
2371                // No transforms registered — pass through all sections unchanged.
2372                let passthrough: HashMap<String, crate::transforms::TransformSection> = sections;
2373                serde_json::json!({ "sections": passthrough })
2374            };
2375
2376            let rpc_response = JsonRpcResponse {
2377                jsonrpc: "2.0".to_string(),
2378                id: request.id,
2379                result: Some(rpc_result),
2380                error: None,
2381            };
2382            let _ = client.send_response(&rpc_response).await;
2383        }
2384
2385        method => {
2386            warn!(
2387                method = method,
2388                "unhandled request method in session event loop"
2389            );
2390            let _ = send_error_response(
2391                client,
2392                request.id,
2393                error_codes::METHOD_NOT_FOUND,
2394                &format!("unknown method: {method}"),
2395            )
2396            .await;
2397        }
2398    }
2399}
2400
2401async fn send_error_response(
2402    client: &Client,
2403    id: u64,
2404    code: i32,
2405    message: &str,
2406) -> Result<(), Error> {
2407    let response = JsonRpcResponse {
2408        jsonrpc: "2.0".to_string(),
2409        id,
2410        result: None,
2411        error: Some(crate::JsonRpcError {
2412            code,
2413            message: message.to_string(),
2414            data: None,
2415        }),
2416    };
2417    client.send_response(&response).await
2418}
2419
2420/// Inject `action: "transform"` sections into a `SystemMessageConfig`,
2421/// forcing `mode: "customize"` (required by the CLI for transforms to fire).
2422/// Preserves any existing caller-provided section overrides.
2423fn apply_transform_sections(
2424    sys_msg: &mut SystemMessageConfig,
2425    transforms: &dyn SystemMessageTransform,
2426) {
2427    sys_msg.mode = Some("customize".to_string());
2428    let sections = sys_msg.sections.get_or_insert_with(HashMap::new);
2429    for id in transforms.section_ids() {
2430        sections.entry(id).or_insert_with(|| SectionOverride {
2431            action: Some("transform".to_string()),
2432            content: None,
2433        });
2434    }
2435}
2436
2437fn inject_transform_sections(config: &mut SessionConfig, transforms: &dyn SystemMessageTransform) {
2438    let sys_msg = config.system_message.get_or_insert_with(Default::default);
2439    apply_transform_sections(sys_msg, transforms);
2440}
2441
2442fn inject_transform_sections_resume(
2443    config: &mut ResumeSessionConfig,
2444    transforms: &dyn SystemMessageTransform,
2445) {
2446    let sys_msg = config.system_message.get_or_insert_with(Default::default);
2447    apply_transform_sections(sys_msg, transforms);
2448}
2449
2450#[cfg(test)]
2451mod tests {
2452    use serde_json::json;
2453
2454    use super::notification_permission_payload;
2455    use crate::handler::PermissionResult;
2456
2457    #[test]
2458    fn notification_payload_suppresses_no_result() {
2459        assert!(notification_permission_payload(&PermissionResult::NoResult).is_none());
2460    }
2461
2462    #[test]
2463    fn notification_payload_serializes_decisions() {
2464        assert_eq!(
2465            notification_permission_payload(&PermissionResult::approve_once()),
2466            Some(json!({ "kind": "approve-once" }))
2467        );
2468        assert_eq!(
2469            notification_permission_payload(&PermissionResult::reject(None)),
2470            Some(json!({ "kind": "reject" }))
2471        );
2472        assert_eq!(
2473            notification_permission_payload(&PermissionResult::reject(Some("bad".to_string()))),
2474            Some(json!({ "kind": "reject", "feedback": "bad" }))
2475        );
2476        assert_eq!(
2477            notification_permission_payload(&PermissionResult::user_not_available()),
2478            Some(json!({ "kind": "user-not-available" }))
2479        );
2480    }
2481}