1mod builder;
40mod error;
41mod execute_request;
42mod tools;
43
44use std::sync::Arc;
45
46use async_trait::async_trait;
47pub use builder::AgentBuilder;
48pub use execute_request::ExecuteRequestBuilder;
49use tokio::sync::mpsc;
50
51use bamboo_engine::session_app::errors::{SessionLoadError, SessionSaveError};
52use bamboo_engine::session_app::repository::SessionAccess;
53use bamboo_engine::session_app::respond::{
54 submit_pending_response, PERMISSION_REEXECUTE_METADATA_KEY,
55};
56use bamboo_engine::session_app::types::RespondInput;
57
58pub use tokio_util::sync::CancellationToken;
61pub use tools::{
62 builtin_tool_names, builtin_tool_specs, BuiltinTool, ToolSpec, CANONICAL_TOOL_NAMES,
63};
64
65pub use error::SdkError;
66
67pub use bamboo_agent_core::{
70 AgentError, AgentEvent, Message, MessageContent, PendingQuestion, Role, Session,
71 TokenBudgetUsage, TokenUsage,
72};
73pub use bamboo_domain::{TaskItem, TaskItemStatus, TaskList};
74pub use bamboo_engine::session_app::respond::PlanModeTransition;
75pub use bamboo_engine::ExecuteRequest;
76pub use bamboo_llm::LLMProvider;
77pub use bamboo_mcp::manager::McpServerManager;
78pub use bamboo_mcp::McpServerConfig;
79pub use bamboo_storage::SessionIndexEntry;
80pub use bamboo_tools::permission::{PermissionChecker, PermissionType};
81pub use bamboo_tools::{BuiltinToolExecutor, BuiltinToolExecutorBuilder, ToolOutputManager};
82
83const EVENT_CHANNEL_CAPACITY: usize = 256;
85
86#[derive(Clone)]
91pub struct Agent {
92 inner: bamboo_engine::Agent,
93 system_prompt: Option<String>,
96 model: Option<String>,
98 session_store: Option<Arc<bamboo_storage::SessionStoreV2>>,
103 permission_checker: Option<Arc<dyn bamboo_tools::permission::PermissionChecker>>,
109}
110
111impl Agent {
112 pub fn builder() -> AgentBuilder {
114 AgentBuilder::new()
115 }
116
117 pub fn from_runtime(inner: bamboo_engine::Agent) -> Self {
120 Self {
121 inner,
122 system_prompt: None,
123 model: None,
124 session_store: None,
125 permission_checker: None,
126 }
127 }
128
129 pub(crate) fn from_runtime_with_config(
132 inner: bamboo_engine::Agent,
133 system_prompt: Option<String>,
134 model: Option<String>,
135 session_store: Option<Arc<bamboo_storage::SessionStoreV2>>,
136 permission_checker: Option<Arc<dyn bamboo_tools::permission::PermissionChecker>>,
137 ) -> Self {
138 Self {
139 inner,
140 system_prompt,
141 model,
142 session_store,
143 permission_checker,
144 }
145 }
146
147 pub async fn run(
158 &self,
159 session: &mut Session,
160 input: impl Into<String>,
161 ) -> Result<(), AgentError> {
162 session.add_message(Message::user(input.into()));
163 self.run_session(session).await
164 }
165
166 pub async fn run_with_cancel(
170 &self,
171 session: &mut Session,
172 input: impl Into<String>,
173 cancel_token: CancellationToken,
174 ) -> Result<(), AgentError> {
175 session.add_message(Message::user(input.into()));
176 self.run_session_with_cancel(session, cancel_token).await
177 }
178
179 pub async fn run_session(&self, session: &mut Session) -> Result<(), AgentError> {
194 self.run_session_with_cancel(session, CancellationToken::new())
195 .await
196 }
197
198 pub async fn run_session_with_cancel(
202 &self,
203 session: &mut Session,
204 cancel_token: CancellationToken,
205 ) -> Result<(), AgentError> {
206 let (event_tx, mut event_rx) = mpsc::channel::<AgentEvent>(EVENT_CHANNEL_CAPACITY);
207
208 let drain = tokio::spawn(async move { while event_rx.recv().await.is_some() {} });
210
211 let result = self.execute_internal(session, event_tx, cancel_token).await;
212
213 drain.abort();
217 result
218 }
219
220 pub fn run_stream(
224 &self,
225 mut session: Session,
226 input: impl Into<String>,
227 ) -> mpsc::Receiver<AgentEvent> {
228 session.add_message(Message::user(input.into()));
229 self.run_stream_session(session)
230 }
231
232 pub fn run_stream_cancellable(
237 &self,
238 mut session: Session,
239 input: impl Into<String>,
240 ) -> (mpsc::Receiver<AgentEvent>, CancellationToken) {
241 session.add_message(Message::user(input.into()));
242 self.run_stream_session_cancellable(session)
243 }
244
245 pub fn run_stream_session(&self, session: Session) -> mpsc::Receiver<AgentEvent> {
248 self.run_stream_session_with_cancel(session, CancellationToken::new())
249 }
250
251 pub fn run_stream_session_cancellable(
254 &self,
255 session: Session,
256 ) -> (mpsc::Receiver<AgentEvent>, CancellationToken) {
257 let cancel_token = CancellationToken::new();
258 let rx = self.run_stream_session_with_cancel(session, cancel_token.clone());
259 (rx, cancel_token)
260 }
261
262 pub fn run_stream_session_with_cancel(
266 &self,
267 mut session: Session,
268 cancel_token: CancellationToken,
269 ) -> mpsc::Receiver<AgentEvent> {
270 let (event_tx, event_rx) = mpsc::channel::<AgentEvent>(EVENT_CHANNEL_CAPACITY);
271 let agent = self.clone();
272
273 tokio::spawn(async move {
274 if let Err(error) = agent
275 .execute_internal(&mut session, event_tx, cancel_token)
276 .await
277 {
278 tracing::warn!("Agent::run_stream execution failed: {error}");
279 }
280 });
281
282 event_rx
283 }
284
285 pub async fn execute(
303 &self,
304 session: &mut Session,
305 request: ExecuteRequest,
306 ) -> Result<(), AgentError> {
307 self.inner.execute(session, request).await
308 }
309
310 async fn execute_internal(
314 &self,
315 session: &mut Session,
316 event_tx: mpsc::Sender<AgentEvent>,
317 cancel_token: CancellationToken,
318 ) -> Result<(), AgentError> {
319 self.reexecute_approved_tool_if_pending(session, &event_tx)
328 .await;
329
330 bamboo_engine::session_app::execution_prep::prepare_session_for_execution(
337 session,
338 self.system_prompt.as_deref(),
339 self.model.as_deref(),
340 );
341
342 let initial_message = session
345 .messages
346 .iter()
347 .rev()
348 .find(|m| matches!(m.role, Role::User))
349 .map(|m| m.content.clone())
350 .unwrap_or_default();
351
352 let mut builder = ExecuteRequestBuilder::new(initial_message, event_tx, cancel_token);
356 if let Some(model) = self.model.clone() {
357 builder = builder.model(model);
358 }
359
360 self.inner.execute(session, builder.build()).await
361 }
362
363 async fn reexecute_approved_tool_if_pending(
385 &self,
386 session: &mut Session,
387 event_tx: &mpsc::Sender<AgentEvent>,
388 ) {
389 let Some(tool_call_id) = session.metadata.remove(PERMISSION_REEXECUTE_METADATA_KEY) else {
390 return;
391 };
392
393 let Some(tool_call) = find_pending_tool_call(session, &tool_call_id) else {
394 tracing::warn!(
395 session_id = %session.id,
396 tool_call_id = %tool_call_id,
397 "Permission re-exec marker set but tool call not found in history"
398 );
399 return;
400 };
401
402 let executor = self.inner.default_tools();
403 let tool_name = tool_call.function.name.clone();
404 let is_mutating = bamboo_tools::orchestrator::classify_tool(&tool_name)
405 == bamboo_tools::orchestrator::ToolMutability::Mutating;
406
407 let mut emitter = bamboo_tools::ToolEmitter::new(&tool_call.id, &tool_name, is_mutating);
412 emitter.set_auto_approved(true);
413 let _ = event_tx
414 .send(emitter.begin().clone().into_agent_event())
415 .await;
416
417 let exec_result = {
418 let ctx = bamboo_agent_core::tools::ToolExecutionContext {
419 session_id: Some(session.id.as_str()),
420 tool_call_id: tool_call_id.as_str(),
421 event_tx: Some(event_tx),
422 available_tool_schemas: None,
423 bypass_permissions: false,
424 can_async_resume: false,
425 bash_completion_sink: None,
426 pre_parsed_args: None,
427 };
428 executor.execute_with_context(&tool_call, ctx).await
429 };
430
431 let (content, success) = match exec_result {
432 Ok(tool_result) => {
433 let _ = event_tx
434 .send(
435 emitter
436 .finish(Some("Re-executed after approval".to_string()))
437 .clone()
438 .into_agent_event(),
439 )
440 .await;
441 let _ = event_tx
442 .send(AgentEvent::ToolComplete {
443 tool_call_id: tool_call.id.clone(),
444 result: tool_result.clone(),
445 })
446 .await;
447 (tool_result.result, tool_result.success)
448 }
449 Err(error) => {
450 let message = format!("Tool re-execution after approval failed: {error}");
451 let _ = event_tx
452 .send(emitter.error(message.clone()).clone().into_agent_event())
453 .await;
454 (message, false)
455 }
456 };
457
458 tracing::info!(
459 session_id = %session.id,
460 tool_name = %tool_name,
461 tool_call_id = %tool_call_id,
462 success,
463 "Re-executed approved tool after permission grant"
464 );
465 apply_tool_result(session, &tool_call_id, content, success);
466
467 if let Err(error) = self.persistence().save_runtime_session(session).await {
468 tracing::warn!(
469 session_id = %session.id,
470 %error,
471 "Failed to persist session after tool re-execution (loop's own save will retry)"
472 );
473 }
474 }
475
476 pub fn storage(&self) -> &Arc<dyn bamboo_agent_core::storage::Storage> {
478 self.inner.storage()
479 }
480
481 pub fn persistence(&self) -> &Arc<dyn bamboo_domain::RuntimeSessionPersistence> {
483 self.inner.persistence()
484 }
485
486 pub async fn answer(
535 &self,
536 session_id: impl Into<String>,
537 response: impl Into<String>,
538 ) -> Result<AnswerOutcome, SdkError> {
539 let input = RespondInput {
540 session_id: session_id.into(),
541 user_response: response.into(),
542 model: None,
543 model_ref: None,
544 provider: None,
545 reasoning_effort: None,
546 };
547 let (session, response, plan_mode_transition, permission_grants) =
548 submit_pending_response(self, input).await?;
549
550 if let Some(checker) = &self.permission_checker {
551 for (perm_type, resource) in &permission_grants {
552 checker.grant_session_permission(*perm_type, resource.clone());
553 }
554 }
555
556 Ok(AnswerOutcome {
557 session,
558 response,
559 plan_mode_transition,
560 permission_grants,
561 })
562 }
563
564 pub async fn resume(&self, session: &mut Session) -> Result<(), AgentError> {
573 self.run_session(session).await
574 }
575
576 pub async fn resume_with_cancel(
579 &self,
580 session: &mut Session,
581 cancel_token: CancellationToken,
582 ) -> Result<(), AgentError> {
583 self.run_session_with_cancel(session, cancel_token).await
584 }
585
586 pub fn resume_stream(&self, session: Session) -> mpsc::Receiver<AgentEvent> {
590 self.run_stream_session(session)
591 }
592
593 pub fn resume_stream_cancellable(
596 &self,
597 session: Session,
598 ) -> (mpsc::Receiver<AgentEvent>, CancellationToken) {
599 self.run_stream_session_cancellable(session)
600 }
601
602 pub async fn answer_and_resume_stream(
606 &self,
607 session_id: impl Into<String>,
608 response: impl Into<String>,
609 ) -> Result<mpsc::Receiver<AgentEvent>, SdkError> {
610 let outcome = self.answer(session_id, response).await?;
611 Ok(self.resume_stream(outcome.session))
612 }
613
614 pub fn answer_child_approval(
651 &self,
652 child_session_id: impl AsRef<str>,
653 request_id: impl AsRef<str>,
654 approved: bool,
655 ) -> bool {
656 bamboo_engine::external_agents::live::deliver_approval_checked(
657 child_session_id.as_ref(),
658 request_id.as_ref(),
659 approved,
660 )
661 }
662
663 pub async fn list_sessions(&self) -> Result<Vec<bamboo_storage::SessionIndexEntry>, SdkError> {
674 let store = self.session_store.as_ref().ok_or_else(|| {
675 SdkError::Unsupported(
676 "list_sessions requires an Agent built via with_defaults_for_data_dir".to_string(),
677 )
678 })?;
679 Ok(store.list_index_entries().await)
680 }
681
682 pub async fn get_session(&self, session_id: &str) -> Result<Option<Session>, SdkError> {
685 self.storage()
686 .load_session(session_id)
687 .await
688 .map_err(SdkError::Io)
689 }
690
691 pub async fn session_history(&self, session_id: &str) -> Result<Vec<Message>, SdkError> {
694 self.get_session(session_id)
695 .await?
696 .map(|session| session.messages)
697 .ok_or_else(|| SdkError::SessionNotFound(session_id.to_string()))
698 }
699
700 pub async fn delete_session(&self, session_id: &str) -> Result<bool, SdkError> {
702 self.storage()
703 .delete_session(session_id)
704 .await
705 .map_err(SdkError::Io)
706 }
707}
708
709#[derive(Debug)]
713pub struct AnswerOutcome {
714 pub session: Session,
717 pub response: String,
719 pub plan_mode_transition: Option<PlanModeTransition>,
722 pub permission_grants: Vec<(bamboo_tools::permission::PermissionType, String)>,
727}
728
729#[async_trait]
736impl SessionAccess for Agent {
737 async fn load_session(&self, id: &str) -> Result<Option<Session>, SessionLoadError> {
738 self.storage()
739 .load_session(id)
740 .await
741 .map_err(|e| SessionLoadError::StorageError(e.to_string()))
742 }
743
744 async fn load_or_create(&self, id: &str, model: &str) -> Result<Session, SessionLoadError> {
745 match SessionAccess::load_session(self, id).await? {
746 Some(session) => Ok(session),
747 None => Ok(Session::new(id.to_string(), model.to_string())),
748 }
749 }
750
751 async fn load_merged(&self, id: &str) -> Result<Option<Session>, SessionLoadError> {
752 SessionAccess::load_session(self, id).await
754 }
755
756 async fn save_session(&self, session: &mut Session) -> Result<(), SessionSaveError> {
757 self.persistence()
758 .save_runtime_session(session)
759 .await
760 .map_err(|e| SessionSaveError::StorageError(e.to_string()))
761 }
762
763 async fn save_and_cache(&self, session: &mut Session) -> Result<(), SessionSaveError> {
764 SessionAccess::save_session(self, session).await
765 }
766}
767
768fn find_pending_tool_call(
771 session: &Session,
772 tool_call_id: &str,
773) -> Option<bamboo_agent_core::tools::ToolCall> {
774 session.messages.iter().find_map(|message| {
775 message
776 .tool_calls
777 .as_ref()
778 .and_then(|calls| calls.iter().find(|call| call.id == tool_call_id).cloned())
779 })
780}
781
782fn apply_tool_result(session: &mut Session, tool_call_id: &str, content: String, success: bool) {
785 for message in &mut session.messages {
786 if message.tool_call_id.as_deref() == Some(tool_call_id) {
787 message.content = content;
788 message.tool_success = Some(success);
789 return;
790 }
791 }
792}
793
794#[cfg(test)]
795mod approval_and_session_tests {
796 use super::*;
797 use bamboo_tools::permission::{
798 PermissionChecker, PermissionContext, PermissionError, PermissionMode, PermissionType,
799 };
800 use std::sync::Mutex as StdMutex;
801
802 async fn build_test_agent(data_dir: std::path::PathBuf) -> Agent {
806 let config_json = r#"{
807 "provider": "anthropic",
808 "providers": {
809 "anthropic": { "api_key": "test-key", "model": "claude-test" }
810 }
811 }"#;
812 std::fs::write(data_dir.join("config.json"), config_json).expect("write config");
813
814 AgentBuilder::new()
815 .model("claude-test")
816 .instruction("test agent")
817 .with_defaults_for_data_dir(data_dir)
818 .await
819 .expect("defaults should assemble")
820 .build()
821 .expect("agent should build")
822 }
823
824 fn seed_session_with_pending_question(
825 session_id: &str,
826 options: Vec<String>,
827 allow_custom: bool,
828 ) -> Session {
829 let mut session = Session::new(session_id.to_string(), "claude-test".to_string());
830 session.set_pending_question(
831 "call-1".to_string(),
832 "ConclusionWithOptions".to_string(),
833 "Pick one".to_string(),
834 options,
835 allow_custom,
836 );
837 session
838 }
839
840 #[tokio::test]
841 async fn answer_resolves_pending_question_and_persists() {
842 let tmp = tempfile::tempdir().expect("tempdir");
843 let agent = build_test_agent(tmp.path().to_path_buf()).await;
844
845 let session = seed_session_with_pending_question(
846 "sess-answer-ok",
847 vec!["A".to_string(), "B".to_string()],
848 false,
849 );
850 agent
851 .storage()
852 .save_session(&session)
853 .await
854 .expect("seed session");
855
856 let outcome = agent
857 .answer("sess-answer-ok", "A")
858 .await
859 .expect("answer should succeed");
860 assert_eq!(outcome.response, "A");
861 assert!(outcome.session.pending_question.is_none());
862 assert!(outcome.permission_grants.is_empty());
863
864 let reloaded = agent
866 .storage()
867 .load_session("sess-answer-ok")
868 .await
869 .expect("load")
870 .expect("present");
871 assert!(reloaded.pending_question.is_none());
872 assert!(reloaded
873 .messages
874 .iter()
875 .any(|m| m.tool_call_id.as_deref() == Some("call-1")
876 && m.content.contains("Selected response: A")));
877 }
878
879 #[tokio::test]
880 async fn answer_rejects_response_outside_fixed_options() {
881 let tmp = tempfile::tempdir().expect("tempdir");
882 let agent = build_test_agent(tmp.path().to_path_buf()).await;
883
884 let session = seed_session_with_pending_question(
885 "sess-answer-invalid",
886 vec!["A".to_string(), "B".to_string()],
887 false,
888 );
889 agent
890 .storage()
891 .save_session(&session)
892 .await
893 .expect("seed session");
894
895 let error = agent
896 .answer("sess-answer-invalid", "not-an-option")
897 .await
898 .expect_err("response outside options should be rejected");
899 assert!(matches!(error, SdkError::InvalidResponse(_)));
900 }
901
902 #[tokio::test]
903 async fn answer_errors_when_no_pending_question() {
904 let tmp = tempfile::tempdir().expect("tempdir");
905 let agent = build_test_agent(tmp.path().to_path_buf()).await;
906
907 let session = Session::new("sess-no-pending".to_string(), "claude-test".to_string());
908 agent
909 .storage()
910 .save_session(&session)
911 .await
912 .expect("seed session");
913
914 let error = agent
915 .answer("sess-no-pending", "anything")
916 .await
917 .expect_err("no pending question should error");
918 assert!(matches!(error, SdkError::NoPendingQuestion));
919 }
920
921 #[tokio::test]
922 async fn answer_errors_when_session_missing() {
923 let tmp = tempfile::tempdir().expect("tempdir");
924 let agent = build_test_agent(tmp.path().to_path_buf()).await;
925
926 let error = agent
927 .answer("does-not-exist", "anything")
928 .await
929 .expect_err("missing session should error");
930 assert!(matches!(error, SdkError::SessionNotFound(id) if id == "does-not-exist"));
931 }
932
933 #[tokio::test]
934 async fn session_ergonomics_list_get_history_delete_round_trip() {
935 let tmp = tempfile::tempdir().expect("tempdir");
936 let agent = build_test_agent(tmp.path().to_path_buf()).await;
937
938 let mut session_a = Session::new("sess-a".to_string(), "claude-test".to_string());
939 session_a.add_message(Message::user("hello"));
940 agent
941 .storage()
942 .save_session(&session_a)
943 .await
944 .expect("save a");
945
946 let session_b = Session::new("sess-b".to_string(), "claude-test".to_string());
947 agent
948 .storage()
949 .save_session(&session_b)
950 .await
951 .expect("save b");
952
953 let listed = agent.list_sessions().await.expect("list_sessions");
954 let ids: Vec<&str> = listed.iter().map(|entry| entry.id.as_str()).collect();
955 assert!(ids.contains(&"sess-a"));
956 assert!(ids.contains(&"sess-b"));
957
958 let history = agent
959 .session_history("sess-a")
960 .await
961 .expect("session_history");
962 assert_eq!(history.len(), 1);
963 assert_eq!(history[0].content, "hello");
964
965 let missing_history = agent.session_history("does-not-exist").await;
966 assert!(matches!(
967 missing_history,
968 Err(SdkError::SessionNotFound(id)) if id == "does-not-exist"
969 ));
970
971 let deleted = agent.delete_session("sess-a").await.expect("delete");
972 assert!(deleted);
973 assert!(agent
974 .get_session("sess-a")
975 .await
976 .expect("get_session")
977 .is_none());
978 }
979
980 #[derive(Default)]
986 struct RecordingPermissionChecker {
987 grants: StdMutex<Vec<(PermissionType, String)>>,
988 }
989
990 #[async_trait]
991 impl PermissionChecker for RecordingPermissionChecker {
992 async fn needs_confirmation(&self, _perm_type: PermissionType, _resource: &str) -> bool {
993 false
994 }
995
996 async fn request_confirmation(
997 &self,
998 _ctx: PermissionContext,
999 ) -> Result<bool, PermissionError> {
1000 Ok(true)
1001 }
1002
1003 fn grant_session_permission(&self, perm_type: PermissionType, resource: String) {
1004 self.grants.lock().unwrap().push((perm_type, resource));
1005 }
1006
1007 fn set_permission_mode(&self, _mode: PermissionMode) {}
1008 }
1009
1010 #[tokio::test]
1011 async fn answer_applies_permission_grants_to_configured_checker() {
1012 let tmp = tempfile::tempdir().expect("tempdir");
1013 let config_json = r#"{
1014 "provider": "anthropic",
1015 "providers": {
1016 "anthropic": { "api_key": "test-key", "model": "claude-test" }
1017 }
1018 }"#;
1019 std::fs::write(tmp.path().join("config.json"), config_json).expect("write config");
1020
1021 let checker = Arc::new(RecordingPermissionChecker::default());
1022 let agent = AgentBuilder::new()
1023 .model("claude-test")
1024 .permission_checker(checker.clone())
1025 .with_defaults_for_data_dir(tmp.path().to_path_buf())
1026 .await
1027 .expect("defaults should assemble")
1028 .build()
1029 .expect("agent should build");
1030
1031 let mut session = Session::new("sess-permission".to_string(), "claude-test".to_string());
1036 session.set_pending_question(
1037 "call-perm-1".to_string(),
1038 "Write".to_string(),
1039 "Permission required".to_string(),
1040 vec!["Approve".to_string(), "Deny".to_string()],
1041 false,
1042 );
1043 session.add_message(Message::tool_result(
1044 "call-perm-1",
1045 serde_json::json!({
1046 "status": "awaiting_permission_approval",
1047 "question": "Permission required",
1048 "permission_type": "write_file",
1049 "resource": "/tmp/example.txt",
1050 "options": ["Approve", "Deny"],
1051 "allow_custom": false,
1052 })
1053 .to_string(),
1054 ));
1055 agent
1056 .storage()
1057 .save_session(&session)
1058 .await
1059 .expect("seed session");
1060
1061 let outcome = agent
1062 .answer("sess-permission", "Approve")
1063 .await
1064 .expect("answer should succeed");
1065 assert_eq!(
1066 outcome.permission_grants,
1067 vec![(PermissionType::WriteFile, "/tmp/example.txt".to_string())]
1068 );
1069
1070 let recorded = checker.grants.lock().unwrap();
1071 assert_eq!(
1072 *recorded,
1073 vec![(PermissionType::WriteFile, "/tmp/example.txt".to_string())]
1074 );
1075 }
1076
1077 #[tokio::test]
1078 async fn list_sessions_unsupported_without_defaults_for_data_dir() {
1079 let tmp = tempfile::tempdir().expect("tempdir");
1084 let agent = build_test_agent(tmp.path().to_path_buf()).await;
1085 let bare = Agent::from_runtime_with_config(
1086 agent.inner.clone(),
1089 None,
1090 None,
1091 None,
1092 None,
1093 );
1094 let result = bare.list_sessions().await;
1095 assert!(matches!(result, Err(SdkError::Unsupported(_))));
1096 }
1097}
1098
1099#[cfg(test)]
1100mod reexecute_and_child_approval_tests {
1101 use super::*;
1102 use bamboo_agent_core::tools::{FunctionCall, Tool, ToolCall, ToolCtx, ToolError, ToolOutcome};
1103 use std::sync::atomic::{AtomicUsize, Ordering};
1104
1105 struct RealOutputTool {
1110 calls: AtomicUsize,
1111 }
1112
1113 impl RealOutputTool {
1114 fn new() -> Self {
1115 Self {
1116 calls: AtomicUsize::new(0),
1117 }
1118 }
1119 }
1120
1121 #[async_trait]
1122 impl Tool for RealOutputTool {
1123 fn name(&self) -> &str {
1124 "real_output_tool"
1125 }
1126
1127 fn description(&self) -> &str {
1128 "test-only tool that returns a distinctive real result"
1129 }
1130
1131 fn parameters_schema(&self) -> serde_json::Value {
1132 serde_json::json!({ "type": "object", "properties": {} })
1133 }
1134
1135 async fn invoke(
1136 &self,
1137 _args: serde_json::Value,
1138 _ctx: ToolCtx,
1139 ) -> Result<ToolOutcome, ToolError> {
1140 let n = self.calls.fetch_add(1, Ordering::SeqCst);
1141 Ok(ToolOutcome::Completed(
1142 bamboo_agent_core::tools::ToolResult::text(true, format!("REAL TOOL OUTPUT #{n}")),
1143 ))
1144 }
1145 }
1146
1147 async fn build_test_agent_with_tool(
1148 data_dir: std::path::PathBuf,
1149 tool: Arc<RealOutputTool>,
1150 ) -> Agent {
1151 let config_json = r#"{
1152 "provider": "anthropic",
1153 "providers": {
1154 "anthropic": { "api_key": "test-key", "model": "claude-test" }
1155 }
1156 }"#;
1157 std::fs::write(data_dir.join("config.json"), config_json).expect("write config");
1158
1159 AgentBuilder::new()
1160 .model("claude-test")
1161 .instruction("test agent")
1162 .tool_shared(tool)
1163 .with_defaults_for_data_dir(data_dir)
1164 .await
1165 .expect("defaults should assemble")
1166 .build()
1167 .expect("agent should build")
1168 }
1169
1170 fn seed_gated_tool_session(session_id: &str, tool_call_id: &str) -> Session {
1175 let mut session = Session::new(session_id.to_string(), "claude-test".to_string());
1176 session.add_message(Message::assistant(
1177 "",
1178 Some(vec![ToolCall {
1179 id: tool_call_id.to_string(),
1180 tool_type: "function".to_string(),
1181 function: FunctionCall {
1182 name: "real_output_tool".to_string(),
1183 arguments: "{}".to_string(),
1184 },
1185 }]),
1186 ));
1187 session.set_pending_question(
1188 tool_call_id.to_string(),
1189 "real_output_tool".to_string(),
1190 "Permission required".to_string(),
1191 vec!["Approve".to_string(), "Deny".to_string()],
1192 false,
1193 );
1194 session.add_message(Message::tool_result(
1195 tool_call_id,
1196 serde_json::json!({
1197 "status": "awaiting_permission_approval",
1198 "question": "Permission required",
1199 "permission_type": "write_file",
1200 "resource": "/tmp/example.txt",
1201 "options": ["Approve", "Deny"],
1202 "allow_custom": false,
1203 })
1204 .to_string(),
1205 ));
1206 session
1207 }
1208
1209 #[tokio::test]
1210 async fn approve_marks_session_for_reexecution() {
1211 let tmp = tempfile::tempdir().expect("tempdir");
1212 let tool = Arc::new(RealOutputTool::new());
1213 let agent = build_test_agent_with_tool(tmp.path().to_path_buf(), tool).await;
1214
1215 let session = seed_gated_tool_session("sess-mark", "call-mark-1");
1216 agent
1217 .storage()
1218 .save_session(&session)
1219 .await
1220 .expect("seed session");
1221
1222 let outcome = agent
1223 .answer("sess-mark", "Approve")
1224 .await
1225 .expect("answer should succeed");
1226
1227 assert_eq!(
1228 outcome
1229 .session
1230 .metadata
1231 .get(PERMISSION_REEXECUTE_METADATA_KEY)
1232 .map(String::as_str),
1233 Some("call-mark-1"),
1234 "approving a permission prompt must stamp the re-exec marker"
1235 );
1236 }
1237
1238 #[tokio::test]
1239 async fn deny_does_not_mark_session_for_reexecution() {
1240 let tmp = tempfile::tempdir().expect("tempdir");
1241 let tool = Arc::new(RealOutputTool::new());
1242 let agent = build_test_agent_with_tool(tmp.path().to_path_buf(), tool).await;
1243
1244 let session = seed_gated_tool_session("sess-deny", "call-deny-1");
1245 agent
1246 .storage()
1247 .save_session(&session)
1248 .await
1249 .expect("seed session");
1250
1251 let outcome = agent
1252 .answer("sess-deny", "Deny")
1253 .await
1254 .expect("answer should succeed");
1255
1256 assert!(outcome.permission_grants.is_empty());
1257 assert!(!outcome
1258 .session
1259 .metadata
1260 .contains_key(PERMISSION_REEXECUTE_METADATA_KEY));
1261 let tool_message = outcome
1264 .session
1265 .messages
1266 .iter()
1267 .find(|m| m.tool_call_id.as_deref() == Some("call-deny-1"))
1268 .expect("tool result message present");
1269 assert_eq!(tool_message.content, "Selected response: Deny");
1270 }
1271
1272 #[tokio::test]
1273 async fn approve_then_reexecute_runs_real_tool_and_overwrites_placeholder() {
1274 let tmp = tempfile::tempdir().expect("tempdir");
1275 let tool = Arc::new(RealOutputTool::new());
1276 let agent = build_test_agent_with_tool(tmp.path().to_path_buf(), tool.clone()).await;
1277
1278 let session = seed_gated_tool_session("sess-reexec", "call-reexec-1");
1279 agent
1280 .storage()
1281 .save_session(&session)
1282 .await
1283 .expect("seed session");
1284
1285 let outcome = agent
1286 .answer("sess-reexec", "Approve")
1287 .await
1288 .expect("answer should succeed");
1289 let mut session = outcome.session;
1290
1291 let placeholder = session
1294 .messages
1295 .iter()
1296 .find(|m| m.tool_call_id.as_deref() == Some("call-reexec-1"))
1297 .expect("tool result message present");
1298 assert_eq!(placeholder.content, "Selected response: Approve");
1299 assert_eq!(tool.calls.load(Ordering::SeqCst), 0);
1300
1301 let (event_tx, mut event_rx) = mpsc::channel::<AgentEvent>(16);
1304 agent
1305 .reexecute_approved_tool_if_pending(&mut session, &event_tx)
1306 .await;
1307 drop(event_tx);
1308
1309 assert_eq!(tool.calls.load(Ordering::SeqCst), 1);
1312 let real_result = session
1313 .messages
1314 .iter()
1315 .find(|m| m.tool_call_id.as_deref() == Some("call-reexec-1"))
1316 .expect("tool result message present");
1317 assert_eq!(real_result.content, "REAL TOOL OUTPUT #0");
1318 assert_eq!(real_result.tool_success, Some(true));
1319 assert!(
1320 !session
1321 .metadata
1322 .contains_key(PERMISSION_REEXECUTE_METADATA_KEY),
1323 "the marker must be consumed (removed) after re-execution"
1324 );
1325
1326 let mut saw_tool_complete = false;
1329 while let Ok(event) = event_rx.try_recv() {
1330 if let AgentEvent::ToolComplete { tool_call_id, .. } = event {
1331 assert_eq!(tool_call_id, "call-reexec-1");
1332 saw_tool_complete = true;
1333 }
1334 }
1335 assert!(saw_tool_complete, "expected a ToolComplete event");
1336
1337 let reloaded = agent
1339 .storage()
1340 .load_session("sess-reexec")
1341 .await
1342 .expect("load")
1343 .expect("present");
1344 let reloaded_result = reloaded
1345 .messages
1346 .iter()
1347 .find(|m| m.tool_call_id.as_deref() == Some("call-reexec-1"))
1348 .expect("tool result message present");
1349 assert_eq!(reloaded_result.content, "REAL TOOL OUTPUT #0");
1350 }
1351
1352 #[tokio::test]
1353 async fn reexecute_is_noop_without_pending_marker() {
1354 let tmp = tempfile::tempdir().expect("tempdir");
1355 let tool = Arc::new(RealOutputTool::new());
1356 let agent = build_test_agent_with_tool(tmp.path().to_path_buf(), tool.clone()).await;
1357
1358 let mut session = Session::new("sess-noop".to_string(), "claude-test".to_string());
1359 session.add_message(Message::user("hi"));
1360
1361 let (event_tx, _event_rx) = mpsc::channel::<AgentEvent>(16);
1362 agent
1363 .reexecute_approved_tool_if_pending(&mut session, &event_tx)
1364 .await;
1365
1366 assert_eq!(tool.calls.load(Ordering::SeqCst), 0);
1367 assert_eq!(session.messages.len(), 1);
1368 }
1369
1370 #[tokio::test]
1371 async fn reexecute_warns_and_clears_marker_when_tool_call_missing() {
1372 let tmp = tempfile::tempdir().expect("tempdir");
1373 let tool = Arc::new(RealOutputTool::new());
1374 let agent = build_test_agent_with_tool(tmp.path().to_path_buf(), tool.clone()).await;
1375
1376 let mut session = Session::new("sess-missing".to_string(), "claude-test".to_string());
1377 session.metadata.insert(
1379 PERMISSION_REEXECUTE_METADATA_KEY.to_string(),
1380 "ghost-call".to_string(),
1381 );
1382
1383 let (event_tx, mut event_rx) = mpsc::channel::<AgentEvent>(16);
1384 agent
1385 .reexecute_approved_tool_if_pending(&mut session, &event_tx)
1386 .await;
1387 drop(event_tx);
1388
1389 assert_eq!(tool.calls.load(Ordering::SeqCst), 0);
1390 assert!(event_rx.try_recv().is_err(), "no events should be emitted");
1391 assert!(
1392 !session
1393 .metadata
1394 .contains_key(PERMISSION_REEXECUTE_METADATA_KEY),
1395 "the marker is removed even when the tool call can't be found, so a \
1396 missing/pruned call can't wedge every future execution"
1397 );
1398 }
1399
1400 #[tokio::test]
1401 async fn answer_child_approval_delivers_only_genuinely_pending_requests() {
1402 let tmp = tempfile::tempdir().expect("tempdir");
1403 let tool = Arc::new(RealOutputTool::new());
1404 let agent = build_test_agent_with_tool(tmp.path().to_path_buf(), tool).await;
1405
1406 assert!(!agent.answer_child_approval("child-x", "req-unknown", true));
1408
1409 let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
1415 let _live_guard = bamboo_engine::external_agents::live::register("child-x", tx);
1416 bamboo_engine::external_agents::live::register_pending_approval("child-x", "req-1");
1417
1418 assert!(agent.answer_child_approval("child-x", "req-1", true));
1419 match rx.try_recv() {
1420 Ok(bamboo_subagent::proto::ParentFrame::ApprovalReply { id, approved }) => {
1421 assert_eq!(id, "req-1");
1422 assert!(approved);
1423 }
1424 other => panic!("expected an ApprovalReply frame, got {other:?}"),
1425 }
1426
1427 assert!(!agent.answer_child_approval("child-x", "req-1", true));
1429 }
1430}