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 for_dispatch_maps_flags_onto_context() {
733 let (tx, _rx) = mpsc::channel(1);
734 let ctx = ToolExecutionContext::for_dispatch(
735 "s1",
736 "root-s1",
737 "call-1",
738 &tx,
739 &[],
740 ToolExecutionSessionFlags {
741 bypass_permissions: true,
742 auto_approve_permissions: false,
743 plan_read_only: false,
744 },
745 true,
746 None,
747 None,
748 );
749 assert_eq!(ctx.session_id, Some("s1"));
750 assert_eq!(ctx.root_session_id(), Some("root-s1"));
751 assert!(ctx.bypass_permissions);
752 assert!(!ctx.auto_approve_permissions);
753 assert!(ctx.can_async_resume);
754 assert!(ctx.pre_parsed_args.is_none());
755
756 let owned = ctx.to_tool_ctx();
757 assert!(owned.bypass_permissions);
758 assert!(!owned.auto_approve_permissions);
759 assert!(!owned.plan_read_only);
760 }
761
762 #[test]
763 fn for_dispatch_threads_pre_parsed_args() {
764 let (tx, _rx) = mpsc::channel(1);
765 let parsed = serde_json::json!({"v": "x"});
766 let ctx = ToolExecutionContext::for_dispatch(
767 "s1",
768 "root-s1",
769 "call-1",
770 &tx,
771 &[],
772 ToolExecutionSessionFlags::default(),
773 false,
774 None,
775 Some(&parsed),
776 );
777 assert_eq!(ctx.pre_parsed_args, Some(&parsed));
778 }
779}
780
781#[cfg(test)]
782mod tests {
783 use super::*;
784
785 #[tokio::test]
786 async fn emit_does_not_block_when_channel_is_full() {
787 let (tx, mut rx) = mpsc::channel(1);
788 tx.send(AgentEvent::Token {
789 content: "full".to_string(),
790 })
791 .await
792 .unwrap();
793 let ctx = ToolExecutionContext {
794 executing_supervisor: None,
795 session_id: Some("session_1"),
796 root_session_id: None,
797 tool_call_id: "call_1",
798 event_tx: Some(&tx),
799 available_tool_schemas: None,
800 bypass_permissions: false,
801 auto_approve_permissions: false,
802 plan_read_only: false,
803 can_async_resume: false,
804 bash_completion_sink: None,
805 pre_parsed_args: None,
806 };
807
808 tokio::time::timeout(
809 std::time::Duration::from_millis(100),
810 ctx.emit(AgentEvent::Token {
811 content: "next".to_string(),
812 }),
813 )
814 .await
815 .expect("emit should not block on full channel");
816
817 let first = rx.recv().await.unwrap();
818 match first {
819 AgentEvent::Token { content } => assert_eq!(content, "full"),
820 other => panic!("unexpected event: {other:?}"),
821 }
822 }
823
824 #[tokio::test]
825 async fn emit_converts_token_to_tool_token() {
826 let (tx, mut rx) = mpsc::channel(10);
827 let ctx = ToolExecutionContext {
828 executing_supervisor: None,
829 session_id: Some("session_1"),
830 root_session_id: None,
831 tool_call_id: "call_123",
832 event_tx: Some(&tx),
833 available_tool_schemas: None,
834 bypass_permissions: false,
835 auto_approve_permissions: false,
836 plan_read_only: false,
837 can_async_resume: false,
838 bash_completion_sink: None,
839 pre_parsed_args: None,
840 };
841
842 ctx.emit(AgentEvent::Token {
843 content: "test content".to_string(),
844 })
845 .await;
846
847 let event = rx.recv().await.unwrap();
848 match event {
849 AgentEvent::ToolToken {
850 tool_call_id,
851 content,
852 } => {
853 assert_eq!(tool_call_id, "call_123");
854 assert_eq!(content, "test content");
855 }
856 other => panic!("Expected ToolToken, got: {other:?}"),
857 }
858 }
859
860 #[tokio::test]
861 async fn emit_passes_through_non_token_events() {
862 let (tx, mut rx) = mpsc::channel(10);
863 let ctx = ToolExecutionContext {
864 executing_supervisor: None,
865 session_id: Some("session_1"),
866 root_session_id: None,
867 tool_call_id: "call_456",
868 event_tx: Some(&tx),
869 available_tool_schemas: None,
870 bypass_permissions: false,
871 auto_approve_permissions: false,
872 plan_read_only: false,
873 can_async_resume: false,
874 bash_completion_sink: None,
875 pre_parsed_args: None,
876 };
877
878 ctx.emit(AgentEvent::ToolToken {
880 tool_call_id: "other".to_string(),
881 content: "direct tool token".to_string(),
882 })
883 .await;
884
885 let event = rx.recv().await.unwrap();
886 match event {
887 AgentEvent::ToolToken { content, .. } => {
888 assert_eq!(content, "direct tool token");
889 }
890 other => panic!("Expected ToolToken, got: {other:?}"),
891 }
892 }
893
894 #[tokio::test]
895 async fn emit_does_nothing_when_no_sender() {
896 let ctx = ToolExecutionContext::none("call_789");
897
898 ctx.emit(AgentEvent::Token {
900 content: "test".to_string(),
901 })
902 .await;
903
904 }
906
907 #[tokio::test]
908 async fn emit_tool_token_convenience_method() {
909 let (tx, mut rx) = mpsc::channel(10);
910 let ctx = ToolExecutionContext {
911 executing_supervisor: None,
912 session_id: None,
913 root_session_id: None,
914 tool_call_id: "call_abc",
915 event_tx: Some(&tx),
916 available_tool_schemas: None,
917 bypass_permissions: false,
918 auto_approve_permissions: false,
919 plan_read_only: false,
920 can_async_resume: false,
921 bash_completion_sink: None,
922 pre_parsed_args: None,
923 };
924
925 ctx.emit_tool_token("convenient output").await;
926
927 let event = rx.recv().await.unwrap();
928 match event {
929 AgentEvent::ToolToken {
930 tool_call_id,
931 content,
932 } => {
933 assert_eq!(tool_call_id, "call_abc");
934 assert_eq!(content, "convenient output");
935 }
936 other => panic!("Expected ToolToken, got: {other:?}"),
937 }
938 }
939
940 #[tokio::test]
941 async fn emit_tool_token_with_no_sender_does_nothing() {
942 let ctx = ToolExecutionContext::none("call_def");
943
944 ctx.emit_tool_token("test").await;
946
947 }
949
950 #[test]
951 fn none_creates_context_with_no_optional_fields() {
952 let ctx = ToolExecutionContext::none("call_xyz");
953
954 assert_eq!(ctx.session_id, None);
955 assert_eq!(ctx.tool_call_id, "call_xyz");
956 assert!(ctx.event_tx.is_none());
957 }
958
959 #[test]
960 fn cloned_sender_returns_none_when_no_sender() {
961 let ctx = ToolExecutionContext::none("call_test");
962 assert!(ctx.cloned_sender().is_none());
963 }
964
965 #[tokio::test]
966 async fn cloned_sender_returns_clone_when_sender_present() {
967 let (tx, _rx) = mpsc::channel(10);
968 let ctx = ToolExecutionContext {
969 executing_supervisor: None,
970 session_id: None,
971 root_session_id: None,
972 tool_call_id: "call_clone",
973 event_tx: Some(&tx),
974 available_tool_schemas: None,
975 bypass_permissions: false,
976 auto_approve_permissions: false,
977 plan_read_only: false,
978 can_async_resume: false,
979 bash_completion_sink: None,
980 pre_parsed_args: None,
981 };
982
983 let cloned = ctx.cloned_sender();
984 assert!(cloned.is_some());
985
986 cloned
988 .unwrap()
989 .send(AgentEvent::Token {
990 content: "test".to_string(),
991 })
992 .await
993 .unwrap();
994 }
995
996 #[tokio::test]
997 async fn emit_handles_multiple_sequential_calls() {
998 let (tx, mut rx) = mpsc::channel(10);
999 let ctx = ToolExecutionContext {
1000 executing_supervisor: None,
1001 session_id: Some("session_multi"),
1002 root_session_id: None,
1003 tool_call_id: "call_multi",
1004 event_tx: Some(&tx),
1005 available_tool_schemas: None,
1006 bypass_permissions: false,
1007 auto_approve_permissions: false,
1008 plan_read_only: false,
1009 can_async_resume: false,
1010 bash_completion_sink: None,
1011 pre_parsed_args: None,
1012 };
1013
1014 for i in 0..5 {
1015 ctx.emit(AgentEvent::Token {
1016 content: format!("message {}", i),
1017 })
1018 .await;
1019 }
1020
1021 for i in 0..5 {
1022 let event = rx.recv().await.unwrap();
1023 match event {
1024 AgentEvent::ToolToken { content, .. } => {
1025 assert_eq!(content, format!("message {}", i));
1026 }
1027 other => panic!("Expected ToolToken, got: {other:?}"),
1028 }
1029 }
1030 }
1031
1032 #[test]
1033 fn context_is_clone_and_copy() {
1034 let (tx, _rx) = mpsc::channel(10);
1035 let ctx = ToolExecutionContext {
1036 executing_supervisor: None,
1037 session_id: Some("session_copy"),
1038 root_session_id: None,
1039 tool_call_id: "call_copy",
1040 event_tx: Some(&tx),
1041 available_tool_schemas: None,
1042 bypass_permissions: false,
1043 auto_approve_permissions: false,
1044 plan_read_only: false,
1045 can_async_resume: false,
1046 bash_completion_sink: None,
1047 pre_parsed_args: None,
1048 };
1049
1050 let _cloned = ctx;
1052
1053 let copied = ctx;
1055
1056 assert_eq!(copied.tool_call_id, "call_copy");
1058 }
1059
1060 #[test]
1061 fn context_is_debug() {
1062 let ctx = ToolExecutionContext::none("call_debug");
1063 let debug_str = format!("{:?}", ctx);
1064 assert!(debug_str.contains("call_debug"));
1065 }
1066
1067 #[tokio::test]
1068 async fn emit_with_empty_tool_call_id() {
1069 let (tx, mut rx) = mpsc::channel(10);
1070 let ctx = ToolExecutionContext {
1071 executing_supervisor: None,
1072 session_id: None,
1073 root_session_id: None,
1074 tool_call_id: "",
1075 event_tx: Some(&tx),
1076 available_tool_schemas: None,
1077 bypass_permissions: false,
1078 auto_approve_permissions: false,
1079 plan_read_only: false,
1080 can_async_resume: false,
1081 bash_completion_sink: None,
1082 pre_parsed_args: None,
1083 };
1084
1085 ctx.emit(AgentEvent::Token {
1086 content: "test".to_string(),
1087 })
1088 .await;
1089
1090 let event = rx.recv().await.unwrap();
1091 match event {
1092 AgentEvent::ToolToken { tool_call_id, .. } => {
1093 assert_eq!(tool_call_id, "");
1094 }
1095 other => panic!("Expected ToolToken, got: {other:?}"),
1096 }
1097 }
1098
1099 #[tokio::test]
1100 async fn emit_with_unicode_content() {
1101 let (tx, mut rx) = mpsc::channel(10);
1102 let ctx = ToolExecutionContext {
1103 executing_supervisor: None,
1104 session_id: Some("会话"),
1105 root_session_id: None,
1106 tool_call_id: "调用_123",
1107 event_tx: Some(&tx),
1108 available_tool_schemas: None,
1109 bypass_permissions: false,
1110 auto_approve_permissions: false,
1111 plan_read_only: false,
1112 can_async_resume: false,
1113 bash_completion_sink: None,
1114 pre_parsed_args: None,
1115 };
1116
1117 ctx.emit(AgentEvent::Token {
1118 content: "测试内容 🎯".to_string(),
1119 })
1120 .await;
1121
1122 let event = rx.recv().await.unwrap();
1123 match event {
1124 AgentEvent::ToolToken {
1125 tool_call_id,
1126 content,
1127 } => {
1128 assert_eq!(tool_call_id, "调用_123");
1129 assert_eq!(content, "测试内容 🎯");
1130 }
1131 other => panic!("Expected ToolToken, got: {other:?}"),
1132 }
1133 }
1134
1135 #[tokio::test]
1136 async fn emit_with_special_characters_in_tool_call_id() {
1137 let (tx, mut rx) = mpsc::channel(10);
1138 let ctx = ToolExecutionContext {
1139 executing_supervisor: None,
1140 session_id: None,
1141 root_session_id: None,
1142 tool_call_id: "call-with_special.chars:123",
1143 event_tx: Some(&tx),
1144 available_tool_schemas: None,
1145 bypass_permissions: false,
1146 auto_approve_permissions: false,
1147 plan_read_only: false,
1148 can_async_resume: false,
1149 bash_completion_sink: None,
1150 pre_parsed_args: None,
1151 };
1152
1153 ctx.emit(AgentEvent::Token {
1154 content: "test".to_string(),
1155 })
1156 .await;
1157
1158 let event = rx.recv().await.unwrap();
1159 match event {
1160 AgentEvent::ToolToken { tool_call_id, .. } => {
1161 assert_eq!(tool_call_id, "call-with_special.chars:123");
1162 }
1163 other => panic!("Expected ToolToken, got: {other:?}"),
1164 }
1165 }
1166
1167 #[tokio::test]
1168 async fn emit_tool_token_with_string_content() {
1169 let (tx, mut rx) = mpsc::channel(10);
1170 let ctx = ToolExecutionContext {
1171 executing_supervisor: None,
1172 session_id: None,
1173 root_session_id: None,
1174 tool_call_id: "call_string",
1175 event_tx: Some(&tx),
1176 available_tool_schemas: None,
1177 bypass_permissions: false,
1178 auto_approve_permissions: false,
1179 plan_read_only: false,
1180 can_async_resume: false,
1181 bash_completion_sink: None,
1182 pre_parsed_args: None,
1183 };
1184
1185 let content = String::from("owned string");
1186 ctx.emit_tool_token(content).await;
1187
1188 let event = rx.recv().await.unwrap();
1189 match event {
1190 AgentEvent::ToolToken { content, .. } => {
1191 assert_eq!(content, "owned string");
1192 }
1193 other => panic!("Expected ToolToken, got: {other:?}"),
1194 }
1195 }
1196
1197 #[tokio::test]
1198 async fn emit_tool_token_with_str_content() {
1199 let (tx, mut rx) = mpsc::channel(10);
1200 let ctx = ToolExecutionContext {
1201 executing_supervisor: None,
1202 session_id: None,
1203 root_session_id: None,
1204 tool_call_id: "call_str",
1205 event_tx: Some(&tx),
1206 available_tool_schemas: None,
1207 bypass_permissions: false,
1208 auto_approve_permissions: false,
1209 plan_read_only: false,
1210 can_async_resume: false,
1211 bash_completion_sink: None,
1212 pre_parsed_args: None,
1213 };
1214
1215 ctx.emit_tool_token("string slice").await;
1216
1217 let event = rx.recv().await.unwrap();
1218 match event {
1219 AgentEvent::ToolToken { content, .. } => {
1220 assert_eq!(content, "string slice");
1221 }
1222 other => panic!("Expected ToolToken, got: {other:?}"),
1223 }
1224 }
1225}