1use std::collections::HashMap;
2use std::path::{Path, PathBuf};
3use std::sync::Arc;
4use std::time::{Duration, Instant};
5
6use parking_lot::Mutex as ParkingLotMutex;
7use serde_json::Value;
8use tokio::sync::oneshot;
9use tokio::task::JoinHandle;
10use tokio_util::sync::CancellationToken;
11use tracing::{Instrument, warn};
12
13use crate::canvas::CanvasHandler;
14use crate::generated::api_types::{
15 LogRequest, ModelSwitchToRequest, OpenCanvasInstance, PermissionDecisionRequest,
16 RegisterEventInterestParams, ToolsGetCurrentMetadataResult, rpc_methods,
17};
18use crate::generated::session_events::{
19 CommandExecuteData, ElicitationRequestedData, ExternalToolRequestedData, McpOauthRequiredData,
20 SessionCanvasClosedData, SessionErrorData, SessionEventType, SessionIdleData, SessionMode,
21};
22use crate::handler::{
23 AutoModeSwitchHandler, AutoModeSwitchResponse, ElicitationHandler, ExitPlanModeHandler,
24 McpAuthHandler, McpAuthRequest, McpAuthResult, PermissionHandler, PermissionResult,
25 UserInputHandler, UserInputResponse,
26};
27use crate::hooks::SessionHooks;
28use crate::provider_token::BearerTokenProvider;
29use crate::session_fs::SessionFsProvider;
30use crate::trace_context::inject_trace_context;
31use crate::transforms::SystemMessageTransform;
32use crate::types::{
33 CommandContext, CommandDefinition, CommandHandler, CreateSessionResult, ElicitationRequest,
34 ElicitationResult, ExitPlanModeData, GetMessagesResponse, MessageOptions,
35 PermissionRequestData, RequestId, ResumeSessionConfig, ResumeSessionResult, SectionOverride,
36 SessionCapabilities, SessionConfig, SessionEvent, SessionId, SetModelOptions,
37 SystemMessageConfig, ToolInvocation, ToolResult, ToolResultExpanded, TraceContext,
38 UiInputOptions, ensure_attachment_display_names,
39};
40use crate::{
41 Client, Error, ErrorKind, JsonRpcResponse, SessionErrorKind, SessionEventNotification,
42 error_codes,
43};
44
45const TOOL_SEARCH_TOOL_NAME: &str = "tool_search_tool";
49
50#[derive(Clone)]
58pub(crate) struct SessionHandlers {
59 pub permission: Option<Arc<dyn PermissionHandler>>,
60 pub managed_settings_enabled: bool,
61 pub elicitation: Option<Arc<dyn ElicitationHandler>>,
62 pub mcp_auth: Option<Arc<dyn McpAuthHandler>>,
63 pub user_input: Option<Arc<dyn UserInputHandler>>,
64 pub exit_plan_mode: Option<Arc<dyn ExitPlanModeHandler>>,
65 pub auto_mode_switch: Option<Arc<dyn AutoModeSwitchHandler>>,
66 pub tools: Arc<HashMap<String, Arc<dyn crate::tool::ToolHandler>>>,
67}
68
69fn has_managed_settings(
70 enable_managed_settings: Option<bool>,
71 managed_settings: Option<&crate::types::ManagedSettings>,
72) -> bool {
73 enable_managed_settings == Some(true) || managed_settings.is_some()
74}
75
76struct IdleWaiter {
78 tx: oneshot::Sender<Result<Option<SessionEvent>, Error>>,
79 last_assistant_message: Option<SessionEvent>,
80 started_at: Instant,
81 first_assistant_message_seen: bool,
82}
83
84struct WaiterGuard {
96 slot: Arc<ParkingLotMutex<Option<IdleWaiter>>>,
97}
98
99impl Drop for WaiterGuard {
100 fn drop(&mut self) {
101 self.slot.lock().take();
102 }
103}
104
105struct PendingSessionRegistration {
106 client: Client,
107 session_id: SessionId,
108 shutdown: CancellationToken,
109 disarmed: bool,
110}
111
112impl PendingSessionRegistration {
113 fn new(client: Client, session_id: SessionId, shutdown: CancellationToken) -> Self {
114 Self {
115 client,
116 session_id,
117 shutdown,
118 disarmed: false,
119 }
120 }
121
122 async fn cleanup(mut self, event_loop: JoinHandle<()>) {
123 self.shutdown.cancel();
124 let _ = event_loop.await;
125 self.client.unregister_session(&self.session_id);
126 self.disarmed = true;
127 }
128
129 fn disarm(&mut self) {
130 self.disarmed = true;
131 }
132}
133
134impl Drop for PendingSessionRegistration {
135 fn drop(&mut self) {
136 if !self.disarmed {
137 self.shutdown.cancel();
138 self.client.unregister_session(&self.session_id);
139 }
140 }
141}
142
143pub struct Session {
156 id: SessionId,
157 cwd: PathBuf,
158 workspace_path: Option<PathBuf>,
159 remote_url: Option<String>,
160 client: Client,
161 event_loop: ParkingLotMutex<Option<JoinHandle<()>>>,
166 shutdown: CancellationToken,
180 idle_waiter: Arc<ParkingLotMutex<Option<IdleWaiter>>>,
187 capabilities: Arc<parking_lot::RwLock<SessionCapabilities>>,
189 open_canvases: Arc<parking_lot::RwLock<Vec<OpenCanvasInstance>>>,
191 event_tx: tokio::sync::broadcast::Sender<SessionEvent>,
193 github_token_registration:
194 ParkingLotMutex<Option<crate::github_token::GitHubTokenRegistration>>,
195}
196
197impl Session {
198 pub fn id(&self) -> &SessionId {
200 &self.id
201 }
202
203 pub fn cwd(&self) -> &PathBuf {
205 &self.cwd
206 }
207
208 pub fn workspace_path(&self) -> Option<&Path> {
210 self.workspace_path.as_deref()
211 }
212
213 pub fn remote_url(&self) -> Option<&str> {
215 self.remote_url.as_deref()
216 }
217
218 pub fn capabilities(&self) -> SessionCapabilities {
223 self.capabilities.read().clone()
224 }
225
226 pub fn open_canvases(&self) -> Vec<OpenCanvasInstance> {
229 self.open_canvases.read().clone()
230 }
231
232 pub fn cancellation_token(&self) -> CancellationToken {
259 self.shutdown.child_token()
260 }
261
262 pub fn subscribe(&self) -> crate::subscription::EventSubscription {
301 crate::subscription::EventSubscription::new(self.event_tx.subscribe())
302 }
303
304 pub fn client(&self) -> &Client {
306 &self.client
307 }
308
309 pub fn rpc(&self) -> crate::generated::rpc::SessionRpc<'_> {
320 crate::generated::rpc::SessionRpc { session: self }
321 }
322
323 pub async fn stop_event_loop(&self) {
331 self.shutdown.cancel();
332 let handle = self.event_loop.lock().take();
333 if let Some(handle) = handle {
334 let _ = handle.await;
335 }
336 if let Some(waiter) = self.idle_waiter.lock().take() {
338 let _ = waiter.tx.send(Err(
339 ErrorKind::Session(SessionErrorKind::EventLoopClosed).into()
340 ));
341 }
342 }
343
344 pub async fn send(&self, opts: impl Into<MessageOptions>) -> Result<String, Error> {
369 if self.idle_waiter.lock().is_some() {
370 return Err(ErrorKind::Session(SessionErrorKind::SendWhileWaiting).into());
371 }
372 self.send_inner(opts.into()).await
373 }
374
375 async fn send_inner(&self, opts: MessageOptions) -> Result<String, Error> {
376 let mut params = serde_json::json!({
377 "sessionId": self.id,
378 "prompt": opts.prompt,
379 });
380 if let Some(m) = opts.mode {
381 params["mode"] = serde_json::to_value(m)?;
382 }
383 if let Some(am) = opts.agent_mode {
384 params["agentMode"] = serde_json::to_value(am)?;
385 }
386 if let Some(mut a) = opts.attachments {
387 ensure_attachment_display_names(&mut a);
388 params["attachments"] = serde_json::to_value(a)?;
389 }
390 if let Some(headers) = opts.request_headers
391 && !headers.is_empty()
392 {
393 params["requestHeaders"] = serde_json::to_value(headers)?;
394 }
395 if let Some(display_prompt) = opts.display_prompt {
396 params["displayPrompt"] = serde_json::to_value(display_prompt)?;
397 }
398 let trace_ctx = if opts.traceparent.is_some() || opts.tracestate.is_some() {
399 TraceContext {
400 traceparent: opts.traceparent,
401 tracestate: opts.tracestate,
402 }
403 } else {
404 self.client.resolve_trace_context().await
405 };
406 inject_trace_context(&mut params, &trace_ctx);
407 let rpc_start = Instant::now();
408 let result = self.client.call("session.send", Some(params)).await?;
409 let message_id = result
410 .get("messageId")
411 .and_then(|v| v.as_str())
412 .map(|s| s.to_string())
413 .unwrap_or_default();
414 tracing::debug!(
415 elapsed_ms = rpc_start.elapsed().as_millis(),
416 session_id = %self.id,
417 message_id = %message_id,
418 "Session::send completed successfully"
419 );
420 Ok(message_id)
421 }
422
423 pub async fn send_and_wait(
443 &self,
444 opts: impl Into<MessageOptions>,
445 ) -> Result<Option<SessionEvent>, Error> {
446 let total_start = Instant::now();
447 let opts = opts.into();
448 let timeout_duration = opts.wait_timeout.unwrap_or(Duration::from_secs(60));
449 let (tx, rx) = oneshot::channel();
450
451 {
452 let mut guard = self.idle_waiter.lock();
453 if guard.is_some() {
454 return Err(ErrorKind::Session(SessionErrorKind::SendWhileWaiting).into());
455 }
456 *guard = Some(IdleWaiter {
457 tx,
458 last_assistant_message: None,
459 started_at: total_start,
460 first_assistant_message_seen: false,
461 });
462 }
463
464 let _waiter_guard = WaiterGuard {
469 slot: self.idle_waiter.clone(),
470 };
471
472 let result = tokio::time::timeout(timeout_duration, async {
473 self.send_inner(opts).await?;
474 match rx.await {
475 Ok(result) => result,
476 Err(_) => Err(ErrorKind::Session(SessionErrorKind::EventLoopClosed).into()),
477 }
478 })
479 .await;
480
481 match result {
482 Ok(inner) => {
483 tracing::debug!(
484 elapsed_ms = total_start.elapsed().as_millis(),
485 session_id = %self.id,
486 completed_by = if inner.is_ok() { "idle" } else { "error" },
487 "Session::send_and_wait complete"
488 );
489 inner
490 }
491 Err(_) => {
492 tracing::warn!(
493 elapsed_ms = total_start.elapsed().as_millis(),
494 session_id = %self.id,
495 completed_by = "timeout",
496 "Session::send_and_wait failed"
497 );
498 Err(ErrorKind::Session(SessionErrorKind::Timeout(timeout_duration)).into())
499 }
500 }
501 }
502
503 pub async fn get_events(&self) -> Result<Vec<SessionEvent>, Error> {
505 let result = self
506 .client
507 .call(
508 "session.getMessages",
509 Some(serde_json::json!({ "sessionId": self.id })),
510 )
511 .await?;
512 let response: GetMessagesResponse = serde_json::from_value(result)?;
513 Ok(response.events)
514 }
515
516 #[deprecated(since = "0.1.0", note = "Use `get_events()` instead")]
518 pub async fn get_messages(&self) -> Result<Vec<SessionEvent>, Error> {
519 self.get_events().await
520 }
521
522 pub async fn abort(&self) -> Result<(), Error> {
530 self.client
531 .call(
532 "session.abort",
533 Some(serde_json::json!({ "sessionId": self.id })),
534 )
535 .await?;
536 Ok(())
537 }
538
539 pub async fn set_model(&self, model: &str, opts: Option<SetModelOptions>) -> Result<(), Error> {
543 let opts = opts.unwrap_or_default();
544 let request = ModelSwitchToRequest {
545 compaction_decision: None,
546 context_tier: opts.context_tier,
547 defer_if_model_change_queued: None,
548 model_capabilities: opts.model_capabilities,
549 model_change_scope: None,
550 model_id: model.to_string(),
551 picker_persistence: None,
552 reasoning_effort: opts.reasoning_effort,
553 reasoning_summary: opts.reasoning_summary,
554 repo_scope: None,
555 require_available: None,
556 run_compaction_preflight: None,
557 source: None,
558 verbosity: None,
559 };
560 self.rpc().model().switch_to(request).await?;
561 Ok(())
562 }
563
564 pub async fn disconnect(&self) -> Result<(), Error> {
581 self.client
582 .call(
583 "session.destroy",
584 Some(serde_json::json!({ "sessionId": self.id })),
585 )
586 .await?;
587 self.stop_event_loop().await;
588 self.client.unregister_session(&self.id);
589 self.github_token_registration.lock().take();
590 Ok(())
591 }
592
593 #[deprecated(since = "0.1.0", note = "Use `disconnect()` instead")]
598 pub async fn destroy(&self) -> Result<(), Error> {
599 self.disconnect().await
600 }
601
602 pub async fn log(
606 &self,
607 message: &str,
608 opts: Option<crate::types::LogOptions>,
609 ) -> Result<(), Error> {
610 let opts = opts.unwrap_or_default();
611 let level = match opts.level {
612 Some(level) => Some(serde_json::from_value(serde_json::to_value(level)?)?),
613 None => None,
614 };
615 let request = LogRequest {
616 message: message.to_string(),
617 level,
618 ephemeral: opts.ephemeral,
619 r#type: None,
620 tip: None,
621 url: None,
622 };
623 self.rpc().log(request).await?;
624 Ok(())
625 }
626
627 pub fn ui(&self) -> SessionUi<'_> {
633 SessionUi { session: self }
634 }
635
636 fn assert_elicitation(&self) -> Result<(), Error> {
638 if self
639 .capabilities
640 .read()
641 .ui
642 .as_ref()
643 .and_then(|u| u.elicitation)
644 != Some(true)
645 {
646 return Err(ErrorKind::Session(SessionErrorKind::ElicitationNotSupported).into());
647 }
648 Ok(())
649 }
650}
651
652impl Drop for Session {
653 fn drop(&mut self) {
654 self.shutdown.cancel();
666 self.client.unregister_session(&self.id);
667 self.github_token_registration.lock().take();
668 }
669}
670
671pub struct SessionUi<'a> {
678 session: &'a Session,
679}
680
681impl<'a> SessionUi<'a> {
682 pub async fn elicitation(
690 &self,
691 message: &str,
692 schema: Value,
693 ) -> Result<ElicitationResult, Error> {
694 self.session.assert_elicitation()?;
695 let result = self
696 .session
697 .client
698 .call(
699 "session.ui.elicitation",
700 Some(serde_json::json!({
701 "sessionId": self.session.id,
702 "message": message,
703 "requestedSchema": schema,
704 })),
705 )
706 .await?;
707 let elicitation: ElicitationResult = serde_json::from_value(result)?;
708 Ok(elicitation)
709 }
710
711 pub async fn confirm(&self, message: &str) -> Result<bool, Error> {
715 self.session.assert_elicitation()?;
716 let schema = serde_json::json!({
717 "type": "object",
718 "properties": {
719 "confirmed": {
720 "type": "boolean",
721 "default": true,
722 }
723 },
724 "required": ["confirmed"]
725 });
726 let result = self.elicitation(message, schema).await?;
727 Ok(result.action == "accept"
728 && result
729 .content
730 .and_then(|c| c.get("confirmed").and_then(|v| v.as_bool()))
731 == Some(true))
732 }
733
734 pub async fn select(&self, message: &str, options: &[&str]) -> Result<Option<String>, Error> {
738 self.session.assert_elicitation()?;
739 let schema = serde_json::json!({
740 "type": "object",
741 "properties": {
742 "selection": {
743 "type": "string",
744 "enum": options,
745 }
746 },
747 "required": ["selection"]
748 });
749 let result = self.elicitation(message, schema).await?;
750 if result.action != "accept" {
751 return Ok(None);
752 }
753 let selection = result.content.and_then(|c| {
754 c.get("selection")
755 .and_then(|v| v.as_str())
756 .map(String::from)
757 });
758 Ok(selection)
759 }
760
761 pub async fn input(
766 &self,
767 message: &str,
768 options: Option<&UiInputOptions<'_>>,
769 ) -> Result<Option<String>, Error> {
770 self.session.assert_elicitation()?;
771 let mut field = serde_json::json!({ "type": "string" });
772 if let Some(opts) = options {
773 if let Some(title) = opts.title {
774 field["title"] = Value::String(title.to_string());
775 }
776 if let Some(desc) = opts.description {
777 field["description"] = Value::String(desc.to_string());
778 }
779 if let Some(min) = opts.min_length {
780 field["minLength"] = Value::Number(min.into());
781 }
782 if let Some(max) = opts.max_length {
783 field["maxLength"] = Value::Number(max.into());
784 }
785 if let Some(fmt) = &opts.format {
786 field["format"] = Value::String(fmt.as_str().to_string());
787 }
788 if let Some(default) = opts.default {
789 field["default"] = Value::String(default.to_string());
790 }
791 }
792 let schema = serde_json::json!({
793 "type": "object",
794 "properties": { "value": field },
795 "required": ["value"]
796 });
797 let result = self.elicitation(message, schema).await?;
798 if result.action != "accept" {
799 return Ok(None);
800 }
801 let value = result
802 .content
803 .and_then(|c| c.get("value").and_then(|v| v.as_str()).map(String::from));
804 Ok(value)
805 }
806}
807
808impl Client {
809 pub async fn create_session(&self, mut config: SessionConfig) -> Result<Session, Error> {
831 let total_start = Instant::now();
832 let caller_session_id = config.session_id.clone();
840 let use_server_generated_id = config.cloud.is_some() && caller_session_id.is_none();
841 let local_session_id: Option<SessionId> = if use_server_generated_id {
842 None
843 } else {
844 Some(
845 caller_session_id
846 .clone()
847 .unwrap_or_else(|| SessionId::new(uuid::Uuid::new_v4().to_string())),
848 )
849 };
850 if config.hooks_handler.is_some() && config.hooks.is_none() {
851 config.hooks = Some(true);
852 }
853 if let Some(transforms) = config.system_message_transform.clone() {
854 inject_transform_sections(&mut config, transforms.as_ref());
855 }
856 let mode = self.inner.mode;
857 if mode == crate::ClientMode::Empty && config.available_tools.is_none() {
858 return Err(Error::with_message(
859 ErrorKind::InvalidConfig,
860 "ClientMode::Empty requires available_tools to be set on the session config. \
861 Use ToolSet to specify which tools the session may use (e.g. \
862 ToolSet::new().add_builtin_many(BUILTIN_TOOLS_ISOLATED)).",
863 ));
864 }
865 crate::mode::validate_tool_filter_list(
866 "available_tools",
867 config.available_tools.as_deref(),
868 )?;
869 crate::mode::validate_tool_filter_list("excluded_tools", config.excluded_tools.as_deref())?;
870 config.system_message =
871 crate::mode::system_message_for_mode(mode, config.system_message.take());
872 config.memory = crate::mode::memory_for_mode(mode, config.memory.take());
873 config.enable_experimental_mode =
874 crate::mode::experimental_mode_for_mode(mode, config.enable_experimental_mode);
875 if mode == crate::ClientMode::Empty {
876 if config.enable_session_telemetry.is_none() {
877 config.enable_session_telemetry = Some(false);
878 }
879 if config.skip_embedding_retrieval.is_none() {
880 config.skip_embedding_retrieval = Some(true);
881 }
882 if config.enable_on_demand_instruction_discovery.is_none() {
883 config.enable_on_demand_instruction_discovery = Some(false);
884 }
885 if config.enable_file_hooks.is_none() {
886 config.enable_file_hooks = Some(false);
887 }
888 if config.enable_host_git_operations.is_none() {
889 config.enable_host_git_operations = Some(false);
890 }
891 if config.enable_session_store.is_none() {
892 config.enable_session_store = Some(false);
893 }
894 if config.enable_skills.is_none() {
895 config.enable_skills = Some(false);
896 }
897 }
898 if mode == crate::ClientMode::Empty && config.mcp_oauth_token_storage.is_none() {
899 config.mcp_oauth_token_storage = Some("in-memory".into());
900 }
901 if mode == crate::ClientMode::Empty && config.embedding_cache_storage.is_none() {
902 config.embedding_cache_storage = Some("in-memory".into());
903 }
904 config.custom_agents_local_only =
905 crate::mode::resolve_custom_agents_local_only(mode, config.custom_agents_local_only);
906 let opt_skip_custom_instructions = config.skip_custom_instructions;
907 let opt_custom_agents_local_only = config.custom_agents_local_only;
908 let opt_coauthor_enabled = config.coauthor_enabled;
909 let opt_manage_schedule_enabled = config.manage_schedule_enabled;
910 let opt_included_builtin_skills = config.included_builtin_skills.take();
911 let (mut wire, mut runtime) = config.into_wire(local_session_id.clone())?;
912 wire.enable_github_telemetry_forwarding =
913 self.inner.on_github_telemetry.is_some().then_some(true);
914
915 let permission_handler = crate::permission::resolve_handler(
916 runtime.permission_handler.take(),
917 runtime.permission_policy.take(),
918 );
919 let handlers = SessionHandlers {
920 permission: permission_handler,
921 managed_settings_enabled: has_managed_settings(
922 wire.enable_managed_settings,
923 wire.managed_settings.as_ref(),
924 ),
925 elicitation: runtime.elicitation_handler.take(),
926 mcp_auth: runtime.mcp_auth_handler.take(),
927 user_input: runtime.user_input_handler.take(),
928 exit_plan_mode: runtime.exit_plan_mode_handler.take(),
929 auto_mode_switch: runtime.auto_mode_switch_handler.take(),
930 tools: Arc::new(std::mem::take(&mut runtime.tool_handlers)),
931 };
932 let hooks = runtime.hooks_handler.take();
933 let transforms = runtime.system_message_transform.take();
934 let tools_count = wire.tools.as_ref().map_or(0, Vec::len);
935 let commands_count = runtime.commands.as_ref().map_or(0, Vec::len);
936 let has_hooks = hooks.is_some();
937 let command_handlers = build_command_handler_map(runtime.commands.as_deref());
938 let canvas_handler = runtime.canvas_handler.take();
939 let session_fs_provider = runtime.session_fs_provider.take();
940 let bearer_token_providers = std::mem::take(&mut runtime.bearer_token_providers);
941 let github_token_registration = runtime
942 .github_token_provider
943 .take()
944 .map(|provider| self.register_github_token_provider(provider));
945 wire.github_token_provider_registration_id = github_token_registration
946 .as_ref()
947 .map(|registration| registration.id().to_string());
948 let has_mcp_auth_handler = handlers.mcp_auth.is_some();
949 if self.inner.session_fs_configured && session_fs_provider.is_none() {
950 return Err(ErrorKind::Session(SessionErrorKind::SessionFsProviderRequired).into());
951 }
952 if self.inner.session_fs_sqlite_declared
953 && let Some(ref provider) = session_fs_provider
954 && provider.sqlite().is_none()
955 {
956 return Err(Error::with_message(
957 ErrorKind::InvalidConfig,
958 "SessionFs capabilities declare SQLite support but the provider \
959 does not implement SessionFsSqliteProvider",
960 ));
961 }
962
963 let mut params = serde_json::to_value(&wire)?;
964 let trace_ctx = self.resolve_trace_context().await;
965 inject_trace_context(&mut params, &trace_ctx);
966
967 let setup_start = Instant::now();
968 let capabilities = Arc::new(parking_lot::RwLock::new(SessionCapabilities::default()));
969 let idle_waiter = Arc::new(ParkingLotMutex::new(None));
970 let open_canvases = Arc::new(parking_lot::RwLock::new(Vec::new()));
971 let shutdown = CancellationToken::new();
972 let (event_tx, _) = tokio::sync::broadcast::channel(512);
973
974 let inline_stash: Arc<
980 ParkingLotMutex<Option<(SessionId, crate::router::SessionChannels)>>,
981 > = Arc::new(ParkingLotMutex::new(None));
982
983 let inline_callback: Option<crate::jsonrpc::InlineResponseCallback> = if let Some(ref sid) =
984 local_session_id
985 {
986 let channels = self.register_session(sid);
987 *inline_stash.lock() = Some((sid.clone(), channels));
988 None
989 } else {
990 let client = self.clone();
991 let stash = inline_stash.clone();
992 let expected = caller_session_id.clone();
993 Some(Box::new(move |response| {
994 let result = response.result.as_ref().ok_or_else(|| {
995 Error::with_message(ErrorKind::Json, "session.create response had no result")
996 })?;
997 let parsed: CreateSessionResult =
998 serde_json::from_value(result.clone()).map_err(Error::from)?;
999 if let Some(requested) = expected.as_ref()
1000 && parsed.session_id != *requested
1001 {
1002 return Err(ErrorKind::Session(SessionErrorKind::SessionIdMismatch {
1003 requested: requested.clone(),
1004 returned: parsed.session_id,
1005 })
1006 .into());
1007 }
1008 let channels = client.register_session(&parsed.session_id);
1009 *stash.lock() = Some((parsed.session_id, channels));
1010 Ok(())
1011 }))
1012 };
1013
1014 let rpc_start = Instant::now();
1015 let result = match self
1016 .call_with_inline_callback("session.create", Some(params), inline_callback)
1017 .await
1018 {
1019 Ok(result) => result,
1020 Err(error) => {
1021 if let Some((id, _channels)) = inline_stash.lock().take() {
1022 self.unregister_session(&id);
1023 }
1024 return Err(error);
1025 }
1026 };
1027 tracing::debug!(
1028 elapsed_ms = rpc_start.elapsed().as_millis(),
1029 "Client::create_session session creation request completed successfully"
1030 );
1031 let create_result: CreateSessionResult = match serde_json::from_value(result) {
1032 Ok(result) => result,
1033 Err(error) => {
1034 if let Some((id, _channels)) = inline_stash.lock().take() {
1035 self.unregister_session(&id);
1036 }
1037 return Err(error.into());
1038 }
1039 };
1040
1041 if let Some(ref requested) = local_session_id
1042 && create_result.session_id != *requested
1043 {
1044 if let Some((id, _channels)) = inline_stash.lock().take() {
1045 self.unregister_session(&id);
1046 }
1047 return Err(ErrorKind::Session(SessionErrorKind::SessionIdMismatch {
1048 requested: requested.clone(),
1049 returned: create_result.session_id.clone(),
1050 })
1051 .into());
1052 }
1053
1054 let (session_id, channels) = inline_stash
1055 .lock()
1056 .take()
1057 .expect("session registration must have populated stash on success");
1058 let event_loop = spawn_event_loop(
1059 session_id.clone(),
1060 self.clone(),
1061 handlers,
1062 hooks,
1063 transforms,
1064 command_handlers,
1065 canvas_handler,
1066 session_fs_provider,
1067 bearer_token_providers,
1068 channels,
1069 idle_waiter.clone(),
1070 capabilities.clone(),
1071 open_canvases.clone(),
1072 event_tx.clone(),
1073 shutdown.clone(),
1074 );
1075 tracing::debug!(
1076 elapsed_ms = setup_start.elapsed().as_millis(),
1077 session_id = %session_id,
1078 tools_count,
1079 commands_count,
1080 has_hooks,
1081 "Client::create_session local setup complete"
1082 );
1083 *capabilities.write() = create_result.capabilities.unwrap_or_default();
1084 if has_mcp_auth_handler {
1085 register_mcp_auth_interest(self, &session_id).await?;
1086 }
1087
1088 tracing::debug!(
1089 elapsed_ms = total_start.elapsed().as_millis(),
1090 session_id = %session_id,
1091 "Client::create_session complete"
1092 );
1093 let session = Session {
1094 id: session_id,
1095 cwd: self.cwd().clone(),
1096 workspace_path: create_result.workspace_path,
1097 remote_url: create_result.remote_url,
1098 client: self.clone(),
1099 event_loop: ParkingLotMutex::new(Some(event_loop)),
1100 shutdown,
1101 idle_waiter,
1102 capabilities,
1103 open_canvases,
1104 event_tx,
1105 github_token_registration: ParkingLotMutex::new(github_token_registration),
1106 };
1107 apply_mode_post_create_patch(
1108 &session,
1109 mode,
1110 opt_skip_custom_instructions,
1111 opt_custom_agents_local_only,
1112 opt_coauthor_enabled,
1113 opt_manage_schedule_enabled,
1114 opt_included_builtin_skills,
1115 )
1116 .await?;
1117 if let Some(registration) = session.github_token_registration.lock().as_ref() {
1118 registration.claim(session.id.clone());
1119 } else {
1120 self.retire_github_token_provider(&session.id);
1121 }
1122 Ok(session)
1123 }
1124
1125 pub async fn resume_session(&self, mut config: ResumeSessionConfig) -> Result<Session, Error> {
1136 let total_start = Instant::now();
1137 let session_id = config.session_id.clone();
1138 if config.hooks_handler.is_some() && config.hooks.is_none() {
1139 config.hooks = Some(true);
1140 }
1141 if let Some(transforms) = config.system_message_transform.clone() {
1142 inject_transform_sections_resume(&mut config, transforms.as_ref());
1143 }
1144 let mode = self.inner.mode;
1145 if mode == crate::ClientMode::Empty && config.available_tools.is_none() {
1146 return Err(Error::with_message(
1147 ErrorKind::InvalidConfig,
1148 "ClientMode::Empty requires available_tools to be set on the session config. \
1149 Use ToolSet to specify which tools the session may use (e.g. \
1150 ToolSet::new().add_builtin_many(BUILTIN_TOOLS_ISOLATED)).",
1151 ));
1152 }
1153 crate::mode::validate_tool_filter_list(
1154 "available_tools",
1155 config.available_tools.as_deref(),
1156 )?;
1157 crate::mode::validate_tool_filter_list("excluded_tools", config.excluded_tools.as_deref())?;
1158 config.system_message =
1159 crate::mode::system_message_for_mode(mode, config.system_message.take());
1160 config.memory = crate::mode::memory_for_mode(mode, config.memory.take());
1161 config.enable_experimental_mode =
1162 crate::mode::experimental_mode_for_mode(mode, config.enable_experimental_mode);
1163 if mode == crate::ClientMode::Empty {
1164 if config.enable_session_telemetry.is_none() {
1165 config.enable_session_telemetry = Some(false);
1166 }
1167 if config.skip_embedding_retrieval.is_none() {
1168 config.skip_embedding_retrieval = Some(true);
1169 }
1170 if config.enable_on_demand_instruction_discovery.is_none() {
1171 config.enable_on_demand_instruction_discovery = Some(false);
1172 }
1173 if config.enable_file_hooks.is_none() {
1174 config.enable_file_hooks = Some(false);
1175 }
1176 if config.enable_host_git_operations.is_none() {
1177 config.enable_host_git_operations = Some(false);
1178 }
1179 if config.enable_session_store.is_none() {
1180 config.enable_session_store = Some(false);
1181 }
1182 if config.enable_skills.is_none() {
1183 config.enable_skills = Some(false);
1184 }
1185 }
1186 if mode == crate::ClientMode::Empty && config.mcp_oauth_token_storage.is_none() {
1187 config.mcp_oauth_token_storage = Some("in-memory".into());
1188 }
1189 if mode == crate::ClientMode::Empty && config.embedding_cache_storage.is_none() {
1190 config.embedding_cache_storage = Some("in-memory".into());
1191 }
1192 config.custom_agents_local_only =
1193 crate::mode::resolve_custom_agents_local_only(mode, config.custom_agents_local_only);
1194 let opt_skip_custom_instructions = config.skip_custom_instructions;
1195 let opt_custom_agents_local_only = config.custom_agents_local_only;
1196 let opt_coauthor_enabled = config.coauthor_enabled;
1197 let opt_manage_schedule_enabled = config.manage_schedule_enabled;
1198 let opt_included_builtin_skills = config.included_builtin_skills.take();
1199 let (mut wire, mut runtime) = config.into_wire()?;
1200 wire.enable_github_telemetry_forwarding =
1201 self.inner.on_github_telemetry.is_some().then_some(true);
1202
1203 let permission_handler = crate::permission::resolve_handler(
1204 runtime.permission_handler.take(),
1205 runtime.permission_policy.take(),
1206 );
1207 let handlers = SessionHandlers {
1208 permission: permission_handler,
1209 managed_settings_enabled: has_managed_settings(
1210 wire.enable_managed_settings,
1211 wire.managed_settings.as_ref(),
1212 ),
1213 elicitation: runtime.elicitation_handler.take(),
1214 mcp_auth: runtime.mcp_auth_handler.take(),
1215 user_input: runtime.user_input_handler.take(),
1216 exit_plan_mode: runtime.exit_plan_mode_handler.take(),
1217 auto_mode_switch: runtime.auto_mode_switch_handler.take(),
1218 tools: Arc::new(std::mem::take(&mut runtime.tool_handlers)),
1219 };
1220 let hooks = runtime.hooks_handler.take();
1221 let transforms = runtime.system_message_transform.take();
1222 let tools_count = wire.tools.as_ref().map_or(0, Vec::len);
1223 let commands_count = runtime.commands.as_ref().map_or(0, Vec::len);
1224 let has_hooks = hooks.is_some();
1225 let command_handlers = build_command_handler_map(runtime.commands.as_deref());
1226 let canvas_handler = runtime.canvas_handler.take();
1227 let session_fs_provider = runtime.session_fs_provider.take();
1228 let bearer_token_providers = std::mem::take(&mut runtime.bearer_token_providers);
1229 let github_token_registration = runtime
1230 .github_token_provider
1231 .take()
1232 .map(|provider| self.register_github_token_provider(provider));
1233 wire.github_token_provider_registration_id = github_token_registration
1234 .as_ref()
1235 .map(|registration| registration.id().to_string());
1236 let has_mcp_auth_handler = handlers.mcp_auth.is_some();
1237 if self.inner.session_fs_configured && session_fs_provider.is_none() {
1238 return Err(ErrorKind::Session(SessionErrorKind::SessionFsProviderRequired).into());
1239 }
1240 if self.inner.session_fs_sqlite_declared
1241 && let Some(ref provider) = session_fs_provider
1242 && provider.sqlite().is_none()
1243 {
1244 return Err(Error::with_message(
1245 ErrorKind::InvalidConfig,
1246 "SessionFs capabilities declare SQLite support but the provider \
1247 does not implement SessionFsSqliteProvider",
1248 ));
1249 }
1250
1251 let mut params = serde_json::to_value(&wire)?;
1252 let trace_ctx = self.resolve_trace_context().await;
1253 inject_trace_context(&mut params, &trace_ctx);
1254
1255 let capabilities = Arc::new(parking_lot::RwLock::new(SessionCapabilities::default()));
1256 let setup_start = Instant::now();
1257 let channels = self.register_session(&session_id);
1258 let idle_waiter = Arc::new(ParkingLotMutex::new(None));
1259 let open_canvases = Arc::new(parking_lot::RwLock::new(Vec::new()));
1260 let shutdown = CancellationToken::new();
1261 let (event_tx, _) = tokio::sync::broadcast::channel(512);
1262 let event_loop = spawn_event_loop(
1263 session_id.clone(),
1264 self.clone(),
1265 handlers,
1266 hooks,
1267 transforms,
1268 command_handlers,
1269 canvas_handler,
1270 session_fs_provider,
1271 bearer_token_providers,
1272 channels,
1273 idle_waiter.clone(),
1274 capabilities.clone(),
1275 open_canvases.clone(),
1276 event_tx.clone(),
1277 shutdown.clone(),
1278 );
1279 let mut registration =
1280 PendingSessionRegistration::new(self.clone(), session_id.clone(), shutdown.clone());
1281 tracing::debug!(
1282 elapsed_ms = setup_start.elapsed().as_millis(),
1283 session_id = %session_id,
1284 tools_count,
1285 commands_count,
1286 has_hooks,
1287 "Client::resume_session local setup complete"
1288 );
1289
1290 let rpc_start = Instant::now();
1291 let result = match self.call("session.resume", Some(params)).await {
1292 Ok(result) => result,
1293 Err(error) => {
1294 registration.cleanup(event_loop).await;
1295 return Err(error);
1296 }
1297 };
1298 tracing::debug!(
1299 elapsed_ms = rpc_start.elapsed().as_millis(),
1300 session_id = %session_id,
1301 "Client::resume_session session resume request completed successfully"
1302 );
1303
1304 let resume_result: ResumeSessionResult = match serde_json::from_value(result) {
1305 Ok(result) => result,
1306 Err(error) => {
1307 registration.cleanup(event_loop).await;
1308 return Err(error.into());
1309 }
1310 };
1311 let cli_session_id = resume_result
1312 .session_id
1313 .clone()
1314 .unwrap_or_else(|| session_id.clone());
1315 if cli_session_id != session_id {
1316 registration.cleanup(event_loop).await;
1317 return Err(ErrorKind::Session(SessionErrorKind::SessionIdMismatch {
1318 requested: session_id,
1319 returned: cli_session_id,
1320 })
1321 .into());
1322 }
1323 if has_mcp_auth_handler {
1324 register_mcp_auth_interest(self, &session_id).await?;
1325 }
1326
1327 let skills_reload_start = Instant::now();
1329 if let Err(e) = self
1330 .call(
1331 "session.skills.reload",
1332 Some(serde_json::json!({ "sessionId": session_id })),
1333 )
1334 .await
1335 {
1336 warn!(
1337 elapsed_ms = skills_reload_start.elapsed().as_millis(),
1338 session_id = %session_id,
1339 error = %e,
1340 "Client::resume_session skills reload request failed"
1341 );
1342 } else {
1343 tracing::debug!(
1344 elapsed_ms = skills_reload_start.elapsed().as_millis(),
1345 session_id = %session_id,
1346 "Client::resume_session skills reload request completed successfully"
1347 );
1348 }
1349
1350 *capabilities.write() = resume_result.capabilities.unwrap_or_default();
1351 {
1356 let mut snapshots = open_canvases.write();
1357 for snapshot in resume_result.open_canvases.unwrap_or_default() {
1358 upsert_open_canvas_snapshot(&mut snapshots, snapshot);
1359 }
1360 }
1361
1362 tracing::debug!(
1363 elapsed_ms = total_start.elapsed().as_millis(),
1364 session_id = %session_id,
1365 "Client::resume_session complete"
1366 );
1367 registration.disarm();
1368 let session = Session {
1369 id: session_id,
1370 cwd: self.cwd().clone(),
1371 workspace_path: resume_result.workspace_path,
1372 remote_url: resume_result.remote_url,
1373 client: self.clone(),
1374 event_loop: ParkingLotMutex::new(Some(event_loop)),
1375 shutdown,
1376 idle_waiter,
1377 capabilities,
1378 open_canvases,
1379 event_tx,
1380 github_token_registration: ParkingLotMutex::new(github_token_registration),
1381 };
1382 apply_mode_post_create_patch(
1383 &session,
1384 mode,
1385 opt_skip_custom_instructions,
1386 opt_custom_agents_local_only,
1387 opt_coauthor_enabled,
1388 opt_manage_schedule_enabled,
1389 opt_included_builtin_skills,
1390 )
1391 .await?;
1392 if let Some(registration) = session.github_token_registration.lock().as_ref() {
1393 registration.claim(session.id.clone());
1394 } else {
1395 self.retire_github_token_provider(&session.id);
1396 }
1397 Ok(session)
1398 }
1399}
1400
1401type CommandHandlerMap = HashMap<String, Arc<dyn CommandHandler>>;
1402
1403async fn apply_mode_post_create_patch(
1404 session: &Session,
1405 mode: crate::ClientMode,
1406 opt_skip_custom_instructions: Option<bool>,
1407 opt_custom_agents_local_only: Option<bool>,
1408 opt_coauthor_enabled: Option<bool>,
1409 opt_manage_schedule_enabled: Option<bool>,
1410 opt_included_builtin_skills: Option<Vec<String>>,
1411) -> Result<(), Error> {
1412 let Some(patch) = build_mode_post_create_patch(
1413 mode,
1414 opt_skip_custom_instructions,
1415 opt_custom_agents_local_only,
1416 opt_coauthor_enabled,
1417 opt_manage_schedule_enabled,
1418 opt_included_builtin_skills,
1419 ) else {
1420 return Ok(());
1421 };
1422 if let Err(error) = session.rpc().options().update(patch).await {
1423 let _ = session.disconnect().await;
1424 return Err(error);
1425 }
1426 Ok(())
1427}
1428
1429fn build_mode_post_create_patch(
1439 mode: crate::ClientMode,
1440 opt_skip_custom_instructions: Option<bool>,
1441 opt_custom_agents_local_only: Option<bool>,
1442 opt_coauthor_enabled: Option<bool>,
1443 opt_manage_schedule_enabled: Option<bool>,
1444 opt_included_builtin_skills: Option<Vec<String>>,
1445) -> Option<crate::generated::api_types::SessionUpdateOptionsParams> {
1446 use crate::generated::api_types::SessionUpdateOptionsParams;
1447 let mut patch = SessionUpdateOptionsParams::default();
1448 let should_send = if mode == crate::ClientMode::Empty {
1449 patch.skip_custom_instructions = Some(opt_skip_custom_instructions.unwrap_or(true));
1450 patch.custom_agents_local_only = Some(opt_custom_agents_local_only.unwrap_or(true));
1451 patch.coauthor_enabled = Some(opt_coauthor_enabled.unwrap_or(false));
1452 patch.manage_schedule_enabled = Some(opt_manage_schedule_enabled.unwrap_or(false));
1453 patch.installed_plugins = Some(Vec::new());
1454 patch.included_builtin_skills = Some(opt_included_builtin_skills.unwrap_or_default());
1455 true
1456 } else {
1457 let mut any = false;
1458 if let Some(v) = opt_skip_custom_instructions {
1459 patch.skip_custom_instructions = Some(v);
1460 any = true;
1461 }
1462 if let Some(v) = opt_custom_agents_local_only {
1463 patch.custom_agents_local_only = Some(v);
1464 any = true;
1465 }
1466 if let Some(v) = opt_coauthor_enabled {
1467 patch.coauthor_enabled = Some(v);
1468 any = true;
1469 }
1470 if let Some(v) = opt_manage_schedule_enabled {
1471 patch.manage_schedule_enabled = Some(v);
1472 any = true;
1473 }
1474 if let Some(v) = opt_included_builtin_skills {
1475 patch.included_builtin_skills = Some(v);
1476 any = true;
1477 }
1478 any
1479 };
1480 if !should_send {
1481 return None;
1482 }
1483 Some(patch)
1484}
1485
1486fn build_command_handler_map(commands: Option<&[CommandDefinition]>) -> Arc<CommandHandlerMap> {
1487 let map = match commands {
1488 Some(commands) => commands
1489 .iter()
1490 .filter(|cmd| !cmd.name.is_empty())
1491 .map(|cmd| (cmd.name.clone(), cmd.handler.clone()))
1492 .collect(),
1493 None => HashMap::new(),
1494 };
1495 Arc::new(map)
1496}
1497
1498fn upsert_open_canvas_snapshot(
1499 snapshots: &mut Vec<OpenCanvasInstance>,
1500 snapshot: OpenCanvasInstance,
1501) {
1502 if let Some(existing) = snapshots
1503 .iter_mut()
1504 .find(|open| open.instance_id == snapshot.instance_id)
1505 {
1506 *existing = snapshot;
1507 } else {
1508 snapshots.push(snapshot);
1509 }
1510}
1511
1512fn remove_open_canvas_snapshot(snapshots: &mut Vec<OpenCanvasInstance>, instance_id: &str) {
1513 snapshots.retain(|open| open.instance_id != instance_id);
1514}
1515
1516#[allow(clippy::too_many_arguments)]
1517fn spawn_event_loop(
1518 session_id: SessionId,
1519 client: Client,
1520 handlers: SessionHandlers,
1521 hooks: Option<Arc<dyn SessionHooks>>,
1522 transforms: Option<Arc<dyn SystemMessageTransform>>,
1523 command_handlers: Arc<CommandHandlerMap>,
1524 canvas_handler: Option<Arc<dyn CanvasHandler>>,
1525 session_fs_provider: Option<Arc<dyn SessionFsProvider>>,
1526 bearer_token_providers: HashMap<String, Arc<dyn BearerTokenProvider>>,
1527 channels: crate::router::SessionChannels,
1528 idle_waiter: Arc<ParkingLotMutex<Option<IdleWaiter>>>,
1529 capabilities: Arc<parking_lot::RwLock<SessionCapabilities>>,
1530 open_canvases: Arc<parking_lot::RwLock<Vec<OpenCanvasInstance>>>,
1531 event_tx: tokio::sync::broadcast::Sender<SessionEvent>,
1532 shutdown: CancellationToken,
1533) -> JoinHandle<()> {
1534 let crate::router::SessionChannels {
1535 mut notifications,
1536 mut requests,
1537 } = channels;
1538
1539 let span = tracing::error_span!("session_event_loop", session_id = %session_id);
1540 tokio::spawn(
1541 async move {
1542 loop {
1543 tokio::select! {
1567 _ = shutdown.cancelled() => break,
1568 Some(notification) = notifications.recv() => {
1569 handle_notification(
1570 &session_id, &client, &handlers, &command_handlers, notification, &idle_waiter, &capabilities, &open_canvases, &event_tx, &shutdown,
1571 ).await;
1572 }
1573 Some(request) = requests.recv() => {
1574 let span = tracing::error_span!("session_request_handler", session_id = %session_id);
1578 let session_id = session_id.clone();
1579 let client = client.clone();
1580 let handlers = handlers.clone();
1581 let hooks = hooks.clone();
1582 let transforms = transforms.clone();
1583 let canvas_handler = canvas_handler.clone();
1584 let session_fs_provider = session_fs_provider.clone();
1585 let bearer_token_providers = bearer_token_providers.clone();
1586 tokio::spawn(
1587 async move {
1588 let ctx = RequestDispatchContext {
1589 client: &client,
1590 handlers: &handlers,
1591 hooks: hooks.as_deref(),
1592 transforms: transforms.as_deref(),
1593 canvas_handler: canvas_handler.as_ref(),
1594 session_fs_provider: session_fs_provider.as_ref(),
1595 bearer_token_providers: &bearer_token_providers,
1596 };
1597 handle_request(&session_id, ctx, request).await;
1598 }
1599 .instrument(span),
1600 );
1601 }
1602 else => break,
1603 }
1604 }
1605 if let Some(waiter) = idle_waiter.lock().take() {
1608 let _ = waiter
1609 .tx
1610 .send(Err(ErrorKind::Session(SessionErrorKind::EventLoopClosed).into()));
1611 }
1612 }
1613 .instrument(span),
1614 )
1615}
1616
1617fn extract_request_id(data: &Value) -> Option<RequestId> {
1618 data.get("requestId")
1619 .and_then(|v| v.as_str())
1620 .filter(|s| !s.is_empty())
1621 .map(RequestId::new)
1622}
1623
1624fn permission_request_data(
1625 event_data: &Value,
1626 managed_settings_enabled: bool,
1627) -> PermissionRequestData {
1628 let request_data = event_data
1629 .get("permissionRequest")
1630 .cloned()
1631 .unwrap_or_else(|| event_data.clone());
1632 let managed_approval_required = match request_data.get("managedApprovalRequired") {
1633 None => None,
1634 Some(Value::Bool(value)) => Some(*value),
1635 Some(_) => Some(true),
1636 };
1637 match serde_json::from_value::<PermissionRequestData>(request_data) {
1638 Ok(mut data) => {
1639 data.extra = event_data.clone();
1640 data.managed_settings_enabled = managed_settings_enabled;
1641 data
1642 }
1643 Err(_) => PermissionRequestData {
1644 kind: None,
1645 tool_call_id: None,
1646 managed_approval_required,
1647 managed_settings_enabled,
1648 extra: event_data.clone(),
1649 },
1650 }
1651}
1652
1653fn permission_response_params(
1661 session_id: &SessionId,
1662 request_id: &RequestId,
1663 result: &PermissionResult,
1664) -> Option<Value> {
1665 let (decision, decision_context) = match result {
1666 PermissionResult::Decision { decision, context } => (decision, context.clone()),
1667 PermissionResult::NoResult => return None,
1668 };
1669 let mut params = serde_json::to_value(PermissionDecisionRequest {
1670 decision_context,
1671 request_id: request_id.clone(),
1672 result: decision.clone(),
1673 })
1674 .expect("serializing permission response should succeed");
1675 params["sessionId"] =
1676 serde_json::to_value(session_id).expect("serializing session ID should succeed");
1677 Some(params)
1678}
1679
1680async fn register_mcp_auth_interest(client: &Client, session_id: &SessionId) -> Result<(), Error> {
1681 let mut params = serde_json::to_value(RegisterEventInterestParams {
1682 event_type: "mcp.oauth_required".to_string(),
1683 })?;
1684 params["sessionId"] = Value::String(session_id.to_string());
1685 client
1686 .call(rpc_methods::SESSION_EVENTLOG_REGISTERINTEREST, Some(params))
1687 .await?;
1688 Ok(())
1689}
1690
1691fn tool_failure_result(message: impl Into<String>) -> ToolResult {
1692 let message = message.into();
1693 ToolResult::Expanded(ToolResultExpanded {
1694 text_result_for_llm: message.clone(),
1695 result_type: "failure".to_string(),
1696 binary_results_for_llm: None,
1697 session_log: None,
1698 error: Some(message),
1699 tool_telemetry: None,
1700 tool_references: None,
1701 })
1702}
1703
1704fn is_autopilot_continuation_idle(event: &SessionEvent) -> bool {
1705 event
1706 .typed_data::<SessionIdleData>()
1707 .is_some_and(|data| data.mode == Some(SessionMode::Autopilot))
1708}
1709
1710#[allow(clippy::too_many_arguments)]
1712async fn handle_notification(
1713 session_id: &SessionId,
1714 client: &Client,
1715 handlers: &SessionHandlers,
1716 command_handlers: &Arc<CommandHandlerMap>,
1717 notification: SessionEventNotification,
1718 idle_waiter: &Arc<ParkingLotMutex<Option<IdleWaiter>>>,
1719 capabilities: &Arc<parking_lot::RwLock<SessionCapabilities>>,
1720 open_canvases: &Arc<parking_lot::RwLock<Vec<OpenCanvasInstance>>>,
1721 event_tx: &tokio::sync::broadcast::Sender<SessionEvent>,
1722 shutdown: &CancellationToken,
1723) {
1724 let dispatch_start = Instant::now();
1725 let event = notification.event.clone();
1726 let event_type = event.parsed_type();
1727 if event_type == SessionEventType::PermissionRequested {
1728 tracing::debug!(
1729 session_id = %session_id,
1730 event_type = %event.event_type,
1731 "Session::handle_notification permission request received"
1732 );
1733 }
1734
1735 match event_type {
1738 SessionEventType::AssistantMessage
1739 | SessionEventType::SessionIdle
1740 | SessionEventType::SessionError => {
1741 let mut guard = idle_waiter.lock();
1742 if let Some(waiter) = guard.as_mut() {
1743 match event_type {
1744 SessionEventType::AssistantMessage => {
1745 if !waiter.first_assistant_message_seen {
1746 waiter.first_assistant_message_seen = true;
1747 tracing::debug!(
1748 elapsed_ms = waiter.started_at.elapsed().as_millis(),
1749 session_id = %session_id,
1750 "Session::send_and_wait first assistant message"
1751 );
1752 }
1753 waiter.last_assistant_message = Some(event.clone());
1754 }
1755 SessionEventType::SessionIdle if is_autopilot_continuation_idle(&event) => {}
1756 SessionEventType::SessionIdle | SessionEventType::SessionError => {
1757 if let Some(waiter) = guard.take() {
1758 if event_type == SessionEventType::SessionIdle {
1759 tracing::debug!(
1760 elapsed_ms = waiter.started_at.elapsed().as_millis(),
1761 session_id = %session_id,
1762 "Session::send_and_wait idle received"
1763 );
1764 let _ = waiter.tx.send(Ok(waiter.last_assistant_message));
1765 } else {
1766 let error_msg = event
1767 .typed_data::<SessionErrorData>()
1768 .map(|d| d.message)
1769 .or_else(|| {
1770 event
1771 .data
1772 .get("message")
1773 .and_then(|v| v.as_str())
1774 .map(|s| s.to_string())
1775 })
1776 .unwrap_or_else(|| "session error".to_string());
1777 let _ = waiter.tx.send(Err(Error::with_message(
1778 ErrorKind::Session(SessionErrorKind::AgentError),
1779 error_msg,
1780 )));
1781 }
1782 }
1783 }
1784 _ => {}
1785 }
1786 }
1787 }
1788 _ => {}
1789 }
1790
1791 if event_type == SessionEventType::CapabilitiesChanged {
1795 match serde_json::from_value::<SessionCapabilities>(notification.event.data.clone()) {
1796 Ok(changed) => *capabilities.write() = changed,
1797 Err(e) => warn!(error = %e, "failed to deserialize capabilities.changed payload"),
1798 }
1799 }
1800 if event_type == SessionEventType::SessionCanvasOpened {
1801 match serde_json::from_value::<OpenCanvasInstance>(notification.event.data.clone()) {
1802 Ok(open_canvas) => {
1803 upsert_open_canvas_snapshot(&mut open_canvases.write(), open_canvas);
1804 }
1805 Err(e) => warn!(error = %e, "failed to deserialize session.canvas.opened payload"),
1806 }
1807 }
1808 if event_type == SessionEventType::SessionCanvasClosed {
1809 match serde_json::from_value::<SessionCanvasClosedData>(notification.event.data.clone()) {
1810 Ok(closed) => {
1811 if closed.instance_id.is_empty() {
1812 warn!("failed to deserialize session.canvas.closed payload");
1813 } else {
1814 remove_open_canvas_snapshot(&mut open_canvases.write(), &closed.instance_id);
1815 }
1816 }
1817 Err(e) => warn!(error = %e, "failed to deserialize session.canvas.closed payload"),
1818 }
1819 }
1820
1821 let _ = event_tx.send(event.clone());
1825
1826 tracing::debug!(
1827 elapsed_ms = dispatch_start.elapsed().as_millis(),
1828 session_id = %session_id,
1829 event_type = %notification.event.event_type,
1830 "Session::handle_notification dispatch"
1831 );
1832
1833 match event_type {
1836 SessionEventType::PermissionRequested => {
1837 let Some(request_id) = extract_request_id(¬ification.event.data) else {
1838 return;
1839 };
1840 if notification
1844 .event
1845 .data
1846 .get("resolvedByHook")
1847 .and_then(|v| v.as_bool())
1848 .unwrap_or(false)
1849 {
1850 return;
1851 }
1852 let Some(permission_handler) = handlers.permission.clone() else {
1856 return;
1857 };
1858 let client = client.clone();
1859 let sid = session_id.clone();
1860 let shutdown = shutdown.clone();
1861 let data = permission_request_data(
1862 ¬ification.event.data,
1863 handlers.managed_settings_enabled,
1864 );
1865 let span = tracing::error_span!(
1866 "permission_request_handler",
1867 session_id = %sid,
1868 request_id = %request_id
1869 );
1870 tokio::spawn(
1871 async move {
1872 let handler_start = Instant::now();
1873 let result = permission_handler
1874 .handle(sid.clone(), request_id.clone(), data)
1875 .await;
1876 tracing::debug!(
1877 elapsed_ms = handler_start.elapsed().as_millis(),
1878 session_id = %sid,
1879 request_id = %request_id,
1880 "PermissionHandler::handle dispatch"
1881 );
1882 let Some(params) = permission_response_params(&sid, &request_id, &result)
1883 else {
1884 return;
1888 };
1889 let rpc_start = Instant::now();
1890 let method =
1891 rpc_methods::SESSION_PERMISSIONS_HANDLEPENDINGPERMISSIONREQUEST;
1892 tokio::select! {
1893 biased;
1894 response = client.call(method, Some(params)) => {
1895 match response {
1896 Ok(_) => tracing::debug!(
1897 elapsed_ms = rpc_start.elapsed().as_millis(),
1898 session_id = %sid,
1899 request_id = %request_id,
1900 method,
1901 "Session::handle_notification response sent successfully"
1902 ),
1903 Err(error) => warn!(
1904 error = %error,
1905 session_id = %sid,
1906 request_id = %request_id,
1907 method,
1908 "failed to deliver permission decision back to the runtime"
1909 ),
1910 }
1911 }
1912 _ = shutdown.cancelled() => {
1913 warn!(
1914 elapsed_ms = rpc_start.elapsed().as_millis(),
1915 session_id = %sid,
1916 request_id = %request_id,
1917 method,
1918 delivery_outcome = "unknown",
1919 "permission confirmation acknowledgement wait cancelled during session shutdown"
1920 );
1921 }
1922 }
1923 }
1924 .instrument(span),
1925 );
1926 }
1927 SessionEventType::ExternalToolRequested => {
1928 let Some(request_id) = extract_request_id(¬ification.event.data) else {
1929 return;
1930 };
1931 let data: ExternalToolRequestedData =
1932 match serde_json::from_value(notification.event.data.clone()) {
1933 Ok(d) => d,
1934 Err(e) => {
1935 warn!(error = %e, "failed to deserialize external_tool.requested");
1936 let client = client.clone();
1937 let sid = session_id.clone();
1938 let span = tracing::error_span!(
1939 "external_tool_deserialize_error",
1940 session_id = %sid,
1941 request_id = %request_id
1942 );
1943 tokio::spawn(
1944 async move {
1945 let rpc_start = Instant::now();
1946 let _ = client
1947 .call(
1948 "session.tools.handlePendingToolCall",
1949 Some(serde_json::json!({
1950 "sessionId": sid,
1951 "requestId": request_id,
1952 "error": format!("Failed to deserialize tool request: {e}"),
1953 })),
1954 )
1955 .await;
1956 tracing::debug!(
1957 elapsed_ms = rpc_start.elapsed().as_millis(),
1958 session_id = %sid,
1959 request_id = %request_id,
1960 "Session::handle_notification response sent successfully"
1961 );
1962 }
1963 .instrument(span),
1964 );
1965 return;
1966 }
1967 };
1968 let tool_handler = if data.tool_name.is_empty() {
1972 None
1973 } else {
1974 handlers.tools.get(&data.tool_name).cloned()
1975 };
1976 let Some(tool_handler) = tool_handler else {
1977 return;
1978 };
1979 let client = client.clone();
1980 let sid = session_id.clone();
1981 let span = tracing::error_span!(
1982 "external_tool_handler",
1983 session_id = %sid,
1984 request_id = %request_id
1985 );
1986 tokio::spawn(
1987 async move {
1988 if data.tool_call_id.is_empty() {
1993 let error_msg = "Missing toolCallId";
1994 let rpc_start = Instant::now();
1995 let _ = client
1996 .call(
1997 "session.tools.handlePendingToolCall",
1998 Some(serde_json::json!({
1999 "sessionId": sid,
2000 "requestId": request_id,
2001 "error": error_msg,
2002 })),
2003 )
2004 .await;
2005 tracing::debug!(
2006 elapsed_ms = rpc_start.elapsed().as_millis(),
2007 session_id = %sid,
2008 request_id = %request_id,
2009 "Session::handle_notification response sent successfully"
2010 );
2011 return;
2012 }
2013 let tool_call_id = data.tool_call_id.clone();
2014 let tool_name = data.tool_name.clone();
2015 let available_tools = if tool_name == TOOL_SEARCH_TOOL_NAME {
2022 match client
2023 .call(
2024 rpc_methods::SESSION_TOOLS_GETCURRENTMETADATA,
2025 Some(serde_json::json!({ "sessionId": sid })),
2026 )
2027 .await
2028 {
2029 Ok(value) => {
2030 serde_json::from_value::<ToolsGetCurrentMetadataResult>(value)
2031 .ok()
2032 .and_then(|result| result.tools)
2033 }
2034 Err(_) => None,
2035 }
2036 } else {
2037 None
2038 };
2039 let invocation = ToolInvocation {
2040 session_id: sid.clone(),
2041 tool_call_id: data.tool_call_id,
2042 tool_name: data.tool_name,
2043 arguments: data
2044 .arguments
2045 .unwrap_or(Value::Object(serde_json::Map::new())),
2046 available_tools,
2047 traceparent: data.traceparent,
2048 tracestate: data.tracestate,
2049 };
2050 let handler_start = Instant::now();
2051 let tool_result = match tool_handler.call(invocation).await {
2052 Ok(r) => r,
2053 Err(e) => tool_failure_result(e.to_string()),
2054 };
2055 tracing::debug!(
2056 elapsed_ms = handler_start.elapsed().as_millis(),
2057 session_id = %sid,
2058 request_id = %request_id,
2059 tool_call_id = %tool_call_id,
2060 tool_name = %tool_name,
2061 "ToolHandler::call dispatch"
2062 );
2063 let result_value = serde_json::to_value(tool_result).unwrap_or(Value::Null);
2064 let rpc_start = Instant::now();
2065 let _ = client
2066 .call(
2067 "session.tools.handlePendingToolCall",
2068 Some(serde_json::json!({
2069 "sessionId": sid,
2070 "requestId": request_id,
2071 "result": result_value,
2072 })),
2073 )
2074 .await;
2075 tracing::debug!(
2076 elapsed_ms = rpc_start.elapsed().as_millis(),
2077 session_id = %sid,
2078 request_id = %request_id,
2079 tool_call_id = %tool_call_id,
2080 tool_name = %tool_name,
2081 "Session::handle_notification response sent successfully"
2082 );
2083 }
2084 .instrument(span),
2085 );
2086 }
2087 SessionEventType::UserInputRequested => {
2088 }
2095 SessionEventType::ElicitationRequested => {
2096 let Some(request_id) = extract_request_id(¬ification.event.data) else {
2097 return;
2098 };
2099 let Some(elicitation_handler) = handlers.elicitation.clone() else {
2103 return;
2104 };
2105 let elicitation_data: ElicitationRequestedData =
2106 match serde_json::from_value(notification.event.data.clone()) {
2107 Ok(d) => d,
2108 Err(e) => {
2109 warn!(error = %e, "failed to deserialize elicitation request");
2110 return;
2111 }
2112 };
2113 let request = ElicitationRequest {
2114 message: elicitation_data.message,
2115 requested_schema: elicitation_data
2116 .requested_schema
2117 .map(|s| serde_json::to_value(s).unwrap_or(Value::Null)),
2118 mode: elicitation_data.mode.map(|m| match m {
2119 crate::generated::session_events::ElicitationRequestedMode::Form => {
2120 crate::types::ElicitationMode::Form
2121 }
2122 crate::generated::session_events::ElicitationRequestedMode::Url => {
2123 crate::types::ElicitationMode::Url
2124 }
2125 _ => crate::types::ElicitationMode::Unknown,
2126 }),
2127 elicitation_source: elicitation_data.elicitation_source,
2128 url: elicitation_data.url,
2129 };
2130 let client = client.clone();
2131 let sid = session_id.clone();
2132 let span = tracing::error_span!(
2133 "elicitation_request_handler",
2134 session_id = %sid,
2135 request_id = %request_id
2136 );
2137 tokio::spawn(
2138 async move {
2139 let cancel = ElicitationResult {
2140 action: "cancel".to_string(),
2141 content: None,
2142 };
2143 let handler_task = tokio::spawn({
2145 let sid = sid.clone();
2146 let request_id = request_id.clone();
2147 let span = tracing::error_span!(
2148 "elicitation_callback",
2149 session_id = %sid,
2150 request_id = %request_id
2151 );
2152 async move {
2153 let handler_start = Instant::now();
2154 let response = elicitation_handler
2155 .handle(sid.clone(), request_id.clone(), request)
2156 .await;
2157 tracing::debug!(
2158 elapsed_ms = handler_start.elapsed().as_millis(),
2159 session_id = %sid,
2160 request_id = %request_id,
2161 "ElicitationHandler::handle dispatch"
2162 );
2163 response
2164 }
2165 .instrument(span)
2166 });
2167 let result = match handler_task.await {
2168 Ok(r) => r,
2169 Err(_) => cancel.clone(),
2170 };
2171 let rpc_start = Instant::now();
2172 if let Err(e) = client
2173 .call(
2174 "session.ui.handlePendingElicitation",
2175 Some(serde_json::json!({
2176 "sessionId": sid,
2177 "requestId": request_id,
2178 "result": result,
2179 })),
2180 )
2181 .await
2182 {
2183 warn!(error = %e, "handlePendingElicitation failed, sending cancel");
2185 let _ = client
2186 .call(
2187 "session.ui.handlePendingElicitation",
2188 Some(serde_json::json!({
2189 "sessionId": sid,
2190 "requestId": request_id,
2191 "result": cancel,
2192 })),
2193 )
2194 .await;
2195 } else {
2196 tracing::debug!(
2197 elapsed_ms = rpc_start.elapsed().as_millis(),
2198 session_id = %sid,
2199 request_id = %request_id,
2200 "Session::handle_notification response sent successfully"
2201 );
2202 }
2203 }
2204 .instrument(span),
2205 );
2206 }
2207 SessionEventType::McpOauthRequired => {
2208 let Some(request_id) = extract_request_id(¬ification.event.data) else {
2209 return;
2210 };
2211 let Some(mcp_auth_handler) = handlers.mcp_auth.clone() else {
2212 warn!(
2213 session_id = %session_id,
2214 request_id = %request_id,
2215 "received MCP OAuth request without a registered MCP auth handler"
2216 );
2217 return;
2218 };
2219 let data: McpOauthRequiredData =
2220 match serde_json::from_value(notification.event.data.clone()) {
2221 Ok(d) => d,
2222 Err(e) => {
2223 warn!(error = %e, "failed to deserialize MCP OAuth request");
2224 return;
2225 }
2226 };
2227 let request = McpAuthRequest {
2228 request_id: request_id.clone(),
2229 server_name: data.server_name,
2230 server_url: data.server_url,
2231 reason: data.reason,
2232 www_authenticate_params: data.www_authenticate_params,
2233 resource_metadata: data.resource_metadata,
2234 static_client_config: data.static_client_config,
2235 };
2236 let client = client.clone();
2237 let sid = session_id.clone();
2238 let span = tracing::error_span!(
2239 "mcp_auth_request_handler",
2240 session_id = %sid,
2241 request_id = %request_id
2242 );
2243 tokio::spawn(
2244 async move {
2245 let cancel = McpAuthResult::Cancelled;
2246 let handler_task = tokio::spawn({
2247 let sid = sid.clone();
2248 let request_id = request_id.clone();
2249 let span = tracing::error_span!(
2250 "mcp_auth_callback",
2251 session_id = %sid,
2252 request_id = %request_id
2253 );
2254 async move {
2255 let handler_start = Instant::now();
2256 let response = mcp_auth_handler
2257 .handle(sid.clone(), request_id.clone(), request)
2258 .await;
2259 tracing::debug!(
2260 elapsed_ms = handler_start.elapsed().as_millis(),
2261 session_id = %sid,
2262 request_id = %request_id,
2263 "McpAuthHandler::handle dispatch"
2264 );
2265 response
2266 }
2267 .instrument(span)
2268 });
2269 let result = match handler_task.await {
2270 Ok(result) => result,
2271 Err(_) => cancel,
2272 };
2273 let rpc_start = Instant::now();
2274 let _ = client
2275 .call(
2276 "session.mcp.oauth.handlePendingRequest",
2277 Some(serde_json::json!({
2278 "sessionId": sid,
2279 "requestId": request_id,
2280 "result": result.into_wire(),
2281 })),
2282 )
2283 .await;
2284 tracing::debug!(
2285 elapsed_ms = rpc_start.elapsed().as_millis(),
2286 "Session::handle_notification MCP auth response sent"
2287 );
2288 }
2289 .instrument(span),
2290 );
2291 }
2292 SessionEventType::CommandExecute => {
2293 let data: CommandExecuteData =
2294 match serde_json::from_value(notification.event.data.clone()) {
2295 Ok(d) => d,
2296 Err(e) => {
2297 warn!(error = %e, "failed to deserialize command.execute");
2298 return;
2299 }
2300 };
2301 let client = client.clone();
2302 let command_handlers = command_handlers.clone();
2303 let sid = session_id.clone();
2304 let span = tracing::error_span!("command_handler", session_id = %sid);
2305 tokio::spawn(
2306 async move {
2307 let request_id = data.request_id;
2308 let ack_error = match command_handlers.get(&data.command_name).cloned() {
2309 None => Some(format!("Unknown command: {}", data.command_name)),
2310 Some(handler) => {
2311 let command_name = data.command_name.clone();
2312 let ctx = CommandContext {
2313 session_id: sid.clone(),
2314 command: data.command,
2315 command_name: data.command_name,
2316 args: data.args,
2317 };
2318 let handler_start = Instant::now();
2319 let result = handler.on_command(ctx).await;
2320 tracing::debug!(
2321 elapsed_ms = handler_start.elapsed().as_millis(),
2322 session_id = %sid,
2323 request_id = %request_id,
2324 command_name = %command_name,
2325 "CommandHandler::call dispatch"
2326 );
2327 match result {
2328 Ok(()) => None,
2329 Err(e) => Some(e.to_string()),
2330 }
2331 }
2332 };
2333 let mut params = serde_json::json!({
2334 "sessionId": sid,
2335 "requestId": request_id,
2336 });
2337 if let Some(error_msg) = ack_error {
2338 params["error"] = serde_json::Value::String(error_msg);
2339 }
2340 let rpc_start = Instant::now();
2341 let _ = client
2342 .call("session.commands.handlePendingCommand", Some(params))
2343 .await;
2344 tracing::debug!(
2345 elapsed_ms = rpc_start.elapsed().as_millis(),
2346 session_id = %sid,
2347 request_id = %request_id,
2348 "Session::handle_notification response sent successfully"
2349 );
2350 }
2351 .instrument(span),
2352 );
2353 }
2354 _ => {}
2355 }
2356}
2357
2358struct RequestDispatchContext<'a> {
2359 client: &'a Client,
2360 handlers: &'a SessionHandlers,
2361 hooks: Option<&'a dyn SessionHooks>,
2362 transforms: Option<&'a dyn SystemMessageTransform>,
2363 canvas_handler: Option<&'a Arc<dyn CanvasHandler>>,
2364 session_fs_provider: Option<&'a Arc<dyn SessionFsProvider>>,
2365 bearer_token_providers: &'a HashMap<String, Arc<dyn BearerTokenProvider>>,
2366}
2367
2368async fn handle_request(
2370 session_id: &SessionId,
2371 ctx: RequestDispatchContext<'_>,
2372 request: crate::JsonRpcRequest,
2373) {
2374 let sid = session_id.clone();
2375 let client = ctx.client;
2376 let handlers = ctx.handlers;
2377 let hooks = ctx.hooks;
2378 let transforms = ctx.transforms;
2379 let canvas_handler = ctx.canvas_handler;
2380 let session_fs_provider = ctx.session_fs_provider;
2381 let bearer_token_providers = ctx.bearer_token_providers;
2382
2383 if request.method.starts_with("sessionFs.") {
2384 crate::session_fs_dispatch::dispatch(client, session_fs_provider, request).await;
2385 return;
2386 }
2387
2388 if request.method.starts_with("canvas.") {
2389 crate::canvas_dispatch::dispatch(client, canvas_handler, request).await;
2390 return;
2391 }
2392
2393 if request.method == crate::generated::api_types::rpc_methods::PROVIDERTOKEN_GETTOKEN {
2394 crate::provider_token_dispatch::dispatch(client, bearer_token_providers, request).await;
2395 return;
2396 }
2397
2398 match request.method.as_str() {
2399 "hooks.invoke" => {
2400 let params = request.params.as_ref();
2401 let hook_type = params
2402 .and_then(|p| p.get("hookType"))
2403 .and_then(|v| v.as_str())
2404 .unwrap_or("");
2405 let input = params
2406 .and_then(|p| p.get("input"))
2407 .cloned()
2408 .unwrap_or(Value::Object(Default::default()));
2409
2410 let rpc_result = if let Some(hooks) = hooks {
2411 match crate::hooks::dispatch_hook(hooks, &sid, hook_type, input).await {
2412 Ok(output) => output,
2413 Err(e) => {
2414 warn!(error = %e, hook_type = hook_type, "hook dispatch failed");
2415 serde_json::json!({ "output": {} })
2416 }
2417 }
2418 } else {
2419 serde_json::json!({ "output": {} })
2420 };
2421
2422 let rpc_response = JsonRpcResponse {
2423 jsonrpc: "2.0".to_string(),
2424 id: request.id,
2425 result: Some(rpc_result),
2426 error: None,
2427 };
2428 let _ = client.send_response(&rpc_response).await;
2429 }
2430
2431 "userInput.request" => {
2432 let params = request.params.as_ref();
2433 let Some(question) = params
2434 .and_then(|p| p.get("question"))
2435 .and_then(|v| v.as_str())
2436 else {
2437 warn!("userInput.request missing 'question' field");
2438 let rpc_response = JsonRpcResponse {
2439 jsonrpc: "2.0".to_string(),
2440 id: request.id,
2441 result: None,
2442 error: Some(crate::JsonRpcError {
2443 code: error_codes::INVALID_PARAMS,
2444 message: "missing required field: question".to_string(),
2445 data: None,
2446 }),
2447 };
2448 let _ = client.send_response(&rpc_response).await;
2449 return;
2450 };
2451 let question = question.to_string();
2452 let choices = params
2453 .and_then(|p| p.get("choices"))
2454 .and_then(|v| v.as_array())
2455 .map(|arr| {
2456 arr.iter()
2457 .filter_map(|v| v.as_str().map(|s| s.to_string()))
2458 .collect()
2459 });
2460 let allow_freeform = params
2461 .and_then(|p| p.get("allowFreeform"))
2462 .and_then(|v| v.as_bool());
2463
2464 let handler_start = Instant::now();
2465 let response = if let Some(user_input_handler) = handlers.user_input.as_ref() {
2466 user_input_handler
2467 .handle(sid.clone(), question, choices, allow_freeform)
2468 .await
2469 } else {
2470 None
2471 };
2472 tracing::debug!(
2473 elapsed_ms = handler_start.elapsed().as_millis(),
2474 session_id = %sid,
2475 "UserInputHandler::handle dispatch"
2476 );
2477
2478 let rpc_result = match response {
2479 Some(UserInputResponse {
2480 answer,
2481 was_freeform,
2482 }) => serde_json::json!({
2483 "answer": answer,
2484 "wasFreeform": was_freeform,
2485 }),
2486 None => serde_json::json!({ "noResponse": true }),
2487 };
2488 let rpc_response = JsonRpcResponse {
2489 jsonrpc: "2.0".to_string(),
2490 id: request.id,
2491 result: Some(rpc_result),
2492 error: None,
2493 };
2494 let _ = client.send_response(&rpc_response).await;
2495 }
2496
2497 "exitPlanMode.request" => {
2498 let params = request
2499 .params
2500 .as_ref()
2501 .cloned()
2502 .unwrap_or(Value::Object(serde_json::Map::new()));
2503 let data: ExitPlanModeData = match serde_json::from_value(params) {
2504 Ok(d) => d,
2505 Err(e) => {
2506 warn!(error = %e, "failed to deserialize exitPlanMode.request params, using defaults");
2507 ExitPlanModeData::default()
2508 }
2509 };
2510
2511 let rpc_result = if let Some(exit_plan_handler) = handlers.exit_plan_mode.as_ref() {
2512 let result = exit_plan_handler.handle(sid, data).await;
2513 serde_json::to_value(result).expect("ExitPlanModeResult serialization cannot fail")
2514 } else {
2515 serde_json::json!({ "approved": true })
2516 };
2517 let rpc_response = JsonRpcResponse {
2518 jsonrpc: "2.0".to_string(),
2519 id: request.id,
2520 result: Some(rpc_result),
2521 error: None,
2522 };
2523 let _ = client.send_response(&rpc_response).await;
2524 }
2525
2526 "autoModeSwitch.request" => {
2527 let error_code = request
2528 .params
2529 .as_ref()
2530 .and_then(|p| p.get("errorCode"))
2531 .and_then(|v| v.as_str())
2532 .map(|s| s.to_string());
2533 let retry_after_seconds = request
2534 .params
2535 .as_ref()
2536 .and_then(|p| p.get("retryAfterSeconds"))
2537 .and_then(|v| v.as_f64());
2538
2539 let answer = if let Some(auto_mode_handler) = handlers.auto_mode_switch.as_ref() {
2540 auto_mode_handler
2541 .handle(sid, error_code, retry_after_seconds)
2542 .await
2543 } else {
2544 AutoModeSwitchResponse::No
2545 };
2546 let rpc_response = JsonRpcResponse {
2547 jsonrpc: "2.0".to_string(),
2548 id: request.id,
2549 result: Some(serde_json::json!({ "response": answer })),
2550 error: None,
2551 };
2552 let _ = client.send_response(&rpc_response).await;
2553 }
2554
2555 "systemMessage.transform" => {
2556 let params = request.params.as_ref();
2557 let sections: HashMap<String, crate::transforms::TransformSection> =
2558 match params.and_then(|p| p.get("sections")) {
2559 Some(v) => match serde_json::from_value(v.clone()) {
2560 Ok(s) => s,
2561 Err(e) => {
2562 let _ = send_error_response(
2563 client,
2564 request.id,
2565 error_codes::INVALID_PARAMS,
2566 &format!("invalid sections: {e}"),
2567 )
2568 .await;
2569 return;
2570 }
2571 },
2572 None => {
2573 let _ = send_error_response(
2574 client,
2575 request.id,
2576 error_codes::INVALID_PARAMS,
2577 "missing sections parameter",
2578 )
2579 .await;
2580 return;
2581 }
2582 };
2583
2584 let rpc_result = if let Some(transforms) = transforms {
2585 let transform_start = Instant::now();
2586 let response =
2587 crate::transforms::dispatch_transform(transforms, &sid, sections).await;
2588 tracing::debug!(
2589 elapsed_ms = transform_start.elapsed().as_millis(),
2590 session_id = %sid,
2591 "SystemMessageTransform::transform_section dispatch"
2592 );
2593 match serde_json::to_value(response) {
2594 Ok(v) => v,
2595 Err(e) => {
2596 warn!(error = %e, "failed to serialize transform response");
2597 serde_json::json!({ "sections": {} })
2598 }
2599 }
2600 } else {
2601 let passthrough: HashMap<String, crate::transforms::TransformSection> = sections;
2603 serde_json::json!({ "sections": passthrough })
2604 };
2605
2606 let rpc_response = JsonRpcResponse {
2607 jsonrpc: "2.0".to_string(),
2608 id: request.id,
2609 result: Some(rpc_result),
2610 error: None,
2611 };
2612 let _ = client.send_response(&rpc_response).await;
2613 }
2614
2615 method => {
2616 warn!(
2617 method = method,
2618 "unhandled request method in session event loop"
2619 );
2620 let _ = send_error_response(
2621 client,
2622 request.id,
2623 error_codes::METHOD_NOT_FOUND,
2624 &format!("unknown method: {method}"),
2625 )
2626 .await;
2627 }
2628 }
2629}
2630
2631async fn send_error_response(
2632 client: &Client,
2633 id: u64,
2634 code: i32,
2635 message: &str,
2636) -> Result<(), Error> {
2637 let response = JsonRpcResponse {
2638 jsonrpc: "2.0".to_string(),
2639 id,
2640 result: None,
2641 error: Some(crate::JsonRpcError {
2642 code,
2643 message: message.to_string(),
2644 data: None,
2645 }),
2646 };
2647 client.send_response(&response).await
2648}
2649
2650fn apply_transform_sections(
2654 sys_msg: &mut SystemMessageConfig,
2655 transforms: &dyn SystemMessageTransform,
2656) {
2657 sys_msg.mode = Some("customize".to_string());
2658 let sections = sys_msg.sections.get_or_insert_with(HashMap::new);
2659 for id in transforms.section_ids() {
2660 sections.entry(id).or_insert_with(|| SectionOverride {
2661 action: Some("transform".to_string()),
2662 content: None,
2663 });
2664 }
2665}
2666
2667fn inject_transform_sections(config: &mut SessionConfig, transforms: &dyn SystemMessageTransform) {
2668 let sys_msg = config.system_message.get_or_insert_with(Default::default);
2669 apply_transform_sections(sys_msg, transforms);
2670}
2671
2672fn inject_transform_sections_resume(
2673 config: &mut ResumeSessionConfig,
2674 transforms: &dyn SystemMessageTransform,
2675) {
2676 let sys_msg = config.system_message.get_or_insert_with(Default::default);
2677 apply_transform_sections(sys_msg, transforms);
2678}
2679
2680#[cfg(test)]
2681mod tests {
2682 use serde_json::json;
2683
2684 use super::{
2685 build_mode_post_create_patch, has_managed_settings, is_autopilot_continuation_idle,
2686 permission_request_data, permission_response_params,
2687 };
2688 use crate::handler::PermissionResult;
2689 use crate::types::{
2690 PermissionDecisionContext, PermissionDecisionOutcome, PermissionDecisionSource,
2691 PermissionDecisionSurface, RequestId, SessionEvent, SessionId,
2692 };
2693
2694 #[test]
2695 fn identifies_only_autopilot_continuation_idles() {
2696 let mut event = SessionEvent {
2697 id: "event-1".to_string(),
2698 timestamp: "2026-01-01T00:00:00Z".to_string(),
2699 parent_id: None,
2700 ephemeral: None,
2701 agent_id: None,
2702 debug_cli_received_at_ms: None,
2703 debug_ws_forwarded_at_ms: None,
2704 event_type: "session.idle".to_string(),
2705 data: json!({ "mode": "autopilot" }),
2706 };
2707
2708 assert!(is_autopilot_continuation_idle(&event));
2709
2710 event.data = json!({ "mode": "interactive" });
2711 assert!(!is_autopilot_continuation_idle(&event));
2712
2713 event.data = json!({});
2714 assert!(!is_autopilot_continuation_idle(&event));
2715 }
2716
2717 #[test]
2718 fn empty_mode_post_patch_sets_empty_included_builtin_skills() {
2719 let patch =
2720 build_mode_post_create_patch(crate::ClientMode::Empty, None, None, None, None, None)
2721 .expect("empty mode always sends a patch");
2722 assert_eq!(
2723 patch.included_builtin_skills,
2724 Some(Vec::new()),
2725 "empty mode must fail closed with an empty includedBuiltinSkills list"
2726 );
2727 assert_eq!(patch.installed_plugins.as_ref().map(|p| p.len()), Some(0));
2728 let value = serde_json::to_value(&patch).expect("serialize patch");
2730 assert_eq!(value["includedBuiltinSkills"], serde_json::json!([]));
2731 }
2732
2733 #[test]
2734 fn empty_mode_post_patch_preserves_explicit_builtin_skill_allowlist() {
2735 let patch = build_mode_post_create_patch(
2736 crate::ClientMode::Empty,
2737 Some(false),
2738 Some(false),
2739 Some(true),
2740 Some(true),
2741 Some(vec!["code-review".to_string()]),
2742 )
2743 .expect("empty mode always sends a patch");
2744 assert_eq!(
2745 patch.included_builtin_skills,
2746 Some(vec!["code-review".to_string()])
2747 );
2748 }
2749
2750 #[test]
2751 fn copilot_cli_mode_does_not_inject_included_builtin_skills() {
2752 assert!(
2754 build_mode_post_create_patch(
2755 crate::ClientMode::CopilotCli,
2756 None,
2757 None,
2758 None,
2759 None,
2760 None
2761 )
2762 .is_none()
2763 );
2764 let patch = build_mode_post_create_patch(
2766 crate::ClientMode::CopilotCli,
2767 Some(true),
2768 None,
2769 None,
2770 None,
2771 None,
2772 )
2773 .expect("a set field triggers a patch");
2774 assert_eq!(patch.included_builtin_skills, None);
2775 assert!(patch.installed_plugins.is_none());
2776 let value = serde_json::to_value(&patch).expect("serialize patch");
2777 assert!(value.get("includedBuiltinSkills").is_none());
2778
2779 let patch = build_mode_post_create_patch(
2780 crate::ClientMode::CopilotCli,
2781 None,
2782 None,
2783 None,
2784 None,
2785 Some(vec!["code-review".to_string()]),
2786 )
2787 .expect("an explicit allowlist triggers a patch");
2788 assert_eq!(
2789 patch.included_builtin_skills,
2790 Some(vec!["code-review".to_string()])
2791 );
2792 }
2793
2794 #[test]
2795 fn direct_injection_enables_managed_safeguards() {
2796 let settings = crate::types::ManagedSettings::default();
2797 assert!(has_managed_settings(None, Some(&settings)));
2798 assert!(!has_managed_settings(None, None));
2799 }
2800
2801 fn attribution_context() -> PermissionDecisionContext {
2802 PermissionDecisionContext {
2803 outcome: PermissionDecisionOutcome::AutoApproved,
2804 response_capability: None,
2805 source: PermissionDecisionSource::AssistedApproval,
2806 surface: PermissionDecisionSurface::CopilotApp,
2807 }
2808 }
2809
2810 #[test]
2811 fn response_params_omit_decision_context_without_attribution() {
2812 for (result, expected) in [
2813 (
2814 PermissionResult::approve_once(),
2815 json!({ "kind": "approve-once" }),
2816 ),
2817 (PermissionResult::reject(None), json!({ "kind": "reject" })),
2818 (
2819 PermissionResult::reject(Some("bad".to_string())),
2820 json!({ "kind": "reject", "feedback": "bad" }),
2821 ),
2822 (
2823 PermissionResult::user_not_available(),
2824 json!({ "kind": "user-not-available" }),
2825 ),
2826 ] {
2827 let params = permission_response_params(
2828 &SessionId::from("session-1"),
2829 &RequestId::from("permission-1"),
2830 &result,
2831 )
2832 .unwrap();
2833 assert_eq!(
2834 params,
2835 json!({
2836 "sessionId": "session-1",
2837 "requestId": "permission-1",
2838 "result": expected,
2839 })
2840 );
2841 }
2842 }
2843
2844 #[test]
2845 fn response_params_forward_decision_context_alongside_result() {
2846 let params = permission_response_params(
2847 &SessionId::from("session-1"),
2848 &RequestId::from("permission-1"),
2849 &PermissionResult::approve_once().with_context(attribution_context()),
2850 )
2851 .unwrap();
2852 assert_eq!(
2853 params,
2854 json!({
2855 "sessionId": "session-1",
2856 "requestId": "permission-1",
2857 "result": { "kind": "approve-once" },
2858 "decisionContext": {
2859 "outcome": "auto_approved",
2860 "source": "assisted_approval",
2861 "surface": "copilot_app",
2862 },
2863 })
2864 );
2865 assert!(params["result"].get("decisionContext").is_none());
2867 }
2868
2869 #[test]
2870 fn response_params_suppressed_for_no_result() {
2871 assert!(
2872 permission_response_params(
2873 &SessionId::from("session-1"),
2874 &RequestId::from("permission-1"),
2875 &PermissionResult::NoResult,
2876 )
2877 .is_none()
2878 );
2879 }
2880
2881 #[test]
2882 fn with_context_is_a_no_op_on_no_result() {
2883 let result = PermissionResult::no_result().with_context(attribution_context());
2884 assert!(matches!(result, PermissionResult::NoResult));
2885 }
2886
2887 #[test]
2888 fn with_context_replaces_rather_than_nests() {
2889 let result = PermissionResult::approve_once()
2890 .with_context(attribution_context())
2891 .with_context(PermissionDecisionContext {
2892 outcome: PermissionDecisionOutcome::PromptedUser,
2893 response_capability: None,
2894 source: PermissionDecisionSource::HumanResponse,
2895 surface: PermissionDecisionSurface::Sdk,
2896 });
2897 let params = permission_response_params(
2898 &SessionId::from("session-1"),
2899 &RequestId::from("permission-1"),
2900 &result,
2901 )
2902 .unwrap();
2903 assert_eq!(
2904 params["decisionContext"],
2905 json!({
2906 "outcome": "prompted_user",
2907 "source": "human_response",
2908 "surface": "sdk",
2909 })
2910 );
2911 }
2912
2913 #[test]
2914 fn permission_request_data_reads_nested_managed_approval_metadata() {
2915 let data = permission_request_data(
2916 &json!({
2917 "requestId": "permission-1",
2918 "permissionRequest": {
2919 "kind": "read",
2920 "managedApprovalRequired": true,
2921 "path": "/workspace/file.txt"
2922 }
2923 }),
2924 false,
2925 );
2926
2927 assert_eq!(data.managed_approval_required, Some(true));
2928 assert_eq!(
2929 data.extra["permissionRequest"]["path"],
2930 "/workspace/file.txt"
2931 );
2932 }
2933
2934 #[test]
2935 fn permission_request_data_preserves_managed_flag_when_other_fields_are_malformed() {
2936 let data = permission_request_data(
2937 &json!({
2938 "requestId": "permission-1",
2939 "permissionRequest": {
2940 "kind": "read",
2941 "managedApprovalRequired": true,
2942 "toolCallId": 42
2943 }
2944 }),
2945 false,
2946 );
2947
2948 assert_eq!(data.managed_approval_required, Some(true));
2949 assert_eq!(data.extra["requestId"], "permission-1");
2950 }
2951
2952 #[test]
2953 fn permission_request_data_fails_closed_for_malformed_managed_flag() {
2954 let data = permission_request_data(
2955 &json!({
2956 "requestId": "permission-1",
2957 "permissionRequest": {
2958 "kind": "read",
2959 "managedApprovalRequired": "yes",
2960 "path": "/workspace/file.txt"
2961 }
2962 }),
2963 false,
2964 );
2965
2966 assert_eq!(data.managed_approval_required, Some(true));
2967 }
2968
2969 #[test]
2970 fn permission_request_data_preserves_valid_false_managed_flag() {
2971 let data = permission_request_data(
2972 &json!({
2973 "requestId": "permission-1",
2974 "permissionRequest": {
2975 "kind": "read",
2976 "managedApprovalRequired": false,
2977 "path": "/workspace/file.txt"
2978 }
2979 }),
2980 false,
2981 );
2982
2983 assert_eq!(data.managed_approval_required, Some(false));
2984 }
2985}