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.platform_store 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::platform_store::PlatformStore;
2115 use crate::session_file::{FileInfo, FileStat, SessionFile};
2116 use crate::session_task::SessionTaskRegistry;
2117 use crate::traits::{SessionFileSystem, SessionScheduleStore};
2118 use crate::typed_id::{HarnessId, SessionId};
2119 use crate::{AgentId, KeyInfo, PlatformMessage, 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 TestPlatformStore {
2495 sent_messages: Mutex<Vec<String>>,
2496 }
2497
2498 #[async_trait]
2499 impl PlatformStore for TestPlatformStore {
2500 async fn list_harnesses(&self) -> crate::Result<Vec<crate::Harness>> {
2501 Ok(Vec::new())
2502 }
2503 async fn get_harness(&self, _id: HarnessId) -> crate::Result<Option<crate::Harness>> {
2504 Ok(None)
2505 }
2506 async fn create_harness(
2507 &self,
2508 _name: &str,
2509 _display_name: Option<&str>,
2510 _description: Option<&str>,
2511 _system_prompt: Option<&str>,
2512 _parent_harness_id: Option<HarnessId>,
2513 _capabilities: &[String],
2514 ) -> crate::Result<crate::Harness> {
2515 unreachable!()
2516 }
2517 async fn update_harness(
2518 &self,
2519 _id: HarnessId,
2520 _name: Option<&str>,
2521 _display_name: Option<&str>,
2522 _description: Option<&str>,
2523 _system_prompt: Option<&str>,
2524 _parent_harness_id: Option<Option<HarnessId>>,
2525 ) -> crate::Result<crate::Harness> {
2526 unreachable!()
2527 }
2528 async fn delete_harness(&self, _id: HarnessId) -> crate::Result<()> {
2529 Ok(())
2530 }
2531 async fn copy_harness(
2532 &self,
2533 _id: HarnessId,
2534 _new_name: Option<&str>,
2535 ) -> crate::Result<crate::Harness> {
2536 unreachable!()
2537 }
2538 async fn list_agents(&self) -> crate::Result<Vec<crate::Agent>> {
2539 Ok(Vec::new())
2540 }
2541 async fn get_agent_by_id(&self, _id: AgentId) -> crate::Result<Option<crate::Agent>> {
2542 Ok(None)
2543 }
2544 async fn create_agent(
2545 &self,
2546 _name: &str,
2547 _display_name: Option<&str>,
2548 _description: Option<&str>,
2549 _system_prompt: &str,
2550 _capabilities: &[String],
2551 ) -> crate::Result<crate::Agent> {
2552 unreachable!()
2553 }
2554 async fn update_agent(
2555 &self,
2556 _id: AgentId,
2557 _name: Option<&str>,
2558 _display_name: Option<&str>,
2559 _description: Option<&str>,
2560 _system_prompt: Option<&str>,
2561 ) -> crate::Result<crate::Agent> {
2562 unreachable!()
2563 }
2564 async fn delete_agent(&self, _id: AgentId) -> crate::Result<()> {
2565 Ok(())
2566 }
2567 async fn list_apps(
2568 &self,
2569 _search: Option<&str>,
2570 _include_archived: bool,
2571 ) -> crate::Result<Vec<crate::App>> {
2572 Ok(Vec::new())
2573 }
2574 async fn get_app(&self, _id: crate::AppId) -> crate::Result<Option<crate::App>> {
2575 Ok(None)
2576 }
2577 async fn create_app(
2578 &self,
2579 _name: &str,
2580 _description: Option<&str>,
2581 _harness_id: HarnessId,
2582 _agent_id: Option<AgentId>,
2583 _agent_identity_id: Option<crate::AgentIdentityId>,
2584 _channel_type: Option<crate::ChannelType>,
2585 _channel_config: Option<&serde_json::Value>,
2586 ) -> crate::Result<crate::App> {
2587 unreachable!()
2588 }
2589 async fn update_app(
2590 &self,
2591 _id: crate::AppId,
2592 _name: Option<&str>,
2593 _description: Option<&str>,
2594 _harness_id: Option<HarnessId>,
2595 _agent_id: Option<AgentId>,
2596 _agent_identity_id: Option<Option<crate::AgentIdentityId>>,
2597 ) -> crate::Result<crate::App> {
2598 unreachable!()
2599 }
2600 async fn delete_app(&self, _id: crate::AppId) -> crate::Result<()> {
2601 Ok(())
2602 }
2603 async fn destroy_app(&self, _id: crate::AppId) -> crate::Result<()> {
2604 Ok(())
2605 }
2606 async fn publish_app(&self, _id: crate::AppId) -> crate::Result<crate::App> {
2607 unreachable!()
2608 }
2609 async fn unpublish_app(&self, _id: crate::AppId) -> crate::Result<crate::App> {
2610 unreachable!()
2611 }
2612 async fn add_app_channel(
2613 &self,
2614 _app_id: crate::AppId,
2615 _channel_type: crate::ChannelType,
2616 _channel_config: Option<&serde_json::Value>,
2617 _enabled: Option<bool>,
2618 ) -> crate::Result<crate::AppChannel> {
2619 unreachable!()
2620 }
2621 async fn update_app_channel(
2622 &self,
2623 _app_id: crate::AppId,
2624 _channel_id: crate::AppChannelId,
2625 _channel_type: Option<crate::ChannelType>,
2626 _channel_config: Option<&serde_json::Value>,
2627 _enabled: Option<bool>,
2628 ) -> crate::Result<crate::AppChannel> {
2629 unreachable!()
2630 }
2631 async fn delete_app_channel(
2632 &self,
2633 _app_id: crate::AppId,
2634 _channel_id: crate::AppChannelId,
2635 ) -> crate::Result<()> {
2636 Ok(())
2637 }
2638 async fn list_sessions(
2639 &self,
2640 _limit: Option<usize>,
2641 _agent_id: Option<AgentId>,
2642 ) -> crate::Result<Vec<crate::Session>> {
2643 Ok(Vec::new())
2644 }
2645 async fn create_session(
2646 &self,
2647 _harness_id: HarnessId,
2648 _agent_id: Option<AgentId>,
2649 _title: Option<&str>,
2650 _locale: Option<&str>,
2651 _blueprint_id: Option<&str>,
2652 _blueprint_config: Option<&serde_json::Value>,
2653 _parent_session_id: Option<SessionId>,
2654 ) -> crate::Result<crate::Session> {
2655 unreachable!()
2656 }
2657 async fn get_session_by_id(&self, _id: SessionId) -> crate::Result<Option<crate::Session>> {
2658 Ok(None)
2659 }
2660 async fn add_agent_session_participant(
2661 &self,
2662 _session_id: SessionId,
2663 _agent_id: AgentId,
2664 ) -> crate::Result<crate::SessionParticipant> {
2665 unreachable!()
2666 }
2667 async fn get_session_context_report(
2668 &self,
2669 id: SessionId,
2670 ) -> crate::Result<crate::SessionContextReport> {
2671 Ok(crate::SessionContextReport {
2672 session_id: id.to_string(),
2673 model: "llmsim".to_string(),
2674 context_window_tokens: None,
2675 estimated_input_tokens: 0,
2676 sections: vec![],
2677 contributions: vec![],
2678 cumulative_usage: None,
2679 })
2680 }
2681 async fn delete_session(&self, _id: SessionId) -> crate::Result<()> {
2682 Ok(())
2683 }
2684 async fn send_message(&self, _session_id: SessionId, content: &str) -> crate::Result<()> {
2685 self.sent_messages.lock().unwrap().push(content.to_string());
2686 Ok(())
2687 }
2688 async fn get_messages(
2689 &self,
2690 _session_id: SessionId,
2691 _limit: Option<usize>,
2692 ) -> crate::Result<Vec<PlatformMessage>> {
2693 Ok(Vec::new())
2694 }
2695 async fn wait_for_idle(
2696 &self,
2697 _session_id: SessionId,
2698 _timeout_secs: Option<u64>,
2699 ) -> crate::Result<String> {
2700 Ok("idle".to_string())
2701 }
2702 async fn list_capabilities(
2703 &self,
2704 _search: Option<&str>,
2705 ) -> crate::Result<Vec<crate::CapabilityInfo>> {
2706 Ok(Vec::new())
2707 }
2708 fn base_url(&self) -> &str {
2709 "http://localhost:9300"
2710 }
2711 }
2712
2713 #[derive(Default)]
2714 struct NoopStorageStore;
2715
2716 #[derive(Default)]
2717 struct TestScheduleStore {
2718 schedules: Mutex<Vec<crate::session_schedule::SessionSchedule>>,
2719 }
2720
2721 #[async_trait]
2722 impl crate::traits::SessionScheduleStore for TestScheduleStore {
2723 async fn create_schedule(
2724 &self,
2725 session_id: SessionId,
2726 description: String,
2727 cron_expression: Option<String>,
2728 scheduled_at: Option<chrono::DateTime<chrono::Utc>>,
2729 timezone: String,
2730 ) -> crate::Result<crate::session_schedule::SessionSchedule> {
2731 let schedule = crate::session_schedule::SessionSchedule {
2732 id: crate::typed_id::ScheduleId::new(),
2733 session_id,
2734 owner_principal_id: crate::PrincipalId::from_seed(1),
2735 resolved_owner_user_id: None,
2736 owner: None,
2737 effective_owner: None,
2738 description,
2739 cron_expression: cron_expression.clone(),
2740 scheduled_at,
2741 timezone,
2742 enabled: true,
2743 schedule_type: crate::session_schedule::SessionSchedule::derive_type(
2744 &cron_expression,
2745 ),
2746 next_trigger_at: Some(chrono::Utc::now() + chrono::Duration::minutes(10)),
2747 last_triggered_at: None,
2748 trigger_count: 0,
2749 created_at: chrono::Utc::now(),
2750 updated_at: chrono::Utc::now(),
2751 };
2752 self.schedules.lock().unwrap().push(schedule.clone());
2753 Ok(schedule)
2754 }
2755
2756 async fn cancel_schedule(
2757 &self,
2758 _session_id: SessionId,
2759 schedule_id: crate::ScheduleId,
2760 ) -> crate::Result<crate::session_schedule::SessionSchedule> {
2761 let mut schedules = self.schedules.lock().unwrap();
2762 let schedule = schedules
2763 .iter_mut()
2764 .find(|schedule| schedule.id == schedule_id)
2765 .ok_or_else(|| crate::AgentLoopError::tool("Schedule not found".to_string()))?;
2766 schedule.enabled = false;
2767 Ok(schedule.clone())
2768 }
2769
2770 async fn list_schedules(
2771 &self,
2772 session_id: SessionId,
2773 ) -> crate::Result<Vec<crate::session_schedule::SessionSchedule>> {
2774 Ok(self
2775 .schedules
2776 .lock()
2777 .unwrap()
2778 .iter()
2779 .filter(|schedule| schedule.session_id == session_id)
2780 .cloned()
2781 .collect())
2782 }
2783
2784 async fn count_active_schedules(&self, session_id: SessionId) -> crate::Result<u32> {
2785 Ok(self
2786 .schedules
2787 .lock()
2788 .unwrap()
2789 .iter()
2790 .filter(|schedule| schedule.session_id == session_id && schedule.enabled)
2791 .count() as u32)
2792 }
2793
2794 async fn count_active_org_schedules(&self) -> crate::Result<u32> {
2795 Ok(self
2797 .schedules
2798 .lock()
2799 .unwrap()
2800 .iter()
2801 .filter(|schedule| schedule.enabled)
2802 .count() as u32)
2803 }
2804 }
2805
2806 #[async_trait]
2807 impl crate::traits::SessionStorageStore for NoopStorageStore {
2808 async fn set_value(
2809 &self,
2810 _session_id: SessionId,
2811 _key: &str,
2812 _value: &str,
2813 ) -> crate::Result<()> {
2814 Ok(())
2815 }
2816 async fn get_value(
2817 &self,
2818 _session_id: SessionId,
2819 _key: &str,
2820 ) -> crate::Result<Option<String>> {
2821 Ok(None)
2822 }
2823 async fn delete_value(&self, _session_id: SessionId, _key: &str) -> crate::Result<bool> {
2824 Ok(false)
2825 }
2826 async fn list_keys(&self, _session_id: SessionId) -> crate::Result<Vec<KeyInfo>> {
2827 Ok(Vec::new())
2828 }
2829 async fn set_secret(
2830 &self,
2831 _session_id: SessionId,
2832 _name: &str,
2833 _value: &str,
2834 ) -> crate::Result<()> {
2835 Ok(())
2836 }
2837 async fn get_secret(
2838 &self,
2839 _session_id: SessionId,
2840 _name: &str,
2841 ) -> crate::Result<Option<String>> {
2842 Ok(None)
2843 }
2844 async fn delete_secret(&self, _session_id: SessionId, _name: &str) -> crate::Result<bool> {
2845 Ok(false)
2846 }
2847 async fn list_secrets(&self, _session_id: SessionId) -> crate::Result<Vec<SecretInfo>> {
2848 Ok(Vec::new())
2849 }
2850 }
2851
2852 #[tokio::test]
2853 async fn test_echo_tool() {
2854 let tool = EchoTool;
2855
2856 let result = tool
2857 .execute(serde_json::json!({"message": "Hello, world!"}))
2858 .await;
2859
2860 if let ToolExecutionResult::Success(value) = result {
2861 assert_eq!(
2862 value.get("echoed").unwrap().as_str().unwrap(),
2863 "Hello, world!"
2864 );
2865 assert_eq!(value.get("length").unwrap().as_u64().unwrap(), 13);
2866 } else {
2867 panic!("Expected success");
2868 }
2869 }
2870
2871 #[tokio::test]
2872 async fn test_failing_tool_with_tool_error() {
2873 let tool = FailingTool::with_tool_error("Something went wrong");
2874
2875 let result = tool.execute(serde_json::json!({})).await;
2876
2877 if let ToolExecutionResult::ToolError(msg) = result {
2878 assert_eq!(msg, "Something went wrong");
2879 } else {
2880 panic!("Expected tool error");
2881 }
2882 }
2883
2884 #[tokio::test]
2885 async fn test_failing_tool_with_internal_error() {
2886 let tool = FailingTool::with_internal_error("Database connection failed");
2887
2888 let result = tool.execute(serde_json::json!({})).await;
2889
2890 if let ToolExecutionResult::InternalError(err) = result {
2891 assert_eq!(err.message, "Database connection failed");
2892 } else {
2893 panic!("Expected internal error");
2894 }
2895 }
2896
2897 #[tokio::test]
2898 async fn test_tool_result_conversion() {
2899 let result = ToolExecutionResult::success(serde_json::json!({"value": 42}));
2901 let tool_result = result.into_tool_result("call_1", "test_tool");
2902 assert!(tool_result.error.is_none());
2903 assert_eq!(tool_result.result.unwrap()["value"], 42);
2904
2905 let result = ToolExecutionResult::tool_error("Invalid input");
2907 let tool_result = result.into_tool_result("call_2", "test_tool");
2908 assert_eq!(tool_result.error.as_deref(), Some("Invalid input"));
2909 assert_eq!(
2910 tool_result.result.unwrap(),
2911 serde_json::json!({"error": "Invalid input"})
2912 );
2913
2914 let result = ToolExecutionResult::internal_error_msg("Secret database error");
2916 let tool_result = result.into_tool_result("call_3", "test_tool");
2917 assert_eq!(
2918 tool_result.error.as_deref(),
2919 Some("An internal error occurred while executing the tool")
2920 );
2921 assert_eq!(
2922 tool_result.result.unwrap(),
2923 serde_json::json!({"error": "An internal error occurred while executing the tool"})
2924 );
2925 }
2926
2927 #[tokio::test]
2928 async fn test_tool_registry() {
2929 let mut registry = ToolRegistry::new();
2930 registry.register(GetCurrentTimeTool);
2931 registry.register(EchoTool);
2932
2933 assert_eq!(registry.len(), 2);
2934 assert!(registry.has("get_current_time"));
2935 assert!(registry.has("echo"));
2936 assert!(!registry.has("nonexistent"));
2937
2938 let definitions = registry.tool_definitions();
2939 assert_eq!(definitions.len(), 2);
2940 }
2941
2942 #[tokio::test]
2943 async fn test_tool_registry_builder() {
2944 let registry = ToolRegistry::builder()
2945 .tool(GetCurrentTimeTool)
2946 .tool(EchoTool)
2947 .build();
2948
2949 assert_eq!(registry.len(), 2);
2950 }
2951
2952 #[test]
2953 fn test_tool_display_name_in_definition() {
2954 let tool = GetCurrentTimeTool;
2956 assert_eq!(tool.display_name(), Some("Get Current Time"));
2957
2958 let def = tool.to_definition();
2959 assert_eq!(def.display_name(), Some("Get Current Time"));
2960 }
2961
2962 #[test]
2963 fn test_success_with_raw_output_object_preserves_shape() {
2964 let res = ToolExecutionResult::success_with_raw_output(
2965 serde_json::json!({"stdout": "hello"}),
2966 "raw stdout bytes".to_string(),
2967 );
2968 let tr = res.into_tool_result("call_1", "demo");
2969 assert_eq!(tr.result.as_ref().unwrap()["stdout"], "hello");
2970 assert!(
2971 tr.result
2972 .as_ref()
2973 .unwrap()
2974 .as_object()
2975 .unwrap()
2976 .get("_raw_output")
2977 .is_none(),
2978 "sidecar key must not leak to the LLM-visible result"
2979 );
2980 assert_eq!(tr.raw_output.as_deref(), Some("raw stdout bytes"));
2981 }
2982
2983 #[test]
2984 fn test_success_with_raw_output_scalar_unwraps_to_string() {
2985 let res = ToolExecutionResult::success_with_raw_output(
2986 "compact summary".to_string(),
2987 "full output bytes".to_string(),
2988 );
2989 let tr = res.into_tool_result("call_1", "demo");
2990 assert_eq!(
2991 tr.result,
2992 Some(serde_json::Value::String("compact summary".into()))
2993 );
2994 assert_eq!(tr.raw_output.as_deref(), Some("full output bytes"));
2995 }
2996
2997 #[test]
2998 fn test_success_result_with_raw_output_scalar_key_is_not_unwrapped() {
2999 let res = ToolExecutionResult::success(
3000 serde_json::json!({"_raw_output_scalar": "user_value", "kept": true}),
3001 );
3002 let tr = res.into_tool_result("call_1", "demo");
3003 assert_eq!(
3004 tr.result,
3005 Some(serde_json::json!({"_raw_output_scalar": "user_value", "kept": true}))
3006 );
3007 assert_eq!(tr.raw_output, None);
3008 }
3009
3010 #[test]
3011 fn test_success_result_with_only_raw_output_scalar_key_is_not_unwrapped() {
3012 let res = ToolExecutionResult::success(serde_json::json!({"_raw_output_scalar": "v"}));
3015 let tr = res.into_tool_result("call_1", "demo");
3016 assert_eq!(
3017 tr.result,
3018 Some(serde_json::json!({"_raw_output_scalar": "v"}))
3019 );
3020 assert_eq!(tr.raw_output, None);
3021 }
3022
3023 #[test]
3024 fn test_echo_tool_display_name() {
3025 let tool = EchoTool;
3026 assert_eq!(tool.display_name(), Some("Echo"));
3027
3028 let def = tool.to_definition();
3029 assert_eq!(def.display_name(), Some("Echo"));
3030 }
3031
3032 #[test]
3033 fn test_all_default_tools_have_display_names() {
3034 let registry = ToolRegistry::with_defaults();
3035 let definitions = registry.tool_definitions();
3036
3037 for def in &definitions {
3038 assert!(
3039 def.display_name().is_some(),
3040 "Tool '{}' should have a display_name",
3041 def.name()
3042 );
3043 }
3044 }
3045
3046 #[tokio::test]
3047 async fn test_tool_registry_as_executor() {
3048 let mut registry = ToolRegistry::new();
3049 registry.register(EchoTool);
3050
3051 let tool_call = ToolCall {
3052 id: "call_1".to_string(),
3053 name: "echo".to_string(),
3054 arguments: serde_json::json!({"message": "test"}),
3055 };
3056
3057 let tool_def = registry.get("echo").unwrap().to_definition();
3058 let result = registry.execute(&tool_call, &tool_def).await.unwrap();
3059
3060 assert!(result.error.is_none());
3061 assert_eq!(result.result.unwrap()["echoed"], "test");
3062 }
3063
3064 #[test]
3065 fn test_tool_to_definition() {
3066 let tool = GetCurrentTimeTool;
3067 let def = tool.to_definition();
3068
3069 let ToolDefinition::Builtin(builtin) = def else {
3070 panic!("expected Builtin variant");
3071 };
3072 assert_eq!(builtin.name, "get_current_time");
3073 assert_eq!(builtin.policy, ToolPolicy::Auto);
3074 }
3075
3076 #[test]
3077 fn test_with_defaults_has_expected_tools() {
3078 let registry = ToolRegistry::with_defaults();
3079
3080 assert!(
3082 registry.has("get_current_time"),
3083 "should have get_current_time"
3084 );
3085 assert!(registry.has("echo"), "should have echo");
3086 assert!(
3089 !registry.has("spawn_background"),
3090 "spawn_background must NOT be in defaults — it comes from the \
3091 background_execution capability"
3092 );
3093 assert!(
3094 registry.has("report_progress"),
3095 "should have report_progress"
3096 );
3097
3098 assert!(registry.has("add"), "should have add");
3100 assert!(registry.has("subtract"), "should have subtract");
3101 assert!(registry.has("multiply"), "should have multiply");
3102 assert!(registry.has("divide"), "should have divide");
3103
3104 assert!(registry.has("get_weather"), "should have get_weather");
3106 assert!(registry.has("get_forecast"), "should have get_forecast");
3107
3108 assert!(registry.has("write_todos"), "should have write_todos");
3110
3111 assert!(registry.has("read_file"), "should have read_file");
3113 assert!(registry.has("write_file"), "should have write_file");
3114 assert!(registry.has("edit_file"), "should have edit_file");
3115 assert!(registry.has("list_directory"), "should have list_directory");
3116 assert!(registry.has("grep_files"), "should have grep_files");
3117 assert!(registry.has("delete_file"), "should have delete_file");
3118 assert!(registry.has("stat_file"), "should have stat_file");
3119
3120 assert!(registry.has("web_fetch"), "should have web_fetch");
3122
3123 assert_eq!(registry.len(), 18, "should have 18 default tools");
3125 }
3126
3127 #[tokio::test]
3128 async fn test_with_defaults_tools_are_executable() {
3129 let registry = ToolRegistry::with_defaults();
3130
3131 let tool_call = ToolCall {
3133 id: "call_1".to_string(),
3134 name: "echo".to_string(),
3135 arguments: serde_json::json!({"message": "hello from defaults"}),
3136 };
3137
3138 let tool_def = registry.get("echo").unwrap().to_definition();
3139 let result = registry.execute(&tool_call, &tool_def).await.unwrap();
3140
3141 assert!(result.error.is_none());
3142 assert_eq!(result.result.unwrap()["echoed"], "hello from defaults");
3143 }
3144
3145 #[tokio::test]
3146 async fn test_with_defaults_math_tools() {
3147 let registry = ToolRegistry::with_defaults();
3148
3149 let tool_call = ToolCall {
3151 id: "call_add".to_string(),
3152 name: "add".to_string(),
3153 arguments: serde_json::json!({"a": 5, "b": 3}),
3154 };
3155
3156 let tool_def = registry.get("add").unwrap().to_definition();
3157 let result = registry.execute(&tool_call, &tool_def).await.unwrap();
3158
3159 assert!(result.error.is_none());
3160 assert_eq!(result.result.unwrap()["result"].as_f64().unwrap(), 8.0);
3162 }
3163
3164 #[test]
3168 fn test_with_defaults_excludes_capability_only_tools() {
3169 let registry = ToolRegistry::with_defaults();
3170
3171 assert!(
3173 !registry.has("bash"),
3174 "bash must not be in defaults — it comes from bashkit_shell capability"
3175 );
3176 assert!(
3178 !registry.has("kv_store"),
3179 "kv_store must not be in defaults — it comes from session_storage capability"
3180 );
3181 assert!(
3185 !registry.has("spawn_background"),
3186 "spawn_background must not be in defaults — it comes from the \
3187 background_execution capability (auto-activated by tool hints)"
3188 );
3189 }
3190
3191 #[tokio::test]
3192 async fn test_spawn_background_executes_and_signals_session() {
3193 let session_id = SessionId::new();
3194 let file_store = Arc::new(TestFileStore::default());
3195 let platform_store = Arc::new(TestPlatformStore::default());
3196 let storage_store = Arc::new(NoopStorageStore);
3197 let task_registry = Arc::new(InMemoryTaskRegistry::default());
3198 let tool_registry = ToolRegistry::builder()
3199 .tool(SpawnBackgroundTool)
3200 .tool(TestBackgroundTool)
3201 .build();
3202
3203 let context = ToolContext::with_stores(session_id, file_store.clone(), storage_store)
3204 .with_tool_registry(Arc::new(tool_registry))
3205 .with_platform_store(platform_store.clone())
3206 .with_session_task_registry(task_registry.clone());
3207
3208 let tool = SpawnBackgroundTool;
3209 let result = tool
3210 .execute_with_context(
3211 json!({
3212 "tool": "test_background",
3213 "args": { "summary": "Background complete" }
3214 }),
3215 &context,
3216 )
3217 .await;
3218
3219 let ToolExecutionResult::Success(value) = result else {
3220 panic!("spawn_background should succeed");
3221 };
3222 let run_id = value["run_id"].as_str().unwrap().to_string();
3223 let task_id = value["task_id"].as_str().unwrap().to_string();
3224
3225 tokio::time::timeout(std::time::Duration::from_secs(2), async {
3226 loop {
3227 if let Ok(Some(task)) = task_registry.get(session_id, &task_id).await
3228 && task.state == crate::session_task::SessionTaskState::Succeeded
3229 {
3230 break task;
3231 }
3232 tokio::time::sleep(std::time::Duration::from_millis(20)).await;
3233 }
3234 })
3235 .await
3236 .expect("background run should complete");
3237 let _ = run_id; let messages = platform_store.sent_messages.lock().unwrap().clone();
3240 assert_eq!(messages.len(), 1);
3241 assert!(messages[0].contains("Background run completed"));
3242
3243 let log_file = file_store
3244 .read_file(session_id, &format!("/.background/{run_id}/output.log"))
3245 .await
3246 .unwrap()
3247 .expect("log file");
3248 assert!(
3249 log_file
3250 .content
3251 .as_deref()
3252 .unwrap_or_default()
3253 .contains("hello from background")
3254 );
3255 }
3256
3257 #[tokio::test]
3258 async fn test_spawn_background_persists_failure_artifacts() {
3259 let session_id = SessionId::new();
3260 let file_store = Arc::new(TestFileStore::default());
3261 let storage_store = Arc::new(NoopStorageStore);
3262 let task_registry = Arc::new(InMemoryTaskRegistry::default());
3263 let tool_registry = ToolRegistry::builder()
3264 .tool(SpawnBackgroundTool)
3265 .tool(TestFailingBackgroundTool)
3266 .build();
3267
3268 let context = ToolContext::with_stores(session_id, file_store.clone(), storage_store)
3269 .with_tool_registry(Arc::new(tool_registry))
3270 .with_session_task_registry(task_registry.clone());
3271
3272 let result = SpawnBackgroundTool
3273 .execute_with_context(
3274 json!({
3275 "tool": "test_background_fail",
3276 "args": {}
3277 }),
3278 &context,
3279 )
3280 .await;
3281
3282 let ToolExecutionResult::Success(value) = result else {
3283 panic!("spawn_background should succeed");
3284 };
3285 let run_id = value["run_id"].as_str().unwrap().to_string();
3286 let task_id = value["task_id"].as_str().unwrap().to_string();
3287
3288 tokio::time::timeout(std::time::Duration::from_secs(2), async {
3289 loop {
3290 if let Ok(Some(task)) = task_registry.get(session_id, &task_id).await
3291 && task.state == crate::session_task::SessionTaskState::Failed
3292 {
3293 break task;
3294 }
3295 tokio::time::sleep(std::time::Duration::from_millis(20)).await;
3296 }
3297 })
3298 .await
3299 .expect("background run should fail");
3300 let _ = run_id;
3301
3302 let log_file = file_store
3303 .read_file(session_id, &format!("/.background/{run_id}/output.log"))
3304 .await
3305 .unwrap()
3306 .expect("log file");
3307 assert!(
3308 log_file
3309 .content
3310 .as_deref()
3311 .unwrap_or_default()
3312 .contains("background failed")
3313 );
3314
3315 let result_file = file_store
3316 .read_file(session_id, &format!("/.background/{run_id}/result.json"))
3317 .await
3318 .unwrap()
3319 .expect("result file");
3320 let result_json: Value =
3321 serde_json::from_str(result_file.content.as_deref().unwrap_or_default())
3322 .expect("valid json");
3323 assert_eq!(result_json["status"], "failed");
3324 assert_eq!(result_json["error"], "boom");
3325 }
3326
3327 #[tokio::test]
3328 async fn test_spawn_background_rejects_when_session_active_run_limit_reached() {
3329 let session_id = SessionId::new();
3330 let file_store = Arc::new(TestFileStore::default());
3331 let storage_store = Arc::new(NoopStorageStore);
3332 let task_registry = Arc::new(InMemoryTaskRegistry::default());
3333 let release = StdArc::new(AtomicBool::new(false));
3334 let tool_registry = ToolRegistry::builder()
3335 .tool(SpawnBackgroundTool)
3336 .tool(BlockingBackgroundTool {
3337 release: release.clone(),
3338 })
3339 .build();
3340
3341 let context = ToolContext::with_stores(session_id, file_store, storage_store)
3342 .with_tool_registry(Arc::new(tool_registry))
3343 .with_session_task_registry(task_registry.clone());
3344
3345 let mut task_ids = Vec::new();
3346 for _ in 0..MAX_ACTIVE_BACKGROUND_RUNS_PER_SESSION {
3347 let result = SpawnBackgroundTool
3348 .execute_with_context(
3349 json!({
3350 "tool": "test_background_blocking",
3351 "args": {}
3352 }),
3353 &context,
3354 )
3355 .await;
3356
3357 let ToolExecutionResult::Success(value) = result else {
3358 panic!("background run below the session limit should start");
3359 };
3360 task_ids.push(value["task_id"].as_str().unwrap().to_string());
3361 }
3362
3363 tokio::time::timeout(std::time::Duration::from_secs(2), async {
3365 loop {
3366 let running = task_registry
3367 .list(
3368 session_id,
3369 Some(&crate::session_task::SessionTaskFilter {
3370 kind: Some(crate::session_task::TASK_KIND_BACKGROUND_TOOL.to_string()),
3371 state: Some(crate::session_task::SessionTaskState::Running),
3372 }),
3373 )
3374 .await
3375 .unwrap();
3376 if running.len() == MAX_ACTIVE_BACKGROUND_RUNS_PER_SESSION {
3377 break;
3378 }
3379 tokio::time::sleep(std::time::Duration::from_millis(20)).await;
3380 }
3381 })
3382 .await
3383 .expect("background runs should become running");
3384
3385 let result = SpawnBackgroundTool
3386 .execute_with_context(
3387 json!({
3388 "tool": "test_background_blocking",
3389 "args": {}
3390 }),
3391 &context,
3392 )
3393 .await;
3394
3395 let ToolExecutionResult::ToolError(message) = result else {
3396 release.store(true, Ordering::SeqCst);
3397 panic!("spawn_background should reject once the session limit is reached");
3398 };
3399 assert!(message.contains("active background runs per session"));
3400
3401 release.store(true, Ordering::SeqCst);
3402 tokio::time::timeout(std::time::Duration::from_secs(2), async {
3403 for task_id in task_ids {
3404 loop {
3405 if let Ok(Some(task)) = task_registry.get(session_id, &task_id).await
3406 && task.state == crate::session_task::SessionTaskState::Succeeded
3407 {
3408 break;
3409 }
3410 tokio::time::sleep(std::time::Duration::from_millis(20)).await;
3411 }
3412 }
3413 })
3414 .await
3415 .expect("blocking background runs should complete after release");
3416
3417 tokio::time::timeout(std::time::Duration::from_secs(1), async {
3422 loop {
3423 if !has_session_background_permits(session_id) {
3424 break;
3425 }
3426 tokio::time::sleep(std::time::Duration::from_millis(5)).await;
3427 }
3428 })
3429 .await
3430 .expect("completed background runs should prune their per-session permit cache entry");
3431 }
3432
3433 #[tokio::test]
3434 async fn test_spawn_background_requires_task_registry() {
3435 let session_id = SessionId::new();
3436 let file_store = Arc::new(TestFileStore::default());
3437 let storage_store = Arc::new(NoopStorageStore);
3438 let tool_registry = ToolRegistry::builder()
3439 .tool(SpawnBackgroundTool)
3440 .tool(TestBackgroundTool)
3441 .build();
3442
3443 let context = ToolContext::with_stores(session_id, file_store, storage_store)
3445 .with_tool_registry(Arc::new(tool_registry));
3446
3447 let result = SpawnBackgroundTool
3448 .execute_with_context(
3449 json!({
3450 "tool": "test_background",
3451 "args": {}
3452 }),
3453 &context,
3454 )
3455 .await;
3456
3457 let ToolExecutionResult::ToolError(message) = result else {
3458 panic!("spawn_background should reject missing task registry");
3459 };
3460 assert!(message.contains("Session task registry not available"));
3461 }
3462
3463 #[tokio::test]
3464 async fn test_spawn_background_requires_file_store() {
3465 let session_id = SessionId::new();
3466 let storage_store = Arc::new(NoopStorageStore);
3467 let task_registry = Arc::new(InMemoryTaskRegistry::default());
3468 let tool_registry = ToolRegistry::builder()
3469 .tool(SpawnBackgroundTool)
3470 .tool(TestBackgroundTool)
3471 .build();
3472
3473 let context = ToolContext::with_storage_store(session_id, storage_store)
3475 .with_tool_registry(Arc::new(tool_registry))
3476 .with_session_task_registry(task_registry);
3477
3478 let result = SpawnBackgroundTool
3479 .execute_with_context(
3480 json!({
3481 "tool": "test_background",
3482 "args": {}
3483 }),
3484 &context,
3485 )
3486 .await;
3487
3488 let ToolExecutionResult::ToolError(message) = result else {
3489 panic!("spawn_background should reject missing file store");
3490 };
3491 assert!(message.contains("Session file store not available"));
3492 }
3493
3494 #[tokio::test]
3495 async fn test_spawn_background_caps_output_log_size() {
3496 let session_id = SessionId::new();
3497 let file_store = Arc::new(TestFileStore::default());
3498 let storage_store = Arc::new(NoopStorageStore);
3499 let task_registry = Arc::new(InMemoryTaskRegistry::default());
3500 let tool_registry = ToolRegistry::builder()
3501 .tool(SpawnBackgroundTool)
3502 .tool(TestLargeOutputBackgroundTool)
3503 .build();
3504
3505 let context = ToolContext::with_stores(session_id, file_store.clone(), storage_store)
3506 .with_tool_registry(Arc::new(tool_registry))
3507 .with_session_task_registry(task_registry.clone());
3508
3509 let result = SpawnBackgroundTool
3510 .execute_with_context(
3511 json!({
3512 "tool": "test_background_large_output",
3513 "args": {}
3514 }),
3515 &context,
3516 )
3517 .await;
3518
3519 let ToolExecutionResult::Success(value) = result else {
3520 panic!("spawn_background should succeed");
3521 };
3522 let run_id = value["run_id"].as_str().unwrap().to_string();
3523 let task_id = value["task_id"].as_str().unwrap().to_string();
3524
3525 tokio::time::timeout(std::time::Duration::from_secs(2), async {
3526 loop {
3527 if let Ok(Some(task)) = task_registry.get(session_id, &task_id).await
3528 && task.state == crate::session_task::SessionTaskState::Succeeded
3529 {
3530 break;
3531 }
3532 tokio::time::sleep(std::time::Duration::from_millis(20)).await;
3533 }
3534 })
3535 .await
3536 .expect("background run should complete");
3537 let _ = run_id;
3538
3539 let log_content = file_store
3540 .read_file(session_id, &format!("/.background/{run_id}/output.log"))
3541 .await
3542 .unwrap()
3543 .expect("log file")
3544 .content
3545 .unwrap_or_default();
3546
3547 assert!(log_content.contains("[system] background output truncated"));
3548 assert!(log_content.chars().count() <= MAX_BACKGROUND_OUTPUT_LOG_CHARS + 128);
3549 }
3550
3551 #[tokio::test]
3552 async fn test_spawn_background_can_create_scheduled_monitor() {
3553 let session_id = SessionId::new();
3554 let schedule_store = Arc::new(TestScheduleStore::default());
3555 let storage_store = Arc::new(NoopStorageStore);
3556 let tool_registry = ToolRegistry::builder()
3557 .tool(SpawnBackgroundTool)
3558 .tool(TestBackgroundTool)
3559 .build();
3560
3561 let context = ToolContext::with_storage_store(session_id, storage_store)
3562 .with_tool_registry(Arc::new(tool_registry))
3563 .with_schedule_store(schedule_store.clone());
3564
3565 let result = SpawnBackgroundTool
3566 .execute_with_context(
3567 json!({
3568 "tool": "test_background",
3569 "title": "Watch PR 1319",
3570 "args": { "summary": "Background complete" },
3571 "schedule": {
3572 "cron_expression": "*/10 * * * *",
3573 "timezone": "America/Chicago"
3574 }
3575 }),
3576 &context,
3577 )
3578 .await;
3579
3580 let ToolExecutionResult::Success(value) = result else {
3581 panic!("spawn_background should create a schedule: {result:?}");
3582 };
3583
3584 assert_eq!(value["status"], "scheduled");
3585 assert_eq!(value["title"], "Watch PR 1319");
3586 assert_eq!(value["cron_expression"], "*/10 * * * *");
3587 assert_eq!(value["timezone"], "America/Chicago");
3588
3589 let schedules = schedule_store.list_schedules(session_id).await.unwrap();
3590 assert_eq!(schedules.len(), 1);
3591 assert_eq!(
3592 schedules[0].cron_expression.as_deref(),
3593 Some("*/10 * * * *")
3594 );
3595 assert!(schedules[0].description.contains("Monitor: Watch PR 1319"));
3596 assert!(
3597 schedules[0]
3598 .description
3599 .contains("\"summary\": \"Background complete\"")
3600 );
3601 }
3602
3603 #[tokio::test]
3604 async fn test_spawn_background_rejects_invalid_scheduled_at() {
3605 let session_id = SessionId::new();
3606 let storage_store = Arc::new(NoopStorageStore);
3607 let tool_registry = ToolRegistry::builder()
3608 .tool(SpawnBackgroundTool)
3609 .tool(TestBackgroundTool)
3610 .build();
3611 let context = ToolContext::with_storage_store(session_id, storage_store)
3612 .with_tool_registry(Arc::new(tool_registry));
3613
3614 let result = SpawnBackgroundTool
3615 .execute_with_context(
3616 json!({
3617 "tool": "test_background",
3618 "args": {},
3619 "schedule": {
3620 "scheduled_at": "tomorrow at noon"
3621 }
3622 }),
3623 &context,
3624 )
3625 .await;
3626
3627 let ToolExecutionResult::ToolError(message) = result else {
3628 panic!("spawn_background should reject invalid scheduled_at");
3629 };
3630 assert!(message.contains("scheduled_at must be RFC3339"));
3631 }
3632
3633 #[tokio::test]
3634 async fn test_spawn_background_rejects_ambiguous_schedule_shape() {
3635 let session_id = SessionId::new();
3636 let storage_store = Arc::new(NoopStorageStore);
3637 let tool_registry = ToolRegistry::builder()
3638 .tool(SpawnBackgroundTool)
3639 .tool(TestBackgroundTool)
3640 .build();
3641 let context = ToolContext::with_storage_store(session_id, storage_store)
3642 .with_tool_registry(Arc::new(tool_registry));
3643
3644 let result = SpawnBackgroundTool
3645 .execute_with_context(
3646 json!({
3647 "tool": "test_background",
3648 "args": {},
3649 "schedule": {
3650 "cron_expression": "*/10 * * * *",
3651 "scheduled_at": "2026-04-16T15:30:00Z"
3652 }
3653 }),
3654 &context,
3655 )
3656 .await;
3657
3658 let ToolExecutionResult::ToolError(message) = result else {
3659 panic!("spawn_background should reject ambiguous schedule shape");
3660 };
3661 assert!(message.contains("must not include both cron_expression and scheduled_at"));
3662 }
3663
3664 #[test]
3669 fn test_is_canceled_outcome_detects_sentinel() {
3670 let sentinel: std::result::Result<BackgroundOutcome, ToolExecutionResult> = Err(
3672 ToolExecutionResult::ToolError(BACKGROUND_CANCEL_SENTINEL.to_string()),
3673 );
3674 assert!(is_canceled_outcome(&sentinel));
3675 }
3676
3677 #[test]
3678 fn test_is_canceled_outcome_does_not_match_other_errors() {
3679 let other_err: std::result::Result<BackgroundOutcome, ToolExecutionResult> =
3680 Err(ToolExecutionResult::ToolError("boom".to_string()));
3681 assert!(!is_canceled_outcome(&other_err));
3682
3683 let success: std::result::Result<BackgroundOutcome, ToolExecutionResult> =
3684 Ok(BackgroundOutcome {
3685 summary: "ok".to_string(),
3686 result: json!({"ok": true}),
3687 raw_output: None,
3688 });
3689 assert!(!is_canceled_outcome(&success));
3690 }
3691
3692 #[derive(Default)]
3695 struct SleepingBackgroundTool;
3696
3697 #[async_trait]
3698 impl BackgroundExecutableTool for SleepingBackgroundTool {
3699 async fn execute_background(
3700 &self,
3701 _arguments: Value,
3702 _context: ToolContext,
3703 sink: Arc<dyn BackgroundEventSink>,
3704 ) -> std::result::Result<BackgroundOutcome, ToolExecutionResult> {
3705 sink.status("Sleeping forever")
3706 .await
3707 .map_err(ToolExecutionResult::internal_error)?;
3708 tokio::time::sleep(std::time::Duration::from_secs(3600)).await;
3710 Ok(BackgroundOutcome {
3711 summary: "should not reach here".to_string(),
3712 result: json!({}),
3713 raw_output: None,
3714 })
3715 }
3716 }
3717
3718 #[async_trait]
3719 impl Tool for SleepingBackgroundTool {
3720 fn name(&self) -> &str {
3721 "test_background_sleeping"
3722 }
3723
3724 fn display_name(&self) -> Option<&str> {
3725 Some("Test Background Sleeping")
3726 }
3727
3728 fn description(&self) -> &str {
3729 "background test tool that sleeps indefinitely"
3730 }
3731
3732 fn parameters_schema(&self) -> Value {
3733 json!({
3734 "type": "object",
3735 "properties": {}
3736 })
3737 }
3738
3739 async fn execute(&self, _arguments: Value) -> ToolExecutionResult {
3740 ToolExecutionResult::tool_error("foreground unsupported")
3741 }
3742
3743 fn hints(&self) -> ToolHints {
3744 ToolHints::default().with_supports_background(true)
3745 }
3746
3747 fn as_background_executable(&self) -> Option<&dyn BackgroundExecutableTool> {
3748 Some(self)
3749 }
3750 }
3751
3752 #[derive(Default)]
3756 struct InMemoryTaskRegistry {
3757 tasks: Mutex<HashMap<String, crate::session_task::SessionTask>>,
3758 }
3759
3760 #[async_trait]
3761 impl crate::session_task::SessionTaskRegistry for InMemoryTaskRegistry {
3762 async fn create(
3763 &self,
3764 input: crate::session_task::CreateSessionTask,
3765 ) -> crate::Result<crate::session_task::SessionTask> {
3766 let mut tasks = self.tasks.lock().unwrap();
3767 if let Some(id) = &input.id
3768 && let Some(existing) = tasks.get(id)
3769 {
3770 return Ok(existing.clone());
3771 }
3772 let task = crate::session_task::new_session_task(input, chrono::Utc::now());
3773 tasks.insert(task.id.clone(), task.clone());
3774 Ok(task)
3775 }
3776
3777 async fn update(
3778 &self,
3779 _session_id: SessionId,
3780 task_id: &str,
3781 update: crate::session_task::SessionTaskUpdate,
3782 ) -> crate::Result<Option<crate::session_task::SessionTask>> {
3783 let mut tasks = self.tasks.lock().unwrap();
3784 let Some(task) = tasks.get_mut(task_id) else {
3785 return Ok(None);
3786 };
3787 crate::session_task::apply_task_update(task, update, chrono::Utc::now());
3788 Ok(Some(task.clone()))
3789 }
3790
3791 async fn get(
3792 &self,
3793 _session_id: SessionId,
3794 task_id: &str,
3795 ) -> crate::Result<Option<crate::session_task::SessionTask>> {
3796 Ok(self.tasks.lock().unwrap().get(task_id).cloned())
3797 }
3798
3799 async fn list(
3800 &self,
3801 _session_id: SessionId,
3802 filter: Option<&crate::session_task::SessionTaskFilter>,
3803 ) -> crate::Result<Vec<crate::session_task::SessionTask>> {
3804 let tasks = self.tasks.lock().unwrap();
3805 Ok(tasks
3806 .values()
3807 .filter(|task| {
3808 filter.is_none_or(|f| {
3809 f.kind.as_deref().is_none_or(|kind| task.kind == kind)
3810 && f.state.is_none_or(|state| task.state == state)
3811 })
3812 })
3813 .cloned()
3814 .collect())
3815 }
3816
3817 async fn request_cancel(
3818 &self,
3819 _session_id: SessionId,
3820 task_id: &str,
3821 ) -> crate::Result<Option<crate::session_task::SessionTask>> {
3822 let mut tasks = self.tasks.lock().unwrap();
3823 let Some(task) = tasks.get_mut(task_id) else {
3824 return Ok(None);
3825 };
3826 task.cancel_requested_at
3827 .get_or_insert_with(chrono::Utc::now);
3828 task.updated_at = chrono::Utc::now();
3829 Ok(Some(task.clone()))
3830 }
3831
3832 async fn record_message(
3833 &self,
3834 _session_id: SessionId,
3835 task_id: &str,
3836 message: crate::session_task::NewTaskMessage,
3837 ) -> crate::Result<crate::session_task::TaskMessage> {
3838 let tasks = self.tasks.lock().unwrap();
3839 let _task = tasks
3840 .get(task_id)
3841 .ok_or_else(|| crate::AgentLoopError::tool(format!("no task {task_id}")))?;
3842 Ok(crate::session_task::TaskMessage {
3843 id: crate::session_task::generate_task_message_id(),
3844 task_id: task_id.to_string(),
3845 direction: message.direction,
3846 content: message.content,
3847 in_reply_to: message.in_reply_to,
3848 created_at: chrono::Utc::now(),
3849 })
3850 }
3851
3852 async fn list_messages(
3853 &self,
3854 _session_id: SessionId,
3855 _task_id: &str,
3856 _limit: Option<u32>,
3857 _after_id: Option<&str>,
3858 ) -> crate::Result<Vec<crate::session_task::TaskMessage>> {
3859 Ok(Vec::new())
3860 }
3861 }
3862
3863 #[tokio::test]
3866 async fn test_cancel_background_run_via_task_registry() {
3867 let session_id = SessionId::new();
3868 let file_store = Arc::new(TestFileStore::default());
3869 let storage_store = Arc::new(NoopStorageStore);
3870 let task_registry = Arc::new(InMemoryTaskRegistry::default());
3871
3872 let tool_registry = ToolRegistry::builder()
3873 .tool(SpawnBackgroundTool)
3874 .tool(SleepingBackgroundTool)
3875 .build();
3876
3877 let context = ToolContext::with_stores(session_id, file_store.clone(), storage_store)
3878 .with_tool_registry(Arc::new(tool_registry))
3879 .with_session_task_registry(task_registry.clone());
3880
3881 let result = SpawnBackgroundTool
3882 .execute_with_context(
3883 json!({
3884 "tool": "test_background_sleeping",
3885 "args": {},
3886 "signal_on_completion": false
3887 }),
3888 &context,
3889 )
3890 .await;
3891
3892 let ToolExecutionResult::Success(value) = result else {
3893 panic!("spawn_background should succeed");
3894 };
3895 let run_id = value["run_id"].as_str().unwrap().to_string();
3896 let task_id = value["task_id"].as_str().unwrap().to_string();
3897
3898 tokio::time::timeout(std::time::Duration::from_secs(5), async {
3900 loop {
3901 if let Ok(Some(task)) = task_registry.get(session_id, &task_id).await
3903 && task.heartbeat_at.is_some()
3904 {
3905 break;
3906 }
3907 tokio::time::sleep(std::time::Duration::from_millis(50)).await;
3908 }
3909 })
3910 .await
3911 .expect("background run should start and send at least one heartbeat");
3912
3913 task_registry
3915 .request_cancel(session_id, &task_id)
3916 .await
3917 .expect("request_cancel should succeed");
3918
3919 tokio::time::timeout(std::time::Duration::from_secs(10), async {
3921 loop {
3922 if let Ok(Some(task)) = task_registry.get(session_id, &task_id).await
3923 && task.state == crate::session_task::SessionTaskState::Canceled
3924 {
3925 break task;
3926 }
3927 tokio::time::sleep(std::time::Duration::from_millis(50)).await;
3928 }
3929 })
3930 .await
3931 .expect("background task should reach Canceled state");
3932
3933 let result_file = file_store
3935 .read_file(session_id, &format!("/.background/{run_id}/result.json"))
3936 .await
3937 .unwrap()
3938 .expect("result.json should exist");
3939 let result_json: Value =
3940 serde_json::from_str(result_file.content.as_deref().unwrap_or_default())
3941 .expect("valid json");
3942 assert_eq!(result_json["status"], "canceled");
3943
3944 let log_file = file_store
3945 .read_file(session_id, &format!("/.background/{run_id}/output.log"))
3946 .await
3947 .unwrap()
3948 .expect("output.log should exist");
3949 assert!(
3950 log_file
3951 .content
3952 .as_deref()
3953 .unwrap_or_default()
3954 .contains("Canceled by request.")
3955 );
3956 }
3957
3958 fn make_reattach_task(spec: serde_json::Value) -> crate::session_task::SessionTask {
3963 use crate::session_task::{SessionTaskState, TaskLinks, TaskWakePolicy};
3964 crate::session_task::SessionTask {
3965 id: "t-reattach".to_string(),
3966 session_id: SessionId::new(),
3967 root_session_id: None,
3968 kind: crate::session_task::TASK_KIND_BACKGROUND_TOOL.to_string(),
3969 display_name: "Reattach test".to_string(),
3970 spec,
3971 state: SessionTaskState::Running,
3972 state_detail: None,
3973 progress: None,
3974 input_request: None,
3975 cancel_requested_at: None,
3976 summary: None,
3977 result_path: None,
3978 artifacts: vec![],
3979 error: None,
3980 attempt: 2,
3981 worker_id: None,
3982 heartbeat_at: None,
3983 links: TaskLinks::default(),
3984 wake_policy: TaskWakePolicy::Silent,
3985 created_at: chrono::Utc::now(),
3986 started_at: None,
3987 finished_at: None,
3988 updated_at: chrono::Utc::now(),
3989 }
3990 }
3991
3992 #[tokio::test]
3993 async fn reattach_fails_with_missing_file_store() {
3994 let session_id = SessionId::new();
3995 let task_registry = Arc::new(InMemoryTaskRegistry::default());
3997 let context =
3998 crate::traits::ToolContext::new(session_id).with_session_task_registry(task_registry);
3999 let task = make_reattach_task(serde_json::json!({
4000 "tool": "get_current_time",
4001 "arguments": {},
4002 "reattachable": true,
4003 "signal_on_completion": true,
4004 }));
4005 let err = reattach_background_run(&task, &context)
4006 .await
4007 .expect_err("should fail without file store");
4008 assert!(
4009 err.to_string().contains("file store"),
4010 "error should mention file store, got: {err}"
4011 );
4012 }
4013
4014 #[tokio::test]
4015 async fn reattach_fails_with_missing_task_registry() {
4016 let session_id = SessionId::new();
4017 let file_store = Arc::new(TestFileStore::default());
4018 let storage_store = Arc::new(NoopStorageStore);
4019 let context =
4021 crate::traits::ToolContext::with_stores(session_id, file_store, storage_store);
4022 let task = make_reattach_task(serde_json::json!({
4023 "tool": "get_current_time",
4024 "arguments": {},
4025 "reattachable": true,
4026 "signal_on_completion": true,
4027 }));
4028 let err = reattach_background_run(&task, &context)
4029 .await
4030 .expect_err("should fail without task registry");
4031 assert!(
4032 err.to_string().contains("task registry"),
4033 "error should mention task registry, got: {err}"
4034 );
4035 }
4036
4037 #[tokio::test]
4038 async fn reattach_fails_with_unknown_tool_name() {
4039 let session_id = SessionId::new();
4040 let file_store = Arc::new(TestFileStore::default());
4041 let storage_store = Arc::new(NoopStorageStore);
4042 let task_registry = Arc::new(InMemoryTaskRegistry::default());
4043 let context =
4044 crate::traits::ToolContext::with_stores(session_id, file_store, storage_store)
4045 .with_session_task_registry(task_registry);
4046 let task = make_reattach_task(serde_json::json!({
4048 "tool": "test_background",
4049 "arguments": {},
4050 "reattachable": true,
4051 "signal_on_completion": true,
4052 }));
4053 let err = reattach_background_run(&task, &context)
4054 .await
4055 .expect_err("should fail for unknown tool");
4056 assert!(
4057 err.to_string().contains("not found in built-in registry"),
4058 "error should mention built-in registry, got: {err}"
4059 );
4060 }
4061
4062 #[tokio::test]
4063 async fn reattach_fails_with_missing_tool_spec_field() {
4064 let session_id = SessionId::new();
4065 let file_store = Arc::new(TestFileStore::default());
4066 let storage_store = Arc::new(NoopStorageStore);
4067 let task_registry = Arc::new(InMemoryTaskRegistry::default());
4068 let context =
4069 crate::traits::ToolContext::with_stores(session_id, file_store, storage_store)
4070 .with_session_task_registry(task_registry);
4071 let task = make_reattach_task(serde_json::json!({ "reattachable": true }));
4073 let err = reattach_background_run(&task, &context)
4074 .await
4075 .expect_err("should fail with missing tool field");
4076 assert!(
4077 err.to_string().contains("missing 'tool' field"),
4078 "error should mention missing tool field, got: {err}"
4079 );
4080 }
4081}