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 startup_tasks: Arc<StartupTasks>,
262 disarmed: bool,
263}
264
265enum PendingSessionId {
273 Known(SessionId, crate::router::RegistrationToken),
276 Deferred(Arc<ParkingLotMutex<Option<(SessionId, crate::router::SessionRegistration)>>>),
279}
280
281impl PendingSessionRegistration {
282 fn new(
283 client: Client,
284 session_id: SessionId,
285 token: crate::router::RegistrationToken,
286 shutdown: CancellationToken,
287 external_tools_shutdown: CancellationToken,
288 ) -> Self {
289 Self {
290 client,
291 session_id: PendingSessionId::Known(session_id, token),
292 shutdown,
293 external_tools_shutdown,
294 startup_tasks: Arc::new(StartupTasks::default()),
295 disarmed: false,
296 }
297 }
298
299 fn deferred(
301 client: Client,
302 stash: Arc<ParkingLotMutex<Option<(SessionId, crate::router::SessionRegistration)>>>,
303 shutdown: CancellationToken,
304 external_tools_shutdown: CancellationToken,
305 ) -> (Self, crate::jsonrpc::InlineResponseCallback) {
306 let guard = Self {
307 client: client.clone(),
308 session_id: PendingSessionId::Deferred(stash.clone()),
309 shutdown: shutdown.clone(),
310 external_tools_shutdown,
311 startup_tasks: Arc::new(StartupTasks::default()),
312 disarmed: false,
313 };
314 let callback: crate::jsonrpc::InlineResponseCallback = Box::new(move |response| {
315 let result = response.result.as_ref().ok_or_else(|| {
316 Error::with_message(ErrorKind::Json, "session.create response had no result")
317 })?;
318 let parsed: CreateSessionResult = serde_json::from_value(result.clone())?;
319 let mut stashed = stash.lock();
322 if shutdown.is_cancelled() {
323 return Err(ErrorKind::Session(SessionErrorKind::EventLoopClosed).into());
324 }
325 let registration = client.register_session(&parsed.session_id);
326 *stashed = Some((parsed.session_id, registration));
327 Ok(())
328 });
329 (guard, callback)
330 }
331
332 fn registered_id(&self) -> Option<SessionId> {
333 match &self.session_id {
334 PendingSessionId::Known(id, _) => Some(id.clone()),
335 PendingSessionId::Deferred(stash) => stash.lock().as_ref().map(|(id, _)| id.clone()),
336 }
337 }
338
339 fn resolve_to(&mut self, session_id: SessionId, token: crate::router::RegistrationToken) {
343 self.session_id = PendingSessionId::Known(session_id, token);
344 }
345
346 fn cancel(&self) {
347 self.external_tools_shutdown.cancel();
348 self.shutdown.cancel();
349 self.startup_tasks.abort();
351 }
352
353 async fn cleanup(mut self, event_loop: JoinHandle<()>) {
354 self.cancel();
355 let _ = event_loop.await;
356 if let Some(id) = self.registered_id() {
357 if let PendingSessionId::Known(_, token) = self.session_id {
358 self.client.unregister_session_owned(&id, token);
359 } else if let PendingSessionId::Deferred(stash) = &self.session_id
360 && let Some((id, registration)) = stash.lock().as_ref()
361 {
362 self.client.unregister_session_owned(id, registration.token);
363 }
364 }
365 self.disarmed = true;
366 }
367
368 fn disarm(&mut self) {
369 self.startup_tasks.disarm();
370 self.disarmed = true;
371 }
372}
373
374impl Drop for PendingSessionRegistration {
375 fn drop(&mut self) {
376 if !self.disarmed {
377 self.cancel();
378 if let Some(id) = self.registered_id() {
379 if let PendingSessionId::Known(_, token) = self.session_id {
380 self.client.unregister_session_owned(&id, token);
381 } else if let PendingSessionId::Deferred(stash) = &self.session_id
382 && let Some((id, registration)) = stash.lock().as_ref()
383 {
384 self.client.unregister_session_owned(id, registration.token);
385 }
386 }
387 }
388 }
389}
390
391type EventLoopSpawner = Box<
392 dyn FnOnce(
393 SessionId,
394 crate::router::SessionChannels,
395 Option<Arc<StartupTasks>>,
396 ) -> JoinHandle<()>
397 + Send,
398>;
399
400struct StartupTasks {
401 state: ParkingLotMutex<StartupTaskState>,
402}
403
404enum StartupTaskState {
405 Pending(Vec<tokio::task::AbortHandle>),
406 Disarmed,
407 Aborted,
408}
409
410impl Default for StartupTasks {
411 fn default() -> Self {
412 Self {
413 state: ParkingLotMutex::new(StartupTaskState::Pending(Vec::new())),
414 }
415 }
416}
417
418impl StartupTasks {
419 fn abort(&self) {
420 if let StartupTaskState::Pending(tasks) =
421 std::mem::replace(&mut *self.state.lock(), StartupTaskState::Aborted)
422 {
423 for handle in tasks {
426 handle.abort();
427 }
428 }
429 }
430
431 fn disarm(&self) {
432 *self.state.lock() = StartupTaskState::Disarmed;
433 }
434}
435
436fn spawn_startup_tracked<F>(
437 future: F,
438 startup_tasks: Option<&Arc<StartupTasks>>,
439) -> Option<JoinHandle<F::Output>>
440where
441 F: std::future::Future + Send + 'static,
442 F::Output: Send + 'static,
443{
444 let mut state = startup_tasks.map(|tasks| tasks.state.lock());
447 if matches!(state.as_deref(), Some(StartupTaskState::Aborted)) {
448 return None;
449 }
450 let task = tokio::spawn(future);
451 if let Some(StartupTaskState::Pending(tasks)) = state.as_deref_mut() {
452 tasks.push(task.abort_handle());
453 }
454 Some(task)
455}
456
457enum CreateEventLoop {
461 Running {
462 event_loop: JoinHandle<()>,
463 token: crate::router::RegistrationToken,
464 },
465 Deferred(EventLoopSpawner),
466}
467
468impl CreateEventLoop {
469 async fn cleanup(self, pending_registration: PendingSessionRegistration) {
470 match self {
471 Self::Running { event_loop, .. } => {
472 pending_registration.cleanup(event_loop).await;
473 }
474 Self::Deferred(_) => drop(pending_registration),
475 }
476 }
477}
478
479pub struct Session {
492 id: SessionId,
493 cwd: PathBuf,
494 workspace_path: Option<PathBuf>,
495 remote_url: Option<String>,
496 client: Client,
497 event_loop: ParkingLotMutex<Option<JoinHandle<()>>>,
502 shutdown: CancellationToken,
516 external_tools_shutdown: CancellationToken,
519 idle_waiter: Arc<ParkingLotMutex<Option<IdleWaiter>>>,
526 capabilities: Arc<parking_lot::RwLock<SessionCapabilities>>,
528 open_canvases: Arc<parking_lot::RwLock<Vec<OpenCanvasInstance>>>,
530 event_tx: tokio::sync::broadcast::Sender<SessionEvent>,
532 resume_bootstrap: Option<Arc<crate::subscription::ResumeBootstrap>>,
535 github_token_registration:
536 ParkingLotMutex<Option<crate::github_token::GitHubTokenRegistration>>,
537 registration_token: crate::router::RegistrationToken,
539}
540
541impl Session {
542 pub fn id(&self) -> &SessionId {
544 &self.id
545 }
546
547 pub fn cwd(&self) -> &PathBuf {
549 &self.cwd
550 }
551
552 pub fn workspace_path(&self) -> Option<&Path> {
554 self.workspace_path.as_deref()
555 }
556
557 pub fn remote_url(&self) -> Option<&str> {
559 self.remote_url.as_deref()
560 }
561
562 pub fn capabilities(&self) -> SessionCapabilities {
567 self.capabilities.read().clone()
568 }
569
570 pub fn open_canvases(&self) -> Vec<OpenCanvasInstance> {
573 self.open_canvases.read().clone()
574 }
575
576 pub fn cancellation_token(&self) -> CancellationToken {
603 self.shutdown.child_token()
604 }
605
606 pub fn subscribe(&self) -> crate::subscription::EventSubscription {
661 match &self.resume_bootstrap {
662 Some(bootstrap) => bootstrap.subscribe(&self.event_tx),
663 None => crate::subscription::EventSubscription::new(self.event_tx.subscribe()),
664 }
665 }
666
667 pub fn client(&self) -> &Client {
669 &self.client
670 }
671
672 pub fn rpc(&self) -> crate::generated::rpc::SessionRpc<'_> {
683 crate::generated::rpc::SessionRpc { session: self }
684 }
685
686 pub async fn stop_event_loop(&self) {
699 self.shutdown.cancel();
700 let handle = self.event_loop.lock().take();
701 if let Some(handle) = handle {
702 let _ = handle.await;
703 }
704 if let Some(waiter) = self.idle_waiter.lock().take() {
706 let _ = waiter.tx.send(Err(
707 ErrorKind::Session(SessionErrorKind::EventLoopClosed).into()
708 ));
709 }
710 }
711
712 pub async fn send(&self, opts: impl Into<MessageOptions>) -> Result<String, Error> {
737 if self.idle_waiter.lock().is_some() {
738 return Err(ErrorKind::Session(SessionErrorKind::SendWhileWaiting).into());
739 }
740 self.send_inner(opts.into()).await
741 }
742
743 async fn send_inner(&self, opts: MessageOptions) -> Result<String, Error> {
744 let mut params = serde_json::json!({
745 "sessionId": self.id,
746 "prompt": opts.prompt,
747 });
748 if let Some(source) = opts.source {
749 params["source"] = serde_json::to_value(source)?;
750 }
751 if let Some(m) = opts.mode {
752 params["mode"] = serde_json::to_value(m)?;
753 }
754 if let Some(am) = opts.agent_mode {
755 params["agentMode"] = serde_json::to_value(am)?;
756 }
757 if let Some(mut a) = opts.attachments {
758 ensure_attachment_display_names(&mut a);
759 params["attachments"] = serde_json::to_value(a)?;
760 }
761 if let Some(headers) = opts.request_headers
762 && !headers.is_empty()
763 {
764 params["requestHeaders"] = serde_json::to_value(headers)?;
765 }
766 if let Some(display_prompt) = opts.display_prompt {
767 params["displayPrompt"] = serde_json::to_value(display_prompt)?;
768 }
769 if let Some(schema) = opts.response_schema {
770 params["responseFormat"] = serde_json::json!({
771 "type": "json_schema",
772 "jsonSchema": { "name": "response", "strict": true, "schema": schema }
773 });
774 }
775 let trace_ctx = if opts.traceparent.is_some() || opts.tracestate.is_some() {
776 TraceContext {
777 traceparent: opts.traceparent,
778 tracestate: opts.tracestate,
779 }
780 } else {
781 self.client.resolve_trace_context().await
782 };
783 inject_trace_context(&mut params, &trace_ctx);
784 let rpc_start = Instant::now();
785 let result = self.client.call("session.send", Some(params)).await?;
786 let message_id = result
787 .get("messageId")
788 .and_then(|v| v.as_str())
789 .map(|s| s.to_string())
790 .unwrap_or_default();
791 tracing::debug!(
792 elapsed_ms = rpc_start.elapsed().as_millis(),
793 session_id = %self.id,
794 message_id = %message_id,
795 "Session::send completed successfully"
796 );
797 Ok(message_id)
798 }
799
800 pub async fn send_and_wait(
828 &self,
829 opts: impl Into<MessageOptions>,
830 ) -> Result<Option<SessionEvent>, Error> {
831 let total_start = Instant::now();
832 let opts = opts.into();
833 if opts.response_schema.is_some() {
834 return self.send_and_wait_structured(opts).await.map(Some);
835 }
836 let timeout_duration = opts.wait_timeout.unwrap_or(Duration::from_secs(60));
837 let (tx, rx) = oneshot::channel();
838
839 {
840 let mut guard = self.idle_waiter.lock();
841 if guard.is_some() {
842 return Err(ErrorKind::Session(SessionErrorKind::SendWhileWaiting).into());
843 }
844 *guard = Some(IdleWaiter {
845 tx,
846 last_assistant_message: None,
847 started_at: total_start,
848 first_assistant_message_seen: false,
849 });
850 }
851
852 let _waiter_guard = WaiterGuard {
857 slot: self.idle_waiter.clone(),
858 };
859
860 let result = tokio::time::timeout(timeout_duration, async {
861 self.send_inner(opts).await?;
862 match rx.await {
863 Ok(result) => result,
864 Err(_) => Err(ErrorKind::Session(SessionErrorKind::EventLoopClosed).into()),
865 }
866 })
867 .await;
868
869 match result {
870 Ok(inner) => {
871 tracing::debug!(
872 elapsed_ms = total_start.elapsed().as_millis(),
873 session_id = %self.id,
874 completed_by = if inner.is_ok() { "idle" } else { "error" },
875 "Session::send_and_wait complete"
876 );
877 inner
878 }
879 Err(_) => {
880 tracing::warn!(
881 elapsed_ms = total_start.elapsed().as_millis(),
882 session_id = %self.id,
883 completed_by = "timeout",
884 "Session::send_and_wait failed"
885 );
886 Err(ErrorKind::Session(SessionErrorKind::Timeout(timeout_duration)).into())
887 }
888 }
889 }
890
891 #[cfg(feature = "derive")]
899 pub async fn send_and_wait_typed<T>(&self, opts: impl Into<MessageOptions>) -> Result<T, Error>
900 where
901 T: schemars::JsonSchema + serde::de::DeserializeOwned,
902 {
903 let mut opts = opts.into();
904 if opts.response_schema.is_some()
905 || opts.mode == Some(crate::types::DeliveryMode::Immediate)
906 {
907 return Err(Error::with_message(
908 ErrorKind::InvalidConfig,
909 "typed structured output cannot specify a response schema or immediate delivery",
910 ));
911 }
912 opts.response_schema = Some(crate::tool::schema_for::<T>());
913 let event = self.send_and_wait_structured(opts).await?;
914 let data: crate::session_events::AssistantMessageData = serde_json::from_value(event.data)?;
915 serde_json::from_str::<Option<T>>(&data.content)?.ok_or_else(|| {
916 structured_output_error("structured response was JSON null, not a result")
917 })
918 }
919
920 async fn send_and_wait_structured(&self, opts: MessageOptions) -> Result<SessionEvent, Error> {
921 let duration = opts.wait_timeout.unwrap_or(Duration::from_secs(60));
922 let mut events = crate::subscription::EventSubscription::new(self.event_tx.subscribe());
924 let wait = async {
925 let mut admission = Box::pin(self.send(opts));
926 let mut pending = Vec::new();
927 let message_id = loop {
928 tokio::select! {
929 result = &mut admission => break result?,
930 event = events.recv() => {
931 let event = event.map_err(|err| structured_output_error(err.to_string()))?;
932 if is_structured_output_event(&event) {
933 pending.push(event);
934 }
935 }
936 _ = self.shutdown.cancelled() =>
937 return Err(structured_output_error("session closed before structured output completed")),
938 }
939 };
940 let mut state = StructuredOutputState {
941 message_id,
942 started: false,
943 final_message: None,
944 };
945 for event in pending {
946 if let Some(result) = state.observe(event)? {
947 return Ok(result);
948 }
949 }
950 loop {
951 tokio::select! {
952 event = events.recv() => {
953 let event = event.map_err(|err| structured_output_error(err.to_string()))?;
954 if let Some(result) = state.observe(event)? {
955 return Ok(result);
956 }
957 }
958 _ = self.shutdown.cancelled() =>
959 return Err(structured_output_error("session closed before structured output completed")),
960 }
961 }
962 };
963 tokio::time::timeout(duration, wait)
964 .await
965 .map_err(|_| Error::from(ErrorKind::Session(SessionErrorKind::Timeout(duration))))?
966 }
967
968 pub async fn get_events(&self) -> Result<Vec<SessionEvent>, Error> {
970 let result = self
971 .client
972 .call(
973 "session.getMessages",
974 Some(serde_json::json!({ "sessionId": self.id })),
975 )
976 .await?;
977 let response: GetMessagesResponse = serde_json::from_value(result)?;
978 Ok(response.events)
979 }
980
981 #[deprecated(since = "0.1.0", note = "Use `get_events()` instead")]
983 pub async fn get_messages(&self) -> Result<Vec<SessionEvent>, Error> {
984 self.get_events().await
985 }
986
987 pub async fn abort(&self) -> Result<(), Error> {
995 self.client
996 .call(
997 "session.abort",
998 Some(serde_json::json!({ "sessionId": self.id })),
999 )
1000 .await?;
1001 Ok(())
1002 }
1003
1004 pub async fn set_model(&self, model: &str, opts: Option<SetModelOptions>) -> Result<(), Error> {
1008 let opts = opts.unwrap_or_default();
1009 let auto_tier = opts.auto_tier.clone();
1010 let request = ModelSwitchToRequest {
1011 auto_tier: match &auto_tier {
1012 Some(AutoTierPreference::Tier(tier)) => Some(tier.clone()),
1013 _ => None,
1014 },
1015 compaction_decision: None,
1016 context_tier: opts.context_tier,
1017 defer_if_model_change_queued: None,
1018 model_capabilities: opts.model_capabilities,
1019 model_change_scope: None,
1020 model_id: model.to_string(),
1021 picker_persistence: None,
1022 reasoning_effort: opts.reasoning_effort,
1023 reasoning_summary: opts.reasoning_summary,
1024 repo_scope: None,
1025 require_available: None,
1026 run_compaction_preflight: None,
1027 source: None,
1028 verbosity: None,
1029 };
1030
1031 if matches!(auto_tier, Some(AutoTierPreference::Reset)) {
1032 let mut wire_params = serde_json::to_value(request)?;
1036 wire_params["sessionId"] = serde_json::Value::String(self.id.to_string());
1037 wire_params["autoTier"] = serde_json::Value::Null;
1038 self.client
1039 .call("session.model.switchTo", Some(wire_params))
1040 .await?;
1041 return Ok(());
1042 }
1043
1044 self.rpc().model().switch_to(request).await?;
1045 Ok(())
1046 }
1047
1048 pub async fn set_auto_tier(
1077 &self,
1078 auto_tier: Option<AutoTier>,
1079 ) -> Result<ModelSwitchAutoTierResult, Error> {
1080 self.rpc()
1081 .model()
1082 .switch_auto_tier(ModelSwitchAutoTierRequest {
1083 auto_tier,
1084 source: None,
1085 })
1086 .await
1087 }
1088
1089 pub async fn disconnect(&self) -> Result<(), Error> {
1106 self.client.detach_session(&self.id).await?;
1107 self.external_tools_shutdown.cancel();
1108 self.stop_event_loop().await;
1109 self.github_token_registration.lock().take();
1110 self.client
1111 .unregister_session_owned(&self.id, self.registration_token);
1112 Ok(())
1113 }
1114
1115 #[deprecated(since = "0.1.0", note = "Use `disconnect()` instead")]
1118 pub async fn destroy(&self) -> Result<(), Error> {
1119 self.disconnect().await
1120 }
1121
1122 pub async fn log(
1126 &self,
1127 message: &str,
1128 opts: Option<crate::types::LogOptions>,
1129 ) -> Result<(), Error> {
1130 let opts = opts.unwrap_or_default();
1131 let level = match opts.level {
1132 Some(level) => Some(serde_json::from_value(serde_json::to_value(level)?)?),
1133 None => None,
1134 };
1135 let request = LogRequest {
1136 message: message.to_string(),
1137 level,
1138 ephemeral: opts.ephemeral,
1139 r#type: None,
1140 tip: None,
1141 url: None,
1142 };
1143 self.rpc().log(request).await?;
1144 Ok(())
1145 }
1146
1147 pub fn ui(&self) -> SessionUi<'_> {
1153 SessionUi { session: self }
1154 }
1155
1156 fn assert_elicitation(&self) -> Result<(), Error> {
1158 if self
1159 .capabilities
1160 .read()
1161 .ui
1162 .as_ref()
1163 .and_then(|u| u.elicitation)
1164 != Some(true)
1165 {
1166 return Err(ErrorKind::Session(SessionErrorKind::ElicitationNotSupported).into());
1167 }
1168 Ok(())
1169 }
1170}
1171
1172impl Drop for Session {
1173 fn drop(&mut self) {
1174 self.shutdown.cancel();
1186 self.external_tools_shutdown.cancel();
1187 self.github_token_registration.lock().take();
1188 self.client
1189 .unregister_session_owned(&self.id, self.registration_token);
1190 }
1191}
1192
1193pub struct SessionUi<'a> {
1200 session: &'a Session,
1201}
1202
1203impl<'a> SessionUi<'a> {
1204 pub async fn elicitation(
1212 &self,
1213 message: &str,
1214 schema: Value,
1215 ) -> Result<ElicitationResult, Error> {
1216 self.session.assert_elicitation()?;
1217 let result = self
1218 .session
1219 .client
1220 .call(
1221 "session.ui.elicitation",
1222 Some(serde_json::json!({
1223 "sessionId": self.session.id,
1224 "message": message,
1225 "requestedSchema": schema,
1226 })),
1227 )
1228 .await?;
1229 let elicitation: ElicitationResult = serde_json::from_value(result)?;
1230 Ok(elicitation)
1231 }
1232
1233 pub async fn confirm(&self, message: &str) -> Result<bool, Error> {
1237 self.session.assert_elicitation()?;
1238 let schema = serde_json::json!({
1239 "type": "object",
1240 "properties": {
1241 "confirmed": {
1242 "type": "boolean",
1243 "default": true,
1244 }
1245 },
1246 "required": ["confirmed"]
1247 });
1248 let result = self.elicitation(message, schema).await?;
1249 Ok(result.action == "accept"
1250 && result
1251 .content
1252 .and_then(|c| c.get("confirmed").and_then(|v| v.as_bool()))
1253 == Some(true))
1254 }
1255
1256 pub async fn select(&self, message: &str, options: &[&str]) -> Result<Option<String>, Error> {
1260 self.session.assert_elicitation()?;
1261 let schema = serde_json::json!({
1262 "type": "object",
1263 "properties": {
1264 "selection": {
1265 "type": "string",
1266 "enum": options,
1267 }
1268 },
1269 "required": ["selection"]
1270 });
1271 let result = self.elicitation(message, schema).await?;
1272 if result.action != "accept" {
1273 return Ok(None);
1274 }
1275 let selection = result.content.and_then(|c| {
1276 c.get("selection")
1277 .and_then(|v| v.as_str())
1278 .map(String::from)
1279 });
1280 Ok(selection)
1281 }
1282
1283 pub async fn input(
1288 &self,
1289 message: &str,
1290 options: Option<&UiInputOptions<'_>>,
1291 ) -> Result<Option<String>, Error> {
1292 self.session.assert_elicitation()?;
1293 let mut field = serde_json::json!({ "type": "string" });
1294 if let Some(opts) = options {
1295 if let Some(title) = opts.title {
1296 field["title"] = Value::String(title.to_string());
1297 }
1298 if let Some(desc) = opts.description {
1299 field["description"] = Value::String(desc.to_string());
1300 }
1301 if let Some(min) = opts.min_length {
1302 field["minLength"] = Value::Number(min.into());
1303 }
1304 if let Some(max) = opts.max_length {
1305 field["maxLength"] = Value::Number(max.into());
1306 }
1307 if let Some(fmt) = &opts.format {
1308 field["format"] = Value::String(fmt.as_str().to_string());
1309 }
1310 if let Some(default) = opts.default {
1311 field["default"] = Value::String(default.to_string());
1312 }
1313 }
1314 let schema = serde_json::json!({
1315 "type": "object",
1316 "properties": { "value": field },
1317 "required": ["value"]
1318 });
1319 let result = self.elicitation(message, schema).await?;
1320 if result.action != "accept" {
1321 return Ok(None);
1322 }
1323 let value = result
1324 .content
1325 .and_then(|c| c.get("value").and_then(|v| v.as_str()).map(String::from));
1326 Ok(value)
1327 }
1328}
1329
1330impl Client {
1331 pub fn prepare_session(&self, config: SessionConfig) -> Result<PreparedSession, Error> {
1393 let capacity = resolve_event_buffer_capacity(config.event_buffer_capacity)?;
1394 Ok(PreparedSession::new(
1395 self.clone(),
1396 PreparedKind::Create(Box::new(config)),
1397 capacity,
1398 ))
1399 }
1400
1401 pub fn prepare_resume_session(
1416 &self,
1417 config: ResumeSessionConfig,
1418 ) -> Result<PreparedSession, Error> {
1419 let capacity = resolve_event_buffer_capacity(config.event_buffer_capacity)?;
1420 Ok(PreparedSession::new(
1421 self.clone(),
1422 PreparedKind::Resume(Box::new(config)),
1423 capacity,
1424 ))
1425 }
1426
1427 pub async fn create_session(&self, config: SessionConfig) -> Result<Session, Error> {
1459 self.prepare_session(config)?.start().await
1460 }
1461
1462 pub async fn resume_session(&self, config: ResumeSessionConfig) -> Result<Session, Error> {
1486 self.prepare_resume_session(config)?.start().await
1487 }
1488
1489 async fn start_prepared_create(
1490 &self,
1491 mut config: SessionConfig,
1492 event_tx: tokio::sync::broadcast::Sender<SessionEvent>,
1493 shutdown: CancellationToken,
1494 ) -> Result<Session, Error> {
1495 let total_start = Instant::now();
1496 let caller_session_id = config.session_id.clone();
1501 let use_server_generated_id = config.cloud.is_some() && caller_session_id.is_none();
1502 let local_session_id: Option<SessionId> = if use_server_generated_id {
1503 None
1504 } else {
1505 Some(
1506 caller_session_id
1507 .clone()
1508 .unwrap_or_else(|| SessionId::new(uuid::Uuid::new_v4().to_string())),
1509 )
1510 };
1511 if config.hooks_handler.is_some() && config.hooks.is_none() {
1512 config.hooks = Some(true);
1513 }
1514 if let Some(transforms) = config.system_message_transform.clone() {
1515 inject_transform_sections(&mut config, transforms.as_ref());
1516 }
1517 let mode = self.inner.mode;
1518 if mode == crate::ClientMode::Empty && config.available_tools.is_none() {
1519 return Err(Error::with_message(
1520 ErrorKind::InvalidConfig,
1521 "ClientMode::Empty requires available_tools to be set on the session config. \
1522 Use ToolSet to specify which tools the session may use (e.g. \
1523 ToolSet::new().add_builtin_many(BUILTIN_TOOLS_ISOLATED)).",
1524 ));
1525 }
1526 crate::mode::validate_tool_filter_list(
1527 "available_tools",
1528 config.available_tools.as_deref(),
1529 )?;
1530 crate::mode::validate_tool_filter_list("excluded_tools", config.excluded_tools.as_deref())?;
1531 config.system_message =
1532 crate::mode::system_message_for_mode(mode, config.system_message.take());
1533 config.memory = crate::mode::memory_for_mode(mode, config.memory.take());
1534 config.enable_experimental_mode =
1535 crate::mode::experimental_mode_for_mode(mode, config.enable_experimental_mode);
1536 if mode == crate::ClientMode::Empty {
1537 if config.enable_session_telemetry.is_none() {
1538 config.enable_session_telemetry = Some(false);
1539 }
1540 if config.skip_embedding_retrieval.is_none() {
1541 config.skip_embedding_retrieval = Some(true);
1542 }
1543 if config.enable_on_demand_instruction_discovery.is_none() {
1544 config.enable_on_demand_instruction_discovery = Some(false);
1545 }
1546 if config.enable_file_hooks.is_none() {
1547 config.enable_file_hooks = Some(false);
1548 }
1549 if config.enable_host_git_operations.is_none() {
1550 config.enable_host_git_operations = Some(false);
1551 }
1552 if config.enable_session_store.is_none() {
1553 config.enable_session_store = Some(false);
1554 }
1555 if config.enable_skills.is_none() {
1556 config.enable_skills = Some(false);
1557 }
1558 }
1559 if mode == crate::ClientMode::Empty && config.mcp_oauth_token_storage.is_none() {
1560 config.mcp_oauth_token_storage = Some("in-memory".into());
1561 }
1562 if mode == crate::ClientMode::Empty && config.embedding_cache_storage.is_none() {
1563 config.embedding_cache_storage = Some("in-memory".into());
1564 }
1565 config.custom_agents_local_only =
1566 crate::mode::resolve_custom_agents_local_only(mode, config.custom_agents_local_only);
1567 let opt_skip_custom_instructions = config.skip_custom_instructions;
1568 let opt_custom_agents_local_only = config.custom_agents_local_only;
1569 let opt_coauthor_enabled = config.coauthor_enabled;
1570 let opt_manage_schedule_enabled = config.manage_schedule_enabled;
1571 let opt_included_builtin_skills = config.included_builtin_skills.take();
1572 let (mut wire, mut runtime) = config.into_wire(local_session_id.clone())?;
1573 wire.enable_github_telemetry_forwarding =
1574 self.inner.on_github_telemetry.is_some().then_some(true);
1575
1576 let permission_handler = crate::permission::resolve_handler(
1577 runtime.permission_handler.take(),
1578 runtime.permission_policy.take(),
1579 );
1580 let handlers = SessionHandlers {
1581 permission: permission_handler,
1582 managed_settings_enabled: has_managed_settings(
1583 wire.enable_managed_settings,
1584 wire.managed_settings.as_ref(),
1585 ),
1586 elicitation: runtime.elicitation_handler.take(),
1587 mcp_auth: runtime.mcp_auth_handler.take(),
1588 user_input: runtime.user_input_handler.take(),
1589 exit_plan_mode: runtime.exit_plan_mode_handler.take(),
1590 auto_mode_switch: runtime.auto_mode_switch_handler.take(),
1591 tools: Arc::new(std::mem::take(&mut runtime.tool_handlers)),
1592 };
1593 let hooks = runtime.hooks_handler.take();
1594 let transforms = runtime.system_message_transform.take();
1595 let tools_count = wire.tools.as_ref().map_or(0, Vec::len);
1596 let commands_count = runtime.commands.as_ref().map_or(0, Vec::len);
1597 let has_hooks = hooks.is_some();
1598 let command_handlers = build_command_handler_map(runtime.commands.as_deref());
1599 let canvas_handler = runtime.canvas_handler.take();
1600 let session_fs_provider = runtime.session_fs_provider.take();
1601 let bearer_token_providers = std::mem::take(&mut runtime.bearer_token_providers);
1602 let github_token_registration = runtime
1603 .github_token_provider
1604 .take()
1605 .map(|provider| self.register_github_token_provider(provider));
1606 wire.github_token_provider_registration_id = github_token_registration
1607 .as_ref()
1608 .map(|registration| registration.id().to_string());
1609 let has_mcp_auth_handler = handlers.mcp_auth.is_some();
1610 if self.inner.session_fs_configured && session_fs_provider.is_none() {
1611 return Err(ErrorKind::Session(SessionErrorKind::SessionFsProviderRequired).into());
1612 }
1613 if self.inner.session_fs_sqlite_declared
1614 && let Some(ref provider) = session_fs_provider
1615 && provider.sqlite().is_none()
1616 {
1617 return Err(Error::with_message(
1618 ErrorKind::InvalidConfig,
1619 "SessionFs capabilities declare SQLite support but the provider \
1620 does not implement SessionFsSqliteProvider",
1621 ));
1622 }
1623
1624 let mut params = serde_json::to_value(&wire)?;
1625 let trace_ctx = self.resolve_trace_context().await;
1626 inject_trace_context(&mut params, &trace_ctx);
1627
1628 let setup_start = Instant::now();
1629 let capabilities = Arc::new(parking_lot::RwLock::new(SessionCapabilities::default()));
1630 let idle_waiter = Arc::new(ParkingLotMutex::new(None));
1631 let open_canvases = Arc::new(parking_lot::RwLock::new(Vec::new()));
1632 let external_tools_shutdown = self.inner.rpc.connection_closed_token();
1633
1634 let spawn_loop: EventLoopSpawner = {
1635 let client = self.clone();
1636 let idle_waiter = idle_waiter.clone();
1637 let capabilities = capabilities.clone();
1638 let open_canvases = open_canvases.clone();
1639 let event_tx = event_tx.clone();
1640 let shutdown = shutdown.clone();
1641 let external_tools_shutdown = external_tools_shutdown.clone();
1642 Box::new(
1643 move |session_id: SessionId,
1644 channels: crate::router::SessionChannels,
1645 startup_tasks: Option<Arc<StartupTasks>>| {
1646 spawn_event_loop(
1647 session_id,
1648 client,
1649 handlers,
1650 hooks,
1651 transforms,
1652 command_handlers,
1653 canvas_handler,
1654 session_fs_provider,
1655 bearer_token_providers,
1656 channels,
1657 idle_waiter,
1658 capabilities,
1659 open_canvases,
1660 event_tx,
1661 None,
1662 shutdown,
1663 external_tools_shutdown,
1664 startup_tasks,
1665 )
1666 },
1667 )
1668 };
1669
1670 let (inline_callback, mut pending_registration, event_loop, inline_stash) =
1673 if let Some(ref sid) = local_session_id {
1674 let registration = self.register_session(sid);
1675 let token = registration.token;
1676 let guard = PendingSessionRegistration::new(
1677 self.clone(),
1678 sid.clone(),
1679 token,
1680 shutdown.clone(),
1681 external_tools_shutdown.clone(),
1682 );
1683 let event_loop = spawn_loop(
1684 sid.clone(),
1685 registration.channels,
1686 Some(guard.startup_tasks.clone()),
1687 );
1688 (
1689 None,
1690 guard,
1691 CreateEventLoop::Running { event_loop, token },
1692 None,
1693 )
1694 } else {
1695 let inline_stash = Arc::new(ParkingLotMutex::new(None));
1696 let (guard, callback) = PendingSessionRegistration::deferred(
1697 self.clone(),
1698 inline_stash.clone(),
1699 shutdown.clone(),
1700 external_tools_shutdown.clone(),
1701 );
1702 (
1703 Some(callback),
1704 guard,
1705 CreateEventLoop::Deferred(spawn_loop),
1706 Some(inline_stash),
1707 )
1708 };
1709
1710 let rpc_start = Instant::now();
1711 let result = match self
1712 .call_with_inline_callback("session.create", Some(params), inline_callback)
1713 .await
1714 {
1715 Ok(result) => result,
1716 Err(error) => {
1717 event_loop.cleanup(pending_registration).await;
1718 return Err(error);
1719 }
1720 };
1721 tracing::debug!(
1722 elapsed_ms = rpc_start.elapsed().as_millis(),
1723 "Client::create_session session creation request completed successfully"
1724 );
1725 let create_result: CreateSessionResult = match serde_json::from_value(result) {
1726 Ok(result) => result,
1727 Err(error) => {
1728 event_loop.cleanup(pending_registration).await;
1729 return Err(error.into());
1730 }
1731 };
1732
1733 if let Some(ref requested) = local_session_id
1734 && create_result.session_id != *requested
1735 {
1736 event_loop.cleanup(pending_registration).await;
1737 return Err(ErrorKind::Session(SessionErrorKind::SessionIdMismatch {
1738 requested: requested.clone(),
1739 returned: create_result.session_id.clone(),
1740 })
1741 .into());
1742 }
1743
1744 let (session_id, event_loop, registration_token) = match event_loop {
1745 CreateEventLoop::Running { event_loop, token } => {
1746 (create_result.session_id.clone(), event_loop, token)
1747 }
1748 CreateEventLoop::Deferred(spawn_loop) => {
1749 let (session_id, registration) = inline_stash
1750 .expect("deferred create has a response stash")
1751 .lock()
1752 .take()
1753 .expect("session registration must have populated stash on success");
1754 pending_registration.resolve_to(session_id.clone(), registration.token);
1755 let event_loop = spawn_loop(
1756 session_id.clone(),
1757 registration.channels,
1758 Some(pending_registration.startup_tasks.clone()),
1759 );
1760 (session_id, event_loop, registration.token)
1761 }
1762 };
1763 tracing::debug!(
1764 elapsed_ms = setup_start.elapsed().as_millis(),
1765 session_id = %session_id,
1766 tools_count,
1767 commands_count,
1768 has_hooks,
1769 "Client::create_session local setup complete"
1770 );
1771 *capabilities.write() = create_result.capabilities.unwrap_or_default();
1772 if has_mcp_auth_handler
1773 && let Err(error) = register_mcp_auth_interest(self, &session_id).await
1774 {
1775 pending_registration.cleanup(event_loop).await;
1776 return Err(error);
1777 }
1778
1779 tracing::debug!(
1780 elapsed_ms = total_start.elapsed().as_millis(),
1781 session_id = %session_id,
1782 "Client::create_session complete"
1783 );
1784 let session = Session {
1785 id: session_id,
1786 cwd: self.cwd().clone(),
1787 workspace_path: create_result.workspace_path,
1788 remote_url: create_result.remote_url,
1789 client: self.clone(),
1790 event_loop: ParkingLotMutex::new(Some(event_loop)),
1791 shutdown,
1792 external_tools_shutdown,
1793 idle_waiter,
1794 capabilities,
1795 open_canvases,
1796 event_tx,
1797 resume_bootstrap: None,
1798 github_token_registration: ParkingLotMutex::new(github_token_registration),
1799 registration_token,
1800 };
1801 let patch = build_mode_post_create_patch(
1802 mode,
1803 opt_skip_custom_instructions,
1804 opt_custom_agents_local_only,
1805 opt_coauthor_enabled,
1806 opt_manage_schedule_enabled,
1807 opt_included_builtin_skills,
1808 );
1809 finish_session_setup(session, pending_registration, patch).await
1810 }
1811
1812 async fn start_prepared_resume(
1823 &self,
1824 mut config: ResumeSessionConfig,
1825 event_tx: tokio::sync::broadcast::Sender<SessionEvent>,
1826 shutdown: CancellationToken,
1827 ) -> Result<Session, Error> {
1828 let total_start = Instant::now();
1829 let session_id = config.session_id.clone();
1830 if config.hooks_handler.is_some() && config.hooks.is_none() {
1831 config.hooks = Some(true);
1832 }
1833 if let Some(transforms) = config.system_message_transform.clone() {
1834 inject_transform_sections_resume(&mut config, transforms.as_ref());
1835 }
1836 let mode = self.inner.mode;
1837 if mode == crate::ClientMode::Empty && config.available_tools.is_none() {
1838 return Err(Error::with_message(
1839 ErrorKind::InvalidConfig,
1840 "ClientMode::Empty requires available_tools to be set on the session config. \
1841 Use ToolSet to specify which tools the session may use (e.g. \
1842 ToolSet::new().add_builtin_many(BUILTIN_TOOLS_ISOLATED)).",
1843 ));
1844 }
1845 crate::mode::validate_tool_filter_list(
1846 "available_tools",
1847 config.available_tools.as_deref(),
1848 )?;
1849 crate::mode::validate_tool_filter_list("excluded_tools", config.excluded_tools.as_deref())?;
1850 config.system_message =
1851 crate::mode::system_message_for_mode(mode, config.system_message.take());
1852 config.memory = crate::mode::memory_for_mode(mode, config.memory.take());
1853 config.enable_experimental_mode =
1854 crate::mode::experimental_mode_for_mode(mode, config.enable_experimental_mode);
1855 if mode == crate::ClientMode::Empty {
1856 if config.enable_session_telemetry.is_none() {
1857 config.enable_session_telemetry = Some(false);
1858 }
1859 if config.skip_embedding_retrieval.is_none() {
1860 config.skip_embedding_retrieval = Some(true);
1861 }
1862 if config.enable_on_demand_instruction_discovery.is_none() {
1863 config.enable_on_demand_instruction_discovery = Some(false);
1864 }
1865 if config.enable_file_hooks.is_none() {
1866 config.enable_file_hooks = Some(false);
1867 }
1868 if config.enable_host_git_operations.is_none() {
1869 config.enable_host_git_operations = Some(false);
1870 }
1871 if config.enable_session_store.is_none() {
1872 config.enable_session_store = Some(false);
1873 }
1874 if config.enable_skills.is_none() {
1875 config.enable_skills = Some(false);
1876 }
1877 }
1878 if mode == crate::ClientMode::Empty && config.mcp_oauth_token_storage.is_none() {
1879 config.mcp_oauth_token_storage = Some("in-memory".into());
1880 }
1881 if mode == crate::ClientMode::Empty && config.embedding_cache_storage.is_none() {
1882 config.embedding_cache_storage = Some("in-memory".into());
1883 }
1884 config.custom_agents_local_only =
1885 crate::mode::resolve_custom_agents_local_only(mode, config.custom_agents_local_only);
1886 let opt_skip_custom_instructions = config.skip_custom_instructions;
1887 let opt_custom_agents_local_only = config.custom_agents_local_only;
1888 let opt_coauthor_enabled = config.coauthor_enabled;
1889 let opt_manage_schedule_enabled = config.manage_schedule_enabled;
1890 let opt_included_builtin_skills = config.included_builtin_skills.take();
1891 let (mut wire, mut runtime) = config.into_wire()?;
1892 wire.enable_github_telemetry_forwarding =
1893 self.inner.on_github_telemetry.is_some().then_some(true);
1894
1895 let permission_handler = crate::permission::resolve_handler(
1896 runtime.permission_handler.take(),
1897 runtime.permission_policy.take(),
1898 );
1899 let handlers = SessionHandlers {
1900 permission: permission_handler,
1901 managed_settings_enabled: has_managed_settings(
1902 wire.enable_managed_settings,
1903 wire.managed_settings.as_ref(),
1904 ),
1905 elicitation: runtime.elicitation_handler.take(),
1906 mcp_auth: runtime.mcp_auth_handler.take(),
1907 user_input: runtime.user_input_handler.take(),
1908 exit_plan_mode: runtime.exit_plan_mode_handler.take(),
1909 auto_mode_switch: runtime.auto_mode_switch_handler.take(),
1910 tools: Arc::new(std::mem::take(&mut runtime.tool_handlers)),
1911 };
1912 let hooks = runtime.hooks_handler.take();
1913 let transforms = runtime.system_message_transform.take();
1914 let tools_count = wire.tools.as_ref().map_or(0, Vec::len);
1915 let commands_count = runtime.commands.as_ref().map_or(0, Vec::len);
1916 let has_hooks = hooks.is_some();
1917 let command_handlers = build_command_handler_map(runtime.commands.as_deref());
1918 let canvas_handler = runtime.canvas_handler.take();
1919 let session_fs_provider = runtime.session_fs_provider.take();
1920 let bearer_token_providers = std::mem::take(&mut runtime.bearer_token_providers);
1921 let github_token_registration = runtime
1922 .github_token_provider
1923 .take()
1924 .map(|provider| self.register_github_token_provider(provider));
1925 wire.github_token_provider_registration_id = github_token_registration
1926 .as_ref()
1927 .map(|registration| registration.id().to_string());
1928 let has_mcp_auth_handler = handlers.mcp_auth.is_some();
1929 if self.inner.session_fs_configured && session_fs_provider.is_none() {
1930 return Err(ErrorKind::Session(SessionErrorKind::SessionFsProviderRequired).into());
1931 }
1932 if self.inner.session_fs_sqlite_declared
1933 && let Some(ref provider) = session_fs_provider
1934 && provider.sqlite().is_none()
1935 {
1936 return Err(Error::with_message(
1937 ErrorKind::InvalidConfig,
1938 "SessionFs capabilities declare SQLite support but the provider \
1939 does not implement SessionFsSqliteProvider",
1940 ));
1941 }
1942
1943 let mut params = serde_json::to_value(&wire)?;
1944 let trace_ctx = self.resolve_trace_context().await;
1945 inject_trace_context(&mut params, &trace_ctx);
1946
1947 let capabilities = Arc::new(parking_lot::RwLock::new(SessionCapabilities::default()));
1948 let setup_start = Instant::now();
1949 let resume_bootstrap = (event_tx.receiver_count() == 0)
1951 .then(|| crate::subscription::ResumeBootstrap::new(&event_tx));
1952 let registration = self.register_session(&session_id);
1953 let registration_token = registration.token;
1954 let channels = registration.channels;
1955 let idle_waiter = Arc::new(ParkingLotMutex::new(None));
1956 let open_canvases = Arc::new(parking_lot::RwLock::new(Vec::new()));
1957 let external_tools_shutdown = self.inner.rpc.connection_closed_token();
1958 let registration = PendingSessionRegistration::new(
1959 self.clone(),
1960 session_id.clone(),
1961 registration_token,
1962 shutdown.clone(),
1963 external_tools_shutdown.clone(),
1964 );
1965 let event_loop = spawn_event_loop(
1966 session_id.clone(),
1967 self.clone(),
1968 handlers,
1969 hooks,
1970 transforms,
1971 command_handlers,
1972 canvas_handler,
1973 session_fs_provider,
1974 bearer_token_providers,
1975 channels,
1976 idle_waiter.clone(),
1977 capabilities.clone(),
1978 open_canvases.clone(),
1979 event_tx.clone(),
1980 resume_bootstrap.clone(),
1981 shutdown.clone(),
1982 external_tools_shutdown.clone(),
1983 Some(registration.startup_tasks.clone()),
1984 );
1985 tracing::debug!(
1986 elapsed_ms = setup_start.elapsed().as_millis(),
1987 session_id = %session_id,
1988 tools_count,
1989 commands_count,
1990 has_hooks,
1991 "Client::resume_session local setup complete"
1992 );
1993
1994 let rpc_start = Instant::now();
1995 let result = match self.call("session.resume", Some(params)).await {
1996 Ok(result) => result,
1997 Err(error) => {
1998 registration.cleanup(event_loop).await;
1999 return Err(error);
2000 }
2001 };
2002 tracing::debug!(
2003 elapsed_ms = rpc_start.elapsed().as_millis(),
2004 session_id = %session_id,
2005 "Client::resume_session session resume request completed successfully"
2006 );
2007
2008 let resume_result: ResumeSessionResult = match serde_json::from_value(result) {
2009 Ok(result) => result,
2010 Err(error) => {
2011 registration.cleanup(event_loop).await;
2012 return Err(error.into());
2013 }
2014 };
2015 let cli_session_id = resume_result
2016 .session_id
2017 .clone()
2018 .unwrap_or_else(|| session_id.clone());
2019 if cli_session_id != session_id {
2020 registration.cleanup(event_loop).await;
2021 return Err(ErrorKind::Session(SessionErrorKind::SessionIdMismatch {
2022 requested: session_id,
2023 returned: cli_session_id,
2024 })
2025 .into());
2026 }
2027 if has_mcp_auth_handler
2028 && let Err(error) = register_mcp_auth_interest(self, &session_id).await
2029 {
2030 registration.cleanup(event_loop).await;
2031 return Err(error);
2032 }
2033 let skills_reload_start = Instant::now();
2035 if let Err(e) = self
2036 .call(
2037 "session.skills.reload",
2038 Some(serde_json::json!({ "sessionId": session_id })),
2039 )
2040 .await
2041 {
2042 warn!(
2043 elapsed_ms = skills_reload_start.elapsed().as_millis(),
2044 session_id = %session_id,
2045 error = %e,
2046 "Client::resume_session skills reload request failed"
2047 );
2048 } else {
2049 tracing::debug!(
2050 elapsed_ms = skills_reload_start.elapsed().as_millis(),
2051 session_id = %session_id,
2052 "Client::resume_session skills reload request completed successfully"
2053 );
2054 }
2055
2056 *capabilities.write() = resume_result.capabilities.unwrap_or_default();
2057 {
2062 let mut snapshots = open_canvases.write();
2063 for snapshot in resume_result.open_canvases.unwrap_or_default() {
2064 upsert_open_canvas_snapshot(&mut snapshots, snapshot);
2065 }
2066 }
2067
2068 tracing::debug!(
2069 elapsed_ms = total_start.elapsed().as_millis(),
2070 session_id = %session_id,
2071 "Client::resume_session complete"
2072 );
2073 let session = Session {
2074 id: session_id,
2075 cwd: self.cwd().clone(),
2076 workspace_path: resume_result.workspace_path,
2077 remote_url: resume_result.remote_url,
2078 client: self.clone(),
2079 event_loop: ParkingLotMutex::new(Some(event_loop)),
2080 shutdown,
2081 external_tools_shutdown,
2082 idle_waiter,
2083 capabilities,
2084 open_canvases,
2085 event_tx,
2086 resume_bootstrap,
2087 github_token_registration: ParkingLotMutex::new(github_token_registration),
2088 registration_token,
2089 };
2090 let patch = build_mode_post_create_patch(
2091 mode,
2092 opt_skip_custom_instructions,
2093 opt_custom_agents_local_only,
2094 opt_coauthor_enabled,
2095 opt_manage_schedule_enabled,
2096 opt_included_builtin_skills,
2097 );
2098 finish_session_setup(session, registration, patch).await
2099 }
2100}
2101
2102#[must_use = "a PreparedSession does nothing until started"]
2166pub struct PreparedSession {
2167 client: Client,
2168 kind: PreparedKind,
2169 event_tx: tokio::sync::broadcast::Sender<SessionEvent>,
2170 shutdown: CancellationToken,
2171}
2172
2173enum PreparedKind {
2176 Create(Box<SessionConfig>),
2177 Resume(Box<ResumeSessionConfig>),
2178}
2179
2180impl PreparedSession {
2181 fn new(client: Client, kind: PreparedKind, event_buffer_capacity: usize) -> Self {
2182 let (event_tx, _) = tokio::sync::broadcast::channel(event_buffer_capacity);
2183 Self {
2184 client,
2185 kind,
2186 event_tx,
2187 shutdown: CancellationToken::new(),
2188 }
2189 }
2190
2191 pub fn subscribe(&self) -> crate::subscription::EventSubscription {
2209 crate::subscription::EventSubscription::new(self.event_tx.subscribe())
2210 }
2211
2212 pub async fn start(self) -> Result<Session, Error> {
2230 let Self {
2231 client,
2232 kind,
2233 event_tx,
2234 shutdown,
2235 } = self;
2236 match kind {
2237 PreparedKind::Create(config) => {
2238 client
2239 .start_prepared_create(*config, event_tx, shutdown)
2240 .await
2241 }
2242 PreparedKind::Resume(config) => {
2243 client
2244 .start_prepared_resume(*config, event_tx, shutdown)
2245 .await
2246 }
2247 }
2248 }
2249}
2250
2251type CommandHandlerMap = HashMap<String, Arc<dyn CommandHandler>>;
2252
2253async fn finish_session_setup(
2254 session: Session,
2255 mut registration: PendingSessionRegistration,
2256 patch: Option<crate::generated::api_types::SessionUpdateOptionsParams>,
2257) -> Result<Session, Error> {
2258 if let Some(patch) = patch
2259 && let Err(error) = session.rpc().options().update(patch).await
2260 {
2261 registration.cancel();
2262 let _ = session.disconnect().await;
2263 return Err(error);
2264 }
2265 if let Some(provider) = session.github_token_registration.lock().as_ref() {
2266 provider.claim(session.id.clone());
2267 } else {
2268 session.client.retire_github_token_provider(&session.id);
2269 }
2270 registration.disarm();
2271 Ok(session)
2272}
2273
2274fn build_mode_post_create_patch(
2284 mode: crate::ClientMode,
2285 opt_skip_custom_instructions: Option<bool>,
2286 opt_custom_agents_local_only: Option<bool>,
2287 opt_coauthor_enabled: Option<bool>,
2288 opt_manage_schedule_enabled: Option<bool>,
2289 opt_included_builtin_skills: Option<Vec<String>>,
2290) -> Option<crate::generated::api_types::SessionUpdateOptionsParams> {
2291 use crate::generated::api_types::SessionUpdateOptionsParams;
2292 let mut patch = SessionUpdateOptionsParams::default();
2293 let should_send = if mode == crate::ClientMode::Empty {
2294 patch.skip_custom_instructions = Some(opt_skip_custom_instructions.unwrap_or(true));
2295 patch.custom_agents_local_only = Some(opt_custom_agents_local_only.unwrap_or(true));
2296 patch.coauthor_enabled = Some(opt_coauthor_enabled.unwrap_or(false));
2297 patch.manage_schedule_enabled = Some(opt_manage_schedule_enabled.unwrap_or(false));
2298 patch.installed_plugins = Some(Vec::new());
2299 patch.included_builtin_skills = Some(opt_included_builtin_skills.unwrap_or_default());
2300 true
2301 } else {
2302 let mut any = false;
2303 if let Some(v) = opt_skip_custom_instructions {
2304 patch.skip_custom_instructions = Some(v);
2305 any = true;
2306 }
2307 if let Some(v) = opt_custom_agents_local_only {
2308 patch.custom_agents_local_only = Some(v);
2309 any = true;
2310 }
2311 if let Some(v) = opt_coauthor_enabled {
2312 patch.coauthor_enabled = Some(v);
2313 any = true;
2314 }
2315 if let Some(v) = opt_manage_schedule_enabled {
2316 patch.manage_schedule_enabled = Some(v);
2317 any = true;
2318 }
2319 if let Some(v) = opt_included_builtin_skills {
2320 patch.included_builtin_skills = Some(v);
2321 any = true;
2322 }
2323 any
2324 };
2325 if !should_send {
2326 return None;
2327 }
2328 Some(patch)
2329}
2330
2331fn build_command_handler_map(commands: Option<&[CommandDefinition]>) -> Arc<CommandHandlerMap> {
2332 let map = match commands {
2333 Some(commands) => commands
2334 .iter()
2335 .filter(|cmd| !cmd.name.is_empty())
2336 .map(|cmd| (cmd.name.clone(), cmd.handler.clone()))
2337 .collect(),
2338 None => HashMap::new(),
2339 };
2340 Arc::new(map)
2341}
2342
2343fn upsert_open_canvas_snapshot(
2344 snapshots: &mut Vec<OpenCanvasInstance>,
2345 snapshot: OpenCanvasInstance,
2346) {
2347 if let Some(existing) = snapshots
2348 .iter_mut()
2349 .find(|open| open.instance_id == snapshot.instance_id)
2350 {
2351 *existing = snapshot;
2352 } else {
2353 snapshots.push(snapshot);
2354 }
2355}
2356
2357fn remove_open_canvas_snapshot(snapshots: &mut Vec<OpenCanvasInstance>, instance_id: &str) {
2358 snapshots.retain(|open| open.instance_id != instance_id);
2359}
2360
2361#[allow(clippy::too_many_arguments)]
2362fn spawn_event_loop(
2363 session_id: SessionId,
2364 client: Client,
2365 handlers: SessionHandlers,
2366 hooks: Option<Arc<dyn SessionHooks>>,
2367 transforms: Option<Arc<dyn SystemMessageTransform>>,
2368 command_handlers: Arc<CommandHandlerMap>,
2369 canvas_handler: Option<Arc<dyn CanvasHandler>>,
2370 session_fs_provider: Option<Arc<dyn SessionFsProvider>>,
2371 bearer_token_providers: HashMap<String, Arc<dyn BearerTokenProvider>>,
2372 channels: crate::router::SessionChannels,
2373 idle_waiter: Arc<ParkingLotMutex<Option<IdleWaiter>>>,
2374 capabilities: Arc<parking_lot::RwLock<SessionCapabilities>>,
2375 open_canvases: Arc<parking_lot::RwLock<Vec<OpenCanvasInstance>>>,
2376 event_tx: tokio::sync::broadcast::Sender<SessionEvent>,
2377 resume_bootstrap: Option<Arc<crate::subscription::ResumeBootstrap>>,
2378 shutdown: CancellationToken,
2379 external_tools_shutdown: CancellationToken,
2380 startup_tasks: Option<Arc<StartupTasks>>,
2381) -> JoinHandle<()> {
2382 let crate::router::SessionChannels {
2383 mut notifications,
2384 mut requests,
2385 } = channels;
2386 let pending_external_tools: PendingExternalTools =
2387 Arc::new(ParkingLotMutex::new(HashMap::new()));
2388 let bootstrap_cleanup = resume_bootstrap
2391 .as_ref()
2392 .map(|bootstrap| bootstrap.cleanup_guard());
2393
2394 let span = tracing::error_span!("session_event_loop", session_id = %session_id);
2395 tokio::spawn(
2396 async move {
2397 let _bootstrap_cleanup = bootstrap_cleanup;
2398 loop {
2399 tokio::select! {
2423 _ = shutdown.cancelled() => break,
2424 Some(notification) = notifications.recv() => {
2425 handle_notification(
2426 &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(),
2427 ).await;
2428 }
2429 Some(request) = requests.recv() => {
2430 let span = tracing::error_span!("session_request_handler", session_id = %session_id);
2434 let session_id = session_id.clone();
2435 let client = client.clone();
2436 let handlers = handlers.clone();
2437 let hooks = hooks.clone();
2438 let transforms = transforms.clone();
2439 let canvas_handler = canvas_handler.clone();
2440 let session_fs_provider = session_fs_provider.clone();
2441 let bearer_token_providers = bearer_token_providers.clone();
2442 let request_id = request.id;
2443 let method = request.method.clone();
2444 let _ = spawn_startup_tracked(
2445 async move {
2446 let ctx = RequestDispatchContext {
2447 client: &client,
2448 handlers: &handlers,
2449 hooks: hooks.as_deref(),
2450 transforms: transforms.as_deref(),
2451 canvas_handler: canvas_handler.as_ref(),
2452 session_fs_provider: session_fs_provider.as_ref(),
2453 bearer_token_providers: &bearer_token_providers,
2454 };
2455 let dispatch = handle_request(&session_id, ctx, request);
2456 if AssertUnwindSafe(dispatch).catch_unwind().await.is_err() {
2457 error!(method = %method, "request handler panicked");
2460 let _ = send_error_response(
2461 &client,
2462 request_id,
2463 error_codes::INTERNAL_ERROR,
2464 "request handler panicked",
2465 )
2466 .await;
2467 }
2468 }
2469 .instrument(span),
2470 startup_tasks.as_ref(),
2471 );
2472 }
2473 else => break,
2474 }
2475 }
2476 if let Some(waiter) = idle_waiter.lock().take() {
2479 let _ = waiter
2480 .tx
2481 .send(Err(ErrorKind::Session(SessionErrorKind::EventLoopClosed).into()));
2482 }
2483 }
2484 .instrument(span),
2485 )
2486}
2487
2488fn extract_request_id(data: &Value) -> Option<RequestId> {
2489 data.get("requestId")
2490 .and_then(|v| v.as_str())
2491 .filter(|s| !s.is_empty())
2492 .map(RequestId::new)
2493}
2494
2495fn permission_request_data(
2496 event_data: &Value,
2497 managed_settings_enabled: bool,
2498) -> PermissionRequestData {
2499 let request_data = event_data
2500 .get("permissionRequest")
2501 .cloned()
2502 .unwrap_or_else(|| event_data.clone());
2503 let managed_approval_required = match request_data.get("managedApprovalRequired") {
2504 None => None,
2505 Some(Value::Bool(value)) => Some(*value),
2506 Some(_) => Some(true),
2507 };
2508 match serde_json::from_value::<PermissionRequestData>(request_data) {
2509 Ok(mut data) => {
2510 data.extra = event_data.clone();
2511 data.managed_settings_enabled = managed_settings_enabled;
2512 data
2513 }
2514 Err(_) => PermissionRequestData {
2515 kind: None,
2516 tool_call_id: None,
2517 managed_approval_required,
2518 managed_settings_enabled,
2519 extra: event_data.clone(),
2520 },
2521 }
2522}
2523
2524fn permission_response_params(
2532 session_id: &SessionId,
2533 request_id: &RequestId,
2534 result: &PermissionResult,
2535) -> Option<Value> {
2536 let (decision, decision_context) = match result {
2537 PermissionResult::Decision { decision, context } => (decision, context.clone()),
2538 PermissionResult::NoResult => return None,
2539 };
2540 let mut params = serde_json::to_value(PermissionDecisionRequest {
2541 decision_context,
2542 request_id: request_id.clone(),
2543 result: decision.clone(),
2544 })
2545 .expect("serializing permission response should succeed");
2546 params["sessionId"] =
2547 serde_json::to_value(session_id).expect("serializing session ID should succeed");
2548 Some(params)
2549}
2550
2551async fn register_mcp_auth_interest(client: &Client, session_id: &SessionId) -> Result<(), Error> {
2552 let mut params = serde_json::to_value(RegisterEventInterestParams {
2553 event_type: "mcp.oauth_required".to_string(),
2554 })?;
2555 params["sessionId"] = Value::String(session_id.to_string());
2556 client
2557 .call(rpc_methods::SESSION_EVENTLOG_REGISTERINTEREST, Some(params))
2558 .await?;
2559 Ok(())
2560}
2561
2562fn tool_failure_result(message: impl Into<String>) -> ToolResult {
2563 let message = message.into();
2564 ToolResult::Expanded(ToolResultExpanded {
2565 text_result_for_llm: message.clone(),
2566 result_type: "failure".to_string(),
2567 binary_results_for_llm: None,
2568 session_log: None,
2569 error: Some(message),
2570 tool_telemetry: None,
2571 tool_references: None,
2572 })
2573}
2574
2575fn is_autopilot_continuation_idle(event: &SessionEvent) -> bool {
2576 event
2577 .typed_data::<SessionIdleData>()
2578 .is_some_and(|data| data.mode == Some(SessionMode::Autopilot))
2579}
2580
2581#[allow(clippy::too_many_arguments)]
2583async fn handle_notification(
2584 session_id: &SessionId,
2585 client: &Client,
2586 handlers: &SessionHandlers,
2587 command_handlers: &Arc<CommandHandlerMap>,
2588 notification: SessionEventNotification,
2589 idle_waiter: &Arc<ParkingLotMutex<Option<IdleWaiter>>>,
2590 capabilities: &Arc<parking_lot::RwLock<SessionCapabilities>>,
2591 open_canvases: &Arc<parking_lot::RwLock<Vec<OpenCanvasInstance>>>,
2592 event_tx: &tokio::sync::broadcast::Sender<SessionEvent>,
2593 resume_bootstrap: Option<&Arc<crate::subscription::ResumeBootstrap>>,
2594 shutdown: &CancellationToken,
2595 external_tools_shutdown: &CancellationToken,
2596 pending_external_tools: &PendingExternalTools,
2597 startup_tasks: Option<&Arc<StartupTasks>>,
2598) {
2599 let dispatch_start = Instant::now();
2600 let event = ¬ification.event;
2601 let event_type = event.parsed_type();
2602 if event_type == SessionEventType::PermissionRequested {
2603 tracing::debug!(
2604 session_id = %session_id,
2605 event_type = %event.event_type,
2606 "Session::handle_notification permission request received"
2607 );
2608 }
2609
2610 let mut completion = None;
2613 match event_type {
2614 SessionEventType::AssistantMessage
2615 | SessionEventType::SessionIdle
2616 | SessionEventType::SessionError
2617 if is_root_agent_event(event) =>
2618 {
2619 let mut guard = idle_waiter.lock();
2620 if let Some(waiter) = guard.as_mut() {
2621 match event_type {
2622 SessionEventType::AssistantMessage => {
2623 if !waiter.first_assistant_message_seen {
2624 waiter.first_assistant_message_seen = true;
2625 tracing::debug!(
2626 elapsed_ms = waiter.started_at.elapsed().as_millis(),
2627 session_id = %session_id,
2628 "Session::send_and_wait first assistant message"
2629 );
2630 }
2631 waiter.last_assistant_message = Some(event.clone());
2632 }
2633 SessionEventType::SessionIdle if is_autopilot_continuation_idle(event) => {}
2634 SessionEventType::SessionIdle | SessionEventType::SessionError => {
2635 if let Some(waiter) = guard.take() {
2636 if event_type == SessionEventType::SessionIdle {
2637 tracing::debug!(
2638 elapsed_ms = waiter.started_at.elapsed().as_millis(),
2639 session_id = %session_id,
2640 "Session::send_and_wait idle received"
2641 );
2642 completion = Some((waiter.tx, Ok(waiter.last_assistant_message)));
2643 } else {
2644 let error_msg = event
2645 .typed_data::<SessionErrorData>()
2646 .map(|d| d.message)
2647 .or_else(|| {
2648 event
2649 .data
2650 .get("message")
2651 .and_then(|v| v.as_str())
2652 .map(|s| s.to_string())
2653 })
2654 .unwrap_or_else(|| "session error".to_string());
2655 completion = Some((
2656 waiter.tx,
2657 Err(Error::with_message(
2658 ErrorKind::Session(SessionErrorKind::AgentError),
2659 error_msg,
2660 )),
2661 ));
2662 }
2663 }
2664 }
2665 _ => {}
2666 }
2667 }
2668 }
2669 _ => {}
2670 }
2671
2672 if event_type == SessionEventType::CapabilitiesChanged {
2676 match serde_json::from_value::<SessionCapabilities>(notification.event.data.clone()) {
2677 Ok(changed) => *capabilities.write() = changed,
2678 Err(e) => warn!(error = %e, "failed to deserialize capabilities.changed payload"),
2679 }
2680 }
2681 if event_type == SessionEventType::SessionCanvasOpened {
2682 match serde_json::from_value::<OpenCanvasInstance>(notification.event.data.clone()) {
2683 Ok(open_canvas) => {
2684 upsert_open_canvas_snapshot(&mut open_canvases.write(), open_canvas);
2685 }
2686 Err(e) => warn!(error = %e, "failed to deserialize session.canvas.opened payload"),
2687 }
2688 }
2689 if event_type == SessionEventType::SessionCanvasClosed {
2690 match serde_json::from_value::<SessionCanvasClosedData>(notification.event.data.clone()) {
2691 Ok(closed) => {
2692 if closed.instance_id.is_empty() {
2693 warn!("failed to deserialize session.canvas.closed payload");
2694 } else {
2695 remove_open_canvas_snapshot(&mut open_canvases.write(), &closed.instance_id);
2696 }
2697 }
2698 Err(e) => warn!(error = %e, "failed to deserialize session.canvas.closed payload"),
2699 }
2700 }
2701
2702 if let Some(bootstrap) = resume_bootstrap {
2703 bootstrap.publish(event_tx, event.clone());
2704 } else {
2705 let _ = event_tx.send(event.clone());
2707 }
2708 if let Some((tx, result)) = completion {
2709 let _ = tx.send(result);
2710 }
2711
2712 tracing::debug!(
2713 elapsed_ms = dispatch_start.elapsed().as_millis(),
2714 session_id = %session_id,
2715 event_type = %notification.event.event_type,
2716 "Session::handle_notification dispatch"
2717 );
2718
2719 match event_type {
2722 SessionEventType::ExternalToolCompleted => {
2723 if let Some(request_id) = extract_request_id(¬ification.event.data)
2724 && let Some(token) = pending_external_tools.lock().remove(&request_id)
2725 {
2726 token.cancel();
2727 }
2728 }
2729 SessionEventType::PermissionRequested => {
2730 let Some(request_id) = extract_request_id(¬ification.event.data) else {
2731 return;
2732 };
2733 if notification
2737 .event
2738 .data
2739 .get("resolvedByHook")
2740 .and_then(|v| v.as_bool())
2741 .unwrap_or(false)
2742 {
2743 return;
2744 }
2745 let Some(permission_handler) = handlers.permission.clone() else {
2749 return;
2750 };
2751 let client = client.clone();
2752 let sid = session_id.clone();
2753 let shutdown = shutdown.clone();
2754 let data = permission_request_data(
2755 ¬ification.event.data,
2756 handlers.managed_settings_enabled,
2757 );
2758 let span = tracing::error_span!(
2759 "permission_request_handler",
2760 session_id = %sid,
2761 request_id = %request_id
2762 );
2763 let _ = spawn_startup_tracked(
2764 async move {
2765 let handler_start = Instant::now();
2766 let result = permission_handler
2767 .handle(sid.clone(), request_id.clone(), data)
2768 .await;
2769 tracing::debug!(
2770 elapsed_ms = handler_start.elapsed().as_millis(),
2771 session_id = %sid,
2772 request_id = %request_id,
2773 "PermissionHandler::handle dispatch"
2774 );
2775 let Some(params) = permission_response_params(&sid, &request_id, &result)
2776 else {
2777 return;
2781 };
2782 let rpc_start = Instant::now();
2783 let method =
2784 rpc_methods::SESSION_PERMISSIONS_HANDLEPENDINGPERMISSIONREQUEST;
2785 tokio::select! {
2786 biased;
2787 response = client.call(method, Some(params)) => {
2788 match response {
2789 Ok(_) => tracing::debug!(
2790 elapsed_ms = rpc_start.elapsed().as_millis(),
2791 session_id = %sid,
2792 request_id = %request_id,
2793 method,
2794 "Session::handle_notification response sent successfully"
2795 ),
2796 Err(error) => warn!(
2797 error = %error,
2798 session_id = %sid,
2799 request_id = %request_id,
2800 method,
2801 "failed to deliver permission decision back to the runtime"
2802 ),
2803 }
2804 }
2805 _ = shutdown.cancelled() => {
2806 warn!(
2807 elapsed_ms = rpc_start.elapsed().as_millis(),
2808 session_id = %sid,
2809 request_id = %request_id,
2810 method,
2811 delivery_outcome = "unknown",
2812 "permission confirmation acknowledgement wait cancelled during session shutdown"
2813 );
2814 }
2815 }
2816 }
2817 .instrument(span),
2818 startup_tasks,
2819 );
2820 }
2821 SessionEventType::ExternalToolRequested => {
2822 let Some(request_id) = extract_request_id(¬ification.event.data) else {
2823 return;
2824 };
2825 let data: ExternalToolRequestedData =
2826 match serde_json::from_value(notification.event.data.clone()) {
2827 Ok(d) => d,
2828 Err(e) => {
2829 warn!(error = %e, "failed to deserialize external_tool.requested");
2830 let client = client.clone();
2831 let sid = session_id.clone();
2832 let span = tracing::error_span!(
2833 "external_tool_deserialize_error",
2834 session_id = %sid,
2835 request_id = %request_id
2836 );
2837 let _ = spawn_startup_tracked(
2838 async move {
2839 let rpc_start = Instant::now();
2840 let _ = client
2841 .call(
2842 "session.tools.handlePendingToolCall",
2843 Some(serde_json::json!({
2844 "sessionId": sid,
2845 "requestId": request_id,
2846 "error": format!("Failed to deserialize tool request: {e}"),
2847 })),
2848 )
2849 .await;
2850 tracing::debug!(
2851 elapsed_ms = rpc_start.elapsed().as_millis(),
2852 session_id = %sid,
2853 request_id = %request_id,
2854 "Session::handle_notification response sent successfully"
2855 );
2856 }
2857 .instrument(span),
2858 startup_tasks,
2859 );
2860 return;
2861 }
2862 };
2863 let tool_handler = if data.tool_name.is_empty() {
2867 None
2868 } else {
2869 handlers.tools.get(&data.tool_name).cloned()
2870 };
2871 let Some(tool_handler) = tool_handler else {
2872 return;
2873 };
2874 let cancellation = Arc::new(external_tools_shutdown.child_token());
2875 {
2876 let mut pending = pending_external_tools.lock();
2877 if external_tools_shutdown.is_cancelled() || pending.contains_key(&request_id) {
2878 return;
2879 }
2880 pending.insert(request_id.clone(), cancellation.clone());
2881 }
2882 let client = client.clone();
2883 let sid = session_id.clone();
2884 let pending_external_tools = pending_external_tools.clone();
2885 let guard_request_id = request_id.clone();
2886 let guard_cancellation = cancellation.clone();
2887 let span = tracing::error_span!(
2888 "external_tool_handler",
2889 session_id = %sid,
2890 request_id = %request_id
2891 );
2892 let _ = spawn_startup_tracked(
2893 async move {
2894 let guard = PendingExternalToolGuard {
2895 request_id: guard_request_id,
2896 token: guard_cancellation,
2897 pending: pending_external_tools,
2898 };
2899 if cancellation.is_cancelled() {
2900 return;
2901 }
2902 if data.tool_call_id.is_empty() {
2907 if !guard.claim() {
2908 return;
2909 }
2910 let error_msg = "Missing toolCallId";
2911 let rpc_start = Instant::now();
2912 let _ = client
2913 .call(
2914 "session.tools.handlePendingToolCall",
2915 Some(serde_json::json!({
2916 "sessionId": sid,
2917 "requestId": request_id,
2918 "error": error_msg,
2919 })),
2920 )
2921 .await;
2922 tracing::debug!(
2923 elapsed_ms = rpc_start.elapsed().as_millis(),
2924 session_id = %sid,
2925 request_id = %request_id,
2926 "Session::handle_notification response sent successfully"
2927 );
2928 return;
2929 }
2930 let tool_call_id = data.tool_call_id.clone();
2931 let tool_name = data.tool_name.clone();
2932 let available_tools = if tool_name == TOOL_SEARCH_TOOL_NAME {
2939 let metadata_result = tokio::select! {
2940 biased;
2941 _ = cancellation.cancelled() => return,
2942 result = client.call(
2943 rpc_methods::SESSION_TOOLS_GETCURRENTMETADATA,
2944 Some(serde_json::json!({ "sessionId": sid })),
2945 ) => result,
2946 };
2947 match metadata_result {
2948 Ok(value) => {
2949 serde_json::from_value::<ToolsGetCurrentMetadataResult>(value)
2950 .ok()
2951 .and_then(|result| result.tools)
2952 }
2953 Err(_) => None,
2954 }
2955 } else {
2956 None
2957 };
2958 let invocation = ToolInvocation {
2959 session_id: sid.clone(),
2960 tool_call_id: data.tool_call_id,
2961 tool_name: data.tool_name,
2962 arguments: data
2963 .arguments
2964 .unwrap_or(Value::Object(serde_json::Map::new())),
2965 available_tools,
2966 traceparent: data.traceparent,
2967 tracestate: data.tracestate,
2968 };
2969 let handler_start = Instant::now();
2970 let tool_result = tokio::select! {
2971 biased;
2972 _ = cancellation.cancelled() => return,
2973 result = tool_handler.call(invocation) => match result {
2974 Ok(r) => r,
2975 Err(e) => tool_failure_result(e.to_string()),
2976 },
2977 };
2978 tracing::debug!(
2979 elapsed_ms = handler_start.elapsed().as_millis(),
2980 session_id = %sid,
2981 request_id = %request_id,
2982 tool_call_id = %tool_call_id,
2983 tool_name = %tool_name,
2984 "ToolHandler::call dispatch"
2985 );
2986 if !guard.claim() {
2987 return;
2988 }
2989 let result_value = serde_json::to_value(tool_result).unwrap_or(Value::Null);
2990 let rpc_start = Instant::now();
2991 let _ = client
2992 .call(
2993 "session.tools.handlePendingToolCall",
2994 Some(serde_json::json!({
2995 "sessionId": sid,
2996 "requestId": request_id,
2997 "result": result_value,
2998 })),
2999 )
3000 .await;
3001 tracing::debug!(
3002 elapsed_ms = rpc_start.elapsed().as_millis(),
3003 session_id = %sid,
3004 request_id = %request_id,
3005 tool_call_id = %tool_call_id,
3006 tool_name = %tool_name,
3007 "Session::handle_notification response sent successfully"
3008 );
3009 }
3010 .instrument(span),
3011 startup_tasks,
3012 );
3013 }
3014 SessionEventType::UserInputRequested => {
3015 }
3022 SessionEventType::ElicitationRequested => {
3023 let Some(request_id) = extract_request_id(¬ification.event.data) else {
3024 return;
3025 };
3026 let Some(elicitation_handler) = handlers.elicitation.clone() else {
3030 return;
3031 };
3032 let elicitation_data: ElicitationRequestedData =
3033 match serde_json::from_value(notification.event.data.clone()) {
3034 Ok(d) => d,
3035 Err(e) => {
3036 warn!(error = %e, "failed to deserialize elicitation request");
3037 return;
3038 }
3039 };
3040 let request = ElicitationRequest {
3041 message: elicitation_data.message,
3042 requested_schema: elicitation_data
3043 .requested_schema
3044 .map(|s| serde_json::to_value(s).unwrap_or(Value::Null)),
3045 mode: elicitation_data.mode.map(|m| match m {
3046 crate::generated::session_events::ElicitationRequestedMode::Form => {
3047 crate::types::ElicitationMode::Form
3048 }
3049 crate::generated::session_events::ElicitationRequestedMode::Url => {
3050 crate::types::ElicitationMode::Url
3051 }
3052 _ => crate::types::ElicitationMode::Unknown,
3053 }),
3054 elicitation_source: elicitation_data.elicitation_source,
3055 url: elicitation_data.url,
3056 };
3057 let client = client.clone();
3058 let sid = session_id.clone();
3059 let span = tracing::error_span!(
3060 "elicitation_request_handler",
3061 session_id = %sid,
3062 request_id = %request_id
3063 );
3064 let nested_startup_tasks = startup_tasks.cloned();
3065 let _ = spawn_startup_tracked(
3066 async move {
3067 let cancel = ElicitationResult {
3068 action: "cancel".to_string(),
3069 content: None,
3070 };
3071 let handler_task = spawn_startup_tracked(
3073 {
3074 let sid = sid.clone();
3075 let request_id = request_id.clone();
3076 let span = tracing::error_span!(
3077 "elicitation_callback",
3078 session_id = %sid,
3079 request_id = %request_id
3080 );
3081 async move {
3082 let handler_start = Instant::now();
3083 let response = elicitation_handler
3084 .handle(sid.clone(), request_id.clone(), request)
3085 .await;
3086 tracing::debug!(
3087 elapsed_ms = handler_start.elapsed().as_millis(),
3088 session_id = %sid,
3089 request_id = %request_id,
3090 "ElicitationHandler::handle dispatch"
3091 );
3092 response
3093 }
3094 .instrument(span)
3095 },
3096 nested_startup_tasks.as_ref(),
3097 );
3098 let Some(handler_task) = handler_task else {
3099 return;
3101 };
3102 let result = match handler_task.await {
3103 Ok(r) => r,
3104 Err(error) if error.is_cancelled() => return,
3106 Err(_) => cancel.clone(),
3107 };
3108 let rpc_start = Instant::now();
3109 if let Err(e) = client
3110 .call(
3111 "session.ui.handlePendingElicitation",
3112 Some(serde_json::json!({
3113 "sessionId": sid,
3114 "requestId": request_id,
3115 "result": result,
3116 })),
3117 )
3118 .await
3119 {
3120 warn!(error = %e, "handlePendingElicitation failed, sending cancel");
3122 let _ = client
3123 .call(
3124 "session.ui.handlePendingElicitation",
3125 Some(serde_json::json!({
3126 "sessionId": sid,
3127 "requestId": request_id,
3128 "result": cancel,
3129 })),
3130 )
3131 .await;
3132 } else {
3133 tracing::debug!(
3134 elapsed_ms = rpc_start.elapsed().as_millis(),
3135 session_id = %sid,
3136 request_id = %request_id,
3137 "Session::handle_notification response sent successfully"
3138 );
3139 }
3140 }
3141 .instrument(span),
3142 startup_tasks,
3143 );
3144 }
3145 SessionEventType::McpOauthRequired => {
3146 let Some(request_id) = extract_request_id(¬ification.event.data) else {
3147 return;
3148 };
3149 let Some(mcp_auth_handler) = handlers.mcp_auth.clone() else {
3150 warn!(
3151 session_id = %session_id,
3152 request_id = %request_id,
3153 "received MCP OAuth request without a registered MCP auth handler"
3154 );
3155 return;
3156 };
3157 let data: McpOauthRequiredData =
3158 match serde_json::from_value(notification.event.data.clone()) {
3159 Ok(d) => d,
3160 Err(e) => {
3161 warn!(error = %e, "failed to deserialize MCP OAuth request");
3162 return;
3163 }
3164 };
3165 let request = McpAuthRequest {
3166 request_id: request_id.clone(),
3167 server_name: data.server_name,
3168 server_url: data.server_url,
3169 reason: data.reason,
3170 www_authenticate_params: data.www_authenticate_params,
3171 resource_metadata: data.resource_metadata,
3172 static_client_config: data.static_client_config,
3173 };
3174 let client = client.clone();
3175 let sid = session_id.clone();
3176 let span = tracing::error_span!(
3177 "mcp_auth_request_handler",
3178 session_id = %sid,
3179 request_id = %request_id
3180 );
3181 let nested_startup_tasks = startup_tasks.cloned();
3182 let _ = spawn_startup_tracked(
3183 async move {
3184 let cancel = McpAuthResult::Cancelled;
3185 let handler_task = spawn_startup_tracked(
3186 {
3187 let sid = sid.clone();
3188 let request_id = request_id.clone();
3189 let span = tracing::error_span!(
3190 "mcp_auth_callback",
3191 session_id = %sid,
3192 request_id = %request_id
3193 );
3194 async move {
3195 let handler_start = Instant::now();
3196 let response = mcp_auth_handler
3197 .handle(sid.clone(), request_id.clone(), request)
3198 .await;
3199 tracing::debug!(
3200 elapsed_ms = handler_start.elapsed().as_millis(),
3201 session_id = %sid,
3202 request_id = %request_id,
3203 "McpAuthHandler::handle dispatch"
3204 );
3205 response
3206 }
3207 .instrument(span)
3208 },
3209 nested_startup_tasks.as_ref(),
3210 );
3211 let Some(handler_task) = handler_task else {
3212 return;
3214 };
3215 let result = match handler_task.await {
3216 Ok(result) => result,
3217 Err(error) if error.is_cancelled() => return,
3219 Err(_) => cancel,
3220 };
3221 let rpc_start = Instant::now();
3222 let _ = client
3223 .call(
3224 "session.mcp.oauth.handlePendingRequest",
3225 Some(serde_json::json!({
3226 "sessionId": sid,
3227 "requestId": request_id,
3228 "result": result.into_wire(),
3229 })),
3230 )
3231 .await;
3232 tracing::debug!(
3233 elapsed_ms = rpc_start.elapsed().as_millis(),
3234 "Session::handle_notification MCP auth response sent"
3235 );
3236 }
3237 .instrument(span),
3238 startup_tasks,
3239 );
3240 }
3241 SessionEventType::CommandExecute => {
3242 let data: CommandExecuteData =
3243 match serde_json::from_value(notification.event.data.clone()) {
3244 Ok(d) => d,
3245 Err(e) => {
3246 warn!(error = %e, "failed to deserialize command.execute");
3247 return;
3248 }
3249 };
3250 let client = client.clone();
3251 let command_handlers = command_handlers.clone();
3252 let sid = session_id.clone();
3253 let span = tracing::error_span!("command_handler", session_id = %sid);
3254 let _ = spawn_startup_tracked(
3255 async move {
3256 let request_id = data.request_id;
3257 let ack_error = match command_handlers.get(&data.command_name).cloned() {
3258 None => Some(format!("Unknown command: {}", data.command_name)),
3259 Some(handler) => {
3260 let command_name = data.command_name.clone();
3261 let ctx = CommandContext {
3262 session_id: sid.clone(),
3263 command: data.command,
3264 command_name: data.command_name,
3265 args: data.args,
3266 };
3267 let handler_start = Instant::now();
3268 let result = handler.on_command(ctx).await;
3269 tracing::debug!(
3270 elapsed_ms = handler_start.elapsed().as_millis(),
3271 session_id = %sid,
3272 request_id = %request_id,
3273 command_name = %command_name,
3274 "CommandHandler::call dispatch"
3275 );
3276 match result {
3277 Ok(()) => None,
3278 Err(e) => Some(e.to_string()),
3279 }
3280 }
3281 };
3282 let mut params = serde_json::json!({
3283 "sessionId": sid,
3284 "requestId": request_id,
3285 });
3286 if let Some(error_msg) = ack_error {
3287 params["error"] = serde_json::Value::String(error_msg);
3288 }
3289 let rpc_start = Instant::now();
3290 let _ = client
3291 .call("session.commands.handlePendingCommand", Some(params))
3292 .await;
3293 tracing::debug!(
3294 elapsed_ms = rpc_start.elapsed().as_millis(),
3295 session_id = %sid,
3296 request_id = %request_id,
3297 "Session::handle_notification response sent successfully"
3298 );
3299 }
3300 .instrument(span),
3301 startup_tasks,
3302 );
3303 }
3304 _ => {}
3305 }
3306}
3307
3308struct RequestDispatchContext<'a> {
3309 client: &'a Client,
3310 handlers: &'a SessionHandlers,
3311 hooks: Option<&'a dyn SessionHooks>,
3312 transforms: Option<&'a dyn SystemMessageTransform>,
3313 canvas_handler: Option<&'a Arc<dyn CanvasHandler>>,
3314 session_fs_provider: Option<&'a Arc<dyn SessionFsProvider>>,
3315 bearer_token_providers: &'a HashMap<String, Arc<dyn BearerTokenProvider>>,
3316}
3317
3318async fn handle_request(
3320 session_id: &SessionId,
3321 ctx: RequestDispatchContext<'_>,
3322 request: crate::JsonRpcRequest,
3323) {
3324 let sid = session_id.clone();
3325 let client = ctx.client;
3326 let handlers = ctx.handlers;
3327 let hooks = ctx.hooks;
3328 let transforms = ctx.transforms;
3329 let canvas_handler = ctx.canvas_handler;
3330 let session_fs_provider = ctx.session_fs_provider;
3331 let bearer_token_providers = ctx.bearer_token_providers;
3332
3333 if request.method.starts_with("sessionFs.") {
3334 crate::session_fs_dispatch::dispatch(client, session_fs_provider, request).await;
3335 return;
3336 }
3337
3338 if request.method.starts_with("canvas.") {
3339 crate::canvas_dispatch::dispatch(client, canvas_handler, request).await;
3340 return;
3341 }
3342
3343 if request.method == crate::generated::api_types::rpc_methods::PROVIDERTOKEN_GETTOKEN {
3344 crate::provider_token_dispatch::dispatch(client, bearer_token_providers, request).await;
3345 return;
3346 }
3347
3348 match request.method.as_str() {
3349 "hooks.invoke" => {
3350 let params = request.params.as_ref();
3351 let hook_type = params
3352 .and_then(|p| p.get("hookType"))
3353 .and_then(|v| v.as_str())
3354 .unwrap_or("");
3355 let input = params
3356 .and_then(|p| p.get("input"))
3357 .cloned()
3358 .unwrap_or(Value::Object(Default::default()));
3359
3360 let rpc_result = if let Some(hooks) = hooks {
3361 match crate::hooks::dispatch_hook(hooks, &sid, hook_type, input).await {
3362 Ok(output) => output,
3363 Err(e) => {
3364 warn!(error = %e, hook_type = hook_type, "hook dispatch failed");
3365 serde_json::json!({ "output": {} })
3366 }
3367 }
3368 } else {
3369 serde_json::json!({ "output": {} })
3370 };
3371
3372 let rpc_response = JsonRpcResponse {
3373 jsonrpc: "2.0".to_string(),
3374 id: request.id,
3375 result: Some(rpc_result),
3376 error: None,
3377 };
3378 let _ = client.send_response(&rpc_response).await;
3379 }
3380
3381 "userInput.request" => {
3382 let params = request.params.as_ref();
3383 let Some(question) = params
3384 .and_then(|p| p.get("question"))
3385 .and_then(|v| v.as_str())
3386 else {
3387 warn!("userInput.request missing 'question' field");
3388 let rpc_response = JsonRpcResponse {
3389 jsonrpc: "2.0".to_string(),
3390 id: request.id,
3391 result: None,
3392 error: Some(crate::JsonRpcError {
3393 code: error_codes::INVALID_PARAMS,
3394 message: "missing required field: question".to_string(),
3395 data: None,
3396 }),
3397 };
3398 let _ = client.send_response(&rpc_response).await;
3399 return;
3400 };
3401 let question = question.to_string();
3402 let choices = params
3403 .and_then(|p| p.get("choices"))
3404 .and_then(|v| v.as_array())
3405 .map(|arr| {
3406 arr.iter()
3407 .filter_map(|v| v.as_str().map(|s| s.to_string()))
3408 .collect()
3409 });
3410 let allow_freeform = params
3411 .and_then(|p| p.get("allowFreeform"))
3412 .and_then(|v| v.as_bool());
3413
3414 let handler_start = Instant::now();
3415 let response = if let Some(user_input_handler) = handlers.user_input.as_ref() {
3416 user_input_handler
3417 .handle(sid.clone(), question, choices, allow_freeform)
3418 .await
3419 } else {
3420 None
3421 };
3422 tracing::debug!(
3423 elapsed_ms = handler_start.elapsed().as_millis(),
3424 session_id = %sid,
3425 "UserInputHandler::handle dispatch"
3426 );
3427
3428 let rpc_result = match response {
3429 Some(UserInputResponse {
3430 answer,
3431 was_freeform,
3432 }) => serde_json::json!({
3433 "answer": answer,
3434 "wasFreeform": was_freeform,
3435 }),
3436 None => serde_json::json!({ "noResponse": true }),
3437 };
3438 let rpc_response = JsonRpcResponse {
3439 jsonrpc: "2.0".to_string(),
3440 id: request.id,
3441 result: Some(rpc_result),
3442 error: None,
3443 };
3444 let _ = client.send_response(&rpc_response).await;
3445 }
3446
3447 "exitPlanMode.request" => {
3448 let params = request
3449 .params
3450 .as_ref()
3451 .cloned()
3452 .unwrap_or(Value::Object(serde_json::Map::new()));
3453 let data: ExitPlanModeData = match serde_json::from_value(params) {
3454 Ok(d) => d,
3455 Err(e) => {
3456 warn!(error = %e, "failed to deserialize exitPlanMode.request params, using defaults");
3457 ExitPlanModeData::default()
3458 }
3459 };
3460
3461 let rpc_result = if let Some(exit_plan_handler) = handlers.exit_plan_mode.as_ref() {
3462 let result = exit_plan_handler.handle(sid, data).await;
3463 serde_json::to_value(result).expect("ExitPlanModeResult serialization cannot fail")
3464 } else {
3465 serde_json::json!({ "approved": true })
3466 };
3467 let rpc_response = JsonRpcResponse {
3468 jsonrpc: "2.0".to_string(),
3469 id: request.id,
3470 result: Some(rpc_result),
3471 error: None,
3472 };
3473 let _ = client.send_response(&rpc_response).await;
3474 }
3475
3476 "autoModeSwitch.request" => {
3477 let error_code = request
3478 .params
3479 .as_ref()
3480 .and_then(|p| p.get("errorCode"))
3481 .and_then(|v| v.as_str())
3482 .map(|s| s.to_string());
3483 let retry_after_seconds = request
3484 .params
3485 .as_ref()
3486 .and_then(|p| p.get("retryAfterSeconds"))
3487 .and_then(|v| v.as_f64());
3488
3489 let answer = if let Some(auto_mode_handler) = handlers.auto_mode_switch.as_ref() {
3490 auto_mode_handler
3491 .handle(sid, error_code, retry_after_seconds)
3492 .await
3493 } else {
3494 AutoModeSwitchResponse::No
3495 };
3496 let rpc_response = JsonRpcResponse {
3497 jsonrpc: "2.0".to_string(),
3498 id: request.id,
3499 result: Some(serde_json::json!({ "response": answer })),
3500 error: None,
3501 };
3502 let _ = client.send_response(&rpc_response).await;
3503 }
3504
3505 "systemMessage.transform" => {
3506 let params = request.params.as_ref();
3507 let sections: HashMap<String, crate::transforms::TransformSection> =
3508 match params.and_then(|p| p.get("sections")) {
3509 Some(v) => match serde_json::from_value(v.clone()) {
3510 Ok(s) => s,
3511 Err(e) => {
3512 let _ = send_error_response(
3513 client,
3514 request.id,
3515 error_codes::INVALID_PARAMS,
3516 &format!("invalid sections: {e}"),
3517 )
3518 .await;
3519 return;
3520 }
3521 },
3522 None => {
3523 let _ = send_error_response(
3524 client,
3525 request.id,
3526 error_codes::INVALID_PARAMS,
3527 "missing sections parameter",
3528 )
3529 .await;
3530 return;
3531 }
3532 };
3533
3534 let rpc_result = if let Some(transforms) = transforms {
3535 let transform_start = Instant::now();
3536 let response =
3537 crate::transforms::dispatch_transform(transforms, &sid, sections).await;
3538 tracing::debug!(
3539 elapsed_ms = transform_start.elapsed().as_millis(),
3540 session_id = %sid,
3541 "SystemMessageTransform::transform_section dispatch"
3542 );
3543 match serde_json::to_value(response) {
3544 Ok(v) => v,
3545 Err(e) => {
3546 warn!(error = %e, "failed to serialize transform response");
3547 serde_json::json!({ "sections": {} })
3548 }
3549 }
3550 } else {
3551 let passthrough: HashMap<String, crate::transforms::TransformSection> = sections;
3553 serde_json::json!({ "sections": passthrough })
3554 };
3555
3556 let rpc_response = JsonRpcResponse {
3557 jsonrpc: "2.0".to_string(),
3558 id: request.id,
3559 result: Some(rpc_result),
3560 error: None,
3561 };
3562 let _ = client.send_response(&rpc_response).await;
3563 }
3564
3565 method => {
3566 warn!(
3567 method = method,
3568 "unhandled request method in session event loop"
3569 );
3570 let _ = send_error_response(
3571 client,
3572 request.id,
3573 error_codes::METHOD_NOT_FOUND,
3574 &format!("unknown method: {method}"),
3575 )
3576 .await;
3577 }
3578 }
3579}
3580
3581async fn send_error_response(
3582 client: &Client,
3583 id: u64,
3584 code: i32,
3585 message: &str,
3586) -> Result<(), Error> {
3587 let response = JsonRpcResponse {
3588 jsonrpc: "2.0".to_string(),
3589 id,
3590 result: None,
3591 error: Some(crate::JsonRpcError {
3592 code,
3593 message: message.to_string(),
3594 data: None,
3595 }),
3596 };
3597 client.send_response(&response).await
3598}
3599
3600fn apply_transform_sections(
3604 sys_msg: &mut SystemMessageConfig,
3605 transforms: &dyn SystemMessageTransform,
3606) {
3607 sys_msg.mode = Some("customize".to_string());
3608 let sections = sys_msg.sections.get_or_insert_with(HashMap::new);
3609 for id in transforms.section_ids() {
3610 sections.entry(id).or_insert_with(|| SectionOverride {
3611 action: Some("transform".to_string()),
3612 content: None,
3613 });
3614 }
3615}
3616
3617fn inject_transform_sections(config: &mut SessionConfig, transforms: &dyn SystemMessageTransform) {
3618 let sys_msg = config.system_message.get_or_insert_with(Default::default);
3619 apply_transform_sections(sys_msg, transforms);
3620}
3621
3622fn inject_transform_sections_resume(
3623 config: &mut ResumeSessionConfig,
3624 transforms: &dyn SystemMessageTransform,
3625) {
3626 let sys_msg = config.system_message.get_or_insert_with(Default::default);
3627 apply_transform_sections(sys_msg, transforms);
3628}
3629
3630#[cfg(test)]
3631mod tests {
3632 use serde_json::json;
3633
3634 use super::{
3635 build_mode_post_create_patch, has_managed_settings, is_autopilot_continuation_idle,
3636 permission_request_data, permission_response_params,
3637 };
3638 use crate::handler::PermissionResult;
3639 use crate::types::{
3640 PermissionDecisionContext, PermissionDecisionOutcome, PermissionDecisionSource,
3641 PermissionDecisionSurface, RequestId, SessionEvent, SessionId,
3642 };
3643
3644 #[test]
3645 fn identifies_only_autopilot_continuation_idles() {
3646 let mut event = SessionEvent {
3647 id: "event-1".to_string(),
3648 timestamp: "2026-01-01T00:00:00Z".to_string(),
3649 parent_id: None,
3650 ephemeral: None,
3651 agent_id: None,
3652 debug_cli_received_at_ms: None,
3653 debug_ws_forwarded_at_ms: None,
3654 event_type: "session.idle".to_string(),
3655 data: json!({ "mode": "autopilot" }),
3656 };
3657
3658 assert!(is_autopilot_continuation_idle(&event));
3659
3660 event.data = json!({ "mode": "interactive" });
3661 assert!(!is_autopilot_continuation_idle(&event));
3662
3663 event.data = json!({});
3664 assert!(!is_autopilot_continuation_idle(&event));
3665 }
3666
3667 #[test]
3668 fn empty_mode_post_patch_sets_empty_included_builtin_skills() {
3669 let patch =
3670 build_mode_post_create_patch(crate::ClientMode::Empty, None, None, None, None, None)
3671 .expect("empty mode always sends a patch");
3672 assert_eq!(
3673 patch.included_builtin_skills,
3674 Some(Vec::new()),
3675 "empty mode must fail closed with an empty includedBuiltinSkills list"
3676 );
3677 assert_eq!(patch.installed_plugins.as_ref().map(|p| p.len()), Some(0));
3678 let value = serde_json::to_value(&patch).expect("serialize patch");
3680 assert_eq!(value["includedBuiltinSkills"], serde_json::json!([]));
3681 }
3682
3683 #[test]
3684 fn empty_mode_post_patch_preserves_explicit_builtin_skill_allowlist() {
3685 let patch = build_mode_post_create_patch(
3686 crate::ClientMode::Empty,
3687 Some(false),
3688 Some(false),
3689 Some(true),
3690 Some(true),
3691 Some(vec!["code-review".to_string()]),
3692 )
3693 .expect("empty mode always sends a patch");
3694 assert_eq!(
3695 patch.included_builtin_skills,
3696 Some(vec!["code-review".to_string()])
3697 );
3698 }
3699
3700 #[test]
3701 fn copilot_cli_mode_does_not_inject_included_builtin_skills() {
3702 assert!(
3704 build_mode_post_create_patch(
3705 crate::ClientMode::CopilotCli,
3706 None,
3707 None,
3708 None,
3709 None,
3710 None
3711 )
3712 .is_none()
3713 );
3714 let patch = build_mode_post_create_patch(
3716 crate::ClientMode::CopilotCli,
3717 Some(true),
3718 None,
3719 None,
3720 None,
3721 None,
3722 )
3723 .expect("a set field triggers a patch");
3724 assert_eq!(patch.included_builtin_skills, None);
3725 assert!(patch.installed_plugins.is_none());
3726 let value = serde_json::to_value(&patch).expect("serialize patch");
3727 assert!(value.get("includedBuiltinSkills").is_none());
3728
3729 let patch = build_mode_post_create_patch(
3730 crate::ClientMode::CopilotCli,
3731 None,
3732 None,
3733 None,
3734 None,
3735 Some(vec!["code-review".to_string()]),
3736 )
3737 .expect("an explicit allowlist triggers a patch");
3738 assert_eq!(
3739 patch.included_builtin_skills,
3740 Some(vec!["code-review".to_string()])
3741 );
3742 }
3743
3744 #[test]
3745 fn direct_injection_enables_managed_safeguards() {
3746 let settings = crate::types::ManagedSettings::default();
3747 assert!(has_managed_settings(None, Some(&settings)));
3748 assert!(!has_managed_settings(None, None));
3749 }
3750
3751 fn attribution_context() -> PermissionDecisionContext {
3752 PermissionDecisionContext {
3753 outcome: PermissionDecisionOutcome::AutoApproved,
3754 response_capability: None,
3755 source: PermissionDecisionSource::AssistedApproval,
3756 surface: PermissionDecisionSurface::CopilotApp,
3757 }
3758 }
3759
3760 #[test]
3761 fn response_params_omit_decision_context_without_attribution() {
3762 for (result, expected) in [
3763 (
3764 PermissionResult::approve_once(),
3765 json!({ "kind": "approve-once" }),
3766 ),
3767 (PermissionResult::reject(None), json!({ "kind": "reject" })),
3768 (
3769 PermissionResult::reject(Some("bad".to_string())),
3770 json!({ "kind": "reject", "feedback": "bad" }),
3771 ),
3772 (
3773 PermissionResult::user_not_available(),
3774 json!({ "kind": "user-not-available" }),
3775 ),
3776 ] {
3777 let params = permission_response_params(
3778 &SessionId::from("session-1"),
3779 &RequestId::from("permission-1"),
3780 &result,
3781 )
3782 .unwrap();
3783 assert_eq!(
3784 params,
3785 json!({
3786 "sessionId": "session-1",
3787 "requestId": "permission-1",
3788 "result": expected,
3789 })
3790 );
3791 }
3792 }
3793
3794 #[test]
3795 fn response_params_forward_decision_context_alongside_result() {
3796 let params = permission_response_params(
3797 &SessionId::from("session-1"),
3798 &RequestId::from("permission-1"),
3799 &PermissionResult::approve_once().with_context(attribution_context()),
3800 )
3801 .unwrap();
3802 assert_eq!(
3803 params,
3804 json!({
3805 "sessionId": "session-1",
3806 "requestId": "permission-1",
3807 "result": { "kind": "approve-once" },
3808 "decisionContext": {
3809 "outcome": "auto_approved",
3810 "source": "assisted_approval",
3811 "surface": "copilot_app",
3812 },
3813 })
3814 );
3815 assert!(params["result"].get("decisionContext").is_none());
3817 }
3818
3819 #[test]
3820 fn response_params_suppressed_for_no_result() {
3821 assert!(
3822 permission_response_params(
3823 &SessionId::from("session-1"),
3824 &RequestId::from("permission-1"),
3825 &PermissionResult::NoResult,
3826 )
3827 .is_none()
3828 );
3829 }
3830
3831 #[test]
3832 fn with_context_is_a_no_op_on_no_result() {
3833 let result = PermissionResult::no_result().with_context(attribution_context());
3834 assert!(matches!(result, PermissionResult::NoResult));
3835 }
3836
3837 #[test]
3838 fn with_context_replaces_rather_than_nests() {
3839 let result = PermissionResult::approve_once()
3840 .with_context(attribution_context())
3841 .with_context(PermissionDecisionContext {
3842 outcome: PermissionDecisionOutcome::PromptedUser,
3843 response_capability: None,
3844 source: PermissionDecisionSource::HumanResponse,
3845 surface: PermissionDecisionSurface::Sdk,
3846 });
3847 let params = permission_response_params(
3848 &SessionId::from("session-1"),
3849 &RequestId::from("permission-1"),
3850 &result,
3851 )
3852 .unwrap();
3853 assert_eq!(
3854 params["decisionContext"],
3855 json!({
3856 "outcome": "prompted_user",
3857 "source": "human_response",
3858 "surface": "sdk",
3859 })
3860 );
3861 }
3862
3863 #[test]
3864 fn permission_request_data_reads_nested_managed_approval_metadata() {
3865 let data = permission_request_data(
3866 &json!({
3867 "requestId": "permission-1",
3868 "permissionRequest": {
3869 "kind": "read",
3870 "managedApprovalRequired": true,
3871 "path": "/workspace/file.txt"
3872 }
3873 }),
3874 false,
3875 );
3876
3877 assert_eq!(data.managed_approval_required, Some(true));
3878 assert_eq!(
3879 data.extra["permissionRequest"]["path"],
3880 "/workspace/file.txt"
3881 );
3882 }
3883
3884 #[test]
3885 fn permission_request_data_preserves_managed_flag_when_other_fields_are_malformed() {
3886 let data = permission_request_data(
3887 &json!({
3888 "requestId": "permission-1",
3889 "permissionRequest": {
3890 "kind": "read",
3891 "managedApprovalRequired": true,
3892 "toolCallId": 42
3893 }
3894 }),
3895 false,
3896 );
3897
3898 assert_eq!(data.managed_approval_required, Some(true));
3899 assert_eq!(data.extra["requestId"], "permission-1");
3900 }
3901
3902 #[test]
3903 fn permission_request_data_fails_closed_for_malformed_managed_flag() {
3904 let data = permission_request_data(
3905 &json!({
3906 "requestId": "permission-1",
3907 "permissionRequest": {
3908 "kind": "read",
3909 "managedApprovalRequired": "yes",
3910 "path": "/workspace/file.txt"
3911 }
3912 }),
3913 false,
3914 );
3915
3916 assert_eq!(data.managed_approval_required, Some(true));
3917 }
3918
3919 #[test]
3920 fn permission_request_data_preserves_valid_false_managed_flag() {
3921 let data = permission_request_data(
3922 &json!({
3923 "requestId": "permission-1",
3924 "permissionRequest": {
3925 "kind": "read",
3926 "managedApprovalRequired": false,
3927 "path": "/workspace/file.txt"
3928 }
3929 }),
3930 false,
3931 );
3932
3933 assert_eq!(data.managed_approval_required, Some(false));
3934 }
3935}
3936
3937#[cfg(test)]
3938mod startup_tasks_tests;