1use std::collections::HashMap;
2use std::path::{Path, PathBuf};
3use std::sync::Arc;
4use std::time::{Duration, Instant};
5
6use parking_lot::Mutex as ParkingLotMutex;
7use serde_json::Value;
8use tokio::sync::oneshot;
9use tokio::task::JoinHandle;
10use tokio_util::sync::CancellationToken;
11use tracing::{Instrument, warn};
12
13use crate::canvas::CanvasHandler;
14use crate::generated::api_types::{
15 LogRequest, ModelSwitchToRequest, OpenCanvasInstance, RegisterEventInterestParams, rpc_methods,
16};
17use crate::generated::session_events::{
18 CommandExecuteData, ElicitationRequestedData, ExternalToolRequestedData, McpOauthRequiredData,
19 SessionCanvasClosedData, SessionErrorData, SessionEventType,
20};
21use crate::handler::{
22 AutoModeSwitchHandler, AutoModeSwitchResponse, ElicitationHandler, ExitPlanModeHandler,
23 McpAuthHandler, McpAuthRequest, McpAuthResult, PermissionHandler, PermissionResult,
24 UserInputHandler, UserInputResponse,
25};
26use crate::hooks::SessionHooks;
27use crate::provider_token::BearerTokenProvider;
28use crate::session_fs::SessionFsProvider;
29use crate::trace_context::inject_trace_context;
30use crate::transforms::SystemMessageTransform;
31use crate::types::{
32 CommandContext, CommandDefinition, CommandHandler, CreateSessionResult, ElicitationRequest,
33 ElicitationResult, ExitPlanModeData, GetMessagesResponse, MessageOptions,
34 PermissionRequestData, RequestId, ResumeSessionConfig, ResumeSessionResult, SectionOverride,
35 SessionCapabilities, SessionConfig, SessionEvent, SessionId, SetModelOptions,
36 SystemMessageConfig, ToolInvocation, ToolResult, ToolResultExpanded, TraceContext,
37 UiInputOptions, ensure_attachment_display_names,
38};
39use crate::{
40 Client, Error, ErrorKind, JsonRpcResponse, SessionErrorKind, SessionEventNotification,
41 error_codes,
42};
43
44#[derive(Clone)]
52pub(crate) struct SessionHandlers {
53 pub permission: Option<Arc<dyn PermissionHandler>>,
54 pub elicitation: Option<Arc<dyn ElicitationHandler>>,
55 pub mcp_auth: Option<Arc<dyn McpAuthHandler>>,
56 pub user_input: Option<Arc<dyn UserInputHandler>>,
57 pub exit_plan_mode: Option<Arc<dyn ExitPlanModeHandler>>,
58 pub auto_mode_switch: Option<Arc<dyn AutoModeSwitchHandler>>,
59 pub tools: Arc<HashMap<String, Arc<dyn crate::tool::ToolHandler>>>,
60}
61
62struct IdleWaiter {
64 tx: oneshot::Sender<Result<Option<SessionEvent>, Error>>,
65 last_assistant_message: Option<SessionEvent>,
66 started_at: Instant,
67 first_assistant_message_seen: bool,
68}
69
70struct WaiterGuard {
82 slot: Arc<ParkingLotMutex<Option<IdleWaiter>>>,
83}
84
85impl Drop for WaiterGuard {
86 fn drop(&mut self) {
87 self.slot.lock().take();
88 }
89}
90
91struct PendingSessionRegistration {
92 client: Client,
93 session_id: SessionId,
94 shutdown: CancellationToken,
95 disarmed: bool,
96}
97
98impl PendingSessionRegistration {
99 fn new(client: Client, session_id: SessionId, shutdown: CancellationToken) -> Self {
100 Self {
101 client,
102 session_id,
103 shutdown,
104 disarmed: false,
105 }
106 }
107
108 async fn cleanup(mut self, event_loop: JoinHandle<()>) {
109 self.shutdown.cancel();
110 let _ = event_loop.await;
111 self.client.unregister_session(&self.session_id);
112 self.disarmed = true;
113 }
114
115 fn disarm(&mut self) {
116 self.disarmed = true;
117 }
118}
119
120impl Drop for PendingSessionRegistration {
121 fn drop(&mut self) {
122 if !self.disarmed {
123 self.shutdown.cancel();
124 self.client.unregister_session(&self.session_id);
125 }
126 }
127}
128
129pub struct Session {
142 id: SessionId,
143 cwd: PathBuf,
144 workspace_path: Option<PathBuf>,
145 remote_url: Option<String>,
146 client: Client,
147 event_loop: ParkingLotMutex<Option<JoinHandle<()>>>,
152 shutdown: CancellationToken,
166 idle_waiter: Arc<ParkingLotMutex<Option<IdleWaiter>>>,
173 capabilities: Arc<parking_lot::RwLock<SessionCapabilities>>,
175 open_canvases: Arc<parking_lot::RwLock<Vec<OpenCanvasInstance>>>,
177 event_tx: tokio::sync::broadcast::Sender<SessionEvent>,
179}
180
181impl Session {
182 pub fn id(&self) -> &SessionId {
184 &self.id
185 }
186
187 pub fn cwd(&self) -> &PathBuf {
189 &self.cwd
190 }
191
192 pub fn workspace_path(&self) -> Option<&Path> {
194 self.workspace_path.as_deref()
195 }
196
197 pub fn remote_url(&self) -> Option<&str> {
199 self.remote_url.as_deref()
200 }
201
202 pub fn capabilities(&self) -> SessionCapabilities {
207 self.capabilities.read().clone()
208 }
209
210 pub fn open_canvases(&self) -> Vec<OpenCanvasInstance> {
213 self.open_canvases.read().clone()
214 }
215
216 pub fn cancellation_token(&self) -> CancellationToken {
243 self.shutdown.child_token()
244 }
245
246 pub fn subscribe(&self) -> crate::subscription::EventSubscription {
285 crate::subscription::EventSubscription::new(self.event_tx.subscribe())
286 }
287
288 pub fn client(&self) -> &Client {
290 &self.client
291 }
292
293 pub fn rpc(&self) -> crate::generated::rpc::SessionRpc<'_> {
304 crate::generated::rpc::SessionRpc { session: self }
305 }
306
307 pub async fn stop_event_loop(&self) {
315 self.shutdown.cancel();
316 let handle = self.event_loop.lock().take();
317 if let Some(handle) = handle {
318 let _ = handle.await;
319 }
320 if let Some(waiter) = self.idle_waiter.lock().take() {
322 let _ = waiter.tx.send(Err(
323 ErrorKind::Session(SessionErrorKind::EventLoopClosed).into()
324 ));
325 }
326 }
327
328 pub async fn send(&self, opts: impl Into<MessageOptions>) -> Result<String, Error> {
353 if self.idle_waiter.lock().is_some() {
354 return Err(ErrorKind::Session(SessionErrorKind::SendWhileWaiting).into());
355 }
356 self.send_inner(opts.into()).await
357 }
358
359 async fn send_inner(&self, opts: MessageOptions) -> Result<String, Error> {
360 let mut params = serde_json::json!({
361 "sessionId": self.id,
362 "prompt": opts.prompt,
363 });
364 if let Some(m) = opts.mode {
365 params["mode"] = serde_json::to_value(m)?;
366 }
367 if let Some(am) = opts.agent_mode {
368 params["agentMode"] = serde_json::to_value(am)?;
369 }
370 if let Some(mut a) = opts.attachments {
371 ensure_attachment_display_names(&mut a);
372 params["attachments"] = serde_json::to_value(a)?;
373 }
374 if let Some(headers) = opts.request_headers
375 && !headers.is_empty()
376 {
377 params["requestHeaders"] = serde_json::to_value(headers)?;
378 }
379 if let Some(display_prompt) = opts.display_prompt {
380 params["displayPrompt"] = serde_json::to_value(display_prompt)?;
381 }
382 let trace_ctx = if opts.traceparent.is_some() || opts.tracestate.is_some() {
383 TraceContext {
384 traceparent: opts.traceparent,
385 tracestate: opts.tracestate,
386 }
387 } else {
388 self.client.resolve_trace_context().await
389 };
390 inject_trace_context(&mut params, &trace_ctx);
391 let rpc_start = Instant::now();
392 let result = self.client.call("session.send", Some(params)).await?;
393 let message_id = result
394 .get("messageId")
395 .and_then(|v| v.as_str())
396 .map(|s| s.to_string())
397 .unwrap_or_default();
398 tracing::debug!(
399 elapsed_ms = rpc_start.elapsed().as_millis(),
400 session_id = %self.id,
401 message_id = %message_id,
402 "Session::send completed successfully"
403 );
404 Ok(message_id)
405 }
406
407 pub async fn send_and_wait(
427 &self,
428 opts: impl Into<MessageOptions>,
429 ) -> Result<Option<SessionEvent>, Error> {
430 let total_start = Instant::now();
431 let opts = opts.into();
432 let timeout_duration = opts.wait_timeout.unwrap_or(Duration::from_secs(60));
433 let (tx, rx) = oneshot::channel();
434
435 {
436 let mut guard = self.idle_waiter.lock();
437 if guard.is_some() {
438 return Err(ErrorKind::Session(SessionErrorKind::SendWhileWaiting).into());
439 }
440 *guard = Some(IdleWaiter {
441 tx,
442 last_assistant_message: None,
443 started_at: total_start,
444 first_assistant_message_seen: false,
445 });
446 }
447
448 let _waiter_guard = WaiterGuard {
453 slot: self.idle_waiter.clone(),
454 };
455
456 let result = tokio::time::timeout(timeout_duration, async {
457 self.send_inner(opts).await?;
458 match rx.await {
459 Ok(result) => result,
460 Err(_) => Err(ErrorKind::Session(SessionErrorKind::EventLoopClosed).into()),
461 }
462 })
463 .await;
464
465 match result {
466 Ok(inner) => {
467 tracing::debug!(
468 elapsed_ms = total_start.elapsed().as_millis(),
469 session_id = %self.id,
470 completed_by = if inner.is_ok() { "idle" } else { "error" },
471 "Session::send_and_wait complete"
472 );
473 inner
474 }
475 Err(_) => {
476 tracing::warn!(
477 elapsed_ms = total_start.elapsed().as_millis(),
478 session_id = %self.id,
479 completed_by = "timeout",
480 "Session::send_and_wait failed"
481 );
482 Err(ErrorKind::Session(SessionErrorKind::Timeout(timeout_duration)).into())
483 }
484 }
485 }
486
487 pub async fn get_events(&self) -> Result<Vec<SessionEvent>, Error> {
489 let result = self
490 .client
491 .call(
492 "session.getMessages",
493 Some(serde_json::json!({ "sessionId": self.id })),
494 )
495 .await?;
496 let response: GetMessagesResponse = serde_json::from_value(result)?;
497 Ok(response.events)
498 }
499
500 #[deprecated(since = "0.1.0", note = "Use `get_events()` instead")]
502 pub async fn get_messages(&self) -> Result<Vec<SessionEvent>, Error> {
503 self.get_events().await
504 }
505
506 pub async fn abort(&self) -> Result<(), Error> {
514 self.client
515 .call(
516 "session.abort",
517 Some(serde_json::json!({ "sessionId": self.id })),
518 )
519 .await?;
520 Ok(())
521 }
522
523 pub async fn set_model(&self, model: &str, opts: Option<SetModelOptions>) -> Result<(), Error> {
527 let opts = opts.unwrap_or_default();
528 let request = ModelSwitchToRequest {
529 model_id: model.to_string(),
530 reasoning_effort: opts.reasoning_effort,
531 reasoning_summary: opts.reasoning_summary,
532 verbosity: None,
533 context_tier: opts.context_tier,
534 model_capabilities: opts.model_capabilities,
535 };
536 self.rpc().model().switch_to(request).await?;
537 Ok(())
538 }
539
540 pub async fn disconnect(&self) -> Result<(), Error> {
557 self.client
558 .call(
559 "session.destroy",
560 Some(serde_json::json!({ "sessionId": self.id })),
561 )
562 .await?;
563 self.stop_event_loop().await;
564 self.client.unregister_session(&self.id);
565 Ok(())
566 }
567
568 #[deprecated(since = "0.1.0", note = "Use `disconnect()` instead")]
573 pub async fn destroy(&self) -> Result<(), Error> {
574 self.disconnect().await
575 }
576
577 pub async fn log(
581 &self,
582 message: &str,
583 opts: Option<crate::types::LogOptions>,
584 ) -> Result<(), Error> {
585 let opts = opts.unwrap_or_default();
586 let level = match opts.level {
587 Some(level) => Some(serde_json::from_value(serde_json::to_value(level)?)?),
588 None => None,
589 };
590 let request = LogRequest {
591 message: message.to_string(),
592 level,
593 ephemeral: opts.ephemeral,
594 r#type: None,
595 tip: None,
596 url: None,
597 };
598 self.rpc().log(request).await?;
599 Ok(())
600 }
601
602 pub fn ui(&self) -> SessionUi<'_> {
608 SessionUi { session: self }
609 }
610
611 fn assert_elicitation(&self) -> Result<(), Error> {
613 if self
614 .capabilities
615 .read()
616 .ui
617 .as_ref()
618 .and_then(|u| u.elicitation)
619 != Some(true)
620 {
621 return Err(ErrorKind::Session(SessionErrorKind::ElicitationNotSupported).into());
622 }
623 Ok(())
624 }
625}
626
627impl Drop for Session {
628 fn drop(&mut self) {
629 self.shutdown.cancel();
641 self.client.unregister_session(&self.id);
642 }
643}
644
645pub struct SessionUi<'a> {
652 session: &'a Session,
653}
654
655impl<'a> SessionUi<'a> {
656 pub async fn elicitation(
664 &self,
665 message: &str,
666 schema: Value,
667 ) -> Result<ElicitationResult, Error> {
668 self.session.assert_elicitation()?;
669 let result = self
670 .session
671 .client
672 .call(
673 "session.ui.elicitation",
674 Some(serde_json::json!({
675 "sessionId": self.session.id,
676 "message": message,
677 "requestedSchema": schema,
678 })),
679 )
680 .await?;
681 let elicitation: ElicitationResult = serde_json::from_value(result)?;
682 Ok(elicitation)
683 }
684
685 pub async fn confirm(&self, message: &str) -> Result<bool, Error> {
689 self.session.assert_elicitation()?;
690 let schema = serde_json::json!({
691 "type": "object",
692 "properties": {
693 "confirmed": {
694 "type": "boolean",
695 "default": true,
696 }
697 },
698 "required": ["confirmed"]
699 });
700 let result = self.elicitation(message, schema).await?;
701 Ok(result.action == "accept"
702 && result
703 .content
704 .and_then(|c| c.get("confirmed").and_then(|v| v.as_bool()))
705 == Some(true))
706 }
707
708 pub async fn select(&self, message: &str, options: &[&str]) -> Result<Option<String>, Error> {
712 self.session.assert_elicitation()?;
713 let schema = serde_json::json!({
714 "type": "object",
715 "properties": {
716 "selection": {
717 "type": "string",
718 "enum": options,
719 }
720 },
721 "required": ["selection"]
722 });
723 let result = self.elicitation(message, schema).await?;
724 if result.action != "accept" {
725 return Ok(None);
726 }
727 let selection = result.content.and_then(|c| {
728 c.get("selection")
729 .and_then(|v| v.as_str())
730 .map(String::from)
731 });
732 Ok(selection)
733 }
734
735 pub async fn input(
740 &self,
741 message: &str,
742 options: Option<&UiInputOptions<'_>>,
743 ) -> Result<Option<String>, Error> {
744 self.session.assert_elicitation()?;
745 let mut field = serde_json::json!({ "type": "string" });
746 if let Some(opts) = options {
747 if let Some(title) = opts.title {
748 field["title"] = Value::String(title.to_string());
749 }
750 if let Some(desc) = opts.description {
751 field["description"] = Value::String(desc.to_string());
752 }
753 if let Some(min) = opts.min_length {
754 field["minLength"] = Value::Number(min.into());
755 }
756 if let Some(max) = opts.max_length {
757 field["maxLength"] = Value::Number(max.into());
758 }
759 if let Some(fmt) = &opts.format {
760 field["format"] = Value::String(fmt.as_str().to_string());
761 }
762 if let Some(default) = opts.default {
763 field["default"] = Value::String(default.to_string());
764 }
765 }
766 let schema = serde_json::json!({
767 "type": "object",
768 "properties": { "value": field },
769 "required": ["value"]
770 });
771 let result = self.elicitation(message, schema).await?;
772 if result.action != "accept" {
773 return Ok(None);
774 }
775 let value = result
776 .content
777 .and_then(|c| c.get("value").and_then(|v| v.as_str()).map(String::from));
778 Ok(value)
779 }
780}
781
782impl Client {
783 pub async fn create_session(&self, mut config: SessionConfig) -> Result<Session, Error> {
805 let total_start = Instant::now();
806 let caller_session_id = config.session_id.clone();
814 let use_server_generated_id = config.cloud.is_some() && caller_session_id.is_none();
815 let local_session_id: Option<SessionId> = if use_server_generated_id {
816 None
817 } else {
818 Some(
819 caller_session_id
820 .clone()
821 .unwrap_or_else(|| SessionId::new(uuid::Uuid::new_v4().to_string())),
822 )
823 };
824 if config.hooks_handler.is_some() && config.hooks.is_none() {
825 config.hooks = Some(true);
826 }
827 if let Some(transforms) = config.system_message_transform.clone() {
828 inject_transform_sections(&mut config, transforms.as_ref());
829 }
830 let mode = self.inner.mode;
831 if mode == crate::ClientMode::Empty && config.available_tools.is_none() {
832 return Err(Error::with_message(
833 ErrorKind::InvalidConfig,
834 "ClientMode::Empty requires available_tools to be set on the session config. \
835 Use ToolSet to specify which tools the session may use (e.g. \
836 ToolSet::new().add_builtin_many(BUILTIN_TOOLS_ISOLATED)).",
837 ));
838 }
839 crate::mode::validate_tool_filter_list(
840 "available_tools",
841 config.available_tools.as_deref(),
842 )?;
843 crate::mode::validate_tool_filter_list("excluded_tools", config.excluded_tools.as_deref())?;
844 config.system_message =
845 crate::mode::system_message_for_mode(mode, config.system_message.take());
846 config.memory = crate::mode::memory_for_mode(mode, config.memory.take());
847 if mode == crate::ClientMode::Empty {
848 if config.enable_session_telemetry.is_none() {
849 config.enable_session_telemetry = Some(false);
850 }
851 if config.skip_embedding_retrieval.is_none() {
852 config.skip_embedding_retrieval = Some(true);
853 }
854 if config.enable_on_demand_instruction_discovery.is_none() {
855 config.enable_on_demand_instruction_discovery = Some(false);
856 }
857 if config.enable_file_hooks.is_none() {
858 config.enable_file_hooks = Some(false);
859 }
860 if config.enable_host_git_operations.is_none() {
861 config.enable_host_git_operations = Some(false);
862 }
863 if config.enable_session_store.is_none() {
864 config.enable_session_store = Some(false);
865 }
866 if config.enable_skills.is_none() {
867 config.enable_skills = Some(false);
868 }
869 }
870 if mode == crate::ClientMode::Empty && config.mcp_oauth_token_storage.is_none() {
871 config.mcp_oauth_token_storage = Some("in-memory".into());
872 }
873 if mode == crate::ClientMode::Empty && config.embedding_cache_storage.is_none() {
874 config.embedding_cache_storage = Some("in-memory".into());
875 }
876 let opt_skip_custom_instructions = config.skip_custom_instructions;
877 let opt_custom_agents_local_only = config.custom_agents_local_only;
878 let opt_coauthor_enabled = config.coauthor_enabled;
879 let opt_manage_schedule_enabled = config.manage_schedule_enabled;
880 let (mut wire, mut runtime) = config.into_wire(local_session_id.clone())?;
881 wire.enable_github_telemetry_forwarding =
882 self.inner.on_github_telemetry.is_some().then_some(true);
883
884 let permission_handler = crate::permission::resolve_handler(
885 runtime.permission_handler.take(),
886 runtime.permission_policy.take(),
887 );
888 let handlers = SessionHandlers {
889 permission: permission_handler,
890 elicitation: runtime.elicitation_handler.take(),
891 mcp_auth: runtime.mcp_auth_handler.take(),
892 user_input: runtime.user_input_handler.take(),
893 exit_plan_mode: runtime.exit_plan_mode_handler.take(),
894 auto_mode_switch: runtime.auto_mode_switch_handler.take(),
895 tools: Arc::new(std::mem::take(&mut runtime.tool_handlers)),
896 };
897 let hooks = runtime.hooks_handler.take();
898 let transforms = runtime.system_message_transform.take();
899 let tools_count = wire.tools.as_ref().map_or(0, Vec::len);
900 let commands_count = runtime.commands.as_ref().map_or(0, Vec::len);
901 let has_hooks = hooks.is_some();
902 let command_handlers = build_command_handler_map(runtime.commands.as_deref());
903 let canvas_handler = runtime.canvas_handler.take();
904 let session_fs_provider = runtime.session_fs_provider.take();
905 let bearer_token_providers = std::mem::take(&mut runtime.bearer_token_providers);
906 let has_mcp_auth_handler = handlers.mcp_auth.is_some();
907 if self.inner.session_fs_configured && session_fs_provider.is_none() {
908 return Err(ErrorKind::Session(SessionErrorKind::SessionFsProviderRequired).into());
909 }
910 if self.inner.session_fs_sqlite_declared
911 && let Some(ref provider) = session_fs_provider
912 && provider.sqlite().is_none()
913 {
914 return Err(Error::with_message(
915 ErrorKind::InvalidConfig,
916 "SessionFs capabilities declare SQLite support but the provider \
917 does not implement SessionFsSqliteProvider",
918 ));
919 }
920
921 let mut params = serde_json::to_value(&wire)?;
922 let trace_ctx = self.resolve_trace_context().await;
923 inject_trace_context(&mut params, &trace_ctx);
924
925 let setup_start = Instant::now();
926 let capabilities = Arc::new(parking_lot::RwLock::new(SessionCapabilities::default()));
927 let idle_waiter = Arc::new(ParkingLotMutex::new(None));
928 let open_canvases = Arc::new(parking_lot::RwLock::new(Vec::new()));
929 let shutdown = CancellationToken::new();
930 let (event_tx, _) = tokio::sync::broadcast::channel(512);
931
932 let inline_stash: Arc<
938 ParkingLotMutex<Option<(SessionId, crate::router::SessionChannels)>>,
939 > = Arc::new(ParkingLotMutex::new(None));
940
941 let inline_callback: Option<crate::jsonrpc::InlineResponseCallback> = if let Some(ref sid) =
942 local_session_id
943 {
944 let channels = self.register_session(sid);
945 *inline_stash.lock() = Some((sid.clone(), channels));
946 None
947 } else {
948 let client = self.clone();
949 let stash = inline_stash.clone();
950 let expected = caller_session_id.clone();
951 Some(Box::new(move |response| {
952 let result = response.result.as_ref().ok_or_else(|| {
953 Error::with_message(ErrorKind::Json, "session.create response had no result")
954 })?;
955 let parsed: CreateSessionResult =
956 serde_json::from_value(result.clone()).map_err(Error::from)?;
957 if let Some(requested) = expected.as_ref()
958 && parsed.session_id != *requested
959 {
960 return Err(ErrorKind::Session(SessionErrorKind::SessionIdMismatch {
961 requested: requested.clone(),
962 returned: parsed.session_id,
963 })
964 .into());
965 }
966 let channels = client.register_session(&parsed.session_id);
967 *stash.lock() = Some((parsed.session_id, channels));
968 Ok(())
969 }))
970 };
971
972 let rpc_start = Instant::now();
973 let result = match self
974 .call_with_inline_callback("session.create", Some(params), inline_callback)
975 .await
976 {
977 Ok(result) => result,
978 Err(error) => {
979 if let Some((id, _channels)) = inline_stash.lock().take() {
980 self.unregister_session(&id);
981 }
982 return Err(error);
983 }
984 };
985 tracing::debug!(
986 elapsed_ms = rpc_start.elapsed().as_millis(),
987 "Client::create_session session creation request completed successfully"
988 );
989 let create_result: CreateSessionResult = match serde_json::from_value(result) {
990 Ok(result) => result,
991 Err(error) => {
992 if let Some((id, _channels)) = inline_stash.lock().take() {
993 self.unregister_session(&id);
994 }
995 return Err(error.into());
996 }
997 };
998
999 if let Some(ref requested) = local_session_id
1000 && create_result.session_id != *requested
1001 {
1002 if let Some((id, _channels)) = inline_stash.lock().take() {
1003 self.unregister_session(&id);
1004 }
1005 return Err(ErrorKind::Session(SessionErrorKind::SessionIdMismatch {
1006 requested: requested.clone(),
1007 returned: create_result.session_id.clone(),
1008 })
1009 .into());
1010 }
1011
1012 let (session_id, channels) = inline_stash
1013 .lock()
1014 .take()
1015 .expect("session registration must have populated stash on success");
1016 let event_loop = spawn_event_loop(
1017 session_id.clone(),
1018 self.clone(),
1019 handlers,
1020 hooks,
1021 transforms,
1022 command_handlers,
1023 canvas_handler,
1024 session_fs_provider,
1025 bearer_token_providers,
1026 channels,
1027 idle_waiter.clone(),
1028 capabilities.clone(),
1029 open_canvases.clone(),
1030 event_tx.clone(),
1031 shutdown.clone(),
1032 );
1033 tracing::debug!(
1034 elapsed_ms = setup_start.elapsed().as_millis(),
1035 session_id = %session_id,
1036 tools_count,
1037 commands_count,
1038 has_hooks,
1039 "Client::create_session local setup complete"
1040 );
1041 *capabilities.write() = create_result.capabilities.unwrap_or_default();
1042 if has_mcp_auth_handler {
1043 register_mcp_auth_interest(self, &session_id).await?;
1044 }
1045
1046 tracing::debug!(
1047 elapsed_ms = total_start.elapsed().as_millis(),
1048 session_id = %session_id,
1049 "Client::create_session complete"
1050 );
1051 let session = Session {
1052 id: session_id,
1053 cwd: self.cwd().clone(),
1054 workspace_path: create_result.workspace_path,
1055 remote_url: create_result.remote_url,
1056 client: self.clone(),
1057 event_loop: ParkingLotMutex::new(Some(event_loop)),
1058 shutdown,
1059 idle_waiter,
1060 capabilities,
1061 open_canvases,
1062 event_tx,
1063 };
1064 apply_mode_post_create_patch(
1065 &session,
1066 mode,
1067 opt_skip_custom_instructions,
1068 opt_custom_agents_local_only,
1069 opt_coauthor_enabled,
1070 opt_manage_schedule_enabled,
1071 )
1072 .await?;
1073 Ok(session)
1074 }
1075
1076 pub async fn resume_session(&self, mut config: ResumeSessionConfig) -> Result<Session, Error> {
1087 let total_start = Instant::now();
1088 let session_id = config.session_id.clone();
1089 if config.hooks_handler.is_some() && config.hooks.is_none() {
1090 config.hooks = Some(true);
1091 }
1092 if let Some(transforms) = config.system_message_transform.clone() {
1093 inject_transform_sections_resume(&mut config, transforms.as_ref());
1094 }
1095 let mode = self.inner.mode;
1096 if mode == crate::ClientMode::Empty && config.available_tools.is_none() {
1097 return Err(Error::with_message(
1098 ErrorKind::InvalidConfig,
1099 "ClientMode::Empty requires available_tools to be set on the session config. \
1100 Use ToolSet to specify which tools the session may use (e.g. \
1101 ToolSet::new().add_builtin_many(BUILTIN_TOOLS_ISOLATED)).",
1102 ));
1103 }
1104 crate::mode::validate_tool_filter_list(
1105 "available_tools",
1106 config.available_tools.as_deref(),
1107 )?;
1108 crate::mode::validate_tool_filter_list("excluded_tools", config.excluded_tools.as_deref())?;
1109 config.system_message =
1110 crate::mode::system_message_for_mode(mode, config.system_message.take());
1111 config.memory = crate::mode::memory_for_mode(mode, config.memory.take());
1112 if mode == crate::ClientMode::Empty {
1113 if config.enable_session_telemetry.is_none() {
1114 config.enable_session_telemetry = Some(false);
1115 }
1116 if config.skip_embedding_retrieval.is_none() {
1117 config.skip_embedding_retrieval = Some(true);
1118 }
1119 if config.enable_on_demand_instruction_discovery.is_none() {
1120 config.enable_on_demand_instruction_discovery = Some(false);
1121 }
1122 if config.enable_file_hooks.is_none() {
1123 config.enable_file_hooks = Some(false);
1124 }
1125 if config.enable_host_git_operations.is_none() {
1126 config.enable_host_git_operations = Some(false);
1127 }
1128 if config.enable_session_store.is_none() {
1129 config.enable_session_store = Some(false);
1130 }
1131 if config.enable_skills.is_none() {
1132 config.enable_skills = Some(false);
1133 }
1134 }
1135 if mode == crate::ClientMode::Empty && config.mcp_oauth_token_storage.is_none() {
1136 config.mcp_oauth_token_storage = Some("in-memory".into());
1137 }
1138 if mode == crate::ClientMode::Empty && config.embedding_cache_storage.is_none() {
1139 config.embedding_cache_storage = Some("in-memory".into());
1140 }
1141 let opt_skip_custom_instructions = config.skip_custom_instructions;
1142 let opt_custom_agents_local_only = config.custom_agents_local_only;
1143 let opt_coauthor_enabled = config.coauthor_enabled;
1144 let opt_manage_schedule_enabled = config.manage_schedule_enabled;
1145 let (mut wire, mut runtime) = config.into_wire()?;
1146 wire.enable_github_telemetry_forwarding =
1147 self.inner.on_github_telemetry.is_some().then_some(true);
1148
1149 let permission_handler = crate::permission::resolve_handler(
1150 runtime.permission_handler.take(),
1151 runtime.permission_policy.take(),
1152 );
1153 let handlers = SessionHandlers {
1154 permission: permission_handler,
1155 elicitation: runtime.elicitation_handler.take(),
1156 mcp_auth: runtime.mcp_auth_handler.take(),
1157 user_input: runtime.user_input_handler.take(),
1158 exit_plan_mode: runtime.exit_plan_mode_handler.take(),
1159 auto_mode_switch: runtime.auto_mode_switch_handler.take(),
1160 tools: Arc::new(std::mem::take(&mut runtime.tool_handlers)),
1161 };
1162 let hooks = runtime.hooks_handler.take();
1163 let transforms = runtime.system_message_transform.take();
1164 let tools_count = wire.tools.as_ref().map_or(0, Vec::len);
1165 let commands_count = runtime.commands.as_ref().map_or(0, Vec::len);
1166 let has_hooks = hooks.is_some();
1167 let command_handlers = build_command_handler_map(runtime.commands.as_deref());
1168 let canvas_handler = runtime.canvas_handler.take();
1169 let session_fs_provider = runtime.session_fs_provider.take();
1170 let bearer_token_providers = std::mem::take(&mut runtime.bearer_token_providers);
1171 let has_mcp_auth_handler = handlers.mcp_auth.is_some();
1172 if self.inner.session_fs_configured && session_fs_provider.is_none() {
1173 return Err(ErrorKind::Session(SessionErrorKind::SessionFsProviderRequired).into());
1174 }
1175 if self.inner.session_fs_sqlite_declared
1176 && let Some(ref provider) = session_fs_provider
1177 && provider.sqlite().is_none()
1178 {
1179 return Err(Error::with_message(
1180 ErrorKind::InvalidConfig,
1181 "SessionFs capabilities declare SQLite support but the provider \
1182 does not implement SessionFsSqliteProvider",
1183 ));
1184 }
1185
1186 let mut params = serde_json::to_value(&wire)?;
1187 let trace_ctx = self.resolve_trace_context().await;
1188 inject_trace_context(&mut params, &trace_ctx);
1189
1190 let capabilities = Arc::new(parking_lot::RwLock::new(SessionCapabilities::default()));
1191 let setup_start = Instant::now();
1192 let channels = self.register_session(&session_id);
1193 let idle_waiter = Arc::new(ParkingLotMutex::new(None));
1194 let open_canvases = Arc::new(parking_lot::RwLock::new(Vec::new()));
1195 let shutdown = CancellationToken::new();
1196 let (event_tx, _) = tokio::sync::broadcast::channel(512);
1197 let event_loop = spawn_event_loop(
1198 session_id.clone(),
1199 self.clone(),
1200 handlers,
1201 hooks,
1202 transforms,
1203 command_handlers,
1204 canvas_handler,
1205 session_fs_provider,
1206 bearer_token_providers,
1207 channels,
1208 idle_waiter.clone(),
1209 capabilities.clone(),
1210 open_canvases.clone(),
1211 event_tx.clone(),
1212 shutdown.clone(),
1213 );
1214 let mut registration =
1215 PendingSessionRegistration::new(self.clone(), session_id.clone(), shutdown.clone());
1216 tracing::debug!(
1217 elapsed_ms = setup_start.elapsed().as_millis(),
1218 session_id = %session_id,
1219 tools_count,
1220 commands_count,
1221 has_hooks,
1222 "Client::resume_session local setup complete"
1223 );
1224
1225 let rpc_start = Instant::now();
1226 let result = match self.call("session.resume", Some(params)).await {
1227 Ok(result) => result,
1228 Err(error) => {
1229 registration.cleanup(event_loop).await;
1230 return Err(error);
1231 }
1232 };
1233 tracing::debug!(
1234 elapsed_ms = rpc_start.elapsed().as_millis(),
1235 session_id = %session_id,
1236 "Client::resume_session session resume request completed successfully"
1237 );
1238
1239 let resume_result: ResumeSessionResult = match serde_json::from_value(result) {
1240 Ok(result) => result,
1241 Err(error) => {
1242 registration.cleanup(event_loop).await;
1243 return Err(error.into());
1244 }
1245 };
1246 let cli_session_id = resume_result
1247 .session_id
1248 .clone()
1249 .unwrap_or_else(|| session_id.clone());
1250 if cli_session_id != session_id {
1251 registration.cleanup(event_loop).await;
1252 return Err(ErrorKind::Session(SessionErrorKind::SessionIdMismatch {
1253 requested: session_id,
1254 returned: cli_session_id,
1255 })
1256 .into());
1257 }
1258 if has_mcp_auth_handler {
1259 register_mcp_auth_interest(self, &session_id).await?;
1260 }
1261
1262 let skills_reload_start = Instant::now();
1264 if let Err(e) = self
1265 .call(
1266 "session.skills.reload",
1267 Some(serde_json::json!({ "sessionId": session_id })),
1268 )
1269 .await
1270 {
1271 warn!(
1272 elapsed_ms = skills_reload_start.elapsed().as_millis(),
1273 session_id = %session_id,
1274 error = %e,
1275 "Client::resume_session skills reload request failed"
1276 );
1277 } else {
1278 tracing::debug!(
1279 elapsed_ms = skills_reload_start.elapsed().as_millis(),
1280 session_id = %session_id,
1281 "Client::resume_session skills reload request completed successfully"
1282 );
1283 }
1284
1285 *capabilities.write() = resume_result.capabilities.unwrap_or_default();
1286 {
1291 let mut snapshots = open_canvases.write();
1292 for snapshot in resume_result.open_canvases.unwrap_or_default() {
1293 upsert_open_canvas_snapshot(&mut snapshots, snapshot);
1294 }
1295 }
1296
1297 tracing::debug!(
1298 elapsed_ms = total_start.elapsed().as_millis(),
1299 session_id = %session_id,
1300 "Client::resume_session complete"
1301 );
1302 registration.disarm();
1303 let session = Session {
1304 id: session_id,
1305 cwd: self.cwd().clone(),
1306 workspace_path: resume_result.workspace_path,
1307 remote_url: resume_result.remote_url,
1308 client: self.clone(),
1309 event_loop: ParkingLotMutex::new(Some(event_loop)),
1310 shutdown,
1311 idle_waiter,
1312 capabilities,
1313 open_canvases,
1314 event_tx,
1315 };
1316 apply_mode_post_create_patch(
1317 &session,
1318 mode,
1319 opt_skip_custom_instructions,
1320 opt_custom_agents_local_only,
1321 opt_coauthor_enabled,
1322 opt_manage_schedule_enabled,
1323 )
1324 .await?;
1325 Ok(session)
1326 }
1327}
1328
1329type CommandHandlerMap = HashMap<String, Arc<dyn CommandHandler>>;
1330
1331async fn apply_mode_post_create_patch(
1332 session: &Session,
1333 mode: crate::ClientMode,
1334 opt_skip_custom_instructions: Option<bool>,
1335 opt_custom_agents_local_only: Option<bool>,
1336 opt_coauthor_enabled: Option<bool>,
1337 opt_manage_schedule_enabled: Option<bool>,
1338) -> Result<(), Error> {
1339 use crate::generated::api_types::SessionUpdateOptionsParams;
1340 let mut patch = SessionUpdateOptionsParams::default();
1341 let should_send = if mode == crate::ClientMode::Empty {
1342 patch.skip_custom_instructions = Some(opt_skip_custom_instructions.unwrap_or(true));
1343 patch.custom_agents_local_only = Some(opt_custom_agents_local_only.unwrap_or(true));
1344 patch.coauthor_enabled = Some(opt_coauthor_enabled.unwrap_or(false));
1345 patch.manage_schedule_enabled = Some(opt_manage_schedule_enabled.unwrap_or(false));
1346 patch.installed_plugins = Some(Vec::new());
1347 true
1348 } else {
1349 let mut any = false;
1350 if let Some(v) = opt_skip_custom_instructions {
1351 patch.skip_custom_instructions = Some(v);
1352 any = true;
1353 }
1354 if let Some(v) = opt_custom_agents_local_only {
1355 patch.custom_agents_local_only = Some(v);
1356 any = true;
1357 }
1358 if let Some(v) = opt_coauthor_enabled {
1359 patch.coauthor_enabled = Some(v);
1360 any = true;
1361 }
1362 if let Some(v) = opt_manage_schedule_enabled {
1363 patch.manage_schedule_enabled = Some(v);
1364 any = true;
1365 }
1366 any
1367 };
1368 if !should_send {
1369 return Ok(());
1370 }
1371 if let Err(error) = session.rpc().options().update(patch).await {
1372 let _ = session.disconnect().await;
1373 return Err(error);
1374 }
1375 Ok(())
1376}
1377
1378fn build_command_handler_map(commands: Option<&[CommandDefinition]>) -> Arc<CommandHandlerMap> {
1379 let map = match commands {
1380 Some(commands) => commands
1381 .iter()
1382 .filter(|cmd| !cmd.name.is_empty())
1383 .map(|cmd| (cmd.name.clone(), cmd.handler.clone()))
1384 .collect(),
1385 None => HashMap::new(),
1386 };
1387 Arc::new(map)
1388}
1389
1390fn upsert_open_canvas_snapshot(
1391 snapshots: &mut Vec<OpenCanvasInstance>,
1392 snapshot: OpenCanvasInstance,
1393) {
1394 if let Some(existing) = snapshots
1395 .iter_mut()
1396 .find(|open| open.instance_id == snapshot.instance_id)
1397 {
1398 *existing = snapshot;
1399 } else {
1400 snapshots.push(snapshot);
1401 }
1402}
1403
1404fn remove_open_canvas_snapshot(snapshots: &mut Vec<OpenCanvasInstance>, instance_id: &str) {
1405 snapshots.retain(|open| open.instance_id != instance_id);
1406}
1407
1408#[allow(clippy::too_many_arguments)]
1409fn spawn_event_loop(
1410 session_id: SessionId,
1411 client: Client,
1412 handlers: SessionHandlers,
1413 hooks: Option<Arc<dyn SessionHooks>>,
1414 transforms: Option<Arc<dyn SystemMessageTransform>>,
1415 command_handlers: Arc<CommandHandlerMap>,
1416 canvas_handler: Option<Arc<dyn CanvasHandler>>,
1417 session_fs_provider: Option<Arc<dyn SessionFsProvider>>,
1418 bearer_token_providers: HashMap<String, Arc<dyn BearerTokenProvider>>,
1419 channels: crate::router::SessionChannels,
1420 idle_waiter: Arc<ParkingLotMutex<Option<IdleWaiter>>>,
1421 capabilities: Arc<parking_lot::RwLock<SessionCapabilities>>,
1422 open_canvases: Arc<parking_lot::RwLock<Vec<OpenCanvasInstance>>>,
1423 event_tx: tokio::sync::broadcast::Sender<SessionEvent>,
1424 shutdown: CancellationToken,
1425) -> JoinHandle<()> {
1426 let crate::router::SessionChannels {
1427 mut notifications,
1428 mut requests,
1429 } = channels;
1430
1431 let span = tracing::error_span!("session_event_loop", session_id = %session_id);
1432 tokio::spawn(
1433 async move {
1434 loop {
1435 tokio::select! {
1446 _ = shutdown.cancelled() => break,
1447 Some(notification) = notifications.recv() => {
1448 handle_notification(
1449 &session_id, &client, &handlers, &command_handlers, notification, &idle_waiter, &capabilities, &open_canvases, &event_tx,
1450 ).await;
1451 }
1452 Some(request) = requests.recv() => {
1453 let ctx = RequestDispatchContext {
1454 client: &client,
1455 handlers: &handlers,
1456 hooks: hooks.as_deref(),
1457 transforms: transforms.as_deref(),
1458 canvas_handler: canvas_handler.as_ref(),
1459 session_fs_provider: session_fs_provider.as_ref(),
1460 bearer_token_providers: &bearer_token_providers,
1461 };
1462 handle_request(&session_id, ctx, request).await;
1463 }
1464 else => break,
1465 }
1466 }
1467 if let Some(waiter) = idle_waiter.lock().take() {
1470 let _ = waiter
1471 .tx
1472 .send(Err(ErrorKind::Session(SessionErrorKind::EventLoopClosed).into()));
1473 }
1474 }
1475 .instrument(span),
1476 )
1477}
1478
1479fn extract_request_id(data: &Value) -> Option<RequestId> {
1480 data.get("requestId")
1481 .and_then(|v| v.as_str())
1482 .filter(|s| !s.is_empty())
1483 .map(RequestId::new)
1484}
1485
1486fn notification_permission_payload(result: &PermissionResult) -> Option<Value> {
1491 match result {
1492 PermissionResult::NoResult => None,
1493 PermissionResult::Decision(decision) => Some(
1494 serde_json::to_value(decision).expect("serializing permission decision should succeed"),
1495 ),
1496 }
1497}
1498
1499async fn register_mcp_auth_interest(client: &Client, session_id: &SessionId) -> Result<(), Error> {
1500 let mut params = serde_json::to_value(RegisterEventInterestParams {
1501 event_type: "mcp.oauth_required".to_string(),
1502 })?;
1503 params["sessionId"] = Value::String(session_id.to_string());
1504 client
1505 .call(rpc_methods::SESSION_EVENTLOG_REGISTERINTEREST, Some(params))
1506 .await?;
1507 Ok(())
1508}
1509
1510fn tool_failure_result(message: impl Into<String>) -> ToolResult {
1511 let message = message.into();
1512 ToolResult::Expanded(ToolResultExpanded {
1513 text_result_for_llm: message.clone(),
1514 result_type: "failure".to_string(),
1515 binary_results_for_llm: None,
1516 session_log: None,
1517 error: Some(message),
1518 tool_telemetry: None,
1519 })
1520}
1521
1522#[allow(clippy::too_many_arguments)]
1524async fn handle_notification(
1525 session_id: &SessionId,
1526 client: &Client,
1527 handlers: &SessionHandlers,
1528 command_handlers: &Arc<CommandHandlerMap>,
1529 notification: SessionEventNotification,
1530 idle_waiter: &Arc<ParkingLotMutex<Option<IdleWaiter>>>,
1531 capabilities: &Arc<parking_lot::RwLock<SessionCapabilities>>,
1532 open_canvases: &Arc<parking_lot::RwLock<Vec<OpenCanvasInstance>>>,
1533 event_tx: &tokio::sync::broadcast::Sender<SessionEvent>,
1534) {
1535 let dispatch_start = Instant::now();
1536 let event = notification.event.clone();
1537 let event_type = event.parsed_type();
1538 if event_type == SessionEventType::PermissionRequested {
1539 tracing::debug!(
1540 session_id = %session_id,
1541 event_type = %event.event_type,
1542 "Session::handle_notification permission request received"
1543 );
1544 }
1545
1546 match event_type {
1549 SessionEventType::AssistantMessage
1550 | SessionEventType::SessionIdle
1551 | SessionEventType::SessionError => {
1552 let mut guard = idle_waiter.lock();
1553 if let Some(waiter) = guard.as_mut() {
1554 match event_type {
1555 SessionEventType::AssistantMessage => {
1556 if !waiter.first_assistant_message_seen {
1557 waiter.first_assistant_message_seen = true;
1558 tracing::debug!(
1559 elapsed_ms = waiter.started_at.elapsed().as_millis(),
1560 session_id = %session_id,
1561 "Session::send_and_wait first assistant message"
1562 );
1563 }
1564 waiter.last_assistant_message = Some(event.clone());
1565 }
1566 SessionEventType::SessionIdle | SessionEventType::SessionError => {
1567 if let Some(waiter) = guard.take() {
1568 if event_type == SessionEventType::SessionIdle {
1569 tracing::debug!(
1570 elapsed_ms = waiter.started_at.elapsed().as_millis(),
1571 session_id = %session_id,
1572 "Session::send_and_wait idle received"
1573 );
1574 let _ = waiter.tx.send(Ok(waiter.last_assistant_message));
1575 } else {
1576 let error_msg = event
1577 .typed_data::<SessionErrorData>()
1578 .map(|d| d.message)
1579 .or_else(|| {
1580 event
1581 .data
1582 .get("message")
1583 .and_then(|v| v.as_str())
1584 .map(|s| s.to_string())
1585 })
1586 .unwrap_or_else(|| "session error".to_string());
1587 let _ = waiter.tx.send(Err(Error::with_message(
1588 ErrorKind::Session(SessionErrorKind::AgentError),
1589 error_msg,
1590 )));
1591 }
1592 }
1593 }
1594 _ => {}
1595 }
1596 }
1597 }
1598 _ => {}
1599 }
1600
1601 if event_type == SessionEventType::CapabilitiesChanged {
1605 match serde_json::from_value::<SessionCapabilities>(notification.event.data.clone()) {
1606 Ok(changed) => *capabilities.write() = changed,
1607 Err(e) => warn!(error = %e, "failed to deserialize capabilities.changed payload"),
1608 }
1609 }
1610 if event_type == SessionEventType::SessionCanvasOpened {
1611 match serde_json::from_value::<OpenCanvasInstance>(notification.event.data.clone()) {
1612 Ok(open_canvas) => {
1613 upsert_open_canvas_snapshot(&mut open_canvases.write(), open_canvas);
1614 }
1615 Err(e) => warn!(error = %e, "failed to deserialize session.canvas.opened payload"),
1616 }
1617 }
1618 if event_type == SessionEventType::SessionCanvasClosed {
1619 match serde_json::from_value::<SessionCanvasClosedData>(notification.event.data.clone()) {
1620 Ok(closed) => {
1621 if closed.instance_id.is_empty() {
1622 warn!("failed to deserialize session.canvas.closed payload");
1623 } else {
1624 remove_open_canvas_snapshot(&mut open_canvases.write(), &closed.instance_id);
1625 }
1626 }
1627 Err(e) => warn!(error = %e, "failed to deserialize session.canvas.closed payload"),
1628 }
1629 }
1630
1631 let _ = event_tx.send(event.clone());
1635
1636 tracing::debug!(
1637 elapsed_ms = dispatch_start.elapsed().as_millis(),
1638 session_id = %session_id,
1639 event_type = %notification.event.event_type,
1640 "Session::handle_notification dispatch"
1641 );
1642
1643 match event_type {
1646 SessionEventType::PermissionRequested => {
1647 let Some(request_id) = extract_request_id(¬ification.event.data) else {
1648 return;
1649 };
1650 if notification
1654 .event
1655 .data
1656 .get("resolvedByHook")
1657 .and_then(|v| v.as_bool())
1658 .unwrap_or(false)
1659 {
1660 return;
1661 }
1662 let Some(permission_handler) = handlers.permission.clone() else {
1666 return;
1667 };
1668 let client = client.clone();
1669 let sid = session_id.clone();
1670 let data: PermissionRequestData =
1671 serde_json::from_value(notification.event.data.clone()).unwrap_or_else(|_| {
1672 PermissionRequestData {
1673 kind: None,
1674 tool_call_id: None,
1675 extra: notification.event.data.clone(),
1676 }
1677 });
1678 let span = tracing::error_span!(
1679 "permission_request_handler",
1680 session_id = %sid,
1681 request_id = %request_id
1682 );
1683 tokio::spawn(
1684 async move {
1685 let handler_start = Instant::now();
1686 let result = permission_handler
1687 .handle(sid.clone(), request_id.clone(), data)
1688 .await;
1689 tracing::debug!(
1690 elapsed_ms = handler_start.elapsed().as_millis(),
1691 session_id = %sid,
1692 request_id = %request_id,
1693 "PermissionHandler::handle dispatch"
1694 );
1695 let Some(result_value) = notification_permission_payload(&result) else {
1696 return;
1700 };
1701 let rpc_start = Instant::now();
1702 let _ = client
1703 .call(
1704 "session.permissions.handlePendingPermissionRequest",
1705 Some(serde_json::json!({
1706 "sessionId": sid,
1707 "requestId": request_id,
1708 "result": result_value,
1709 })),
1710 )
1711 .await;
1712 tracing::debug!(
1713 elapsed_ms = rpc_start.elapsed().as_millis(),
1714 session_id = %sid,
1715 request_id = %request_id,
1716 "Session::handle_notification response sent successfully"
1717 );
1718 }
1719 .instrument(span),
1720 );
1721 }
1722 SessionEventType::ExternalToolRequested => {
1723 let Some(request_id) = extract_request_id(¬ification.event.data) else {
1724 return;
1725 };
1726 let data: ExternalToolRequestedData =
1727 match serde_json::from_value(notification.event.data.clone()) {
1728 Ok(d) => d,
1729 Err(e) => {
1730 warn!(error = %e, "failed to deserialize external_tool.requested");
1731 let client = client.clone();
1732 let sid = session_id.clone();
1733 let span = tracing::error_span!(
1734 "external_tool_deserialize_error",
1735 session_id = %sid,
1736 request_id = %request_id
1737 );
1738 tokio::spawn(
1739 async move {
1740 let rpc_start = Instant::now();
1741 let _ = client
1742 .call(
1743 "session.tools.handlePendingToolCall",
1744 Some(serde_json::json!({
1745 "sessionId": sid,
1746 "requestId": request_id,
1747 "error": format!("Failed to deserialize tool request: {e}"),
1748 })),
1749 )
1750 .await;
1751 tracing::debug!(
1752 elapsed_ms = rpc_start.elapsed().as_millis(),
1753 session_id = %sid,
1754 request_id = %request_id,
1755 "Session::handle_notification response sent successfully"
1756 );
1757 }
1758 .instrument(span),
1759 );
1760 return;
1761 }
1762 };
1763 let tool_handler = if data.tool_name.is_empty() {
1767 None
1768 } else {
1769 handlers.tools.get(&data.tool_name).cloned()
1770 };
1771 let Some(tool_handler) = tool_handler else {
1772 return;
1773 };
1774 let client = client.clone();
1775 let sid = session_id.clone();
1776 let span = tracing::error_span!(
1777 "external_tool_handler",
1778 session_id = %sid,
1779 request_id = %request_id
1780 );
1781 tokio::spawn(
1782 async move {
1783 if data.tool_call_id.is_empty() {
1788 let error_msg = "Missing toolCallId";
1789 let rpc_start = Instant::now();
1790 let _ = client
1791 .call(
1792 "session.tools.handlePendingToolCall",
1793 Some(serde_json::json!({
1794 "sessionId": sid,
1795 "requestId": request_id,
1796 "error": error_msg,
1797 })),
1798 )
1799 .await;
1800 tracing::debug!(
1801 elapsed_ms = rpc_start.elapsed().as_millis(),
1802 session_id = %sid,
1803 request_id = %request_id,
1804 "Session::handle_notification response sent successfully"
1805 );
1806 return;
1807 }
1808 let tool_call_id = data.tool_call_id.clone();
1809 let tool_name = data.tool_name.clone();
1810 let invocation = ToolInvocation {
1811 session_id: sid.clone(),
1812 tool_call_id: data.tool_call_id,
1813 tool_name: data.tool_name,
1814 arguments: data
1815 .arguments
1816 .unwrap_or(Value::Object(serde_json::Map::new())),
1817 traceparent: data.traceparent,
1818 tracestate: data.tracestate,
1819 };
1820 let handler_start = Instant::now();
1821 let tool_result = match tool_handler.call(invocation).await {
1822 Ok(r) => r,
1823 Err(e) => tool_failure_result(e.to_string()),
1824 };
1825 tracing::debug!(
1826 elapsed_ms = handler_start.elapsed().as_millis(),
1827 session_id = %sid,
1828 request_id = %request_id,
1829 tool_call_id = %tool_call_id,
1830 tool_name = %tool_name,
1831 "ToolHandler::call dispatch"
1832 );
1833 let result_value = serde_json::to_value(tool_result).unwrap_or(Value::Null);
1834 let rpc_start = Instant::now();
1835 let _ = client
1836 .call(
1837 "session.tools.handlePendingToolCall",
1838 Some(serde_json::json!({
1839 "sessionId": sid,
1840 "requestId": request_id,
1841 "result": result_value,
1842 })),
1843 )
1844 .await;
1845 tracing::debug!(
1846 elapsed_ms = rpc_start.elapsed().as_millis(),
1847 session_id = %sid,
1848 request_id = %request_id,
1849 tool_call_id = %tool_call_id,
1850 tool_name = %tool_name,
1851 "Session::handle_notification response sent successfully"
1852 );
1853 }
1854 .instrument(span),
1855 );
1856 }
1857 SessionEventType::UserInputRequested => {
1858 }
1865 SessionEventType::ElicitationRequested => {
1866 let Some(request_id) = extract_request_id(¬ification.event.data) else {
1867 return;
1868 };
1869 let Some(elicitation_handler) = handlers.elicitation.clone() else {
1873 return;
1874 };
1875 let elicitation_data: ElicitationRequestedData =
1876 match serde_json::from_value(notification.event.data.clone()) {
1877 Ok(d) => d,
1878 Err(e) => {
1879 warn!(error = %e, "failed to deserialize elicitation request");
1880 return;
1881 }
1882 };
1883 let request = ElicitationRequest {
1884 message: elicitation_data.message,
1885 requested_schema: elicitation_data
1886 .requested_schema
1887 .map(|s| serde_json::to_value(s).unwrap_or(Value::Null)),
1888 mode: elicitation_data.mode.map(|m| match m {
1889 crate::generated::session_events::ElicitationRequestedMode::Form => {
1890 crate::types::ElicitationMode::Form
1891 }
1892 crate::generated::session_events::ElicitationRequestedMode::Url => {
1893 crate::types::ElicitationMode::Url
1894 }
1895 _ => crate::types::ElicitationMode::Unknown,
1896 }),
1897 elicitation_source: elicitation_data.elicitation_source,
1898 url: elicitation_data.url,
1899 };
1900 let client = client.clone();
1901 let sid = session_id.clone();
1902 let span = tracing::error_span!(
1903 "elicitation_request_handler",
1904 session_id = %sid,
1905 request_id = %request_id
1906 );
1907 tokio::spawn(
1908 async move {
1909 let cancel = ElicitationResult {
1910 action: "cancel".to_string(),
1911 content: None,
1912 };
1913 let handler_task = tokio::spawn({
1915 let sid = sid.clone();
1916 let request_id = request_id.clone();
1917 let span = tracing::error_span!(
1918 "elicitation_callback",
1919 session_id = %sid,
1920 request_id = %request_id
1921 );
1922 async move {
1923 let handler_start = Instant::now();
1924 let response = elicitation_handler
1925 .handle(sid.clone(), request_id.clone(), request)
1926 .await;
1927 tracing::debug!(
1928 elapsed_ms = handler_start.elapsed().as_millis(),
1929 session_id = %sid,
1930 request_id = %request_id,
1931 "ElicitationHandler::handle dispatch"
1932 );
1933 response
1934 }
1935 .instrument(span)
1936 });
1937 let result = match handler_task.await {
1938 Ok(r) => r,
1939 Err(_) => cancel.clone(),
1940 };
1941 let rpc_start = Instant::now();
1942 if let Err(e) = client
1943 .call(
1944 "session.ui.handlePendingElicitation",
1945 Some(serde_json::json!({
1946 "sessionId": sid,
1947 "requestId": request_id,
1948 "result": result,
1949 })),
1950 )
1951 .await
1952 {
1953 warn!(error = %e, "handlePendingElicitation failed, sending cancel");
1955 let _ = client
1956 .call(
1957 "session.ui.handlePendingElicitation",
1958 Some(serde_json::json!({
1959 "sessionId": sid,
1960 "requestId": request_id,
1961 "result": cancel,
1962 })),
1963 )
1964 .await;
1965 } else {
1966 tracing::debug!(
1967 elapsed_ms = rpc_start.elapsed().as_millis(),
1968 session_id = %sid,
1969 request_id = %request_id,
1970 "Session::handle_notification response sent successfully"
1971 );
1972 }
1973 }
1974 .instrument(span),
1975 );
1976 }
1977 SessionEventType::McpOauthRequired => {
1978 let Some(request_id) = extract_request_id(¬ification.event.data) else {
1979 return;
1980 };
1981 let Some(mcp_auth_handler) = handlers.mcp_auth.clone() else {
1982 warn!(
1983 session_id = %session_id,
1984 request_id = %request_id,
1985 "received MCP OAuth request without a registered MCP auth handler"
1986 );
1987 return;
1988 };
1989 let data: McpOauthRequiredData =
1990 match serde_json::from_value(notification.event.data.clone()) {
1991 Ok(d) => d,
1992 Err(e) => {
1993 warn!(error = %e, "failed to deserialize MCP OAuth request");
1994 return;
1995 }
1996 };
1997 let request = McpAuthRequest {
1998 request_id: request_id.clone(),
1999 server_name: data.server_name,
2000 server_url: data.server_url,
2001 reason: data.reason,
2002 www_authenticate_params: data.www_authenticate_params,
2003 resource_metadata: data.resource_metadata,
2004 static_client_config: data.static_client_config,
2005 };
2006 let client = client.clone();
2007 let sid = session_id.clone();
2008 let span = tracing::error_span!(
2009 "mcp_auth_request_handler",
2010 session_id = %sid,
2011 request_id = %request_id
2012 );
2013 tokio::spawn(
2014 async move {
2015 let cancel = McpAuthResult::Cancelled;
2016 let handler_task = tokio::spawn({
2017 let sid = sid.clone();
2018 let request_id = request_id.clone();
2019 let span = tracing::error_span!(
2020 "mcp_auth_callback",
2021 session_id = %sid,
2022 request_id = %request_id
2023 );
2024 async move {
2025 let handler_start = Instant::now();
2026 let response = mcp_auth_handler
2027 .handle(sid.clone(), request_id.clone(), request)
2028 .await;
2029 tracing::debug!(
2030 elapsed_ms = handler_start.elapsed().as_millis(),
2031 session_id = %sid,
2032 request_id = %request_id,
2033 "McpAuthHandler::handle dispatch"
2034 );
2035 response
2036 }
2037 .instrument(span)
2038 });
2039 let result = match handler_task.await {
2040 Ok(result) => result,
2041 Err(_) => cancel,
2042 };
2043 let rpc_start = Instant::now();
2044 let _ = client
2045 .call(
2046 "session.mcp.oauth.handlePendingRequest",
2047 Some(serde_json::json!({
2048 "sessionId": sid,
2049 "requestId": request_id,
2050 "result": result.into_wire(),
2051 })),
2052 )
2053 .await;
2054 tracing::debug!(
2055 elapsed_ms = rpc_start.elapsed().as_millis(),
2056 "Session::handle_notification MCP auth response sent"
2057 );
2058 }
2059 .instrument(span),
2060 );
2061 }
2062 SessionEventType::CommandExecute => {
2063 let data: CommandExecuteData =
2064 match serde_json::from_value(notification.event.data.clone()) {
2065 Ok(d) => d,
2066 Err(e) => {
2067 warn!(error = %e, "failed to deserialize command.execute");
2068 return;
2069 }
2070 };
2071 let client = client.clone();
2072 let command_handlers = command_handlers.clone();
2073 let sid = session_id.clone();
2074 let span = tracing::error_span!("command_handler", session_id = %sid);
2075 tokio::spawn(
2076 async move {
2077 let request_id = data.request_id;
2078 let ack_error = match command_handlers.get(&data.command_name).cloned() {
2079 None => Some(format!("Unknown command: {}", data.command_name)),
2080 Some(handler) => {
2081 let command_name = data.command_name.clone();
2082 let ctx = CommandContext {
2083 session_id: sid.clone(),
2084 command: data.command,
2085 command_name: data.command_name,
2086 args: data.args,
2087 };
2088 let handler_start = Instant::now();
2089 let result = handler.on_command(ctx).await;
2090 tracing::debug!(
2091 elapsed_ms = handler_start.elapsed().as_millis(),
2092 session_id = %sid,
2093 request_id = %request_id,
2094 command_name = %command_name,
2095 "CommandHandler::call dispatch"
2096 );
2097 match result {
2098 Ok(()) => None,
2099 Err(e) => Some(e.to_string()),
2100 }
2101 }
2102 };
2103 let mut params = serde_json::json!({
2104 "sessionId": sid,
2105 "requestId": request_id,
2106 });
2107 if let Some(error_msg) = ack_error {
2108 params["error"] = serde_json::Value::String(error_msg);
2109 }
2110 let rpc_start = Instant::now();
2111 let _ = client
2112 .call("session.commands.handlePendingCommand", Some(params))
2113 .await;
2114 tracing::debug!(
2115 elapsed_ms = rpc_start.elapsed().as_millis(),
2116 session_id = %sid,
2117 request_id = %request_id,
2118 "Session::handle_notification response sent successfully"
2119 );
2120 }
2121 .instrument(span),
2122 );
2123 }
2124 _ => {}
2125 }
2126}
2127
2128struct RequestDispatchContext<'a> {
2129 client: &'a Client,
2130 handlers: &'a SessionHandlers,
2131 hooks: Option<&'a dyn SessionHooks>,
2132 transforms: Option<&'a dyn SystemMessageTransform>,
2133 canvas_handler: Option<&'a Arc<dyn CanvasHandler>>,
2134 session_fs_provider: Option<&'a Arc<dyn SessionFsProvider>>,
2135 bearer_token_providers: &'a HashMap<String, Arc<dyn BearerTokenProvider>>,
2136}
2137
2138async fn handle_request(
2140 session_id: &SessionId,
2141 ctx: RequestDispatchContext<'_>,
2142 request: crate::JsonRpcRequest,
2143) {
2144 let sid = session_id.clone();
2145 let client = ctx.client;
2146 let handlers = ctx.handlers;
2147 let hooks = ctx.hooks;
2148 let transforms = ctx.transforms;
2149 let canvas_handler = ctx.canvas_handler;
2150 let session_fs_provider = ctx.session_fs_provider;
2151 let bearer_token_providers = ctx.bearer_token_providers;
2152
2153 if request.method.starts_with("sessionFs.") {
2154 crate::session_fs_dispatch::dispatch(client, session_fs_provider, request).await;
2155 return;
2156 }
2157
2158 if request.method.starts_with("canvas.") {
2159 crate::canvas_dispatch::dispatch(client, canvas_handler, request).await;
2160 return;
2161 }
2162
2163 if request.method == crate::generated::api_types::rpc_methods::PROVIDERTOKEN_GETTOKEN {
2164 crate::provider_token_dispatch::dispatch(client, bearer_token_providers, request).await;
2165 return;
2166 }
2167
2168 match request.method.as_str() {
2169 "hooks.invoke" => {
2170 let params = request.params.as_ref();
2171 let hook_type = params
2172 .and_then(|p| p.get("hookType"))
2173 .and_then(|v| v.as_str())
2174 .unwrap_or("");
2175 let input = params
2176 .and_then(|p| p.get("input"))
2177 .cloned()
2178 .unwrap_or(Value::Object(Default::default()));
2179
2180 let rpc_result = if let Some(hooks) = hooks {
2181 match crate::hooks::dispatch_hook(hooks, &sid, hook_type, input).await {
2182 Ok(output) => output,
2183 Err(e) => {
2184 warn!(error = %e, hook_type = hook_type, "hook dispatch failed");
2185 serde_json::json!({ "output": {} })
2186 }
2187 }
2188 } else {
2189 serde_json::json!({ "output": {} })
2190 };
2191
2192 let rpc_response = JsonRpcResponse {
2193 jsonrpc: "2.0".to_string(),
2194 id: request.id,
2195 result: Some(rpc_result),
2196 error: None,
2197 };
2198 let _ = client.send_response(&rpc_response).await;
2199 }
2200
2201 "userInput.request" => {
2202 let params = request.params.as_ref();
2203 let Some(question) = params
2204 .and_then(|p| p.get("question"))
2205 .and_then(|v| v.as_str())
2206 else {
2207 warn!("userInput.request missing 'question' field");
2208 let rpc_response = JsonRpcResponse {
2209 jsonrpc: "2.0".to_string(),
2210 id: request.id,
2211 result: None,
2212 error: Some(crate::JsonRpcError {
2213 code: error_codes::INVALID_PARAMS,
2214 message: "missing required field: question".to_string(),
2215 data: None,
2216 }),
2217 };
2218 let _ = client.send_response(&rpc_response).await;
2219 return;
2220 };
2221 let question = question.to_string();
2222 let choices = params
2223 .and_then(|p| p.get("choices"))
2224 .and_then(|v| v.as_array())
2225 .map(|arr| {
2226 arr.iter()
2227 .filter_map(|v| v.as_str().map(|s| s.to_string()))
2228 .collect()
2229 });
2230 let allow_freeform = params
2231 .and_then(|p| p.get("allowFreeform"))
2232 .and_then(|v| v.as_bool());
2233
2234 let handler_start = Instant::now();
2235 let response = if let Some(user_input_handler) = handlers.user_input.as_ref() {
2236 user_input_handler
2237 .handle(sid.clone(), question, choices, allow_freeform)
2238 .await
2239 } else {
2240 None
2241 };
2242 tracing::debug!(
2243 elapsed_ms = handler_start.elapsed().as_millis(),
2244 session_id = %sid,
2245 "UserInputHandler::handle dispatch"
2246 );
2247
2248 let rpc_result = match response {
2249 Some(UserInputResponse {
2250 answer,
2251 was_freeform,
2252 }) => serde_json::json!({
2253 "answer": answer,
2254 "wasFreeform": was_freeform,
2255 }),
2256 None => serde_json::json!({ "noResponse": true }),
2257 };
2258 let rpc_response = JsonRpcResponse {
2259 jsonrpc: "2.0".to_string(),
2260 id: request.id,
2261 result: Some(rpc_result),
2262 error: None,
2263 };
2264 let _ = client.send_response(&rpc_response).await;
2265 }
2266
2267 "exitPlanMode.request" => {
2268 let params = request
2269 .params
2270 .as_ref()
2271 .cloned()
2272 .unwrap_or(Value::Object(serde_json::Map::new()));
2273 let data: ExitPlanModeData = match serde_json::from_value(params) {
2274 Ok(d) => d,
2275 Err(e) => {
2276 warn!(error = %e, "failed to deserialize exitPlanMode.request params, using defaults");
2277 ExitPlanModeData::default()
2278 }
2279 };
2280
2281 let rpc_result = if let Some(exit_plan_handler) = handlers.exit_plan_mode.as_ref() {
2282 let result = exit_plan_handler.handle(sid, data).await;
2283 serde_json::to_value(result).expect("ExitPlanModeResult serialization cannot fail")
2284 } else {
2285 serde_json::json!({ "approved": true })
2286 };
2287 let rpc_response = JsonRpcResponse {
2288 jsonrpc: "2.0".to_string(),
2289 id: request.id,
2290 result: Some(rpc_result),
2291 error: None,
2292 };
2293 let _ = client.send_response(&rpc_response).await;
2294 }
2295
2296 "autoModeSwitch.request" => {
2297 let error_code = request
2298 .params
2299 .as_ref()
2300 .and_then(|p| p.get("errorCode"))
2301 .and_then(|v| v.as_str())
2302 .map(|s| s.to_string());
2303 let retry_after_seconds = request
2304 .params
2305 .as_ref()
2306 .and_then(|p| p.get("retryAfterSeconds"))
2307 .and_then(|v| v.as_f64());
2308
2309 let answer = if let Some(auto_mode_handler) = handlers.auto_mode_switch.as_ref() {
2310 auto_mode_handler
2311 .handle(sid, error_code, retry_after_seconds)
2312 .await
2313 } else {
2314 AutoModeSwitchResponse::No
2315 };
2316 let rpc_response = JsonRpcResponse {
2317 jsonrpc: "2.0".to_string(),
2318 id: request.id,
2319 result: Some(serde_json::json!({ "response": answer })),
2320 error: None,
2321 };
2322 let _ = client.send_response(&rpc_response).await;
2323 }
2324
2325 "systemMessage.transform" => {
2326 let params = request.params.as_ref();
2327 let sections: HashMap<String, crate::transforms::TransformSection> =
2328 match params.and_then(|p| p.get("sections")) {
2329 Some(v) => match serde_json::from_value(v.clone()) {
2330 Ok(s) => s,
2331 Err(e) => {
2332 let _ = send_error_response(
2333 client,
2334 request.id,
2335 error_codes::INVALID_PARAMS,
2336 &format!("invalid sections: {e}"),
2337 )
2338 .await;
2339 return;
2340 }
2341 },
2342 None => {
2343 let _ = send_error_response(
2344 client,
2345 request.id,
2346 error_codes::INVALID_PARAMS,
2347 "missing sections parameter",
2348 )
2349 .await;
2350 return;
2351 }
2352 };
2353
2354 let rpc_result = if let Some(transforms) = transforms {
2355 let transform_start = Instant::now();
2356 let response =
2357 crate::transforms::dispatch_transform(transforms, &sid, sections).await;
2358 tracing::debug!(
2359 elapsed_ms = transform_start.elapsed().as_millis(),
2360 session_id = %sid,
2361 "SystemMessageTransform::transform_section dispatch"
2362 );
2363 match serde_json::to_value(response) {
2364 Ok(v) => v,
2365 Err(e) => {
2366 warn!(error = %e, "failed to serialize transform response");
2367 serde_json::json!({ "sections": {} })
2368 }
2369 }
2370 } else {
2371 let passthrough: HashMap<String, crate::transforms::TransformSection> = sections;
2373 serde_json::json!({ "sections": passthrough })
2374 };
2375
2376 let rpc_response = JsonRpcResponse {
2377 jsonrpc: "2.0".to_string(),
2378 id: request.id,
2379 result: Some(rpc_result),
2380 error: None,
2381 };
2382 let _ = client.send_response(&rpc_response).await;
2383 }
2384
2385 method => {
2386 warn!(
2387 method = method,
2388 "unhandled request method in session event loop"
2389 );
2390 let _ = send_error_response(
2391 client,
2392 request.id,
2393 error_codes::METHOD_NOT_FOUND,
2394 &format!("unknown method: {method}"),
2395 )
2396 .await;
2397 }
2398 }
2399}
2400
2401async fn send_error_response(
2402 client: &Client,
2403 id: u64,
2404 code: i32,
2405 message: &str,
2406) -> Result<(), Error> {
2407 let response = JsonRpcResponse {
2408 jsonrpc: "2.0".to_string(),
2409 id,
2410 result: None,
2411 error: Some(crate::JsonRpcError {
2412 code,
2413 message: message.to_string(),
2414 data: None,
2415 }),
2416 };
2417 client.send_response(&response).await
2418}
2419
2420fn apply_transform_sections(
2424 sys_msg: &mut SystemMessageConfig,
2425 transforms: &dyn SystemMessageTransform,
2426) {
2427 sys_msg.mode = Some("customize".to_string());
2428 let sections = sys_msg.sections.get_or_insert_with(HashMap::new);
2429 for id in transforms.section_ids() {
2430 sections.entry(id).or_insert_with(|| SectionOverride {
2431 action: Some("transform".to_string()),
2432 content: None,
2433 });
2434 }
2435}
2436
2437fn inject_transform_sections(config: &mut SessionConfig, transforms: &dyn SystemMessageTransform) {
2438 let sys_msg = config.system_message.get_or_insert_with(Default::default);
2439 apply_transform_sections(sys_msg, transforms);
2440}
2441
2442fn inject_transform_sections_resume(
2443 config: &mut ResumeSessionConfig,
2444 transforms: &dyn SystemMessageTransform,
2445) {
2446 let sys_msg = config.system_message.get_or_insert_with(Default::default);
2447 apply_transform_sections(sys_msg, transforms);
2448}
2449
2450#[cfg(test)]
2451mod tests {
2452 use serde_json::json;
2453
2454 use super::notification_permission_payload;
2455 use crate::handler::PermissionResult;
2456
2457 #[test]
2458 fn notification_payload_suppresses_no_result() {
2459 assert!(notification_permission_payload(&PermissionResult::NoResult).is_none());
2460 }
2461
2462 #[test]
2463 fn notification_payload_serializes_decisions() {
2464 assert_eq!(
2465 notification_permission_payload(&PermissionResult::approve_once()),
2466 Some(json!({ "kind": "approve-once" }))
2467 );
2468 assert_eq!(
2469 notification_permission_payload(&PermissionResult::reject(None)),
2470 Some(json!({ "kind": "reject" }))
2471 );
2472 assert_eq!(
2473 notification_permission_payload(&PermissionResult::reject(Some("bad".to_string()))),
2474 Some(json!({ "kind": "reject", "feedback": "bad" }))
2475 );
2476 assert_eq!(
2477 notification_permission_payload(&PermissionResult::user_not_available()),
2478 Some(json!({ "kind": "user-not-available" }))
2479 );
2480 }
2481}