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