1use async_trait::async_trait;
14use serde_json::{Value, json};
15use std::collections::HashMap;
16use std::sync::{Arc, OnceLock};
17use tracing::error;
18
19use crate::background::{
20 BackgroundEventSink, BackgroundExecutableTool, BackgroundOutcome, BackgroundProgress,
21};
22use crate::tool_types::{
23 BuiltinTool, DeferrablePolicy, ToolCall, ToolDefinition, ToolHints, ToolPolicy, ToolResult,
24};
25use crate::traits::{ToolContext, ToolContextService, ToolContextServices};
26use crate::typed_id::SessionId;
27use tokio::sync::{OwnedSemaphorePermit, Semaphore};
28
29use crate::error::Result;
30use crate::traits::ToolExecutor;
31
32pub const MAX_ACTIVE_BACKGROUND_RUNS_PER_SESSION: usize = 5;
37
38const MAX_ACTIVE_BACKGROUND_RUNS_PER_WORKER: usize = 64;
43static ACTIVE_BACKGROUND_RUNS_PER_WORKER: Semaphore =
44 Semaphore::const_new(MAX_ACTIVE_BACKGROUND_RUNS_PER_WORKER);
45
46static SESSION_BACKGROUND_PERMITS: OnceLock<std::sync::Mutex<HashMap<SessionId, Arc<Semaphore>>>> =
52 OnceLock::new();
53
54struct SessionBackgroundPermit {
55 session_id: SessionId,
56 semaphore: Arc<Semaphore>,
57 permit: Option<OwnedSemaphorePermit>,
58}
59
60pub(crate) struct BackgroundRunPermit {
65 _worker: tokio::sync::SemaphorePermit<'static>,
66 _session: SessionBackgroundPermit,
67}
68
69impl Drop for SessionBackgroundPermit {
70 fn drop(&mut self) {
71 drop(self.permit.take());
72
73 let Some(permits) = SESSION_BACKGROUND_PERMITS.get() else {
74 return;
75 };
76 let mut permits = permits.lock().unwrap();
77 let should_remove = permits.get(&self.session_id).is_some_and(|current| {
78 Arc::ptr_eq(current, &self.semaphore)
79 && self.semaphore.available_permits() == MAX_ACTIVE_BACKGROUND_RUNS_PER_SESSION
80 && Arc::strong_count(&self.semaphore) == 2
81 });
82 if should_remove {
83 permits.remove(&self.session_id);
84 }
85 }
86}
87
88fn try_acquire_session_background_permit(
89 session_id: SessionId,
90) -> std::result::Result<SessionBackgroundPermit, tokio::sync::TryAcquireError> {
91 let permits = SESSION_BACKGROUND_PERMITS.get_or_init(Default::default);
92 let mut permits = permits.lock().unwrap();
93 let semaphore = permits
94 .entry(session_id)
95 .or_insert_with(|| Arc::new(Semaphore::new(MAX_ACTIVE_BACKGROUND_RUNS_PER_SESSION)))
96 .clone();
97 let permit = semaphore.clone().try_acquire_owned()?;
98
99 Ok(SessionBackgroundPermit {
100 session_id,
101 semaphore,
102 permit: Some(permit),
103 })
104}
105
106pub(crate) fn try_acquire_background_run_permit(
107 session_id: SessionId,
108) -> std::result::Result<BackgroundRunPermit, String> {
109 let worker = ACTIVE_BACKGROUND_RUNS_PER_WORKER.try_acquire().map_err(|_| {
110 format!(
111 "Worker is already running the maximum {MAX_ACTIVE_BACKGROUND_RUNS_PER_WORKER} active background runs. Try again after an existing run finishes."
112 )
113 })?;
114
115 let session = match try_acquire_session_background_permit(session_id) {
116 Ok(permit) => permit,
117 Err(_) => {
118 drop(worker);
119 return Err(format!(
120 "Maximum {MAX_ACTIVE_BACKGROUND_RUNS_PER_SESSION} active background runs per session. Wait for an existing run to finish before starting another."
121 ));
122 }
123 };
124
125 Ok(BackgroundRunPermit {
126 _worker: worker,
127 _session: session,
128 })
129}
130
131#[cfg(test)]
132fn has_session_background_permits(session_id: SessionId) -> bool {
133 SESSION_BACKGROUND_PERMITS
134 .get()
135 .and_then(|permits| permits.lock().unwrap().get(&session_id).cloned())
136 .is_some()
137}
138
139pub use crate::tool_types::ToolResultImage;
150
151#[derive(Debug)]
167pub enum ToolExecutionResult {
168 Success(Value),
170
171 SuccessWithImages {
175 result: Value,
176 images: Vec<ToolResultImage>,
177 },
178
179 ToolError(String),
184
185 InternalError(ToolInternalError),
191
192 ConnectionRequired {
199 provider: String,
201 },
202}
203
204impl ToolExecutionResult {
205 pub fn success(value: impl Into<Value>) -> Self {
207 ToolExecutionResult::Success(value.into())
208 }
209
210 pub fn success_with_raw_output(value: impl Into<Value>, raw_output: String) -> Self {
213 let mut value = value.into();
214 match value.as_object_mut() {
218 Some(obj) => {
219 obj.insert("_raw_output".to_string(), Value::String(raw_output));
220 }
221 None => {
222 value = serde_json::json!({
223 "_raw_output_scalar": value,
224 "_raw_output": raw_output,
225 });
226 }
227 }
228 ToolExecutionResult::Success(value)
229 }
230
231 pub fn success_with_images(value: impl Into<Value>, images: Vec<ToolResultImage>) -> Self {
233 ToolExecutionResult::SuccessWithImages {
234 result: value.into(),
235 images,
236 }
237 }
238
239 pub fn tool_error(message: impl Into<String>) -> Self {
241 ToolExecutionResult::ToolError(message.into())
242 }
243
244 pub fn internal_error(error: impl std::error::Error + Send + Sync + 'static) -> Self {
246 ToolExecutionResult::InternalError(ToolInternalError::new(error))
247 }
248
249 pub fn internal_error_msg(message: impl Into<String>) -> Self {
251 ToolExecutionResult::InternalError(ToolInternalError::from_message(message))
252 }
253
254 pub fn connection_required(provider: impl Into<String>) -> Self {
256 ToolExecutionResult::ConnectionRequired {
257 provider: provider.into(),
258 }
259 }
260
261 pub fn is_success(&self) -> bool {
263 matches!(
264 self,
265 ToolExecutionResult::Success(_) | ToolExecutionResult::SuccessWithImages { .. }
266 )
267 }
268
269 pub fn is_error(&self) -> bool {
271 matches!(
272 self,
273 ToolExecutionResult::ToolError(_) | ToolExecutionResult::InternalError(_)
274 )
275 }
276
277 pub fn is_connection_required(&self) -> bool {
279 matches!(self, ToolExecutionResult::ConnectionRequired { .. })
280 }
281
282 pub fn into_tool_result(self, tool_call_id: &str, tool_name: &str) -> ToolResult {
290 match self {
291 ToolExecutionResult::Success(mut value) => {
292 let raw_output = value
294 .as_object_mut()
295 .and_then(|obj| obj.remove("_raw_output"))
296 .and_then(|v| v.as_str().map(|s| s.to_string()));
297 let result_value = if let Some(obj) = value.as_object_mut() {
300 let is_scalar_carrier = raw_output.is_some()
301 && obj.len() == 1
302 && obj.contains_key("_raw_output_scalar");
303 if is_scalar_carrier {
304 obj.remove("_raw_output_scalar").unwrap_or(Value::Null)
305 } else {
306 value
307 }
308 } else {
309 value
310 };
311 ToolResult {
312 tool_call_id: tool_call_id.to_string(),
313 result: Some(result_value),
314 images: None,
315 error: None,
316 connection_required: None,
317 raw_output,
318 }
319 }
320 ToolExecutionResult::SuccessWithImages { result, images } => ToolResult {
321 tool_call_id: tool_call_id.to_string(),
322 result: Some(result),
323 images: if images.is_empty() {
324 None
325 } else {
326 Some(images)
327 },
328 error: None,
329 connection_required: None,
330 raw_output: None,
331 },
332 ToolExecutionResult::ToolError(message) => ToolResult {
333 tool_call_id: tool_call_id.to_string(),
334 result: Some(serde_json::json!({ "error": &message })),
335 images: None,
336 error: Some(message),
337 connection_required: None,
338 raw_output: None,
339 },
340 ToolExecutionResult::InternalError(err) => {
341 error!(
343 tool_name = %tool_name,
344 tool_call_id = %tool_call_id,
345 error = %err.message,
346 error_chain = %err.chain_string(),
347 "Tool internal error (details hidden from LLM)"
348 );
349
350 let generic_msg = "An internal error occurred while executing the tool";
352 ToolResult {
353 tool_call_id: tool_call_id.to_string(),
354 result: Some(serde_json::json!({
355 "error": generic_msg
356 })),
357 images: None,
358 error: Some(generic_msg.to_string()),
359 connection_required: None,
360 raw_output: None,
361 }
362 }
363 ToolExecutionResult::ConnectionRequired { ref provider } => ToolResult {
364 tool_call_id: tool_call_id.to_string(),
365 result: Some(serde_json::json!({
366 "connection_required": provider,
367 })),
368 images: None,
369 error: None,
370 connection_required: Some(provider.clone()),
371 raw_output: None,
372 },
373 }
374 }
375}
376
377#[derive(Debug)]
379pub struct ToolInternalError {
380 pub message: String,
382 pub source: Option<Box<dyn std::error::Error + Send + Sync>>,
384}
385
386impl ToolInternalError {
387 pub fn new(error: impl std::error::Error + Send + Sync + 'static) -> Self {
389 Self {
390 message: error.to_string(),
391 source: Some(Box::new(error)),
392 }
393 }
394
395 pub fn from_message(message: impl Into<String>) -> Self {
397 Self {
398 message: message.into(),
399 source: None,
400 }
401 }
402
403 pub fn chain_string(&self) -> String {
404 let mut parts = vec![self.message.clone()];
405 let mut current = <Self as std::error::Error>::source(self);
406 while let Some(source) = current {
407 let message = source.to_string();
408 if parts.last() != Some(&message) {
409 parts.push(message);
410 }
411 current = source.source();
412 }
413 parts.join(": ")
414 }
415}
416
417impl std::fmt::Display for ToolInternalError {
418 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
419 write!(f, "{}", self.message)
420 }
421}
422
423impl std::error::Error for ToolInternalError {
424 fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
425 self.source
426 .as_ref()
427 .map(|e| e.as_ref() as &(dyn std::error::Error + 'static))
428 }
429}
430
431#[async_trait]
480pub trait Tool: Send + Sync {
481 fn name(&self) -> &str;
486
487 fn display_name(&self) -> Option<&str> {
494 None
495 }
496
497 fn description(&self) -> &str;
502
503 fn parameters_schema(&self) -> Value;
509
510 async fn execute(&self, arguments: Value) -> ToolExecutionResult;
521
522 async fn execute_with_context(
539 &self,
540 arguments: Value,
541 _context: &ToolContext,
542 ) -> ToolExecutionResult {
543 self.execute(arguments).await
545 }
546
547 fn requires_context(&self) -> bool {
552 false
553 }
554
555 fn required_context_services(&self) -> &'static [ToolContextService] {
560 &[]
561 }
562
563 fn policy(&self) -> ToolPolicy {
568 ToolPolicy::Auto
569 }
570
571 fn hints(&self) -> ToolHints {
576 ToolHints::default()
577 }
578
579 fn narrate(
589 &self,
590 _tool_call: &crate::tool_types::ToolCall,
591 _phase: crate::tool_narration::ToolNarrationPhase,
592 _locale: Option<&str>,
593 _ctx: crate::tool_narration::ToolNarrationContext<'_>,
594 ) -> Option<String> {
595 None
596 }
597
598 fn as_background_executable(&self) -> Option<&dyn BackgroundExecutableTool> {
601 None
602 }
603
604 fn deferrable_policy(&self) -> DeferrablePolicy {
609 DeferrablePolicy::default()
610 }
611
612 fn to_definition(&self) -> ToolDefinition {
617 ToolDefinition::Builtin(BuiltinTool {
618 name: self.name().to_string(),
619 display_name: self.display_name().map(|s| s.to_string()),
620 description: self.description().to_string(),
621 parameters: self.parameters_schema(),
622 policy: self.policy(),
623 category: None,
624 deferrable: self.deferrable_policy(),
625 hints: self.hints(),
626 full_parameters: None,
627 })
628 }
629}
630
631#[derive(Default, Clone)]
658pub struct ToolRegistry {
659 tools: HashMap<String, Arc<dyn Tool>>,
660}
661
662impl ToolRegistry {
663 pub fn new() -> Self {
665 Self {
666 tools: HashMap::new(),
667 }
668 }
669
670 pub fn with_defaults() -> Self {
684 use crate::capabilities::{
685 AddTool, DeleteFileTool, DivideTool, EditFileTool, GetCurrentTimeTool, GetForecastTool,
686 GetWeatherTool, GrepFilesTool, ListDirectoryTool, MultiplyTool, ReadFileTool,
687 StatFileTool, SubtractTool, WriteFileTool, WriteTodosTool,
688 };
689 use crate::progress_reporting::ReportProgressTool;
690
691 let builder = ToolRegistry::builder()
692 .tool(GetCurrentTimeTool)
693 .tool(EchoTool)
694 .tool(ReportProgressTool)
704 .tool(AddTool)
706 .tool(SubtractTool)
707 .tool(MultiplyTool)
708 .tool(DivideTool)
709 .tool(GetWeatherTool)
711 .tool(GetForecastTool)
712 .tool(WriteTodosTool)
714 .tool(ReadFileTool)
716 .tool(WriteFileTool)
717 .tool(EditFileTool)
718 .tool(ListDirectoryTool)
719 .tool(GrepFilesTool)
720 .tool(DeleteFileTool)
721 .tool(StatFileTool);
722
723 #[cfg(feature = "web-fetch")]
726 let builder = builder.tool(crate::capabilities::WebFetchTool::default());
727
728 builder.build()
729 }
730
731 pub fn with_monitor_probe_defaults() -> Self {
739 use crate::capabilities::{
740 AddTool, DivideTool, GetCurrentTimeTool, MultiplyTool, SubtractTool,
741 };
742
743 ToolRegistry::builder()
744 .tool(GetCurrentTimeTool)
745 .tool(EchoTool)
746 .tool(AddTool)
747 .tool(SubtractTool)
748 .tool(MultiplyTool)
749 .tool(DivideTool)
750 .build()
751 }
752
753 pub fn register(&mut self, tool: impl Tool + 'static) {
757 self.tools.insert(tool.name().to_string(), Arc::new(tool));
758 }
759
760 pub fn register_boxed(&mut self, tool: Box<dyn Tool>) {
762 self.tools.insert(tool.name().to_string(), Arc::from(tool));
763 }
764
765 pub fn register_arc(&mut self, tool: Arc<dyn Tool>) {
767 self.tools.insert(tool.name().to_string(), tool);
768 }
769
770 pub fn get(&self, name: &str) -> Option<&Arc<dyn Tool>> {
772 self.tools.get(name)
773 }
774
775 pub fn has(&self, name: &str) -> bool {
777 self.tools.contains_key(name)
778 }
779
780 pub fn len(&self) -> usize {
782 self.tools.len()
783 }
784
785 pub fn is_empty(&self) -> bool {
787 self.tools.is_empty()
788 }
789
790 pub fn tool_names(&self) -> Vec<&str> {
792 self.tools.keys().map(|s| s.as_str()).collect()
793 }
794
795 pub fn tool_definitions(&self) -> Vec<ToolDefinition> {
800 self.tools.values().map(|t| t.to_definition()).collect()
801 }
802
803 pub fn validate_context_services(&self, services: &ToolContextServices) -> Result<()> {
806 let mut tools: Vec<_> = self.tools.values().collect();
807 tools.sort_by_key(|tool| tool.name());
808 for tool in tools {
809 for service in tool.required_context_services() {
810 if !services.provides(*service) {
811 return Err(crate::error::AgentLoopError::config(format!(
812 "tool \"{}\" requires unavailable ToolContext service {}",
813 tool.name(),
814 service.name(),
815 )));
816 }
817 }
818 }
819 Ok(())
820 }
821
822 pub fn unregister(&mut self, name: &str) -> Option<Arc<dyn Tool>> {
824 self.tools.remove(name)
825 }
826
827 pub fn clear(&mut self) {
829 self.tools.clear();
830 }
831
832 pub fn builder() -> ToolRegistryBuilder {
834 ToolRegistryBuilder::new()
835 }
836}
837
838impl std::fmt::Debug for ToolRegistry {
839 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
840 f.debug_struct("ToolRegistry")
841 .field("tools", &self.tool_names())
842 .finish()
843 }
844}
845
846#[async_trait]
847impl ToolExecutor for ToolRegistry {
848 async fn execute(
849 &self,
850 tool_call: &ToolCall,
851 _tool_def: &ToolDefinition,
852 ) -> Result<ToolResult> {
853 let tool = self.tools.get(&tool_call.name).ok_or_else(|| {
854 crate::error::AgentLoopError::tool(format!("Tool not found: {}", tool_call.name))
855 })?;
856
857 let result = tool.execute(tool_call.arguments.clone()).await;
858 Ok(result.into_tool_result(&tool_call.id, &tool_call.name))
859 }
860
861 async fn execute_with_context(
862 &self,
863 tool_call: &ToolCall,
864 _tool_def: &ToolDefinition,
865 context: &ToolContext,
866 ) -> Result<ToolResult> {
867 let tool = self.tools.get(&tool_call.name).ok_or_else(|| {
868 crate::error::AgentLoopError::tool(format!("Tool not found: {}", tool_call.name))
869 })?;
870
871 let result = tool
874 .execute_with_context(tool_call.arguments.clone(), context)
875 .await;
876 Ok(result.into_tool_result(&tool_call.id, &tool_call.name))
877 }
878}
879
880pub struct ToolRegistryBuilder {
895 registry: ToolRegistry,
896}
897
898impl ToolRegistryBuilder {
899 pub fn new() -> Self {
901 Self {
902 registry: ToolRegistry::new(),
903 }
904 }
905
906 pub fn tool(mut self, tool: impl Tool + 'static) -> Self {
908 self.registry.register(tool);
909 self
910 }
911
912 pub fn tool_boxed(mut self, tool: Box<dyn Tool>) -> Self {
914 self.registry.register_boxed(tool);
915 self
916 }
917
918 pub fn tool_arc(mut self, tool: Arc<dyn Tool>) -> Self {
920 self.registry.register_arc(tool);
921 self
922 }
923
924 pub fn build(self) -> ToolRegistry {
926 self.registry
927 }
928}
929
930impl Default for ToolRegistryBuilder {
931 fn default() -> Self {
932 Self::new()
933 }
934}
935
936pub struct EchoTool;
942
943#[async_trait]
944impl Tool for EchoTool {
945 fn name(&self) -> &str {
946 "echo"
947 }
948
949 fn display_name(&self) -> Option<&str> {
950 Some("Echo")
951 }
952
953 fn description(&self) -> &str {
954 "Echo back the provided message. Useful for testing tool execution."
955 }
956
957 fn parameters_schema(&self) -> Value {
958 serde_json::json!({
959 "type": "object",
960 "properties": {
961 "message": {
962 "type": "string",
963 "description": "The message to echo back"
964 }
965 },
966 "required": ["message"],
967 "additionalProperties": false
968 })
969 }
970
971 fn hints(&self) -> ToolHints {
972 ToolHints::default()
973 .with_readonly(true)
974 .with_idempotent(true)
975 }
976
977 async fn execute(&self, arguments: Value) -> ToolExecutionResult {
978 let message = arguments
979 .get("message")
980 .and_then(|v| v.as_str())
981 .unwrap_or("");
982
983 ToolExecutionResult::success(serde_json::json!({
984 "echoed": message,
985 "length": message.len()
986 }))
987 }
988}
989
990pub struct SpawnBackgroundTool;
992
993#[derive(Debug, Clone)]
994struct BackgroundScheduleRequest {
995 cron_expression: Option<String>,
996 scheduled_at: Option<chrono::DateTime<chrono::Utc>>,
997 timezone: String,
998}
999
1000fn parse_background_schedule(
1001 arguments: &Value,
1002) -> std::result::Result<Option<BackgroundScheduleRequest>, String> {
1003 let Some(schedule) = arguments.get("schedule") else {
1004 return Ok(None);
1005 };
1006 let Some(schedule) = schedule.as_object() else {
1007 return Err("schedule must be an object".to_string());
1008 };
1009
1010 let cron_expression = schedule
1011 .get("cron_expression")
1012 .and_then(Value::as_str)
1013 .map(str::trim)
1014 .filter(|value| !value.is_empty())
1015 .map(ToString::to_string);
1016 let scheduled_at = match schedule.get("scheduled_at").and_then(Value::as_str) {
1017 Some(value) => {
1018 let value = value.trim();
1019 if value.is_empty() {
1020 None
1021 } else {
1022 Some(
1023 chrono::DateTime::parse_from_rfc3339(value)
1024 .map_err(|_| "scheduled_at must be RFC3339".to_string())?
1025 .with_timezone(&chrono::Utc),
1026 )
1027 }
1028 }
1029 None => None,
1030 };
1031
1032 match (cron_expression.is_some(), scheduled_at.is_some()) {
1033 (false, false) => {
1034 return Err(
1035 "schedule must include exactly one of cron_expression (recurring) or scheduled_at (one-shot)"
1036 .to_string(),
1037 );
1038 }
1039 (true, true) => {
1040 return Err(
1041 "schedule must not include both cron_expression and scheduled_at; provide exactly one"
1042 .to_string(),
1043 );
1044 }
1045 _ => {}
1046 }
1047
1048 let timezone = schedule
1049 .get("timezone")
1050 .and_then(Value::as_str)
1051 .map(str::trim)
1052 .filter(|value| !value.is_empty())
1053 .unwrap_or("UTC")
1054 .to_string();
1055
1056 Ok(Some(BackgroundScheduleRequest {
1057 cron_expression,
1058 scheduled_at,
1059 timezone,
1060 }))
1061}
1062
1063fn build_background_schedule_description(
1064 tool_name: &str,
1065 tool_args: &Value,
1066 title: &str,
1067 signal_on_completion: bool,
1068) -> String {
1069 let payload = json!({
1070 "tool": tool_name,
1071 "title": title,
1072 "signal_on_completion": signal_on_completion,
1073 "args": tool_args,
1074 });
1075 let payload_json =
1076 serde_json::to_string_pretty(&payload).unwrap_or_else(|_| payload.to_string());
1077
1078 format!(
1079 "Monitor: {title}\n\n\
1080This scheduled monitor fired. Start the background run now.\n\n\
1081spawn_background payload:\n{payload_json}"
1082 )
1083}
1084
1085#[async_trait]
1086impl Tool for SpawnBackgroundTool {
1087 fn narrate(
1088 &self,
1089 tool_call: &crate::tool_types::ToolCall,
1090 phase: crate::tool_narration::ToolNarrationPhase,
1091 locale: Option<&str>,
1092 _ctx: crate::tool_narration::ToolNarrationContext<'_>,
1093 ) -> Option<String> {
1094 Some(crate::tool_narration::narrate_spawn_background(
1095 &tool_call.arguments,
1096 phase,
1097 locale,
1098 ))
1099 }
1100
1101 fn name(&self) -> &str {
1102 "spawn_background"
1103 }
1104
1105 fn display_name(&self) -> Option<&str> {
1106 Some("Spawn Background")
1107 }
1108
1109 fn description(&self) -> &str {
1110 "Run a background-capable built-in tool asynchronously. Returns immediately and signals the session when the background run completes."
1111 }
1112
1113 fn parameters_schema(&self) -> Value {
1114 json!({
1115 "type": "object",
1116 "properties": {
1117 "tool": {
1118 "type": "string",
1119 "description": "Name of the built-in tool to execute in the background"
1120 },
1121 "args": {
1122 "type": "object",
1123 "description": "Arguments to pass to the target tool"
1124 },
1125 "title": {
1126 "type": "string",
1127 "description": "Optional human-readable label for the background run"
1128 },
1129 "schedule": {
1130 "type": "object",
1131 "description": "Optional session schedule. When provided, this creates a scheduled monitor instead of starting the run immediately.",
1132 "properties": {
1133 "cron_expression": {
1134 "type": "string",
1135 "description": "Standard 5-field cron expression for recurring runs (e.g. '*/10 * * * *' for every 10 minutes)"
1136 },
1137 "scheduled_at": {
1138 "type": "string",
1139 "description": "ISO 8601 datetime for a one-shot run (e.g. '2026-04-16T15:30:00Z')"
1140 },
1141 "timezone": {
1142 "type": "string",
1143 "description": "IANA timezone for the schedule. Default: UTC"
1144 }
1145 },
1146 "additionalProperties": false
1147 },
1148 "signal_on_completion": {
1149 "type": "boolean",
1150 "description": "Send a synthetic user message back to the session when the run completes",
1151 "default": true
1152 }
1153 },
1154 "required": ["tool", "args"],
1155 "additionalProperties": false
1156 })
1157 }
1158
1159 async fn execute(&self, _arguments: Value) -> ToolExecutionResult {
1160 ToolExecutionResult::tool_error(
1161 "spawn_background requires context. This tool must be executed with session context.",
1162 )
1163 }
1164
1165 async fn execute_with_context(
1166 &self,
1167 arguments: Value,
1168 context: &ToolContext,
1169 ) -> ToolExecutionResult {
1170 let tool_name = match arguments.get("tool").and_then(|v| v.as_str()) {
1171 Some(name) if !name.trim().is_empty() => name.trim(),
1172 _ => return ToolExecutionResult::tool_error("Missing required parameter: tool"),
1173 };
1174 let tool_args = match arguments.get("args") {
1175 Some(args) if args.is_object() => args.clone(),
1176 _ => {
1177 return ToolExecutionResult::tool_error(
1178 "Missing required parameter: args (object expected)",
1179 );
1180 }
1181 };
1182 let signal_on_completion = arguments
1183 .get("signal_on_completion")
1184 .and_then(|v| v.as_bool())
1185 .unwrap_or(true);
1186 let schedule_request = match parse_background_schedule(&arguments) {
1187 Ok(schedule) => schedule,
1188 Err(message) => return ToolExecutionResult::tool_error(message),
1189 };
1190
1191 let Some(tool_registry) = &context.tool_registry else {
1192 return ToolExecutionResult::tool_error(
1193 "Tool registry not available in this context. spawn_background requires worker-side tool execution.",
1194 );
1195 };
1196
1197 let Some(tool) = tool_registry.get(tool_name).cloned() else {
1198 return ToolExecutionResult::tool_error(format!("Unknown tool: {tool_name}"));
1199 };
1200 if tool_name == self.name() {
1201 return ToolExecutionResult::tool_error(
1202 "spawn_background cannot target itself recursively",
1203 );
1204 }
1205 if tool.hints().supports_background != Some(true) {
1206 return ToolExecutionResult::tool_error(format!(
1207 "Tool does not support background execution: {tool_name}"
1208 ));
1209 }
1210 if tool.as_background_executable().is_none() {
1211 return ToolExecutionResult::tool_error(format!(
1212 "Tool declared background support but has no background executor: {tool_name}"
1213 ));
1214 }
1215 let title = arguments
1216 .get("title")
1217 .and_then(|v| v.as_str())
1218 .map(str::trim)
1219 .filter(|s| !s.is_empty())
1220 .map(|s| s.to_string())
1221 .unwrap_or_else(|| {
1222 tool.display_name()
1223 .map(ToString::to_string)
1224 .unwrap_or_else(|| format!("Background {tool_name}"))
1225 });
1226
1227 if let Some(schedule_request) = schedule_request {
1228 let Some(schedule_store) = &context.schedule_store else {
1229 return ToolExecutionResult::tool_error(
1230 "Schedule store not available in this context. Scheduled monitors require session schedules.",
1231 );
1232 };
1233
1234 let description = build_background_schedule_description(
1235 tool_name,
1236 &tool_args,
1237 &title,
1238 signal_on_completion,
1239 );
1240
1241 return match schedule_store
1242 .create_schedule_enforcing_limits(
1243 context.session_id,
1244 description,
1245 schedule_request.cron_expression.clone(),
1246 schedule_request.scheduled_at,
1247 schedule_request.timezone.clone(),
1248 )
1249 .await
1250 {
1251 Ok(schedule) => {
1252 let mut monitor_task_id: Option<String> = None;
1256 if let Some(ref task_registry) = context.session_task_registry {
1257 let spec = json!({
1258 "tool": tool_name,
1259 "arguments": &tool_args,
1260 "schedule_id": schedule.id.to_string(),
1261 "schedule_type": schedule.schedule_type,
1262 "cron_expression": schedule.cron_expression,
1263 "scheduled_at": schedule.scheduled_at,
1264 "timezone": schedule.timezone,
1265 "signal_on_completion": signal_on_completion,
1266 });
1267 match task_registry
1268 .create(crate::session_task::CreateSessionTask {
1269 session_id: context.session_id,
1270 id: None,
1271 kind: crate::session_task::TASK_KIND_MONITOR.to_string(),
1272 display_name: title.clone(),
1273 spec,
1274 state: crate::session_task::SessionTaskState::Running,
1275 links: crate::session_task::TaskLinks::default(),
1276 wake_policy: crate::session_task::TaskWakePolicy::Silent,
1279 })
1280 .await
1281 {
1282 Ok(task) => {
1283 monitor_task_id = Some(task.id);
1284 }
1285 Err(e) => {
1286 tracing::warn!(
1287 session_id = %context.session_id,
1288 schedule_id = %schedule.id,
1289 error = %e,
1290 "Failed to create monitor task for schedule (best-effort)"
1291 );
1292 }
1293 }
1294 }
1295 ToolExecutionResult::success(json!({
1296 "created": true,
1297 "status": "scheduled",
1298 "title": title,
1299 "tool": tool_name,
1300 "signal_on_completion": signal_on_completion,
1301 "schedule_id": schedule.id.to_string(),
1302 "schedule_type": schedule.schedule_type,
1303 "cron_expression": schedule.cron_expression,
1304 "scheduled_at": schedule.scheduled_at,
1305 "timezone": schedule.timezone,
1306 "next_trigger_at": schedule.next_trigger_at,
1307 "enabled": schedule.enabled,
1308 "task_id": monitor_task_id,
1309 }))
1310 }
1311 Err(crate::session_schedule::ScheduleLimitError::Store(err)) => {
1312 ToolExecutionResult::internal_error(err)
1313 }
1314 Err(crate::session_schedule::ScheduleLimitError::Rejected(msg)) => {
1315 ToolExecutionResult::tool_error(msg)
1316 }
1317 };
1318 }
1319
1320 let Some(task_registry) = &context.session_task_registry else {
1321 return ToolExecutionResult::tool_error(
1322 "Session task registry not available in this context. Background runs require task tracking.",
1323 );
1324 };
1325 if context.file_store.is_none() {
1326 return ToolExecutionResult::tool_error(
1327 "Session file store not available in this context. spawn_background requires artifact persistence.",
1328 );
1329 }
1330
1331 let background_run_permit = match try_acquire_background_run_permit(context.session_id) {
1332 Ok(permit) => permit,
1333 Err(message) => return ToolExecutionResult::tool_error(message),
1334 };
1335
1336 let run_id = format!("bg_{}", uuid::Uuid::now_v7().simple());
1337 let artifact_dir = format!("/.background/{run_id}");
1338 let log_path = format!("{artifact_dir}/output.log");
1339 let result_path = format!("{artifact_dir}/result.json");
1340
1341 let (task_id, task_attempt): (Option<String>, i32) = match task_registry
1344 .create(crate::session_task::CreateSessionTask {
1345 session_id: context.session_id,
1346 id: None,
1347 kind: crate::session_task::TASK_KIND_BACKGROUND_TOOL.to_string(),
1348 display_name: title.clone(),
1349 spec: json!({
1350 "tool": tool_name,
1351 "arguments": &tool_args,
1352 "reattachable": tool.hints().idempotent.unwrap_or(false)
1355 || tool.hints().readonly.unwrap_or(false),
1356 "signal_on_completion": signal_on_completion,
1358 }),
1359 state: crate::session_task::SessionTaskState::Running,
1360 links: crate::session_task::TaskLinks::default(),
1361 wake_policy: crate::session_task::TaskWakePolicy::Silent,
1362 })
1363 .await
1364 {
1365 Ok(task) => (Some(task.id), task.attempt),
1366 Err(e) => {
1367 return ToolExecutionResult::internal_error_msg(format!(
1368 "Failed to create background run task: {e}"
1369 ));
1370 }
1371 };
1372
1373 let background_context = context.clone().with_tool_registry(tool_registry.clone());
1374 let sink = Arc::new(SessionBackgroundSink::new(
1375 background_context.clone(),
1376 run_id.clone(),
1377 title.clone(),
1378 tool_name.to_string(),
1379 log_path.clone(),
1380 result_path.clone(),
1381 signal_on_completion,
1382 task_id.clone(),
1383 ));
1384 let run_id_for_task = run_id.clone();
1385 let tool_for_task = tool.clone();
1386 let tool_name_for_task = tool_name.to_string();
1387
1388 let cancel_registry = context.session_task_registry.clone();
1390 let cancel_session_id = context.session_id;
1391 let cancel_task_id = task_id.clone();
1392 let cancel_task_attempt = task_attempt;
1394
1395 tokio::spawn(async move {
1396 let _background_run_permit = background_run_permit;
1397 let _ = sink.status("Starting").await;
1398
1399 let outcome: std::result::Result<BackgroundOutcome, ToolExecutionResult> = match (
1407 cancel_registry.as_ref(),
1408 cancel_task_id.as_deref(),
1409 ) {
1410 (Some(registry), Some(task_id_str)) => {
1411 let registry = registry.clone();
1412 let task_id_str = task_id_str.to_string();
1413 let tool_fut = async {
1414 match tool_for_task.as_background_executable() {
1415 Some(background_tool) => {
1416 background_tool
1417 .execute_background(tool_args, background_context, sink.clone())
1418 .await
1419 }
1420 None => Err(ToolExecutionResult::tool_error(format!(
1421 "Tool declared background support but has no background executor: {}",
1422 tool_name_for_task
1423 ))),
1424 }
1425 };
1426 let watch_fut = async {
1427 loop {
1428 tokio::time::sleep(std::time::Duration::from_secs(2)).await;
1429 let _ = registry
1433 .update(
1434 cancel_session_id,
1435 &task_id_str,
1436 crate::session_task::SessionTaskUpdate {
1437 heartbeat_at: Some(chrono::Utc::now()),
1438 expected_attempt: Some(cancel_task_attempt),
1439 ..Default::default()
1440 },
1441 )
1442 .await;
1443 if let Ok(Some(task)) =
1445 registry.get(cancel_session_id, &task_id_str).await
1446 && task.cancel_requested_at.is_some()
1447 {
1448 break;
1449 }
1450 }
1451 };
1452 tokio::select! {
1453 result = tool_fut => result,
1454 () = watch_fut => {
1455 Err(ToolExecutionResult::ToolError(
1457 BACKGROUND_CANCEL_SENTINEL.to_string(),
1458 ))
1459 }
1460 }
1461 }
1462 _ => match tool_for_task.as_background_executable() {
1464 Some(background_tool) => {
1465 background_tool
1466 .execute_background(tool_args, background_context, sink.clone())
1467 .await
1468 }
1469 None => Err(ToolExecutionResult::tool_error(format!(
1470 "Tool declared background support but has no background executor: {}",
1471 tool_name_for_task
1472 ))),
1473 },
1474 };
1475
1476 let finalize_result = if is_canceled_outcome(&outcome) {
1479 sink.finalize_canceled().await
1480 } else {
1481 sink.finalize(outcome).await
1482 };
1483 if let Err(err) = finalize_result {
1484 tracing::warn!(
1485 run_id = run_id_for_task,
1486 error = %err,
1487 "Background run finalization failed"
1488 );
1489 }
1490 });
1491
1492 ToolExecutionResult::success(json!({
1493 "run_id": run_id,
1494 "resource_id": run_id,
1495 "task_id": task_id,
1496 "title": title,
1497 "tool": tool_name,
1498 "status": "running",
1499 "signal_on_completion": signal_on_completion,
1500 "artifact_dir": artifact_dir,
1501 "log_path": log_path,
1502 "result_path": result_path
1503 }))
1504 }
1505
1506 fn requires_context(&self) -> bool {
1507 true
1508 }
1509}
1510
1511#[derive(Debug, Default)]
1512struct SessionBackgroundState {
1513 status_text: String,
1514 progress: Option<BackgroundProgress>,
1515 output_tail: String,
1516 output_log: String,
1517 output_log_chars: usize,
1518 output_log_truncated: bool,
1519}
1520
1521const MAX_BACKGROUND_OUTPUT_LOG_CHARS: usize = 256 * 1024;
1522
1523struct SessionBackgroundSink {
1524 context: ToolContext,
1525 run_id: String,
1526 display_name: String,
1527 tool_name: String,
1528 log_path: String,
1529 result_path: String,
1530 signal_on_completion: bool,
1531 task_id: Option<String>,
1533 state: tokio::sync::Mutex<SessionBackgroundState>,
1534}
1535
1536impl SessionBackgroundSink {
1537 #[allow(clippy::too_many_arguments)]
1538 fn new(
1539 context: ToolContext,
1540 run_id: String,
1541 display_name: String,
1542 tool_name: String,
1543 log_path: String,
1544 result_path: String,
1545 signal_on_completion: bool,
1546 task_id: Option<String>,
1547 ) -> Self {
1548 Self {
1549 context,
1550 run_id,
1551 display_name,
1552 tool_name,
1553 log_path,
1554 result_path,
1555 signal_on_completion,
1556 task_id,
1557 state: tokio::sync::Mutex::new(SessionBackgroundState {
1558 status_text: "Queued".to_string(),
1559 ..Default::default()
1560 }),
1561 }
1562 }
1563
1564 async fn mirror_task(&self, update: crate::session_task::SessionTaskUpdate) {
1566 let (Some(registry), Some(task_id)) = (&self.context.session_task_registry, &self.task_id)
1567 else {
1568 return;
1569 };
1570 let _ = registry
1571 .update(self.context.session_id, task_id, update)
1572 .await;
1573 }
1574
1575 async fn finalize_canceled(&self) -> Result<()> {
1576 let output_log = {
1577 let state = self.state.lock().await;
1578 let mut log = Self::final_output_log(&state);
1579 log.push_str("\nCanceled by request.\n");
1580 log
1581 };
1582 self.write_text_file(&self.log_path, &output_log).await?;
1583 let result_json = serde_json::to_string_pretty(&serde_json::json!({"status": "canceled"}))
1584 .unwrap_or_else(|_| r#"{"status":"canceled"}"#.to_string());
1585 self.write_text_file(&self.result_path, &result_json)
1586 .await?;
1587
1588 let mut state = self.state.lock().await;
1589 state.status_text = "Canceled".to_string();
1590 drop(state);
1591
1592 self.mirror_task(crate::session_task::SessionTaskUpdate {
1593 state: Some(crate::session_task::SessionTaskState::Canceled),
1594 summary: Some("Canceled by request.".to_string()),
1595 result_path: Some(self.result_path.clone()),
1596 ..Default::default()
1597 })
1598 .await;
1599 if self.signal_on_completion {
1600 self.signal_session("canceled", "Canceled by request.")
1601 .await?;
1602 }
1603 Ok(())
1604 }
1605
1606 async fn finalize(
1607 &self,
1608 outcome: std::result::Result<BackgroundOutcome, ToolExecutionResult>,
1609 ) -> Result<()> {
1610 match outcome {
1611 Ok(outcome) => {
1612 let output_log = if let Some(raw_output) = &outcome.raw_output {
1613 raw_output.clone()
1614 } else {
1615 let state = self.state.lock().await;
1616 Self::final_output_log(&state)
1617 };
1618 self.write_text_file(&self.log_path, &output_log).await?;
1619 let result_json = serde_json::to_string_pretty(&outcome.result)
1620 .unwrap_or_else(|_| outcome.result.to_string());
1621 self.write_text_file(&self.result_path, &result_json)
1622 .await?;
1623
1624 let mut state = self.state.lock().await;
1625 state.status_text = "Completed".to_string();
1626 drop(state);
1627 self.mirror_task(crate::session_task::SessionTaskUpdate {
1628 state: Some(crate::session_task::SessionTaskState::Succeeded),
1629 summary: Some(outcome.summary.clone()),
1630 result_path: Some(self.result_path.clone()),
1631 ..Default::default()
1632 })
1633 .await;
1634 if self.signal_on_completion {
1635 self.signal_session("completed", &outcome.summary).await?;
1636 }
1637 }
1638 Err(err) => {
1639 let message = match err {
1640 ToolExecutionResult::ToolError(msg) => msg,
1641 ToolExecutionResult::InternalError(inner) => inner.message,
1642 ToolExecutionResult::ConnectionRequired { provider } => {
1643 format!("Background tool requires connection setup: {provider}")
1644 }
1645 ToolExecutionResult::Success(_)
1646 | ToolExecutionResult::SuccessWithImages { .. } => {
1647 "Background run ended unexpectedly".to_string()
1648 }
1649 };
1650 let output_log = {
1651 let state = self.state.lock().await;
1652 Self::final_output_log(&state)
1653 };
1654 self.write_text_file(&self.log_path, &output_log).await?;
1655 let error_json = serde_json::to_string_pretty(&json!({
1656 "status": "failed",
1657 "error": &message,
1658 }))
1659 .unwrap_or_else(|_| {
1660 json!({
1661 "status": "failed",
1662 "error": &message,
1663 })
1664 .to_string()
1665 });
1666 self.write_text_file(&self.result_path, &error_json).await?;
1667 let mut state = self.state.lock().await;
1668 state.status_text = "Failed".to_string();
1669 drop(state);
1670 self.mirror_task(crate::session_task::SessionTaskUpdate {
1671 state: Some(crate::session_task::SessionTaskState::Failed),
1672 summary: Some(message.clone()),
1673 result_path: Some(self.result_path.clone()),
1674 error: Some(crate::session_task::TaskError {
1675 kind: "error".to_string(),
1676 message: message.clone(),
1677 }),
1678 ..Default::default()
1679 })
1680 .await;
1681 if self.signal_on_completion {
1682 self.signal_session("failed", &message).await?;
1683 }
1684 }
1685 }
1686
1687 Ok(())
1688 }
1689
1690 async fn signal_session(&self, status: &str, summary: &str) -> Result<()> {
1691 let Some(platform_store) = &self.context.subagent_delegate else {
1696 tracing::warn!(
1697 run_id = %self.run_id,
1698 tool = %self.tool_name,
1699 session_id = %self.context.session_id,
1700 %status,
1701 "background run finished but no platform store is configured; \
1702 the session will not be woken (see everruns-local's wake routing)"
1703 );
1704 return Ok(());
1705 };
1706 let message = format!(
1707 "Background run {status}.\n- run_id: {}\n- title: {}\n- tool: {}\n- summary: {}\n- result_path: {}\n- log_path: {}",
1708 self.run_id,
1709 self.display_name,
1710 self.tool_name,
1711 summary,
1712 self.result_path,
1713 self.log_path
1714 );
1715 platform_store
1716 .send_message(self.context.session_id, &message)
1717 .await
1718 }
1719
1720 async fn write_text_file(&self, path: &str, content: &str) -> Result<()> {
1721 let file_store = self.context.file_store.as_ref().ok_or_else(|| {
1722 anyhow::anyhow!(
1723 "background run {} cannot persist artifact {} because no session file store is configured",
1724 self.run_id,
1725 path
1726 )
1727 })?;
1728
1729 ensure_directory(file_store.as_ref(), self.context.session_id, "/.background").await?;
1730 let run_dir = format!("/.background/{}", self.run_id);
1731 ensure_directory(file_store.as_ref(), self.context.session_id, &run_dir).await?;
1732 file_store
1733 .write_file(self.context.session_id, path, content, "text")
1734 .await?;
1735 Ok(())
1736 }
1737}
1738
1739#[async_trait]
1740impl BackgroundEventSink for SessionBackgroundSink {
1741 async fn status(&self, message: &str) -> Result<()> {
1742 let mut state = self.state.lock().await;
1743 state.status_text = message.to_string();
1744 drop(state);
1745 self.mirror_task(crate::session_task::SessionTaskUpdate {
1746 state_detail: Some(message.to_string()),
1747 ..Default::default()
1748 })
1749 .await;
1750 Ok(())
1751 }
1752
1753 async fn output(&self, stream: &str, delta: &str) -> Result<()> {
1754 let mut state = self.state.lock().await;
1755 if !delta.is_empty() {
1756 let prefix = format!("[{stream}] ");
1757 state.output_tail.push_str(&prefix);
1758 state.output_tail.push_str(delta);
1759 Self::append_to_output_log(&mut state, &prefix, delta);
1760 if state.output_tail.chars().count() > 2048 {
1761 state.output_tail = state
1762 .output_tail
1763 .chars()
1764 .rev()
1765 .take(2048)
1766 .collect::<Vec<_>>()
1767 .into_iter()
1768 .rev()
1769 .collect();
1770 }
1771 }
1772 Ok(())
1773 }
1774
1775 async fn progress(&self, progress: BackgroundProgress) -> Result<()> {
1776 let mut state = self.state.lock().await;
1777 state.progress = Some(progress.clone());
1778 drop(state);
1779 self.mirror_task(crate::session_task::SessionTaskUpdate {
1780 progress: Some(progress),
1781 ..Default::default()
1782 })
1783 .await;
1784 Ok(())
1785 }
1786}
1787
1788impl SessionBackgroundSink {
1789 fn append_to_output_log(state: &mut SessionBackgroundState, prefix: &str, delta: &str) {
1790 if state.output_log_chars >= MAX_BACKGROUND_OUTPUT_LOG_CHARS {
1791 state.output_log_truncated = true;
1792 return;
1793 }
1794
1795 let chunk = format!("{prefix}{delta}");
1796 let remaining = MAX_BACKGROUND_OUTPUT_LOG_CHARS - state.output_log_chars;
1797 let chunk_chars = chunk.chars().count();
1798
1799 if chunk_chars <= remaining {
1800 state.output_log.push_str(&chunk);
1801 state.output_log_chars += chunk_chars;
1802 return;
1803 }
1804
1805 let truncated_chunk: String = chunk.chars().take(remaining).collect();
1806 state.output_log.push_str(&truncated_chunk);
1807 state.output_log_chars += truncated_chunk.chars().count();
1808 state.output_log_truncated = true;
1809 }
1810
1811 fn final_output_log(state: &SessionBackgroundState) -> String {
1812 if !state.output_log_truncated {
1813 return state.output_log.clone();
1814 }
1815
1816 format!(
1817 "{}\n[system] background output truncated at {} characters\n",
1818 state.output_log, MAX_BACKGROUND_OUTPUT_LOG_CHARS
1819 )
1820 }
1821}
1822
1823const BACKGROUND_CANCEL_SENTINEL: &str = "__everruns_background_cancel__";
1827
1828fn is_canceled_outcome(
1832 outcome: &std::result::Result<BackgroundOutcome, ToolExecutionResult>,
1833) -> bool {
1834 matches!(outcome, Err(ToolExecutionResult::ToolError(msg)) if msg == BACKGROUND_CANCEL_SENTINEL)
1835}
1836
1837pub(crate) async fn reattach_background_run(
1853 task: &crate::session_task::SessionTask,
1854 context: &crate::traits::ToolContext,
1855) -> crate::error::Result<()> {
1856 if context.file_store.is_none() {
1859 return Err(crate::error::AgentLoopError::tool(
1860 "file store not available; cannot re-attach background run",
1861 ));
1862 }
1863 if context.session_task_registry.is_none() {
1864 return Err(crate::error::AgentLoopError::tool(
1865 "task registry not available; cannot re-attach background run",
1866 ));
1867 }
1868
1869 let tool_name: String = task
1870 .spec
1871 .get("tool")
1872 .and_then(|v| v.as_str())
1873 .filter(|s| !s.is_empty())
1874 .map(str::to_owned)
1875 .ok_or_else(|| {
1876 crate::error::AgentLoopError::tool(
1877 "background_tool spec missing 'tool' field; cannot re-attach",
1878 )
1879 })?;
1880
1881 let tool_args = task
1882 .spec
1883 .get("arguments")
1884 .cloned()
1885 .unwrap_or_else(|| serde_json::Value::Object(Default::default()));
1886
1887 let registry = std::sync::Arc::new(ToolRegistry::with_defaults());
1888
1889 let Some(tool) = registry.get(&tool_name).cloned() else {
1890 return Err(crate::error::AgentLoopError::tool(format!(
1891 "tool '{tool_name}' not found in built-in registry; cannot re-attach"
1892 )));
1893 };
1894
1895 if tool.as_background_executable().is_none() {
1896 return Err(crate::error::AgentLoopError::tool(format!(
1897 "tool '{tool_name}' does not support background execution; cannot re-attach"
1898 )));
1899 }
1900
1901 let hints = tool.hints();
1904 if !hints.idempotent.unwrap_or(false) && !hints.readonly.unwrap_or(false) {
1905 return Err(crate::error::AgentLoopError::tool(format!(
1906 "tool '{tool_name}' is not idempotent or readonly; re-attach declined",
1907 )));
1908 }
1909
1910 let background_run_permit = try_acquire_background_run_permit(task.session_id)
1913 .map_err(crate::error::AgentLoopError::tool)?;
1914
1915 let signal_on_completion = task
1918 .spec
1919 .get("signal_on_completion")
1920 .and_then(|v| v.as_bool())
1921 .unwrap_or(true);
1922
1923 let run_id = format!("bg_{}", uuid::Uuid::now_v7().simple());
1924 let artifact_dir = format!("/.background/{run_id}");
1925 let log_path = format!("{artifact_dir}/output.log");
1926 let result_path = format!("{artifact_dir}/result.json");
1927
1928 let task_id = task.id.clone();
1929 let task_attempt = task.attempt;
1930 let session_id = task.session_id;
1931
1932 let sink_context = context.clone().with_tool_registry(registry);
1933 let sink = std::sync::Arc::new(SessionBackgroundSink::new(
1934 sink_context.clone(),
1935 run_id.clone(),
1936 task.display_name.clone(),
1937 tool_name.to_string(),
1938 log_path,
1939 result_path,
1940 signal_on_completion,
1941 Some(task_id.clone()),
1942 ));
1943
1944 let cancel_registry = context.session_task_registry.clone();
1945 let run_id_for_log = run_id.clone();
1946
1947 tokio::spawn(async move {
1948 let _background_run_permit = background_run_permit;
1950 let _ = sink.status("Re-attaching").await;
1951
1952 let outcome: std::result::Result<BackgroundOutcome, ToolExecutionResult> =
1953 match (cancel_registry.as_ref(), Some(task_id.as_str())) {
1954 (Some(registry), Some(task_id_str)) => {
1955 let registry = registry.clone();
1956 let task_id_str = task_id_str.to_string();
1957 let tool_fut = async {
1958 match tool.as_background_executable() {
1959 Some(bg) => {
1960 bg.execute_background(tool_args, sink_context.clone(), sink.clone())
1961 .await
1962 }
1963 None => Err(ToolExecutionResult::tool_error(format!(
1964 "tool '{tool_name}' lost background support during re-attach"
1965 ))),
1966 }
1967 };
1968 let watch_fut = async {
1969 loop {
1970 tokio::time::sleep(std::time::Duration::from_secs(2)).await;
1971 let _ = registry
1972 .update(
1973 session_id,
1974 &task_id_str,
1975 crate::session_task::SessionTaskUpdate {
1976 heartbeat_at: Some(chrono::Utc::now()),
1977 expected_attempt: Some(task_attempt),
1978 ..Default::default()
1979 },
1980 )
1981 .await;
1982 if let Ok(Some(t)) = registry.get(session_id, &task_id_str).await
1983 && t.cancel_requested_at.is_some()
1984 {
1985 break;
1986 }
1987 }
1988 };
1989 tokio::select! {
1990 result = tool_fut => result,
1991 () = watch_fut => Err(ToolExecutionResult::ToolError(
1992 BACKGROUND_CANCEL_SENTINEL.to_string(),
1993 )),
1994 }
1995 }
1996 _ => match tool.as_background_executable() {
1997 Some(bg) => {
1998 bg.execute_background(tool_args, sink_context, sink.clone())
1999 .await
2000 }
2001 None => Err(ToolExecutionResult::tool_error(format!(
2002 "tool '{tool_name}' lost background support during re-attach"
2003 ))),
2004 },
2005 };
2006
2007 let finalize_result = if is_canceled_outcome(&outcome) {
2008 sink.finalize_canceled().await
2009 } else {
2010 sink.finalize(outcome).await
2011 };
2012 if let Err(err) = finalize_result {
2013 tracing::warn!(
2014 run_id = run_id_for_log,
2015 error = %err,
2016 "Background run re-attach finalization failed"
2017 );
2018 }
2019 });
2020
2021 Ok(())
2022}
2023
2024async fn ensure_directory(
2025 file_store: &dyn crate::traits::SessionFileSystem,
2026 session_id: crate::SessionId,
2027 path: &str,
2028) -> Result<()> {
2029 if let Some(entry) = file_store.stat_file(session_id, path).await? {
2030 if entry.is_directory {
2031 return Ok(());
2032 }
2033 return Err(anyhow::anyhow!("path exists but is not a directory: {path}").into());
2034 }
2035 let _ = file_store.create_directory(session_id, path).await?;
2036 Ok(())
2037}
2038
2039pub struct FailingTool {
2041 error_message: String,
2042 use_internal_error: bool,
2043}
2044
2045impl FailingTool {
2046 pub fn with_tool_error(message: impl Into<String>) -> Self {
2048 Self {
2049 error_message: message.into(),
2050 use_internal_error: false,
2051 }
2052 }
2053
2054 pub fn with_internal_error(message: impl Into<String>) -> Self {
2056 Self {
2057 error_message: message.into(),
2058 use_internal_error: true,
2059 }
2060 }
2061}
2062
2063impl Default for FailingTool {
2064 fn default() -> Self {
2065 Self::with_tool_error("Tool execution failed")
2066 }
2067}
2068
2069#[async_trait]
2070impl Tool for FailingTool {
2071 fn name(&self) -> &str {
2072 "failing_tool"
2073 }
2074
2075 fn display_name(&self) -> Option<&str> {
2076 Some("Failing Tool")
2077 }
2078
2079 fn description(&self) -> &str {
2080 "A tool that always fails (for testing error handling)"
2081 }
2082
2083 fn parameters_schema(&self) -> Value {
2084 serde_json::json!({
2085 "type": "object",
2086 "properties": {},
2087 "additionalProperties": false
2088 })
2089 }
2090
2091 fn hints(&self) -> ToolHints {
2092 ToolHints::default()
2093 .with_readonly(true)
2094 .with_idempotent(true)
2095 }
2096
2097 async fn execute(&self, _arguments: Value) -> ToolExecutionResult {
2098 if self.use_internal_error {
2099 ToolExecutionResult::internal_error_msg(&self.error_message)
2100 } else {
2101 ToolExecutionResult::tool_error(&self.error_message)
2102 }
2103 }
2104}
2105
2106#[cfg(test)]
2111mod tests {
2112 use super::*;
2113 use crate::capabilities::GetCurrentTimeTool;
2114 use crate::session_file::{FileInfo, FileStat, SessionFile};
2115 use crate::session_task::SessionTaskRegistry;
2116 use crate::subagent_delegation::SubagentSessionDelegate;
2117 use crate::traits::{SessionFileSystem, SessionScheduleStore};
2118 use crate::typed_id::{HarnessId, SessionId};
2119 use crate::{KeyInfo, SecretInfo};
2120 use async_trait::async_trait;
2121 use std::sync::{
2122 Arc as StdArc, Mutex,
2123 atomic::{AtomicBool, Ordering},
2124 };
2125
2126 #[derive(Default)]
2127 struct TestBackgroundTool;
2128
2129 #[async_trait]
2130 impl BackgroundExecutableTool for TestBackgroundTool {
2131 async fn execute_background(
2132 &self,
2133 arguments: Value,
2134 _context: ToolContext,
2135 sink: Arc<dyn BackgroundEventSink>,
2136 ) -> std::result::Result<BackgroundOutcome, ToolExecutionResult> {
2137 sink.status("Waiting for test result")
2138 .await
2139 .map_err(ToolExecutionResult::internal_error)?;
2140 sink.output("stdout", "hello from background")
2141 .await
2142 .map_err(ToolExecutionResult::internal_error)?;
2143 sink.progress(BackgroundProgress {
2144 current: Some(1),
2145 total: Some(1),
2146 unit: Some("step".to_string()),
2147 label: Some("done".to_string()),
2148 })
2149 .await
2150 .map_err(ToolExecutionResult::internal_error)?;
2151
2152 Ok(BackgroundOutcome {
2153 summary: arguments["summary"].as_str().unwrap_or("done").to_string(),
2154 result: json!({"ok": true}),
2155 raw_output: None,
2156 })
2157 }
2158 }
2159
2160 #[async_trait]
2161 impl Tool for TestBackgroundTool {
2162 fn name(&self) -> &str {
2163 "test_background"
2164 }
2165
2166 fn display_name(&self) -> Option<&str> {
2167 Some("Test Background")
2168 }
2169
2170 fn description(&self) -> &str {
2171 "test tool"
2172 }
2173
2174 fn parameters_schema(&self) -> Value {
2175 json!({
2176 "type": "object",
2177 "properties": {
2178 "summary": { "type": "string" }
2179 }
2180 })
2181 }
2182
2183 async fn execute(&self, _arguments: Value) -> ToolExecutionResult {
2184 ToolExecutionResult::tool_error("foreground unsupported")
2185 }
2186
2187 fn hints(&self) -> ToolHints {
2188 ToolHints::default().with_supports_background(true)
2189 }
2190
2191 fn as_background_executable(&self) -> Option<&dyn BackgroundExecutableTool> {
2192 Some(self)
2193 }
2194 }
2195
2196 #[derive(Default)]
2197 struct TestFailingBackgroundTool;
2198
2199 #[async_trait]
2200 impl BackgroundExecutableTool for TestFailingBackgroundTool {
2201 async fn execute_background(
2202 &self,
2203 _arguments: Value,
2204 _context: ToolContext,
2205 sink: Arc<dyn BackgroundEventSink>,
2206 ) -> std::result::Result<BackgroundOutcome, ToolExecutionResult> {
2207 sink.status("Running failing test")
2208 .await
2209 .map_err(ToolExecutionResult::internal_error)?;
2210 sink.output("stderr", "background failed")
2211 .await
2212 .map_err(ToolExecutionResult::internal_error)?;
2213 Err(ToolExecutionResult::tool_error("boom"))
2214 }
2215 }
2216
2217 #[async_trait]
2218 impl Tool for TestFailingBackgroundTool {
2219 fn name(&self) -> &str {
2220 "test_background_fail"
2221 }
2222
2223 fn display_name(&self) -> Option<&str> {
2224 Some("Test Background Fail")
2225 }
2226
2227 fn description(&self) -> &str {
2228 "failing background test tool"
2229 }
2230
2231 fn parameters_schema(&self) -> Value {
2232 json!({
2233 "type": "object",
2234 "properties": {}
2235 })
2236 }
2237
2238 async fn execute(&self, _arguments: Value) -> ToolExecutionResult {
2239 ToolExecutionResult::tool_error("foreground unsupported")
2240 }
2241
2242 fn hints(&self) -> ToolHints {
2243 ToolHints::default().with_supports_background(true)
2244 }
2245
2246 fn as_background_executable(&self) -> Option<&dyn BackgroundExecutableTool> {
2247 Some(self)
2248 }
2249 }
2250
2251 #[derive(Default)]
2252 struct TestLargeOutputBackgroundTool;
2253
2254 #[async_trait]
2255 impl BackgroundExecutableTool for TestLargeOutputBackgroundTool {
2256 async fn execute_background(
2257 &self,
2258 _arguments: Value,
2259 _context: ToolContext,
2260 sink: Arc<dyn BackgroundEventSink>,
2261 ) -> std::result::Result<BackgroundOutcome, ToolExecutionResult> {
2262 let large_chunk = "x".repeat(MAX_BACKGROUND_OUTPUT_LOG_CHARS + 4096);
2263 sink.output("stdout", &large_chunk)
2264 .await
2265 .map_err(ToolExecutionResult::internal_error)?;
2266 Ok(BackgroundOutcome {
2267 summary: "large output complete".to_string(),
2268 result: json!({"ok": true}),
2269 raw_output: None,
2270 })
2271 }
2272 }
2273
2274 #[async_trait]
2275 impl Tool for TestLargeOutputBackgroundTool {
2276 fn name(&self) -> &str {
2277 "test_background_large_output"
2278 }
2279
2280 fn display_name(&self) -> Option<&str> {
2281 Some("Test Background Large Output")
2282 }
2283
2284 fn description(&self) -> &str {
2285 "background test tool with huge output"
2286 }
2287
2288 fn parameters_schema(&self) -> Value {
2289 json!({
2290 "type": "object",
2291 "properties": {}
2292 })
2293 }
2294
2295 async fn execute(&self, _arguments: Value) -> ToolExecutionResult {
2296 ToolExecutionResult::tool_error("foreground unsupported")
2297 }
2298
2299 fn hints(&self) -> ToolHints {
2300 ToolHints::default().with_supports_background(true)
2301 }
2302
2303 fn as_background_executable(&self) -> Option<&dyn BackgroundExecutableTool> {
2304 Some(self)
2305 }
2306 }
2307
2308 struct BlockingBackgroundTool {
2309 release: StdArc<AtomicBool>,
2310 }
2311
2312 #[async_trait]
2313 impl BackgroundExecutableTool for BlockingBackgroundTool {
2314 async fn execute_background(
2315 &self,
2316 _arguments: Value,
2317 _context: ToolContext,
2318 sink: Arc<dyn BackgroundEventSink>,
2319 ) -> std::result::Result<BackgroundOutcome, ToolExecutionResult> {
2320 sink.status("Blocking until released")
2321 .await
2322 .map_err(ToolExecutionResult::internal_error)?;
2323 while !self.release.load(Ordering::SeqCst) {
2324 tokio::time::sleep(std::time::Duration::from_millis(10)).await;
2325 }
2326 Ok(BackgroundOutcome {
2327 summary: "released".to_string(),
2328 result: json!({"ok": true}),
2329 raw_output: None,
2330 })
2331 }
2332 }
2333
2334 #[async_trait]
2335 impl Tool for BlockingBackgroundTool {
2336 fn name(&self) -> &str {
2337 "test_background_blocking"
2338 }
2339
2340 fn display_name(&self) -> Option<&str> {
2341 Some("Test Background Blocking")
2342 }
2343
2344 fn description(&self) -> &str {
2345 "background test tool that waits for test release"
2346 }
2347
2348 fn parameters_schema(&self) -> Value {
2349 json!({
2350 "type": "object",
2351 "properties": {}
2352 })
2353 }
2354
2355 async fn execute(&self, _arguments: Value) -> ToolExecutionResult {
2356 ToolExecutionResult::tool_error("foreground unsupported")
2357 }
2358
2359 fn hints(&self) -> ToolHints {
2360 ToolHints::default().with_supports_background(true)
2361 }
2362
2363 fn as_background_executable(&self) -> Option<&dyn BackgroundExecutableTool> {
2364 Some(self)
2365 }
2366 }
2367
2368 #[derive(Default)]
2369 struct TestFileStore {
2370 files: Mutex<HashMap<String, SessionFile>>,
2371 }
2372
2373 #[async_trait]
2374 impl crate::traits::SessionFileSystem for TestFileStore {
2375 fn is_mount_resolver(&self) -> bool {
2376 false
2377 }
2378
2379 async fn read_file(
2380 &self,
2381 _session_id: SessionId,
2382 path: &str,
2383 ) -> crate::Result<Option<SessionFile>> {
2384 Ok(self.files.lock().unwrap().get(path).cloned())
2385 }
2386
2387 async fn write_file(
2388 &self,
2389 session_id: SessionId,
2390 path: &str,
2391 content: &str,
2392 encoding: &str,
2393 ) -> crate::Result<SessionFile> {
2394 let now = chrono::Utc::now();
2395 let file = SessionFile {
2396 id: uuid::Uuid::now_v7(),
2397 session_id: session_id.uuid(),
2398 path: path.to_string(),
2399 name: FileInfo::name_from_path(path),
2400 content: Some(content.to_string()),
2401 encoding: encoding.to_string(),
2402 is_directory: false,
2403 is_readonly: false,
2404 size_bytes: content.len() as i64,
2405 created_at: now,
2406 updated_at: now,
2407 };
2408 self.files
2409 .lock()
2410 .unwrap()
2411 .insert(path.to_string(), file.clone());
2412 Ok(file)
2413 }
2414
2415 async fn delete_file(
2416 &self,
2417 _session_id: SessionId,
2418 _path: &str,
2419 _recursive: bool,
2420 ) -> crate::Result<bool> {
2421 Ok(false)
2422 }
2423
2424 async fn list_directory(
2425 &self,
2426 _session_id: SessionId,
2427 _path: &str,
2428 ) -> crate::Result<Vec<FileInfo>> {
2429 Ok(Vec::new())
2430 }
2431
2432 async fn stat_file(
2433 &self,
2434 _session_id: SessionId,
2435 path: &str,
2436 ) -> crate::Result<Option<FileStat>> {
2437 let file = self.files.lock().unwrap().get(path).cloned();
2438 Ok(file.map(|entry| FileStat {
2439 path: entry.path,
2440 name: entry.name,
2441 is_directory: entry.is_directory,
2442 is_readonly: entry.is_readonly,
2443 size_bytes: entry.size_bytes,
2444 created_at: entry.created_at,
2445 updated_at: entry.updated_at,
2446 }))
2447 }
2448
2449 async fn grep_files(
2450 &self,
2451 _session_id: SessionId,
2452 _pattern: &str,
2453 _path_pattern: Option<&str>,
2454 ) -> crate::Result<Vec<crate::session_file::GrepMatch>> {
2455 Ok(Vec::new())
2456 }
2457
2458 async fn create_directory(
2459 &self,
2460 session_id: SessionId,
2461 path: &str,
2462 ) -> crate::Result<FileInfo> {
2463 let now = chrono::Utc::now();
2464 let id = uuid::Uuid::now_v7();
2465 let dir = SessionFile {
2466 id,
2467 session_id: session_id.uuid(),
2468 path: path.to_string(),
2469 name: FileInfo::name_from_path(path),
2470 content: None,
2471 encoding: "text".to_string(),
2472 is_directory: true,
2473 is_readonly: false,
2474 size_bytes: 0,
2475 created_at: now,
2476 updated_at: now,
2477 };
2478 self.files.lock().unwrap().insert(path.to_string(), dir);
2479 Ok(FileInfo {
2480 id,
2481 session_id: session_id.uuid(),
2482 path: path.to_string(),
2483 name: FileInfo::name_from_path(path),
2484 is_directory: true,
2485 is_readonly: false,
2486 size_bytes: 0,
2487 created_at: now,
2488 updated_at: now,
2489 })
2490 }
2491 }
2492
2493 #[derive(Default)]
2494 struct TestSubagentDelegate {
2495 sent_messages: Mutex<Vec<String>>,
2496 }
2497
2498 #[async_trait]
2499 #[async_trait]
2500 impl SubagentSessionDelegate for TestSubagentDelegate {
2501 async fn get_agent_by_id(
2502 &self,
2503 _id: crate::typed_id::AgentId,
2504 ) -> crate::Result<Option<crate::Agent>> {
2505 Ok(None)
2506 }
2507 async fn add_agent_session_participant(
2508 &self,
2509 _session_id: SessionId,
2510 _agent_id: crate::typed_id::AgentId,
2511 ) -> crate::Result<crate::session::SessionParticipant> {
2512 Err(crate::error::AgentLoopError::tool(
2513 "test store does not add participants",
2514 ))
2515 }
2516 async fn get_harness(&self, _id: HarnessId) -> crate::Result<Option<crate::Harness>> {
2517 Ok(None)
2518 }
2519 async fn create_session_with_options(
2520 &self,
2521 _request: crate::subagent_delegation::PlatformCreateSessionRequest,
2522 ) -> crate::Result<crate::Session> {
2523 Err(crate::error::AgentLoopError::tool(
2524 "test store does not create sessions",
2525 ))
2526 }
2527 async fn get_session_by_id(&self, _id: SessionId) -> crate::Result<Option<crate::Session>> {
2528 Ok(None)
2529 }
2530 async fn send_message(&self, _session_id: SessionId, content: &str) -> crate::Result<()> {
2531 self.sent_messages.lock().unwrap().push(content.to_string());
2532 Ok(())
2533 }
2534 async fn get_messages(
2535 &self,
2536 _session_id: SessionId,
2537 _limit: Option<usize>,
2538 ) -> crate::Result<Vec<crate::subagent_delegation::PlatformMessage>> {
2539 Ok(Vec::new())
2540 }
2541 async fn wait_for_idle(
2542 &self,
2543 _session_id: SessionId,
2544 _timeout_secs: Option<u64>,
2545 ) -> crate::Result<String> {
2546 Ok("idle".to_string())
2547 }
2548 }
2549
2550 #[derive(Default)]
2551 struct NoopStorageStore;
2552
2553 #[derive(Default)]
2554 struct TestScheduleStore {
2555 schedules: Mutex<Vec<crate::session_schedule::SessionSchedule>>,
2556 }
2557
2558 #[async_trait]
2559 impl crate::traits::SessionScheduleStore for TestScheduleStore {
2560 async fn create_schedule(
2561 &self,
2562 session_id: SessionId,
2563 description: String,
2564 cron_expression: Option<String>,
2565 scheduled_at: Option<chrono::DateTime<chrono::Utc>>,
2566 timezone: String,
2567 ) -> crate::Result<crate::session_schedule::SessionSchedule> {
2568 let schedule = crate::session_schedule::SessionSchedule {
2569 id: crate::typed_id::ScheduleId::new(),
2570 session_id,
2571 owner_principal_id: crate::PrincipalId::from_seed(1),
2572 resolved_owner_user_id: None,
2573 owner: None,
2574 effective_owner: None,
2575 description,
2576 cron_expression: cron_expression.clone(),
2577 scheduled_at,
2578 timezone,
2579 enabled: true,
2580 schedule_type: crate::session_schedule::SessionSchedule::derive_type(
2581 &cron_expression,
2582 ),
2583 next_trigger_at: Some(chrono::Utc::now() + chrono::Duration::minutes(10)),
2584 last_triggered_at: None,
2585 trigger_count: 0,
2586 created_at: chrono::Utc::now(),
2587 updated_at: chrono::Utc::now(),
2588 };
2589 self.schedules.lock().unwrap().push(schedule.clone());
2590 Ok(schedule)
2591 }
2592
2593 async fn cancel_schedule(
2594 &self,
2595 _session_id: SessionId,
2596 schedule_id: crate::ScheduleId,
2597 ) -> crate::Result<crate::session_schedule::SessionSchedule> {
2598 let mut schedules = self.schedules.lock().unwrap();
2599 let schedule = schedules
2600 .iter_mut()
2601 .find(|schedule| schedule.id == schedule_id)
2602 .ok_or_else(|| crate::AgentLoopError::tool("Schedule not found".to_string()))?;
2603 schedule.enabled = false;
2604 Ok(schedule.clone())
2605 }
2606
2607 async fn list_schedules(
2608 &self,
2609 session_id: SessionId,
2610 ) -> crate::Result<Vec<crate::session_schedule::SessionSchedule>> {
2611 Ok(self
2612 .schedules
2613 .lock()
2614 .unwrap()
2615 .iter()
2616 .filter(|schedule| schedule.session_id == session_id)
2617 .cloned()
2618 .collect())
2619 }
2620
2621 async fn count_active_schedules(&self, session_id: SessionId) -> crate::Result<u32> {
2622 Ok(self
2623 .schedules
2624 .lock()
2625 .unwrap()
2626 .iter()
2627 .filter(|schedule| schedule.session_id == session_id && schedule.enabled)
2628 .count() as u32)
2629 }
2630
2631 async fn count_active_org_schedules(&self) -> crate::Result<u32> {
2632 Ok(self
2634 .schedules
2635 .lock()
2636 .unwrap()
2637 .iter()
2638 .filter(|schedule| schedule.enabled)
2639 .count() as u32)
2640 }
2641 }
2642
2643 #[async_trait]
2644 impl crate::traits::SessionStorageStore for NoopStorageStore {
2645 async fn set_value(
2646 &self,
2647 _session_id: SessionId,
2648 _key: &str,
2649 _value: &str,
2650 ) -> crate::Result<()> {
2651 Ok(())
2652 }
2653 async fn get_value(
2654 &self,
2655 _session_id: SessionId,
2656 _key: &str,
2657 ) -> crate::Result<Option<String>> {
2658 Ok(None)
2659 }
2660 async fn delete_value(&self, _session_id: SessionId, _key: &str) -> crate::Result<bool> {
2661 Ok(false)
2662 }
2663 async fn list_keys(&self, _session_id: SessionId) -> crate::Result<Vec<KeyInfo>> {
2664 Ok(Vec::new())
2665 }
2666 async fn set_secret(
2667 &self,
2668 _session_id: SessionId,
2669 _name: &str,
2670 _value: &str,
2671 ) -> crate::Result<()> {
2672 Ok(())
2673 }
2674 async fn get_secret(
2675 &self,
2676 _session_id: SessionId,
2677 _name: &str,
2678 ) -> crate::Result<Option<String>> {
2679 Ok(None)
2680 }
2681 async fn delete_secret(&self, _session_id: SessionId, _name: &str) -> crate::Result<bool> {
2682 Ok(false)
2683 }
2684 async fn list_secrets(&self, _session_id: SessionId) -> crate::Result<Vec<SecretInfo>> {
2685 Ok(Vec::new())
2686 }
2687 }
2688
2689 #[tokio::test]
2690 async fn test_echo_tool() {
2691 let tool = EchoTool;
2692
2693 let result = tool
2694 .execute(serde_json::json!({"message": "Hello, world!"}))
2695 .await;
2696
2697 if let ToolExecutionResult::Success(value) = result {
2698 assert_eq!(
2699 value.get("echoed").unwrap().as_str().unwrap(),
2700 "Hello, world!"
2701 );
2702 assert_eq!(value.get("length").unwrap().as_u64().unwrap(), 13);
2703 } else {
2704 panic!("Expected success");
2705 }
2706 }
2707
2708 #[tokio::test]
2709 async fn test_failing_tool_with_tool_error() {
2710 let tool = FailingTool::with_tool_error("Something went wrong");
2711
2712 let result = tool.execute(serde_json::json!({})).await;
2713
2714 if let ToolExecutionResult::ToolError(msg) = result {
2715 assert_eq!(msg, "Something went wrong");
2716 } else {
2717 panic!("Expected tool error");
2718 }
2719 }
2720
2721 #[tokio::test]
2722 async fn test_failing_tool_with_internal_error() {
2723 let tool = FailingTool::with_internal_error("Database connection failed");
2724
2725 let result = tool.execute(serde_json::json!({})).await;
2726
2727 if let ToolExecutionResult::InternalError(err) = result {
2728 assert_eq!(err.message, "Database connection failed");
2729 } else {
2730 panic!("Expected internal error");
2731 }
2732 }
2733
2734 #[tokio::test]
2735 async fn test_tool_result_conversion() {
2736 let result = ToolExecutionResult::success(serde_json::json!({"value": 42}));
2738 let tool_result = result.into_tool_result("call_1", "test_tool");
2739 assert!(tool_result.error.is_none());
2740 assert_eq!(tool_result.result.unwrap()["value"], 42);
2741
2742 let result = ToolExecutionResult::tool_error("Invalid input");
2744 let tool_result = result.into_tool_result("call_2", "test_tool");
2745 assert_eq!(tool_result.error.as_deref(), Some("Invalid input"));
2746 assert_eq!(
2747 tool_result.result.unwrap(),
2748 serde_json::json!({"error": "Invalid input"})
2749 );
2750
2751 let result = ToolExecutionResult::internal_error_msg("Secret database error");
2753 let tool_result = result.into_tool_result("call_3", "test_tool");
2754 assert_eq!(
2755 tool_result.error.as_deref(),
2756 Some("An internal error occurred while executing the tool")
2757 );
2758 assert_eq!(
2759 tool_result.result.unwrap(),
2760 serde_json::json!({"error": "An internal error occurred while executing the tool"})
2761 );
2762 }
2763
2764 #[tokio::test]
2765 async fn test_tool_registry() {
2766 let mut registry = ToolRegistry::new();
2767 registry.register(GetCurrentTimeTool);
2768 registry.register(EchoTool);
2769
2770 assert_eq!(registry.len(), 2);
2771 assert!(registry.has("get_current_time"));
2772 assert!(registry.has("echo"));
2773 assert!(!registry.has("nonexistent"));
2774
2775 let definitions = registry.tool_definitions();
2776 assert_eq!(definitions.len(), 2);
2777 }
2778
2779 #[tokio::test]
2780 async fn test_tool_registry_builder() {
2781 let registry = ToolRegistry::builder()
2782 .tool(GetCurrentTimeTool)
2783 .tool(EchoTool)
2784 .build();
2785
2786 assert_eq!(registry.len(), 2);
2787 }
2788
2789 #[test]
2790 fn test_tool_display_name_in_definition() {
2791 let tool = GetCurrentTimeTool;
2793 assert_eq!(tool.display_name(), Some("Get Current Time"));
2794
2795 let def = tool.to_definition();
2796 assert_eq!(def.display_name(), Some("Get Current Time"));
2797 }
2798
2799 #[test]
2800 fn test_success_with_raw_output_object_preserves_shape() {
2801 let res = ToolExecutionResult::success_with_raw_output(
2802 serde_json::json!({"stdout": "hello"}),
2803 "raw stdout bytes".to_string(),
2804 );
2805 let tr = res.into_tool_result("call_1", "demo");
2806 assert_eq!(tr.result.as_ref().unwrap()["stdout"], "hello");
2807 assert!(
2808 tr.result
2809 .as_ref()
2810 .unwrap()
2811 .as_object()
2812 .unwrap()
2813 .get("_raw_output")
2814 .is_none(),
2815 "sidecar key must not leak to the LLM-visible result"
2816 );
2817 assert_eq!(tr.raw_output.as_deref(), Some("raw stdout bytes"));
2818 }
2819
2820 #[test]
2821 fn test_success_with_raw_output_scalar_unwraps_to_string() {
2822 let res = ToolExecutionResult::success_with_raw_output(
2823 "compact summary".to_string(),
2824 "full output bytes".to_string(),
2825 );
2826 let tr = res.into_tool_result("call_1", "demo");
2827 assert_eq!(
2828 tr.result,
2829 Some(serde_json::Value::String("compact summary".into()))
2830 );
2831 assert_eq!(tr.raw_output.as_deref(), Some("full output bytes"));
2832 }
2833
2834 #[test]
2835 fn test_success_result_with_raw_output_scalar_key_is_not_unwrapped() {
2836 let res = ToolExecutionResult::success(
2837 serde_json::json!({"_raw_output_scalar": "user_value", "kept": true}),
2838 );
2839 let tr = res.into_tool_result("call_1", "demo");
2840 assert_eq!(
2841 tr.result,
2842 Some(serde_json::json!({"_raw_output_scalar": "user_value", "kept": true}))
2843 );
2844 assert_eq!(tr.raw_output, None);
2845 }
2846
2847 #[test]
2848 fn test_success_result_with_only_raw_output_scalar_key_is_not_unwrapped() {
2849 let res = ToolExecutionResult::success(serde_json::json!({"_raw_output_scalar": "v"}));
2852 let tr = res.into_tool_result("call_1", "demo");
2853 assert_eq!(
2854 tr.result,
2855 Some(serde_json::json!({"_raw_output_scalar": "v"}))
2856 );
2857 assert_eq!(tr.raw_output, None);
2858 }
2859
2860 #[test]
2861 fn test_echo_tool_display_name() {
2862 let tool = EchoTool;
2863 assert_eq!(tool.display_name(), Some("Echo"));
2864
2865 let def = tool.to_definition();
2866 assert_eq!(def.display_name(), Some("Echo"));
2867 }
2868
2869 #[test]
2870 fn test_all_default_tools_have_display_names() {
2871 let registry = ToolRegistry::with_defaults();
2872 let definitions = registry.tool_definitions();
2873
2874 for def in &definitions {
2875 assert!(
2876 def.display_name().is_some(),
2877 "Tool '{}' should have a display_name",
2878 def.name()
2879 );
2880 }
2881 }
2882
2883 #[tokio::test]
2884 async fn test_tool_registry_as_executor() {
2885 let mut registry = ToolRegistry::new();
2886 registry.register(EchoTool);
2887
2888 let tool_call = ToolCall {
2889 id: "call_1".to_string(),
2890 name: "echo".to_string(),
2891 arguments: serde_json::json!({"message": "test"}),
2892 };
2893
2894 let tool_def = registry.get("echo").unwrap().to_definition();
2895 let result = registry.execute(&tool_call, &tool_def).await.unwrap();
2896
2897 assert!(result.error.is_none());
2898 assert_eq!(result.result.unwrap()["echoed"], "test");
2899 }
2900
2901 #[test]
2902 fn test_tool_to_definition() {
2903 let tool = GetCurrentTimeTool;
2904 let def = tool.to_definition();
2905
2906 let ToolDefinition::Builtin(builtin) = def else {
2907 panic!("expected Builtin variant");
2908 };
2909 assert_eq!(builtin.name, "get_current_time");
2910 assert_eq!(builtin.policy, ToolPolicy::Auto);
2911 }
2912
2913 #[test]
2914 fn test_with_defaults_has_expected_tools() {
2915 let registry = ToolRegistry::with_defaults();
2916
2917 assert!(
2919 registry.has("get_current_time"),
2920 "should have get_current_time"
2921 );
2922 assert!(registry.has("echo"), "should have echo");
2923 assert!(
2926 !registry.has("spawn_background"),
2927 "spawn_background must NOT be in defaults — it comes from the \
2928 background_execution capability"
2929 );
2930 assert!(
2931 registry.has("report_progress"),
2932 "should have report_progress"
2933 );
2934
2935 assert!(registry.has("add"), "should have add");
2937 assert!(registry.has("subtract"), "should have subtract");
2938 assert!(registry.has("multiply"), "should have multiply");
2939 assert!(registry.has("divide"), "should have divide");
2940
2941 assert!(registry.has("get_weather"), "should have get_weather");
2943 assert!(registry.has("get_forecast"), "should have get_forecast");
2944
2945 assert!(registry.has("write_todos"), "should have write_todos");
2947
2948 assert!(registry.has("read_file"), "should have read_file");
2950 assert!(registry.has("write_file"), "should have write_file");
2951 assert!(registry.has("edit_file"), "should have edit_file");
2952 assert!(registry.has("list_directory"), "should have list_directory");
2953 assert!(registry.has("grep_files"), "should have grep_files");
2954 assert!(registry.has("delete_file"), "should have delete_file");
2955 assert!(registry.has("stat_file"), "should have stat_file");
2956
2957 assert!(registry.has("web_fetch"), "should have web_fetch");
2959
2960 assert_eq!(registry.len(), 18, "should have 18 default tools");
2962 }
2963
2964 #[tokio::test]
2965 async fn test_with_defaults_tools_are_executable() {
2966 let registry = ToolRegistry::with_defaults();
2967
2968 let tool_call = ToolCall {
2970 id: "call_1".to_string(),
2971 name: "echo".to_string(),
2972 arguments: serde_json::json!({"message": "hello from defaults"}),
2973 };
2974
2975 let tool_def = registry.get("echo").unwrap().to_definition();
2976 let result = registry.execute(&tool_call, &tool_def).await.unwrap();
2977
2978 assert!(result.error.is_none());
2979 assert_eq!(result.result.unwrap()["echoed"], "hello from defaults");
2980 }
2981
2982 #[tokio::test]
2983 async fn test_with_defaults_math_tools() {
2984 let registry = ToolRegistry::with_defaults();
2985
2986 let tool_call = ToolCall {
2988 id: "call_add".to_string(),
2989 name: "add".to_string(),
2990 arguments: serde_json::json!({"a": 5, "b": 3}),
2991 };
2992
2993 let tool_def = registry.get("add").unwrap().to_definition();
2994 let result = registry.execute(&tool_call, &tool_def).await.unwrap();
2995
2996 assert!(result.error.is_none());
2997 assert_eq!(result.result.unwrap()["result"].as_f64().unwrap(), 8.0);
2999 }
3000
3001 #[test]
3005 fn test_with_defaults_excludes_capability_only_tools() {
3006 let registry = ToolRegistry::with_defaults();
3007
3008 assert!(
3010 !registry.has("bash"),
3011 "bash must not be in defaults — it comes from bashkit_shell capability"
3012 );
3013 assert!(
3015 !registry.has("kv_store"),
3016 "kv_store must not be in defaults — it comes from session_storage capability"
3017 );
3018 assert!(
3022 !registry.has("spawn_background"),
3023 "spawn_background must not be in defaults — it comes from the \
3024 background_execution capability (auto-activated by tool hints)"
3025 );
3026 }
3027
3028 #[tokio::test]
3029 async fn test_spawn_background_executes_and_signals_session() {
3030 let session_id = SessionId::new();
3031 let file_store = Arc::new(TestFileStore::default());
3032 let platform_store = Arc::new(TestSubagentDelegate::default());
3033 let storage_store = Arc::new(NoopStorageStore);
3034 let task_registry = Arc::new(InMemoryTaskRegistry::default());
3035 let tool_registry = ToolRegistry::builder()
3036 .tool(SpawnBackgroundTool)
3037 .tool(TestBackgroundTool)
3038 .build();
3039
3040 let context = ToolContext::with_stores(session_id, file_store.clone(), storage_store)
3041 .with_tool_registry(Arc::new(tool_registry))
3042 .with_subagent_delegate(platform_store.clone())
3043 .with_session_task_registry(task_registry.clone());
3044
3045 let tool = SpawnBackgroundTool;
3046 let result = tool
3047 .execute_with_context(
3048 json!({
3049 "tool": "test_background",
3050 "args": { "summary": "Background complete" }
3051 }),
3052 &context,
3053 )
3054 .await;
3055
3056 let ToolExecutionResult::Success(value) = result else {
3057 panic!("spawn_background should succeed");
3058 };
3059 let run_id = value["run_id"].as_str().unwrap().to_string();
3060 let task_id = value["task_id"].as_str().unwrap().to_string();
3061
3062 tokio::time::timeout(std::time::Duration::from_secs(2), async {
3063 loop {
3064 if let Ok(Some(task)) = task_registry.get(session_id, &task_id).await
3065 && task.state == crate::session_task::SessionTaskState::Succeeded
3066 {
3067 break task;
3068 }
3069 tokio::time::sleep(std::time::Duration::from_millis(20)).await;
3070 }
3071 })
3072 .await
3073 .expect("background run should complete");
3074 let _ = run_id; let messages = platform_store.sent_messages.lock().unwrap().clone();
3077 assert_eq!(messages.len(), 1);
3078 assert!(messages[0].contains("Background run completed"));
3079
3080 let log_file = file_store
3081 .read_file(session_id, &format!("/.background/{run_id}/output.log"))
3082 .await
3083 .unwrap()
3084 .expect("log file");
3085 assert!(
3086 log_file
3087 .content
3088 .as_deref()
3089 .unwrap_or_default()
3090 .contains("hello from background")
3091 );
3092 }
3093
3094 #[tokio::test]
3095 async fn test_spawn_background_persists_failure_artifacts() {
3096 let session_id = SessionId::new();
3097 let file_store = Arc::new(TestFileStore::default());
3098 let storage_store = Arc::new(NoopStorageStore);
3099 let task_registry = Arc::new(InMemoryTaskRegistry::default());
3100 let tool_registry = ToolRegistry::builder()
3101 .tool(SpawnBackgroundTool)
3102 .tool(TestFailingBackgroundTool)
3103 .build();
3104
3105 let context = ToolContext::with_stores(session_id, file_store.clone(), storage_store)
3106 .with_tool_registry(Arc::new(tool_registry))
3107 .with_session_task_registry(task_registry.clone());
3108
3109 let result = SpawnBackgroundTool
3110 .execute_with_context(
3111 json!({
3112 "tool": "test_background_fail",
3113 "args": {}
3114 }),
3115 &context,
3116 )
3117 .await;
3118
3119 let ToolExecutionResult::Success(value) = result else {
3120 panic!("spawn_background should succeed");
3121 };
3122 let run_id = value["run_id"].as_str().unwrap().to_string();
3123 let task_id = value["task_id"].as_str().unwrap().to_string();
3124
3125 tokio::time::timeout(std::time::Duration::from_secs(2), async {
3126 loop {
3127 if let Ok(Some(task)) = task_registry.get(session_id, &task_id).await
3128 && task.state == crate::session_task::SessionTaskState::Failed
3129 {
3130 break task;
3131 }
3132 tokio::time::sleep(std::time::Duration::from_millis(20)).await;
3133 }
3134 })
3135 .await
3136 .expect("background run should fail");
3137 let _ = run_id;
3138
3139 let log_file = file_store
3140 .read_file(session_id, &format!("/.background/{run_id}/output.log"))
3141 .await
3142 .unwrap()
3143 .expect("log file");
3144 assert!(
3145 log_file
3146 .content
3147 .as_deref()
3148 .unwrap_or_default()
3149 .contains("background failed")
3150 );
3151
3152 let result_file = file_store
3153 .read_file(session_id, &format!("/.background/{run_id}/result.json"))
3154 .await
3155 .unwrap()
3156 .expect("result file");
3157 let result_json: Value =
3158 serde_json::from_str(result_file.content.as_deref().unwrap_or_default())
3159 .expect("valid json");
3160 assert_eq!(result_json["status"], "failed");
3161 assert_eq!(result_json["error"], "boom");
3162 }
3163
3164 #[tokio::test]
3165 async fn test_spawn_background_rejects_when_session_active_run_limit_reached() {
3166 let session_id = SessionId::new();
3167 let file_store = Arc::new(TestFileStore::default());
3168 let storage_store = Arc::new(NoopStorageStore);
3169 let task_registry = Arc::new(InMemoryTaskRegistry::default());
3170 let release = StdArc::new(AtomicBool::new(false));
3171 let tool_registry = ToolRegistry::builder()
3172 .tool(SpawnBackgroundTool)
3173 .tool(BlockingBackgroundTool {
3174 release: release.clone(),
3175 })
3176 .build();
3177
3178 let context = ToolContext::with_stores(session_id, file_store, storage_store)
3179 .with_tool_registry(Arc::new(tool_registry))
3180 .with_session_task_registry(task_registry.clone());
3181
3182 let mut task_ids = Vec::new();
3183 for _ in 0..MAX_ACTIVE_BACKGROUND_RUNS_PER_SESSION {
3184 let result = SpawnBackgroundTool
3185 .execute_with_context(
3186 json!({
3187 "tool": "test_background_blocking",
3188 "args": {}
3189 }),
3190 &context,
3191 )
3192 .await;
3193
3194 let ToolExecutionResult::Success(value) = result else {
3195 panic!("background run below the session limit should start");
3196 };
3197 task_ids.push(value["task_id"].as_str().unwrap().to_string());
3198 }
3199
3200 tokio::time::timeout(std::time::Duration::from_secs(2), async {
3202 loop {
3203 let running = task_registry
3204 .list(
3205 session_id,
3206 Some(&crate::session_task::SessionTaskFilter {
3207 kind: Some(crate::session_task::TASK_KIND_BACKGROUND_TOOL.to_string()),
3208 state: Some(crate::session_task::SessionTaskState::Running),
3209 }),
3210 )
3211 .await
3212 .unwrap();
3213 if running.len() == MAX_ACTIVE_BACKGROUND_RUNS_PER_SESSION {
3214 break;
3215 }
3216 tokio::time::sleep(std::time::Duration::from_millis(20)).await;
3217 }
3218 })
3219 .await
3220 .expect("background runs should become running");
3221
3222 let result = SpawnBackgroundTool
3223 .execute_with_context(
3224 json!({
3225 "tool": "test_background_blocking",
3226 "args": {}
3227 }),
3228 &context,
3229 )
3230 .await;
3231
3232 let ToolExecutionResult::ToolError(message) = result else {
3233 release.store(true, Ordering::SeqCst);
3234 panic!("spawn_background should reject once the session limit is reached");
3235 };
3236 assert!(message.contains("active background runs per session"));
3237
3238 release.store(true, Ordering::SeqCst);
3239 tokio::time::timeout(std::time::Duration::from_secs(2), async {
3240 for task_id in task_ids {
3241 loop {
3242 if let Ok(Some(task)) = task_registry.get(session_id, &task_id).await
3243 && task.state == crate::session_task::SessionTaskState::Succeeded
3244 {
3245 break;
3246 }
3247 tokio::time::sleep(std::time::Duration::from_millis(20)).await;
3248 }
3249 }
3250 })
3251 .await
3252 .expect("blocking background runs should complete after release");
3253
3254 tokio::time::timeout(std::time::Duration::from_secs(1), async {
3259 loop {
3260 if !has_session_background_permits(session_id) {
3261 break;
3262 }
3263 tokio::time::sleep(std::time::Duration::from_millis(5)).await;
3264 }
3265 })
3266 .await
3267 .expect("completed background runs should prune their per-session permit cache entry");
3268 }
3269
3270 #[tokio::test]
3271 async fn test_spawn_background_requires_task_registry() {
3272 let session_id = SessionId::new();
3273 let file_store = Arc::new(TestFileStore::default());
3274 let storage_store = Arc::new(NoopStorageStore);
3275 let tool_registry = ToolRegistry::builder()
3276 .tool(SpawnBackgroundTool)
3277 .tool(TestBackgroundTool)
3278 .build();
3279
3280 let context = ToolContext::with_stores(session_id, file_store, storage_store)
3282 .with_tool_registry(Arc::new(tool_registry));
3283
3284 let result = SpawnBackgroundTool
3285 .execute_with_context(
3286 json!({
3287 "tool": "test_background",
3288 "args": {}
3289 }),
3290 &context,
3291 )
3292 .await;
3293
3294 let ToolExecutionResult::ToolError(message) = result else {
3295 panic!("spawn_background should reject missing task registry");
3296 };
3297 assert!(message.contains("Session task registry not available"));
3298 }
3299
3300 #[tokio::test]
3301 async fn test_spawn_background_requires_file_store() {
3302 let session_id = SessionId::new();
3303 let storage_store = Arc::new(NoopStorageStore);
3304 let task_registry = Arc::new(InMemoryTaskRegistry::default());
3305 let tool_registry = ToolRegistry::builder()
3306 .tool(SpawnBackgroundTool)
3307 .tool(TestBackgroundTool)
3308 .build();
3309
3310 let context = ToolContext::with_storage_store(session_id, storage_store)
3312 .with_tool_registry(Arc::new(tool_registry))
3313 .with_session_task_registry(task_registry);
3314
3315 let result = SpawnBackgroundTool
3316 .execute_with_context(
3317 json!({
3318 "tool": "test_background",
3319 "args": {}
3320 }),
3321 &context,
3322 )
3323 .await;
3324
3325 let ToolExecutionResult::ToolError(message) = result else {
3326 panic!("spawn_background should reject missing file store");
3327 };
3328 assert!(message.contains("Session file store not available"));
3329 }
3330
3331 #[tokio::test]
3332 async fn test_spawn_background_caps_output_log_size() {
3333 let session_id = SessionId::new();
3334 let file_store = Arc::new(TestFileStore::default());
3335 let storage_store = Arc::new(NoopStorageStore);
3336 let task_registry = Arc::new(InMemoryTaskRegistry::default());
3337 let tool_registry = ToolRegistry::builder()
3338 .tool(SpawnBackgroundTool)
3339 .tool(TestLargeOutputBackgroundTool)
3340 .build();
3341
3342 let context = ToolContext::with_stores(session_id, file_store.clone(), storage_store)
3343 .with_tool_registry(Arc::new(tool_registry))
3344 .with_session_task_registry(task_registry.clone());
3345
3346 let result = SpawnBackgroundTool
3347 .execute_with_context(
3348 json!({
3349 "tool": "test_background_large_output",
3350 "args": {}
3351 }),
3352 &context,
3353 )
3354 .await;
3355
3356 let ToolExecutionResult::Success(value) = result else {
3357 panic!("spawn_background should succeed");
3358 };
3359 let run_id = value["run_id"].as_str().unwrap().to_string();
3360 let task_id = value["task_id"].as_str().unwrap().to_string();
3361
3362 tokio::time::timeout(std::time::Duration::from_secs(2), async {
3363 loop {
3364 if let Ok(Some(task)) = task_registry.get(session_id, &task_id).await
3365 && task.state == crate::session_task::SessionTaskState::Succeeded
3366 {
3367 break;
3368 }
3369 tokio::time::sleep(std::time::Duration::from_millis(20)).await;
3370 }
3371 })
3372 .await
3373 .expect("background run should complete");
3374 let _ = run_id;
3375
3376 let log_content = file_store
3377 .read_file(session_id, &format!("/.background/{run_id}/output.log"))
3378 .await
3379 .unwrap()
3380 .expect("log file")
3381 .content
3382 .unwrap_or_default();
3383
3384 assert!(log_content.contains("[system] background output truncated"));
3385 assert!(log_content.chars().count() <= MAX_BACKGROUND_OUTPUT_LOG_CHARS + 128);
3386 }
3387
3388 #[tokio::test]
3389 async fn test_spawn_background_can_create_scheduled_monitor() {
3390 let session_id = SessionId::new();
3391 let schedule_store = Arc::new(TestScheduleStore::default());
3392 let storage_store = Arc::new(NoopStorageStore);
3393 let tool_registry = ToolRegistry::builder()
3394 .tool(SpawnBackgroundTool)
3395 .tool(TestBackgroundTool)
3396 .build();
3397
3398 let context = ToolContext::with_storage_store(session_id, storage_store)
3399 .with_tool_registry(Arc::new(tool_registry))
3400 .with_schedule_store(schedule_store.clone());
3401
3402 let result = SpawnBackgroundTool
3403 .execute_with_context(
3404 json!({
3405 "tool": "test_background",
3406 "title": "Watch PR 1319",
3407 "args": { "summary": "Background complete" },
3408 "schedule": {
3409 "cron_expression": "*/10 * * * *",
3410 "timezone": "America/Chicago"
3411 }
3412 }),
3413 &context,
3414 )
3415 .await;
3416
3417 let ToolExecutionResult::Success(value) = result else {
3418 panic!("spawn_background should create a schedule: {result:?}");
3419 };
3420
3421 assert_eq!(value["status"], "scheduled");
3422 assert_eq!(value["title"], "Watch PR 1319");
3423 assert_eq!(value["cron_expression"], "*/10 * * * *");
3424 assert_eq!(value["timezone"], "America/Chicago");
3425
3426 let schedules = schedule_store.list_schedules(session_id).await.unwrap();
3427 assert_eq!(schedules.len(), 1);
3428 assert_eq!(
3429 schedules[0].cron_expression.as_deref(),
3430 Some("*/10 * * * *")
3431 );
3432 assert!(schedules[0].description.contains("Monitor: Watch PR 1319"));
3433 assert!(
3434 schedules[0]
3435 .description
3436 .contains("\"summary\": \"Background complete\"")
3437 );
3438 }
3439
3440 #[tokio::test]
3441 async fn test_spawn_background_rejects_invalid_scheduled_at() {
3442 let session_id = SessionId::new();
3443 let storage_store = Arc::new(NoopStorageStore);
3444 let tool_registry = ToolRegistry::builder()
3445 .tool(SpawnBackgroundTool)
3446 .tool(TestBackgroundTool)
3447 .build();
3448 let context = ToolContext::with_storage_store(session_id, storage_store)
3449 .with_tool_registry(Arc::new(tool_registry));
3450
3451 let result = SpawnBackgroundTool
3452 .execute_with_context(
3453 json!({
3454 "tool": "test_background",
3455 "args": {},
3456 "schedule": {
3457 "scheduled_at": "tomorrow at noon"
3458 }
3459 }),
3460 &context,
3461 )
3462 .await;
3463
3464 let ToolExecutionResult::ToolError(message) = result else {
3465 panic!("spawn_background should reject invalid scheduled_at");
3466 };
3467 assert!(message.contains("scheduled_at must be RFC3339"));
3468 }
3469
3470 #[tokio::test]
3471 async fn test_spawn_background_rejects_ambiguous_schedule_shape() {
3472 let session_id = SessionId::new();
3473 let storage_store = Arc::new(NoopStorageStore);
3474 let tool_registry = ToolRegistry::builder()
3475 .tool(SpawnBackgroundTool)
3476 .tool(TestBackgroundTool)
3477 .build();
3478 let context = ToolContext::with_storage_store(session_id, storage_store)
3479 .with_tool_registry(Arc::new(tool_registry));
3480
3481 let result = SpawnBackgroundTool
3482 .execute_with_context(
3483 json!({
3484 "tool": "test_background",
3485 "args": {},
3486 "schedule": {
3487 "cron_expression": "*/10 * * * *",
3488 "scheduled_at": "2026-04-16T15:30:00Z"
3489 }
3490 }),
3491 &context,
3492 )
3493 .await;
3494
3495 let ToolExecutionResult::ToolError(message) = result else {
3496 panic!("spawn_background should reject ambiguous schedule shape");
3497 };
3498 assert!(message.contains("must not include both cron_expression and scheduled_at"));
3499 }
3500
3501 #[test]
3506 fn test_is_canceled_outcome_detects_sentinel() {
3507 let sentinel: std::result::Result<BackgroundOutcome, ToolExecutionResult> = Err(
3509 ToolExecutionResult::ToolError(BACKGROUND_CANCEL_SENTINEL.to_string()),
3510 );
3511 assert!(is_canceled_outcome(&sentinel));
3512 }
3513
3514 #[test]
3515 fn test_is_canceled_outcome_does_not_match_other_errors() {
3516 let other_err: std::result::Result<BackgroundOutcome, ToolExecutionResult> =
3517 Err(ToolExecutionResult::ToolError("boom".to_string()));
3518 assert!(!is_canceled_outcome(&other_err));
3519
3520 let success: std::result::Result<BackgroundOutcome, ToolExecutionResult> =
3521 Ok(BackgroundOutcome {
3522 summary: "ok".to_string(),
3523 result: json!({"ok": true}),
3524 raw_output: None,
3525 });
3526 assert!(!is_canceled_outcome(&success));
3527 }
3528
3529 #[derive(Default)]
3532 struct SleepingBackgroundTool;
3533
3534 #[async_trait]
3535 impl BackgroundExecutableTool for SleepingBackgroundTool {
3536 async fn execute_background(
3537 &self,
3538 _arguments: Value,
3539 _context: ToolContext,
3540 sink: Arc<dyn BackgroundEventSink>,
3541 ) -> std::result::Result<BackgroundOutcome, ToolExecutionResult> {
3542 sink.status("Sleeping forever")
3543 .await
3544 .map_err(ToolExecutionResult::internal_error)?;
3545 tokio::time::sleep(std::time::Duration::from_secs(3600)).await;
3547 Ok(BackgroundOutcome {
3548 summary: "should not reach here".to_string(),
3549 result: json!({}),
3550 raw_output: None,
3551 })
3552 }
3553 }
3554
3555 #[async_trait]
3556 impl Tool for SleepingBackgroundTool {
3557 fn name(&self) -> &str {
3558 "test_background_sleeping"
3559 }
3560
3561 fn display_name(&self) -> Option<&str> {
3562 Some("Test Background Sleeping")
3563 }
3564
3565 fn description(&self) -> &str {
3566 "background test tool that sleeps indefinitely"
3567 }
3568
3569 fn parameters_schema(&self) -> Value {
3570 json!({
3571 "type": "object",
3572 "properties": {}
3573 })
3574 }
3575
3576 async fn execute(&self, _arguments: Value) -> ToolExecutionResult {
3577 ToolExecutionResult::tool_error("foreground unsupported")
3578 }
3579
3580 fn hints(&self) -> ToolHints {
3581 ToolHints::default().with_supports_background(true)
3582 }
3583
3584 fn as_background_executable(&self) -> Option<&dyn BackgroundExecutableTool> {
3585 Some(self)
3586 }
3587 }
3588
3589 #[derive(Default)]
3593 struct InMemoryTaskRegistry {
3594 tasks: Mutex<HashMap<String, crate::session_task::SessionTask>>,
3595 }
3596
3597 #[async_trait]
3598 impl crate::session_task::SessionTaskRegistry for InMemoryTaskRegistry {
3599 async fn create(
3600 &self,
3601 input: crate::session_task::CreateSessionTask,
3602 ) -> crate::Result<crate::session_task::SessionTask> {
3603 let mut tasks = self.tasks.lock().unwrap();
3604 if let Some(id) = &input.id
3605 && let Some(existing) = tasks.get(id)
3606 {
3607 return Ok(existing.clone());
3608 }
3609 let task = crate::session_task::new_session_task(input, chrono::Utc::now());
3610 tasks.insert(task.id.clone(), task.clone());
3611 Ok(task)
3612 }
3613
3614 async fn update(
3615 &self,
3616 _session_id: SessionId,
3617 task_id: &str,
3618 update: crate::session_task::SessionTaskUpdate,
3619 ) -> crate::Result<Option<crate::session_task::SessionTask>> {
3620 let mut tasks = self.tasks.lock().unwrap();
3621 let Some(task) = tasks.get_mut(task_id) else {
3622 return Ok(None);
3623 };
3624 crate::session_task::apply_task_update(task, update, chrono::Utc::now());
3625 Ok(Some(task.clone()))
3626 }
3627
3628 async fn get(
3629 &self,
3630 _session_id: SessionId,
3631 task_id: &str,
3632 ) -> crate::Result<Option<crate::session_task::SessionTask>> {
3633 Ok(self.tasks.lock().unwrap().get(task_id).cloned())
3634 }
3635
3636 async fn list(
3637 &self,
3638 _session_id: SessionId,
3639 filter: Option<&crate::session_task::SessionTaskFilter>,
3640 ) -> crate::Result<Vec<crate::session_task::SessionTask>> {
3641 let tasks = self.tasks.lock().unwrap();
3642 Ok(tasks
3643 .values()
3644 .filter(|task| {
3645 filter.is_none_or(|f| {
3646 f.kind.as_deref().is_none_or(|kind| task.kind == kind)
3647 && f.state.is_none_or(|state| task.state == state)
3648 })
3649 })
3650 .cloned()
3651 .collect())
3652 }
3653
3654 async fn request_cancel(
3655 &self,
3656 _session_id: SessionId,
3657 task_id: &str,
3658 ) -> crate::Result<Option<crate::session_task::SessionTask>> {
3659 let mut tasks = self.tasks.lock().unwrap();
3660 let Some(task) = tasks.get_mut(task_id) else {
3661 return Ok(None);
3662 };
3663 task.cancel_requested_at
3664 .get_or_insert_with(chrono::Utc::now);
3665 task.updated_at = chrono::Utc::now();
3666 Ok(Some(task.clone()))
3667 }
3668
3669 async fn record_message(
3670 &self,
3671 _session_id: SessionId,
3672 task_id: &str,
3673 message: crate::session_task::NewTaskMessage,
3674 ) -> crate::Result<crate::session_task::TaskMessage> {
3675 let tasks = self.tasks.lock().unwrap();
3676 let _task = tasks
3677 .get(task_id)
3678 .ok_or_else(|| crate::AgentLoopError::tool(format!("no task {task_id}")))?;
3679 Ok(crate::session_task::TaskMessage {
3680 id: crate::session_task::generate_task_message_id(),
3681 task_id: task_id.to_string(),
3682 direction: message.direction,
3683 content: message.content,
3684 in_reply_to: message.in_reply_to,
3685 created_at: chrono::Utc::now(),
3686 })
3687 }
3688
3689 async fn list_messages(
3690 &self,
3691 _session_id: SessionId,
3692 _task_id: &str,
3693 _limit: Option<u32>,
3694 _after_id: Option<&str>,
3695 ) -> crate::Result<Vec<crate::session_task::TaskMessage>> {
3696 Ok(Vec::new())
3697 }
3698 }
3699
3700 #[tokio::test]
3703 async fn test_cancel_background_run_via_task_registry() {
3704 let session_id = SessionId::new();
3705 let file_store = Arc::new(TestFileStore::default());
3706 let storage_store = Arc::new(NoopStorageStore);
3707 let task_registry = Arc::new(InMemoryTaskRegistry::default());
3708
3709 let tool_registry = ToolRegistry::builder()
3710 .tool(SpawnBackgroundTool)
3711 .tool(SleepingBackgroundTool)
3712 .build();
3713
3714 let context = ToolContext::with_stores(session_id, file_store.clone(), storage_store)
3715 .with_tool_registry(Arc::new(tool_registry))
3716 .with_session_task_registry(task_registry.clone());
3717
3718 let result = SpawnBackgroundTool
3719 .execute_with_context(
3720 json!({
3721 "tool": "test_background_sleeping",
3722 "args": {},
3723 "signal_on_completion": false
3724 }),
3725 &context,
3726 )
3727 .await;
3728
3729 let ToolExecutionResult::Success(value) = result else {
3730 panic!("spawn_background should succeed");
3731 };
3732 let run_id = value["run_id"].as_str().unwrap().to_string();
3733 let task_id = value["task_id"].as_str().unwrap().to_string();
3734
3735 tokio::time::timeout(std::time::Duration::from_secs(5), async {
3737 loop {
3738 if let Ok(Some(task)) = task_registry.get(session_id, &task_id).await
3740 && task.heartbeat_at.is_some()
3741 {
3742 break;
3743 }
3744 tokio::time::sleep(std::time::Duration::from_millis(50)).await;
3745 }
3746 })
3747 .await
3748 .expect("background run should start and send at least one heartbeat");
3749
3750 task_registry
3752 .request_cancel(session_id, &task_id)
3753 .await
3754 .expect("request_cancel should succeed");
3755
3756 tokio::time::timeout(std::time::Duration::from_secs(10), async {
3758 loop {
3759 if let Ok(Some(task)) = task_registry.get(session_id, &task_id).await
3760 && task.state == crate::session_task::SessionTaskState::Canceled
3761 {
3762 break task;
3763 }
3764 tokio::time::sleep(std::time::Duration::from_millis(50)).await;
3765 }
3766 })
3767 .await
3768 .expect("background task should reach Canceled state");
3769
3770 let result_file = file_store
3772 .read_file(session_id, &format!("/.background/{run_id}/result.json"))
3773 .await
3774 .unwrap()
3775 .expect("result.json should exist");
3776 let result_json: Value =
3777 serde_json::from_str(result_file.content.as_deref().unwrap_or_default())
3778 .expect("valid json");
3779 assert_eq!(result_json["status"], "canceled");
3780
3781 let log_file = file_store
3782 .read_file(session_id, &format!("/.background/{run_id}/output.log"))
3783 .await
3784 .unwrap()
3785 .expect("output.log should exist");
3786 assert!(
3787 log_file
3788 .content
3789 .as_deref()
3790 .unwrap_or_default()
3791 .contains("Canceled by request.")
3792 );
3793 }
3794
3795 fn make_reattach_task(spec: serde_json::Value) -> crate::session_task::SessionTask {
3800 use crate::session_task::{SessionTaskState, TaskLinks, TaskWakePolicy};
3801 crate::session_task::SessionTask {
3802 id: "t-reattach".to_string(),
3803 session_id: SessionId::new(),
3804 root_session_id: None,
3805 kind: crate::session_task::TASK_KIND_BACKGROUND_TOOL.to_string(),
3806 display_name: "Reattach test".to_string(),
3807 spec,
3808 state: SessionTaskState::Running,
3809 state_detail: None,
3810 progress: None,
3811 input_request: None,
3812 cancel_requested_at: None,
3813 summary: None,
3814 result_path: None,
3815 artifacts: vec![],
3816 error: None,
3817 attempt: 2,
3818 worker_id: None,
3819 heartbeat_at: None,
3820 links: TaskLinks::default(),
3821 wake_policy: TaskWakePolicy::Silent,
3822 created_at: chrono::Utc::now(),
3823 started_at: None,
3824 finished_at: None,
3825 updated_at: chrono::Utc::now(),
3826 }
3827 }
3828
3829 #[tokio::test]
3830 async fn reattach_fails_with_missing_file_store() {
3831 let session_id = SessionId::new();
3832 let task_registry = Arc::new(InMemoryTaskRegistry::default());
3834 let context =
3835 crate::traits::ToolContext::new(session_id).with_session_task_registry(task_registry);
3836 let task = make_reattach_task(serde_json::json!({
3837 "tool": "get_current_time",
3838 "arguments": {},
3839 "reattachable": true,
3840 "signal_on_completion": true,
3841 }));
3842 let err = reattach_background_run(&task, &context)
3843 .await
3844 .expect_err("should fail without file store");
3845 assert!(
3846 err.to_string().contains("file store"),
3847 "error should mention file store, got: {err}"
3848 );
3849 }
3850
3851 #[tokio::test]
3852 async fn reattach_fails_with_missing_task_registry() {
3853 let session_id = SessionId::new();
3854 let file_store = Arc::new(TestFileStore::default());
3855 let storage_store = Arc::new(NoopStorageStore);
3856 let context =
3858 crate::traits::ToolContext::with_stores(session_id, file_store, storage_store);
3859 let task = make_reattach_task(serde_json::json!({
3860 "tool": "get_current_time",
3861 "arguments": {},
3862 "reattachable": true,
3863 "signal_on_completion": true,
3864 }));
3865 let err = reattach_background_run(&task, &context)
3866 .await
3867 .expect_err("should fail without task registry");
3868 assert!(
3869 err.to_string().contains("task registry"),
3870 "error should mention task registry, got: {err}"
3871 );
3872 }
3873
3874 #[tokio::test]
3875 async fn reattach_fails_with_unknown_tool_name() {
3876 let session_id = SessionId::new();
3877 let file_store = Arc::new(TestFileStore::default());
3878 let storage_store = Arc::new(NoopStorageStore);
3879 let task_registry = Arc::new(InMemoryTaskRegistry::default());
3880 let context =
3881 crate::traits::ToolContext::with_stores(session_id, file_store, storage_store)
3882 .with_session_task_registry(task_registry);
3883 let task = make_reattach_task(serde_json::json!({
3885 "tool": "test_background",
3886 "arguments": {},
3887 "reattachable": true,
3888 "signal_on_completion": true,
3889 }));
3890 let err = reattach_background_run(&task, &context)
3891 .await
3892 .expect_err("should fail for unknown tool");
3893 assert!(
3894 err.to_string().contains("not found in built-in registry"),
3895 "error should mention built-in registry, got: {err}"
3896 );
3897 }
3898
3899 #[tokio::test]
3900 async fn reattach_fails_with_missing_tool_spec_field() {
3901 let session_id = SessionId::new();
3902 let file_store = Arc::new(TestFileStore::default());
3903 let storage_store = Arc::new(NoopStorageStore);
3904 let task_registry = Arc::new(InMemoryTaskRegistry::default());
3905 let context =
3906 crate::traits::ToolContext::with_stores(session_id, file_store, storage_store)
3907 .with_session_task_registry(task_registry);
3908 let task = make_reattach_task(serde_json::json!({ "reattachable": true }));
3910 let err = reattach_background_run(&task, &context)
3911 .await
3912 .expect_err("should fail with missing tool field");
3913 assert!(
3914 err.to_string().contains("missing 'tool' field"),
3915 "error should mention missing tool field, got: {err}"
3916 );
3917 }
3918}