1use std::sync::Arc;
9
10use tokio::sync::mpsc;
11
12use serde_json::Value;
13
14use crate::tools::{BashCompletionSink, ToolCall, ToolSchema};
15use crate::{AgentEvent, Session};
16use bamboo_domain::{
17 PermissionMode, SessionAuthorityIdentity, SessionKind, SupervisorReference,
18 DEFAULT_SUPERVISOR_SESSION_ID,
19};
20use uuid::Uuid;
21
22#[derive(Clone, Copy, Debug, PartialEq, Eq)]
28pub struct ExecutingSupervisorObservation {
29 incarnation_id: Uuid,
30}
31
32const SUPERVISOR_PERMISSION_REPLAY_METADATA_KEY: &str = "permission.executing_supervisor.v1";
34
35#[derive(serde::Serialize, serde::Deserialize)]
36#[serde(deny_unknown_fields)]
37struct SupervisorPermissionReplayRecord {
38 version: u8,
39 incarnation_id: Uuid,
40 session_id: String,
41 result_message_id: String,
42 request_generation: String,
43 tool_call: ToolCall,
44 execution_name: String,
45}
46
47impl ExecutingSupervisorObservation {
48 pub const PERMISSION_REPLAY_METADATA_KEY: &'static str =
49 SUPERVISOR_PERMISSION_REPLAY_METADATA_KEY;
50
51 pub fn capture_from_executing_session(session: &Session) -> Option<Self> {
55 let SessionAuthorityIdentity::Supervisor { incarnation_id } = &session.authority_identity
56 else {
57 return None;
58 };
59 (session.id == DEFAULT_SUPERVISOR_SESSION_ID
60 && session.kind == SessionKind::Root
61 && session.root_session_id == session.id
62 && session.parent_session_id.is_none()
63 && session.spawn_depth == 0
64 && !incarnation_id.is_nil())
65 .then_some(Self {
66 incarnation_id: *incarnation_id,
67 })
68 }
69
70 pub fn supervisor_reference(self) -> SupervisorReference {
73 SupervisorReference {
74 session_id: DEFAULT_SUPERVISOR_SESSION_ID.to_string(),
75 incarnation_id: self.incarnation_id,
76 }
77 }
78
79 pub fn permission_replay_record(
83 self,
84 session_id: &str,
85 result_message_id: &str,
86 tool_call: &ToolCall,
87 execution_name: &str,
88 request_generation: &str,
89 ) -> Value {
90 serde_json::to_value(SupervisorPermissionReplayRecord {
91 version: 1,
92 incarnation_id: self.incarnation_id,
93 session_id: session_id.to_string(),
94 result_message_id: result_message_id.to_string(),
95 request_generation: request_generation.to_string(),
96 tool_call: tool_call.clone(),
97 execution_name: execution_name.to_string(),
98 })
99 .expect("Supervisor permission binding serializes")
100 }
101
102 pub fn restore_permission_replay_record(
107 session: &Session,
108 result_index: usize,
109 tool_call: &ToolCall,
110 execution_name: &str,
111 request_generation: Option<&str>,
112 ) -> Result<Option<Self>, &'static str> {
113 let message = session
114 .messages
115 .get(result_index)
116 .ok_or("result occurrence missing")?;
117 if serde_json::from_str::<Value>(&message.content)
118 .ok()
119 .is_some_and(|payload| {
120 payload
121 .get(SUPERVISOR_PERMISSION_REPLAY_METADATA_KEY)
122 .is_some()
123 })
124 {
125 return Err("Supervisor authority is not accepted from a result payload");
126 }
127 let Some(value) = message
128 .metadata
129 .as_ref()
130 .and_then(|metadata| metadata.get(SUPERVISOR_PERMISSION_REPLAY_METADATA_KEY))
131 else {
132 return Ok(None);
133 };
134 let record: SupervisorPermissionReplayRecord = serde_json::from_value(value.clone())
135 .map_err(|_| "Supervisor permission binding is malformed")?;
136 let latest_result = session
137 .messages
138 .iter()
139 .rposition(|message| message.tool_call_id.as_deref() == Some(tool_call.id.as_str()));
140 let newer_call = session.messages[result_index + 1..].iter().any(|message| {
141 message
142 .tool_calls
143 .as_ref()
144 .is_some_and(|calls| calls.iter().any(|call| call.id == tool_call.id))
145 });
146 let preceding_call = session.messages[..result_index]
147 .iter()
148 .rev()
149 .find(|message| {
150 message
151 .tool_calls
152 .as_ref()
153 .is_some_and(|calls| calls.iter().any(|call| call.id == tool_call.id))
154 });
155 let exact_call = preceding_call.is_some_and(|message| {
156 message.role == bamboo_domain::Role::Assistant
157 && message.tool_calls.as_ref().is_some_and(|calls| {
158 let matching: Vec<_> = calls
159 .iter()
160 .filter(|call| call.id == tool_call.id)
161 .collect();
162 matching.len() == 1 && matching[0] == tool_call
163 })
164 });
165 if record.version != 1
166 || record.incarnation_id.is_nil()
167 || session.id != DEFAULT_SUPERVISOR_SESSION_ID
168 || session.kind != SessionKind::Root
169 || session.root_session_id != session.id
170 || session.parent_session_id.is_some()
171 || session.spawn_depth != 0
172 || session.authority_identity
173 != (SessionAuthorityIdentity::Supervisor {
174 incarnation_id: record.incarnation_id,
175 })
176 || message.role != bamboo_domain::Role::Tool
177 || latest_result != Some(result_index)
178 || newer_call
179 || !exact_call
180 || record.session_id != session.id
181 || record.result_message_id != message.id
182 || record.tool_call != *tool_call
183 || record.execution_name != execution_name
184 || execution_name.trim().is_empty()
185 || record.request_generation.trim().is_empty()
186 || Some(record.request_generation.as_str()) != request_generation
187 {
188 return Err("Supervisor permission binding does not match the current operation");
189 }
190 Ok(Some(Self {
191 incarnation_id: record.incarnation_id,
192 }))
193 }
194
195 pub(super) fn for_caller(self, session_id: Option<&str>) -> Option<Self> {
196 (session_id == Some(DEFAULT_SUPERVISOR_SESSION_ID)).then_some(self)
197 }
198}
199
200#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
210pub struct ToolExecutionSessionFlags {
211 pub bypass_permissions: bool,
214 pub auto_approve_permissions: bool,
217 pub plan_read_only: bool,
220}
221
222impl ToolExecutionSessionFlags {
223 pub fn from_session(session: &Session) -> Self {
226 Self::from_session_and_configured_mode(session, PermissionMode::Default)
227 }
228
229 pub fn from_session_and_configured_mode(
234 session: &Session,
235 configured_mode: PermissionMode,
236 ) -> Self {
237 let requested = session
238 .agent_runtime_state
239 .as_ref()
240 .map(|state| state.effective_permission_mode())
241 .unwrap_or_default();
242 let configured_mode = if session
243 .agent_runtime_state
244 .as_ref()
245 .is_some_and(|state| state.plan_mode.is_some())
246 {
247 PermissionMode::Plan
248 } else {
249 configured_mode
250 };
251 let resolution = bamboo_domain::resolve_permission_mode(requested, configured_mode);
252 Self {
253 bypass_permissions: resolution.bypass_permissions(),
254 auto_approve_permissions: resolution.auto_approve_permissions(),
255 plan_read_only: resolution.effective == PermissionMode::Plan,
256 }
257 }
258}
259
260#[derive(Clone, Copy)]
271pub struct ToolExecutionContext<'a> {
272 pub session_id: Option<&'a str>,
274 pub executing_supervisor: Option<ExecutingSupervisorObservation>,
276 pub root_session_id: Option<&'a str>,
280 pub tool_call_id: &'a str,
282 pub event_tx: Option<&'a mpsc::Sender<AgentEvent>>,
284 pub available_tool_schemas: Option<&'a [ToolSchema]>,
286 pub bypass_permissions: bool,
290 pub auto_approve_permissions: bool,
293 pub plan_read_only: bool,
295 pub can_async_resume: bool,
305 pub bash_completion_sink: Option<&'a Arc<dyn BashCompletionSink>>,
316 pub pre_parsed_args: Option<&'a Value>,
326}
327
328impl std::fmt::Debug for ToolExecutionContext<'_> {
331 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
332 f.debug_struct("ToolExecutionContext")
333 .field("session_id", &self.session_id)
334 .field("executing_supervisor", &self.executing_supervisor)
335 .field("root_session_id", &self.root_session_id)
336 .field("tool_call_id", &self.tool_call_id)
337 .field("event_tx", &self.event_tx)
338 .field("available_tool_schemas", &self.available_tool_schemas)
339 .field("bypass_permissions", &self.bypass_permissions)
340 .field("auto_approve_permissions", &self.auto_approve_permissions)
341 .field("plan_read_only", &self.plan_read_only)
342 .field("can_async_resume", &self.can_async_resume)
343 .field("bash_completion_sink", &self.bash_completion_sink.is_some())
344 .field("pre_parsed_args", &self.pre_parsed_args)
345 .finish()
346 }
347}
348
349impl<'a> ToolExecutionContext<'a> {
350 pub fn none(tool_call_id: &'a str) -> Self {
351 Self {
352 session_id: None,
353 executing_supervisor: None,
354 root_session_id: None,
355 tool_call_id,
356 event_tx: None,
357 available_tool_schemas: None,
358 bypass_permissions: false,
359 auto_approve_permissions: false,
360 plan_read_only: false,
361 can_async_resume: false,
362 bash_completion_sink: None,
363 pre_parsed_args: None,
364 }
365 }
366
367 #[allow(clippy::too_many_arguments)]
375 pub fn for_dispatch(
376 session_id: &'a str,
377 root_session_id: &'a str,
378 tool_call_id: &'a str,
379 event_tx: &'a mpsc::Sender<AgentEvent>,
380 available_tool_schemas: &'a [ToolSchema],
381 flags: ToolExecutionSessionFlags,
382 can_async_resume: bool,
387 bash_completion_sink: Option<&'a Arc<dyn BashCompletionSink>>,
392 pre_parsed_args: Option<&'a Value>,
400 ) -> Self {
401 Self {
402 session_id: Some(session_id),
403 executing_supervisor: None,
404 root_session_id: Some(if root_session_id.trim().is_empty() {
405 session_id
406 } else {
407 root_session_id
408 }),
409 tool_call_id,
410 event_tx: Some(event_tx),
411 available_tool_schemas: Some(available_tool_schemas),
412 bypass_permissions: flags.bypass_permissions,
413 auto_approve_permissions: flags.auto_approve_permissions,
414 plan_read_only: flags.plan_read_only,
415 can_async_resume,
416 bash_completion_sink,
417 pre_parsed_args,
418 }
419 }
420
421 pub fn with_executing_supervisor(
424 mut self,
425 observation: Option<ExecutingSupervisorObservation>,
426 ) -> Self {
427 self.executing_supervisor =
428 observation.and_then(|observation| observation.for_caller(self.session_id));
429 self
430 }
431
432 pub fn cloned_sender(&self) -> Option<mpsc::Sender<AgentEvent>> {
434 self.event_tx.cloned()
435 }
436
437 pub fn cloned_bash_completion_sink(&self) -> Option<Arc<dyn BashCompletionSink>> {
442 self.bash_completion_sink.map(Arc::clone)
443 }
444
445 pub fn root_session_id(&self) -> Option<&'a str> {
447 self.root_session_id
448 }
449
450 pub fn to_tool_ctx(&self) -> crate::tools::ToolCtx {
457 crate::tools::ToolCtx {
458 session_id: self.session_id.map(Arc::from),
459 executing_supervisor: self
460 .executing_supervisor
461 .and_then(|observation| observation.for_caller(self.session_id)),
462 tool_call_id: Arc::from(self.tool_call_id),
463 event_tx: self.event_tx.cloned(),
464 available_tool_schemas: self
465 .available_tool_schemas
466 .map(Arc::from)
467 .unwrap_or_else(|| Arc::from(Vec::new())),
468 bypass_permissions: self.bypass_permissions,
469 auto_approve_permissions: self.auto_approve_permissions,
470 plan_read_only: self.plan_read_only,
471 can_async_resume: self.can_async_resume,
472 async_completion_sink: None,
473 bash_completion_sink: self.bash_completion_sink.map(Arc::clone),
474 }
475 }
476
477 pub async fn emit(&self, event: AgentEvent) {
479 if let Some(tx) = self.event_tx {
480 let event = match event {
484 AgentEvent::Token { content } => AgentEvent::ToolToken {
485 tool_call_id: self.tool_call_id.to_string(),
486 content,
487 },
488 other => other,
489 };
490 let _ = tx.try_send(event);
491 }
492 }
493
494 pub async fn emit_tool_token(&self, content: impl Into<String>) {
496 self.emit(AgentEvent::ToolToken {
497 tool_call_id: self.tool_call_id.to_string(),
498 content: content.into(),
499 })
500 .await;
501 }
502}
503
504#[cfg(test)]
505mod supervisor_observation_tests {
506 use super::*;
507
508 fn supervisor() -> Session {
509 let mut session = Session::new(DEFAULT_SUPERVISOR_SESSION_ID, "model");
510 session.authority_identity = SessionAuthorityIdentity::Supervisor {
511 incarnation_id: Uuid::new_v4(),
512 };
513 session
514 }
515
516 #[test]
517 fn supervisor_observation_requires_typed_canonical_root_lineage() {
518 let valid = supervisor();
519 let observed = ExecutingSupervisorObservation::capture_from_executing_session(&valid)
520 .expect("canonical executing Supervisor");
521 assert_eq!(observed.supervisor_reference().session_id, valid.id);
522 for mutation in [
523 |s: &mut Session| s.authority_identity = SessionAuthorityIdentity::Ordinary,
524 |s: &mut Session| s.id = "other".into(),
525 |s: &mut Session| s.root_session_id = "other".into(),
526 |s: &mut Session| s.parent_session_id = Some("parent".into()),
527 |s: &mut Session| s.spawn_depth = 1,
528 |s: &mut Session| s.kind = SessionKind::Child,
529 |s: &mut Session| {
530 s.authority_identity = SessionAuthorityIdentity::Supervisor {
531 incarnation_id: Uuid::nil(),
532 }
533 },
534 ] {
535 let mut malformed = valid.clone();
536 mutation(&mut malformed);
537 assert_eq!(
538 ExecutingSupervisorObservation::capture_from_executing_session(&malformed),
539 None,
540 "malformed Session: {malformed:?}"
541 );
542 }
543 let mut ordinary = Session::new(DEFAULT_SUPERVISOR_SESSION_ID, "model");
544 ordinary.metadata.insert(
545 "authority_identity".into(),
546 serde_json::to_string(&valid.authority_identity).unwrap(),
547 );
548 ordinary.metadata.insert("supervisor".into(), "true".into());
549 assert_eq!(
550 ExecutingSupervisorObservation::capture_from_executing_session(&ordinary),
551 None,
552 "reserved ID and forged metadata cannot supply a typed identity"
553 );
554 }
555
556 #[test]
557 fn supervisor_observation_is_not_minted_by_flags_arguments_or_generic_contexts() {
558 let session = supervisor();
559 let (tx, _rx) = mpsc::channel(1);
560 let forged_args = serde_json::json!({
561 "executing_supervisor": session.authority_identity,
562 "session_id": DEFAULT_SUPERVISOR_SESSION_ID,
563 });
564 let ctx = ToolExecutionContext::for_dispatch(
565 &session.id,
566 &session.root_session_id,
567 "call",
568 &tx,
569 &[],
570 ToolExecutionSessionFlags {
571 bypass_permissions: true,
572 auto_approve_permissions: true,
573 plan_read_only: true,
574 },
575 false,
576 None,
577 Some(&forged_args),
578 );
579 assert_eq!(ctx.executing_supervisor, None);
580 assert_eq!(ctx.to_tool_ctx().executing_supervisor, None);
581 assert_eq!(
582 ToolExecutionContext::none("none").executing_supervisor,
583 None
584 );
585 assert_eq!(
586 crate::tools::ToolCtx::none("owned").executing_supervisor,
587 None
588 );
589
590 let observation = ExecutingSupervisorObservation::capture_from_executing_session(&session);
591 let captured = ctx.with_executing_supervisor(observation);
592 let owned = captured.to_tool_ctx().clone();
593 assert_eq!(owned.executing_supervisor_for(&session.id), observation);
594 assert_eq!(owned.executing_supervisor_for("child"), None);
595 assert_eq!(
596 ToolExecutionContext::none("missing")
597 .with_executing_supervisor(observation)
598 .executing_supervisor,
599 None
600 );
601 let mismatched = ToolExecutionContext {
603 session_id: Some("child"),
604 ..captured
605 };
606 assert_eq!(mismatched.to_tool_ctx().executing_supervisor, None);
607 }
608}
609
610#[cfg(test)]
611mod session_flags_tests {
612 use super::*;
613 use bamboo_domain::{AgentRuntimeState, SessionPermissionMode};
614
615 #[test]
616 fn dispatch_context_falls_back_from_empty_legacy_root_to_session_id() {
617 let (event_tx, _event_rx) = mpsc::channel(1);
618 let ctx = ToolExecutionContext::for_dispatch(
619 "session-id",
620 " ",
621 "call-id",
622 &event_tx,
623 &[],
624 ToolExecutionSessionFlags::default(),
625 false,
626 None,
627 None,
628 );
629
630 assert_eq!(ctx.root_session_id(), Some("session-id"));
631 }
632
633 #[test]
634 fn from_session_defaults_false_without_runtime_state() {
635 let session = Session::new("s-none", "test-model");
636 assert_eq!(
637 ToolExecutionSessionFlags::from_session(&session),
638 ToolExecutionSessionFlags {
639 bypass_permissions: false,
640 auto_approve_permissions: false,
641 plan_read_only: false,
642 }
643 );
644 }
645
646 #[test]
647 fn from_session_reads_bypass_from_runtime_state() {
648 let mut session = Session::new("s-bypass", "test-model");
649 let mut runtime = AgentRuntimeState::new("run-1");
650 runtime.bypass_permissions = true;
651 session.agent_runtime_state = Some(runtime);
652 assert!(ToolExecutionSessionFlags::from_session(&session).bypass_permissions);
653 assert!(!ToolExecutionSessionFlags::from_session(&session).auto_approve_permissions);
654 }
655
656 #[test]
657 fn from_session_distinguishes_auto_from_legacy_bypass() {
658 let mut session = Session::new("s-auto", "test-model");
659 let mut runtime = AgentRuntimeState::new("run-1");
660 runtime.set_permission_mode(SessionPermissionMode::Auto);
661 session.agent_runtime_state = Some(runtime);
662
663 let flags = ToolExecutionSessionFlags::from_session(&session);
664 assert!(!flags.bypass_permissions);
665 assert!(flags.auto_approve_permissions);
666 }
667
668 #[test]
669 fn configured_auto_applies_to_default_but_not_explicit_bypass() {
670 let default_session = Session::new("s-global-auto", "test-model");
671 let flags = ToolExecutionSessionFlags::from_session_and_configured_mode(
672 &default_session,
673 PermissionMode::Auto,
674 );
675 assert_eq!(
676 flags,
677 ToolExecutionSessionFlags {
678 bypass_permissions: false,
679 auto_approve_permissions: true,
680 plan_read_only: false,
681 }
682 );
683
684 let mut bypass_session = default_session;
685 bypass_session
686 .agent_runtime_state
687 .get_or_insert_default()
688 .set_permission_mode(SessionPermissionMode::Bypass);
689 let flags = ToolExecutionSessionFlags::from_session_and_configured_mode(
690 &bypass_session,
691 PermissionMode::Auto,
692 );
693 assert!(flags.bypass_permissions);
694 assert!(!flags.auto_approve_permissions);
695 }
696
697 #[test]
698 fn configured_plan_preserves_no_prompt_but_sets_read_only_gate() {
699 let mut session = Session::new("s-plan-auto", "test-model");
700 session
701 .agent_runtime_state
702 .get_or_insert_default()
703 .set_permission_mode(SessionPermissionMode::Auto);
704 let flags = ToolExecutionSessionFlags::from_session_and_configured_mode(
705 &session,
706 PermissionMode::Plan,
707 );
708 assert!(!flags.bypass_permissions);
709 assert!(flags.auto_approve_permissions);
710 assert!(flags.plan_read_only);
711 }
712
713 #[test]
714 fn persisted_plan_overlay_is_honored_without_config_reconstruction() {
715 let mut session = Session::new("s-persisted-plan-auto", "test-model");
716 let runtime = session.agent_runtime_state.get_or_insert_default();
717 runtime.set_permission_mode(SessionPermissionMode::Auto);
718 runtime.plan_mode = Some(bamboo_domain::PlanModeState {
719 entered_at: chrono::Utc::now(),
720 pre_permission_mode: "auto".to_string(),
721 plan_file_path: None,
722 status: bamboo_domain::PlanModeStatus::Designing,
723 });
724
725 let flags = ToolExecutionSessionFlags::from_session(&session);
726 assert!(flags.plan_read_only);
727 assert!(flags.auto_approve_permissions);
728 assert!(!flags.bypass_permissions);
729 }
730
731 #[test]
732 fn typed_read_only_child_does_not_reactivate_the_legacy_plan_name_gate() {
733 for requested in [SessionPermissionMode::Auto, SessionPermissionMode::Bypass] {
734 let mut session = Session::new("s-read-only-child", "test-model");
735 let runtime = session.agent_runtime_state.get_or_insert_default();
736 runtime.set_permission_mode(requested);
737 runtime.read_only = true;
738
739 let flags = ToolExecutionSessionFlags::from_session_and_configured_mode(
744 &session,
745 PermissionMode::Auto,
746 );
747
748 assert!(!flags.plan_read_only);
749 assert_eq!(
750 flags.bypass_permissions,
751 requested == SessionPermissionMode::Bypass
752 );
753 assert_eq!(
754 flags.auto_approve_permissions,
755 requested == SessionPermissionMode::Auto
756 );
757 }
758 }
759
760 #[test]
761 fn for_dispatch_maps_flags_onto_context() {
762 let (tx, _rx) = mpsc::channel(1);
763 let ctx = ToolExecutionContext::for_dispatch(
764 "s1",
765 "root-s1",
766 "call-1",
767 &tx,
768 &[],
769 ToolExecutionSessionFlags {
770 bypass_permissions: true,
771 auto_approve_permissions: false,
772 plan_read_only: false,
773 },
774 true,
775 None,
776 None,
777 );
778 assert_eq!(ctx.session_id, Some("s1"));
779 assert_eq!(ctx.root_session_id(), Some("root-s1"));
780 assert!(ctx.bypass_permissions);
781 assert!(!ctx.auto_approve_permissions);
782 assert!(ctx.can_async_resume);
783 assert!(ctx.pre_parsed_args.is_none());
784
785 let owned = ctx.to_tool_ctx();
786 assert!(owned.bypass_permissions);
787 assert!(!owned.auto_approve_permissions);
788 assert!(!owned.plan_read_only);
789 }
790
791 #[test]
792 fn for_dispatch_threads_pre_parsed_args() {
793 let (tx, _rx) = mpsc::channel(1);
794 let parsed = serde_json::json!({"v": "x"});
795 let ctx = ToolExecutionContext::for_dispatch(
796 "s1",
797 "root-s1",
798 "call-1",
799 &tx,
800 &[],
801 ToolExecutionSessionFlags::default(),
802 false,
803 None,
804 Some(&parsed),
805 );
806 assert_eq!(ctx.pre_parsed_args, Some(&parsed));
807 }
808}
809
810#[cfg(test)]
811mod tests {
812 use super::*;
813
814 #[tokio::test]
815 async fn emit_does_not_block_when_channel_is_full() {
816 let (tx, mut rx) = mpsc::channel(1);
817 tx.send(AgentEvent::Token {
818 content: "full".to_string(),
819 })
820 .await
821 .unwrap();
822 let ctx = ToolExecutionContext {
823 executing_supervisor: None,
824 session_id: Some("session_1"),
825 root_session_id: None,
826 tool_call_id: "call_1",
827 event_tx: Some(&tx),
828 available_tool_schemas: None,
829 bypass_permissions: false,
830 auto_approve_permissions: false,
831 plan_read_only: false,
832 can_async_resume: false,
833 bash_completion_sink: None,
834 pre_parsed_args: None,
835 };
836
837 tokio::time::timeout(
838 std::time::Duration::from_millis(100),
839 ctx.emit(AgentEvent::Token {
840 content: "next".to_string(),
841 }),
842 )
843 .await
844 .expect("emit should not block on full channel");
845
846 let first = rx.recv().await.unwrap();
847 match first {
848 AgentEvent::Token { content } => assert_eq!(content, "full"),
849 other => panic!("unexpected event: {other:?}"),
850 }
851 }
852
853 #[tokio::test]
854 async fn emit_converts_token_to_tool_token() {
855 let (tx, mut rx) = mpsc::channel(10);
856 let ctx = ToolExecutionContext {
857 executing_supervisor: None,
858 session_id: Some("session_1"),
859 root_session_id: None,
860 tool_call_id: "call_123",
861 event_tx: Some(&tx),
862 available_tool_schemas: None,
863 bypass_permissions: false,
864 auto_approve_permissions: false,
865 plan_read_only: false,
866 can_async_resume: false,
867 bash_completion_sink: None,
868 pre_parsed_args: None,
869 };
870
871 ctx.emit(AgentEvent::Token {
872 content: "test content".to_string(),
873 })
874 .await;
875
876 let event = rx.recv().await.unwrap();
877 match event {
878 AgentEvent::ToolToken {
879 tool_call_id,
880 content,
881 } => {
882 assert_eq!(tool_call_id, "call_123");
883 assert_eq!(content, "test content");
884 }
885 other => panic!("Expected ToolToken, got: {other:?}"),
886 }
887 }
888
889 #[tokio::test]
890 async fn emit_passes_through_non_token_events() {
891 let (tx, mut rx) = mpsc::channel(10);
892 let ctx = ToolExecutionContext {
893 executing_supervisor: None,
894 session_id: Some("session_1"),
895 root_session_id: None,
896 tool_call_id: "call_456",
897 event_tx: Some(&tx),
898 available_tool_schemas: None,
899 bypass_permissions: false,
900 auto_approve_permissions: false,
901 plan_read_only: false,
902 can_async_resume: false,
903 bash_completion_sink: None,
904 pre_parsed_args: None,
905 };
906
907 ctx.emit(AgentEvent::ToolToken {
909 tool_call_id: "other".to_string(),
910 content: "direct tool token".to_string(),
911 })
912 .await;
913
914 let event = rx.recv().await.unwrap();
915 match event {
916 AgentEvent::ToolToken { content, .. } => {
917 assert_eq!(content, "direct tool token");
918 }
919 other => panic!("Expected ToolToken, got: {other:?}"),
920 }
921 }
922
923 #[tokio::test]
924 async fn emit_does_nothing_when_no_sender() {
925 let ctx = ToolExecutionContext::none("call_789");
926
927 ctx.emit(AgentEvent::Token {
929 content: "test".to_string(),
930 })
931 .await;
932
933 }
935
936 #[tokio::test]
937 async fn emit_tool_token_convenience_method() {
938 let (tx, mut rx) = mpsc::channel(10);
939 let ctx = ToolExecutionContext {
940 executing_supervisor: None,
941 session_id: None,
942 root_session_id: None,
943 tool_call_id: "call_abc",
944 event_tx: Some(&tx),
945 available_tool_schemas: None,
946 bypass_permissions: false,
947 auto_approve_permissions: false,
948 plan_read_only: false,
949 can_async_resume: false,
950 bash_completion_sink: None,
951 pre_parsed_args: None,
952 };
953
954 ctx.emit_tool_token("convenient output").await;
955
956 let event = rx.recv().await.unwrap();
957 match event {
958 AgentEvent::ToolToken {
959 tool_call_id,
960 content,
961 } => {
962 assert_eq!(tool_call_id, "call_abc");
963 assert_eq!(content, "convenient output");
964 }
965 other => panic!("Expected ToolToken, got: {other:?}"),
966 }
967 }
968
969 #[tokio::test]
970 async fn emit_tool_token_with_no_sender_does_nothing() {
971 let ctx = ToolExecutionContext::none("call_def");
972
973 ctx.emit_tool_token("test").await;
975
976 }
978
979 #[test]
980 fn none_creates_context_with_no_optional_fields() {
981 let ctx = ToolExecutionContext::none("call_xyz");
982
983 assert_eq!(ctx.session_id, None);
984 assert_eq!(ctx.tool_call_id, "call_xyz");
985 assert!(ctx.event_tx.is_none());
986 }
987
988 #[test]
989 fn cloned_sender_returns_none_when_no_sender() {
990 let ctx = ToolExecutionContext::none("call_test");
991 assert!(ctx.cloned_sender().is_none());
992 }
993
994 #[tokio::test]
995 async fn cloned_sender_returns_clone_when_sender_present() {
996 let (tx, _rx) = mpsc::channel(10);
997 let ctx = ToolExecutionContext {
998 executing_supervisor: None,
999 session_id: None,
1000 root_session_id: None,
1001 tool_call_id: "call_clone",
1002 event_tx: Some(&tx),
1003 available_tool_schemas: None,
1004 bypass_permissions: false,
1005 auto_approve_permissions: false,
1006 plan_read_only: false,
1007 can_async_resume: false,
1008 bash_completion_sink: None,
1009 pre_parsed_args: None,
1010 };
1011
1012 let cloned = ctx.cloned_sender();
1013 assert!(cloned.is_some());
1014
1015 cloned
1017 .unwrap()
1018 .send(AgentEvent::Token {
1019 content: "test".to_string(),
1020 })
1021 .await
1022 .unwrap();
1023 }
1024
1025 #[tokio::test]
1026 async fn emit_handles_multiple_sequential_calls() {
1027 let (tx, mut rx) = mpsc::channel(10);
1028 let ctx = ToolExecutionContext {
1029 executing_supervisor: None,
1030 session_id: Some("session_multi"),
1031 root_session_id: None,
1032 tool_call_id: "call_multi",
1033 event_tx: Some(&tx),
1034 available_tool_schemas: None,
1035 bypass_permissions: false,
1036 auto_approve_permissions: false,
1037 plan_read_only: false,
1038 can_async_resume: false,
1039 bash_completion_sink: None,
1040 pre_parsed_args: None,
1041 };
1042
1043 for i in 0..5 {
1044 ctx.emit(AgentEvent::Token {
1045 content: format!("message {}", i),
1046 })
1047 .await;
1048 }
1049
1050 for i in 0..5 {
1051 let event = rx.recv().await.unwrap();
1052 match event {
1053 AgentEvent::ToolToken { content, .. } => {
1054 assert_eq!(content, format!("message {}", i));
1055 }
1056 other => panic!("Expected ToolToken, got: {other:?}"),
1057 }
1058 }
1059 }
1060
1061 #[test]
1062 fn context_is_clone_and_copy() {
1063 let (tx, _rx) = mpsc::channel(10);
1064 let ctx = ToolExecutionContext {
1065 executing_supervisor: None,
1066 session_id: Some("session_copy"),
1067 root_session_id: None,
1068 tool_call_id: "call_copy",
1069 event_tx: Some(&tx),
1070 available_tool_schemas: None,
1071 bypass_permissions: false,
1072 auto_approve_permissions: false,
1073 plan_read_only: false,
1074 can_async_resume: false,
1075 bash_completion_sink: None,
1076 pre_parsed_args: None,
1077 };
1078
1079 let _cloned = ctx;
1081
1082 let copied = ctx;
1084
1085 assert_eq!(copied.tool_call_id, "call_copy");
1087 }
1088
1089 #[test]
1090 fn context_is_debug() {
1091 let ctx = ToolExecutionContext::none("call_debug");
1092 let debug_str = format!("{:?}", ctx);
1093 assert!(debug_str.contains("call_debug"));
1094 }
1095
1096 #[tokio::test]
1097 async fn emit_with_empty_tool_call_id() {
1098 let (tx, mut rx) = mpsc::channel(10);
1099 let ctx = ToolExecutionContext {
1100 executing_supervisor: None,
1101 session_id: None,
1102 root_session_id: None,
1103 tool_call_id: "",
1104 event_tx: Some(&tx),
1105 available_tool_schemas: None,
1106 bypass_permissions: false,
1107 auto_approve_permissions: false,
1108 plan_read_only: false,
1109 can_async_resume: false,
1110 bash_completion_sink: None,
1111 pre_parsed_args: None,
1112 };
1113
1114 ctx.emit(AgentEvent::Token {
1115 content: "test".to_string(),
1116 })
1117 .await;
1118
1119 let event = rx.recv().await.unwrap();
1120 match event {
1121 AgentEvent::ToolToken { tool_call_id, .. } => {
1122 assert_eq!(tool_call_id, "");
1123 }
1124 other => panic!("Expected ToolToken, got: {other:?}"),
1125 }
1126 }
1127
1128 #[tokio::test]
1129 async fn emit_with_unicode_content() {
1130 let (tx, mut rx) = mpsc::channel(10);
1131 let ctx = ToolExecutionContext {
1132 executing_supervisor: None,
1133 session_id: Some("会话"),
1134 root_session_id: None,
1135 tool_call_id: "调用_123",
1136 event_tx: Some(&tx),
1137 available_tool_schemas: None,
1138 bypass_permissions: false,
1139 auto_approve_permissions: false,
1140 plan_read_only: false,
1141 can_async_resume: false,
1142 bash_completion_sink: None,
1143 pre_parsed_args: None,
1144 };
1145
1146 ctx.emit(AgentEvent::Token {
1147 content: "测试内容 🎯".to_string(),
1148 })
1149 .await;
1150
1151 let event = rx.recv().await.unwrap();
1152 match event {
1153 AgentEvent::ToolToken {
1154 tool_call_id,
1155 content,
1156 } => {
1157 assert_eq!(tool_call_id, "调用_123");
1158 assert_eq!(content, "测试内容 🎯");
1159 }
1160 other => panic!("Expected ToolToken, got: {other:?}"),
1161 }
1162 }
1163
1164 #[tokio::test]
1165 async fn emit_with_special_characters_in_tool_call_id() {
1166 let (tx, mut rx) = mpsc::channel(10);
1167 let ctx = ToolExecutionContext {
1168 executing_supervisor: None,
1169 session_id: None,
1170 root_session_id: None,
1171 tool_call_id: "call-with_special.chars:123",
1172 event_tx: Some(&tx),
1173 available_tool_schemas: None,
1174 bypass_permissions: false,
1175 auto_approve_permissions: false,
1176 plan_read_only: false,
1177 can_async_resume: false,
1178 bash_completion_sink: None,
1179 pre_parsed_args: None,
1180 };
1181
1182 ctx.emit(AgentEvent::Token {
1183 content: "test".to_string(),
1184 })
1185 .await;
1186
1187 let event = rx.recv().await.unwrap();
1188 match event {
1189 AgentEvent::ToolToken { tool_call_id, .. } => {
1190 assert_eq!(tool_call_id, "call-with_special.chars:123");
1191 }
1192 other => panic!("Expected ToolToken, got: {other:?}"),
1193 }
1194 }
1195
1196 #[tokio::test]
1197 async fn emit_tool_token_with_string_content() {
1198 let (tx, mut rx) = mpsc::channel(10);
1199 let ctx = ToolExecutionContext {
1200 executing_supervisor: None,
1201 session_id: None,
1202 root_session_id: None,
1203 tool_call_id: "call_string",
1204 event_tx: Some(&tx),
1205 available_tool_schemas: None,
1206 bypass_permissions: false,
1207 auto_approve_permissions: false,
1208 plan_read_only: false,
1209 can_async_resume: false,
1210 bash_completion_sink: None,
1211 pre_parsed_args: None,
1212 };
1213
1214 let content = String::from("owned string");
1215 ctx.emit_tool_token(content).await;
1216
1217 let event = rx.recv().await.unwrap();
1218 match event {
1219 AgentEvent::ToolToken { content, .. } => {
1220 assert_eq!(content, "owned string");
1221 }
1222 other => panic!("Expected ToolToken, got: {other:?}"),
1223 }
1224 }
1225
1226 #[tokio::test]
1227 async fn emit_tool_token_with_str_content() {
1228 let (tx, mut rx) = mpsc::channel(10);
1229 let ctx = ToolExecutionContext {
1230 executing_supervisor: None,
1231 session_id: None,
1232 root_session_id: None,
1233 tool_call_id: "call_str",
1234 event_tx: Some(&tx),
1235 available_tool_schemas: None,
1236 bypass_permissions: false,
1237 auto_approve_permissions: false,
1238 plan_read_only: false,
1239 can_async_resume: false,
1240 bash_completion_sink: None,
1241 pre_parsed_args: None,
1242 };
1243
1244 ctx.emit_tool_token("string slice").await;
1245
1246 let event = rx.recv().await.unwrap();
1247 match event {
1248 AgentEvent::ToolToken { content, .. } => {
1249 assert_eq!(content, "string slice");
1250 }
1251 other => panic!("Expected ToolToken, got: {other:?}"),
1252 }
1253 }
1254}