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