1use std::collections::HashMap;
2use std::panic::AssertUnwindSafe;
3use std::path::{Path, PathBuf};
4use std::sync::Arc;
5use std::time::{Duration, Instant};
6
7use futures_util::FutureExt;
8use parking_lot::Mutex as ParkingLotMutex;
9use serde_json::Value;
10use tokio::sync::oneshot;
11use tokio::task::JoinHandle;
12use tokio_util::sync::CancellationToken;
13use tracing::{Instrument, error, warn};
14
15use crate::canvas::CanvasHandler;
16use crate::generated::api_types::{
17 LogRequest, ModelSwitchAutoTierRequest, ModelSwitchAutoTierResult, ModelSwitchToRequest,
18 OpenCanvasInstance, PermissionDecisionRequest, RegisterEventInterestParams,
19 ToolsGetCurrentMetadataResult, rpc_methods,
20};
21use crate::generated::session_events::{
22 CommandExecuteData, ElicitationRequestedData, ExternalToolRequestedData, McpOauthRequiredData,
23 SessionCanvasClosedData, SessionErrorData, SessionEventType, SessionIdleData, SessionMode,
24};
25use crate::handler::{
26 AutoModeSwitchHandler, AutoModeSwitchResponse, ElicitationHandler, ExitPlanModeHandler,
27 McpAuthHandler, McpAuthRequest, McpAuthResult, PermissionHandler, PermissionResult,
28 UserInputHandler, UserInputResponse,
29};
30use crate::hooks::SessionHooks;
31use crate::provider_token::BearerTokenProvider;
32use crate::session_fs::SessionFsProvider;
33use crate::trace_context::inject_trace_context;
34use crate::transforms::SystemMessageTransform;
35use crate::types::{
36 AutoTier, AutoTierPreference, CommandContext, CommandDefinition, CommandHandler,
37 CreateSessionResult, ElicitationRequest, ElicitationResult, ExitPlanModeData,
38 GetMessagesResponse, MessageOptions, PermissionRequestData, RequestId, ResumeSessionConfig,
39 ResumeSessionResult, SectionOverride, SessionCapabilities, SessionConfig, SessionEvent,
40 SessionId, SetModelOptions, SystemMessageConfig, ToolInvocation, ToolResult,
41 ToolResultExpanded, TraceContext, UiInputOptions, ensure_attachment_display_names,
42};
43use crate::{
44 Client, Error, ErrorKind, JsonRpcResponse, SessionErrorKind, SessionEventNotification,
45 error_codes,
46};
47
48const TOOL_SEARCH_TOOL_NAME: &str = "tool_search_tool";
52
53pub const DEFAULT_EVENT_BUFFER_CAPACITY: usize = 512;
61
62fn resolve_event_buffer_capacity(capacity: Option<usize>) -> Result<usize, Error> {
68 match capacity {
69 Some(0) => Err(Error::with_message(
70 ErrorKind::InvalidConfig,
71 "event_buffer_capacity must be greater than zero",
72 )),
73 Some(capacity) => Ok(capacity),
74 None => Ok(DEFAULT_EVENT_BUFFER_CAPACITY),
75 }
76}
77
78#[derive(Clone)]
86pub(crate) struct SessionHandlers {
87 pub permission: Option<Arc<dyn PermissionHandler>>,
88 pub managed_settings_enabled: bool,
89 pub elicitation: Option<Arc<dyn ElicitationHandler>>,
90 pub mcp_auth: Option<Arc<dyn McpAuthHandler>>,
91 pub user_input: Option<Arc<dyn UserInputHandler>>,
92 pub exit_plan_mode: Option<Arc<dyn ExitPlanModeHandler>>,
93 pub auto_mode_switch: Option<Arc<dyn AutoModeSwitchHandler>>,
94 pub tools: Arc<HashMap<String, Arc<dyn crate::tool::ToolHandler>>>,
95}
96
97type PendingExternalTools = Arc<ParkingLotMutex<HashMap<RequestId, Arc<CancellationToken>>>>;
98
99struct PendingExternalToolGuard {
100 request_id: RequestId,
101 token: Arc<CancellationToken>,
102 pending: PendingExternalTools,
103}
104
105impl Drop for PendingExternalToolGuard {
106 fn drop(&mut self) {
107 let mut pending = self.pending.lock();
108 if pending
109 .get(&self.request_id)
110 .is_some_and(|token| Arc::ptr_eq(token, &self.token))
111 {
112 pending.remove(&self.request_id);
113 }
114 }
115}
116
117impl PendingExternalToolGuard {
118 fn claim(&self) -> bool {
119 let mut pending = self.pending.lock();
120 if pending
121 .get(&self.request_id)
122 .is_some_and(|token| Arc::ptr_eq(token, &self.token))
123 {
124 pending.remove(&self.request_id);
125 true
126 } else {
127 false
128 }
129 }
130}
131
132fn has_managed_settings(
133 enable_managed_settings: Option<bool>,
134 managed_settings: Option<&crate::types::ManagedSettings>,
135) -> bool {
136 enable_managed_settings == Some(true) || managed_settings.is_some()
137}
138
139struct IdleWaiter {
141 tx: oneshot::Sender<Result<Option<SessionEvent>, Error>>,
142 last_assistant_message: Option<SessionEvent>,
143 started_at: Instant,
144 first_assistant_message_seen: bool,
145}
146
147fn structured_output_error(message: impl Into<String>) -> Error {
148 Error::with_message(
149 ErrorKind::Session(SessionErrorKind::AgentError),
150 message.into(),
151 )
152}
153
154fn is_structured_output_event(event: &SessionEvent) -> bool {
155 matches!(
156 event.parsed_type(),
157 SessionEventType::UserMessage
158 | SessionEventType::AssistantMessage
159 | SessionEventType::SessionIdle
160 | SessionEventType::SessionError
161 )
162}
163
164struct StructuredOutputState {
165 message_id: String,
166 started: bool,
167 final_message: Option<SessionEvent>,
168}
169
170impl StructuredOutputState {
171 fn observe(&mut self, event: SessionEvent) -> Result<Option<SessionEvent>, Error> {
172 if event.agent_id.as_deref().is_some_and(|id| !id.is_empty()) {
173 return Ok(None);
174 }
175 match event.parsed_type() {
176 SessionEventType::UserMessage => {
177 let data: crate::session_events::UserMessageData =
178 serde_json::from_value(event.data)?;
179 if data.message_id.as_deref() == Some(self.message_id.as_str()) {
180 self.started = true;
181 }
182 }
183 SessionEventType::AssistantMessage => {
184 let data: crate::session_events::AssistantMessageData =
185 serde_json::from_value(event.data.clone())?;
186 if data.originating_message_id.as_deref() == Some(self.message_id.as_str()) {
187 self.started = true;
188 self.final_message =
189 if data.tool_requests.is_some_and(|tools| !tools.is_empty()) {
190 None
191 } else {
192 Some(event)
193 };
194 }
195 }
196 SessionEventType::SessionIdle if self.started => {
197 let data: SessionIdleData = serde_json::from_value(event.data)?;
198 if data.mode == Some(SessionMode::Autopilot) {
199 return Ok(None);
200 }
201 if data.aborted == Some(true) {
202 return Err(structured_output_error(
203 "session aborted before structured output completed",
204 ));
205 }
206 let result = self.final_message.take().ok_or_else(|| {
207 structured_output_error("run completed without a structured assistant response")
208 })?;
209 let data: crate::session_events::AssistantMessageData =
210 serde_json::from_value(result.data.clone())?;
211 if data.content.trim().is_empty() {
212 return Err(structured_output_error(
213 "run completed without a structured assistant response",
214 ));
215 }
216 return Ok(Some(result));
217 }
218 SessionEventType::SessionError if self.started => {
219 let data: SessionErrorData = serde_json::from_value(event.data)?;
220 return Err(structured_output_error(format!(
221 "session error: {}",
222 data.message
223 )));
224 }
225 _ => {}
226 }
227 Ok(None)
228 }
229}
230
231struct WaiterGuard {
243 slot: Arc<ParkingLotMutex<Option<IdleWaiter>>>,
244}
245
246impl Drop for WaiterGuard {
247 fn drop(&mut self) {
248 self.slot.lock().take();
249 }
250}
251
252struct PendingSessionRegistration {
253 client: Client,
254 session_id: PendingSessionId,
255 shutdown: CancellationToken,
256 external_tools_shutdown: CancellationToken,
257 disarmed: bool,
258}
259
260enum PendingSessionId {
268 Known(SessionId, crate::router::RegistrationToken),
271 Deferred(Arc<ParkingLotMutex<Option<(SessionId, crate::router::SessionRegistration)>>>),
274}
275
276impl PendingSessionRegistration {
277 fn new(
278 client: Client,
279 session_id: SessionId,
280 token: crate::router::RegistrationToken,
281 shutdown: CancellationToken,
282 external_tools_shutdown: CancellationToken,
283 ) -> Self {
284 Self {
285 client,
286 session_id: PendingSessionId::Known(session_id, token),
287 shutdown,
288 external_tools_shutdown,
289 disarmed: false,
290 }
291 }
292
293 fn deferred(
295 client: Client,
296 stash: Arc<ParkingLotMutex<Option<(SessionId, crate::router::SessionRegistration)>>>,
297 shutdown: CancellationToken,
298 external_tools_shutdown: CancellationToken,
299 ) -> Self {
300 Self {
301 client,
302 session_id: PendingSessionId::Deferred(stash),
303 shutdown,
304 external_tools_shutdown,
305 disarmed: false,
306 }
307 }
308
309 fn registered_id(&self) -> Option<SessionId> {
310 match &self.session_id {
311 PendingSessionId::Known(id, _) => Some(id.clone()),
312 PendingSessionId::Deferred(stash) => stash.lock().as_ref().map(|(id, _)| id.clone()),
313 }
314 }
315
316 fn resolve_to(&mut self, session_id: SessionId, token: crate::router::RegistrationToken) {
320 self.session_id = PendingSessionId::Known(session_id, token);
321 }
322
323 async fn cleanup(mut self, event_loop: JoinHandle<()>) {
324 self.external_tools_shutdown.cancel();
325 self.shutdown.cancel();
326 let _ = event_loop.await;
327 if let Some(id) = self.registered_id() {
328 if let PendingSessionId::Known(_, token) = self.session_id {
329 self.client.unregister_session_owned(&id, token);
330 } else if let PendingSessionId::Deferred(stash) = &self.session_id
331 && let Some((id, registration)) = stash.lock().as_ref()
332 {
333 self.client.unregister_session_owned(id, registration.token);
334 }
335 }
336 self.disarmed = true;
337 }
338
339 fn disarm(&mut self) {
340 self.disarmed = true;
341 }
342}
343
344impl Drop for PendingSessionRegistration {
345 fn drop(&mut self) {
346 if !self.disarmed {
347 self.external_tools_shutdown.cancel();
348 self.shutdown.cancel();
349 if let Some(id) = self.registered_id() {
350 if let PendingSessionId::Known(_, token) = self.session_id {
351 self.client.unregister_session_owned(&id, token);
352 } else if let PendingSessionId::Deferred(stash) = &self.session_id
353 && let Some((id, registration)) = stash.lock().as_ref()
354 {
355 self.client.unregister_session_owned(id, registration.token);
356 }
357 }
358 }
359 }
360}
361
362pub struct Session {
375 id: SessionId,
376 cwd: PathBuf,
377 workspace_path: Option<PathBuf>,
378 remote_url: Option<String>,
379 client: Client,
380 event_loop: ParkingLotMutex<Option<JoinHandle<()>>>,
385 shutdown: CancellationToken,
399 external_tools_shutdown: CancellationToken,
402 idle_waiter: Arc<ParkingLotMutex<Option<IdleWaiter>>>,
409 capabilities: Arc<parking_lot::RwLock<SessionCapabilities>>,
411 open_canvases: Arc<parking_lot::RwLock<Vec<OpenCanvasInstance>>>,
413 event_tx: tokio::sync::broadcast::Sender<SessionEvent>,
415 github_token_registration:
416 ParkingLotMutex<Option<crate::github_token::GitHubTokenRegistration>>,
417 registration_token: crate::router::RegistrationToken,
419}
420
421impl Session {
422 pub fn id(&self) -> &SessionId {
424 &self.id
425 }
426
427 pub fn cwd(&self) -> &PathBuf {
429 &self.cwd
430 }
431
432 pub fn workspace_path(&self) -> Option<&Path> {
434 self.workspace_path.as_deref()
435 }
436
437 pub fn remote_url(&self) -> Option<&str> {
439 self.remote_url.as_deref()
440 }
441
442 pub fn capabilities(&self) -> SessionCapabilities {
447 self.capabilities.read().clone()
448 }
449
450 pub fn open_canvases(&self) -> Vec<OpenCanvasInstance> {
453 self.open_canvases.read().clone()
454 }
455
456 pub fn cancellation_token(&self) -> CancellationToken {
483 self.shutdown.child_token()
484 }
485
486 pub fn subscribe(&self) -> crate::subscription::EventSubscription {
525 crate::subscription::EventSubscription::new(self.event_tx.subscribe())
526 }
527
528 pub fn client(&self) -> &Client {
530 &self.client
531 }
532
533 pub fn rpc(&self) -> crate::generated::rpc::SessionRpc<'_> {
544 crate::generated::rpc::SessionRpc { session: self }
545 }
546
547 pub async fn stop_event_loop(&self) {
560 self.shutdown.cancel();
561 let handle = self.event_loop.lock().take();
562 if let Some(handle) = handle {
563 let _ = handle.await;
564 }
565 if let Some(waiter) = self.idle_waiter.lock().take() {
567 let _ = waiter.tx.send(Err(
568 ErrorKind::Session(SessionErrorKind::EventLoopClosed).into()
569 ));
570 }
571 }
572
573 pub async fn send(&self, opts: impl Into<MessageOptions>) -> Result<String, Error> {
598 if self.idle_waiter.lock().is_some() {
599 return Err(ErrorKind::Session(SessionErrorKind::SendWhileWaiting).into());
600 }
601 self.send_inner(opts.into()).await
602 }
603
604 async fn send_inner(&self, opts: MessageOptions) -> Result<String, Error> {
605 let mut params = serde_json::json!({
606 "sessionId": self.id,
607 "prompt": opts.prompt,
608 });
609 if let Some(source) = opts.source {
610 params["source"] = serde_json::to_value(source)?;
611 }
612 if let Some(m) = opts.mode {
613 params["mode"] = serde_json::to_value(m)?;
614 }
615 if let Some(am) = opts.agent_mode {
616 params["agentMode"] = serde_json::to_value(am)?;
617 }
618 if let Some(mut a) = opts.attachments {
619 ensure_attachment_display_names(&mut a);
620 params["attachments"] = serde_json::to_value(a)?;
621 }
622 if let Some(headers) = opts.request_headers
623 && !headers.is_empty()
624 {
625 params["requestHeaders"] = serde_json::to_value(headers)?;
626 }
627 if let Some(display_prompt) = opts.display_prompt {
628 params["displayPrompt"] = serde_json::to_value(display_prompt)?;
629 }
630 if let Some(schema) = opts.response_schema {
631 params["responseFormat"] = serde_json::json!({
632 "type": "json_schema",
633 "jsonSchema": { "name": "response", "strict": true, "schema": schema }
634 });
635 }
636 let trace_ctx = if opts.traceparent.is_some() || opts.tracestate.is_some() {
637 TraceContext {
638 traceparent: opts.traceparent,
639 tracestate: opts.tracestate,
640 }
641 } else {
642 self.client.resolve_trace_context().await
643 };
644 inject_trace_context(&mut params, &trace_ctx);
645 let rpc_start = Instant::now();
646 let result = self.client.call("session.send", Some(params)).await?;
647 let message_id = result
648 .get("messageId")
649 .and_then(|v| v.as_str())
650 .map(|s| s.to_string())
651 .unwrap_or_default();
652 tracing::debug!(
653 elapsed_ms = rpc_start.elapsed().as_millis(),
654 session_id = %self.id,
655 message_id = %message_id,
656 "Session::send completed successfully"
657 );
658 Ok(message_id)
659 }
660
661 pub async fn send_and_wait(
683 &self,
684 opts: impl Into<MessageOptions>,
685 ) -> Result<Option<SessionEvent>, Error> {
686 let total_start = Instant::now();
687 let opts = opts.into();
688 if opts.response_schema.is_some() {
689 return self.send_and_wait_structured(opts).await.map(Some);
690 }
691 let timeout_duration = opts.wait_timeout.unwrap_or(Duration::from_secs(60));
692 let (tx, rx) = oneshot::channel();
693
694 {
695 let mut guard = self.idle_waiter.lock();
696 if guard.is_some() {
697 return Err(ErrorKind::Session(SessionErrorKind::SendWhileWaiting).into());
698 }
699 *guard = Some(IdleWaiter {
700 tx,
701 last_assistant_message: None,
702 started_at: total_start,
703 first_assistant_message_seen: false,
704 });
705 }
706
707 let _waiter_guard = WaiterGuard {
712 slot: self.idle_waiter.clone(),
713 };
714
715 let result = tokio::time::timeout(timeout_duration, async {
716 self.send_inner(opts).await?;
717 match rx.await {
718 Ok(result) => result,
719 Err(_) => Err(ErrorKind::Session(SessionErrorKind::EventLoopClosed).into()),
720 }
721 })
722 .await;
723
724 match result {
725 Ok(inner) => {
726 tracing::debug!(
727 elapsed_ms = total_start.elapsed().as_millis(),
728 session_id = %self.id,
729 completed_by = if inner.is_ok() { "idle" } else { "error" },
730 "Session::send_and_wait complete"
731 );
732 inner
733 }
734 Err(_) => {
735 tracing::warn!(
736 elapsed_ms = total_start.elapsed().as_millis(),
737 session_id = %self.id,
738 completed_by = "timeout",
739 "Session::send_and_wait failed"
740 );
741 Err(ErrorKind::Session(SessionErrorKind::Timeout(timeout_duration)).into())
742 }
743 }
744 }
745
746 #[cfg(feature = "derive")]
754 pub async fn send_and_wait_typed<T>(&self, opts: impl Into<MessageOptions>) -> Result<T, Error>
755 where
756 T: schemars::JsonSchema + serde::de::DeserializeOwned,
757 {
758 let mut opts = opts.into();
759 if opts.response_schema.is_some()
760 || opts.mode == Some(crate::types::DeliveryMode::Immediate)
761 {
762 return Err(Error::with_message(
763 ErrorKind::InvalidConfig,
764 "typed structured output cannot specify a response schema or immediate delivery",
765 ));
766 }
767 opts.response_schema = Some(crate::tool::schema_for::<T>());
768 let event = self.send_and_wait_structured(opts).await?;
769 let data: crate::session_events::AssistantMessageData = serde_json::from_value(event.data)?;
770 serde_json::from_str::<Option<T>>(&data.content)?.ok_or_else(|| {
771 structured_output_error("structured response was JSON null, not a result")
772 })
773 }
774
775 async fn send_and_wait_structured(&self, opts: MessageOptions) -> Result<SessionEvent, Error> {
776 let duration = opts.wait_timeout.unwrap_or(Duration::from_secs(60));
777 let mut events = self.subscribe();
778 let wait = async {
779 let mut admission = Box::pin(self.send(opts));
780 let mut pending = Vec::new();
781 let message_id = loop {
782 tokio::select! {
783 result = &mut admission => break result?,
784 event = events.recv() => {
785 let event = event.map_err(|err| structured_output_error(err.to_string()))?;
786 if is_structured_output_event(&event) {
787 pending.push(event);
788 }
789 }
790 _ = self.shutdown.cancelled() =>
791 return Err(structured_output_error("session closed before structured output completed")),
792 }
793 };
794 let mut state = StructuredOutputState {
795 message_id,
796 started: false,
797 final_message: None,
798 };
799 for event in pending {
800 if let Some(result) = state.observe(event)? {
801 return Ok(result);
802 }
803 }
804 loop {
805 tokio::select! {
806 event = events.recv() => {
807 let event = event.map_err(|err| structured_output_error(err.to_string()))?;
808 if let Some(result) = state.observe(event)? {
809 return Ok(result);
810 }
811 }
812 _ = self.shutdown.cancelled() =>
813 return Err(structured_output_error("session closed before structured output completed")),
814 }
815 }
816 };
817 tokio::time::timeout(duration, wait)
818 .await
819 .map_err(|_| Error::from(ErrorKind::Session(SessionErrorKind::Timeout(duration))))?
820 }
821
822 pub async fn get_events(&self) -> Result<Vec<SessionEvent>, Error> {
824 let result = self
825 .client
826 .call(
827 "session.getMessages",
828 Some(serde_json::json!({ "sessionId": self.id })),
829 )
830 .await?;
831 let response: GetMessagesResponse = serde_json::from_value(result)?;
832 Ok(response.events)
833 }
834
835 #[deprecated(since = "0.1.0", note = "Use `get_events()` instead")]
837 pub async fn get_messages(&self) -> Result<Vec<SessionEvent>, Error> {
838 self.get_events().await
839 }
840
841 pub async fn abort(&self) -> Result<(), Error> {
849 self.client
850 .call(
851 "session.abort",
852 Some(serde_json::json!({ "sessionId": self.id })),
853 )
854 .await?;
855 Ok(())
856 }
857
858 pub async fn set_model(&self, model: &str, opts: Option<SetModelOptions>) -> Result<(), Error> {
862 let opts = opts.unwrap_or_default();
863 let auto_tier = opts.auto_tier.clone();
864 let request = ModelSwitchToRequest {
865 auto_tier: match &auto_tier {
866 Some(AutoTierPreference::Tier(tier)) => Some(tier.clone()),
867 _ => None,
868 },
869 compaction_decision: None,
870 context_tier: opts.context_tier,
871 defer_if_model_change_queued: None,
872 model_capabilities: opts.model_capabilities,
873 model_change_scope: None,
874 model_id: model.to_string(),
875 picker_persistence: None,
876 reasoning_effort: opts.reasoning_effort,
877 reasoning_summary: opts.reasoning_summary,
878 repo_scope: None,
879 require_available: None,
880 run_compaction_preflight: None,
881 source: None,
882 verbosity: None,
883 };
884
885 if matches!(auto_tier, Some(AutoTierPreference::Reset)) {
886 let mut wire_params = serde_json::to_value(request)?;
890 wire_params["sessionId"] = serde_json::Value::String(self.id.to_string());
891 wire_params["autoTier"] = serde_json::Value::Null;
892 self.client
893 .call("session.model.switchTo", Some(wire_params))
894 .await?;
895 return Ok(());
896 }
897
898 self.rpc().model().switch_to(request).await?;
899 Ok(())
900 }
901
902 pub async fn set_auto_tier(
931 &self,
932 auto_tier: Option<AutoTier>,
933 ) -> Result<ModelSwitchAutoTierResult, Error> {
934 self.rpc()
935 .model()
936 .switch_auto_tier(ModelSwitchAutoTierRequest {
937 auto_tier,
938 source: None,
939 })
940 .await
941 }
942
943 pub async fn disconnect(&self) -> Result<(), Error> {
960 self.client.detach_session(&self.id).await?;
961 self.external_tools_shutdown.cancel();
962 self.stop_event_loop().await;
963 self.github_token_registration.lock().take();
964 self.client
965 .unregister_session_owned(&self.id, self.registration_token);
966 Ok(())
967 }
968
969 #[deprecated(since = "0.1.0", note = "Use `disconnect()` instead")]
972 pub async fn destroy(&self) -> Result<(), Error> {
973 self.disconnect().await
974 }
975
976 pub async fn log(
980 &self,
981 message: &str,
982 opts: Option<crate::types::LogOptions>,
983 ) -> Result<(), Error> {
984 let opts = opts.unwrap_or_default();
985 let level = match opts.level {
986 Some(level) => Some(serde_json::from_value(serde_json::to_value(level)?)?),
987 None => None,
988 };
989 let request = LogRequest {
990 message: message.to_string(),
991 level,
992 ephemeral: opts.ephemeral,
993 r#type: None,
994 tip: None,
995 url: None,
996 };
997 self.rpc().log(request).await?;
998 Ok(())
999 }
1000
1001 pub fn ui(&self) -> SessionUi<'_> {
1007 SessionUi { session: self }
1008 }
1009
1010 fn assert_elicitation(&self) -> Result<(), Error> {
1012 if self
1013 .capabilities
1014 .read()
1015 .ui
1016 .as_ref()
1017 .and_then(|u| u.elicitation)
1018 != Some(true)
1019 {
1020 return Err(ErrorKind::Session(SessionErrorKind::ElicitationNotSupported).into());
1021 }
1022 Ok(())
1023 }
1024}
1025
1026impl Drop for Session {
1027 fn drop(&mut self) {
1028 self.shutdown.cancel();
1040 self.external_tools_shutdown.cancel();
1041 self.github_token_registration.lock().take();
1042 self.client
1043 .unregister_session_owned(&self.id, self.registration_token);
1044 }
1045}
1046
1047pub struct SessionUi<'a> {
1054 session: &'a Session,
1055}
1056
1057impl<'a> SessionUi<'a> {
1058 pub async fn elicitation(
1066 &self,
1067 message: &str,
1068 schema: Value,
1069 ) -> Result<ElicitationResult, Error> {
1070 self.session.assert_elicitation()?;
1071 let result = self
1072 .session
1073 .client
1074 .call(
1075 "session.ui.elicitation",
1076 Some(serde_json::json!({
1077 "sessionId": self.session.id,
1078 "message": message,
1079 "requestedSchema": schema,
1080 })),
1081 )
1082 .await?;
1083 let elicitation: ElicitationResult = serde_json::from_value(result)?;
1084 Ok(elicitation)
1085 }
1086
1087 pub async fn confirm(&self, message: &str) -> Result<bool, Error> {
1091 self.session.assert_elicitation()?;
1092 let schema = serde_json::json!({
1093 "type": "object",
1094 "properties": {
1095 "confirmed": {
1096 "type": "boolean",
1097 "default": true,
1098 }
1099 },
1100 "required": ["confirmed"]
1101 });
1102 let result = self.elicitation(message, schema).await?;
1103 Ok(result.action == "accept"
1104 && result
1105 .content
1106 .and_then(|c| c.get("confirmed").and_then(|v| v.as_bool()))
1107 == Some(true))
1108 }
1109
1110 pub async fn select(&self, message: &str, options: &[&str]) -> Result<Option<String>, Error> {
1114 self.session.assert_elicitation()?;
1115 let schema = serde_json::json!({
1116 "type": "object",
1117 "properties": {
1118 "selection": {
1119 "type": "string",
1120 "enum": options,
1121 }
1122 },
1123 "required": ["selection"]
1124 });
1125 let result = self.elicitation(message, schema).await?;
1126 if result.action != "accept" {
1127 return Ok(None);
1128 }
1129 let selection = result.content.and_then(|c| {
1130 c.get("selection")
1131 .and_then(|v| v.as_str())
1132 .map(String::from)
1133 });
1134 Ok(selection)
1135 }
1136
1137 pub async fn input(
1142 &self,
1143 message: &str,
1144 options: Option<&UiInputOptions<'_>>,
1145 ) -> Result<Option<String>, Error> {
1146 self.session.assert_elicitation()?;
1147 let mut field = serde_json::json!({ "type": "string" });
1148 if let Some(opts) = options {
1149 if let Some(title) = opts.title {
1150 field["title"] = Value::String(title.to_string());
1151 }
1152 if let Some(desc) = opts.description {
1153 field["description"] = Value::String(desc.to_string());
1154 }
1155 if let Some(min) = opts.min_length {
1156 field["minLength"] = Value::Number(min.into());
1157 }
1158 if let Some(max) = opts.max_length {
1159 field["maxLength"] = Value::Number(max.into());
1160 }
1161 if let Some(fmt) = &opts.format {
1162 field["format"] = Value::String(fmt.as_str().to_string());
1163 }
1164 if let Some(default) = opts.default {
1165 field["default"] = Value::String(default.to_string());
1166 }
1167 }
1168 let schema = serde_json::json!({
1169 "type": "object",
1170 "properties": { "value": field },
1171 "required": ["value"]
1172 });
1173 let result = self.elicitation(message, schema).await?;
1174 if result.action != "accept" {
1175 return Ok(None);
1176 }
1177 let value = result
1178 .content
1179 .and_then(|c| c.get("value").and_then(|v| v.as_str()).map(String::from));
1180 Ok(value)
1181 }
1182}
1183
1184impl Client {
1185 pub fn prepare_session(&self, config: SessionConfig) -> Result<PreparedSession, Error> {
1247 let capacity = resolve_event_buffer_capacity(config.event_buffer_capacity)?;
1248 Ok(PreparedSession::new(
1249 self.clone(),
1250 PreparedKind::Create(Box::new(config)),
1251 capacity,
1252 ))
1253 }
1254
1255 pub fn prepare_resume_session(
1270 &self,
1271 config: ResumeSessionConfig,
1272 ) -> Result<PreparedSession, Error> {
1273 let capacity = resolve_event_buffer_capacity(config.event_buffer_capacity)?;
1274 Ok(PreparedSession::new(
1275 self.clone(),
1276 PreparedKind::Resume(Box::new(config)),
1277 capacity,
1278 ))
1279 }
1280
1281 pub async fn create_session(&self, config: SessionConfig) -> Result<Session, Error> {
1313 self.prepare_session(config)?.start().await
1314 }
1315
1316 pub async fn resume_session(&self, config: ResumeSessionConfig) -> Result<Session, Error> {
1335 self.prepare_resume_session(config)?.start().await
1336 }
1337
1338 async fn start_prepared_create(
1339 &self,
1340 mut config: SessionConfig,
1341 event_tx: tokio::sync::broadcast::Sender<SessionEvent>,
1342 shutdown: CancellationToken,
1343 ) -> Result<Session, Error> {
1344 let total_start = Instant::now();
1345 let caller_session_id = config.session_id.clone();
1353 let use_server_generated_id = config.cloud.is_some() && caller_session_id.is_none();
1354 let local_session_id: Option<SessionId> = if use_server_generated_id {
1355 None
1356 } else {
1357 Some(
1358 caller_session_id
1359 .clone()
1360 .unwrap_or_else(|| SessionId::new(uuid::Uuid::new_v4().to_string())),
1361 )
1362 };
1363 if config.hooks_handler.is_some() && config.hooks.is_none() {
1364 config.hooks = Some(true);
1365 }
1366 if let Some(transforms) = config.system_message_transform.clone() {
1367 inject_transform_sections(&mut config, transforms.as_ref());
1368 }
1369 let mode = self.inner.mode;
1370 if mode == crate::ClientMode::Empty && config.available_tools.is_none() {
1371 return Err(Error::with_message(
1372 ErrorKind::InvalidConfig,
1373 "ClientMode::Empty requires available_tools to be set on the session config. \
1374 Use ToolSet to specify which tools the session may use (e.g. \
1375 ToolSet::new().add_builtin_many(BUILTIN_TOOLS_ISOLATED)).",
1376 ));
1377 }
1378 crate::mode::validate_tool_filter_list(
1379 "available_tools",
1380 config.available_tools.as_deref(),
1381 )?;
1382 crate::mode::validate_tool_filter_list("excluded_tools", config.excluded_tools.as_deref())?;
1383 config.system_message =
1384 crate::mode::system_message_for_mode(mode, config.system_message.take());
1385 config.memory = crate::mode::memory_for_mode(mode, config.memory.take());
1386 config.enable_experimental_mode =
1387 crate::mode::experimental_mode_for_mode(mode, config.enable_experimental_mode);
1388 if mode == crate::ClientMode::Empty {
1389 if config.enable_session_telemetry.is_none() {
1390 config.enable_session_telemetry = Some(false);
1391 }
1392 if config.skip_embedding_retrieval.is_none() {
1393 config.skip_embedding_retrieval = Some(true);
1394 }
1395 if config.enable_on_demand_instruction_discovery.is_none() {
1396 config.enable_on_demand_instruction_discovery = Some(false);
1397 }
1398 if config.enable_file_hooks.is_none() {
1399 config.enable_file_hooks = Some(false);
1400 }
1401 if config.enable_host_git_operations.is_none() {
1402 config.enable_host_git_operations = Some(false);
1403 }
1404 if config.enable_session_store.is_none() {
1405 config.enable_session_store = Some(false);
1406 }
1407 if config.enable_skills.is_none() {
1408 config.enable_skills = Some(false);
1409 }
1410 }
1411 if mode == crate::ClientMode::Empty && config.mcp_oauth_token_storage.is_none() {
1412 config.mcp_oauth_token_storage = Some("in-memory".into());
1413 }
1414 if mode == crate::ClientMode::Empty && config.embedding_cache_storage.is_none() {
1415 config.embedding_cache_storage = Some("in-memory".into());
1416 }
1417 config.custom_agents_local_only =
1418 crate::mode::resolve_custom_agents_local_only(mode, config.custom_agents_local_only);
1419 let opt_skip_custom_instructions = config.skip_custom_instructions;
1420 let opt_custom_agents_local_only = config.custom_agents_local_only;
1421 let opt_coauthor_enabled = config.coauthor_enabled;
1422 let opt_manage_schedule_enabled = config.manage_schedule_enabled;
1423 let opt_included_builtin_skills = config.included_builtin_skills.take();
1424 let (mut wire, mut runtime) = config.into_wire(local_session_id.clone())?;
1425 wire.enable_github_telemetry_forwarding =
1426 self.inner.on_github_telemetry.is_some().then_some(true);
1427
1428 let permission_handler = crate::permission::resolve_handler(
1429 runtime.permission_handler.take(),
1430 runtime.permission_policy.take(),
1431 );
1432 let handlers = SessionHandlers {
1433 permission: permission_handler,
1434 managed_settings_enabled: has_managed_settings(
1435 wire.enable_managed_settings,
1436 wire.managed_settings.as_ref(),
1437 ),
1438 elicitation: runtime.elicitation_handler.take(),
1439 mcp_auth: runtime.mcp_auth_handler.take(),
1440 user_input: runtime.user_input_handler.take(),
1441 exit_plan_mode: runtime.exit_plan_mode_handler.take(),
1442 auto_mode_switch: runtime.auto_mode_switch_handler.take(),
1443 tools: Arc::new(std::mem::take(&mut runtime.tool_handlers)),
1444 };
1445 let hooks = runtime.hooks_handler.take();
1446 let transforms = runtime.system_message_transform.take();
1447 let tools_count = wire.tools.as_ref().map_or(0, Vec::len);
1448 let commands_count = runtime.commands.as_ref().map_or(0, Vec::len);
1449 let has_hooks = hooks.is_some();
1450 let command_handlers = build_command_handler_map(runtime.commands.as_deref());
1451 let canvas_handler = runtime.canvas_handler.take();
1452 let session_fs_provider = runtime.session_fs_provider.take();
1453 let bearer_token_providers = std::mem::take(&mut runtime.bearer_token_providers);
1454 let github_token_registration = runtime
1455 .github_token_provider
1456 .take()
1457 .map(|provider| self.register_github_token_provider(provider));
1458 wire.github_token_provider_registration_id = github_token_registration
1459 .as_ref()
1460 .map(|registration| registration.id().to_string());
1461 let has_mcp_auth_handler = handlers.mcp_auth.is_some();
1462 if self.inner.session_fs_configured && session_fs_provider.is_none() {
1463 return Err(ErrorKind::Session(SessionErrorKind::SessionFsProviderRequired).into());
1464 }
1465 if self.inner.session_fs_sqlite_declared
1466 && let Some(ref provider) = session_fs_provider
1467 && provider.sqlite().is_none()
1468 {
1469 return Err(Error::with_message(
1470 ErrorKind::InvalidConfig,
1471 "SessionFs capabilities declare SQLite support but the provider \
1472 does not implement SessionFsSqliteProvider",
1473 ));
1474 }
1475
1476 let mut params = serde_json::to_value(&wire)?;
1477 let trace_ctx = self.resolve_trace_context().await;
1478 inject_trace_context(&mut params, &trace_ctx);
1479
1480 let setup_start = Instant::now();
1481 let capabilities = Arc::new(parking_lot::RwLock::new(SessionCapabilities::default()));
1482 let idle_waiter = Arc::new(ParkingLotMutex::new(None));
1483 let open_canvases = Arc::new(parking_lot::RwLock::new(Vec::new()));
1484 let external_tools_shutdown = self.inner.rpc.connection_closed_token();
1485
1486 let inline_stash: Arc<
1492 ParkingLotMutex<Option<(SessionId, crate::router::SessionRegistration)>>,
1493 > = Arc::new(ParkingLotMutex::new(None));
1494
1495 let inline_callback: Option<crate::jsonrpc::InlineResponseCallback> = if let Some(ref sid) =
1496 local_session_id
1497 {
1498 let channels = self.register_session(sid);
1499 *inline_stash.lock() = Some((sid.clone(), channels));
1500 None
1501 } else {
1502 let client = self.clone();
1503 let stash = inline_stash.clone();
1504 let expected = caller_session_id.clone();
1505 Some(Box::new(move |response| {
1506 let result = response.result.as_ref().ok_or_else(|| {
1507 Error::with_message(ErrorKind::Json, "session.create response had no result")
1508 })?;
1509 let parsed: CreateSessionResult =
1510 serde_json::from_value(result.clone()).map_err(Error::from)?;
1511 if let Some(requested) = expected.as_ref()
1512 && parsed.session_id != *requested
1513 {
1514 return Err(ErrorKind::Session(SessionErrorKind::SessionIdMismatch {
1515 requested: requested.clone(),
1516 returned: parsed.session_id,
1517 })
1518 .into());
1519 }
1520 let mut stashed = stash.lock();
1528 let registration = client.register_session(&parsed.session_id);
1529 *stashed = Some((parsed.session_id, registration));
1530 Ok(())
1531 }))
1532 };
1533
1534 let mut pending_registration = match local_session_id {
1540 Some(ref sid) => {
1541 let token = inline_stash
1542 .lock()
1543 .as_ref()
1544 .expect("session registration must exist")
1545 .1
1546 .token;
1547 PendingSessionRegistration::new(
1548 self.clone(),
1549 sid.clone(),
1550 token,
1551 shutdown.clone(),
1552 external_tools_shutdown.clone(),
1553 )
1554 }
1555 None => PendingSessionRegistration::deferred(
1556 self.clone(),
1557 inline_stash.clone(),
1558 shutdown.clone(),
1559 external_tools_shutdown.clone(),
1560 ),
1561 };
1562
1563 let rpc_start = Instant::now();
1564 let result = self
1565 .call_with_inline_callback("session.create", Some(params), inline_callback)
1566 .await?;
1567 tracing::debug!(
1568 elapsed_ms = rpc_start.elapsed().as_millis(),
1569 "Client::create_session session creation request completed successfully"
1570 );
1571 let create_result: CreateSessionResult = serde_json::from_value(result)?;
1572
1573 if let Some(ref requested) = local_session_id
1574 && create_result.session_id != *requested
1575 {
1576 return Err(ErrorKind::Session(SessionErrorKind::SessionIdMismatch {
1577 requested: requested.clone(),
1578 returned: create_result.session_id.clone(),
1579 })
1580 .into());
1581 }
1582
1583 let (session_id, registration) = inline_stash
1584 .lock()
1585 .take()
1586 .expect("session registration must have populated stash on success");
1587 let channels = registration.channels;
1588 let registration_token = registration.token;
1589 pending_registration.resolve_to(session_id.clone(), registration_token);
1590 let event_loop = spawn_event_loop(
1591 session_id.clone(),
1592 self.clone(),
1593 handlers,
1594 hooks,
1595 transforms,
1596 command_handlers,
1597 canvas_handler,
1598 session_fs_provider,
1599 bearer_token_providers,
1600 channels,
1601 idle_waiter.clone(),
1602 capabilities.clone(),
1603 open_canvases.clone(),
1604 event_tx.clone(),
1605 shutdown.clone(),
1606 external_tools_shutdown.clone(),
1607 );
1608 tracing::debug!(
1609 elapsed_ms = setup_start.elapsed().as_millis(),
1610 session_id = %session_id,
1611 tools_count,
1612 commands_count,
1613 has_hooks,
1614 "Client::create_session local setup complete"
1615 );
1616 *capabilities.write() = create_result.capabilities.unwrap_or_default();
1617 if has_mcp_auth_handler
1618 && let Err(error) = register_mcp_auth_interest(self, &session_id).await
1619 {
1620 pending_registration.cleanup(event_loop).await;
1621 return Err(error);
1622 }
1623
1624 tracing::debug!(
1625 elapsed_ms = total_start.elapsed().as_millis(),
1626 session_id = %session_id,
1627 "Client::create_session complete"
1628 );
1629 pending_registration.disarm();
1630 let session = Session {
1631 id: session_id,
1632 cwd: self.cwd().clone(),
1633 workspace_path: create_result.workspace_path,
1634 remote_url: create_result.remote_url,
1635 client: self.clone(),
1636 event_loop: ParkingLotMutex::new(Some(event_loop)),
1637 shutdown,
1638 external_tools_shutdown,
1639 idle_waiter,
1640 capabilities,
1641 open_canvases,
1642 event_tx,
1643 github_token_registration: ParkingLotMutex::new(github_token_registration),
1644 registration_token,
1645 };
1646 apply_mode_post_create_patch(
1647 &session,
1648 mode,
1649 opt_skip_custom_instructions,
1650 opt_custom_agents_local_only,
1651 opt_coauthor_enabled,
1652 opt_manage_schedule_enabled,
1653 opt_included_builtin_skills,
1654 )
1655 .await?;
1656 if let Some(registration) = session.github_token_registration.lock().as_ref() {
1657 registration.claim(session.id.clone());
1658 } else {
1659 self.retire_github_token_provider(&session.id);
1660 }
1661 Ok(session)
1662 }
1663
1664 async fn start_prepared_resume(
1675 &self,
1676 mut config: ResumeSessionConfig,
1677 event_tx: tokio::sync::broadcast::Sender<SessionEvent>,
1678 shutdown: CancellationToken,
1679 ) -> Result<Session, Error> {
1680 let total_start = Instant::now();
1681 let session_id = config.session_id.clone();
1682 if config.hooks_handler.is_some() && config.hooks.is_none() {
1683 config.hooks = Some(true);
1684 }
1685 if let Some(transforms) = config.system_message_transform.clone() {
1686 inject_transform_sections_resume(&mut config, transforms.as_ref());
1687 }
1688 let mode = self.inner.mode;
1689 if mode == crate::ClientMode::Empty && config.available_tools.is_none() {
1690 return Err(Error::with_message(
1691 ErrorKind::InvalidConfig,
1692 "ClientMode::Empty requires available_tools to be set on the session config. \
1693 Use ToolSet to specify which tools the session may use (e.g. \
1694 ToolSet::new().add_builtin_many(BUILTIN_TOOLS_ISOLATED)).",
1695 ));
1696 }
1697 crate::mode::validate_tool_filter_list(
1698 "available_tools",
1699 config.available_tools.as_deref(),
1700 )?;
1701 crate::mode::validate_tool_filter_list("excluded_tools", config.excluded_tools.as_deref())?;
1702 config.system_message =
1703 crate::mode::system_message_for_mode(mode, config.system_message.take());
1704 config.memory = crate::mode::memory_for_mode(mode, config.memory.take());
1705 config.enable_experimental_mode =
1706 crate::mode::experimental_mode_for_mode(mode, config.enable_experimental_mode);
1707 if mode == crate::ClientMode::Empty {
1708 if config.enable_session_telemetry.is_none() {
1709 config.enable_session_telemetry = Some(false);
1710 }
1711 if config.skip_embedding_retrieval.is_none() {
1712 config.skip_embedding_retrieval = Some(true);
1713 }
1714 if config.enable_on_demand_instruction_discovery.is_none() {
1715 config.enable_on_demand_instruction_discovery = Some(false);
1716 }
1717 if config.enable_file_hooks.is_none() {
1718 config.enable_file_hooks = Some(false);
1719 }
1720 if config.enable_host_git_operations.is_none() {
1721 config.enable_host_git_operations = Some(false);
1722 }
1723 if config.enable_session_store.is_none() {
1724 config.enable_session_store = Some(false);
1725 }
1726 if config.enable_skills.is_none() {
1727 config.enable_skills = Some(false);
1728 }
1729 }
1730 if mode == crate::ClientMode::Empty && config.mcp_oauth_token_storage.is_none() {
1731 config.mcp_oauth_token_storage = Some("in-memory".into());
1732 }
1733 if mode == crate::ClientMode::Empty && config.embedding_cache_storage.is_none() {
1734 config.embedding_cache_storage = Some("in-memory".into());
1735 }
1736 config.custom_agents_local_only =
1737 crate::mode::resolve_custom_agents_local_only(mode, config.custom_agents_local_only);
1738 let opt_skip_custom_instructions = config.skip_custom_instructions;
1739 let opt_custom_agents_local_only = config.custom_agents_local_only;
1740 let opt_coauthor_enabled = config.coauthor_enabled;
1741 let opt_manage_schedule_enabled = config.manage_schedule_enabled;
1742 let opt_included_builtin_skills = config.included_builtin_skills.take();
1743 let (mut wire, mut runtime) = config.into_wire()?;
1744 wire.enable_github_telemetry_forwarding =
1745 self.inner.on_github_telemetry.is_some().then_some(true);
1746
1747 let permission_handler = crate::permission::resolve_handler(
1748 runtime.permission_handler.take(),
1749 runtime.permission_policy.take(),
1750 );
1751 let handlers = SessionHandlers {
1752 permission: permission_handler,
1753 managed_settings_enabled: has_managed_settings(
1754 wire.enable_managed_settings,
1755 wire.managed_settings.as_ref(),
1756 ),
1757 elicitation: runtime.elicitation_handler.take(),
1758 mcp_auth: runtime.mcp_auth_handler.take(),
1759 user_input: runtime.user_input_handler.take(),
1760 exit_plan_mode: runtime.exit_plan_mode_handler.take(),
1761 auto_mode_switch: runtime.auto_mode_switch_handler.take(),
1762 tools: Arc::new(std::mem::take(&mut runtime.tool_handlers)),
1763 };
1764 let hooks = runtime.hooks_handler.take();
1765 let transforms = runtime.system_message_transform.take();
1766 let tools_count = wire.tools.as_ref().map_or(0, Vec::len);
1767 let commands_count = runtime.commands.as_ref().map_or(0, Vec::len);
1768 let has_hooks = hooks.is_some();
1769 let command_handlers = build_command_handler_map(runtime.commands.as_deref());
1770 let canvas_handler = runtime.canvas_handler.take();
1771 let session_fs_provider = runtime.session_fs_provider.take();
1772 let bearer_token_providers = std::mem::take(&mut runtime.bearer_token_providers);
1773 let github_token_registration = runtime
1774 .github_token_provider
1775 .take()
1776 .map(|provider| self.register_github_token_provider(provider));
1777 wire.github_token_provider_registration_id = github_token_registration
1778 .as_ref()
1779 .map(|registration| registration.id().to_string());
1780 let has_mcp_auth_handler = handlers.mcp_auth.is_some();
1781 if self.inner.session_fs_configured && session_fs_provider.is_none() {
1782 return Err(ErrorKind::Session(SessionErrorKind::SessionFsProviderRequired).into());
1783 }
1784 if self.inner.session_fs_sqlite_declared
1785 && let Some(ref provider) = session_fs_provider
1786 && provider.sqlite().is_none()
1787 {
1788 return Err(Error::with_message(
1789 ErrorKind::InvalidConfig,
1790 "SessionFs capabilities declare SQLite support but the provider \
1791 does not implement SessionFsSqliteProvider",
1792 ));
1793 }
1794
1795 let mut params = serde_json::to_value(&wire)?;
1796 let trace_ctx = self.resolve_trace_context().await;
1797 inject_trace_context(&mut params, &trace_ctx);
1798
1799 let capabilities = Arc::new(parking_lot::RwLock::new(SessionCapabilities::default()));
1800 let setup_start = Instant::now();
1801 let registration = self.register_session(&session_id);
1802 let registration_token = registration.token;
1803 let channels = registration.channels;
1804 let idle_waiter = Arc::new(ParkingLotMutex::new(None));
1805 let open_canvases = Arc::new(parking_lot::RwLock::new(Vec::new()));
1806 let external_tools_shutdown = self.inner.rpc.connection_closed_token();
1807 let event_loop = spawn_event_loop(
1808 session_id.clone(),
1809 self.clone(),
1810 handlers,
1811 hooks,
1812 transforms,
1813 command_handlers,
1814 canvas_handler,
1815 session_fs_provider,
1816 bearer_token_providers,
1817 channels,
1818 idle_waiter.clone(),
1819 capabilities.clone(),
1820 open_canvases.clone(),
1821 event_tx.clone(),
1822 shutdown.clone(),
1823 external_tools_shutdown.clone(),
1824 );
1825 let mut registration = PendingSessionRegistration::new(
1826 self.clone(),
1827 session_id.clone(),
1828 registration_token,
1829 shutdown.clone(),
1830 external_tools_shutdown.clone(),
1831 );
1832 tracing::debug!(
1833 elapsed_ms = setup_start.elapsed().as_millis(),
1834 session_id = %session_id,
1835 tools_count,
1836 commands_count,
1837 has_hooks,
1838 "Client::resume_session local setup complete"
1839 );
1840
1841 let rpc_start = Instant::now();
1842 let result = match self.call("session.resume", Some(params)).await {
1843 Ok(result) => result,
1844 Err(error) => {
1845 registration.cleanup(event_loop).await;
1846 return Err(error);
1847 }
1848 };
1849 tracing::debug!(
1850 elapsed_ms = rpc_start.elapsed().as_millis(),
1851 session_id = %session_id,
1852 "Client::resume_session session resume request completed successfully"
1853 );
1854
1855 let resume_result: ResumeSessionResult = match serde_json::from_value(result) {
1856 Ok(result) => result,
1857 Err(error) => {
1858 registration.cleanup(event_loop).await;
1859 return Err(error.into());
1860 }
1861 };
1862 let cli_session_id = resume_result
1863 .session_id
1864 .clone()
1865 .unwrap_or_else(|| session_id.clone());
1866 if cli_session_id != session_id {
1867 registration.cleanup(event_loop).await;
1868 return Err(ErrorKind::Session(SessionErrorKind::SessionIdMismatch {
1869 requested: session_id,
1870 returned: cli_session_id,
1871 })
1872 .into());
1873 }
1874 if has_mcp_auth_handler
1875 && let Err(error) = register_mcp_auth_interest(self, &session_id).await
1876 {
1877 registration.cleanup(event_loop).await;
1878 return Err(error);
1879 }
1880 let skills_reload_start = Instant::now();
1882 if let Err(e) = self
1883 .call(
1884 "session.skills.reload",
1885 Some(serde_json::json!({ "sessionId": session_id })),
1886 )
1887 .await
1888 {
1889 warn!(
1890 elapsed_ms = skills_reload_start.elapsed().as_millis(),
1891 session_id = %session_id,
1892 error = %e,
1893 "Client::resume_session skills reload request failed"
1894 );
1895 } else {
1896 tracing::debug!(
1897 elapsed_ms = skills_reload_start.elapsed().as_millis(),
1898 session_id = %session_id,
1899 "Client::resume_session skills reload request completed successfully"
1900 );
1901 }
1902
1903 *capabilities.write() = resume_result.capabilities.unwrap_or_default();
1904 {
1909 let mut snapshots = open_canvases.write();
1910 for snapshot in resume_result.open_canvases.unwrap_or_default() {
1911 upsert_open_canvas_snapshot(&mut snapshots, snapshot);
1912 }
1913 }
1914
1915 tracing::debug!(
1916 elapsed_ms = total_start.elapsed().as_millis(),
1917 session_id = %session_id,
1918 "Client::resume_session complete"
1919 );
1920 registration.disarm();
1921 let session = Session {
1922 id: session_id,
1923 cwd: self.cwd().clone(),
1924 workspace_path: resume_result.workspace_path,
1925 remote_url: resume_result.remote_url,
1926 client: self.clone(),
1927 event_loop: ParkingLotMutex::new(Some(event_loop)),
1928 shutdown,
1929 external_tools_shutdown,
1930 idle_waiter,
1931 capabilities,
1932 open_canvases,
1933 event_tx,
1934 github_token_registration: ParkingLotMutex::new(github_token_registration),
1935 registration_token,
1936 };
1937 apply_mode_post_create_patch(
1938 &session,
1939 mode,
1940 opt_skip_custom_instructions,
1941 opt_custom_agents_local_only,
1942 opt_coauthor_enabled,
1943 opt_manage_schedule_enabled,
1944 opt_included_builtin_skills,
1945 )
1946 .await?;
1947 if let Some(registration) = session.github_token_registration.lock().as_ref() {
1948 registration.claim(session.id.clone());
1949 } else {
1950 self.retire_github_token_provider(&session.id);
1951 }
1952 Ok(session)
1953 }
1954}
1955
1956#[must_use = "a PreparedSession does nothing until started"]
2014pub struct PreparedSession {
2015 client: Client,
2016 kind: PreparedKind,
2017 event_tx: tokio::sync::broadcast::Sender<SessionEvent>,
2018 shutdown: CancellationToken,
2019}
2020
2021enum PreparedKind {
2024 Create(Box<SessionConfig>),
2025 Resume(Box<ResumeSessionConfig>),
2026}
2027
2028impl PreparedSession {
2029 fn new(client: Client, kind: PreparedKind, event_buffer_capacity: usize) -> Self {
2030 let (event_tx, _) = tokio::sync::broadcast::channel(event_buffer_capacity);
2031 Self {
2032 client,
2033 kind,
2034 event_tx,
2035 shutdown: CancellationToken::new(),
2036 }
2037 }
2038
2039 pub fn subscribe(&self) -> crate::subscription::EventSubscription {
2057 crate::subscription::EventSubscription::new(self.event_tx.subscribe())
2058 }
2059
2060 pub async fn start(self) -> Result<Session, Error> {
2078 let Self {
2079 client,
2080 kind,
2081 event_tx,
2082 shutdown,
2083 } = self;
2084 match kind {
2085 PreparedKind::Create(config) => {
2086 client
2087 .start_prepared_create(*config, event_tx, shutdown)
2088 .await
2089 }
2090 PreparedKind::Resume(config) => {
2091 client
2092 .start_prepared_resume(*config, event_tx, shutdown)
2093 .await
2094 }
2095 }
2096 }
2097}
2098
2099type CommandHandlerMap = HashMap<String, Arc<dyn CommandHandler>>;
2100
2101async fn apply_mode_post_create_patch(
2102 session: &Session,
2103 mode: crate::ClientMode,
2104 opt_skip_custom_instructions: Option<bool>,
2105 opt_custom_agents_local_only: Option<bool>,
2106 opt_coauthor_enabled: Option<bool>,
2107 opt_manage_schedule_enabled: Option<bool>,
2108 opt_included_builtin_skills: Option<Vec<String>>,
2109) -> Result<(), Error> {
2110 let Some(patch) = build_mode_post_create_patch(
2111 mode,
2112 opt_skip_custom_instructions,
2113 opt_custom_agents_local_only,
2114 opt_coauthor_enabled,
2115 opt_manage_schedule_enabled,
2116 opt_included_builtin_skills,
2117 ) else {
2118 return Ok(());
2119 };
2120 if let Err(error) = session.rpc().options().update(patch).await {
2121 let _ = session.disconnect().await;
2122 return Err(error);
2123 }
2124 Ok(())
2125}
2126
2127fn build_mode_post_create_patch(
2137 mode: crate::ClientMode,
2138 opt_skip_custom_instructions: Option<bool>,
2139 opt_custom_agents_local_only: Option<bool>,
2140 opt_coauthor_enabled: Option<bool>,
2141 opt_manage_schedule_enabled: Option<bool>,
2142 opt_included_builtin_skills: Option<Vec<String>>,
2143) -> Option<crate::generated::api_types::SessionUpdateOptionsParams> {
2144 use crate::generated::api_types::SessionUpdateOptionsParams;
2145 let mut patch = SessionUpdateOptionsParams::default();
2146 let should_send = if mode == crate::ClientMode::Empty {
2147 patch.skip_custom_instructions = Some(opt_skip_custom_instructions.unwrap_or(true));
2148 patch.custom_agents_local_only = Some(opt_custom_agents_local_only.unwrap_or(true));
2149 patch.coauthor_enabled = Some(opt_coauthor_enabled.unwrap_or(false));
2150 patch.manage_schedule_enabled = Some(opt_manage_schedule_enabled.unwrap_or(false));
2151 patch.installed_plugins = Some(Vec::new());
2152 patch.included_builtin_skills = Some(opt_included_builtin_skills.unwrap_or_default());
2153 true
2154 } else {
2155 let mut any = false;
2156 if let Some(v) = opt_skip_custom_instructions {
2157 patch.skip_custom_instructions = Some(v);
2158 any = true;
2159 }
2160 if let Some(v) = opt_custom_agents_local_only {
2161 patch.custom_agents_local_only = Some(v);
2162 any = true;
2163 }
2164 if let Some(v) = opt_coauthor_enabled {
2165 patch.coauthor_enabled = Some(v);
2166 any = true;
2167 }
2168 if let Some(v) = opt_manage_schedule_enabled {
2169 patch.manage_schedule_enabled = Some(v);
2170 any = true;
2171 }
2172 if let Some(v) = opt_included_builtin_skills {
2173 patch.included_builtin_skills = Some(v);
2174 any = true;
2175 }
2176 any
2177 };
2178 if !should_send {
2179 return None;
2180 }
2181 Some(patch)
2182}
2183
2184fn build_command_handler_map(commands: Option<&[CommandDefinition]>) -> Arc<CommandHandlerMap> {
2185 let map = match commands {
2186 Some(commands) => commands
2187 .iter()
2188 .filter(|cmd| !cmd.name.is_empty())
2189 .map(|cmd| (cmd.name.clone(), cmd.handler.clone()))
2190 .collect(),
2191 None => HashMap::new(),
2192 };
2193 Arc::new(map)
2194}
2195
2196fn upsert_open_canvas_snapshot(
2197 snapshots: &mut Vec<OpenCanvasInstance>,
2198 snapshot: OpenCanvasInstance,
2199) {
2200 if let Some(existing) = snapshots
2201 .iter_mut()
2202 .find(|open| open.instance_id == snapshot.instance_id)
2203 {
2204 *existing = snapshot;
2205 } else {
2206 snapshots.push(snapshot);
2207 }
2208}
2209
2210fn remove_open_canvas_snapshot(snapshots: &mut Vec<OpenCanvasInstance>, instance_id: &str) {
2211 snapshots.retain(|open| open.instance_id != instance_id);
2212}
2213
2214#[allow(clippy::too_many_arguments)]
2215fn spawn_event_loop(
2216 session_id: SessionId,
2217 client: Client,
2218 handlers: SessionHandlers,
2219 hooks: Option<Arc<dyn SessionHooks>>,
2220 transforms: Option<Arc<dyn SystemMessageTransform>>,
2221 command_handlers: Arc<CommandHandlerMap>,
2222 canvas_handler: Option<Arc<dyn CanvasHandler>>,
2223 session_fs_provider: Option<Arc<dyn SessionFsProvider>>,
2224 bearer_token_providers: HashMap<String, Arc<dyn BearerTokenProvider>>,
2225 channels: crate::router::SessionChannels,
2226 idle_waiter: Arc<ParkingLotMutex<Option<IdleWaiter>>>,
2227 capabilities: Arc<parking_lot::RwLock<SessionCapabilities>>,
2228 open_canvases: Arc<parking_lot::RwLock<Vec<OpenCanvasInstance>>>,
2229 event_tx: tokio::sync::broadcast::Sender<SessionEvent>,
2230 shutdown: CancellationToken,
2231 external_tools_shutdown: CancellationToken,
2232) -> JoinHandle<()> {
2233 let crate::router::SessionChannels {
2234 mut notifications,
2235 mut requests,
2236 } = channels;
2237 let pending_external_tools: PendingExternalTools =
2238 Arc::new(ParkingLotMutex::new(HashMap::new()));
2239
2240 let span = tracing::error_span!("session_event_loop", session_id = %session_id);
2241 tokio::spawn(
2242 async move {
2243 loop {
2244 tokio::select! {
2268 _ = shutdown.cancelled() => break,
2269 Some(notification) = notifications.recv() => {
2270 handle_notification(
2271 &session_id, &client, &handlers, &command_handlers, notification, &idle_waiter, &capabilities, &open_canvases, &event_tx, &shutdown, &external_tools_shutdown, &pending_external_tools,
2272 ).await;
2273 }
2274 Some(request) = requests.recv() => {
2275 let span = tracing::error_span!("session_request_handler", session_id = %session_id);
2279 let session_id = session_id.clone();
2280 let client = client.clone();
2281 let handlers = handlers.clone();
2282 let hooks = hooks.clone();
2283 let transforms = transforms.clone();
2284 let canvas_handler = canvas_handler.clone();
2285 let session_fs_provider = session_fs_provider.clone();
2286 let bearer_token_providers = bearer_token_providers.clone();
2287 let request_id = request.id;
2288 let method = request.method.clone();
2289 tokio::spawn(
2290 async move {
2291 let ctx = RequestDispatchContext {
2292 client: &client,
2293 handlers: &handlers,
2294 hooks: hooks.as_deref(),
2295 transforms: transforms.as_deref(),
2296 canvas_handler: canvas_handler.as_ref(),
2297 session_fs_provider: session_fs_provider.as_ref(),
2298 bearer_token_providers: &bearer_token_providers,
2299 };
2300 let dispatch = handle_request(&session_id, ctx, request);
2301 if AssertUnwindSafe(dispatch).catch_unwind().await.is_err() {
2302 error!(method = %method, "request handler panicked");
2305 let _ = send_error_response(
2306 &client,
2307 request_id,
2308 error_codes::INTERNAL_ERROR,
2309 "request handler panicked",
2310 )
2311 .await;
2312 }
2313 }
2314 .instrument(span),
2315 );
2316 }
2317 else => break,
2318 }
2319 }
2320 if let Some(waiter) = idle_waiter.lock().take() {
2323 let _ = waiter
2324 .tx
2325 .send(Err(ErrorKind::Session(SessionErrorKind::EventLoopClosed).into()));
2326 }
2327 }
2328 .instrument(span),
2329 )
2330}
2331
2332fn extract_request_id(data: &Value) -> Option<RequestId> {
2333 data.get("requestId")
2334 .and_then(|v| v.as_str())
2335 .filter(|s| !s.is_empty())
2336 .map(RequestId::new)
2337}
2338
2339fn permission_request_data(
2340 event_data: &Value,
2341 managed_settings_enabled: bool,
2342) -> PermissionRequestData {
2343 let request_data = event_data
2344 .get("permissionRequest")
2345 .cloned()
2346 .unwrap_or_else(|| event_data.clone());
2347 let managed_approval_required = match request_data.get("managedApprovalRequired") {
2348 None => None,
2349 Some(Value::Bool(value)) => Some(*value),
2350 Some(_) => Some(true),
2351 };
2352 match serde_json::from_value::<PermissionRequestData>(request_data) {
2353 Ok(mut data) => {
2354 data.extra = event_data.clone();
2355 data.managed_settings_enabled = managed_settings_enabled;
2356 data
2357 }
2358 Err(_) => PermissionRequestData {
2359 kind: None,
2360 tool_call_id: None,
2361 managed_approval_required,
2362 managed_settings_enabled,
2363 extra: event_data.clone(),
2364 },
2365 }
2366}
2367
2368fn permission_response_params(
2376 session_id: &SessionId,
2377 request_id: &RequestId,
2378 result: &PermissionResult,
2379) -> Option<Value> {
2380 let (decision, decision_context) = match result {
2381 PermissionResult::Decision { decision, context } => (decision, context.clone()),
2382 PermissionResult::NoResult => return None,
2383 };
2384 let mut params = serde_json::to_value(PermissionDecisionRequest {
2385 decision_context,
2386 request_id: request_id.clone(),
2387 result: decision.clone(),
2388 })
2389 .expect("serializing permission response should succeed");
2390 params["sessionId"] =
2391 serde_json::to_value(session_id).expect("serializing session ID should succeed");
2392 Some(params)
2393}
2394
2395async fn register_mcp_auth_interest(client: &Client, session_id: &SessionId) -> Result<(), Error> {
2396 let mut params = serde_json::to_value(RegisterEventInterestParams {
2397 event_type: "mcp.oauth_required".to_string(),
2398 })?;
2399 params["sessionId"] = Value::String(session_id.to_string());
2400 client
2401 .call(rpc_methods::SESSION_EVENTLOG_REGISTERINTEREST, Some(params))
2402 .await?;
2403 Ok(())
2404}
2405
2406fn tool_failure_result(message: impl Into<String>) -> ToolResult {
2407 let message = message.into();
2408 ToolResult::Expanded(ToolResultExpanded {
2409 text_result_for_llm: message.clone(),
2410 result_type: "failure".to_string(),
2411 binary_results_for_llm: None,
2412 session_log: None,
2413 error: Some(message),
2414 tool_telemetry: None,
2415 tool_references: None,
2416 })
2417}
2418
2419fn is_autopilot_continuation_idle(event: &SessionEvent) -> bool {
2420 event
2421 .typed_data::<SessionIdleData>()
2422 .is_some_and(|data| data.mode == Some(SessionMode::Autopilot))
2423}
2424
2425#[allow(clippy::too_many_arguments)]
2427async fn handle_notification(
2428 session_id: &SessionId,
2429 client: &Client,
2430 handlers: &SessionHandlers,
2431 command_handlers: &Arc<CommandHandlerMap>,
2432 notification: SessionEventNotification,
2433 idle_waiter: &Arc<ParkingLotMutex<Option<IdleWaiter>>>,
2434 capabilities: &Arc<parking_lot::RwLock<SessionCapabilities>>,
2435 open_canvases: &Arc<parking_lot::RwLock<Vec<OpenCanvasInstance>>>,
2436 event_tx: &tokio::sync::broadcast::Sender<SessionEvent>,
2437 shutdown: &CancellationToken,
2438 external_tools_shutdown: &CancellationToken,
2439 pending_external_tools: &PendingExternalTools,
2440) {
2441 let dispatch_start = Instant::now();
2442 let event = ¬ification.event;
2443 let event_type = event.parsed_type();
2444 if event_type == SessionEventType::PermissionRequested {
2445 tracing::debug!(
2446 session_id = %session_id,
2447 event_type = %event.event_type,
2448 "Session::handle_notification permission request received"
2449 );
2450 }
2451
2452 match event_type {
2455 SessionEventType::AssistantMessage
2456 | SessionEventType::SessionIdle
2457 | SessionEventType::SessionError => {
2458 let mut guard = idle_waiter.lock();
2459 if let Some(waiter) = guard.as_mut() {
2460 match event_type {
2461 SessionEventType::AssistantMessage => {
2462 if !waiter.first_assistant_message_seen {
2463 waiter.first_assistant_message_seen = true;
2464 tracing::debug!(
2465 elapsed_ms = waiter.started_at.elapsed().as_millis(),
2466 session_id = %session_id,
2467 "Session::send_and_wait first assistant message"
2468 );
2469 }
2470 waiter.last_assistant_message = Some(event.clone());
2471 }
2472 SessionEventType::SessionIdle if is_autopilot_continuation_idle(event) => {}
2473 SessionEventType::SessionIdle | SessionEventType::SessionError => {
2474 if let Some(waiter) = guard.take() {
2475 if event_type == SessionEventType::SessionIdle {
2476 tracing::debug!(
2477 elapsed_ms = waiter.started_at.elapsed().as_millis(),
2478 session_id = %session_id,
2479 "Session::send_and_wait idle received"
2480 );
2481 let _ = waiter.tx.send(Ok(waiter.last_assistant_message));
2482 } else {
2483 let error_msg = event
2484 .typed_data::<SessionErrorData>()
2485 .map(|d| d.message)
2486 .or_else(|| {
2487 event
2488 .data
2489 .get("message")
2490 .and_then(|v| v.as_str())
2491 .map(|s| s.to_string())
2492 })
2493 .unwrap_or_else(|| "session error".to_string());
2494 let _ = waiter.tx.send(Err(Error::with_message(
2495 ErrorKind::Session(SessionErrorKind::AgentError),
2496 error_msg,
2497 )));
2498 }
2499 }
2500 }
2501 _ => {}
2502 }
2503 }
2504 }
2505 _ => {}
2506 }
2507
2508 if event_type == SessionEventType::CapabilitiesChanged {
2512 match serde_json::from_value::<SessionCapabilities>(notification.event.data.clone()) {
2513 Ok(changed) => *capabilities.write() = changed,
2514 Err(e) => warn!(error = %e, "failed to deserialize capabilities.changed payload"),
2515 }
2516 }
2517 if event_type == SessionEventType::SessionCanvasOpened {
2518 match serde_json::from_value::<OpenCanvasInstance>(notification.event.data.clone()) {
2519 Ok(open_canvas) => {
2520 upsert_open_canvas_snapshot(&mut open_canvases.write(), open_canvas);
2521 }
2522 Err(e) => warn!(error = %e, "failed to deserialize session.canvas.opened payload"),
2523 }
2524 }
2525 if event_type == SessionEventType::SessionCanvasClosed {
2526 match serde_json::from_value::<SessionCanvasClosedData>(notification.event.data.clone()) {
2527 Ok(closed) => {
2528 if closed.instance_id.is_empty() {
2529 warn!("failed to deserialize session.canvas.closed payload");
2530 } else {
2531 remove_open_canvas_snapshot(&mut open_canvases.write(), &closed.instance_id);
2532 }
2533 }
2534 Err(e) => warn!(error = %e, "failed to deserialize session.canvas.closed payload"),
2535 }
2536 }
2537
2538 let _ = event_tx.send(event.clone());
2542
2543 tracing::debug!(
2544 elapsed_ms = dispatch_start.elapsed().as_millis(),
2545 session_id = %session_id,
2546 event_type = %notification.event.event_type,
2547 "Session::handle_notification dispatch"
2548 );
2549
2550 match event_type {
2553 SessionEventType::ExternalToolCompleted => {
2554 if let Some(request_id) = extract_request_id(¬ification.event.data)
2555 && let Some(token) = pending_external_tools.lock().remove(&request_id)
2556 {
2557 token.cancel();
2558 }
2559 }
2560 SessionEventType::PermissionRequested => {
2561 let Some(request_id) = extract_request_id(¬ification.event.data) else {
2562 return;
2563 };
2564 if notification
2568 .event
2569 .data
2570 .get("resolvedByHook")
2571 .and_then(|v| v.as_bool())
2572 .unwrap_or(false)
2573 {
2574 return;
2575 }
2576 let Some(permission_handler) = handlers.permission.clone() else {
2580 return;
2581 };
2582 let client = client.clone();
2583 let sid = session_id.clone();
2584 let shutdown = shutdown.clone();
2585 let data = permission_request_data(
2586 ¬ification.event.data,
2587 handlers.managed_settings_enabled,
2588 );
2589 let span = tracing::error_span!(
2590 "permission_request_handler",
2591 session_id = %sid,
2592 request_id = %request_id
2593 );
2594 tokio::spawn(
2595 async move {
2596 let handler_start = Instant::now();
2597 let result = permission_handler
2598 .handle(sid.clone(), request_id.clone(), data)
2599 .await;
2600 tracing::debug!(
2601 elapsed_ms = handler_start.elapsed().as_millis(),
2602 session_id = %sid,
2603 request_id = %request_id,
2604 "PermissionHandler::handle dispatch"
2605 );
2606 let Some(params) = permission_response_params(&sid, &request_id, &result)
2607 else {
2608 return;
2612 };
2613 let rpc_start = Instant::now();
2614 let method =
2615 rpc_methods::SESSION_PERMISSIONS_HANDLEPENDINGPERMISSIONREQUEST;
2616 tokio::select! {
2617 biased;
2618 response = client.call(method, Some(params)) => {
2619 match response {
2620 Ok(_) => tracing::debug!(
2621 elapsed_ms = rpc_start.elapsed().as_millis(),
2622 session_id = %sid,
2623 request_id = %request_id,
2624 method,
2625 "Session::handle_notification response sent successfully"
2626 ),
2627 Err(error) => warn!(
2628 error = %error,
2629 session_id = %sid,
2630 request_id = %request_id,
2631 method,
2632 "failed to deliver permission decision back to the runtime"
2633 ),
2634 }
2635 }
2636 _ = shutdown.cancelled() => {
2637 warn!(
2638 elapsed_ms = rpc_start.elapsed().as_millis(),
2639 session_id = %sid,
2640 request_id = %request_id,
2641 method,
2642 delivery_outcome = "unknown",
2643 "permission confirmation acknowledgement wait cancelled during session shutdown"
2644 );
2645 }
2646 }
2647 }
2648 .instrument(span),
2649 );
2650 }
2651 SessionEventType::ExternalToolRequested => {
2652 let Some(request_id) = extract_request_id(¬ification.event.data) else {
2653 return;
2654 };
2655 let data: ExternalToolRequestedData =
2656 match serde_json::from_value(notification.event.data.clone()) {
2657 Ok(d) => d,
2658 Err(e) => {
2659 warn!(error = %e, "failed to deserialize external_tool.requested");
2660 let client = client.clone();
2661 let sid = session_id.clone();
2662 let span = tracing::error_span!(
2663 "external_tool_deserialize_error",
2664 session_id = %sid,
2665 request_id = %request_id
2666 );
2667 tokio::spawn(
2668 async move {
2669 let rpc_start = Instant::now();
2670 let _ = client
2671 .call(
2672 "session.tools.handlePendingToolCall",
2673 Some(serde_json::json!({
2674 "sessionId": sid,
2675 "requestId": request_id,
2676 "error": format!("Failed to deserialize tool request: {e}"),
2677 })),
2678 )
2679 .await;
2680 tracing::debug!(
2681 elapsed_ms = rpc_start.elapsed().as_millis(),
2682 session_id = %sid,
2683 request_id = %request_id,
2684 "Session::handle_notification response sent successfully"
2685 );
2686 }
2687 .instrument(span),
2688 );
2689 return;
2690 }
2691 };
2692 let tool_handler = if data.tool_name.is_empty() {
2696 None
2697 } else {
2698 handlers.tools.get(&data.tool_name).cloned()
2699 };
2700 let Some(tool_handler) = tool_handler else {
2701 return;
2702 };
2703 let cancellation = Arc::new(external_tools_shutdown.child_token());
2704 {
2705 let mut pending = pending_external_tools.lock();
2706 if external_tools_shutdown.is_cancelled() || pending.contains_key(&request_id) {
2707 return;
2708 }
2709 pending.insert(request_id.clone(), cancellation.clone());
2710 }
2711 let client = client.clone();
2712 let sid = session_id.clone();
2713 let pending_external_tools = pending_external_tools.clone();
2714 let guard_request_id = request_id.clone();
2715 let guard_cancellation = cancellation.clone();
2716 let span = tracing::error_span!(
2717 "external_tool_handler",
2718 session_id = %sid,
2719 request_id = %request_id
2720 );
2721 tokio::spawn(
2722 async move {
2723 let guard = PendingExternalToolGuard {
2724 request_id: guard_request_id,
2725 token: guard_cancellation,
2726 pending: pending_external_tools,
2727 };
2728 if cancellation.is_cancelled() {
2729 return;
2730 }
2731 if data.tool_call_id.is_empty() {
2736 if !guard.claim() {
2737 return;
2738 }
2739 let error_msg = "Missing toolCallId";
2740 let rpc_start = Instant::now();
2741 let _ = client
2742 .call(
2743 "session.tools.handlePendingToolCall",
2744 Some(serde_json::json!({
2745 "sessionId": sid,
2746 "requestId": request_id,
2747 "error": error_msg,
2748 })),
2749 )
2750 .await;
2751 tracing::debug!(
2752 elapsed_ms = rpc_start.elapsed().as_millis(),
2753 session_id = %sid,
2754 request_id = %request_id,
2755 "Session::handle_notification response sent successfully"
2756 );
2757 return;
2758 }
2759 let tool_call_id = data.tool_call_id.clone();
2760 let tool_name = data.tool_name.clone();
2761 let available_tools = if tool_name == TOOL_SEARCH_TOOL_NAME {
2768 let metadata_result = tokio::select! {
2769 biased;
2770 _ = cancellation.cancelled() => return,
2771 result = client.call(
2772 rpc_methods::SESSION_TOOLS_GETCURRENTMETADATA,
2773 Some(serde_json::json!({ "sessionId": sid })),
2774 ) => result,
2775 };
2776 match metadata_result {
2777 Ok(value) => {
2778 serde_json::from_value::<ToolsGetCurrentMetadataResult>(value)
2779 .ok()
2780 .and_then(|result| result.tools)
2781 }
2782 Err(_) => None,
2783 }
2784 } else {
2785 None
2786 };
2787 let invocation = ToolInvocation {
2788 session_id: sid.clone(),
2789 tool_call_id: data.tool_call_id,
2790 tool_name: data.tool_name,
2791 arguments: data
2792 .arguments
2793 .unwrap_or(Value::Object(serde_json::Map::new())),
2794 available_tools,
2795 traceparent: data.traceparent,
2796 tracestate: data.tracestate,
2797 };
2798 let handler_start = Instant::now();
2799 let tool_result = tokio::select! {
2800 biased;
2801 _ = cancellation.cancelled() => return,
2802 result = tool_handler.call(invocation) => match result {
2803 Ok(r) => r,
2804 Err(e) => tool_failure_result(e.to_string()),
2805 },
2806 };
2807 tracing::debug!(
2808 elapsed_ms = handler_start.elapsed().as_millis(),
2809 session_id = %sid,
2810 request_id = %request_id,
2811 tool_call_id = %tool_call_id,
2812 tool_name = %tool_name,
2813 "ToolHandler::call dispatch"
2814 );
2815 if !guard.claim() {
2816 return;
2817 }
2818 let result_value = serde_json::to_value(tool_result).unwrap_or(Value::Null);
2819 let rpc_start = Instant::now();
2820 let _ = client
2821 .call(
2822 "session.tools.handlePendingToolCall",
2823 Some(serde_json::json!({
2824 "sessionId": sid,
2825 "requestId": request_id,
2826 "result": result_value,
2827 })),
2828 )
2829 .await;
2830 tracing::debug!(
2831 elapsed_ms = rpc_start.elapsed().as_millis(),
2832 session_id = %sid,
2833 request_id = %request_id,
2834 tool_call_id = %tool_call_id,
2835 tool_name = %tool_name,
2836 "Session::handle_notification response sent successfully"
2837 );
2838 }
2839 .instrument(span),
2840 );
2841 }
2842 SessionEventType::UserInputRequested => {
2843 }
2850 SessionEventType::ElicitationRequested => {
2851 let Some(request_id) = extract_request_id(¬ification.event.data) else {
2852 return;
2853 };
2854 let Some(elicitation_handler) = handlers.elicitation.clone() else {
2858 return;
2859 };
2860 let elicitation_data: ElicitationRequestedData =
2861 match serde_json::from_value(notification.event.data.clone()) {
2862 Ok(d) => d,
2863 Err(e) => {
2864 warn!(error = %e, "failed to deserialize elicitation request");
2865 return;
2866 }
2867 };
2868 let request = ElicitationRequest {
2869 message: elicitation_data.message,
2870 requested_schema: elicitation_data
2871 .requested_schema
2872 .map(|s| serde_json::to_value(s).unwrap_or(Value::Null)),
2873 mode: elicitation_data.mode.map(|m| match m {
2874 crate::generated::session_events::ElicitationRequestedMode::Form => {
2875 crate::types::ElicitationMode::Form
2876 }
2877 crate::generated::session_events::ElicitationRequestedMode::Url => {
2878 crate::types::ElicitationMode::Url
2879 }
2880 _ => crate::types::ElicitationMode::Unknown,
2881 }),
2882 elicitation_source: elicitation_data.elicitation_source,
2883 url: elicitation_data.url,
2884 };
2885 let client = client.clone();
2886 let sid = session_id.clone();
2887 let span = tracing::error_span!(
2888 "elicitation_request_handler",
2889 session_id = %sid,
2890 request_id = %request_id
2891 );
2892 tokio::spawn(
2893 async move {
2894 let cancel = ElicitationResult {
2895 action: "cancel".to_string(),
2896 content: None,
2897 };
2898 let handler_task = tokio::spawn({
2900 let sid = sid.clone();
2901 let request_id = request_id.clone();
2902 let span = tracing::error_span!(
2903 "elicitation_callback",
2904 session_id = %sid,
2905 request_id = %request_id
2906 );
2907 async move {
2908 let handler_start = Instant::now();
2909 let response = elicitation_handler
2910 .handle(sid.clone(), request_id.clone(), request)
2911 .await;
2912 tracing::debug!(
2913 elapsed_ms = handler_start.elapsed().as_millis(),
2914 session_id = %sid,
2915 request_id = %request_id,
2916 "ElicitationHandler::handle dispatch"
2917 );
2918 response
2919 }
2920 .instrument(span)
2921 });
2922 let result = match handler_task.await {
2923 Ok(r) => r,
2924 Err(_) => cancel.clone(),
2925 };
2926 let rpc_start = Instant::now();
2927 if let Err(e) = client
2928 .call(
2929 "session.ui.handlePendingElicitation",
2930 Some(serde_json::json!({
2931 "sessionId": sid,
2932 "requestId": request_id,
2933 "result": result,
2934 })),
2935 )
2936 .await
2937 {
2938 warn!(error = %e, "handlePendingElicitation failed, sending cancel");
2940 let _ = client
2941 .call(
2942 "session.ui.handlePendingElicitation",
2943 Some(serde_json::json!({
2944 "sessionId": sid,
2945 "requestId": request_id,
2946 "result": cancel,
2947 })),
2948 )
2949 .await;
2950 } else {
2951 tracing::debug!(
2952 elapsed_ms = rpc_start.elapsed().as_millis(),
2953 session_id = %sid,
2954 request_id = %request_id,
2955 "Session::handle_notification response sent successfully"
2956 );
2957 }
2958 }
2959 .instrument(span),
2960 );
2961 }
2962 SessionEventType::McpOauthRequired => {
2963 let Some(request_id) = extract_request_id(¬ification.event.data) else {
2964 return;
2965 };
2966 let Some(mcp_auth_handler) = handlers.mcp_auth.clone() else {
2967 warn!(
2968 session_id = %session_id,
2969 request_id = %request_id,
2970 "received MCP OAuth request without a registered MCP auth handler"
2971 );
2972 return;
2973 };
2974 let data: McpOauthRequiredData =
2975 match serde_json::from_value(notification.event.data.clone()) {
2976 Ok(d) => d,
2977 Err(e) => {
2978 warn!(error = %e, "failed to deserialize MCP OAuth request");
2979 return;
2980 }
2981 };
2982 let request = McpAuthRequest {
2983 request_id: request_id.clone(),
2984 server_name: data.server_name,
2985 server_url: data.server_url,
2986 reason: data.reason,
2987 www_authenticate_params: data.www_authenticate_params,
2988 resource_metadata: data.resource_metadata,
2989 static_client_config: data.static_client_config,
2990 };
2991 let client = client.clone();
2992 let sid = session_id.clone();
2993 let span = tracing::error_span!(
2994 "mcp_auth_request_handler",
2995 session_id = %sid,
2996 request_id = %request_id
2997 );
2998 tokio::spawn(
2999 async move {
3000 let cancel = McpAuthResult::Cancelled;
3001 let handler_task = tokio::spawn({
3002 let sid = sid.clone();
3003 let request_id = request_id.clone();
3004 let span = tracing::error_span!(
3005 "mcp_auth_callback",
3006 session_id = %sid,
3007 request_id = %request_id
3008 );
3009 async move {
3010 let handler_start = Instant::now();
3011 let response = mcp_auth_handler
3012 .handle(sid.clone(), request_id.clone(), request)
3013 .await;
3014 tracing::debug!(
3015 elapsed_ms = handler_start.elapsed().as_millis(),
3016 session_id = %sid,
3017 request_id = %request_id,
3018 "McpAuthHandler::handle dispatch"
3019 );
3020 response
3021 }
3022 .instrument(span)
3023 });
3024 let result = match handler_task.await {
3025 Ok(result) => result,
3026 Err(_) => cancel,
3027 };
3028 let rpc_start = Instant::now();
3029 let _ = client
3030 .call(
3031 "session.mcp.oauth.handlePendingRequest",
3032 Some(serde_json::json!({
3033 "sessionId": sid,
3034 "requestId": request_id,
3035 "result": result.into_wire(),
3036 })),
3037 )
3038 .await;
3039 tracing::debug!(
3040 elapsed_ms = rpc_start.elapsed().as_millis(),
3041 "Session::handle_notification MCP auth response sent"
3042 );
3043 }
3044 .instrument(span),
3045 );
3046 }
3047 SessionEventType::CommandExecute => {
3048 let data: CommandExecuteData =
3049 match serde_json::from_value(notification.event.data.clone()) {
3050 Ok(d) => d,
3051 Err(e) => {
3052 warn!(error = %e, "failed to deserialize command.execute");
3053 return;
3054 }
3055 };
3056 let client = client.clone();
3057 let command_handlers = command_handlers.clone();
3058 let sid = session_id.clone();
3059 let span = tracing::error_span!("command_handler", session_id = %sid);
3060 tokio::spawn(
3061 async move {
3062 let request_id = data.request_id;
3063 let ack_error = match command_handlers.get(&data.command_name).cloned() {
3064 None => Some(format!("Unknown command: {}", data.command_name)),
3065 Some(handler) => {
3066 let command_name = data.command_name.clone();
3067 let ctx = CommandContext {
3068 session_id: sid.clone(),
3069 command: data.command,
3070 command_name: data.command_name,
3071 args: data.args,
3072 };
3073 let handler_start = Instant::now();
3074 let result = handler.on_command(ctx).await;
3075 tracing::debug!(
3076 elapsed_ms = handler_start.elapsed().as_millis(),
3077 session_id = %sid,
3078 request_id = %request_id,
3079 command_name = %command_name,
3080 "CommandHandler::call dispatch"
3081 );
3082 match result {
3083 Ok(()) => None,
3084 Err(e) => Some(e.to_string()),
3085 }
3086 }
3087 };
3088 let mut params = serde_json::json!({
3089 "sessionId": sid,
3090 "requestId": request_id,
3091 });
3092 if let Some(error_msg) = ack_error {
3093 params["error"] = serde_json::Value::String(error_msg);
3094 }
3095 let rpc_start = Instant::now();
3096 let _ = client
3097 .call("session.commands.handlePendingCommand", Some(params))
3098 .await;
3099 tracing::debug!(
3100 elapsed_ms = rpc_start.elapsed().as_millis(),
3101 session_id = %sid,
3102 request_id = %request_id,
3103 "Session::handle_notification response sent successfully"
3104 );
3105 }
3106 .instrument(span),
3107 );
3108 }
3109 _ => {}
3110 }
3111}
3112
3113struct RequestDispatchContext<'a> {
3114 client: &'a Client,
3115 handlers: &'a SessionHandlers,
3116 hooks: Option<&'a dyn SessionHooks>,
3117 transforms: Option<&'a dyn SystemMessageTransform>,
3118 canvas_handler: Option<&'a Arc<dyn CanvasHandler>>,
3119 session_fs_provider: Option<&'a Arc<dyn SessionFsProvider>>,
3120 bearer_token_providers: &'a HashMap<String, Arc<dyn BearerTokenProvider>>,
3121}
3122
3123async fn handle_request(
3125 session_id: &SessionId,
3126 ctx: RequestDispatchContext<'_>,
3127 request: crate::JsonRpcRequest,
3128) {
3129 let sid = session_id.clone();
3130 let client = ctx.client;
3131 let handlers = ctx.handlers;
3132 let hooks = ctx.hooks;
3133 let transforms = ctx.transforms;
3134 let canvas_handler = ctx.canvas_handler;
3135 let session_fs_provider = ctx.session_fs_provider;
3136 let bearer_token_providers = ctx.bearer_token_providers;
3137
3138 if request.method.starts_with("sessionFs.") {
3139 crate::session_fs_dispatch::dispatch(client, session_fs_provider, request).await;
3140 return;
3141 }
3142
3143 if request.method.starts_with("canvas.") {
3144 crate::canvas_dispatch::dispatch(client, canvas_handler, request).await;
3145 return;
3146 }
3147
3148 if request.method == crate::generated::api_types::rpc_methods::PROVIDERTOKEN_GETTOKEN {
3149 crate::provider_token_dispatch::dispatch(client, bearer_token_providers, request).await;
3150 return;
3151 }
3152
3153 match request.method.as_str() {
3154 "hooks.invoke" => {
3155 let params = request.params.as_ref();
3156 let hook_type = params
3157 .and_then(|p| p.get("hookType"))
3158 .and_then(|v| v.as_str())
3159 .unwrap_or("");
3160 let input = params
3161 .and_then(|p| p.get("input"))
3162 .cloned()
3163 .unwrap_or(Value::Object(Default::default()));
3164
3165 let rpc_result = if let Some(hooks) = hooks {
3166 match crate::hooks::dispatch_hook(hooks, &sid, hook_type, input).await {
3167 Ok(output) => output,
3168 Err(e) => {
3169 warn!(error = %e, hook_type = hook_type, "hook dispatch failed");
3170 serde_json::json!({ "output": {} })
3171 }
3172 }
3173 } else {
3174 serde_json::json!({ "output": {} })
3175 };
3176
3177 let rpc_response = JsonRpcResponse {
3178 jsonrpc: "2.0".to_string(),
3179 id: request.id,
3180 result: Some(rpc_result),
3181 error: None,
3182 };
3183 let _ = client.send_response(&rpc_response).await;
3184 }
3185
3186 "userInput.request" => {
3187 let params = request.params.as_ref();
3188 let Some(question) = params
3189 .and_then(|p| p.get("question"))
3190 .and_then(|v| v.as_str())
3191 else {
3192 warn!("userInput.request missing 'question' field");
3193 let rpc_response = JsonRpcResponse {
3194 jsonrpc: "2.0".to_string(),
3195 id: request.id,
3196 result: None,
3197 error: Some(crate::JsonRpcError {
3198 code: error_codes::INVALID_PARAMS,
3199 message: "missing required field: question".to_string(),
3200 data: None,
3201 }),
3202 };
3203 let _ = client.send_response(&rpc_response).await;
3204 return;
3205 };
3206 let question = question.to_string();
3207 let choices = params
3208 .and_then(|p| p.get("choices"))
3209 .and_then(|v| v.as_array())
3210 .map(|arr| {
3211 arr.iter()
3212 .filter_map(|v| v.as_str().map(|s| s.to_string()))
3213 .collect()
3214 });
3215 let allow_freeform = params
3216 .and_then(|p| p.get("allowFreeform"))
3217 .and_then(|v| v.as_bool());
3218
3219 let handler_start = Instant::now();
3220 let response = if let Some(user_input_handler) = handlers.user_input.as_ref() {
3221 user_input_handler
3222 .handle(sid.clone(), question, choices, allow_freeform)
3223 .await
3224 } else {
3225 None
3226 };
3227 tracing::debug!(
3228 elapsed_ms = handler_start.elapsed().as_millis(),
3229 session_id = %sid,
3230 "UserInputHandler::handle dispatch"
3231 );
3232
3233 let rpc_result = match response {
3234 Some(UserInputResponse {
3235 answer,
3236 was_freeform,
3237 }) => serde_json::json!({
3238 "answer": answer,
3239 "wasFreeform": was_freeform,
3240 }),
3241 None => serde_json::json!({ "noResponse": true }),
3242 };
3243 let rpc_response = JsonRpcResponse {
3244 jsonrpc: "2.0".to_string(),
3245 id: request.id,
3246 result: Some(rpc_result),
3247 error: None,
3248 };
3249 let _ = client.send_response(&rpc_response).await;
3250 }
3251
3252 "exitPlanMode.request" => {
3253 let params = request
3254 .params
3255 .as_ref()
3256 .cloned()
3257 .unwrap_or(Value::Object(serde_json::Map::new()));
3258 let data: ExitPlanModeData = match serde_json::from_value(params) {
3259 Ok(d) => d,
3260 Err(e) => {
3261 warn!(error = %e, "failed to deserialize exitPlanMode.request params, using defaults");
3262 ExitPlanModeData::default()
3263 }
3264 };
3265
3266 let rpc_result = if let Some(exit_plan_handler) = handlers.exit_plan_mode.as_ref() {
3267 let result = exit_plan_handler.handle(sid, data).await;
3268 serde_json::to_value(result).expect("ExitPlanModeResult serialization cannot fail")
3269 } else {
3270 serde_json::json!({ "approved": true })
3271 };
3272 let rpc_response = JsonRpcResponse {
3273 jsonrpc: "2.0".to_string(),
3274 id: request.id,
3275 result: Some(rpc_result),
3276 error: None,
3277 };
3278 let _ = client.send_response(&rpc_response).await;
3279 }
3280
3281 "autoModeSwitch.request" => {
3282 let error_code = request
3283 .params
3284 .as_ref()
3285 .and_then(|p| p.get("errorCode"))
3286 .and_then(|v| v.as_str())
3287 .map(|s| s.to_string());
3288 let retry_after_seconds = request
3289 .params
3290 .as_ref()
3291 .and_then(|p| p.get("retryAfterSeconds"))
3292 .and_then(|v| v.as_f64());
3293
3294 let answer = if let Some(auto_mode_handler) = handlers.auto_mode_switch.as_ref() {
3295 auto_mode_handler
3296 .handle(sid, error_code, retry_after_seconds)
3297 .await
3298 } else {
3299 AutoModeSwitchResponse::No
3300 };
3301 let rpc_response = JsonRpcResponse {
3302 jsonrpc: "2.0".to_string(),
3303 id: request.id,
3304 result: Some(serde_json::json!({ "response": answer })),
3305 error: None,
3306 };
3307 let _ = client.send_response(&rpc_response).await;
3308 }
3309
3310 "systemMessage.transform" => {
3311 let params = request.params.as_ref();
3312 let sections: HashMap<String, crate::transforms::TransformSection> =
3313 match params.and_then(|p| p.get("sections")) {
3314 Some(v) => match serde_json::from_value(v.clone()) {
3315 Ok(s) => s,
3316 Err(e) => {
3317 let _ = send_error_response(
3318 client,
3319 request.id,
3320 error_codes::INVALID_PARAMS,
3321 &format!("invalid sections: {e}"),
3322 )
3323 .await;
3324 return;
3325 }
3326 },
3327 None => {
3328 let _ = send_error_response(
3329 client,
3330 request.id,
3331 error_codes::INVALID_PARAMS,
3332 "missing sections parameter",
3333 )
3334 .await;
3335 return;
3336 }
3337 };
3338
3339 let rpc_result = if let Some(transforms) = transforms {
3340 let transform_start = Instant::now();
3341 let response =
3342 crate::transforms::dispatch_transform(transforms, &sid, sections).await;
3343 tracing::debug!(
3344 elapsed_ms = transform_start.elapsed().as_millis(),
3345 session_id = %sid,
3346 "SystemMessageTransform::transform_section dispatch"
3347 );
3348 match serde_json::to_value(response) {
3349 Ok(v) => v,
3350 Err(e) => {
3351 warn!(error = %e, "failed to serialize transform response");
3352 serde_json::json!({ "sections": {} })
3353 }
3354 }
3355 } else {
3356 let passthrough: HashMap<String, crate::transforms::TransformSection> = sections;
3358 serde_json::json!({ "sections": passthrough })
3359 };
3360
3361 let rpc_response = JsonRpcResponse {
3362 jsonrpc: "2.0".to_string(),
3363 id: request.id,
3364 result: Some(rpc_result),
3365 error: None,
3366 };
3367 let _ = client.send_response(&rpc_response).await;
3368 }
3369
3370 method => {
3371 warn!(
3372 method = method,
3373 "unhandled request method in session event loop"
3374 );
3375 let _ = send_error_response(
3376 client,
3377 request.id,
3378 error_codes::METHOD_NOT_FOUND,
3379 &format!("unknown method: {method}"),
3380 )
3381 .await;
3382 }
3383 }
3384}
3385
3386async fn send_error_response(
3387 client: &Client,
3388 id: u64,
3389 code: i32,
3390 message: &str,
3391) -> Result<(), Error> {
3392 let response = JsonRpcResponse {
3393 jsonrpc: "2.0".to_string(),
3394 id,
3395 result: None,
3396 error: Some(crate::JsonRpcError {
3397 code,
3398 message: message.to_string(),
3399 data: None,
3400 }),
3401 };
3402 client.send_response(&response).await
3403}
3404
3405fn apply_transform_sections(
3409 sys_msg: &mut SystemMessageConfig,
3410 transforms: &dyn SystemMessageTransform,
3411) {
3412 sys_msg.mode = Some("customize".to_string());
3413 let sections = sys_msg.sections.get_or_insert_with(HashMap::new);
3414 for id in transforms.section_ids() {
3415 sections.entry(id).or_insert_with(|| SectionOverride {
3416 action: Some("transform".to_string()),
3417 content: None,
3418 });
3419 }
3420}
3421
3422fn inject_transform_sections(config: &mut SessionConfig, transforms: &dyn SystemMessageTransform) {
3423 let sys_msg = config.system_message.get_or_insert_with(Default::default);
3424 apply_transform_sections(sys_msg, transforms);
3425}
3426
3427fn inject_transform_sections_resume(
3428 config: &mut ResumeSessionConfig,
3429 transforms: &dyn SystemMessageTransform,
3430) {
3431 let sys_msg = config.system_message.get_or_insert_with(Default::default);
3432 apply_transform_sections(sys_msg, transforms);
3433}
3434
3435#[cfg(test)]
3436mod tests {
3437 use serde_json::json;
3438
3439 use super::{
3440 build_mode_post_create_patch, has_managed_settings, is_autopilot_continuation_idle,
3441 permission_request_data, permission_response_params,
3442 };
3443 use crate::handler::PermissionResult;
3444 use crate::types::{
3445 PermissionDecisionContext, PermissionDecisionOutcome, PermissionDecisionSource,
3446 PermissionDecisionSurface, RequestId, SessionEvent, SessionId,
3447 };
3448
3449 #[test]
3450 fn identifies_only_autopilot_continuation_idles() {
3451 let mut event = SessionEvent {
3452 id: "event-1".to_string(),
3453 timestamp: "2026-01-01T00:00:00Z".to_string(),
3454 parent_id: None,
3455 ephemeral: None,
3456 agent_id: None,
3457 debug_cli_received_at_ms: None,
3458 debug_ws_forwarded_at_ms: None,
3459 event_type: "session.idle".to_string(),
3460 data: json!({ "mode": "autopilot" }),
3461 };
3462
3463 assert!(is_autopilot_continuation_idle(&event));
3464
3465 event.data = json!({ "mode": "interactive" });
3466 assert!(!is_autopilot_continuation_idle(&event));
3467
3468 event.data = json!({});
3469 assert!(!is_autopilot_continuation_idle(&event));
3470 }
3471
3472 #[test]
3473 fn empty_mode_post_patch_sets_empty_included_builtin_skills() {
3474 let patch =
3475 build_mode_post_create_patch(crate::ClientMode::Empty, None, None, None, None, None)
3476 .expect("empty mode always sends a patch");
3477 assert_eq!(
3478 patch.included_builtin_skills,
3479 Some(Vec::new()),
3480 "empty mode must fail closed with an empty includedBuiltinSkills list"
3481 );
3482 assert_eq!(patch.installed_plugins.as_ref().map(|p| p.len()), Some(0));
3483 let value = serde_json::to_value(&patch).expect("serialize patch");
3485 assert_eq!(value["includedBuiltinSkills"], serde_json::json!([]));
3486 }
3487
3488 #[test]
3489 fn empty_mode_post_patch_preserves_explicit_builtin_skill_allowlist() {
3490 let patch = build_mode_post_create_patch(
3491 crate::ClientMode::Empty,
3492 Some(false),
3493 Some(false),
3494 Some(true),
3495 Some(true),
3496 Some(vec!["code-review".to_string()]),
3497 )
3498 .expect("empty mode always sends a patch");
3499 assert_eq!(
3500 patch.included_builtin_skills,
3501 Some(vec!["code-review".to_string()])
3502 );
3503 }
3504
3505 #[test]
3506 fn copilot_cli_mode_does_not_inject_included_builtin_skills() {
3507 assert!(
3509 build_mode_post_create_patch(
3510 crate::ClientMode::CopilotCli,
3511 None,
3512 None,
3513 None,
3514 None,
3515 None
3516 )
3517 .is_none()
3518 );
3519 let patch = build_mode_post_create_patch(
3521 crate::ClientMode::CopilotCli,
3522 Some(true),
3523 None,
3524 None,
3525 None,
3526 None,
3527 )
3528 .expect("a set field triggers a patch");
3529 assert_eq!(patch.included_builtin_skills, None);
3530 assert!(patch.installed_plugins.is_none());
3531 let value = serde_json::to_value(&patch).expect("serialize patch");
3532 assert!(value.get("includedBuiltinSkills").is_none());
3533
3534 let patch = build_mode_post_create_patch(
3535 crate::ClientMode::CopilotCli,
3536 None,
3537 None,
3538 None,
3539 None,
3540 Some(vec!["code-review".to_string()]),
3541 )
3542 .expect("an explicit allowlist triggers a patch");
3543 assert_eq!(
3544 patch.included_builtin_skills,
3545 Some(vec!["code-review".to_string()])
3546 );
3547 }
3548
3549 #[test]
3550 fn direct_injection_enables_managed_safeguards() {
3551 let settings = crate::types::ManagedSettings::default();
3552 assert!(has_managed_settings(None, Some(&settings)));
3553 assert!(!has_managed_settings(None, None));
3554 }
3555
3556 fn attribution_context() -> PermissionDecisionContext {
3557 PermissionDecisionContext {
3558 outcome: PermissionDecisionOutcome::AutoApproved,
3559 response_capability: None,
3560 source: PermissionDecisionSource::AssistedApproval,
3561 surface: PermissionDecisionSurface::CopilotApp,
3562 }
3563 }
3564
3565 #[test]
3566 fn response_params_omit_decision_context_without_attribution() {
3567 for (result, expected) in [
3568 (
3569 PermissionResult::approve_once(),
3570 json!({ "kind": "approve-once" }),
3571 ),
3572 (PermissionResult::reject(None), json!({ "kind": "reject" })),
3573 (
3574 PermissionResult::reject(Some("bad".to_string())),
3575 json!({ "kind": "reject", "feedback": "bad" }),
3576 ),
3577 (
3578 PermissionResult::user_not_available(),
3579 json!({ "kind": "user-not-available" }),
3580 ),
3581 ] {
3582 let params = permission_response_params(
3583 &SessionId::from("session-1"),
3584 &RequestId::from("permission-1"),
3585 &result,
3586 )
3587 .unwrap();
3588 assert_eq!(
3589 params,
3590 json!({
3591 "sessionId": "session-1",
3592 "requestId": "permission-1",
3593 "result": expected,
3594 })
3595 );
3596 }
3597 }
3598
3599 #[test]
3600 fn response_params_forward_decision_context_alongside_result() {
3601 let params = permission_response_params(
3602 &SessionId::from("session-1"),
3603 &RequestId::from("permission-1"),
3604 &PermissionResult::approve_once().with_context(attribution_context()),
3605 )
3606 .unwrap();
3607 assert_eq!(
3608 params,
3609 json!({
3610 "sessionId": "session-1",
3611 "requestId": "permission-1",
3612 "result": { "kind": "approve-once" },
3613 "decisionContext": {
3614 "outcome": "auto_approved",
3615 "source": "assisted_approval",
3616 "surface": "copilot_app",
3617 },
3618 })
3619 );
3620 assert!(params["result"].get("decisionContext").is_none());
3622 }
3623
3624 #[test]
3625 fn response_params_suppressed_for_no_result() {
3626 assert!(
3627 permission_response_params(
3628 &SessionId::from("session-1"),
3629 &RequestId::from("permission-1"),
3630 &PermissionResult::NoResult,
3631 )
3632 .is_none()
3633 );
3634 }
3635
3636 #[test]
3637 fn with_context_is_a_no_op_on_no_result() {
3638 let result = PermissionResult::no_result().with_context(attribution_context());
3639 assert!(matches!(result, PermissionResult::NoResult));
3640 }
3641
3642 #[test]
3643 fn with_context_replaces_rather_than_nests() {
3644 let result = PermissionResult::approve_once()
3645 .with_context(attribution_context())
3646 .with_context(PermissionDecisionContext {
3647 outcome: PermissionDecisionOutcome::PromptedUser,
3648 response_capability: None,
3649 source: PermissionDecisionSource::HumanResponse,
3650 surface: PermissionDecisionSurface::Sdk,
3651 });
3652 let params = permission_response_params(
3653 &SessionId::from("session-1"),
3654 &RequestId::from("permission-1"),
3655 &result,
3656 )
3657 .unwrap();
3658 assert_eq!(
3659 params["decisionContext"],
3660 json!({
3661 "outcome": "prompted_user",
3662 "source": "human_response",
3663 "surface": "sdk",
3664 })
3665 );
3666 }
3667
3668 #[test]
3669 fn permission_request_data_reads_nested_managed_approval_metadata() {
3670 let data = permission_request_data(
3671 &json!({
3672 "requestId": "permission-1",
3673 "permissionRequest": {
3674 "kind": "read",
3675 "managedApprovalRequired": true,
3676 "path": "/workspace/file.txt"
3677 }
3678 }),
3679 false,
3680 );
3681
3682 assert_eq!(data.managed_approval_required, Some(true));
3683 assert_eq!(
3684 data.extra["permissionRequest"]["path"],
3685 "/workspace/file.txt"
3686 );
3687 }
3688
3689 #[test]
3690 fn permission_request_data_preserves_managed_flag_when_other_fields_are_malformed() {
3691 let data = permission_request_data(
3692 &json!({
3693 "requestId": "permission-1",
3694 "permissionRequest": {
3695 "kind": "read",
3696 "managedApprovalRequired": true,
3697 "toolCallId": 42
3698 }
3699 }),
3700 false,
3701 );
3702
3703 assert_eq!(data.managed_approval_required, Some(true));
3704 assert_eq!(data.extra["requestId"], "permission-1");
3705 }
3706
3707 #[test]
3708 fn permission_request_data_fails_closed_for_malformed_managed_flag() {
3709 let data = permission_request_data(
3710 &json!({
3711 "requestId": "permission-1",
3712 "permissionRequest": {
3713 "kind": "read",
3714 "managedApprovalRequired": "yes",
3715 "path": "/workspace/file.txt"
3716 }
3717 }),
3718 false,
3719 );
3720
3721 assert_eq!(data.managed_approval_required, Some(true));
3722 }
3723
3724 #[test]
3725 fn permission_request_data_preserves_valid_false_managed_flag() {
3726 let data = permission_request_data(
3727 &json!({
3728 "requestId": "permission-1",
3729 "permissionRequest": {
3730 "kind": "read",
3731 "managedApprovalRequired": false,
3732 "path": "/workspace/file.txt"
3733 }
3734 }),
3735 false,
3736 );
3737
3738 assert_eq!(data.managed_approval_required, Some(false));
3739 }
3740}