1use std::collections::HashMap;
2use std::panic::AssertUnwindSafe;
3use std::path::{Path, PathBuf};
4use std::sync::Arc;
5use std::time::{Duration, Instant};
6
7use futures_util::FutureExt;
8use parking_lot::Mutex as ParkingLotMutex;
9use serde_json::Value;
10use tokio::sync::oneshot;
11use tokio::task::JoinHandle;
12use tokio_util::sync::CancellationToken;
13use tracing::{Instrument, error, warn};
14
15use crate::canvas::CanvasHandler;
16use crate::generated::api_types::{
17 LogRequest, ModelSwitchAutoTierRequest, ModelSwitchAutoTierResult, ModelSwitchToRequest,
18 OpenCanvasInstance, PermissionDecisionRequest, RegisterEventInterestParams,
19 ToolsGetCurrentMetadataResult, rpc_methods,
20};
21use crate::generated::session_events::{
22 CommandExecuteData, ElicitationRequestedData, ExternalToolRequestedData, McpOauthRequiredData,
23 SessionCanvasClosedData, SessionErrorData, SessionEventType, SessionIdleData, SessionMode,
24};
25use crate::handler::{
26 AutoModeSwitchHandler, AutoModeSwitchResponse, ElicitationHandler, ExitPlanModeHandler,
27 McpAuthHandler, McpAuthRequest, McpAuthResult, PermissionHandler, PermissionResult,
28 UserInputHandler, UserInputResponse,
29};
30use crate::hooks::SessionHooks;
31use crate::provider_token::BearerTokenProvider;
32use crate::session_fs::SessionFsProvider;
33use crate::trace_context::inject_trace_context;
34use crate::transforms::SystemMessageTransform;
35use crate::types::{
36 AutoTier, AutoTierPreference, CommandContext, CommandDefinition, CommandHandler,
37 CreateSessionResult, ElicitationRequest, ElicitationResult, ExitPlanModeData,
38 GetMessagesResponse, MessageOptions, PermissionRequestData, RequestId, ResumeSessionConfig,
39 ResumeSessionResult, SectionOverride, SessionCapabilities, SessionConfig, SessionEvent,
40 SessionId, SetModelOptions, SystemMessageConfig, ToolInvocation, ToolResult,
41 ToolResultExpanded, TraceContext, UiInputOptions, ensure_attachment_display_names,
42};
43use crate::{
44 Client, Error, ErrorKind, JsonRpcResponse, SessionErrorKind, SessionEventNotification,
45 error_codes,
46};
47
48const TOOL_SEARCH_TOOL_NAME: &str = "tool_search_tool";
52
53pub const DEFAULT_EVENT_BUFFER_CAPACITY: usize = 512;
61
62fn resolve_event_buffer_capacity(capacity: Option<usize>) -> Result<usize, Error> {
68 match capacity {
69 Some(0) => Err(Error::with_message(
70 ErrorKind::InvalidConfig,
71 "event_buffer_capacity must be greater than zero",
72 )),
73 Some(capacity) => Ok(capacity),
74 None => Ok(DEFAULT_EVENT_BUFFER_CAPACITY),
75 }
76}
77
78#[derive(Clone)]
86pub(crate) struct SessionHandlers {
87 pub permission: Option<Arc<dyn PermissionHandler>>,
88 pub managed_settings_enabled: bool,
89 pub elicitation: Option<Arc<dyn ElicitationHandler>>,
90 pub mcp_auth: Option<Arc<dyn McpAuthHandler>>,
91 pub user_input: Option<Arc<dyn UserInputHandler>>,
92 pub exit_plan_mode: Option<Arc<dyn ExitPlanModeHandler>>,
93 pub auto_mode_switch: Option<Arc<dyn AutoModeSwitchHandler>>,
94 pub tools: Arc<HashMap<String, Arc<dyn crate::tool::ToolHandler>>>,
95}
96
97type PendingExternalTools = Arc<ParkingLotMutex<HashMap<RequestId, Arc<CancellationToken>>>>;
98
99struct PendingExternalToolGuard {
100 request_id: RequestId,
101 token: Arc<CancellationToken>,
102 pending: PendingExternalTools,
103}
104
105impl Drop for PendingExternalToolGuard {
106 fn drop(&mut self) {
107 let mut pending = self.pending.lock();
108 if pending
109 .get(&self.request_id)
110 .is_some_and(|token| Arc::ptr_eq(token, &self.token))
111 {
112 pending.remove(&self.request_id);
113 }
114 }
115}
116
117impl PendingExternalToolGuard {
118 fn claim(&self) -> bool {
119 let mut pending = self.pending.lock();
120 if pending
121 .get(&self.request_id)
122 .is_some_and(|token| Arc::ptr_eq(token, &self.token))
123 {
124 pending.remove(&self.request_id);
125 true
126 } else {
127 false
128 }
129 }
130}
131
132fn has_managed_settings(
133 enable_managed_settings: Option<bool>,
134 managed_settings: Option<&crate::types::ManagedSettings>,
135) -> bool {
136 enable_managed_settings == Some(true) || managed_settings.is_some()
137}
138
139struct IdleWaiter {
141 tx: oneshot::Sender<Result<Option<SessionEvent>, Error>>,
142 last_assistant_message: Option<SessionEvent>,
143 started_at: Instant,
144 first_assistant_message_seen: bool,
145}
146
147fn structured_output_error(message: impl Into<String>) -> Error {
148 Error::with_message(
149 ErrorKind::Session(SessionErrorKind::AgentError),
150 message.into(),
151 )
152}
153
154fn is_structured_output_event(event: &SessionEvent) -> bool {
155 matches!(
156 event.parsed_type(),
157 SessionEventType::UserMessage
158 | SessionEventType::AssistantMessage
159 | SessionEventType::SessionIdle
160 | SessionEventType::SessionError
161 )
162}
163
164fn is_root_agent_event(event: &SessionEvent) -> bool {
165 event.agent_id.as_deref().is_none_or(str::is_empty)
166}
167
168struct StructuredOutputState {
169 message_id: String,
170 started: bool,
171 final_message: Option<SessionEvent>,
172}
173
174impl StructuredOutputState {
175 fn observe(&mut self, event: SessionEvent) -> Result<Option<SessionEvent>, Error> {
176 if !is_root_agent_event(&event) {
177 return Ok(None);
178 }
179 match event.parsed_type() {
180 SessionEventType::UserMessage => {
181 let data: crate::session_events::UserMessageData =
182 serde_json::from_value(event.data)?;
183 if data.message_id.as_deref() == Some(self.message_id.as_str()) {
184 self.started = true;
185 }
186 }
187 SessionEventType::AssistantMessage => {
188 let data: crate::session_events::AssistantMessageData =
189 serde_json::from_value(event.data.clone())?;
190 if data.originating_message_id.as_deref() == Some(self.message_id.as_str()) {
191 self.started = true;
192 self.final_message =
193 if data.tool_requests.is_some_and(|tools| !tools.is_empty()) {
194 None
195 } else {
196 Some(event)
197 };
198 }
199 }
200 SessionEventType::SessionIdle if self.started => {
201 let data: SessionIdleData = serde_json::from_value(event.data)?;
202 if data.mode == Some(SessionMode::Autopilot) {
203 return Ok(None);
204 }
205 if data.aborted == Some(true) {
206 return Err(structured_output_error(
207 "session aborted before structured output completed",
208 ));
209 }
210 let result = self.final_message.take().ok_or_else(|| {
211 structured_output_error("run completed without a structured assistant response")
212 })?;
213 let data: crate::session_events::AssistantMessageData =
214 serde_json::from_value(result.data.clone())?;
215 if data.content.trim().is_empty() {
216 return Err(structured_output_error(
217 "run completed without a structured assistant response",
218 ));
219 }
220 return Ok(Some(result));
221 }
222 SessionEventType::SessionError if self.started => {
223 let data: SessionErrorData = serde_json::from_value(event.data)?;
224 return Err(structured_output_error(format!(
225 "session error: {}",
226 data.message
227 )));
228 }
229 _ => {}
230 }
231 Ok(None)
232 }
233}
234
235struct WaiterGuard {
247 slot: Arc<ParkingLotMutex<Option<IdleWaiter>>>,
248}
249
250impl Drop for WaiterGuard {
251 fn drop(&mut self) {
252 self.slot.lock().take();
253 }
254}
255
256struct PendingSessionRegistration {
257 client: Client,
258 session_id: PendingSessionId,
259 shutdown: CancellationToken,
260 external_tools_shutdown: CancellationToken,
261 disarmed: bool,
262}
263
264enum PendingSessionId {
272 Known(SessionId, crate::router::RegistrationToken),
275 Deferred(Arc<ParkingLotMutex<Option<(SessionId, crate::router::SessionRegistration)>>>),
278}
279
280impl PendingSessionRegistration {
281 fn new(
282 client: Client,
283 session_id: SessionId,
284 token: crate::router::RegistrationToken,
285 shutdown: CancellationToken,
286 external_tools_shutdown: CancellationToken,
287 ) -> Self {
288 Self {
289 client,
290 session_id: PendingSessionId::Known(session_id, token),
291 shutdown,
292 external_tools_shutdown,
293 disarmed: false,
294 }
295 }
296
297 fn deferred(
299 client: Client,
300 stash: Arc<ParkingLotMutex<Option<(SessionId, crate::router::SessionRegistration)>>>,
301 shutdown: CancellationToken,
302 external_tools_shutdown: CancellationToken,
303 ) -> Self {
304 Self {
305 client,
306 session_id: PendingSessionId::Deferred(stash),
307 shutdown,
308 external_tools_shutdown,
309 disarmed: false,
310 }
311 }
312
313 fn registered_id(&self) -> Option<SessionId> {
314 match &self.session_id {
315 PendingSessionId::Known(id, _) => Some(id.clone()),
316 PendingSessionId::Deferred(stash) => stash.lock().as_ref().map(|(id, _)| id.clone()),
317 }
318 }
319
320 fn resolve_to(&mut self, session_id: SessionId, token: crate::router::RegistrationToken) {
324 self.session_id = PendingSessionId::Known(session_id, token);
325 }
326
327 async fn cleanup(mut self, event_loop: JoinHandle<()>) {
328 self.external_tools_shutdown.cancel();
329 self.shutdown.cancel();
330 let _ = event_loop.await;
331 if let Some(id) = self.registered_id() {
332 if let PendingSessionId::Known(_, token) = self.session_id {
333 self.client.unregister_session_owned(&id, token);
334 } else if let PendingSessionId::Deferred(stash) = &self.session_id
335 && let Some((id, registration)) = stash.lock().as_ref()
336 {
337 self.client.unregister_session_owned(id, registration.token);
338 }
339 }
340 self.disarmed = true;
341 }
342
343 fn disarm(&mut self) {
344 self.disarmed = true;
345 }
346}
347
348impl Drop for PendingSessionRegistration {
349 fn drop(&mut self) {
350 if !self.disarmed {
351 self.external_tools_shutdown.cancel();
352 self.shutdown.cancel();
353 if let Some(id) = self.registered_id() {
354 if let PendingSessionId::Known(_, token) = self.session_id {
355 self.client.unregister_session_owned(&id, token);
356 } else if let PendingSessionId::Deferred(stash) = &self.session_id
357 && let Some((id, registration)) = stash.lock().as_ref()
358 {
359 self.client.unregister_session_owned(id, registration.token);
360 }
361 }
362 }
363 }
364}
365
366type EventLoopSpawner = Box<
367 dyn FnOnce(
368 SessionId,
369 crate::router::SessionChannels,
370 Option<Arc<StartupTasks>>,
371 ) -> JoinHandle<()>
372 + Send,
373>;
374
375struct StartupTasks {
376 state: ParkingLotMutex<StartupTaskState>,
377}
378
379enum StartupTaskState {
380 Pending {
381 tasks: Vec<JoinHandle<()>>,
382 nested: Vec<tokio::task::AbortHandle>,
383 },
384 Disarmed,
385 Aborted,
386}
387
388impl Default for StartupTasks {
389 fn default() -> Self {
390 Self {
391 state: ParkingLotMutex::new(StartupTaskState::Pending {
392 tasks: Vec::new(),
393 nested: Vec::new(),
394 }),
395 }
396 }
397}
398
399impl StartupTasks {
400 fn track(&self, task: JoinHandle<()>) {
401 match &mut *self.state.lock() {
402 StartupTaskState::Pending { tasks, .. } => tasks.push(task),
403 StartupTaskState::Disarmed => {}
404 StartupTaskState::Aborted => task.abort(),
405 }
406 }
407
408 fn track_nested<T>(&self, task: &JoinHandle<T>) {
409 match &mut *self.state.lock() {
410 StartupTaskState::Pending { nested, .. } => nested.push(task.abort_handle()),
411 StartupTaskState::Disarmed => {}
412 StartupTaskState::Aborted => task.abort(),
413 }
414 }
415
416 fn abort(&self) {
417 if let StartupTaskState::Pending { tasks, nested } =
418 std::mem::replace(&mut *self.state.lock(), StartupTaskState::Aborted)
419 {
420 for handle in nested {
421 handle.abort();
422 }
423 for handle in tasks {
424 handle.abort();
425 }
426 }
427 }
428
429 fn disarm(&self) {
430 *self.state.lock() = StartupTaskState::Disarmed;
431 }
432}
433
434fn spawn_startup_tracked(
435 future: impl std::future::Future<Output = ()> + Send + 'static,
436 startup_tasks: Option<&Arc<StartupTasks>>,
437) {
438 let task = tokio::spawn(future);
439 if let Some(startup_tasks) = startup_tasks {
440 startup_tasks.track(task);
441 }
442}
443
444struct RunningCreateLoop {
445 event_loop: Option<JoinHandle<()>>,
446 startup_tasks: Arc<StartupTasks>,
447}
448
449impl RunningCreateLoop {
450 fn into_handle(mut self) -> JoinHandle<()> {
451 self.event_loop
452 .take()
453 .expect("running create loop has a handle")
454 }
455}
456
457impl Drop for RunningCreateLoop {
458 fn drop(&mut self) {
459 if let Some(event_loop) = self.event_loop.take() {
460 self.startup_tasks.abort();
461 event_loop.abort();
462 }
463 }
464}
465
466enum CreateEventLoop {
469 Running {
470 running: RunningCreateLoop,
471 token: crate::router::RegistrationToken,
472 },
473 Deferred(EventLoopSpawner),
474}
475
476impl CreateEventLoop {
477 async fn cleanup(self, pending_registration: PendingSessionRegistration) {
478 match self {
479 Self::Running { running, .. } => {
480 pending_registration.shutdown.cancel();
481 running.startup_tasks.abort();
483 pending_registration.cleanup(running.into_handle()).await;
484 }
485 Self::Deferred(_) => drop(pending_registration),
486 }
487 }
488}
489
490pub struct Session {
503 id: SessionId,
504 cwd: PathBuf,
505 workspace_path: Option<PathBuf>,
506 remote_url: Option<String>,
507 client: Client,
508 event_loop: ParkingLotMutex<Option<JoinHandle<()>>>,
513 shutdown: CancellationToken,
527 external_tools_shutdown: CancellationToken,
530 idle_waiter: Arc<ParkingLotMutex<Option<IdleWaiter>>>,
537 capabilities: Arc<parking_lot::RwLock<SessionCapabilities>>,
539 open_canvases: Arc<parking_lot::RwLock<Vec<OpenCanvasInstance>>>,
541 event_tx: tokio::sync::broadcast::Sender<SessionEvent>,
543 resume_bootstrap: Option<Arc<crate::subscription::ResumeBootstrap>>,
546 github_token_registration:
547 ParkingLotMutex<Option<crate::github_token::GitHubTokenRegistration>>,
548 registration_token: crate::router::RegistrationToken,
550}
551
552impl Session {
553 pub fn id(&self) -> &SessionId {
555 &self.id
556 }
557
558 pub fn cwd(&self) -> &PathBuf {
560 &self.cwd
561 }
562
563 pub fn workspace_path(&self) -> Option<&Path> {
565 self.workspace_path.as_deref()
566 }
567
568 pub fn remote_url(&self) -> Option<&str> {
570 self.remote_url.as_deref()
571 }
572
573 pub fn capabilities(&self) -> SessionCapabilities {
578 self.capabilities.read().clone()
579 }
580
581 pub fn open_canvases(&self) -> Vec<OpenCanvasInstance> {
584 self.open_canvases.read().clone()
585 }
586
587 pub fn cancellation_token(&self) -> CancellationToken {
614 self.shutdown.child_token()
615 }
616
617 pub fn subscribe(&self) -> crate::subscription::EventSubscription {
672 match &self.resume_bootstrap {
673 Some(bootstrap) => bootstrap.subscribe(&self.event_tx),
674 None => crate::subscription::EventSubscription::new(self.event_tx.subscribe()),
675 }
676 }
677
678 pub fn client(&self) -> &Client {
680 &self.client
681 }
682
683 pub fn rpc(&self) -> crate::generated::rpc::SessionRpc<'_> {
694 crate::generated::rpc::SessionRpc { session: self }
695 }
696
697 pub async fn stop_event_loop(&self) {
710 self.shutdown.cancel();
711 let handle = self.event_loop.lock().take();
712 if let Some(handle) = handle {
713 let _ = handle.await;
714 }
715 if let Some(waiter) = self.idle_waiter.lock().take() {
717 let _ = waiter.tx.send(Err(
718 ErrorKind::Session(SessionErrorKind::EventLoopClosed).into()
719 ));
720 }
721 }
722
723 pub async fn send(&self, opts: impl Into<MessageOptions>) -> Result<String, Error> {
748 if self.idle_waiter.lock().is_some() {
749 return Err(ErrorKind::Session(SessionErrorKind::SendWhileWaiting).into());
750 }
751 self.send_inner(opts.into()).await
752 }
753
754 async fn send_inner(&self, opts: MessageOptions) -> Result<String, Error> {
755 let mut params = serde_json::json!({
756 "sessionId": self.id,
757 "prompt": opts.prompt,
758 });
759 if let Some(source) = opts.source {
760 params["source"] = serde_json::to_value(source)?;
761 }
762 if let Some(m) = opts.mode {
763 params["mode"] = serde_json::to_value(m)?;
764 }
765 if let Some(am) = opts.agent_mode {
766 params["agentMode"] = serde_json::to_value(am)?;
767 }
768 if let Some(mut a) = opts.attachments {
769 ensure_attachment_display_names(&mut a);
770 params["attachments"] = serde_json::to_value(a)?;
771 }
772 if let Some(headers) = opts.request_headers
773 && !headers.is_empty()
774 {
775 params["requestHeaders"] = serde_json::to_value(headers)?;
776 }
777 if let Some(display_prompt) = opts.display_prompt {
778 params["displayPrompt"] = serde_json::to_value(display_prompt)?;
779 }
780 if let Some(schema) = opts.response_schema {
781 params["responseFormat"] = serde_json::json!({
782 "type": "json_schema",
783 "jsonSchema": { "name": "response", "strict": true, "schema": schema }
784 });
785 }
786 let trace_ctx = if opts.traceparent.is_some() || opts.tracestate.is_some() {
787 TraceContext {
788 traceparent: opts.traceparent,
789 tracestate: opts.tracestate,
790 }
791 } else {
792 self.client.resolve_trace_context().await
793 };
794 inject_trace_context(&mut params, &trace_ctx);
795 let rpc_start = Instant::now();
796 let result = self.client.call("session.send", Some(params)).await?;
797 let message_id = result
798 .get("messageId")
799 .and_then(|v| v.as_str())
800 .map(|s| s.to_string())
801 .unwrap_or_default();
802 tracing::debug!(
803 elapsed_ms = rpc_start.elapsed().as_millis(),
804 session_id = %self.id,
805 message_id = %message_id,
806 "Session::send completed successfully"
807 );
808 Ok(message_id)
809 }
810
811 pub async fn send_and_wait(
839 &self,
840 opts: impl Into<MessageOptions>,
841 ) -> Result<Option<SessionEvent>, Error> {
842 let total_start = Instant::now();
843 let opts = opts.into();
844 if opts.response_schema.is_some() {
845 return self.send_and_wait_structured(opts).await.map(Some);
846 }
847 let timeout_duration = opts.wait_timeout.unwrap_or(Duration::from_secs(60));
848 let (tx, rx) = oneshot::channel();
849
850 {
851 let mut guard = self.idle_waiter.lock();
852 if guard.is_some() {
853 return Err(ErrorKind::Session(SessionErrorKind::SendWhileWaiting).into());
854 }
855 *guard = Some(IdleWaiter {
856 tx,
857 last_assistant_message: None,
858 started_at: total_start,
859 first_assistant_message_seen: false,
860 });
861 }
862
863 let _waiter_guard = WaiterGuard {
868 slot: self.idle_waiter.clone(),
869 };
870
871 let result = tokio::time::timeout(timeout_duration, async {
872 self.send_inner(opts).await?;
873 match rx.await {
874 Ok(result) => result,
875 Err(_) => Err(ErrorKind::Session(SessionErrorKind::EventLoopClosed).into()),
876 }
877 })
878 .await;
879
880 match result {
881 Ok(inner) => {
882 tracing::debug!(
883 elapsed_ms = total_start.elapsed().as_millis(),
884 session_id = %self.id,
885 completed_by = if inner.is_ok() { "idle" } else { "error" },
886 "Session::send_and_wait complete"
887 );
888 inner
889 }
890 Err(_) => {
891 tracing::warn!(
892 elapsed_ms = total_start.elapsed().as_millis(),
893 session_id = %self.id,
894 completed_by = "timeout",
895 "Session::send_and_wait failed"
896 );
897 Err(ErrorKind::Session(SessionErrorKind::Timeout(timeout_duration)).into())
898 }
899 }
900 }
901
902 #[cfg(feature = "derive")]
910 pub async fn send_and_wait_typed<T>(&self, opts: impl Into<MessageOptions>) -> Result<T, Error>
911 where
912 T: schemars::JsonSchema + serde::de::DeserializeOwned,
913 {
914 let mut opts = opts.into();
915 if opts.response_schema.is_some()
916 || opts.mode == Some(crate::types::DeliveryMode::Immediate)
917 {
918 return Err(Error::with_message(
919 ErrorKind::InvalidConfig,
920 "typed structured output cannot specify a response schema or immediate delivery",
921 ));
922 }
923 opts.response_schema = Some(crate::tool::schema_for::<T>());
924 let event = self.send_and_wait_structured(opts).await?;
925 let data: crate::session_events::AssistantMessageData = serde_json::from_value(event.data)?;
926 serde_json::from_str::<Option<T>>(&data.content)?.ok_or_else(|| {
927 structured_output_error("structured response was JSON null, not a result")
928 })
929 }
930
931 async fn send_and_wait_structured(&self, opts: MessageOptions) -> Result<SessionEvent, Error> {
932 let duration = opts.wait_timeout.unwrap_or(Duration::from_secs(60));
933 let mut events = crate::subscription::EventSubscription::new(self.event_tx.subscribe());
935 let wait = async {
936 let mut admission = Box::pin(self.send(opts));
937 let mut pending = Vec::new();
938 let message_id = loop {
939 tokio::select! {
940 result = &mut admission => break result?,
941 event = events.recv() => {
942 let event = event.map_err(|err| structured_output_error(err.to_string()))?;
943 if is_structured_output_event(&event) {
944 pending.push(event);
945 }
946 }
947 _ = self.shutdown.cancelled() =>
948 return Err(structured_output_error("session closed before structured output completed")),
949 }
950 };
951 let mut state = StructuredOutputState {
952 message_id,
953 started: false,
954 final_message: None,
955 };
956 for event in pending {
957 if let Some(result) = state.observe(event)? {
958 return Ok(result);
959 }
960 }
961 loop {
962 tokio::select! {
963 event = events.recv() => {
964 let event = event.map_err(|err| structured_output_error(err.to_string()))?;
965 if let Some(result) = state.observe(event)? {
966 return Ok(result);
967 }
968 }
969 _ = self.shutdown.cancelled() =>
970 return Err(structured_output_error("session closed before structured output completed")),
971 }
972 }
973 };
974 tokio::time::timeout(duration, wait)
975 .await
976 .map_err(|_| Error::from(ErrorKind::Session(SessionErrorKind::Timeout(duration))))?
977 }
978
979 pub async fn get_events(&self) -> Result<Vec<SessionEvent>, Error> {
981 let result = self
982 .client
983 .call(
984 "session.getMessages",
985 Some(serde_json::json!({ "sessionId": self.id })),
986 )
987 .await?;
988 let response: GetMessagesResponse = serde_json::from_value(result)?;
989 Ok(response.events)
990 }
991
992 #[deprecated(since = "0.1.0", note = "Use `get_events()` instead")]
994 pub async fn get_messages(&self) -> Result<Vec<SessionEvent>, Error> {
995 self.get_events().await
996 }
997
998 pub async fn abort(&self) -> Result<(), Error> {
1006 self.client
1007 .call(
1008 "session.abort",
1009 Some(serde_json::json!({ "sessionId": self.id })),
1010 )
1011 .await?;
1012 Ok(())
1013 }
1014
1015 pub async fn set_model(&self, model: &str, opts: Option<SetModelOptions>) -> Result<(), Error> {
1019 let opts = opts.unwrap_or_default();
1020 let auto_tier = opts.auto_tier.clone();
1021 let request = ModelSwitchToRequest {
1022 auto_tier: match &auto_tier {
1023 Some(AutoTierPreference::Tier(tier)) => Some(tier.clone()),
1024 _ => None,
1025 },
1026 compaction_decision: None,
1027 context_tier: opts.context_tier,
1028 defer_if_model_change_queued: None,
1029 model_capabilities: opts.model_capabilities,
1030 model_change_scope: None,
1031 model_id: model.to_string(),
1032 picker_persistence: None,
1033 reasoning_effort: opts.reasoning_effort,
1034 reasoning_summary: opts.reasoning_summary,
1035 repo_scope: None,
1036 require_available: None,
1037 run_compaction_preflight: None,
1038 source: None,
1039 verbosity: None,
1040 };
1041
1042 if matches!(auto_tier, Some(AutoTierPreference::Reset)) {
1043 let mut wire_params = serde_json::to_value(request)?;
1047 wire_params["sessionId"] = serde_json::Value::String(self.id.to_string());
1048 wire_params["autoTier"] = serde_json::Value::Null;
1049 self.client
1050 .call("session.model.switchTo", Some(wire_params))
1051 .await?;
1052 return Ok(());
1053 }
1054
1055 self.rpc().model().switch_to(request).await?;
1056 Ok(())
1057 }
1058
1059 pub async fn set_auto_tier(
1088 &self,
1089 auto_tier: Option<AutoTier>,
1090 ) -> Result<ModelSwitchAutoTierResult, Error> {
1091 self.rpc()
1092 .model()
1093 .switch_auto_tier(ModelSwitchAutoTierRequest {
1094 auto_tier,
1095 source: None,
1096 })
1097 .await
1098 }
1099
1100 pub async fn disconnect(&self) -> Result<(), Error> {
1117 self.client.detach_session(&self.id).await?;
1118 self.external_tools_shutdown.cancel();
1119 self.stop_event_loop().await;
1120 self.github_token_registration.lock().take();
1121 self.client
1122 .unregister_session_owned(&self.id, self.registration_token);
1123 Ok(())
1124 }
1125
1126 #[deprecated(since = "0.1.0", note = "Use `disconnect()` instead")]
1129 pub async fn destroy(&self) -> Result<(), Error> {
1130 self.disconnect().await
1131 }
1132
1133 pub async fn log(
1137 &self,
1138 message: &str,
1139 opts: Option<crate::types::LogOptions>,
1140 ) -> Result<(), Error> {
1141 let opts = opts.unwrap_or_default();
1142 let level = match opts.level {
1143 Some(level) => Some(serde_json::from_value(serde_json::to_value(level)?)?),
1144 None => None,
1145 };
1146 let request = LogRequest {
1147 message: message.to_string(),
1148 level,
1149 ephemeral: opts.ephemeral,
1150 r#type: None,
1151 tip: None,
1152 url: None,
1153 };
1154 self.rpc().log(request).await?;
1155 Ok(())
1156 }
1157
1158 pub fn ui(&self) -> SessionUi<'_> {
1164 SessionUi { session: self }
1165 }
1166
1167 fn assert_elicitation(&self) -> Result<(), Error> {
1169 if self
1170 .capabilities
1171 .read()
1172 .ui
1173 .as_ref()
1174 .and_then(|u| u.elicitation)
1175 != Some(true)
1176 {
1177 return Err(ErrorKind::Session(SessionErrorKind::ElicitationNotSupported).into());
1178 }
1179 Ok(())
1180 }
1181}
1182
1183impl Drop for Session {
1184 fn drop(&mut self) {
1185 self.shutdown.cancel();
1197 self.external_tools_shutdown.cancel();
1198 self.github_token_registration.lock().take();
1199 self.client
1200 .unregister_session_owned(&self.id, self.registration_token);
1201 }
1202}
1203
1204pub struct SessionUi<'a> {
1211 session: &'a Session,
1212}
1213
1214impl<'a> SessionUi<'a> {
1215 pub async fn elicitation(
1223 &self,
1224 message: &str,
1225 schema: Value,
1226 ) -> Result<ElicitationResult, Error> {
1227 self.session.assert_elicitation()?;
1228 let result = self
1229 .session
1230 .client
1231 .call(
1232 "session.ui.elicitation",
1233 Some(serde_json::json!({
1234 "sessionId": self.session.id,
1235 "message": message,
1236 "requestedSchema": schema,
1237 })),
1238 )
1239 .await?;
1240 let elicitation: ElicitationResult = serde_json::from_value(result)?;
1241 Ok(elicitation)
1242 }
1243
1244 pub async fn confirm(&self, message: &str) -> Result<bool, Error> {
1248 self.session.assert_elicitation()?;
1249 let schema = serde_json::json!({
1250 "type": "object",
1251 "properties": {
1252 "confirmed": {
1253 "type": "boolean",
1254 "default": true,
1255 }
1256 },
1257 "required": ["confirmed"]
1258 });
1259 let result = self.elicitation(message, schema).await?;
1260 Ok(result.action == "accept"
1261 && result
1262 .content
1263 .and_then(|c| c.get("confirmed").and_then(|v| v.as_bool()))
1264 == Some(true))
1265 }
1266
1267 pub async fn select(&self, message: &str, options: &[&str]) -> Result<Option<String>, Error> {
1271 self.session.assert_elicitation()?;
1272 let schema = serde_json::json!({
1273 "type": "object",
1274 "properties": {
1275 "selection": {
1276 "type": "string",
1277 "enum": options,
1278 }
1279 },
1280 "required": ["selection"]
1281 });
1282 let result = self.elicitation(message, schema).await?;
1283 if result.action != "accept" {
1284 return Ok(None);
1285 }
1286 let selection = result.content.and_then(|c| {
1287 c.get("selection")
1288 .and_then(|v| v.as_str())
1289 .map(String::from)
1290 });
1291 Ok(selection)
1292 }
1293
1294 pub async fn input(
1299 &self,
1300 message: &str,
1301 options: Option<&UiInputOptions<'_>>,
1302 ) -> Result<Option<String>, Error> {
1303 self.session.assert_elicitation()?;
1304 let mut field = serde_json::json!({ "type": "string" });
1305 if let Some(opts) = options {
1306 if let Some(title) = opts.title {
1307 field["title"] = Value::String(title.to_string());
1308 }
1309 if let Some(desc) = opts.description {
1310 field["description"] = Value::String(desc.to_string());
1311 }
1312 if let Some(min) = opts.min_length {
1313 field["minLength"] = Value::Number(min.into());
1314 }
1315 if let Some(max) = opts.max_length {
1316 field["maxLength"] = Value::Number(max.into());
1317 }
1318 if let Some(fmt) = &opts.format {
1319 field["format"] = Value::String(fmt.as_str().to_string());
1320 }
1321 if let Some(default) = opts.default {
1322 field["default"] = Value::String(default.to_string());
1323 }
1324 }
1325 let schema = serde_json::json!({
1326 "type": "object",
1327 "properties": { "value": field },
1328 "required": ["value"]
1329 });
1330 let result = self.elicitation(message, schema).await?;
1331 if result.action != "accept" {
1332 return Ok(None);
1333 }
1334 let value = result
1335 .content
1336 .and_then(|c| c.get("value").and_then(|v| v.as_str()).map(String::from));
1337 Ok(value)
1338 }
1339}
1340
1341impl Client {
1342 pub fn prepare_session(&self, config: SessionConfig) -> Result<PreparedSession, Error> {
1404 let capacity = resolve_event_buffer_capacity(config.event_buffer_capacity)?;
1405 Ok(PreparedSession::new(
1406 self.clone(),
1407 PreparedKind::Create(Box::new(config)),
1408 capacity,
1409 ))
1410 }
1411
1412 pub fn prepare_resume_session(
1427 &self,
1428 config: ResumeSessionConfig,
1429 ) -> Result<PreparedSession, Error> {
1430 let capacity = resolve_event_buffer_capacity(config.event_buffer_capacity)?;
1431 Ok(PreparedSession::new(
1432 self.clone(),
1433 PreparedKind::Resume(Box::new(config)),
1434 capacity,
1435 ))
1436 }
1437
1438 pub async fn create_session(&self, config: SessionConfig) -> Result<Session, Error> {
1470 self.prepare_session(config)?.start().await
1471 }
1472
1473 pub async fn resume_session(&self, config: ResumeSessionConfig) -> Result<Session, Error> {
1497 self.prepare_resume_session(config)?.start().await
1498 }
1499
1500 async fn start_prepared_create(
1501 &self,
1502 mut config: SessionConfig,
1503 event_tx: tokio::sync::broadcast::Sender<SessionEvent>,
1504 shutdown: CancellationToken,
1505 ) -> Result<Session, Error> {
1506 let total_start = Instant::now();
1507 let caller_session_id = config.session_id.clone();
1512 let use_server_generated_id = config.cloud.is_some() && caller_session_id.is_none();
1513 let local_session_id: Option<SessionId> = if use_server_generated_id {
1514 None
1515 } else {
1516 Some(
1517 caller_session_id
1518 .clone()
1519 .unwrap_or_else(|| SessionId::new(uuid::Uuid::new_v4().to_string())),
1520 )
1521 };
1522 if config.hooks_handler.is_some() && config.hooks.is_none() {
1523 config.hooks = Some(true);
1524 }
1525 if let Some(transforms) = config.system_message_transform.clone() {
1526 inject_transform_sections(&mut config, transforms.as_ref());
1527 }
1528 let mode = self.inner.mode;
1529 if mode == crate::ClientMode::Empty && config.available_tools.is_none() {
1530 return Err(Error::with_message(
1531 ErrorKind::InvalidConfig,
1532 "ClientMode::Empty requires available_tools to be set on the session config. \
1533 Use ToolSet to specify which tools the session may use (e.g. \
1534 ToolSet::new().add_builtin_many(BUILTIN_TOOLS_ISOLATED)).",
1535 ));
1536 }
1537 crate::mode::validate_tool_filter_list(
1538 "available_tools",
1539 config.available_tools.as_deref(),
1540 )?;
1541 crate::mode::validate_tool_filter_list("excluded_tools", config.excluded_tools.as_deref())?;
1542 config.system_message =
1543 crate::mode::system_message_for_mode(mode, config.system_message.take());
1544 config.memory = crate::mode::memory_for_mode(mode, config.memory.take());
1545 config.enable_experimental_mode =
1546 crate::mode::experimental_mode_for_mode(mode, config.enable_experimental_mode);
1547 if mode == crate::ClientMode::Empty {
1548 if config.enable_session_telemetry.is_none() {
1549 config.enable_session_telemetry = Some(false);
1550 }
1551 if config.skip_embedding_retrieval.is_none() {
1552 config.skip_embedding_retrieval = Some(true);
1553 }
1554 if config.enable_on_demand_instruction_discovery.is_none() {
1555 config.enable_on_demand_instruction_discovery = Some(false);
1556 }
1557 if config.enable_file_hooks.is_none() {
1558 config.enable_file_hooks = Some(false);
1559 }
1560 if config.enable_host_git_operations.is_none() {
1561 config.enable_host_git_operations = Some(false);
1562 }
1563 if config.enable_session_store.is_none() {
1564 config.enable_session_store = Some(false);
1565 }
1566 if config.enable_skills.is_none() {
1567 config.enable_skills = Some(false);
1568 }
1569 }
1570 if mode == crate::ClientMode::Empty && config.mcp_oauth_token_storage.is_none() {
1571 config.mcp_oauth_token_storage = Some("in-memory".into());
1572 }
1573 if mode == crate::ClientMode::Empty && config.embedding_cache_storage.is_none() {
1574 config.embedding_cache_storage = Some("in-memory".into());
1575 }
1576 config.custom_agents_local_only =
1577 crate::mode::resolve_custom_agents_local_only(mode, config.custom_agents_local_only);
1578 let opt_skip_custom_instructions = config.skip_custom_instructions;
1579 let opt_custom_agents_local_only = config.custom_agents_local_only;
1580 let opt_coauthor_enabled = config.coauthor_enabled;
1581 let opt_manage_schedule_enabled = config.manage_schedule_enabled;
1582 let opt_included_builtin_skills = config.included_builtin_skills.take();
1583 let (mut wire, mut runtime) = config.into_wire(local_session_id.clone())?;
1584 wire.enable_github_telemetry_forwarding =
1585 self.inner.on_github_telemetry.is_some().then_some(true);
1586
1587 let permission_handler = crate::permission::resolve_handler(
1588 runtime.permission_handler.take(),
1589 runtime.permission_policy.take(),
1590 );
1591 let handlers = SessionHandlers {
1592 permission: permission_handler,
1593 managed_settings_enabled: has_managed_settings(
1594 wire.enable_managed_settings,
1595 wire.managed_settings.as_ref(),
1596 ),
1597 elicitation: runtime.elicitation_handler.take(),
1598 mcp_auth: runtime.mcp_auth_handler.take(),
1599 user_input: runtime.user_input_handler.take(),
1600 exit_plan_mode: runtime.exit_plan_mode_handler.take(),
1601 auto_mode_switch: runtime.auto_mode_switch_handler.take(),
1602 tools: Arc::new(std::mem::take(&mut runtime.tool_handlers)),
1603 };
1604 let hooks = runtime.hooks_handler.take();
1605 let transforms = runtime.system_message_transform.take();
1606 let tools_count = wire.tools.as_ref().map_or(0, Vec::len);
1607 let commands_count = runtime.commands.as_ref().map_or(0, Vec::len);
1608 let has_hooks = hooks.is_some();
1609 let command_handlers = build_command_handler_map(runtime.commands.as_deref());
1610 let canvas_handler = runtime.canvas_handler.take();
1611 let session_fs_provider = runtime.session_fs_provider.take();
1612 let bearer_token_providers = std::mem::take(&mut runtime.bearer_token_providers);
1613 let github_token_registration = runtime
1614 .github_token_provider
1615 .take()
1616 .map(|provider| self.register_github_token_provider(provider));
1617 wire.github_token_provider_registration_id = github_token_registration
1618 .as_ref()
1619 .map(|registration| registration.id().to_string());
1620 let has_mcp_auth_handler = handlers.mcp_auth.is_some();
1621 if self.inner.session_fs_configured && session_fs_provider.is_none() {
1622 return Err(ErrorKind::Session(SessionErrorKind::SessionFsProviderRequired).into());
1623 }
1624 if self.inner.session_fs_sqlite_declared
1625 && let Some(ref provider) = session_fs_provider
1626 && provider.sqlite().is_none()
1627 {
1628 return Err(Error::with_message(
1629 ErrorKind::InvalidConfig,
1630 "SessionFs capabilities declare SQLite support but the provider \
1631 does not implement SessionFsSqliteProvider",
1632 ));
1633 }
1634
1635 let mut params = serde_json::to_value(&wire)?;
1636 let trace_ctx = self.resolve_trace_context().await;
1637 inject_trace_context(&mut params, &trace_ctx);
1638
1639 let setup_start = Instant::now();
1640 let capabilities = Arc::new(parking_lot::RwLock::new(SessionCapabilities::default()));
1641 let idle_waiter = Arc::new(ParkingLotMutex::new(None));
1642 let open_canvases = Arc::new(parking_lot::RwLock::new(Vec::new()));
1643 let external_tools_shutdown = self.inner.rpc.connection_closed_token();
1644
1645 let spawn_loop: EventLoopSpawner = {
1646 let client = self.clone();
1647 let idle_waiter = idle_waiter.clone();
1648 let capabilities = capabilities.clone();
1649 let open_canvases = open_canvases.clone();
1650 let event_tx = event_tx.clone();
1651 let shutdown = shutdown.clone();
1652 let external_tools_shutdown = external_tools_shutdown.clone();
1653 Box::new(
1654 move |session_id: SessionId,
1655 channels: crate::router::SessionChannels,
1656 startup_tasks: Option<Arc<StartupTasks>>| {
1657 spawn_event_loop(
1658 session_id,
1659 client,
1660 handlers,
1661 hooks,
1662 transforms,
1663 command_handlers,
1664 canvas_handler,
1665 session_fs_provider,
1666 bearer_token_providers,
1667 channels,
1668 idle_waiter,
1669 capabilities,
1670 open_canvases,
1671 event_tx,
1672 None,
1673 shutdown,
1674 external_tools_shutdown,
1675 startup_tasks,
1676 )
1677 },
1678 )
1679 };
1680
1681 let (inline_callback, mut pending_registration, event_loop, inline_stash) =
1684 if let Some(ref sid) = local_session_id {
1685 let registration = self.register_session(sid);
1686 let token = registration.token;
1687 let startup_tasks = Arc::new(StartupTasks::default());
1688 let event_loop = spawn_loop(
1689 sid.clone(),
1690 registration.channels,
1691 Some(startup_tasks.clone()),
1692 );
1693 let guard = PendingSessionRegistration::new(
1694 self.clone(),
1695 sid.clone(),
1696 token,
1697 shutdown.clone(),
1698 external_tools_shutdown.clone(),
1699 );
1700 (
1701 None,
1702 guard,
1703 CreateEventLoop::Running {
1704 running: RunningCreateLoop {
1705 event_loop: Some(event_loop),
1706 startup_tasks,
1707 },
1708 token,
1709 },
1710 None,
1711 )
1712 } else {
1713 let inline_stash = Arc::new(ParkingLotMutex::new(None));
1714 let client = self.clone();
1715 let stash = inline_stash.clone();
1716 let expected = caller_session_id.clone();
1717 let callback: crate::jsonrpc::InlineResponseCallback = Box::new(move |response| {
1718 let result = response.result.as_ref().ok_or_else(|| {
1719 Error::with_message(
1720 ErrorKind::Json,
1721 "session.create response had no result",
1722 )
1723 })?;
1724 let parsed: CreateSessionResult =
1725 serde_json::from_value(result.clone()).map_err(Error::from)?;
1726 if let Some(requested) = expected.as_ref()
1727 && parsed.session_id != *requested
1728 {
1729 return Err(ErrorKind::Session(SessionErrorKind::SessionIdMismatch {
1730 requested: requested.clone(),
1731 returned: parsed.session_id,
1732 })
1733 .into());
1734 }
1735 let mut stashed = stash.lock();
1743 let registration = client.register_session(&parsed.session_id);
1744 *stashed = Some((parsed.session_id, registration));
1745 Ok(())
1746 });
1747 let guard = PendingSessionRegistration::deferred(
1748 self.clone(),
1749 inline_stash.clone(),
1750 shutdown.clone(),
1751 external_tools_shutdown.clone(),
1752 );
1753 (
1754 Some(callback),
1755 guard,
1756 CreateEventLoop::Deferred(spawn_loop),
1757 Some(inline_stash),
1758 )
1759 };
1760
1761 let rpc_start = Instant::now();
1762 let result = match self
1763 .call_with_inline_callback("session.create", Some(params), inline_callback)
1764 .await
1765 {
1766 Ok(result) => result,
1767 Err(error) => {
1768 event_loop.cleanup(pending_registration).await;
1769 return Err(error);
1770 }
1771 };
1772 tracing::debug!(
1773 elapsed_ms = rpc_start.elapsed().as_millis(),
1774 "Client::create_session session creation request completed successfully"
1775 );
1776 let create_result: CreateSessionResult = match serde_json::from_value(result) {
1777 Ok(result) => result,
1778 Err(error) => {
1779 event_loop.cleanup(pending_registration).await;
1780 return Err(error.into());
1781 }
1782 };
1783
1784 if let Some(ref requested) = local_session_id
1785 && create_result.session_id != *requested
1786 {
1787 event_loop.cleanup(pending_registration).await;
1788 return Err(ErrorKind::Session(SessionErrorKind::SessionIdMismatch {
1789 requested: requested.clone(),
1790 returned: create_result.session_id.clone(),
1791 })
1792 .into());
1793 }
1794
1795 let (session_id, event_loop, registration_token) = match event_loop {
1796 CreateEventLoop::Running { running, token } => {
1797 running.startup_tasks.disarm();
1798 (
1799 create_result.session_id.clone(),
1800 running.into_handle(),
1801 token,
1802 )
1803 }
1804 CreateEventLoop::Deferred(spawn_loop) => {
1805 let (session_id, registration) = inline_stash
1806 .expect("deferred create has a response stash")
1807 .lock()
1808 .take()
1809 .expect("session registration must have populated stash on success");
1810 pending_registration.resolve_to(session_id.clone(), registration.token);
1811 let event_loop = spawn_loop(session_id.clone(), registration.channels, None);
1812 (session_id, event_loop, registration.token)
1813 }
1814 };
1815 tracing::debug!(
1816 elapsed_ms = setup_start.elapsed().as_millis(),
1817 session_id = %session_id,
1818 tools_count,
1819 commands_count,
1820 has_hooks,
1821 "Client::create_session local setup complete"
1822 );
1823 *capabilities.write() = create_result.capabilities.unwrap_or_default();
1824 if has_mcp_auth_handler
1825 && let Err(error) = register_mcp_auth_interest(self, &session_id).await
1826 {
1827 pending_registration.cleanup(event_loop).await;
1828 return Err(error);
1829 }
1830
1831 tracing::debug!(
1832 elapsed_ms = total_start.elapsed().as_millis(),
1833 session_id = %session_id,
1834 "Client::create_session complete"
1835 );
1836 pending_registration.disarm();
1837 let session = Session {
1838 id: session_id,
1839 cwd: self.cwd().clone(),
1840 workspace_path: create_result.workspace_path,
1841 remote_url: create_result.remote_url,
1842 client: self.clone(),
1843 event_loop: ParkingLotMutex::new(Some(event_loop)),
1844 shutdown,
1845 external_tools_shutdown,
1846 idle_waiter,
1847 capabilities,
1848 open_canvases,
1849 event_tx,
1850 resume_bootstrap: None,
1851 github_token_registration: ParkingLotMutex::new(github_token_registration),
1852 registration_token,
1853 };
1854 apply_mode_post_create_patch(
1855 &session,
1856 mode,
1857 opt_skip_custom_instructions,
1858 opt_custom_agents_local_only,
1859 opt_coauthor_enabled,
1860 opt_manage_schedule_enabled,
1861 opt_included_builtin_skills,
1862 )
1863 .await?;
1864 if let Some(registration) = session.github_token_registration.lock().as_ref() {
1865 registration.claim(session.id.clone());
1866 } else {
1867 self.retire_github_token_provider(&session.id);
1868 }
1869 Ok(session)
1870 }
1871
1872 async fn start_prepared_resume(
1883 &self,
1884 mut config: ResumeSessionConfig,
1885 event_tx: tokio::sync::broadcast::Sender<SessionEvent>,
1886 shutdown: CancellationToken,
1887 ) -> Result<Session, Error> {
1888 let total_start = Instant::now();
1889 let session_id = config.session_id.clone();
1890 if config.hooks_handler.is_some() && config.hooks.is_none() {
1891 config.hooks = Some(true);
1892 }
1893 if let Some(transforms) = config.system_message_transform.clone() {
1894 inject_transform_sections_resume(&mut config, transforms.as_ref());
1895 }
1896 let mode = self.inner.mode;
1897 if mode == crate::ClientMode::Empty && config.available_tools.is_none() {
1898 return Err(Error::with_message(
1899 ErrorKind::InvalidConfig,
1900 "ClientMode::Empty requires available_tools to be set on the session config. \
1901 Use ToolSet to specify which tools the session may use (e.g. \
1902 ToolSet::new().add_builtin_many(BUILTIN_TOOLS_ISOLATED)).",
1903 ));
1904 }
1905 crate::mode::validate_tool_filter_list(
1906 "available_tools",
1907 config.available_tools.as_deref(),
1908 )?;
1909 crate::mode::validate_tool_filter_list("excluded_tools", config.excluded_tools.as_deref())?;
1910 config.system_message =
1911 crate::mode::system_message_for_mode(mode, config.system_message.take());
1912 config.memory = crate::mode::memory_for_mode(mode, config.memory.take());
1913 config.enable_experimental_mode =
1914 crate::mode::experimental_mode_for_mode(mode, config.enable_experimental_mode);
1915 if mode == crate::ClientMode::Empty {
1916 if config.enable_session_telemetry.is_none() {
1917 config.enable_session_telemetry = Some(false);
1918 }
1919 if config.skip_embedding_retrieval.is_none() {
1920 config.skip_embedding_retrieval = Some(true);
1921 }
1922 if config.enable_on_demand_instruction_discovery.is_none() {
1923 config.enable_on_demand_instruction_discovery = Some(false);
1924 }
1925 if config.enable_file_hooks.is_none() {
1926 config.enable_file_hooks = Some(false);
1927 }
1928 if config.enable_host_git_operations.is_none() {
1929 config.enable_host_git_operations = Some(false);
1930 }
1931 if config.enable_session_store.is_none() {
1932 config.enable_session_store = Some(false);
1933 }
1934 if config.enable_skills.is_none() {
1935 config.enable_skills = Some(false);
1936 }
1937 }
1938 if mode == crate::ClientMode::Empty && config.mcp_oauth_token_storage.is_none() {
1939 config.mcp_oauth_token_storage = Some("in-memory".into());
1940 }
1941 if mode == crate::ClientMode::Empty && config.embedding_cache_storage.is_none() {
1942 config.embedding_cache_storage = Some("in-memory".into());
1943 }
1944 config.custom_agents_local_only =
1945 crate::mode::resolve_custom_agents_local_only(mode, config.custom_agents_local_only);
1946 let opt_skip_custom_instructions = config.skip_custom_instructions;
1947 let opt_custom_agents_local_only = config.custom_agents_local_only;
1948 let opt_coauthor_enabled = config.coauthor_enabled;
1949 let opt_manage_schedule_enabled = config.manage_schedule_enabled;
1950 let opt_included_builtin_skills = config.included_builtin_skills.take();
1951 let (mut wire, mut runtime) = config.into_wire()?;
1952 wire.enable_github_telemetry_forwarding =
1953 self.inner.on_github_telemetry.is_some().then_some(true);
1954
1955 let permission_handler = crate::permission::resolve_handler(
1956 runtime.permission_handler.take(),
1957 runtime.permission_policy.take(),
1958 );
1959 let handlers = SessionHandlers {
1960 permission: permission_handler,
1961 managed_settings_enabled: has_managed_settings(
1962 wire.enable_managed_settings,
1963 wire.managed_settings.as_ref(),
1964 ),
1965 elicitation: runtime.elicitation_handler.take(),
1966 mcp_auth: runtime.mcp_auth_handler.take(),
1967 user_input: runtime.user_input_handler.take(),
1968 exit_plan_mode: runtime.exit_plan_mode_handler.take(),
1969 auto_mode_switch: runtime.auto_mode_switch_handler.take(),
1970 tools: Arc::new(std::mem::take(&mut runtime.tool_handlers)),
1971 };
1972 let hooks = runtime.hooks_handler.take();
1973 let transforms = runtime.system_message_transform.take();
1974 let tools_count = wire.tools.as_ref().map_or(0, Vec::len);
1975 let commands_count = runtime.commands.as_ref().map_or(0, Vec::len);
1976 let has_hooks = hooks.is_some();
1977 let command_handlers = build_command_handler_map(runtime.commands.as_deref());
1978 let canvas_handler = runtime.canvas_handler.take();
1979 let session_fs_provider = runtime.session_fs_provider.take();
1980 let bearer_token_providers = std::mem::take(&mut runtime.bearer_token_providers);
1981 let github_token_registration = runtime
1982 .github_token_provider
1983 .take()
1984 .map(|provider| self.register_github_token_provider(provider));
1985 wire.github_token_provider_registration_id = github_token_registration
1986 .as_ref()
1987 .map(|registration| registration.id().to_string());
1988 let has_mcp_auth_handler = handlers.mcp_auth.is_some();
1989 if self.inner.session_fs_configured && session_fs_provider.is_none() {
1990 return Err(ErrorKind::Session(SessionErrorKind::SessionFsProviderRequired).into());
1991 }
1992 if self.inner.session_fs_sqlite_declared
1993 && let Some(ref provider) = session_fs_provider
1994 && provider.sqlite().is_none()
1995 {
1996 return Err(Error::with_message(
1997 ErrorKind::InvalidConfig,
1998 "SessionFs capabilities declare SQLite support but the provider \
1999 does not implement SessionFsSqliteProvider",
2000 ));
2001 }
2002
2003 let mut params = serde_json::to_value(&wire)?;
2004 let trace_ctx = self.resolve_trace_context().await;
2005 inject_trace_context(&mut params, &trace_ctx);
2006
2007 let capabilities = Arc::new(parking_lot::RwLock::new(SessionCapabilities::default()));
2008 let setup_start = Instant::now();
2009 let resume_bootstrap = (event_tx.receiver_count() == 0)
2011 .then(|| crate::subscription::ResumeBootstrap::new(&event_tx));
2012 let registration = self.register_session(&session_id);
2013 let registration_token = registration.token;
2014 let channels = registration.channels;
2015 let idle_waiter = Arc::new(ParkingLotMutex::new(None));
2016 let open_canvases = Arc::new(parking_lot::RwLock::new(Vec::new()));
2017 let external_tools_shutdown = self.inner.rpc.connection_closed_token();
2018 let event_loop = spawn_event_loop(
2019 session_id.clone(),
2020 self.clone(),
2021 handlers,
2022 hooks,
2023 transforms,
2024 command_handlers,
2025 canvas_handler,
2026 session_fs_provider,
2027 bearer_token_providers,
2028 channels,
2029 idle_waiter.clone(),
2030 capabilities.clone(),
2031 open_canvases.clone(),
2032 event_tx.clone(),
2033 resume_bootstrap.clone(),
2034 shutdown.clone(),
2035 external_tools_shutdown.clone(),
2036 None,
2037 );
2038 let mut registration = PendingSessionRegistration::new(
2039 self.clone(),
2040 session_id.clone(),
2041 registration_token,
2042 shutdown.clone(),
2043 external_tools_shutdown.clone(),
2044 );
2045 tracing::debug!(
2046 elapsed_ms = setup_start.elapsed().as_millis(),
2047 session_id = %session_id,
2048 tools_count,
2049 commands_count,
2050 has_hooks,
2051 "Client::resume_session local setup complete"
2052 );
2053
2054 let rpc_start = Instant::now();
2055 let result = match self.call("session.resume", Some(params)).await {
2056 Ok(result) => result,
2057 Err(error) => {
2058 registration.cleanup(event_loop).await;
2059 return Err(error);
2060 }
2061 };
2062 tracing::debug!(
2063 elapsed_ms = rpc_start.elapsed().as_millis(),
2064 session_id = %session_id,
2065 "Client::resume_session session resume request completed successfully"
2066 );
2067
2068 let resume_result: ResumeSessionResult = match serde_json::from_value(result) {
2069 Ok(result) => result,
2070 Err(error) => {
2071 registration.cleanup(event_loop).await;
2072 return Err(error.into());
2073 }
2074 };
2075 let cli_session_id = resume_result
2076 .session_id
2077 .clone()
2078 .unwrap_or_else(|| session_id.clone());
2079 if cli_session_id != session_id {
2080 registration.cleanup(event_loop).await;
2081 return Err(ErrorKind::Session(SessionErrorKind::SessionIdMismatch {
2082 requested: session_id,
2083 returned: cli_session_id,
2084 })
2085 .into());
2086 }
2087 if has_mcp_auth_handler
2088 && let Err(error) = register_mcp_auth_interest(self, &session_id).await
2089 {
2090 registration.cleanup(event_loop).await;
2091 return Err(error);
2092 }
2093 let skills_reload_start = Instant::now();
2095 if let Err(e) = self
2096 .call(
2097 "session.skills.reload",
2098 Some(serde_json::json!({ "sessionId": session_id })),
2099 )
2100 .await
2101 {
2102 warn!(
2103 elapsed_ms = skills_reload_start.elapsed().as_millis(),
2104 session_id = %session_id,
2105 error = %e,
2106 "Client::resume_session skills reload request failed"
2107 );
2108 } else {
2109 tracing::debug!(
2110 elapsed_ms = skills_reload_start.elapsed().as_millis(),
2111 session_id = %session_id,
2112 "Client::resume_session skills reload request completed successfully"
2113 );
2114 }
2115
2116 *capabilities.write() = resume_result.capabilities.unwrap_or_default();
2117 {
2122 let mut snapshots = open_canvases.write();
2123 for snapshot in resume_result.open_canvases.unwrap_or_default() {
2124 upsert_open_canvas_snapshot(&mut snapshots, snapshot);
2125 }
2126 }
2127
2128 tracing::debug!(
2129 elapsed_ms = total_start.elapsed().as_millis(),
2130 session_id = %session_id,
2131 "Client::resume_session complete"
2132 );
2133 registration.disarm();
2134 let session = Session {
2135 id: session_id,
2136 cwd: self.cwd().clone(),
2137 workspace_path: resume_result.workspace_path,
2138 remote_url: resume_result.remote_url,
2139 client: self.clone(),
2140 event_loop: ParkingLotMutex::new(Some(event_loop)),
2141 shutdown,
2142 external_tools_shutdown,
2143 idle_waiter,
2144 capabilities,
2145 open_canvases,
2146 event_tx,
2147 resume_bootstrap,
2148 github_token_registration: ParkingLotMutex::new(github_token_registration),
2149 registration_token,
2150 };
2151 apply_mode_post_create_patch(
2152 &session,
2153 mode,
2154 opt_skip_custom_instructions,
2155 opt_custom_agents_local_only,
2156 opt_coauthor_enabled,
2157 opt_manage_schedule_enabled,
2158 opt_included_builtin_skills,
2159 )
2160 .await?;
2161 if let Some(registration) = session.github_token_registration.lock().as_ref() {
2162 registration.claim(session.id.clone());
2163 } else {
2164 self.retire_github_token_provider(&session.id);
2165 }
2166 Ok(session)
2167 }
2168}
2169
2170#[must_use = "a PreparedSession does nothing until started"]
2232pub struct PreparedSession {
2233 client: Client,
2234 kind: PreparedKind,
2235 event_tx: tokio::sync::broadcast::Sender<SessionEvent>,
2236 shutdown: CancellationToken,
2237}
2238
2239enum PreparedKind {
2242 Create(Box<SessionConfig>),
2243 Resume(Box<ResumeSessionConfig>),
2244}
2245
2246impl PreparedSession {
2247 fn new(client: Client, kind: PreparedKind, event_buffer_capacity: usize) -> Self {
2248 let (event_tx, _) = tokio::sync::broadcast::channel(event_buffer_capacity);
2249 Self {
2250 client,
2251 kind,
2252 event_tx,
2253 shutdown: CancellationToken::new(),
2254 }
2255 }
2256
2257 pub fn subscribe(&self) -> crate::subscription::EventSubscription {
2275 crate::subscription::EventSubscription::new(self.event_tx.subscribe())
2276 }
2277
2278 pub async fn start(self) -> Result<Session, Error> {
2296 let Self {
2297 client,
2298 kind,
2299 event_tx,
2300 shutdown,
2301 } = self;
2302 match kind {
2303 PreparedKind::Create(config) => {
2304 client
2305 .start_prepared_create(*config, event_tx, shutdown)
2306 .await
2307 }
2308 PreparedKind::Resume(config) => {
2309 client
2310 .start_prepared_resume(*config, event_tx, shutdown)
2311 .await
2312 }
2313 }
2314 }
2315}
2316
2317type CommandHandlerMap = HashMap<String, Arc<dyn CommandHandler>>;
2318
2319async fn apply_mode_post_create_patch(
2320 session: &Session,
2321 mode: crate::ClientMode,
2322 opt_skip_custom_instructions: Option<bool>,
2323 opt_custom_agents_local_only: Option<bool>,
2324 opt_coauthor_enabled: Option<bool>,
2325 opt_manage_schedule_enabled: Option<bool>,
2326 opt_included_builtin_skills: Option<Vec<String>>,
2327) -> Result<(), Error> {
2328 let Some(patch) = build_mode_post_create_patch(
2329 mode,
2330 opt_skip_custom_instructions,
2331 opt_custom_agents_local_only,
2332 opt_coauthor_enabled,
2333 opt_manage_schedule_enabled,
2334 opt_included_builtin_skills,
2335 ) else {
2336 return Ok(());
2337 };
2338 if let Err(error) = session.rpc().options().update(patch).await {
2339 let _ = session.disconnect().await;
2340 return Err(error);
2341 }
2342 Ok(())
2343}
2344
2345fn build_mode_post_create_patch(
2355 mode: crate::ClientMode,
2356 opt_skip_custom_instructions: Option<bool>,
2357 opt_custom_agents_local_only: Option<bool>,
2358 opt_coauthor_enabled: Option<bool>,
2359 opt_manage_schedule_enabled: Option<bool>,
2360 opt_included_builtin_skills: Option<Vec<String>>,
2361) -> Option<crate::generated::api_types::SessionUpdateOptionsParams> {
2362 use crate::generated::api_types::SessionUpdateOptionsParams;
2363 let mut patch = SessionUpdateOptionsParams::default();
2364 let should_send = if mode == crate::ClientMode::Empty {
2365 patch.skip_custom_instructions = Some(opt_skip_custom_instructions.unwrap_or(true));
2366 patch.custom_agents_local_only = Some(opt_custom_agents_local_only.unwrap_or(true));
2367 patch.coauthor_enabled = Some(opt_coauthor_enabled.unwrap_or(false));
2368 patch.manage_schedule_enabled = Some(opt_manage_schedule_enabled.unwrap_or(false));
2369 patch.installed_plugins = Some(Vec::new());
2370 patch.included_builtin_skills = Some(opt_included_builtin_skills.unwrap_or_default());
2371 true
2372 } else {
2373 let mut any = false;
2374 if let Some(v) = opt_skip_custom_instructions {
2375 patch.skip_custom_instructions = Some(v);
2376 any = true;
2377 }
2378 if let Some(v) = opt_custom_agents_local_only {
2379 patch.custom_agents_local_only = Some(v);
2380 any = true;
2381 }
2382 if let Some(v) = opt_coauthor_enabled {
2383 patch.coauthor_enabled = Some(v);
2384 any = true;
2385 }
2386 if let Some(v) = opt_manage_schedule_enabled {
2387 patch.manage_schedule_enabled = Some(v);
2388 any = true;
2389 }
2390 if let Some(v) = opt_included_builtin_skills {
2391 patch.included_builtin_skills = Some(v);
2392 any = true;
2393 }
2394 any
2395 };
2396 if !should_send {
2397 return None;
2398 }
2399 Some(patch)
2400}
2401
2402fn build_command_handler_map(commands: Option<&[CommandDefinition]>) -> Arc<CommandHandlerMap> {
2403 let map = match commands {
2404 Some(commands) => commands
2405 .iter()
2406 .filter(|cmd| !cmd.name.is_empty())
2407 .map(|cmd| (cmd.name.clone(), cmd.handler.clone()))
2408 .collect(),
2409 None => HashMap::new(),
2410 };
2411 Arc::new(map)
2412}
2413
2414fn upsert_open_canvas_snapshot(
2415 snapshots: &mut Vec<OpenCanvasInstance>,
2416 snapshot: OpenCanvasInstance,
2417) {
2418 if let Some(existing) = snapshots
2419 .iter_mut()
2420 .find(|open| open.instance_id == snapshot.instance_id)
2421 {
2422 *existing = snapshot;
2423 } else {
2424 snapshots.push(snapshot);
2425 }
2426}
2427
2428fn remove_open_canvas_snapshot(snapshots: &mut Vec<OpenCanvasInstance>, instance_id: &str) {
2429 snapshots.retain(|open| open.instance_id != instance_id);
2430}
2431
2432#[allow(clippy::too_many_arguments)]
2433fn spawn_event_loop(
2434 session_id: SessionId,
2435 client: Client,
2436 handlers: SessionHandlers,
2437 hooks: Option<Arc<dyn SessionHooks>>,
2438 transforms: Option<Arc<dyn SystemMessageTransform>>,
2439 command_handlers: Arc<CommandHandlerMap>,
2440 canvas_handler: Option<Arc<dyn CanvasHandler>>,
2441 session_fs_provider: Option<Arc<dyn SessionFsProvider>>,
2442 bearer_token_providers: HashMap<String, Arc<dyn BearerTokenProvider>>,
2443 channels: crate::router::SessionChannels,
2444 idle_waiter: Arc<ParkingLotMutex<Option<IdleWaiter>>>,
2445 capabilities: Arc<parking_lot::RwLock<SessionCapabilities>>,
2446 open_canvases: Arc<parking_lot::RwLock<Vec<OpenCanvasInstance>>>,
2447 event_tx: tokio::sync::broadcast::Sender<SessionEvent>,
2448 resume_bootstrap: Option<Arc<crate::subscription::ResumeBootstrap>>,
2449 shutdown: CancellationToken,
2450 external_tools_shutdown: CancellationToken,
2451 startup_tasks: Option<Arc<StartupTasks>>,
2452) -> JoinHandle<()> {
2453 let crate::router::SessionChannels {
2454 mut notifications,
2455 mut requests,
2456 } = channels;
2457 let pending_external_tools: PendingExternalTools =
2458 Arc::new(ParkingLotMutex::new(HashMap::new()));
2459 let bootstrap_cleanup = resume_bootstrap
2462 .as_ref()
2463 .map(|bootstrap| bootstrap.cleanup_guard());
2464
2465 let span = tracing::error_span!("session_event_loop", session_id = %session_id);
2466 tokio::spawn(
2467 async move {
2468 let _bootstrap_cleanup = bootstrap_cleanup;
2469 loop {
2470 tokio::select! {
2494 _ = shutdown.cancelled() => break,
2495 Some(notification) = notifications.recv() => {
2496 handle_notification(
2497 &session_id, &client, &handlers, &command_handlers, notification, &idle_waiter, &capabilities, &open_canvases, &event_tx, resume_bootstrap.as_ref(), &shutdown, &external_tools_shutdown, &pending_external_tools, startup_tasks.as_ref(),
2498 ).await;
2499 }
2500 Some(request) = requests.recv() => {
2501 let span = tracing::error_span!("session_request_handler", session_id = %session_id);
2505 let session_id = session_id.clone();
2506 let client = client.clone();
2507 let handlers = handlers.clone();
2508 let hooks = hooks.clone();
2509 let transforms = transforms.clone();
2510 let canvas_handler = canvas_handler.clone();
2511 let session_fs_provider = session_fs_provider.clone();
2512 let bearer_token_providers = bearer_token_providers.clone();
2513 let request_id = request.id;
2514 let method = request.method.clone();
2515 spawn_startup_tracked(
2516 async move {
2517 let ctx = RequestDispatchContext {
2518 client: &client,
2519 handlers: &handlers,
2520 hooks: hooks.as_deref(),
2521 transforms: transforms.as_deref(),
2522 canvas_handler: canvas_handler.as_ref(),
2523 session_fs_provider: session_fs_provider.as_ref(),
2524 bearer_token_providers: &bearer_token_providers,
2525 };
2526 let dispatch = handle_request(&session_id, ctx, request);
2527 if AssertUnwindSafe(dispatch).catch_unwind().await.is_err() {
2528 error!(method = %method, "request handler panicked");
2531 let _ = send_error_response(
2532 &client,
2533 request_id,
2534 error_codes::INTERNAL_ERROR,
2535 "request handler panicked",
2536 )
2537 .await;
2538 }
2539 }
2540 .instrument(span),
2541 startup_tasks.as_ref(),
2542 );
2543 }
2544 else => break,
2545 }
2546 }
2547 if let Some(waiter) = idle_waiter.lock().take() {
2550 let _ = waiter
2551 .tx
2552 .send(Err(ErrorKind::Session(SessionErrorKind::EventLoopClosed).into()));
2553 }
2554 }
2555 .instrument(span),
2556 )
2557}
2558
2559fn extract_request_id(data: &Value) -> Option<RequestId> {
2560 data.get("requestId")
2561 .and_then(|v| v.as_str())
2562 .filter(|s| !s.is_empty())
2563 .map(RequestId::new)
2564}
2565
2566fn permission_request_data(
2567 event_data: &Value,
2568 managed_settings_enabled: bool,
2569) -> PermissionRequestData {
2570 let request_data = event_data
2571 .get("permissionRequest")
2572 .cloned()
2573 .unwrap_or_else(|| event_data.clone());
2574 let managed_approval_required = match request_data.get("managedApprovalRequired") {
2575 None => None,
2576 Some(Value::Bool(value)) => Some(*value),
2577 Some(_) => Some(true),
2578 };
2579 match serde_json::from_value::<PermissionRequestData>(request_data) {
2580 Ok(mut data) => {
2581 data.extra = event_data.clone();
2582 data.managed_settings_enabled = managed_settings_enabled;
2583 data
2584 }
2585 Err(_) => PermissionRequestData {
2586 kind: None,
2587 tool_call_id: None,
2588 managed_approval_required,
2589 managed_settings_enabled,
2590 extra: event_data.clone(),
2591 },
2592 }
2593}
2594
2595fn permission_response_params(
2603 session_id: &SessionId,
2604 request_id: &RequestId,
2605 result: &PermissionResult,
2606) -> Option<Value> {
2607 let (decision, decision_context) = match result {
2608 PermissionResult::Decision { decision, context } => (decision, context.clone()),
2609 PermissionResult::NoResult => return None,
2610 };
2611 let mut params = serde_json::to_value(PermissionDecisionRequest {
2612 decision_context,
2613 request_id: request_id.clone(),
2614 result: decision.clone(),
2615 })
2616 .expect("serializing permission response should succeed");
2617 params["sessionId"] =
2618 serde_json::to_value(session_id).expect("serializing session ID should succeed");
2619 Some(params)
2620}
2621
2622async fn register_mcp_auth_interest(client: &Client, session_id: &SessionId) -> Result<(), Error> {
2623 let mut params = serde_json::to_value(RegisterEventInterestParams {
2624 event_type: "mcp.oauth_required".to_string(),
2625 })?;
2626 params["sessionId"] = Value::String(session_id.to_string());
2627 client
2628 .call(rpc_methods::SESSION_EVENTLOG_REGISTERINTEREST, Some(params))
2629 .await?;
2630 Ok(())
2631}
2632
2633fn tool_failure_result(message: impl Into<String>) -> ToolResult {
2634 let message = message.into();
2635 ToolResult::Expanded(ToolResultExpanded {
2636 text_result_for_llm: message.clone(),
2637 result_type: "failure".to_string(),
2638 binary_results_for_llm: None,
2639 session_log: None,
2640 error: Some(message),
2641 tool_telemetry: None,
2642 tool_references: None,
2643 })
2644}
2645
2646fn is_autopilot_continuation_idle(event: &SessionEvent) -> bool {
2647 event
2648 .typed_data::<SessionIdleData>()
2649 .is_some_and(|data| data.mode == Some(SessionMode::Autopilot))
2650}
2651
2652#[allow(clippy::too_many_arguments)]
2654async fn handle_notification(
2655 session_id: &SessionId,
2656 client: &Client,
2657 handlers: &SessionHandlers,
2658 command_handlers: &Arc<CommandHandlerMap>,
2659 notification: SessionEventNotification,
2660 idle_waiter: &Arc<ParkingLotMutex<Option<IdleWaiter>>>,
2661 capabilities: &Arc<parking_lot::RwLock<SessionCapabilities>>,
2662 open_canvases: &Arc<parking_lot::RwLock<Vec<OpenCanvasInstance>>>,
2663 event_tx: &tokio::sync::broadcast::Sender<SessionEvent>,
2664 resume_bootstrap: Option<&Arc<crate::subscription::ResumeBootstrap>>,
2665 shutdown: &CancellationToken,
2666 external_tools_shutdown: &CancellationToken,
2667 pending_external_tools: &PendingExternalTools,
2668 startup_tasks: Option<&Arc<StartupTasks>>,
2669) {
2670 let dispatch_start = Instant::now();
2671 let event = ¬ification.event;
2672 let event_type = event.parsed_type();
2673 if event_type == SessionEventType::PermissionRequested {
2674 tracing::debug!(
2675 session_id = %session_id,
2676 event_type = %event.event_type,
2677 "Session::handle_notification permission request received"
2678 );
2679 }
2680
2681 let mut completion = None;
2684 match event_type {
2685 SessionEventType::AssistantMessage
2686 | SessionEventType::SessionIdle
2687 | SessionEventType::SessionError
2688 if is_root_agent_event(event) =>
2689 {
2690 let mut guard = idle_waiter.lock();
2691 if let Some(waiter) = guard.as_mut() {
2692 match event_type {
2693 SessionEventType::AssistantMessage => {
2694 if !waiter.first_assistant_message_seen {
2695 waiter.first_assistant_message_seen = true;
2696 tracing::debug!(
2697 elapsed_ms = waiter.started_at.elapsed().as_millis(),
2698 session_id = %session_id,
2699 "Session::send_and_wait first assistant message"
2700 );
2701 }
2702 waiter.last_assistant_message = Some(event.clone());
2703 }
2704 SessionEventType::SessionIdle if is_autopilot_continuation_idle(event) => {}
2705 SessionEventType::SessionIdle | SessionEventType::SessionError => {
2706 if let Some(waiter) = guard.take() {
2707 if event_type == SessionEventType::SessionIdle {
2708 tracing::debug!(
2709 elapsed_ms = waiter.started_at.elapsed().as_millis(),
2710 session_id = %session_id,
2711 "Session::send_and_wait idle received"
2712 );
2713 completion = Some((waiter.tx, Ok(waiter.last_assistant_message)));
2714 } else {
2715 let error_msg = event
2716 .typed_data::<SessionErrorData>()
2717 .map(|d| d.message)
2718 .or_else(|| {
2719 event
2720 .data
2721 .get("message")
2722 .and_then(|v| v.as_str())
2723 .map(|s| s.to_string())
2724 })
2725 .unwrap_or_else(|| "session error".to_string());
2726 completion = Some((
2727 waiter.tx,
2728 Err(Error::with_message(
2729 ErrorKind::Session(SessionErrorKind::AgentError),
2730 error_msg,
2731 )),
2732 ));
2733 }
2734 }
2735 }
2736 _ => {}
2737 }
2738 }
2739 }
2740 _ => {}
2741 }
2742
2743 if event_type == SessionEventType::CapabilitiesChanged {
2747 match serde_json::from_value::<SessionCapabilities>(notification.event.data.clone()) {
2748 Ok(changed) => *capabilities.write() = changed,
2749 Err(e) => warn!(error = %e, "failed to deserialize capabilities.changed payload"),
2750 }
2751 }
2752 if event_type == SessionEventType::SessionCanvasOpened {
2753 match serde_json::from_value::<OpenCanvasInstance>(notification.event.data.clone()) {
2754 Ok(open_canvas) => {
2755 upsert_open_canvas_snapshot(&mut open_canvases.write(), open_canvas);
2756 }
2757 Err(e) => warn!(error = %e, "failed to deserialize session.canvas.opened payload"),
2758 }
2759 }
2760 if event_type == SessionEventType::SessionCanvasClosed {
2761 match serde_json::from_value::<SessionCanvasClosedData>(notification.event.data.clone()) {
2762 Ok(closed) => {
2763 if closed.instance_id.is_empty() {
2764 warn!("failed to deserialize session.canvas.closed payload");
2765 } else {
2766 remove_open_canvas_snapshot(&mut open_canvases.write(), &closed.instance_id);
2767 }
2768 }
2769 Err(e) => warn!(error = %e, "failed to deserialize session.canvas.closed payload"),
2770 }
2771 }
2772
2773 if let Some(bootstrap) = resume_bootstrap {
2774 bootstrap.publish(event_tx, event.clone());
2775 } else {
2776 let _ = event_tx.send(event.clone());
2778 }
2779 if let Some((tx, result)) = completion {
2780 let _ = tx.send(result);
2781 }
2782
2783 tracing::debug!(
2784 elapsed_ms = dispatch_start.elapsed().as_millis(),
2785 session_id = %session_id,
2786 event_type = %notification.event.event_type,
2787 "Session::handle_notification dispatch"
2788 );
2789
2790 match event_type {
2793 SessionEventType::ExternalToolCompleted => {
2794 if let Some(request_id) = extract_request_id(¬ification.event.data)
2795 && let Some(token) = pending_external_tools.lock().remove(&request_id)
2796 {
2797 token.cancel();
2798 }
2799 }
2800 SessionEventType::PermissionRequested => {
2801 let Some(request_id) = extract_request_id(¬ification.event.data) else {
2802 return;
2803 };
2804 if notification
2808 .event
2809 .data
2810 .get("resolvedByHook")
2811 .and_then(|v| v.as_bool())
2812 .unwrap_or(false)
2813 {
2814 return;
2815 }
2816 let Some(permission_handler) = handlers.permission.clone() else {
2820 return;
2821 };
2822 let client = client.clone();
2823 let sid = session_id.clone();
2824 let shutdown = shutdown.clone();
2825 let data = permission_request_data(
2826 ¬ification.event.data,
2827 handlers.managed_settings_enabled,
2828 );
2829 let span = tracing::error_span!(
2830 "permission_request_handler",
2831 session_id = %sid,
2832 request_id = %request_id
2833 );
2834 spawn_startup_tracked(
2835 async move {
2836 let handler_start = Instant::now();
2837 let result = permission_handler
2838 .handle(sid.clone(), request_id.clone(), data)
2839 .await;
2840 tracing::debug!(
2841 elapsed_ms = handler_start.elapsed().as_millis(),
2842 session_id = %sid,
2843 request_id = %request_id,
2844 "PermissionHandler::handle dispatch"
2845 );
2846 let Some(params) = permission_response_params(&sid, &request_id, &result)
2847 else {
2848 return;
2852 };
2853 let rpc_start = Instant::now();
2854 let method =
2855 rpc_methods::SESSION_PERMISSIONS_HANDLEPENDINGPERMISSIONREQUEST;
2856 tokio::select! {
2857 biased;
2858 response = client.call(method, Some(params)) => {
2859 match response {
2860 Ok(_) => tracing::debug!(
2861 elapsed_ms = rpc_start.elapsed().as_millis(),
2862 session_id = %sid,
2863 request_id = %request_id,
2864 method,
2865 "Session::handle_notification response sent successfully"
2866 ),
2867 Err(error) => warn!(
2868 error = %error,
2869 session_id = %sid,
2870 request_id = %request_id,
2871 method,
2872 "failed to deliver permission decision back to the runtime"
2873 ),
2874 }
2875 }
2876 _ = shutdown.cancelled() => {
2877 warn!(
2878 elapsed_ms = rpc_start.elapsed().as_millis(),
2879 session_id = %sid,
2880 request_id = %request_id,
2881 method,
2882 delivery_outcome = "unknown",
2883 "permission confirmation acknowledgement wait cancelled during session shutdown"
2884 );
2885 }
2886 }
2887 }
2888 .instrument(span),
2889 startup_tasks,
2890 );
2891 }
2892 SessionEventType::ExternalToolRequested => {
2893 let Some(request_id) = extract_request_id(¬ification.event.data) else {
2894 return;
2895 };
2896 let data: ExternalToolRequestedData =
2897 match serde_json::from_value(notification.event.data.clone()) {
2898 Ok(d) => d,
2899 Err(e) => {
2900 warn!(error = %e, "failed to deserialize external_tool.requested");
2901 let client = client.clone();
2902 let sid = session_id.clone();
2903 let span = tracing::error_span!(
2904 "external_tool_deserialize_error",
2905 session_id = %sid,
2906 request_id = %request_id
2907 );
2908 spawn_startup_tracked(
2909 async move {
2910 let rpc_start = Instant::now();
2911 let _ = client
2912 .call(
2913 "session.tools.handlePendingToolCall",
2914 Some(serde_json::json!({
2915 "sessionId": sid,
2916 "requestId": request_id,
2917 "error": format!("Failed to deserialize tool request: {e}"),
2918 })),
2919 )
2920 .await;
2921 tracing::debug!(
2922 elapsed_ms = rpc_start.elapsed().as_millis(),
2923 session_id = %sid,
2924 request_id = %request_id,
2925 "Session::handle_notification response sent successfully"
2926 );
2927 }
2928 .instrument(span),
2929 startup_tasks,
2930 );
2931 return;
2932 }
2933 };
2934 let tool_handler = if data.tool_name.is_empty() {
2938 None
2939 } else {
2940 handlers.tools.get(&data.tool_name).cloned()
2941 };
2942 let Some(tool_handler) = tool_handler else {
2943 return;
2944 };
2945 let cancellation = Arc::new(external_tools_shutdown.child_token());
2946 {
2947 let mut pending = pending_external_tools.lock();
2948 if external_tools_shutdown.is_cancelled() || pending.contains_key(&request_id) {
2949 return;
2950 }
2951 pending.insert(request_id.clone(), cancellation.clone());
2952 }
2953 let client = client.clone();
2954 let sid = session_id.clone();
2955 let pending_external_tools = pending_external_tools.clone();
2956 let guard_request_id = request_id.clone();
2957 let guard_cancellation = cancellation.clone();
2958 let span = tracing::error_span!(
2959 "external_tool_handler",
2960 session_id = %sid,
2961 request_id = %request_id
2962 );
2963 spawn_startup_tracked(
2964 async move {
2965 let guard = PendingExternalToolGuard {
2966 request_id: guard_request_id,
2967 token: guard_cancellation,
2968 pending: pending_external_tools,
2969 };
2970 if cancellation.is_cancelled() {
2971 return;
2972 }
2973 if data.tool_call_id.is_empty() {
2978 if !guard.claim() {
2979 return;
2980 }
2981 let error_msg = "Missing toolCallId";
2982 let rpc_start = Instant::now();
2983 let _ = client
2984 .call(
2985 "session.tools.handlePendingToolCall",
2986 Some(serde_json::json!({
2987 "sessionId": sid,
2988 "requestId": request_id,
2989 "error": error_msg,
2990 })),
2991 )
2992 .await;
2993 tracing::debug!(
2994 elapsed_ms = rpc_start.elapsed().as_millis(),
2995 session_id = %sid,
2996 request_id = %request_id,
2997 "Session::handle_notification response sent successfully"
2998 );
2999 return;
3000 }
3001 let tool_call_id = data.tool_call_id.clone();
3002 let tool_name = data.tool_name.clone();
3003 let available_tools = if tool_name == TOOL_SEARCH_TOOL_NAME {
3010 let metadata_result = tokio::select! {
3011 biased;
3012 _ = cancellation.cancelled() => return,
3013 result = client.call(
3014 rpc_methods::SESSION_TOOLS_GETCURRENTMETADATA,
3015 Some(serde_json::json!({ "sessionId": sid })),
3016 ) => result,
3017 };
3018 match metadata_result {
3019 Ok(value) => {
3020 serde_json::from_value::<ToolsGetCurrentMetadataResult>(value)
3021 .ok()
3022 .and_then(|result| result.tools)
3023 }
3024 Err(_) => None,
3025 }
3026 } else {
3027 None
3028 };
3029 let invocation = ToolInvocation {
3030 session_id: sid.clone(),
3031 tool_call_id: data.tool_call_id,
3032 tool_name: data.tool_name,
3033 arguments: data
3034 .arguments
3035 .unwrap_or(Value::Object(serde_json::Map::new())),
3036 available_tools,
3037 traceparent: data.traceparent,
3038 tracestate: data.tracestate,
3039 };
3040 let handler_start = Instant::now();
3041 let tool_result = tokio::select! {
3042 biased;
3043 _ = cancellation.cancelled() => return,
3044 result = tool_handler.call(invocation) => match result {
3045 Ok(r) => r,
3046 Err(e) => tool_failure_result(e.to_string()),
3047 },
3048 };
3049 tracing::debug!(
3050 elapsed_ms = handler_start.elapsed().as_millis(),
3051 session_id = %sid,
3052 request_id = %request_id,
3053 tool_call_id = %tool_call_id,
3054 tool_name = %tool_name,
3055 "ToolHandler::call dispatch"
3056 );
3057 if !guard.claim() {
3058 return;
3059 }
3060 let result_value = serde_json::to_value(tool_result).unwrap_or(Value::Null);
3061 let rpc_start = Instant::now();
3062 let _ = client
3063 .call(
3064 "session.tools.handlePendingToolCall",
3065 Some(serde_json::json!({
3066 "sessionId": sid,
3067 "requestId": request_id,
3068 "result": result_value,
3069 })),
3070 )
3071 .await;
3072 tracing::debug!(
3073 elapsed_ms = rpc_start.elapsed().as_millis(),
3074 session_id = %sid,
3075 request_id = %request_id,
3076 tool_call_id = %tool_call_id,
3077 tool_name = %tool_name,
3078 "Session::handle_notification response sent successfully"
3079 );
3080 }
3081 .instrument(span),
3082 startup_tasks,
3083 );
3084 }
3085 SessionEventType::UserInputRequested => {
3086 }
3093 SessionEventType::ElicitationRequested => {
3094 let Some(request_id) = extract_request_id(¬ification.event.data) else {
3095 return;
3096 };
3097 let Some(elicitation_handler) = handlers.elicitation.clone() else {
3101 return;
3102 };
3103 let elicitation_data: ElicitationRequestedData =
3104 match serde_json::from_value(notification.event.data.clone()) {
3105 Ok(d) => d,
3106 Err(e) => {
3107 warn!(error = %e, "failed to deserialize elicitation request");
3108 return;
3109 }
3110 };
3111 let request = ElicitationRequest {
3112 message: elicitation_data.message,
3113 requested_schema: elicitation_data
3114 .requested_schema
3115 .map(|s| serde_json::to_value(s).unwrap_or(Value::Null)),
3116 mode: elicitation_data.mode.map(|m| match m {
3117 crate::generated::session_events::ElicitationRequestedMode::Form => {
3118 crate::types::ElicitationMode::Form
3119 }
3120 crate::generated::session_events::ElicitationRequestedMode::Url => {
3121 crate::types::ElicitationMode::Url
3122 }
3123 _ => crate::types::ElicitationMode::Unknown,
3124 }),
3125 elicitation_source: elicitation_data.elicitation_source,
3126 url: elicitation_data.url,
3127 };
3128 let client = client.clone();
3129 let sid = session_id.clone();
3130 let span = tracing::error_span!(
3131 "elicitation_request_handler",
3132 session_id = %sid,
3133 request_id = %request_id
3134 );
3135 let nested_startup_tasks = startup_tasks.cloned();
3136 spawn_startup_tracked(
3137 async move {
3138 let cancel = ElicitationResult {
3139 action: "cancel".to_string(),
3140 content: None,
3141 };
3142 let handler_task = tokio::spawn({
3144 let sid = sid.clone();
3145 let request_id = request_id.clone();
3146 let span = tracing::error_span!(
3147 "elicitation_callback",
3148 session_id = %sid,
3149 request_id = %request_id
3150 );
3151 async move {
3152 let handler_start = Instant::now();
3153 let response = elicitation_handler
3154 .handle(sid.clone(), request_id.clone(), request)
3155 .await;
3156 tracing::debug!(
3157 elapsed_ms = handler_start.elapsed().as_millis(),
3158 session_id = %sid,
3159 request_id = %request_id,
3160 "ElicitationHandler::handle dispatch"
3161 );
3162 response
3163 }
3164 .instrument(span)
3165 });
3166 if let Some(tasks) = nested_startup_tasks.as_ref() {
3167 tasks.track_nested(&handler_task);
3168 }
3169 let result = match handler_task.await {
3170 Ok(r) => r,
3171 Err(_) => cancel.clone(),
3172 };
3173 let rpc_start = Instant::now();
3174 if let Err(e) = client
3175 .call(
3176 "session.ui.handlePendingElicitation",
3177 Some(serde_json::json!({
3178 "sessionId": sid,
3179 "requestId": request_id,
3180 "result": result,
3181 })),
3182 )
3183 .await
3184 {
3185 warn!(error = %e, "handlePendingElicitation failed, sending cancel");
3187 let _ = client
3188 .call(
3189 "session.ui.handlePendingElicitation",
3190 Some(serde_json::json!({
3191 "sessionId": sid,
3192 "requestId": request_id,
3193 "result": cancel,
3194 })),
3195 )
3196 .await;
3197 } else {
3198 tracing::debug!(
3199 elapsed_ms = rpc_start.elapsed().as_millis(),
3200 session_id = %sid,
3201 request_id = %request_id,
3202 "Session::handle_notification response sent successfully"
3203 );
3204 }
3205 }
3206 .instrument(span),
3207 startup_tasks,
3208 );
3209 }
3210 SessionEventType::McpOauthRequired => {
3211 let Some(request_id) = extract_request_id(¬ification.event.data) else {
3212 return;
3213 };
3214 let Some(mcp_auth_handler) = handlers.mcp_auth.clone() else {
3215 warn!(
3216 session_id = %session_id,
3217 request_id = %request_id,
3218 "received MCP OAuth request without a registered MCP auth handler"
3219 );
3220 return;
3221 };
3222 let data: McpOauthRequiredData =
3223 match serde_json::from_value(notification.event.data.clone()) {
3224 Ok(d) => d,
3225 Err(e) => {
3226 warn!(error = %e, "failed to deserialize MCP OAuth request");
3227 return;
3228 }
3229 };
3230 let request = McpAuthRequest {
3231 request_id: request_id.clone(),
3232 server_name: data.server_name,
3233 server_url: data.server_url,
3234 reason: data.reason,
3235 www_authenticate_params: data.www_authenticate_params,
3236 resource_metadata: data.resource_metadata,
3237 static_client_config: data.static_client_config,
3238 };
3239 let client = client.clone();
3240 let sid = session_id.clone();
3241 let span = tracing::error_span!(
3242 "mcp_auth_request_handler",
3243 session_id = %sid,
3244 request_id = %request_id
3245 );
3246 let nested_startup_tasks = startup_tasks.cloned();
3247 spawn_startup_tracked(
3248 async move {
3249 let cancel = McpAuthResult::Cancelled;
3250 let handler_task = tokio::spawn({
3251 let sid = sid.clone();
3252 let request_id = request_id.clone();
3253 let span = tracing::error_span!(
3254 "mcp_auth_callback",
3255 session_id = %sid,
3256 request_id = %request_id
3257 );
3258 async move {
3259 let handler_start = Instant::now();
3260 let response = mcp_auth_handler
3261 .handle(sid.clone(), request_id.clone(), request)
3262 .await;
3263 tracing::debug!(
3264 elapsed_ms = handler_start.elapsed().as_millis(),
3265 session_id = %sid,
3266 request_id = %request_id,
3267 "McpAuthHandler::handle dispatch"
3268 );
3269 response
3270 }
3271 .instrument(span)
3272 });
3273 if let Some(tasks) = nested_startup_tasks.as_ref() {
3274 tasks.track_nested(&handler_task);
3275 }
3276 let result = match handler_task.await {
3277 Ok(result) => result,
3278 Err(_) => cancel,
3279 };
3280 let rpc_start = Instant::now();
3281 let _ = client
3282 .call(
3283 "session.mcp.oauth.handlePendingRequest",
3284 Some(serde_json::json!({
3285 "sessionId": sid,
3286 "requestId": request_id,
3287 "result": result.into_wire(),
3288 })),
3289 )
3290 .await;
3291 tracing::debug!(
3292 elapsed_ms = rpc_start.elapsed().as_millis(),
3293 "Session::handle_notification MCP auth response sent"
3294 );
3295 }
3296 .instrument(span),
3297 startup_tasks,
3298 );
3299 }
3300 SessionEventType::CommandExecute => {
3301 let data: CommandExecuteData =
3302 match serde_json::from_value(notification.event.data.clone()) {
3303 Ok(d) => d,
3304 Err(e) => {
3305 warn!(error = %e, "failed to deserialize command.execute");
3306 return;
3307 }
3308 };
3309 let client = client.clone();
3310 let command_handlers = command_handlers.clone();
3311 let sid = session_id.clone();
3312 let span = tracing::error_span!("command_handler", session_id = %sid);
3313 spawn_startup_tracked(
3314 async move {
3315 let request_id = data.request_id;
3316 let ack_error = match command_handlers.get(&data.command_name).cloned() {
3317 None => Some(format!("Unknown command: {}", data.command_name)),
3318 Some(handler) => {
3319 let command_name = data.command_name.clone();
3320 let ctx = CommandContext {
3321 session_id: sid.clone(),
3322 command: data.command,
3323 command_name: data.command_name,
3324 args: data.args,
3325 };
3326 let handler_start = Instant::now();
3327 let result = handler.on_command(ctx).await;
3328 tracing::debug!(
3329 elapsed_ms = handler_start.elapsed().as_millis(),
3330 session_id = %sid,
3331 request_id = %request_id,
3332 command_name = %command_name,
3333 "CommandHandler::call dispatch"
3334 );
3335 match result {
3336 Ok(()) => None,
3337 Err(e) => Some(e.to_string()),
3338 }
3339 }
3340 };
3341 let mut params = serde_json::json!({
3342 "sessionId": sid,
3343 "requestId": request_id,
3344 });
3345 if let Some(error_msg) = ack_error {
3346 params["error"] = serde_json::Value::String(error_msg);
3347 }
3348 let rpc_start = Instant::now();
3349 let _ = client
3350 .call("session.commands.handlePendingCommand", Some(params))
3351 .await;
3352 tracing::debug!(
3353 elapsed_ms = rpc_start.elapsed().as_millis(),
3354 session_id = %sid,
3355 request_id = %request_id,
3356 "Session::handle_notification response sent successfully"
3357 );
3358 }
3359 .instrument(span),
3360 startup_tasks,
3361 );
3362 }
3363 _ => {}
3364 }
3365}
3366
3367struct RequestDispatchContext<'a> {
3368 client: &'a Client,
3369 handlers: &'a SessionHandlers,
3370 hooks: Option<&'a dyn SessionHooks>,
3371 transforms: Option<&'a dyn SystemMessageTransform>,
3372 canvas_handler: Option<&'a Arc<dyn CanvasHandler>>,
3373 session_fs_provider: Option<&'a Arc<dyn SessionFsProvider>>,
3374 bearer_token_providers: &'a HashMap<String, Arc<dyn BearerTokenProvider>>,
3375}
3376
3377async fn handle_request(
3379 session_id: &SessionId,
3380 ctx: RequestDispatchContext<'_>,
3381 request: crate::JsonRpcRequest,
3382) {
3383 let sid = session_id.clone();
3384 let client = ctx.client;
3385 let handlers = ctx.handlers;
3386 let hooks = ctx.hooks;
3387 let transforms = ctx.transforms;
3388 let canvas_handler = ctx.canvas_handler;
3389 let session_fs_provider = ctx.session_fs_provider;
3390 let bearer_token_providers = ctx.bearer_token_providers;
3391
3392 if request.method.starts_with("sessionFs.") {
3393 crate::session_fs_dispatch::dispatch(client, session_fs_provider, request).await;
3394 return;
3395 }
3396
3397 if request.method.starts_with("canvas.") {
3398 crate::canvas_dispatch::dispatch(client, canvas_handler, request).await;
3399 return;
3400 }
3401
3402 if request.method == crate::generated::api_types::rpc_methods::PROVIDERTOKEN_GETTOKEN {
3403 crate::provider_token_dispatch::dispatch(client, bearer_token_providers, request).await;
3404 return;
3405 }
3406
3407 match request.method.as_str() {
3408 "hooks.invoke" => {
3409 let params = request.params.as_ref();
3410 let hook_type = params
3411 .and_then(|p| p.get("hookType"))
3412 .and_then(|v| v.as_str())
3413 .unwrap_or("");
3414 let input = params
3415 .and_then(|p| p.get("input"))
3416 .cloned()
3417 .unwrap_or(Value::Object(Default::default()));
3418
3419 let rpc_result = if let Some(hooks) = hooks {
3420 match crate::hooks::dispatch_hook(hooks, &sid, hook_type, input).await {
3421 Ok(output) => output,
3422 Err(e) => {
3423 warn!(error = %e, hook_type = hook_type, "hook dispatch failed");
3424 serde_json::json!({ "output": {} })
3425 }
3426 }
3427 } else {
3428 serde_json::json!({ "output": {} })
3429 };
3430
3431 let rpc_response = JsonRpcResponse {
3432 jsonrpc: "2.0".to_string(),
3433 id: request.id,
3434 result: Some(rpc_result),
3435 error: None,
3436 };
3437 let _ = client.send_response(&rpc_response).await;
3438 }
3439
3440 "userInput.request" => {
3441 let params = request.params.as_ref();
3442 let Some(question) = params
3443 .and_then(|p| p.get("question"))
3444 .and_then(|v| v.as_str())
3445 else {
3446 warn!("userInput.request missing 'question' field");
3447 let rpc_response = JsonRpcResponse {
3448 jsonrpc: "2.0".to_string(),
3449 id: request.id,
3450 result: None,
3451 error: Some(crate::JsonRpcError {
3452 code: error_codes::INVALID_PARAMS,
3453 message: "missing required field: question".to_string(),
3454 data: None,
3455 }),
3456 };
3457 let _ = client.send_response(&rpc_response).await;
3458 return;
3459 };
3460 let question = question.to_string();
3461 let choices = params
3462 .and_then(|p| p.get("choices"))
3463 .and_then(|v| v.as_array())
3464 .map(|arr| {
3465 arr.iter()
3466 .filter_map(|v| v.as_str().map(|s| s.to_string()))
3467 .collect()
3468 });
3469 let allow_freeform = params
3470 .and_then(|p| p.get("allowFreeform"))
3471 .and_then(|v| v.as_bool());
3472
3473 let handler_start = Instant::now();
3474 let response = if let Some(user_input_handler) = handlers.user_input.as_ref() {
3475 user_input_handler
3476 .handle(sid.clone(), question, choices, allow_freeform)
3477 .await
3478 } else {
3479 None
3480 };
3481 tracing::debug!(
3482 elapsed_ms = handler_start.elapsed().as_millis(),
3483 session_id = %sid,
3484 "UserInputHandler::handle dispatch"
3485 );
3486
3487 let rpc_result = match response {
3488 Some(UserInputResponse {
3489 answer,
3490 was_freeform,
3491 }) => serde_json::json!({
3492 "answer": answer,
3493 "wasFreeform": was_freeform,
3494 }),
3495 None => serde_json::json!({ "noResponse": true }),
3496 };
3497 let rpc_response = JsonRpcResponse {
3498 jsonrpc: "2.0".to_string(),
3499 id: request.id,
3500 result: Some(rpc_result),
3501 error: None,
3502 };
3503 let _ = client.send_response(&rpc_response).await;
3504 }
3505
3506 "exitPlanMode.request" => {
3507 let params = request
3508 .params
3509 .as_ref()
3510 .cloned()
3511 .unwrap_or(Value::Object(serde_json::Map::new()));
3512 let data: ExitPlanModeData = match serde_json::from_value(params) {
3513 Ok(d) => d,
3514 Err(e) => {
3515 warn!(error = %e, "failed to deserialize exitPlanMode.request params, using defaults");
3516 ExitPlanModeData::default()
3517 }
3518 };
3519
3520 let rpc_result = if let Some(exit_plan_handler) = handlers.exit_plan_mode.as_ref() {
3521 let result = exit_plan_handler.handle(sid, data).await;
3522 serde_json::to_value(result).expect("ExitPlanModeResult serialization cannot fail")
3523 } else {
3524 serde_json::json!({ "approved": true })
3525 };
3526 let rpc_response = JsonRpcResponse {
3527 jsonrpc: "2.0".to_string(),
3528 id: request.id,
3529 result: Some(rpc_result),
3530 error: None,
3531 };
3532 let _ = client.send_response(&rpc_response).await;
3533 }
3534
3535 "autoModeSwitch.request" => {
3536 let error_code = request
3537 .params
3538 .as_ref()
3539 .and_then(|p| p.get("errorCode"))
3540 .and_then(|v| v.as_str())
3541 .map(|s| s.to_string());
3542 let retry_after_seconds = request
3543 .params
3544 .as_ref()
3545 .and_then(|p| p.get("retryAfterSeconds"))
3546 .and_then(|v| v.as_f64());
3547
3548 let answer = if let Some(auto_mode_handler) = handlers.auto_mode_switch.as_ref() {
3549 auto_mode_handler
3550 .handle(sid, error_code, retry_after_seconds)
3551 .await
3552 } else {
3553 AutoModeSwitchResponse::No
3554 };
3555 let rpc_response = JsonRpcResponse {
3556 jsonrpc: "2.0".to_string(),
3557 id: request.id,
3558 result: Some(serde_json::json!({ "response": answer })),
3559 error: None,
3560 };
3561 let _ = client.send_response(&rpc_response).await;
3562 }
3563
3564 "systemMessage.transform" => {
3565 let params = request.params.as_ref();
3566 let sections: HashMap<String, crate::transforms::TransformSection> =
3567 match params.and_then(|p| p.get("sections")) {
3568 Some(v) => match serde_json::from_value(v.clone()) {
3569 Ok(s) => s,
3570 Err(e) => {
3571 let _ = send_error_response(
3572 client,
3573 request.id,
3574 error_codes::INVALID_PARAMS,
3575 &format!("invalid sections: {e}"),
3576 )
3577 .await;
3578 return;
3579 }
3580 },
3581 None => {
3582 let _ = send_error_response(
3583 client,
3584 request.id,
3585 error_codes::INVALID_PARAMS,
3586 "missing sections parameter",
3587 )
3588 .await;
3589 return;
3590 }
3591 };
3592
3593 let rpc_result = if let Some(transforms) = transforms {
3594 let transform_start = Instant::now();
3595 let response =
3596 crate::transforms::dispatch_transform(transforms, &sid, sections).await;
3597 tracing::debug!(
3598 elapsed_ms = transform_start.elapsed().as_millis(),
3599 session_id = %sid,
3600 "SystemMessageTransform::transform_section dispatch"
3601 );
3602 match serde_json::to_value(response) {
3603 Ok(v) => v,
3604 Err(e) => {
3605 warn!(error = %e, "failed to serialize transform response");
3606 serde_json::json!({ "sections": {} })
3607 }
3608 }
3609 } else {
3610 let passthrough: HashMap<String, crate::transforms::TransformSection> = sections;
3612 serde_json::json!({ "sections": passthrough })
3613 };
3614
3615 let rpc_response = JsonRpcResponse {
3616 jsonrpc: "2.0".to_string(),
3617 id: request.id,
3618 result: Some(rpc_result),
3619 error: None,
3620 };
3621 let _ = client.send_response(&rpc_response).await;
3622 }
3623
3624 method => {
3625 warn!(
3626 method = method,
3627 "unhandled request method in session event loop"
3628 );
3629 let _ = send_error_response(
3630 client,
3631 request.id,
3632 error_codes::METHOD_NOT_FOUND,
3633 &format!("unknown method: {method}"),
3634 )
3635 .await;
3636 }
3637 }
3638}
3639
3640async fn send_error_response(
3641 client: &Client,
3642 id: u64,
3643 code: i32,
3644 message: &str,
3645) -> Result<(), Error> {
3646 let response = JsonRpcResponse {
3647 jsonrpc: "2.0".to_string(),
3648 id,
3649 result: None,
3650 error: Some(crate::JsonRpcError {
3651 code,
3652 message: message.to_string(),
3653 data: None,
3654 }),
3655 };
3656 client.send_response(&response).await
3657}
3658
3659fn apply_transform_sections(
3663 sys_msg: &mut SystemMessageConfig,
3664 transforms: &dyn SystemMessageTransform,
3665) {
3666 sys_msg.mode = Some("customize".to_string());
3667 let sections = sys_msg.sections.get_or_insert_with(HashMap::new);
3668 for id in transforms.section_ids() {
3669 sections.entry(id).or_insert_with(|| SectionOverride {
3670 action: Some("transform".to_string()),
3671 content: None,
3672 });
3673 }
3674}
3675
3676fn inject_transform_sections(config: &mut SessionConfig, transforms: &dyn SystemMessageTransform) {
3677 let sys_msg = config.system_message.get_or_insert_with(Default::default);
3678 apply_transform_sections(sys_msg, transforms);
3679}
3680
3681fn inject_transform_sections_resume(
3682 config: &mut ResumeSessionConfig,
3683 transforms: &dyn SystemMessageTransform,
3684) {
3685 let sys_msg = config.system_message.get_or_insert_with(Default::default);
3686 apply_transform_sections(sys_msg, transforms);
3687}
3688
3689#[cfg(test)]
3690mod tests {
3691 use serde_json::json;
3692
3693 use super::{
3694 build_mode_post_create_patch, has_managed_settings, is_autopilot_continuation_idle,
3695 permission_request_data, permission_response_params,
3696 };
3697 use crate::handler::PermissionResult;
3698 use crate::types::{
3699 PermissionDecisionContext, PermissionDecisionOutcome, PermissionDecisionSource,
3700 PermissionDecisionSurface, RequestId, SessionEvent, SessionId,
3701 };
3702
3703 #[test]
3704 fn identifies_only_autopilot_continuation_idles() {
3705 let mut event = SessionEvent {
3706 id: "event-1".to_string(),
3707 timestamp: "2026-01-01T00:00:00Z".to_string(),
3708 parent_id: None,
3709 ephemeral: None,
3710 agent_id: None,
3711 debug_cli_received_at_ms: None,
3712 debug_ws_forwarded_at_ms: None,
3713 event_type: "session.idle".to_string(),
3714 data: json!({ "mode": "autopilot" }),
3715 };
3716
3717 assert!(is_autopilot_continuation_idle(&event));
3718
3719 event.data = json!({ "mode": "interactive" });
3720 assert!(!is_autopilot_continuation_idle(&event));
3721
3722 event.data = json!({});
3723 assert!(!is_autopilot_continuation_idle(&event));
3724 }
3725
3726 #[test]
3727 fn empty_mode_post_patch_sets_empty_included_builtin_skills() {
3728 let patch =
3729 build_mode_post_create_patch(crate::ClientMode::Empty, None, None, None, None, None)
3730 .expect("empty mode always sends a patch");
3731 assert_eq!(
3732 patch.included_builtin_skills,
3733 Some(Vec::new()),
3734 "empty mode must fail closed with an empty includedBuiltinSkills list"
3735 );
3736 assert_eq!(patch.installed_plugins.as_ref().map(|p| p.len()), Some(0));
3737 let value = serde_json::to_value(&patch).expect("serialize patch");
3739 assert_eq!(value["includedBuiltinSkills"], serde_json::json!([]));
3740 }
3741
3742 #[test]
3743 fn empty_mode_post_patch_preserves_explicit_builtin_skill_allowlist() {
3744 let patch = build_mode_post_create_patch(
3745 crate::ClientMode::Empty,
3746 Some(false),
3747 Some(false),
3748 Some(true),
3749 Some(true),
3750 Some(vec!["code-review".to_string()]),
3751 )
3752 .expect("empty mode always sends a patch");
3753 assert_eq!(
3754 patch.included_builtin_skills,
3755 Some(vec!["code-review".to_string()])
3756 );
3757 }
3758
3759 #[test]
3760 fn copilot_cli_mode_does_not_inject_included_builtin_skills() {
3761 assert!(
3763 build_mode_post_create_patch(
3764 crate::ClientMode::CopilotCli,
3765 None,
3766 None,
3767 None,
3768 None,
3769 None
3770 )
3771 .is_none()
3772 );
3773 let patch = build_mode_post_create_patch(
3775 crate::ClientMode::CopilotCli,
3776 Some(true),
3777 None,
3778 None,
3779 None,
3780 None,
3781 )
3782 .expect("a set field triggers a patch");
3783 assert_eq!(patch.included_builtin_skills, None);
3784 assert!(patch.installed_plugins.is_none());
3785 let value = serde_json::to_value(&patch).expect("serialize patch");
3786 assert!(value.get("includedBuiltinSkills").is_none());
3787
3788 let patch = build_mode_post_create_patch(
3789 crate::ClientMode::CopilotCli,
3790 None,
3791 None,
3792 None,
3793 None,
3794 Some(vec!["code-review".to_string()]),
3795 )
3796 .expect("an explicit allowlist triggers a patch");
3797 assert_eq!(
3798 patch.included_builtin_skills,
3799 Some(vec!["code-review".to_string()])
3800 );
3801 }
3802
3803 #[test]
3804 fn direct_injection_enables_managed_safeguards() {
3805 let settings = crate::types::ManagedSettings::default();
3806 assert!(has_managed_settings(None, Some(&settings)));
3807 assert!(!has_managed_settings(None, None));
3808 }
3809
3810 fn attribution_context() -> PermissionDecisionContext {
3811 PermissionDecisionContext {
3812 outcome: PermissionDecisionOutcome::AutoApproved,
3813 response_capability: None,
3814 source: PermissionDecisionSource::AssistedApproval,
3815 surface: PermissionDecisionSurface::CopilotApp,
3816 }
3817 }
3818
3819 #[test]
3820 fn response_params_omit_decision_context_without_attribution() {
3821 for (result, expected) in [
3822 (
3823 PermissionResult::approve_once(),
3824 json!({ "kind": "approve-once" }),
3825 ),
3826 (PermissionResult::reject(None), json!({ "kind": "reject" })),
3827 (
3828 PermissionResult::reject(Some("bad".to_string())),
3829 json!({ "kind": "reject", "feedback": "bad" }),
3830 ),
3831 (
3832 PermissionResult::user_not_available(),
3833 json!({ "kind": "user-not-available" }),
3834 ),
3835 ] {
3836 let params = permission_response_params(
3837 &SessionId::from("session-1"),
3838 &RequestId::from("permission-1"),
3839 &result,
3840 )
3841 .unwrap();
3842 assert_eq!(
3843 params,
3844 json!({
3845 "sessionId": "session-1",
3846 "requestId": "permission-1",
3847 "result": expected,
3848 })
3849 );
3850 }
3851 }
3852
3853 #[test]
3854 fn response_params_forward_decision_context_alongside_result() {
3855 let params = permission_response_params(
3856 &SessionId::from("session-1"),
3857 &RequestId::from("permission-1"),
3858 &PermissionResult::approve_once().with_context(attribution_context()),
3859 )
3860 .unwrap();
3861 assert_eq!(
3862 params,
3863 json!({
3864 "sessionId": "session-1",
3865 "requestId": "permission-1",
3866 "result": { "kind": "approve-once" },
3867 "decisionContext": {
3868 "outcome": "auto_approved",
3869 "source": "assisted_approval",
3870 "surface": "copilot_app",
3871 },
3872 })
3873 );
3874 assert!(params["result"].get("decisionContext").is_none());
3876 }
3877
3878 #[test]
3879 fn response_params_suppressed_for_no_result() {
3880 assert!(
3881 permission_response_params(
3882 &SessionId::from("session-1"),
3883 &RequestId::from("permission-1"),
3884 &PermissionResult::NoResult,
3885 )
3886 .is_none()
3887 );
3888 }
3889
3890 #[test]
3891 fn with_context_is_a_no_op_on_no_result() {
3892 let result = PermissionResult::no_result().with_context(attribution_context());
3893 assert!(matches!(result, PermissionResult::NoResult));
3894 }
3895
3896 #[test]
3897 fn with_context_replaces_rather_than_nests() {
3898 let result = PermissionResult::approve_once()
3899 .with_context(attribution_context())
3900 .with_context(PermissionDecisionContext {
3901 outcome: PermissionDecisionOutcome::PromptedUser,
3902 response_capability: None,
3903 source: PermissionDecisionSource::HumanResponse,
3904 surface: PermissionDecisionSurface::Sdk,
3905 });
3906 let params = permission_response_params(
3907 &SessionId::from("session-1"),
3908 &RequestId::from("permission-1"),
3909 &result,
3910 )
3911 .unwrap();
3912 assert_eq!(
3913 params["decisionContext"],
3914 json!({
3915 "outcome": "prompted_user",
3916 "source": "human_response",
3917 "surface": "sdk",
3918 })
3919 );
3920 }
3921
3922 #[test]
3923 fn permission_request_data_reads_nested_managed_approval_metadata() {
3924 let data = permission_request_data(
3925 &json!({
3926 "requestId": "permission-1",
3927 "permissionRequest": {
3928 "kind": "read",
3929 "managedApprovalRequired": true,
3930 "path": "/workspace/file.txt"
3931 }
3932 }),
3933 false,
3934 );
3935
3936 assert_eq!(data.managed_approval_required, Some(true));
3937 assert_eq!(
3938 data.extra["permissionRequest"]["path"],
3939 "/workspace/file.txt"
3940 );
3941 }
3942
3943 #[test]
3944 fn permission_request_data_preserves_managed_flag_when_other_fields_are_malformed() {
3945 let data = permission_request_data(
3946 &json!({
3947 "requestId": "permission-1",
3948 "permissionRequest": {
3949 "kind": "read",
3950 "managedApprovalRequired": true,
3951 "toolCallId": 42
3952 }
3953 }),
3954 false,
3955 );
3956
3957 assert_eq!(data.managed_approval_required, Some(true));
3958 assert_eq!(data.extra["requestId"], "permission-1");
3959 }
3960
3961 #[test]
3962 fn permission_request_data_fails_closed_for_malformed_managed_flag() {
3963 let data = permission_request_data(
3964 &json!({
3965 "requestId": "permission-1",
3966 "permissionRequest": {
3967 "kind": "read",
3968 "managedApprovalRequired": "yes",
3969 "path": "/workspace/file.txt"
3970 }
3971 }),
3972 false,
3973 );
3974
3975 assert_eq!(data.managed_approval_required, Some(true));
3976 }
3977
3978 #[test]
3979 fn permission_request_data_preserves_valid_false_managed_flag() {
3980 let data = permission_request_data(
3981 &json!({
3982 "requestId": "permission-1",
3983 "permissionRequest": {
3984 "kind": "read",
3985 "managedApprovalRequired": false,
3986 "path": "/workspace/file.txt"
3987 }
3988 }),
3989 false,
3990 );
3991
3992 assert_eq!(data.managed_approval_required, Some(false));
3993 }
3994}